Scripts for form-factor dataset creation

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}")