jax-md

import jax import jax.numpy as jnp from jax_md import simulate, energy, space, quantity # Set up a simple cubic lattice of particles dimension = 2 box_size = 5.0 displacement_fn, shift_fn = space.periodic(box_size) # Create positions on a lattice N = 9 lattice = space.initialize_canonical_lattice(N, box_size=box_size, dim=dimension, center=[0.0, 0.0]) positions = lattice['position'] # Define a soft sphere interaction energy_fn = energy.soft_sphere(displacement_fn, sigma=1.0, epsilon=1.0) # Initialize neighbor list neighbor_fn = partition.neighbor_list(displacement_fn, box_size, r_cutoff=2.5, capacity_multiplier=1.2) neighbor_list = neighbor_fn.allocate(positions) # Wrap energy function with neighbor list def total_energy(R, **kwargs): return energy_fn(R, neighbor_list.idx) # Run simulation with NVE init, apply = simulate.nve(energy_fn, shift_fn, dt=0.001, T0=1.0) state = init(jax.random.PRNGKey(0), positions, neighbor_list=neighbor_list) # Simulate for 100 steps for i in range(100): state = apply(state, neighbor_list=neighbor_list) if i % 10 == 0: print(f'Step {i}, KE={quantity.kinetic_energy(state.velocity, state.mass):.3f}') print('Quickstart complete.')