msmjax.calculators

Code for high-level interface.

References

[1] Hardy, D. J.; Wolff, M. A.; Xia, J.; Schulten, K.; Skeel, R. D. Multilevel Summation with B-Spline Interpolation for Pairwise Interactions in Molecular Dynamics Simulations. J. Chem. Phys. 2016, 144 (11), 114112. https://doi.org/10.1063/1.4943868.

[2] Hardy, D. J. Multilevel Summation for the Fast Evaluation of Forces for the Simulation of Biomolecules (PhD thesis), University of Illinois at Urbana-Champaign, 2006.

class msmjax.calculators.CustomJSONEncoder(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, sort_keys=False, indent=None, separators=None, default=None)[source]

Class that extends JSONEncoder to handle different data types.

From Clinamen2: R. Wanzenböck, F. Buchner, P. Kovács, G. K. H. Madsen, and J. Carrete, “Clinamen2: Functional-style evolutionary optimization in Python for atomistic structure searches,” Computer Physics Communications, vol. 297, p. 109065, Apr. 2024, doi: 10.1016/j.cpc.2023.109065.

default(o)[source]

Return a json-izable version of o or delegate on the base class.

class msmjax.calculators.CustomJSONDecoder(*args, **kwargs)[source]

Class that extends the JSONDecoder to handle different data types.

From Clinamen2: R. Wanzenböck, F. Buchner, P. Kovács, G. K. H. Madsen, and J. Carrete, “Clinamen2: Functional-style evolutionary optimization in Python for atomistic structure searches,” Computer Physics Communications, vol. 297, p. 109065, Apr. 2024, doi: 10.1016/j.cpc.2023.109065.

object_hook(o)[source]

Reencode numpy arrays from dictionary.

class msmjax.calculators.MSMParams(p, mu, max_splitting_level, max_grid_level, cutoffs, cell, cell_mode, pbc, dynamic_cell, supercell_diag, use_neighborlist, neighborlist_prefactor, grids_defined_on_unitcube, grid_shapes, grid_spacings, stencil_extents_from_center, convolution_methods, n_dim, version='2025.10.1', info=<factory>)[source]

Class for holding all settings required for MSM with B-spline interpolation for Coulomb potentials.

While it is possible to manually instantiate this class, this requires a lot of parameters, and might leave it in an inconsistent state. The recommended way is to use set_up_msm_params() instead.

Parameters:

  • p (int) – Interpolation order in the convention of the article (Ref. [1]).

  • mu (int) – Adjustable parameter in the quasi-interpolation scheme that controls its accuracy at grid points, see Ref. [1].

  • max_splitting_level (int) – The level of the highest term in the splitting of the Coulomb kernel into partial kernels.

  • max_grid_level (int) – The level of the highest grid included in the evaluation. Can be lower than max_splitting_level.

  • cutoffs (Sequence[float]) – Sequence of cutoff radii of the partial kernels at all levels.

  • cell (ndarray) – Simulation cell. The exact interpretation depends on the dynamic_cell parameter. If dynamic_cell=False, it is the actual fixed cell for which the model is evaluated. If dynamic_cell=True, it is merely a reference cell.

  • cell_mode (Literal['ortho', 'triclinic']) – A string specifying assumptions on the shape of the unit cell. Either the cell is assumed orthorhombic and axis-aligned, in which case only its diagonal is considered, reducing computational cost, or a general triclinic one.

  • pbc (Sequence[bool]) – One boolean per direction signaling periodicity.

  • dynamic_cell (bool) – Whether the cell can dynamically change.

  • supercell_diag (Sequence[int]) – An optional sequence of positive integers, one per direction, indicating that the short-range contribution should be evaluated in a supercell. See msmjax.core.shortrange.make_eval_pair_pot() for more details.

  • use_neighborlist (bool) – Whether to use a neighbor list.

  • neighborlist_prefactor (float) – Set to 0.5 if the neighbor lists you are going to pass contain duplicates, and to 1.0 if they are duplicate-free. Setting to other values is possible, but unlikely to make sense.

  • grids_defined_on_unitcube (bool) – Whether the grid spacings are given on the unit cube or in dimensionful length units.

  • grid_shapes (Sequence[None | tuple[int, ...]]) – Sequence of tuples of integers indicating the shape of grid at each level.

  • grid_spacings (Sequence[None | ndarray]) – Sequence of 1-d arrays (one value per direction) indicating the grid spacings at each level. If grids_defined_on_unitcube = True, they must be given in fractional coordinates.

  • stencil_extents_from_center (Sequence[None | tuple[int, ...]])

  • convolution_methods (Sequence[Optional[Literal['scipy-direct', 'scipy-fft']]]) – Algorithms to use for convolution in the grid potential calculation. One value per grid level (i.e., different algorithms can be used at different levels). See msmjax.utils.general.ConvMeth and msmjax.core.longrange.special_periodic_convolve_scipy() for more details.

  • n_dim (int) – Spatial dimension.

  • version (str) – Package version.

  • info (dict) – Any additional information you want to store for later reference.

p: int
mu: int
max_splitting_level: int
max_grid_level: int
cutoffs: Sequence[float]
cell: ndarray
cell_mode: Literal['ortho', 'triclinic']
pbc: Sequence[bool]
dynamic_cell: bool
supercell_diag: Sequence[int]
use_neighborlist: bool
neighborlist_prefactor: float
grids_defined_on_unitcube: bool
grid_shapes: Sequence[None | tuple[int, ...]]
grid_spacings: Sequence[None | ndarray]
stencil_extents_from_center: Sequence[None | tuple[int, ...]]
convolution_methods: Sequence[Optional[Literal['scipy-direct', 'scipy-fft']]]
n_dim: int
version: str = '2025.10.1'
info: dict
save_json(filename, **kwargs)[source]
Return type:

None

Parameters:

filename (str | Path)

classmethod load_json(filename)[source]
Parameters:

filename (str | Path)

msmjax.calculators.find_stencil_extents_all_levels(cell, level_one_spacings, level_zero_cutoff, n_levels_intermed, include_toplevel, grid_shape_toplevel=None)[source]

Determine required extents of the kernel stencils at all grid levels.

See Ref. [1] for the criteria used.

Parameters:
Return type:

list[None | tuple[int, ...]]

Returns:

List of tuples of integers indicating kernel stencil extents, one for each level, with a None placeholder at level zero, where there is no grid. See msmjax.kernels.make_construct_stencils() for the precise meaning of the returned stencil extents.

msmjax.calculators.set_up_msm_params(cell, level_one_spacings, level_zero_cutoff, p, pbc, cell_mode, dynamic_cell, mu=None, n_particles=None, max_splitting_level=None, supercell_diag=None, use_neighborlist=False, neighborlist_prefactor=None, intermediate_kernel_stencil_extents=None, convolution_methods='scipy-fft')[source]

Set up an MSMParams instance.

This is the recommended way for doing so, with a higher-level interface requiring fewer parameters, and with more error checking, than manually instantiating MSMParams.

The required arguments can be grouped into the main MSM algorithm parameters on the one hand (level_zero_cutoff, level_one_spacings, p), and specification of system geometry (in a wider sense) on the other hand (cell, pbc, cell_mode, dynamic_cell).

Parameters:

  • cell (jax.typing.ArrayLike) – Array representing unit cell, shape (n_dim, n_dim). Rows correspond to different cell vectors, columns correspond to different cartesian coordinates. The cell is required for setup even in non-periodic cases, as well as when the cell can change (dynamic_cell = True), because the cell together with the grid spacings determines the numbers of grid points (which need to be fixed at setup time). In case the cell can change during simulation ( dynamic_cell = True), it is only the ratio between the lengths of the cell vectors and the grid spacings that matters for setting up the number of grid points. For proper automatic determination of intermediate_kernel_stencil_extents in cases with periodicity, it is advised to supply the cell and the spacings in the most compressed/tilted state that the model is intended to be able to handle.

  • level_one_spacings (jax.typing.ArrayLike) – Level-one grid spacings, one value per direction.

  • level_zero_cutoff (float) – Cutoff radius of the lowest-level partial kernel \(k_0(r)\).

  • p (int) – Interpolation order in the convention of the article (Ref. [1]).

  • pbc (Sequence[bool]) – One boolean per direction signaling periodicity.

  • cell_mode (Literal['ortho', 'triclinic']) – A string specifying assumptions on the shape of the unit cell. Either the cell is assumed orthorhombic and axis-aligned, in which case only its diagonal is considered, reducing computational cost, or a general triclinic one.

  • dynamic_cell (bool) – Whether the cell can dynamically change, or the one given at setup time remains fixed.

  • mu (int | None) – Adjustable parameter in the quasi-interpolation scheme that controls its accuracy at grid points. Chosen based on heuristics from Ref. [1] if not given.

  • n_particles (int | None) – Number of particles. Required in the non-periodic case if the number of levels is to be chosen automatically.

  • max_splitting_level (int | None) – The level of the highest term in the splitting of the Coulomb kernel into partial kernels. Note that, in cases with periodicity, this is not the same as the level of the highest grid included in the evaluation. Determined automatically by default. In cases without periodicity only, you can try manually tuning it, which may or may not improve performance.

  • supercell_diag (Optional[Sequence[int]]) – An optional sequence of positive integers, one per direction, indicating that the short-range contribution should be evaluated in a supercell. Does not work in combination with use_neighborlist = True. See msmjax.core.shortrange.make_eval_pair_pot() for more details.

  • use_neighborlist (bool) – Whether to use a neighbor list.

  • neighborlist_prefactor (float | None) – Set to 0.5 if the neighbor lists you are going to pass contain duplicates, and to 1.0 if they are duplicate-free. Setting to other values is possible, but unlikely to make sense.

  • intermediate_kernel_stencil_extents (tuple[int, ...] | None) – See the extents_from_center_intermed parameter of msmjax.kernels.make_construct_stencils().

  • convolution_methods (Union[Literal['scipy-direct', 'scipy-fft'], Sequence[Literal['scipy-direct', 'scipy-fft']]]) – Algorithms to use for convolution in the grid potential calculation. One value per grid level (i.e., different algorithms can be used at different levels). See msmjax.utils.general.ConvMeth and msmjax.core.longrange.special_periodic_convolve_scipy() for more details.

Returns:

An MSMParams instance.

msmjax.calculators.create_msm(params, part='total', use_custom_derivatives_for_longrange=True, extra_uncharged_interaction=None)[source]

From an MSMParams instance, create electrostatics evaluation functions.

Parameters:

  • params (MSMParams) – The MSMParams instance from which to set up the evaluation functions.

  • part (Literal['total', 'shortrange', 'longrange']) – Whether to calculate the total energy, or just the short-range (\(U^0\)) or long-range (\(U^{1+}\)) contribution.

  • use_custom_derivatives_for_longrange (bool) – Whether to use the custom optimized JVP for the long-range part.

  • extra_uncharged_interaction (Optional[Callable[[jax.typing.ArrayLike], Array]]) – See msmjax.core.shortrange.make_eval_pair_pot().

Return type:

dict[str, Callable]

Returns:

A dictionary of functions that evaluate various different electrostatic quantities. Their exact argument structure depends on the settings in params. In all cases, they take positions and charges as their first two arguments. If params specifies a dynamic cell, they have an additional cell argument, and if params specifies that a neighbor list is to be used, they have an additional neighbor list argument.

msmjax.calculators.check_cutoffs_and_get_actual_spacings(cell, params)[source]

For a given cell, check if MSM cutoff and kernel stencil size settings are appropriate, and calculate the actual grid spacings.

If ok, the spacings are returned. Otherwise, an error is intentionally raised to signal that an inappropriate settings-cell combination was encountered.

Parameters:

  • cell (jax.typing.ArrayLike) – The cell for which to check the cutoff and kernel stencil sizes.

  • params (MSMParams) – The parameter instance specifying the cutoff and the kernel stencil sizes.

Return type:

ndarray

Returns:

Array of actual grid spacings for the given cell, one value per direction.

Raises:

  • ValueError – If the cutoff is too small for the minimum-image convention to work in the given cell. Only raised if there is periodicity.

  • ValueError – If any of the kernel stencils specified by params is too small.