From c7cbd93918707ebfefa23c05df029eeda00d859d Mon Sep 17 00:00:00 2001 From: jung235 Date: Wed, 6 Mar 2024 17:42:11 +0900 Subject: [PATCH 1/5] feat: enable interacting particle generation in `sde` --- pydiffuser/models/core/sde.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/pydiffuser/models/core/sde.py b/pydiffuser/models/core/sde.py index 607467b..85e6142 100644 --- a/pydiffuser/models/core/sde.py +++ b/pydiffuser/models/core/sde.py @@ -73,9 +73,7 @@ def __init__( raise KeyError(f"Unsupported potential {potential} is encountered") super(OverdampedLangevin, self).__init__() - self.potential = ( - FIELD_REGISTRY[potential] if potential is not None else potential - ) + self.potential = potential self.potential_params = potential_params self.external_force = external_force self.friction_coefficient = friction_coefficient @@ -92,6 +90,8 @@ def generate( ) -> Ensemble: ens = super().generate(realization, length, dimension, dt, **generate_kwargs) realization, length, dimension, dt = list(self.generate_info.values())[:4] + if self.interacting: + raise RuntimeError("Interacting particles are not supported.") x = self._get_initial_position() # realization x 1 x dimension dx = jnp.zeros((realization, (length - 1), dimension)) # init @@ -99,10 +99,7 @@ def generate( if self.generate_hooks is not None: for hook in self.generate_hooks: - if self.interacting: - pass # TODO - else: - out: Array = self._load_hook(hook)() + out: Array = self._load_hook(hook)() if dx.shape == dx_shape: dx += out else: @@ -157,7 +154,7 @@ def get_diff_from_potential(self, r: Array, remaining_terms: Array) -> Array: return ( r - coeff - * grad(self.potential)(r, **self.potential_params) + * grad(FIELD_REGISTRY[self.potential])(r, **self.potential_params) # type: ignore[index] * self.generate_info["dt"] + remaining_terms # constant with respect to `r` ) From cd43466d2cc8237d057de8d3be45332c74123d97 Mon Sep 17 00:00:00 2001 From: jung235 Date: Wed, 6 Mar 2024 17:43:44 +0900 Subject: [PATCH 2/5] feat: add injection function for init args --- pydiffuser/models/core/base.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/pydiffuser/models/core/base.py b/pydiffuser/models/core/base.py index b342e1e..1572ee1 100644 --- a/pydiffuser/models/core/base.py +++ b/pydiffuser/models/core/base.py @@ -38,7 +38,7 @@ def from_config(cls, config: BaseDiffusionConfig) -> BaseDiffusion: if MODEL_KWARGS in params: params.pop() model = cls(*[getattr(config, param) for param in params[1:]]) - model.config = config + model._inject_init_args(config) return model @property @@ -70,7 +70,10 @@ def pre_generate(self, *generate_args) -> None: f"Unsupported dimension {dimension} is encountered" ) if self.interacting: - raise NotImplementedError("Interactive particles are unsupported") # TODO + logger.debug( + f"Generating interacting particles from `{self.__class__.__name__}`. " + "The calculation will be significantly slower than with non-interacting particles." + ) return @property @@ -85,6 +88,13 @@ def generate_info(self) -> Dict[str, Any]: pass return info + def _inject_init_args(self, config: BaseDiffusionConfig) -> None: + for param in vars(self): + if param in config: + setattr(self, param, getattr(config, param)) + self.config = config + return + def _stash_generate_args(self, *user_args) -> None: params = list(signature(self.generate).parameters.keys()) defaults = list(signature(self.generate).parameters.values()) From d6b6c3c64233226db5d343a9c2190b35ed0de884 Mon Sep 17 00:00:00 2001 From: jung235 Date: Wed, 6 Mar 2024 17:44:06 +0900 Subject: [PATCH 3/5] feat: add vicsek model --- pydiffuser/models/vicsek.py | 170 ++++++++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 pydiffuser/models/vicsek.py diff --git a/pydiffuser/models/vicsek.py b/pydiffuser/models/vicsek.py new file mode 100644 index 0000000..e6815d1 --- /dev/null +++ b/pydiffuser/models/vicsek.py @@ -0,0 +1,170 @@ +import math +from typing import Optional, Tuple + +import jax.numpy as jnp +from jax import lax +from numpy.random import rand + +from pydiffuser.exceptions import InvalidDimensionError +from pydiffuser.models.core import OverdampedLangevin, OverdampedLangevinConfig +from pydiffuser.tracer import Ensemble +from pydiffuser.typing import LongLongPosType, LongPosType +from pydiffuser.utils import jitted + + +class VicsekModelConfig(OverdampedLangevinConfig): + name: str = "vicsek" + + def __init__( + self, + boxsize: float = 40.0, + interaction_strength: float = 1.0, + interaction_radius: float = 2.0, + diffusivity: float = 0.0, + speed: float = 1.0, + rotational_diffusivity: float = 0.1, + angular_velocity: float = 0.0, + **kwargs, + ): + kwargs = { + param: kwargs[param] if param in kwargs else default + for param, default in OverdampedLangevinConfig.show_fields().items() + } + kwargs["diffusivity"] = diffusivity + kwargs["model_alias"] = VicsekModelConfig.name + super(VicsekModelConfig, self).__init__(**kwargs) + + self.boxsize = boxsize + self.interaction_strength = interaction_strength + self.interaction_radius = interaction_radius + self.speed = speed + self.rotational_diffusivity = rotational_diffusivity + self.angular_velocity = angular_velocity + + +class VicsekModel(OverdampedLangevin): + name: str = "vicsek" + + def __init__( + self, + boxsize: float, + interaction_strength: float, + interaction_radius: float, + speed: float, + rotational_diffusivity: float, + angular_velocity: float, + **model_kwargs, + ): + """ + We consider the Vicsek model utilizing active Brownian particles (ABPs) + in a square box of size L x L, subjected to a periodic boundary condition (PBC). + The following equation governs the velocity direction φᵢ of the ith particle: + dφᵢ K ___ + ─── = ω + ─── Σ sin(φⱼ - φᵢ) + √2Dr ηᵢ(t). + dt πR² j + Here, we ignore the excluded volume effect and external force term, which means + U = 0 and F = 0 in `pydiffuser.models.core.sde.OverdampedLangevin`. + For a detailed description of ABPs, see `pydiffuser.models.abp.ActiveBrownianParticle`. + + Args: + boxsize (float): L. + interaction_strength (float): K. + interaction_radius (float): R. + """ + + super(VicsekModel, self).__init__(**model_kwargs) + self.interacting = True + + self.boxsize = boxsize + self.interaction_strength = interaction_strength + self.interaction_radius = interaction_radius + self.speed = speed + self.rotational_diffusivity = rotational_diffusivity + self.angular_velocity = angular_velocity + + def pre_generate(self, *generate_args) -> None: + super().pre_generate(*generate_args) + dimension = self.generate_info["dimension"] + if dimension != 2: + raise InvalidDimensionError( + f"Unsupported dimension {dimension} is encountered" + ) + return + + def generate( + self, + realization: int = 1000, + length: int = 100, + dimension: int = 2, + dt: float = 1.0, + **generate_kwargs, + ) -> Ensemble: + self.pre_generate(realization, length, dimension, dt, *generate_kwargs.values()) + realization, length, dimension, dt = list(self.generate_info.values())[:4] + ens = Ensemble(dt=dt) + + phi = self._get_initial_orientation() # realization x 1 x 1 + x = self._get_initial_position() # realization x 1 x 2 + state = jnp.concatenate((phi, x), axis=-1) # realization x 1 x 3 + + dphi = self.get_diff_from_white_noise( + self.rotational_diffusivity, shape=(realization, (length - 1), 1) + ) + dphi += self.angular_velocity * dt + dx = self.get_diff_from_white_noise( + self.diffusivity, shape=(realization, (length - 1), dimension) + ) + dstate = jnp.concatenate((dphi, dx), axis=-1) # realization x (length - 1) x 3 + + _, stx = lax.scan( + f=self.get_next_state_from_vicsek_interaction, + init=jnp.squeeze(state, axis=1), # realization x 3 + xs=jnp.transpose(dstate, (1, 0, 2)), # (length - 1) x realization x 3 + ) + stx_phi, stx_x = jnp.split(stx, indices_or_sections=[1], axis=-1) + phi = jnp.concatenate((phi, jnp.transpose(stx_phi, (1, 0, 2))), axis=1) + x = jnp.concatenate((x, jnp.transpose(stx_x, (1, 0, 2))), axis=1) + + ens.update_microstate(microstate=x) + for id in range(realization): + ens[id].update_meta_dict(item={"direction": phi[id]}) + ens.update_meta_dict(item={"direction": phi}) + return ens + + def get_next_state_from_vicsek_interaction( + self, + state: LongPosType, + dstate: LongPosType, + ) -> Tuple[LongPosType, LongPosType]: + dt = self.generate_info["dt"] + phi, x = jnp.split(state, indices_or_sections=[1], axis=-1) + dphi, dx = jnp.split(dstate, indices_or_sections=[1], axis=-1) + coeff = self.interaction_strength / (jnp.pi * self.interaction_radius**2) + + dx += jnp.concatenate( + arrays=jitted.polar_to_cartesian(r=self.speed * dt, phi=phi), + axis=1, + ) + next_x = (x + dx) % self.boxsize # PBC + + abs_dx = jnp.abs(x[:, jnp.newaxis] - x[jnp.newaxis, :]) + abs_dx = jnp.where( + abs_dx <= self.boxsize / 2, abs_dx, self.boxsize - abs_dx + ) # PBC + dr = jnp.sqrt(jnp.sum(abs_dx**2, axis=-1)) + mask = jnp.where(dr <= self.interaction_radius, 1, 0) + sine = jnp.sin(phi.T - phi) + dphi_vicsek = coeff * jnp.sum(sine * mask, axis=-1) * dt + next_phi = phi + dphi + dphi_vicsek[:, jnp.newaxis] + + next_state = jnp.concatenate((next_phi, next_x), axis=-1) # realization x 3 + return next_state, next_state + + def _get_initial_position( + self, realization: Optional[int] = None, dimension: Optional[int] = None + ) -> LongLongPosType: + if realization is None or dimension is None: + realization, _, dimension = list(self.generate_info.values())[:3] + shape = (realization, 1, dimension) + x = jitted.get_noise(generator=rand, size=math.prod(shape), shape=shape) # type: ignore[arg-type] + return x * self.boxsize # PBC From 557ee71094035724e9ccefa465c1437c27ffc1b2 Mon Sep 17 00:00:00 2001 From: jung235 Date: Wed, 6 Mar 2024 17:45:09 +0900 Subject: [PATCH 4/5] feat: cli for vicsek model --- pydiffuser/_cli/cli_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pydiffuser/_cli/cli_utils.py b/pydiffuser/_cli/cli_utils.py index 9effddb..5093220 100644 --- a/pydiffuser/_cli/cli_utils.py +++ b/pydiffuser/_cli/cli_utils.py @@ -12,6 +12,7 @@ def get_model_info(): "levy": _, "rtp": _, "smoluchowski": "1d, 2d", + "vicsek": "2d", } info = [("NAME", "MODEL", "CONFIG", "DIMENSION")] info += [ From 52b298b4e71b355736e045056ea5e593ac84b659 Mon Sep 17 00:00:00 2001 From: jung235 Date: Wed, 6 Mar 2024 17:46:57 +0900 Subject: [PATCH 5/5] fix: add vicsek model in constructors --- pydiffuser/__init__.py | 4 ++++ pydiffuser/models/__init__.py | 3 +++ 2 files changed, 7 insertions(+) diff --git a/pydiffuser/__init__.py b/pydiffuser/__init__.py index 8ec01c8..9c4f696 100644 --- a/pydiffuser/__init__.py +++ b/pydiffuser/__init__.py @@ -12,6 +12,8 @@ RunAndTumbleParticleConfig, SmoluchowskiEquation, SmoluchowskiEquationConfig, + VicsekModel, + VicsekModelConfig, ) from .tracer import Ensemble, Trajectory from .utils import load, save @@ -30,6 +32,8 @@ "RunAndTumbleParticleConfig", "SmoluchowskiEquation", "SmoluchowskiEquationConfig", + "VicsekModel", + "VicsekModelConfig", "Ensemble", "Trajectory", "load", diff --git a/pydiffuser/models/__init__.py b/pydiffuser/models/__init__.py index 7d45dd5..2dcbf8a 100644 --- a/pydiffuser/models/__init__.py +++ b/pydiffuser/models/__init__.py @@ -9,6 +9,7 @@ from .levy import LevyWalk, LevyWalkConfig from .rtp import RunAndTumbleParticle, RunAndTumbleParticleConfig from .smoluchowski import SmoluchowskiEquation, SmoluchowskiEquationConfig +from .vicsek import VicsekModel, VicsekModelConfig __all__ = [ "ActiveBrownianParticle", @@ -23,6 +24,8 @@ "RunAndTumbleParticleConfig", "SmoluchowskiEquation", "SmoluchowskiEquationConfig", + "VicsekModel", + "VicsekModelConfig", ]