From 6035fe9e0623bbbbde46cc91c9e16ff12b012aac Mon Sep 17 00:00:00 2001 From: Patrik Kula Date: Wed, 8 Apr 2026 16:44:46 +0200 Subject: [PATCH 1/3] Added script for form factor dataset creation --- developer/make_train_data.py | 343 +++++++++++++++++++++++++++++++++++ 1 file changed, 343 insertions(+) create mode 100644 developer/make_train_data.py diff --git a/developer/make_train_data.py b/developer/make_train_data.py new file mode 100644 index 000000000..3d0192ffd --- /dev/null +++ b/developer/make_train_data.py @@ -0,0 +1,343 @@ +from fairmd.lipids.core import initialize_databank +from fairmd.lipids.databankio import download_resource_from_uri +from fairmd.lipids.molecules import Lipid, Molecule, lipids_set +from fairmd.lipids.api import get_eqtimes, get_thickness, UniverseConstructor#, get_mean_ApL +from fairmd.lipids.analib.maicos import DensityPlanar, FormFactorPlanar, is_system_suitable_4_maicos, first_last_carbon, traj_centering_for_maicos_gromacs, traj_centering_for_maicos_mda_parallel, traj_centering_for_maicos_mda +from fairmd.lipids import FMDL_SIMU_PATH, FMDL_MAICOS_NCORES +from fairmd.lipids.auxiliary import mollib +from maicos.core.base import AnalysisCollection +import logging +import MDAnalysis as mda +import fairmd +import numpy as np +import matplotlib.pyplot as plt +import h5py +import os + + +WATER_TO_LIPID_RATIO_THRESHOLD = 20 +RECOMPUTE = True + +systems = initialize_databank() +logger = logging.getLogger(__name__) + +#### OLDER version get_mean_ApL new one is giving an error - Ill test it and send bug report later +from fairmd.lipids.core import System +import json +def get_mean_ApL(system) -> float: # noqa: N802 (API name) + """ + Calculate average area per lipid for a system. + + :param system: Simulation object. + + :return: area per lipid (Å^2) + """ + path = os.path.join(FMDL_SIMU_PATH, system["path"], "apl.json") + try: + with open(path) as f: + data = json.load(f) + except FileNotFoundError as e: + msg = "Area per lipid data is absent for system #{}".format(system["ID"]) + raise FileNotFoundError(msg) from e + except json.JSONDecodeError as e: + msg = "Area per lipid data for system #{} in {} is invalid.".format(system["ID"], path) + raise ValueError(msg) from e + vals = np.array(list(data.values())) + return vals.mean() + +def is_suitable(system): + flag = True + if "WARNINGS" in system.keys() and type(system["WARNINGS"]) == dict: + flag = False + if not is_system_suitable_4_maicos(system): + print(f"system {system} not suitable for maicos") + flag = False + return flag +# if "PBC" in system["WARNINGS"].keys(): +# continue +# elif "ORIENTATION" in system["WARNINGS"].keys(): +# continue + +def get_scalar_properties(system): + flag = True + try: + ApL = get_mean_ApL(system) + except: + print(f"System {system} - can't load ApL") + flag = False + ApL = -1 + + try: + thickness = get_thickness(system) + except: + print(f"System {system} - can't load thickness") + flag = False + thickness = -1 + return ApL, thickness, flag + +def center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger): + if "gromacs" in system["SOFTWARE"]: + traj_centered = traj_centering_for_maicos_gromacs( + spath, + tpr_name=uc.paths["top"], + trj_name=uc.paths["traj"], + last_atom=last_atom, + g3_atom=g3_atom, + eq_time=eq_time, + recompute=RECOMPUTE, + ) + else: + if FMDL_MAICOS_NCORES != 1: + try: + n_jobs = FMDL_MAICOS_NCORES if FMDL_MAICOS_NCORES is not None else -1 + logger.info(f"Using parallel trajectory centering (n_jobs={n_jobs})") + traj_centered = traj_centering_for_maicos_mda_parallel( + u, + spath, + last_atom, + eq_time, + n_jobs=n_jobs, + recompute=RECOMPUTE, + logger=logger, + show_progress=True, + ) + except ImportError: + logger.warning("joblib not available, falling back to sequential centering") + traj_centered = traj_centering_for_maicos_mda( + u, + spath, + last_atom, + eq_time, + recompute=RECOMPUTE, + logger=logger, + ) + else: + logger.info("Using sequential trajectory centering (FMDL_MAICOS_NCORES=1)") + traj_centered = traj_centering_for_maicos_mda( + u, + spath, + last_atom, + eq_time, + recompute=RECOMPUTE, + logger=logger, + ) + +def separate_lipid_atoms(mapping_dict): + head_atoms = '' + tail_atoms = '' + backbone_atoms = '' + for atom in mapping_dict.keys(): + fragment = mapping_dict[atom]['FRAGMENT'] + atom_name = mapping_dict[atom]['ATOMNAME'] + if fragment == 'headgroup': + head_atoms += atom_name + ' ' + elif fragment == 'sn-1' or fragment == 'sn-2' or fragment == 'tail': + tail_atoms += atom_name + ' ' + elif fragment == 'glycerol backbone': + backbone_atoms += atom_name + ' ' + else: + print(f'What is this? {fragment} - {atom_name}') + return (head_atoms, tail_atoms, backbone_atoms) + +def create_fragment_selectors(lipid_names): + head_selector, tail_selector, backbone_selector = 'name ', 'name ', 'name ' + for lipid in lipid_names: + lipid_class = Lipid(lipid) + lipid_class.register_mapping() + + mapping_dict = lipid_class.mapping_dict + head_atoms, tail_atoms, backbone_atoms = separate_lipid_atoms(mapping_dict) + + head_selector += head_atoms + tail_selector += tail_atoms + backbone_selector += backbone_atoms + + return [head_selector, tail_selector, backbone_selector] + + +class HDF5LipidWriter: + def __init__(self, filename, overwrite_file=False): + self.filename = filename + + if overwrite_file and os.path.exists(self.filename): + print(f"Clearing existing file: {self.filename}") + os.remove(self.filename) + + def _get_file_mode(self): + return 'a' + + def save_system(self, system, scalar_data, form_factor, total_dens, mol_densities, frag_densities): + """ + Opens, writes one system, and closes immediately to prevent data loss. + """ + with h5py.File(self.filename, self._get_file_mode()) as f: + sys_id = str(system["ID"]) + + # Prevent overwriting if system ID already exists + if sys_id in f: + print(f"Warning: System {sys_id} already exists in HDF5. Overwriting.") + del f[sys_id] + + grp = f.create_group(sys_id) + + grp.attrs['path'] = system.get("path", "") + grp.attrs['ApL'] = scalar_data.get('ApL', 0) + grp.attrs['thickness'] = scalar_data.get('thickness', 0) + + axis_grp = grp.create_group('axis') + self._write_dataset(axis_grp, 'q_pos', form_factor[0]) + self._write_dataset(axis_grp, 'r_pos', total_dens[0]) + + ff_grp = grp.create_group('form_factor') + self._write_dataset(ff_grp, 'profile', form_factor[1]) + self._write_dataset(ff_grp, 'dprofile', form_factor[2]) + + td_grp = grp.create_group('density_total') + self._write_dataset(td_grp, 'profile', total_dens[1]) + self._write_dataset(td_grp, 'dprofile', total_dens[2]) + + frag_labels = ['head', 'tail', 'backbone'] + frag_grp = grp.create_group('density_fragments') + for i, (profile, dprofile) in enumerate(frag_densities): + label = frag_labels[i] if i < len(frag_labels) else f"frag_{i}" + sub_grp = frag_grp.create_group(label) + self._write_dataset(sub_grp, 'profile', profile) + self._write_dataset(sub_grp, 'dprofile', dprofile) + + mol_grp = grp.create_group('density_molecules') + for i, (profile, dprofile) in enumerate(mol_densities): + sub_grp = mol_grp.create_group(f"mol_{i}") + self._write_dataset(sub_grp, 'profile', profile) + self._write_dataset(sub_grp, 'dprofile', dprofile) + + + def _write_dataset(self, group, name, data): + if data is not None: + group.create_dataset(name, data=np.array(data), compression="gzip", compression_opts=4) + +writer = HDF5LipidWriter("lipid_FF_dataset.h5") +test_flag = True + +count = 0 +print(f"Number of systems: {len(systems)}") +for system in systems: + if system['TRAJECTORY_SIZE'] > 10**8 and test_flag: #For testing purpouses + continue + + if not is_suitable(system): + continue + + ApL, thickness, flag = get_scalar_properties(system) + #if not flag: + # continue + print(ApL, thickness) + scalar_info = { + 'ApL': ApL, + 'thickness': thickness + } + + water_n = system['COMPOSITION']['SOL']['COUNT'] + molecules = system['COMPOSITION'].keys() + lipid_names = [] + lipid_n = 0 + for molecule in molecules: + if molecule in lipids_set: + lipid_n += sum(system['COMPOSITION'][molecule]['COUNT']) + lipid_names.append(molecule) + + water_to_lipid_ratio = water_n / lipid_n + print(water_to_lipid_ratio) + + + if water_to_lipid_ratio < WATER_TO_LIPID_RATIO_THRESHOLD: + continue + + try: + uc = UniverseConstructor(system) + uc.download_mddata() + except: + continue + + eq_time = float(system["TIMELEFTOUT"]) * 1000 + + last_atom, g3_atom = first_last_carbon(system, logger) + + spath = f'{FMDL_SIMU_PATH}/{system["path"]}' + u = uc.build_universe() + + traj_centered = center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger) + + + u.load_new(traj_centered, format="XTC") + u.guess_TopologyAttrs(force_guess=["elements"]) + mollib.guess_elements(system, u) + + bin_width = 0.3 + + + L_min = u.dimensions[2] + for ts in u.trajectory: + L_min = min(L_min, ts.dimensions[2]) + + base_options = {"unwrap": False, "bin_width": bin_width, "pack": False} + zlim = {"zmin": -L_min / 2, "zmax": L_min / 2} + dens_options = {**zlim, **base_options} + + save_data = [] + + print(f'Calculating form factor') + form_factor = FormFactorPlanar( + atomgroup=u.atoms, + **base_options, + zmin=None, + zmax=None, + ).run() + ff = (form_factor.results.bin_pos, form_factor.results.profile, form_factor.results.dprofile) + + print(f'Calculating total density') + dens_total_runner = DensityPlanar( + u.atoms, + dens="electron", + **dens_options, + ).run() + dens_total = (dens_total_runner.results.bin_pos, dens_total_runner.results.profile, dens_total_runner.results.dprofile) + + molecule_types_selector = [f"resname {system['COMPOSITION'][molecule_type]['NAME']}" for molecule_type in molecules] + dens_molecule = [] + for selector in molecule_types_selector: + print(f'Calculating {selector.split(" ")[1]} density') + molecule_group = u.select_atoms(selector) + dens_molecule_runner = DensityPlanar( + molecule_group, + dens="electron", + **dens_options, + ).run() + dens = (dens_molecule_runner.results.profile, dens_molecule_runner.results.dprofile) + dens_molecule.append(dens) + + fragment_selectors = create_fragment_selectors(lipid_names) + dens_fragment = [] + frag_labels = ['head', 'tail', 'backbone'] + for i, selector in enumerate(fragment_selectors): + print(f'Calculating {frag_labels[i]} density') + fragment_group = u.select_atoms(selector) + dens_fragment_runner = DensityPlanar( + fragment_group, + dens="electron", + **dens_options, + ).run() + dens = (dens_fragment_runner.results.profile, dens_fragment_runner.results.dprofile) + dens_fragment.append(dens) + + writer.save_system( + system=system, + scalar_data=scalar_info, + form_factor=ff, + total_dens=dens_total, + mol_densities=dens_molecule, + frag_densities=dens_fragment + ) + + count += 1 + +print(f"Final number of systems saved into dataset: {count}") From 4cd699ad8026f12c36195377ad287d543ea71b0c Mon Sep 17 00:00:00 2001 From: Patrik Kula Date: Thu, 9 Apr 2026 15:00:27 +0200 Subject: [PATCH 2/3] Added shebang; added file-level, function and class docstrings; linting; removed commented code, dev prints and provisionary get_mean_ApL function --- developer/make_train_data.py | 313 ++++++++++++++++++++++------------- 1 file changed, 202 insertions(+), 111 deletions(-) diff --git a/developer/make_train_data.py b/developer/make_train_data.py index 3d0192ffd..7fd49ad20 100644 --- a/developer/make_train_data.py +++ b/developer/make_train_data.py @@ -1,51 +1,47 @@ -from fairmd.lipids.core import initialize_databank -from fairmd.lipids.databankio import download_resource_from_uri -from fairmd.lipids.molecules import Lipid, Molecule, lipids_set -from fairmd.lipids.api import get_eqtimes, get_thickness, UniverseConstructor#, get_mean_ApL -from fairmd.lipids.analib.maicos import DensityPlanar, FormFactorPlanar, is_system_suitable_4_maicos, first_last_carbon, traj_centering_for_maicos_gromacs, traj_centering_for_maicos_mda_parallel, traj_centering_for_maicos_mda -from fairmd.lipids import FMDL_SIMU_PATH, FMDL_MAICOS_NCORES -from fairmd.lipids.auxiliary import mollib -from maicos.core.base import AnalysisCollection +#!/usr/bin/env python3 + +""" +Script for processing lipid databank systems and saving analysis results to HDF5. + +This script filters systems by water-to-lipid ratio, calculates density +profiles using maicos, and exports the data for machine learning training. +""" + import logging -import MDAnalysis as mda -import fairmd -import numpy as np -import matplotlib.pyplot as plt -import h5py import os +import h5py +import numpy as np -WATER_TO_LIPID_RATIO_THRESHOLD = 20 -RECOMPUTE = True +from fairmd.lipids import FMDL_MAICOS_NCORES, FMDL_SIMU_PATH +from fairmd.lipids.analib.maicos import ( + DensityPlanar, + FormFactorPlanar, + first_last_carbon, + is_system_suitable_4_maicos, + traj_centering_for_maicos_gromacs, + traj_centering_for_maicos_mda, + traj_centering_for_maicos_mda_parallel, +) +from fairmd.lipids.api import UniverseConstructor, get_thickness, get_mean_ApL +from fairmd.lipids.auxiliary import mollib +from fairmd.lipids.core import initialize_databank +from fairmd.lipids.molecules import Lipid, lipids_set -systems = initialize_databank() -logger = logging.getLogger(__name__) -#### OLDER version get_mean_ApL new one is giving an error - Ill test it and send bug report later -from fairmd.lipids.core import System -import json -def get_mean_ApL(system) -> float: # noqa: N802 (API name) +def is_suitable(system): """ - Calculate average area per lipid for a system. + Checks if a simulation system is suitable for analysis with maicos. - :param system: Simulation object. + Checks for the presence of a 'WARNINGS' dictionary in the system metadata + and verifies compatibility using the internal fairmd suitability check. - :return: area per lipid (Å^2) - """ - path = os.path.join(FMDL_SIMU_PATH, system["path"], "apl.json") - try: - with open(path) as f: - data = json.load(f) - except FileNotFoundError as e: - msg = "Area per lipid data is absent for system #{}".format(system["ID"]) - raise FileNotFoundError(msg) from e - except json.JSONDecodeError as e: - msg = "Area per lipid data for system #{} in {} is invalid.".format(system["ID"], path) - raise ValueError(msg) from e - vals = np.array(list(data.values())) - return vals.mean() + Args: + system (dict): The system dictionary from the databank. -def is_suitable(system): + Returns: + bool: True if the system is suitable for analysis, False otherwise. + """ flag = True if "WARNINGS" in system.keys() and type(system["WARNINGS"]) == dict: flag = False @@ -53,12 +49,25 @@ def is_suitable(system): print(f"system {system} not suitable for maicos") flag = False return flag -# if "PBC" in system["WARNINGS"].keys(): -# continue -# elif "ORIENTATION" in system["WARNINGS"].keys(): -# continue + def get_scalar_properties(system): + """ + Retrieves scalar physical properties (ApL and thickness) for a given system. + + Attempts to load the mean Area per Lipid (ApL) and bilayer thickness. + If a property cannot be loaded, it assigns a default value of -1 and + logs a warning. + + Args: + system (dict): The system dictionary from the databank. + + Returns: + tuple: A tuple containing: + - ApL (float): Average area per lipid (Ų). + - thickness (float): Bilayer thickness (Å). + - flag (bool): True if all properties were loaded successfully. + """ flag = True try: ApL = get_mean_ApL(system) @@ -75,7 +84,27 @@ def get_scalar_properties(system): thickness = -1 return ApL, thickness, flag + def center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger): + """ + Centers the simulation trajectory for analysis, handling different software backends. + + Coordinates trajectory centering using either Gromacs commands or MDAnalysis + (sequential or parallel) based on the simulation metadata and environment + configuration. + + Args: + u (MDAnalysis.Universe): The initial MDAnalysis universe. + uc (UniverseConstructor): Object containing simulation path and topology info. + spath (str): Full path to the simulation directory. + last_atom (int/str): Index or name of the last atom for centering reference. + g3_atom (int/str): Index or name of the glycerol 3 atom for orientation. + eq_time (float): Equilibration time to skip in milliseconds. + logger (logging.Logger): Logger instance for status and error reporting. + + Returns: + str: The file path to the newly created centered trajectory file + """ if "gromacs" in system["SOFTWARE"]: traj_centered = traj_centering_for_maicos_gromacs( spath, @@ -121,30 +150,60 @@ def center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger): recompute=RECOMPUTE, logger=logger, ) + return traj_centered + def separate_lipid_atoms(mapping_dict): - head_atoms = '' - tail_atoms = '' - backbone_atoms = '' + """ + Groups lipid atom names into headgroup, tail, and backbone fragments. + + Parses a mapping dictionary (usually from a Lipid class instance) to + categorize atoms based on their structural fragment. + + Args: + mapping_dict (dict): Dictionary mapping atom IDs to names and fragments. + + Returns: + tuple: A triplet of space-separated strings (head_atoms, tail_atoms, backbone_atoms) + containing the atom names for each respective fragment. + """ + head_atoms = "" + tail_atoms = "" + backbone_atoms = "" for atom in mapping_dict.keys(): - fragment = mapping_dict[atom]['FRAGMENT'] - atom_name = mapping_dict[atom]['ATOMNAME'] - if fragment == 'headgroup': - head_atoms += atom_name + ' ' - elif fragment == 'sn-1' or fragment == 'sn-2' or fragment == 'tail': - tail_atoms += atom_name + ' ' - elif fragment == 'glycerol backbone': - backbone_atoms += atom_name + ' ' + fragment = mapping_dict[atom]["FRAGMENT"] + atom_name = mapping_dict[atom]["ATOMNAME"] + if fragment == "headgroup": + head_atoms += atom_name + " " + elif fragment == "sn-1" or fragment == "sn-2" or fragment == "tail": + tail_atoms += atom_name + " " + elif fragment == "glycerol backbone": + backbone_atoms += atom_name + " " else: - print(f'What is this? {fragment} - {atom_name}') + print(f"Invalid atom - {atom} - {atom_name} - {fragment}") return (head_atoms, tail_atoms, backbone_atoms) + def create_fragment_selectors(lipid_names): - head_selector, tail_selector, backbone_selector = 'name ', 'name ', 'name ' + """ + Creates MDAnalysis atom selection strings for lipid fragments across a system. + + Iterates through a list of lipid names, registers their fragment mappings, + and constructs 'name ...' strings used to select specific fragments in + a simulation. + + Args: + lipid_names (list of str): List of lipid molecule names present in the system. + + Returns: + list of str: A list containing three selection strings in the order: + [head_selector, tail_selector, backbone_selector]. + """ + head_selector, tail_selector, backbone_selector = "name ", "name ", "name " for lipid in lipid_names: lipid_class = Lipid(lipid) lipid_class.register_mapping() - + mapping_dict = lipid_class.mapping_dict head_atoms, tail_atoms, backbone_atoms = separate_lipid_atoms(mapping_dict) @@ -156,98 +215,128 @@ def create_fragment_selectors(lipid_names): class HDF5LipidWriter: + """ + A handler for saving lipid simulation analysis results into HDF5 format. + + This class manages the hierarchical storage of form factors, density + profiles, and scalar properties. It is designed for long-running + processes by opening and flushing to the file for every system processed. + """ def __init__(self, filename, overwrite_file=False): + """ + Initializes the writer and optionally clears the existing file. + + Args: + filename (str): Path to the output .h5 file. + overwrite_file (bool): If True, deletes the existing file on initialization. + """ self.filename = filename if overwrite_file and os.path.exists(self.filename): print(f"Clearing existing file: {self.filename}") os.remove(self.filename) - def _get_file_mode(self): - return 'a' - def save_system(self, system, scalar_data, form_factor, total_dens, mol_densities, frag_densities): """ - Opens, writes one system, and closes immediately to prevent data loss. + Saves a single system's results to the HDF5 file. + + Organizes data into groups for metadata, axes, form factors, and + various electron density profiles (total, fragment-based, and molecule-based). + If the system ID already exists, it is overwritten to ensure data integrity. + + Args: + system (dict): System metadata from the databank. + scalar_data (dict): Dictionary containing 'ApL' and 'thickness'. + form_factor (tuple): (q_pos, profile, dprofile) for the form factor. + total_dens (tuple): (r_pos, profile, dprofile) for total electron density. + mol_densities (list of tuples): List of (profile, dprofile) for each molecule type. + frag_densities (list of tuples): List of (profile, dprofile) for lipid fragments. """ - with h5py.File(self.filename, self._get_file_mode()) as f: + with h5py.File(self.filename, "a") as f: sys_id = str(system["ID"]) - - # Prevent overwriting if system ID already exists + if sys_id in f: print(f"Warning: System {sys_id} already exists in HDF5. Overwriting.") del f[sys_id] - + grp = f.create_group(sys_id) - grp.attrs['path'] = system.get("path", "") - grp.attrs['ApL'] = scalar_data.get('ApL', 0) - grp.attrs['thickness'] = scalar_data.get('thickness', 0) + grp.attrs["path"] = system.get("path", "") + grp.attrs["ApL"] = scalar_data.get("ApL", 0) + grp.attrs["thickness"] = scalar_data.get("thickness", 0) + + axis_grp = grp.create_group("axis") + self._write_dataset(axis_grp, "q_pos", form_factor[0]) + self._write_dataset(axis_grp, "r_pos", total_dens[0]) - axis_grp = grp.create_group('axis') - self._write_dataset(axis_grp, 'q_pos', form_factor[0]) - self._write_dataset(axis_grp, 'r_pos', total_dens[0]) - - ff_grp = grp.create_group('form_factor') - self._write_dataset(ff_grp, 'profile', form_factor[1]) - self._write_dataset(ff_grp, 'dprofile', form_factor[2]) + ff_grp = grp.create_group("form_factor") + self._write_dataset(ff_grp, "profile", form_factor[1]) + self._write_dataset(ff_grp, "dprofile", form_factor[2]) - td_grp = grp.create_group('density_total') - self._write_dataset(td_grp, 'profile', total_dens[1]) - self._write_dataset(td_grp, 'dprofile', total_dens[2]) + td_grp = grp.create_group("density_total") + self._write_dataset(td_grp, "profile", total_dens[1]) + self._write_dataset(td_grp, "dprofile", total_dens[2]) - frag_labels = ['head', 'tail', 'backbone'] - frag_grp = grp.create_group('density_fragments') + frag_labels = ["head", "tail", "backbone"] + frag_grp = grp.create_group("density_fragments") for i, (profile, dprofile) in enumerate(frag_densities): label = frag_labels[i] if i < len(frag_labels) else f"frag_{i}" sub_grp = frag_grp.create_group(label) - self._write_dataset(sub_grp, 'profile', profile) - self._write_dataset(sub_grp, 'dprofile', dprofile) + self._write_dataset(sub_grp, "profile", profile) + self._write_dataset(sub_grp, "dprofile", dprofile) - mol_grp = grp.create_group('density_molecules') + mol_grp = grp.create_group("density_molecules") for i, (profile, dprofile) in enumerate(mol_densities): sub_grp = mol_grp.create_group(f"mol_{i}") - self._write_dataset(sub_grp, 'profile', profile) - self._write_dataset(sub_grp, 'dprofile', dprofile) - + self._write_dataset(sub_grp, "profile", profile) + self._write_dataset(sub_grp, "dprofile", dprofile) def _write_dataset(self, group, name, data): + """ + Helper method to write a NumPy array to an HDF5 group with compression. + + Args: + group (h5py.Group): The parent group to write into. + name (str): The name of the dataset. + data (array-like): The numerical data to save. + """ if data is not None: group.create_dataset(name, data=np.array(data), compression="gzip", compression_opts=4) + +WATER_TO_LIPID_RATIO_THRESHOLD = 20 +RECOMPUTE = True + +systems = initialize_databank() +logger = logging.getLogger(__name__) + writer = HDF5LipidWriter("lipid_FF_dataset.h5") test_flag = True count = 0 print(f"Number of systems: {len(systems)}") for system in systems: - if system['TRAJECTORY_SIZE'] > 10**8 and test_flag: #For testing purpouses + if system["TRAJECTORY_SIZE"] > 10**8 and test_flag: # For testing purpouses continue if not is_suitable(system): continue - + ApL, thickness, flag = get_scalar_properties(system) - #if not flag: + # if not flag: # continue - print(ApL, thickness) - scalar_info = { - 'ApL': ApL, - 'thickness': thickness - } - - water_n = system['COMPOSITION']['SOL']['COUNT'] - molecules = system['COMPOSITION'].keys() + scalar_info = {"ApL": ApL, "thickness": thickness} + + water_n = system["COMPOSITION"]["SOL"]["COUNT"] + molecules = system["COMPOSITION"].keys() lipid_names = [] lipid_n = 0 for molecule in molecules: if molecule in lipids_set: - lipid_n += sum(system['COMPOSITION'][molecule]['COUNT']) + lipid_n += sum(system["COMPOSITION"][molecule]["COUNT"]) lipid_names.append(molecule) water_to_lipid_ratio = water_n / lipid_n - print(water_to_lipid_ratio) - if water_to_lipid_ratio < WATER_TO_LIPID_RATIO_THRESHOLD: continue @@ -262,20 +351,18 @@ def _write_dataset(self, group, name, data): last_atom, g3_atom = first_last_carbon(system, logger) - spath = f'{FMDL_SIMU_PATH}/{system["path"]}' + spath = f"{FMDL_SIMU_PATH}/{system['path']}" u = uc.build_universe() traj_centered = center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger) - u.load_new(traj_centered, format="XTC") u.guess_TopologyAttrs(force_guess=["elements"]) mollib.guess_elements(system, u) bin_width = 0.3 - - L_min = u.dimensions[2] + L_min = u.dimensions[2] for ts in u.trajectory: L_min = min(L_min, ts.dimensions[2]) @@ -285,7 +372,7 @@ def _write_dataset(self, group, name, data): save_data = [] - print(f'Calculating form factor') + print("Calculating form factor") form_factor = FormFactorPlanar( atomgroup=u.atoms, **base_options, @@ -294,18 +381,22 @@ def _write_dataset(self, group, name, data): ).run() ff = (form_factor.results.bin_pos, form_factor.results.profile, form_factor.results.dprofile) - print(f'Calculating total density') + print("Calculating total density") dens_total_runner = DensityPlanar( u.atoms, dens="electron", **dens_options, ).run() - dens_total = (dens_total_runner.results.bin_pos, dens_total_runner.results.profile, dens_total_runner.results.dprofile) + dens_total = ( + dens_total_runner.results.bin_pos, + dens_total_runner.results.profile, + dens_total_runner.results.dprofile, + ) molecule_types_selector = [f"resname {system['COMPOSITION'][molecule_type]['NAME']}" for molecule_type in molecules] dens_molecule = [] for selector in molecule_types_selector: - print(f'Calculating {selector.split(" ")[1]} density') + print(f"Calculating {selector.split(' ')[1]} density") molecule_group = u.select_atoms(selector) dens_molecule_runner = DensityPlanar( molecule_group, @@ -317,9 +408,9 @@ def _write_dataset(self, group, name, data): fragment_selectors = create_fragment_selectors(lipid_names) dens_fragment = [] - frag_labels = ['head', 'tail', 'backbone'] + frag_labels = ["head", "tail", "backbone"] for i, selector in enumerate(fragment_selectors): - print(f'Calculating {frag_labels[i]} density') + print(f"Calculating {frag_labels[i]} density") fragment_group = u.select_atoms(selector) dens_fragment_runner = DensityPlanar( fragment_group, @@ -335,9 +426,9 @@ def _write_dataset(self, group, name, data): form_factor=ff, total_dens=dens_total, mol_densities=dens_molecule, - frag_densities=dens_fragment + frag_densities=dens_fragment, ) count += 1 - + print(f"Final number of systems saved into dataset: {count}") From 6e4dafabaab3068a4d158b99dc2505a319923f35 Mon Sep 17 00:00:00 2001 From: Alexey Nesterenko Date: Sat, 11 Apr 2026 23:25:57 +0200 Subject: [PATCH 3/3] clean a bit --- developer/make_train_data.py | 356 ++++++++++++++++++----------------- 1 file changed, 181 insertions(+), 175 deletions(-) mode change 100644 => 100755 developer/make_train_data.py diff --git a/developer/make_train_data.py b/developer/make_train_data.py old mode 100644 new mode 100755 index 7fd49ad20..c4a6a784d --- a/developer/make_train_data.py +++ b/developer/make_train_data.py @@ -23,15 +23,15 @@ traj_centering_for_maicos_mda, traj_centering_for_maicos_mda_parallel, ) -from fairmd.lipids.api import UniverseConstructor, get_thickness, get_mean_ApL +from fairmd.lipids.api import UniverseConstructor, get_mean_ApL, get_thickness from fairmd.lipids.auxiliary import mollib -from fairmd.lipids.core import initialize_databank -from fairmd.lipids.molecules import Lipid, lipids_set +from fairmd.lipids.core import System, initialize_databank +from fairmd.lipids.molecules import Lipid -def is_suitable(system): +def is_suitable(system: System) -> bool: """ - Checks if a simulation system is suitable for analysis with maicos. + Check if a simulation system is suitable for analysis with maicos. Checks for the presence of a 'WARNINGS' dictionary in the system metadata and verifies compatibility using the internal fairmd suitability check. @@ -43,7 +43,7 @@ def is_suitable(system): bool: True if the system is suitable for analysis, False otherwise. """ flag = True - if "WARNINGS" in system.keys() and type(system["WARNINGS"]) == dict: + if "WARNINGS" in system and isinstance(system["WARNINGS"], dict): flag = False if not is_system_suitable_4_maicos(system): print(f"system {system} not suitable for maicos") @@ -51,12 +51,12 @@ def is_suitable(system): return flag -def get_scalar_properties(system): +def get_scalar_properties(system: System) -> tuple[float, float, bool]: """ - Retrieves scalar physical properties (ApL and thickness) for a given system. + Retrieve scalar physical properties (ApL and thickness) for a given system. - Attempts to load the mean Area per Lipid (ApL) and bilayer thickness. - If a property cannot be loaded, it assigns a default value of -1 and + Attempts to load the mean Area per Lipid (ApL) and bilayer thickness. + If a property cannot be loaded, it assigns a default value of -1 and logs a warning. Args: @@ -66,31 +66,40 @@ def get_scalar_properties(system): tuple: A tuple containing: - ApL (float): Average area per lipid (Ų). - thickness (float): Bilayer thickness (Å). - - flag (bool): True if all properties were loaded successfully. + - no_error_flag (bool): True if all properties were loaded successfully. """ - flag = True + no_error_flag = True try: - ApL = get_mean_ApL(system) + apl = get_mean_ApL(system) except: print(f"System {system} - can't load ApL") - flag = False - ApL = -1 + no_error_flag = False + apl = -1 try: thickness = get_thickness(system) except: print(f"System {system} - can't load thickness") - flag = False + no_error_flag = False thickness = -1 - return ApL, thickness, flag - - -def center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger): + return apl, thickness, no_error_flag + + +def center_trajectory( + system: System, + uc: UniverseConstructor, + last_atom, + g3_atom, + eq_time, + logger, + *, + recompute: bool = False, +) -> str: """ Centers the simulation trajectory for analysis, handling different software backends. - Coordinates trajectory centering using either Gromacs commands or MDAnalysis - (sequential or parallel) based on the simulation metadata and environment + Coordinates trajectory centering using either Gromacs commands or MDAnalysis + (sequential or parallel) based on the simulation metadata and environment configuration. Args: @@ -105,6 +114,8 @@ def center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger): Returns: str: The file path to the newly created centered trajectory file """ + u = uc.build_universe() + spath = os.path.join(FMDL_SIMU_PATH, system["path"]) if "gromacs" in system["SOFTWARE"]: traj_centered = traj_centering_for_maicos_gromacs( spath, @@ -113,51 +124,53 @@ def center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger): last_atom=last_atom, g3_atom=g3_atom, eq_time=eq_time, - recompute=RECOMPUTE, + recompute=recompute, ) - else: - if FMDL_MAICOS_NCORES != 1: - try: - n_jobs = FMDL_MAICOS_NCORES if FMDL_MAICOS_NCORES is not None else -1 - logger.info(f"Using parallel trajectory centering (n_jobs={n_jobs})") - traj_centered = traj_centering_for_maicos_mda_parallel( - u, - spath, - last_atom, - eq_time, - n_jobs=n_jobs, - recompute=RECOMPUTE, - logger=logger, - show_progress=True, - ) - except ImportError: - logger.warning("joblib not available, falling back to sequential centering") - traj_centered = traj_centering_for_maicos_mda( - u, - spath, - last_atom, - eq_time, - recompute=RECOMPUTE, - logger=logger, - ) - else: - logger.info("Using sequential trajectory centering (FMDL_MAICOS_NCORES=1)") + elif FMDL_MAICOS_NCORES != 1: + try: + n_jobs = FMDL_MAICOS_NCORES if FMDL_MAICOS_NCORES is not None else -1 + logger.info(f"Using parallel trajectory centering (n_jobs={n_jobs})") + traj_centered = traj_centering_for_maicos_mda_parallel( + u, + spath, + last_atom, + eq_time, + n_jobs=n_jobs, + recompute=recompute, + logger=logger, + show_progress=True, + ) + except ImportError: + logger.warning("joblib not available, falling back to sequential centering") traj_centered = traj_centering_for_maicos_mda( u, spath, last_atom, eq_time, - recompute=RECOMPUTE, + recompute=recompute, logger=logger, ) - return traj_centered + else: + logger.info("Using sequential trajectory centering (FMDL_MAICOS_NCORES=1)") + traj_centered = traj_centering_for_maicos_mda( + u, + spath, + last_atom, + eq_time, + recompute=recompute, + logger=logger, + ) + u.load_new(traj_centered, format="XTC") + u.guess_TopologyAttrs(force_guess=["elements"]) + mollib.guess_elements(system, u) + return u def separate_lipid_atoms(mapping_dict): """ - Groups lipid atom names into headgroup, tail, and backbone fragments. + Group lipid atom names into headgroup, tail, and backbone fragments. - Parses a mapping dictionary (usually from a Lipid class instance) to + Parses a mapping dictionary (usually from a Lipid class instance) to categorize atoms based on their structural fragment. Args: @@ -170,7 +183,7 @@ def separate_lipid_atoms(mapping_dict): head_atoms = "" tail_atoms = "" backbone_atoms = "" - for atom in mapping_dict.keys(): + for atom in mapping_dict: fragment = mapping_dict[atom]["FRAGMENT"] atom_name = mapping_dict[atom]["ATOMNAME"] if fragment == "headgroup": @@ -186,10 +199,10 @@ def separate_lipid_atoms(mapping_dict): def create_fragment_selectors(lipid_names): """ - Creates MDAnalysis atom selection strings for lipid fragments across a system. + Create MDAnalysis atom selection strings for lipid fragments across a system. - Iterates through a list of lipid names, registers their fragment mappings, - and constructs 'name ...' strings used to select specific fragments in + Iterates through a list of lipid names, registers their fragment mappings, + and constructs 'name ...' strings used to select specific fragments in a simulation. Args: @@ -216,15 +229,16 @@ def create_fragment_selectors(lipid_names): class HDF5LipidWriter: """ - A handler for saving lipid simulation analysis results into HDF5 format. + Save lipid simulation analysis results into HDF5 format. - This class manages the hierarchical storage of form factors, density - profiles, and scalar properties. It is designed for long-running + This class manages the hierarchical storage of form factors, density + profiles, and scalar properties. It is designed for long-running processes by opening and flushing to the file for every system processed. """ - def __init__(self, filename, overwrite_file=False): + + def __init__(self, filename: str, *, overwrite_file: bool = False) -> None: """ - Initializes the writer and optionally clears the existing file. + Initialize the writer and optionally clears the existing file. Args: filename (str): Path to the output .h5 file. @@ -238,9 +252,9 @@ def __init__(self, filename, overwrite_file=False): def save_system(self, system, scalar_data, form_factor, total_dens, mol_densities, frag_densities): """ - Saves a single system's results to the HDF5 file. + Save a single system's results to the HDF5 file. - Organizes data into groups for metadata, axes, form factors, and + Organizes data into groups for metadata, axes, form factors, and various electron density profiles (total, fragment-based, and molecule-based). If the system ID already exists, it is overwritten to ensure data integrity. @@ -293,7 +307,7 @@ def save_system(self, system, scalar_data, form_factor, total_dens, mol_densitie def _write_dataset(self, group, name, data): """ - Helper method to write a NumPy array to an HDF5 group with compression. + Write a NumPy array to an HDF5 group with compression. Args: group (h5py.Group): The parent group to write into. @@ -304,131 +318,123 @@ def _write_dataset(self, group, name, data): group.create_dataset(name, data=np.array(data), compression="gzip", compression_opts=4) -WATER_TO_LIPID_RATIO_THRESHOLD = 20 -RECOMPUTE = True - -systems = initialize_databank() -logger = logging.getLogger(__name__) - -writer = HDF5LipidWriter("lipid_FF_dataset.h5") -test_flag = True +def recompute_extended_ff_dataset( + h5fpath: str, *, + hydration_threshold: int = 20, + recompute_centering: bool = True, + small_trajs_only: bool = False, +) -> None: + """ + Recompute the extended form factor dataset for all systems in the databank. -count = 0 -print(f"Number of systems: {len(systems)}") -for system in systems: - if system["TRAJECTORY_SIZE"] > 10**8 and test_flag: # For testing purpouses - continue + This function iterates through all systems, checks their suitability, and + processes them to extract form factors and density profiles. The results + are saved into an HDF5 file using the HDF5LipidWriter class. Systems that + do not meet the criteria (e.g., water-to-lipid ratio) are skipped. + """ + systems = initialize_databank() + logger = logging.getLogger(__name__) + writer = HDF5LipidWriter(h5fpath) - if not is_suitable(system): - continue + count = 0 + print(f"Number of systems: {len(systems)}") + for system in systems: + if system["TRAJECTORY_SIZE"] > 10**8 and small_trajs_only: # For testing purpouses + continue - ApL, thickness, flag = get_scalar_properties(system) - # if not flag: - # continue - scalar_info = {"ApL": ApL, "thickness": thickness} + if not is_suitable(system): + continue - water_n = system["COMPOSITION"]["SOL"]["COUNT"] - molecules = system["COMPOSITION"].keys() - lipid_names = [] - lipid_n = 0 - for molecule in molecules: - if molecule in lipids_set: - lipid_n += sum(system["COMPOSITION"][molecule]["COUNT"]) - lipid_names.append(molecule) + ApL, thickness, flag = get_scalar_properties(system) + scalar_info = {"ApL": ApL, "thickness": thickness} - water_to_lipid_ratio = water_n / lipid_n + if system.get_hydration(basis="number") < hydration_threshold: + continue - if water_to_lipid_ratio < WATER_TO_LIPID_RATIO_THRESHOLD: - continue + try: + uc = UniverseConstructor(system) + uc.download_mddata() + except: + continue - try: - uc = UniverseConstructor(system) - uc.download_mddata() - except: - continue + eq_time = float(system["TIMELEFTOUT"]) * 1000 + last_atom, g3_atom = first_last_carbon(system, logger) - eq_time = float(system["TIMELEFTOUT"]) * 1000 + u = center_trajectory(system, uc, last_atom, g3_atom, eq_time, logger, recompute=recompute_centering) - last_atom, g3_atom = first_last_carbon(system, logger) + bin_width = 0.3 - spath = f"{FMDL_SIMU_PATH}/{system['path']}" - u = uc.build_universe() + L_min = u.dimensions[2] + for ts in u.trajectory: + L_min = min(L_min, ts.dimensions[2]) - traj_centered = center_trajectory(u, uc, spath, last_atom, g3_atom, eq_time, logger) + base_options = {"unwrap": False, "bin_width": bin_width, "pack": False} + zlim = {"zmin": -L_min / 2, "zmax": L_min / 2} + dens_options = {**zlim, **base_options} - u.load_new(traj_centered, format="XTC") - u.guess_TopologyAttrs(force_guess=["elements"]) - mollib.guess_elements(system, u) - - bin_width = 0.3 - - L_min = u.dimensions[2] - for ts in u.trajectory: - L_min = min(L_min, ts.dimensions[2]) - - base_options = {"unwrap": False, "bin_width": bin_width, "pack": False} - zlim = {"zmin": -L_min / 2, "zmax": L_min / 2} - dens_options = {**zlim, **base_options} - - save_data = [] - - print("Calculating form factor") - form_factor = FormFactorPlanar( - atomgroup=u.atoms, - **base_options, - zmin=None, - zmax=None, - ).run() - ff = (form_factor.results.bin_pos, form_factor.results.profile, form_factor.results.dprofile) - - print("Calculating total density") - dens_total_runner = DensityPlanar( - u.atoms, - dens="electron", - **dens_options, - ).run() - dens_total = ( - dens_total_runner.results.bin_pos, - dens_total_runner.results.profile, - dens_total_runner.results.dprofile, - ) - - molecule_types_selector = [f"resname {system['COMPOSITION'][molecule_type]['NAME']}" for molecule_type in molecules] - dens_molecule = [] - for selector in molecule_types_selector: - print(f"Calculating {selector.split(' ')[1]} density") - molecule_group = u.select_atoms(selector) - dens_molecule_runner = DensityPlanar( - molecule_group, - dens="electron", - **dens_options, + print("Calculating form factor") + form_factor = FormFactorPlanar( + atomgroup=u.atoms, + **base_options, + zmin=None, + zmax=None, ).run() - dens = (dens_molecule_runner.results.profile, dens_molecule_runner.results.dprofile) - dens_molecule.append(dens) - - fragment_selectors = create_fragment_selectors(lipid_names) - dens_fragment = [] - frag_labels = ["head", "tail", "backbone"] - for i, selector in enumerate(fragment_selectors): - print(f"Calculating {frag_labels[i]} density") - fragment_group = u.select_atoms(selector) - dens_fragment_runner = DensityPlanar( - fragment_group, + ff = (form_factor.results.bin_pos, form_factor.results.profile, form_factor.results.dprofile) + + print("Calculating total density") + dens_total_runner = DensityPlanar( + u.atoms, dens="electron", **dens_options, ).run() - dens = (dens_fragment_runner.results.profile, dens_fragment_runner.results.dprofile) - dens_fragment.append(dens) + dens_total = ( + dens_total_runner.results.bin_pos, + dens_total_runner.results.profile, + dens_total_runner.results.dprofile, + ) + + molecule_types_selector = [ + f"resname {system['COMPOSITION'][molkey]['NAME']}" for molkey in system.content + ] + dens_molecule = [] + for selector in molecule_types_selector: + print(f"Calculating {selector.split(' ')[1]} density") + molecule_group = u.select_atoms(selector) + dens_molecule_runner = DensityPlanar( + molecule_group, + dens="electron", + **dens_options, + ).run() + dens = (dens_molecule_runner.results.profile, dens_molecule_runner.results.dprofile) + dens_molecule.append(dens) + + fragment_selectors = create_fragment_selectors(system.lipids.keys()) + dens_fragment = [] + frag_labels = ["head", "tail", "backbone"] + for i, selector in enumerate(fragment_selectors): + print(f"Calculating {frag_labels[i]} density") + fragment_group = u.select_atoms(selector) + dens_fragment_runner = DensityPlanar( + fragment_group, + dens="electron", + **dens_options, + ).run() + dens = (dens_fragment_runner.results.profile, dens_fragment_runner.results.dprofile) + dens_fragment.append(dens) + + writer.save_system( + system=system, + scalar_data=scalar_info, + form_factor=ff, + total_dens=dens_total, + mol_densities=dens_molecule, + frag_densities=dens_fragment, + ) - writer.save_system( - system=system, - scalar_data=scalar_info, - form_factor=ff, - total_dens=dens_total, - mol_densities=dens_molecule, - frag_densities=dens_fragment, - ) + count += 1 - count += 1 + print(f"Final number of systems saved into dataset: {count}") -print(f"Final number of systems saved into dataset: {count}") +if __name__ == "__main__": + h5fpath = "lipid_dataset_extended.h5" + recompute_extended_ff_dataset(h5fpath, hydration_threshold=0, recompute_centering=False, small_trajs_only=False) \ No newline at end of file