diff --git a/examples/gauss.py b/examples/gauss.py new file mode 100644 index 0000000..17d2bc3 --- /dev/null +++ b/examples/gauss.py @@ -0,0 +1,125 @@ +"""Demonstrates generating a Gaussian map for Chignolin. + +NOTE: You must have JAX installed to run this script. +""" + +from typing import ( + Tuple, + Any, +) +import re +from pathlib import Path +import numpy as np +import mdtraj as md # type: ignore [import-untyped] + +# tools for preparing map optimization call +from aggforce import ( + LinearMap, + guess_pairwise_constraints, + joptgauss_map, + project_forces, +) + + +def get_data() -> Tuple[np.ndarray, np.ndarray, md.Trajectory, float]: + r"""Return data for analysis. + + This is currently grabs a group of numpy coordinate and force files, stacks them, + and then along with a pdb-derived mdtraj trajectory and kbt value returns them. + + Note that we must manually provide a value for KbT in appropriate units. + + Returns + ------- + A tuple of the following: + coordinates array + array of positions as a function of time (shape should be + (n_frames,n_sites,n_dims)). Should correspond to the same frames + as the forces array. + forces array + array of forces as a function of time (shape should be + (n_frames,n_sites,n_dims)). Should correspond to the same frames + as the coordinates array. + mdtraj.Trajectory + mdtraj trajectory corresponding to the sites in the coordinates and + forces array. We use it to make the configurational map by + considering the atom names, although the method used to generate the + configurational map may be modified. It does not need more than one + frame (it can be generated from a pdb). + KbT (float) + Boltzmann's constant times the temperature of the reference + trajectory. See code for units. + """ + kbt = 0.6955215 # kbt for 350K in kcal/mol, known a priori + + force_list = [ + np.load(str(name))["Fs"] for name in Path().glob("record_*_prod_*.npz") + ] + coord_list = [ + np.load(str(name))["coords"] for name in Path().glob("record_*_prod_*.npz") + ] + forces = np.vstack(force_list) + coords = np.vstack(coord_list) + pdb = md.load("data/cln025.pdb") + return (coords, forces, pdb, kbt) + + +def gen_config_map(pdb: md.Trajectory, string: str) -> LinearMap: + """Create the configurational map. + + This is needed as it defines constraints which dictate which force maps are + feasible. + + We here generate a (usually carbon alpha) configurational map using mdtraj's + topology. The map could also be specified externally. + + Arguments: + --------- + pdb (mdtraj.Trajectory): + Trajectory object describing the fine-grained (e.g. atomistic) + resolution. + string (string): + Regex string which is compared against the str() of the topology.atoms + entry--- if matched that atom is retained in the configurational map. + + Returns: + ------- + A LinearMap object which characterizes the configurational map. There are + multiple ways to initialize this object; see the main code for more details. + """ + inds = [] + atomlist = list(pdb.topology.atoms) + # record which atoms match the string via str casing, e.g., which are carbon alphas. + for ind, a in enumerate(atomlist): + if re.search(string, str(a)): + inds.append([ind]) + return LinearMap(inds, n_fg_sites=pdb.xyz.shape[1]) + + +def main() -> Any: + """Create Gaussian map.""" + coords, forces, pdb, kbt = get_data() + # cmap is the configurational coarse-grained map + cmap = gen_config_map(pdb, "CA$") + # guess molecular constraints + constraints = guess_pairwise_constraints(coords[0:10], threshold=1e-3) + + optim_results = project_forces( + coords=coords, + forces=forces, + coord_map=cmap, + constrained_inds=constraints, + kbt=kbt, + method=joptgauss_map, + l2_regularization=1e1, + var=0.01, + ) + + # this is the optimized trajectory map + return optim_results["tmap"] + + +if __name__ == "__main__": + """Create Gaussian map.""" + main() + # we do nothing with the output