Pydiffuser is a numerical simulation framework for nonequilibrium statistical physics based on JAX.
This package mainly aims:
- to share code to implement a numerical simulation on physical models written in various forms of stochastic differential equations.
- to revisit recent research highlights in nonequilibrium statistical physics.
- to reduce the repeated code on time-series data analysis, e.g., statistical analysis of single-particle trajectory for SPT experiments.
- to provide the skeleton of stochastic modeling for anyone interested in stochastic processes.
Python 3.10+, jax>=0.4.18, and jaxlib>=0.4.18.
From PyPI
$ pip install pydiffuserIf properly installed, you can run:
$ pydiffuser --version
pydiffuser, version 0.0.3Pydiffuser provides various stochastic models that implement a numerical simulation based on the Monte Carlo method.
All Pydiffuser's models inherit an abstract class pydiffuser.models.BaseDiffusion and initiate the simulation after a method generate is called.
For the simplest case, you can produce a non-interacting Brownian motion at low Reynolds numbers as follows:
from pydiffuser.models import BrownianMotion
from pydiffuser.tracer import Ensemble, Trajectory
model = BrownianMotion()
ensemble: Ensemble = model.generate()
tracer: Trajectory = ensemble[0] # the 0th particleRelevant stochastic observables, such as mean-squared displacement and normalized velocity autocorrelation function, can be calculated through the methods of Trajectory and Ensemble.
tamsd = tracer.get_mean_squared_displacement(lagtime=1, rolling=True)
eamsd = ensemble.get_mean_squared_displacement(lagtime=1, rolling=False)
eatamsd = ensemble.get_mean_squared_displacement(lagtime=1, rolling=True)You can visualize the trajectory using matplotlib:
It is obtained by matplotlib.pyplot.plot(tracer.position_x1, tracer.position_x2).
List all stochastic models supported by Pydiffuser.
$ pydiffuser model list
NAME MODEL CONFIG DIMENSION
abp ActiveBrownianParticle ActiveBrownianParticleConfig 2d
aoup ActiveOUParticle ActiveOUParticleConfig 1d, 2d, 3d
bm BrownianMotion BrownianMotionConfig 1d, 2d, 3d
levy LevyWalk LevyWalkConfig 1d, 2d, 3d
mips PhaseSeparation PhaseSeparationConfig 1d, 2d, 3d
rtp RunAndTumbleParticle RunAndTumbleParticleConfig 1d, 2d, 3d
smoluchowski SmoluchowskiEquation SmoluchowskiEquationConfig 1d, 2d
vicsek VicsekModel VicsekModelConfig 2d When generating
═════════════════════════════════════════════════════════════════════════════════════════════════
Model Method Running time [s] for N x L =
10² x 10⁵ 10³ x 10⁴ 10⁴ x 10³
─────────────────────────────────────────────────────────────────────────────────────────────────
`loop` [*] 3.62 (0.19) 3.45 (0.23) 3.37 (0.21)
─────────────────────────────────────────────────────────────────────────────────────────────────
`abp` `generate` 1.95 (0.14) 1.74 (0.12) 1.59 (0.11)
`aoup` `generate` 1.61 (0.08) 1.61 (0.15) 1.55 (0.09)
`bm` `generate` 1.45 (0.11) 1.46 (0.13) 1.46 (0.14)
`smoluchowski` `generate` 1.71 (0.12) 1.67 (0.15) 1.64 (0.13)
─────────────────────────────────────────────────────────────────────────────────────────────────
`bm` `create` 1440.72 (158.16) 964.90 (83.06) 1195.41 (94.28)
═════════════════════════════════════════════════════════════════════════════════════════════════[*]
def loop(N: int, L: int) -> float:
"""Even the most straightforward loop requires over 3 [s] for all (N, L) conditions.
"""
t1 = time.time()
xes = []
for _ in range(N):
x = []
for _ in range(1, L):
x.append([])
xes.append(x)
t2 = time.time()
return t2 - t1The represented running time is a mean
class pydiffuser.tracer.Trajectory ∈ class pydiffuser.tracer.Ensemble
get_incrementsget_displacement_momentget_mean_squared_displacementget_cosine_momentget_velocity_autocorrelationget_real_time
The above methods are defined in both Trajectory and Ensemble to enhance transparency.
Using Trajectory, the statistical analysis of single-particle trajectory can be accelerated.
We introduce a configuration to deal with extensive parameter manipulation.
For instance, see config.json, which contains all parameters demanded to instantiate pydiffuser.ActiveBrownianParticle.
Every JSON of the configurations listed in CLI can be obtained as follows:
import pydiffuser as pyd
from pydiffuser.models import ActiveBrownianParticle, ActiveBrownianParticleConfig
config = ActiveBrownianParticleConfig()
config.to_json(json_path=<JSON_PATH>)We suggest a research pipeline:
┌────┐ ┌─────────────────────┐ ┌───────────────┐ ┌──────────┐ ┌────────────┐
│JSON├──>──┤`BaseDiffusionConfig`├──>──┤`BaseDiffusion`├──>──┤`Ensemble`├──>──┤NPY | PICKLE│
└────┘ [1] └─────────────────────┘ [2] └───────────────┘ [3] └──────────┘ [4] └────────────┘It can be automized as follows:
config = ActiveBrownianParticleConfig.from_json(json_path=<JSON_PATH>) # [1]
model = ActiveBrownianParticle.from_config(config=config) # [2]
ensemble = model.generate() # [3]
ensemble.to_npy(npy_path=<NPY_PATH>) # [4]You can save and load any picklable object through pydiffuser.save and pydiffuser.load.
MODEL_PATH = "model.pickle"
pyd.save(obj=model, pickle_path=MODEL_PATH) # Here, <PICKLE_PATH> = MODEL_PATH
model = pyd.load(pickle_path=MODEL_PATH)Hyperdiffusion of Poissonian run-and-tumble particles in two dimensions
@misc{jung2023pydiffuser,
title = {Pydiffuser: a simulation framework for nonequilibrium statistical physics},
author = {Jung, Yurim},
year = {2023},
note = {doi: 10.5281/zenodo.10017027},
}