Note
Go to the end to download the full example code.
Nonreciprocal ABPs at large scale — 3 000 particles¶
Infer three pairwise interaction kernels — repulsion, alignment, and pursuit — in a system of 3 000 active Brownian particles whose interactions are nonreciprocal and vision-gated.
This example scales up the ABP pair-interaction demo to large particle
numbers using truncated-range neighbor lists (CSR format) and the
simulate_chunked() helper that periodically
rebuilds the neighbor list during simulation.
It demonstrates:
Building a parametric simulation force from composable geometric primitives in
SFI.bases.pairs— includingvision_gate()for nonreciprocal coupling andparticle_heading()for pursuit.Chunked simulation with periodic neighbor-list rebuilds via
simulate_chunked().Constructing an overcomplete inference basis with
scalar_pair_basis()and the same geometric building blocks.Linear force inference on a subset of trajectory chunks.
Recovering all three interaction kernels vs the true model.
Tags
synthetic · overdamped · multi-particle · linear · interactions · large-scale · nonreciprocal
System: nonreciprocal, vision-gated ABPs¶
Each particle has state \((x, y, \theta)\). The deterministic force contains three pairwise interactions, each modulated by a vision cone \(v(\delta) = (1 + \cos\delta)/2\) that depends on the bearing angle \(\delta\) from the focal particle’s heading to the neighbor — making the interactions nonreciprocal:
Self-propulsion: \(c_0\,\hat{e}_\theta\)
Repulsion: \(-\varepsilon\, e^{-r/R_0}\,\hat{r}_{ij}\)
Alignment (vision-gated): \(A\,v(\delta)\,\sin(\Delta\theta)\, e^{-r/L_a}\)
Pursuit (vision-gated): \(P\,v(\delta)\,\hat{e}_{\theta_j}\, e^{-r/L_p}\) — steers the focal particle toward the neighbor’s heading.
The system lives in a periodic box with ~3 000 particles at the same density as the 60-particle ABP demo.
N_particles = 3_000
density = 60 / (30.0 * 30.0) # same as abp_align_demo
area = N_particles / density
Lx = Ly = float(np.sqrt(area)) # ≈ 387
Nsteps = 2000
dt_sim = 0.02
D_iso = 0.05
seed = 0
box = jnp.array([Lx, Ly])
box_np = np.array([Lx, Ly])
# Interaction range and neighbor-list parameters
cutoff = 12.0 # truncation radius
rebuild_every = 1 # rebuild neighbor list every step (fast with cKDTree)
# True model parameters
c0_true = 1.0 # self-propulsion
eps_true = 2.0 # repulsion strength
A_true = 0.3 # alignment strength
P_true = 1.5 # pursuit strength
R0_true = 1.0 # repulsion length scale
La_true = 2.0 # alignment length scale
Lp_true = 4.0 # pursuit length scale
theta_F_exact = dict(c0=c0_true, eps=eps_true, R0=R0_true,
A=A_true, La=La_true, P=P_true, Lp=Lp_true)
print(f"N = {N_particles}, box = {Lx:.1f}×{Ly:.1f}, cutoff = {cutoff}")
N = 3000, box = 212.1×212.1, cutoff = 12.0
Building the parametric simulation force¶
We compose the simulation force from geometric primitives in
SFI.bases.pairs. Each primitive is a building block — a
radial kernel, a direction vector, a gating function — that can
be multiplied together and dispatched over pairs.
from SFI.bases.pairs import (
angle_coupling,
exp_poly_kernels,
heading_vector,
pair_direction,
parametric_radial_kernel,
particle_heading,
scalar_pair_basis,
vision_gate,
)
from SFI.statefunc import Basis
dim = 3 # (x, y, θ) per particle
# Geometric primitives
B_heading = heading_vector(dim=dim, angle_index=2)
e_ij = pair_direction(
dim=dim, box="extras", spatial_dims=slice(0, 2),
embed_dim=dim, embed_axes=[0, 1],
)
g_align = angle_coupling(jnp.sin, dim=dim, angle_index=2)
e_j = particle_heading(1, dim=dim, angle_index=2)
v = vision_gate(
lambda d: (1 + jnp.cos(d)) / 2,
dim=dim, angle_index=2,
box="extras", spatial_dims=slice(0, 2),
)
# Parametric radial kernels
k_repel = parametric_radial_kernel(
lambda r, p: -p["eps"] * jnp.exp(-r / p["R0"]),
params={"eps": (), "R0": ()},
dim=dim, box="extras", spatial_dims=slice(0, 2),
)
k_align = parametric_radial_kernel(
lambda r, p: p["A"] * jnp.exp(-r / p["La"]),
params={"A": (), "La": ()},
dim=dim, box="extras", spatial_dims=slice(0, 2),
)
k_pursuit = parametric_radial_kernel(
lambda r, p: p["P"] * jnp.exp(-r / p["Lp"]),
params={"P": (), "Lp": ()},
dim=dim, box="extras", spatial_dims=slice(0, 2),
)
# CSR dispatch keys (neighbor list stored in extras)
csr_kw = dict(indptr_key="indptr", indices_key="indices")
# Compose the full force: self-propulsion + repulsion + alignment + pursuit
F_sim = (
B_heading.to_psf(coeff_key="c0")
+ (k_repel * e_ij).dispatch_pairs_from_extras(**csr_kw, return_as="psf")
+ (k_align * v * g_align).dispatch_pairs_from_extras(**csr_kw, return_as="psf")
+ (k_pursuit * v * e_j).dispatch_pairs_from_extras(**csr_kw, return_as="psf")
)
print(f"Force primitives: heading, repulsion, "
f"vision-gated alignment, vision-gated pursuit")
Force primitives: heading, repulsion, vision-gated alignment, vision-gated pursuit
Chunked simulation with neighbor rebuilds¶
At this scale (N = 3 000) a full N² pair list is infeasible.
simulate_chunked() builds a CSR
neighbor list using scipy.spatial.cKDTree() (≈ 0.05 s per
rebuild) and rebuilds it periodically. Here we rebuild every
step because the collective flock motion can exceed any practical
Verlet skin in a single time step.
The save_every parameter decouples the output chunk size (50
frames) from the rebuild frequency (1 step).
from SFI.langevin import OverdampedProcess
from SFI.langevin.chunked import simulate_chunked
from SFI.utils.neighbors import make_neighbor_extras
# 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)
# Build initial neighbor list
nbr0 = make_neighbor_extras(np.asarray(x0[:, :2]), cutoff, box_np)
extras0 = {"box": box}
extras0.update(nbr0)
proc = OverdampedProcess(F_sim, D=D_iso, extras_global=extras0)
proc.set_params(theta_F=theta_F_exact)
proc.initialize(x0)
# Run chunked simulation
print(f"Simulating {Nsteps} steps with neighbor rebuild every step ...")
t0 = time.perf_counter()
key, sub = random.split(key)
coll = simulate_chunked(
proc, dt=dt_sim, Nsteps=Nsteps, key=sub,
cutoff=cutoff, box=box_np,
skin=0.0, rebuild_every=rebuild_every,
save_every=50,
spatial_dims=slice(0, 2),
nnz_safety=3.0, verbose=True,
)
sim_time = time.perf_counter() - t0
n_chunks = len(coll.datasets)
print(f"Simulation done in {sim_time:.0f}s ({n_chunks} chunks)")
Simulating 2000 steps with neighbor rebuild every step ...
SFI/trajectory/collection.py:641: UserWarning: Very short trajectory (T=1). Most inference methods need T >> 1 for meaningful results.
ds = TrajectoryDataset.from_arrays(
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 1: 50/2000 steps
max displacement = 0.250 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 2: 100/2000 steps
max displacement = 0.249 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 3: 150/2000 steps
max displacement = 0.282 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 4: 200/2000 steps
max displacement = 0.277 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 5: 250/2000 steps
max displacement = 0.376 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 6: 300/2000 steps
max displacement = 0.577 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 7: 350/2000 steps
max displacement = 0.881 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 8: 400/2000 steps
max displacement = 1.076 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 9: 450/2000 steps
max displacement = 1.545 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 10: 500/2000 steps
max displacement = 1.500 (skin/2 = 0.000)
max_nnz grew to 818490
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 11: 550/2000 steps
max displacement = 1.600 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 12: 600/2000 steps
max displacement = 4.049 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 13: 650/2000 steps
max displacement = 2.803 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 14: 700/2000 steps
max displacement = 2.812 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 15: 750/2000 steps
max displacement = 2.717 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 16: 800/2000 steps
max displacement = 1.393 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 17: 850/2000 steps
max displacement = 1.356 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 18: 900/2000 steps
max displacement = 1.489 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 19: 950/2000 steps
max displacement = 1.258 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 20: 1000/2000 steps
max displacement = 1.214 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 21: 1050/2000 steps
max displacement = 1.181 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 22: 1100/2000 steps
max displacement = 1.226 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 23: 1150/2000 steps
max displacement = 1.338 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 24: 1200/2000 steps
max displacement = 1.589 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 25: 1250/2000 steps
max displacement = 1.416 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 26: 1300/2000 steps
max displacement = 1.145 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 27: 1350/2000 steps
max displacement = 1.173 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 28: 1400/2000 steps
max displacement = 1.057 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 29: 1450/2000 steps
max displacement = 1.085 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 30: 1500/2000 steps
max displacement = 0.785 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 31: 1550/2000 steps
max displacement = 0.902 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 32: 1600/2000 steps
max displacement = 1.190 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 33: 1650/2000 steps
max displacement = 1.135 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 34: 1700/2000 steps
max displacement = 1.117 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 35: 1750/2000 steps
max displacement = 1.072 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 36: 1800/2000 steps
max displacement = 0.860 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 37: 1850/2000 steps
max displacement = 0.888 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 38: 1900/2000 steps
max displacement = 1.163 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 39: 1950/2000 steps
max displacement = 1.182 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 40: 2000/2000 steps
Simulation done in 2511s (40 chunks)
Simulation snapshot¶
Final-frame snapshot with 3 000 particles coloured by heading angle \(\theta\). Collective flocking structures emerge from the pursuit and alignment interactions.
Xfull = coll.to_array(axis="time") # (T, N, 3)
Simulation movie¶
Full-frame animation of the flock dynamics showing the emergence and persistence of collective structures.
Building the overcomplete inference basis¶
For inference we build an overcomplete basis from the same geometric
primitives. Instead of parametric kernels we use a grid of
exponential-polynomial basis functions via
exp_poly_kernels(), combined with the same
vision gate \(v\), alignment coupling \(g\), and heading
vector \(\hat{e}_j\).
kernels = exp_poly_kernels(degrees=[0, 1], lengths=[0.5, 1.0, 2.0, 4.0])
phi_r = scalar_pair_basis(kernels, dim=dim, box="extras",
spatial_dims=slice(0, 2))
B_repel = (phi_r * e_ij).dispatch_pairs_from_extras(
**csr_kw, return_as="basis")
B_align = (phi_r * v * g_align).dispatch_pairs_from_extras(
**csr_kw, return_as="basis")
B_pursuit = (phi_r * v * e_j).dispatch_pairs_from_extras(
**csr_kw, return_as="basis")
B_full = Basis.stack([B_heading, B_repel, B_align, B_pursuit])
n_heading = B_heading.n_features
n_repel = B_repel.n_features
n_align = B_align.n_features
n_pursuit = B_pursuit.n_features
Linear force inference on mid-trajectory chunks¶
We select 8 evenly-spaced chunks from the middle of the trajectory (after the system has relaxed). For each chunk we build a fresh CSR neighbor list from the first frame.
from SFI import OverdampedLangevinInference
from SFI.trajectory import TrajectoryCollection
INFER_CHUNKS = list(range(12, 28, 2)) # [12, 14, 16, 18, 20, 22, 24, 26]
datasets = []
for ci in INFER_CHUNKS:
Xc = np.asarray(coll.to_arrays(dataset=ci)[1]) # (t, X, mask) -> X
X0 = Xc[0]
nbr = make_neighbor_extras(X0[:, :2], cutoff, box_np)
eg = {"box": box}
eg.update(nbr)
ds = TrajectoryCollection.from_arrays(
X=Xc, dt=dt_sim, extras_global=eg,
).datasets[0]
datasets.append(ds)
print(f" chunk {ci}: nnz = {int(nbr['indptr'][-1])}")
coll_infer = TrajectoryCollection(
datasets=datasets,
weights=jnp.ones(len(datasets), dtype=jnp.float32),
).with_weights("pool")
# Inference
inf = OverdampedLangevinInference(coll_infer)
inf.compute_diffusion_constant(method="WeakNoise")
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 error metrics
X0_ref = np.asarray(coll_infer.to_arrays(dataset=0)[1][0]) # X, first frame
nbr_ref = make_neighbor_extras(X0_ref[:, :2], cutoff, box_np)
extras_ref = {"box": box}
extras_ref.update(nbr_ref)
proc_ref = OverdampedProcess(F_sim, D=D_iso, extras_global=extras_ref)
proc_ref.set_params(theta_F=theta_F_exact)
proc_ref.initialize(jnp.array(X0_ref))
inf.compare_to_exact(model_exact=proc_ref, maxpoints=2000)
inf.print_report()
# NMSE(force) is already printed by inf.print_report(); keep the scalar
# only for the kernel-recovery figure title below.
nmse = float(inf.NMSE_force)
# Report inferred self-propulsion (true-vs-inferred narration)
c_heading, _ = inf.coeff_block((0, n_heading))
chunk 12: nnz = 415632
chunk 14: nnz = 340048
chunk 16: nnz = 234986
chunk 18: nnz = 240874
chunk 20: nnz = 216122
chunk 22: nnz = 218036
chunk 24: nnz = 261440
chunk 26: nnz = 191516
--- StochasticForceInference Report ---
Average diffusion tensor:
[[ 5.0915726e-02 -6.5736566e-04 -9.1523943e-05]
[-6.5736566e-04 5.1229026e-02 -5.7592075e-05]
[-9.1523951e-05 -5.7592068e-05 4.9817488e-02]]
Measurement noise tensor:
[[-1.0783467e-01 9.3433745e-02 6.5726155e-05]
[ 9.3433760e-02 -1.5797479e-01 -1.2734835e-04]
[ 6.5726155e-05 -1.2734835e-04 -1.3002133e-05]]
Force estimated information: 72313200.0
Force: estimated normalized mean squared error (sampling only): 1.1062799186607579e-07
Normalized MSE (force): 0.0013
Normalized MSE (diffusion): 0.0004
Force Coefficient Table
──────────────────────────────────────────────────────────────────────────────────────────
# Label Coefficient Std.Err SNR Sig
──────────────────────────────────────────────────────────────────────────────────────────
0 e_heading 1.70803e+00 3.05078e-03 559.9 ***
1 (r^0·exp(-r/0.5))·r̂_ij -6.31290e-01 1.23689e-02 51.0 **
2 (r^0·exp(-r/1))·r̂_ij -4.53953e-01 4.96479e-03 91.4 **
3 (r^0·exp(-r/2))·r̂_ij -1.23375e-01 7.48174e-04 164.9 ***
4 (r^0·exp(-r/4))·r̂_ij 7.79701e-02 5.39988e-04 144.4 ***
5 (r^1·exp(-r/0.5))·r̂_ij -3.05828e+00 3.89760e-02 78.5 **
6 (r^1·exp(-r/1))·r̂_ij -3.63216e-01 7.58078e-03 47.9 **
7 (r^1·exp(-r/2))·r̂_ij 1.11130e-01 1.70465e-03 65.2 **
8 (r^1·exp(-r/4))·r̂_ij -3.24837e-02 2.83821e-04 114.5 ***
9 ((r^0·exp(-r/0.5))·vision)·g 1.72235e+00 6.82426e-02 25.2 **
10 ((r^0·exp(-r/1))·vision)·g -1.67757e+00 2.29336e-02 73.1 **
11 ((r^0·exp(-r/2))·vision)·g 4.58102e-01 2.58290e-03 177.4 ***
12 ((r^0·exp(-r/4))·vision)·g -1.49466e-01 2.21182e-03 67.6 **
13 ((r^1·exp(-r/0.5))·vision)·g 1.46566e+00 1.80297e-01 8.1 *
14 ((r^1·exp(-r/1))·vision)·g 4.85935e-01 2.74626e-02 17.7 **
15 ((r^1·exp(-r/2))·vision)·g 1.00282e-02 5.43485e-03 1.8 ·
16 ((r^1·exp(-r/4))·vision)·g 1.41672e-02 1.01886e-03 13.9 **
17 ((r^0·exp(-r/0.5))·vision)·ê_θ[1] -1.23483e+00 1.69453e-02 72.9 **
18 ((r^0·exp(-r/1))·vision)·ê_θ[1] -8.24450e-01 4.30121e-03 191.7 ***
19 ((r^0·exp(-r/2))·vision)·ê_θ[1] 9.29590e-02 1.75395e-03 53.0 **
20 ((r^0·exp(-r/4))·vision)·ê_θ[1] 5.40698e-01 8.76343e-04 617.0 ***
21 ((r^1·exp(-r/0.5))·vision)·ê_θ[1] 1.04261e+01 5.58161e-02 186.8 ***
22 ((r^1·exp(-r/1))·vision)·ê_θ[1] -1.94002e+00 1.08098e-02 179.5 ***
23 ((r^1·exp(-r/2))·vision)·ê_θ[1] 1.02251e+00 2.41261e-03 423.8 ***
24 ((r^1·exp(-r/4))·vision)·ê_θ[1] -2.27031e-02 3.91993e-04 57.9 **
──────────────────────────────────────────────────────────────────────────────────────────
25/25 basis functions in support, sig: 24* / 23** / 10*** (|SNR| ≥ 2 / 10 / 100)
(Std.err. reflects sampling error only; discretization bias is not included.)
Recovered interaction kernels¶
Each interaction kernel is reconstructed as
\(\sum_k c_k \phi_k(r)\) using the inferred coefficients
and the basis functions. Shaded bands show 95 % confidence
intervals derived from the coefficient covariance
(see kernel_predict_ci()).
from SFI.inference import kernel_predict_ci
# Per-group coefficients and covariance sub-blocks (no hand-sliced
# offsets into the flat covariance — coeff_block returns both).
i0_repel = n_heading
i1_repel = n_heading + n_repel
i0_align = i1_repel
i1_align = i0_align + n_align
i0_pursuit = i1_align
i1_pursuit = i0_pursuit + n_pursuit
c_repel, cov_repel = inf.coeff_block((i0_repel, i1_repel))
c_align, cov_align = inf.coeff_block((i0_align, i1_align))
c_pursuit, cov_pursuit = inf.coeff_block((i0_pursuit, i1_pursuit))
# True coefficient vectors
# Kernel order from exp_poly_kernels(degrees=[0,1], lengths=[0.5,1,2,4]):
# 0: r⁰·exp(-r/0.5) 1: r⁰·exp(-r/1) 2: r⁰·exp(-r/2) 3: r⁰·exp(-r/4)
# 4: r¹·exp(-r/0.5) 5: r¹·exp(-r/1) 6: r¹·exp(-r/2) 7: r¹·exp(-r/4)
idx_repel = 1 # r⁰·exp(-r/R₀) with R₀=1
idx_align = 2 # r⁰·exp(-r/Lₐ) with Lₐ=2
idx_pursuit = 3 # r⁰·exp(-r/Lₚ) with Lₚ=4
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
true_c_pursuit = np.zeros(n_pursuit)
true_c_pursuit[idx_pursuit] = P_true
# Compute kernel profiles with confidence intervals
r_eval = np.linspace(0.01, 8.0, 200)
ci_repel = kernel_predict_ci(r_eval, kernels, c_repel, cov_repel)
ci_align = kernel_predict_ci(r_eval, kernels, c_align, cov_align)
ci_pursuit = kernel_predict_ci(r_eval, kernels, c_pursuit, cov_pursuit)
# True kernel profiles
true_repel = true_c_repel @ ci_repel["phi"]
true_align = true_c_align @ ci_align["phi"]
true_pursuit = true_c_pursuit @ ci_pursuit["phi"]
Bootstrap simulation¶
We simulate a trajectory from the inferred model starting from a relaxed mid-trajectory frame and using chunked integration with periodic neighbor rebuilds.
key_boot = random.PRNGKey(seed + 77)
proc_boot = OverdampedProcess(
inf.force_inferred._psf, inf.diffusion_inferred._psf,
)
proc_boot.set_params(
theta_F=inf.force_inferred.params,
theta_D=inf.diffusion_inferred.params,
)
# Start from a relaxed mid-trajectory frame
mid_frame = Xfull.shape[0] // 2
x0_boot = jnp.array(Xfull[mid_frame])
nbr_boot = make_neighbor_extras(np.asarray(x0_boot[:, :2]),
cutoff, box_np)
extras_boot = {"box": box}
extras_boot.update(nbr_boot)
proc_boot.set_extras(extras_global=extras_boot)
proc_boot.initialize(x0_boot)
boot_steps = 200
print(f"Bootstrap: {boot_steps} steps from mid-trajectory frame {mid_frame} ...")
key_boot, sub_boot = random.split(key_boot)
coll_boot = simulate_chunked(
proc_boot, dt=dt_sim, Nsteps=boot_steps, key=sub_boot,
cutoff=cutoff, box=box_np,
skin=0.0, rebuild_every=rebuild_every,
save_every=50,
spatial_dims=slice(0, 2),
nnz_safety=3.0, verbose=True,
)
Xboot = coll_boot.to_array(axis="time")
print(f"Bootstrap done: {Xboot.shape[0]} frames")
Bootstrap: 200 steps from mid-trajectory frame 1000 ...
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 1: 50/200 steps
max displacement = 1.235 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 2: 100/200 steps
max displacement = 1.166 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 3: 150/200 steps
max displacement = 1.202 (skin/2 = 0.000)
merged 50 sub-chunks → X shape (50, 3000, 3)
chunk 4: 200/200 steps
Bootstrap done: 200 frames
Thumbnail¶
Dedicated single-panel figure for the gallery thumbnail.
stamp_output()
examples/gallery/abp_nonreciprocal_demo.py:520: UserWarning: The figure layout has changed to tight
plt.tight_layout()
[Generated: 2026-06-30 12:52]
Total running time of the script: (48 minutes 38.677 seconds)