Load data from numpy arrays

.gitignore

@@ -6,4 +6,7 @@
 # pip install folders
 countoscope.egg-info/
 build/
+
+# example script outputs
 examples/example_xyz.png
+examples/example_numpy.png

countoscope/numpytrajectory.py

@@ -0,0 +1,58 @@
+import numpy as np
+
+class NumpyFrame():
+    def __init__(self, positions):
+        assert len(positions.shape) == 2
+        self.positions = positions
+
+class NumpyTrajectory:
+    """a class that performs like a chemfiles trajectory"""
+    def __init__(self, array, column_labels):
+        self.array = array
+
+        self.x_column = None
+        self.y_column = None
+        self.z_column = None
+        self.t_column = None
+        self.dimension = 2
+
+        for column_index, column_header in enumerate(column_labels):
+            if column_header == 'x':
+                self.x_column = column_index
+            if column_header == 'y':
+                self.y_column = column_index
+            if column_header == 'z':
+                self.z_column = column_index
+                self.dimension = 3
+            if column_header == 't':
+                self.t_column = column_index
+
+        assert self.t_column is not None, 'column_headers must include "t"'
+        assert self.x_column is not None, 'column_headers must include "x"'
+        assert self.y_column is not None, 'column_headers must include "y"'
+
+        # ensure timestep is zero-based
+        self.array[:, self.t_column] -= self.array[:, self.t_column].min()
+        self.nsteps = int(self.array[:, self.t_column].max() + 1)
+
+    def get_system_size(self):
+        if self.dimension == 2:
+            return np.array([
+                self.array[:, self.x_column].max(),
+                self.array[:, self.y_column].max(),
+            ])
+        elif self.dimension == 3:
+            return np.array([
+                self.array[:, self.x_column].max(),
+                self.array[:, self.y_column].max(),
+                self.array[:, self.z_column].max(),
+            ])
+
+    def read_step(self, timestep):
+        row_indices_at_timestep = self.array[:, self.t_column] == timestep
+        array_at_timestep = self.array[row_indices_at_timestep, :]
+
+        if self.dimension == 2:
+            return NumpyFrame(array_at_timestep[:, [self.x_column, self.y_column]])
+        elif self.dimension == 3:
+            return NumpyFrame(array_at_timestep[:, [self.x_column, self.y_column, self.z_column]])

countoscope/trajectorytool.py

@@ -9,34 +9,47 @@
 import os
 import numpy as np
 from chemfiles import Trajectory as chem_traj
-
+from .numpytrajectory import NumpyTrajectory
 
 class TrajectoryTool():
     r"""Class for tools based on trajectories."""
 
     def __init__(self,
-                 trajectory_file: str,
+                 trajectory_file: str = None,
+                 trajectory_array: np.ndarray = None,
+                 trajectory_labels: list = None,
                  topology_file: str = None,
                  system_size: np.array = None,
                  dimension: int = None,
                  symmetric_system: bool = False,
                  *args,
                  **kwargs
                  ):
-        self.trajectory_file = trajectory_file
-        self.topology_file = topology_file
         self.system_size = system_size
         self.dimension = dimension
         self.symmetric_system = symmetric_system
-        self.prepare_trajectories()
 
-    def prepare_trajectories(self):
-        self.import_trajectory()
-        self.read_information_trajectory()
-        self.detect_system_dimension()
-        self.detect_system_size()
+        if trajectory_file is not None:
+            self.trajectory_file = trajectory_file
+            self.topology_file = topology_file
+            self.import_chemfiles_trajectory()
+            self.read_information_chemfiles_trajectory()
+            self.detect_system_dimension()
+
+        elif trajectory_array is not None:
+            assert trajectory_labels is not None, 'when using trajectory_array, trajectory_array_column_headers must be supplied'
+            self.import_numpy_trajectory(trajectory_array, trajectory_labels)
+            self.read_information_numpy_trajectory()
+            self.dimension = self.trajectory.dimension
+
+        else:
+            raise Exception('Either trajectory_file or trajectory_array must be supplied')
+
+        self.detect_system_boundaries()
 
-    def import_trajectory(self):
+        assert self.nb_steps > 0, f'found {self.nb_steps} steps'
+
+    def import_chemfiles_trajectory(self):
         """Import trajectory file using Chemfiles"""
         # Make sure that the trajectory file exists.
         assert os.path.exists(self.trajectory_file),  \
@@ -47,7 +60,10 @@ def import_trajectory(self):
         if self.topology_file is not None:
             self.trajectory.set_topology(self.topology_file)
 
-    def read_information_trajectory(self):
+    def import_numpy_trajectory(self, trajectory_array, trajectory_array_column_headers):
+        self.trajectory = NumpyTrajectory(trajectory_array, trajectory_array_column_headers)
+
+    def read_information_chemfiles_trajectory(self):
         """Read basic information from trajectory"""
         self.nb_steps = self.trajectory.nsteps
         # If box_size was not provided, try reading it from trajectory.
@@ -56,6 +72,11 @@ def read_information_trajectory(self):
         # If box_size is 0, return a warning
         assert np.sum(np.array(self.system_size)) > 0, \
             """Error: No box size in trajectory, provide system_size"""
+
+    def read_information_numpy_trajectory(self):
+        """Read basic information from trajectory"""
+        self.nb_steps    = self.trajectory.nsteps
+        self.system_size = self.trajectory.get_system_size()
 
     def detect_system_dimension(self):
         """Attempt to guess the dimension from the trajectory.
@@ -89,7 +110,7 @@ def detect_system_dimension(self):
             assert (self.dimension == 2) | (self.dimension == 3), \
                 """ERROR: Unsuported dimension. Must be 2 or 3."""
 
-    def detect_system_size(self):
+    def detect_system_boundaries(self):
         """From the box size, estimate the system lower and higher coordinate."""
         system_boundaries = []
         for system_length in self.system_size:

examples/example_numpy.py

@@ -0,0 +1,33 @@
+from countoscope import Countoscope
+import matplotlib.pyplot as plt
+import numpy as np
+
+"""
+run this example from within the examples folder with
+python example_numpy.py
+"""
+
+# load dataset
+# this example is an array of rows of (x, y, t) - eg from Trackpy
+data = np.load('example_dataset.npy')
+column_labels = ['x', 'y', 't']
+
+# set up the Countoscope
+countoscope = Countoscope(
+    trajectory_array  = data,
+    trajectory_labels = column_labels, # the column labels
+    box_size          = np.array([0.5, 0.5]),
+)
+
+# do the counting
+countoscope.count()
+
+# calcuate <dN^2(t)>
+delta_N2 = countoscope.evaluate_deltaN2()
+
+# plot
+t = range(0, len(delta_N2))
+plt.plot(t[1:], delta_N2[1:])
+plt.loglog()
+plt.savefig('example_numpy.png')
+plt.show()

tests/test_numpy_trajs.py

@@ -0,0 +1,43 @@
+from countoscope import Countoscope
+import numpy as np
+
+def test_load_numpy_trajs():
+    trajs = np.array([
+    #    x    y    t
+        [0.0, 0.0, 0],
+        [1.0, 0.0, 0],
+        [0.0, 1.0, 1],
+        [1.0, 1.0, 1],
+    ])
+
+    countoscope = Countoscope(
+        box_size          = np.array([0.5, 0.5]),
+        trajectory_array  = trajs,
+        trajectory_labels = ['x', 'y', 't'],
+    )
+    assert countoscope.nb_steps == 2
+    assert np.all(countoscope.trajectory.read_step(0).positions == [[0.0, 0.0], [1.0, 0.0]])
+
+    countoscope.count()
+    assert np.all(countoscope.number_matrix == [
+        [[1, 0],
+         [1, 0]],
+        [[0, 1],
+         [0, 1]],
+    ])
+
+def test_load_numpy_trajs_3D():
+    trajs = np.array([
+    #    x    y    z    t
+        [0.0, 0.0, 1.0, 0],
+        [1.0, 0.0, 0.5, 0],
+        [0.0, 1.0, 0.5, 1],
+        [1.0, 1.0, 0.0, 1],
+    ])
+
+    countoscope = Countoscope(
+        box_size          = np.array([0.5, 0.5, 0.5]),
+        trajectory_array  = trajs,
+        trajectory_labels = ['x', 'y', 'z', 't'],
+    )
+    countoscope.count()