From 9f4e5ddf9036391855d767887e06e4c56fab52dd Mon Sep 17 00:00:00 2001 From: Carlo Camilloni Date: Thu, 16 Apr 2026 17:13:13 +0200 Subject: [PATCH 1/4] mego: faster --- src/multiego/contacts.py | 67 ++++++++++++++++++++-------------------- 1 file changed, 33 insertions(+), 34 deletions(-) diff --git a/src/multiego/contacts.py b/src/multiego/contacts.py index 828e50a7..bc5170e9 100644 --- a/src/multiego/contacts.py +++ b/src/multiego/contacts.py @@ -19,7 +19,7 @@ Private helpers are prefixed with ``_`` and are not part of the public API. """ -from . import fileio as io +from . import fileio from . import type_definitions from . import _term from .model_config import config @@ -49,7 +49,7 @@ def _index_training_topology(topology, custom_dict): - the topology DataFrame (used later in ``init_meGO_matrices`` to check that all non-hydrogen atom names have been mapped to multi-eGO sb_types) - the ``molecules_idx_sbtype_dictionary`` (used by - ``io.read_molecular_contacts`` to remap raw atom indices to sb_type + ``fileio.read_molecular_contacts`` to remap raw atom indices to sb_type labels) LJ parameter lookup and all ``sbtype_*`` dict extractions are intentionally @@ -139,13 +139,8 @@ def _get_lj_params(topology): DataFrame with columns ``ai`` (atom type string), ``c6``, and ``c12``, indexed 0...N-1 for each atom in the topology. """ - lj_params = pd.DataFrame( - { - "ai": [atom.atom_type for atom in topology.atoms], - "c6": [atom.sigma * 0.1 for atom in topology.atoms], - "c12": [atom.epsilon * 4.184 for atom in topology.atoms], - } - ) + rows = [(a.atom_type, a.sigma * 0.1, a.epsilon * 4.184) for a in topology.atoms] + lj_params = pd.DataFrame(rows, columns=["ai", "c6", "c12"]) return lj_params @@ -225,20 +220,21 @@ def _get_lj14_pairs(topology): DataFrame with columns ``ai``, ``aj``, ``epsilon``, and ``sigma``, concatenated across all molecules. """ - lj14_pairs = pd.DataFrame() - for mol, top in topology.molecules.items(): - pair14 = [ - { - "ai": pair.atom1.type, - "aj": pair.atom2.type, - "c6": pair.type.sigma * 0.1, - "c12": pair.type.epsilon * 4.184, - } - for pair in top[0].adjusts - ] - lj14_pairs = pd.concat([lj14_pairs, pd.DataFrame(pair14)]) - - lj14_pairs = lj14_pairs.reset_index(drop=True) + frames = [ + pd.DataFrame( + [ + { + "ai": pair.atom1.type, + "aj": pair.atom2.type, + "c6": pair.type.sigma * 0.1, + "c12": pair.type.epsilon * 4.184, + } + for pair in top[0].adjusts + ] + ) + for _mol, top in topology.molecules.items() + ] + lj14_pairs = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame(columns=["ai", "aj", "c6", "c12"]) lj14_pairs["epsilon"] = np.where( lj14_pairs["c6"] > 0, lj14_pairs["c6"] ** 2 / (4 * lj14_pairs["c12"]), @@ -488,8 +484,8 @@ def _symmetrize_reference_contacts(contact_matrix, symmetry, topology_df): # ── Group by canonical pair and average ──────────────────────────────── ai_str = contact_matrix["rc_ai"].astype(str) aj_str = contact_matrix["rc_aj"].astype(str) - canon_ai = ai_str.map(lambda x: canonical_map.get(x, x)) - canon_aj = aj_str.map(lambda x: canonical_map.get(x, x)) + canon_ai = ai_str.map(canonical_map).fillna(ai_str) + canon_aj = aj_str.map(canonical_map).fillna(aj_str) group_key = canon_ai.values + "|" + canon_aj.values + "|" + contact_matrix["rc_same_chain"].astype(str).values contact_matrix = contact_matrix.copy() @@ -568,7 +564,8 @@ def _load_reference_matrix(reference, ensemble, args): topol = parmed.load_file(topology_path) lj_data = _get_lj_params(topol) - lj_data_dict = {str(key): val for key, val in zip(lj_data["ai"], lj_data[["c6", "c12"]].values)} + c6_map = dict(zip(lj_data["ai"].astype(str), lj_data["c6"])) + c12_map = dict(zip(lj_data["ai"].astype(str), lj_data["c12"])) ensemble.topology_dataframe["c6"] = lj_data["c6"].to_numpy() ensemble.topology_dataframe["c12"] = lj_data["c12"].to_numpy() @@ -586,7 +583,7 @@ def _load_reference_matrix(reference, ensemble, args): path = matrix_paths[0] name = _path_to_matrix_name(path, args.inputs_dir) - contact_matrix = io.read_molecular_contacts( + contact_matrix = fileio.read_molecular_contacts( path, ensemble.molecules_idx_sbtype_dictionary, reference["reference"], path.endswith(".h5") ) @@ -597,11 +594,11 @@ def _load_reference_matrix(reference, ensemble, args): if args.symmetry: contact_matrix = _symmetrize_reference_contacts(contact_matrix, args.symmetry, ensemble.topology_dataframe) - contact_matrix["c6_i"] = [lj_data_dict[x][0] for x in contact_matrix["rc_ai"]] - contact_matrix["c6_j"] = [lj_data_dict[x][0] for x in contact_matrix["rc_aj"]] + contact_matrix["c6_i"] = contact_matrix["rc_ai"].map(c6_map) + contact_matrix["c6_j"] = contact_matrix["rc_aj"].map(c6_map) contact_matrix["c6"] = np.sqrt(contact_matrix["c6_i"] * contact_matrix["c6_j"]) - contact_matrix["c12_i"] = [lj_data_dict[x][1] for x in contact_matrix["rc_ai"]] - contact_matrix["c12_j"] = [lj_data_dict[x][1] for x in contact_matrix["rc_aj"]] + contact_matrix["c12_i"] = contact_matrix["rc_ai"].map(c12_map) + contact_matrix["c12_j"] = contact_matrix["rc_aj"].map(c12_map) contact_matrix["c12"] = np.sqrt(contact_matrix["c12_i"] * contact_matrix["c12_j"]) # Combination-rule defaults; overridden below where explicit entries exist. @@ -731,7 +728,7 @@ def _load_train_matrix( name = f"{args.system}/{reference['reference']}/{simulation}/{reference['matrix']}".replace("/", "_") if train_name not in computed_contact_matrices: - train_contact_matrices_general[train_name] = io.read_molecular_contacts( + train_contact_matrices_general[train_name] = fileio.read_molecular_contacts( path, ensemble.molecules_idx_sbtype_dictionary, simulation, path.endswith(".h5") ) computed_contact_matrices.add(train_name) @@ -796,7 +793,7 @@ def init_meGO_matrices(ensemble, args, custom_dict): train_contact_matrices_general = {} computed_contact_matrices = set() # set for O(1) membership checks computed_train_topologies = {} # simulation_path → (topology_df, sbtype_dict) - train_topology_dataframe = pd.DataFrame() + train_topology_frames = [] for reference in args.input_refs: # Skip if this (reference, matrix) combination has already been loaded — @@ -837,7 +834,7 @@ def init_meGO_matrices(ensemble, args, custom_dict): computed_train_topologies, ) train_contact_matrices[name] = contact_matrix - train_topology_dataframe = pd.concat([train_topology_dataframe, topology_df], axis=0, ignore_index=True) + train_topology_frames.append(topology_df) ensemble.train_matrix_tuples.append((name, ref_name)) et = time.time() _term.timing(et - st) @@ -846,6 +843,8 @@ def init_meGO_matrices(ensemble, args, custom_dict): del train_contact_matrices_general del computed_train_topologies + train_topology_dataframe = pd.concat(train_topology_frames, axis=0, ignore_index=True) + # Verify that all non-hydrogen atom names in the training topologies have # been mapped to multi-eGO sb_types. Unmapped names indicate missing entries # in from_ff_to_multiego or custom_dict. From 25284005854b3ed4b03d067ddb4924cc83d68f69 Mon Sep 17 00:00:00 2001 From: Carlo Camilloni Date: Thu, 16 Apr 2026 17:23:04 +0200 Subject: [PATCH 2/4] fileio: cleaning --- src/multiego/fileio.py | 213 ++++++++++++----------------------------- 1 file changed, 61 insertions(+), 152 deletions(-) diff --git a/src/multiego/fileio.py b/src/multiego/fileio.py index bab0584e..9efb3131 100644 --- a/src/multiego/fileio.py +++ b/src/multiego/fileio.py @@ -1,4 +1,5 @@ import importlib.metadata +import itertools import numpy as np import pandas as pd import glob @@ -26,12 +27,8 @@ def strip_gz_h5_suffix(filename): - str: The filename without the '.gz' suffix, if it was originally present. Otherwise, the original filename is returned. """ - if filename.endswith(".gz"): + if filename.endswith((".gz", ".h5")): return filename[:-3] - - if filename.endswith(".h5"): - return filename[:-3] - return filename @@ -56,33 +53,15 @@ def check_matrix_compatibility(input_path): Returns: - None: The function returns None but raises an error if incompatible files are found. """ - matrix_paths = glob.glob(f"{input_path}.ndx") - matrix_paths_gz = glob.glob(f"{input_path}.ndx.gz") - matrix_paths_h5 = glob.glob(f"{input_path}.ndx.h5") - stripped_matrix_paths_gz_set = set(map(strip_gz_h5_suffix, matrix_paths_gz)) - stripped_matrix_paths_h5_set = set(map(strip_gz_h5_suffix, matrix_paths_h5)) - matrix_paths_set = set(matrix_paths) - - # Find intersection of the two sets - common_files = matrix_paths_set.intersection(stripped_matrix_paths_gz_set) - - # Check if there are any common elements and raise an error if there are - if common_files: - raise ValueError(f"Error: Some files have both text and gz versions: {common_files}") - - # Find intersection of the two sets - common_files = matrix_paths_set.intersection(stripped_matrix_paths_h5_set) - - # Check if there are any common elements and raise an error if there are - if common_files: - raise ValueError(f"Error: Some files have both text and hdf5 versions: {common_files}") - - # Find intersection of the two sets - common_files = stripped_matrix_paths_gz_set.intersection(stripped_matrix_paths_h5_set) - - # Check if there are any common elements and raise an error if there are - if common_files: - raise ValueError(f"Error: Some files have both gz and hdf5 versions: {common_files}") + format_sets = { + "text (.ndx)": set(glob.glob(f"{input_path}.ndx")), + "gz (.ndx.gz)": set(map(strip_gz_h5_suffix, glob.glob(f"{input_path}.ndx.gz"))), + "hdf5 (.h5)": set(map(strip_gz_h5_suffix, glob.glob(f"{input_path}.ndx.h5"))), + } + for (n1, s1), (n2, s2) in itertools.combinations(format_sets.items(), 2): + overlap = s1 & s2 + if overlap: + raise ValueError(f"Error: Some files exist in both {n1} and {n2} formats: {overlap}") def check_mat_name(mat_name, ref): @@ -194,18 +173,13 @@ def parse_symmetry_list(symmetry_list): symmetry = [] for line in symmetry_list: - if "#" in line: - line = line[: line.index("#")] - line = line.replace("\n", "") - line = line.strip() + line = line.split("#", 1)[0].strip() if not line: continue - line = line.split(" ") - line = [x for x in line if x] - if len(line) < 3: + parts = line.split() + if len(parts) < 3: continue - - symmetry.append(line) + symmetry.append(parts) return symmetry @@ -236,21 +210,19 @@ def read_molecular_contacts(path, ensemble_molecules_idx_sbtype_dictionary, simu contact_matrix["learned"] = contact_matrix["learned"].astype(bool) - # Validation checks using `query` for more efficient conditional filtering - if contact_matrix.query("probability < 0 or probability > 1").shape[0] > 0: + # Validation — extract arrays once for all checks (avoids 5 separate DataFrame scans) + probs = contact_matrix["probability"].to_numpy() + dists = contact_matrix["distance"].to_numpy() + cuts = contact_matrix["cutoff"].to_numpy() + if np.any((probs < 0) | (probs > 1)): raise ValueError("ERROR: Probabilities should be between 0 and 1.") - - if contact_matrix.query("distance < 0 or distance > cutoff").shape[0] > 0: - raise ValueError("ERROR: Distances should be between 0 and cutoff.") - - if contact_matrix.query("cutoff < 0").shape[0] > 0: + if np.any(cuts < 0): raise ValueError("ERROR: Cutoff values cannot be negative.") - - # Check for NaN or infinite values in critical columns - if contact_matrix[["probability", "distance", "cutoff"]].isnull().any().any(): + if np.any((dists < 0) | (dists > cuts)): + raise ValueError("ERROR: Distances should be between 0 and cutoff.") + if np.any(np.isnan(probs) | np.isnan(dists) | np.isnan(cuts)): raise ValueError("ERROR: The matrix contains NaN values.") - - if np.isinf(contact_matrix[["probability", "distance", "cutoff"]].values).any(): + if np.any(np.isinf(probs) | np.isinf(dists) | np.isinf(cuts)): raise ValueError("ERROR: The matrix contains INF values.") molecule_names_dictionary = { @@ -420,36 +392,36 @@ def write_output_readme(meGO_LJ, parameters, output_dir, stat_str): def print_stats(meGO_LJ): # it would be nice to cycle over molecule types and print an half matrix with all the relevant information - intrad_contacts = len(meGO_LJ.loc[(meGO_LJ["same_chain"])]) - interm_contacts = len(meGO_LJ.loc[~(meGO_LJ["same_chain"])]) - intrad_a_contacts = len(meGO_LJ.loc[(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)]) - interm_a_contacts = len(meGO_LJ.loc[~(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)]) + intra = meGO_LJ[meGO_LJ["same_chain"]] + inter = meGO_LJ[~meGO_LJ["same_chain"]] + intra_att = intra[intra["epsilon"] > 0.0] + inter_att = inter[inter["epsilon"] > 0.0] + + intrad_contacts = len(intra) + interm_contacts = len(inter) + intrad_a_contacts = len(intra_att) + interm_a_contacts = len(inter_att) intrad_r_contacts = intrad_contacts - intrad_a_contacts interm_r_contacts = interm_contacts - interm_a_contacts - intrad_a_ave_contacts = 0.000 - intrad_a_min_contacts = 0.000 - intrad_a_max_contacts = 0.000 - intrad_a_s_min_contacts = 0.000 - intrad_a_s_max_contacts = 0.000 - interm_a_ave_contacts = 0.000 - interm_a_min_contacts = 0.000 - interm_a_max_contacts = 0.000 - interm_a_s_min_contacts = 0.000 - interm_a_s_max_contacts = 0.000 + + intrad_a_ave_contacts = intrad_a_min_contacts = intrad_a_max_contacts = 0.0 + intrad_a_s_min_contacts = intrad_a_s_max_contacts = 0.0 + interm_a_ave_contacts = interm_a_min_contacts = interm_a_max_contacts = 0.0 + interm_a_s_min_contacts = interm_a_s_max_contacts = 0.0 if intrad_a_contacts > 0: - intrad_a_ave_contacts = meGO_LJ["epsilon"].loc[(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].mean() - intrad_a_min_contacts = meGO_LJ["epsilon"].loc[(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].min() - intrad_a_max_contacts = meGO_LJ["epsilon"].loc[(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].max() - intrad_a_s_min_contacts = meGO_LJ["sigma"].loc[(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].min() - intrad_a_s_max_contacts = meGO_LJ["sigma"].loc[(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].max() + intrad_a_ave_contacts = intra_att["epsilon"].mean() + intrad_a_min_contacts = intra_att["epsilon"].min() + intrad_a_max_contacts = intra_att["epsilon"].max() + intrad_a_s_min_contacts = intra_att["sigma"].min() + intrad_a_s_max_contacts = intra_att["sigma"].max() if interm_a_contacts > 0: - interm_a_ave_contacts = meGO_LJ["epsilon"].loc[~(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].mean() - interm_a_min_contacts = meGO_LJ["epsilon"].loc[~(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].min() - interm_a_max_contacts = meGO_LJ["epsilon"].loc[~(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].max() - interm_a_s_min_contacts = meGO_LJ["sigma"].loc[~(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].min() - interm_a_s_max_contacts = meGO_LJ["sigma"].loc[~(meGO_LJ["same_chain"]) & (meGO_LJ["epsilon"] > 0.0)].max() + interm_a_ave_contacts = inter_att["epsilon"].mean() + interm_a_min_contacts = inter_att["epsilon"].min() + interm_a_max_contacts = inter_att["epsilon"].max() + interm_a_s_min_contacts = inter_att["sigma"].min() + interm_a_s_max_contacts = inter_att["sigma"].max() stat_str = f""" \t- LJ parameterization completed for a total of {len(meGO_LJ)} contacts. @@ -567,10 +539,7 @@ def make_header(parameters, write_header): header += ";\t- {:<26} :\n".format(parameter) for line in value: header += f";\t - {' '.join(line)}\n" - elif isinstance(value, list): - value = np.array(value, dtype=str) - header += ";\t- {:<26} = {:<20}\n".format(parameter, ", ".join(value)) - elif type(value) is np.ndarray: + elif isinstance(value, (list, np.ndarray)): value = np.array(value, dtype=str) header += ";\t- {:<26} = {:<20}\n".format(parameter, ", ".join(value)) elif isinstance(value, dict): @@ -691,14 +660,7 @@ def create_output_directories(parameters, out_dir): output_folder : str The path to the output directory """ - if not os.path.exists(parameters.outputs_dir) and not os.path.isdir(parameters.outputs_dir): - os.mkdir(parameters.outputs_dir) - if not os.path.exists(f"{parameters.outputs_dir}/{parameters.system}") and not os.path.isdir( - f"{parameters.outputs_dir}/{parameters.system}" - ): - os.mkdir(f"{parameters.outputs_dir}/{parameters.system}") - if not os.path.isdir(out_dir) and not os.path.exists(out_dir): - os.mkdir(out_dir) + os.makedirs(out_dir, exist_ok=True) def check_files_existence(args): @@ -826,70 +788,17 @@ def sort_LJ(meGO_ensemble, meGO_LJ): meGO_LJ = meGO_LJ[(meGO_LJ["ai"].cat.codes <= meGO_LJ["aj"].cat.codes)].copy() - ( - meGO_LJ["ai"], - meGO_LJ["aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["number_ai"], - meGO_LJ["number_aj"], - ) = np.where( - (meGO_LJ["molecule_name_ai"].astype(str) <= meGO_LJ["molecule_name_aj"].astype(str)), - [ - meGO_LJ["ai"], - meGO_LJ["aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["number_ai"], - meGO_LJ["number_aj"], - ], - [ - meGO_LJ["aj"], - meGO_LJ["ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["number_aj"], - meGO_LJ["number_ai"], - ], - ) - - ( - meGO_LJ["ai"], - meGO_LJ["aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["number_ai"], - meGO_LJ["number_aj"], - ) = np.where( - meGO_LJ["molecule_name_ai"] == meGO_LJ["molecule_name_aj"], - np.where( - meGO_LJ["number_ai"] <= meGO_LJ["number_aj"], - [ - meGO_LJ["ai"], - meGO_LJ["aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["number_ai"], - meGO_LJ["number_aj"], - ], - [ - meGO_LJ["aj"], - meGO_LJ["ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["number_aj"], - meGO_LJ["number_ai"], - ], - ), - [ - meGO_LJ["ai"], - meGO_LJ["aj"], - meGO_LJ["molecule_name_ai"], - meGO_LJ["molecule_name_aj"], - meGO_LJ["number_ai"], - meGO_LJ["number_aj"], - ], + # Canonical ordering: swap (ai, aj) so that molecule_name_ai <= molecule_name_aj, + # and within the same molecule so that number_ai <= number_aj. + cols_i = ["ai", "molecule_name_ai", "number_ai"] + cols_j = ["aj", "molecule_name_aj", "number_aj"] + swap_mol = meGO_LJ["molecule_name_ai"].astype(str) > meGO_LJ["molecule_name_aj"].astype(str) + swap_same = (meGO_LJ["molecule_name_ai"] == meGO_LJ["molecule_name_aj"]) & ( + meGO_LJ["number_ai"] > meGO_LJ["number_aj"] ) + swap = swap_mol | swap_same + if swap.any(): + meGO_LJ.loc[swap, cols_i + cols_j] = meGO_LJ.loc[swap, cols_j + cols_i].values meGO_LJ.sort_values(by=["molecule_name_ai", "molecule_name_aj", "number_ai", "number_aj"], inplace=True) From 9d09a1f4646a03afbad91692388a6972b9a78742 Mon Sep 17 00:00:00 2001 From: Carlo Camilloni Date: Thu, 16 Apr 2026 20:48:14 +0200 Subject: [PATCH 3/4] codeql config --- .github/codeql-config-python.yml | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 .github/codeql-config-python.yml diff --git a/.github/codeql-config-python.yml b/.github/codeql-config-python.yml new file mode 100644 index 00000000..bd11a704 --- /dev/null +++ b/.github/codeql-config-python.yml @@ -0,0 +1,2 @@ +paths: + - src From 7b5ad1aff2926d34ef68f3535f0ee5397ffa994a Mon Sep 17 00:00:00 2001 From: Carlo Camilloni Date: Thu, 16 Apr 2026 20:53:30 +0200 Subject: [PATCH 4/4] lj: cleanup --- src/multiego/lj.py | 66 +++++++++++++++++----------------------------- 1 file changed, 24 insertions(+), 42 deletions(-) diff --git a/src/multiego/lj.py b/src/multiego/lj.py index 311f7d87..9c4cd890 100644 --- a/src/multiego/lj.py +++ b/src/multiego/lj.py @@ -3,7 +3,7 @@ from . import bonded from . import _term from .model_config import config -from . import fileio as io +from . import fileio import numpy as np import pandas as pd @@ -88,6 +88,9 @@ def set_sig_epsilon(meGO_LJ, parameters): meGO_LJ["epsilon"] = meGO_LJ["epsilon_prior"] meGO_LJ["sigma"] = meGO_LJ["sigma_prior"] + # Precompute shared sub-expression (used in both attractive and repulsive conditions) + eps_prior_pos = np.maximum(0.0, meGO_LJ["epsilon_prior"]) + # Attractive interactions # These are defined only if the training probability is greater than MD_threshold and # by comparing them with RC_probabilities so that the resulting epsilon is between eps_min and eps_0 @@ -96,8 +99,8 @@ def set_sig_epsilon(meGO_LJ, parameters): > meGO_LJ["limit_rc_att"] * np.maximum(meGO_LJ["rc_probability"], meGO_LJ["rc_threshold"]) ) & (meGO_LJ["probability"] > meGO_LJ["md_threshold"]) with np.errstate(divide="ignore"): - meGO_LJ.loc[condition, "epsilon"] = np.maximum(0.0, meGO_LJ["epsilon_prior"]) - ( - (meGO_LJ["epsilon_0"] - np.maximum(0.0, meGO_LJ["epsilon_prior"])) / np.log(meGO_LJ["rc_threshold"]) + meGO_LJ.loc[condition, "epsilon"] = eps_prior_pos - ( + (meGO_LJ["epsilon_0"] - eps_prior_pos) / np.log(meGO_LJ["rc_threshold"]) ) * (np.log(meGO_LJ["probability"] / (np.maximum(meGO_LJ["rc_probability"], meGO_LJ["rc_threshold"])))) meGO_LJ.loc[condition, "learned"] = 1 meGO_LJ.loc[condition, "sigma"] = meGO_LJ["distance"] / 2.0 ** (1.0 / 6.0) @@ -112,9 +115,10 @@ def set_sig_epsilon(meGO_LJ, parameters): & (meGO_LJ["probability"] > meGO_LJ["md_threshold"]) & (meGO_LJ["rc_probability"] > meGO_LJ["md_threshold"]) ) - meGO_LJ.loc[condition, "epsilon"] = (-meGO_LJ["rep"] * (meGO_LJ["distance"] / meGO_LJ["rc_distance"]) ** 12).clip( - lower=-20 * meGO_LJ["rep"], upper=-0.05 * meGO_LJ["rep"] - )[condition] + sub = meGO_LJ.loc[condition] + meGO_LJ.loc[condition, "epsilon"] = (-sub["rep"] * (sub["distance"] / sub["rc_distance"]) ** 12).clip( + lower=-20 * sub["rep"], upper=-0.05 * sub["rep"] + ) meGO_LJ.loc[condition, "learned"] = 1 # for repulsive interactions reset sigma to its effective value for consistent merging @@ -263,7 +267,7 @@ def init_LJ_datasets(meGO_ensemble, matrices, args): chunks = [] for name, ref_name in meGO_ensemble.train_matrix_tuples: - if ref_name not in matrices["reference_matrices"].keys(): + if ref_name not in matrices["reference_matrices"]: raise RuntimeError( f'Encountered error while trying to find {ref_name} in reference matrices {matrices["reference_matrices"].keys()}' ) @@ -347,7 +351,7 @@ def init_LJ_datasets(meGO_ensemble, matrices, args): * train_dataset["aj"].map(meGO_ensemble.sbtype_mg_c6_dict) ) ) ** (1 / 12) - train_dataset["mg_sigma"] = pd.Series(pairwise_mg_sigma) + train_dataset["mg_sigma"] = pairwise_mg_sigma # default (mg) epsilon pairwise_mg_epsilon = ( @@ -360,7 +364,7 @@ def init_LJ_datasets(meGO_ensemble, matrices, args): * train_dataset["aj"].map(meGO_ensemble.sbtype_mg_c12_dict) ) ) - train_dataset["mg_epsilon"] = pd.Series(pairwise_mg_epsilon) + train_dataset["mg_epsilon"] = pairwise_mg_epsilon # Apply special non-local interaction rules from type_definitions. # @@ -461,6 +465,12 @@ def init_LJ_datasets(meGO_ensemble, matrices, args): return train_dataset +def _add_c6_c12(df): + """Compute c6/c12 columns in-place from epsilon and sigma.""" + df["c6"] = np.where(df["epsilon"] < 0.0, 0.0, 4 * df["epsilon"] * df["sigma"] ** 6) + df["c12"] = np.where(df["epsilon"] < 0.0, -df["epsilon"], 4 * df["epsilon"] * df["sigma"] ** 12) + + def generate_LJ(meGO_ensemble, train_dataset, parameters): """ Generates production LJ interactions from the training dataset. @@ -564,7 +574,7 @@ def generate_LJ(meGO_ensemble, train_dataset, parameters): meGO_LJ = meGO_LJ[needed_fields] - stat_str = io.print_stats(meGO_LJ) + stat_str = fileio.print_stats(meGO_LJ) # --- identify duplicates ON ORIGINAL DATA --- dup_mask = meGO_LJ.duplicated(subset=["ai", "aj"], keep=False) @@ -596,26 +606,7 @@ def generate_LJ(meGO_ensemble, train_dataset, parameters): meGO_LJ = meGO_LJ.loc[(~meGO_LJ["same_chain"])] # Add MG prior interactions for pairs not learned in any other way - needed_fields_mg = [ - "molecule_name_ai", - "ai", - "molecule_name_aj", - "aj", - "probability", - "same_chain", - "source", - "reference", - "rc_probability", - "sigma", - "epsilon", - "bond_distance", - "rep", - "mg_sigma", - "mg_epsilon", - "md_threshold", - "rc_threshold", - "learned", - ] + needed_fields_mg = [f for f in needed_fields if f not in {"sigma_prior", "epsilon_prior"}] basic_LJ = mg.generate_MG_LJ(meGO_ensemble)[needed_fields_mg] meGO_LJ = pd.concat([meGO_LJ, basic_LJ]) @@ -627,7 +618,7 @@ def generate_LJ(meGO_ensemble, train_dataset, parameters): "molecule_name_ai": "molecule_name_aj", "molecule_name_aj": "molecule_name_ai", } - ).copy() + ) meGO_LJ = pd.concat([meGO_LJ, inverse_meGO_LJ], axis=0, sort=False, ignore_index=True) meGO_LJ["ai"] = meGO_LJ["ai"].astype("category") @@ -638,17 +629,8 @@ def generate_LJ(meGO_ensemble, train_dataset, parameters): meGO_LJ.sort_values(by=["ai", "aj", "learned", "sigma"], ascending=[True, True, False, True], inplace=True) meGO_LJ = meGO_LJ.drop_duplicates(subset=["ai", "aj"], keep="first") - meGO_LJ["c6"] = np.where(meGO_LJ["epsilon"] < 0.0, 0.0, 4 * meGO_LJ["epsilon"] * (meGO_LJ["sigma"] ** 6)) - meGO_LJ["c12"] = np.where( - meGO_LJ["epsilon"] < 0.0, -meGO_LJ["epsilon"], 4 * meGO_LJ["epsilon"] * (meGO_LJ["sigma"] ** 12) - ) - - meGO_LJ_14["c6"] = np.where( - meGO_LJ_14["epsilon"] < 0.0, 0.0, 4 * meGO_LJ_14["epsilon"] * (meGO_LJ_14["sigma"] ** 6) - ) - meGO_LJ_14["c12"] = np.where( - meGO_LJ_14["epsilon"] < 0.0, -meGO_LJ_14["epsilon"], 4 * meGO_LJ_14["epsilon"] * (meGO_LJ_14["sigma"] ** 12) - ) + _add_c6_c12(meGO_LJ) + _add_c6_c12(meGO_LJ_14) et = time.time() _term.timing(et - st)