diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 54a04f3..2d3b8fe 100644
--- a/.pre-commit-config.yaml
+++ b/.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
diff --git a/pydiffuser/__init__.py b/pydiffuser/__init__.py
index 9c4f696..4c52e71 100644
--- a/pydiffuser/__init__.py
+++ b/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",
diff --git a/pydiffuser/_cli/cli_utils.py b/pydiffuser/_cli/cli_utils.py
index 5093220..83357ad 100644
--- a/pydiffuser/_cli/cli_utils.py
+++ b/pydiffuser/_cli/cli_utils.py
@@ -10,6 +10,7 @@ def get_model_info():
         "aoup": _,
         "bm": _,
         "levy": _,
+        "mips": _,
         "rtp": _,
         "smoluchowski": "1d, 2d",
         "vicsek": "2d",
diff --git a/pydiffuser/mech/fields.py b/pydiffuser/mech/fields.py
index efefe20..db5fe67 100644
--- a/pydiffuser/mech/fields.py
+++ b/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,
+    )
diff --git a/pydiffuser/models/__init__.py b/pydiffuser/models/__init__.py
index 2dcbf8a..d8a6168 100644
--- a/pydiffuser/models/__init__.py
+++ b/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",
diff --git a/pydiffuser/models/abp.py b/pydiffuser/models/abp.py
index 37dc28b..b6a7fdd 100644
--- a/pydiffuser/models/abp.py
+++ b/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.
diff --git a/pydiffuser/models/aoup.py b/pydiffuser/models/aoup.py
index 36409db..f53161f 100644
--- a/pydiffuser/models/aoup.py
+++ b/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),
diff --git a/pydiffuser/models/core/base.py b/pydiffuser/models/core/base.py
index 1572ee1..c82b2dd 100644
--- a/pydiffuser/models/core/base.py
+++ b/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
diff --git a/pydiffuser/models/core/ctrw.py b/pydiffuser/models/core/ctrw.py
index e0d29d0..6f3ae01 100644
--- a/pydiffuser/models/core/ctrw.py
+++ b/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`.
diff --git a/pydiffuser/models/core/sde.py b/pydiffuser/models/core/sde.py
index 85e6142..2340915 100644
--- a/pydiffuser/models/core/sde.py
+++ b/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),
diff --git a/pydiffuser/models/mips.py b/pydiffuser/models/mips.py
new file mode 100755
index 0000000..c00545b
--- /dev/null
+++ b/pydiffuser/models/mips.py
@@ -0,0 +1,171 @@
+import math
+from functools import partial
+from typing import Any, Dict, Optional, Tuple
+
+import jax.numpy as jnp
+from jax import Array, grad, lax, vmap
+from numpy.random import rand
+
+from pydiffuser.mech.fields import FIELD_REGISTRY, get_static_argsigs, wca_potential
+from pydiffuser.models.aoup import ActiveOUParticle, ActiveOUParticleConfig
+from pydiffuser.tracer import Ensemble
+from pydiffuser.typing import LongLongPosType, LongPosType, PosType
+from pydiffuser.utils import jitted
+
+DEFAULT_POTENTIAL = wca_potential.__name__
+DEFAULT_POTENTIAL_PARAMS = get_static_argsigs(wca_potential, 3, 4)
+
+
+class PhaseSeparationConfig(ActiveOUParticleConfig):
+    name: str = "mips"
+
+    def __init__(
+        self,
+        boxsize: float = 100.0,
+        mean_diameter: float = 2.0,
+        potential: str = DEFAULT_POTENTIAL,
+        potential_params: Dict[str, Any] = DEFAULT_POTENTIAL_PARAMS,
+        diffusivity: float = 1.0,
+        drift_coefficient: float = 1.0,
+        diffusion_coefficient: float = 0.1,
+        **kwargs,
+    ):
+        kwargs = {
+            param: kwargs[param] if param in kwargs else default
+            for param, default in ActiveOUParticleConfig.show_fields().items()
+        }
+        kwargs["potential"] = potential
+        kwargs["potential_params"] = potential_params
+        kwargs["diffusivity"] = diffusivity
+        kwargs["drift_coefficient"] = drift_coefficient
+        kwargs["diffusion_coefficient"] = diffusion_coefficient
+        kwargs["model_alias"] = PhaseSeparationConfig.name
+        super(PhaseSeparationConfig, self).__init__(**kwargs)
+
+        self.boxsize = boxsize
+        self.mean_diameter = mean_diameter
+
+
+class PhaseSeparation(ActiveOUParticle):
+    name: str = "mips"
+
+    def __init__(
+        self,
+        boxsize: float,
+        mean_diameter: float,
+        potential: str,
+        potential_params: Dict[str, Any],
+        diffusivity: float,
+        drift_coefficient: float,
+        diffusion_coefficient: float,
+        **model_kwargs,
+    ):
+        """
+        We consider AOUPs interacting via the Weeks-Chandler-Andersen (WCA) potential
+        in a square box of unit cell size L, subjected to a periodic boundary condition (PBC).
+        Note that the particles can exhibit motility-induced phase separation (MIPS).
+        The following equation governs the ith particle:
+        ```
+            drᵢ     1              1      __
+            ─── = - ─ Σ ∇ᵢU(rᵢⱼ) + ─ F + √2D ξᵢ(t) + pᵢ,
+            dt      γ j            γ
+        ```
+        where rᵢⱼ = |rⱼ - rᵢ| is the magnitude of the relative position vector.
+        Here, U is the purely repulsive WCA potential defined in `pydiffuser.mech.fields`.
+        For a detailed description of pᵢ, see `pydiffuser.models.aoup.ActiveOUParticle`.
+
+        Args:
+            boxsize (float): L.
+            mean_diameter (float): The mean diameter of all particles.
+        """
+
+        super(PhaseSeparation, self).__init__(
+            potential=potential,
+            potential_params=potential_params,
+            diffusivity=diffusivity,
+            drift_coefficient=drift_coefficient,
+            diffusion_coefficient=diffusion_coefficient,
+            **model_kwargs,
+        )
+        self.interacting = True
+        self.precision_x64 = True
+
+        self.boxsize = boxsize
+        self.mean_diameter = mean_diameter
+
+    def generate(
+        self,
+        realization: int = 50,
+        length: int = 1000,
+        dimension: int = 2,
+        dt: float = 0.01,
+        **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)
+
+        hook = self.generate_hooks[-1]  # type: ignore[index]
+        assert (
+            hook == self.ornstein_uhlenbeck_process.__name__
+            and self.potential == DEFAULT_POTENTIAL
+        ), "Only AOUPs imposed on the WCA potential are supported"
+
+        sigma = jitted.get_noise(rand, size=realization) * self.mean_diameter * 2
+        x = self._get_initial_position()  # realization x 1 x dimension
+        dx: Array = self._load_hook(hook)()  # realization x (length - 1) x dimension
+        dx_shape = dx.shape
+
+        # add terms given in the Langevin eq.
+        if self.diffusivity:
+            dx += self.get_diff_from_white_noise(self.diffusivity, dx_shape)
+
+        if self.external_force:
+            dx += self.get_diff_from_const_force(self.external_force, dx_shape)
+
+        _, stx = lax.scan(
+            f=partial(self.get_next_position_from_pairwise_potential, sigma=sigma),
+            init=jnp.squeeze(x, axis=1),  # realization x dimension
+            xs=jnp.transpose(dx, (1, 0, 2)),  # (length - 1) x realization x dimension
+        )
+        x = jnp.concatenate((x, jnp.transpose(stx, (1, 0, 2))), axis=1)
+
+        ens.update_microstate(microstate=x)
+        for id in range(realization):
+            ens[id].update_meta_dict(item={"diameter": sigma[id]})
+        ens.update_meta_dict(item={"diameter": sigma})
+        return ens
+
+    def get_next_position_from_pairwise_potential(
+        self, x: LongPosType, dx: LongPosType, sigma: PosType
+    ) -> Tuple[LongPosType, LongPosType]:
+        dt = self.generate_info["dt"]
+        potential_fn = partial(
+            FIELD_REGISTRY[self.potential],
+            boxsize=self.boxsize,  # PBC
+            **self.potential_params,
+        )
+        sij = (sigma[:, jnp.newaxis] + sigma[jnp.newaxis, :]) / 2
+
+        vmap_fn = vmap(
+            vmap(grad(potential_fn, argnums=0), in_axes=(None, 0, 0)),
+            in_axes=(0, None, 0),
+        )
+        del_fn = vmap_fn(x, x, sij)  # realization x realization x dimension
+        del_fn = jnp.nan_to_num(del_fn, nan=0.0)
+
+        # sum over all j to calculate Σⱼ∇ᵢU(rᵢⱼ)
+        sum_del_fn = jnp.sum(del_fn, axis=1)  # realization x dimension
+        dx_wca = -1 / self.friction_coefficient * sum_del_fn * dt
+
+        next_x = (x + dx + dx_wca) % self.boxsize  # PBC
+        return next_x, next_x
+
+    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
diff --git a/pydiffuser/models/vicsek.py b/pydiffuser/models/vicsek.py
index e6815d1..1421e4f 100644
--- a/pydiffuser/models/vicsek.py
+++ b/pydiffuser/models/vicsek.py
@@ -59,9 +59,11 @@ def __init__(
         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`.
@@ -147,11 +149,10 @@ def get_next_state_from_vicsek_interaction(
         )
         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))
+        dx_vec = x[:, jnp.newaxis] - x[jnp.newaxis, :]
+        dx_vec = dx_vec - self.boxsize * jnp.round(dx_vec / self.boxsize)  # PBC
+        dr = jnp.linalg.norm(dx_vec, 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
diff --git a/pydiffuser/tracer/ensemble.py b/pydiffuser/tracer/ensemble.py
index 2eccea2..307ea03 100644
--- a/pydiffuser/tracer/ensemble.py
+++ b/pydiffuser/tracer/ensemble.py
@@ -88,7 +88,10 @@ def __delitem__(self, *args, **kwargs):
     def __setitem__(self, id: int, tracer: Trajectory) -> None:
         if not isinstance(tracer, Trajectory):
             raise ValueError("Only `Trajectory` can be an element of `Ensemble`")
+        if jnp.any(jnp.isnan(tracer.position)):
+            raise ValueError("`Trajectory` contains NaN values")
         assert tracer.dt == self.dt
+        
         if id >= self.N:
             if id > self.N:
                 logger.warning(
diff --git a/requirements.txt b/requirements.txt
index 4f3cf6e..b524c26 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 # dev
-black==23.3.0
+black>=24.3.0
 ruff==0.0.285
 mypy==1.4.1
 pre-commit>=2.21.0