feat: add the MIPS model of AOUPs subjected to the WCA potential

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ repos:
     hooks:
       - id: ruff
   - repo: https://github.com/psf/black
-    rev: 23.3.0
+    rev: 24.3.0
     hooks:
       - id: black
   - repo: https://github.com/pre-commit/mirrors-mypy

pydiffuser/__init__.py

@@ -8,6 +8,8 @@
     BrownianMotionConfig,
     LevyWalk,
     LevyWalkConfig,
+    PhaseSeparation,
+    PhaseSeparationConfig,
     RunAndTumbleParticle,
     RunAndTumbleParticleConfig,
     SmoluchowskiEquation,
@@ -28,6 +30,8 @@
     "BrownianMotionConfig",
     "LevyWalk",
     "LevyWalkConfig",
+    "PhaseSeparation",
+    "PhaseSeparationConfig",
     "RunAndTumbleParticle",
     "RunAndTumbleParticleConfig",
     "SmoluchowskiEquation",

pydiffuser/_cli/cli_utils.py

@@ -10,6 +10,7 @@ def get_model_info():
         "aoup": _,
         "bm": _,
         "levy": _,
+        "mips": _,
         "rtp": _,
         "smoluchowski": "1d, 2d",
         "vicsek": "2d",

pydiffuser/mech/fields.py

@@ -6,7 +6,7 @@
 from jax import Array, jit
 from jaxlib.xla_extension import PjitFunction
 
-from pydiffuser.typing import ArrayType, ConstType
+from pydiffuser.typing import ArrayType, ConstType, PosType
 
 FIELD_REGISTRY = {}
 
@@ -24,11 +24,14 @@ def decorator() -> Any:
 
 
 def get_static_argsigs(
-    potential: Any, static_args_start: int = 1
+    potential: Any, static_args_start: int = 1, static_args_end: int | None = None
 ) -> Dict[str, ConstType]:
-    i = static_args_start
-    keys = list(signature(potential).parameters.keys())[i:]
-    values = list(signature(potential).parameters.values())[i:]
+    sigs = signature(potential).parameters
+    i = static_args_start if static_args_start else 0
+    j = static_args_end if static_args_end else len(sigs)
+
+    keys = list(sigs.keys())[i:j]
+    values = list(sigs.values())[i:j]
     return {k: v.default for k, v in zip(keys, values, strict=True)}
 
 
@@ -52,5 +55,41 @@ def lennard_jones_potential():
     pass
 
 
-def weeks_chandler_andersen_potential():
-    pass
+@register
+@partial(jit, static_argnames=("epsilon", "boxsize"))
+def wca_potential(
+    ri: PosType,
+    rj: PosType,
+    sigma: ConstType,
+    epsilon: ConstType = 10.0,
+    boxsize: ConstType | None = None,
+) -> Array:
+    """The Weeks-Chandler-Andersen potential is given as
+    ```
+                     σᵢⱼ       σᵢⱼ     1
+        U(rᵢⱼ) = 4ε[(───)¹² - (───)⁶ + ─]
+                     rᵢⱼ       rᵢⱼ     4
+    ```
+    for rᵢⱼ < 2¹ᐟ⁶σᵢⱼ and U(rᵢⱼ) = 0 otherwise.
+    Here, we have rᵢⱼ = |rᵢ - rⱼ| and σᵢⱼ = (σᵢ + σⱼ) / 2,
+    where σᵢ represents the diameter of the ith particle.
+
+    Args:
+        ri (PosType): rᵢ.
+        rj (PosType): rⱼ.
+        sigma (ConstType): σᵢⱼ.
+        epsilon (ConstType, optional): ε.
+        boxsize (ConstType | None, optional): Size of the unit cell.
+            If not None, a periodic boundary condition (PBC) is enforced.
+    """
+
+    rij_vec = ri - rj
+    if boxsize:
+        rij_vec = rij_vec - boxsize * jnp.round(rij_vec / boxsize)
+    rij = jnp.linalg.norm(rij_vec)
+
+    return jnp.where(
+        rij < 2 ** (1 / 6) * sigma,
+        4 * epsilon * ((sigma / rij) ** 12 - (sigma / rij) ** 6 + 1 / 4),
+        0,
+    )

pydiffuser/models/__init__.py

@@ -7,6 +7,7 @@
 from .bm import BrownianMotion, BrownianMotionConfig
 from .core.base import BaseDiffusion
 from .levy import LevyWalk, LevyWalkConfig
+from .mips import PhaseSeparation, PhaseSeparationConfig
 from .rtp import RunAndTumbleParticle, RunAndTumbleParticleConfig
 from .smoluchowski import SmoluchowskiEquation, SmoluchowskiEquationConfig
 from .vicsek import VicsekModel, VicsekModelConfig
@@ -20,6 +21,8 @@
     "BrownianMotionConfig",
     "LevyWalk",
     "LevyWalkConfig",
+    "PhaseSeparation",
+    "PhaseSeparationConfig",
     "RunAndTumbleParticle",
     "RunAndTumbleParticleConfig",
     "SmoluchowskiEquation",

pydiffuser/models/abp.py

@@ -55,9 +55,11 @@ def __init__(
     ):
         """
         Consider an overdamped Langevin equation for orientation φ:
+        ```
             dφ        ___
             ── = ω + √2Dr η(t),
             dt
+        ```
         where ω is a constant angular velocity,
         Dr is a rotational diffusion coefficient,
         and η(t) denotes a Gaussian white noise with zero mean and unit variance.

pydiffuser/models/aoup.py

@@ -1,9 +1,12 @@
+from functools import partial
 from typing import List
 
 import jax.numpy as jnp
 from jax import Array
+from numpy.random import uniform
 
 from pydiffuser.models.core import OverdampedLangevin, OverdampedLangevinConfig
+from pydiffuser.utils import jitted
 
 _GENERATE_HOOKS = ["ornstein_uhlenbeck_process"]
 
@@ -50,9 +53,11 @@ def __init__(
     ):
         """
         Consider an Ornstein-Uhlenbeck process for a self-propulsion velocity p:
+        ```
             dp              ____
             ── = - μ x p + √2Dou Γ(t),
             dt
+        ```
         where Γ(t) is a Gaussian white noise with zero mean and unit variance.
         Note that p is coupled with the overdamped Langevin equation
             written in `pydiffuser.models.core.sde.OverdampedLangevin`.
@@ -72,7 +77,12 @@ def __init__(
 
     def ornstein_uhlenbeck_process(self) -> Array:
         realization, length, dimension, _ = self.generate_info.values()
-        p = jnp.ones((realization, 1, dimension))  # init
+        # TODO patternize
+        p = jitted.get_noise(  # init
+            generator=partial(uniform, -1, 1),
+            size=realization * dimension,
+            shape=(realization, 1, dimension),
+        )
         dp = self.get_diff_from_white_noise(
             diffusivity=self.diffusion_coefficient,
             shape=(realization, (length - 2), dimension),

pydiffuser/models/core/base.py

@@ -6,6 +6,7 @@
 from inspect import signature
 from typing import Any, Callable, Dict, Optional
 
+import jax
 import jax.numpy as jnp
 from numpy.random import rand, randint
 
@@ -27,6 +28,7 @@ def __init__(self) -> None:
         self.config: Optional[BaseDiffusionConfig] = None
         self._state_dict: Dict[str, Any] = {}
         self._interacting: bool = False
+        self._precision_x64: bool = False
 
     @classmethod
     def from_config(cls, config: BaseDiffusionConfig) -> BaseDiffusion:
@@ -49,6 +51,14 @@ def interacting(self):
     def interacting(self, interacting: bool):
         self._interacting = interacting
 
+    @property
+    def precision_x64(self):
+        return self._precision_x64
+
+    @precision_x64.setter
+    def precision_x64(self, precision_x64: bool):
+        self._precision_x64 = precision_x64
+
     @abc.abstractmethod
     def generate(
         self,
@@ -74,6 +84,12 @@ def pre_generate(self, *generate_args) -> None:
                 f"Generating interacting particles from `{self.__class__.__name__}`. "
                 "The calculation will be significantly slower than with non-interacting particles."
             )
+
+        jax.config.update("jax_enable_x64", self.precision_x64)
+        if self.precision_x64:
+            logger.debug(
+                f"The simulation launched from `{self.__class__.__name__}` requires x64 precision."
+            )
         return
 
     @property

pydiffuser/models/core/ctrw.py

@@ -84,8 +84,7 @@ def get_orientation_steps(self, realization: int, capacity: int) -> Array:
     @staticmethod
     def slice(arr: Array, threshold: ConstType, axis: int = 0) -> Array:
         """For a 2darray with shape `realization` x -1 (here, `arr.T`),
-        _summary_
-
+        ```
         □□□□□...□□│□□□□□□□            □□□□□...□□ppppppp│
         □□□□□...□□│□□□□□□□            □□□□□...□□ppppppp│
         □□□□□...  └──┐                □□□□□...□□□□□pppp│
@@ -94,6 +93,8 @@ def slice(arr: Array, threshold: ConstType, axis: int = 0) -> Array:
         □□□□□...□□□□□│□□□□            □□□□□...□□□□□pppp│
         □□□□□...□□□□□│□□□□            □□□□□...□□□□□pppp│
                      └───> threshold                   └> padding
+        ```
+        _summary_
 
         Args:
             arr (Array): 2darray with shape -1 x `realization`.

pydiffuser/models/core/sde.py

@@ -50,9 +50,11 @@ def __init__(
     ):
         """
         Consider an overdamped Langevin equation in d dimensions:
+        ```
             dr     1         1      __
             ── = - ─ ∇U(r) + ─ F + √2D ξ(t) + p,
             dt     γ         γ
+        ```
         where γ is a friction coefficient,
         D is a translational diffusion coefficient,
         p is a user-defined stochastic variable (e.g., self-propulsion velocity),