Example: temperature and density control via simple rescaling.

This script builds a small periodic system of spheres, initialized at a low temperature, and then runs dynamics while imposing a time-dependent packing fraction using jd.utils.control_nvt_density_rollout. This function can modulate both the temperature and density of the system using very simple rescaling. It comes in two variants, one with _rollout and one without. The _rollout variant provides access to the intermediate trajectory data, whereas the other does not. Although it may be convenient to get the trajectory data, the non _rollout option is higher performance.

In either, we can set a target temperature and/or density and the control_nvt_density function will use a linear scheduler to reach the desired set point after a set number of steps.

You can also control either variable using a custom scheduler. Here, we show an example of that. We modulate the density according to a single sine wave.

In this example we: - Modulate density (packing fraction) with a prescribed schedule. - Do not thermostat (temperature is allowed to fluctuate in response to compression).

Imports

import jax
import jax.numpy as jnp
import jaxdem as jd
from functools import partial
from jaxdem.utils.randomSphereConfiguration import random_sphere_configuration

# We need to enable double precision for accuracy in cold systems
jax.config.update("jax_enable_x64", True)

Parameters

Background: We build a single system in the same spirit as examples/jam_spheres.py: - choose a polydisperse radius distribution - generate a random, overlap-free configuration at the desired packing fraction phi - create a periodic System with a simple spring contact model

N = 100
phi = 0.7
dim = 2
e_int = 1.0
dt = 1e-2

can_rotate = False  # smooth spheres cannot rotate
subtract_drift = True
seed = 0


def build_microstate(i):
    mass = 1.0
    e_int = 1.0
    dt = 1e-2
    if dim == 2:
        cr = [0.5, 0.5]
        sr = [1.0, 1.4]
    else:
        cr = [1.0]
        sr = [1.0]
    particle_radii = jd.utils.dispersity.get_polydisperse_radii(N, cr, sr)
    pos, box_size = random_sphere_configuration(particle_radii.tolist(), phi, dim, seed)
    state = jd.State.create(pos=pos, rad=particle_radii, mass=jnp.ones(N) * mass)
    mats = [jd.Material.create("elastic", young=e_int, poisson=0.5, density=1.0)]
    matcher = jd.MaterialMatchmaker.create("harmonic")
    mat_table = jd.MaterialTable.from_materials(mats, matcher=matcher)
    system = jd.System.create(
        state_shape=state.shape,
        dt=dt,
        linear_integrator_type="verlet",
        rotation_integrator_type="",
        domain_type="periodic",
        force_model_type="spring",
        collider_type="naive",
        mat_table=mat_table,
        domain_kw={
            "box_size": box_size,
        },
    )
    return state, system

Run dynamics while imposing a sine-wave compression

We impose a single-period sinusoidal modulation of the packing fraction over the protocol duration (a compression/decompression cycle). We intentionally do not set a temperature schedule here, so temperature is free to fluctuate.

initial_temperature = 1e-4
packing_fraction_amplitude = -0.05
n_steps = 10_000
save_stride = 100
n_snapshots = int(n_steps) // int(save_stride)

state, system = build_microstate(0)
state = jd.utils.thermal.set_temperature(
    state,
    initial_temperature,
    can_rotate=can_rotate,
    subtract_drift=subtract_drift,
    seed=0,
)


def sine_dens_schedule(k, K, start, target):
    x = k / jnp.maximum(K, 1)  # 0..1
    return start + packing_fraction_amplitude * jnp.sin(
        2.0 * jnp.pi * x
    )  # one full period


state, system, (traj_state, traj_system) = jd.utils.control_nvt_density_rollout(
    state,
    system,
    n=n_snapshots,
    stride=save_stride,
    rescale_every=100,
    # NOTE: Density control is enabled when either packing_fraction_target or
    # packing_fraction_delta is provided. Our schedule ignores the "target", so we
    # pass delta=0.0 simply to enable control.
    packing_fraction_delta=0.0,
    density_schedule=sine_dens_schedule,
    can_rotate=can_rotate,
    subtract_drift=subtract_drift,
)

temperature = jax.vmap(
    partial(
        jd.utils.thermal.compute_temperature,
        can_rotate=can_rotate,
        subtract_drift=subtract_drift,
    )
)(traj_state)
phi = jax.vmap(partial(jd.utils.packingUtils.compute_packing_fraction))(
    traj_state, traj_system
)
pe = jax.vmap(partial(jd.utils.thermal.compute_potential_energy))(
    traj_state, traj_system
)

print("The recorded temperature profile is: ", temperature)
print("The recorded density profile is: ", phi)
print("The recorded potential-energy profile is: ", pe)

The recorded temperature profile is:  [9.33094880e-05 8.54299606e-05 7.92524407e-05 8.20803919e-05
 8.03058175e-05 7.47673854e-05 7.28302215e-05 6.88730304e-05
 6.88260067e-05 6.84916193e-05 6.46841546e-05 6.32966682e-05
 6.02580825e-05 6.22173641e-05 6.12192442e-05 6.14830213e-05
 5.78662698e-05 5.91876929e-05 5.79813470e-05 5.87028133e-05
 5.76878835e-05 5.62916500e-05 5.44034181e-05 5.68672405e-05
 5.47027158e-05 5.57097421e-05 5.66834385e-05 5.56166591e-05
 5.43716042e-05 5.63501712e-05 5.73616693e-05 5.80161530e-05
 5.80360636e-05 5.98364010e-05 5.82441273e-05 6.24505974e-05
 6.34642374e-05 6.46894773e-05 6.89427242e-05 6.61038315e-05
 7.68571387e-05 8.04094459e-05 8.25110437e-05 8.42384683e-05
 9.65663115e-05 9.76201259e-05 9.65047239e-05 1.07747395e-04
 1.06847420e-04 1.24332693e-04 1.20392706e-04 1.06917759e-04
 1.32589163e-04 1.44302912e-04 1.58059136e-04 1.53538712e-04
 1.64589427e-04 1.68039739e-04 1.73181837e-04 1.82230499e-04
 1.91113632e-04 1.79369924e-04 2.01646374e-04 1.93249082e-04
 1.95918536e-04 2.11428481e-04 2.11015468e-04 2.25335581e-04
 2.34052620e-04 2.28650407e-04 2.33119047e-04 2.50956933e-04
 2.40704451e-04 2.48065738e-04 2.38804493e-04 2.58868263e-04
 2.31551259e-04 2.51513230e-04 2.55727586e-04 2.46193591e-04
 2.39233619e-04 2.50534961e-04 2.48516603e-04 2.30711919e-04
 2.11763392e-04 2.20826807e-04 2.25878088e-04 2.15812260e-04
 2.10747852e-04 2.05534982e-04 1.89852405e-04 1.93547179e-04
 1.81418717e-04 1.82200193e-04 1.79665174e-04 1.68763432e-04
 1.68375303e-04 1.65295054e-04 1.47087529e-04 1.44035309e-04]
The recorded density profile is:  [0.69686047 0.69373334 0.69063093 0.68756551 0.68454915 0.68159377
 0.67871104 0.67591232 0.67320866 0.67061074 0.6681288  0.66577264
 0.66355157 0.66147434 0.65954915 0.6577836  0.65618467 0.65475865
 0.65351118 0.65244717 0.65157084 0.65088564 0.65039426 0.65009866
 0.65       0.65009866 0.65039426 0.65088564 0.65157084 0.65244717
 0.65351118 0.65475865 0.65618467 0.6577836  0.65954915 0.66147434
 0.66355157 0.66577264 0.6681288  0.67061074 0.67320866 0.67591232
 0.67871104 0.68159377 0.68454915 0.68756551 0.69063093 0.69373334
 0.69686047 0.7        0.70313953 0.70626666 0.70936907 0.71243449
 0.71545085 0.71840623 0.72128896 0.72408768 0.72679134 0.72938926
 0.7318712  0.73422736 0.73644843 0.73852566 0.74045085 0.7422164
 0.74381533 0.74524135 0.74648882 0.74755283 0.74842916 0.74911436
 0.74960574 0.74990134 0.75       0.74990134 0.74960574 0.74911436
 0.74842916 0.74755283 0.74648882 0.74524135 0.74381533 0.7422164
 0.74045085 0.73852566 0.73644843 0.73422736 0.7318712  0.72938926
 0.72679134 0.72408768 0.72128896 0.71840623 0.71545085 0.71243449
 0.70936907 0.70626666 0.70313953 0.7       ]
The recorded potential-energy profile is:  [3.88969557e-04 6.70618384e-04 7.73829276e-04 2.78877300e-04
 2.04135159e-04 4.69208843e-04 4.22557970e-04 5.32432893e-04
 2.85927967e-04 1.56704030e-04 3.17987822e-04 2.84772328e-04
 4.53435619e-04 1.57670761e-04 1.80266617e-04 8.05369632e-05
 3.18523630e-04 1.24708923e-04 1.80375335e-04 7.99434744e-05
 1.44568309e-04 2.40849454e-04 3.84671129e-04 1.27968354e-04
 3.35431622e-04 2.41481662e-04 1.55949524e-04 2.96404350e-04
 4.89880427e-04 3.55546019e-04 3.29331659e-04 3.63363003e-04
 4.80349272e-04 4.39836152e-04 8.28276173e-04 6.43951856e-04
 8.18779692e-04 1.02876840e-03 9.33525598e-04 1.71130993e-03
 1.06247530e-03 1.13495653e-03 1.49037768e-03 1.91289714e-03
 1.08687248e-03 1.39110820e-03 2.10523769e-03 1.53595136e-03
 2.38418070e-03 1.10687141e-03 2.13864711e-03 4.49425970e-03
 2.94673838e-03 2.57624700e-03 1.78849918e-03 2.94655913e-03
 2.60440874e-03 3.01231656e-03 3.35075680e-03 3.15564129e-03
 2.88303535e-03 4.97897966e-03 3.57014286e-03 5.29044353e-03
 5.97484450e-03 5.19426529e-03 5.97341725e-03 5.26937424e-03
 4.88784626e-03 5.91549957e-03 5.87366757e-03 4.38652813e-03
 5.62865790e-03 5.03027252e-03 5.99305802e-03 3.96828265e-03
 6.52836574e-03 4.34066323e-03 3.67954250e-03 4.27079426e-03
 4.57087994e-03 3.07606987e-03 2.87676407e-03 4.05687824e-03
 5.14636436e-03 3.56317729e-03 2.52290993e-03 2.87423814e-03
 2.64808066e-03 2.48493170e-03 3.26968822e-03 2.15473683e-03
 2.54610752e-03 1.83347851e-03 1.53444990e-03 2.02834742e-03
 1.54556216e-03 1.23564536e-03 2.28640783e-03 1.92976099e-03]