diff --git a/.gitignore b/.gitignore index 69a01de..1d3dbd6 100644 --- a/.gitignore +++ b/.gitignore @@ -185,6 +185,13 @@ memory/ static/iridium/* scripts/ benchmarks/ +datasets/ GDB20* +TMCONF40/ +TMCONF5/ +TMCONF_Econf.ods +tmconf5.tar docs/ + +TM_CONFORMER_BENCHMARK.md diff --git a/README.md b/README.md index 5ce6c4c..c0f22fa 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ ensemble.to_xyz("output.xyz") ### Named Presets -Six use-case presets are available out of the box: +Seven use-case presets are available out of the box: ```python from openconf import generate_conformers @@ -49,12 +49,19 @@ ensemble = generate_conformers(mol, preset="ensemble") # property predictio ensemble = generate_conformers(mol, preset="spectroscopic") # NMR / IR / VCD ensemble = generate_conformers(mol, preset="docking") # docking pose recovery ensemble = generate_conformers(mol, preset="macrocycle") # macrocyclic ring systems +ensemble = generate_conformers(mol, preset="transition_metal") # organometallics ``` For FEP-style analogue generation from a fixed pose, see [`generate_conformers_from_pose`](#5-analogue--fep-style-r-group-exploration) below. For macrocyclic ring systems, use the `"macrocycle"` preset, which enables low-mode following and a wide energy window to capture the full range of ring-pucker conformations. +For transition-metal complexes with 3D input coordinates, use the +`"transition_metal"` preset. Other presets also add a transition-metal move +budget automatically when a metal atom is detected; set +`ConformerConfig(auto_transition_metal_moves=False)` to preserve an exact custom +move schedule. + ### Custom Configuration For full control, pass a `ConformerConfig` directly. You can also use @@ -82,7 +89,7 @@ ensemble = generate_conformers(mol, config=config) ## Use-Case Examples -The right configuration depends on the downstream task. Four named presets +The right configuration depends on the downstream task. Named presets cover the most common workflows: ```python @@ -97,7 +104,7 @@ config.max_out = 200 # override a single field ensemble = generate_conformers(mol, config=config) ``` -Available presets: `"rapid"`, `"ensemble"`, `"spectroscopic"`, `"docking"`, `"analogue"`, `"macrocycle"`. +Available presets: `"rapid"`, `"ensemble"`, `"spectroscopic"`, `"docking"`, `"analogue"`, `"macrocycle"`, `"transition_metal"`. Below are representative wall-clock timings measured on a single CPU core (Apple M3 Pro), mean over 3 runs. diff --git a/openconf/api.py b/openconf/api.py index 3f69e46..d25827f 100644 --- a/openconf/api.py +++ b/openconf/api.py @@ -11,6 +11,7 @@ from .exceptions import OpenConfRuntimeError, OpenConfValueError from .perceive import ( StereoSignature, + _is_metal, build_rotor_model, conformer_matches_specified_stereochemistry, prepare_molecule, @@ -24,6 +25,38 @@ # Gas constant in kcal/(mol·K); kT at 298.15 K ≈ 0.5924 kcal/mol. _R_KCAL_PER_MOL_K = 1.987204e-3 +_AUTO_TM_MOVE_BUDGET = 0.35 +_AUTO_TM_MOVE_WEIGHTS = { + "tm_ligand_rotate": 0.75, + "tm_haptic_rotate": 0.25, +} +_TM_MOVE_TYPES = frozenset(_AUTO_TM_MOVE_WEIGHTS) + + +def _has_transition_metal_move_budget(config: ConformerConfig) -> bool: + """Return whether config already assigns positive budget to TM moves.""" + return any(config.move_probs.get(move_type, 0.0) > 0.0 for move_type in _TM_MOVE_TYPES) + + +def _with_auto_transition_metal_sampling( + config: ConformerConfig, + *, + has_metal: bool, + has_metal_input: bool, +) -> ConformerConfig: + """Return config with automatic TM sampling overlay when appropriate.""" + if not has_metal or not config.auto_transition_metal_moves or _has_transition_metal_move_budget(config): + return config + + total = sum(config.move_probs.values()) + move_probs = {move_type: prob * (1.0 - _AUTO_TM_MOVE_BUDGET) for move_type, prob in config.move_probs.items()} + for move_type, weight in _AUTO_TM_MOVE_WEIGHTS.items(): + move_probs[move_type] = move_probs.get(move_type, 0.0) + total * _AUTO_TM_MOVE_BUDGET * weight + + updates: dict[str, object] = {"move_probs": move_probs} + if has_metal_input and config.tm_seed_move_attempts == 0: + updates["tm_seed_move_attempts"] = 2 + return dataclasses.replace(config, **updates) def _filter_stereochemistry_consistent_conformers( @@ -355,8 +388,9 @@ def generate_conformers( method: generation method; `"hybrid"` is the default and recommended config: configuration options; defaults are used when omitted preset: named use-case preset; one of `"rapid"`, - `"ensemble"`, `"spectroscopic"`, `"docking"`. Mutually - exclusive with *config*; raises ValueError if both are supplied. + `"ensemble"`, `"spectroscopic"`, `"docking"`, `"analogue"`, + `"macrocycle"`, or `"transition_metal"`. Mutually exclusive with + *config*; raises ValueError if both are supplied. torsion_library: torsion library override; when omitted, uses the bundled cached CrystalFF-derived library. add_hs: add explicit hydrogens before embedding; set to @@ -397,8 +431,14 @@ def generate_conformers( if method == "hybrid": low_flex_tuning = get_runtime_tuning().low_flex_path + has_metal = any(_is_metal(atom) for atom in mol.GetAtoms()) + has_metal_input = has_metal and mol.GetNumConformers() > 0 + config = _with_auto_transition_metal_sampling(config, has_metal=has_metal, has_metal_input=has_metal_input) + has_tm_sampling_profile = has_metal and config.move_probs.get("tm_ligand_rotate", 0.0) > 0.0 use_low_flex_path = ( config.constraint_spec is None + and not has_metal_input + and not has_tm_sampling_profile and rotor_model.n_rotatable <= low_flex_tuning.max_rotatable and (low_flex_tuning.allow_macrocycles or not rotor_model.ring_info.get("has_macrocycle")) and (low_flex_tuning.allow_rings or not rotor_model.ring_info.get("ring_sizes")) diff --git a/openconf/config.py b/openconf/config.py index 496d2d7..d3b01b3 100644 --- a/openconf/config.py +++ b/openconf/config.py @@ -6,7 +6,15 @@ from .exceptions import OpenConfValueError from .tuning import get_default_move_probs -ConformerPreset = Literal["rapid", "ensemble", "spectroscopic", "docking", "analogue", "macrocycle"] +ConformerPreset = Literal[ + "rapid", + "ensemble", + "spectroscopic", + "docking", + "analogue", + "macrocycle", + "transition_metal", +] _SUPPORTED_MOVE_TYPES = frozenset( { "single_rotor", @@ -17,6 +25,8 @@ "crankshaft", "ring_kic", "amide_flip", + "tm_ligand_rotate", + "tm_haptic_rotate", } ) _SUPPORTED_PARENT_STRATEGIES = frozenset({"softmax", "uniform", "best"}) @@ -130,7 +140,7 @@ class ConformerConfig: values flatten the distribution (more exploration); smaller values concentrate sampling on the lowest-energy pool members. Default 2.0 matches typical MCMM practice and is unrelated to physical temperature. - final_select: How the final conformer set is chosen. + final_select: How final conformer set is chosen. skip_clash_check: skip the pre-minimization clash check entirely. When False (default), use a fast numpy-based clash filter that avoids expensive minimization of heavily clashed structures. For large or @@ -191,6 +201,10 @@ class ConformerConfig: preserved. collect_stats: record stage timings and counters for benchmark analysis and attach them to the returned ensemble. + auto_transition_metal_moves: add transition-metal move budget + automatically when molecule contains metal atoms and configured + move probabilities do not already include TM-specific moves. + Set to False to preserve exact custom move probabilities. torsion_multitry_attempts: pre-minimization torsion proposal count to try for clash-filtered torsion moves. The candidate with the lowest clash score is kept, reducing wasted minimizations on crowded @@ -230,6 +244,11 @@ class ConformerConfig: skipped for low mode following but still enter the MCMM pool normally. Default 1 (only the lowest-energy seed) keeps the Hessian cost fixed regardless of total seed count. + tm_seed_move_attempts: number of seed conformers to attempt per + available TM-specific move type when starting from input + coordinates. This supplements the protected input seed with + early metal-ligand rotation and haptic-rotation variants while + reusing the normal move and minimization machinery. """ max_out: int = 200 @@ -270,6 +289,7 @@ class ConformerConfig: patience: int = 150 auto_tune_large_flexible: bool = True collect_stats: bool = False + auto_transition_metal_moves: bool = True torsion_multitry_attempts: int = 4 use_low_mode_following: bool = False low_mode_eigenvalue_threshold: float = 100.0 @@ -278,6 +298,7 @@ class ConformerConfig: low_mode_scan_energy_threshold: float = 2390.0 low_mode_scan_max_steps: int = 10 low_mode_n_source_seeds: int = 1 + tm_seed_move_attempts: int = 0 def __post_init__(self) -> None: _require_at_least("max_out", self.max_out, 1) @@ -309,6 +330,7 @@ def __post_init__(self) -> None: _require_greater_than("low_mode_scan_energy_threshold", self.low_mode_scan_energy_threshold, 0.0) _require_at_least("low_mode_scan_max_steps", self.low_mode_scan_max_steps, 1) _require_at_least("low_mode_n_source_seeds", self.low_mode_n_source_seeds, 1) + _require_at_least("tm_seed_move_attempts", self.tm_seed_move_attempts, 0) _validate_move_probs(self.move_probs) if self.minimizer != "rdkit_mmff": @@ -364,9 +386,15 @@ def preset_config(preset: ConformerPreset) -> ConformerConfig: and low-mode following enabled to supplement ETKDG seeds with Hessian-guided scan points along collective ring-puckering modes. + - `"transition_metal"` — Organometallic and coordination complexes. + Keeps wide energy filtering, uses uniform parent sampling, enables + low-mode following, and allocates move budget to metal-ligand fragment + rotations while relying on shared metal-shell constraints. + Args: preset: one of `"rapid"`, `"ensemble"`, `"spectroscopic"`, - `"docking"`, `"analogue"`, `"macrocycle"`. + `"docking"`, `"analogue"`, `"macrocycle"`, + `"transition_metal"`. Returns: Configuration for requested use case @@ -384,6 +412,9 @@ def preset_config(preset: ConformerPreset) -> ConformerConfig: >>> config = preset_config("macrocycle") >>> config.energy_window_kcal 100.0 + >>> config = preset_config("transition_metal") + >>> config.use_low_mode_following + True """ match preset: case "rapid": @@ -468,8 +499,37 @@ def preset_config(preset: ConformerPreset) -> ConformerConfig: parent_strategy="softmax", final_select="diverse", ) + case "transition_metal": + return ConformerConfig( + max_out=100, + pool_max=500, + n_steps=250, + energy_window_kcal=100.0, + move_probs={ + "single_rotor": 0.25, + "multi_rotor": 0.15, + "correlated": 0.1, + "ring_flip": 0.05, + "tm_ligand_rotate": 0.25, + "tm_haptic_rotate": 0.1, + }, + use_low_mode_following=True, + low_mode_max_modes=3, + low_mode_scan_max_steps=4, + low_mode_n_source_seeds=1, + tm_seed_move_attempts=2, + seed_n_per_rotor=1, + seed_prune_rms_thresh=1.5, + prism_max_deviation=0.005, + do_final_refine=False, + minimize_batch_size=1, + parent_strategy="uniform", + final_select="diverse", + adaptive_moves=True, + ) case _: raise OpenConfValueError( f"Unknown preset {preset!r}. " - "Choose from: 'rapid', 'ensemble', 'spectroscopic', 'docking', 'analogue', 'macrocycle'." + "Choose from: 'rapid', 'ensemble', 'spectroscopic', 'docking', 'analogue', 'macrocycle', " + "'transition_metal'." ) diff --git a/openconf/constraints.py b/openconf/constraints.py new file mode 100644 index 0000000..4a44fc8 --- /dev/null +++ b/openconf/constraints.py @@ -0,0 +1,254 @@ +"""Reusable geometry constraints for conformer relaxation.""" + +from dataclasses import dataclass + +import numpy as np +from rdkit import Chem + +_DEFAULT_POSITION_FORCE_CONSTANT = 1000.0 +_METAL_POSITION_FORCE_CONSTANT = 1e4 +_METAL_LIGAND_DISTANCE_TOLERANCE = 0.05 +_METAL_LIGAND_DISTANCE_FORCE_CONSTANT = 100000.0 + + +@dataclass(frozen=True) +class PositionConstraint: + """Harmonic position constraint on an atom. + + Attributes: + atom_idx: constrained atom index + position: reference Cartesian position + force_constant: force constant in kcal/mol/A^2 + """ + + atom_idx: int + position: tuple[float, float, float] + force_constant: float = _DEFAULT_POSITION_FORCE_CONSTANT + + +@dataclass(frozen=True) +class DistanceConstraint: + """Harmonic distance-window constraint between two atoms. + + Attributes: + atom_i: first atom index + atom_j: second atom index + min_distance: lower distance bound in Angstrom + max_distance: upper distance bound in Angstrom + force_constant: force constant in kcal/mol/A^2 + """ + + atom_i: int + atom_j: int + min_distance: float + max_distance: float + force_constant: float + + +@dataclass(frozen=True) +class ConstraintModel: + """Geometry constraints shared by all minimization paths. + + Attributes: + position_constraints: atom position constraints + distance_constraints: atom-pair distance constraints + """ + + position_constraints: tuple[PositionConstraint, ...] = () + distance_constraints: tuple[DistanceConstraint, ...] = () + + @property + def constrained_atoms(self) -> frozenset[int]: + """Atom indices with position constraints.""" + return frozenset(constraint.atom_idx for constraint in self.position_constraints) + + @classmethod + def empty(cls) -> "ConstraintModel": + """Return model with no constraints.""" + return cls() + + @classmethod + def from_atom_positions( + cls, + mol: Chem.Mol, + atom_indices: frozenset[int], + force_constant: float = _DEFAULT_POSITION_FORCE_CONSTANT, + ) -> "ConstraintModel": + """Build position constraints from first conformer. + + Args: + mol: molecule containing reference conformer + atom_indices: atoms to constrain + force_constant: force constant in kcal/mol/A^2 + + Returns: + Constraint model with atom position constraints + """ + if not atom_indices or mol.GetNumConformers() == 0: + return cls.empty() + + conf = mol.GetConformer(mol.GetConformers()[0].GetId()) + constraints = [] + for atom_idx in sorted(atom_indices): + pos = conf.GetAtomPosition(int(atom_idx)) + constraints.append( + PositionConstraint( + atom_idx=int(atom_idx), + position=(float(pos.x), float(pos.y), float(pos.z)), + force_constant=force_constant, + ) + ) + return cls(position_constraints=tuple(constraints)) + + @classmethod + def from_metal_shell( + cls, + mol: Chem.Mol, + metal_atom_indices: frozenset[int], + position_force_constant: float = _METAL_POSITION_FORCE_CONSTANT, + distance_tolerance: float = _METAL_LIGAND_DISTANCE_TOLERANCE, + distance_force_constant: float = _METAL_LIGAND_DISTANCE_FORCE_CONSTANT, + ) -> "ConstraintModel": + """Build metal position and first-shell distance constraints. + + Args: + mol: molecule containing reference conformer + metal_atom_indices: metal atom indices + position_force_constant: metal position force constant + distance_tolerance: allowed metal-ligand distance tolerance + distance_force_constant: metal-ligand distance force constant + + Returns: + Constraint model for metal centers and direct ligands + """ + if not metal_atom_indices or mol.GetNumConformers() == 0: + return cls.empty() + + conf = mol.GetConformer(mol.GetConformers()[0].GetId()) + positions = cls.from_atom_positions(mol, metal_atom_indices, position_force_constant).position_constraints + distances: list[DistanceConstraint] = [] + for metal_idx in sorted(metal_atom_indices): + metal_pos = conf.GetAtomPosition(int(metal_idx)) + for neighbor in mol.GetAtomWithIdx(int(metal_idx)).GetNeighbors(): + ligand_idx = int(neighbor.GetIdx()) + ligand_pos = conf.GetAtomPosition(ligand_idx) + distance = float(metal_pos.Distance(ligand_pos)) + if distance <= 0.0: + continue + distances.append( + DistanceConstraint( + atom_i=int(metal_idx), + atom_j=ligand_idx, + min_distance=max(0.0, distance - distance_tolerance), + max_distance=distance + distance_tolerance, + force_constant=distance_force_constant, + ) + ) + return cls(position_constraints=positions, distance_constraints=tuple(distances)) + + def combine(self, other: "ConstraintModel") -> "ConstraintModel": + """Return merged constraints, with later position constraints winning. + + Args: + other: constraints to merge after self + + Returns: + Combined constraint model + """ + positions = {constraint.atom_idx: constraint for constraint in self.position_constraints} + positions.update({constraint.atom_idx: constraint for constraint in other.position_constraints}) + distances = { + (min(c.atom_i, c.atom_j), max(c.atom_i, c.atom_j), c.min_distance, c.max_distance): c + for c in self.distance_constraints + } + distances.update( + { + (min(c.atom_i, c.atom_j), max(c.atom_i, c.atom_j), c.min_distance, c.max_distance): c + for c in other.distance_constraints + } + ) + return ConstraintModel( + position_constraints=tuple(positions[idx] for idx in sorted(positions)), + distance_constraints=tuple(distances[key] for key in sorted(distances)), + ) + + def reset_positions(self, mol: Chem.Mol, conf_id: int) -> None: + """Snap constrained atoms and distances back to reference geometry. + + Args: + mol: molecule containing conformer + conf_id: conformer ID to update + """ + if not self.position_constraints and not self.distance_constraints: + return + conf = mol.GetConformer(int(conf_id)) + position_atoms = {constraint.atom_idx for constraint in self.position_constraints} + for constraint in self.position_constraints: + conf.SetAtomPosition(constraint.atom_idx, constraint.position) + for constraint in self.distance_constraints: + atom_i = int(constraint.atom_i) + atom_j = int(constraint.atom_j) + if atom_i in position_atoms and atom_j in position_atoms: + continue + pos_i = conf.GetAtomPosition(atom_i) + pos_j = conf.GetAtomPosition(atom_j) + vec = np.array([pos_j.x - pos_i.x, pos_j.y - pos_i.y, pos_j.z - pos_i.z]) + dist = float(np.linalg.norm(vec)) + if constraint.min_distance <= dist <= constraint.max_distance or dist <= 1e-12: + continue + target = 0.5 * (constraint.min_distance + constraint.max_distance) + scaled = vec * (target / dist) + if atom_i in position_atoms: + conf.SetAtomPosition(atom_j, (pos_i.x + scaled[0], pos_i.y + scaled[1], pos_i.z + scaled[2])) + elif atom_j in position_atoms: + conf.SetAtomPosition(atom_i, (pos_j.x - scaled[0], pos_j.y - scaled[1], pos_j.z - scaled[2])) + + def max_position_drift(self, mol: Chem.Mol, conf_id: int) -> float: + """Return maximum drift of position-constrained atoms. + + Args: + mol: molecule containing conformer + conf_id: conformer ID to inspect + + Returns: + Maximum distance from reference position + """ + if not self.position_constraints: + return 0.0 + conf = mol.GetConformer(int(conf_id)) + max_drift = 0.0 + for constraint in self.position_constraints: + pos = conf.GetAtomPosition(constraint.atom_idx) + ref = np.array(constraint.position) + cur = np.array([pos.x, pos.y, pos.z]) + max_drift = max(max_drift, float(np.linalg.norm(cur - ref))) + return max_drift + + +def add_constraints_to_force_field(ff: object, constraints: ConstraintModel, family: str) -> None: + """Add constraints to RDKit force field object when supported. + + Args: + ff: RDKit force field object + constraints: constraints to apply + family: force-field family, either `"MMFF"` or `"UFF"` + """ + if not constraints.position_constraints and not constraints.distance_constraints: + return + + position_method = getattr(ff, f"{family}AddPositionConstraint", None) + if position_method is not None: + for constraint in constraints.position_constraints: + position_method(int(constraint.atom_idx), 0.0, float(constraint.force_constant)) + + distance_method = getattr(ff, f"{family}AddDistanceConstraint", None) + if distance_method is not None: + for constraint in constraints.distance_constraints: + distance_method( + int(constraint.atom_i), + int(constraint.atom_j), + False, + float(constraint.min_distance), + float(constraint.max_distance), + float(constraint.force_constant), + ) diff --git a/openconf/data/runtime_tuning.json b/openconf/data/runtime_tuning.json index 82d7b3c..c1e4cac 100644 --- a/openconf/data/runtime_tuning.json +++ b/openconf/data/runtime_tuning.json @@ -48,7 +48,9 @@ "ring_flip": 0.09, "crankshaft": 0.05, "ring_kic": 0.05, - "amide_flip": 0.05 + "amide_flip": 0.05, + "tm_ligand_rotate": 0.0, + "tm_haptic_rotate": 0.0 }, "forced_periodic_move": "global_shake", "constrained_fallback_move": "single_rotor", @@ -62,7 +64,9 @@ "ring_flip": "single_rotor", "crankshaft": "single_rotor", "ring_kic": "crankshaft", - "amide_flip": "single_rotor" + "amide_flip": "single_rotor", + "tm_ligand_rotate": "single_rotor", + "tm_haptic_rotate": "tm_ligand_rotate" } }, "clash_check": { @@ -70,7 +74,9 @@ "ring_flip", "crankshaft", "ring_kic", - "amide_flip" + "amide_flip", + "tm_ligand_rotate", + "tm_haptic_rotate" ] }, "low_flex_path": { diff --git a/openconf/pool.py b/openconf/pool.py index 3b7aa39..d96947f 100644 --- a/openconf/pool.py +++ b/openconf/pool.py @@ -11,14 +11,17 @@ from .config import ConformerConfig from .dedupe import prism_dedupe -type Pool = ConformerPool - def _energy_or_inf(energy: float | None) -> float: """Return a finite sentinel for missing energies.""" return energy if energy is not None else float("inf") +def _is_protected(record: "ConformerRecord") -> bool: + """Return whether pool operations must keep conformer when possible.""" + return bool(record.tags.get("protected", False)) + + def _build_3d_descriptors( mol: Chem.Mol, conf_id: int, @@ -55,6 +58,19 @@ def _build_3d_descriptors( ] +def _build_selection_features(mol: Chem.Mol, conf_ids: list[int]) -> np.ndarray: + """Build final-selection feature matrix for conformers.""" + radii = rdFreeSASA.classifyAtoms(mol) + polar_atom_indices = np.array( + [atom.GetIdx() for atom in mol.GetAtoms() if atom.GetAtomicNum() not in {1, 6}], + dtype=int, + ) + features = [] + for cid in conf_ids: + features.append(_build_3d_descriptors(mol, cid, radii, polar_atom_indices)) + return np.asarray(features, dtype=float) + + def _softmax_parent_weights(energies: np.ndarray, temperature: float) -> np.ndarray: """Return numerically stable parent-selection weights.""" finite = np.isfinite(energies) @@ -90,19 +106,13 @@ def _pick_diverse_maxmin( conf_ids: conformer IDs to choose from energies: energies corresponding to each conformer k: conformer count to select - Returns: Selected conformer IDs, at most `k` """ if k >= len(conf_ids): return conf_ids - radii = rdFreeSASA.classifyAtoms(mol) - polar_atom_indices = np.array( - [atom.GetIdx() for atom in mol.GetAtoms() if atom.GetAtomicNum() not in {1, 6}], - dtype=int, - ) - features = np.array([_build_3d_descriptors(mol, cid, radii, polar_atom_indices) for cid in conf_ids]) + features = _build_selection_features(mol, conf_ids) # Normalize: zero mean, unit variance (avoid div-by-zero on flat features). mean = features.mean(axis=0) @@ -154,7 +164,7 @@ class ConformerRecord: class ParentSampler: """Cached parent selector for a conformer pool.""" - pool: Pool + pool: "ConformerPool" _dirty: bool = field(default=True, init=False, repr=False) _records_ref: dict[int, ConformerRecord] | None = field(default=None, init=False, repr=False) _records_version: int = field(default=-1, init=False, repr=False) @@ -311,11 +321,15 @@ def insert( assert self.config.pool_max is not None if self.size >= self.config.pool_max: # Refresh cached worst if stale (O(N) scan, amortised rare). + if self._worst_id is not None and self._worst_id not in self.records: + self._worst_dirty = True + if not self._worst_dirty and self._worst_id is not None and _is_protected(self.records[self._worst_id]): + self._worst_dirty = True if self._worst_dirty or self._worst_id is None: - self._worst_id = max( - self.records, - key=lambda cid: _energy_or_inf(self.records[cid].energy_kcal), - ) + removable = [cid for cid, record in self.records.items() if not _is_protected(record)] + if not removable: + return False + self._worst_id = max(removable, key=lambda cid: _energy_or_inf(self.records[cid].energy_kcal)) self._worst_energy = _energy_or_inf(self.records[self._worst_id].energy_kcal) self._worst_dirty = False @@ -325,18 +339,24 @@ def insert( # Remove worst (O(1) now that we know _worst_id). del self.records[self._worst_id] self.mol.RemoveConformer(self._worst_id) + self._worst_id = None self._worst_dirty = True # next worst needs a fresh scan # Add new conformer. + record_tags = tags or {} self.records[conf_id] = ConformerRecord( conf_id=conf_id, energy_kcal=energy, source=source, - tags=tags or {}, + tags=record_tags, ) # Update cached worst in O(1) if the new conformer is the new worst. - if not self._worst_dirty and (self._worst_id is None or energy > self._worst_energy): + if ( + not record_tags.get("protected", False) + and not self._worst_dirty + and (self._worst_id is None or energy > self._worst_energy) + ): self._worst_energy = energy self._worst_id = conf_id @@ -371,7 +391,8 @@ def dedupe(self) -> int: # Remove duplicates keep_set = set(keep_ids) - to_remove = [cid for cid in conf_ids if cid not in keep_set] + protected = {cid for cid, record in self.records.items() if _is_protected(record)} + to_remove = [cid for cid in conf_ids if cid not in keep_set and cid not in protected] for cid in to_remove: del self.records[cid] @@ -403,19 +424,30 @@ def select_final(self) -> list[int]: conf_ids = self.conf_ids energies = self.energies + protected_ids = [cid for cid in conf_ids if _is_protected(self.records[cid])] + if len(protected_ids) >= self.config.max_out: + sorted_protected = sorted( + protected_ids, + key=lambda cid: _energy_or_inf(self.records[cid].energy_kcal), + ) + return sorted_protected[: self.config.max_out] # Take lowest-energy conformers if self.config.final_select == "energy": - sorted_pairs = sorted(zip(energies, conf_ids, strict=True)) - return [cid for _, cid in sorted_pairs[: self.config.max_out]] + unprotected_pairs = sorted( + (energy, cid) for energy, cid in zip(energies, conf_ids, strict=True) if cid not in protected_ids + ) + n_needed = self.config.max_out - len(protected_ids) + return protected_ids + [cid for _, cid in unprotected_pairs[:n_needed]] # Diverse final selection (greedy MaxMin on shape descriptors) - return _pick_diverse_maxmin( + selected = _pick_diverse_maxmin( self.mol, - conf_ids, - energies, - k=self.config.max_out, + [cid for cid in conf_ids if cid not in protected_ids], + [energy for energy, cid in zip(energies, conf_ids, strict=True) if cid not in protected_ids], + k=self.config.max_out - len(protected_ids), ) + return protected_ids + selected def get_parent(self, strategy: str = "softmax") -> int | None: """Select a parent conformer for mutation. diff --git a/openconf/propose/hybrid.py b/openconf/propose/hybrid.py index 8c57f56..77644d8 100644 --- a/openconf/propose/hybrid.py +++ b/openconf/propose/hybrid.py @@ -10,18 +10,12 @@ from rdkit.Geometry import rdGeometry from ..config import ConformerConfig, ConstraintSpec +from ..constraints import ConstraintModel, add_constraints_to_force_field from ..dedupe import prism_dedupe from ..exceptions import OpenConfValueError from ..perceive import RotorModel, _is_metal, filter_constrained_rotors from ..pool import ConformerPool -from ..relax import ( - _METAL_LIGAND_DISTANCE_FORCE_CONSTANT, - _METAL_LIGAND_DISTANCE_TOLERANCE, - _METAL_POSITION_FORCE_CONSTANT, - _metal_ligand_reference_distances, - get_minimizer, - minimize_confs_mmff, -) +from ..relax import get_minimizer, minimize_confs_mmff from ..torsionlib import TorsionLibrary, get_default_torsion_library from ..tuning import ( get_runtime_tuning, @@ -52,6 +46,10 @@ ] _TORSION_MOVE_TYPES = frozenset({"single_rotor", "multi_rotor", "correlated", "global_shake"}) +_TM_SEED_MOVE_TYPES = ( + "tm_haptic_rotate", + "tm_ligand_rotate", +) def _resolve_runtime_tuned_config(config: ConformerConfig, rotor_model: RotorModel) -> tuple[ConformerConfig, bool]: @@ -136,27 +134,32 @@ def __init__( self.stats = stats self._metal_atom_indices: frozenset[int] = frozenset(a.GetIdx() for a in mol.GetAtoms() if _is_metal(a)) - self._metal_ref_pos: dict[int, np.ndarray] = {} - if self._metal_atom_indices and mol.GetNumConformers() > 0: - ref_conf = mol.GetConformer(mol.GetConformers()[0].GetId()) - for idx in self._metal_atom_indices: - self._metal_ref_pos[int(idx)] = np.array(ref_conf.GetAtomPosition(int(idx))) - self._metal_ligand_ref_distances: dict[tuple[int, int], float] = _metal_ligand_reference_distances( - mol, - self._metal_atom_indices, + metal_constraints = ConstraintModel.from_metal_shell(mol, self._metal_atom_indices) + pose_constraints = ( + ConstraintModel.from_atom_positions( + mol, + constraint_spec.constrained_atoms, + force_constant=constraint_spec.position_force_constant, + ) + if constraint_spec is not None + else ConstraintModel.empty() ) + self.constraint_model = metal_constraints.combine(pose_constraints) + self._has_position_constraints = bool(self.constraint_model.position_constraints) self.fast_minimizer = get_minimizer( config.minimizer, max_iters=config.fast_minimization_iters, dielectric=config.fast_dielectric, metal_atom_indices=self._metal_atom_indices, + constraint_model=self.constraint_model, ) self.full_minimizer = get_minimizer( config.minimizer, max_iters=config.max_minimization_iters, dielectric=config.final_dielectric, metal_atom_indices=self._metal_atom_indices, + constraint_model=self.constraint_model, ) self.fast_minimizer.prepare(mol) self.full_minimizer.prepare(mol) @@ -174,13 +177,6 @@ def __init__( if self._staging_mmff_props is not None: self._staging_mmff_props.SetMMFFDielectricConstant(config.fast_dielectric) - # snap constrained atoms back exactly after each minimization to eliminate restraint drift - self._constrained_ref_pos: dict[int, np.ndarray] = {} - if constraint_spec is not None and mol.GetNumConformers() > 0: - ref_conf = mol.GetConformer(mol.GetConformers()[0].GetId()) - for idx in constraint_spec.constrained_atoms: - self._constrained_ref_pos[idx] = np.array(ref_conf.GetAtomPosition(idx)) - if config.random_seed is not None: random.seed(config.random_seed) np.random.seed(config.random_seed) @@ -204,6 +200,10 @@ def _increment_stat(self, key: str, amount: int = 1) -> None: if self.stats is not None: self.stats[key] = int(self.stats.get(key, 0)) + amount + def _increment_move_stat(self, prefix: str, move_type: str) -> None: + """Increment move-specific counter when instrumentation is enabled.""" + self._increment_stat(f"{prefix}_{move_type}") + def _add_time_stat(self, key: str, elapsed_s: float) -> None: """Accumulate a timing stat when instrumentation is enabled.""" if self.stats is not None: @@ -319,8 +319,8 @@ def generate_seeds(self, n_seeds: int, prune_rms_thresh: float | None = None) -> return seed_results - def _reset_constrained_positions(self, mol: Chem.Mol, conf_id: int) -> None: - """Snap constrained atoms back to their exact reference coordinates. + def _reset_constraint_positions(self, mol: Chem.Mol, conf_id: int) -> None: + """Snap position-constrained atoms back to exact reference coordinates. Called after every constrained minimization to eliminate any residual drift that the position restraints did not fully suppress. @@ -329,48 +329,17 @@ def _reset_constrained_positions(self, mol: Chem.Mol, conf_id: int) -> None: mol: molecule containing conformer conf_id: conformer ID to update in place """ - if not self._constrained_ref_pos: - return - conf = mol.GetConformer(conf_id) - for idx, pos in self._constrained_ref_pos.items(): - conf.SetAtomPosition(idx, pos.tolist()) - - def _reset_metal_positions(self, mol: Chem.Mol, conf_id: int) -> None: - """Snap metal centers back to reference coordinates when available. - - Args: - mol: molecule containing conformer - conf_id: conformer ID to update - """ - if not self._metal_ref_pos: - return - conf = mol.GetConformer(conf_id) - for idx, pos in self._metal_ref_pos.items(): - conf.SetAtomPosition(idx, pos.tolist()) - - def _add_metal_uff_constraints(self, ff) -> None: - """Add metal position and reference-shell distance constraints to UFF force field.""" - for m_idx in self._metal_atom_indices: - ff.UFFAddPositionConstraint(int(m_idx), 0.0, _METAL_POSITION_FORCE_CONSTANT) - for (m_idx, nb_idx), distance in self._metal_ligand_ref_distances.items(): - ff.UFFAddDistanceConstraint( - int(m_idx), - int(nb_idx), - False, - max(0.0, distance - _METAL_LIGAND_DISTANCE_TOLERANCE), - distance + _METAL_LIGAND_DISTANCE_TOLERANCE, - _METAL_LIGAND_DISTANCE_FORCE_CONSTANT, - ) + self.constraint_model.reset_positions(mol, conf_id) def _minimize_uff_single(self, mol: Chem.Mol, conf_id: int, max_its: int) -> float: """UFF-minimize one conformer with metal position/distance constraints.""" ff = AllChem.UFFGetMoleculeForceField(mol, confId=int(conf_id)) if ff is None: return float("inf") - self._add_metal_uff_constraints(ff) + add_constraints_to_force_field(ff, self.constraint_model, "UFF") try: ff.Minimize(maxIts=max_its) - self._reset_metal_positions(mol, conf_id) + self._reset_constraint_positions(mol, conf_id) return float(ff.CalcEnergy()) except (ValueError, RuntimeError): return float("inf") @@ -439,7 +408,6 @@ def _minimize_constrained(self, mol: Chem.Mol, conf_id: int, use_fast: bool = Tr Returns: Energy in kcal/mol after minimization, or inf on failure """ - assert self.constraint_spec is not None minimizer = self.fast_minimizer if use_fast else self.full_minimizer props = getattr(minimizer, "_mmff_props", None) @@ -448,15 +416,14 @@ def _minimize_constrained(self, mol: Chem.Mol, conf_id: int, use_fast: bool = Tr ff = AllChem.MMFFGetMoleculeForceField(mol, props, confId=int(conf_id)) if ff is None: return float("inf") - for idx in self.constraint_spec.constrained_atoms: - ff.MMFFAddPositionConstraint(idx, 0.0, self.constraint_spec.position_force_constant) + add_constraints_to_force_field(ff, self.constraint_model, "MMFF") ff.Minimize(maxIts=int(minimizer.max_iters)) energy = float(ff.CalcEnergy()) else: ff = AllChem.UFFGetMoleculeForceField(mol, confId=int(conf_id)) if ff is None: return float("inf") - self._add_metal_uff_constraints(ff) + add_constraints_to_force_field(ff, self.constraint_model, "UFF") ff.Minimize(maxIts=int(minimizer.max_iters)) energy = float(ff.CalcEnergy()) except (ValueError, RuntimeError): @@ -464,8 +431,7 @@ def _minimize_constrained(self, mol: Chem.Mol, conf_id: int, use_fast: bool = Tr # Snap constrained atoms back to exact reference coordinates, eliminating # any residual drift that the position restraints did not fully suppress. - self._reset_constrained_positions(mol, conf_id) - self._reset_metal_positions(mol, conf_id) + self._reset_constraint_positions(mol, conf_id) return energy def _propose_constrained(self, pool: ConformerPool, step: int) -> tuple[int, float, str] | None: @@ -490,6 +456,78 @@ def _propose_constrained(self, pool: ConformerPool, step: int) -> tuple[int, flo return (new_conf_id, energy, f"hybrid_{move_type}") + def _available_tm_seed_move_types(self) -> list[str]: + """Return TM-specific move types available for seed augmentation.""" + availability = { + "tm_ligand_rotate": bool(self._moves.tm_ligand_rotations), + "tm_haptic_rotate": bool(self._moves.tm_haptic_rotations), + } + return [move_type for move_type in _TM_SEED_MOVE_TYPES if availability[move_type]] + + def _generate_move_seeds( + self, + source_conf_id: int, + attempts_per_type: int, + move_types: list[str], + attempt_stat_key: str, + ) -> list[tuple[int, float]]: + """Generate minimized seeds by copying source conformer and applying moves. + + Args: + source_conf_id: conformer ID to copy before applying each move + attempts_per_type: number of attempts per move type + move_types: move types to apply + attempt_stat_key: stat counter incremented for attempted seeds + + Returns: + Minimized seed conformer identifiers and energies + """ + if attempts_per_type <= 0 or not move_types: + return [] + + seeds: list[tuple[int, float]] = [] + for move_type in move_types: + operator = self._move_operators[move_type] + for _ in range(attempts_per_type): + self._increment_stat(attempt_stat_key) + self._increment_move_stat("seed_move_attempted", move_type) + new_conf_id = _copy_conformer(self.mol, source_conf_id) + operator(new_conf_id) + + if self._clash_checker.has_clash( + new_conf_id, + move_type, + skip_check=self.config.skip_clash_check, + ): + self.mol.RemoveConformer(new_conf_id) + continue + + self._increment_stat("n_minimization_calls") + energy = self.fast_minimizer.minimize(self.mol, new_conf_id) + if np.isfinite(energy): + self._increment_move_stat("seed_move_generated", move_type) + seeds.append((new_conf_id, energy)) + else: + self.mol.RemoveConformer(new_conf_id) + return seeds + + def generate_tm_move_seeds(self, source_conf_id: int, attempts_per_type: int) -> list[tuple[int, float]]: + """Generate input-derived TM move seeds using existing move operators. + + Args: + source_conf_id: conformer ID to copy before applying each move + attempts_per_type: number of attempts per available TM move type + + Returns: + Minimized seed conformer identifiers and energies + """ + return self._generate_move_seeds( + source_conf_id, + attempts_per_type, + self._available_tm_seed_move_types(), + "n_tm_move_seed_attempts", + ) + def _select_move_type(self, step: int) -> str: """Select move type based on probabilities and step count. @@ -502,18 +540,20 @@ def _select_move_type(self, step: int) -> str: forced = resolve_forced_move( step, self.config.shake_period, - constrained=self.constraint_spec is not None, + constrained=self._has_position_constraints, ) if forced is not None: return forced probs = resolve_move_probabilities( self._current_move_probs, - constrained=self.constraint_spec is not None, + constrained=self._has_position_constraints, has_ring_flips=bool(self.rotor_model.ring_flips), has_crankshaft=bool(self._moves.crankable_rings), has_kic=bool(self._moves.macro_kic_data), has_amide=bool(self._moves.amide_flips), + has_tm_moves=bool(self._moves.tm_ligand_rotations), + has_tm_haptic=bool(self._moves.tm_haptic_rotations), has_rotors=bool(self.rotor_model.rotors), ) @@ -608,6 +648,7 @@ def record_dedupe_outcome(self, surviving_ids: set[int]) -> None: self._move_attempts[move_type] = self._move_attempts.get(move_type, 0.0) + 1.0 if cid in surviving_ids: self._move_rewards[move_type] = self._move_rewards.get(move_type, 0.0) + 1.0 + self._increment_move_stat("move_survived_dedupe", move_type) self._pending_tags.clear() if self.config.adaptive_moves: @@ -672,6 +713,7 @@ def _generate_candidate(self, pool: ConformerPool, step: int) -> tuple[int, str] self._increment_stat("n_candidate_attempts") move_select_start = time.perf_counter() move_type = self._select_move_type(step) + self._increment_move_stat("move_selected", move_type) self._add_time_stat("move_selection_time_s", time.perf_counter() - move_select_start) if self._use_torsion_multitry(move_type): @@ -698,6 +740,7 @@ def _generate_candidate(self, pool: ConformerPool, step: int) -> tuple[int, str] self._add_time_stat("clash_check_time_s", time.perf_counter() - clash_start) self._increment_stat("n_candidates_passed_clash") + self._increment_move_stat("move_clash_passed", move_type) return (new_conf_id, move_type) def propose(self, pool: ConformerPool, step: int) -> tuple[int, float, str] | None: @@ -745,8 +788,8 @@ def propose_batch(self, pool: ConformerPool, step: int) -> list[tuple[int, float Returns: Accepted conformer IDs, energies, and sources """ - # Constrained mode: per-conformer MMFF with position restraints. - if self.constraint_spec is not None: + # Constraint mode: per-conformer minimization with explicit restraints. + if self.constraint_spec is not None or self.constraint_model.position_constraints: results: list[tuple[int, float, str]] = [] for i in range(self.config.minimize_batch_size): result = self._propose_constrained(pool, step + i) @@ -868,8 +911,7 @@ def full_refine_final_constrained( if ff is None: energies.append(float("inf")) continue - for idx in self.constraint_spec.constrained_atoms: - ff.MMFFAddPositionConstraint(idx, 0.0, self.constraint_spec.position_force_constant) + add_constraints_to_force_field(ff, self.constraint_model, "MMFF") ff.Minimize(maxIts=int(max_iters)) energies.append(float(ff.CalcEnergy())) else: @@ -877,14 +919,13 @@ def full_refine_final_constrained( if ff is None: energies.append(float("inf")) continue - self._add_metal_uff_constraints(ff) + add_constraints_to_force_field(ff, self.constraint_model, "UFF") ff.Minimize(maxIts=int(max_iters)) energies.append(float(ff.CalcEnergy())) except (ValueError, RuntimeError): energies.append(float("inf")) continue - self._reset_constrained_positions(mol, cid) - self._reset_metal_positions(mol, cid) + self._reset_constraint_positions(mol, cid) return energies @@ -909,6 +950,8 @@ def run_hybrid_generation( total_start = time.perf_counter() stats = new_generation_stats() if config.collect_stats else {} constraint_spec = config.constraint_spec + has_metal_input = any(_is_metal(atom) for atom in mol.GetAtoms()) and mol.GetNumConformers() > 0 + use_input_seed = constraint_spec is not None or has_metal_input # Filter rotors before building the proposer so _rotor_angles is computed # only for free rotors. @@ -925,9 +968,11 @@ def run_hybrid_generation( prune_rms_thresh=_resolve_seed_prune_rms_thresh(mol, rotor_model, effective_config), reason="seed_pose", ) - if constraint_spec is not None + if use_input_seed else resolve_seed_plan(mol, rotor_model, effective_config) ) + if use_input_seed and constraint_spec is None: + seed_plan = dataclasses.replace(seed_plan, reason="seed_input_metal") populate_effective_config_stats( stats, config=effective_config, @@ -943,15 +988,18 @@ def run_hybrid_generation( pool = ConformerPool(mol, effective_config) seed_start = time.perf_counter() - if constraint_spec is not None: + if use_input_seed: existing_ids = [c.GetId() for c in mol.GetConformers()] if not existing_ids: raise OpenConfValueError( - "Constrained conformer generation requires a starting conformer. " - "Use generate_conformers_from_pose to supply one." + "Input-seeded conformer generation requires a starting conformer. Supply a molecule with coordinates." ) seeds = proposer.seed_from_conformer(existing_ids[0]) - seed_source = "seed_pose" + if not seeds and has_metal_input and constraint_spec is None: + seeds = [(existing_ids[0], float("inf"))] + if stats: + stats["n_unminimized_input_seeds"] = 1 + seed_source = "seed_pose" if constraint_spec is not None else "seed_input" input_positions = None else: input_positions = mol.GetConformers()[0].GetPositions().copy() if mol.GetNumConformers() > 0 else None @@ -962,7 +1010,8 @@ def run_hybrid_generation( stats["n_seed_conformers"] = len(seeds) for conf_id, energy in seeds: - pool.insert(conf_id, energy, source=seed_source) + tags = {"protected": True} if use_input_seed else None + pool.insert(conf_id, energy, source=seed_source, tags=tags) if input_positions is not None: for conf_id, energy in proposer.seed_from_input_conformer(input_positions): @@ -970,38 +1019,51 @@ def run_hybrid_generation( if effective_config.use_low_mode_following and seeds: lm_ff_props = getattr(proposer.fast_minimizer, "_mmff_props", None) - if lm_ff_props is not None: - lm_start = time.perf_counter() - n_lm_accepted = 0 - source_seeds = sorted(seeds, key=lambda x: x[1])[: effective_config.low_mode_n_source_seeds] - # Materialize the low-mode children for every source seed before any - # pool insertion. The source seeds are themselves pool members, so - # inserting one seed's children could evict a later, not-yet-processed - # source seed and leave generate_low_mode_seeds referencing a freed - # conformer (Bad Conformer Id, or a hard abort during MMFF setup). - lm_children: list[tuple[int, float]] = [] - for seed_conf_id, _ in source_seeds: - lm_children.extend( - generate_low_mode_seeds( - mol, - lm_ff_props, - seed_conf_id, - proposer.fast_minimizer, - eigenvalue_threshold=effective_config.low_mode_eigenvalue_threshold, - max_modes=effective_config.low_mode_max_modes, - scan_step_size=effective_config.low_mode_scan_step_size, - scan_energy_threshold=effective_config.low_mode_scan_energy_threshold, - scan_max_steps=effective_config.low_mode_scan_max_steps, - ) + lm_start = time.perf_counter() + n_lm_accepted = 0 + source_seeds = sorted(seeds, key=lambda x: x[1])[: effective_config.low_mode_n_source_seeds] + # Materialize the low-mode children for every source seed before any + # pool insertion. The source seeds are themselves pool members, so + # inserting one seed's children could evict a later, not-yet-processed + # source seed and leave generate_low_mode_seeds referencing a freed + # conformer (Bad Conformer Id, or a hard abort during MMFF setup). + lm_children: list[tuple[int, float]] = [] + for seed_conf_id, _ in source_seeds: + lm_children.extend( + generate_low_mode_seeds( + mol, + lm_ff_props, + seed_conf_id, + proposer.fast_minimizer, + eigenvalue_threshold=effective_config.low_mode_eigenvalue_threshold, + max_modes=effective_config.low_mode_max_modes, + scan_step_size=effective_config.low_mode_scan_step_size, + scan_energy_threshold=effective_config.low_mode_scan_energy_threshold, + scan_max_steps=effective_config.low_mode_scan_max_steps, + constraints=proposer.constraint_model, ) - for lm_conf_id, lm_energy in lm_children: - if pool.insert(lm_conf_id, lm_energy, source="low_mode"): - n_lm_accepted += 1 - else: - mol.RemoveConformer(lm_conf_id) - if stats: - stats["low_mode_time_s"] = time.perf_counter() - lm_start - stats["n_low_mode_seeds"] = n_lm_accepted + ) + for lm_conf_id, lm_energy in lm_children: + if pool.insert(lm_conf_id, lm_energy, source="low_mode"): + n_lm_accepted += 1 + else: + mol.RemoveConformer(lm_conf_id) + if stats: + stats["low_mode_time_s"] = time.perf_counter() - lm_start + stats["n_low_mode_seeds"] = n_lm_accepted + + if has_metal_input and constraint_spec is None and effective_config.tm_seed_move_attempts > 0 and seeds: + tm_seed_start = time.perf_counter() + n_tm_move_accepted = 0 + tm_move_seeds = proposer.generate_tm_move_seeds(seeds[0][0], effective_config.tm_seed_move_attempts) + for tm_conf_id, tm_energy in tm_move_seeds: + if pool.insert(tm_conf_id, tm_energy, source="seed_tm_move"): + n_tm_move_accepted += 1 + else: + mol.RemoveConformer(tm_conf_id) + if stats: + stats["tm_move_seed_time_s"] = time.perf_counter() - tm_seed_start + stats["n_tm_move_seeds"] = n_tm_move_accepted batch_size = effective_config.minimize_batch_size step = 0 @@ -1030,6 +1092,7 @@ def run_hybrid_generation( if stats: stats["n_pool_accepts"] = int(stats["n_pool_accepts"]) + 1 move_type = source.removeprefix("hybrid_") if source.startswith("hybrid_") else source + proposer._increment_move_stat("move_pool_accepted", move_type) proposer.record_accepted(conf_id, move_type) else: if stats: diff --git a/openconf/propose/low_mode.py b/openconf/propose/low_mode.py index 6d8528c..5673c29 100644 --- a/openconf/propose/low_mode.py +++ b/openconf/propose/low_mode.py @@ -12,6 +12,8 @@ from rdkit import Chem from rdkit.Chem import AllChem +from ..constraints import ConstraintModel, add_constraints_to_force_field + if TYPE_CHECKING: from ..relax import Minimizer @@ -23,13 +25,46 @@ _DEFAULT_SCAN_MAX_STEPS: int = 10 +def _force_field( + mol: Chem.Mol, + ff_props: object | None, + conf_id: int, + constraints: ConstraintModel, +) -> object | None: + """Build constrained MMFF or UFF force field. + + Args: + mol: molecule containing conformer + ff_props: MMFF properties, or None to use UFF + conf_id: conformer ID + constraints: geometry constraints to apply + + Returns: + RDKit force field object, or None when setup fails + """ + try: + if ff_props is not None: + ff = AllChem.MMFFGetMoleculeForceField(mol, ff_props, confId=int(conf_id)) + family = "MMFF" + else: + ff = AllChem.UFFGetMoleculeForceField(mol, confId=int(conf_id)) + family = "UFF" + if ff is None: + return None + add_constraints_to_force_field(ff, constraints, family) + return ff + except (RuntimeError, ValueError): + return None + + def _compute_hessian( mol: Chem.Mol, - ff_props: object, + ff_props: object | None, conf_id: int, step: float = _DEFAULT_FD_STEP, + constraints: ConstraintModel | None = None, ) -> np.ndarray: - """Compute numerical Hessian via central differences of MMFF gradients. + """Compute numerical Hessian via central differences of force gradients. Perturbs each Cartesian coordinate of each atom by ±step, evaluates the MMFF gradient at each displaced geometry, and assembles the @@ -37,16 +72,18 @@ def _compute_hessian( Args: mol: molecule containing the conformer - ff_props: pre-prepared MMFFMoleculeProperties used to build force fields + ff_props: pre-prepared MMFFMoleculeProperties, or None to use UFF at displaced geometries conf_id: ID of the conformer at which to evaluate the Hessian; should be at or near a local minimum for meaningful low modes step: finite difference displacement in Å + constraints: geometry constraints to apply during gradient evaluation Returns: Symmetric 3N×3N Hessian matrix in kcal/mol/Ų """ conf = mol.GetConformer(conf_id) + constraints = constraints or ConstraintModel.empty() n_atoms = mol.GetNumAtoms() n_dof = 3 * n_atoms pos0 = conf.GetPositions().copy() @@ -60,7 +97,7 @@ def _compute_hessian( fwd = original.copy() fwd[coord_i] += step conf.SetAtomPosition(atom_i, fwd.tolist()) - ff_fwd = AllChem.MMFFGetMoleculeForceField(mol, ff_props, confId=conf_id) + ff_fwd = _force_field(mol, ff_props, conf_id, constraints) if ff_fwd is None: conf.SetAtomPosition(atom_i, original.tolist()) continue @@ -69,7 +106,7 @@ def _compute_hessian( bwd = original.copy() bwd[coord_i] -= step conf.SetAtomPosition(atom_i, bwd.tolist()) - ff_bwd = AllChem.MMFFGetMoleculeForceField(mol, ff_props, confId=conf_id) + ff_bwd = _force_field(mol, ff_props, conf_id, constraints) if ff_bwd is None: conf.SetAtomPosition(atom_i, original.tolist()) continue @@ -131,9 +168,9 @@ def _build_vibrational_basis(mol: Chem.Mol, conf_id: int) -> np.ndarray: for k in range(n_atoms): sqrt_mk = sqrt_masses[k] rx, ry, rz = r[k] - rot[0, 3 * k : 3 * k + 3] = sqrt_mk * np.array([0.0, rz, -ry]) # r x e_x - rot[1, 3 * k : 3 * k + 3] = sqrt_mk * np.array([-rz, 0.0, rx]) # r x e_y - rot[2, 3 * k : 3 * k + 3] = sqrt_mk * np.array([ry, -rx, 0.0]) # r x e_z + rot[0, 3 * k : 3 * k + 3] = sqrt_mk * np.array([0.0, rz, -ry]) # r x e_x + rot[1, 3 * k : 3 * k + 3] = sqrt_mk * np.array([-rz, 0.0, rx]) # r x e_y + rot[2, 3 * k : 3 * k + 3] = sqrt_mk * np.array([ry, -rx, 0.0]) # r x e_z # Threshold scales with sqrt(total_mass) since mass-weighted norms grow # as sqrt(M_total * r^2). Drops near-zero vectors for linear molecules. @@ -191,8 +228,8 @@ def _select_low_modes( # d_vib @ eigenvectors: mass-weighted eigenvectors L_mw # Unweight to Cartesian: L_cart[i] = L_mw[i] / sqrt(m_i) - vecs_mw = d_vib @ eigenvectors # (3N, n_vib) in mass-weighted space - vecs = vecs_mw / sqrt_masses[:, None] # (3N, n_vib) in Cartesian space + vecs_mw = d_vib @ eigenvectors # (3N, n_vib) in mass-weighted space + vecs = vecs_mw / sqrt_masses[:, None] # (3N, n_vib) in Cartesian space # Renormalise: mass-unweighting changes column norms col_norms = np.linalg.norm(vecs, axis=0, keepdims=True) @@ -207,12 +244,13 @@ def _select_low_modes( def _scan_along_mode( mol: Chem.Mol, - ff_props: object, + ff_props: object | None, start_conf_id: int, direction: np.ndarray, step_size: float, energy_threshold: float, max_steps: int, + constraints: ConstraintModel | None = None, ) -> np.ndarray: """Scan from a minimized conformer along a unit direction vector. @@ -226,7 +264,7 @@ def _scan_along_mode( Args: mol: molecule to scan; receives a temporary conformer during the call - ff_props: pre-prepared MMFFMoleculeProperties for energy evaluation + ff_props: pre-prepared MMFFMoleculeProperties, or None to use UFF start_conf_id: ID of the minimized starting conformer direction: unit displacement direction of shape (n_atoms, 3) in Å; each step moves the geometry by step_size × direction in 3N space @@ -234,15 +272,17 @@ def _scan_along_mode( energy_threshold: maximum allowed per-step energy increase (kcal/mol); scanning stops when ΔE in a single step exceeds this value max_steps: maximum number of steps regardless of energy criterion + constraints: geometry constraints to apply during energy evaluation Returns: Positions of shape (n_atoms, 3) at the last accepted scan point """ n_atoms = mol.GetNumAtoms() src_conf = mol.GetConformer(start_conf_id) + constraints = constraints or ConstraintModel.empty() start_pos = src_conf.GetPositions().copy() - ff0 = AllChem.MMFFGetMoleculeForceField(mol, ff_props, confId=start_conf_id) + ff0 = _force_field(mol, ff_props, start_conf_id, constraints) if ff0 is None: return start_pos prev_energy = float(ff0.CalcEnergy()) @@ -260,7 +300,7 @@ def _scan_along_mode( for i in range(n_atoms): working.SetAtomPosition(i, current_pos[i].tolist()) - ff = AllChem.MMFFGetMoleculeForceField(mol, ff_props, confId=working_id) + ff = _force_field(mol, ff_props, working_id, constraints) if ff is None: break @@ -278,7 +318,7 @@ def _scan_along_mode( def generate_low_mode_seeds( mol: Chem.Mol, - ff_props: object, + ff_props: object | None, conf_id: int, minimizer: "Minimizer", *, @@ -288,10 +328,11 @@ def generate_low_mode_seeds( scan_energy_threshold: float = _DEFAULT_SCAN_ENERGY_THRESHOLD, scan_max_steps: int = _DEFAULT_SCAN_MAX_STEPS, fd_step: float = _DEFAULT_FD_STEP, + constraints: ConstraintModel | None = None, ) -> list[tuple[int, float]]: """Generate conformers by scanning along low-frequency Hessian eigenvectors. - Numerically evaluates the MMFF Hessian at the given minimized conformer, + Numerically evaluates the MMFF or UFF Hessian at the given minimized conformer, identifies eigenvectors with eigenvalues below eigenvalue_threshold (soft conformational modes), and for each such mode scans in both the positive and negative directions using discrete steps of scan_step_size Å. Scanning along @@ -314,7 +355,7 @@ def generate_low_mode_seeds( Args: mol: molecule containing the conformer; receives new conformers in place - ff_props: pre-prepared MMFFMoleculeProperties used to build force fields + ff_props: pre-prepared MMFFMoleculeProperties, or None to use UFF conf_id: ID of a minimized conformer to compute modes from minimizer: minimizer applied to each scan endpoint eigenvalue_threshold: Hessian eigenvalue cutoff in mass-weighted @@ -331,6 +372,8 @@ def generate_low_mode_seeds( scan_max_steps: upper bound on scan steps per direction regardless of the energy criterion; acts as a safety cap on total displacement fd_step: finite difference step size for numerical Hessian in Å + constraints: geometry constraints to apply during Hessian and scan + energy evaluation Returns: Pairs of (conformer_id, energy_kcal_mol) for each successfully @@ -338,7 +381,8 @@ def generate_low_mode_seeds( senses); empty when no low modes satisfy the threshold or all minimizations fail """ - hessian = _compute_hessian(mol, ff_props, conf_id, fd_step) + constraints = constraints or ConstraintModel.empty() + hessian = _compute_hessian(mol, ff_props, conf_id, fd_step, constraints=constraints) low_vecs = _select_low_modes(hessian, mol, conf_id, eigenvalue_threshold, max_modes) if low_vecs.shape[1] == 0: return [] @@ -359,6 +403,7 @@ def generate_low_mode_seeds( scan_step_size, scan_energy_threshold, scan_max_steps, + constraints=constraints, ) if np.allclose(final_pos, start_pos, atol=1e-8): diff --git a/openconf/propose/moves.py b/openconf/propose/moves.py index c604a70..e23df55 100644 --- a/openconf/propose/moves.py +++ b/openconf/propose/moves.py @@ -26,6 +26,8 @@ type Config = ConformerConfig type MoveOperator = Callable[[int], None] type RingFlipArc = tuple[int, int, list[int]] +type TMLigandRotation = tuple[int, int, np.ndarray] +type TMHapticRotation = tuple[int, np.ndarray, np.ndarray] type TorsionAngleLibrary = TorsionLibrary @@ -111,6 +113,8 @@ def __init__( for ring_flip in rotor_model.ring_flips ] self.amide_flips = self._build_amide_flips(mol) + self.tm_ligand_rotations = self._build_tm_ligand_rotations(mol) + self.tm_haptic_rotations = self._build_tm_haptic_rotations(mol) self.operators: dict[str, MoveOperator] = { "single_rotor": self.apply_single_rotor_move, "multi_rotor": self.apply_multi_rotor_move, @@ -120,6 +124,8 @@ def __init__( "crankshaft": self.apply_crankshaft_move, "ring_kic": self.apply_ring_kic_move, "amide_flip": self.apply_amide_flip_move, + "tm_ligand_rotate": self.apply_tm_ligand_rotate_move, + "tm_haptic_rotate": self.apply_tm_haptic_rotate_move, } @staticmethod @@ -206,6 +212,106 @@ def _build_amide_flips(self, mol: Chem.Mol) -> list[tuple[int, int, np.ndarray]] amide_flips.append((c_idx, n_idx, moving_arr)) return amide_flips + def _build_tm_ligand_rotations(self, mol: Chem.Mol) -> list[TMLigandRotation]: + """Precompute ligand fragments rotatable around metal-donor axes.""" + rotations: list[TMLigandRotation] = [] + for atom in mol.GetAtoms(): + if not _is_metal(atom): + continue + metal_idx = atom.GetIdx() + for neighbor in atom.GetNeighbors(): + donor_idx = neighbor.GetIdx() + if _is_metal(neighbor): + continue + moving = self._tm_ligand_fragment(mol, metal_idx, donor_idx) + if len(moving) <= 1: + continue + moving_arr = np.fromiter(moving, dtype=np.int64, count=len(moving)) + rotations.append((metal_idx, donor_idx, moving_arr)) + return rotations + + @staticmethod + def _reachable_without(mol: Chem.Mol, start_atoms: set[int], blocked_atoms: set[int]) -> frozenset[int]: + """Return non-metal atoms reachable from starts without crossing blocked atoms.""" + visited: set[int] = set(blocked_atoms) + moving: set[int] = set() + stack = list(start_atoms) + while stack: + cur = stack.pop() + if cur in visited: + continue + visited.add(cur) + if _is_metal(mol.GetAtomWithIdx(cur)): + continue + moving.add(cur) + for neighbor in mol.GetAtomWithIdx(cur).GetNeighbors(): + idx = neighbor.GetIdx() + if idx in visited: + continue + if _is_metal(neighbor): + continue + stack.append(idx) + return frozenset(moving) + + @staticmethod + def _tm_ligand_fragment(mol: Chem.Mol, metal_idx: int, donor_idx: int) -> frozenset[int]: + """Return atoms reachable from donor without crossing metal centers.""" + return MoveExecutor._reachable_without(mol, {donor_idx}, {metal_idx}) + + @staticmethod + def _metal_edge_count(mol: Chem.Mol, atom_indices: frozenset[int]) -> int: + """Count metal-ligand edges touching atom_indices.""" + atom_set = set(atom_indices) + count = 0 + for idx in atom_indices: + atom = mol.GetAtomWithIdx(int(idx)) + for neighbor in atom.GetNeighbors(): + if _is_metal(neighbor) and neighbor.GetIdx() not in atom_set: + count += 1 + return count + + @staticmethod + def _connected_components(mol: Chem.Mol, atom_indices: set[int]) -> list[frozenset[int]]: + """Return graph components induced by atom_indices.""" + remaining = set(atom_indices) + components: list[frozenset[int]] = [] + while remaining: + start = remaining.pop() + component = {start} + stack = [start] + while stack: + cur = stack.pop() + for neighbor in mol.GetAtomWithIdx(cur).GetNeighbors(): + idx = neighbor.GetIdx() + if idx not in remaining: + continue + remaining.remove(idx) + component.add(idx) + stack.append(idx) + components.append(frozenset(component)) + return components + + def _build_tm_haptic_rotations(self, mol: Chem.Mol) -> list[TMHapticRotation]: + """Precompute eta-bound ring-like ligands for rotation about metal-centroid axes.""" + rotations: list[TMHapticRotation] = [] + for atom in mol.GetAtoms(): + if not _is_metal(atom): + continue + metal_idx = atom.GetIdx() + ligand_neighbors = {neighbor.GetIdx() for neighbor in atom.GetNeighbors() if not _is_metal(neighbor)} + for component in self._connected_components(mol, ligand_neighbors): + if len(component) < 3: + continue + moving = self._reachable_without(mol, set(component), {metal_idx}) + if len(moving) <= len(component): + continue + if self._metal_edge_count(mol, moving) != len(component): + continue + haptic_arr = np.fromiter(component, dtype=np.int64, count=len(component)) + moving_arr = np.fromiter(moving, dtype=np.int64, count=len(moving)) + rotations.append((metal_idx, haptic_arr, moving_arr)) + return rotations + def sample_angle(self, rotor_idx: int) -> float: """Sample a torsion angle for a rotor.""" angles_arr, weights_arr = self._rotor_angles[rotor_idx] @@ -340,6 +446,63 @@ def apply_amide_flip_move(self, conf_id: int) -> None: ) conf.SetPositions(all_pos) + def apply_tm_ligand_rotate_move(self, conf_id: int) -> None: + """Rotate one ligand fragment about a metal-donor axis.""" + if not self.tm_ligand_rotations: + return + + metal_idx, donor_idx, moving_arr = random.choice(self.tm_ligand_rotations) + conf = self.mol.GetConformer(conf_id) + all_pos = conf.GetPositions() + axis_origin = all_pos[metal_idx] + axis_vec = all_pos[donor_idx] - axis_origin + axis_norm = float(np.linalg.norm(axis_vec)) + if axis_norm < 1e-6: + return + axis_unit = axis_vec / axis_norm + + base = random.uniform(30.0, 90.0) * random.choice([-1.0, 1.0]) + theta = np.deg2rad(base + np.random.normal(0.0, self.config.torsion_jitter_deg)) + _rodrigues_rotate( + all_pos, + moving_arr, + axis_origin, + float(axis_unit[0]), + float(axis_unit[1]), + float(axis_unit[2]), + theta, + ) + conf.SetPositions(all_pos) + + def apply_tm_haptic_rotate_move(self, conf_id: int) -> None: + """Rotate eta-bound ligand around metal-ligand-centroid axis.""" + if not self.tm_haptic_rotations: + return + + metal_idx, haptic_arr, moving_arr = random.choice(self.tm_haptic_rotations) + conf = self.mol.GetConformer(conf_id) + all_pos = conf.GetPositions() + centroid = all_pos[haptic_arr].mean(axis=0) + axis_origin = all_pos[metal_idx] + axis_vec = centroid - axis_origin + axis_norm = float(np.linalg.norm(axis_vec)) + if axis_norm < 1e-6: + return + axis_unit = axis_vec / axis_norm + + base = random.uniform(30.0, 120.0) * random.choice([-1.0, 1.0]) + theta = np.deg2rad(base + np.random.normal(0.0, self.config.torsion_jitter_deg)) + _rodrigues_rotate( + all_pos, + moving_arr, + axis_origin, + float(axis_unit[0]), + float(axis_unit[1]), + float(axis_unit[2]), + theta, + ) + conf.SetPositions(all_pos) + def apply_ring_flip_move(self, conf_id: int) -> None: """Apply ring flip or stereo-preserving ring pucker move.""" if not self.rotor_model.ring_flips: diff --git a/openconf/propose/stats.py b/openconf/propose/stats.py index c8196c8..ed39e05 100644 --- a/openconf/propose/stats.py +++ b/openconf/propose/stats.py @@ -25,6 +25,7 @@ def new_generation_stats() -> dict[str, GenerationStat]: "effective_dedupe_period": 0, "effective_minimize_batch_size": 0, "n_seed_conformers": 0, + "n_unminimized_input_seeds": 0, "n_steps_executed": 0, "n_batches": 0, "n_candidate_attempts": 0, @@ -39,10 +40,13 @@ def new_generation_stats() -> dict[str, GenerationStat]: "n_dedupe_removed": 0, "n_final_selected": 0, "n_low_mode_seeds": 0, + "n_tm_move_seed_attempts": 0, + "n_tm_move_seeds": 0, "seed_time_s": 0.0, "seed_embedding_time_s": 0.0, "seed_minimization_time_s": 0.0, "low_mode_time_s": 0.0, + "tm_move_seed_time_s": 0.0, "proposal_stage_time_s": 0.0, "parent_selection_time_s": 0.0, "move_selection_time_s": 0.0, diff --git a/openconf/relax.py b/openconf/relax.py index f812d37..07233e6 100644 --- a/openconf/relax.py +++ b/openconf/relax.py @@ -6,41 +6,9 @@ from rdkit import Chem from rdkit.Chem import AllChem +from .constraints import ConstraintModel, add_constraints_to_force_field from .exceptions import OpenConfValueError -_METAL_POSITION_FORCE_CONSTANT = 1e4 -_METAL_LIGAND_DISTANCE_TOLERANCE = 0.05 -_METAL_LIGAND_DISTANCE_FORCE_CONSTANT = 100000.0 - - -def _metal_ligand_reference_distances( - mol: Chem.Mol, - metal_atom_indices: frozenset[int], -) -> dict[tuple[int, int], float]: - """Measure reference metal-ligand distances from first conformer. - - Args: - mol: molecule with reference conformer - metal_atom_indices: metal-center atom indices - - Returns: - Reference distances keyed by metal-neighbor atom index pairs - """ - if not metal_atom_indices or mol.GetNumConformers() == 0: - return {} - - conf = mol.GetConformer(mol.GetConformers()[0].GetId()) - distances: dict[tuple[int, int], float] = {} - for m_idx in metal_atom_indices: - metal_pos = conf.GetAtomPosition(int(m_idx)) - for nb in mol.GetAtomWithIdx(int(m_idx)).GetNeighbors(): - nb_idx = int(nb.GetIdx()) - nb_pos = conf.GetAtomPosition(nb_idx) - distance = float(metal_pos.Distance(nb_pos)) - if distance > 0.0: - distances[(int(m_idx), nb_idx)] = distance - return distances - class Minimizer(Protocol): """Protocol for conformer minimizers.""" @@ -76,10 +44,7 @@ class RDKitMMFFMinimizer: unavailable these are pinned via UFFAddPositionConstraint so that untyped metals (lanthanides, actinides, …) do not drift freely during UFF minimization. - _metal_ref_positions: reference coordinates for metal centers when - input molecule already has coordinates. - _metal_ligand_ref_distances: reference distances from metal centers to - directly connected ligating atoms when input molecule has coordinates. + constraint_model: geometry constraints to apply during minimization. """ max_iters: int = 500 @@ -88,9 +53,9 @@ class RDKitMMFFMinimizer: variant: str = "MMFF94s" dielectric: float = 4.0 metal_atom_indices: frozenset[int] = field(default_factory=frozenset) + constraint_model: ConstraintModel = field(default_factory=ConstraintModel.empty) _mmff_props: object = field(default=None, init=False, repr=False) - _metal_ref_positions: dict[int, tuple[float, float, float]] = field(default_factory=dict, init=False, repr=False) - _metal_ligand_ref_distances: dict[tuple[int, int], float] = field(default_factory=dict, init=False, repr=False) + _prepared_constraints: ConstraintModel = field(default_factory=ConstraintModel.empty, init=False, repr=False) def prepare(self, mol: Chem.Mol) -> None: """Cache MMFF properties for the molecule. @@ -103,40 +68,17 @@ def prepare(self, mol: Chem.Mol) -> None: self._mmff_props = AllChem.MMFFGetMoleculeProperties(mol, mmffVariant=self.variant) if self._mmff_props is not None: self._mmff_props.SetMMFFDielectricConstant(self.dielectric) - self._metal_ref_positions = {} - if self.metal_atom_indices and mol.GetNumConformers() > 0: - conf = mol.GetConformer(mol.GetConformers()[0].GetId()) - for idx in self.metal_atom_indices: - pos = conf.GetAtomPosition(int(idx)) - self._metal_ref_positions[int(idx)] = (float(pos.x), float(pos.y), float(pos.z)) - self._metal_ligand_ref_distances = _metal_ligand_reference_distances(mol, self.metal_atom_indices) - - def _add_metal_uff_constraints(self, ff) -> None: - """Add metal position and reference-shell distance constraints to UFF force field.""" - for m_idx in self.metal_atom_indices: - ff.UFFAddPositionConstraint(int(m_idx), 0.0, _METAL_POSITION_FORCE_CONSTANT) - for (m_idx, nb_idx), distance in self._metal_ligand_ref_distances.items(): - ff.UFFAddDistanceConstraint( - int(m_idx), - int(nb_idx), - False, - max(0.0, distance - _METAL_LIGAND_DISTANCE_TOLERANCE), - distance + _METAL_LIGAND_DISTANCE_TOLERANCE, - _METAL_LIGAND_DISTANCE_FORCE_CONSTANT, - ) - - def _reset_metal_positions(self, mol: Chem.Mol, conf_id: int) -> None: - """Snap metal centers back to reference coordinates when available. + metal_constraints = ConstraintModel.from_metal_shell(mol, self.metal_atom_indices) + self._prepared_constraints = metal_constraints.combine(self.constraint_model) + + def _reset_constrained_positions(self, mol: Chem.Mol, conf_id: int) -> None: + """Snap position-constrained atoms back to reference coordinates. Args: mol: molecule containing conformer conf_id: conformer ID to update """ - if not self._metal_ref_positions: - return - conf = mol.GetConformer(int(conf_id)) - for idx, pos in self._metal_ref_positions.items(): - conf.SetAtomPosition(idx, pos) + self._prepared_constraints.reset_positions(mol, conf_id) def minimize(self, mol: Chem.Mol, conf_id: int) -> float: """Minimize conformer in place and return energy in kcal/mol. @@ -151,14 +93,15 @@ def minimize(self, mol: Chem.Mol, conf_id: int) -> float: try: if self._mmff_props is not None: ff = AllChem.MMFFGetMoleculeForceField(mol, self._mmff_props, confId=int(conf_id)) + family = "MMFF" else: ff = AllChem.UFFGetMoleculeForceField(mol, confId=int(conf_id)) + family = "UFF" if ff is None: return float("inf") - if self._mmff_props is None and self.metal_atom_indices: - self._add_metal_uff_constraints(ff) + add_constraints_to_force_field(ff, self._prepared_constraints, family) ff.Minimize(maxIts=int(self.max_iters)) - self._reset_metal_positions(mol, conf_id) + self._reset_constrained_positions(mol, conf_id) return float(ff.CalcEnergy()) except (ValueError, RuntimeError): return float("inf") @@ -210,12 +153,18 @@ def minimize_confs_mmff( ] -def get_minimizer(name: str = "rdkit_mmff", metal_atom_indices: frozenset[int] = frozenset(), **kwargs) -> Minimizer: +def get_minimizer( + name: str = "rdkit_mmff", + metal_atom_indices: frozenset[int] = frozenset(), + constraint_model: ConstraintModel | None = None, + **kwargs, +) -> Minimizer: """Get a minimizer by name. Args: name: minimizer name; only `"rdkit_mmff"` is currently supported metal_atom_indices: metal atoms to exclude from MMFF typing + constraint_model: geometry constraints to apply during minimization **kwargs: additional arguments for minimizer Returns: @@ -225,5 +174,9 @@ def get_minimizer(name: str = "rdkit_mmff", metal_atom_indices: frozenset[int] = OpenConfValueError: unknown minimizer name """ if name == "rdkit_mmff": - return RDKitMMFFMinimizer(metal_atom_indices=metal_atom_indices, **kwargs) + return RDKitMMFFMinimizer( + metal_atom_indices=metal_atom_indices, + constraint_model=constraint_model or ConstraintModel.empty(), + **kwargs, + ) raise OpenConfValueError(f"Unknown minimizer: {name}") diff --git a/openconf/tuning.py b/openconf/tuning.py index 0b73a28..46f952f 100644 --- a/openconf/tuning.py +++ b/openconf/tuning.py @@ -209,6 +209,8 @@ def resolve_move_probabilities( has_crankshaft: bool, has_kic: bool = False, has_amide: bool = False, + has_tm_moves: bool = False, + has_tm_haptic: bool = False, has_rotors: bool = True, ) -> dict[str, float]: """Return move probabilities adjusted for current availability constraints.""" @@ -226,6 +228,8 @@ def resolve_move_probabilities( "crankshaft": has_crankshaft, "ring_kic": has_kic, "amide_flip": has_amide, + "tm_ligand_rotate": has_tm_moves, + "tm_haptic_rotate": has_tm_haptic, } for move_type, is_available in availability.items(): if is_available or move_type not in probs: diff --git a/tests/test_basic.py b/tests/test_basic.py index 3125257..6fdc7ad 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -187,6 +187,91 @@ def test_runtime_tuning_policy_helpers(): assert is_clash_exempt_move("single_rotor") is False +def test_runtime_tuning_keeps_tm_moves_without_rotors(): + """TM ligand moves can receive budget when no torsion rotors are available.""" + from openconf.tuning import resolve_move_probabilities + + resolved = resolve_move_probabilities( + {"single_rotor": 0.5, "tm_ligand_rotate": 0.5}, + constrained=True, + has_ring_flips=False, + has_crankshaft=False, + has_tm_moves=True, + has_rotors=False, + ) + + positive = {move_type: prob for move_type, prob in resolved.items() if prob > 0.0} + assert positive == {"tm_ligand_rotate": pytest.approx(1.0)} + + +def test_runtime_tuning_falls_back_from_unavailable_tm_moves(): + """Unavailable TM-specialized moves fall back to simpler metal-ligand moves.""" + from openconf.tuning import resolve_move_probabilities + + resolved = resolve_move_probabilities( + {"tm_haptic_rotate": 0.3, "tm_ligand_rotate": 0.7}, + constrained=True, + has_ring_flips=False, + has_crankshaft=False, + has_tm_moves=True, + has_tm_haptic=False, + ) + + positive = {move_type: prob for move_type, prob in resolved.items() if prob > 0.0} + assert positive == {"tm_ligand_rotate": pytest.approx(1.0)} + + +def test_transition_metal_preset_config_values(): + """Transition-metal preset enables TM-specific exploration knobs.""" + from openconf.config import preset_config + + config = preset_config("transition_metal") + + assert config.use_low_mode_following is True + assert config.parent_strategy == "uniform" + assert config.do_final_refine is False + assert config.minimize_batch_size == 1 + assert config.prism_max_deviation == pytest.approx(0.005) + assert config.tm_seed_move_attempts == 2 + assert config.final_select == "diverse" + assert config.move_probs["tm_ligand_rotate"] == pytest.approx(0.25) + assert config.move_probs["tm_haptic_rotate"] == pytest.approx(0.1) + + +def test_auto_transition_metal_sampling_adds_tm_budget_for_metal_inputs(): + """Metal inputs receive automatic TM move budget unless already configured.""" + from openconf import ConformerConfig + from openconf.api import _with_auto_transition_metal_sampling + + config = ConformerConfig(move_probs={"single_rotor": 1.0}, tm_seed_move_attempts=0) + resolved = _with_auto_transition_metal_sampling(config, has_metal=True, has_metal_input=True) + + assert resolved is not config + assert resolved.move_probs["single_rotor"] == pytest.approx(0.65) + assert resolved.move_probs["tm_ligand_rotate"] == pytest.approx(0.2625) + assert resolved.move_probs["tm_haptic_rotate"] == pytest.approx(0.0875) + assert sum(resolved.move_probs.values()) == pytest.approx(1.0) + assert resolved.tm_seed_move_attempts == 2 + assert config.move_probs == {"single_rotor": 1.0} + assert config.tm_seed_move_attempts == 0 + + +def test_auto_transition_metal_sampling_respects_opt_out_and_existing_tm_budget(): + """Automatic TM move overlay preserves explicit user choices.""" + from openconf import ConformerConfig + from openconf.api import _with_auto_transition_metal_sampling + + opt_out = ConformerConfig( + move_probs={"single_rotor": 1.0}, + auto_transition_metal_moves=False, + ) + assert _with_auto_transition_metal_sampling(opt_out, has_metal=True, has_metal_input=True) is opt_out + + existing_tm = ConformerConfig(move_probs={"single_rotor": 0.5, "tm_ligand_rotate": 0.5}) + assert _with_auto_transition_metal_sampling(existing_tm, has_metal=True, has_metal_input=True) is existing_tm + assert _with_auto_transition_metal_sampling(existing_tm, has_metal=False, has_metal_input=False) is existing_tm + + def test_generate_candidate_uses_operator_table_and_clash_checker(monkeypatch): """Candidate generation dispatches through the operator table and clash helper.""" from rdkit.Chem import AllChem @@ -415,6 +500,53 @@ def test_zero_energy_is_preserved(): assert pool.get_parent(strategy="best") == 0 +def test_pool_overflow_preserves_protected_conformer(): + """Pool eviction should not remove protected input conformers.""" + from openconf import ConformerConfig + from openconf.pool import ConformerPool + + mol = Chem.AddHs(Chem.MolFromSmiles("CC")) + for _ in range(3): + mol.AddConformer(Chem.Conformer(mol.GetNumAtoms()), assignId=True) + + pool = ConformerPool(mol=mol, config=ConformerConfig(max_out=1, pool_max=2, energy_window_kcal=20.0)) + assert pool.insert(0, 10.0, source="seed", tags={"protected": True}) + assert pool.insert(1, 1.0, source="seed") + assert pool.insert(2, 0.5, source="hybrid") + + assert 0 in pool.records + assert 1 not in pool.records + assert 2 in pool.records + + +def test_diverse_final_selection_preserves_protected_seed(): + """Diverse final selection should work with protected conformers.""" + from pathlib import Path + + from openconf import ConformerConfig + from openconf.io import read_xyz + from openconf.pool import ConformerPool + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + base_conf = mol.GetConformer(0) + for offset in (0.15, -0.15): + conf = Chem.Conformer(base_conf) + pos = conf.GetAtomPosition(0) + conf.SetAtomPosition(0, (pos.x + offset, pos.y, pos.z)) + mol.AddConformer(conf, assignId=True) + + config = ConformerConfig(max_out=2, pool_max=3, final_select="diverse") + pool = ConformerPool(mol=mol, config=config) + assert pool.insert(0, 10.0, source="seed", tags={"protected": True}) + assert pool.insert(1, 1.0, source="hybrid") + assert pool.insert(2, 2.0, source="hybrid") + + selected = pool.select_final() + + assert len(selected) == 2 + assert selected[0] == 0 + + @pytest.mark.parametrize( ("kwargs", "match"), [ @@ -459,7 +591,15 @@ def test_hybrid_collect_stats_populates_proposal_breakdown(): """Flexible molecules record nonzero proposal-stage timing components.""" from openconf import ConformerConfig, generate_conformers - config = ConformerConfig(max_out=5, n_seeds=4, n_steps=20, pool_max=30, random_seed=5, collect_stats=True) + config = ConformerConfig( + max_out=5, + n_seeds=4, + n_steps=20, + pool_max=30, + random_seed=5, + collect_stats=True, + move_probs={"single_rotor": 1.0}, + ) ensemble = generate_conformers("CCCCc1ccccc1", config=config) assert float(ensemble.generation_stats["proposal_stage_time_s"]) > 0.0 @@ -470,6 +610,9 @@ def test_hybrid_collect_stats_populates_proposal_breakdown(): assert float(ensemble.generation_stats["batch_staging_time_s"]) >= 0.0 assert float(ensemble.generation_stats["batch_commit_time_s"]) >= 0.0 assert int(ensemble.generation_stats["n_steps_executed"]) > 0 + assert int(ensemble.generation_stats["move_selected_single_rotor"]) > 0 + assert int(ensemble.generation_stats["move_clash_passed_single_rotor"]) > 0 + assert int(ensemble.generation_stats["move_pool_accepted_single_rotor"]) > 0 def test_large_flexible_defaults_are_topology_tuned(): @@ -1225,3 +1368,226 @@ def test_xyz_metals_remain_pinned_during_generation(xyz_name: str): for (i, j), distance in ref_distances.items() ) assert max_shell_delta < 0.35 + + +def test_tm_ligand_rotate_moves_ligand_and_keeps_metal_fixed(): + """TM ligand rotation changes donor-side atoms while leaving metal fixed.""" + from pathlib import Path + + import numpy as np + + from openconf.config import ConformerConfig + from openconf.io import read_xyz + from openconf.perceive import _is_metal, build_rotor_model + from openconf.propose.hybrid import HybridProposer, _copy_conformer + from openconf.torsionlib import TorsionLibrary + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + rotor_model = build_rotor_model(mol) + proposer = HybridProposer(mol, rotor_model, TorsionLibrary(), ConformerConfig(random_seed=2)) + assert proposer._moves.tm_ligand_rotations + + parent_id = mol.GetConformer(0).GetId() + new_id = _copy_conformer(mol, parent_id) + before = mol.GetConformer(new_id).GetPositions().copy() + + proposer._moves.apply_tm_ligand_rotate_move(new_id) + + after = mol.GetConformer(new_id).GetPositions() + metal_indices = [atom.GetIdx() for atom in mol.GetAtoms() if _is_metal(atom)] + nonmetal_indices = [atom.GetIdx() for atom in mol.GetAtoms() if not _is_metal(atom)] + + assert np.allclose(after[metal_indices], before[metal_indices], atol=1e-8) + assert max(float(np.linalg.norm(after[idx] - before[idx])) for idx in nonmetal_indices) > 0.05 + + +def test_tm_specific_move_sites_detected_from_xyz(): + """TM move chemotypes are detected conservatively from XYZ connectivity.""" + from pathlib import Path + + from openconf.config import ConformerConfig + from openconf.io import read_xyz + from openconf.perceive import build_rotor_model + from openconf.propose.hybrid import HybridProposer + from openconf.torsionlib import TorsionLibrary + + ru_mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + ru_moves = HybridProposer( + ru_mol, + build_rotor_model(ru_mol), + TorsionLibrary(), + ConformerConfig(random_seed=0), + )._moves + + assert ru_moves.tm_ligand_rotations + assert len(ru_moves.tm_haptic_rotations) == 1 + + dppf_mol = read_xyz(Path(__file__).resolve().parent / "data" / "dppf_ni_ph_aniline.xyz") + dppf_moves = HybridProposer( + dppf_mol, + build_rotor_model(dppf_mol), + TorsionLibrary(), + ConformerConfig(random_seed=0), + )._moves + + assert not dppf_moves.tm_haptic_rotations + assert dppf_moves.tm_ligand_rotations + + +def test_tm_haptic_rotate_moves_eta_ligand(): + """Haptic ligand rotation moves eta-bound ligand while keeping metal fixed.""" + from pathlib import Path + + import numpy as np + + from openconf.config import ConformerConfig + from openconf.io import read_xyz + from openconf.perceive import _is_metal, build_rotor_model + from openconf.propose.hybrid import HybridProposer, _copy_conformer + from openconf.torsionlib import TorsionLibrary + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + proposer = HybridProposer(mol, build_rotor_model(mol), TorsionLibrary(), ConformerConfig(random_seed=4)) + assert proposer._moves.tm_haptic_rotations + + parent_id = mol.GetConformer(0).GetId() + new_id = _copy_conformer(mol, parent_id) + before = mol.GetConformer(new_id).GetPositions().copy() + metal_indices = [atom.GetIdx() for atom in mol.GetAtoms() if _is_metal(atom)] + _metal_idx, _haptic_arr, moving_arr = proposer._moves.tm_haptic_rotations[0] + + proposer._moves.apply_tm_haptic_rotate_move(new_id) + + after = mol.GetConformer(new_id).GetPositions() + assert np.allclose(after[metal_indices], before[metal_indices], atol=1e-8) + assert max(float(np.linalg.norm(after[int(idx)] - before[int(idx)])) for idx in moving_arr) > 0.05 + + +def test_xyz_metal_input_seed_is_retained(): + """Metal XYZ generation should preserve input conformer basin.""" + from pathlib import Path + + from rdkit import Chem + from rdkit.Chem import rdMolAlign + + from openconf import ConformerConfig, generate_conformers + from openconf.io import read_xyz + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + ref_mol = Chem.Mol(mol) + heavy_map = [(atom.GetIdx(), atom.GetIdx()) for atom in mol.GetAtoms() if atom.GetAtomicNum() > 1] + + config = ConformerConfig( + max_out=3, + n_steps=0, + n_seeds=1, + random_seed=1, + num_threads=1, + skip_clash_check=True, + seed_minimization_iters=20, + collect_stats=True, + ) + ensemble = generate_conformers(mol, config=config, add_hs=False) + + best_rmsd = min( + rdMolAlign.AlignMol(Chem.Mol(ensemble.mol), ref_mol, prbCid=conf_id, refCid=0, atomMap=heavy_map) + for conf_id in ensemble.conf_ids + ) + assert best_rmsd < 0.8 + assert ensemble.generation_stats["seed_plan_reason"] == "seed_input_metal" + + +def test_xyz_metal_unminimized_input_seed_fallback(monkeypatch): + """Metal XYZ input seed should survive force-field setup failure.""" + from pathlib import Path + + from openconf import ConformerConfig, generate_conformers + from openconf.io import read_xyz + from openconf.propose import hybrid + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + + def fail_seed_from_conformer(self: hybrid.HybridProposer, conf_id: int) -> list[tuple[int, float]]: + return [] + + monkeypatch.setattr(hybrid.HybridProposer, "seed_from_conformer", fail_seed_from_conformer) + + config = ConformerConfig( + max_out=1, + n_steps=0, + n_seeds=1, + random_seed=1, + num_threads=1, + do_final_refine=False, + collect_stats=True, + ) + ensemble = generate_conformers(mol, config=config, add_hs=False) + + assert ensemble.n_conformers == 1 + assert ensemble.energies == [float("inf")] + assert ensemble.generation_stats["seed_plan_reason"] == "seed_input_metal" + assert int(ensemble.generation_stats["n_unminimized_input_seeds"]) == 1 + + +def test_xyz_metal_low_mode_runs_with_uff_backend(): + """Low-mode seeding should run for metal XYZ inputs without MMFF typing.""" + from pathlib import Path + + from rdkit.Chem import AllChem + + from openconf import ConformerConfig, generate_conformers + from openconf.io import read_xyz + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + assert AllChem.MMFFGetMoleculeProperties(mol, mmffVariant="MMFF94s") is None + + config = ConformerConfig( + max_out=4, + n_steps=0, + n_seeds=1, + random_seed=1, + num_threads=1, + skip_clash_check=True, + seed_minimization_iters=5, + fast_minimization_iters=5, + use_low_mode_following=True, + low_mode_max_modes=1, + low_mode_scan_max_steps=1, + low_mode_n_source_seeds=1, + collect_stats=True, + ) + ensemble = generate_conformers(mol, config=config, add_hs=False) + + assert ensemble.n_conformers > 0 + assert "low_mode_time_s" in ensemble.generation_stats + + +def test_xyz_metal_tm_move_seeds_are_generated(): + """TM move seeding should augment input-seeded metal generation.""" + from pathlib import Path + + from openconf import ConformerConfig, generate_conformers + from openconf.io import read_xyz + + mol = read_xyz(Path(__file__).resolve().parent / "data" / "ru_pcymene_aqua.xyz") + config = ConformerConfig( + max_out=8, + n_steps=0, + n_seeds=1, + random_seed=2, + num_threads=1, + skip_clash_check=True, + seed_minimization_iters=5, + fast_minimization_iters=5, + use_low_mode_following=False, + tm_seed_move_attempts=1, + energy_window_kcal=100.0, + collect_stats=True, + ) + + ensemble = generate_conformers(mol, config=config, add_hs=False) + + assert ensemble.generation_stats["seed_plan_reason"] == "seed_input_metal" + assert int(ensemble.generation_stats["n_tm_move_seed_attempts"]) > 0 + assert int(ensemble.generation_stats["n_tm_move_seeds"]) > 0 + assert ensemble.n_conformers > 1 diff --git a/tests/test_constrained.py b/tests/test_constrained.py index acfb118..f9b53ce 100644 --- a/tests/test_constrained.py +++ b/tests/test_constrained.py @@ -239,6 +239,59 @@ def test_filter_constrained_rotors_free_side_reoriented(): ) +def test_filtered_constrained_torsions_do_not_translate_core(): + """Filtered torsion moves leave constrained atom coordinates unchanged.""" + from rdkit.Chem import rdMolTransforms + + from openconf import build_rotor_model, filter_constrained_rotors, prepare_molecule + + mol = prepare_molecule(Chem.MolFromSmiles("CCCCc1ccccc1")) + AllChem.EmbedMolecule(mol, randomSeed=0) + constrained = frozenset(range(4, 10)) # benzene ring atoms + filtered = filter_constrained_rotors(build_rotor_model(mol), constrained) + + assert filtered.rotors + constrained_indices = sorted(constrained) + for rotor in filtered.rotors: + trial = Chem.Mol(mol) + conf = trial.GetConformer(0) + before = conf.GetPositions().copy() + angle = rdMolTransforms.GetDihedralDeg(conf, *rotor.dihedral_atoms) + + rdMolTransforms.SetDihedralDeg(conf, *rotor.dihedral_atoms, angle + 73.0) + + after = conf.GetPositions() + assert np.allclose(after[constrained_indices], before[constrained_indices], atol=1e-8), ( + f"Rotor {rotor.atom_idxs} translated constrained atoms" + ) + + +def test_boundary_constrained_torsion_moves_only_free_fragment(): + """Boundary torsion with constrained axis atoms moves free distal fragment.""" + from rdkit.Chem import rdMolTransforms + + from openconf import build_rotor_model, filter_constrained_rotors, prepare_molecule + + mol = prepare_molecule(Chem.MolFromSmiles("CCCCc1ccccc1")) + AllChem.EmbedMolecule(mol, randomSeed=0) + constrained = frozenset([3, 4, 5, 6, 7, 8, 9]) + filtered = filter_constrained_rotors(build_rotor_model(mol), constrained) + boundary_rotor = next((rotor for rotor in filtered.rotors if set(rotor.atom_idxs) == {3, 4}), None) + + assert boundary_rotor is not None + conf = mol.GetConformer(0) + before = conf.GetPositions().copy() + angle = rdMolTransforms.GetDihedralDeg(conf, *boundary_rotor.dihedral_atoms) + + rdMolTransforms.SetDihedralDeg(conf, *boundary_rotor.dihedral_atoms, angle + 90.0) + + after = conf.GetPositions() + constrained_indices = sorted(constrained) + free_chain_indices = [0, 1, 2] + assert np.allclose(after[constrained_indices], before[constrained_indices], atol=1e-8) + assert max(float(np.linalg.norm(after[idx] - before[idx])) for idx in free_chain_indices) > 0.05 + + def test_filter_constrained_rotors_no_constrained_atoms(): """Empty constraint set leaves the rotor model unchanged.""" from openconf import build_rotor_model, filter_constrained_rotors, prepare_molecule