Create full unit velocity trajectory for tests

* Contains mass, velocities, positions, and volume

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