Note
Go to the end to download the full example code.
Batch Sharding Across GPUs#
This guide creates periodic sphere systems, then splits a batch across devices.
Create One Setup#
We first create a function to create the system we want to study:
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
import jaxdem as jdem
box_length = (64 * (4 * jnp.pi / 3) / 0.5) ** (1 / 3)
def create_setup(key):
state = jdem.utils.grid_state(n_per_axis=(4, 4, 4), spacing=3.0, key=key)
state.pos_c %= box_length
state.vel = 0.5 * jdem.utils.unit(state.vel)
system = jdem.System.create(
state=state,
key=key,
dt=1e-2,
domain_type="periodic",
domain_kw={"box_size": jnp.full(3, box_length)},
collider_type="neighbor_list",
collider_kw={
"cutoff": 2.0,
"skin": 0.2,
"max_neighbors": 32,
"secondary_collider_type": "cell_list",
"secondary_collider_kw": {"cell_size": 2.2, "search_range": 1},
},
)
return state, system
state, system = create_setup(jax.random.key(0))
print("One setup:", state.pos.shape)
One setup: (64, 3)
Create a Batch with vmap#
vmap calls the setup function for each random key. All systems have the
same array shapes and static configuration. Each system has its own
velocities and dynamic data:
batch_size = 16
keys = jax.random.split(jax.random.key(42), batch_size)
batch_state, batch_system = jax.vmap(create_setup)(keys)
print("Batch positions:", batch_state.pos.shape)
Batch positions: (16, 64, 3)
Create the Batch Directly on the Devices#
For a larger setup, give jit the output sharding. This replaces batch
creation followed by transfer. Only the small array of keys needs placement
before the call. JAX creates the output arrays with the requested sharding.
create_batch = jax.jit(jax.vmap(create_setup), out_shardings=sharding)
batch_state, batch_system = create_batch(jax.device_put(keys, sharding))
jax.block_until_ready((batch_state, batch_system))
(State(pos_c=Array([[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
...,
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]]], dtype=float64), pos_p=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), vel=Array([[[ 0.15753041, 0.07031648, 0.4692971 ],
[-0.32753323, 0.14538199, 0.34869193],
[-0.39629476, 0.09985669, -0.28806094],
...,
[-0.31460058, 0.31412474, -0.22880586],
[ 0.38054994, -0.11393869, 0.30364407],
[ 0.286559 , 0.14705313, -0.38243865]],
[[ 0.34285425, -0.17136487, 0.3210686 ],
[ 0.05863437, 0.36868646, -0.33261435],
[ 0.25925375, 0.39720178, 0.15817153],
...,
[-0.27531116, -0.15378584, -0.38801247],
[ 0.46023066, 0.07653712, 0.17980492],
[ 0.32328841, 0.19027425, 0.33057573]],
[[ 0.39362887, 0.29042362, 0.10349122],
[-0.35437242, 0.00738209, 0.35265521],
[-0.33823182, -0.32720676, -0.16892297],
...,
[-0.39911524, -0.26077192, -0.15068189],
[-0.27568693, -0.33960797, 0.24220475],
[ 0.29387099, 0.32506734, 0.24077182]],
...,
[[-0.18185362, -0.29092591, 0.36371882],
[ 0.10193624, 0.36413808, 0.32712759],
[-0.36241623, 0.23197993, -0.25463658],
...,
[-0.40625891, 0.02559584, -0.29034213],
[ 0.18584723, -0.28800981, -0.36402081],
[ 0.30860499, 0.03356622, -0.39196463]],
[[ 0.03868465, 0.14877224, -0.4757839 ],
[ 0.34479777, -0.35126312, 0.08791313],
[-0.04986618, -0.47207151, -0.15704094],
...,
[-0.05254389, -0.00220851, -0.49722657],
[-0.29890101, 0.33263489, 0.22363411],
[ 0.38267832, 0.1077892 , -0.30321411]],
[[ 0.1294673 , -0.23336296, 0.42282378],
[-0.2702412 , 0.04145824, -0.4186298 ],
[-0.36310967, 0.06745739, -0.33704728],
...,
[-0.45162629, -0.04229712, 0.21034411],
[-0.40266001, -0.29547025, -0.02371165],
[-0.31867352, -0.22004124, -0.31627368]]], dtype=float64), force=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), q=Quaternion(w=Array([[[1.],
[1.],
[1.],
...,
[1.],
[1.],
[1.]],
[[1.],
[1.],
[1.],
...,
[1.],
[1.],
[1.]],
[[1.],
[1.],
[1.],
...,
[1.],
[1.],
[1.]],
...,
[[1.],
[1.],
[1.],
...,
[1.],
[1.],
[1.]],
[[1.],
[1.],
[1.],
...,
[1.],
[1.],
[1.]],
[[1.],
[1.],
[1.],
...,
[1.],
[1.],
[1.]]], dtype=float64), xyz=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64)), ang_vel=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), torque=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), rad=Array([[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]], dtype=float64), _rad=Array([[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]], dtype=float64), volume=Array([[4.1887902, 4.1887902, 4.1887902, ..., 4.1887902, 4.1887902,
4.1887902],
[4.1887902, 4.1887902, 4.1887902, ..., 4.1887902, 4.1887902,
4.1887902],
[4.1887902, 4.1887902, 4.1887902, ..., 4.1887902, 4.1887902,
4.1887902],
...,
[4.1887902, 4.1887902, 4.1887902, ..., 4.1887902, 4.1887902,
4.1887902],
[4.1887902, 4.1887902, 4.1887902, ..., 4.1887902, 4.1887902,
4.1887902],
[4.1887902, 4.1887902, 4.1887902, ..., 4.1887902, 4.1887902,
4.1887902]], dtype=float64), mass=Array([[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]], dtype=float64), inertia=Array([[[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
...,
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4]],
[[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
...,
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4]],
[[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
...,
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4]],
...,
[[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
...,
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4]],
[[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
...,
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4]],
[[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
...,
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4],
[0.4, 0.4, 0.4]]], dtype=float64), clump_id=Array([[ 0, 1, 2, ..., 61, 62, 63],
[ 0, 1, 2, ..., 61, 62, 63],
[ 0, 1, 2, ..., 61, 62, 63],
...,
[ 0, 1, 2, ..., 61, 62, 63],
[ 0, 1, 2, ..., 61, 62, 63],
[ 0, 1, 2, ..., 61, 62, 63]], dtype=int64), bond_id=Array([[[-1],
[-1],
[-1],
...,
[-1],
[-1],
[-1]],
[[-1],
[-1],
[-1],
...,
[-1],
[-1],
[-1]],
[[-1],
[-1],
[-1],
...,
[-1],
[-1],
[-1]],
...,
[[-1],
[-1],
[-1],
...,
[-1],
[-1],
[-1]],
[[-1],
[-1],
[-1],
...,
[-1],
[-1],
[-1]],
[[-1],
[-1],
[-1],
...,
[-1],
[-1],
[-1]]], dtype=int64), mat_id=Array([[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]], dtype=int64), species_id=Array([[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]], dtype=int64), fixed=Array([[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
...,
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False]], dtype=bool), facet_id=Array([[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1],
...,
[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1]], dtype=int64), facet_vertices=Array([[[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1],
...,
[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1]],
[[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1],
...,
[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1]],
[[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1],
...,
[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1]],
...,
[[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1],
...,
[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1]],
[[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1],
...,
[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1]],
[[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1],
...,
[-1, -1, -1],
[-1, -1, -1],
[-1, -1, -1]]], dtype=int64), facet_ids_extra=Array([], shape=(16, 64, 0), dtype=int64), facet_vertices_extra=Array([], shape=(16, 64, 0, 3), dtype=int64), _pos_p_rot=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64)), System(linear_integrator=VelocityVerlet(), rotation_integrator=VelocityVerletSpiral(), collider=NeighborList(overflow=Array([False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False], dtype=bool), secondary_collider=DynamicCellList(overflow=Array([False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False], dtype=bool), neighbor_mask=Array([[[-1, -1, -1],
[-1, -1, 0],
[-1, -1, 1],
...,
[ 1, 1, -1],
[ 1, 1, 0],
[ 1, 1, 1]],
[[-1, -1, -1],
[-1, -1, 0],
[-1, -1, 1],
...,
[ 1, 1, -1],
[ 1, 1, 0],
[ 1, 1, 1]],
[[-1, -1, -1],
[-1, -1, 0],
[-1, -1, 1],
...,
[ 1, 1, -1],
[ 1, 1, 0],
[ 1, 1, 1]],
...,
[[-1, -1, -1],
[-1, -1, 0],
[-1, -1, 1],
...,
[ 1, 1, -1],
[ 1, 1, 0],
[ 1, 1, 1]],
[[-1, -1, -1],
[-1, -1, 0],
[-1, -1, 1],
...,
[ 1, 1, -1],
[ 1, 1, 0],
[ 1, 1, 1]],
[[-1, -1, -1],
[-1, -1, 0],
[-1, -1, 1],
...,
[ 1, 1, -1],
[ 1, 1, 0],
[ 1, 1, 1]]], dtype=int64), cell_size=Array([2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2,
2.2, 2.2, 2.2], dtype=float64)), neighbor_list=Array([[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1],
...,
[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1],
[-1, -1, -1, ..., -1, -1, -1]], dtype=int64), old_pos=Array([[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
...,
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]],
[[0. , 0. , 0. ],
[0. , 0. , 3. ],
[0. , 0. , 6. ],
...,
[0.87606962, 0.87606962, 3. ],
[0.87606962, 0.87606962, 6. ],
[0.87606962, 0.87606962, 0.87606962]]], dtype=float64), n_build_times=Array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int64), cutoff=Array([2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2.], dtype=float64), skin=Array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2,
0.2, 0.2, 0.2], dtype=float64), max_neighbors=32, row_offsets=Array([[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
...,
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0],
[0, 0, 0, ..., 0, 0, 0]], dtype=int64), history=Array([], shape=(16, 2048, 0), dtype=float64), metric_snapshot=Array([[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.]], dtype=float64), physical_cutoff_snapshot=Array([-1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1.,
-1., -1., -1.], dtype=float64), skin_snapshot=Array([-1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1.,
-1., -1., -1.], dtype=float64), facet_id_snapshot=Array([-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2], dtype=int64), invalidated=Array([ True, True, True, True, True, True, True, True, True,
True, True, True, True, True, True, True], dtype=bool)), domain=PeriodicDomain(box_size=Array([[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038],
[8.12393038, 8.12393038, 8.12393038]], dtype=float64), inv_box_size=Array([[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313],
[0.12309313, 0.12309313, 0.12309313]], dtype=float64), anchor=Array([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]], dtype=float64), box_rate=Array([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]], dtype=float64)), force_manager=ForceManager(gravity=Array([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]], dtype=float64), external_force=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), external_force_com=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), external_torque=Array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
...,
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.],
...,
[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float64), is_com_force=(), force_functions=(), energy_functions=()), bonded_force_model=None, force_model=SpringForce(laws=()), mat_table=MaterialTable(props={'density': Array([[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27],
[0.27]], dtype=float64), 'poisson': Array([[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3],
[0.3]], dtype=float64), 'young': Array([[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.],
[10000.]], dtype=float64)}, pair={'density_eff': Array([[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]],
[[0.27]]], dtype=float64), 'poisson_eff': Array([[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]],
[[0.3]]], dtype=float64), 'young_eff': Array([[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]],
[[10000.]]], dtype=float64)}, matcher=HarmonicMaterialMatchmaker()), dt=Array([0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01,
0.01, 0.01, 0.01, 0.01, 0.01], dtype=float64), time=Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64), dim=Array([3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3], dtype=int64), step_count=Array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int64), key=Array((16,), dtype=key<fry>) overlaying:
[[1832780943 270669613]
[ 64467757 2916123636]
[2465931498 255383827]
[3134548294 894150801]
[2954079971 3276725750]
[2765691542 824333390]
[2768684296 3055579793]
[2547012911 1371500959]
[1016697191 2390192106]
[1128875147 2463678267]
[1039196627 1683848162]
[ 246739928 3519402408]
[3114009986 1419417030]
[3514951389 229662949]
[2526883203 3973959769]
[ 991576401 3935454969]], interact_same_bond_id=Array([False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False], dtype=bool), search_overflow=Array([False, False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False], dtype=bool), user_pre_step_actions=<PjitFunction of <function _save_state_system at 0x7e66757b9bc0>>, user_post_step_actions=<PjitFunction of <function _save_state_system at 0x7e66757b9bc0>>, minimizer=CustomGradientTransformation(init=<function fire.<locals>.init at 0x7e660c92ca90>, update=<function CustomGradientTransformation.__new__.<locals>.update at 0x7e660c92cb40>), target_fn=None))
Take One Step#
The existing System methods handle the batch axis and keep positions and velocities on their shards:
batch_state, batch_system = jdem.System.initialize(batch_state, batch_system)
batch_state, batch_system = jdem.System.step(batch_state, batch_system, n=1)
jax.block_until_ready((batch_state, batch_system))
print("Step counts:", batch_system.step_count)
ParaView rendering of this run’s VTKWriter output.
Step counts: [1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]
Total running time of the script: (0 minutes 13.353 seconds)