qdiffusivity 0.1.0__tar.gz → 0.2.0__tar.gz

{qdiffusivity-0.1.0 → qdiffusivity-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: qdiffusivity
-Version: 0.1.0
+Version: 0.2.0
 Summary: Quantile-based estimation of local diffusivities in molecular dynamics simulations of nanoconfined liquids
 Author-email: Shern Ren Tee <shernren@gmail.com>
 Maintainer-email: Shern Ren Tee <shernren@gmail.com>

{qdiffusivity-0.1.0 → qdiffusivity-0.2.0}/qdiffusivity/__init__.py

@@ -9,12 +9,14 @@ from importlib.metadata import version
 
 from .binned import (
     LocalDiffusivityQBinned,
-    TransverseDensityQBinned,
+    TransverseMassDensityQBinned,
+    TransverseNumDensityQBinned,
     cic_assign,
     resolve_bins,
 )
 from .density import (
-    TransverseDensityQKDE,
+    TransverseMassDensityQKDE,
+    TransverseNumDensityQKDE,
     epanechnikov_kernel,
     kde_1d,
     select_bandwidth,
@@ -32,10 +34,12 @@ from .diffusivity import (
 __version__ = version("qdiffusivity")
 
 __all__ = [
-    "TransverseDensityQBinned",
-    "TransverseDensityQKDE",
     "LocalDiffusivityQBinned",
     "LocalDiffusivityQKDE",
+    "TransverseMassDensityQBinned",
+    "TransverseMassDensityQKDE",
+    "TransverseNumDensityQBinned",
+    "TransverseNumDensityQKDE",
     "build_cdf",
     "cic_assign",
     "epanechnikov_kernel",

{qdiffusivity-0.1.0 → qdiffusivity-0.2.0}/qdiffusivity/binned.py

@@ -30,8 +30,8 @@ from __future__ import annotations
 
 import numpy as np
 from MDAnalysis.analysis.base import AnalysisBase
-from MDAnalysis.transformations import NoJump
 
+from .density import _AMU_PER_ANGSTROM3_TO_G_PER_CM3
 from .diffusivity import build_cdf
 
 
@@ -138,14 +138,10 @@ def _bin_centers_from_edges(edges):
     return 0.5 * (edges[:-1] + edges[1:])
 
 
-class TransverseDensityQBinned(AnalysisBase):
-    r"""CDF-binned transverse number-density profile.
+class _TransverseDensityQBinnedBase(AnalysisBase):
+    r"""Shared base for CDF-binned transverse density profiles.
 
-    Pools per-frame positions of ``atomgroup`` along ``dim`` (the
-    confined axis), builds a CDF-uniformised u-space, and assigns each
-    sample to u-space bins via cloud-in-cell (CIC) or hard assignment.
-    The bin population is converted to a number density via the
-    cross-sectional area and the number of analysis frames.
+    Subclasses set ``_mass_density`` to select number or mass density.
 
     Parameters
     ----------
@@ -175,15 +171,12 @@ class TransverseDensityQBinned(AnalysisBase):
     u_centers : numpy.ndarray
         ``(n_bins,)`` bin centres in u-space.
     density : numpy.ndarray
-        ``(n_bins,)`` number density (particles per volume), computed
-        as ``(N_total / (n_frames_used * A)) * (bin_population /
-        bin_width_in_u) * rho(z_center)`` — i.e. the CDF bin population
-        rescaled by the local Jacobian :math:`\rho = du/dz`.
+        ``(n_bins,)`` density profile.  For number density:
+        particles/ų.  For mass density: g/cm³.
     n_per_bin : numpy.ndarray
         ``(n_bins,)`` raw (weighted) population per bin.
     n_eff : numpy.ndarray
-        ``(n_bins,)`` effective sample size (sum of CIC weights, or
-        integer count for hard assignment).
+        ``(n_bins,)`` effective sample size.
     bin_edges_u : numpy.ndarray
         ``(n_bins + 1,)`` u-space bin edges.
     n_total : int
@@ -194,29 +187,10 @@ class TransverseDensityQBinned(AnalysisBase):
         Mean number of samples per frame.
     P, P_inv, rho : callables
         CDF, inverse-CDF and equilibrium density closures.
-
-    Notes
-    -----
-    Confined-axis positions are wrapped into the primary box cell along
-    ``dim`` before pooling.  For ``grouping="residues"`` the
-    centre-of-mass is computed with the atoms' topology masses.
-
-    Examples
-    --------
-    ::
-
-        import MDAnalysis as mda
-        from qdiffusivity import TransverseDensityQBinned
-
-        u = mda.Universe("topology.data", "trajectory.xtc")
-        ag = u.select_atoms("type 1 2")
-        binned = TransverseDensityQBinned(
-            ag, dim=2, z_bot=10.0, z_top=90.0, bins=30,
-        )
-        binned.run()
-        # binned.density is in particles/Å^3.
     """
 
+    _mass_density = False
+
     def __init__(
         self,
         atomgroup,
@@ -242,13 +216,13 @@ class TransverseDensityQBinned(AnalysisBase):
         self._grouping = grouping
 
     def _prepare(self):
+        self._masses = self._ag.masses.astype(np.float64)
         if self._grouping == "residues":
             residx = self._ag.resindices
             self._res_unique, self._res_inv = np.unique(
                 residx, return_inverse=True
             )
             self._n_res = self._res_unique.size
-            self._masses = self._ag.masses.astype(np.float64)
             self._res_mass = np.zeros(self._n_res, dtype=np.float64)
             np.add.at(self._res_mass, self._res_inv, self._masses)
         self._u_frames = []
@@ -282,7 +256,14 @@ class TransverseDensityQBinned(AnalysisBase):
             self.n_total / self.n_frames_used if self.n_frames_used > 0 else 0.0
         )
 
-        self.P, self.P_inv, self.rho, _, _, _ = build_cdf(z_frames)
+        (
+            self.P,
+            self.P_inv,
+            self.rho,
+            self.rho_prime,
+            _,
+            _,
+        ) = build_cdf(z_frames)
         self.u_centers = _bin_centers_from_edges(self._edges_u)
         self.z_centers = self.P_inv(self.u_centers)
         n_bins = self._n_bins
@@ -302,8 +283,6 @@ class TransverseDensityQBinned(AnalysisBase):
             np.add.at(pop, k1, w1)
             self.n_eff = pop.copy()
         else:
-            # Hard assignment via digitize.  Samples exactly on the
-            # upper edge (u == 1) are folded into the last bin.
             idx = np.clip(
                 np.digitize(u_samples, self._edges_u) - 1,
                 0,
@@ -314,26 +293,109 @@ class TransverseDensityQBinned(AnalysisBase):
 
         self.n_per_bin = pop
 
-        # Convert bin population to a number density.
-        #   n(z) = (N_total / (n_frames * A)) * rho(z)
-        # because in u-space the population per bin is ~N_total * du,
-        # and du = rho(z) * dz, so the z-space number density is
-        # (N_total / (n_frames * A)) * (pop / N_total) / dz
-        # = pop / (n_frames * A * dz), with dz = du / rho(z).
-        A = 1.0
+        # Cross-sectional area.
         dims = self._ag.universe.dimensions
-        if dims is not None:
-            para = [i for i in (0, 1, 2) if i != self._dim]
-            A = float(dims[para[0]] * dims[para[1]])
+        para = [i for i in (0, 1, 2) if i != self._dim]
+        A = float(dims[para[0]] * dims[para[1]])
+
         rho_centers = self.rho(self.z_centers)
         du = np.diff(self._edges_u)
         with np.errstate(divide="ignore", invalid="ignore"):
             dz = np.where(rho_centers > 0, du / rho_centers, 0.0)
-            self.density = np.where(
-                dz > 0,
-                pop / (self.n_frames_used * A * dz),
-                0.0,
-            )
+
+            if self._mass_density:
+                # Mass-weighted population: each sample's CIC/hard
+                # weight is multiplied by its mass.  For CIC we
+                # re-accumulate with mass weights; for hard assignment
+                # we use np.bincount with weights.
+                if self._use_cic:
+                    mass_per_sample = (
+                        np.tile(self._masses, self.n_frames_used)
+                        if self._grouping == "atoms"
+                        else np.repeat(self._res_mass, self.n_frames_used)
+                    )
+                    mass_pop = np.zeros(n_bins, dtype=np.float64)
+                    np.add.at(mass_pop, k0, w0 * mass_per_sample)
+                    np.add.at(mass_pop, k1, w1 * mass_per_sample)
+                else:
+                    mass_per_sample = (
+                        np.tile(self._masses, self.n_frames_used)
+                        if self._grouping == "atoms"
+                        else np.repeat(self._res_mass, self.n_frames_used)
+                    )
+                    mass_pop = np.bincount(
+                        idx,
+                        weights=mass_per_sample,
+                        minlength=n_bins,
+                    ).astype(np.float64)
+                # Mass density (amu/Å^3) -> g/cm^3.
+                self.density = np.where(
+                    dz > 0,
+                    mass_pop
+                    / (self.n_frames_used * A * dz)
+                    / _AMU_PER_ANGSTROM3_TO_G_PER_CM3,
+                    0.0,
+                )
+            else:
+                self.density = np.where(
+                    dz > 0,
+                    pop / (self.n_frames_used * A * dz),
+                    0.0,
+                )
+
+
+class TransverseNumDensityQBinned(_TransverseDensityQBinnedBase):
+    r"""CDF-binned transverse number-density profile.
+
+    Returns number density in particles/ų.  See
+    :class:`_TransverseDensityQBinnedBase` for parameters and attributes.
+
+    Examples
+    --------
+    ::
+
+        import MDAnalysis as mda
+        from qdiffusivity import TransverseNumDensityQBinned
+
+        u = mda.Universe("topology.data", "trajectory.xtc")
+        ag = u.select_atoms("type 1 2")
+        binned = TransverseNumDensityQBinned(
+            ag, dim=2, z_bot=10.0, z_top=90.0, bins=30,
+        )
+        binned.run()
+        # binned.density is in particles/Å^3.
+    """
+
+    _mass_density = False
+
+
+class TransverseMassDensityQBinned(_TransverseDensityQBinnedBase):
+    r"""CDF-binned transverse mass-density profile.
+
+    Returns mass density in g/cm³ (matching MDAnalysis's
+    :class:`~MDAnalysis.analysis.lineardensity.LinearDensity`).  Each
+    sample is weighted by its topology mass (per atom for
+    ``grouping="atoms"``; per-residue total mass for
+    ``grouping="residues"``).  See
+    :class:`_TransverseDensityQBinnedBase` for parameters and attributes.
+
+    Examples
+    --------
+    ::
+
+        import MDAnalysis as mda
+        from qdiffusivity import TransverseMassDensityQBinned
+
+        u = mda.Universe("topology.data", "trajectory.xtc")
+        ag = u.select_atoms("type 1 2")
+        binned = TransverseMassDensityQBinned(
+            ag, dim=2, z_bot=10.0, z_top=90.0, bins=30,
+        )
+        binned.run()
+        # binned.density is in g/cm^3.
+    """
+
+    _mass_density = True
 
 
 class LocalDiffusivityQBinned(AnalysisBase):
@@ -373,6 +435,17 @@ class LocalDiffusivityQBinned(AnalysisBase):
         :math:`\frac{\Delta t}{2}\Phi(z)^2` (with
         :math:`\Phi = D\,\rho'/\rho` in the isothermal convention) from
         the perpendicular estimator.  Default ``False``.
+    density_result : TransverseNumDensityQKDE or None, optional
+        A pre-computed density result whose CDF closures
+        (:attr:`P`, :attr:`P_inv`, :attr:`rho`, :attr:`rho_prime`)
+        are reused for the u-space mapping.  If ``None`` (default),
+        a :class:`TransverseNumDensityQKDE` is run internally in
+        ``_prepare()`` to build the CDF automatically.  Passing a
+        pre-computed result enables a **two-pass parallelization**
+        strategy: run the density analysis first (parallelizable via
+        split-apply-combine), then pass its result to the diffusivity
+        analysis (each frame's u-mapping is then stateless and
+        parallelizable).
 
     Attributes
     ----------
@@ -407,14 +480,18 @@ class LocalDiffusivityQBinned(AnalysisBase):
         Time step used in the estimator.
     P, P_inv, rho : callables
         CDF, inverse-CDF and equilibrium density closures.
+    rho_prime : callable
+        Spatial derivative of the equilibrium density.
 
     Notes
     -----
-    Parallel (x/y) positions are unwrapped with MDAnalysis's ``NoJump``
-    transformation so that multi-frame displacements across periodic
-    boundaries are captured.  The transformation is attached in place
-    (guarded so it is only attached once per Universe) and the full
-    trajectory is iterated from frame 0 (``NoJump`` is stateful).
+    Multi-frame displacements across periodic boundaries are computed
+    with the **minimum-image convention** applied directly to the
+    stored wrapped positions.  This is a stateless, frame-local
+    operation — only consecutive-frame positions are needed — so the
+    estimator is compatible with split-apply-combine parallelization
+    (each worker block can be processed independently without missing
+    prior image corrections).
 
     Examples
     --------
@@ -439,6 +516,7 @@ class LocalDiffusivityQBinned(AnalysisBase):
         bins=30,
         dt=None,
         ito_correction: bool = False,
+        density_result=None,
     ):
         super().__init__(atomgroup.universe.trajectory)
         self._ag = atomgroup
@@ -449,8 +527,26 @@ class LocalDiffusivityQBinned(AnalysisBase):
         self._n_bins, self._edges_u, self._use_cic = resolve_bins(bins)
         self._dt = dt
         self._ito_correction = bool(ito_correction)
+        self._density_result = density_result
 
     def _prepare(self):
+        if self._density_result is not None:
+            dr = self._density_result
+            self.P = dr.P
+            self.P_inv = dr.P_inv
+            self.rho = dr.rho
+            self.rho_prime = dr.rho_prime
+        else:
+            dens = TransverseNumDensityQBinned(
+                self._ag,
+                dim=self._dim,
+                bins=self._n_bins,
+            )
+            dens.run()
+            self.P = dens.P
+            self.P_inv = dens.P_inv
+            self.rho = dens.rho
+            self.rho_prime = dens.rho_prime
         self._pos_frames = []
 
     def _single_frame(self):
@@ -470,40 +566,44 @@ class LocalDiffusivityQBinned(AnalysisBase):
         if dt <= 0:
             raise ValueError(f"dt must be positive, got {dt}")
 
-        # Unwrap parallel coordinates via NoJump.
-        u = self._ag.universe
-        if not getattr(u, "_qdiff_nojump_applied", False):
-            u.trajectory[0]
-            u.trajectory.add_transformations(NoJump(self._ag))
-            u._qdiff_nojump_applied = True
-
-        pos_unwrapped = np.empty_like(pos)
-        Lz = None
-        k = 0
-        for idx in range(0, u.trajectory.n_frames):
-            u.trajectory[idx]
-            if k >= n_frames:
-                break
-            pos_unwrapped[k] = self._ag.positions
-            Lz = u.dimensions[self._dim]
-            k += 1
-        if Lz is not None and Lz > 0:
-            pos_unwrapped[:, :, self._dim] %= Lz
-
-        # Build CDF from pooled wrapped-z.
-        z_pooled = pos_unwrapped[:, :, self._dim].ravel()
-        self.P, self.P_inv, self.rho, self.rho_prime, _, _ = build_cdf(z_pooled)
+        # Box lengths from the last analysed frame (constant-box / NVT
+        # assumption).  Minimum-image correction is applied to the
+        # *displacement* of consecutive frames, not to the absolute
+        # positions — stateless and parallelization-friendly.
+        box = self._ts.dimensions
+        Lz = float(box[self._dim])
+        Lpara = {d: float(box[d]) for d in self._para_dims}
+
+        # Wrap the confined coordinate into [0, Lz] for CDF construction.
+        if Lz > 0:
+            pos[:, :, self._dim] %= Lz
+
+        # Build CDF from pooled wrapped-z — unless a pre-computed
+        # CDF was supplied (or auto-run in _prepare).
+        if self._density_result is None:
+            z_pooled = pos[:, :, self._dim].ravel()
+            (
+                self.P,
+                self.P_inv,
+                self.rho,
+                self.rho_prime,
+                _,
+                _,
+            ) = build_cdf(z_pooled)
 
         self.u_centers = _bin_centers_from_edges(self._edges_u)
         self.z_centers = self.P_inv(self.u_centers)
         n_bins = self._n_bins
 
         # u-positions of all increments (starting frame of each pair).
-        z_start = pos_unwrapped[:-1, :, self._dim]
+        z_start = pos[:-1, :, self._dim]
         u_start = self.P(z_start.ravel())
 
         # Perpendicular: z-space local estimator (Δz)²/(2Δt).
-        dz = np.diff(pos_unwrapped[:, :, self._dim], axis=0)
+        # Minimum-image correction on the wrapped-z displacement.
+        dz = np.diff(pos[:, :, self._dim], axis=0)
+        if Lz > 0:
+            dz -= np.rint(dz / Lz) * Lz
         d_perp_local = (dz**2) / (2.0 * dt)
         self.n_increments = int(d_perp_local.size)
 
@@ -532,10 +632,14 @@ class LocalDiffusivityQBinned(AnalysisBase):
 
         self.D_perp = np.clip(self.D_perp, 0, None)
 
-        # Parallel: Δx² + Δy², only z_start mapped to u.
+        # Parallel: Δx² + Δy², only z_start mapped to u.  Minimum-image
+        # correction on each parallel displacement.
         d_para_acc = np.zeros_like(dz)
         for d in self._para_dims:
-            diff_d = np.diff(pos_unwrapped[:, :, d], axis=0)
+            diff_d = np.diff(pos[:, :, d], axis=0)
+            Ld = Lpara[d]
+            if Ld > 0:
+                diff_d -= np.rint(diff_d / Ld) * Ld
             d_para_acc += diff_d**2
         d_para_local = d_para_acc / (4.0 * dt)