jaxdem.forces#

Interface for defining force laws and their corresponding potential energy.

Classes

ForceModel([required_material_properties, laws])

Abstract base class for defining inter-particle force laws and their potential energies.

SpringForce([required_material_properties, laws])

A ForceModel implementation for a linear spring-like interaction between particles.
class jaxdem.forces.ForceModel(required_material_properties: Tuple[str, ...] = (), laws: Tuple[ForceModel, ...] = ())[source][source]#

Bases: Factory[ForceModel], ABC

Abstract base class for defining inter-particle force laws and their potential energies.

Concrete subclasses implement specific force and energy models, such as linear springs, Hertzian contacts, etc.

Notes

  • Implementations should be JIT-compilable.

  • The force() and energy() methods should correctly handle the case where i and j refer to the same particle (i.e., i == j). There is no guarantee that self-interaction calls will not occur.

Example

To define a custom force model, inherit from ForceModel and implement its abstract methods:

>>> @ForceModel.register("myCustomForce")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomForce(ForceModel):
        ...
required_material_properties: Tuple[str, ...]#

A static tuple of strings specifying the material properties required by this force model.

These properties (e.g., ‘young_eff’, ‘restitution’) must be present in the System.mat_table for the model to function correctly. This is used for validation.

laws: Tuple[ForceModel, ...]#

A static tuple of other ForceModel instances that compose this force model.

This allows for creating composite force models (e.g., a total force being the sum of a spring force and a damping force).

abstractmethod static force(i: int, j: int, state: State, system: System) → Array[source][source]#

Compute the force vector acting on particle \(i\) due to particle \(j\).

Parameters:

  • i (int) – Index of the first particle (on which the force is acting).

  • j (int) – Index of the second particle (which is exerting the force).

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Force vector acting on particle \(i\) due to particle \(j\). Shape (dim,).

Return type:

jax.Array

Raises:

NotImplementedError – This is an abstract method and must be implemented by subclasses.

Example

This method is called internally by a Collider when computing total forces:

>>> force_on_particle_0_from_1 = system.force_model.force(0, 1, state, system)
abstractmethod static energy(i: int, j: int, state: State, system: System) → Array[source][source]#

Compute the potential energy of the interaction between particle \(i\) and particle \(j\).

Parameters:

  • i (int) – Index of the first particle.

  • j (int) – Index of the second particle.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Scalar JAX array representing the potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array

Raises:

NotImplementedError – This is an abstract method and must be implemented by subclasses.

Example

This method is typically called internally by a Collider when computing the total potential energy of the system:

>>> energy_0_1 = system.force_model.energy(0, 1, state, system)
classmethod create(key: str, /, **kw: Any) → T[source]#

Creates and returns an instance of a registered subclass.

This method looks up the subclass associated with the given key in the factory’s registry and then calls its constructor with the provided arguments.

Parameters:

  • key (str) – The registration key of the subclass to be created.

  • **kw (Any) – Arbitrary keyword arguments to be passed directly to the constructor of the registered subclass.

Returns:

An instance of the registered subclass.

Return type:

T

Raises:

  • KeyError – If the provided key is not found in the factory’s registry.

  • TypeError – If the provided **kw arguments do not match the signature of the registered subclass’s constructor.

Example

Given Foo factory and Bar registered:

>>> bar_instance = Foo.create("bar", value=42)
>>> print(bar_instance)
Bar(value=42)
classmethod register(key: str | None = None) → Callable[[Type[T]], Type[T]][source]#

Registers a subclass with the factory’s registry.

This method returns a decorator that can be applied to a class to register it under a specific key.

Parameters:

key (str or None, optional) – The string key under which to register the subclass. If None, the lowercase name of the subclass itself will be used as the key.

Returns:

A decorator function that takes a class and registers it, returning the class unchanged.

Return type:

Callable[[Type[T]], Type[T]]

Raises:

ValueError – If the provided key (or the default class name) is already registered in the factory’s registry.

Example

Register a class named “MyComponent” under the key “mycomp”:

>>> @MyFactory.register("mycomp")
>>> class MyComponent:
>>>     ...

Register a class named “DefaultComponent” using its own name as the key:

>>> @MyFactory.register()
>>> class DefaultComponent:
>>>     ...
class jaxdem.forces.SpringForce(required_material_properties: Tuple[str, ...] = ('young_eff',), laws: Tuple[ForceModel, ...] = ())[source][source]#

Bases: ForceModel

A ForceModel implementation for a linear spring-like interaction between particles.

Notes

  • The ‘effective Young’s modulus’ (\(k_{eff,\; ij}\)) is retrieved from the System.mat_table.young_eff based on the material IDs of the interacting particles.

  • The force is zero if \(i == j\).

  • A small epsilon is added to the squared distance (\(r^2\)) before taking the square root to prevent division by zero or NaN issues when particles are perfectly co-located.

The penetration \(\delta\) (overlap) between two particles \(i\) and \(j\) is:

\[\delta = (R_i + R_j) - r\]

where \(R_i\) and \(R_j\) are the radii of particles \(i\) and \(j\) respectively, and \(r = ||r_{ij}||\) is the distance between their centers.

The scalar overlap \(s\) is defined as:

\[s = \max\left(0, \frac{R_i + R_j}{r} - 1\right)\]

The force \(F_{ij}\) acting on particle \(i\) due to particle \(j\) is:

\[F_{ij} = k_{eff,\; ij} \cdot s \cdot r_{ij}\]

The potential energy \(E_{ij}\) of the interaction is:

\[E_{ij} = \frac{1}{2} k_{eff,\; ij} \cdot s^2\]

where \(k_{eff,\; ij}\) is the effective Young’s modulus for the particle pair.

Example

To use SpringForce in a simulation, specify it as the force_model when creating your System:

>>> system = jaxdem.System.create(dim=3, force_model_type="spring", force_model_kw={})

For this force model, the typical MaterialMatchmaker: type is “harmonic”.

required_material_properties: Tuple[str, ...]#

A static tuple of strings specifying the material properties required by this force model.

These properties (e.g., ‘young_eff’, ‘restitution’) must be present in the System.mat_table for the model to function correctly. This is used for validation.

classmethod create(key: str, /, **kw: Any) → T[source]#

Creates and returns an instance of a registered subclass.

This method looks up the subclass associated with the given key in the factory’s registry and then calls its constructor with the provided arguments.

Parameters:

  • key (str) – The registration key of the subclass to be created.

  • **kw (Any) – Arbitrary keyword arguments to be passed directly to the constructor of the registered subclass.

Returns:

An instance of the registered subclass.

Return type:

T

Raises:

  • KeyError – If the provided key is not found in the factory’s registry.

  • TypeError – If the provided **kw arguments do not match the signature of the registered subclass’s constructor.

Example

Given Foo factory and Bar registered:

>>> bar_instance = Foo.create("bar", value=42)
>>> print(bar_instance)
Bar(value=42)
static force(i: int, j: int, state: State, system: System) → Array[source][source]#

Compute linear spring-like interaction force acting on particle \(i\) due to particle \(j\).

Returns zero when \(i = j\).

Parameters:

  • i (int) – Index of the first particle.

  • j (int) – Index of the second particle.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Force vector acting on particle \(i\) due to particle \(j\).

Return type:

jax.Array

laws: Tuple['ForceModel', ...]#

A static tuple of other ForceModel instances that compose this force model.

This allows for creating composite force models (e.g., a total force being the sum of a spring force and a damping force).

classmethod register(key: str | None = None) → Callable[[Type[T]], Type[T]][source]#

Registers a subclass with the factory’s registry.

This method returns a decorator that can be applied to a class to register it under a specific key.

Parameters:

key (str or None, optional) – The string key under which to register the subclass. If None, the lowercase name of the subclass itself will be used as the key.

Returns:

A decorator function that takes a class and registers it, returning the class unchanged.

Return type:

Callable[[Type[T]], Type[T]]

Raises:

ValueError – If the provided key (or the default class name) is already registered in the factory’s registry.

Example

Register a class named “MyComponent” under the key “mycomp”:

>>> @MyFactory.register("mycomp")
>>> class MyComponent:
>>>     ...

Register a class named “DefaultComponent” using its own name as the key:

>>> @MyFactory.register()
>>> class DefaultComponent:
>>>     ...
static energy(i: int, j: int, state: State, system: System) → Array[source][source]#

Compute linear spring-like interaction potential energy between particle \(i\) and particle \(j\).

Returns zero when \(i = j\).

Parameters:

  • i (int) – Index of the first particle.

  • j (int) – Index of the second particle.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Scalar JAX array representing the potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array