def test_chunked_trajectory_with_lag(self): data = np.vstack((self.d, self.d)) reader = FragmentedTrajectoryReader([self.d, self.d]) for lag in [0, 1, 3]: for stride in [1, 3, 5]: for chunksize in [1, 34, 53, 72]: reader.chunksize = chunksize if lag > 0: collected = None collected_lagged = None for itraj, X, Y in reader.iterator(stride=stride, lag=lag): collected = X if collected is None else np.vstack( (collected, X)) collected_lagged = Y if collected_lagged is None else np.vstack( (collected_lagged, Y)) np.testing.assert_array_almost_equal( data[::stride][0:len(collected_lagged)], collected) np.testing.assert_array_almost_equal( data[lag::stride], collected_lagged) else: collected = None for itraj, X in reader.iterator(stride=stride): collected = X if collected is None else np.vstack( (collected, X)) np.testing.assert_array_almost_equal( data[::stride], collected)
def test_cols(self): dim = 5 arr = np.arange(60).reshape(-1, dim) data = [(arr, arr), arr, (arr, arr, arr)] reader = FragmentedTrajectoryReader(data) cols = (0, 3) for itraj, x in reader.iterator(chunk=0, return_trajindex=True, cols=cols): if isinstance(data[itraj], tuple): syn_traj = np.concatenate(data[itraj]) else: syn_traj = data[itraj] np.testing.assert_equal(x, syn_traj[:, cols])
def test_full_trajectory_stridden_with_lag(self): reader = FragmentedTrajectoryReader([self.d, self.d]) data = np.vstack((self.d, self.d)) for lag in [1, 5, 7]: for stride in [1, 3, 5, 7, 13, 20]: reader.chunksize = 0 X, Y = None, None # not chunked for itraj, X, Y in reader.iterator(stride=stride, lag=lag): pass np.testing.assert_array_almost_equal(data[::stride][0:len(Y)], X) np.testing.assert_array_almost_equal(data[lag::stride], Y)