Force Models.#

A ForceModel defines the pairwise interaction law between two particles. It is evaluated by the collider for every interacting pair and returns a force vector and torque.

This guide covers:

  • The available (registered) force models.

  • How materials supply the parameters each model needs.

  • Combining several laws with LawCombiner.

  • Species-wise interactions with ForceRouter.

Selecting a Force Model#

The force model is chosen via the force_model_type string when creating a System:

import jax.numpy as jnp
import jaxdem as jdem

state = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0]]),
    rad=jnp.array([1.0, 1.0]),
)
system = jdem.System.create(state.shape, force_model_type="spring")
print("Force model:", type(system.force_model).__name__)
Force model: SpringForce

Available Force Models#

JaxDEM ships with several pairwise force models registered in the ForceModel factory. You select one via force_model_type when calling create().

system_spring = jdem.System.create(state.shape, force_model_type="spring")
print("spring →", type(system_spring.force_model).__name__)

# WCA/LJ models require an ``epsilon_eff`` material table.
lj_mat = jdem.MaterialTable.from_materials(
    [jdem.Material.create("lj", density=1.0, epsilon=1.0)]
)
system_wca = jdem.System.create(state.shape, force_model_type="wca", mat_table=lj_mat)
print("wca    →", type(system_wca.force_model).__name__)

system_lj = jdem.System.create(
    state.shape, force_model_type="lennardjones", mat_table=lj_mat
)
print("lj     →", type(system_lj.force_model).__name__)
spring → SpringForce
wca    → WCA
lj     → LennardJones

The following ForceModels are registered:

print("ForceModels:", list(jdem.ForceModel._registry.keys()))
ForceModels: ['lawcombiner', 'forcerouter', 'spring', 'wca', 'lennardjones', 'wca_shifted', 'hertz', 'cundallstrack']

Materials and Force Requirements#

Each force model declares the material-pair properties it needs. For example, "spring" requires young_eff while "wca" requires epsilon_eff. These effective properties are automatically computed from per-material scalars by a MaterialMatchmaker.

When you pass mat_table=None (the default), create() builds a single-material table with sensible defaults. For custom materials, build the table yourself:

mat_a = jdem.Material.create("lj", density=1.0, epsilon=1.0)
mat_b = jdem.Material.create("lj", density=1.0, epsilon=2.0)
mat_table = jdem.MaterialTable.from_materials([mat_a, mat_b])

print("epsilon per material:", mat_table.epsilon)
print("epsilon_eff (pair table):\n", mat_table.epsilon_eff)
epsilon per material: [1. 2.]
epsilon_eff (pair table):
 [[1.         1.33333333]
 [1.33333333 2.        ]]

Pass the table to create(), and assign per-particle material IDs via mat_id in the state:

state_2mat = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0]]),
    rad=jnp.array([1.0, 1.0]),
    mat_id=jnp.array([0, 1]),  # particle 0 → mat_a, particle 1 → mat_b
)
system_2mat = jdem.System.create(
    state_2mat.shape,
    force_model_type="wca",
    mat_table=mat_table,
)

state_2mat, system_2mat = system_2mat.step(state_2mat, system_2mat)
print("Force model:", type(system_2mat.force_model).__name__)
Force model: WCA

Combining Force Models#

LawCombiner sums several elementary force laws into one composite model. You pass it a tuple of ForceModel instances via the laws field. Both forces and energies are added together.

combined = jdem.LawCombiner(
    laws=(jdem.ForceModel.create("spring"), jdem.ForceModel.create("wca"))
)
print("Combined laws:", [type(l).__name__ for l in combined.laws])
Combined laws: ['SpringForce', 'WCA']

To use a combined model, pass it directly to create() via the force_model_kw argument. Make sure the material table provides all the properties required by the child laws:

mat = jdem.Material.create("elastic", density=1.0, young=1e4, poisson=0.3)
mat_lj = jdem.Material.create("lj", density=1.0, epsilon=1.0)

# Build a table that has both young_eff and epsilon_eff.
mat_table_both = jdem.MaterialTable.from_materials([mat, mat_lj])
print("Table has young_eff:", hasattr(mat_table_both, "young_eff"))
print("Table has epsilon_eff:", hasattr(mat_table_both, "epsilon_eff"))

state_combined = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0]]),
    rad=jnp.array([1.0, 1.0]),
    mat_id=jnp.array([0, 1]),  # particle 0 → mat_a, particle 1 → mat_b
)
system_combined = jdem.System.create(
    state.shape,
    force_model_type="lawcombiner",
    force_model_kw={
        "laws": (jdem.ForceModel.create("spring"), jdem.ForceModel.create("wca"))
    },
    mat_table=mat_table_both,
)

state_combined, system_combined = system_combined.step(state_combined, system_combined)
print("System force model:", type(system_combined.force_model).__name__)
Table has young_eff: True
Table has epsilon_eff: True
System force model: LawCombiner

Species-Wise Interactions#

In many simulations different particle species should interact under different force laws (e.g. soft repulsion between species A-B but a stiff spring between A-A). The ForceRouter provides this via a species-pair lookup table.

Use from_dict() to build the table. The keys are (species_i, species_j) tuples and the values are ForceModel instances. Pairs not specified default to zero interaction.

# Two species: 0 and 1
router = jdem.ForceRouter.from_dict(
    S=2,
    mapping={
        (0, 0): jdem.ForceModel.create("spring"),  # A-A: spring
        (1, 1): jdem.ForceModel.create("wca"),  # B-B: WCA
        (0, 1): jdem.LawCombiner(
            laws=(jdem.ForceModel.create("spring"), jdem.ForceModel.create("wca"))
        ),  # A-B: both
    },
)
print("Router table shape:", len(router.table), "x", len(router.table[0]))
Router table shape: 2 x 2

Assign species IDs via species_id in the state, and pass the router table to create() using the force_model_type / force_model_kw pattern:

state_species = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0], [3.0, 0.0]]),
    rad=jnp.array([1.0, 1.0, 1.0]),
    species_id=jnp.array([0, 0, 1]),  # first two are species 0, last is species 1
)

system_species = jdem.System.create(
    state_species.shape,
    force_model_type="forcerouter",
    force_model_kw={"table": router.table},
    mat_table=jdem.MaterialTable.from_materials([mat, mat_lj]),
)

state_species, system_species = system_species.step(state_species, system_species)
print("Active force model:", type(system_species.force_model).__name__)
Active force model: ForceRouter

Note

species_id selects the force law, while mat_id selects the material parameters (stiffness, epsilon, etc.). They are independent — you can have species 0 and 1 share the same material but use different force laws, or vice-versa.

Total running time of the script: (0 minutes 1.843 seconds)