def test_io_txt(): # 1d a = np.arange(0, 3) fn = pj(testdir, 'a1d.txt') write_read_check(fn, a) # 2d shape = (3, 5) a = np.arange(0, np.prod(shape)).reshape(shape) fn = pj(testdir, 'a2d.txt') write_read_check(fn, a) # 3d, store array along axis 0,1,2,-1 shape = (3, 5, 7) a = np.arange(0, np.prod(shape)).reshape(shape) # ignore file header if shape != None for sh in [None, shape]: fn = pj(testdir, 'a3d0.txt') write_read_check(fn, a, axis=0, shape=sh) fn = pj(testdir, 'a3d1.txt') write_read_check(fn, a, axis=1, shape=sh) fn = pj(testdir, 'a3d2.txt') write_read_check(fn, a, axis=2, shape=sh) fn = pj(testdir, 'a3dm1.txt') write_read_check(fn, a, axis=-1, shape=sh) # API shape = (3, 5) arr = np.arange(0, np.prod(shape)).reshape(shape) fn = pj(testdir, 'a2d_api.txt') fh = open(fn, 'wb') for ii in range(3): fh.write('@@ some comment\n'.encode()) np.savetxt(fh, arr) fh.close() a = arrayio.readtxt(fn, shape=shape, axis=-1, comments='@@') assert (a == arr).all() txt = "1.0 2.0 3\n4 5 6\n" arr = arrayio.readtxt(StringIO(txt), shape=(2, 3), axis=-1, dtype=float) assert arr.dtype == np.array([1.0]).dtype # Apparently in Python 2.7: # float('1.0') -> 1.0 # int('1.0') -> ValueError: invalid literal for int() with base 10: '1.0' # int(float('1.0')) -> 1 # We need to (ab)use converters. Ugh. arr = arrayio.readtxt(StringIO(txt), shape=(2, 3), axis=-1, dtype=int, converters=dict( (x, lambda a: int(float(a))) for x in [0, 1, 2])) assert arr.dtype == np.array([1]).dtype
def test_io_txt(): # 1d a = np.arange(0, 3) fn = pj(testdir, 'a1d.txt') write_read_check(fn, a) # 2d shape = (3, 5) a = np.arange(0, np.prod(shape)).reshape(shape) fn = pj(testdir, 'a2d.txt') write_read_check(fn, a) # 3d, store array along axis 0,1,2,-1 shape = (3, 5, 7) a = np.arange(0, np.prod(shape)).reshape(shape) # ignore file header if shape != None for sh in [None, shape]: fn = pj(testdir, 'a3d0.txt') write_read_check(fn, a, axis=0, shape=sh) fn = pj(testdir, 'a3d1.txt') write_read_check(fn, a, axis=1, shape=sh) fn = pj(testdir, 'a3d2.txt') write_read_check(fn, a, axis=2, shape=sh) fn = pj(testdir, 'a3dm1.txt') write_read_check(fn, a, axis=-1, shape=sh) # API shape = (3, 5) arr = np.arange(0, np.prod(shape)).reshape(shape) fn = pj(testdir, 'a2d_api.txt') fh = open(fn, 'w') fh.write('@@ some comment\n') fh.write('@@ some comment\n') fh.write('@@ some comment\n') np.savetxt(fh, arr) fh.close() a = arrayio.readtxt(fn, shape=shape, axis=-1, comments='@@') assert (a == arr).all() txt = "1.0 2.0 3\n4 5 6\n" arr = arrayio.readtxt(StringIO(txt), shape=(2,3), axis=-1, dtype=float) assert arr.dtype == np.array([1.0]).dtype # Apparently in Python 2.7: # float('1.0') -> 1.0 # int('1.0') -> ValueError: invalid literal for int() with base 10: '1.0' # int(float('1.0')) -> 1 # We need to (ab)use converters. Ugh. arr = arrayio.readtxt(StringIO(txt), shape=(2,3), axis=-1, dtype=int, converters=dict((x,lambda a: int(float(a))) for x in [0,1,2])) assert arr.dtype == np.array([1]).dtype
def load_old(fn): # Load old ref data and reshape: # (natoms, 3, nstep) -> (nstep, natoms, 3) arr = arrayio.readtxt(fn) arr2 = np.rollaxis(arr, 2, 0) for istep in range(arr.shape[-1]): assert (arr[..., istep] == arr2[istep, ...]).all() return arr2
def test_traj_from_txt(): shape = (10, 20, 30) arr3d_orig = np.random.rand(*shape) axis = 0 # general stuff for axis != 0, here for axis=0 we have written_shape==shape shape_2d_chunk = shape[:axis] + shape[(axis + 1):] written_shape = (shape[axis], ) + shape_2d_chunk print("axis, written_shape:", axis, written_shape) fn = pj(testdir, 'arr_test_traj_from_txt_axis%i.txt' % axis) arrayio.writetxt(fn, arr3d_orig, axis=axis, header=False) with open(fn) as fd: arr3d = parse.traj_from_txt(fd.read(), axis=axis, shape=written_shape) fd.seek(0) # test if the "old" way of reading also works arr3d_readtxt = arrayio.readtxt(fd, shape=written_shape, axis=axis) assert (arr3d_readtxt == arr3d_orig).all() # now test traj_from_txt result assert arr3d.shape == arr3d_orig.shape, \ ("axis={}, shapes: read={} written={} orig={}".format(axis, arr3d.shape, written_shape, arr3d_orig.shape)) assert (arr3d == arr3d_orig).all()
def test_traj_from_txt(): shape = (10,20,30) arr3d_orig = np.random.rand(*shape) axis = 0 # general stuff for axis != 0, here for axis=0 we have written_shape==shape shape_2d_chunk = shape[:axis] + shape[(axis+1):] written_shape = (shape[axis],) + shape_2d_chunk print "axis, written_shape:", axis, written_shape fn = pj(testdir, 'arr_test_traj_from_txt_axis%i.txt' %axis) arrayio.writetxt(fn, arr3d_orig, axis=axis, header=False) with open(fn) as fd: arr3d = parse.traj_from_txt(fd.read(), axis=axis, shape=written_shape) fd.seek(0) # test if the "old" way of reading also works arr3d_readtxt = arrayio.readtxt(fd, shape=written_shape, axis=axis) assert (arr3d_readtxt == arr3d_orig).all() # now test traj_from_txt result assert arr3d.shape == arr3d_orig.shape, \ ("axis={0}, shapes: read={1} written={2} orig={3}".format(axis, arr3d.shape, written_shape, arr3d_orig.shape)) assert (arr3d == arr3d_orig).all()
def write_read_check(fn, arr, axis=-1, shape=None): print fn + ' ...' arrayio.writetxt(fn, arr, axis=axis) a = arrayio.readtxt(fn, axis=axis, shape=shape) assert (a == arr).all()
def write_read_check(fn, arr, axis=-1, shape=None): print("fn={} axis={} shape={}".format(fn, axis, shape)) arrayio.writetxt(fn, arr, axis=axis) a = arrayio.readtxt(fn, axis=axis, shape=shape) assert (a == arr).all()