def memmap_chunk(filename, mode="r+", **kwargs):
"""
Read-write memmap to chunk ID out of NID.
ID and NID are normally taken from arrayjob.get_ID() and arrayjob.get_NID(),
but can be specified as keyword arguments for testing.
Assuming a Numpy array has already been saved to file.
>>> np.save("test.npy", np.arange(10))
Return chunk 1 out of [0, 1, 2], [3, 4, 5], [6, 7], [8, 9].
>>> c = memmap_chunk("test.npy", ID=1, NID=4)
>>> c
memmap([3, 4, 5])
The offset attribute is the original index of the first item in the chunk.
>>> c.offset
3
Clean up after the test.
>>> del c
>>> os.remove("test.npy")
"""
myID = kwargs.get("ID", ID)
myNID = kwargs.get("NID", NID)
# deferred import so arrayjob can be used without load_memmap_offset
from cgp.utils.load_memmap_offset import open_memmap, memmap_chunk_ind
r = open_memmap(filename, "r")
n = r.shape[0]
i = np.array_split(range(n), myNID)[myID]
if len(i) > 0:
result, offset = memmap_chunk_ind(filename, i, mode=mode)
result.offset = offset
return result
else:
return np.empty(0, r.dtype)
def __missing__(self, key):
"""Memory-map existing {key}.npy on first lookup of key."""
# deferred import so arrayjob can be used without load_memmap_offset
from cgp.utils.load_memmap_offset import open_memmap
return open_memmap("%s/%s.npy" % (self.pardir, key), **self.kwargs)
