Setting up and running a calculation

import numpy as onp
from msmjax.calculators import set_up_msm_params

msm_params = set_up_msm_params(
    cell=onp.diag([10., 10., 10.]),
    level_one_spacings=onp.array([1., 1., 1.]),
    level_zero_cutoff=3.0,
    p=4,
    pbc=(False, False, False),
    cell_mode="ortho",
    dynamic_cell=False,
)

from msmjax.calculators import create_msm

evaluation_fns = create_msm(msm_params)
# jit for performance
eval_energy_and_forces = jax.jit(
    evaluation_fns["energy_and_forces"]
)
energy, forces = eval_energy_and_forces(
    positions, charges
)

from msmjax.calculators import MSMParams

msm_params.save_json("params.json")
msm_params_restored = MSMParams.load_json("params.json")