Release GEOM scaffold substituent LoRA rank-16 checkpoint after 25 epochs (part 2)

geom/scaffold_murcko_substituent_lora_r16/runtime/src/morph/utils/metrics/tensor.py

@@ -0,0 +1,336 @@
+"""Geometric tensor metrics consumed directly from ``ChemBatch``.
+
+These run in every representation (no RDKit, no validity conditioning). Target
+stats (per-element-pair distance histogram, RoG histogram) are computed once at
+preprocessing time on training-set coords in Angstroms and stored in
+``Distributions``. Generated coords are expected to arrive in Angstroms too —
+``sample()`` in the LightningModule already multiplies by ``coordinate_std``
+before returning the final ``ChemBatch``.
+
+Also hosts the shared bin edges + histogram accumulators because both this
+module and ``preprocessing.py`` must agree on them exactly.
+"""
+
+from dataclasses import dataclass
+
+import torch
+from torchmetrics import Metric
+
+from morph.dataset.core.atomic_data import ChemBatch
+
+# Shared bin edges (Angstroms). Used both at preprocessing time to compute
+# target histograms and at metric time to bin generated samples — they must
+# match exactly. Defaults are generous enough for QM9 / GEOM / TMQM; values
+# above the upper bound are counted in an overflow accumulator (see the binning
+# helpers below) so silent clamping is visible at validation time.
+DIST_RANGE = (0.0, 15.0)
+DIST_N_BINS = 150
+RG_RANGE = (0.0, 8.0)
+RG_N_BINS = 80
+
+
+@dataclass(frozen=True)
+class BatchGeometry:
+    """Geometry shared by all tensor metrics for one generated batch."""
+
+    num_graphs: int
+    node_counts: torch.Tensor
+    pair_distances: torch.Tensor
+    pair_atom_lo: torch.Tensor
+    pair_atom_hi: torch.Tensor
+    pair_graph: torch.Tensor
+    radii_of_gyration: torch.Tensor
+
+
+def compute_batch_geometry(batch: ChemBatch) -> BatchGeometry:
+    """Compute pair distances and radii of gyration once for all metrics."""
+    rows, cols = batch.edge_index
+    upper = rows < cols
+    rows, cols = rows[upper], cols[upper]
+    pair_distances = (batch.x[rows] - batch.x[cols]).norm(dim=-1)
+    atom_i, atom_j = batch.a[rows], batch.a[cols]
+
+    counts = batch.num_nodes_per_graph
+    centers = batch.x.new_zeros((batch.num_graphs, batch.x.shape[-1]))
+    centers.index_add_(0, batch.batch, batch.x)
+    centers = centers / counts.clamp(min=1).unsqueeze(1)
+    squared_radius = (batch.x - centers[batch.batch]).square().sum(dim=-1)
+    mean_squared_radius = batch.x.new_zeros(batch.num_graphs)
+    mean_squared_radius.index_add_(0, batch.batch, squared_radius)
+    mean_squared_radius = mean_squared_radius / counts.clamp(min=1)
+
+    return BatchGeometry(
+        num_graphs=batch.num_graphs,
+        node_counts=counts,
+        pair_distances=pair_distances,
+        pair_atom_lo=torch.minimum(atom_i, atom_j),
+        pair_atom_hi=torch.maximum(atom_i, atom_j),
+        pair_graph=batch.batch[rows],
+        radii_of_gyration=mean_squared_radius.sqrt(),
+    )
+
+
+def dist_edges() -> torch.Tensor:
+    return torch.linspace(DIST_RANGE[0], DIST_RANGE[1], DIST_N_BINS + 1)
+
+
+def rg_edges() -> torch.Tensor:
+    return torch.linspace(RG_RANGE[0], RG_RANGE[1], RG_N_BINS + 1)
+
+
+def accumulate_pairwise_distance_hist(
+    hist: torch.Tensor,
+    coord: torch.Tensor,
+    atom_idx: torch.Tensor,
+    edges: torch.Tensor,
+) -> tuple[int, int]:
+    """Add upper-triangular pair distances to ``hist`` (A, A, B), in place.
+
+    The pair key is ``(min(a_i, a_j), max(a_i, a_j))`` so the histogram is
+    naturally invariant to atom ordering.
+
+    Returns ``(n_overflow, n_total)`` for the pairs in this molecule, so callers
+    can track the fraction of distances that fell beyond the histogram range
+    (silent clamping otherwise hides "the cloud is way too large").
+    """
+    n = coord.shape[0]
+    if n < 2:
+        return 0, 0
+    i_idx, j_idx = torch.triu_indices(n, n, offset=1, device=coord.device)
+    dists = (coord[i_idx] - coord[j_idx]).norm(dim=-1)
+    lo = torch.minimum(atom_idx[i_idx], atom_idx[j_idx])
+    hi = torch.maximum(atom_idx[i_idx], atom_idx[j_idx])
+    edges_d = edges.to(coord.device)
+    n_overflow = int((dists >= edges_d[-1]).sum().item())
+    bin_idx = torch.bucketize(dists, edges_d).sub_(1).clamp_(0, hist.shape[-1] - 1)
+    a = hist.shape[0]
+    flat = (lo.long() * a + hi.long()) * hist.shape[-1] + bin_idx
+    hist.view(-1).index_add_(
+        0, flat, torch.ones_like(flat, dtype=hist.dtype, device=coord.device)
+    )
+    return n_overflow, int(dists.numel())
+
+
+def accumulate_rog_hist(
+    hist: torch.Tensor, coord: torch.Tensor, edges: torch.Tensor
+) -> tuple[int, int]:
+    """Add one molecule's radius of gyration into ``hist`` (B,), in place.
+
+    Returns ``(n_overflow, n_total)`` (0 or 1 each) for the same reason as
+    ``accumulate_pairwise_distance_hist``.
+    """
+    if coord.shape[0] == 0:
+        return 0, 0
+    centered = coord - coord.mean(dim=0)
+    rog = (centered.pow(2).sum(dim=-1).mean()).sqrt()
+    edges_d = edges.to(coord.device)
+    overflow = int((rog >= edges_d[-1]).item())
+    bin_idx = (
+        torch.bucketize(rog.unsqueeze(0), edges_d).sub_(1).clamp_(0, hist.shape[0] - 1)
+    )
+    hist.index_add_(0, bin_idx, torch.ones(1, dtype=hist.dtype, device=coord.device))
+    return overflow, 1
+
+
+class BatchMetric(Metric):
+    """Base class for metrics that consume a ``ChemBatch`` directly."""
+
+    def update(self, batch: ChemBatch) -> None:
+        """Accumulate tensor statistics from molecular batches."""
+        raise NotImplementedError
+
+    def compute(self) -> torch.Tensor:
+        """Reduce the accumulated tensor statistics to a metric value."""
+        raise NotImplementedError
+
+
+class PairwiseDistanceL1(BatchMetric):
+    """Mean per-element-pair L1 distance between gen and target distance histograms.
+
+    For every ordered (lo, hi) atom-type pair the *target* has meaningful mass
+    for, we normalize gen and target slices independently and report
+    ``|gen - target|.sum()``. The mean over those pairs is a 0-2 metric where 0
+    means perfect distributional match.
+
+    Note: gen-empty pairs are *not* excluded. If the target has C-N distances
+    but the model never generates a C-N pair, the L1 against zeros is 1.0 and
+    that pair counts in the average — which is what we want.
+    """
+
+    def __init__(
+        self, target_hist: torch.Tensor, min_pair_count: float = 10.0, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.register_buffer("target_hist", target_hist.detach().to(torch.float32))
+        self.register_buffer("edges", dist_edges())
+        self.min_pair_count = min_pair_count
+        self.add_state(
+            "gen_hist",
+            default=torch.zeros_like(self.target_hist),
+            dist_reduce_fx="sum",
+        )
+
+    def update(self, geometry: BatchGeometry) -> None:
+        """Accumulate pair-distance histograms by unordered atom-type pair."""
+        distances = geometry.pair_distances
+        if distances.numel() == 0:
+            return
+        edges = self.edges.to(distances.device)
+        bins = (
+            torch.bucketize(distances, edges)
+            .sub_(1)
+            .clamp_(0, self.gen_hist.shape[-1] - 1)
+        )
+        num_atom_types = self.gen_hist.shape[0]
+        flat = (
+            geometry.pair_atom_lo.long() * num_atom_types + geometry.pair_atom_hi.long()
+        ) * self.gen_hist.shape[-1] + bins
+        self.gen_hist.view(-1).index_add_(
+            0, flat, torch.ones_like(flat, dtype=self.gen_hist.dtype)
+        )
+
+    def compute(self) -> torch.Tensor:
+        """Average normalized histogram L1 errors over populated target pairs."""
+        gen = self.gen_hist  # (A, A, B)
+        tgt = self.target_hist
+        pair_t = tgt.sum(dim=-1)
+        # Target-populated only: a pair the training set knows about counts in
+        # the metric even if the model generates zero of them (gen_n = 0 →
+        # L1 == tgt_n.sum() == 1, which is maximal pain).
+        populated = pair_t >= self.min_pair_count
+        if not populated.any():
+            return torch.tensor(0.0, device=gen.device)
+        gen_n = gen / gen.sum(dim=-1).unsqueeze(-1).clamp(min=1.0)
+        tgt_n = tgt / pair_t.unsqueeze(-1).clamp(min=1.0)
+        l1 = (gen_n - tgt_n).abs().sum(dim=-1)  # (A, A)
+        return l1[populated].mean()
+
+
+class MinDistanceViolation(BatchMetric):
+    """Fraction of generated molecules whose smallest pairwise distance is below
+    a physical threshold (default 0.5 Å). Catches collapse and overlap.
+    """
+
+    def __init__(self, threshold: float = 0.5, **kwargs):
+        super().__init__(**kwargs)
+        self.threshold = float(threshold)
+        self.add_state("violations", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, geometry: BatchGeometry) -> None:
+        """Count graphs with an atom pair below the distance threshold."""
+        valid = geometry.node_counts >= 2
+        minima = geometry.pair_distances.new_full((geometry.num_graphs,), float("inf"))
+        minima.scatter_reduce_(
+            0, geometry.pair_graph, geometry.pair_distances, reduce="amin"
+        )
+        self.violations += ((minima < self.threshold) & valid).sum()
+        self.total += valid.sum()
+
+    def compute(self) -> torch.Tensor:
+        """Return the violation fraction among graphs with at least two atoms."""
+        if self.total <= 0:
+            return torch.tensor(0.0, device=self.violations.device)
+        return self.violations / self.total
+
+
+class RangeOverflow(BatchMetric):
+    """Fraction of generated values that landed outside the histogram range.
+
+    If this is non-trivial (say >1%), the fixed range in this module is too
+    small for the dataset and we're silently clamping into the last bin —
+    revisit ``DIST_RANGE`` / ``RG_RANGE`` or switch to target-derived edges.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.register_buffer("dist_edges_buf", dist_edges())
+        self.register_buffer("rog_edges_buf", rg_edges())
+        self.add_state("dist_over", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("dist_total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("rog_over", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("rog_total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, geometry: BatchGeometry) -> None:
+        """Count distances and radii outside their histogram ranges."""
+        self.dist_over += (geometry.pair_distances >= self.dist_edges_buf[-1]).sum()
+        self.dist_total += geometry.pair_distances.numel()
+        valid = geometry.node_counts > 0
+        self.rog_over += (
+            (geometry.radii_of_gyration >= self.rog_edges_buf[-1]) & valid
+        ).sum()
+        self.rog_total += valid.sum()
+
+    def compute(self) -> torch.Tensor:
+        """Return the larger of the distance and radius overflow fractions."""
+        d = self.dist_over / self.dist_total.clamp(min=1.0)
+        r = self.rog_over / self.rog_total.clamp(min=1.0)
+        return torch.maximum(d, r)
+
+
+class RoGL1(BatchMetric):
+    """L1 distance between gen and target normalized radius-of-gyration histograms."""
+
+    def __init__(self, target_hist: torch.Tensor, **kwargs):
+        super().__init__(**kwargs)
+        self.register_buffer("target_hist", target_hist.detach().to(torch.float32))
+        self.register_buffer("edges", rg_edges())
+        self.add_state(
+            "gen_hist",
+            default=torch.zeros_like(self.target_hist),
+            dist_reduce_fx="sum",
+        )
+
+    def update(self, geometry: BatchGeometry) -> None:
+        """Accumulate radii of gyration for nonempty graphs."""
+        valid = geometry.node_counts > 0
+        radii = geometry.radii_of_gyration[valid]
+        if radii.numel() == 0:
+            return
+        bins = (
+            torch.bucketize(radii, self.edges.to(radii.device))
+            .sub_(1)
+            .clamp_(0, self.gen_hist.shape[0] - 1)
+        )
+        self.gen_hist.index_add_(
+            0, bins, torch.ones_like(bins, dtype=self.gen_hist.dtype)
+        )
+
+    def compute(self) -> torch.Tensor:
+        """Return the L1 distance between normalized radius histograms."""
+        gen = self.gen_hist
+        tgt = self.target_hist
+        if gen.sum() <= 0 or tgt.sum() <= 0:
+            return torch.tensor(0.0, device=gen.device)
+        gen_n = gen / gen.sum().clamp(min=1.0)
+        tgt_n = tgt / tgt.sum().clamp(min=1.0)
+        return (gen_n - tgt_n).abs().sum()
+
+
+def build_batch_metrics(distributions) -> dict[str, BatchMetric]:
+    """Construct the geometric batch-side metric set from training target stats.
+
+    Atom-count / atom-type / edge / charge KLs live in the distribution
+    collection (see ``init_metrics``), not here — this set is purely geometric.
+
+    Returns an empty dict if Distributions lacks either geometric target stat
+    (e.g. an older cache built before these were added). Both are checked
+    independently in case one is added without the other later.
+    """
+    if (
+        distributions.pairwise_distance_histogram is None
+        or distributions.radius_of_gyration_histogram is None
+    ):
+        return {}
+    return {
+        "min_dist_violation": MinDistanceViolation(
+            threshold=0.5, sync_on_compute=False
+        ),
+        "range_overflow": RangeOverflow(sync_on_compute=False),
+        "pair_dist_l1": PairwiseDistanceL1(
+            distributions.pairwise_distance_histogram, sync_on_compute=False
+        ),
+        "rog_l1": RoGL1(
+            distributions.radius_of_gyration_histogram, sync_on_compute=False
+        ),
+    }

geom/scaffold_murcko_substituent_lora_r16/runtime/src/morph/utils/rdkit_utils.py

@@ -0,0 +1,291 @@
+"""RDKit molecule construction, conversion, and checks.
+
+Upstream attribution is recorded in docs/development.md.
+"""
+
+import threading
+
+import numpy as np
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem
+from rdkit.Chem.Descriptors import NumRadicalElectrons
+from scipy.spatial.transform import Rotation
+
+# Silence RDKit's noisy validity/kekulize warnings. These fire once per
+# generated molecule during validation and, on a shared Lustre stdout,
+# can block every Python thread for minutes -- including the TorchElastic
+# rendezvous keep-alive thread, which then times out and kills the job.
+RDLogger.DisableLog("rdApp.*")
+
+ArrT = np.ndarray
+
+
+class PeriodicTable:
+    """Thread-safe singleton wrapper for RDKit periodic-table queries."""
+
+    _instance = None
+
+    def __new__(cls):
+        """Return the shared periodic-table wrapper instance."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+
+        return cls._instance
+
+    def __init__(self):
+        self._table = Chem.GetPeriodicTable()
+
+        # Just to be certain that vocab objects are thread safe
+        self._pt_lock = threading.Lock()
+
+    def atomic_from_symbol(self, symbol: str) -> int:
+        """Look up an atomic number under the periodic-table lock."""
+        with self._pt_lock:
+            symbol = symbol.upper() if len(symbol) == 1 else symbol
+            atomic = self._table.GetAtomicNumber(symbol)
+
+        return atomic
+
+    def symbol_from_atomic(self, atomic_num: int) -> str:
+        """Look up an element symbol under the periodic-table lock."""
+        with self._pt_lock:
+            token = self._table.GetElementSymbol(atomic_num)
+
+        return token
+
+    def valence(self, atom: str | int) -> int:
+        """Look up the default valence for an element symbol or atomic number."""
+        with self._pt_lock:
+            valence = self._table.GetDefaultValence(atom)
+
+        return valence
+
+
+PT = PeriodicTable()
+
+IDX_BOND_MAP = {
+    0: "NO_BOND",
+    1: Chem.BondType.SINGLE,
+    2: Chem.BondType.DOUBLE,
+    3: Chem.BondType.TRIPLE,
+    4: Chem.BondType.AROMATIC,
+}
+BOND_IDX_MAP = {bond: idx for idx, bond in IDX_BOND_MAP.items()}
+
+# IDX_CHARGE_MAP = {0: 0, 1: 1, 2: 2, 3: 3, 4: -1, 5: -2, 6: -3}
+# CHARGE_IDX_MAP = {charge: idx for idx, charge in IDX_CHARGE_MAP.items()}
+
+
+# TODO merge these with check functions in other files
+
+
+def tensors_to_rdkit_mol(
+    atom_types: list[str],
+    coords: np.ndarray,
+    charge_types: list[int],
+    edge_types: list[Chem.BondType],
+    edge_index: list[tuple[int, int]],
+    sanitize: bool = True,
+    implicit_hydrogens: bool = False,
+):
+    mol = Chem.EditableMol(Chem.Mol())
+
+    # Add atoms with formal charges
+    for idx, atomic in enumerate(atom_types):
+        atom = Chem.Atom(atomic)
+        charge = charge_types[idx]
+        atom.SetFormalCharge(charge)
+        atom.SetNoImplicit(not implicit_hydrogens)
+        mol.AddAtom(atom)
+
+    mol = mol.GetMol()
+
+    # Add 3D coords
+    conf = Chem.Conformer(len(atom_types))
+
+    conf.SetPositions(np.ascontiguousarray(coords, dtype=np.float64))
+
+    mol.AddConformer(conf, assignId=True)
+
+    mol = Chem.EditableMol(mol)
+    for edge, edge_type in zip(edge_index, edge_types, strict=True):
+        start, end = edge
+        mol.AddBond(start, end, edge_type)
+
+    try:
+        mol = mol.GetMol()
+        # DetermineBonds(mol)
+        for atom in mol.GetAtoms():
+            atom.UpdatePropertyCache(strict=False)
+    except Exception:
+        return None
+
+    if sanitize:
+        return sanitize_mol_correctly(mol, implicit_hydrogens=implicit_hydrogens)
+
+    return mol
+
+
+def sanitize_mol_correctly(mol: Chem.Mol, *, implicit_hydrogens: bool | None = False):
+    """Sanitize in place; None preserves each atom's existing hydrogen policy."""
+    for a in mol.GetAtoms():
+        if implicit_hydrogens is not None:
+            a.SetNoImplicit(not implicit_hydrogens)
+        if a.HasProp("_MolFileHCount"):
+            a.ClearProp("_MolFileHCount")
+
+    flags = Chem.SanitizeFlags.SANITIZE_ALL
+    if all(a.GetNoImplicit() for a in mol.GetAtoms()):
+        flags &= ~Chem.SanitizeFlags.SANITIZE_ADJUSTHS
+
+    # Explicit-H graphs must not acquire extra H; implicit-H graphs need ADJUSTHS.
+    err = Chem.SanitizeMol(
+        mol,
+        sanitizeOps=flags,
+        catchErrors=True,
+    )
+    if err:  # nonzero bitmask means some step failed
+        return None
+    return mol
+
+
+def mol_is_valid(
+    mol: Chem.Mol, allow_radical: bool = False, allow_charged: bool = False
+) -> bool:
+    # Validity is a predicate: do not let sanitization mutate the molecule that
+    # subsequent metrics, serialization, or dataset processing will consume.
+    if mol is None or mol.GetNumAtoms() == 0:
+        return False
+    sanitized_mol = sanitize_mol_correctly(Chem.Mol(mol), implicit_hydrogens=None)
+    return not (
+        sanitized_mol is None
+        or len(Chem.GetMolFrags(sanitized_mol)) > 1
+        or (not allow_radical and NumRadicalElectrons(sanitized_mol) != 0)
+        or (not allow_charged and Chem.GetFormalCharge(sanitized_mol) != 0)
+    )
+
+
+def calc_energy(mol: Chem.Mol, per_atom: bool = False, level: str = "MMFF94") -> float:
+    """Calculate the energy for an RDKit molecule using the MMFF forcefield
+
+    Evaluate the first conformer on a copy, leaving the original unchanged.
+
+    Args:
+        mol (Chem.Mol): RDKit molecule
+        per_atom (bool): Whether to normalise by the number of atoms.
+        level: Force-field variant; currently only "MMFF94" is supported.
+
+    Returns:
+        float: Energy of the molecule or None if the energy could not be calculated
+    """
+
+    mol_copy = Chem.Mol(mol)
+
+    try:
+        if level == "MMFF94":
+            mmff_props = AllChem.MMFFGetMoleculeProperties(
+                mol_copy, mmffVariant="MMFF94"
+            )
+            ff = AllChem.MMFFGetMoleculeForceField(mol_copy, mmff_props, confId=0)
+            energy = ff.CalcEnergy()
+            energy = energy / mol.GetNumAtoms() if per_atom else energy
+    except Exception:
+        energy = None
+
+    return energy
+
+
+def optimise_mol(
+    mol: Chem.rdchem.Mol, level: str = "MMFF94", max_iters: int = 1000
+) -> Chem.rdchem.Mol:
+    """Optimise the conformation of an RDKit molecule
+
+    Optimize the first conformer on a copy, leaving the original unchanged.
+
+    Args:
+        mol (Chem.Mol): RDKit molecule
+        max_iters (int): Maximum iterations for conformer optimization.
+        level: Force-field variant; currently only "MMFF94" is supported.
+
+    Returns:
+        Chem.Mol: Optimized copy, or None if the force-field call raises an error.
+    """
+
+    mol_copy = Chem.Mol(mol)
+    try:
+        if level == "MMFF94":
+            AllChem.MMFFOptimizeMolecule(mol_copy, maxIters=max_iters)
+    except Exception:
+        return None
+
+    return mol_copy
+
+
+def conf_distance(
+    mol1: Chem.rdchem.Mol, mol2: Chem.rdchem.Mol, fix_order: bool = True
+) -> float:
+    """Approximately align two molecules and then calculate RMSD between them
+
+    Alignment and distance use the default conformer of each molecule.
+
+    Args:
+        mol1 (Chem.Mol): First molecule to align
+        mol2 (Chem.Mol): Second molecule to align
+        fix_order (bool): Whether to fix the atom order of the molecules
+
+    Returns:
+        float: RMSD between molecules after approximate alignment
+    """
+
+    assert len(mol1.GetAtoms()) == len(mol2.GetAtoms())
+
+    if not fix_order:
+        raise NotImplementedError
+
+    coords1 = np.array(mol1.GetConformer().GetPositions())
+    coords2 = np.array(mol2.GetConformer().GetPositions())
+
+    # Firstly, centre both molecules
+    coords1 = coords1 - (coords1.sum(axis=0) / coords1.shape[0])
+    coords2 = coords2 - (coords2.sum(axis=0) / coords2.shape[0])
+
+    if coords1.shape[0] == 1:
+        # catch the case where the molecule is a single atom
+        return 0.0
+
+    else:
+        # Find the best rotation alignment between the centred mols
+        rotation, _ = Rotation.align_vectors(coords1, coords2)
+        aligned_coords2 = rotation.apply(coords2)
+
+        sqrd_dists = (coords1 - aligned_coords2) ** 2
+        rmsd = np.sqrt(sqrd_dists.sum(axis=1).mean())
+        return rmsd
+
+
+def smiles_from_mol(
+    mol: Chem.rdchem.Mol, canonical: bool = True, explicit_hs: bool = False
+) -> str | None:
+    """Create a SMILES string from a molecule
+
+    Args:
+        mol (Chem.Mol): RDKit molecule object
+        canonical (bool): Whether to create a canonical SMILES, default True
+        explicit_hs (bool): Add explicit hydrogen atoms before creating the SMILES.
+            If False, remove hydrogen atoms. Explicit and implicit hydrogen
+            representations can produce different strings.
+
+    Returns:
+        str: SMILES string which could be None if the SMILES generation failed
+    """
+
+    if mol is None:
+        return None
+
+    try:
+        mol = Chem.AddHs(mol) if explicit_hs else Chem.RemoveHs(mol)
+        smiles = Chem.MolToSmiles(mol, canonical=canonical)
+    except Exception:
+        smiles = None
+
+    return smiles

geom/scaffold_murcko_substituent_lora_r16/runtime/src/morph/utils/run_config.py

@@ -0,0 +1,192 @@
+"""Runtime settings, Hydra/Lightning configuration, callbacks, and run IDs."""
+
+import faulthandler
+import os
+import sys
+
+import hydra
+import pytorch_lightning as pl
+import torch
+import torch.distributed as dist
+import wandb
+from omegaconf import DictConfig
+from pytorch_lightning import Callback
+from pytorch_lightning.callbacks import (
+    EarlyStopping,
+    ModelCheckpoint,
+)
+from pytorch_lightning.loggers import WandbLogger
+from rdkit import RDLogger
+
+
+def configure_runtime(*, seed: int | None = None, diagnostics: bool = False) -> None:
+    """Set arithmetic, chemistry logging, and optional reproducibility diagnostics."""
+    torch.set_float32_matmul_precision("medium")
+    RDLogger.DisableLog("rdApp.*")
+    if seed is not None:
+        pl.seed_everything(seed)
+    if diagnostics and os.environ.get("MORPH_STALL_DUMP"):
+        faulthandler.dump_traceback_later(
+            int(os.environ["MORPH_STALL_DUMP"]), repeat=True, file=sys.stderr
+        )
+        faulthandler.enable(file=sys.stderr, all_threads=True)
+    if diagnostics and os.environ.get("MORPH_FORCE_MATH_SDP"):
+        torch.backends.cuda.enable_flash_sdp(False)
+        torch.backends.cuda.enable_mem_efficient_sdp(False)
+        torch.backends.cuda.enable_math_sdp(True)
+
+
+class ManualAccumulationWandbLogger(WandbLogger):
+    """Report manual-optimization updates on W&B's standard step axis."""
+
+    def __init__(self, *args, accumulation_steps: int = 1, **kwargs):
+        if accumulation_steps <= 0:
+            raise ValueError("accumulation_steps must be positive")
+        self._accumulation_steps = accumulation_steps
+        super().__init__(*args, **kwargs)
+
+    def log_metrics(self, metrics, step=None):
+        """Log metrics against completed updates rather than microbatches."""
+        # Lightning's logger counter advances per batch when its Trainer-side
+        # accumulation is one. Manual optimization accumulates inside the
+        # module, so convert that counter back to completed model updates.
+        if step is not None:
+            step //= self._accumulation_steps
+        return super().log_metrics(metrics, step=step)
+
+
+def validation_batches_per_rank(
+    num_samples: int, batch_size: int, world_size: int
+) -> int:
+    """Return the per-rank batch limit for an exact global sample count."""
+    if num_samples <= 0 or batch_size <= 0 or world_size <= 0:
+        raise ValueError(
+            "validation sample count, batch size and world size must be positive"
+        )
+    samples_per_distributed_batch = batch_size * world_size
+    if num_samples % samples_per_distributed_batch:
+        raise ValueError(
+            f"num_val_samples ({num_samples}) must be divisible by val batch size "
+            f"({batch_size}) * world size ({world_size})"
+        )
+    return num_samples // samples_per_distributed_batch
+
+
+def configure_manual_gradient_accumulation(module, trainer_kwargs: dict) -> None:
+    """Move Lightning's accumulation setting onto a manual-optimization module.
+
+    Lightning rejects ``Trainer(accumulate_grad_batches != 1)`` for manual
+    optimization. The optimizer mixin implements the equivalent accumulation,
+    while the Trainer must therefore retain its neutral value of one.
+    """
+    steps = int(trainer_kwargs.get("accumulate_grad_batches", 1))
+    if steps <= 0:
+        raise ValueError("accumulate_grad_batches must be positive")
+    if module.automatic_optimization:
+        return
+    module.manual_gradient_accumulation_steps = steps
+    trainer_kwargs["accumulate_grad_batches"] = 1
+
+
+def gather_rank_outputs(outputs: list) -> list | None:
+    """Gather picklable rank outputs on rank zero without replicating them."""
+    if not (dist.is_available() and dist.is_initialized()):
+        return outputs
+
+    is_global_zero = dist.get_rank() == 0
+    gathered = [None] * dist.get_world_size() if is_global_zero else None
+    dist.gather_object(outputs, gathered, dst=0)
+    if not is_global_zero:
+        return None
+    return [output for rank_outputs in gathered for output in rank_outputs]
+
+
+def build_callbacks(cfg: DictConfig) -> list[Callback]:
+    callbacks: list[Callback] = []
+
+    if "early_stopping" in cfg.callbacks:
+        hydra.utils.log.info("Adding callback <EarlyStopping>")
+        es_cfg = cfg.callbacks.early_stopping
+        # stopping_threshold (optional): halt the moment the monitored metric
+        # crosses this value (>= for mode=max), independent of the plateau path.
+        # Used to stop a run once validity reaches a target (e.g. 0.8).
+        callbacks.append(
+            EarlyStopping(
+                monitor=cfg.callbacks.monitor_metric,
+                mode=cfg.callbacks.monitor_metric_mode,
+                patience=es_cfg.patience,
+                verbose=es_cfg.verbose,
+                stopping_threshold=es_cfg.get("stopping_threshold", None),
+            )
+        )
+
+    if "model_checkpoints" in cfg.callbacks:
+        hydra.utils.log.info("Adding callback <ModelCheckpoint>")
+        callbacks.append(
+            ModelCheckpoint(
+                monitor=cfg.callbacks.monitor_metric,
+                mode=cfg.callbacks.monitor_metric_mode,
+                every_n_epochs=cfg.callbacks.model_checkpoints.get("every_n_epochs", 1),
+                save_top_k=cfg.callbacks.model_checkpoints.save_top_k,
+                verbose=cfg.callbacks.model_checkpoints.verbose,
+                save_last=cfg.callbacks.model_checkpoints.save_last,
+            )
+        )
+
+    if "every_n_epochs_checkpoint" in cfg.callbacks:
+        cb_cfg = cfg.callbacks.every_n_epochs_checkpoint
+        hydra.utils.log.info(
+            "Adding callback <ModelCheckpoint> for every "
+            f"{cb_cfg.every_n_epochs} epochs"
+        )
+        # Keep best-by-validation separate from latest-for-resume. Lightning only saves
+        # train-epoch checkpoints when save_on_train_epoch_end=True if validation is not
+        # run every epoch.
+        callbacks.append(
+            ModelCheckpoint(
+                dirpath="every_n_epochs",
+                every_n_epochs=cb_cfg.every_n_epochs,
+                save_top_k=cb_cfg.save_top_k,
+                verbose=cb_cfg.verbose,
+                save_last=cb_cfg.save_last,
+                save_on_train_epoch_end=cb_cfg.get("save_on_train_epoch_end", None),
+            )
+        )
+
+    # Task-specific callbacks are built generically: any ``callbacks.<name>`` that
+    # carries a ``_target_`` is instantiated as-is. This module therefore imports
+    # none of them, which matters because the per-task eval callbacks ship with
+    # their own recipe (edit / multimode) and may not be installed at
+    # all -- a hard import here would make this core module depend on them.
+    # Defaults live on the callback classes, so the config need only override.
+    for entry in (cfg.callbacks or {}).values():
+        if not isinstance(entry, DictConfig) or "_target_" not in entry:
+            continue  # the explicitly-handled callbacks above
+        hydra.utils.log.info(f"Adding callback <{entry._target_.rsplit('.', 1)[-1]}>")
+        callbacks.append(hydra.utils.instantiate(entry, _convert_="all"))
+
+    return callbacks
+
+
+def bootstrap_run_id() -> None:
+    """Set the shared run IDs before Hydra composes ``hydra.run.dir``.
+
+    A single process generates a W&B-style ID. Multi-process ``torchrun`` must
+    receive the ID in its environment so every rank shares one output directory.
+    """
+    if os.environ.get("MORPH_RUN_ID"):
+        os.environ.setdefault("WANDB_RUN_ID", os.environ["MORPH_RUN_ID"])
+        return
+
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size > 1:
+        raise RuntimeError(
+            "Multi-process runs require MORPH_RUN_ID in the environment (use the same "
+            "value as WANDB_RUN_ID). Example:\n"
+            "  export MORPH_RUN_ID=<shared-id>\n"
+            '  export WANDB_RUN_ID="$MORPH_RUN_ID"'
+        )
+
+    rid = wandb.util.generate_id()
+    os.environ["MORPH_RUN_ID"] = rid
+    os.environ["WANDB_RUN_ID"] = rid

geom/scaffold_murcko_substituent_lora_r16/runtime/src/morph/utils/vocab_utils.py

@@ -0,0 +1,177 @@
+"""Token vocabulary helpers and loss-weight computation."""
+
+from copy import deepcopy
+
+import torch
+from rdkit import Chem
+
+
+def z_to_atom_types(z):
+    """Convert the atomic numbers to atom symbols with rdkit."""
+
+    table = Chem.GetPeriodicTable()
+    return [table.GetElementSymbol(z_i) for z_i in z]
+
+
+def token_to_index(token_list, token: str):
+    return token_list.index(token)
+
+
+def index_to_token(token_list, index: int):
+    return token_list[index]
+
+
+def compute_token_weights(
+    token_list: list[str],
+    distribution: torch.Tensor,
+    special_token_names: list[str],
+    weight_alpha: float = 1.0,
+    type_loss_token_weights: str = "uniform",
+) -> torch.Tensor:
+    """
+    Compute token weights with special handling for special tokens.
+
+    This function can be used for both node tokens (atom types) and edge tokens.
+    It computes inverse frequency weights, normalizes regular tokens, and assigns
+    appropriate weights to special tokens.
+
+    Args:
+        token_list: List of all tokens (e.g., atom types or edge types)
+        distribution: Distribution of token types from training data.
+            MUST have the same length as token_list and be aligned by index.
+        special_token_names: List of special token names to handle specially
+            (e.g., ["<NO_BOND>"] for edges)
+        weight_alpha: Alpha parameter for weight scaling (default: 1.0)
+        type_loss_token_weights: "uniform" or "training" - if "uniform", returns
+            uniform weights (all 1.0)
+
+    Returns:
+        Weights tensor with same shape as distribution
+    """
+    # Validate inputs
+    if len(distribution) != len(token_list):
+        raise ValueError(
+            f"Distribution length ({len(distribution)}) must match "
+            f"token_list length ({len(token_list)})"
+        )
+
+    # Get indices for special tokens (with validation)
+    special_token_indices = set()
+    for token_name in special_token_names:
+        if token_name not in token_list:
+            print(f"Warning: Special token '{token_name}' not in token_list, skipping")
+            continue
+        idx = token_to_index(token_list, token_name)
+        special_token_indices.add(idx)
+
+    # --- Handle uniform weights case early ---
+    if type_loss_token_weights == "uniform":
+        return torch.ones_like(distribution)
+
+    # --- 1. Initial Weight Calculation (Inverse Frequency) ---
+    epsilon = 1e-8
+    weights = 1.0 / (distribution + epsilon)
+    weights = weights**weight_alpha  # Apply alpha scaling
+
+    # --- 2. Isolate & Normalize REGULAR Tokens ---
+    all_indices = set(range(len(token_list)))
+    regular_token_indices = list(all_indices - special_token_indices)
+
+    if not regular_token_indices:
+        return torch.ones_like(weights)
+
+    # Get weights for only the regular tokens
+    regular_weights = weights[regular_token_indices]
+
+    # Normalize *only* the regular weights to have a mean of 1.0
+    mean_regular_weight = regular_weights.mean()
+    if mean_regular_weight > 0:
+        regular_weights = regular_weights / mean_regular_weight
+
+    # --- 3. Build Final Weights Tensor ---
+    # Start with the full weights (including special tokens computed from distribution)
+    final_weights = weights.clone()
+
+    # Assign normalized regular weights to their correct positions
+    final_weights[regular_token_indices] = regular_weights
+
+    # Handle special tokens: use their actual inverse frequency weight,
+    # normalized the same way as regular tokens
+    for token_name in special_token_names:
+        if token_name not in token_list:
+            continue
+        token_idx = token_to_index(token_list, token_name)
+        special_weight = weights[token_idx]
+        if mean_regular_weight > 0:
+            special_weight = special_weight / mean_regular_weight
+        final_weights[token_idx] = special_weight
+
+    # set the minimum weight to 1.0
+    # final_weights = final_weights / final_weights.min()
+
+    return final_weights
+
+
+def init_uniform_prior(distributions):
+    """Return a new Distributions with uniform atom, edge, and charge type
+        distributions.
+
+    All other fields (n_atoms_distribution, coordinate_std) are copied from the
+    original. The original object is not modified.
+    """
+    prior = deepcopy(distributions)
+    for attr in (
+        "atom_type_distribution",
+        "edge_type_distribution",
+        "charge_type_distribution",
+    ):
+        dist = torch.ones_like(getattr(prior, attr))
+        setattr(prior, attr, dist / dist.sum())
+    return prior
+
+
+_PRIOR_MODES = ("uniform", "marginal", "absorbing")
+
+
+def build_token_prior(
+    distributions,
+    *,
+    edge_mode: str = "uniform",
+    atom_mode: str = "uniform",
+    charge_mode: str = "uniform",
+):
+    """Return a new Distributions with per-token prior modes.
+
+    Generalises :func:`init_uniform_prior` (``all modes == "uniform"`` reproduces
+    it exactly). The prior is the flow-matching *source*: it seeds
+    ``sample_prior_graph`` and the interpolator, so changing it changes the whole
+    transport path the denoiser is trained on.
+
+    Modes per token channel:
+      * ``uniform``   -- maximum entropy (the historical default; for edges this
+                        is why the prior graph starts massively over-bonded).
+      * ``marginal``  -- the data marginal already on ``distributions`` (DiGress-
+                        style; the source is close to the data, so far fewer edits).
+      * ``absorbing`` -- one-hot on token index 0 (``<NO_BOND>`` for edges): an
+                        empty-graph / masked source with zero deletions to make.
+    """
+    prior = deepcopy(distributions)
+    for attr, mode in (
+        ("atom_type_distribution", atom_mode),
+        ("edge_type_distribution", edge_mode),
+        ("charge_type_distribution", charge_mode),
+    ):
+        cur = getattr(prior, attr)
+        if mode == "uniform":
+            dist = torch.ones_like(cur)
+        elif mode == "marginal":
+            dist = cur.clamp_min(0).clone()  # keep the data marginal already present
+        elif mode == "absorbing":
+            dist = torch.zeros_like(cur)
+            dist[0] = 1.0
+        else:
+            raise ValueError(
+                f"prior mode for {attr} must be one of {_PRIOR_MODES}, got {mode!r}"
+            )
+        setattr(prior, attr, dist / dist.sum())
+    return prior

geom/scaffold_murcko_substituent_lora_r16/training_config.yaml

@@ -0,0 +1,479 @@
+representation: molecule
+seed: 42
+heads:
+  heads:
+    _target_: morph.model.layers.heads.MultiHeadModule
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    heads_configs:
+      graph_aggregation: mean
+      node_heads:
+        atom_type_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.atom_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        charge_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: ${len:${data.vocab.charge_tokens}}
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: state
+        ins_rate_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_rate: ${model.head_initialization.ins_rate}
+        do_del_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_del_probability}
+        do_sub_a_head:
+          input_dim: ${model.head_node_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_node_adapter_dim}
+          group: edit
+          initial_probability: ${model.head_initialization.do_sub_a_probability}
+      edge_heads:
+        edge_type_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: ${len:${data.vocab.edge_tokens}}
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+        do_sub_e_head:
+          input_dim: ${model.head_edge_dim}
+          output_dim: 1
+          adapter_dim: ${model.head_edge_adapter_dim}
+          group: state
+          initial_probability: ${model.head_initialization.do_sub_e_probability}
+  ins_edge_head:
+    _target_: morph.model.layers.heads.InsertionEdgeHead
+    _recursive_: false
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_edge_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    n_edge_types: ${len:${data.vocab.edge_tokens}}
+    n_atom_types: ${len:${data.vocab.atom_tokens}}
+    n_charge_types: ${len:${data.vocab.charge_tokens}}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+insertion_head:
+  mode: unimodal
+  dimension: 3
+  num_components: 1
+  model:
+    _target_: morph.model.layers.heads.UnimodalInsertionHead
+    hidden_dim: ${model.head_node_dim}
+    readout_dim: ${model.head_insertion_adapter_dim}
+    output_init_std: ${model.head_initialization.output_std}
+    independent_adapters: ${oc.select:model.head_independent_adapters,false}
+    N_a: ${len:${data.vocab.atom_tokens}}
+    N_c: ${len:${data.vocab.charge_tokens}}
+model:
+  module:
+    optimizer_config:
+      implementation: torch
+      optimizer:
+        execution: ${oc.env:MORPH_MUON_EXECUTION,native}
+        lr: 0.001
+        momentum: 0.95
+        weight_decay: 0.0
+        nesterov: true
+        ns_steps: 5
+        eps: 1.0e-07
+        adjust_lr_fn: original
+      aux_optimizer:
+        lr: 0.0001
+        betas:
+        - 0.9
+        - 0.95
+        eps: 1.0e-10
+        weight_decay: 0.0
+      gradient_clip_val: 1.0
+      gradient_clip_algorithm: norm
+      scheduler:
+        _target_: morph.utils.lr_schedulers.MultiGroupCosineWarmupLR
+        warmup_steps: 2000
+        total_steps: 200000
+        min_lr_fraction: 0.05
+        last_epoch: -1
+      interval: step
+      monitor: null
+    loss_weights:
+      do_sub_a: 0.2
+      sub_a_class: 0.03
+      do_sub_e: 0.25
+      sub_e_class: 0.5
+      do_del: 0.4
+      ins_rate: 0.15
+      ins_state: 0.01
+      ins_e: 0.4
+      x: 2.0
+      c: 0.5
+      ins_e_ii: 0.03
+    _target_: morph.model.finetune.scaffold.LightningModuleRatesScaffoldFinetune
+    insertion_head_mode: ${insertion_head.mode}
+    ins_noise_scale: ${data.ins_noise_scale}
+    ema_decay: 0.999
+    ema_warmup_steps: 5000
+    use_ema_for_eval: true
+    time_dist: ${if:${eq:${model.sample_time_strategy}, log}, ${model.beta_sched},
+      ${model.linear_sched}}
+    integrator: ${model.integrator}
+    vocab: ${data.vocab}
+    use_learnable_loss_weights: false
+    use_time_weights: false
+    model:
+      _target_: morph.model.finetune.scaffold.ScaffoldConditionalDiTBackboneWithHeads
+      _recursive_: false
+      atom_type_embedding_args: ${embeddings.atom_type_embedding}
+      edge_type_embedding_args: ${embeddings.edge_type_embedding}
+      time_embedding_args: ${embeddings.time_embedding}
+      node_count_embedding_args: ${embeddings.node_count_embedding}
+      incident_bond_embedding: false
+      backbone_model_args: ${model.backbone}
+      heads_args: ${heads.heads}
+      insertion_head_args: ${insertion_head.model}
+      ins_edge_head_args: ${heads.ins_edge_head}
+      xattn_encoder_args: ${model.xattn_encoder}
+      use_distance_bias: true
+    pretrained_ckpt: /cluster/project/krause/frankem/morph-chem/outputs/hf_geom_qm9_20261009/stage/geom/explicit_h/epoch_100.ckpt
+    pretrained_prefer_ema: true
+    freeze_pretrained_backbone: true
+    freeze_output_heads: false
+    exclude_scaffold_from_sub_loss: true
+    exclude_scaffold_from_x_loss: false
+    freeze_keep_positions: false
+    freeze_keep_charges: true
+  rotate: true
+  n_augmentations: 1
+  integrator: ${integrator.integrator}
+  beta_sched:
+    _target_: torch.distributions.Beta
+    concentration1: 2.0
+    concentration0: 1.0
+  linear_sched:
+    _target_: torch.distributions.Uniform
+    low: 0.0
+    high: 1.0
+  sample_time_strategy: linear
+  head_node_dim: ${model.backbone.out_node_nf}
+  head_edge_dim: ${model.backbone.out_edge_nf}
+  head_node_adapter_dim: 256
+  head_edge_adapter_dim: 256
+  head_insertion_adapter_dim: 256
+  head_insertion_edge_dim: 256
+  head_independent_adapters: true
+  head_initialization:
+    output_std: 0.001
+    coordinate_output_std: 0.0
+    do_del_probability: 0.5
+    do_sub_a_probability: 0.5
+    do_sub_e_probability: 0.5
+    ins_rate: 0.6931471805599453
+  token_weighting:
+    type_loss_token_weights: uniform
+    weight_alpha: 0.4
+  lora:
+    enabled: true
+    rank: 16
+    alpha: 16.0
+    dropout: 0.0
+    targets:
+    - qkv
+    - proj
+    - mlp.gate_value
+    - mlp.out_proj
+    freeze_base: true
+  backbone:
+    _target_: morph.model.backbones.conditional_dit.ConditionalDiTBackbone
+    _recursive_: false
+    d_input: ${embeddings.atom_type_embedding.out_dim}
+    d_cond: ${oc.eval:"${embeddings.time_embedding.out_dim} + ${embeddings.node_count_embedding.out_dim}"}
+    d_model: 512
+    out_node_nf: 256
+    out_edge_nf: 256
+    num_layers: 8
+    nhead: 16
+    mlp_ratio: 4.0
+    ffn_hidden_dim: null
+    ffn_multiple_of: 64
+    qk_norm: true
+    coordinate_output_init_std: ${model.head_initialization.coordinate_output_std}
+    center_coordinate_output: true
+    in_edge_nf: ${embeddings.edge_type_embedding.out_dim}
+    edge_pair_dim: 256
+    pair_bias_hidden_dim: 32
+    num_atom_types: ${len:${data.vocab.atom_tokens}}
+    lora: ${model.lora}
+    rbf_embedding_args:
+      _target_: morph.model.layers.embedding.RBFEmbedding
+      num_rbf: 128
+      rbf_dmax: 3.5
+      out_dim: 64
+      dropout: 0.0
+      trainable: false
+    ctx_dim: 128
+    use_cross_attn: true
+    use_node_cond: true
+    use_edge_cond: true
+  xattn_encoder:
+    _target_: morph.model.finetune.base.LightweightContextEncoder
+    num_atom_types: ${len:${data.vocab.atom_tokens}}
+    num_charges: ${len:${data.vocab.charge_tokens}}
+    atom_embed_dim: 64
+    charge_embed_dim: 16
+    pos_embed_dim: 32
+    ctx_dim: ${model.backbone.ctx_dim}
+    num_attn_layers: 4
+    nhead: 8
+    mlp_ratio: 4.0
+    dropout: 0.0
+    use_pair_bias: true
+    num_edge_types: ${len:${data.vocab.edge_tokens}}
+    edge_embed_dim: 32
+    num_rbf: 32
+    rbf_dmax: 5.0
+    pair_bias_hidden_dim: 32
+schedule:
+  _target_: morph.flow_matching.schedules.SmoothstepSchedule
+  shift: 0.8
+trainer:
+  trainer:
+    deterministic: warn
+    devices: 2
+    accelerator: gpu
+    strategy: ddp_find_unused_parameters_false
+    precision: bf16-mixed
+    max_epochs: 50
+    accumulate_grad_batches: 1
+    num_sanity_val_steps: 0
+    gradient_clip_val: null
+    gradient_clip_algorithm: norm
+    check_val_every_n_epoch: 1
+    num_val_samples: 4096
+    limit_val_batches: null
+    limit_predict_batches: 4
+    enable_progress_bar: false
+    max_time:
+      hours: 0
+      seconds: 169927
+  do_train: true
+  do_validate: false
+  do_predict: false
+  checkpoint_path: null
+  final_checkpoint_path: final.ckpt
+callbacks:
+  monitor_metric: ${if:${oc.eval:"'${representation}' in ('geometric_graph', 'molecule')"},val/validity,val/batch/pair_dist_l1}
+  monitor_metric_mode: ${if:${oc.eval:"'${representation}' in ('geometric_graph',
+    'molecule')"},max,min}
+  model_checkpoints:
+    save_top_k: 1
+    save_last: false
+    verbose: false
+    every_n_epochs: 5
+  every_n_epochs_checkpoint:
+    every_n_epochs: 5
+    save_top_k: -1
+    save_last: link
+    verbose: false
+    save_on_train_epoch_end: true
+  metric_audit:
+    _target_: run_audit.ValidationAudit
+    expected_samples: 4096
+    expected_scaffolds: 512
+logging:
+  enabled: true
+  project: morph-geom-scaffold-decoration
+  name: ${oc.env:MORPH_RUN_ID,base}
+  save_dir: .
+  offline: false
+  group: geom-e100-murcko-original-50ep-20261010
+  tags:
+  - substituent-insertion
+  - rank-16
+  - scaffold-decoration
+  - bemis-murcko
+  - geom-e100
+  - lora
+  - trainable-heads
+  - movable-scaffold
+  - jitter-0.1A
+  - scaffold-metrics
+integrator:
+  integrator:
+    _target_: morph.flow_matching.integration.RateIntegrator
+    vocab: ${data.vocab}
+    distributions: None
+    insertion_head_mode: ${insertion_head.mode}
+    insertion_dimension: ${insertion_head.dimension}
+    insertion_num_components: ${insertion_head.num_components}
+    time_strategy: log
+    time_beta_alpha: 2.0
+    time_beta_beta: 1.0
+    num_integration_steps: 100
+    ins_noise_scale: ${data.ins_noise_scale}
+    renoise_insertions: true
+    exact_schedule_integration: true
+    del_schedule: ${data.interpolator.del_schedule}
+    ins_schedule: ${data.interpolator.ins_schedule}
+    sub_schedule: ${data.interpolator.sub_schedule}
+    sub_e_schedule: ${data.interpolator.sub_e_schedule}
+data:
+  root_path: ${oc.env:PROJECT_ROOT}/data
+  preprocess_workers: 10
+  preprocess_chunk_size: 128
+  statistics_workers: ${data.preprocess_workers}
+  hydrogen_mode: explicit
+  cache_name: ${data.dataset_name}${if:${eq:${data.hydrogen_mode},implicit},_implicit_h,''}
+  source_dir: ${oc.env:GEOM_SOURCE_DIR,null}
+  max_atoms: 72
+  allowed_atom_symbols:
+  - H
+  - C
+  - 'N'
+  - O
+  - S
+  - F
+  - Cl
+  - Br
+  - P
+  - I
+  n_atoms_strategy: uniform
+  uniform_n_atoms_padding: 5
+  optimal_transport: equivariant
+  ins_noise_scale: 0.65
+  move_noise_scale: 0.5
+  allow_charged: true
+  preprocessing:
+    _target_: morph.dataset.core.preprocessing.Preprocessing
+    _recursive_: false
+    root: /cluster/project/krause/frankem/morph/data/geom
+    num_workers: ${data.statistics_workers}
+    train_dataset: ${data.datamodule.datasets.train}
+    atom_tokens_path: ${data.preprocessing.root}/processed/atom_tokens.txt
+    edge_tokens_path: ${data.preprocessing.root}/processed/edge_tokens.txt
+    charge_tokens_path: ${data.preprocessing.root}/processed/charge_tokens.txt
+    distributions_path: ${data.preprocessing.root}/processed/distributions.pt
+  vocab: null
+  interpolator:
+    _target_: morph.flow_matching.interpolation.Interpolator
+    distributions: null
+    optimal_transport: ${data.optimal_transport}
+    ins_noise_scale: ${data.ins_noise_scale}
+    move_noise_scale: ${data.move_noise_scale}
+    atom_type_noise_prob: 0.0
+    pre_align_iters: 2
+    c_move: 1.0
+    c_sub: 0.0
+    c_ins: 6.125
+    c_del: 6.125
+    del_schedule: ${schedule}
+    ins_schedule: ${schedule}
+    sub_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+    sub_e_schedule:
+      _target_: morph.flow_matching.schedules.SmoothstepSchedule
+      shift: 1.5
+    jitter_scale: ${data.jitter_scale}
+    insertion_assignment: ${data.insertion_assignment}
+  datamodule:
+    _target_: morph.dataset.conditional.scaffold.datamodule.ScaffoldLightningDataModule
+    vocab: ${data.vocab}
+    distributions: null
+    representation: ${representation}
+    datasets:
+      train:
+        _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+      val:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+      test:
+      - _target_: ${data.dataset_target}
+        root: ${data.preprocessing.root}
+        preprocess_workers: ${data.preprocess_workers}
+        preprocess_chunk_size: ${data.preprocess_chunk_size}
+        hydrogen_mode: ${data.hydrogen_mode}
+        source_dir: ${data.source_dir}
+        max_atoms: ${data.max_atoms}
+        allowed_atom_symbols: ${data.allowed_atom_symbols}
+        vocab: ${data.vocab}
+        distributions: null
+    interpolator: ${data.interpolator}
+    rotate: ${model.rotate}
+    n_augmentations: ${model.n_augmentations}
+    num_workers:
+      train: 20
+      val: 0
+      test: 0
+    batch_size:
+      train: 320
+      val: 32
+      test: 32
+    n_atoms_strategy: ${data.n_atoms_strategy}
+    uniform_n_atoms_padding: ${data.uniform_n_atoms_padding}
+    time_dist: ${model.module.time_dist}
+    scaffold_groups_dir: ${data.scaffold_groups_dir}
+    scaffold_kind: ${data.scaffold_kind}
+    enumerate_slicings: ${data.enumerate_slicings}
+    eval_panel: ${data.eval_panel}
+    train_prefetch_factor: 1
+  dataset_name: geom
+  dataset_target: morph.dataset.unconditional.geom.FlowMatchingGEOMDataset
+  scaffold_groups_dir: ${data.preprocessing.root}/processed/scaffold_groups
+  scaffold_kind: murcko
+  jitter_scale: 0.1
+  insertion_assignment: substituent
+  enumerate_slicings: false
+  eval_panel:
+    directory: /cluster/project/krause/frankem/morph-chem/logs/training/geom-scaffold-substituent-lora-r16-20261012T095326Z/panels
+    n_scaffolds: 512
+    samples_per_scaffold: 8
+    seed: 42
+embeddings:
+  atom_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.atom_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  edge_type_embedding:
+    _target_: morph.model.layers.embedding.Embedding
+    num_embeddings: ${len:${data.vocab.edge_tokens}}
+    embedding_dim: 64
+    out_dim: 128
+  time_embedding:
+    _target_: morph.model.layers.embedding.TimeEmbedding
+    embedding_dim: 128
+    out_dim: 128
+  node_count_embedding:
+    _target_: morph.model.layers.embedding.ExtrapolatableCountEmbedding
+    embedding_dim: 64
+    out_dim: 128
+    min_period: 128.0
+    max_period: 1024.0
+    max_count: 96.0

geom/scaffold_murcko_substituent_lora_r16/training_source.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2ee1bf919ad17d664f0614252f43bdcf39b2474e5a1069a37429a520adb8a2f
+size 510894