Dask timeseries prototype

.gitignore

@@ -52,3 +52,6 @@ benchmarks/results
 .idea
 .vscode
 *.lock
+
+# dev
+tmp/

package/MDAnalysis/analysis/dasktimeseries.py

@@ -0,0 +1,49 @@
+from .results import Results, ResultsGroup
+import dask.array as da
+import numpy as np
+
+
+class DaskTimeSeriesAnalysisBase:
+
+    def __init__(self, dask_timeseries, verbose=False, **kwargs):
+        self._dts = dask_timeseries
+        self._verbose = verbose
+        self.results = Results()
+
+    def _prepare(self):
+        pass  # pylint: disable=unnecessary-pass
+
+    def _compute(self):
+        pass
+
+    def _conclude(self):
+        pass  # pylint: disable=unnecessary-pass
+
+    def run(self):
+        self._prepare()
+        self._compute()
+        self._conclude()
+        return self
+
+
+class DaskRMSF(DaskTimeSeriesAnalysisBase):
+    def __init__(self, dask_timeseries, verbose=False, **kwargs):
+        super().__init__(dask_timeseries, verbose=verbose, **kwargs)
+
+    def _prepare(self):
+        n_atoms = len(self._dts[0])
+        self.results["rmsf"] = np.zeros((n_atoms, 3))
+
+    def _compute(self):
+        positions = self._dts
+        mean_positions = positions.mean(axis=0)
+        subtracted_positions = positions - mean_positions
+        squared_deviations = subtracted_positions**2
+        avg_squared_deviations = squared_deviations.mean(axis=0)
+        sqrt_avg_squared_deviations = da.sqrt(avg_squared_deviations)
+        self.results.rmsf = da.sqrt(
+            (sqrt_avg_squared_deviations**2).sum(axis=1)
+        ).compute()
+
+    def _conclude(self):
+        pass

package/MDAnalysis/coordinates/H5MD.py

@@ -208,7 +208,9 @@
 
 """
 import warnings
-
+import dask.array as da
+from .. import units
+from typing import Any, Union, Optional, List, Dict
 import numpy as np
 import MDAnalysis as mda
 from . import base, core
@@ -810,6 +812,57 @@
         kwargs.setdefault('forces', self.has_forces)
         return H5MDWriter(filename, n_atoms, **kwargs)
 
+    def dask_timeseries(self, asel: Optional['AtomGroup']=None,
+                   atomgroup: Optional['Atomgroup']=None,
+                   start: Optional[int]=None, stop: Optional[int]=None,
+                   step: Optional[int]=None,
+                   order: Optional[str]='fac') -> np.ndarray:
+        if asel is not None:
+            warnings.warn(
+                "asel argument to timeseries will be renamed to"
+                "'atomgroup' in 3.0, see #3911",
+                category=DeprecationWarning)
+            if atomgroup:
+                raise ValueError("Cannot provide both asel and atomgroup kwargs")
+            atomgroup = asel
+        start, stop, step = self.check_slice_indices(start, stop, step)
+        nframes = len(range(start, stop, step))
+
+        if atomgroup is not None:
+            if len(atomgroup) == 0:
+                raise ValueError(
+                    "Timeseries requires at least one atom to analyze")
+            atom_numbers = atomgroup.indices
+            natoms = len(atom_numbers)
+        else:
+            natoms = self.n_atoms
+            atom_numbers = np.arange(natoms)
+
+        coordinates = da.from_array(self._particle_group['position']['value'],)[start:stop:step, atom_numbers, :]
+
+        # switch axes around
+        default_order = 'fac'
+        if order != default_order:
+            try:
+                newidx = [default_order.index(i) for i in order]
+            except ValueError:
+                raise ValueError(f"Unrecognized order key in {order}, "
+                                 "must be permutation of 'fac'")
+
+            try:
+                coordinates = da.moveaxis(coordinates, newidx, [0, 1, 2])
+            except ValueError:
+                errmsg = ("Repeated or missing keys passed to argument "
+                          f"`order`: {order}, each key must be used once")
+                raise ValueError(errmsg)
+
+        f = units.get_conversion_factor('length',
+                                        self.units['length'], 'Angstrom')
+        coordinates *= f
+
+        return coordinates
+
+
     @property
     def has_positions(self):
         """``True`` if 'position' group is in trajectory."""

tmp/env.yaml

@@ -0,0 +1,43 @@
+name: dask-timeseries-dev
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - chemfiles>=0.10
+  - codecov
+  - cython
+  - dask
+  - docutils
+  - fasteners
+  - griddataformats
+  - gsd
+  - h5py>=2.10
+  - hypothesis
+  - ipykernel
+  - joblib>=0.12
+  - mdanalysis-sphinx-theme >=1.3.0
+  - matplotlib>=3.2.2
+  - mmtf-python
+  - mock
+  - networkx
+  - numpy>=1.23.2
+  - pytest
+  - python==3.10
+  - pytng>=0.2.3
+  - scikit-learn
+  - scipy
+  - pip
+  - sphinx <7.0
+  - tidynamics>=1.0.0
+  - tqdm>=4.43.0
+  - sphinxcontrib-bibtex
+  - mdaencore
+  - waterdynamics
+  - pip:
+      - mdahole2
+      - pathsimanalysis
+      - duecredit
+      - parmed
+      - sphinx-sitemap
+      - packaging
+      - pyedr>=0.7.0