.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "gallery/abp_align_demo.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_gallery_abp_align_demo.py: Aligning active Brownian particles — generic pairs API ======================================================== Infer pairwise interaction forces in a system of aligning active Brownian particles (ABPs) using **generic pair-interaction building blocks** from :mod:`SFI.bases.pairs`. First we simulate the ABP system using the pre-built ABP model, then save and reload the trajectory to demonstrate the I/O round-trip, and finally construct a generic pair-interaction basis for linear inference. This example demonstrates: 1. Simulating a multi-particle ABP system with an example-local ABP helper (thin wrapper around :mod:`SFI.bases.pairs`). 2. Saving and loading trajectory data (CSV round-trip). 3. Constructing a generic **Basis** from :func:`~SFI.bases.pairs.heading_vector`, :func:`~SFI.bases.pairs.radial_pair_basis`, and :func:`~SFI.bases.pairs.angular_pair_basis`. 4. **Linear** force inference using that basis. 5. Recovering the learned interaction kernels vs the true model. .. rubric:: Tags synthetic · overdamped · multi-particle · linear · interactions .. GENERATED FROM PYTHON SOURCE LINES 30-31 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 72-85 System: aligning active Brownian particles -------------------------------------------- Each particle has position :math:`(x, y)` and heading :math:`\theta` — three degrees of freedom. The force includes: - **Self-propulsion**: :math:`c_0\,\hat{e}_\theta` - **Pairwise repulsion**: short-range, radially repulsive: :math:`\varepsilon\, e^{-r/R_0}\,\hat{r}_{ij}` - **Pairwise alignment**: angular torque: :math:`A\,\sin(\Delta\theta)\,e^{-r/L_0}` Positions obey periodic boundary conditions. .. GENERATED FROM PYTHON SOURCE LINES 85-104 .. code-block:: Python N_particles = 60 Lx, Ly = 30.0, 30.0 # periodic box dimensions Nsteps = 2000 # simulation length dt_sim = 0.02 # integration time step D_iso = 0.05 # isotropic diffusion coefficient seed = 0 box = jnp.array([Lx, Ly]) # True model parameters c0_true = 1.0 # self-propulsion speed eps_true = 2.0 # repulsion strength A_true = 0.5 # alignment strength R0_true = 1.0 # repulsion length scale L0_true = 2.0 # alignment length scale theta_F_exact = dict(c0=c0_true, eps=eps_true, A=A_true, R0=R0_true, L0=L0_true) .. GENERATED FROM PYTHON SOURCE LINES 105-113 Simulating the true model ---------------------------- We build the ABP model from the helpers in ``_gallery_utils/abp.py`` (a thin composition of :mod:`SFI.bases.pairs` primitives). The result is a parametric state function (PSF) with named parameters ``c0``, ``eps``, ``A``, ``R0``, ``L0``. .. GENERATED FROM PYTHON SOURCE LINES 113-135 .. code-block:: Python from _gallery_utils.abp import make_abp_align_psf from SFI.langevin import OverdampedProcess F_psf = make_abp_align_psf(dim=3) proc = OverdampedProcess(F_psf, D=D_iso, extras_global={"box": box}) proc.set_params(theta_F=theta_F_exact) # Random initial conditions: uniform in box, random headings key = random.PRNGKey(seed) key, kx, kth = random.split(key, 3) X0_xy = random.uniform(kx, (N_particles, 2)) * jnp.array([Lx, Ly]) TH0 = random.uniform(kth, (N_particles,), minval=-jnp.pi, maxval=jnp.pi) x0 = jnp.concatenate([X0_xy, TH0[:, None]], axis=1) proc.initialize(x0) key, sub = random.split(key) coll = proc.simulate(dt=dt_sim, Nsteps=Nsteps, key=sub) print(f"Trajectory: {coll.T} frames, " f"{coll.N} particles, dim={coll.d}") .. rst-class:: sphx-glr-script-out .. code-block:: none Trajectory: 2000 frames, 60 particles, dim=3 .. GENERATED FROM PYTHON SOURCE LINES 136-142 Simulation movie ------------------ Animation of the simulated ABP system. Particles are coloured by heading angle :math:`\theta`; white arrows show the propulsion direction. .. GENERATED FROM PYTHON SOURCE LINES 142-145 .. code-block:: Python Xfull = np.asarray(coll.X) # full trajectory array (T × N × 3) .. video:: /gallery/images/sphx_glr_abp_align_demo_001.mp4 :class: sphx-glr-single-img :height: 675 :width: 675 :autoplay: :loop: .. GENERATED FROM PYTHON SOURCE LINES 184-191 Saving and loading trajectory data ------------------------------------- SFI trajectory collections can be serialised to CSV, Parquet, or HDF5 via :meth:`~SFI.trajectory.TrajectoryCollection.save`. Global extras — here the periodic box — are embedded in a YAML header and survive the round-trip. .. GENERATED FROM PYTHON SOURCE LINES 191-209 .. code-block:: Python import os, tempfile from SFI.trajectory import TrajectoryCollection csv_path = os.path.join(tempfile.gettempdir(), "sfi_abp_demo.csv") coll.save(csv_path, format="csv") with open(csv_path) as f: for line in f.readlines()[:10]: print(line.rstrip()) print(" ...") # Reload into a fresh collection and verify coll_reloaded = TrajectoryCollection.load(csv_path) print(f"\nOriginal: T={coll.T}, N={coll.N}, d={coll.d} → Loaded: T={coll_reloaded.T}, N={coll_reloaded.N}, d={coll_reloaded.d}") max_err = float(np.max(np.abs(coll_reloaded.X - coll.X))) print(f"Max |ΔX| = {max_err:.1e} (numerical round-trip error)") .. rst-class:: sphx-glr-script-out .. code-block:: none # --- # kind: overdamped # dimension: 3 # pdepth: 1 # num_particles: 60 # x0: # - - 0.21881461143493652 # - 0.6267356872558594 # - 2.5286648273468018 # - - 17.442794799804688 ... Original: T=2000, N=60, d=3 → Loaded: T=2000, N=60, d=3 Max |ΔX| = 7.5e-09 (numerical round-trip error) .. GENERATED FROM PYTHON SOURCE LINES 210-216 Particle snapshot ------------------- Final-frame snapshot with heading arrows (left) and individual particle trajectories (right). Positions are wrapped into the periodic box; particles are coloured by heading angle. .. GENERATED FROM PYTHON SOURCE LINES 216-219 .. code-block:: Python n_show = min(5, N_particles) # trajectories to display .. image-sg:: /gallery/images/sphx_glr_abp_align_demo_002.png :alt: Aligning ABPs (N = 60), Snapshot (t = 40), Trajectories (5 particles) :srcset: /gallery/images/sphx_glr_abp_align_demo_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 236-253 Building the pair-interaction basis ---------------------------------------- For **inference**, we construct a generic basis from building blocks in :mod:`SFI.bases.pairs`. The basis is deliberately over-complete — it contains many more radial kernels than the true model uses — so the linear solver must pick out the right combination. 1. **Self-propulsion** — :func:`~SFI.bases.pairs.heading_vector` gives the unit heading :math:`(\cos\theta, \sin\theta, 0)`. 2. **Radial repulsion** — :func:`~SFI.bases.pairs.radial_pair_basis` with exponential- polynomial kernels. 3. **Angular alignment** — :func:`~SFI.bases.pairs.angular_pair_basis` with a :math:`\sin(\Delta\theta)` coupling. .. GENERATED FROM PYTHON SOURCE LINES 253-303 .. code-block:: Python from SFI.bases.pairs import ( angular_pair_basis, exp_poly_kernels, heading_vector, radial_pair_basis, ) from SFI.statefunc import Basis dim = 3 # (x, y, θ) per particle # Radial kernels: r^n * exp(-r/L) for n ∈ {0,1}, L ∈ {0.5, 1, 2, 4} repel_kernels = exp_poly_kernels(degrees=[0, 1], lengths=[0.5, 1.0, 2.0, 4.0]) align_kernels = exp_poly_kernels(degrees=[0, 1], lengths=[0.5, 1.0, 2.0, 4.0]) # Self-propulsion: (cos θ, sin θ, 0) B_heading = heading_vector(dim=dim, angle_index=2) # Radial pair repulsion → dispatched Basis (acts in xy-plane) inter_repel = radial_pair_basis( repel_kernels, dim=dim, box="extras", spatial_dims=slice(0, 2), # xy positions for distance embed_dim=dim, embed_axes=[0, 1], # embed radial force into xy ) B_repel = inter_repel.dispatch_pairs( symmetric=True, exclude_self=True, owners="focal", reducer="sum", return_as="basis", ) # Angular alignment → dispatched Basis (acts on θ) inter_align = angular_pair_basis( align_kernels, jnp.sin, dim=dim, angle_index=2, output_index=2, # coupling: sin(Δθ) → force on θ box="extras", spatial_dims=slice(0, 2), ) B_align = inter_align.dispatch_pairs( symmetric=True, exclude_self=True, owners="focal", reducer="sum", return_as="basis", ) # Stack into one combined basis for inference B_full = Basis.stack([B_heading, B_repel, B_align]) n_heading = B_heading.n_features n_repel = B_repel.n_features n_align = B_align.n_features .. GENERATED FROM PYTHON SOURCE LINES 304-310 Linear force inference ------------------------- Standard :meth:`~SFI.inference.OverdampedLangevinInference.infer_force_linear` recovers the interaction coefficients in a single linear solve — no gradient-based optimisation needed. .. GENERATED FROM PYTHON SOURCE LINES 310-332 .. code-block:: Python from SFI import OverdampedLangevinInference from SFI.utils import plotting # Create inference object from the trajectory collection inf = OverdampedLangevinInference(coll) # Estimate diffusion coefficient inf.compute_diffusion_constant(method="WeakNoise") # Solve for force coefficients (one-shot linear solve) inf.infer_force_linear(B_full, M_mode="Ito", G_mode="rectangle") # Estimate coefficient uncertainty (populates force_coefficients_covariance) inf.compute_force_error() # Compare to exact model for quantitative error metrics inf.compare_to_exact(model_exact=proc, maxpoints=2000) inf.print_report() nmse_abp = float(inf.NMSE_force) # reused in the force-scatter title .. rst-class:: sphx-glr-script-out .. code-block:: none --- StochasticForceInference Report --- Average diffusion tensor: [[ 4.9648169e-02 -1.6120239e-05 -1.0172471e-04] [-1.6120215e-05 4.9877696e-02 -1.6907154e-04] [-1.0172469e-04 -1.6907144e-04 4.9801379e-02]] Measurement noise tensor: [[-1.5219006e-04 -8.3711362e-05 5.3390534e-07] [-8.3711362e-05 -2.2979152e-04 -3.4853206e-06] [ 5.3390534e-07 -3.4853210e-06 -8.2759789e-06]] Force estimated information: 12177.7060546875 Force: estimated normalized mean squared error (sampling only): 0.0005337358762452214 Normalized MSE (force): 0.0007 Normalized MSE (diffusion): 0.0000 Force Coefficient Table ────────────────────────────────────────────────────────────────────────── # Label Coefficient Std.Err SNR Sig ────────────────────────────────────────────────────────────────────────── 0 e_heading 9.95722e-01 6.56048e-03 151.8 *** 1 r^0·exp(-r/0.5) 3.34897e+00 1.93552e+00 1.7 · 2 r^0·exp(-r/1) -4.99166e+00 2.76279e+00 1.8 · 3 r^0·exp(-r/2) -1.84840e+00 5.16968e-01 3.6 * 4 r^0·exp(-r/4) 3.16612e-02 1.45463e-01 0.2 · 5 r^1·exp(-r/0.5) 4.37969e+00 3.73906e+00 1.2 · 6 r^1·exp(-r/1) 2.38924e+00 4.90655e-01 4.9 * 7 r^1·exp(-r/2) 1.11856e-01 4.70822e-02 2.4 * 8 r^1·exp(-r/4) -2.23570e-02 1.33291e-02 1.7 · 9 g·r^0·exp(-r/0.5) -4.35410e+01 3.46116e+00 12.6 ** 10 g·r^0·exp(-r/1) 2.50350e+01 4.92693e+00 5.1 * 11 g·r^0·exp(-r/2) 2.55251e+01 4.97670e-01 51.3 ** 12 g·r^0·exp(-r/4) 3.83326e+00 8.76765e-02 43.7 ** 13 g·r^1·exp(-r/0.5) -5.32409e+01 7.19212e+00 7.4 * 14 g·r^1·exp(-r/1) -3.07139e+01 1.00916e+00 30.4 ** 15 g·r^1·exp(-r/2) -4.90880e+00 6.85575e-02 71.6 ** 16 g·r^1·exp(-r/4) -1.75143e-01 8.78098e-03 19.9 ** ────────────────────────────────────────────────────────────────────────── 17/17 basis functions in support, sig: 12* / 7** / 1*** (|SNR| ≥ 2 / 10 / 100) (Std.err. reflects sampling error only; discretization bias is not included.) .. GENERATED FROM PYTHON SOURCE LINES 333-339 Force scatter: exact vs inferred ----------------------------------- All force components (translational and angular) along the trajectory, plotted as a scatter. Points near the diagonal indicate good inference. .. GENERATED FROM PYTHON SOURCE LINES 339-346 .. code-block:: Python # The force is a multi-particle interacting field (per-frame neighbour # structure + the ``box`` extra), so evaluate it per frame on ``coll.X`` # rather than via force_comparison_arrays (which flattens to single points). Xe = proc.force_sf(coll.X, extras={"box": box}) Xi = inf.force_inferred(coll.X, extras={"box": box}) .. image-sg:: /gallery/images/sphx_glr_abp_align_demo_003.png :alt: $r=0.9995$ MSE=0.0, Force scatter (NMSE = 0.001) :srcset: /gallery/images/sphx_glr_abp_align_demo_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 355-363 Recovered interaction kernels -------------------------------- The inferred coefficients weight the radial kernels. Summing :math:`\sum_k c_k \phi_k(r)` reconstructs the effective radial interaction shape. Shaded bands show 95 % confidence intervals derived from the coefficient covariance (see :func:`~SFI.inference.kernel_predict_ci`). .. GENERATED FROM PYTHON SOURCE LINES 363-392 .. code-block:: Python from SFI.inference import kernel_predict_ci # Extract inferred coefficients and covariance sub-blocks for each basis block c_heading, _ = inf.coeff_block(B_heading, field='force') c_repel, cov_repel = inf.coeff_block(B_repel, field='force') c_align, cov_align = inf.coeff_block(B_align, field='force') # True coefficient vectors in the generic-basis coordinate system # Kernel ordering from exp_poly_kernels(degrees=[0,1], lengths=[0.5,1,2,4]): # k=0: r⁰·e⁻ʳ/0.5 k=1: r⁰·e⁻ʳ/1 k=2: r⁰·e⁻ʳ/2 k=3: r⁰·e⁻ʳ/4 # k=4: r¹·e⁻ʳ/0.5 k=5: r¹·e⁻ʳ/1 k=6: r¹·e⁻ʳ/2 k=7: r¹·e⁻ʳ/4 idx_repel = 1 # r⁰·e⁻ʳ/R₀ with R₀=1.0 idx_align = 2 # sin(Δθ)·r⁰·e⁻ʳ/L₀ with L₀=2.0 true_c_repel = np.zeros(n_repel) true_c_repel[idx_repel] = -eps_true true_c_align = np.zeros(n_align) true_c_align[idx_align] = A_true # Compute kernel profiles with confidence intervals r_eval = np.linspace(0.01, 8.0, 200) ci_repel = kernel_predict_ci(r_eval, repel_kernels, c_repel, cov_repel) ci_align = kernel_predict_ci(r_eval, align_kernels, c_align, cov_align) # True kernel profiles (for comparison) true_repel = true_c_repel @ ci_repel["phi"] true_align = true_c_align @ ci_align["phi"] .. image-sg:: /gallery/images/sphx_glr_abp_align_demo_004.png :alt: Interaction kernel recovery, Radial repulsion kernel, Angular alignment kernel :srcset: /gallery/images/sphx_glr_abp_align_demo_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 416-421 Inferred self-propulsion ~~~~~~~~~~~~~~~~~~~~~~~~~~ The heading coefficient recovers the active thrust magnitude :math:`c_0`: .. GENERATED FROM PYTHON SOURCE LINES 423-429 Bootstrapped trajectory ------------------------- Simulating from the inferred model and comparing emergent collective behaviour is a powerful qualitative validation for multi-particle systems. .. GENERATED FROM PYTHON SOURCE LINES 429-438 .. code-block:: Python # Run a bootstrap simulation from the inferred coefficients key_boot = random.PRNGKey(seed + 77) proc_boot = inf.simulate_bootstrapped_trajectory(key_boot, simulate=False) proc_boot.set_extras(extras_global={"box": box}) proc_boot.initialize(x0) # same initial conditions key_boot, sub_boot = random.split(key_boot) coll_boot = proc_boot.simulate(dt=dt_sim, Nsteps=Nsteps, key=sub_boot) .. image-sg:: /gallery/images/sphx_glr_abp_align_demo_005.png :alt: Bootstrap validation, Original (final frame), Bootstrapped (inferred model) :srcset: /gallery/images/sphx_glr_abp_align_demo_005.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 455-461 Bootstrap comparison movie ---------------------------- Side-by-side animation of the original trajectory (left) and a trajectory simulated from the *inferred* model (right). Both start from the same initial conditions. .. GENERATED FROM PYTHON SOURCE LINES 461-464 .. code-block:: Python Xboot_full = np.asarray(coll_boot.X) .. video:: /gallery/images/sphx_glr_abp_align_demo_006.mp4 :class: sphx-glr-single-img :height: 900 :width: 1800 :autoplay: :loop: .. GENERATED FROM PYTHON SOURCE LINES 529-532 Both simulations start from the same initial conditions. Visual similarity confirms the inferred model captures the collective dynamics. .. GENERATED FROM PYTHON SOURCE LINES 534-547 Diagnostics ----------- :func:`~SFI.diagnostics.assess` computes standardised Euler residuals pooled across **all particles and spatial components** (:math:`N_\mathrm{obs} \approx N_\mathrm{particles} \times T \times d`), giving high statistical power even from a short trajectory. All three force channels — translational :math:`(x,y)` and angular :math:`(\theta)` — are pooled into a single residual vector. The linear basis is over-complete by design, so the NMSE should be small but the coefficient z-scores for unused kernels will be near zero. .. GENERATED FROM PYTHON SOURCE LINES 547-554 .. code-block:: Python from SFI.diagnostics import assess, plot_summary report = assess(inf, level="standard") report.print_summary() .. image-sg:: /gallery/images/sphx_glr_abp_align_demo_007.png :alt: Aligning ABPs — diagnostics (linear model), Q--Q plot, Residual histogram, Residual ACF :srcset: /gallery/images/sphx_glr_abp_align_demo_007.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none === SFI diagnostics report === backend : linear regime : OD n_obs : 119940 n_particles: 60 d: 3 level : standard -- Residuals -- mean = +0.0003 std = 0.9988 skew = -0.005 kurt-3 = -0.000 (n=359820) ✓ normality ks stat=0.00143 p=0.454 ✓ autocorr ljung_box stat=14.9 p=0.783 ✓ autocorr ljung_box_squared stat=21 p=0.4 predicted NMSE = 0.000534 realised NMSE = 0 χ² z = -1.00 (|z|>5 ⇒ bias) -- Flags -- (no issues at α = 0.01) .. GENERATED FROM PYTHON SOURCE LINES 560-564 Thumbnail --------- Dedicated single-panel figure for the gallery thumbnail. .. GENERATED FROM PYTHON SOURCE LINES 564-567 .. code-block:: Python stamp_output() .. image-sg:: /gallery/images/sphx_glr_abp_align_demo_008.png :alt: abp align demo :srcset: /gallery/images/sphx_glr_abp_align_demo_008.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none examples/gallery/abp_align_demo.py:570: UserWarning: The figure layout has changed to tight plt.tight_layout() [Generated: 2026-06-30 12:03] .. rst-class:: sphx-glr-timing **Total running time of the script:** (3 minutes 24.697 seconds) .. rst-class:: sphx-glr-example-tags 🏷 Tags: synthetic, overdamped, multi-particle, linear, interactions .. _sphx_glr_download_gallery_abp_align_demo.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: abp_align_demo.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: abp_align_demo.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: abp_align_demo.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_