D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/src/utils/splitting.py

@@ -0,0 +1,250 @@
+"""Molecular split strategies and auditable split reports."""
+
+import csv
+import json
+import os
+import random
+import uuid
+from pathlib import Path
+
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem.Scaffolds import MurckoScaffold
+
+
+SPLIT_REPORT_SCHEMA_VERSION = 1
+SPLIT_NAMES = ("train", "val", "test")
+DEFAULT_SPLIT_FRACTIONS = (0.8, 0.1, 0.1)
+
+
+def murcko_scaffold(smiles, include_chirality=False):
+    molecule = Chem.MolFromSmiles(str(smiles))
+    if molecule is None:
+        raise ValueError(
+            f"RDKit could not parse SMILES {smiles!r} while computing a Murcko scaffold. "
+            "Fix the molecule or ensure invalid graph rows are filtered before splitting."
+        )
+    scaffold = MurckoScaffold.GetScaffoldForMol(molecule)
+    return Chem.MolToSmiles(
+        scaffold,
+        canonical=True,
+        isomericSmiles=bool(include_chirality),
+    )
+
+
+def scaffold_split_indices(
+    smiles,
+    seed=42,
+    fractions=DEFAULT_SPLIT_FRACTIONS,
+    include_chirality=False,
+):
+    """Assign whole Murcko-scaffold groups to reproducible 80/10/10 splits."""
+    fractions = tuple(float(value) for value in fractions)
+    if len(fractions) != 3 or any(value <= 0 for value in fractions):
+        raise ValueError(
+            f"split fractions must contain three positive values, got {fractions!r}."
+        )
+    if not np.isclose(sum(fractions), 1.0):
+        raise ValueError(
+            f"split fractions must sum to 1.0, got {fractions!r} "
+            f"(sum={sum(fractions)!r})."
+        )
+    groups = {}
+    row_scaffolds = []
+    for index, value in enumerate(smiles):
+        scaffold = murcko_scaffold(value, include_chirality=include_chirality)
+        row_scaffolds.append(scaffold)
+        groups.setdefault(scaffold, []).append(index)
+    if len(groups) < 3:
+        raise ValueError(
+            f"Scaffold splitting requires at least 3 distinct Murcko scaffolds, "
+            f"but found {len(groups)} across {len(row_scaffolds)} usable rows. "
+            "Add structurally diverse molecules or use split_strategy='random' or "
+            "a predefined 'set' column."
+        )
+
+    rng = random.Random(seed)
+    grouped = list(groups.items())
+    rng.shuffle(grouped)
+    grouped.sort(key=lambda item: -len(item[1]))
+    target_counts = np.asarray(fractions) * len(row_scaffolds)
+    assigned_groups = {name: [] for name in SPLIT_NAMES}
+    assigned_indices = {name: [] for name in SPLIT_NAMES}
+
+    for scaffold, indices in grouped:
+        scores = []
+        for position, name in enumerate(SPLIT_NAMES):
+            projected = np.asarray(
+                [len(assigned_indices[item]) for item in SPLIT_NAMES],
+                dtype=float,
+            )
+            projected[position] += len(indices)
+            scores.append(float(np.abs(projected - target_counts).sum()))
+        selected = SPLIT_NAMES[min(range(3), key=lambda index: (scores[index], index))]
+        assigned_groups[selected].append(scaffold)
+        assigned_indices[selected].extend(indices)
+
+    for empty_name in [
+        name for name in SPLIT_NAMES if not assigned_indices[name]
+    ]:
+        donors = [
+            name for name in SPLIT_NAMES if len(assigned_groups[name]) > 1
+        ]
+        if not donors:
+            raise ValueError(
+                "Scaffold grouping could not produce non-empty train, val, and test "
+                "splits without breaking a scaffold group. Use a more diverse dataset."
+            )
+        donor = max(donors, key=lambda name: len(assigned_indices[name]))
+        scaffold = min(
+            assigned_groups[donor],
+            key=lambda value: (len(groups[value]), value),
+        )
+        assigned_groups[donor].remove(scaffold)
+        assigned_groups[empty_name].append(scaffold)
+        moved = groups[scaffold]
+        assigned_indices[donor] = [
+            index for index in assigned_indices[donor] if index not in set(moved)
+        ]
+        assigned_indices[empty_name].extend(moved)
+
+    result = tuple(
+        np.asarray(sorted(assigned_indices[name]), dtype=int)
+        for name in SPLIT_NAMES
+    )
+    scaffold_sets = [
+        {row_scaffolds[index] for index in indices} for indices in result
+    ]
+    overlap = (
+        scaffold_sets[0] & scaffold_sets[1]
+        | scaffold_sets[0] & scaffold_sets[2]
+        | scaffold_sets[1] & scaffold_sets[2]
+    )
+    if overlap:
+        raise RuntimeError(
+            f"Internal scaffold split error: scaffolds crossed split boundaries: "
+            f"{sorted(overlap)!r}."
+        )
+    return result, row_scaffolds
+
+
+def _target_stats(values):
+    array = np.asarray(values, dtype=float)
+    result = []
+    for column in range(array.shape[1]):
+        usable = array[:, column]
+        usable = usable[~np.isnan(usable)]
+        result.append(
+            {
+                "count": int(array.shape[0]),
+                "usable_count": int(usable.size),
+                "missing_count": int(array.shape[0] - usable.size),
+                "min": float(usable.min()) if usable.size else None,
+                "max": float(usable.max()) if usable.size else None,
+                "mean": float(usable.mean()) if usable.size else None,
+                "std": float(usable.std(ddof=1)) if usable.size > 1 else None,
+            }
+        )
+    return result
+
+
+def write_split_report(
+    config,
+    strategy,
+    indices_by_split,
+    row_indices,
+    target_values,
+    target_names,
+    smiles,
+    scaffold_column,
+    include_chirality=False,
+):
+    """Write atomic JSON summary and row-level CSV split assignments."""
+    scaffolds = [
+        murcko_scaffold(value, include_chirality=include_chirality)
+        for value in smiles
+    ]
+    scaffold_sets = {
+        name: {scaffolds[index] for index in indices}
+        for name, indices in zip(SPLIT_NAMES, indices_by_split)
+    }
+    overlap = sorted(
+        scaffold_sets["train"] & scaffold_sets["val"]
+        | scaffold_sets["train"] & scaffold_sets["test"]
+        | scaffold_sets["val"] & scaffold_sets["test"]
+    )
+    if strategy == "scaffold" and overlap:
+        raise RuntimeError(
+            f"Scaffold leakage detected after scaffold splitting: {overlap!r}."
+        )
+
+    split_target_summary = {}
+    for name, indices in zip(SPLIT_NAMES, indices_by_split):
+        stats = _target_stats(np.asarray(target_values)[indices])
+        split_target_summary[name] = {
+            target: value for target, value in zip(target_names, stats)
+        }
+    report = {
+        "split_report_schema_version": SPLIT_REPORT_SCHEMA_VERSION,
+        "strategy": strategy,
+        "split_random_seed": config.get("split_random_seed", 42),
+        "fractions": dict(zip(SPLIT_NAMES, DEFAULT_SPLIT_FRACTIONS)),
+        "scaffold_column": scaffold_column,
+        "scaffold_include_chirality": bool(include_chirality),
+        "counts": {
+            name: int(len(indices))
+            for name, indices in zip(SPLIT_NAMES, indices_by_split)
+        },
+        "scaffold_counts": {
+            name: len(scaffold_sets[name]) for name in SPLIT_NAMES
+        },
+        "scaffold_overlap_count": len(overlap),
+        "scaffold_overlap": overlap,
+        "target_summary": split_target_summary,
+    }
+
+    model_path = Path(config["MODEL_PATH"])
+    model_path.mkdir(parents=True, exist_ok=True)
+    json_path = model_path / "split_report.json"
+    csv_path = model_path / "split_assignments.csv"
+    token = uuid.uuid4().hex
+    json_staging = json_path.with_name(f".{json_path.name}.{token}.tmp")
+    csv_staging = csv_path.with_name(f".{csv_path.name}.{token}.tmp")
+    try:
+        json_staging.write_text(
+            json.dumps(report, indent=2, sort_keys=True),
+            encoding="utf-8",
+        )
+        assignment = {}
+        for name, indices in zip(SPLIT_NAMES, indices_by_split):
+            for index in indices:
+                assignment[int(index)] = name
+        with open(csv_staging, "w", encoding="utf-8", newline="") as stream:
+            writer = csv.DictWriter(
+                stream,
+                fieldnames=[
+                    "dataset_index",
+                    "original_row_index",
+                    "split",
+                    "scaffold",
+                    "smiles",
+                ],
+            )
+            writer.writeheader()
+            for index in range(len(smiles)):
+                writer.writerow(
+                    {
+                        "dataset_index": index,
+                        "original_row_index": int(row_indices[index]),
+                        "split": assignment[index],
+                        "scaffold": scaffolds[index],
+                        "smiles": smiles[index],
+                    }
+                )
+        os.replace(json_staging, json_path)
+        os.replace(csv_staging, csv_path)
+    finally:
+        for staging in (json_staging, csv_staging):
+            if staging.exists():
+                staging.unlink()
+    return report

D4CMPP2/src/utils/supportfile_saver.py

@@ -0,0 +1,94 @@
+import pandas as pd
+import yaml
+import hashlib
+import inspect
+import json
+from . import PATH
+
+def save_funtionalgroup_csv(file_path,save_path):
+    """
+    save the functional group csv file
+    """
+    df = pd.read_csv(file_path)
+    df.to_csv(save_path+'/functional_group.csv',index=False)
+
+def save_network_refer(file_path,save_path):
+    """
+    save the network refer yaml file
+    """
+    with open(file_path, "r") as file:
+        network_refer = yaml.safe_load(file)
+    with open(save_path+'/network_refer.yaml', "w") as file:
+        yaml.dump(network_refer, file)
+
+def save_network_module(file_path,save_path):
+    """
+    save the network module
+    """
+    with open(file_path, "r") as file:
+        network_module = file.read()
+    with open(save_path+'/network.py', "w") as file:
+        file.write(network_module)
+
+def save_registered_network(network_class, save_path):
+    source_path = inspect.getsourcefile(network_class)
+    if source_path is None:
+        raise ValueError(
+            f"Cannot save network {network_class.__qualname__!r}: its source file "
+            "cannot be located. Define custom networks in an importable .py module."
+        )
+    with open(source_path, "r", encoding="utf-8") as stream:
+        source = stream.read()
+    class_name = network_class.__qualname__
+    if "." in class_name:
+        raise ValueError(
+            f"Cannot save nested network class {class_name!r}; define it at module scope."
+        )
+    with open(save_path + '/network.py', "w", encoding="utf-8") as file:
+        file.write(source)
+        file.write(f"\n\n# Saved D4CMPP2 entry point\nnetwork = {class_name}\n")
+    digest = hashlib.sha256(source.encode("utf-8")).hexdigest()
+    with open(save_path + "/network_identity.json", "w", encoding="utf-8") as file:
+        json.dump(
+            {
+                "module": network_class.__module__,
+                "class": class_name,
+                "source_sha256": digest,
+                "snapshot": "network.py",
+            },
+            file,
+            indent=2,
+            sort_keys=True,
+        )
+
+def save_config(config):
+    """
+    save the configuration file
+    """
+    with open(config['MODEL_PATH']+'/config.yaml', "w") as file:
+        yaml.dump(config, file)
+
+def save_additional_files(config):
+    """
+    save the additional files
+    """
+    if "sculptor_s" in config:
+        FRAG_REF = PATH.get_frag_ref_path(config)
+        save_funtionalgroup_csv(FRAG_REF,config['MODEL_PATH'])
+
+    NET_REF = PATH.get_network_refer(config)
+    save_network_refer(NET_REF,config['MODEL_PATH'])
+
+    try:
+        from D4CMPP2.networks.registry import get_model
+
+        definition = get_model(config.get("network_id", config["network"]))
+    except ValueError:
+        definition = None
+    if definition is None:
+        NET_MODULE = PATH.get_NET_DIR(config)+'/'+config['network']+'.py'
+        save_network_module(NET_MODULE,config['MODEL_PATH'])
+    else:
+        save_registered_network(definition.network, config['MODEL_PATH'])
+    
+    save_config(config)

D4CMPP2/src/utils/tools.py

@@ -0,0 +1,156 @@
+
+from rdkit.Chem.Draw import rdMolDraw2D
+import matplotlib.pyplot as plt
+import rdkit.Chem as Chem
+from rdkit.Chem import Draw
+
+import io
+import PIL
+import numpy as np
+from PIL import Image
+
+
+def _figure_to_pil(fig):
+    """Render a Matplotlib figure to an RGB Pillow image."""
+
+    fig.canvas.draw()
+    width, height = fig.canvas.get_width_height()
+    rgba = np.asarray(fig.canvas.buffer_rgba())
+    return Image.fromarray(rgba.reshape(height, width, 4), mode="RGBA").convert("RGB")
+
+
+def plot_learning_curve(train_losses,val_losses):
+    "Draw the learning curve"
+    train_losses = np.array(train_losses)
+    val_losses = np.array(val_losses)
+
+    plt.plot(train_losses, label='Training loss', c= 'k')
+    plt.plot(val_losses, label='Validation loss', c= 'r')
+    y_max = np.concatenate([train_losses, val_losses]).mean() * 3
+    plt.xlabel('Epochs')
+    plt.ylabel('Loss')
+    plt.ylim(0, y_max)
+    plt.legend()
+    plt.title('Learning curve')
+
+    fig = plt.gcf()
+    pil_image = _figure_to_pil(fig)
+    plt.close()
+    return pil_image
+
+def plot_prediction(targets, prediction_df):
+    "Draw the prediction plot"
+    train_df = prediction_df[prediction_df['set']=="train"]
+    val_df = prediction_df[prediction_df['set']=="val"]
+    test_df = prediction_df[prediction_df['set']=="test"]
+    
+    target_nums = len(targets)
+    row_num = 1
+    col_num = 1
+    while True:
+        if row_num*col_num>=target_nums:
+            break
+        if row_num==col_num:
+            row_num+=1
+        else:
+            col_num+=1
+    plt.subplots(row_num,col_num,figsize=(col_num*6,row_num*5))
+    for i in range(target_nums):
+        plt.subplot(row_num,col_num,i+1)
+        plt.scatter(train_df[f"{targets[i]}_true"],train_df[f"{targets[i]}_pred"],label="train",c='k',s=0.5,alpha=0.3)
+        plt.scatter(val_df[f"{targets[i]}_true"],val_df[f"{targets[i]}_pred"],label="val",c='r',s=0.5,alpha=0.8)
+        plt.scatter(test_df[f"{targets[i]}_true"],test_df[f"{targets[i]}_pred"],label="test",c='b',s=0.5,alpha=0.8)
+        plt.xlabel("target")
+        plt.ylabel("prediction")
+        plt.title(targets[i])
+        plt.legend()
+
+    fig = plt.gcf()
+    pil_image = _figure_to_pil(fig)
+    plt.close()
+    return pil_image    
+    
+
+    
+def mol_with_atom_index( mol ):
+    "Return the mol object with atom index"
+    if type(mol) is str:
+        mol = Chem.MolFromSmiles(mol)
+    atoms = mol.GetNumAtoms()
+    for idx in range( atoms ):
+        mol.GetAtomWithIdx( idx ).SetProp( 'molAtomMapNumber', str( mol.GetAtomWithIdx( idx ).GetIdx() ) )
+    return mol
+
+def showAtomHighlight(mol,atms_list,atom_with_index=True):
+    """Draw the molecule with atom highlight
+    Args:
+        mol: str or rdkit.Chem.Mol
+            the molecule to draw
+        atms_list: list of list of int e.g [[0,1,2],[3,4,5]]
+            the list of atoms to highlight. the list of atoms will be drawn with different colors.
+        atom_with_index: bool
+            if True, the atom index will be shown
+
+    """
+    if type(mol) is str:
+        mol = Chem.MolFromSmiles(mol)
+    atom_num = mol.GetNumAtoms()
+    if atom_with_index:
+        mol=mol_with_atom_index(mol)
+    hit_ats=[]
+    hit_bonds = []
+    hit_ats_colormap={}
+    hit_bonds_colormap={}
+    
+    colormap=[(1.0, 0.75, 0.79),
+               (1.0, 1.0, 0.6),
+               (0.74, 0.99, 0.79),
+               (0.68, 0.85, 0.90),
+               (0.87, 0.63, 0.87),
+               (0.90, 0.90, 0.98),
+               (0.82, 0.94, 0.75),
+               (0.56, 0.93, 0.56),
+               (0.50, 1.0, 0.83),
+               (0.80, 0.60, 1.0),
+               (1.0, 0.85, 0.73),
+               (1.0, 1.0, 0.6),
+               (0.53, 0.81, 0.92),
+               (1.0, 0.75, 0.79),
+               (1.0, 0.94, 0.84)
+              ]
+    for count,atms in enumerate(atms_list):
+        atms = [int(x) for x in atms]
+        hit_ats=hit_ats+atms
+        for i in range(len(atms)):
+            hit_ats_colormap[atms[i]]=colormap[count%len(colormap)]
+            for j in range(i+1,len(atms)):
+                bond=mol.GetBondBetweenAtoms(atms[i],atms[j])
+                if bond:
+                    bond_index=bond.GetIdx()
+                    hit_bonds.append(bond_index)
+                    hit_bonds_colormap[bond_index]=colormap[count%len(colormap)]
+                    
+                    
+    d = rdMolDraw2D.MolDraw2DCairo(1000, 800)
+    d.DrawMolecule(mol,highlightAtoms = hit_ats, highlightAtomColors=hit_ats_colormap,
+                   highlightBonds=hit_bonds, highlightBondColors=hit_bonds_colormap)
+    d.FinishDrawing()
+    d = d.GetDrawingText()
+    
+    img_buf = io.BytesIO(d)
+    plot(img_buf)
+    
+def drawSmiles(smiles):
+    img_buf = io.BytesIO()
+    img = Draw.MolToImage(Chem.MolFromSmiles(smiles), size=(300,300))
+    img.save(img_buf, format='png')
+    img_buf.seek(0)
+    plot(img_buf)
+
+def plot(buf):
+    image = PIL.Image.open(buf)
+
+    plt.imshow(image)
+    plt.axis('off')
+    plt.show()
+

D4CMPP2/src/utils/transfer_learning.py

@@ -0,0 +1,111 @@
+"""Transfer-learning state selection and audit reporting."""
+
+import hashlib
+import json
+import os
+import uuid
+from pathlib import Path
+
+import torch
+
+
+TRANSFER_REPORT_SCHEMA_VERSION = 1
+
+
+def _tensor_description(tensor):
+    return {
+        "shape": list(tensor.shape),
+        "dtype": str(tensor.dtype),
+    }
+
+
+def select_compatible_state(source_state, target_state):
+    """Return name/shape-compatible source state and a complete audit report."""
+    if not isinstance(source_state, dict):
+        raise TypeError(
+            "Transfer checkpoint must contain a model state-dict mapping parameter "
+            f"names to tensors, got {type(source_state).__name__}."
+        )
+
+    non_tensors = [
+        name for name, value in source_state.items() if not isinstance(value, torch.Tensor)
+    ]
+    if non_tensors:
+        raise TypeError(
+            "Transfer checkpoint contains non-tensor state entries "
+            f"{non_tensors[:10]!r}. Check that final.pth is a plain model state dict."
+        )
+
+    selected = {}
+    loaded = []
+    shape_mismatch = []
+    source_only = []
+    for name, source_value in source_state.items():
+        target_value = target_state.get(name)
+        if target_value is None:
+            source_only.append({"name": name, **_tensor_description(source_value)})
+        elif source_value.shape != target_value.shape:
+            shape_mismatch.append(
+                {
+                    "name": name,
+                    "source_shape": list(source_value.shape),
+                    "target_shape": list(target_value.shape),
+                    "source_dtype": str(source_value.dtype),
+                    "target_dtype": str(target_value.dtype),
+                }
+            )
+        else:
+            selected[name] = source_value
+            loaded.append(
+                {
+                    "name": name,
+                    "shape": list(source_value.shape),
+                    "source_dtype": str(source_value.dtype),
+                    "target_dtype": str(target_value.dtype),
+                }
+            )
+
+    target_only = [
+        {"name": name, **_tensor_description(value)}
+        for name, value in target_state.items()
+        if name not in source_state
+    ]
+    report = {
+        "transfer_report_schema_version": TRANSFER_REPORT_SCHEMA_VERSION,
+        "counts": {
+            "loaded": len(loaded),
+            "shape_mismatch": len(shape_mismatch),
+            "source_only": len(source_only),
+            "target_only": len(target_only),
+        },
+        "loaded": loaded,
+        "shape_mismatch": shape_mismatch,
+        "source_only": source_only,
+        "target_only": target_only,
+    }
+    return selected, report
+
+
+def file_sha256(path):
+    digest = hashlib.sha256()
+    with open(path, "rb") as stream:
+        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def write_transfer_report(report, model_path):
+    """Atomically write the additive transfer audit artifact."""
+    path = Path(model_path) / "transfer_report.json"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    staging = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+    try:
+        staging.write_text(
+            json.dumps(report, indent=2, sort_keys=True),
+            encoding="utf-8",
+        )
+        os.replace(staging, path)
+    finally:
+        if staging.exists():
+            staging.unlink()
+    return str(path)