jax_md.mm_forcefields.reaxff package

Submodules

jax_md.mm_forcefields.reaxff.reaxff_energy module

Contains energy related functions for ReaxFF

Author: Mehmet Cagri Kaymak

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_acks2_charges(species, atom_mask, nbr_inds, hulp2_mat, bond_softness, tapered_dists, idempotential, electronegativity, total_charge, backprop_solve=False, tol=1e-06)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_acks2_coulomb_pot(nbr_inds, atom_mask, effpot, bond_softness)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_bo(nbr_inds, nbr_dist, species, species_AN, force_field)

Usage:

first update/allocate neighborlist will be called then the info will be passed to this function

for now, assume the format is “Dense”

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_boncor_pot(nbr_inds, nbr_mask, species, bo, bopi, bopi2, abo, force_field)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_charge_energy(species, charges, idempotential, electronegativity)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_charge_energy_acks2(charges, effpot)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_coulomb_pot(nbr_inds, atom_mask, hulp2_mat, tapered_dists, charges)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_covbon_pot(nbr_inds, nbr_dist, nbr_mask, species, triple_bond, force_field, tapered_reaxff=False)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_eem_charges(species, atom_mask, nbr_inds, hulp2_mat, tapered_dists, idempotential, electronegativity, init_charges=None, total_charge=0.0, backprop_solve=False, tol=1e-06, max_solver_iter=500)

EEM charge solver.

If max_solver_iter is set to -1, use direct solve.

Returns:

An array of shape [n+1,] where first n entries are the charges and last entry is the electronegativity equalization value.

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_hb_pot(species, hbond_inds, hbond_angles, hbond_mask, close_nbr_inds, far_nbr_inds, bo, force_field, tapered_reaxff=False)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_lonpar_pot(species, atom_mask, abo, force_field)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_ovcor_pot(species, atoms_AN, atom_mask, nbr_inds, nbr_dists, nbr_mask, bo, bopi, bopi2, abo, vlp, force_field)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_reaxff_energy(species, atomic_numbers, nbr_lists, close_nbr_dists, far_nbr_dists, body_3_angles, body_4_angles, hb_ang_dist, force_field, init_charges=None, total_charge=0.0, tol=1e-06, max_solver_iter=500, backprop_solve=False, tors_2013=False, tapered_reaxff=False, solver_model='EEM')

Calculate full ReaxFF potential.

Parameters:

• species (Array) – An ndarray of shape [n, ] for the atom types.

• atomic_numbers (Array) – An ndarray of shape [n, ] for the atomic numbers of the atoms.

• nbr_lists (ReaxFFNeighborLists) – Contains the interaction lists for ReaxFF potential.

• close_nbr_dists (Array) – An ndarray of shape [n,m] for bonded interaction distances

• far_nbr_dists (Array) – An ndarray of shape [n,m] for non-bonded interaction distances (vdw and Coulomb)

• body_3_angles (Array) – Angles and distances for many-body interactions

• body_4_angles (Array) – Angles and distances for many-body interactions

• hb_ang_dist (Array) – Angles and distances for many-body interactions

• force_field (ForceField) – ReaxFF parameters

• init_charges (Array | None) – Initial charges for the iterative solver An ndarray of shape [n, ] or None

• total_charge (float) – Total charge of the system (float)

• tol (float) – Tolarence for the charge solver

• max_solver_iter (int) – Maximum number of solver iterations If set to -1, use direct solve

• backprop_solve (bool) – Control variable to decide whether to do a solve to calculate the gradients of the charges wrt positions. By definition, the gradients should be 0 but if the solver tolerance is high, the gradients might be non-ignorable.

• tors_2013 (bool) – Control variable to decide whether to use more stable version of the torsion interactions

• tapered_reaxff (bool) – Control variable to decide whether to use tapered cutoffs for various bonded interactions, causes computational overhead

• solver_model (str) – Control variable for the solver model (“EEM” or “ACKS”)

Returns:

System energy in kcal/mol

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_torsion_pot(species, body_4_inds, body_4_angles, body_4_mask, nbr_inds, bo, bopi, abo, force_field, tapered_reaxff=False, tors_2013=False)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_valency_pot(species, body_3_inds, body_3_angles, body_3_mask, nbr_inds, vlp, bo, bopi, bopi2, abo, force_field, tapered_reaxff=False)

jax_md.mm_forcefields.reaxff.reaxff_energy.calculate_vdw_pot(species, far_nbr_mask, nbr_inds, dists, tapered_dists, force_field)

jax_md.mm_forcefields.reaxff.reaxff_energy.smooth_lone_pair_casting(number, p_lambda=0.9999, l1=-1.3, l2=-0.3, r1=0.3, r2=1.3)

jax_md.mm_forcefields.reaxff.reaxff_energy.taper(value, low_tap_rad, up_tap_rad)

Decreasing tapering function 1 at low_tap_rad and 0 at up_tap_rad smoothly taper the value in between

jax_md.mm_forcefields.reaxff.reaxff_energy.taper_inc(dist, low_tap_rad=0, up_tap_rad=10)

Increasing tapering function 0 at low_tap_rad and 1 at up_tap_rad smoothly taper the value in between

jax_md.mm_forcefields.reaxff.reaxff_forcefield module

Contains force field related code

Author: Mehmet Cagri Kaymak

class jax_md.mm_forcefields.reaxff.reaxff_forcefield.ForceField(num_atom_types, name_to_index, params_to_indices, body3_indices_src, body3_indices_dst, body4_indices_src, body4_indices_dst, self_energies, shift, low_tap_rad, up_tap_rad, cutoff, cutoff2, hb_close_cutoff, hb_far_cutoff, body2_params_mask, body3_params_mask, body4_params_mask, body34_params_mask, hb_params_mask, global_params, electronegativity, idempotential, gamma, rvdw, p1co, p1co_off, p1co_off_mask, eps, p2co, p2co_off, p2co_off_mask, alf, p3co, p3co_off, p3co_off_mask, vop, amas, rat, rob1, rob1_off, rob1_off_mask, rapt, rob2, rob2_off, rob2_off_mask, vnq, rob3, rob3_off, rob3_off_mask, ptp, pdp, popi, pdo, bop1, bop2, de1, de2, de3, psp, psi, aval, vval3, bo131, bo132, bo133, ovc, v13cor, softcut, softcut_2d, stlp, valf, vval1, vval2, vval4, vkac, th0, vka, vkap, vka3, vka8, vlp1, valp1, vovun, vover, v1, v2, v3, v4, vconj, nphb, rhb, dehb, vhb1, vhb2, vdw_shiedling, trip_stab4, trip_stab5, trip_stab8, trip_stab11, over_coord1, over_coord2, val_par3, val_par15, val_par17, val_par18, val_par20, val_par21, val_par22, val_par31, val_par34, val_par39, par_16, par_6, par_7, par_9, par_10, par_32, par_33, par_24, par_25, par_26, par_28, par_35)

Bases: object

Container for ReaxFF parameters

alf: Array

amas: Array

aval: Array

bo131: Array

bo132: Array

bo133: Array

body2_params_mask: Array

body34_params_mask: Array

body3_indices_dst: tuple

body3_indices_src: tuple

body3_params_mask: Array

body4_indices_dst: tuple

body4_indices_src: tuple

body4_params_mask: Array

bop1: Array

bop2: Array

cutoff: Array

cutoff2: Array

de1: Array

de2: Array

de3: Array

dehb: Array

electronegativity: Array

eps: Array

fill_off_diag()

Fills the off-diagonal entries in the parameter arrays

fill_symm()

Fills the parameter arrays based on the symmetries

gamma: Array

global_params: Array

hb_close_cutoff: Array

hb_far_cutoff: Array

hb_params_mask: Array

idempotential: Array

classmethod init_from_arg_dict(kwargs)

low_tap_rad: Array

name_to_index: dict

nphb: Array

num_atom_types: int

ovc: Array

over_coord1: Array

over_coord2: Array

p1co: Array

p1co_off: Array

p1co_off_mask: Array

p2co: Array

p2co_off: Array

p2co_off_mask: Array

p3co: Array

p3co_off: Array

p3co_off_mask: Array

par_10: Array

par_16: Array

par_24: Array

par_25: Array

par_26: Array

par_28: Array

par_32: Array

par_33: Array

par_35: Array

par_6: Array

par_7: Array

par_9: Array

params_to_indices: dict

pdo: Array

pdp: Array

popi: Array

psi: Array

psp: Array

ptp: Array

rapt: Array

rat: Array

rhb: Array

rob1: Array

rob1_off: Array

rob1_off_mask: Array

rob2: Array

rob2_off: Array

rob2_off_mask: Array

rob3: Array

rob3_off: Array

rob3_off_mask: Array

rvdw: Array

self_energies: Array

set(**kwargs)

shift: Array

softcut: Array

softcut_2d: Array

stlp: Array

th0: Array

trip_stab11: Array

trip_stab4: Array

trip_stab5: Array

trip_stab8: Array

up_tap_rad: Array

v1: Array

v13cor: Array

v2: Array

v3: Array

v4: Array

val_par15: Array

val_par17: Array

val_par18: Array

val_par20: Array

val_par21: Array

val_par22: Array

val_par3: Array

val_par31: Array

val_par34: Array

val_par39: Array

valf: Array

valp1: Array

vconj: Array

vdw_shiedling: Array

vhb1: Array

vhb2: Array

vka: Array

vka3: Array

vka8: Array

vkac: Array

vkap: Array

vlp1: Array

vnq: Array

vop: Array

vover: Array

vovun: Array

vval1: Array

vval2: Array

vval3: Array

vval4: Array

jax_md.mm_forcefields.reaxff.reaxff_helper module

Contains helper functions ReaxFF

Author: Mehmet Cagri Kaymak

jax_md.mm_forcefields.reaxff.reaxff_helper.init_params_for_filler_atom_type(FF_field_dict)

jax_md.mm_forcefields.reaxff.reaxff_helper.read_force_field(force_field_file, cutoff2=0.001, hbond_close_cutoff=0.01, hbond_far_cutoff=7.5, dtype=<class 'jax.numpy.float32'>)

jax_md.mm_forcefields.reaxff.reaxff_helper.safe_sqrt_jvp(primals, tangents)

jax_md.mm_forcefields.reaxff.reaxff_helper.vectorized_cond(pred, true_fun, false_fun, operand)

jax_md.mm_forcefields.reaxff.reaxff_interactions module

Contains interaction list related functions for ReaxFF

Author: Mehmet Cagri Kaymak

class jax_md.mm_forcefields.reaxff.reaxff_interactions.Filtration(candidate_fn, mask_fn, is_dense, idx, did_buffer_overflow)

Bases: object

Allows filtering out the masked indices from the interaction lists for better performance. Supports dense -> dense and sparse -> sparse filtering.

candidate_fn

A function that creates candidate list to be filtered.

mask_fn

A function that decides which indices to mask. returns a boolean array.

is_dense

An boolean flag to show if it is a dense->dense filter.

idx

Index array.

did_buffer_overflow

A boolean specifying whether or not the cell list exceeded the maximum allocated capacity.

allocate(candidate_args, capacity_multiplier=1.25, min_capacity=0)

Initial allocation

Return type:

Filtration

allocate_fixed(candidate_args, capacity=0)

Allocate a fixed size idx array

Return type:

Filtration

candidate_fn: Callable

count(candidate_args)

Count the number of interactions in the filtered array

did_buffer_overflow: Array

idx: Array | None

is_dense: bool

mask_fn: Callable

set(**kwargs)

update(candidate_args)

Updates the filtered index array and overflow flag

Return type:

Filtration

class jax_md.mm_forcefields.reaxff.reaxff_interactions.ReaxFFNeighborListFns(allocate, update)

Bases: object

allocate: Callable

set(**kwargs)

update: Callable

class jax_md.mm_forcefields.reaxff.reaxff_interactions.ReaxFFNeighborLists(close_nbrs, far_nbrs, filter2, filter3, filter34, filter4, filter_hb_close, filter_hb_far, filter_hb, did_buffer_overflow)

Bases: object

Stores the neighbor lists and filters required for ReaxFF

close_nbrs: NeighborList

did_buffer_overflow: Array

far_nbrs: NeighborList

filter2: Filtration

filter3: Filtration

filter34: Filtration

filter4: Filtration

filter_hb: Filtration | None

filter_hb_close: Filtration | None

filter_hb_far: Filtration | None

set(**kwargs)

jax_md.mm_forcefields.reaxff.reaxff_interactions.body_34_candidate_fn(body_3_inds, nbr_inds, neigh_bo, species, cutoff2, param_mask)

jax_md.mm_forcefields.reaxff.reaxff_interactions.body_3_candidate_fn(nbr_inds, neigh_bo, species, cutoff2, param_mask)

Creates full candidate index and value arrays for 3 body interactions, to be used by a filter.

jax_md.mm_forcefields.reaxff.reaxff_interactions.body_4_candidate_fn(body_3_inds, nbr_inds, neigh_bo, species, cutoff2, param_mask)

Creates full candidate index and value arrays for 4 body interactions, to be used by a filter.

jax_md.mm_forcefields.reaxff.reaxff_interactions.calculate_all_4_body_angles(body_4_inds, nbr_inds, nbr_disps)

Calculates the angle related terms for the provided 4-body interaction list

jax_md.mm_forcefields.reaxff.reaxff_interactions.calculate_all_angles_and_distances(R, nbr_lists, map_metric, map_disp)

jax_md.mm_forcefields.reaxff.reaxff_interactions.calculate_all_hbond_angles_and_dists(hbond_inds, close_nbr_disps, far_nbr_disps)

Calculates the angles and distances for the provided hydrogen-bond list

jax_md.mm_forcefields.reaxff.reaxff_interactions.calculate_angle(disp12, disp32)

Assume there are 3 atoms: atom 1,2 and 3 where 2 is the center disp12 = pos1 - pos2 disp32 = pos3 - pos2

jax_md.mm_forcefields.reaxff.reaxff_interactions.filtration(candidate_fn, mask_fn, is_dense=False)

Returns an empty Filtration object to be used

Return type:

Filtration

jax_md.mm_forcefields.reaxff.reaxff_interactions.find_3_body_inter(N, ctr_ind, nbr_pot, nbr_inds, lcl_inds, cutoff2)

Finds all potential 3-body interactions based on bonded potentials

jax_md.mm_forcefields.reaxff.reaxff_interactions.find_body_4_inter(body_3_item, neigh_inds, neigh_bo, cutoff2)

Finds all potential 4 body interactions involving a 3-body item

jax_md.mm_forcefields.reaxff.reaxff_interactions.hbond_candidate_fn(donor_inds, close_nbr_inds, hb_short_inds, far_nbr_inds, hb_long_inds, species, param_mask)

Creates full candidate index and value arrays for hydrogen bonds, to be used by a filter.

jax_md.mm_forcefields.reaxff.reaxff_interactions.reaxff_inter_list(displacement, box, species, species_AN, force_field, total_charge=0.0, tol=1e-06, max_solver_iter=500, backprop_solve=False, tors_2013=False, solver_model='EEM', short_inters_capacity_multiplier=1.2, long_inters_capacity_multiplier=1.2)

Contains all the necessary logic to run a reaxff simulation and allocate, reallocate, update and energy_fn functions.

Parameters:

• species (Array) – An ndarray of shape [n, ] for the atom types.

• species_AN (Array) – An ndarray of shape [n, ] for the atomic numbers of the atoms.

• total_charge (float) – Total charge of the system (float)

• tol (float) – Tolarence for the charge solver

• max_solver_iter (int) – Maximum number of solver iterations

• backprop_solve (bool) – Control variable to decide whether to do a solve to calculate the gradients of the charges wrt positions. By definition, the gradients should be 0 but if the solver tolerance is high, the gradients might be non-ignorable.

• tors_2013 (bool) – Control variable to decide whether to use more stable version of the torsion interactions

• solver_model (str) – Control variable for the solver model (“EEM” or “ACKS”)

• short_inters_capacity_multiplier (float) – capacity multiplier for all short range interactions

• long_inters_capacity_multiplier (float) – capacity multiplier for all long range interactions

Return type:

Tuple[ReaxFFNeighborListFns, Callable]

Returns:

System energy in kcal/mol