Trj slicing and dump colored trj

.gitignore

@@ -7,7 +7,9 @@ __pycache__
 docs/build
 docs/source/_autosummary
 examples/onion_analysis/data.json
-examples/analysis_workflow/water_lens.json
+examples/analysis_workflow/lens.json
+examples/analysis_workflow/onion.json
+examples/analysis_workflow/colored_trj.xyz
 examples/info_gain/trj_*.npy
 tests/systems/coex/.*
 tests/systems/.*

examples/analysis_workflow.py

@@ -4,7 +4,7 @@
 
 import MDAnalysis
 
-from dynsight.trajectory import OnionInsight, Trj
+from dynsight.trajectory import Insight, OnionInsight, Trj
 
 
 def main() -> None:
@@ -15,28 +15,36 @@ def main() -> None:
     universe = MDAnalysis.Universe(
         files_path / "oxygens.gro", files_path / "oxygens.xtc"
     )
-    water_trj = Trj(universe)
+    trj = Trj(universe)
+    n_frames = len(trj.universe.trajectory)
 
     # We want, for instance, compute LENS on this trajectory
     # From here, we work with an Insight, containing data computed from a Trj
-    water_lens = water_trj.get_lens(r_cut=7.5)
+    lens_file = files_path / "lens.json"
+    if lens_file.exists():
+        lens = Insight.load_from_json(files_path / "lens.json")
+    else:
+        lens = trj.get_lens(r_cut=7.5)
+        lens.dump_to_json(lens_file)
 
     # We can do spatial average on the computed LENS
-    water_smooth = water_lens.spatial_average(water_trj, r_cut=7.5)
+    trj_lens = trj.with_slice(slice(1, n_frames, 1))
+    lens_smooth = lens.spatial_average(trj_lens, r_cut=7.5, num_processes=6)
 
     # And we can perform onion-clustering
-    water_onion = water_smooth.get_onion(delta_t=10)
+    lens_onion = lens_smooth.get_onion_smooth(delta_t=10)
 
-    water_onion.plot_output(files_path / "tmp_fig1.png", water_smooth)
-    water_onion.plot_one_trj(
+    lens_onion.plot_output(files_path / "tmp_fig1.png", lens_smooth)
+    lens_onion.plot_one_trj(
         files_path / "tmp_fig2.png",
-        water_smooth,
+        lens_smooth,
         particle_id=1234,
     )
+    lens_onion.dump_colored_trj(trj_lens, files_path / "colored_trj.xyz")
 
     # Save/load the Insight with all the results
-    water_onion.dump_to_json(files_path / "water_lens.json")
-    _ = OnionInsight.load_from_json(files_path / "water_lens.json")
+    lens_onion.dump_to_json(files_path / "onion.json")
+    _ = OnionInsight.load_from_json(files_path / "onion.json")
 
 
 if __name__ == "__main__":

src/dynsight/_internal/analysis/spatial_average.py

@@ -3,69 +3,73 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    import MDAnalysis
+    from MDAnalysis import Universe
+    from MDAnalysis.core.groups import AtomGroup
     from numpy.typing import NDArray
 
 import ctypes
+from collections import defaultdict
 from multiprocessing import Array, Pool
 
 import numpy as np
-from MDAnalysis.analysis.distances import distance_array
+from MDAnalysis.lib.nsgrid import FastNS
 
 array: NDArray[np.float64]
 
 
 def initworker(
-    shared_array: NDArray[np.float64],
+    shared_arr: NDArray[np.float64],
     shape: int,
     dtype: tuple[float, Any],
 ) -> None:
     # The use of global statement is necessary for the correct functioning
     # of the code. Ruff error PLW0603 is therefore ignored.
     global array  # noqa: PLW0603
-    array = np.frombuffer(shared_array, dtype=dtype).reshape(shape)
+    array = np.frombuffer(shared_arr, dtype=dtype).reshape(shape)
 
 
 def processframe(
-    args: tuple[MDAnalysis.Universe, MDAnalysis.AtomGroup, float, int, bool],
+    args: tuple[Universe, AtomGroup, float, int, int, bool],
 ) -> tuple[int, NDArray[np.float64]]:
-    universe, selection, cutoff, frame, is_vector = args
-    universe.trajectory[frame]
-    distances = distance_array(
-        reference=selection.positions,
-        configuration=selection.positions,
-        box=universe.dimensions,
-    )
-    atom_id = np.argsort(distances, axis=1)
-    nn = np.sum(distances < cutoff, axis=1)
+    universe, selection, cutoff, traj_frame, result_index, is_vector = args
 
-    rows = np.arange(distances.shape[0])
-    sp_dict = {row: atom_id[row, : nn[row]] for row in rows}
+    universe.trajectory[traj_frame]
+    positions = selection.positions
+    box = universe.trajectory.ts.dimensions
 
-    if is_vector:
-        sp_array_frame = np.zeros((array.shape[0], array.shape[2]))
-        for key, value in sp_dict.items():
-            if len(value) == 0:
-                continue
-            sp_array_frame[key, :] = np.mean(array[value, frame, :], axis=0)
-        return frame, sp_array_frame
+    ns = FastNS(cutoff, positions, box=box)
+    pairs = ns.self_search().get_pairs()
+
+    neighbor_dict = defaultdict(list)
+    for i, j in pairs:
+        neighbor_dict[i].append(j)
+        neighbor_dict[j].append(i)
 
-    sp_array_frame_1 = np.zeros(array.shape[0])
-    for key, value in sp_dict.items():
-        if len(value) == 0:
-            continue
-        sp_array_frame_1[key] = np.mean(array[value, frame])
+    n_atoms = len(selection)
+    if is_vector:
+        n_features = array.shape[2]
+        sp_array_frame1 = np.zeros((n_atoms, n_features), dtype=array.dtype)
+        for key in range(n_atoms):
+            neighbors = neighbor_dict[key] + [key]
+            sp_array_frame1[key] = np.mean(
+                array[neighbors, result_index, :], axis=0
+            )
+        return result_index, sp_array_frame1
+    sp_array_frame2 = np.zeros(n_atoms, dtype=array.dtype)
+    for key in range(n_atoms):
+        neighbors = neighbor_dict[key] + [key]
+        sp_array_frame2[key] = np.mean(array[neighbors, result_index])
 
-    return frame, sp_array_frame_1
+    return result_index, sp_array_frame2
 
 
 def spatialaverage(
-    universe: MDAnalysis.Universe,
+    universe: Universe,
     descriptor_array: NDArray[np.float64],
     selection: str,
     cutoff: float,
-    traj_cut: int = 0,
-    num_processes: int = 4,
+    trajslice: slice | None = None,
+    num_processes: int = 1,
 ) -> NDArray[np.float64]:
     """Compute spatially averaged descriptor values over neighboring particles.
 
@@ -142,11 +146,11 @@ def spatialaverage(
             descriptor = np.load('descriptor_array.npy')
 
             averaged_values = spatialaverage(
-                                universe=u,
-                                descriptor_array=descriptor,
-                                selection='name CA',
-                                cutoff=5.0,
-                                num_processes=8)
+                universe=u,
+                descriptor_array=descriptor,
+                selection='name CA',
+                cutoff=5.0,
+                num_processes=8)
 
         This example computes the spatial averages of the descriptor values
         for atoms selected as `CA` atoms, within a cutoff of 5.0 units, using 8
@@ -155,6 +159,8 @@ def spatialaverage(
         to set up the Universe.
     """
     selection = universe.select_atoms(selection)
+    if trajslice is None:
+        trajslice = slice(None)
 
     shape = descriptor_array.shape
     dtype = descriptor_array.dtype
@@ -165,45 +171,40 @@ def spatialaverage(
     shared_array_np = np.frombuffer(shared_array, dtype=dtype).reshape(shape)  # type: ignore[call-overload]
 
     np.copyto(shared_array_np, descriptor_array)
+
     two_dim = 2
     three_dim = 3
+    sp_array = np.zeros(descriptor_array.shape)
     if descriptor_array.ndim == two_dim:
-        sp_array_2 = np.zeros(
-            (descriptor_array.shape[0], descriptor_array.shape[1])
-        )
         is_vector = False
     elif descriptor_array.ndim == three_dim:
-        sp_array_3 = np.zeros(
-            (
-                descriptor_array.shape[0],
-                descriptor_array.shape[1],
-                descriptor_array.shape[2],
-            )
-        )
         is_vector = True
     else:
-        error_string = "INVALID ARRAY SHAPE"
-        raise ValueError(error_string)
+        msg = "descriptor_array must have ndim == 2 or ndim == 3."
+        raise ValueError(msg)
 
-    num_frames = len(universe.trajectory) - traj_cut
     pool = Pool(
         processes=num_processes,
         initializer=initworker,
         initargs=(shared_array, shape, dtype),
     )
+
+    frame_indices = list(
+        range(*trajslice.indices(universe.trajectory.n_frames))
+    )
     args = [
-        (universe, selection, cutoff, frame, is_vector)
-        for frame in range(num_frames)
+        (universe, selection, cutoff, traj_frame, i, is_vector)
+        for i, traj_frame in enumerate(frame_indices)
     ]
     results = pool.map(processframe, args)
     pool.close()
     pool.join()
 
     if is_vector:
         for frame, sp_array_frame in results:
-            sp_array_3[:, frame, :] = sp_array_frame
-        return sp_array_3
+            sp_array[:, frame, :] = sp_array_frame
+    else:
+        for frame, sp_array_frame in results:
+            sp_array[:, frame] = sp_array_frame
 
-    for frame, sp_array_frame in results:
-        sp_array_2[:, frame] = sp_array_frame
-    return sp_array_2
+    return sp_array

src/dynsight/_internal/analysis/time_correlations.py

@@ -22,8 +22,6 @@ def self_time_correlation(
     Returns the self-TCF averaged over all the particles, normalized so that
     the value at t = 0 is equal to 1.
 
-    * Author: Matteo Becchi
-
     Parameters:
         data:
             Array of shape (n_particles, n_frames).
@@ -62,8 +60,9 @@ def self_time_correlation(
 
     """
     n_part, n_frames = data.shape
+    data2 = data.copy()
     # Subtract the mean for each particle
-    data = data - np.mean(data, axis=1, keepdims=True)
+    data2 = data2 - np.mean(data2, axis=1, keepdims=True)
 
     if max_delay is None:
         max_delay = n_frames
@@ -79,7 +78,7 @@ def self_time_correlation(
                 tmp[:valid_t],
                 tmp[t_prime : valid_t + t_prime],
             )
-            for tmp in data
+            for tmp in data2
         ]
         correlation[t_prime] = np.mean(corr_sum) / valid_t
         correlation_error[t_prime] = np.std(corr_sum) / (
@@ -105,8 +104,6 @@ def cross_time_correlation(
     For each particles pair, computes the cross time correlation function
     (cross-TCF). Returns the cross-TCF averaged over all the particles pairs.
 
-    * Author: Matteo Becchi
-
     Parameters:
         data:
             Array of shape (n_particles, n_frames).

src/dynsight/_internal/soapify/saponify.py

@@ -22,6 +22,7 @@ def saponify_trajectory(
     soap_respectpbc: bool = True,
     n_core: int = 1,
     centers: str = "all",
+    trajslice: slice | None = None,
 ) -> NDArray[np.float64]:
     """Calculate the SOAP fingerprints for each atom in a MDA universe.
 
@@ -53,6 +54,8 @@ def saponify_trajectory(
             (option passed to the desired SOAP engine). Defaults to True.
         n_core:
             Number of core used for parallel processing. Default to 1.
+        trajslice:
+            The slice of the trajectory to consider. Defaults to slice(None).
 
     Returns:
         The SOAP spectra for all the particles and frames. np.ndarray of shape
@@ -83,6 +86,8 @@ def saponify_trajectory(
             assert np.isclose(
                 np.sum(soap[0]), 8627.847941030795, atol=1e-6, rtol=1e-3)
     """
+    if trajslice is None:
+        trajslice = slice(None)
     sel = universe.select_atoms(selection)
     centers_list_id = sel.select_atoms(centers).indices.tolist()
     centers_list = [
@@ -98,7 +103,7 @@ def saponify_trajectory(
         periodic=soap_respectpbc,
     )
     traj = []
-    for t_s in universe.trajectory:
+    for t_s in universe.trajectory[trajslice]:
         pos = sel.atoms.positions
         symbols = sel.atoms.types
         box = t_s.dimensions