diff --git a/src/molgen3D/data_processing/data_preprocessing.py b/src/molgen3D/data_processing/data_preprocessing.py index 878167c..8c411e3 100644 --- a/src/molgen3D/data_processing/data_preprocessing.py +++ b/src/molgen3D/data_processing/data_preprocessing.py @@ -1,5 +1,4 @@ import argparse -import ast import glob import json import os @@ -7,7 +6,7 @@ import random from collections import defaultdict from multiprocessing import Pool -from typing import Any, Dict, Optional, Set, Tuple, List +from typing import Any, Dict, List, Optional, Set, Tuple import numpy as np from loguru import logger as log from rdkit import Chem, RDLogger @@ -16,12 +15,15 @@ from molgen3D.data_processing.utils import ( JsonlSplitWriter, filter_mols, + get_embedding_func_and_config, + parse_coordinate_ranges, save_processed_pickle, ) from molgen3D.data_processing.smiles_encoder_decoder import ( - encode_cartesian_v2, + BinConfig, encode_cartesian_binned, encode_cartesian_binned_v2, + encode_cartesian_with_config, ) from molgen3D.utils.utils import load_pkl @@ -29,12 +31,16 @@ def read_mol( - args: Tuple[str, int, int, Any, float, List[Tuple[float, float]], bool, str, str] + args: Tuple, ) -> Optional[Tuple[List[str], Dict[str, Any]]]: - mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir, _geom_root = args + mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir, _geom_root = args[:9] + bin_config = args[9] if len(args) > 9 else None + use_isomeric_smiles = args[10] if len(args) > 10 else False try: return _read_mol_impl( - mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir + mol_path, max_confs, precision, embedding_func, bin_size, ranges, do_filter, pickle_dir, + bin_config=bin_config, + use_isomeric_smiles=use_isomeric_smiles, ) except Exception as exc: log.error("Unhandled exception in read_mol | path={} | error={}", mol_path, exc) @@ -50,6 +56,8 @@ def _read_mol_impl( ranges: List[Tuple[float, float]], do_filter: bool, pickle_dir: str, + bin_config: Optional[BinConfig] = None, + use_isomeric_smiles: bool = False, ) -> Tuple[List[str], Dict[str, Any]]: mol_object = load_pkl(mol_path) geom_smiles = mol_object["smiles"] @@ -75,7 +83,9 @@ def _read_mol_impl( continue try: - if embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): embedded_smile, iso_smile = embedding_func(mol, bin_size=bin_size, ranges=ranges) else: embedded_smile, iso_smile = embedding_func(mol, precision=precision) @@ -85,6 +95,7 @@ def _read_mol_impl( continue # Compute nonisomeric SMILES only for conformers that encoded successfully + noniso = None try: noniso = Chem.MolToSmiles(Chem.RemoveHs(mol, sanitize=False), canonical=True, isomericSmiles=False) nonisomeric_smiles.add(noniso) @@ -93,10 +104,11 @@ def _read_mol_impl( except Exception: pass + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) samples.append( json.dumps( { - "canonical_smiles": iso_smile, + "canonical_smiles": canonical_smiles, "embedded_smiles": embedded_smile, }, separators=(",", ":"), @@ -158,19 +170,16 @@ def preprocess( bin_size: float = 0.104, ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", filter_ranges: str = None, + bin_config_path: Optional[str] = None, + use_isomeric_smiles: bool = False, ) -> None: if dest_path is None: raise ValueError("dest_path must be provided for preprocessing output") - embedding_registry = { - "cartesian_v2": encode_cartesian_v2, - "cartesian": encode_cartesian_v2, - "cartesian_binned": encode_cartesian_binned, - "cartesian_binned_v2": encode_cartesian_binned_v2, - } - if embedding_type not in embedding_registry: - raise ValueError(f"Unsupported embedding_type '{embedding_type}'. Options: {sorted(embedding_registry)}") - embedding_func = embedding_registry[embedding_type] + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) overall_total_input_mols = overall_total_confs = overall_total_mols = 0 overall_multi_distinct_graphs = overall_mol_with_dotted_smiles = overall_total_dotted_smiles = 0 @@ -199,13 +208,7 @@ def preprocess( if pickle_paths.size == 0: raise FileNotFoundError(f"No pickle files found under pattern {pickle_glob}") - # Parse ranges string once - try: - parsed_ranges = ast.literal_eval(f"[{ranges}]") - parsed_ranges = [tuple(r) for r in parsed_ranges] - except Exception as e: - log.error(f"Failed to parse ranges: {ranges}. Error: {e}") - parsed_ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + parsed_ranges = parse_coordinate_ranges(ranges) do_filter = False if filter_ranges is not None: @@ -250,6 +253,8 @@ def preprocess( do_filter, split_pickle_dirs[split_name], geom_raw_path, + bin_config, + use_isomeric_smiles, ) for path in mol_paths ] @@ -379,7 +384,8 @@ def preprocess( "--embedding_type", "-et", type=str, - choices=["cartesian", "cartesian_v2", "cartesian_binned", "cartesian_binned_v2"], + choices=["cartesian", "cartesian_v2", "cartesian_binned", "cartesian_binned_v2", + "uniform_binned", "quantile_binned"], default="cartesian_v2", help="Embedding type to use for enrichment.", ) @@ -446,6 +452,25 @@ def preprocess( default=None, help="Filter ranges for binned embedding.", ) + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + + parser.add_argument( + "--isomeric", + action="store_true", + help="Alias for --use_isomeric_smiles.", + ) + parser.add_argument( + "--sort_by", + type=str, + choices=["energy", "weight", "none"], + default="energy", + help="Sort conformers by energy, weight, or keep original order.", + ) args = parser.parse_args() dest_path = osp.join(args.dest, args.run_name) @@ -458,6 +483,8 @@ def preprocess( diagnose=False, ) + use_isomeric = args.isomeric + preprocess( geom_raw_path=args.geom_raw_path, indices_path=args.indices_path, @@ -471,5 +498,7 @@ def preprocess( bin_size=args.bin_size, ranges=args.ranges, filter_ranges=args.filter_ranges, + bin_config_path=args.bin_config_path, + use_isomeric_smiles=use_isomeric, ) \ No newline at end of file diff --git a/src/molgen3D/data_processing/data_preprocessing_revisited.py b/src/molgen3D/data_processing/data_preprocessing_revisited.py new file mode 100644 index 0000000..e521802 --- /dev/null +++ b/src/molgen3D/data_processing/data_preprocessing_revisited.py @@ -0,0 +1,467 @@ +import argparse +import json +import os +import os.path as osp +import pickle +import random +from collections import defaultdict +from multiprocessing import Pool +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np +from loguru import logger as log +from rdkit import Chem, RDLogger +from tqdm.auto import tqdm + +from molgen3D.data_processing.smiles_encoder_decoder import ( + encode_cartesian_binned, + encode_cartesian_binned_v2, + encode_cartesian_with_config, +) +from molgen3D.data_processing.utils import ( + JsonlSplitWriter, + filter_revisited_mols, + get_embedding_func_and_config, + get_revisited_split_path, + parse_coordinate_ranges, + save_processed_pickle, +) + +RDLogger.DisableLog("rdApp.*") + +_REVISITED_DATA: Optional[List] = None + + +def _process_revisited_mol_impl( + idx: int, + max_confs: int, + precision: int, + embedding_func: Any, + bin_size: float, + parsed_ranges: List[Tuple[float, float]], + do_filter: bool, + bin_config: Any, + use_isomeric_smiles: bool, +) -> Tuple[List[str], Dict[str, Any], List[Chem.Mol]]: + smiles, mols = _REVISITED_DATA[idx] + local_failures: Dict[str, int] = defaultdict(int) + + valid = filter_revisited_mols( + smiles=smiles, + mols=mols, + failures=local_failures, + max_confs=max_confs, + ) + + nonisomeric_smiles, dotted_smiles, isomeric_smiles = set(), set(), set() + samples: List[str] = [] + filtered_mols: List[Chem.Mol] = [] + + for mol in valid: + if do_filter: + pos = mol.GetConformer().GetPositions() + out_of_range = False + for i in range(3): + min_val, max_val = parsed_ranges[i] + if np.any(pos[:, i] < min_val) or np.any(pos[:, i] > max_val): + out_of_range = True + break + if out_of_range: + local_failures["coord_out_of_range"] += 1 + continue + + try: + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + embedded_smile, iso_smile = embedding_func( + mol, + bin_size=bin_size, + ranges=parsed_ranges, + ) + else: + embedded_smile, iso_smile = embedding_func(mol, precision=precision) + except Exception as exc: + log.error("Error encoding conformer | smiles={} | failure={}", smiles, exc) + local_failures["encoding_error"] += 1 + continue + + noniso = None + try: + noniso = Chem.MolToSmiles( + Chem.RemoveHs(mol, sanitize=False), + canonical=True, + isomericSmiles=False, + ) + nonisomeric_smiles.add(noniso) + if "." in noniso: + dotted_smiles.add(noniso) + except Exception: + pass + + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) + samples.append( + json.dumps( + { + "canonical_smiles": canonical_smiles, + "embedded_smiles": embedded_smile, + }, + separators=(",", ":"), + ) + + "\n" + ) + isomeric_smiles.add(iso_smile) + filtered_mols.append(mol) + + if not samples: + local_failures["no_samples_after_filtering"] += 1 + + stats = { + "path": smiles, + "geom_smiles": smiles, + "confs_count_pre_filter": len(mols), + "confs_count_post_filter": len(samples), + "nonisomeric_smiles_post_filter": len(nonisomeric_smiles), + "isomeric_smiles_post_filter": isomeric_smiles, + "num_distinct_smiles_with_dot": len(dotted_smiles), + "has_dotted_smiles": bool(dotted_smiles), + "failures": local_failures, + "processed_pickle_path": None, + } + return samples, stats, filtered_mols + + +def _process_revisited_by_idx(args: Tuple) -> Optional[Tuple]: + idx = args[0] + try: + return _process_revisited_mol_impl(idx, *args[1:]) + except Exception as exc: + log.error("Unhandled exception | idx={} | error={}", idx, exc) + return None + + +def preprocess_revisited( + geom_raw_path: str, + embedding_type: str, + num_workers: int = 20, + precision: int = 4, + splits: Optional[str] = None, + dest_path: Optional[str] = None, + max_confs: int = 30, + bin_size: float = 0.104, + ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + filter_ranges: str = None, + bin_config_path: Optional[str] = None, + isomeric: bool = False, + use_centered: bool = True, +) -> None: + global _REVISITED_DATA + + if dest_path is None: + raise ValueError("dest_path must be provided for preprocessing output") + + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) + parsed_ranges = parse_coordinate_ranges(ranges) + + do_filter = False + if filter_ranges is not None: + if isinstance(filter_ranges, str): + do_filter = filter_ranges.lower() in ("true", "1", "yes", "on") + else: + do_filter = bool(filter_ranges) + + requested_splits = [splits] if splits else ["train", "valid", "test"] + + strings_root = osp.join(dest_path, "processed_strings") + split_writers = { + split: JsonlSplitWriter(osp.join(strings_root, split), split, chunk_size=50_000) + for split in ("train", "valid", "test") + } + split_pickle_dirs = { + split: osp.join(dest_path, "processed_pickles", split) for split in ("train", "valid", "test") + } + for split_dir in split_pickle_dirs.values(): + os.makedirs(split_dir, exist_ok=True) + + overall_total_input = overall_total_confs = overall_total_mols = 0 + overall_failure_counts: Dict[str, int] = defaultdict(int) + + for split_name in requested_splits: + data_path = get_revisited_split_path( + geom_raw_path=geom_raw_path, + split_name=split_name, + use_centered=use_centered, + ) + if not osp.exists(data_path): + raise FileNotFoundError(f"Revisited split file not found: {data_path}") + + log.info("Loading {} ...", data_path) + with open(data_path, "rb") as fh: + _REVISITED_DATA = pickle.load(fh) + log.info("Loaded {:,} molecules for split '{}'", len(_REVISITED_DATA), split_name) + + n = len(_REVISITED_DATA) + overall_total_input += n + + job_args = [ + ( + idx, + max_confs, + precision, + embedding_func, + bin_size, + parsed_ranges, + do_filter, + bin_config, + isomeric, + ) + for idx in range(n) + ] + + conf_count_pre = conf_count_post = mol_count_post = 0 + split_num_mol_with_multi_distinct_graphs = split_num_mol_with_dotted_smiles = total_dotted_smiles = 0 + failure_counts: Dict[str, int] = defaultdict(int) + split_geom_to_iso_map: Dict[str, Set[str]] = defaultdict(set) + + geom_to_iso_path = osp.join(dest_path, f"{split_name}_geom_to_isomeric_smiles.jsonl") + geom_to_iso_fh = open(geom_to_iso_path, "w") + try: + with tqdm(total=n, dynamic_ncols=True, mininterval=0.2) as pbar: + with Pool(processes=num_workers) as pool: + chunk_size = max(1, n // max(num_workers * 2, 1)) + for result in pool.imap_unordered( + _process_revisited_by_idx, + job_args, + chunksize=chunk_size, + ): + if result is None: + failure_counts["unhandled_exception"] += 1 + pbar.update() + continue + + samples, stats, filtered_mols = result + split_writers[split_name].write(samples) + + geom_smiles = stats.get("geom_smiles") + if filtered_mols and geom_smiles: + try: + processed_pickle_path = save_processed_pickle( + split_dir=split_pickle_dirs[split_name], + geom_smiles=geom_smiles, + mols=filtered_mols, + ) + stats["processed_pickle_path"] = processed_pickle_path + except Exception as exc: + log.error( + "Failed to write processed pickle | geom_smiles={} | failure={}", + geom_smiles, + exc, + ) + + conf_count_pre += stats["confs_count_pre_filter"] + conf_count_post += stats["confs_count_post_filter"] + overall_total_confs += stats["confs_count_post_filter"] + total_dotted_smiles += stats.get("num_distinct_smiles_with_dot", 0) + + if stats["nonisomeric_smiles_post_filter"] > 1: + split_num_mol_with_multi_distinct_graphs += 1 + if stats.get("has_dotted_smiles", False): + split_num_mol_with_dotted_smiles += 1 + + for reason, count in stats["failures"].items(): + failure_counts[reason] += int(count) + + if stats["confs_count_post_filter"] > 0: + mol_count_post += 1 + overall_total_mols += 1 + + if geom_smiles: + for iso in stats.get("isomeric_smiles_post_filter", ()): + split_geom_to_iso_map[geom_smiles].add(iso) + geom_to_iso_fh.write( + json.dumps( + { + "geom_smiles": geom_smiles, + "isomeric_smiles": iso, + }, + separators=(",", ":"), + ) + + "\n" + ) + + pbar.update() + finally: + geom_to_iso_fh.close() + _REVISITED_DATA = None + + avg_confs = conf_count_post / mol_count_post if mol_count_post else 0.0 + success_rate = mol_count_post / n if n else 0.0 + split_report = { + "split": split_name, + "num_input_molecules": n, + "num_output_molecules": mol_count_post, + "num_input_conformers": conf_count_pre, + "total_conformers_after": conf_count_post, + "avg_conformers_per_molecule_after": avg_confs, + "success_rate": success_rate, + "failure_counts": dict(failure_counts), + "molecules_with_multiple_distinct_graphs": split_num_mol_with_multi_distinct_graphs, + "molecules_with_dotted_smiles": split_num_mol_with_dotted_smiles, + "total_dotted_smiles": total_dotted_smiles, + } + log.info(json.dumps({"split_summary": split_report}, ensure_ascii=False, separators=(",", ":"))) + + for reason, count in failure_counts.items(): + overall_failure_counts[reason] += count + + for writer in split_writers.values(): + writer.close() + + grand_total = sum(writer.total_samples for writer in split_writers.values()) + overall_success_rate = float(overall_total_mols) / max(1, overall_total_input) + run_summary = { + "grand_total_samples_written": grand_total, + "total_input_molecules": overall_total_input, + "molecules_after_filter": overall_total_mols, + "conformers_after_filter": overall_total_confs, + "avg_confs_per_mol_after": float(overall_total_confs) / max(1, overall_total_mols), + "overall_success_rate": overall_success_rate, + "overall_failure_counts": dict(overall_failure_counts), + } + log.info(json.dumps({"run_summary": run_summary}, ensure_ascii=False, separators=(",", ":"))) + + +if __name__ == "__main__": + random.seed(42) + parser = argparse.ArgumentParser() + parser.add_argument( + "--geom_raw_path", + "-p", + type=str, + required=True, + help="Path to the revisited GEOM split pickle files.", + ) + parser.add_argument( + "--dest", + "-d", + type=str, + default="/data/molgen/", + help="Destination directory for processed outputs.", + ) + parser.add_argument( + "--embedding_type", + "-et", + type=str, + choices=[ + "cartesian", + "cartesian_v2", + "cartesian_binned", + "cartesian_binned_v2", + "uniform_binned", + "quantile_binned", + ], + default="cartesian_v2", + help="Embedding type to use for enrichment.", + ) + parser.add_argument( + "--num_workers", + "-nw", + type=int, + default=max(4, os.cpu_count() or 4), + help="Number of worker processes.", + ) + parser.add_argument( + "--precision", + type=int, + default=4, + help="Numeric precision for encoded coordinates.", + ) + parser.add_argument( + "--splits", + type=str, + choices=["train", "valid", "test"], + default=None, + help="Optional single split to process.", + ) + parser.add_argument( + "--run_name", + type=str, + default="", + help="Run name, appended to destination directory.", + ) + parser.add_argument( + "--max_confs", + type=int, + default=30, + help="Maximum number of conformers per molecule.", + ) + parser.add_argument( + "--bin_size", + type=float, + default=0.104, + help="Bin size for binned embedding.", + ) + parser.add_argument( + "--ranges", + type=str, + default="[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + help="Ranges for binned embedding.", + ) + parser.add_argument( + "--filter_ranges", + type=str, + default=None, + help="Filter ranges for binned embedding.", + ) + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + parser.add_argument( + "--use_isomeric_smiles", + action="store_true", + help="If set, writes isomeric SMILES to canonical_smiles; otherwise writes non-isomeric SMILES.", + ) + parser.add_argument( + "--isomeric", + action="store_true", + help="Alias for --use_isomeric_smiles.", + ) + parser.add_argument( + "--use_centered", + action="store_true", + default=False, + help="Load *_data_centered.pickle instead of *_data.pickle.", + ) + args = parser.parse_args() + + dest_path = osp.join(args.dest, args.run_name) + os.makedirs(dest_path, exist_ok=True) + log.add( + osp.join(dest_path, "preprocessing.log"), + mode="w", + enqueue=True, + backtrace=False, + diagnose=False, + ) + + preprocess_revisited( + geom_raw_path=args.geom_raw_path, + embedding_type=args.embedding_type, + dest_path=dest_path, + max_confs=args.max_confs, + num_workers=args.num_workers, + precision=args.precision, + splits=args.splits, + bin_size=args.bin_size, + ranges=args.ranges, + filter_ranges=args.filter_ranges, + bin_config_path=args.bin_config_path, + isomeric=args.isomeric or args.use_isomeric_smiles, + use_centered=args.use_centered, + ) diff --git a/src/molgen3D/data_processing/preprocess_geom_grouped.py b/src/molgen3D/data_processing/preprocess_geom_grouped.py index d2132f8..d44241b 100644 --- a/src/molgen3D/data_processing/preprocess_geom_grouped.py +++ b/src/molgen3D/data_processing/preprocess_geom_grouped.py @@ -1,10 +1,8 @@ import argparse -import ast import glob import json import os import os.path as osp -import pickle import random from collections import defaultdict from multiprocessing import Pool @@ -15,11 +13,18 @@ from rdkit import Chem, RDLogger from tqdm.auto import tqdm -from molgen3D.data_processing.utils import JsonlSplitWriter +from molgen3D.data_processing.utils import ( + JsonlSplitWriter, + copy_single_conformer_mol, + extract_conf_meta, + get_embedding_func_and_config, + parse_coordinate_ranges, + save_grouped_pickle, +) from molgen3D.data_processing.smiles_encoder_decoder import ( encode_cartesian_binned, - encode_cartesian_v2, encode_cartesian_binned_v2, + encode_cartesian_with_config, ) from molgen3D.utils.utils import load_pkl @@ -35,81 +40,6 @@ def infer_geom_id(mol_path: str) -> str: return osp.splitext(osp.basename(mol_path))[0] -def copy_single_conformer_mol(mol: Chem.Mol) -> Chem.Mol: - copied = Chem.Mol(mol) - if copied.GetNumConformers() > 1: - conf = Chem.Conformer(copied.GetConformer(0)) - copied.RemoveAllConformers() - copied.AddConformer(conf, assignId=True) - return copied - - -def _get_prop_float(props: Dict[str, Any], key: str) -> Optional[float]: - if key in props: - try: - return float(props[key]) - except Exception: - return None - return None - - -def extract_conf_meta( - conf_meta: Optional[Dict[str, Any]], mol: Chem.Mol -) -> Tuple[Optional[float], Optional[float], Optional[int]]: - energy = weight = None - conf_id = None - - if conf_meta: - if "totalenergy" in conf_meta: - energy = conf_meta.get("totalenergy") - elif "relativeenergy" in conf_meta: - energy = conf_meta.get("relativeenergy") - - if "boltzmannweight" in conf_meta: - weight = conf_meta.get("boltzmannweight") - elif "weight" in conf_meta: - weight = conf_meta.get("weight") - - if "geom_id" in conf_meta: - conf_id = conf_meta.get("geom_id") - - if mol is not None: - mol_props = mol.GetPropsAsDict() - if energy is None: - energy = _get_prop_float(mol_props, "totalenergy") - if weight is None: - weight = _get_prop_float(mol_props, "boltzmannweight") - if conf_id is None and "geom_id" in mol_props: - try: - conf_id = int(mol_props["geom_id"]) - except Exception: - conf_id = None - - if mol.GetNumConformers() > 0: - conf_props = mol.GetConformer().GetPropsAsDict() - if energy is None: - energy = _get_prop_float(conf_props, "totalenergy") - if weight is None: - weight = _get_prop_float(conf_props, "boltzmannweight") - if conf_id is None and "geom_id" in conf_props: - try: - conf_id = int(conf_props["geom_id"]) - except Exception: - conf_id = None - - try: - energy = float(energy) if energy is not None else None - except Exception: - energy = None - - try: - weight = float(weight) if weight is not None else None - except Exception: - weight = None - - return energy, weight, conf_id - - def filter_conformers_keep_dotted( mol_object: Dict[str, Any], failures: Dict[str, int], @@ -171,6 +101,8 @@ def _read_mol_impl( Dict[str, List[Dict[str, Any]]], ]: mol_path, max_confs, precision, embedding_func, bin_size, ranges, sort_by = args + bin_config = args[7] if len(args) > 7 else None + use_isomeric_smiles = args[8] if len(args) > 8 else False mol_object = load_pkl(mol_path) geom_key = mol_object.get("smiles") geom_id = infer_geom_id(mol_path) @@ -204,7 +136,9 @@ def _read_mol_impl( for _, _, mol, _conf_meta, energy, weight, conf_id in scored_candidates: try: - if embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): embedded_smile, iso_smile = embedding_func( mol, bin_size=bin_size, @@ -217,6 +151,7 @@ def _read_mol_impl( local_failures["encoding_error"] += 1 continue + noniso = None # Compute nonisomeric SMILES only for conformers that encoded successfully try: noniso = Chem.MolToSmiles( @@ -230,7 +165,8 @@ def _read_mol_impl( except Exception: pass - isomeric_smiles.add(iso_smile) + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) + isomeric_smiles.add(canonical_smiles) sample_entry: Dict[str, Any] = { "embedded_smiles": embedded_smile, @@ -240,7 +176,7 @@ def _read_mol_impl( } if conf_id is not None: sample_entry["conf_id"] = conf_id - sample_isomers[iso_smile].append(sample_entry) + sample_isomers[canonical_smiles].append(sample_entry) pickle_entry: Dict[str, Any] = { "mol": copy_single_conformer_mol(mol), @@ -251,7 +187,7 @@ def _read_mol_impl( } if conf_id is not None: pickle_entry["conf_id"] = conf_id - pickle_isomers[iso_smile].append(pickle_entry) + pickle_isomers[canonical_smiles].append(pickle_entry) if len(nonisomeric_smiles) > 1: log.info( @@ -296,14 +232,6 @@ def _read_mol_impl( return json_line, stats, geom_id, isomeric_smiles, pickle_isomers -def save_grouped_pickle(output_path: str, iso_to_confs: Dict[str, List[Dict[str, Any]]]) -> None: - parent = osp.dirname(output_path) - if parent: - os.makedirs(parent, exist_ok=True) - with open(output_path, "wb") as fh: - pickle.dump(iso_to_confs, fh) - - def preprocess( geom_raw_path: str, indices_path: str, @@ -317,19 +245,16 @@ def preprocess( bin_size: float = 0.104, ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", sort_by: str = "energy", + bin_config_path: Optional[str] = None, + isomeric: bool = False, ) -> None: if dest_path is None: raise ValueError("dest_path must be provided for preprocessing output") - embedding_registry = { - "cartesian_v2": encode_cartesian_v2, - "cartesian": encode_cartesian_v2, - "cartesian_binned": encode_cartesian_binned, - "cartesian_binned_v2": encode_cartesian_binned_v2, - } - if embedding_type not in embedding_registry: - raise ValueError(f"Unsupported embedding_type '{embedding_type}'. Options: {sorted(embedding_registry)}") - embedding_func = embedding_registry[embedding_type] + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) overall_total_input_mols = overall_total_confs = overall_total_mols = 0 overall_multi_distinct_graphs = overall_mol_with_dotted_smiles = overall_total_dotted_smiles = 0 @@ -352,12 +277,7 @@ def preprocess( split_indices_array = np.load(indices_path, allow_pickle=True) - try: - parsed_ranges = ast.literal_eval(f"[{ranges}]") - parsed_ranges = [tuple(r) for r in parsed_ranges] - except Exception as exc: - log.error("Failed to parse ranges: {} | failure={}", ranges, exc) - parsed_ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + parsed_ranges = parse_coordinate_ranges(ranges) pickle_glob = osp.join(geom_raw_path, f"{dataset_type}/*.pickle") pickle_paths = np.array(sorted(glob.glob(pickle_glob))) @@ -403,6 +323,8 @@ def preprocess( bin_size, parsed_ranges, sort_by, + bin_config, + isomeric, ) for path in mol_paths ] @@ -553,7 +475,14 @@ def preprocess( "--embedding_type", "-et", type=str, - choices=["cartesian", "cartesian_v2", "cartesian_binned", "cartesian_binned_v2"], + choices=[ + "cartesian", + "cartesian_v2", + "cartesian_binned", + "cartesian_binned_v2", + "uniform_binned", + "quantile_binned", + ], default="cartesian_v2", help="Embedding type to use for enrichment.", ) @@ -621,7 +550,17 @@ def preprocess( default="energy", help="Sort conformers by energy, weight, or keep original order.", ) - + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + parser.add_argument( + "--isomeric", + action="store_true", + help="If set, use isomeric SMILES keys; otherwise use non-isomeric keys when available.", + ) args = parser.parse_args() random.seed(42) @@ -649,4 +588,6 @@ def preprocess( sort_by=args.sort_by, bin_size=args.bin_size, ranges=args.ranges, + bin_config_path=args.bin_config_path, + isomeric=args.isomeric, ) diff --git a/src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py b/src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py new file mode 100644 index 0000000..e91df4c --- /dev/null +++ b/src/molgen3D/data_processing/preprocess_geom_grouped_revisited.py @@ -0,0 +1,521 @@ +import argparse +import json +import os +import os.path as osp +import pickle +import random +from collections import defaultdict +from multiprocessing import Pool +from typing import Any, Dict, List, Optional, Set, Tuple + +from loguru import logger as log +from rdkit import Chem, RDLogger +from tqdm.auto import tqdm + +from molgen3D.data_processing.smiles_encoder_decoder import ( + encode_cartesian_binned, + encode_cartesian_binned_v2, + encode_cartesian_with_config, +) +from molgen3D.data_processing.utils import ( + JsonlSplitWriter, + copy_single_conformer_mol, + extract_conf_meta, + filter_revisited_conformers_keep_dotted, + get_embedding_func_and_config, + get_revisited_split_path, + parse_coordinate_ranges, + save_grouped_pickle, +) + +RDLogger.DisableLog("rdApp.*") + +_REVISITED_DATA: Optional[List] = None + + +def _process_revisited_mol_impl( + idx: int, + max_confs: int, + precision: int, + embedding_func: Any, + bin_size: float, + parsed_ranges: List[Tuple[float, float]], + sort_by: str, + bin_config: Any, + use_isomeric_smiles: bool, +) -> Tuple[ + Optional[str], + Dict[str, Any], + str, + Set[str], + Dict[str, List[Dict[str, Any]]], +]: + smiles, mols = _REVISITED_DATA[idx] + geom_id = str(idx) + local_failures: Dict[str, int] = defaultdict(int) + + candidates = filter_revisited_conformers_keep_dotted( + smiles=smiles, + mols=mols, + failures=local_failures, + ) + + scored_candidates: List[ + Tuple[float, int, Chem.Mol, Optional[float], Optional[float], Optional[int]] + ] = [] + for mol, conf_idx in candidates: + energy, weight, conf_id = extract_conf_meta(None, mol) + if energy is None: + local_failures["missing_energy"] += 1 + if sort_by == "energy": + scored_value = energy if energy is not None else float("inf") + elif sort_by == "weight": + scored_value = -(weight if weight is not None else float("-inf")) + else: + scored_value = float(conf_idx) + scored_candidates.append((scored_value, conf_idx, mol, energy, weight, conf_id)) + + scored_candidates.sort(key=lambda x: (x[0], x[1])) + scored_candidates = scored_candidates[:max_confs] + + nonisomeric_smiles: Set[str] = set() + dotted_smiles: Set[str] = set() + isomeric_smiles: Set[str] = set() + sample_isomers: Dict[str, List[Dict[str, Any]]] = defaultdict(list) + pickle_isomers: Dict[str, List[Dict[str, Any]]] = defaultdict(list) + + for _, _, mol, energy, weight, conf_id in scored_candidates: + try: + if embedding_func is encode_cartesian_with_config: + embedded_smile, iso_smile = embedding_func(mol, bin_config) + elif embedding_func in (encode_cartesian_binned, encode_cartesian_binned_v2): + embedded_smile, iso_smile = embedding_func( + mol, + bin_size=bin_size, + ranges=parsed_ranges, + ) + else: + embedded_smile, iso_smile = embedding_func(mol, precision=precision) + except Exception as exc: + log.error("Error encoding conformer | smiles={} | failure={}", smiles, exc) + local_failures["encoding_error"] += 1 + continue + + noniso = None + try: + noniso = Chem.MolToSmiles( + Chem.RemoveHs(mol, sanitize=False), + canonical=True, + isomericSmiles=False, + ) + nonisomeric_smiles.add(noniso) + if "." in noniso: + dotted_smiles.add(noniso) + except Exception: + pass + + canonical_smiles = iso_smile if use_isomeric_smiles else (noniso or iso_smile) + isomeric_smiles.add(canonical_smiles) + + sample_entry: Dict[str, Any] = { + "embedded_smiles": embedded_smile, + "energy": energy, + "weight": weight, + "geom_id": geom_id, + } + if conf_id is not None: + sample_entry["conf_id"] = conf_id + sample_isomers[canonical_smiles].append(sample_entry) + + pickle_entry: Dict[str, Any] = { + "mol": copy_single_conformer_mol(mol), + "embedded_smiles": embedded_smile, + "energy": energy, + "weight": weight, + "geom_id": geom_id, + } + if conf_id is not None: + pickle_entry["conf_id"] = conf_id + pickle_isomers[canonical_smiles].append(pickle_entry) + + if len(nonisomeric_smiles) > 1: + log.info( + "multiple_distinct_nonisomeric_smiles | path={} | distinct_smiles={}", + smiles, + nonisomeric_smiles, + ) + for dotted in dotted_smiles: + log.info("dot_in_conformer_smiles | path={} | smile={}", smiles, dotted) + + if not sample_isomers: + log.warning("No samples after filtering | path={}", smiles) + local_failures["no_samples_after_filtering"] += 1 + + json_line = None + if sample_isomers: + json_line = ( + json.dumps( + {"geom_key": smiles, "geom_id": geom_id, "isomers": sample_isomers}, + separators=(",", ":"), + ) + + "\n" + ) + + stats = { + "path": smiles, + "geom_smiles": smiles, + "confs_count_pre_filter": len(mols), + "confs_count_post_filter": sum(len(v) for v in sample_isomers.values()), + "nonisomeric_smiles_post_filter": len(nonisomeric_smiles), + "isomeric_smiles_post_filter": isomeric_smiles, + "num_distinct_smiles_with_dot": len(dotted_smiles), + "has_dotted_smiles": bool(dotted_smiles), + "failures": local_failures, + "processed_pickle_path": None, + } + return json_line, stats, geom_id, isomeric_smiles, pickle_isomers + + +def _process_revisited_by_idx(args: Tuple) -> Optional[Tuple]: + idx = args[0] + try: + return _process_revisited_mol_impl(idx, *args[1:]) + except Exception as exc: + log.error("Unhandled exception | idx={} | error={}", idx, exc) + return None + + +def preprocess_revisited( + geom_raw_path: str, + embedding_type: str, + num_workers: int = 20, + precision: int = 4, + splits: Optional[str] = None, + dest_path: Optional[str] = None, + max_confs: int = 30, + bin_size: float = 0.104, + ranges: str = "[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + sort_by: str = "energy", + bin_config_path: Optional[str] = None, + isomeric: bool = False, + use_centered: bool = True, +) -> None: + global _REVISITED_DATA + + if dest_path is None: + raise ValueError("dest_path must be provided for preprocessing output") + + embedding_func, bin_config = get_embedding_func_and_config( + embedding_type=embedding_type, + bin_config_path=bin_config_path, + ) + parsed_ranges = parse_coordinate_ranges(ranges) + requested_splits = [splits] if splits else ["train", "valid", "test"] + + strings_root = osp.join(dest_path, "processed_strings") + split_writers = { + split: JsonlSplitWriter(osp.join(strings_root, split), split, chunk_size=50_000) + for split in ("train", "valid", "test") + } + split_pickle_dirs = { + split: osp.join(dest_path, "processed_pickles", split) for split in ("train", "valid", "test") + } + for split_dir in split_pickle_dirs.values(): + os.makedirs(split_dir, exist_ok=True) + + overall_total_input = overall_total_confs = overall_total_mols = 0 + overall_failure_counts: Dict[str, int] = defaultdict(int) + + for split_name in requested_splits: + data_path = get_revisited_split_path( + geom_raw_path=geom_raw_path, + split_name=split_name, + use_centered=use_centered, + ) + if not osp.exists(data_path): + raise FileNotFoundError(f"Revisited split file not found: {data_path}") + + log.info("Loading {} ...", data_path) + with open(data_path, "rb") as fh: + _REVISITED_DATA = pickle.load(fh) + log.info("Loaded {:,} molecules for split '{}'", len(_REVISITED_DATA), split_name) + + n = len(_REVISITED_DATA) + overall_total_input += n + + job_args = [ + ( + idx, + max_confs, + precision, + embedding_func, + bin_size, + parsed_ranges, + sort_by, + bin_config, + isomeric, + ) + for idx in range(n) + ] + + conf_count_pre = conf_count_post = mol_count_post = 0 + split_num_mol_with_multi_distinct_graphs = split_num_mol_with_dotted_smiles = total_dotted_smiles = 0 + failure_counts: Dict[str, int] = defaultdict(int) + split_geom_to_iso_map: Dict[str, Set[str]] = defaultdict(set) + + geom_to_iso_path = osp.join(dest_path, f"{split_name}_geom_to_isomeric_smiles.jsonl") + iso_to_geom_path = osp.join(dest_path, f"{split_name}_isomeric_to_geom.jsonl") + geom_to_iso_fh = open(geom_to_iso_path, "w") + iso_to_geom_fh = open(iso_to_geom_path, "w") + try: + with tqdm(total=n, dynamic_ncols=True, mininterval=0.2) as pbar: + with Pool(processes=num_workers) as pool: + chunk_size = max(1, n // max(num_workers * 2, 1)) + for result in pool.imap_unordered( + _process_revisited_by_idx, + job_args, + chunksize=chunk_size, + ): + if result is None: + failure_counts["unhandled_exception"] += 1 + pbar.update() + continue + + json_line, stats, geom_id, iso_set, iso_to_pickle_confs = result + + if json_line: + split_writers[split_name].write([json_line]) + + if iso_to_pickle_confs: + processed_pickle_path = osp.join(split_pickle_dirs[split_name], f"{geom_id}.pkl") + try: + save_grouped_pickle(processed_pickle_path, iso_to_pickle_confs) + except Exception as exc: + log.error( + "Failed to write processed pickle | path={} | failure={}", + processed_pickle_path, + exc, + ) + + conf_count_pre += stats["confs_count_pre_filter"] + conf_count_post += stats["confs_count_post_filter"] + overall_total_confs += stats["confs_count_post_filter"] + total_dotted_smiles += stats.get("num_distinct_smiles_with_dot", 0) + + if stats["nonisomeric_smiles_post_filter"] > 1: + split_num_mol_with_multi_distinct_graphs += 1 + if stats.get("has_dotted_smiles", False): + split_num_mol_with_dotted_smiles += 1 + + for reason, count in stats["failures"].items(): + failure_counts[reason] += int(count) + + if stats["confs_count_post_filter"] > 0: + mol_count_post += 1 + overall_total_mols += 1 + + geom_smiles = stats.get("geom_smiles") + if geom_smiles and iso_set: + for iso in iso_set: + split_geom_to_iso_map[geom_smiles].add(iso) + geom_to_iso_fh.write( + json.dumps( + { + "geom_key": geom_smiles, + "geom_id": str(geom_id), + "isomeric_smiles": iso, + }, + separators=(",", ":"), + ) + + "\n" + ) + iso_to_geom_fh.write( + json.dumps( + { + "isomeric_smiles": iso, + "geom_key": geom_smiles, + "geom_id": str(geom_id), + }, + separators=(",", ":"), + ) + + "\n" + ) + + pbar.update() + finally: + geom_to_iso_fh.close() + iso_to_geom_fh.close() + _REVISITED_DATA = None + + avg_confs = conf_count_post / mol_count_post if mol_count_post else 0.0 + success_rate = mol_count_post / n if n else 0.0 + split_report = { + "split": split_name, + "num_input_molecules": n, + "num_output_molecules": mol_count_post, + "num_input_conformers": conf_count_pre, + "total_conformers_after": conf_count_post, + "avg_conformers_per_molecule_after": avg_confs, + "success_rate": success_rate, + "failure_counts": dict(failure_counts), + "molecules_with_multiple_distinct_graphs": split_num_mol_with_multi_distinct_graphs, + "molecules_with_dotted_smiles": split_num_mol_with_dotted_smiles, + "total_dotted_smiles": total_dotted_smiles, + } + log.info(json.dumps({"split_summary": split_report}, ensure_ascii=False, separators=(",", ":"))) + + for reason, count in failure_counts.items(): + overall_failure_counts[reason] += count + + for writer in split_writers.values(): + writer.close() + + grand_total = sum(writer.total_samples for writer in split_writers.values()) + overall_success_rate = float(overall_total_mols) / max(1, overall_total_input) + run_summary = { + "grand_total_samples_written": grand_total, + "total_input_molecules": overall_total_input, + "molecules_after_filter": overall_total_mols, + "conformers_after_filter": overall_total_confs, + "avg_confs_per_mol_after": float(overall_total_confs) / max(1, overall_total_mols), + "overall_success_rate": overall_success_rate, + "overall_failure_counts": dict(overall_failure_counts), + } + log.info(json.dumps({"run_summary": run_summary}, ensure_ascii=False, separators=(",", ":"))) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--geom_raw_path", + "-p", + type=str, + required=True, + help="Path to the revisited GEOM split pickle files.", + ) + parser.add_argument( + "--dest", + "-d", + type=str, + default="/data/molgen/", + help="Destination directory for processed outputs.", + ) + parser.add_argument( + "--embedding_type", + "-et", + type=str, + choices=[ + "cartesian", + "cartesian_v2", + "cartesian_binned", + "cartesian_binned_v2", + "uniform_binned", + "quantile_binned", + ], + default="cartesian_v2", + help="Embedding type to use for enrichment.", + ) + parser.add_argument( + "--num_workers", + "-nw", + type=int, + default=max(4, os.cpu_count() or 4), + help="Number of worker processes.", + ) + parser.add_argument( + "--precision", + type=int, + default=4, + help="Numeric precision for encoded coordinates.", + ) + parser.add_argument( + "--dataset_type", + "-dt", + type=str, + default="drugs", + help="Unused placeholder kept for CLI parity.", + ) + parser.add_argument( + "--splits", + type=str, + choices=["train", "valid", "test"], + default=None, + help="Optional single split to process.", + ) + parser.add_argument( + "--run_name", + type=str, + default="", + help="Run name, appended to destination directory.", + ) + parser.add_argument( + "--max_confs", + type=int, + default=30, + help="Maximum number of conformers per molecule.", + ) + parser.add_argument( + "--bin_size", + type=float, + default=0.104, + help="Bin size for binned embedding.", + ) + parser.add_argument( + "--ranges", + type=str, + default="[-13.0, 13.0], [-13.0, 13.0], [-13.0, 13.0]", + help="Ranges for binned embedding.", + ) + parser.add_argument( + "--sort_by", + type=str, + choices=["energy", "weight", "none"], + default="none", + help="Sort conformers by energy, weight, or keep original order.", + ) + parser.add_argument( + "--bin_config_path", + type=str, + default=None, + help="Path to BinConfig JSON (required for uniform_binned / quantile_binned).", + ) + parser.add_argument( + "--isomeric", + action="store_true", + help="If set, use isomeric SMILES keys; otherwise use non-isomeric keys when available.", + ) + parser.add_argument( + "--use_centered", + action="store_true", + default=False, + help="Load *_data_centered.pickle instead of *_data.pickle.", + ) + args = parser.parse_args() + + random.seed(42) + dest_path = osp.join(args.dest, args.run_name) + os.makedirs(dest_path, exist_ok=True) + enqueue_logs = os.environ.get("LOGURU_ENQUEUE", "1") not in {"0", "false", "False"} + log.add( + osp.join(dest_path, "preprocessing.log"), + mode="w", + enqueue=enqueue_logs, + backtrace=False, + diagnose=False, + ) + + preprocess_revisited( + geom_raw_path=args.geom_raw_path, + embedding_type=args.embedding_type, + dest_path=dest_path, + max_confs=args.max_confs, + num_workers=args.num_workers, + precision=args.precision, + splits=args.splits, + bin_size=args.bin_size, + ranges=args.ranges, + sort_by=args.sort_by, + bin_config_path=args.bin_config_path, + isomeric=args.isomeric, + use_centered=args.use_centered, + ) diff --git a/src/molgen3D/data_processing/smiles_encoder_decoder.py b/src/molgen3D/data_processing/smiles_encoder_decoder.py index bc853ae..3f17218 100644 --- a/src/molgen3D/data_processing/smiles_encoder_decoder.py +++ b/src/molgen3D/data_processing/smiles_encoder_decoder.py @@ -1,6 +1,11 @@ import ast +import json import math import re +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + from rdkit import Chem from rdkit.Chem import AllChem from rdkit.Chem.rdchem import ChiralType @@ -9,7 +14,6 @@ def truncate(x, precision=4): - """Format a float with exactly ``precision`` decimal places (truncation, not rounding).""" if precision < 0: raise ValueError("precision must be non-negative") @@ -62,29 +66,13 @@ def _parse_float_token(token: str) -> float: _ORGANIC_SUBSET = {"B", "C", "N", "O", "P", "S", "F", "Cl", "Br", "I", "b", "c", "n", "o", "p", "s"} def strip_smiles(s: str) -> str: - """ - Normalize enriched SMILES strings into a 'canonical-ish' comparison form. - - Supported inputs: - - Legacy enriched strings: C<...>N<...> (atoms without brackets + coords) - - Current enriched strings: [C]<...>[N]<...> (atoms wrapped in brackets) - - Binned enriched strings: [C]123456789[N]123456789 (atoms wrapped in brackets + digits) - - Plain SMILES: C[NH2+]Cc1... - - Steps: - 1. Remove every <...> coordinate block. - 2. Remove binned coordinate digits: [C]123456789 -> [C]. - 3. Collapse decorative carbon H-counts: [CH3],[CH2],[CH],[cH] -> C/c. - 4. Drop brackets around simple atoms: [C]->C, [c]->c, [N]->N, ... - 5. Keep chemically meaningful brackets: [NH2+], [nH], [H], [Pt+2], etc. - """ - if not s: return "" # Remove any tags that might be present s = s.replace("[CONFORMER]", "").replace("[/CONFORMER]", "") s = s.replace("[SMILES]", "").replace("[/SMILES]", "") + s = re.sub(r"\[SERIALIZATION\].*?\[/SERIALIZATION\]", "", s, flags=re.IGNORECASE | re.DOTALL) s = s.replace(";", "") s = _WHITESPACE_RE.sub('', s) @@ -137,7 +125,6 @@ def _expected_plain_token(atom) -> str: def tokenize_smiles(smiles_str, expected_atom_tokens=None): - """Tokenize a canonical SMILES string into atom/non-atom tokens.""" tokens = [] i = 0 n = len(smiles_str) @@ -236,7 +223,6 @@ def tokenize_smiles(smiles_str, expected_atom_tokens=None): def _format_atom_descriptor(atom, *, allow_chirality: bool = True): - """Return a bracketed atom descriptor that preserves valence information.""" symbol = atom.GetSymbol() aromatic = atom.GetIsAromatic() if aromatic and len(symbol) == 1: @@ -276,9 +262,6 @@ def _format_atom_descriptor(atom, *, allow_chirality: bool = True): def _normalize_atom_descriptor(descriptor: str) -> str: - """ - Collapse decorative hydrogen counts on neutral carbon descriptors. - """ match = _CARBON_DESCRIPTOR_RE.match(descriptor) if not match or match.group("iso"): return descriptor @@ -297,7 +280,6 @@ def _normalize_atom_descriptor(descriptor: str) -> str: def encode_cartesian_v2(mol, precision=4): - """Serialize a 3D RDKit Mol into the enriched text representation.""" mol_no_h = Chem.RemoveHs(mol) if mol_no_h.GetNumConformers() == 0: raise ValueError("Molecule has no conformer / 3D coordinates.") @@ -364,7 +346,6 @@ def encode_cartesian_v2(mol, precision=4): ) def tokenize_enriched(enriched): - """Tokenize the enriched representation back into atoms (with coords) and other tokens.""" tokens = [] pos = 0 for match in _ENRICHED_TOKEN_PATTERN.finditer(enriched): @@ -469,7 +450,6 @@ def tokenize_enriched_v2(enriched, digit_width): def decode_cartesian_v2(enriched_string): - """Reconstruct an RDKit Mol (with conformer) from the enriched string produced by the encoder.""" tokens = tokenize_enriched(enriched_string) smiles_parts = [] @@ -505,7 +485,6 @@ def decode_cartesian_v2(enriched_string): def embed_3d_conformer_from_smiles(smiles, seed=0): - """Generate a 3D conformer for a SMILES, drop implicit hydrogens, and return the resulting Mol.""" mol = Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(f"Could not parse SMILES: {smiles}") @@ -543,7 +522,6 @@ def embed_3d_conformer_from_smiles(smiles, seed=0): def coords_rmsd(mol_a, mol_b): - """Compute RMSD between conformer-0 coordinates assuming identical atom order.""" if mol_a.GetNumAtoms() != mol_b.GetNumAtoms(): raise ValueError("Cannot compare coordinates for molecules with different atom counts.") @@ -566,7 +544,6 @@ def coords_rmsd(mol_a, mol_b): return min(math.sqrt(sse / n), rmsd_rdkit) def get_bins_for_coords(ranges, bin_size=0.104): - """Get bins for coordinates based on the ranges and bin size.""" bins = [] for start, end in ranges: bins.append(np.arange(start, end, bin_size)) @@ -574,60 +551,216 @@ def get_bins_for_coords(ranges, bin_size=0.104): def get_digit_width(bins): - """Determine the zero-padding width based on the maximum number of bins.""" max_bin_len = max(len(b) for b in bins) return max(3, len(str(max_bin_len))) def coords_to_bins(coords, bins): - """Convert coordinates to bins.""" return np.digitize(coords, bins) def bins_to_coords(bin_indices, bins, use_bin_center=False): - """ - Convert bin indices to coordinates by choosing a value within the bin interval for each bin. - - Parameters: - bin_indices (array-like): Indices of the bins. - bins (array-like): The bin edges as used in np.digitize (e.g., output from np.arange). - use_bin_center (bool): If True, use the center of the bin; if False, uniformly sample within the bin. - - Returns: - np.ndarray: Coordinates either as bin centers or randomly sampled within each bin. - """ + step = float(bins[-1] - bins[-2]) if len(bins) > 1 else 1.0 coords = [] for bin_idx in bin_indices: - # Find the left and right edges of the bin if bin_idx <= 0: - left = bins[0] - right = bins[1] if len(bins) > 1 else bins[0] + # Tail low (BIN_L): snap to range start + coords.append(float(bins[0])) elif bin_idx >= len(bins): - left = bins[-1] - right = bins[-1] + (bins[-1] - bins[-2]) if len(bins) > 1 else bins[-1] + 1.0 + # Tail high (BIN_H): snap to range end + coords.append(float(bins[-1] + step)) else: left = bins[bin_idx - 1] right = bins[bin_idx] - # Choose value: random within [left, right) or use center - if use_bin_center: - coord = (left + right) / 2.0 - else: - coord = np.random.uniform(left, right) - coords.append(coord) + if use_bin_center: + coords.append((left + right) / 2.0) + else: + coords.append(np.random.uniform(left, right)) return np.array(coords) +@dataclass +class BinConfig: + mode: str + L: float + H: float + n_bins: int + edges: np.ndarray + digit_width: int = 3 + + def __post_init__(self): + self.digit_width = max(3, len(str(self.n_bins + 1))) + + # -- persistence --------------------------------------------------------- + def save(self, path: str) -> None: + obj = { + "mode": self.mode, + "L": self.L, + "H": self.H, + "n_bins": self.n_bins, + "edges": self.edges.tolist(), + } + with open(path, "w") as f: + json.dump(obj, f, indent=2) + + @classmethod + def load(cls, path: str) -> "BinConfig": + with open(path) as f: + obj = json.load(f) + return cls( + mode=obj["mode"], + L=obj["L"], + H=obj["H"], + n_bins=obj["n_bins"], + edges=np.array(obj["edges"], dtype=np.float64), + ) + + +def fit_uniform_bins( + values: np.ndarray, + n_bins: int = 256, + q_low: float = 0.01, + q_high: float = 0.99, +) -> BinConfig: + L = float(np.quantile(values, q_low)) + H = float(np.quantile(values, q_high)) + edges = np.linspace(L, H, n_bins + 1) + return BinConfig(mode="uniform", L=L, H=H, n_bins=n_bins, edges=edges) + + +def fit_quantile_bins( + values: np.ndarray, + n_bins: int = 256, + q_low: float = 0.01, + q_high: float = 0.99, +) -> BinConfig: + L = float(np.quantile(values, q_low)) + H = float(np.quantile(values, q_high)) + clipped = np.clip(values, L, H) + edges = np.quantile(clipped, np.linspace(0, 1, n_bins + 1)) + return BinConfig(mode="quantile", L=L, H=H, n_bins=n_bins, edges=edges) + + +def _encode_scalar(c: float, config: BinConfig) -> int: + if c < config.L: + return 0 + if c > config.H: + return config.n_bins + 1 + i = int(np.searchsorted(config.edges, c, side="right")) - 1 + i = max(0, min(i, config.n_bins - 1)) + return i + 1 # 1-based interior index + + +def _decode_scalar(idx: int, config: BinConfig) -> float: + if idx <= 0: + return config.L + if idx > config.n_bins: + return config.H + return float((config.edges[idx - 1] + config.edges[idx]) / 2.0) + + +def encode_cartesian_with_config(mol, config: BinConfig): + mol_no_h = Chem.RemoveHs(mol) + if mol_no_h.GetNumConformers() == 0: + raise ValueError("Molecule has no conformer / 3D coordinates.") + + smiles = Chem.MolToSmiles( + mol_no_h, + canonical=True, + isomericSmiles=True, + allHsExplicit=False, + allBondsExplicit=False, + ) + + if not mol_no_h.HasProp("_smilesAtomOutputOrder"): + raise ValueError("Mol is missing _smilesAtomOutputOrder after MolToSmiles.") + + atom_order_raw = mol_no_h.GetProp("_smilesAtomOutputOrder") + atom_order = list(map(int, ast.literal_eval(atom_order_raw))) + + expected_atom_tokens = [ + _expected_plain_token(mol_no_h.GetAtomWithIdx(idx)) for idx in atom_order + ] + + tokens = tokenize_smiles(smiles, expected_atom_tokens=expected_atom_tokens) + dw = config.digit_width + + out_parts = [] + atom_idx_in_smiles = 0 + conformer = mol_no_h.GetConformer() + + for token in tokens: + if token["type"] == "atom": + if atom_idx_in_smiles >= len(atom_order): + raise ValueError("SMILES atom tokens exceed atom order mapping.") + + rd_idx = atom_order[atom_idx_in_smiles] + atom_text = token["text"] + atom_descriptor = atom_text if atom_text.startswith("[") else f"[{atom_text}]" + + pos = conformer.GetAtomPosition(rd_idx) + ix = _encode_scalar(pos.x, config) + iy = _encode_scalar(pos.y, config) + iz = _encode_scalar(pos.z, config) + + out_parts.append( + f"{atom_descriptor}{ix:0{dw}d}{iy:0{dw}d}{iz:0{dw}d};" + ) + atom_idx_in_smiles += 1 + else: + out_parts.append(token["text"]) + + if atom_idx_in_smiles != len(atom_order): + raise ValueError( + f"Atom count mismatch: mapped {atom_idx_in_smiles} atoms " + f"but expected {len(atom_order)}." + ) + + return "".join(out_parts), smiles + + +def decode_cartesian_with_config(enriched_string: str, config: BinConfig): + normalized = enriched_string.replace(";", "") + tokens = tokenize_enriched_v2(normalized, config.digit_width) + + smiles_parts = [] + coords = [] + for token in tokens: + if token["type"] == "atom_with_coords": + desc = token["atom_desc"] + desc_inner = desc[1:-1] + if desc_inner in _ORGANIC_SUBSET: + smiles_parts.append(desc_inner) + else: + smiles_parts.append(desc) + + ix, iy, iz = (int(round(v)) for v in token["coords"]) + x = _decode_scalar(ix, config) + y = _decode_scalar(iy, config) + z = _decode_scalar(iz, config) + coords.append((x, y, z)) + else: + smiles_parts.append(token["text"]) + + smiles = "".join(smiles_parts) + mol = Chem.MolFromSmiles(smiles, sanitize=False) + if mol is None: + raise ValueError(f"Failed to parse rebuilt SMILES: {smiles}") + if mol.GetNumAtoms() != len(coords): + raise ValueError( + f"Atom count mismatch: mol has {mol.GetNumAtoms()} atoms, " + f"coords list has {len(coords)} entries." + ) + + Chem.SanitizeMol(mol) + + conformer = Chem.Conformer(mol.GetNumAtoms()) + for idx, (x, y, z) in enumerate(coords): + conformer.SetAtomPosition(idx, Point3D(x, y, z)) + mol.AddConformer(conformer, assignId=True) + return mol + def encode_cartesian_binned(mol, bin_size, ranges=None): - """ - Serialize a 3D RDKit Mol into an enriched text representation where - the Cartesian coordinates are replaced by bin indices. - - Returns: - enriched_string (str): SMILES-like string with [atom] tokens. - smiles (str): Canonical SMILES of the heavy-atom molecule. - bins (list[np.ndarray]): [bins_x, bins_y, bins_z] used for binning. - ranges (list[tuple[float, float]]): Axis ranges used to construct bins. - """ mol_no_h = Chem.RemoveHs(mol) if mol_no_h.GetNumConformers() == 0: raise ValueError("Molecule has no conformer / 3D coordinates.") @@ -653,7 +786,7 @@ def encode_cartesian_binned(mol, bin_size, ranges=None): tokens = tokenize_smiles(smiles, expected_atom_tokens=expected_atom_tokens) if ranges is None: - ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + ranges = [(-11.0, 11.0), (-11.0, 11.0), (-11.0, 11.0)] if len(ranges) != 3: raise ValueError("ranges must be a sequence of three (start, end) tuples.") bins = get_bins_for_coords(ranges, bin_size=bin_size) @@ -705,16 +838,6 @@ def encode_cartesian_binned(mol, bin_size, ranges=None): def encode_cartesian_binned_v2(mol, bin_size, ranges=None): - """ - Serialize a 3D RDKit Mol into an enriched text representation where - the Cartesian coordinates are replaced by bin indices. - - Returns: - enriched_string (str): SMILES-like string with [atom] tokens. - smiles (str): Canonical SMILES of the heavy-atom molecule. - bins (list[np.ndarray]): [bins_x, bins_y, bins_z] used for binning. - ranges (list[tuple[float, float]]): Axis ranges used to construct bins. - """ mol_no_h = Chem.RemoveHs(mol) if mol_no_h.GetNumConformers() == 0: raise ValueError("Molecule has no conformer / 3D coordinates.") @@ -740,7 +863,7 @@ def encode_cartesian_binned_v2(mol, bin_size, ranges=None): tokens = tokenize_smiles(smiles, expected_atom_tokens=expected_atom_tokens) if ranges is None: - ranges = [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + ranges = [(-11.0, 11.0), (-11.0, 11.0), (-11.0, 11.0)] if len(ranges) != 3: raise ValueError("ranges must be a sequence of three (start, end) tuples.") bins = get_bins_for_coords(ranges, bin_size=bin_size) @@ -777,7 +900,7 @@ def encode_cartesian_binned_v2(mol, bin_size, ranges=None): iy_txt = f"{iy:0{digit_width}d}" iz_txt = f"{iz:0{digit_width}d}" - out_parts.append(f"{atom_descriptor}{ix_txt}{iy_txt}{iz_txt}") + out_parts.append(f"{atom_descriptor}{ix_txt}{iy_txt}{iz_txt};") atom_idx_in_smiles += 1 else: out_parts.append(token["text"]) @@ -792,20 +915,6 @@ def encode_cartesian_binned_v2(mol, bin_size, ranges=None): def decode_cartesian_binned_v2(enriched_string, bins, use_bin_center=True): - """ - Reconstruct an RDKit Mol (with conformer) from a v2 binned enriched string. - - Supports both schemas: - - With semicolons: [c]123123123;[n]456456456; - - Without semicolons: [c]123123123[n]456456456 - - The string must have been produced by ``encode_cartesian_binned_v2`` using - the same set of ``bins`` (one array per axis). - - Semicolons (';') are treated as optional delimiters and are stripped - before parsing, so both ``[c]123123123;`` and ``[c]123123123`` decode - identically here, while upstream code can still see the raw schema. - """ if len(bins) != 3: raise ValueError("bins must be a sequence of three bin arrays (x, y, z).") @@ -855,3 +964,31 @@ def decode_cartesian_binned_v2(enriched_string, bins, use_bin_center=True): conformer.SetAtomPosition(idx, Point3D(x, y, z)) mol.AddConformer(conformer, assignId=True) return mol + + +def decode_conformer_by_serialization( + enriched_string: str, + serialization_tag: str, + *, + bins=None, + uniform_config_path: Optional[str] = None, + quantile_config_path: Optional[str] = None, +): + mode = str(serialization_tag) + if mode == "cartesian": + return decode_cartesian_v2(enriched_string) + if mode == "cartesian_binned": + if bins is None: + raise ValueError("`bins` must be provided for cartesian_binned decoding.") + return decode_cartesian_binned_v2(enriched_string, bins) + if mode == "uniform": + cfg_path = uniform_config_path or str( + Path(__file__).resolve().parents[1] / "config" / "bin_configs" / "uniform_bins.json" + ) + return decode_cartesian_with_config(enriched_string, BinConfig.load(cfg_path)) + if mode == "quantile": + cfg_path = quantile_config_path or str( + Path(__file__).resolve().parents[1] / "config" / "bin_configs" / "quantile_bins.json" + ) + return decode_cartesian_with_config(enriched_string, BinConfig.load(cfg_path)) + raise ValueError(f"Unsupported serialization mode: {serialization_tag}") diff --git a/src/molgen3D/data_processing/utils.py b/src/molgen3D/data_processing/utils.py index 978649f..bfd41dc 100644 --- a/src/molgen3D/data_processing/utils.py +++ b/src/molgen3D/data_processing/utils.py @@ -1,3 +1,4 @@ +import ast import os import os.path as osp import re @@ -7,15 +8,22 @@ import numpy as np import cloudpickle from collections import defaultdict -from typing import Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional, Tuple from pathlib import Path +from loguru import logger as log from rdkit.Chem.rdchem import HybridizationType from rdkit.Chem.rdchem import BondType as BT from rdkit.Chem.rdchem import ChiralType from rdkit import Chem from rdkit.Geometry import Point3D -from molgen3D.data_processing.smiles_encoder_decoder import encode_cartesian_v2 +from molgen3D.data_processing.smiles_encoder_decoder import ( + BinConfig, + encode_cartesian_binned, + encode_cartesian_binned_v2, + encode_cartesian_v2, + encode_cartesian_with_config, +) dihedral_pattern = Chem.MolFromSmarts('[*]~[*]~[*]~[*]') chirality = {ChiralType.CHI_TETRAHEDRAL_CW: -1., @@ -30,6 +38,17 @@ 'Ga': 21, 'Ge': 22, 'As': 23, 'Se': 24, 'Br': 25, 'Ag': 26, 'In': 27, 'Sb': 28, 'I': 29, 'Gd': 30, 'Pt': 31, 'Au': 32, 'Hg': 33, 'Bi': 34} +EMBEDDING_REGISTRY = { + "cartesian_v2": encode_cartesian_v2, + "cartesian": encode_cartesian_v2, + "cartesian_binned": encode_cartesian_binned, + "cartesian_binned_v2": encode_cartesian_binned_v2, + "uniform_binned": encode_cartesian_with_config, + "quantile_binned": encode_cartesian_with_config, +} + +REVISITED_SPLIT_FILE_MAP = {"train": "train", "valid": "val", "test": "test"} + def encode_cartesian_raw(mol, precision=4): """Legacy compatibility wrapper for the enriched representation.""" @@ -103,6 +122,137 @@ def save_processed_pickle( return output_path + +def save_grouped_pickle(output_path: str, iso_to_confs: Dict[str, List[Dict[str, Any]]]) -> None: + parent = osp.dirname(output_path) + if parent: + os.makedirs(parent, exist_ok=True) + with open(output_path, "wb") as fh: + pickle.dump(iso_to_confs, fh) + + +def get_embedding_func_and_config( + embedding_type: str, + bin_config_path: Optional[str] = None, +): + bin_config = None + if embedding_type in ("uniform_binned", "quantile_binned"): + if bin_config_path is None: + raise ValueError( + f"--bin_config_path is required for embedding_type={embedding_type!r}." + ) + bin_config = BinConfig.load(bin_config_path) + log.info( + "Loaded BinConfig from {} | mode={} L={:.4f} H={:.4f} n_bins={}", + bin_config_path, + bin_config.mode, + bin_config.L, + bin_config.H, + bin_config.n_bins, + ) + if embedding_type not in EMBEDDING_REGISTRY: + raise ValueError( + f"Unsupported embedding_type '{embedding_type}'. Options: {sorted(EMBEDDING_REGISTRY)}" + ) + return EMBEDDING_REGISTRY[embedding_type], bin_config + + +def parse_coordinate_ranges(ranges: str) -> List[Tuple[float, float]]: + try: + parsed_ranges = ast.literal_eval(f"[{ranges}]") + return [tuple(r) for r in parsed_ranges] + except Exception as exc: + log.error("Failed to parse ranges: {} | failure={}", ranges, exc) + return [(-13.0, 13.0), (-13.0, 13.0), (-13.0, 13.0)] + + +def get_revisited_split_path( + geom_raw_path: str, + split_name: str, + use_centered: bool = False, +) -> str: + if split_name not in REVISITED_SPLIT_FILE_MAP: + raise ValueError(f"Unsupported revisited split: {split_name!r}") + suffix = "_data_centered.pickle" if use_centered else "_data.pickle" + file_stem = REVISITED_SPLIT_FILE_MAP[split_name] + return osp.join(geom_raw_path, f"{file_stem}{suffix}") + + +def copy_single_conformer_mol(mol: Chem.Mol) -> Chem.Mol: + copied = Chem.Mol(mol) + if copied.GetNumConformers() > 1: + conf = Chem.Conformer(copied.GetConformer(0)) + copied.RemoveAllConformers() + copied.AddConformer(conf, assignId=True) + return copied + + +def _get_prop_float(props: Dict[str, Any], key: str) -> Optional[float]: + if key in props: + try: + return float(props[key]) + except Exception: + return None + return None + + +def extract_conf_meta( + conf_meta: Optional[Dict[str, Any]], + mol: Chem.Mol, +) -> Tuple[Optional[float], Optional[float], Optional[int]]: + energy = weight = None + conf_id = None + + if conf_meta: + if "totalenergy" in conf_meta: + energy = conf_meta.get("totalenergy") + elif "relativeenergy" in conf_meta: + energy = conf_meta.get("relativeenergy") + + if "boltzmannweight" in conf_meta: + weight = conf_meta.get("boltzmannweight") + elif "weight" in conf_meta: + weight = conf_meta.get("weight") + + if "geom_id" in conf_meta: + conf_id = conf_meta.get("geom_id") + + if mol is not None: + mol_props = mol.GetPropsAsDict() + if energy is None: + energy = _get_prop_float(mol_props, "totalenergy") + if weight is None: + weight = _get_prop_float(mol_props, "boltzmannweight") + if conf_id is None and "geom_id" in mol_props: + try: + conf_id = int(mol_props["geom_id"]) + except Exception: + conf_id = None + + if mol.GetNumConformers() > 0: + conf_props = mol.GetConformer().GetPropsAsDict() + if energy is None: + energy = _get_prop_float(conf_props, "totalenergy") + if weight is None: + weight = _get_prop_float(conf_props, "boltzmannweight") + if conf_id is None and "geom_id" in conf_props: + try: + conf_id = int(conf_props["geom_id"]) + except Exception: + conf_id = None + + try: + energy = float(energy) if energy is not None else None + except Exception: + energy = None + + try: + weight = float(weight) if weight is not None else None + except Exception: + weight = None + + return energy, weight, conf_id + class JsonlSplitWriter: def __init__(self, base_dir: str, split_name: str, chunk_size: int = 1_000_000): self.split_name = split_name @@ -170,6 +320,74 @@ def filter_mols( return selected + +def filter_revisited_mols( + smiles: Optional[str], + mols: List[Chem.Mol], + failures: Dict[str, int] = defaultdict(int), + max_confs: int = 30, +) -> List[Chem.Mol]: + if smiles and "." in smiles: + failures["dot_in_smiles"] += 1 + return [] + + mol_from_smiles = Chem.MolFromSmiles(smiles) if smiles else None + if mol_from_smiles is None: + failures["mol_from_smiles_failed"] += 1 + return [] + + num_confs = len(mols) + if max_confs is not None: + num_confs = max_confs + + selected: List[Chem.Mol] = [] + k = 0 + for mol in mols: + if mol is None: + failures["missing_mol"] += 1 + continue + + num_neighbors = [len(a.GetNeighbors()) for a in mol.GetAtoms()] + if not num_neighbors or np.max(num_neighbors) > 4: + failures["large_degree"] += 1 + continue + + selected.append(mol) + k += 1 + if k == num_confs: + break + + return selected + + +def filter_revisited_conformers_keep_dotted( + smiles: Optional[str], + mols: List[Chem.Mol], + failures: Dict[str, int], +) -> List[Tuple[Chem.Mol, int]]: + if smiles and "." in smiles: + failures["dot_in_smiles"] += 1 + + mol_from_smiles = Chem.MolFromSmiles(smiles) if smiles else None + if mol_from_smiles is None: + failures["mol_from_smiles_failed"] += 1 + return [] + + selected: List[Tuple[Chem.Mol, int]] = [] + for idx, mol in enumerate(mols): + if mol is None: + failures["missing_mol"] += 1 + continue + + num_neighbors = [len(a.GetNeighbors()) for a in mol.GetAtoms()] + if not num_neighbors or np.max(num_neighbors) > 4: + failures["large_degree"] += 1 + continue + + selected.append((mol, idx)) + + return selected + def load_pkl(file_path: Path) -> object: if not os.path.exists(file_path): raise FileNotFoundError(f"File {file_path} does not exist.")