diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..43ae0e2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.py[cod] diff --git a/module b/module new file mode 120000 index 0000000..dc428f0 --- /dev/null +++ b/module @@ -0,0 +1 @@ +modules \ No newline at end of file diff --git a/modules b/modules index 78e5ec5..4409f54 160000 --- a/modules +++ b/modules @@ -1 +1 @@ -Subproject commit 78e5ec5d0743420b0eb8e526d80c89a4e734b704 +Subproject commit 4409f540847b21af3dc6f5304d4cec126bc8211b diff --git a/preprocess.py b/preprocess.py index d7657e8..d1cee7d 100755 --- a/preprocess.py +++ b/preprocess.py @@ -1,1106 +1,13 @@ #!/usr/bin/env python3 -import os -import numpy as np -import yaml -import json -import h5py -import mdtraj -from module.torchmd_cg_mappings import CACB_MAP -from module import prior -from module import prior_flex -from module import psfwriter -from module.make_deltaforces import DeltaForces -from module.cg_mapping import CGMapping -import argparse -import traceback -import shutil -import multiprocessing as mp -from tqdm import tqdm -import glob -import pickle -import torch -# this can have a small performance hit as it uses the HDD file_system to share data across different processes, but it doesn't lead to "Too many files open" error, which was limiting the max number of parallel processes to 16. -torch.multiprocessing.set_sharing_strategy('file_system') +"""CLI shim: delegates to the `preprocess` package. -# Raz Dec 30 2024: turns out that when preprocessing the 6.5k dataset, the last 20 pdbs are taking forever to process due to being very big proteins. In addition, some were giving hdf5 errors since the batch_generate job didn't properly finish, and this means we don't have all the frames, and they will be removed later on before training. So here's the best workflow I found to solve it: -# 1) first run just step 1 on all proteins, and set FILTER_NOT_PROCESSED_STEP_ONE = False. if any pdbs are taking forever (generally the last 20), just kill the job -# 2) re-run the pre-processing for a second time with FILTER_NOT_PROCESSED_STEP_ONE = True -FILTER_NOT_PROCESSED_STEP_ONE = False - -# Controls whether to re-generate the priors in step 2 for each of the terms, terms in the list will be loaded from the cache instead of refit -#USE_CACHED_FITS = ['dihedrals', 'angles', 'bonds', 'lj'] -USE_CACHED_FITS = [] - -DEVICE_STEP_3 = 'cpu' -#DEVICE_STEP_3 = 'cuda' - -DO_STEP_1 = True # whether to do step 1. if you got errors in steps 2-3 and want to resume, set this to False -REGEN_CACHE_FILES = True # whether to re-generate cache files - - -def process_init(counter): - """This function sets the worker names such that we can use them to position the tqdm bars""" - with counter.get_lock(): - idx = int(counter.value) - counter.value += 1 - mp.current_process().name = f"PreprocessWorker-{idx}" - - -class CGMappingDef_CA: - def __init__(self): - residues = ["ALA", "CYS", "ASP", "GLU", "PHE", "GLY", "HIS", "ILE", "LYS", "LEU", "MET", "ASN", "PRO", "HYP", "GLN", "ARG", "SER", "THR", "VAL", "TRP", "TYR"] - # For legacy reasons we have a couple extra ambiguous residues (ASX & GLX) in the embedding map but we do not accept these for parsing - embedding_residues = ["ALA", "ARG", "ASN", "ASP", "ASX", "CYS", "GLU", "GLN", "GLX", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"] - self.bead_embeddings = {name: [index + 1] for index, name in enumerate(sorted(embedding_residues))} - - # bead_atom_selection: A list of lists, where each inner list is the names of the atoms that will be combined to form the bead - self.bead_atom_selection = {k: [["CA"]] for k in residues} - # The type names of beads (will become the atom type/element in the cg topology) - self.bead_types = { - "ALA": ["CAA"], - "ARG": ["CAR"], - "ASN": ["CAN"], - "ASP": ["CAD"], - "CYS": ["CAC"], - "GLN": ["CAQ"], - "GLU": ["CAE"], - "GLY": ["CAG"], - "HIS": ["CAH"], - "HSD": ["CAH"], - "ILE": ["CAI"], - "LEU": ["CAL"], - "LYS": ["CAK"], - "MET": ["CAM"], - "PHE": ["CAF"], - "PRO": ["CAP"], - "SER": ["CAS"], - "THR": ["CAT"], - "TRP": ["CAW"], - "TYR": ["CAY"], - "VAL": ["CAV"], - } - # The "atom name" assigned to the beads - self.bead_atom_names = {k: ["CA"] for k in residues} - self.bead_masses = {k: [12.01] for k in residues} - self.bead_backbone_idx = {k: 0 for k in residues} - -class CGMappingDef_CACB: - def __init__(self): - residues = ["ALA", "CYS", "ASP", "GLU", "PHE", "GLY", "HIS", "ILE", "LYS", "LEU", "MET", "ASN", "PRO", "HYP", "GLN", "ARG", "SER", "THR", "VAL", "TRP", "TYR"] - - # bead_atom_selection: A list of lists, where each inner list is the names of the atoms that will be combined to form the bead - self.bead_atom_selection = {k: [["CA"], ["CB"]] for k in residues} - self.bead_atom_selection["GLY"] = [["CA"]] - # The type names of beads (will become the atom type/element in the cg topology) - self.bead_types = { - "ALA": ["CA", "CBA"], - "ARG": ["CA", "CBR"], - "ASN": ["CA", "CBN"], - "ASP": ["CA", "CBD"], - "CYS": ["CA", "CBC"], - "GLN": ["CA", "CBQ"], - "GLU": ["CA", "CBE"], - "GLY": ["CAG"], - "HIS": ["CA", "CBH"], - "HSD": ["CA", "CBH"], - "ILE": ["CA", "CBI"], - "LEU": ["CA", "CBL"], - "LYS": ["CA", "CBK"], - "MET": ["CA", "CBM"], - "PHE": ["CA", "CBF"], - "PRO": ["CA", "CBP"], - "SER": ["CA", "CBS"], - "THR": ["CA", "CBT"], - "TRP": ["CA", "CBW"], - "TYR": ["CA", "CBY"], - "VAL": ["CA", "CBV"], - } - - embedding_map = {k:i for i,k in enumerate(sorted(set.union(*[set(i) for i in self.bead_types.values()])))} - self.bead_embeddings = {k:[embedding_map[i] for i in v] for k, v in self.bead_types.items()} - - # The "atom name" assigned to the beads - self.bead_atom_names = {k: ["CA", "CB"] for k in residues} - self.bead_atom_names["GLY"] = ["CA"] - self.bead_masses = {k: [12.01]*len(v) for k,v in self.bead_types.items()} - self.bead_backbone_idx = {k: 0 for k in residues} - -class PriorBuilder: - def __init__(self): - self.prior_params = dict() - self.priors = None - self.terms = dict() - self.atom_types = set() - self.fit_constraints = True - self.tag_beta_turns = False - self.min_cnt = 0 - - def select_atoms(self, topology): - """Returns tha atom index to be saved for this prior""" - raise NotImplementedError() - - def map_embeddings(self, selected_atoms, trajectory): - """Generates the embeddings array for the selected atoms""" - raise NotImplementedError() - - def write_psf(self, pdb_file, psf_file): - """Write the .psf file describing the course grain geometry""" - raise NotImplementedError() - - def add_molecule(self, mol, traj, cache_dir): - fit_ok_path = os.path.join(cache_dir, "fit_ok.txt") - - if cache_dir and os.path.exists(fit_ok_path): - os.unlink(fit_ok_path) - - for term in self.terms.values(): - term.add_molecule(mol, traj, cache_dir) - self.atom_types = self.atom_types.union(mol.atomtype) - - if cache_dir: - np.save(os.path.join(cache_dir, "atomtype.npy"), mol.atomtype) - with open(fit_ok_path, "wt", encoding="utf-8") as f: - f.write("ok") - - def load_molecule_cache(self, cache_dir): - assert os.path.exists(os.path.join(cache_dir, "fit_ok.txt")) - atomtype = np.load(os.path.join(cache_dir, "atomtype.npy"), allow_pickle=True) - self.atom_types = self.atom_types.union(atomtype) - - for term in self.terms.values(): - term.load_molecule_cache(cache_dir) - - def enable_fit_constraints(self, use_constraints): - self.fit_constraints = use_constraints - self.prior_params["fit_constraints"] = self.fit_constraints - - def enable_bond_tags(self, use_tags): - self.tag_beta_turns = use_tags - self.prior_params["tag_beta_turns"] = self.tag_beta_turns - - def set_min_cnt(self, min_cnt): - assert min_cnt >= 0 - self.min_cnt = min_cnt - self.prior_params["min_cnt"] = self.min_cnt - - def fit(self, temperature, plot_dir=None): - self.init_prior_dict() - assert self.priors is not None - for key, term in self.terms.items(): - if os.path.exists(f"{plot_dir}/prior_{key}.pkl") and (key in USE_CACHED_FITS): - print(f"Used cached fit for {key}...") - with open(f"{plot_dir}/prior_{key}.pkl", "rb") as f: - self.priors[key] = pickle.load(f) - else: - print(f"Fitting {key}...") - self.priors[key] = term.get_param(temperature, plot_dir, self.fit_constraints, self.min_cnt) - # pickle the prior for this term - with open(f"{plot_dir}/prior_{key}.pkl", "wb") as f: - pickle.dump(self.priors[key], f) - - def init_prior_dict(self): - # Define the force field dict - priors = {} - priors['atomtypes'] = sorted(self.atom_types) - priors['bonds'] = {} - priors['angles'] = {} - priors['dihedrals'] = {} - priors['lj'] = {} - # For mass and charge assume everything is a carbon atom - priors['electrostatics'] = {at: {'charge': 0.0} for at in priors['atomtypes']} - # The mass of carbon used here is the from OpenMM/AMBER-14 value - priors['masses'] = {at: 12.01 for at in priors['atomtypes']} - self.priors = priors - - def save_prior(self, output_path, pdbid): - prefix = "" - if pdbid: - prefix = f"{pdbid}_" - with open(os.path.join(output_path, f"{prefix}priors.yaml"), "w") as f: - yaml.dump(self.priors, f) - with open(os.path.join(output_path, f"{prefix}prior_params.json"),"w") as f: - json.dump(self.prior_params, f) - - def make_mol(self, cg_map): - bonds = "bonds" in self.terms - angles = "angles" in self.terms - dihedrals = "dihedrals" in self.terms - return cg_map.to_mol(bonds = bonds, angles = angles, dihedrals = dihedrals) - -class Prior_CA(PriorBuilder): - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA", - "exclusions" : ['bonds'], - "forceterms" : ["bonds"], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - - def build_mapping(self, topology): - return CGMapping(topology, CGMappingDef_CA()) - - def select_atoms(self, topology): - #TODO: Remove this function (replaced by build_mapping) - return topology.select('name CA and protein') - - def map_embeddings(self, selected_atoms, topology): #pyright: ignore[reportIncompatibleMethodOverride] - #TODO: Remove this function (replaced by build_mapping) - standardResidues = {"ALA", "ARG", "ASN", "ASP", "ASX", "CYS", "GLU", "GLN", "GLX", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"} - amino_acid_mapping = {name: index + 1 for index, name in enumerate(sorted(standardResidues))} - - result = [] - for a_idx in selected_atoms: - r_name = topology.atom(a_idx).residue.name - result.append(amino_acid_mapping[r_name]) - return np.array(result, dtype=int) - - def write_psf(self, pdb_file, psf_file): - #TODO: Remove this function (replaced by build_mapping) - bonds = "bonds" in self.terms - angles = "angles" in self.terms - dihedrals = "dihedrals" in self.terms - return psfwriter.pdb2psf_CA(pdb_file, psf_file, bonds = bonds, angles = angles, dihedrals = dihedrals, - tag_beta_turns = self.tag_beta_turns) - -class Prior_CACB(PriorBuilder): - """Implements the torchmd-cg CACB prior""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CACB", - "exclusions" : ['bonds'], - "forceterms" : ["bonds"], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - - def build_mapping(self, topology): - return CGMapping(topology, CGMappingDef_CACB()) - - def select_atoms(self, topology): - #TODO: Remove this function (replaced by build_mapping) - return topology.select('(name CA or name CB) and protein') - - def map_embeddings(self, selected_atoms, topology):#pyright: ignore[reportIncompatibleMethodOverride] - #TODO: Remove this function (replaced by build_mapping) - - # Make a map from embedding name to embedding name number - # e.g. {"CAA":0, "CAC":1, ...} - embedding_map = CACB_MAP - embedding_nums = dict([(k, i) for i, k in enumerate(sorted(set(embedding_map.values())))]) - - result = [] - for a_idx in selected_atoms: - r_name = topology.atom(a_idx).residue.name - a_name = topology.atom(a_idx).name - emb_name = embedding_map[(r_name, a_name)] - result.append(embedding_nums[emb_name]) - return np.array(result, dtype=int) - - def write_psf(self, pdb_file, psf_file): - #TODO: Remove this function (replaced by build_mapping) - bonds = "bonds" in self.terms - angles = "angles" in self.terms - dihedrals = "dihedrals" in self.terms - return psfwriter.pdb2psf_CACB(pdb_file, psf_file, bonds = bonds, angles = angles, dihedrals = dihedrals) - -class Prior_CACB_lj(Prior_CACB): - """torchmd-cg CACB prior with Bonded & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CACB_lj", - "exclusions" : ['bonds'], - "forceterms" : ['bonds', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - -class Prior_CACB_lj_angle_dihedral(Prior_CACB): - """torchmd-cg CACB prior with Bonded, Angle, Dihedral & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CACB_lj_angle_dihedral", - "exclusions" : ['bonds', 'angles', 'dihedrals'], - "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.ParamAngleCalculator() - self.terms["dihedrals"] = prior.ParamDihedralCalculator() - -class Prior_CA_lj(Prior_CA): - """CA prior with Bonded & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj", - "exclusions" : ['bonds'], - "forceterms" : ['bonds', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - -class Prior_CA_lj_angle(Prior_CA): - """CA prior with Bonded, Angle, and RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angle", - "exclusions" : ['bonds', 'angles'], - "forceterms" : ['bonds', 'angles', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms['angles'] = prior.ParamAngleCalculator() - -class Prior_CA_lj_angle_dihedral(Prior_CA): - """torchmd-cg CA prior with Bonded, Angle, Dihedral & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angle_dihedral", - "exclusions" : ['bonds', 'angles', 'dihedrals'], - "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.ParamAngleCalculator() - self.terms["dihedrals"] = prior.ParamDihedralCalculator() - -class Prior_CA_lj_angle_dihedralX(Prior_CA): - """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angle_dihedralX", - "exclusions" : ['bonds', 'angles', 'dihedrals'], - "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.ParamAngleCalculator() - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) - -class Prior_CA_lj_angleXCX_dihedralX(Prior_CA): - """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angleXCX_dihedralX", - "exclusions" : ['bonds', 'angles', 'dihedrals'], - "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.ParamAngleCalculator(center=True) - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) - -class Prior_CA_lj_angleXCX_dihedralX_flex(Prior_CA): - """torchmd-cg CA prior with highly flexible Bonded, Angle, DihedralX & RepulsionCG terms that fit the data. - - """ - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angleXCX_dihedralX_flex", - "exclusions" : ['bonds', 'angles', 'dihedrals'], - "forceterms_nn" : ['bonds', 'angles', 'dihedrals'], - "forceterms_classical": ['repulsioncg'], # changed from lj, would need to re-generated the dataset (Jan 10 2025). repulsioncg is using just the repulsion term from lj. it uses the same parameters as lj, so need to make sure the right function is evaluated. - "external" : True - }) - self.prior_params['forceterms'] = self.prior_params['forceterms_classical'] + self.prior_params['forceterms_nn'] - - self.terms["bonds"] = prior_flex.ParamBondedFlexCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior_flex.ParamAngleFlexCalculator(center=True) - self.terms["dihedrals"] = prior_flex.ParamDihedralFlexCalculator(unified=True) - - # have to override this method since we're saving neural nets as priors - def save_prior(self, output_path, pdbid): - prefix = "" - # if pdbid: - # prefix = f"{pdbid}_" - with open(os.path.join(output_path, f"{prefix}prior_params.json"),"w") as f: - json.dump(self.prior_params, f) - - # print('self.priors', self.priors.keys()) - # remove the dihedrals and bonds from the priors - priorsTruncated = self.priors.copy() - priorsTruncated.pop('dihedrals') - priorsTruncated.pop('bonds') - priorsTruncated.pop('angles') - # print('priorsTruncated', priorsTruncated.keys()) - - # save the classical priors using yaml. this is requires because the classical priors are built from the yaml files - with open(os.path.join(output_path, f"{prefix}priors.yaml"), "w") as f: - yaml.dump(priorsTruncated, f) - - self.priors['terms'] = self.terms - self.priors['prior_params'] = self.prior_params - - # also save with pickle - with open(os.path.join(output_path, f"{prefix}priors.pkl"), "wb") as f: - pickle.dump(self.priors, f) - - - def load_prior_nnets(self, output_path): - # load the prior with pickle - with open(os.path.join(output_path, "priors.pkl"), "rb") as f: - self.priors = pickle.load(f) - - # return self.priors - - # with open(os.path.join(output_path, f"{prefix}priors.pkl"), "wb") as f: - # pickle.dump(self.priors, f) - - - -class Prior_CA_lj_angleXCX_dihedralX_V1(Prior_CA): - """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angleXCX_dihedralX_V1", - "exclusions" : ['bonds', 'angles', '1-4'], - "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.ParamAngleCalculator(center=True) - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) - -class Prior_CA_lj_bondNull_angleXCX_dihedralX(Prior_CA): - """torchmd-cg CA prior with Angle, DihedralX & RepulsionCG terms (+ bond exclusions)""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_bondNull_angleXCX_dihedralX", - "exclusions" : ['bonds', 'angles', '1-4'], - "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], - }) - self.terms["bonds"] = prior.NullParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.ParamAngleCalculator(center=True) - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) - -class Prior_CA_lj_bondNull_angleNull_dihedralX(Prior_CA): - """torchmd-cg CA prior with DihedralX & RepulsionCG terms (+ bond & angle exclusions)""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_bondNull_angleNull_dihedralX", - "exclusions" : ['bonds', 'angles', '1-4'], - "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], - }) - self.terms["bonds"] = prior.NullParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.NullParamAngleCalculator() - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) - -class Prior_CA_lj_bondNull_angleNull_dihedralNull(Prior_CA): - """torchmd-cg CA prior with RepulsionCG terms (+ bond, angle, & dihedral exclusions)""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_bondNull_angleNull_dihedralNull", - "exclusions" : ['bonds', 'angles', '1-4'], - "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], - }) - self.terms["bonds"] = prior.NullParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.NullParamAngleCalculator() - self.terms["dihedrals"] = prior.NullParamDihedralCalculator() - -class Prior_CA_lj_angleNull_dihedralX(Prior_CA): - """torchmd-cg CA prior with Bonded, DihedralX & RepulsionCG terms (+ angle exclusions)""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angleNull_dihedralX", - "exclusions" : ['bonds', 'angles', '1-4'], - "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.NullParamAngleCalculator() - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) - -class Prior_CA_lj_angleNull_dihedralNull(Prior_CA): - """torchmd-cg CA prior with Bonded & RepulsionCG terms (+ angle & dihedral exclusions)""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_angleNull_dihedralNull", - "exclusions" : ['bonds', 'angles', '1-4'], - "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["angles"] = prior.NullParamAngleCalculator() - self.terms["dihedrals"] = prior.NullParamDihedralCalculator() - -class Prior_CA_Majewski2022_v0(Prior_CA): - """torchmd-cg CA prior based on the parameters used in (Majewski 2022) - Note this version (v0) has different lj exclusions than the one used in the paper. - """ - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_Majewski2022_v0", - "exclusions" : ['bonds', 'dihedrals'], - "forceterms" : ['bonds', 'dihedrals', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True, scale=0.5) - -class Prior_CA_Majewski2022_v1(Prior_CA): - """torchmd-cg CA prior based on the parameters used in (Majewski 2022)""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_Majewski2022_v1", - "exclusions" : ['bonds'], - "forceterms" : ['bonds', 'dihedrals', 'repulsioncg'], - }) - self.terms["bonds"] = prior.ParamBondedCalculator() - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6], exclusion_terms={"bonds"}) - self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True, scale=0.5) - -class Prior_CA_null(Prior_CA): - """CA prior with no terms""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_null", - "exclusions" : [], - "forceterms" : [], - }) - self.terms = {} - -class Prior_CA_lj_only(Prior_CA): - """CA prior with just a RepulsionCG term""" - def __init__(self): - super().__init__() - self.prior_params.update({ - "prior_configuration_name": "CA_lj_only", - "exclusions" : [], - "forceterms" : ['RepulsionCG'], - }) - self.terms = {} - self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) - -def slice_to_str(s): - result = [s.start, s.stop, s.step] - result = [str(i) if i is not None else '' for i in result] - return ":".join(result) - -def get_prior_params_path(prior_path): - dir_path, file_name = os.path.split(prior_path) - file_name = file_name.replace("priors.yaml", "prior_params.json") - return os.path.join(dir_path, file_name) - -def load_h5_traj_slice(path, slice): - """Load a slice from a h5 trajectory without reading the entire file into memory""" - base_traj = mdtraj.load_frame(path, 0) - with h5py.File(path) as f: - t_xyz = f["coordinates"][slice][:] #pyright: ignore[reportIndexIssue] - t_time = f["time"][slice][:] #pyright: ignore[reportIndexIssue] - - t_unitcell_lengths = None - t_unitcell_angles = None - if "cell_lengths" in f.keys(): - t_unitcell_lengths = f["cell_lengths"][slice][:] #pyright: ignore[reportIndexIssue] - t_unitcell_angles = f["cell_angles"][slice][:] #pyright: ignore[reportIndexIssue] - - result = mdtraj.Trajectory(t_xyz, base_traj.topology, time=t_time, unitcell_lengths=t_unitcell_lengths, unitcell_angles=t_unitcell_angles) - return result - -class Preprocessor: - def __init__(self, dataset_conf, input_path_map, save_path, prior_builder, prior_file, prior_name, frame_slice, temp, optimize_forces, box, prior_plots, resume_preprocess, num_cores, jobid=None, totalNrJobs=None): - self.dataset_conf = dataset_conf - self.save_path = save_path - self.prior_builder = prior_builder - self.prior_file = prior_file - self.frame_slice = frame_slice - self.temp = temp - self.jobid = jobid - self.totalNrJobs = totalNrJobs - - self.pdbid_list = input_path_map - - if FILTER_NOT_PROCESSED_STEP_ONE: - pdbs_processed_step1 = [f.split('/')[-3] for f in glob.glob(os.path.join(save_path, "*/fit/fit_ok.txt"))] - # remove keys that are not in pdbs_processed_step1 - self.pdbid_list = {k: v for k, v in self.pdbid_list.items() if k in pdbs_processed_step1} - print('%d pdbs left after removing pdbs not processed in step 1' % len(self.pdbid_list)) - - # if os.path.exists(os.path.join(save_path, 'pdb_list.pkl')): - # with open(os.path.join(save_path, 'pdb_list.pkl'), 'rb') as f: - # self.pdbid_list = pickle.load(f) - # else: - # self.pdbid_list = self.get_pdbid_list() - - # if FILTER_NOT_PROCESSED_STEP_ONE: - # pdbs_processed_step1 = [f.split('/')[-3] for f in glob.glob(os.path.join(save_path, "*/fit/fit_ok.txt"))] - # # remove keys that are not in pdbs_processed_step1 - # self.pdbid_list = {k: v for k, v in self.pdbid_list.items() if k in pdbs_processed_step1} - # print('%d pdbs left after removing pdbs not processed in step 1' % len(self.pdbid_list)) - - # # pickle pdb_list - # os.makedirs(save_path, exist_ok=True) - # with open(os.path.join(save_path, 'pdb_list.pkl'), 'wb') as f: - # pickle.dump(self.pdbid_list, f) - - - self.optimize_forces = optimize_forces - self.box = box - self.prior_plots = prior_plots - self.resume_preprocess = resume_preprocess - self.num_cores = num_cores - - print("Input directory paths:", [i["path"] for i in self.dataset_conf]) - print("Save directory path:", self.save_path) - print(f"Temperature: {self.temp}") - print("Frame slice:", slice_to_str(self.frame_slice)) - print("Number of cores used for parallelization:", self.num_cores) - # print("PDB ID list:", self.pdbid_list) - - def step1_threading(self, pdbid): - try: - cache_dir = os.path.join(self.save_path, pdbid, "fit") - if not(self.resume_preprocess and os.path.exists(os.path.join(cache_dir, "fit_ok.txt"))): - # This assumes we've named the processes during initialization - bar_pos = int(mp.current_process().name.split("-")[1]) + 1 - self.process_step1(pdbid, bar_pos) - return [] - except Exception as e: - traceback.print_tb(e.__traceback__) - print(f"{pdbid}:", e) - raise - - def step3_threading(self, pdbid): - try: - # This assumes we've named the processes during initialization - bar_pos = int(mp.current_process().name.split("-")[1]) + 1 - self.process_step3(pdbid, bar_pos) - except Exception as e: - traceback.print_tb(e.__traceback__) - print(f"{pdbid}:", e) - raise - - def preprocess(self): - os.makedirs(os.path.join(self.save_path, "result"), exist_ok=True) - - info_dict = { - "input_paths": [i["path"] for i in self.dataset_conf], - "frame_slice": slice_to_str(self.frame_slice), - "pdbids": list(self.pdbid_list.keys()), - "optimize_forces": self.optimize_forces, - "box": self.box, - "prior_name": prior_name - } - - # If resuming, validate that no paramters that would invalidate the fit cache object have changed - # FIXME: This should also check "--tag-beta-turns" - if self.resume_preprocess: - if os.path.exists(os.path.join(self.save_path, "result/info.json")): - with open(os.path.join(self.save_path, "result/info.json"), "rt", encoding="utf-8") as f: - prevous_info = json.load(f) - for k in ["box", "frame_slice", "optimize_forces", "prior_name"]: - assert info_dict[k] == prevous_info[k], \ - f"Can't resume with different parameters: {k}: {info_dict[k]} != {prevous_info[k]}" - - with open(os.path.join(self.save_path, "result/info.json"), "wt", encoding="utf-8") as f: - json.dump(info_dict, f) - - pdbids = self.pdbid_list - - # Ensure all jobs have the same tqdm lock : https://github.com/tqdm/tqdm/issues/982 - tqdm.get_lock() - - # TODO: Print exceptions in the main thread for legibility - # TODO: Abstract the loop logic instead of repeating it twice - - # Truncate any existing ok_list.txt - with open(os.path.join(self.save_path, "result/ok_list.txt"), "wt", encoding="utf-8") as ok_list: - pass - - - - # Run step 1 in parallel, saving the results to the cache - - if DO_STEP_1: - errorList = {} - thread_counter = mp.Value('i', 0, lock=True) - with tqdm(total=len(pdbids), desc="Processing Step 1", dynamic_ncols=True) as pbar: - with mp.Pool(self.num_cores, initializer=process_init, initargs=(thread_counter,)) as pool: - pending_results = {} - - # Submit tasks and map results to pdbids - for pdbid in pdbids: - result = pool.apply_async(self.step1_threading, args=(pdbid,)) - pending_results[result] = pdbid - - while pending_results: - # Check completed tasks - for result in list(pending_results.keys()): # Iterate over a copy to allow removal - if result.ready(): - try: - result.get() # Retrieve result or raise exception - except Exception as e: - pdbid = pending_results[result] # Get the corresponding pdbid - errorList[pdbid] = str(e) - finally: - del pending_results[result] # Remove completed task - - # Update the progress bar - pbar.n = len(pdbids) - len(pending_results) - pbar.refresh() - - # Wait for 1 second or for the last job to finish - if pending_results: - next(iter(pending_results)).wait(1) - - if len(errorList): - print('errorList', errorList) - print('errorList keys', errorList.keys()) - - if not self.prior_file: - # pickle prior builder - if not os.path.exists(self.save_path + '/prior_builder.pkl') or REGEN_CACHE_FILES: - # Merge cache files back into prior builder - for pdbid in tqdm(pdbids, desc="Merging cache files together"): - cache_dir = os.path.join(self.save_path, pdbid, "fit") - self.prior_builder.load_molecule_cache(cache_dir) - - with open(self.save_path + '/prior_builder.pkl', 'wb') as f: - pickle.dump(self.prior_builder, f) - else: - print("Using cached prior_builder object... ") - - with open(self.save_path + '/prior_builder.pkl', 'rb') as f: - self.prior_builder = pickle.load(f) - - self.process_step2() - - self.prior_builder.save_prior(self.save_path, None) - else: - prior_params_path = get_prior_params_path(prior_file) - shutil.copy(self.prior_file, os.path.join(self.save_path, "priors.yaml")) - shutil.copy(prior_params_path, os.path.join(self.save_path, "prior_params.json")) - - if self.totalNrJobs: - pdblist = list(pdbids.keys()) - pdbidsPerJob = len(pdblist) // self.totalNrJobs + 1 - jobid = self.jobid - assert jobid is not None - if jobid < self.totalNrJobs - 1: - pdbids_c = [pdblist[i] for i in range(jobid * pdbidsPerJob, (jobid + 1) * pdbidsPerJob)] - else: - pdbids_c = [pdblist[i] for i in range(jobid * pdbidsPerJob, len(pdblist))] - - # filter pdbids for this job only - pdbids = {k: v for k, v in pdbids.items() if k in pdbids_c} - - print(f"Step 3: Processing {len(pdbids)} pdbids") - - # Run step 3 in parallel - thread_counter = mp.Value('i', 0, lock=True) - with tqdm(total=len(pdbids), desc="Processing Step 3", dynamic_ncols=True) as pbar: - with mp.Pool(self.num_cores, initializer=process_init, initargs=(thread_counter,)) as pool: - pending_results = [] - - for pdbid in pdbids: - pending_results += [pool.apply_async(self.step3_threading, args=(pdbid,))] - - while pending_results: - # Check for exceptions - [i.get() for i in pending_results if i.ready()] - # Remove finished jobs from the list - pending_results = [i for i in pending_results if not i.ready()] - pbar.n = len(pdbids) - len(pending_results) - pbar.refresh() - # Wait for 1 second or for the last job to finish - if pending_results: - pending_results[0].wait(1) - - # alternatively, cd to the preprocessed_data directory and run this cmd: - # ls */raw/deltaforces.npy | awk '{print substr($1, 1, 4)}' > result/ok_list.txt - with open(os.path.join(self.save_path, "result/ok_list.txt"), "wt", encoding="utf-8") as ok_list: - ok_list.write("\n".join(pdbids)) - - print("Done!") - - def save_data(self, output_path, trajectory, embeddings, forces, pdbid): - # print(f" {pdbid} (coordinates, forces): {trajectory.xyz.shape}, {forces.shape}") - np.save(f"{output_path}/raw/embeddings.npy", embeddings) - np.save(f"{output_path}/raw/forces.npy", forces) - np.save(f"{output_path}/raw/coordinates.npy", trajectory.xyz) - box_path = f"{output_path}/raw/box.npy" - if self.box: - np.save(box_path, trajectory.unitcell_vectors) - elif os.path.exists(box_path): - os.unlink(box_path) - - def process_step1(self, pdbid, bar_position=0): - """Generate the course grained data and topology for the protein, then add it to the prior builder""" - - with tqdm(total=7, position=bar_position, desc=f"{pdbid}: File path setup", dynamic_ncols=True, leave=False) as pbar: - def progress_bar_step(msg): - pbar.update(1) - pbar.set_description_str(f"{pdbid}: {msg}") - - # Set up paths and create directories - output_path = os.path.join(self.save_path, pdbid) - - # TODO: Get rit of the of subdirectories? - os.makedirs(f"{output_path}/raw", exist_ok=True) - os.makedirs(f"{output_path}/processed", exist_ok=True) - - # Find which path this ID belongs to - input_file_path = self.pdbid_list[pdbid] - - progress_bar_step("Loading trajectory") - AAtraj = load_h5_traj_slice(input_file_path, self.frame_slice) - assert AAtraj.xyz is not None # for pyright - AAtraj.xyz *= 10 # convert to angstroms - - progress_bar_step("Building CG mapping") - cg_map = self.prior_builder.build_mapping(AAtraj.topology) - mol = self.prior_builder.make_mol(cg_map) - topology = cg_map.to_mol(bonds=True, angles=True, dihedrals=True) - mol.write(f'{output_path}/processed/{pdbid}_processed.psf') - topology.write(f'{output_path}/processed/topology.psf') # Save the topology for the CG mapping, this is optional but useful for debugging - - # Get the forces - progress_bar_step("Mapping CG forces") - with h5py.File(input_file_path, "r") as f: - forces = f["forces"][self.frame_slice, :, :] #pyright: ignore[reportIndexIssue] - - if self.optimize_forces: - forces = cg_map.cg_optimal_forces(AAtraj, forces) - else: - forces = cg_map.cg_forces(forces) - - assert len(forces) == len(AAtraj) - # Convert from kilojoules/mole/nanometer to kilocalories/mole/angstrom - forces = forces*0.02390057361376673 - - progress_bar_step("Mapping CG coordinates") - xyz = cg_map.cg_positions(AAtraj.xyz) - cg_traj = mdtraj.Trajectory(xyz, topology=cg_map.to_mdtraj()) - if self.box and AAtraj.unitcell_lengths is not None: - cg_traj.unitcell_lengths = AAtraj.unitcell_lengths * 10 - cg_traj.unitcell_angles = AAtraj.unitcell_angles - else: - cg_traj.unitcell_lengths = None - cg_traj.unitcell_angles = None - - # Save the data - progress_bar_step("Saving Data") - self.save_data(output_path, cg_traj, cg_map.embeddings, forces, pdbid) - - # Note: moveaxis creates a view, the original trajectory.xyz is unmodified - assert cg_traj.xyz is not None - mol.coords = np.moveaxis(cg_traj.xyz, 0, -1) - - progress_bar_step("Generating prior fit data") - if not self.prior_file: - cache_dir = os.path.join(output_path, "fit") - os.makedirs(cache_dir, exist_ok=True) - self.prior_builder.add_molecule(mol, cg_traj, cache_dir) - - def process_step2(self): - """Fit the prior forcefield based on accumulated data""" - if self.prior_plots: - plot_dir = os.path.join(self.save_path, "prior_fit_plots") - os.makedirs(plot_dir, exist_ok=True) - else: - plot_dir = None - - self.prior_builder.fit(self.temp, plot_dir=plot_dir) - - def process_step3(self, pdbid, bar_position=0): - """Save prior focefield and generate delta forces data for each protein""" - output_path = os.path.join(self.save_path, pdbid) - - # Remove legacy prior files if they exist - if os.path.exists(f"{output_path}/raw/{pdbid}_priors.yaml"): - os.unlink(f"{output_path}/raw/{pdbid}_priors.yaml") - if os.path.exists(f"{output_path}/raw/{pdbid}_prior_params.json"): - os.unlink(f"{output_path}/raw/{pdbid}_prior_params.json") - - # Generate delta forces for all atom simulation vs. prior FF - coords_npz = f'{output_path}/raw/coordinates.npy' - forces_npz = f'{output_path}/raw/forces.npy' - delta_forces_npz = f'{output_path}/raw/deltaforces.npy' - prior_energy_npz = f'{output_path}/raw/prior_energy.npy' - box_npz = None - if self.box: - box_npz = f"{output_path}/raw/box.npy" - forcefield = os.path.join(self.save_path, "priors.yaml") - psf_file = f'{output_path}/processed/{pdbid}_processed.psf' - prior_params = self.prior_builder.prior_params - - deltaForcesObj = DeltaForces(DEVICE_STEP_3, psf_file, coords_npz, box_npz) - if 'external' in self.prior_builder.prior_params.keys(): - # forceterms = ['bonds', 'angles', 'dihedrals'] - deltaForcesObj.addExternalForces(forcefield, self.prior_builder.priors['bonds'], self.prior_builder.priors['angles'], self.prior_builder.priors['dihedrals'], forceterms=prior_params["forceterms_nn"], bar_position=bar_position) - - # forceterms = ['repulsioncg'] # update them properly in preprocess.py in the _flex class - deltaForcesObj.computePriorForces(forcefield, exclusions=prior_params["exclusions"], - forceterms=prior_params["forceterms_classical"], bar_position=bar_position) - - else: - deltaForcesObj.computePriorForces(forcefield, exclusions=prior_params["exclusions"], - forceterms=prior_params["forceterms"], bar_position=bar_position) - - # load MD forces from forces_npz, compute delta forces, and save them in delta_forces_npz - deltaForcesObj.makeAndSaveDeltaForces(forces_npz, delta_forces_npz, prior_energy_npz) - -def gen_input_mapping(conf): - """Find the list of input files for the passed dataset config""" - pdbid_mapping = dict() - for entry in conf: - input_path = entry["path"] - prefix = entry.get("prefix", "") - suffix = entry.get("suffix", "") - assert os.path.isdir(input_path), f"Input path does not exist: {input_path}" - if "pdbids" in entry: - for dir_name in entry["pdbids"]: - input_h5 = os.path.join(input_path, dir_name, "result", f"output_{dir_name}.h5") - assert os.path.exists(input_h5), "Requested path {input_path}/{dir_name} does not exist" - pdbid_mapping[prefix + dir_name + suffix] = input_h5 - else: - dir_names = os.listdir(input_path) - for dir_name in sorted(dir_names): - input_h5 = os.path.join(input_path, dir_name, "result", f"output_{dir_name}.h5") - if os.path.exists(input_h5): - pdbid_mapping[prefix + dir_name + suffix] = input_h5 - else: - print(f" Skipping \"{dir_name}\" (directory contains no output)") - return pdbid_mapping - -prior_types = { - "CA":Prior_CA, - "CACB":Prior_CACB, - "CACB_lj":Prior_CACB_lj, - "CACB_lj_angle_dihedral":Prior_CACB_lj_angle_dihedral, - "CA_lj":Prior_CA_lj, - "CA_lj_angle":Prior_CA_lj_angle, - "CA_lj_angle_dihedral":Prior_CA_lj_angle_dihedral, - "CA_lj_angle_dihedralX":Prior_CA_lj_angle_dihedralX, - "CA_lj_angleXCX_dihedralX":Prior_CA_lj_angleXCX_dihedralX, - "CA_lj_angleXCX_dihedralX_flex":Prior_CA_lj_angleXCX_dihedralX_flex, - "CA_lj_angleXCX_dihedralX_V1":Prior_CA_lj_angleXCX_dihedralX_V1, - "CA_Majewski2022_v0":Prior_CA_Majewski2022_v0, - "CA_Majewski2022_v1":Prior_CA_Majewski2022_v1, - "CA_lj_bondNull_angleXCX_dihedralX":Prior_CA_lj_bondNull_angleXCX_dihedralX, - "CA_lj_bondNull_angleNull_dihedralX":Prior_CA_lj_bondNull_angleNull_dihedralX, - "CA_lj_bondNull_angleNull_dihedralNull":Prior_CA_lj_bondNull_angleNull_dihedralNull, - "CA_lj_angleNull_dihedralX":Prior_CA_lj_angleNull_dihedralX, - "CA_lj_angleNull_dihedralNull":Prior_CA_lj_angleNull_dihedralNull, - "CA_null":Prior_CA_null, - "CA_lj_only":Prior_CA_lj_only, -} +Full single-file copy (for reference / diff): `preprocess_legacy.py`. +Run from `base_model/` with `module` on PYTHONPATH, same as before: + python preprocess.py ... +or: + python -m preprocess ... +""" +from preprocess.__main__ import main if __name__ == "__main__": - - parser = argparse.ArgumentParser(description="Preprocess data.") - parser.add_argument("input", nargs = "+", help="Input directory path") - parser.add_argument("-o", "--output", required=True, help="Output directory path") - parser.add_argument("--pdbids", nargs="*", help="List of specific PDB IDs to process") - parser.add_argument("--num-frames", "--num_frames", type=int, default=None, help="Number of frames to process") - parser.add_argument("--frame-slice", type=str, default=None, help="Select frames to process using a python slice: start:end:stride") - parser.add_argument("--temp", type=int, default=300, help="Temperature") - parser.add_argument("--prior", type=str, default=None, help="Select the prior forcefield to use, must be one of: " + ", ".join(sorted(prior_types.keys()))) - parser.add_argument("--optimize-forces", action="store_true", help="Use statistically optimal force aggregation (Kramer 2023)") - parser.add_argument("--prior-file", default=None, help="Use PRIOR_FILE instead of fitting a prior") - parser.add_argument('--no-box', default=False, action='store_true', help="Don't use periodic box information") - parser.add_argument('--prior-plots', default=True, action='store_true', help="Save plots of the prior fit functions") - parser.add_argument('--no-prior-plots', dest='prior_plots', action='store_false', help="Don't save plots of the prior fit functions") - parser.add_argument('--no-fit-constraints', default=False, action='store_true', help="Disable range constraints when fitting prior functions") - parser.add_argument('--fit-min-cnt', type=int, default=0, help="Only bins with cnt > min_cnt will be considered when fitting the prior (default 0)") - # parser.add_argument('--tag-beta-turns', default=False, action='store_true', help="Give beta turns a different bond type in the prior") - parser.add_argument('--resume', default=False, action='store_true', help="Resume processing rather than overwriting, all settings must be identical between calls") - parser.add_argument('--num-cores', type=int, default=32, help="Number of cores to be used for parallelization of preprocessing") - parser.add_argument('--jobid', type=int, default=None, help="Integer denoting jobid, if not -1, it will only process a subset of the PDBs") - parser.add_argument('--totalNrJobs', type=int, default=None, help="Integer denoting how many jobs are in total.") - - - args = parser.parse_args() - print(args) - - output_dir = args.output - pdbids = args.pdbids - assert not (args.num_frames and args.frame_slice) - if args.num_frames: - frame_slice = slice(0, args.num_frames) - elif args.frame_slice: - # Convert the arg string into a slice - frame_slice = slice(*[int(i) if i != '' else None for i in args.frame_slice.split(":") ]) - else: - frame_slice = slice(None) - temp = args.temp - optimize_forces = args.optimize_forces - box = not args.no_box - prior_plots = args.prior_plots - prior_name = args.prior - prior_file = args.prior_file - resume_preprocess = args.resume - num_cores = args.num_cores - jobid = args.jobid - totalNrJobs = args.totalNrJobs - - if prior_file: - assert os.path.exists(prior_file), f"Prior file does not exist: {prior_file}" - prior_params_path = get_prior_params_path(prior_file) - with open(prior_params_path, "r", encoding="utf-8") as f: - prior_params = json.load(f) - prior_configuration_name = prior_params["prior_configuration_name"] - if prior_name is None: - prior_name = prior_configuration_name - elif prior_name != prior_configuration_name: - print() - print(f"WARNING: Prior \"{prior_name}\" differs from the one used to build the prior file \"{prior_configuration_name}\"") - print() - - assert prior_name, " You must specify the prior to use with either --prior or --prior-file" - - if prior_name not in prior_types: - raise RuntimeError(f"Unknown prior configuration: {prior_name}") - print(f"Using prior: {prior_name}") - prior_builder = prior_types[prior_name]() # <- () to instantiate the class - prior_builder.enable_fit_constraints(not args.no_fit_constraints) - # prior_builder.enable_bond_tags(args.tag_beta_turns) - prior_builder.enable_bond_tags(False) - prior_builder.set_min_cnt(args.fit_min_cnt) - - if 'external' in prior_builder.prior_params.keys(): - mp.set_start_method('spawn') - - # Set matplotlib to use a thread safe backend (for prior fit plots) - import matplotlib - matplotlib.use('Agg') - - dataset_conf = [] - - for i in args.input: - if os.path.isfile(i): - with open(args.input[0], "r") as f: - dataset_conf += yaml.safe_load(f) - else: - dataset_conf += [{"path": i}] - - input_path_map = gen_input_mapping(dataset_conf) - - if pdbids: - input_path_map = {i: input_path_map[i] for i in pdbids} - - preprocessor = Preprocessor(dataset_conf, input_path_map, output_dir, prior_builder, prior_file, prior_name, frame_slice, temp, optimize_forces, box, prior_plots, resume_preprocess, num_cores, jobid, totalNrJobs) - - preprocessor.preprocess() - + main() diff --git a/preprocess/README.md b/preprocess/README.md new file mode 100644 index 0000000..08a1a2b --- /dev/null +++ b/preprocess/README.md @@ -0,0 +1,54 @@ +# Preprocess Directory + +Refactored preprocessing. `base_model/preprocess.py` only forwards to `preprocess.__main__:main` (same as `python -m preprocess` from `base_model/`). The old single-file script is `preprocess_legacy.py` for reference and diffs. + +## How to run + +From `base_model/` so imports resolve (`preprocess` package and `module`): + +```bash +python preprocess.py -o --prior +# same: +python -m preprocess -o --prior +``` + +Optional YAML (`--config`, e.g. `preprocess/defaults.yaml`): values are applied with `set_defaults` before `parse_args()`, so **any CLI argument you pass overrides** the corresponding YAML default (`config_manager.apply_yaml_defaults_to_argparser`): + +```bash +python -m preprocess --config preprocess/defaults.yaml -o --prior CA_lj +``` + +`input` is one or more directories, or a YAML listing batch paths. + +## How the code is organized + +**Call chain:** + +1. `__main__.py` — argparse, YAML (`config_manager`), inputs (`loaders`), `PriorBuilder`, then `Preprocessor(...).preprocess()`. +2. `runner.py` — `Preprocessor` class holds run state; `preprocess()` calls `pipeline.run_preprocess_pipeline(self)`. +3. `pipeline.py` — ordered stages only: `steps/info.py` → `step1_pool` → `step2_pool` → `job_slice` (optional `--jobid` / `--totalNrJobs`) → `step3_pool` (+ `write_ok_list`). +4. PDB parallelism — `step1_pool` / `step3_pool` use `worker_pool.run_pdb_pool` → `step*_threading` (pool wrapper: resume, tqdm, errors) → `process_step*` (real work). Step 3 pulls `DeltaForces` / worker counts from **`module.*`** (sources under `modules/`). + +**Supporting Files**: + +| Path | Role | +|------|------| +| `paths.py` | `PreprocessPaths`: canonical layout under `-o`. | +| `loaders.py` | H5 batch mapping, trajectory slices. | +| `settings.py` | `PreprocessSettings` (device step 3, resume, cached fits, …). | +| `config_manager.py` / `config_models.py` | YAML merge + CLI wiring. | +| `prior_builder.py` | `PRIOR_TYPES` registry, fit / prior I/O. | +| `trajectory_source.py` | Pluggable pdbid -> path (default H5 batch). | +| `mapping.py` | CG mapping helpers for the pipeline. | + +## Design notes + +- **`Preprocessor`** is the run-scoped “bag” (paths, trajectory, prior, flags, `num_cores`). **`process_step*`** does the work; **`step*_threading`** is the **multiprocessing** pool entrypoint (name is from old file, not special threading). +- **Parallelism:** outer **`mp.Pool` over PDBs**; inner **over frame/chunk workers** inside step 3 (`step3_classical_worker_count` tries to balance the two). **`multiprocessing` is not multi-node / MPI**—use disjoint jobs (e.g. `--jobid` / `--totalNrJobs` for step 3 after global steps 1–2) instead of expecting the pool to coordinate ranks. +- **`PreprocessPaths`**, **`TrajectorySource`**, YAML + CLI: one place for paths, swappable inputs, reproducible defaults. +- **Step 2** is mostly sequential merge + one global prior fit (not a per-PDB process pool); the name **`run_step2_parallel`** is legacy. +- **Tradeoffs:** pool error collection can leave partial outputs; `spawn` vs `fork` by prior type → pickling sensitivity for workers; `prior_builder.pkl` ties resume to class layout. + +## Output layout + +Under `-o`, paths come from `PreprocessPaths`: per-PDB `raw` / `processed` / `fit`; run-level `prior_builder.pkl`, `priors.yaml`, `prior_params.json`; `result/info.json` and `ok_list.txt`. diff --git a/preprocess/__init__.py b/preprocess/__init__.py new file mode 100644 index 0000000..d3391a6 --- /dev/null +++ b/preprocess/__init__.py @@ -0,0 +1,83 @@ +from .config_manager import build_preprocess_settings, load_preprocess_yaml +from .config_models import PreprocessYamlConfig, SettingsSection +from .loaders import ( + BatchGeneratorH5Loader, + gen_input_mapping, + get_prior_params_path, + load_h5_traj_slice, + slice_to_str, +) +from .paths import PreprocessPaths +from .pipeline import run_preprocess_pipeline +from .prior_builder import ( + PRIOR_TYPES, + PriorBuilder, + Prior_CA, + Prior_CA_DNA, + Prior_CA_lj, + Prior_CA_lj_angle, + Prior_CA_lj_angle_dihedral, + Prior_CA_lj_angle_dihedralX, + Prior_CA_lj_angleNull_dihedralNull, + Prior_CA_lj_angleNull_dihedralX, + Prior_CA_lj_angleXCX_dihedralX, + Prior_CA_lj_angleXCX_dihedralX_flex, + Prior_CA_lj_angleXCX_dihedralX_V1, + Prior_CA_lj_bondNull_angleNull_dihedralNull, + Prior_CA_lj_bondNull_angleNull_dihedralX, + Prior_CA_lj_bondNull_angleXCX_dihedralX, + Prior_CA_lj_only, + Prior_CA_Majewski2022_v0, + Prior_CA_Majewski2022_v1, + Prior_CA_null, + Prior_CACB, + Prior_CACB_lj, + Prior_CACB_lj_angle_dihedral, +) +from .runner import Preprocessor +from .settings import PreprocessSettings +from .trajectory_source import H5BatchTrajectorySource, TrajectorySource + +prior_types = PRIOR_TYPES + +__all__ = [ + "PRIOR_TYPES", + "prior_types", + "PriorBuilder", + "Prior_CA", + "Prior_CA_DNA", + "Prior_CACB", + "Prior_CACB_lj", + "Prior_CACB_lj_angle_dihedral", + "Prior_CA_lj", + "Prior_CA_lj_angle", + "Prior_CA_lj_angle_dihedral", + "Prior_CA_lj_angle_dihedralX", + "Prior_CA_lj_angleXCX_dihedralX", + "Prior_CA_lj_angleXCX_dihedralX_flex", + "Prior_CA_lj_angleXCX_dihedralX_V1", + "Prior_CA_Majewski2022_v0", + "Prior_CA_Majewski2022_v1", + "Prior_CA_lj_bondNull_angleXCX_dihedralX", + "Prior_CA_lj_bondNull_angleNull_dihedralX", + "Prior_CA_lj_bondNull_angleNull_dihedralNull", + "Prior_CA_lj_angleNull_dihedralX", + "Prior_CA_lj_angleNull_dihedralNull", + "Prior_CA_null", + "Prior_CA_lj_only", + "Preprocessor", + "PreprocessPaths", + "PreprocessSettings", + "PreprocessYamlConfig", + "SettingsSection", + "TrajectorySource", + "H5BatchTrajectorySource", + "build_preprocess_settings", + "load_preprocess_yaml", + "run_preprocess_pipeline", + "BatchGeneratorH5Loader", + "gen_input_mapping", + "get_prior_params_path", + "load_h5_traj_slice", + "slice_to_str", +] diff --git a/preprocess/__main__.py b/preprocess/__main__.py new file mode 100644 index 0000000..0943d67 --- /dev/null +++ b/preprocess/__main__.py @@ -0,0 +1,183 @@ +"""CLI: `python -m preprocess` from `base_model/` (with `module` on PYTHONPATH).""" + +from __future__ import annotations + +import argparse +import json +import multiprocessing as mp +import os +import sys + +import torch +import yaml + +from .config_manager import apply_yaml_defaults_to_argparser, build_preprocess_settings, load_preprocess_yaml +from .loaders import gen_input_mapping, get_prior_params_path +from .prior_builder import PRIOR_TYPES + + +def _peek_config_path(argv: list[str]) -> str | None: + for i, a in enumerate(argv): + if a == "--config" and i + 1 < len(argv): + return argv[i + 1] + return None + + +def main() -> None: + torch.multiprocessing.set_sharing_strategy("file_system") + + cfg = load_preprocess_yaml(_peek_config_path(sys.argv[1:])) + + parser = argparse.ArgumentParser(description="Preprocess data.") + parser.set_defaults( + filter_not_processed_step_one=False, + use_cached_fits=[], + device_step_3="cpu", + do_step_1=True, + regen_cache_files=True, + resume=False, + prior_plots=True, + no_box=False, + optimize_forces=False, + fit_min_cnt=0, + ) + apply_yaml_defaults_to_argparser(cfg, parser) + + parser.add_argument("--config", help="YAML file with default preprocess settings (CLI overrides file)") + parser.add_argument("input", nargs="+", help="Input directory path") + parser.add_argument("-o", "--output", required=True, help="Output directory path") + parser.add_argument("--pdbids", nargs="*", help="List of specific PDB IDs to process") + parser.add_argument("--num-frames", "--num_frames", type=int, default=None, help="Number of frames to process") + parser.add_argument("--frame-slice", type=str, default=None, help="Select frames using a python slice: start:end:stride") + parser.add_argument("--temp", type=int, help="Temperature") + parser.add_argument( + "--prior", + type=str, + default=None, + help="Select the prior forcefield to use, must be one of: " + ", ".join(sorted(PRIOR_TYPES.keys())), + ) + parser.add_argument("--optimize-forces", action="store_true", help="Use statistically optimal force aggregation (Kramer 2023)") + parser.add_argument("--prior-file", default=None, help="Use PRIOR_FILE instead of fitting a prior") + parser.add_argument("--no-box", action="store_true", help="Don't use periodic box information") + parser.add_argument("--prior-plots", action="store_true", help="Save prior fit plots") + parser.add_argument("--no-prior-plots", dest="prior_plots", action="store_false", help="Do not save prior fit plots") + parser.add_argument("--no-fit-constraints", default=False, action="store_true", help="Disable range constraints when fitting prior functions") + parser.add_argument("--fit-min-cnt", type=int, help="Only bins with cnt > min_cnt when fitting the prior") + parser.add_argument("--resume", action="store_true", help="Resume preprocessing") + parser.add_argument("--no-resume", dest="resume", action="store_false", help="Do not resume") + parser.add_argument("--num-cores", type=int, help="Number of cores for parallel PDB processing") + parser.add_argument("--jobid", type=int, default=None, help="Job id for Step 3 subset (with --totalNrJobs)") + parser.add_argument("--totalNrJobs", type=int, default=None, help="Total array jobs for Step 3 subsetting") + parser.add_argument("--filter-not-processed-step-one", action="store_true", help="Only PDBs with step-1 fit_ok") + parser.add_argument("--skip-step-1", action="store_false", dest="do_step_1", help="Skip step 1 (resume steps 2–3)") + parser.add_argument("--no-regen-cache-files", action="store_false", dest="regen_cache_files", help="Reuse prior_builder.pkl") + parser.add_argument( + "--device-step-3", + type=str, + metavar="DEV", + default=argparse.SUPPRESS, + help="Torch device for step 3 (e.g. cpu, cuda)", + ) + parser.add_argument( + "--use-cached-fits", + nargs="*", + default=argparse.SUPPRESS, + metavar="TERM", + help="Terms to load from cache in step 2 (empty list clears); omit to use YAML default", + ) + + args = parser.parse_args() + print(args) + + output_dir = args.output + pdbids = args.pdbids + assert not (args.num_frames and args.frame_slice) + if args.num_frames: + frame_slice = slice(0, args.num_frames) + elif args.frame_slice: + frame_slice = slice(*[int(i) if i != "" else None for i in args.frame_slice.split(":")]) + else: + frame_slice = slice(None) + temp = args.temp + optimize_forces = args.optimize_forces + box = not args.no_box + prior_plots = args.prior_plots + prior_name = args.prior + prior_file = args.prior_file + resume_preprocess = args.resume + num_cores = args.num_cores + jobid = args.jobid + total_nr_jobs = args.totalNrJobs + + if prior_file: + assert os.path.exists(prior_file), f"Prior file does not exist: {prior_file}" + prior_params_path = get_prior_params_path(prior_file) + with open(prior_params_path, "r", encoding="utf-8") as f: + prior_params = json.load(f) + prior_configuration_name = prior_params["prior_configuration_name"] + if prior_name is None: + prior_name = prior_configuration_name + elif prior_name != prior_configuration_name: + print() + print( + f'WARNING: Prior "{prior_name}" differs from the one used to build the prior file "{prior_configuration_name}"' + ) + print() + + assert prior_name, " You must specify the prior to use with either --prior or --prior-file" + + if prior_name not in PRIOR_TYPES: + raise RuntimeError(f"Unknown prior configuration: {prior_name}") + print(f"Using prior: {prior_name}") + prior_builder = PRIOR_TYPES[prior_name]() + prior_builder.enable_fit_constraints(not args.no_fit_constraints) + prior_builder.enable_bond_tags(False) + prior_builder.set_min_cnt(args.fit_min_cnt) + + if "external" in prior_builder.prior_params.keys(): + mp.set_start_method("spawn") + + import matplotlib + + matplotlib.use("Agg") + + dataset_conf = [] + for i in args.input: + if os.path.isfile(i): + with open(args.input[0], "r") as f: + dataset_conf += yaml.safe_load(f) + else: + dataset_conf += [{"path": i}] + + input_path_map = gen_input_mapping(dataset_conf) + + if pdbids: + input_path_map = {i: input_path_map[i] for i in pdbids} + + from .runner import Preprocessor + + settings = build_preprocess_settings(args, cfg) + preprocessor = Preprocessor( + dataset_conf, + input_path_map, + output_dir, + prior_builder, + prior_file, + prior_name, + frame_slice, + temp, + optimize_forces, + box, + prior_plots, + resume_preprocess, + num_cores, + jobid, + total_nr_jobs, + settings=settings, + ) + + preprocessor.preprocess() + + +if __name__ == "__main__": + main() diff --git a/preprocess/compare_runs.py b/preprocess/compare_runs.py new file mode 100644 index 0000000..b105180 --- /dev/null +++ b/preprocess/compare_runs.py @@ -0,0 +1,172 @@ +"""Compare two preprocess output trees (e.g. pre-refactor vs refactor) for the same pdbs. + +Examples:: + + python -m preprocess.compare_runs /path/to/REF /path/to/NEW --pdbids 1b3t + # Same trajectories/forces/embeddings, ignore priorfit-dependent arrays: + python -m preprocess.compare_runs REF NEW --pdbids 1b3t --geometry-only +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path + +import numpy as np + +# Prior fit + torch step-3: delta/prior_energy differ if priors.yaml differs (cgschnet vs refactor, reruns). +_GEOMETRY_FILES = ("coordinates.npy", "forces.npy", "embeddings.npy", "box.npy") +_RAW_FILES = ( + "deltaforces.npy", + "forces.npy", + "embeddings.npy", + "coordinates.npy", + "prior_energy.npy", + "box.npy", +) + + +def _n_frames_in_info(d: dict) -> int | None: + """int = first n frames; None = all (or unparseable).""" + n = d.get("num_frames") + if n is not None and isinstance(n, (int, float)) and n == n and not isinstance(n, bool): + return int(n) + t = d.get("frame_slice") + if t is None or t == "None" or t == "": + return None + s = str(t) + if s in (":::", "::", ): + return None + m = re.match(r"^0:([0-9]+):", s) + if m: + return int(m.group(1)) + m2 = re.match(r"^:([0-9]+):", s) + if m2: + return int(m2.group(1)) + return None + + +def _is_all_frames_fs(fs) -> bool: + if fs is None or fs == "None" or fs == "": + return True + return str(fs).strip(":") == "" and ":" in str(fs) + + +def _info_frame_match(a: dict, b: dict) -> bool: + """Cgschnet stores `num_frames` only; refactored also stores `frame_slice`.""" + na, nb = _n_frames_in_info(a), _n_frames_in_info(b) + if na is not None or nb is not None: + return na == nb + fa, fb = a.get("frame_slice"), b.get("frame_slice") + if _is_all_frames_fs(fa) and _is_all_frames_fs(fb): + return True + return (fa or "") == (fb or "") + + +def _compare_npy( + a: Path, b: Path, *, rtol: float, atol: float +) -> tuple[str, str]: + if not a.is_file() and not b.is_file(): + return ("skip", "both missing") + if not a.is_file(): + return ("fail", f"ref missing: {a}") + if not b.is_file(): + return ("fail", f"new missing: {b}") + x, y = np.load(a), np.load(b) + if x.shape != y.shape: + return ("fail", f"shape {x.shape} vs {y.shape}") + if x.dtype != y.dtype: + if np.issubdtype(x.dtype, np.floating) and np.issubdtype(y.dtype, np.floating): + pass + else: + return ("fail", f"dtype {x.dtype} vs {y.dtype}") + d = float(np.max(np.abs(x.astype(np.float64) - y.astype(np.float64)))) + if np.allclose(x, y, rtol=rtol, atol=atol): + return ("ok", f"max|diff|={d:.6e} (within rtol/atol)") + return ("fail", f"max_abs_diff={d:.6e} (rtol={rtol} atol={atol})") + + +def compare_roots( + ref_root: Path, + new_root: Path, + pdbids: list[str], + *, + rtol: float = 1e-5, + atol: float = 1e-6, + geometry_only: bool = False, +) -> int: + any_fail = False + files = list(_GEOMETRY_FILES) if geometry_only else list(_RAW_FILES) + if geometry_only: + print( + "compare_runs: geometry-only (coordinates, forces, embeddings, box); " + "skipping deltaforces & prior_energy (depend on fitted priors / step-3).", + flush=True, + ) + for pid in pdbids: + rdir, ndir = ref_root / pid / "raw", new_root / pid / "raw" + print(f"=== {pid} ===", flush=True) + for name in files: + status, msg = _compare_npy(rdir / name, ndir / name, rtol=rtol, atol=atol) + print(f" {name}: {status} — {msg}", flush=True) + if status == "fail": + any_fail = True + ref_info, new_info = ref_root / "result" / "info.json", new_root / "result" / "info.json" + if ref_info.is_file() and new_info.is_file(): + with open(ref_info, encoding="utf-8") as f: + a = json.load(f) + with open(new_info, encoding="utf-8") as f: + b = json.load(f) + for k in ("prior_name", "optimize_forces", "box"): + if a.get(k) != b.get(k): + print(f"WARNING: info.json {k!r} differs: {a.get(k)} vs {b.get(k)}", flush=True) + any_fail = True + if not _info_frame_match(a, b): + print( + f"WARNING: info.json frame selection differs: num_frames/slice " + f"{a.get('num_frames', a.get('frame_slice'))!r} vs {b.get('num_frames', b.get('frame_slice'))!r}", + flush=True, + ) + any_fail = True + if any_fail: + print("compare_runs: FAILED", file=sys.stderr, flush=True) + return 1 + print("compare_runs: all compared arrays match (within tolerances).", flush=True) + return 0 + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Compare preprocess outputs: ref (e.g. 0225_preprocessed) vs new run.", + ) + ap.add_argument("ref_root", type=Path, help="Reference run root (e.g. 0225_preprocessed)") + ap.add_argument("new_root", type=Path, help="New run root to validate") + ap.add_argument( + "--pdbids", type=str, required=True, help="Comma-separated pdbids to compare" + ) + ap.add_argument("--rtol", type=float, default=1e-5) + ap.add_argument("--atol", type=float, default=1e-6) + ap.add_argument( + "--geometry-only", + action="store_true", + help="Only compare coordinates, forces, embeddings, box (excludes delta/prior from prior-fit).", + ) + args = ap.parse_args() + ids = [x.strip() for x in args.pdbids.split(",") if x.strip()] + raise SystemExit( + compare_roots( + args.ref_root.resolve(), + args.new_root.resolve(), + ids, + rtol=args.rtol, + atol=args.atol, + geometry_only=args.geometry_only, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/preprocess/config_manager.py b/preprocess/config_manager.py new file mode 100644 index 0000000..6ffd125 --- /dev/null +++ b/preprocess/config_manager.py @@ -0,0 +1,58 @@ +"""Load merged defaults from optional YAML for CLI argparse.""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from .config_models import PreprocessYamlConfig + + +def load_preprocess_yaml(path: str | None) -> PreprocessYamlConfig: + if not path: + return PreprocessYamlConfig() + p = Path(path) + if not p.is_file(): + raise FileNotFoundError(f"Config file not found: {path}") + with open(p, "r", encoding="utf-8") as f: + raw = yaml.safe_load(f) + if raw is None: + raw = {} + return PreprocessYamlConfig.model_validate(raw) + + +def apply_yaml_defaults_to_argparser(cfg: PreprocessYamlConfig, parser) -> None: + """Call before `parse_args()` so CLI flags override YAML.""" + s = cfg.settings + parser.set_defaults( + num_cores=cfg.num_cores, + resume=cfg.resume, + prior_plots=cfg.prior_plots, + optimize_forces=cfg.optimize_forces, + no_box=cfg.no_box, + temp=cfg.temp, + fit_min_cnt=cfg.fit_min_cnt, + filter_not_processed_step_one=s.filter_not_processed_step_one, + use_cached_fits=list(s.use_cached_fits), + device_step_3=s.device_step_3, + do_step_1=s.do_step_1, + regen_cache_files=s.regen_cache_files, + ) + + +def build_preprocess_settings(args, cfg: PreprocessYamlConfig): + from .settings import PreprocessSettings + + if hasattr(args, "use_cached_fits"): + uc = list(args.use_cached_fits) + else: + uc = list(cfg.settings.use_cached_fits) + dev = getattr(args, "device_step_3", cfg.settings.device_step_3) + return PreprocessSettings( + filter_not_processed_step_one=bool(args.filter_not_processed_step_one), + use_cached_fits=uc, + device_step_3=str(dev), + do_step_1=bool(args.do_step_1), + regen_cache_files=bool(args.regen_cache_files), + ) diff --git a/preprocess/config_models.py b/preprocess/config_models.py new file mode 100644 index 0000000..dd6dd92 --- /dev/null +++ b/preprocess/config_models.py @@ -0,0 +1,30 @@ +"""Pydantic schemas for optional YAML-driven preprocess defaults.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field + + +class SettingsSection(BaseModel): + model_config = ConfigDict(extra="forbid") + + filter_not_processed_step_one: bool = False + use_cached_fits: list[str] = Field(default_factory=list) + device_step_3: str = "cpu" + do_step_1: bool = True + regen_cache_files: bool = True + + +class PreprocessYamlConfig(BaseModel): + """Top-level keys in `preprocess.yaml` (unknown keys ignored).""" + + model_config = ConfigDict(extra="ignore") + + settings: SettingsSection = Field(default_factory=SettingsSection) + num_cores: int = 32 + resume: bool = False + prior_plots: bool = True + optimize_forces: bool = False + no_box: bool = False + temp: int = 300 + fit_min_cnt: int = 0 diff --git a/preprocess/defaults.yaml b/preprocess/defaults.yaml new file mode 100644 index 0000000..12b5180 --- /dev/null +++ b/preprocess/defaults.yaml @@ -0,0 +1,17 @@ +# Optional defaults for `python -m preprocess --config preprocess/defaults.yaml` +# CLI flags override these values. + +settings: + filter_not_processed_step_one: false + use_cached_fits: [] + device_step_3: cpu + do_step_1: true + regen_cache_files: true + +num_cores: 32 +resume: false +prior_plots: true +optimize_forces: false +no_box: false +temp: 300 +fit_min_cnt: 0 diff --git a/preprocess/loaders.py b/preprocess/loaders.py new file mode 100644 index 0000000..0dff42c --- /dev/null +++ b/preprocess/loaders.py @@ -0,0 +1,91 @@ +"""traj discovery and HDF5 slice loading""" + +from __future__ import annotations + +import os + +import h5py +import mdtraj + +from .trajectory_source import H5BatchTrajectorySource + + +def slice_to_str(s: slice) -> str: + result = [s.start, s.stop, s.step] + result = [str(i) if i is not None else "" for i in result] + return ":".join(result) + + +def num_frames_for_info(s: slice | None) -> int | None: + """Cgschnet `result/info.json` uses `num_frames` (int) for `slice(0, N)`; else None for all frames.""" + if s is None or s == slice(None) or (s.start is None and s.stop is None): + return None + st, sp, stp = s.start, s.stop, s.step + if stp not in (None, 1): + return None + if st not in (0, None) or sp is None: + return None + return int(sp) + + +def get_prior_params_path(prior_path: str) -> str: + dir_path, file_name = os.path.split(prior_path) + file_name = file_name.replace("priors.yaml", "prior_params.json") + return os.path.join(dir_path, file_name) + + +def load_h5_traj_slice(path: str, slice_: slice): + """Load a slice from a h5 trajectory without reading the entire file into memory.""" + base_traj = mdtraj.load_frame(path, 0) + with h5py.File(path) as f: + t_xyz = f["coordinates"][slice_][:] # pyright: ignore[reportIndexIssue] + t_time = f["time"][slice_][:] # pyright: ignore[reportIndexIssue] + + t_unitcell_lengths = None + t_unitcell_angles = None + if "cell_lengths" in f.keys(): + t_unitcell_lengths = f["cell_lengths"][slice_][:] # pyright: ignore[reportIndexIssue] + t_unitcell_angles = f["cell_angles"][slice_][:] # pyright: ignore[reportIndexIssue] + + return mdtraj.Trajectory( + t_xyz, + base_traj.topology, + time=t_time, + unitcell_lengths=t_unitcell_lengths, + unitcell_angles=t_unitcell_angles, + ) + + +def gen_input_mapping(conf: list) -> dict: + """Find the list of input files for the passed dataset config.""" + pdbid_mapping: dict = {} + for entry in conf: + input_path = entry["path"] + prefix = entry.get("prefix", "") + suffix = entry.get("suffix", "") + assert os.path.isdir(input_path), f"Input path does not exist: {input_path}" + if "pdbids" in entry: + for dir_name in entry["pdbids"]: + input_h5 = os.path.join(input_path, dir_name, "result", f"output_{dir_name}.h5") + assert os.path.exists(input_h5), "Requested path {input_path}/{dir_name} does not exist" + pdbid_mapping[prefix + dir_name + suffix] = input_h5 + else: + dir_names = os.listdir(input_path) + for dir_name in sorted(dir_names): + input_h5 = os.path.join(input_path, dir_name, "result", f"output_{dir_name}.h5") + if os.path.exists(input_h5): + pdbid_mapping[prefix + dir_name + suffix] = input_h5 + else: + print(f' Skipping "{dir_name}" (directory contains no output)') + return pdbid_mapping + + +class BatchGeneratorH5Loader: + """`TrajectorySource` built from the same YAML/directory discovery as `gen_input_mapping`.""" + + def __init__(self, dataset_conf: list): + self._source = H5BatchTrajectorySource(gen_input_mapping(dataset_conf)) + + @property + def source(self) -> H5BatchTrajectorySource: + return self._source diff --git a/preprocess/mapping.py b/preprocess/mapping.py new file mode 100644 index 0000000..80c569a --- /dev/null +++ b/preprocess/mapping.py @@ -0,0 +1,138 @@ +class CGMappingDef_CA: + def __init__(self): + residues = ["ALA", "CYS", "ASP", "GLU", "PHE", "GLY", "HIS", "ILE", "LYS", "LEU", "MET", "ASN", "PRO", "HYP", "GLN", "ARG", "SER", "THR", "VAL", "TRP", "TYR"] + # For legacy reasons we have a couple extra ambiguous residues (ASX & GLX) in the embedding map but we do not accept these for parsing + embedding_residues = ["ALA", "ARG", "ASN", "ASP", "ASX", "CYS", "GLU", "GLN", "GLX", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"] + self.bead_embeddings = {name: [index + 1] for index, name in enumerate(sorted(embedding_residues))} + + # bead_atom_selection: A list of lists, where each inner list is the names of the atoms that will be combined to form the bead + self.bead_atom_selection = {k: [["CA"]] for k in residues} + # The type names of beads (will become the atom type/element in the cg topology) + self.bead_types = { + "ALA": ["CAA"], + "ARG": ["CAR"], + "ASN": ["CAN"], + "ASP": ["CAD"], + "CYS": ["CAC"], + "GLN": ["CAQ"], + "GLU": ["CAE"], + "GLY": ["CAG"], + "HIS": ["CAH"], + "HSD": ["CAH"], + "ILE": ["CAI"], + "LEU": ["CAL"], + "LYS": ["CAK"], + "MET": ["CAM"], + "PHE": ["CAF"], + "PRO": ["CAP"], + "SER": ["CAS"], + "THR": ["CAT"], + "TRP": ["CAW"], + "TYR": ["CAY"], + "VAL": ["CAV"], + } + # The "atom name" assigned to the beads + self.bead_atom_names = {k: ["CA"] for k in residues} + self.bead_masses = {k: [12.01] for k in residues} + self.bead_backbone_idx = {k: 0 for k in residues} + +class CGMappingDef_CACB: + def __init__(self): + residues = ["ALA", "CYS", "ASP", "GLU", "PHE", "GLY", "HIS", "ILE", "LYS", "LEU", "MET", "ASN", "PRO", "HYP", "GLN", "ARG", "SER", "THR", "VAL", "TRP", "TYR"] + + # bead_atom_selection: A list of lists, where each inner list is the names of the atoms that will be combined to form the bead + self.bead_atom_selection = {k: [["CA"], ["CB"]] for k in residues} + self.bead_atom_selection["GLY"] = [["CA"]] + # The type names of beads (will become the atom type/element in the cg topology) + self.bead_types = { + "ALA": ["CA", "CBA"], + "ARG": ["CA", "CBR"], + "ASN": ["CA", "CBN"], + "ASP": ["CA", "CBD"], + "CYS": ["CA", "CBC"], + "GLN": ["CA", "CBQ"], + "GLU": ["CA", "CBE"], + "GLY": ["CAG"], + "HIS": ["CA", "CBH"], + "HSD": ["CA", "CBH"], + "ILE": ["CA", "CBI"], + "LEU": ["CA", "CBL"], + "LYS": ["CA", "CBK"], + "MET": ["CA", "CBM"], + "PHE": ["CA", "CBF"], + "PRO": ["CA", "CBP"], + "SER": ["CA", "CBS"], + "THR": ["CA", "CBT"], + "TRP": ["CA", "CBW"], + "TYR": ["CA", "CBY"], + "VAL": ["CA", "CBV"], + } + + embedding_map = {k:i for i,k in enumerate(sorted(set.union(*[set(i) for i in self.bead_types.values()])))} + self.bead_embeddings = {k:[embedding_map[i] for i in v] for k, v in self.bead_types.items()} + + # The "atom name" assigned to the beads + self.bead_atom_names = {k: ["CA", "CB"] for k in residues} + self.bead_atom_names["GLY"] = ["CA"] + self.bead_masses = {k: [12.01]*len(v) for k,v in self.bead_types.items()} + self.bead_backbone_idx = {k: 0 for k in residues} + + +class CGMappingDef_CA_DNA(CGMappingDef_CA): + """Protein Cα plus coarse DNA (2 beads/residue: backbone + base) — matches cgschnet `Prior_CA_DNA` mapping.""" + + def __init__(self) -> None: + super().__init__() + dna_residues = ["DA", "DT", "DG", "DC"] + backbone_atoms = [ + "P", + "OP1", + "OP2", + "O5'", + "C5'", + "C4'", + "C3'", + "O3'", + "C1'", + "C2'", + "O4'", + ] + # 5' or fragment models without phosphate: COM of sugar + backbone linkage only (same order as cgschnet, minus P/OP*) + backbone_atoms_no_phosphate = [ + "O5'", + "C5'", + "C4'", + "C3'", + "O3'", + "C1'", + "C2'", + "O4'", + ] + base_atoms = { + "DA": ["N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], + "DT": ["N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6", "C7"], + "DG": ["N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], + "DC": ["N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], + } + for resname in dna_residues: + self.bead_atom_selection[resname] = [backbone_atoms, base_atoms[resname]] + # Tried in order for the first (DBB) bead; see `module.cg_mapping.CGMapping` + self.dna_backbone_atom_candidates = (backbone_atoms, backbone_atoms_no_phosphate) + for resname in dna_residues: + base_code = resname[1] + self.bead_atom_names[resname] = ["DBB", f"DB{base_code}"] + for resname in dna_residues: + base_code = resname[1] + self.bead_types[resname] = ["DBB", f"DB{base_code}"] + base_masses = {"A": 134.1, "T": 125.1, "G": 150.1, "C": 110.1} + for resname in dna_residues: + base_code = resname[1] + self.bead_masses[resname] = [178.08, base_masses[base_code]] + current_max_id = max(self.bead_embeddings.values())[0] + dbb_id = current_max_id + 1 + base_ids = {"A": dbb_id + 1, "T": dbb_id + 2, "G": dbb_id + 3, "C": dbb_id + 4} + for resname in dna_residues: + base_code = resname[1] + self.bead_embeddings[resname] = [dbb_id, base_ids[base_code]] + for resname in dna_residues: + self.bead_backbone_idx[resname] = 0 diff --git a/preprocess/paths.py b/preprocess/paths.py new file mode 100644 index 0000000..326d295 --- /dev/null +++ b/preprocess/paths.py @@ -0,0 +1,54 @@ +"""Filesystem layout for a single preprocess output directory.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class PreprocessPaths: + """All paths under the preprocess run root (`save_path` / `-o` output).""" + + root: Path + + @property + def result_dir(self) -> Path: + return self.root / "result" + + def info_json(self) -> Path: + return self.result_dir / "info.json" + + def ok_list_txt(self) -> Path: + return self.result_dir / "ok_list.txt" + + def pdb_dir(self, pdbid: str) -> Path: + return self.root / pdbid + + def pdb_raw(self, pdbid: str) -> Path: + return self.pdb_dir(pdbid) / "raw" + + def pdb_processed(self, pdbid: str) -> Path: + return self.pdb_dir(pdbid) / "processed" + + def pdb_fit(self, pdbid: str) -> Path: + return self.pdb_dir(pdbid) / "fit" + + def fit_ok(self, pdbid: str) -> Path: + return self.pdb_fit(pdbid) / "fit_ok.txt" + + def prior_builder_pkl(self) -> Path: + return self.root / "prior_builder.pkl" + + def prior_fit_plots_dir(self) -> Path: + return self.root / "prior_fit_plots" + + def priors_yaml(self) -> Path: + return self.root / "priors.yaml" + + def prior_params_json(self) -> Path: + return self.root / "prior_params.json" + + def glob_step1_fit_ok(self): + """Paths `root//fit/fit_ok.txt` (sorted).""" + return sorted(self.root.glob("*/fit/fit_ok.txt")) diff --git a/preprocess/pipeline.py b/preprocess/pipeline.py new file mode 100644 index 0000000..c083603 --- /dev/null +++ b/preprocess/pipeline.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .steps.info import write_result_metadata +from .steps.job_slice import apply_job_slice +from .steps.step1_pool import run_step1_parallel +from .steps.step2_pool import run_step2_parallel +from .steps.step3_pool import run_step3_parallel, write_ok_list + +if TYPE_CHECKING: + from .runner import Preprocessor + + +def run_preprocess_pipeline(pre: Preprocessor) -> None: + write_result_metadata(pre) + run_step1_parallel(pre) + run_step2_parallel(pre) + apply_job_slice(pre) + run_step3_parallel(pre) + write_ok_list(pre) + print("Done!") diff --git a/preprocess/prior_builder.py b/preprocess/prior_builder.py new file mode 100644 index 0000000..3f8fd48 --- /dev/null +++ b/preprocess/prior_builder.py @@ -0,0 +1,774 @@ +import json +import os +import pickle + +import numpy as np +import yaml + +from module import prior +from module import prior_flex +from module import psfwriter +from module.cg_mapping import CGMapping +from module.torchmd_cg_mappings import CACB_MAP + +from .mapping import CGMappingDef_CA, CGMappingDef_CACB, CGMappingDef_CA_DNA + +class PriorBuilder: + def __init__(self): + self.prior_params = dict() + self.priors = None + self.terms = dict() + self.atom_types = set() + self.fit_constraints = True + self.tag_beta_turns = False + self.min_cnt = 0 + + def select_atoms(self, topology): + """Returns tha atom index to be saved for this prior""" + raise NotImplementedError() + + def map_embeddings(self, selected_atoms, trajectory): + """Generates the embeddings array for the selected atoms""" + raise NotImplementedError() + + def write_psf(self, pdb_file, psf_file): + """Write the .psf file describing the course grain geometry""" + raise NotImplementedError() + + def add_molecule(self, mol, traj, cache_dir): + fit_ok_path = os.path.join(cache_dir, "fit_ok.txt") + + if cache_dir and os.path.exists(fit_ok_path): + os.unlink(fit_ok_path) + + for term in self.terms.values(): + term.add_molecule(mol, traj, cache_dir) + self.atom_types = self.atom_types.union(mol.atomtype) + + if cache_dir: + np.save(os.path.join(cache_dir, "atomtype.npy"), mol.atomtype) + with open(fit_ok_path, "wt", encoding="utf-8") as f: + f.write("ok") + + def load_molecule_cache(self, cache_dir): + assert os.path.exists(os.path.join(cache_dir, "fit_ok.txt")) + atomtype = np.load(os.path.join(cache_dir, "atomtype.npy"), allow_pickle=True) + self.atom_types = self.atom_types.union(atomtype) + + for term in self.terms.values(): + term.load_molecule_cache(cache_dir) + + def enable_fit_constraints(self, use_constraints): + self.fit_constraints = use_constraints + self.prior_params["fit_constraints"] = self.fit_constraints + + def enable_bond_tags(self, use_tags): + self.tag_beta_turns = use_tags + self.prior_params["tag_beta_turns"] = self.tag_beta_turns + + def set_min_cnt(self, min_cnt): + assert min_cnt >= 0 + self.min_cnt = min_cnt + self.prior_params["min_cnt"] = self.min_cnt + + def fit(self, temperature, plot_dir=None, use_cached_fits=None): + if use_cached_fits is None: + use_cached_fits = [] + self.init_prior_dict() + assert self.priors is not None + for key, term in self.terms.items(): + cache_pkl = os.path.join(plot_dir, f"prior_{key}.pkl") if plot_dir else None + if cache_pkl and os.path.exists(cache_pkl) and (key in use_cached_fits): + print(f"Used cached fit for {key}...") + with open(cache_pkl, "rb") as f: + self.priors[key] = pickle.load(f) + else: + print(f"Fitting {key}...") + self.priors[key] = term.get_param(temperature, plot_dir, self.fit_constraints, self.min_cnt) + if plot_dir: + with open(os.path.join(plot_dir, f"prior_{key}.pkl"), "wb") as f: + pickle.dump(self.priors[key], f) + + def init_prior_dict(self): + # Define the force field dict + priors = {} + priors['atomtypes'] = sorted(self.atom_types) + priors['bonds'] = {} + priors['angles'] = {} + priors['dihedrals'] = {} + priors['lj'] = {} + # For mass and charge assume everything is a carbon atom + priors['electrostatics'] = {at: {'charge': 0.0} for at in priors['atomtypes']} + # The mass of carbon used here is the from OpenMM/AMBER-14 value + priors['masses'] = {at: 12.01 for at in priors['atomtypes']} + self.priors = priors + + def save_prior(self, output_path, pdbid): + prefix = "" + if pdbid: + prefix = f"{pdbid}_" + with open(os.path.join(output_path, f"{prefix}priors.yaml"), "w") as f: + yaml.dump(self.priors, f) + with open(os.path.join(output_path, f"{prefix}prior_params.json"),"w") as f: + json.dump(self.prior_params, f) + + def make_mol(self, cg_map): + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return cg_map.to_mol(bonds = bonds, angles = angles, dihedrals = dihedrals) + +class Prior_CA(PriorBuilder): + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA", + "exclusions" : ['bonds'], + "forceterms" : ["bonds"], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + + def build_mapping(self, topology): + return CGMapping(topology, CGMappingDef_CA()) + + def select_atoms(self, topology): + #TODO: Remove this function (replaced by build_mapping) + return topology.select('name CA and protein') + + def map_embeddings(self, selected_atoms, topology): #pyright: ignore[reportIncompatibleMethodOverride] + #TODO: Remove this function (replaced by build_mapping) + standardResidues = {"ALA", "ARG", "ASN", "ASP", "ASX", "CYS", "GLU", "GLN", "GLX", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"} + amino_acid_mapping = {name: index + 1 for index, name in enumerate(sorted(standardResidues))} + + result = [] + for a_idx in selected_atoms: + r_name = topology.atom(a_idx).residue.name + result.append(amino_acid_mapping[r_name]) + return np.array(result, dtype=int) + + def write_psf(self, pdb_file, psf_file): + #TODO: Remove this function (replaced by build_mapping) + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return psfwriter.pdb2psf_CA(pdb_file, psf_file, bonds = bonds, angles = angles, dihedrals = dihedrals, + tag_beta_turns = self.tag_beta_turns) + +class Prior_CACB(PriorBuilder): + """Implements the torchmd-cg CACB prior""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CACB", + "exclusions" : ['bonds'], + "forceterms" : ["bonds"], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + + def build_mapping(self, topology): + return CGMapping(topology, CGMappingDef_CACB()) + + def select_atoms(self, topology): + #TODO: Remove this function (replaced by build_mapping) + return topology.select('(name CA or name CB) and protein') + + def map_embeddings(self, selected_atoms, topology):#pyright: ignore[reportIncompatibleMethodOverride] + #TODO: Remove this function (replaced by build_mapping) + + # Make a map from embedding name to embedding name number + # e.g. {"CAA":0, "CAC":1, ...} + embedding_map = CACB_MAP + embedding_nums = dict([(k, i) for i, k in enumerate(sorted(set(embedding_map.values())))]) + + result = [] + for a_idx in selected_atoms: + r_name = topology.atom(a_idx).residue.name + a_name = topology.atom(a_idx).name + emb_name = embedding_map[(r_name, a_name)] + result.append(embedding_nums[emb_name]) + return np.array(result, dtype=int) + + def write_psf(self, pdb_file, psf_file): + #TODO: Remove this function (replaced by build_mapping) + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return psfwriter.pdb2psf_CACB(pdb_file, psf_file, bonds = bonds, angles = angles, dihedrals = dihedrals) + +class Prior_CACB_lj(Prior_CACB): + """torchmd-cg CACB prior with Bonded & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CACB_lj", + "exclusions" : ['bonds'], + "forceterms" : ['bonds', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + +class Prior_CACB_lj_angle_dihedral(Prior_CACB): + """torchmd-cg CACB prior with Bonded, Angle, Dihedral & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CACB_lj_angle_dihedral", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator() + +class Prior_CA_lj(Prior_CA): + """CA prior with Bonded & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj", + "exclusions" : ['bonds'], + "forceterms" : ['bonds', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + +class Prior_CA_lj_angle(Prior_CA): + """CA prior with Bonded, Angle, and RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angle", + "exclusions" : ['bonds', 'angles'], + "forceterms" : ['bonds', 'angles', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms['angles'] = prior.ParamAngleCalculator() + +class Prior_CA_lj_angle_dihedral(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, Dihedral & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angle_dihedral", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator() + +class Prior_CA_lj_angle_dihedralX(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angle_dihedralX", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_angleXCX_dihedralX(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleXCX_dihedralX", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator(center=True) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_angleXCX_dihedralX_flex(Prior_CA): + """torchmd-cg CA prior with highly flexible Bonded, Angle, DihedralX & RepulsionCG terms that fit the data. + + """ + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleXCX_dihedralX_flex", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms_nn" : ['bonds', 'angles', 'dihedrals'], + "forceterms_classical": ['repulsioncg'], # changed from lj, would need to re-generated the dataset (Jan 10 2025). repulsioncg is using just the repulsion term from lj. it uses the same parameters as lj, so need to make sure the right function is evaluated. + "external" : True + }) + self.prior_params['forceterms'] = self.prior_params['forceterms_classical'] + self.prior_params['forceterms_nn'] + + self.terms["bonds"] = prior_flex.ParamBondedFlexCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior_flex.ParamAngleFlexCalculator(center=True) + self.terms["dihedrals"] = prior_flex.ParamDihedralFlexCalculator(unified=True) + + # have to override this method since we're saving neural nets as priors + def save_prior(self, output_path, pdbid): + prefix = "" + # if pdbid: + # prefix = f"{pdbid}_" + with open(os.path.join(output_path, f"{prefix}prior_params.json"),"w") as f: + json.dump(self.prior_params, f) + + # print('self.priors', self.priors.keys()) + # remove the dihedrals and bonds from the priors + priorsTruncated = self.priors.copy() + priorsTruncated.pop('dihedrals') + priorsTruncated.pop('bonds') + priorsTruncated.pop('angles') + # print('priorsTruncated', priorsTruncated.keys()) + + # save the classical priors using yaml. this is requires because the classical priors are built from the yaml files + with open(os.path.join(output_path, f"{prefix}priors.yaml"), "w") as f: + yaml.dump(priorsTruncated, f) + + self.priors['terms'] = self.terms + self.priors['prior_params'] = self.prior_params + + # also save with pickle + with open(os.path.join(output_path, f"{prefix}priors.pkl"), "wb") as f: + pickle.dump(self.priors, f) + + + def load_prior_nnets(self, output_path): + # load the prior with pickle + with open(os.path.join(output_path, "priors.pkl"), "rb") as f: + self.priors = pickle.load(f) + + # return self.priors + + # with open(os.path.join(output_path, f"{prefix}priors.pkl"), "wb") as f: + # pickle.dump(self.priors, f) + + + +class Prior_CA_lj_angleXCX_dihedralX_V1(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleXCX_dihedralX_V1", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator(center=True) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_bondNull_angleXCX_dihedralX(Prior_CA): + """torchmd-cg CA prior with Angle, DihedralX & RepulsionCG terms (+ bond exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_bondNull_angleXCX_dihedralX", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.NullParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator(center=True) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_bondNull_angleNull_dihedralX(Prior_CA): + """torchmd-cg CA prior with DihedralX & RepulsionCG terms (+ bond & angle exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_bondNull_angleNull_dihedralX", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.NullParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_bondNull_angleNull_dihedralNull(Prior_CA): + """torchmd-cg CA prior with RepulsionCG terms (+ bond, angle, & dihedral exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_bondNull_angleNull_dihedralNull", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.NullParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.NullParamDihedralCalculator() + +class Prior_CA_lj_angleNull_dihedralX(Prior_CA): + """torchmd-cg CA prior with Bonded, DihedralX & RepulsionCG terms (+ angle exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleNull_dihedralX", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_angleNull_dihedralNull(Prior_CA): + """torchmd-cg CA prior with Bonded & RepulsionCG terms (+ angle & dihedral exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleNull_dihedralNull", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.NullParamDihedralCalculator() + +class Prior_CA_Majewski2022_v0(Prior_CA): + """torchmd-cg CA prior based on the parameters used in (Majewski 2022) + Note this version (v0) has different lj exclusions than the one used in the paper. + """ + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_Majewski2022_v0", + "exclusions" : ['bonds', 'dihedrals'], + "forceterms" : ['bonds', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True, scale=0.5) + +class Prior_CA_Majewski2022_v1(Prior_CA): + """torchmd-cg CA prior based on the parameters used in (Majewski 2022)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_Majewski2022_v1", + "exclusions" : ['bonds'], + "forceterms" : ['bonds', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6], exclusion_terms={"bonds"}) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True, scale=0.5) + +class Prior_CA_null(Prior_CA): + """CA prior with no terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_null", + "exclusions" : [], + "forceterms" : [], + }) + self.terms = {} + +class Prior_CA_lj_only(Prior_CA): + """CA prior with just a RepulsionCG term""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_only", + "exclusions" : [], + "forceterms" : ['RepulsionCG'], + }) + self.terms = {} + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + + +class Prior_CA_DNA(Prior_CA): + """Protein Cα + DNA 2-bead model; after bonded/angle fit, applies universal DNA stiffness (cgschnet).""" + + def __init__(self) -> None: + super().__init__() + self.prior_params.update( + { + "prior_configuration_name": "CA_DNA", + "exclusions": ["1-2", "1-3"], + "forceterms": ["bonds", "angles"], + } + ) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["angles"] = prior.ParamAngleCalculator(center=False) + self.cg_mapping_def: CGMappingDef_CA_DNA = CGMappingDef_CA_DNA() + + def init_prior_dict(self) -> None: + super().init_prior_dict() + bead_type_to_mass: dict[str, float] = {} + for resname in self.cg_mapping_def.bead_types: + bead_types = self.cg_mapping_def.bead_types[resname] + bead_masses = self.cg_mapping_def.bead_masses.get(resname, [12.01] * len(bead_types)) + for bt, bm in zip(bead_types, bead_masses): + bead_type_to_mass[bt] = float(bm) + for atom_type in self.priors["atomtypes"]: + if atom_type in bead_type_to_mass: + self.priors["masses"][atom_type] = bead_type_to_mass[atom_type] + + def fit( + self, + temperature: float, + plot_dir: str | None = None, + use_cached_fits: list | None = None, + ) -> None: + if use_cached_fits is None: + use_cached_fits = [] + print("[DEBUG] Starting fit for Prior_CA_DNA") + super().fit(temperature, plot_dir, use_cached_fits) + self._apply_universal_cg_parameters() + self._ensure_all_dna_params_exist() + self._validate_physics() + print("[DEBUG] Applied universal CG parameters for mixed protein–DNA systems") + + def _apply_universal_cg_parameters(self) -> None: + UNIVERSAL = { + "DNA_BOND_K": 50.0, + "DNA_ANGLE_K": 30.0, + "PROTEIN_BOND_K": 10.0, + "PROTEIN_ANGLE_K": 2.0, + "DNA_BACKBONE_LENGTH": 5.5, + "DNA_BASE_DISTANCE": 5.75, + "DNA_BACKBONE_ANGLE": 170.0, + } + if "bonds" in self.priors: + for bond_type, params in self.priors["bonds"].items(): + if isinstance(params, dict): + current_k = float(params.get("k0", 1.0)) + current_r0 = float(params.get("req", 3.8)) + elif isinstance(params, (list, tuple)) and len(params) >= 2: + current_k = float(params[0]) + current_r0 = float(params[1]) + else: + print(f"[WARNING] Bond {bond_type} has unexpected format: {params}") + continue + dna_token = any(t in bond_type for t in ("DBB", "DBA", "DBT", "DBG", "DBC")) + if dna_token: + new_k = UNIVERSAL["DNA_BOND_K"] + if "DBB-DBB" in bond_type: + new_r0 = UNIVERSAL["DNA_BACKBONE_LENGTH"] + elif "DBB" in bond_type: + new_r0 = UNIVERSAL["DNA_BASE_DISTANCE"] + else: + new_r0 = current_r0 + self.priors["bonds"][bond_type] = {"k0": float(new_k), "req": float(new_r0)} + elif "CA" in bond_type and current_k < 5.0: + if isinstance(self.priors["bonds"][bond_type], dict): + self.priors["bonds"][bond_type]["k0"] = UNIVERSAL["PROTEIN_BOND_K"] + else: + self.priors["bonds"][bond_type][0] = UNIVERSAL["PROTEIN_BOND_K"] + + if "angles" in self.priors: + for angle_type, value in self.priors["angles"].items(): + if not isinstance(value, dict): + continue + try: + current_k0 = float(value.get("k0", 1.0)) + except (ValueError, TypeError): + continue + dna_kw = ("DBA", "DBB", "DBT", "DBG", "DBC") + if any(kw in angle_type for kw in dna_kw): + if "DBB-DBB-DBB" in angle_type: + value["theta0"] = UNIVERSAL["DNA_BACKBONE_ANGLE"] + value["k0"] = UNIVERSAL["DNA_ANGLE_K"] + else: + value["k0"] = UNIVERSAL["DNA_ANGLE_K"] + elif "CA" in angle_type and current_k0 < 0.5: + value["k0"] = UNIVERSAL["PROTEIN_ANGLE_K"] + + def _has_dna_beads(self) -> bool: + """True if any loaded molecule has DNA CG atom types (DB*), not just protein Cα.""" + for at in self.atom_types: + s = str(at) + if s.startswith("DB") or s.startswith("db"): + return True + return False + + def _ensure_all_dna_params_exist(self) -> None: + if not self._has_dna_beads(): + return + essential_dna_bonds = { + "DBB-DBB": [50.0, 5.5], + "DBB-DBA": [50.0, 5.75], + "DBB-DBT": [50.0, 5.75], + "DBB-DBG": [50.0, 5.75], + "DBB-DBC": [50.0, 5.75], + } + essential_dna_angles = { + "DBB-DBB-DBB": {"k0": 30.0, "theta0": 170.0}, + "DBA-DBB-DBB": {"k0": 25.0, "theta0": 50.0}, + "DBB-DBB-DBC": {"k0": 25.0, "theta0": 83.0}, + "DBB-DBB-DBG": {"k0": 25.0, "theta0": 78.0}, + "DBB-DBB-DBT": {"k0": 25.0, "theta0": 55.0}, + } + for bond_type, params in essential_dna_bonds.items(): + if bond_type not in self.priors.get("bonds", {}): + self.priors.setdefault("bonds", {})[bond_type] = params + for angle_type, params in essential_dna_angles.items(): + if angle_type not in self.priors.get("angles", {}): + self.priors.setdefault("angles", {})[angle_type] = params + + def _validate_physics(self) -> None: + print("\n[PHYSICS VALIDATION]") + print("=" * 50) + dna_stiffness_ok = True + if not self._has_dna_beads(): + print(" (no DNA CG beads in this batch — skip DNA stiffness template checks)\n" + "=" * 50) + return + for bond_type, params in self.priors.get("bonds", {}).items(): + if not any(t in bond_type for t in ("DBB", "DBA", "DBT", "DBG", "DBC")): + continue + try: + if isinstance(params, dict): + k = float(params.get("k0", 0)) + r0 = float(params.get("req", 0)) + elif isinstance(params, (list, tuple)) and len(params) >= 2: + k, r0 = float(params[0]), float(params[1]) + else: + print(f" DNA bond {bond_type}: Unrecognized format {params}") + continue + if k < 20.0: + print(f" DNA bond {bond_type}: k={k:.1f} (might be too soft)") + dna_stiffness_ok = False + if r0 < 4.0 or r0 > 7.0: + print(f" DNA bond {bond_type}: r0={r0:.1f}Å (unusual)") + except (ValueError, TypeError, IndexError): + print(f"Could not parse DNA bond {bond_type}: {params}") + avg_dna_k = 0.0 + avg_p_k = 0.0 + dna_c = 0 + p_c = 0 + for bond_type, params in self.priors.get("bonds", {}).items(): + try: + if isinstance(params, dict): + k = float(params.get("k0", 0)) + elif isinstance(params, (list, tuple)) and len(params) >= 2: + k = float(params[0]) + else: + continue + except (ValueError, TypeError, IndexError): + continue + if any(t in bond_type for t in ("DBB", "DBA", "DBT", "DBG", "DBC")): + avg_dna_k += k + dna_c += 1 + elif "CA" in bond_type: + avg_p_k += k + p_c += 1 + if dna_c > 0 and p_c > 0: + avg_dna_k /= dna_c + avg_p_k /= p_c + ratio = avg_dna_k / avg_p_k if avg_p_k > 0 else 999.0 + print(f" Stiffness ratio (DNA:Protein) = {ratio:.1f}:1") + print(f" Average DNA k = {avg_dna_k:.1f}, Protein k = {avg_p_k:.1f}") + if ratio < 2.0: + print("DNA not stiff enough relative to protein!") + print("\n[ANGLES VALIDATION]") + dna_ang = 0 + tot = 0.0 + for angle_type, value in self.priors.get("angles", {}).items(): + if not isinstance(value, dict): + continue + if not any(t in angle_type for t in ("DBA", "DBB", "DBT", "DBG", "DBC")): + continue + try: + k0 = float(value.get("k0", 0)) + tot += k0 + dna_ang += 1 + if "DBB-DBB-DBB" in angle_type: + t0 = float(value.get("theta0", 0)) + if t0 < 150 or t0 > 190: + print( + f"DNA backbone angle {angle_type}: theta0={t0:.1f}° (should be ~170°)" + ) + except (ValueError, TypeError): + pass + if dna_ang > 0: + print(f" Average DNA angle k0 = {tot / dna_ang:.1f}") + if tot / dna_ang < 10.0: + print("DNA angles might be too soft!") + if dna_stiffness_ok: + print("\n DNA bond parameters look reasonable") + else: + print("\n DNA bond parameters need attention") + print("=" * 50) + + def build_mapping(self, topology): + return CGMapping(topology, self.cg_mapping_def) + + def map_embeddings(self, selected_atoms, topology): # pyright: ignore[reportIncompatibleMethodOverride] + standard = { + "ALA", + "ARG", + "ASN", + "ASP", + "CYS", + "GLU", + "GLN", + "GLY", + "HIS", + "ILE", + "LEU", + "LYS", + "MET", + "PHE", + "PRO", + "SER", + "THR", + "TRP", + "TYR", + "VAL", + } + dna = {"DA", "DT", "DG", "DC"} + all_r = sorted(standard | dna) + residue_map = {res: i + 1 for i, res in enumerate(all_r)} + result: list[int] = [] + for a_idx in selected_atoms: + r_name = "".join(filter(str.isalpha, topology.atom(a_idx).residue.name)) + if r_name not in residue_map: + print(f"[WARNING] Residue {r_name!r} not in mapping!") + result.append(residue_map.get(r_name, 0)) + return np.array(result, dtype=int) + + def write_psf(self, pdb_file, psf_file): + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return psfwriter.pdb2psf_CA( + pdb_file, + psf_file, + bonds=bonds, + angles=angles, + dihedrals=dihedrals, + tag_beta_turns=self.tag_beta_turns, + ) + + +PRIOR_TYPES = { + "CA": Prior_CA, + "CA_DNA": Prior_CA_DNA, + "CACB": Prior_CACB, + "CACB_lj": Prior_CACB_lj, + "CACB_lj_angle_dihedral": Prior_CACB_lj_angle_dihedral, + "CA_lj": Prior_CA_lj, + "CA_lj_angle": Prior_CA_lj_angle, + "CA_lj_angle_dihedral": Prior_CA_lj_angle_dihedral, + "CA_lj_angle_dihedralX": Prior_CA_lj_angle_dihedralX, + "CA_lj_angleXCX_dihedralX": Prior_CA_lj_angleXCX_dihedralX, + "CA_lj_angleXCX_dihedralX_flex": Prior_CA_lj_angleXCX_dihedralX_flex, + "CA_lj_angleXCX_dihedralX_V1": Prior_CA_lj_angleXCX_dihedralX_V1, + "CA_Majewski2022_v0": Prior_CA_Majewski2022_v0, + "CA_Majewski2022_v1": Prior_CA_Majewski2022_v1, + "CA_lj_bondNull_angleXCX_dihedralX": Prior_CA_lj_bondNull_angleXCX_dihedralX, + "CA_lj_bondNull_angleNull_dihedralX": Prior_CA_lj_bondNull_angleNull_dihedralX, + "CA_lj_bondNull_angleNull_dihedralNull": Prior_CA_lj_bondNull_angleNull_dihedralNull, + "CA_lj_angleNull_dihedralX": Prior_CA_lj_angleNull_dihedralX, + "CA_lj_angleNull_dihedralNull": Prior_CA_lj_angleNull_dihedralNull, + "CA_null": Prior_CA_null, + "CA_lj_only": Prior_CA_lj_only, +} diff --git a/preprocess/runner.py b/preprocess/runner.py new file mode 100644 index 0000000..1379e22 --- /dev/null +++ b/preprocess/runner.py @@ -0,0 +1,337 @@ +from __future__ import annotations + +import multiprocessing as mp +import os +import traceback +from collections.abc import Mapping +from pathlib import Path +from typing import Callable + +import h5py +import mdtraj +import numpy as np +from tqdm import tqdm + +from module.frame_utils import step3_classical_worker_count +from module.make_deltaforces import DeltaForces + +from .loaders import load_h5_traj_slice, slice_to_str +from .paths import PreprocessPaths +from .settings import PreprocessSettings +from .trajectory_source import H5BatchTrajectorySource, TrajectorySource + +_KJNM_TO_KCALA = 0.02390057361376673 + + +class Preprocessor: + def __init__( + self, + dataset_conf, + input_path_map: Mapping[str, str], + save_path: str | Path, + prior_builder, + prior_file, + prior_name, + frame_slice, + temp=300, + optimize_forces=False, + box=True, + prior_plots=True, + resume_preprocess=False, + num_cores=32, + jobid=None, + totalNrJobs=None, + settings: PreprocessSettings | None = None, + trajectory_source: TrajectorySource | None = None, + ): + if trajectory_source is not None: + self.trajectory = trajectory_source + else: + self.trajectory = H5BatchTrajectorySource(input_path_map) + + self.dataset_conf = dataset_conf + self.paths = PreprocessPaths(Path(save_path)) + self.save_path = os.fspath(self.paths.root) + self.prior_builder = prior_builder + self.prior_file = prior_file + self.prior_name = prior_name + self.frame_slice = frame_slice + self.temp = temp + self.jobid = jobid + self.totalNrJobs = totalNrJobs + self.settings = settings if settings is not None else PreprocessSettings() + + if self.settings.filter_not_processed_step_one: + done = {p.parent.parent.name for p in self.paths.glob_step1_fit_ok()} + m = {k: v for k, v in self.trajectory.as_dict().items() if k in done} + self.trajectory = H5BatchTrajectorySource(m) + print("%d pdbs left after removing pdbs not processed in step 1" % len(m)) + + self.optimize_forces = optimize_forces + self.box = box + self.prior_plots = prior_plots + self.resume_preprocess = resume_preprocess + self.num_cores = num_cores + + print("Input directory paths:", [i["path"] for i in self.dataset_conf]) + print("Save directory path:", self.save_path) + print(f"Temperature: {self.temp}") + print("Frame slice:", slice_to_str(self.frame_slice)) + print("Number of cores used for parallelization:", self.num_cores) + + def _step3_classical_workers(self, n_frames: int) -> int: + return step3_classical_worker_count( + n_frames, + len(self.trajectory.pdb_ids()), + self.num_cores, + ) + + def step1_threading(self, pdbid): + try: + if not (self.resume_preprocess and self.paths.fit_ok(pdbid).exists()): + bar_pos = _worker_bar_position() + self.process_step1(pdbid, bar_pos) + return [] + except Exception as e: + traceback.print_tb(e.__traceback__) + print(f"{pdbid}:", e) + raise + + def step3_threading(self, pdbid): + try: + bar_pos = _worker_bar_position() + self.process_step3(pdbid, bar_pos) + except Exception as e: + traceback.print_tb(e.__traceback__) + print(f"{pdbid}:", e) + raise + + def preprocess(self): + from .pipeline import run_preprocess_pipeline + + run_preprocess_pipeline(self) + + def save_data(self, pdbid, trajectory, embeddings, forces): + raw = self.paths.pdb_raw(pdbid) + np.save(raw / "embeddings.npy", embeddings) + np.save(raw / "forces.npy", forces) + np.save(raw / "coordinates.npy", trajectory.xyz) + box_path = raw / "box.npy" + if self.box: + np.save(box_path, trajectory.unitcell_vectors) + elif box_path.exists(): + box_path.unlink() + + def process_step1(self, pdbid, bar_position=0): + with tqdm( + total=7, + position=bar_position, + desc=f"{pdbid}: File path setup", + dynamic_ncols=True, + leave=False, + ) as pbar: + + def tick(msg): + pbar.update(1) + pbar.set_description_str(f"{pdbid}: {msg}") + + self.paths.pdb_raw(pdbid).mkdir(parents=True, exist_ok=True) + self.paths.pdb_processed(pdbid).mkdir(parents=True, exist_ok=True) + + AAtraj, path = step1_load_aa(self, pdbid, tick) + cg_map, mol, _topology = step1_write_topology(self, pdbid, AAtraj, tick) + forces = step1_map_forces(self, pdbid, path, AAtraj, cg_map, tick) + cg_traj = step1_cg_traj(self, AAtraj, cg_map, tick) + step1_save_raw(self, pdbid, cg_traj, cg_map, forces, tick) + step1_attach_coords(mol, cg_traj) + step1_prior_cache(self, pdbid, mol, cg_traj, tick) + + def process_step2(self): + if self.prior_plots: + plot_dir = self.paths.prior_fit_plots_dir() + plot_dir.mkdir(parents=True, exist_ok=True) + plot_dir_s = os.fspath(plot_dir) + else: + plot_dir_s = None + + self.prior_builder.fit( + self.temp, + plot_dir=plot_dir_s, + use_cached_fits=self.settings.use_cached_fits, + ) + + def process_step3(self, pdbid, bar_position=0): + step3_remove_stale_prior_files(self, pdbid) + paths = step3_delta_paths(self, pdbid) + step3_run_delta_forces(self, paths, bar_position) + + +# --- Step 1 helpers --- + + +def step1_load_aa(pre: Preprocessor, pdbid: str, tick: Callable[[str], None]) -> tuple[mdtraj.Trajectory, str]: + tick("Loading trajectory") + path = pre.trajectory.h5_path(pdbid) + AAtraj = load_h5_traj_slice(path, pre.frame_slice) + assert AAtraj.xyz is not None + AAtraj.xyz *= 10 + return AAtraj, path + + +def step1_write_topology( + pre: Preprocessor, pdbid: str, AAtraj: mdtraj.Trajectory, tick: Callable[[str], None] +) -> tuple[object, object, object]: + tick("Building CG mapping") + cg_map = pre.prior_builder.build_mapping(AAtraj.topology) + mol = pre.prior_builder.make_mol(cg_map) + topology = cg_map.to_mol(bonds=True, angles=True, dihedrals=True) + mol.write(os.fspath(pre.paths.pdb_processed(pdbid) / f"{pdbid}_processed.psf")) + topology.write(os.fspath(pre.paths.pdb_processed(pdbid) / "topology.psf")) + return cg_map, mol, topology + + +def step1_map_forces( + pre: Preprocessor, + pdbid: str, + path: str, + AAtraj: mdtraj.Trajectory, + cg_map: object, + tick: Callable[[str], None], +) -> np.ndarray: + tick("Mapping CG forces") + with h5py.File(path, "r") as f: + forces = f["forces"][pre.frame_slice, :, :] # pyright: ignore[reportIndexIssue] + if pre.optimize_forces: + forces = cg_map.cg_optimal_forces(AAtraj, forces) + else: + forces = cg_map.cg_forces(forces) + assert len(forces) == len(AAtraj) + return forces * _KJNM_TO_KCALA + + +def step1_cg_traj( + pre: Preprocessor, AAtraj: mdtraj.Trajectory, cg_map: object, tick: Callable[[str], None] +) -> mdtraj.Trajectory: + tick("Mapping CG coordinates") + xyz = cg_map.cg_positions(AAtraj.xyz) + cg_traj = mdtraj.Trajectory(xyz, topology=cg_map.to_mdtraj()) + if pre.box and AAtraj.unitcell_lengths is not None: + cg_traj.unitcell_lengths = AAtraj.unitcell_lengths * 10 + cg_traj.unitcell_angles = AAtraj.unitcell_angles + else: + cg_traj.unitcell_lengths = None + cg_traj.unitcell_angles = None + return cg_traj + + +def step1_save_raw( + pre: Preprocessor, + pdbid: str, + cg_traj: mdtraj.Trajectory, + cg_map: object, + forces: np.ndarray, + tick: Callable[[str], None], +) -> None: + tick("Saving Data") + pre.save_data(pdbid, cg_traj, cg_map.embeddings, forces) + + +def step1_attach_coords(mol: object, cg_traj: mdtraj.Trajectory) -> None: + assert cg_traj.xyz is not None + mol.coords = np.moveaxis(cg_traj.xyz, 0, -1) + + +def step1_prior_cache( + pre: Preprocessor, pdbid: str, mol: object, cg_traj: mdtraj.Trajectory, tick: Callable[[str], None] +) -> None: + tick("Generating prior fit data") + if pre.prior_file: + return + fit_dir = pre.paths.pdb_fit(pdbid) + fit_dir.mkdir(parents=True, exist_ok=True) + pre.prior_builder.add_molecule(mol, cg_traj, os.fspath(fit_dir)) + + +# --- Step 3 helpers --- + + +def step3_remove_stale_prior_files(pre: Preprocessor, pdbid: str) -> None: + raw = pre.paths.pdb_raw(pdbid) + for name in (f"{pdbid}_priors.yaml", f"{pdbid}_prior_params.json"): + p = raw / name + if p.exists(): + p.unlink() + + +def step3_delta_paths(pre: Preprocessor, pdbid: str) -> dict[str, str | None]: + raw = pre.paths.pdb_raw(pdbid) + box_npz = os.fspath(raw / "box.npy") if pre.box else None + return { + "coords": os.fspath(raw / "coordinates.npy"), + "forces": os.fspath(raw / "forces.npy"), + "delta": os.fspath(raw / "deltaforces.npy"), + "prior_e": os.fspath(raw / "prior_energy.npy"), + "box": box_npz, + "ff": os.fspath(pre.paths.priors_yaml()), + "psf": os.fspath(pre.paths.pdb_processed(pdbid) / f"{pdbid}_processed.psf"), + } + + +def _step3_nn_and_classical( + pre: Preprocessor, + dfo: DeltaForces, + paths: dict[str, str | None], + bar_position: int, + classical_workers: int, +) -> None: + pp = pre.prior_builder.prior_params + dfo.addExternalForces( + paths["ff"], + pre.prior_builder.priors["bonds"], + pre.prior_builder.priors["angles"], + pre.prior_builder.priors["dihedrals"], + forceterms=pp["forceterms_nn"], + bar_position=bar_position, + ) + dfo.computePriorForces( + paths["ff"], + exclusions=pp["exclusions"], + forceterms=pp["forceterms_classical"], + bar_position=bar_position, + num_parallel_workers=classical_workers, + ) + + +def step3_run_delta_forces( + pre: Preprocessor, paths: dict[str, str | None], bar_position: int +) -> None: + dfo = DeltaForces( + pre.settings.device_step_3, + paths["psf"], + paths["coords"], + paths["box"], + ) + n_frames = int(dfo.coords.shape[0]) + classical_workers = pre._step3_classical_workers(n_frames) + pp = pre.prior_builder.prior_params + + if "external" in pp.keys(): + _step3_nn_and_classical(pre, dfo, paths, bar_position, classical_workers) + else: + dfo.computePriorForces( + paths["ff"], + exclusions=pp["exclusions"], + forceterms=pp["forceterms"], + bar_position=bar_position, + num_parallel_workers=classical_workers, + ) + + dfo.makeAndSaveDeltaForces(paths["forces"], paths["delta"], paths["prior_e"]) + + +def _worker_bar_position(default: int = 1) -> int: + name = mp.current_process().name + parts = name.split("-") + if len(parts) >= 2 and parts[1].isdigit(): + return int(parts[1]) + 1 + return default diff --git a/preprocess/settings.py b/preprocess/settings.py new file mode 100644 index 0000000..946c5c3 --- /dev/null +++ b/preprocess/settings.py @@ -0,0 +1,15 @@ +"""Runtime toggles for preprocessing (replaces former module-level globals).""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List + + +@dataclass +class PreprocessSettings: + filter_not_processed_step_one: bool = False + use_cached_fits: List[str] = field(default_factory=list) + device_step_3: str = "cpu" + do_step_1: bool = True + regen_cache_files: bool = True diff --git a/preprocess/steps/__init__.py b/preprocess/steps/__init__.py new file mode 100644 index 0000000..2b094d1 --- /dev/null +++ b/preprocess/steps/__init__.py @@ -0,0 +1 @@ +"""Pipeline stage implementations (import submodules directly from `pipeline`).""" diff --git a/preprocess/steps/info.py b/preprocess/steps/info.py new file mode 100644 index 0000000..049cba0 --- /dev/null +++ b/preprocess/steps/info.py @@ -0,0 +1,66 @@ +"""Write `result/info.json`, optional resume checks, touch `ok_list`.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +from tqdm import tqdm + +from ..loaders import num_frames_for_info, slice_to_str + +if TYPE_CHECKING: + from ..runner import Preprocessor + + +def _info_dict(pre: Preprocessor) -> dict: + tag_beta = pre.prior_builder.prior_params.get("tag_beta_turns", False) + return { + "input_paths": [i["path"] for i in pre.dataset_conf], + "num_frames": num_frames_for_info(pre.frame_slice), + "frame_slice": slice_to_str(pre.frame_slice), + "pdbids": pre.trajectory.pdb_ids(), + "optimize_forces": pre.optimize_forces, + "box": pre.box, + "prior_name": pre.prior_name, + "tag_beta_turns": tag_beta, + } + + +def assert_resume_compatible(pre: Preprocessor, info_dict: dict) -> None: + if not pre.resume_preprocess: + return + info_path = pre.paths.info_json() + if not info_path.exists(): + return + with open(info_path, "rt", encoding="utf-8") as f: + previous_info = json.load(f) + for k in ("box", "optimize_forces", "prior_name"): + assert info_dict[k] == previous_info[k], ( + f"Can't resume with different parameters: {k}: {info_dict[k]} != {previous_info[k]}" + ) + if "num_frames" in previous_info and previous_info.get("num_frames") is not None: + assert info_dict.get("num_frames") == previous_info.get("num_frames"), ( + "Can't resume with different num_frames: " + f"{info_dict.get('num_frames')} != {previous_info.get('num_frames')}" + ) + else: + assert info_dict.get("frame_slice") == previous_info.get("frame_slice"), ( + "Can't resume with different frame_slice: " + f"{info_dict.get('frame_slice')!r} != {previous_info.get('frame_slice')!r}" + ) + if "tag_beta_turns" in previous_info: + assert info_dict["tag_beta_turns"] == previous_info["tag_beta_turns"], ( + "Can't resume with different tag_beta_turns: " + f"{info_dict['tag_beta_turns']} != {previous_info['tag_beta_turns']}" + ) + + +def write_result_metadata(pre: Preprocessor) -> None: + pre.paths.result_dir.mkdir(parents=True, exist_ok=True) + info_dict = _info_dict(pre) + assert_resume_compatible(pre, info_dict) + with open(pre.paths.info_json(), "wt", encoding="utf-8") as f: + json.dump(info_dict, f) + tqdm.get_lock() + pre.paths.ok_list_txt().touch() diff --git a/preprocess/steps/job_slice.py b/preprocess/steps/job_slice.py new file mode 100644 index 0000000..8194a29 --- /dev/null +++ b/preprocess/steps/job_slice.py @@ -0,0 +1,23 @@ +"""Restrict PDB set for Step 3 when using job arrays (`--jobid` / `--totalNrJobs`).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ..runner import Preprocessor + + +def apply_job_slice(pre: Preprocessor) -> None: + if not pre.totalNrJobs: + return + pdblist = pre.trajectory.pdb_ids() + pdbids_per_job = len(pdblist) // pre.totalNrJobs + 1 + jobid = pre.jobid + assert jobid is not None + if jobid < pre.totalNrJobs - 1: + lo, hi = jobid * pdbids_per_job, (jobid + 1) * pdbids_per_job + pdbids_c = [pdblist[i] for i in range(lo, hi)] + else: + pdbids_c = [pdblist[i] for i in range(jobid * pdbids_per_job, len(pdblist))] + pre.trajectory = pre.trajectory.filter_to(pdbids_c) diff --git a/preprocess/steps/step1_pool.py b/preprocess/steps/step1_pool.py new file mode 100644 index 0000000..2f2f7a6 --- /dev/null +++ b/preprocess/steps/step1_pool.py @@ -0,0 +1,24 @@ +"""Parallel Step 1 over PDBs.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..worker_pool import run_pdb_pool + +if TYPE_CHECKING: + from ..runner import Preprocessor + + +def run_step1_parallel(pre: Preprocessor) -> None: + if not pre.settings.do_step_1: + return + err = run_pdb_pool( + pre.num_cores, + pre.trajectory.as_dict(), + pre.step1_threading, + "Processing Step 1", + ) + if err: + print("errorList", err) + print("errorList keys", err.keys()) diff --git a/preprocess/steps/step2_pool.py b/preprocess/steps/step2_pool.py new file mode 100644 index 0000000..4d93a78 --- /dev/null +++ b/preprocess/steps/step2_pool.py @@ -0,0 +1,38 @@ +"""Parallel Step 2: merge per-PDB caches, fit prior, or copy fixed prior files.""" + +from __future__ import annotations + +import os +import pickle +import shutil +from typing import TYPE_CHECKING + +from tqdm import tqdm + +from ..loaders import get_prior_params_path + +if TYPE_CHECKING: + from ..runner import Preprocessor + + +def run_step2_parallel(pre: Preprocessor) -> None: + pdb_map = pre.trajectory.as_dict() + if pre.prior_file: + prior_params_path = get_prior_params_path(pre.prior_file) + shutil.copy(pre.prior_file, pre.paths.priors_yaml()) + shutil.copy(prior_params_path, pre.paths.prior_params_json()) + return + + pb_path = pre.paths.prior_builder_pkl() + if not pb_path.exists() or pre.settings.regen_cache_files: + for pdbid in tqdm(pdb_map, desc="Merging cache files together"): + pre.prior_builder.load_molecule_cache(os.fspath(pre.paths.pdb_fit(pdbid))) + with open(pb_path, "wb") as f: + pickle.dump(pre.prior_builder, f) + else: + print("Using cached prior_builder object... ") + with open(pb_path, "rb") as f: + pre.prior_builder = pickle.load(f) + + pre.process_step2() + pre.prior_builder.save_prior(pre.save_path, None) diff --git a/preprocess/steps/step3_pool.py b/preprocess/steps/step3_pool.py new file mode 100644 index 0000000..6fe48f0 --- /dev/null +++ b/preprocess/steps/step3_pool.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..worker_pool import run_pdb_pool + +if TYPE_CHECKING: + from ..runner import Preprocessor + + +def run_step3_parallel(pre: Preprocessor) -> None: + pdbids = pre.trajectory.pdb_ids() + print(f"Step 3: Processing {len(pdbids)} pdbids") + if len(pdbids) <= 1: + for pdbid in pdbids: + pre.step3_threading(pdbid) + return + + err = run_pdb_pool( + pre.num_cores, + pre.trajectory.as_dict(), + pre.step3_threading, + "Processing Step 3", + ) + if err: + print("step3 errorList", err) + print("step3 errorList keys", err.keys()) + + +def write_ok_list(pre: Preprocessor) -> None: + with open(pre.paths.ok_list_txt(), "wt", encoding="utf-8") as ok_list: + ok_list.write("\n".join(pre.trajectory.pdb_ids())) diff --git a/preprocess/stress_harness.py b/preprocess/stress_harness.py new file mode 100644 index 0000000..c481637 --- /dev/null +++ b/preprocess/stress_harness.py @@ -0,0 +1,467 @@ +"""End-to-end preprocess stress driver (parallel pools + inner classical workers + I/O). + +From ``base_model/`` a single default run uses the **protein–DNA** prior (``--prior CA_DNA``) +and **1b3t** from the 0119 batch; outputs are under the system temp only: + + python -m preprocess.stress_harness + +To stress **DNA** mapping (a structure that actually has DA/DT/DG/DC in the H5), run a +second pass with a pdb from your batch, e.g.: + + python -m preprocess.stress_harness --with-dna-pdb 1xyz + +(Replace ``1xyz`` with a real id in ``--seed-input`` that contains DNA in the trajectory.) + +Protein-only prior (no DNA terms path): use ``--prior CA``. + +Override batch dir or pdbs: ``--pdbids 1b3t,1d02`` or ``--all-pdbs``. + +Environment: set ``STRESS_MAX_CORES`` to cap the default core sweep. Use ``--repeat 3`` to +detect flaky timing. If ``Too many open files`` appears, raise ``ulimit -n`` before running +large ``--copies`` / high ``--cores`` sweeps. +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Sequence + +import numpy as np + +from .loaders import gen_input_mapping + +# Default one-command inputs (1b3t under 0119_successes, CA_DNA to match 0225/0120-style runs). +DEFAULT_STRESS_INPUT_DIR = "/media/DATA_18_TB_1/akshitha/0119_successes" +DEFAULT_STRESS_PDBIDS = ("1b3t",) +DEFAULT_STRESS_PRIOR = "CA_DNA" + + +def _first_valid_pdb_h5(input_root: Path) -> tuple[str, Path]: + m = gen_input_mapping([{"path": os.fspath(input_root)}]) + if not m: + raise FileNotFoundError( + f"No result/output_*.h5 under {input_root} (see gen_input_mapping)." + ) + sid = next(iter(m.keys())) + return sid, Path(m[sid]) + + +def _make_multi_pdb_input( + input_root: Path, seed_id: str, h5: Path, n_copies: int, work: Path +) -> Path: + if n_copies < 1: + raise ValueError("n_copies must be >= 1") + if n_copies == 1: + return input_root + h5 = h5.resolve() + for i in range(n_copies): + name = f"stress_{i:04d}" + d = work / name / "result" + d.mkdir(parents=True, exist_ok=True) + link = d / f"output_{name}.h5" + if link.exists() or link.is_symlink(): + link.unlink() + link.symlink_to(h5) + return work + + +def _parse_int_list(s: str) -> list[int]: + return [int(x.strip()) for x in s.split(",") if x.strip()] + + +def _validate_run(out: Path, pdbids: list[str]) -> None: + okp = out / "result" / "ok_list.txt" + if not okp.is_file(): + raise AssertionError(f"Missing {okp}") + lines = [x.strip() for x in okp.read_text(encoding="utf-8").splitlines() if x.strip()] + if set(lines) != set(pdbids): + raise AssertionError(f"ok_list mismatch: {lines} vs {pdbids}") + for pid in pdbids: + dfn = out / pid / "raw" / "deltaforces.npy" + if not dfn.is_file(): + raise AssertionError(f"Missing {dfn}") + + +def _load_delta(path: Path) -> np.ndarray: + return np.load(os.fspath(path)) + + +def _assert_deltas_close(a: Path, b: Path, *, rtol: float, atol: float) -> None: + xa, xb = _load_delta(a), _load_delta(b) + if xa.shape != xb.shape: + raise AssertionError(f"Shape mismatch {a} {xa.shape} vs {b} {xb.shape}") + if not np.allclose(xa, xb, rtol=rtol, atol=atol): + d = float(np.max(np.abs(xa - xb))) + raise AssertionError(f"deltaforces differ: {a} vs {b} max_abs={d}") + + +def _run_subprocess( + base_model: Path, cmd: list[str] +) -> tuple[int, float]: + t0 = time.perf_counter() + p = subprocess.run( + cmd, + cwd=os.fspath(base_model), + env={**os.environ, "PYTHONPATH": os.fspath(base_model)}, + check=False, + ) + return p.returncode, time.perf_counter() - t0 + + +def _cap_cores(cores: Sequence[int], cap: int | None) -> list[int]: + if cap is None: + if os.environ.get("STRESS_MAX_CORES"): + cap = int(os.environ["STRESS_MAX_CORES"]) + else: + cap = (os.cpu_count() or 4) * 2 + return [c for c in sorted(set(cores)) if 1 <= c <= cap] + + +def run_one( + base_model: Path, + input_dir: Path, + out_dir: Path, + *, + prior: str, + num_frames: int, + num_cores: int, + extra: list[str], + subprocess_pdbids: list[str] | None = None, +) -> float: + cmd = [ + sys.executable, + "-m", + "preprocess", + os.fspath(input_dir), + "-o", + os.fspath(out_dir), + "--prior", + prior, + "--num-frames", + str(num_frames), + "--num-cores", + str(num_cores), + "--no-prior-plots", + ] + if subprocess_pdbids: + cmd += ["--pdbids", *subprocess_pdbids] + cmd += extra + print("+", " ".join(cmd), flush=True) + code, wall = _run_subprocess(base_model, cmd) + if code != 0: + raise RuntimeError(f"preprocess failed with {code} after {wall:.2f}s") + return wall + + +def stress_matrix( + base_model: Path, + seed_input: Path, + *, + prior: str = DEFAULT_STRESS_PRIOR, + copies: int = 1, + frames_list: list[int] | None = None, + cores_list: list[int] | None = None, + repeat: int = 1, + verify_parallel: bool = False, + extra: list[str] | None = None, + keep_work: bool = False, + pdbid_filter: list[str] | None = None, + entire_input_dir: bool = False, +) -> None: + """Run preprocess over (frames × cores × repeat) and check outputs; optional 1 vs N-core agreement.""" + if entire_input_dir: + f_pdb: list[str] | None = None + elif pdbid_filter is not None: + f_pdb = list(pdbid_filter) + else: + f_pdb = list(DEFAULT_STRESS_PDBIDS) + if verify_parallel and repeat > 1: + print("verify_parallel: forcing repeat=1 (comparing single runs per (frames, cores)).", flush=True) + repeat = 1 + frames_list = frames_list or [5, 20] + max_cpu = os.cpu_count() or 4 + raw_cores = cores_list or _cap_cores([1, 2, 4, max(8, max_cpu // 2), max_cpu], None) + if verify_parallel: + if 1 not in raw_cores: + raise ValueError("verify_parallel requires --cores to include 1 (serial PDB pool baseline).") + cores_list = [1] + [c for c in sorted(set(raw_cores)) if c != 1] + else: + cores_list = list(raw_cores) + extra = extra or [] + + work_root = ( + Path(tempfile.mkdtemp(prefix="preprocess_stress_in_")) + if copies > 1 + else seed_input + ) + + try: + full_map = gen_input_mapping([{"path": os.fspath(seed_input)}]) + if not full_map: + raise FileNotFoundError( + f"No result/output_*.h5 under {seed_input} (see gen_input_mapping)." + ) + if f_pdb is not None: + missing = set(f_pdb) - set(full_map.keys()) + if missing: + raise ValueError(f"pdbids not found under input: {sorted(missing)}") + seed_id = f_pdb[0] + h5 = Path(full_map[seed_id]) + else: + if len(full_map) > 1: + print( + f"WARNING: processing all {len(full_map)} pdbs under {seed_input}.", + flush=True, + ) + seed_id, h5 = _first_valid_pdb_h5(seed_input) + print(f"Using seed {seed_id!r} -> {h5}", flush=True) + input_dir = ( + _make_multi_pdb_input(seed_input, seed_id, h5, copies, work_root) + if copies > 1 + else seed_input + ) + discovered = list( + gen_input_mapping([{"path": os.fspath(input_dir)}]).keys() + ) + if copies > 1: + subprocess_pdbids: list[str] | None = None + pdbids = discovered + elif f_pdb is not None: + subprocess_pdbids = list(f_pdb) + pdbids = list(f_pdb) + else: + subprocess_pdbids = None + pdbids = discovered + print( + f"PDB ids in run ({len(pdbids)}): {pdbids[:5]}{'...' if len(pdbids) > 5 else ''}", + flush=True, + ) + + for n_fr in frames_list: + if not verify_parallel: + for n_co in cores_list: + times: list[float] = [] + for r in range(repeat): + out = Path( + tempfile.mkdtemp( + prefix=f"out_{n_fr}f_{n_co}c_r{r}_", + ) + ) + try: + t = run_one( + base_model, + input_dir, + out, + prior=prior, + num_frames=n_fr, + num_cores=n_co, + extra=extra, + subprocess_pdbids=subprocess_pdbids, + ) + _validate_run(out, pdbids) + times.append(t) + finally: + shutil.rmtree(out, ignore_errors=True) + mean = statistics.mean(times) + s = f"frames={n_fr} cores={n_co}: " + ( + f"wall_s={times[0]:.2f}" + if repeat == 1 + else f"wall_s mean={mean:.2f} stdev={statistics.pstdev(times):.3f} runs={times}" + ) + print(s, flush=True) + else: + out_ref = Path( + tempfile.mkdtemp(prefix=f"ref_{n_fr}f_1c_"), + ) + try: + t_ref = run_one( + base_model, + input_dir, + out_ref, + prior=prior, + num_frames=n_fr, + num_cores=1, + extra=extra, + subprocess_pdbids=subprocess_pdbids, + ) + _validate_run(out_ref, pdbids) + print( + f"frames={n_fr} cores=1: wall_s={t_ref:.2f} (reference for deltaforces check)", + flush=True, + ) + except Exception: + shutil.rmtree(out_ref, ignore_errors=True) + raise + for n_co in [c for c in cores_list if c != 1]: + out_par = Path( + tempfile.mkdtemp(prefix=f"out_{n_fr}f_{n_co}c_"), + ) + try: + t = run_one( + base_model, + input_dir, + out_par, + prior=prior, + num_frames=n_fr, + num_cores=n_co, + extra=extra, + subprocess_pdbids=subprocess_pdbids, + ) + _validate_run(out_par, pdbids) + for pid in pdbids: + _assert_deltas_close( + out_ref / pid / "raw" / "deltaforces.npy", + out_par / pid / "raw" / "deltaforces.npy", + rtol=1e-5, + atol=1e-6, + ) + print( + f"frames={n_fr} cores={n_co}: wall_s={t:.2f} (deltaforces match ref)", + flush=True, + ) + finally: + shutil.rmtree(out_par, ignore_errors=True) + shutil.rmtree(out_ref, ignore_errors=True) + finally: + if copies > 1 and not keep_work and work_root != seed_input and work_root.exists(): + shutil.rmtree(work_root, ignore_errors=True) + + +def _main() -> None: + ap = argparse.ArgumentParser( + description="Stress-test full preprocess (optional multi-PDB from one H5).", + ) + ap.add_argument( + "seed_input", + type=Path, + nargs="?", + default=Path(DEFAULT_STRESS_INPUT_DIR), + help=( + f"Input batch directory (default: {DEFAULT_STRESS_INPUT_DIR} " + f"— then only {','.join(DEFAULT_STRESS_PDBIDS)} and prior {DEFAULT_STRESS_PRIOR} " + f"unless --all-pdbs or --pdbids)" + ), + ) + ap.add_argument( + "--all-pdbs", + action="store_true", + help="Run on every system under the batch dir (overrides default 1b3t).", + ) + ap.add_argument( + "--pdbids", + type=str, + default="", + help="Comma-separated pdbids, e.g. 1b3t,1d02. Empty uses the default 1b3t (ignored with --all-pdbs).", + ) + ap.add_argument( + "--base-model", + type=Path, + default=Path(__file__).resolve().parent.parent, + help="Path to base_model (default: parent of package preprocess)", + ) + ap.add_argument( + "--prior", + type=str, + default=DEFAULT_STRESS_PRIOR, + help=f"Preprocess --prior (default: {DEFAULT_STRESS_PRIOR} for mixed protein–DNA; use CA for protein only).", + ) + ap.add_argument( + "--with-dna-pdb", + type=str, + default="", + metavar="PDBID", + help="After the default pass, run the same stress matrix for this single pdbid (use a batch entry whose H5 has DNA) to stress DA/DT/DG/DC mapping.", + ) + ap.add_argument( + "--copies", + type=int, + default=1, + help="If >1, build unique pdb dirs in a temp area, all sharing the same H5 (stress step 1/3 pools).", + ) + ap.add_argument( + "--frames", + type=str, + default="5,20", + help="Comma-separated frame counts (e.g. 5,20,100)", + ) + ap.add_argument( + "--cores", + type=str, + default="", + help="Comma-separated core counts; default sweep 1,2,4,… up to STRESS_MAX_CORES or cpu_count", + ) + ap.add_argument("--repeat", type=int, default=1, help="Repeat each (frames,cores) for flakiness") + ap.add_argument( + "--verify-parallel", + action="store_true", + help="After a num-cores=1 run, require deltaforces to match under higher core counts (same data).", + ) + ap.add_argument( + "--keep-work", + action="store_true", + help="If --copies>1, do not delete the mirrored input temp dir (debug).", + ) + ap.add_argument( + "extra", + nargs=argparse.REMAINDER, + help="Extra args to preprocess, after --, e.g. -- --no-box", + ) + args = ap.parse_args() + extra: list[str] = [] + if args.extra and args.extra[0] == "--": + extra = list(args.extra[1:]) + + fr = _parse_int_list(args.frames) + cores: list[int] | None + if args.cores.strip(): + cores = _parse_int_list(args.cores) + else: + cores = None + + with_dna = args.with_dna_pdb.strip() + + if args.all_pdbs: + pfilter: list[str] | None = None + entire = True + elif args.pdbids.strip(): + pfilter = [x.strip() for x in args.pdbids.split(",") if x.strip()] + entire = False + else: + pfilter = None + entire = False + + def _one_matrix(pf: list[str] | None, ent: bool) -> None: + stress_matrix( + args.base_model.resolve(), + args.seed_input.resolve(), + prior=args.prior, + copies=args.copies, + frames_list=fr, + cores_list=cores, + repeat=max(1, args.repeat), + verify_parallel=args.verify_parallel, + extra=extra, + keep_work=args.keep_work, + pdbid_filter=pf, + entire_input_dir=ent, + ) + + _one_matrix(pfilter, entire) + if with_dna and not ent: + print( + f"\n[stress_harness] Second pass: --prior {args.prior} --pdbids {with_dna} (DNA trajectory stress)\n", + flush=True, + ) + _one_matrix([with_dna], False) + + +if __name__ == "__main__": + _main() diff --git a/preprocess/trajectory_source.py b/preprocess/trajectory_source.py new file mode 100644 index 0000000..dad488e --- /dev/null +++ b/preprocess/trajectory_source.py @@ -0,0 +1,46 @@ +"""pluggable trajectory path resolution (H5 batch layout today, can swap to direct westpa if we want""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterator, Mapping + + +class TrajectorySource(ABC): + """Maps preprocess PDB ids to trajectory file paths.""" + + @abstractmethod + def h5_path(self, pdbid: str) -> str: + """Absolute or relative path to the per-structure trajectory (e.g. output_*.h5).""" + + @abstractmethod + def pdb_ids(self) -> list[str]: + """Stable ordering for logging and job slicing (sorted ids).""" + + def items(self) -> Iterator[tuple[str, str]]: + for pid in self.pdb_ids(): + yield pid, self.h5_path(pid) + + def as_dict(self) -> dict[str, str]: + return {pid: self.h5_path(pid) for pid in self.pdb_ids()} + + +class H5BatchTrajectorySource(TrajectorySource): + """Default layout from `gen_input_mapping`: pdbid -> path to HDF5.""" + + def __init__(self, pdbid_to_h5: Mapping[str, str], *, order: list[str] | None = None): + self._m = dict(pdbid_to_h5) + if order is not None: + self._ids = [k for k in order if k in self._m] + else: + self._ids = sorted(self._m.keys()) + + def h5_path(self, pdbid: str) -> str: + return self._m[pdbid] + + def pdb_ids(self) -> list[str]: + return list(self._ids) + + def filter_to(self, pdbids: list[str]) -> H5BatchTrajectorySource: + m = {k: self._m[k] for k in pdbids if k in self._m} + return H5BatchTrajectorySource(m, order=[k for k in pdbids if k in m]) diff --git a/preprocess/worker_pool.py b/preprocess/worker_pool.py new file mode 100644 index 0000000..9232bba --- /dev/null +++ b/preprocess/worker_pool.py @@ -0,0 +1,48 @@ +"""Multiprocessing pool helpers for per-PDB work (Steps 1 and 3).""" + +from __future__ import annotations + +import multiprocessing as mp +from typing import Any, Callable, Dict, Mapping + +from tqdm import tqdm + + +def process_init(counter: Any) -> None: + """Assign stable worker names for tqdm positioning.""" + with counter.get_lock(): + idx = int(counter.value) + counter.value += 1 + mp.current_process().name = f"PreprocessWorker-{idx}" + + +def run_pdb_pool( + num_cores: int, + pdbids: Mapping[str, Any], + task_fn: Callable[[str], None], + desc: str, +) -> Dict[str, str]: + """Run task_fn(pdbid) in parallel; return pdbid -> error message for failures.""" + error_list: Dict[str, str] = {} + n = len(pdbids) + thread_counter = mp.Value("i", 0, lock=True) + tqdm.get_lock() + with tqdm(total=n, desc=desc, dynamic_ncols=True) as pbar: + with mp.Pool(num_cores, initializer=process_init, initargs=(thread_counter,)) as pool: + pending: Dict[Any, str] = {} + for pdbid in pdbids: + pending[pool.apply_async(task_fn, args=(pdbid,))] = pdbid + while pending: + for result in list(pending.keys()): + if result.ready(): + try: + result.get() + except Exception as e: + error_list[pending[result]] = str(e) + finally: + del pending[result] + pbar.n = n - len(pending) + pbar.refresh() + if pending: + next(iter(pending)).wait(1) + return error_list diff --git a/preprocess_legacy.py b/preprocess_legacy.py new file mode 100755 index 0000000..ed38a64 --- /dev/null +++ b/preprocess_legacy.py @@ -0,0 +1,1107 @@ +#!/usr/bin/env python3 +# Legacy monolithic copy of preprocess (pre-package split). Prefer: `python preprocess.py` or `python -m preprocess`. +import os +import numpy as np +import yaml +import json +import h5py +import mdtraj +from module.torchmd_cg_mappings import CACB_MAP +from module import prior +from module import prior_flex +from module import psfwriter +from module.make_deltaforces import DeltaForces +from module.cg_mapping import CGMapping +import argparse +import traceback +import shutil +import multiprocessing as mp +from tqdm import tqdm +import glob +import pickle +import torch +# this can have a small performance hit as it uses the HDD file_system to share data across different processes, but it doesn't lead to "Too many files open" error, which was limiting the max number of parallel processes to 16. +torch.multiprocessing.set_sharing_strategy('file_system') + +# Raz Dec 30 2024: turns out that when preprocessing the 6.5k dataset, the last 20 pdbs are taking forever to process due to being very big proteins. In addition, some were giving hdf5 errors since the batch_generate job didn't properly finish, and this means we don't have all the frames, and they will be removed later on before training. So here's the best workflow I found to solve it: +# 1) first run just step 1 on all proteins, and set FILTER_NOT_PROCESSED_STEP_ONE = False. if any pdbs are taking forever (generally the last 20), just kill the job +# 2) re-run the pre-processing for a second time with FILTER_NOT_PROCESSED_STEP_ONE = True +FILTER_NOT_PROCESSED_STEP_ONE = False + +# Controls whether to re-generate the priors in step 2 for each of the terms, terms in the list will be loaded from the cache instead of refit +#USE_CACHED_FITS = ['dihedrals', 'angles', 'bonds', 'lj'] +USE_CACHED_FITS = [] + +DEVICE_STEP_3 = 'cpu' +#DEVICE_STEP_3 = 'cuda' + +DO_STEP_1 = True # whether to do step 1. if you got errors in steps 2-3 and want to resume, set this to False +REGEN_CACHE_FILES = True # whether to re-generate cache files + + +def process_init(counter): + """This function sets the worker names such that we can use them to position the tqdm bars""" + with counter.get_lock(): + idx = int(counter.value) + counter.value += 1 + mp.current_process().name = f"PreprocessWorker-{idx}" + + +class CGMappingDef_CA: + def __init__(self): + residues = ["ALA", "CYS", "ASP", "GLU", "PHE", "GLY", "HIS", "ILE", "LYS", "LEU", "MET", "ASN", "PRO", "HYP", "GLN", "ARG", "SER", "THR", "VAL", "TRP", "TYR"] + # For legacy reasons we have a couple extra ambiguous residues (ASX & GLX) in the embedding map but we do not accept these for parsing + embedding_residues = ["ALA", "ARG", "ASN", "ASP", "ASX", "CYS", "GLU", "GLN", "GLX", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"] + self.bead_embeddings = {name: [index + 1] for index, name in enumerate(sorted(embedding_residues))} + + # bead_atom_selection: A list of lists, where each inner list is the names of the atoms that will be combined to form the bead + self.bead_atom_selection = {k: [["CA"]] for k in residues} + # The type names of beads (will become the atom type/element in the cg topology) + self.bead_types = { + "ALA": ["CAA"], + "ARG": ["CAR"], + "ASN": ["CAN"], + "ASP": ["CAD"], + "CYS": ["CAC"], + "GLN": ["CAQ"], + "GLU": ["CAE"], + "GLY": ["CAG"], + "HIS": ["CAH"], + "HSD": ["CAH"], + "ILE": ["CAI"], + "LEU": ["CAL"], + "LYS": ["CAK"], + "MET": ["CAM"], + "PHE": ["CAF"], + "PRO": ["CAP"], + "SER": ["CAS"], + "THR": ["CAT"], + "TRP": ["CAW"], + "TYR": ["CAY"], + "VAL": ["CAV"], + } + # The "atom name" assigned to the beads + self.bead_atom_names = {k: ["CA"] for k in residues} + self.bead_masses = {k: [12.01] for k in residues} + self.bead_backbone_idx = {k: 0 for k in residues} + +class CGMappingDef_CACB: + def __init__(self): + residues = ["ALA", "CYS", "ASP", "GLU", "PHE", "GLY", "HIS", "ILE", "LYS", "LEU", "MET", "ASN", "PRO", "HYP", "GLN", "ARG", "SER", "THR", "VAL", "TRP", "TYR"] + + # bead_atom_selection: A list of lists, where each inner list is the names of the atoms that will be combined to form the bead + self.bead_atom_selection = {k: [["CA"], ["CB"]] for k in residues} + self.bead_atom_selection["GLY"] = [["CA"]] + # The type names of beads (will become the atom type/element in the cg topology) + self.bead_types = { + "ALA": ["CA", "CBA"], + "ARG": ["CA", "CBR"], + "ASN": ["CA", "CBN"], + "ASP": ["CA", "CBD"], + "CYS": ["CA", "CBC"], + "GLN": ["CA", "CBQ"], + "GLU": ["CA", "CBE"], + "GLY": ["CAG"], + "HIS": ["CA", "CBH"], + "HSD": ["CA", "CBH"], + "ILE": ["CA", "CBI"], + "LEU": ["CA", "CBL"], + "LYS": ["CA", "CBK"], + "MET": ["CA", "CBM"], + "PHE": ["CA", "CBF"], + "PRO": ["CA", "CBP"], + "SER": ["CA", "CBS"], + "THR": ["CA", "CBT"], + "TRP": ["CA", "CBW"], + "TYR": ["CA", "CBY"], + "VAL": ["CA", "CBV"], + } + + embedding_map = {k:i for i,k in enumerate(sorted(set.union(*[set(i) for i in self.bead_types.values()])))} + self.bead_embeddings = {k:[embedding_map[i] for i in v] for k, v in self.bead_types.items()} + + # The "atom name" assigned to the beads + self.bead_atom_names = {k: ["CA", "CB"] for k in residues} + self.bead_atom_names["GLY"] = ["CA"] + self.bead_masses = {k: [12.01]*len(v) for k,v in self.bead_types.items()} + self.bead_backbone_idx = {k: 0 for k in residues} + +class PriorBuilder: + def __init__(self): + self.prior_params = dict() + self.priors = None + self.terms = dict() + self.atom_types = set() + self.fit_constraints = True + self.tag_beta_turns = False + self.min_cnt = 0 + + def select_atoms(self, topology): + """Returns tha atom index to be saved for this prior""" + raise NotImplementedError() + + def map_embeddings(self, selected_atoms, trajectory): + """Generates the embeddings array for the selected atoms""" + raise NotImplementedError() + + def write_psf(self, pdb_file, psf_file): + """Write the .psf file describing the course grain geometry""" + raise NotImplementedError() + + def add_molecule(self, mol, traj, cache_dir): + fit_ok_path = os.path.join(cache_dir, "fit_ok.txt") + + if cache_dir and os.path.exists(fit_ok_path): + os.unlink(fit_ok_path) + + for term in self.terms.values(): + term.add_molecule(mol, traj, cache_dir) + self.atom_types = self.atom_types.union(mol.atomtype) + + if cache_dir: + np.save(os.path.join(cache_dir, "atomtype.npy"), mol.atomtype) + with open(fit_ok_path, "wt", encoding="utf-8") as f: + f.write("ok") + + def load_molecule_cache(self, cache_dir): + assert os.path.exists(os.path.join(cache_dir, "fit_ok.txt")) + atomtype = np.load(os.path.join(cache_dir, "atomtype.npy"), allow_pickle=True) + self.atom_types = self.atom_types.union(atomtype) + + for term in self.terms.values(): + term.load_molecule_cache(cache_dir) + + def enable_fit_constraints(self, use_constraints): + self.fit_constraints = use_constraints + self.prior_params["fit_constraints"] = self.fit_constraints + + def enable_bond_tags(self, use_tags): + self.tag_beta_turns = use_tags + self.prior_params["tag_beta_turns"] = self.tag_beta_turns + + def set_min_cnt(self, min_cnt): + assert min_cnt >= 0 + self.min_cnt = min_cnt + self.prior_params["min_cnt"] = self.min_cnt + + def fit(self, temperature, plot_dir=None): + self.init_prior_dict() + assert self.priors is not None + for key, term in self.terms.items(): + if os.path.exists(f"{plot_dir}/prior_{key}.pkl") and (key in USE_CACHED_FITS): + print(f"Used cached fit for {key}...") + with open(f"{plot_dir}/prior_{key}.pkl", "rb") as f: + self.priors[key] = pickle.load(f) + else: + print(f"Fitting {key}...") + self.priors[key] = term.get_param(temperature, plot_dir, self.fit_constraints, self.min_cnt) + # pickle the prior for this term + with open(f"{plot_dir}/prior_{key}.pkl", "wb") as f: + pickle.dump(self.priors[key], f) + + def init_prior_dict(self): + # Define the force field dict + priors = {} + priors['atomtypes'] = sorted(self.atom_types) + priors['bonds'] = {} + priors['angles'] = {} + priors['dihedrals'] = {} + priors['lj'] = {} + # For mass and charge assume everything is a carbon atom + priors['electrostatics'] = {at: {'charge': 0.0} for at in priors['atomtypes']} + # The mass of carbon used here is the from OpenMM/AMBER-14 value + priors['masses'] = {at: 12.01 for at in priors['atomtypes']} + self.priors = priors + + def save_prior(self, output_path, pdbid): + prefix = "" + if pdbid: + prefix = f"{pdbid}_" + with open(os.path.join(output_path, f"{prefix}priors.yaml"), "w") as f: + yaml.dump(self.priors, f) + with open(os.path.join(output_path, f"{prefix}prior_params.json"),"w") as f: + json.dump(self.prior_params, f) + + def make_mol(self, cg_map): + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return cg_map.to_mol(bonds = bonds, angles = angles, dihedrals = dihedrals) + +class Prior_CA(PriorBuilder): + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA", + "exclusions" : ['bonds'], + "forceterms" : ["bonds"], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + + def build_mapping(self, topology): + return CGMapping(topology, CGMappingDef_CA()) + + def select_atoms(self, topology): + #TODO: Remove this function (replaced by build_mapping) + return topology.select('name CA and protein') + + def map_embeddings(self, selected_atoms, topology): #pyright: ignore[reportIncompatibleMethodOverride] + #TODO: Remove this function (replaced by build_mapping) + standardResidues = {"ALA", "ARG", "ASN", "ASP", "ASX", "CYS", "GLU", "GLN", "GLX", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"} + amino_acid_mapping = {name: index + 1 for index, name in enumerate(sorted(standardResidues))} + + result = [] + for a_idx in selected_atoms: + r_name = topology.atom(a_idx).residue.name + result.append(amino_acid_mapping[r_name]) + return np.array(result, dtype=int) + + def write_psf(self, pdb_file, psf_file): + #TODO: Remove this function (replaced by build_mapping) + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return psfwriter.pdb2psf_CA(pdb_file, psf_file, bonds = bonds, angles = angles, dihedrals = dihedrals, + tag_beta_turns = self.tag_beta_turns) + +class Prior_CACB(PriorBuilder): + """Implements the torchmd-cg CACB prior""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CACB", + "exclusions" : ['bonds'], + "forceterms" : ["bonds"], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + + def build_mapping(self, topology): + return CGMapping(topology, CGMappingDef_CACB()) + + def select_atoms(self, topology): + #TODO: Remove this function (replaced by build_mapping) + return topology.select('(name CA or name CB) and protein') + + def map_embeddings(self, selected_atoms, topology):#pyright: ignore[reportIncompatibleMethodOverride] + #TODO: Remove this function (replaced by build_mapping) + + # Make a map from embedding name to embedding name number + # e.g. {"CAA":0, "CAC":1, ...} + embedding_map = CACB_MAP + embedding_nums = dict([(k, i) for i, k in enumerate(sorted(set(embedding_map.values())))]) + + result = [] + for a_idx in selected_atoms: + r_name = topology.atom(a_idx).residue.name + a_name = topology.atom(a_idx).name + emb_name = embedding_map[(r_name, a_name)] + result.append(embedding_nums[emb_name]) + return np.array(result, dtype=int) + + def write_psf(self, pdb_file, psf_file): + #TODO: Remove this function (replaced by build_mapping) + bonds = "bonds" in self.terms + angles = "angles" in self.terms + dihedrals = "dihedrals" in self.terms + return psfwriter.pdb2psf_CACB(pdb_file, psf_file, bonds = bonds, angles = angles, dihedrals = dihedrals) + +class Prior_CACB_lj(Prior_CACB): + """torchmd-cg CACB prior with Bonded & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CACB_lj", + "exclusions" : ['bonds'], + "forceterms" : ['bonds', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + +class Prior_CACB_lj_angle_dihedral(Prior_CACB): + """torchmd-cg CACB prior with Bonded, Angle, Dihedral & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CACB_lj_angle_dihedral", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator() + +class Prior_CA_lj(Prior_CA): + """CA prior with Bonded & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj", + "exclusions" : ['bonds'], + "forceterms" : ['bonds', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + +class Prior_CA_lj_angle(Prior_CA): + """CA prior with Bonded, Angle, and RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angle", + "exclusions" : ['bonds', 'angles'], + "forceterms" : ['bonds', 'angles', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms['angles'] = prior.ParamAngleCalculator() + +class Prior_CA_lj_angle_dihedral(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, Dihedral & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angle_dihedral", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator() + +class Prior_CA_lj_angle_dihedralX(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angle_dihedralX", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_angleXCX_dihedralX(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleXCX_dihedralX", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms" : ['bonds', 'angles', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator(center=True) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_angleXCX_dihedralX_flex(Prior_CA): + """torchmd-cg CA prior with highly flexible Bonded, Angle, DihedralX & RepulsionCG terms that fit the data. + + """ + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleXCX_dihedralX_flex", + "exclusions" : ['bonds', 'angles', 'dihedrals'], + "forceterms_nn" : ['bonds', 'angles', 'dihedrals'], + "forceterms_classical": ['repulsioncg'], # changed from lj, would need to re-generated the dataset (Jan 10 2025). repulsioncg is using just the repulsion term from lj. it uses the same parameters as lj, so need to make sure the right function is evaluated. + "external" : True + }) + self.prior_params['forceterms'] = self.prior_params['forceterms_classical'] + self.prior_params['forceterms_nn'] + + self.terms["bonds"] = prior_flex.ParamBondedFlexCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior_flex.ParamAngleFlexCalculator(center=True) + self.terms["dihedrals"] = prior_flex.ParamDihedralFlexCalculator(unified=True) + + # have to override this method since we're saving neural nets as priors + def save_prior(self, output_path, pdbid): + prefix = "" + # if pdbid: + # prefix = f"{pdbid}_" + with open(os.path.join(output_path, f"{prefix}prior_params.json"),"w") as f: + json.dump(self.prior_params, f) + + # print('self.priors', self.priors.keys()) + # remove the dihedrals and bonds from the priors + priorsTruncated = self.priors.copy() + priorsTruncated.pop('dihedrals') + priorsTruncated.pop('bonds') + priorsTruncated.pop('angles') + # print('priorsTruncated', priorsTruncated.keys()) + + # save the classical priors using yaml. this is requires because the classical priors are built from the yaml files + with open(os.path.join(output_path, f"{prefix}priors.yaml"), "w") as f: + yaml.dump(priorsTruncated, f) + + self.priors['terms'] = self.terms + self.priors['prior_params'] = self.prior_params + + # also save with pickle + with open(os.path.join(output_path, f"{prefix}priors.pkl"), "wb") as f: + pickle.dump(self.priors, f) + + + def load_prior_nnets(self, output_path): + # load the prior with pickle + with open(os.path.join(output_path, "priors.pkl"), "rb") as f: + self.priors = pickle.load(f) + + # return self.priors + + # with open(os.path.join(output_path, f"{prefix}priors.pkl"), "wb") as f: + # pickle.dump(self.priors, f) + + + +class Prior_CA_lj_angleXCX_dihedralX_V1(Prior_CA): + """torchmd-cg CA prior with Bonded, Angle, DihedralX & RepulsionCG terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleXCX_dihedralX_V1", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator(center=True) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_bondNull_angleXCX_dihedralX(Prior_CA): + """torchmd-cg CA prior with Angle, DihedralX & RepulsionCG terms (+ bond exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_bondNull_angleXCX_dihedralX", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.NullParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.ParamAngleCalculator(center=True) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_bondNull_angleNull_dihedralX(Prior_CA): + """torchmd-cg CA prior with DihedralX & RepulsionCG terms (+ bond & angle exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_bondNull_angleNull_dihedralX", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.NullParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_bondNull_angleNull_dihedralNull(Prior_CA): + """torchmd-cg CA prior with RepulsionCG terms (+ bond, angle, & dihedral exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_bondNull_angleNull_dihedralNull", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.NullParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.NullParamDihedralCalculator() + +class Prior_CA_lj_angleNull_dihedralX(Prior_CA): + """torchmd-cg CA prior with Bonded, DihedralX & RepulsionCG terms (+ angle exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleNull_dihedralX", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True) + +class Prior_CA_lj_angleNull_dihedralNull(Prior_CA): + """torchmd-cg CA prior with Bonded & RepulsionCG terms (+ angle & dihedral exclusions)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_angleNull_dihedralNull", + "exclusions" : ['bonds', 'angles', '1-4'], + "forceterms" : ['Bonds', 'angles', 'dihedrals', 'RepulsionCG'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["angles"] = prior.NullParamAngleCalculator() + self.terms["dihedrals"] = prior.NullParamDihedralCalculator() + +class Prior_CA_Majewski2022_v0(Prior_CA): + """torchmd-cg CA prior based on the parameters used in (Majewski 2022) + Note this version (v0) has different lj exclusions than the one used in the paper. + """ + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_Majewski2022_v0", + "exclusions" : ['bonds', 'dihedrals'], + "forceterms" : ['bonds', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True, scale=0.5) + +class Prior_CA_Majewski2022_v1(Prior_CA): + """torchmd-cg CA prior based on the parameters used in (Majewski 2022)""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_Majewski2022_v1", + "exclusions" : ['bonds'], + "forceterms" : ['bonds', 'dihedrals', 'repulsioncg'], + }) + self.terms["bonds"] = prior.ParamBondedCalculator() + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6], exclusion_terms={"bonds"}) + self.terms["dihedrals"] = prior.ParamDihedralCalculator(unified=True, scale=0.5) + +class Prior_CA_null(Prior_CA): + """CA prior with no terms""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_null", + "exclusions" : [], + "forceterms" : [], + }) + self.terms = {} + +class Prior_CA_lj_only(Prior_CA): + """CA prior with just a RepulsionCG term""" + def __init__(self): + super().__init__() + self.prior_params.update({ + "prior_configuration_name": "CA_lj_only", + "exclusions" : [], + "forceterms" : ['RepulsionCG'], + }) + self.terms = {} + self.terms["lj"] = prior.ParamNonbondedCalculator(fit_range=[3, 6]) + +def slice_to_str(s): + result = [s.start, s.stop, s.step] + result = [str(i) if i is not None else '' for i in result] + return ":".join(result) + +def get_prior_params_path(prior_path): + dir_path, file_name = os.path.split(prior_path) + file_name = file_name.replace("priors.yaml", "prior_params.json") + return os.path.join(dir_path, file_name) + +def load_h5_traj_slice(path, slice): + """Load a slice from a h5 trajectory without reading the entire file into memory""" + base_traj = mdtraj.load_frame(path, 0) + with h5py.File(path) as f: + t_xyz = f["coordinates"][slice][:] #pyright: ignore[reportIndexIssue] + t_time = f["time"][slice][:] #pyright: ignore[reportIndexIssue] + + t_unitcell_lengths = None + t_unitcell_angles = None + if "cell_lengths" in f.keys(): + t_unitcell_lengths = f["cell_lengths"][slice][:] #pyright: ignore[reportIndexIssue] + t_unitcell_angles = f["cell_angles"][slice][:] #pyright: ignore[reportIndexIssue] + + result = mdtraj.Trajectory(t_xyz, base_traj.topology, time=t_time, unitcell_lengths=t_unitcell_lengths, unitcell_angles=t_unitcell_angles) + return result + +class Preprocessor: + def __init__(self, dataset_conf, input_path_map, save_path, prior_builder, prior_file, prior_name, frame_slice, temp, optimize_forces, box, prior_plots, resume_preprocess, num_cores, jobid=None, totalNrJobs=None): + self.dataset_conf = dataset_conf + self.save_path = save_path + self.prior_builder = prior_builder + self.prior_file = prior_file + self.frame_slice = frame_slice + self.temp = temp + self.jobid = jobid + self.totalNrJobs = totalNrJobs + + self.pdbid_list = input_path_map + + if FILTER_NOT_PROCESSED_STEP_ONE: + pdbs_processed_step1 = [f.split('/')[-3] for f in glob.glob(os.path.join(save_path, "*/fit/fit_ok.txt"))] + # remove keys that are not in pdbs_processed_step1 + self.pdbid_list = {k: v for k, v in self.pdbid_list.items() if k in pdbs_processed_step1} + print('%d pdbs left after removing pdbs not processed in step 1' % len(self.pdbid_list)) + + # if os.path.exists(os.path.join(save_path, 'pdb_list.pkl')): + # with open(os.path.join(save_path, 'pdb_list.pkl'), 'rb') as f: + # self.pdbid_list = pickle.load(f) + # else: + # self.pdbid_list = self.get_pdbid_list() + + # if FILTER_NOT_PROCESSED_STEP_ONE: + # pdbs_processed_step1 = [f.split('/')[-3] for f in glob.glob(os.path.join(save_path, "*/fit/fit_ok.txt"))] + # # remove keys that are not in pdbs_processed_step1 + # self.pdbid_list = {k: v for k, v in self.pdbid_list.items() if k in pdbs_processed_step1} + # print('%d pdbs left after removing pdbs not processed in step 1' % len(self.pdbid_list)) + + # # pickle pdb_list + # os.makedirs(save_path, exist_ok=True) + # with open(os.path.join(save_path, 'pdb_list.pkl'), 'wb') as f: + # pickle.dump(self.pdbid_list, f) + + + self.optimize_forces = optimize_forces + self.box = box + self.prior_plots = prior_plots + self.resume_preprocess = resume_preprocess + self.num_cores = num_cores + + print("Input directory paths:", [i["path"] for i in self.dataset_conf]) + print("Save directory path:", self.save_path) + print(f"Temperature: {self.temp}") + print("Frame slice:", slice_to_str(self.frame_slice)) + print("Number of cores used for parallelization:", self.num_cores) + # print("PDB ID list:", self.pdbid_list) + + def step1_threading(self, pdbid): + try: + cache_dir = os.path.join(self.save_path, pdbid, "fit") + if not(self.resume_preprocess and os.path.exists(os.path.join(cache_dir, "fit_ok.txt"))): + # This assumes we've named the processes during initialization + bar_pos = int(mp.current_process().name.split("-")[1]) + 1 + self.process_step1(pdbid, bar_pos) + return [] + except Exception as e: + traceback.print_tb(e.__traceback__) + print(f"{pdbid}:", e) + raise + + def step3_threading(self, pdbid): + try: + # This assumes we've named the processes during initialization + bar_pos = int(mp.current_process().name.split("-")[1]) + 1 + self.process_step3(pdbid, bar_pos) + except Exception as e: + traceback.print_tb(e.__traceback__) + print(f"{pdbid}:", e) + raise + + def preprocess(self): + os.makedirs(os.path.join(self.save_path, "result"), exist_ok=True) + + info_dict = { + "input_paths": [i["path"] for i in self.dataset_conf], + "frame_slice": slice_to_str(self.frame_slice), + "pdbids": list(self.pdbid_list.keys()), + "optimize_forces": self.optimize_forces, + "box": self.box, + "prior_name": prior_name + } + + # If resuming, validate that no paramters that would invalidate the fit cache object have changed + # FIXME: This should also check "--tag-beta-turns" + if self.resume_preprocess: + if os.path.exists(os.path.join(self.save_path, "result/info.json")): + with open(os.path.join(self.save_path, "result/info.json"), "rt", encoding="utf-8") as f: + prevous_info = json.load(f) + for k in ["box", "frame_slice", "optimize_forces", "prior_name"]: + assert info_dict[k] == prevous_info[k], \ + f"Can't resume with different parameters: {k}: {info_dict[k]} != {prevous_info[k]}" + + with open(os.path.join(self.save_path, "result/info.json"), "wt", encoding="utf-8") as f: + json.dump(info_dict, f) + + pdbids = self.pdbid_list + + # Ensure all jobs have the same tqdm lock : https://github.com/tqdm/tqdm/issues/982 + tqdm.get_lock() + + # TODO: Print exceptions in the main thread for legibility + # TODO: Abstract the loop logic instead of repeating it twice + + # Truncate any existing ok_list.txt + with open(os.path.join(self.save_path, "result/ok_list.txt"), "wt", encoding="utf-8") as ok_list: + pass + + + + # Run step 1 in parallel, saving the results to the cache + + if DO_STEP_1: + errorList = {} + thread_counter = mp.Value('i', 0, lock=True) + with tqdm(total=len(pdbids), desc="Processing Step 1", dynamic_ncols=True) as pbar: + with mp.Pool(self.num_cores, initializer=process_init, initargs=(thread_counter,)) as pool: + pending_results = {} + + # Submit tasks and map results to pdbids + for pdbid in pdbids: + result = pool.apply_async(self.step1_threading, args=(pdbid,)) + pending_results[result] = pdbid + + while pending_results: + # Check completed tasks + for result in list(pending_results.keys()): # Iterate over a copy to allow removal + if result.ready(): + try: + result.get() # Retrieve result or raise exception + except Exception as e: + pdbid = pending_results[result] # Get the corresponding pdbid + errorList[pdbid] = str(e) + finally: + del pending_results[result] # Remove completed task + + # Update the progress bar + pbar.n = len(pdbids) - len(pending_results) + pbar.refresh() + + # Wait for 1 second or for the last job to finish + if pending_results: + next(iter(pending_results)).wait(1) + + if len(errorList): + print('errorList', errorList) + print('errorList keys', errorList.keys()) + + if not self.prior_file: + # pickle prior builder + if not os.path.exists(self.save_path + '/prior_builder.pkl') or REGEN_CACHE_FILES: + # Merge cache files back into prior builder + for pdbid in tqdm(pdbids, desc="Merging cache files together"): + cache_dir = os.path.join(self.save_path, pdbid, "fit") + self.prior_builder.load_molecule_cache(cache_dir) + + with open(self.save_path + '/prior_builder.pkl', 'wb') as f: + pickle.dump(self.prior_builder, f) + else: + print("Using cached prior_builder object... ") + + with open(self.save_path + '/prior_builder.pkl', 'rb') as f: + self.prior_builder = pickle.load(f) + + self.process_step2() + + self.prior_builder.save_prior(self.save_path, None) + else: + prior_params_path = get_prior_params_path(prior_file) + shutil.copy(self.prior_file, os.path.join(self.save_path, "priors.yaml")) + shutil.copy(prior_params_path, os.path.join(self.save_path, "prior_params.json")) + + if self.totalNrJobs: + pdblist = list(pdbids.keys()) + pdbidsPerJob = len(pdblist) // self.totalNrJobs + 1 + jobid = self.jobid + assert jobid is not None + if jobid < self.totalNrJobs - 1: + pdbids_c = [pdblist[i] for i in range(jobid * pdbidsPerJob, (jobid + 1) * pdbidsPerJob)] + else: + pdbids_c = [pdblist[i] for i in range(jobid * pdbidsPerJob, len(pdblist))] + + # filter pdbids for this job only + pdbids = {k: v for k, v in pdbids.items() if k in pdbids_c} + + print(f"Step 3: Processing {len(pdbids)} pdbids") + + # Run step 3 in parallel + thread_counter = mp.Value('i', 0, lock=True) + with tqdm(total=len(pdbids), desc="Processing Step 3", dynamic_ncols=True) as pbar: + with mp.Pool(self.num_cores, initializer=process_init, initargs=(thread_counter,)) as pool: + pending_results = [] + + for pdbid in pdbids: + pending_results += [pool.apply_async(self.step3_threading, args=(pdbid,))] + + while pending_results: + # Check for exceptions + [i.get() for i in pending_results if i.ready()] + # Remove finished jobs from the list + pending_results = [i for i in pending_results if not i.ready()] + pbar.n = len(pdbids) - len(pending_results) + pbar.refresh() + # Wait for 1 second or for the last job to finish + if pending_results: + pending_results[0].wait(1) + + # alternatively, cd to the preprocessed_data directory and run this cmd: + # ls */raw/deltaforces.npy | awk '{print substr($1, 1, 4)}' > result/ok_list.txt + with open(os.path.join(self.save_path, "result/ok_list.txt"), "wt", encoding="utf-8") as ok_list: + ok_list.write("\n".join(pdbids)) + + print("Done!") + + def save_data(self, output_path, trajectory, embeddings, forces, pdbid): + # print(f" {pdbid} (coordinates, forces): {trajectory.xyz.shape}, {forces.shape}") + np.save(f"{output_path}/raw/embeddings.npy", embeddings) + np.save(f"{output_path}/raw/forces.npy", forces) + np.save(f"{output_path}/raw/coordinates.npy", trajectory.xyz) + box_path = f"{output_path}/raw/box.npy" + if self.box: + np.save(box_path, trajectory.unitcell_vectors) + elif os.path.exists(box_path): + os.unlink(box_path) + + def process_step1(self, pdbid, bar_position=0): + """Generate the course grained data and topology for the protein, then add it to the prior builder""" + + with tqdm(total=7, position=bar_position, desc=f"{pdbid}: File path setup", dynamic_ncols=True, leave=False) as pbar: + def progress_bar_step(msg): + pbar.update(1) + pbar.set_description_str(f"{pdbid}: {msg}") + + # Set up paths and create directories + output_path = os.path.join(self.save_path, pdbid) + + # TODO: Get rit of the of subdirectories? + os.makedirs(f"{output_path}/raw", exist_ok=True) + os.makedirs(f"{output_path}/processed", exist_ok=True) + + # Find which path this ID belongs to + input_file_path = self.pdbid_list[pdbid] + + progress_bar_step("Loading trajectory") + AAtraj = load_h5_traj_slice(input_file_path, self.frame_slice) + assert AAtraj.xyz is not None # for pyright + AAtraj.xyz *= 10 # convert to angstroms + + progress_bar_step("Building CG mapping") + cg_map = self.prior_builder.build_mapping(AAtraj.topology) + mol = self.prior_builder.make_mol(cg_map) + topology = cg_map.to_mol(bonds=True, angles=True, dihedrals=True) + mol.write(f'{output_path}/processed/{pdbid}_processed.psf') + topology.write(f'{output_path}/processed/topology.psf') # Save the topology for the CG mapping, this is optional but useful for debugging + + # Get the forces + progress_bar_step("Mapping CG forces") + with h5py.File(input_file_path, "r") as f: + forces = f["forces"][self.frame_slice, :, :] #pyright: ignore[reportIndexIssue] + + if self.optimize_forces: + forces = cg_map.cg_optimal_forces(AAtraj, forces) + else: + forces = cg_map.cg_forces(forces) + + assert len(forces) == len(AAtraj) + # Convert from kilojoules/mole/nanometer to kilocalories/mole/angstrom + forces = forces*0.02390057361376673 + + progress_bar_step("Mapping CG coordinates") + xyz = cg_map.cg_positions(AAtraj.xyz) + cg_traj = mdtraj.Trajectory(xyz, topology=cg_map.to_mdtraj()) + if self.box and AAtraj.unitcell_lengths is not None: + cg_traj.unitcell_lengths = AAtraj.unitcell_lengths * 10 + cg_traj.unitcell_angles = AAtraj.unitcell_angles + else: + cg_traj.unitcell_lengths = None + cg_traj.unitcell_angles = None + + # Save the data + progress_bar_step("Saving Data") + self.save_data(output_path, cg_traj, cg_map.embeddings, forces, pdbid) + + # Note: moveaxis creates a view, the original trajectory.xyz is unmodified + assert cg_traj.xyz is not None + mol.coords = np.moveaxis(cg_traj.xyz, 0, -1) + + progress_bar_step("Generating prior fit data") + if not self.prior_file: + cache_dir = os.path.join(output_path, "fit") + os.makedirs(cache_dir, exist_ok=True) + self.prior_builder.add_molecule(mol, cg_traj, cache_dir) + + def process_step2(self): + """Fit the prior forcefield based on accumulated data""" + if self.prior_plots: + plot_dir = os.path.join(self.save_path, "prior_fit_plots") + os.makedirs(plot_dir, exist_ok=True) + else: + plot_dir = None + + self.prior_builder.fit(self.temp, plot_dir=plot_dir) + + def process_step3(self, pdbid, bar_position=0): + """Save prior focefield and generate delta forces data for each protein""" + output_path = os.path.join(self.save_path, pdbid) + + # Remove legacy prior files if they exist + if os.path.exists(f"{output_path}/raw/{pdbid}_priors.yaml"): + os.unlink(f"{output_path}/raw/{pdbid}_priors.yaml") + if os.path.exists(f"{output_path}/raw/{pdbid}_prior_params.json"): + os.unlink(f"{output_path}/raw/{pdbid}_prior_params.json") + + # Generate delta forces for all atom simulation vs. prior FF + coords_npz = f'{output_path}/raw/coordinates.npy' + forces_npz = f'{output_path}/raw/forces.npy' + delta_forces_npz = f'{output_path}/raw/deltaforces.npy' + prior_energy_npz = f'{output_path}/raw/prior_energy.npy' + box_npz = None + if self.box: + box_npz = f"{output_path}/raw/box.npy" + forcefield = os.path.join(self.save_path, "priors.yaml") + psf_file = f'{output_path}/processed/{pdbid}_processed.psf' + prior_params = self.prior_builder.prior_params + + deltaForcesObj = DeltaForces(DEVICE_STEP_3, psf_file, coords_npz, box_npz) + if 'external' in self.prior_builder.prior_params.keys(): + # forceterms = ['bonds', 'angles', 'dihedrals'] + deltaForcesObj.addExternalForces(forcefield, self.prior_builder.priors['bonds'], self.prior_builder.priors['angles'], self.prior_builder.priors['dihedrals'], forceterms=prior_params["forceterms_nn"], bar_position=bar_position) + + # forceterms = ['repulsioncg'] # update them properly in preprocess.py in the _flex class + deltaForcesObj.computePriorForces(forcefield, exclusions=prior_params["exclusions"], + forceterms=prior_params["forceterms_classical"], bar_position=bar_position) + + else: + deltaForcesObj.computePriorForces(forcefield, exclusions=prior_params["exclusions"], + forceterms=prior_params["forceterms"], bar_position=bar_position) + + # load MD forces from forces_npz, compute delta forces, and save them in delta_forces_npz + deltaForcesObj.makeAndSaveDeltaForces(forces_npz, delta_forces_npz, prior_energy_npz) + +def gen_input_mapping(conf): + """Find the list of input files for the passed dataset config""" + pdbid_mapping = dict() + for entry in conf: + input_path = entry["path"] + prefix = entry.get("prefix", "") + suffix = entry.get("suffix", "") + assert os.path.isdir(input_path), f"Input path does not exist: {input_path}" + if "pdbids" in entry: + for dir_name in entry["pdbids"]: + input_h5 = os.path.join(input_path, dir_name, "result", f"output_{dir_name}.h5") + assert os.path.exists(input_h5), "Requested path {input_path}/{dir_name} does not exist" + pdbid_mapping[prefix + dir_name + suffix] = input_h5 + else: + dir_names = os.listdir(input_path) + for dir_name in sorted(dir_names): + input_h5 = os.path.join(input_path, dir_name, "result", f"output_{dir_name}.h5") + if os.path.exists(input_h5): + pdbid_mapping[prefix + dir_name + suffix] = input_h5 + else: + print(f" Skipping \"{dir_name}\" (directory contains no output)") + return pdbid_mapping + +prior_types = { + "CA":Prior_CA, + "CACB":Prior_CACB, + "CACB_lj":Prior_CACB_lj, + "CACB_lj_angle_dihedral":Prior_CACB_lj_angle_dihedral, + "CA_lj":Prior_CA_lj, + "CA_lj_angle":Prior_CA_lj_angle, + "CA_lj_angle_dihedral":Prior_CA_lj_angle_dihedral, + "CA_lj_angle_dihedralX":Prior_CA_lj_angle_dihedralX, + "CA_lj_angleXCX_dihedralX":Prior_CA_lj_angleXCX_dihedralX, + "CA_lj_angleXCX_dihedralX_flex":Prior_CA_lj_angleXCX_dihedralX_flex, + "CA_lj_angleXCX_dihedralX_V1":Prior_CA_lj_angleXCX_dihedralX_V1, + "CA_Majewski2022_v0":Prior_CA_Majewski2022_v0, + "CA_Majewski2022_v1":Prior_CA_Majewski2022_v1, + "CA_lj_bondNull_angleXCX_dihedralX":Prior_CA_lj_bondNull_angleXCX_dihedralX, + "CA_lj_bondNull_angleNull_dihedralX":Prior_CA_lj_bondNull_angleNull_dihedralX, + "CA_lj_bondNull_angleNull_dihedralNull":Prior_CA_lj_bondNull_angleNull_dihedralNull, + "CA_lj_angleNull_dihedralX":Prior_CA_lj_angleNull_dihedralX, + "CA_lj_angleNull_dihedralNull":Prior_CA_lj_angleNull_dihedralNull, + "CA_null":Prior_CA_null, + "CA_lj_only":Prior_CA_lj_only, +} + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="Preprocess data.") + parser.add_argument("input", nargs = "+", help="Input directory path") + parser.add_argument("-o", "--output", required=True, help="Output directory path") + parser.add_argument("--pdbids", nargs="*", help="List of specific PDB IDs to process") + parser.add_argument("--num-frames", "--num_frames", type=int, default=None, help="Number of frames to process") + parser.add_argument("--frame-slice", type=str, default=None, help="Select frames to process using a python slice: start:end:stride") + parser.add_argument("--temp", type=int, default=300, help="Temperature") + parser.add_argument("--prior", type=str, default=None, help="Select the prior forcefield to use, must be one of: " + ", ".join(sorted(prior_types.keys()))) + parser.add_argument("--optimize-forces", action="store_true", help="Use statistically optimal force aggregation (Kramer 2023)") + parser.add_argument("--prior-file", default=None, help="Use PRIOR_FILE instead of fitting a prior") + parser.add_argument('--no-box', default=False, action='store_true', help="Don't use periodic box information") + parser.add_argument('--prior-plots', default=True, action='store_true', help="Save plots of the prior fit functions") + parser.add_argument('--no-prior-plots', dest='prior_plots', action='store_false', help="Don't save plots of the prior fit functions") + parser.add_argument('--no-fit-constraints', default=False, action='store_true', help="Disable range constraints when fitting prior functions") + parser.add_argument('--fit-min-cnt', type=int, default=0, help="Only bins with cnt > min_cnt will be considered when fitting the prior (default 0)") + # parser.add_argument('--tag-beta-turns', default=False, action='store_true', help="Give beta turns a different bond type in the prior") + parser.add_argument('--resume', default=False, action='store_true', help="Resume processing rather than overwriting, all settings must be identical between calls") + parser.add_argument('--num-cores', type=int, default=32, help="Number of cores to be used for parallelization of preprocessing") + parser.add_argument('--jobid', type=int, default=None, help="Integer denoting jobid, if not -1, it will only process a subset of the PDBs") + parser.add_argument('--totalNrJobs', type=int, default=None, help="Integer denoting how many jobs are in total.") + + + args = parser.parse_args() + print(args) + + output_dir = args.output + pdbids = args.pdbids + assert not (args.num_frames and args.frame_slice) + if args.num_frames: + frame_slice = slice(0, args.num_frames) + elif args.frame_slice: + # Convert the arg string into a slice + frame_slice = slice(*[int(i) if i != '' else None for i in args.frame_slice.split(":") ]) + else: + frame_slice = slice(None) + temp = args.temp + optimize_forces = args.optimize_forces + box = not args.no_box + prior_plots = args.prior_plots + prior_name = args.prior + prior_file = args.prior_file + resume_preprocess = args.resume + num_cores = args.num_cores + jobid = args.jobid + totalNrJobs = args.totalNrJobs + + if prior_file: + assert os.path.exists(prior_file), f"Prior file does not exist: {prior_file}" + prior_params_path = get_prior_params_path(prior_file) + with open(prior_params_path, "r", encoding="utf-8") as f: + prior_params = json.load(f) + prior_configuration_name = prior_params["prior_configuration_name"] + if prior_name is None: + prior_name = prior_configuration_name + elif prior_name != prior_configuration_name: + print() + print(f"WARNING: Prior \"{prior_name}\" differs from the one used to build the prior file \"{prior_configuration_name}\"") + print() + + assert prior_name, " You must specify the prior to use with either --prior or --prior-file" + + if prior_name not in prior_types: + raise RuntimeError(f"Unknown prior configuration: {prior_name}") + print(f"Using prior: {prior_name}") + prior_builder = prior_types[prior_name]() # <- () to instantiate the class + prior_builder.enable_fit_constraints(not args.no_fit_constraints) + # prior_builder.enable_bond_tags(args.tag_beta_turns) + prior_builder.enable_bond_tags(False) + prior_builder.set_min_cnt(args.fit_min_cnt) + + if 'external' in prior_builder.prior_params.keys(): + mp.set_start_method('spawn') + + # Set matplotlib to use a thread safe backend (for prior fit plots) + import matplotlib + matplotlib.use('Agg') + + dataset_conf = [] + + for i in args.input: + if os.path.isfile(i): + with open(args.input[0], "r") as f: + dataset_conf += yaml.safe_load(f) + else: + dataset_conf += [{"path": i}] + + input_path_map = gen_input_mapping(dataset_conf) + + if pdbids: + input_path_map = {i: input_path_map[i] for i in pdbids} + + preprocessor = Preprocessor(dataset_conf, input_path_map, output_dir, prior_builder, prior_file, prior_name, frame_slice, temp, optimize_forces, box, prior_plots, resume_preprocess, num_cores, jobid, totalNrJobs) + + preprocessor.preprocess() + diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..96a5ace --- /dev/null +++ b/pytest.ini @@ -0,0 +1,7 @@ +[pytest] +pythonpath = . +testpaths = tests +markers = + stress: end-to-end preprocess load tests (require data + time) +filterwarnings = + ignore::DeprecationWarning diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..077e4fa --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,2 @@ +-r requirements-preprocess.txt +pytest>=8.0 diff --git a/requirements-preprocess.txt b/requirements-preprocess.txt new file mode 100644 index 0000000..33e475f --- /dev/null +++ b/requirements-preprocess.txt @@ -0,0 +1,2 @@ +pydantic>=2.0,<3 +PyYAML>=6.0 diff --git a/tests/test_frame_utils.py b/tests/test_frame_utils.py new file mode 100644 index 0000000..cf7cd5a --- /dev/null +++ b/tests/test_frame_utils.py @@ -0,0 +1,78 @@ +"""Tests for module.frame_utils (no torch/mdtraj; safe for lightweight CI).""" + +from __future__ import annotations + +import pytest + +from module.frame_utils import split_frame_indices, step3_classical_worker_count + + +def _flatten_in_order(chunks: list[list[int]]) -> list[int]: + out: list[int] = [] + for c in chunks: + out.extend(c) + return out + + +@pytest.mark.parametrize( + "frames,n_workers,expected", + [ + ([], 4, [[]]), + ([0], 8, [[0]]), + ([0, 1, 2, 3], 1, [[0, 1, 2, 3]]), + ([0, 1, 2, 3], 0, [[0, 1, 2, 3]]), + ( + [0, 1, 2, 3, 4, 5, 6, 7], + 3, + [[0, 1, 2], [3, 4, 5], [6, 7]], + ), + ( + list(range(5)), + 2, + [[0, 1, 2], [3, 4]], + ), + ], +) +def test_split_frame_indices_shape( + frames: list[int], n_workers: int, expected: list[list[int]] +) -> None: + assert split_frame_indices(frames, n_workers) == expected + + +def test_split_preserves_all_indices_in_order() -> None: + frames = list(range(23)) + for w in (2, 3, 5, 7, 11): + chunks = split_frame_indices(frames, w) + assert _flatten_in_order(chunks) == frames + if len(frames) > 1 and w > 1: + w_eff = min(w, len(frames)) + assert len(chunks) == w_eff + sizes = [len(c) for c in chunks] + assert max(sizes) - min(sizes) <= 1 + + +@pytest.mark.parametrize( + "n_frames,cpu_count,num_cores,n_pdbs,expected", + [ + (0, 8, 32, 1, 1), + (1, 8, 32, 1, 1), + (10, 8, 32, 1, 8), + (100, 8, 32, 1, 8), + (100, 16, 4, 2, 4), + (100, 8, 32, 2, 2), + (3, 16, 4, 3, 3), + ], +) +def test_step3_classical_worker_count( + n_frames: int, + cpu_count: int, + num_cores: int, + n_pdbs: int, + expected: int, +) -> None: + assert ( + step3_classical_worker_count( + n_frames, n_pdbs, num_cores, cpu_count=cpu_count + ) + == expected + ) diff --git a/tests/test_preprocess_stress.py b/tests/test_preprocess_stress.py new file mode 100644 index 0000000..0269cef --- /dev/null +++ b/tests/test_preprocess_stress.py @@ -0,0 +1,91 @@ +"""Opt-in end-to-end stress: requires real input data and significant time. + +Run examples:: + + STRESS_SEED_INPUT=/path/to/batch pytest tests/test_preprocess_stress.py -m stress -s + + STRESS_SEED_INPUT=... STRESS_FRAMES=5,15 STRESS_COPIES=4 STRESS_CORES=1,4 pytest ... -m stress -s + +Environment (all optional except STRESS_SEED_INPUT): + STRESS_SEED_INPUT — batch directory (default: same hardcoded default as the harness) + STRESS_ALL — if ``1``, all PDBs in the batch + STRESS_FRAMES — comma list (default ``5``) + STRESS_COPIES — duplicate PDB entries sharing one H5 (default ``1``) + STRESS_CORES — comma list, default is harness sweep + STRESS_VERIFY — if ``1``, compare deltaforces: cores=1 vs higher + STRESS_PRIOR — default ``CA_DNA`` (harness default); use ``CA`` for protein-only + STRESS_WITH_DNA_PDB — if set, a second run with this one pdbid (DNA in H5) after the first pass +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +# Module lives under base_model/ (package ``preprocess`` on PYTHONPATH). +_BASE = Path(__file__).resolve().parent.parent + + +@pytest.mark.stress +def test_preprocess_stress_harness() -> None: + from preprocess.stress_harness import ( + DEFAULT_STRESS_INPUT_DIR, + DEFAULT_STRESS_PRIOR, + stress_matrix, + ) + + seed = os.environ.get("STRESS_SEED_INPUT", DEFAULT_STRESS_INPUT_DIR) + if not Path(seed).is_dir(): + pytest.skip("Set STRESS_SEED_INPUT to a valid batch directory (or mount the default).") + + frames = [ + int(x.strip()) + for x in os.environ.get("STRESS_FRAMES", "5").split(",") + if x.strip() + ] + copies = int(os.environ.get("STRESS_COPIES", "1")) + extra = list( + x + for x in os.environ.get("STRESS_EXTRA", "").split() + if x.strip() + ) + ver = os.environ.get("STRESS_VERIFY", "") in ("1", "true", "yes") + cr = os.environ.get("STRESS_CORES", "").strip() + cores = [int(x) for x in cr.split(",") if x.strip()] if cr else None + rpt = int(os.environ.get("STRESS_REPEAT", "1")) + pids = os.environ.get("STRESS_PDBIDS", "").strip() + pfilter = [x.strip() for x in pids.split(",") if x.strip()] if pids else None + all_pdbs = os.environ.get("STRESS_ALL", "").lower() in ("1", "true", "yes") + + pr = os.environ.get("STRESS_PRIOR", DEFAULT_STRESS_PRIOR) + dna_pdb = os.environ.get("STRESS_WITH_DNA_PDB", "").strip() + + stress_matrix( + _BASE, + Path(seed).resolve(), + prior=pr, + copies=copies, + frames_list=frames, + cores_list=cores, + repeat=max(1, rpt), + verify_parallel=ver, + extra=extra, + pdbid_filter=pfilter, + entire_input_dir=all_pdbs, + ) + if dna_pdb and not all_pdbs: + stress_matrix( + _BASE, + Path(seed).resolve(), + prior=pr, + copies=copies, + frames_list=frames, + cores_list=cores, + repeat=max(1, rpt), + verify_parallel=ver, + extra=extra, + pdbid_filter=[dna_pdb], + entire_input_dir=False, + )