diff --git a/transport_analysis/tests/test_viscosity.py b/transport_analysis/tests/test_viscosity.py index 730d5dc..e02f1b8 100644 --- a/transport_analysis/tests/test_viscosity.py +++ b/transport_analysis/tests/test_viscosity.py @@ -42,6 +42,39 @@ def visc_helfand(ag): return vh_t +# Full step trajectory of unit velocities i.e. v = 0 at t = 0, +# v = 1 at t = 1, v = 2 at t = 2, etc. for all components x, y, z +# with mass, positions, and volume +@pytest.fixture(scope="module") +def step_vtraj_full(NSTEP): + # Set up unit velocities + v = np.arange(NSTEP) + velocities = np.vstack([v, v, v]).T + # NSTEP frames x 1 atom x 3 dimensions, where NSTEP = 5001 + velocities_reshape = velocities.reshape([NSTEP, 1, 3]) + + # Positions derived from unit velocity setup + x = np.arange(NSTEP).astype(np.float64) + # Since initial position and velocity are 0 and acceleration is 1, + # position = 1/2t^2 + x *= x / 2 + positions = np.vstack([x, x, x]).T + # NSTEP frames x 1 atom x 3 dimensions, where NSTEP = 5001 + positions_reshape = positions.reshape([NSTEP, 1, 3]) + u = mda.Universe.empty(1, n_frames=NSTEP, velocities=True) + + # volume of 8.0 + dim = [2, 2, 2, 90, 90, 90] + for i, ts in enumerate(u.trajectory): + u.atoms.velocities = velocities_reshape[i] + u.atoms.positions = positions_reshape[i] + set_dimensions(dim)(u.trajectory.ts) + + # mass of 16.0 + u.add_TopologyAttr('masses', [16.0]) + return u + + class TestViscosityHelfand: # fixtures are helpful functions that set up a test # See more at https://docs.pytest.org/en/stable/how-to/fixtures.html