jax_md_mod.custom_quantity.init_pressure

Contents

jax_md_mod.custom_quantity.init_pressure#

init_pressure(energy_fn_template, reference_box=None, include_kinetic=True)[source]#

Initializes a function that computes the pressure for a single state.

This function is applicable to arbitrary many-body interactions, even under periodic boundary conditions. See init_virial_stress_tensor for details.

Parameters:
  • energy_fn_template – A function that takes energy parameters as input and returns an energy function

  • ref_box_tensor – The transformation T of general periodic boundary conditions. If None, box_tensor needs to be provided as ‘box’ during function call, e.g. for NPT ensemble.

  • include_kinetic – Whether kinetic part of stress tensor should be added.

Returns:

A function that takes a simulation state with neighbor list, energy_params and box (if applicable) and returns the instantaneous pressure.