kinetic-st 0.1.0__tar.gz

kinetic_st-0.1.0/PKG-INFO

@@ -0,0 +1,76 @@
+Metadata-Version: 2.4
+Name: kinetic-st
+Version: 0.1.0
+Summary: Kinetic-theory framework for multiscale modelling of spatial transcriptomics
+Author: KineticST Project Team
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/whirlsyu/kinetic_st
+Project-URL: Repository, https://github.com/whirlsyu/kinetic_st.git
+Project-URL: Issues, https://github.com/whirlsyu/kinetic_st/issues
+Keywords: spatial-transcriptomics,kinetic-theory,bioinformatics,single-cell
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Requires-Dist: scikit-learn>=1.3
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: matplotlib>=3.7; extra == "dev"
+Provides-Extra: demo
+Requires-Dist: matplotlib>=3.7; extra == "demo"
+
+# KineticST
+
+Kinetic-theory framework for multiscale modelling of spatial transcriptomics.
+
+Bridges microscopic cell-cell interactions to macroscopic tissue organisation
+via a Boltzmann-like collision operator on latent cell states, integrated on
+a spatial graph with stochastic dynamics and invariant-measure analysis.
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Quick start
+
+```python
+import numpy as np
+from kinetic_st import KineticST, KineticSTConfig, generate_kinetic_simulation
+
+# Synthetic data
+sim = generate_kinetic_simulation(n_types=6, n_spots=1600, seed=42)
+
+# Fit
+model = KineticST(KineticSTConfig(n_steps=150, latent_dim=10, verbose=True))
+model.fit(sim.Y, sim.coords)
+
+# Results
+labels = model.get_domain_labels()
+field = model.get_expression_field()
+print(model.summary())
+```
+
+## Pipeline stages
+
+1. **Latent embedding** — PCA on log-CPM expression to a low-dimensional cell-state space
+2. **Spatial graph** — kNN adjacency from coordinates
+3. **Collision kernel** — ligand-receptor rule with irreversible attractor bias
+4. **Kinetic integration** — deterministic (Boltzmann), stochastic (Euler-Maruyama), or time-sliced solver
+5. **Moment closure** — macroscopic tissue expression fields via local KDE moments
+6. **Invariant measure** — non-equilibrium steady-state estimation (KDE or GMM)
+7. **Spatial domains** — segmentation via local invariant-measure divergence
+8. **Supervised markers** — one-vs-rest z-scores on the dynamics-smoothed field
+
+## License
+
+MIT

kinetic_st-0.1.0/README.md

@@ -0,0 +1,47 @@
+# KineticST
+
+Kinetic-theory framework for multiscale modelling of spatial transcriptomics.
+
+Bridges microscopic cell-cell interactions to macroscopic tissue organisation
+via a Boltzmann-like collision operator on latent cell states, integrated on
+a spatial graph with stochastic dynamics and invariant-measure analysis.
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Quick start
+
+```python
+import numpy as np
+from kinetic_st import KineticST, KineticSTConfig, generate_kinetic_simulation
+
+# Synthetic data
+sim = generate_kinetic_simulation(n_types=6, n_spots=1600, seed=42)
+
+# Fit
+model = KineticST(KineticSTConfig(n_steps=150, latent_dim=10, verbose=True))
+model.fit(sim.Y, sim.coords)
+
+# Results
+labels = model.get_domain_labels()
+field = model.get_expression_field()
+print(model.summary())
+```
+
+## Pipeline stages
+
+1. **Latent embedding** — PCA on log-CPM expression to a low-dimensional cell-state space
+2. **Spatial graph** — kNN adjacency from coordinates
+3. **Collision kernel** — ligand-receptor rule with irreversible attractor bias
+4. **Kinetic integration** — deterministic (Boltzmann), stochastic (Euler-Maruyama), or time-sliced solver
+5. **Moment closure** — macroscopic tissue expression fields via local KDE moments
+6. **Invariant measure** — non-equilibrium steady-state estimation (KDE or GMM)
+7. **Spatial domains** — segmentation via local invariant-measure divergence
+8. **Supervised markers** — one-vs-rest z-scores on the dynamics-smoothed field
+
+## License
+
+MIT

kinetic_st-0.1.0/pyproject.toml

@@ -0,0 +1,45 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kinetic-st"
+version = "0.1.0"
+description = "Kinetic-theory framework for multiscale modelling of spatial transcriptomics"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.9"
+authors = [
+    {name = "KineticST Project Team"},
+]
+keywords = ["spatial-transcriptomics", "kinetic-theory", "bioinformatics", "single-cell"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+dependencies = [
+    "numpy>=1.24",
+    "scipy>=1.10",
+    "scikit-learn>=1.3",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0", "matplotlib>=3.7"]
+demo = ["matplotlib>=3.7"]
+
+[project.urls]
+Homepage = "https://github.com/whirlsyu/kinetic_st"
+Repository = "https://github.com/whirlsyu/kinetic_st.git"
+Issues = "https://github.com/whirlsyu/kinetic_st/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+kinetic_st = ["py.typed"]

kinetic_st-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

kinetic_st-0.1.0/src/kinetic_st/__init__.py

@@ -0,0 +1,34 @@
+"""KineticST: kinetic-theory framework for spatial transcriptomics.
+
+Bridges microscopic cell-cell interactions to macroscopic tissue organisation
+via a Boltzmann-like collision operator on latent cell states, integrated on
+a spatial graph with stochastic dynamics and invariant-measure analysis.
+"""
+
+__version__ = "0.1.0"
+__author__ = "KineticST Project Team"
+
+from kinetic_st.core.latent import LatentEmbedding, fit_latent
+from kinetic_st.core.collision import CollisionKernel, LigandReceptorRule
+from kinetic_st.core.boltzmann import BoltzmannSolver, solve_kinetic
+from kinetic_st.core.fluid_limit import FluidLimiter, compute_moments
+from kinetic_st.core.stochastic_pde import StochasticKineticSolver, InvariantMeasureEstimator
+from kinetic_st.core.time_slicing import TimeSlicer
+from kinetic_st.core.domains import DomainIdentifier
+from kinetic_st.core.wave_turbulence import WaveKineticAnalyzer
+from kinetic_st.core.simulator import TissueSimulator
+from kinetic_st.core.markers import KineticMarkerFinder, find_kinetic_markers, markers_from_model
+from kinetic_st.data.simulation import generate_kinetic_simulation
+from kinetic_st.pipeline import KineticST, KineticSTConfig
+
+__all__ = [
+    "KineticST", "KineticSTConfig",
+    "LatentEmbedding", "fit_latent",
+    "CollisionKernel", "LigandReceptorRule",
+    "BoltzmannSolver", "solve_kinetic",
+    "FluidLimiter", "compute_moments",
+    "StochasticKineticSolver", "InvariantMeasureEstimator",
+    "TimeSlicer", "DomainIdentifier", "WaveKineticAnalyzer",
+    "TissueSimulator", "KineticMarkerFinder", "find_kinetic_markers",
+    "markers_from_model", "generate_kinetic_simulation",
+]

kinetic_st-0.1.0/src/kinetic_st/core/__init__.py

@@ -0,0 +1 @@
+# Core modules

kinetic_st-0.1.0/src/kinetic_st/core/boltzmann.py

@@ -0,0 +1,87 @@
+"""Mesoscopic kinetic equation solver via operator splitting.
+
+Integrates df/dt = jump + (1/tau)*Q(f,f) on a spatial graph, where Q is
+the collision integral realized by CollisionKernel.
+"""
+
+import numpy as np
+from scipy import sparse
+from typing import Optional, Dict, Callable
+from dataclasses import dataclass, field
+
+from kinetic_st.core.collision import CollisionKernel
+
+
+@dataclass
+class KineticConfig:
+    """Solver configuration."""
+    tau: float = 1.0
+    jump_rate: float = 0.1
+    dt: float = 0.05
+    n_steps: int = 200
+    drift_field: Optional[Callable] = None
+    n_pairs_per_spot: int = 4
+    record_every: int = 10
+    seed: int = 0
+
+
+@dataclass
+class KineticTrajectory:
+    """Recorded trajectory."""
+    states: np.ndarray              # (n_spots, d) final states
+    history: np.ndarray             # (T, n_spots, d)
+    times: np.ndarray               # (T,)
+    info: Dict = field(default_factory=dict)
+
+
+class BoltzmannSolver:
+    """Integrate the cell-state kinetic equation via operator splitting.
+
+    Steps per time step: streaming (spatial jump) -> collision -> drift.
+    """
+
+    def __init__(self, kernel: CollisionKernel, config: Optional[KineticConfig] = None):
+        self.kernel = kernel
+        self.config = config or KineticConfig()
+        self.trajectory_ = None
+
+    def solve(self, states0, adjacency, Y_lr=None, verbose=False) -> KineticTrajectory:
+        """Integrate forward in time."""
+        cfg = self.config
+        rng = np.random.RandomState(cfg.seed)
+        states = np.array(states0, dtype=np.float64, copy=True)
+        n, d = states.shape
+
+        A = adjacency.tocsr().astype(np.float64)
+        deg = np.maximum(np.asarray(A.sum(axis=1)).flatten(), 1.0)
+        P = (sparse.diags(1.0 / deg) @ A).tocsr()
+
+        history = [states.copy()]
+        times = [0.0]
+        for step in range(cfg.n_steps):
+            t = (step + 1) * cfg.dt
+            # Streaming
+            states += cfg.dt * cfg.jump_rate * (P @ states - states)
+            # Collision
+            collided = self.kernel.apply_collisions(
+                states, adjacency, Y_lr=Y_lr,
+                n_pairs_per_spot=cfg.n_pairs_per_spot, rng=rng)
+            states += (cfg.dt / cfg.tau) * (collided - states)
+            # Drift
+            if cfg.drift_field is not None:
+                states += cfg.dt * cfg.drift_field(states, t)
+            if (step + 1) % cfg.record_every == 0:
+                history.append(states.copy())
+                times.append(t)
+                if verbose and (step + 1) % (cfg.record_every * 10) == 0:
+                    print(f"  [Boltzmann] t={t:.3f} step={step+1}/{cfg.n_steps}")
+
+        self.trajectory_ = KineticTrajectory(
+            states=states, history=np.array(history), times=np.array(times),
+            info={"n_steps": cfg.n_steps, "dt": cfg.dt, "n_spots": n, "dim": d})
+        return self.trajectory_
+
+
+def solve_kinetic(states0, adjacency, kernel, config=None, Y_lr=None) -> KineticTrajectory:
+    """Functional API for BoltzmannSolver.solve."""
+    return BoltzmannSolver(kernel, config).solve(states0, adjacency, Y_lr=Y_lr)

kinetic_st-0.1.0/src/kinetic_st/core/collision.py

@@ -0,0 +1,106 @@
+"""Ligand-receptor collision operator on latent cell states.
+
+A "collision" is a ligand-receptor-mediated interaction updating two
+neighboring cell states. Irreversibility is baked in via attractor bias
+(gamma), since biological interactions are intrinsically irreversible.
+"""
+
+import numpy as np
+from typing import Optional, Tuple
+from scipy import sparse
+
+
+class LigandReceptorRule:
+    """Collision rule: v' = (1-a)*v + b*v* + g*a(v), and vice versa.
+
+    Parameters
+    ----------
+    alpha : self-relaxation rate (default 0.1)
+    beta : cross-influence rate (default 0.2)
+    gamma : attractor bias, irreversible differentiation (default 0.05)
+    attractors : (n_types, d) cell-type centroids; None = irreversibility off
+    """
+
+    def __init__(self, alpha: float = 0.1, beta: float = 0.2,
+                 gamma: float = 0.05, attractors: Optional[np.ndarray] = None):
+        for v in (alpha, beta, gamma):
+            if not 0.0 <= v <= 1.0:
+                raise ValueError("alpha, beta, gamma must be in [0, 1].")
+        self.alpha, self.beta, self.gamma = alpha, beta, gamma
+        self.attractors = attractors
+
+    def apply(self, v: np.ndarray, vs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Apply collision to cell pairs. v, vs: (m, d) -> vp, vsp: (m, d)."""
+        vp = (1.0 - self.alpha) * v + self.beta * vs
+        vsp = (1.0 - self.alpha) * vs + self.beta * v
+        if self.gamma > 0 and self.attractors is not None:
+            vp += self.gamma * (_nearest_attractor(vp, self.attractors) - vp)
+            vsp += self.gamma * (_nearest_attractor(vsp, self.attractors) - vsp)
+        return vp, vsp
+
+
+class CollisionKernel:
+    """Collision operator combining a rule with a cross-section sigma(v, v*).
+
+    sigma = sigma_LR * sigma_space, where sigma_LR depends on ligand-receptor
+    expression and sigma_space on spatial adjacency.
+
+    Parameters
+    ----------
+    rule : LigandReceptorRule
+    lr_matrix : (n_lr,) per-pair strengths; None = all 1
+    lr_index : (ligand_idx, receptor_idx) for expression-based sigma
+    sigma_scale : overall collision rate scale
+    """
+
+    def __init__(self, rule: LigandReceptorRule,
+                 lr_matrix: Optional[np.ndarray] = None,
+                 lr_index: Optional[Tuple[np.ndarray, np.ndarray]] = None,
+                 sigma_scale: float = 1.0):
+        self.rule = rule
+        self.lr_matrix = lr_matrix
+        self.lr_index = lr_index
+        self.sigma_scale = sigma_scale
+
+    def cross_section(self, states, adjacency, Y_lr=None) -> sparse.spmatrix:
+        """Sparse cross-section matrix over spatial edges."""
+        A = adjacency.tocoo()
+        n = adjacency.shape[0]
+        if self.lr_index is not None and Y_lr is not None:
+            lig_idx, rec_idx = self.lr_index
+            w = np.asarray((Y_lr[A.row][:, lig_idx] * Y_lr[A.col][:, rec_idx]).sum(axis=1)).flatten()
+        else:
+            d2 = np.sum((states[A.row] - states[A.col]) ** 2, axis=1)
+            w = np.exp(-0.5 * d2 / (np.median(d2) + 1e-8))
+        if self.lr_matrix is not None:
+            w *= float(np.mean(self.lr_matrix))
+        w *= self.sigma_scale
+        return sparse.coo_matrix((w, (A.row, A.col)), shape=(n, n)).tocsr()
+
+    def apply_collisions(self, states, adjacency, Y_lr=None,
+                         n_pairs_per_spot=4, rng=None) -> np.ndarray:
+        """One Monte-Carlo collision step on all cells."""
+        rng = rng or np.random.RandomState(0)
+        n, d = states.shape
+        Sigma = self.cross_section(states, adjacency, Y_lr).tocsr()
+        new_states = states.copy()
+        for i in range(n):
+            s, e = Sigma.indptr[i], Sigma.indptr[i + 1]
+            if e == s:
+                continue
+            nbrs, w = Sigma.indices[s:e], Sigma.data[s:e]
+            wsum = w.sum()
+            if wsum <= 0:
+                continue
+            m = min(n_pairs_per_spot, len(nbrs))
+            chosen = rng.choice(nbrs, size=m, replace=True, p=w / wsum)
+            vp, _ = self.rule.apply(
+                np.broadcast_to(states[i], (m, d)), states[chosen])
+            new_states[i] = vp.mean(axis=0)
+        return new_states
+
+
+def _nearest_attractor(states, attractors):
+    """Nearest attractor centroid for each state."""
+    d2 = ((states[:, None, :] - attractors[None, :, :]) ** 2).sum(axis=2)
+    return attractors[np.argmin(d2, axis=1)]

kinetic_st-0.1.0/src/kinetic_st/core/domains.py

@@ -0,0 +1,147 @@
+"""Spatial domain identification via local invariant measure divergence.
+
+A domain is a maximal contiguous region where the local invariant measure
+mu(x, v) is approximately constant. Replaces heuristic graph clustering with
+a dynamics-grounded definition.
+"""
+
+import numpy as np
+from scipy import sparse
+from typing import Optional
+from dataclasses import dataclass
+
+from kinetic_st.core.boltzmann import KineticTrajectory
+from kinetic_st.core.stochastic_pde import InvariantMeasureEstimator
+
+
+@dataclass
+class DomainResult:
+    """Domain identification output."""
+    labels: np.ndarray
+    n_domains: int
+    domain_field: np.ndarray
+    local_log_density: np.ndarray
+    centroids: np.ndarray
+
+
+class DomainIdentifier:
+    """Identify spatial domains from a kinetic trajectory.
+
+    Parameters
+    ----------
+    n_neighbors : spatial neighborhood for local mu_i estimation
+    n_domains : int or "auto" (BIC-like penalty)
+    divergence : "kl", "l2", or "pairwise"
+    representation : invariant-measure representation ("kde" or "gmm")
+    bandwidth : KDE bandwidth
+    n_components : GMM components
+    seed : random seed
+    """
+
+    def __init__(self, n_neighbors: int = 12, n_domains="auto",
+                 divergence: str = "kl", representation: str = "gmm",
+                 bandwidth: Optional[float] = None, n_components: int = 10,
+                 seed: int = 0):
+        self.n_neighbors = n_neighbors
+        self.n_domains = n_domains
+        self.divergence = divergence
+        self.representation = representation
+        self.bandwidth = bandwidth
+        self.n_components = n_components
+        self.seed = seed
+        self.result_ = None
+
+    def fit(self, trajectory: KineticTrajectory, adjacency: sparse.spmatrix,
+            verbose: bool = False) -> DomainResult:
+        """Identify spatial domains."""
+        n = adjacency.shape[0]
+        A = adjacency.tocsr()
+        nbr_lists = _neighbor_lists(A, self.n_neighbors)
+
+        global_est = InvariantMeasureEstimator(
+            representation=self.representation, bandwidth=self.bandwidth,
+            n_components=self.n_components, seed=self.seed).fit(trajectory)
+        self.global_measure_ = global_est
+
+        final_states = trajectory.states
+        d = final_states.shape[1]
+        global_logp = global_est.log_density(final_states)
+
+        if self.divergence == "kl":
+            field = -global_logp
+        elif self.divergence == "l2":
+            local_means = np.array([final_states[idx].mean(axis=0) for idx in nbr_lists])
+            field = np.linalg.norm(local_means - global_est.samples_.mean(axis=0), axis=1)
+        elif self.divergence == "pairwise":
+            field = np.zeros(n)
+            for i in range(n):
+                lp_local = _local_logpdf(final_states[nbr_lists[i]], final_states[i])
+                field[i] = global_logp[i] - lp_local
+        else:
+            raise ValueError(f"Unknown divergence: {self.divergence}")
+
+        field /= field.std() + 1e-8
+        labels = self._segment(field, final_states, A)
+
+        n_dom = int(labels.max() + 1)
+        centroids = np.array([final_states[labels == k].mean(axis=0)
+                              for k in range(n_dom) if (labels == k).any()])
+
+        self.result_ = DomainResult(
+            labels=labels, n_domains=n_dom, domain_field=field,
+            local_log_density=global_logp, centroids=centroids)
+        if verbose:
+            print(f"  [DomainIdentifier] identified {n_dom} domains")
+        return self.result_
+
+    def _segment(self, field, states, A):
+        """Segment tissue into domains via KMeans + spatial contiguity."""
+        n = len(states)
+        if self.n_domains == "auto":
+            from sklearn.cluster import KMeans
+            best_k, best_score = 2, -np.inf
+            for k in range(2, min(15, n // 20) + 1):
+                km = KMeans(n_clusters=k, random_state=self.seed, n_init=3).fit(states)
+                score = np.sum((km.cluster_centers_ - states.mean(axis=0)) ** 2) * k / max(km.inertia_, 1e-8) - 0.05 * k
+                if score > best_score:
+                    best_score, best_k = score, k
+            n_dom = best_k
+        else:
+            n_dom = int(self.n_domains)
+
+        from sklearn.cluster import KMeans
+        km = KMeans(n_clusters=n_dom, random_state=self.seed, n_init=5).fit(states)
+        return _enforce_contiguity(km.labels_, A)
+
+
+def _neighbor_lists(A, k):
+    A = A.tocsr()
+    return [np.concatenate([[i], A.indices[A.indptr[i]:A.indptr[i + 1]][:k]])
+            for i in range(A.shape[0])]
+
+
+def _local_logpdf(samples, x):
+    """Crude local KDE log-density."""
+    d = samples.shape[1]
+    bw = np.std(samples, axis=0) + 1e-3
+    logk = -0.5 * np.sum(((samples - x) / bw) ** 2, axis=1) - np.sum(np.log(bw)) - 0.5 * d * np.log(2 * np.pi)
+    m = logk.max()
+    return float(m + np.log(np.exp(logk - m).mean()))
+
+
+def _enforce_contiguity(labels, A):
+    """Split non-contiguous clusters into connected components."""
+    from scipy.sparse.csgraph import connected_components
+    A = A.tocsr()
+    new_labels = np.full(len(labels), -1, dtype=int)
+    next_id = 0
+    for lab in np.unique(labels):
+        mask = labels == lab
+        idx_global = np.where(mask)[0]
+        if len(idx_global) == 0:
+            continue
+        n_cc, cc_labels = connected_components(A[mask][:, mask], directed=False)
+        for c in range(n_cc):
+            new_labels[idx_global[cc_labels == c]] = next_id
+            next_id += 1
+    return new_labels

kinetic_st-0.1.0/src/kinetic_st/core/fluid_limit.py

@@ -0,0 +1,116 @@
+"""Macroscopic tissue fields via moment closure of the cell-state density.
+
+Moments: density rho, mean state u (expression field), covariance P
+(heterogeneity tensor). A relaxation_indicator reports how close each spot
+is to local quasi-equilibrium (fluid-limit regime).
+"""
+
+import numpy as np
+from scipy import sparse
+from typing import Optional
+from dataclasses import dataclass
+
+from kinetic_st.core.boltzmann import KineticTrajectory
+
+
+@dataclass
+class TissueField:
+    """Macroscopic fields at one time point."""
+    density: np.ndarray
+    mean_state: np.ndarray
+    covariance: np.ndarray
+    heterogeneity: np.ndarray
+    time: float
+
+
+class FluidLimiter:
+    """Compute tissue fields by taking moments of the cell-state density.
+
+    Parameters
+    ----------
+    bandwidth : KDE bandwidth; None = Silverman's rule
+    n_neighbors : spatial neighborhood size for moment support
+    use_local_kde : estimate f_i as Gaussian KDE over neighborhood
+    """
+
+    def __init__(self, bandwidth: Optional[float] = None,
+                 n_neighbors: int = 12, use_local_kde: bool = True):
+        self.bandwidth = bandwidth
+        self.n_neighbors = n_neighbors
+        self.use_local_kde = use_local_kde
+        self.fields_ = None
+
+    def compute(self, trajectory: KineticTrajectory, adjacency: sparse.spmatrix) -> list:
+        """Compute fields for every recorded time."""
+        A = adjacency.tocsr()
+        n = A.shape[0]
+        nbr_lists = _neighbor_lists(A, self.n_neighbors)
+        history, times = trajectory.history, trajectory.times
+        T, _, d = history.shape
+
+        if self.bandwidth is None:
+            all_states = history.reshape(-1, d)
+            self.bandwidth_ = np.median(all_states.std(axis=0) + 1e-8) * (n ** (-1.0 / (d + 4)))
+        else:
+            self.bandwidth_ = self.bandwidth
+
+        fields = []
+        for ti in range(T):
+            states = history[ti]
+            if self.use_local_kde:
+                mean_state, cov_diag = _local_moments(states, nbr_lists, self.bandwidth_)
+            else:
+                mean_state, cov_diag = states.copy(), np.zeros((n, d))
+            covariance = np.array([np.diag(c) for c in cov_diag])
+            fields.append(TissueField(
+                density=np.ones(n), mean_state=mean_state,
+                covariance=covariance, heterogeneity=cov_diag.sum(axis=1),
+                time=float(times[ti])))
+        self.fields_ = fields
+        return fields
+
+    def relaxation_indicator(self, trajectory: KineticTrajectory) -> np.ndarray:
+        """Local-equilibrium indicator per spot (0 = quasi-equilibrium)."""
+        if self.fields_ is None:
+            raise RuntimeError("Call compute() first.")
+        all_h = np.array([f.heterogeneity for f in self.fields_])
+        return (all_h / (all_h.mean() + 1e-8)).mean(axis=0)
+
+
+def compute_moments(states, adjacency, n_neighbors=12) -> TissueField:
+    """One-shot moment computation for a single snapshot."""
+    limiter = FluidLimiter(n_neighbors=n_neighbors, use_local_kde=True)
+    A = adjacency.tocsr()
+    nbr_lists = _neighbor_lists(A, n_neighbors)
+    n, d = states.shape
+    bw = np.median(states.std(axis=0) + 1e-8) * (n ** (-1.0 / (d + 4)))
+    mean_state, cov_diag = _local_moments(states, nbr_lists, bw)
+    return TissueField(
+        density=np.ones(n), mean_state=mean_state,
+        covariance=np.array([np.diag(c) for c in cov_diag]),
+        heterogeneity=cov_diag.sum(axis=1), time=0.0)
+
+
+def _neighbor_lists(A, k):
+    """Self + up to k neighbors per node."""
+    A = A.tocsr()
+    out = []
+    for i in range(A.shape[0]):
+        nbrs = A.indices[A.indptr[i]:A.indptr[i + 1]][:k]
+        out.append(np.concatenate([[i], nbrs]))
+    return out
+
+
+def _local_moments(states, nbr_lists, bandwidth):
+    """Weighted local mean and variance (Gaussian KDE)."""
+    n, d = states.shape
+    mean_state = np.zeros((n, d))
+    cov_diag = np.zeros((n, d))
+    bw2 = 2.0 * bandwidth ** 2 + 1e-12
+    for i in range(n):
+        local = states[nbr_lists[i]]
+        w = np.exp(-np.sum((local - states[i]) ** 2, axis=1) / bw2)
+        w /= w.sum() + 1e-12
+        mean_state[i] = (w[:, None] * local).sum(axis=0)
+        cov_diag[i] = (w[:, None] * (local - mean_state[i]) ** 2).sum(axis=0)
+    return mean_state, cov_diag