From 577e131c76b3e2b0a2b670ed492607880cc7bafe Mon Sep 17 00:00:00 2001
From: Xu Hong Chen <110699064+xhgchen@users.noreply.github.com>
Date: Mon, 7 Aug 2023 14:30:47 -0700
Subject: [PATCH] Create full unit velocity trajectory for tests

* Contains mass, velocities, positions, and volume
---
 transport_analysis/tests/test_viscosity.py | 33 ++++++++++++++++++++++
 1 file changed, 33 insertions(+)

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