class Oracle:
"""Run Validate once for one ligand and return a finite scalar score."""
_AGGREGATORS: ClassVar[frozenset[str]] = frozenset(
{"first", "mean", "max", "min", "median"}
)
_COMPARISONS: ClassVar[Mapping[str, Callable[[float, float], bool]]] = {
"gt": operator.gt,
"ge": operator.ge,
"lt": operator.lt,
"le": operator.le,
}
def __init__(
self,
wrk_dir: str,
system_path: str,
options_path: str,
runner: str = "boltz2",
input_smiles: str | None = None,
input_mol_file: str | None = None,
output_metric: str | None = None,
aggregate: str = "first",
repeats: int = 1,
seed: int | None = None,
scoring_functions: list[str] | None = None,
assess_robustness: bool = True,
assess_bias: bool = False,
protein_training_data_path: str | None = None,
ligand_training_data_path: str | None = None,
bias_training_data_protein_path: str | None = None,
bias_training_data_ligand_path: str | None = None,
bias_query_cache_path: str | None = None,
custom_bias_reference_path: str | None = None,
custom_protein_reference_path: str | None = None,
custom_ligand_reference_path: str | None = None,
bias_release_cutoff: str = "2023-06-01",
bias_protein_similarity_threshold: float = 0.25,
bias_ligand_similarity_threshold: float = 0.35,
bias_chains: list[str] | None = None,
build_bias_training_data: bool = False,
bias_training_components_cif: str | None = None,
conformers: str | None = None,
sdf_file: str | None = None,
reference_path: str | None = None,
pocket_coverage_reference: str | None = None,
reproduction_metrics: list[str] | None = None,
score_components: Mapping[str, float] | None = None,
scoring_function: Callable[[OracleScoreContext], float] | None = None,
score_gates: Sequence[OracleGate] | None = None,
gate_policy: OracleGatePolicy | None = None,
ligand_chain: str | None = None,
):
self.wrk_dir = Path(wrk_dir)
self.system_path = Path(system_path)
self.options_path = Path(options_path)
self.runner = str(runner)
self.run_id = str(uuid4())
self.input_smiles = input_smiles.strip() if input_smiles else None
self.input_mol_file = Path(input_mol_file) if input_mol_file else None
self.output_metric = str(output_metric).strip() if output_metric else None
self.aggregate = str(aggregate).strip().lower()
self.score_components = self._normalize_components(score_components)
self.scoring_function = scoring_function
self.score_gates = tuple(score_gates or ())
self.gate_policy = gate_policy
self.ligand_chain = str(ligand_chain).strip() if ligand_chain else None
effective_scoring_functions = (
scoring_functions
if scoring_functions is not None
else sorted(DEFAULT_SCORING_FUNCTIONS)
)
self.validate_kwargs: dict[str, Any] = {
"repeats": repeats,
"seed": seed,
"scoring_functions": effective_scoring_functions,
"assess_robustness": assess_robustness,
"assess_bias": assess_bias,
"protein_training_data_path": protein_training_data_path,
"ligand_training_data_path": ligand_training_data_path,
"bias_training_data_protein_path": bias_training_data_protein_path,
"bias_training_data_ligand_path": bias_training_data_ligand_path,
"bias_query_cache_path": bias_query_cache_path,
"custom_bias_reference_path": custom_bias_reference_path,
"custom_protein_reference_path": custom_protein_reference_path,
"custom_ligand_reference_path": custom_ligand_reference_path,
"bias_release_cutoff": bias_release_cutoff,
"bias_protein_similarity_threshold": bias_protein_similarity_threshold,
"bias_ligand_similarity_threshold": bias_ligand_similarity_threshold,
"bias_chains": bias_chains,
"build_bias_training_data": build_bias_training_data,
"bias_training_components_cif": bias_training_components_cif,
"conformers": conformers,
"sdf_file": sdf_file,
"reference_path": reference_path,
"pocket_coverage_reference": pocket_coverage_reference,
"reproduction_metrics": reproduction_metrics,
}
self.base_system: dict[str, Any] = {}
self.query_ligand_chain = self.ligand_chain
self.logger = logging.getLogger("cofolder.oracle")
def preflight(self):
"""Validate and describe the oracle request without executing it."""
from dataclasses import replace
from cofolder.recipes.preflight import PreflightReport, prediction_preflight
try:
self._validate_config()
except Exception as exc:
return PreflightReport(
workflow="oracle",
ready=False,
runner=self.runner,
output_dir=str(self.wrk_dir / "results"),
messages=(str(exc),),
)
report, system_obj, _options = prediction_preflight(
workflow="oracle",
system_path=self.system_path,
options_path=self.options_path,
runner_name=self.runner,
repeats=int(self.validate_kwargs["repeats"]),
output_dir=self.wrk_dir,
require_ligand=True,
assess_bias=bool(self.validate_kwargs["assess_bias"]),
use_bias_databases=bool(
self.validate_kwargs["bias_training_data_protein_path"]
or self.validate_kwargs["bias_training_data_ligand_path"]
or (
self.validate_kwargs["protein_training_data_path"] is None
and self.validate_kwargs["ligand_training_data_path"] is None
and self.validate_kwargs["custom_bias_reference_path"] is None
and self.validate_kwargs["custom_protein_reference_path"] is None
and self.validate_kwargs["custom_ligand_reference_path"] is None
and not self.validate_kwargs["build_bias_training_data"]
)
),
protein_database_path=self.validate_kwargs["bias_training_data_protein_path"],
ligand_database_path=self.validate_kwargs["bias_training_data_ligand_path"],
release_cutoff=self.validate_kwargs["bias_release_cutoff"],
bias_chains=self.validate_kwargs["bias_chains"],
bias_query_cache_path=self.validate_kwargs["bias_query_cache_path"],
custom_bias_reference_path=self.validate_kwargs["custom_bias_reference_path"],
protein_similarity_threshold=self.validate_kwargs[
"bias_protein_similarity_threshold"
],
ligand_similarity_threshold=self.validate_kwargs[
"bias_ligand_similarity_threshold"
],
)
if not report.ready:
return report
try:
target = resolve_ligand_target(system_obj, self.ligand_chain)
except Exception as exc:
return replace(report, ready=False, messages=report.messages + (str(exc),))
return replace(
report,
messages=report.messages + (f"selected_ligand_chain={target.chain_ids[0]}",),
)
@staticmethod
def _normalize_components(
components: Mapping[str, float] | None,
) -> dict[str, float]:
normalized: dict[str, float] = {}
for metric, raw_weight in (components or {}).items():
name = str(metric).strip()
if not name:
raise ValueError("score_components metric names must not be empty.")
try:
weight = float(raw_weight)
except (TypeError, ValueError) as exc:
raise ValueError(
f"Weight for score component '{name}' must be a finite number."
) from exc
if not math.isfinite(weight):
raise ValueError(
f"Weight for score component '{name}' must be a finite number."
)
normalized[name] = weight
return normalized
def _validate_config(self) -> None:
if bool(self.input_smiles) == bool(self.input_mol_file):
raise ValueError(
"Provide exactly one input: --input_smiles or --input_mol_file."
)
if self.input_mol_file is not None:
if not self.input_mol_file.exists():
raise ValueError(
f"--input_mol_file does not exist: {self.input_mol_file}"
)
if not self.input_mol_file.is_file():
raise ValueError(
f"--input_mol_file is not a file: {self.input_mol_file}"
)
sources = sum(
(
bool(self.output_metric),
bool(self.score_components),
self.scoring_function is not None,
)
)
if sources != 1:
raise ValueError(
"Provide exactly one score source: output_metric, score_components, or scoring_function."
)
if self.scoring_function is not None and not callable(self.scoring_function):
raise ValueError("scoring_function must be callable.")
if self.aggregate not in self._AGGREGATORS:
raise ValueError(
f"Invalid --aggregate '{self.aggregate}'. Choose from: {sorted(self._AGGREGATORS)}"
)
if any(not isinstance(gate, OracleGate) for gate in self.score_gates):
raise ValueError("score_gates must contain OracleGate instances.")
if self.score_gates and not isinstance(self.gate_policy, OracleGatePolicy):
raise ValueError("gate_policy is required when score_gates are configured.")
if not self.score_gates and self.gate_policy is not None:
raise ValueError("gate_policy requires at least one score gate.")
selectors = [
*self.score_components,
*(gate.metric for gate in self.score_gates),
]
if self.output_metric:
selectors.append(self.output_metric)
for selector in selectors:
self._validate_metric_dependency(selector)
def _validate_metric_dependency(self, selector: str) -> None:
metric = str(selector).rsplit("__", 1)[-1]
scoring = set(self.validate_kwargs.get("scoring_functions") or [])
if metric == "ifp_distance":
if "ifp_distance" not in scoring:
raise ValueError(
f"Metric '{selector}' requires one of scoring functions "
"['ifp_distance']."
)
raise ValueError(
"Metric 'ifp_distance' is vector-valued and cannot be an Oracle scalar; "
"use pocket_coverage_custom for an IFP-reference overlap score."
)
if metric.startswith("bias_"):
if not self.validate_kwargs.get("assess_bias"):
raise ValueError(f"Metric '{selector}' requires assess_bias=True.")
return
required_group: set[str] | None = None
if metric == "sasa":
required_group = {"sasa"}
elif metric == "sasa_norm_heavy":
required_group = {"sasa_normalized"}
elif metric in {"affinity_pred_value", "affinity_probability_binary"}:
required_group = {"affinity_metrics"}
elif metric in {"pIC50", "IC50_M", "pIC50_kcal_per_mol"}:
required_group = {"affinity_metrics_ext"}
elif metric.startswith("pair_chains_iptm_") or metric in {
"chains_ptm",
"ptm",
"iptm",
"confidence_score",
}:
required_group = {"confidence_metrics"}
elif metric in {"pocket_coverage_custom", "pocket_coverage_custom_mean"}:
required_group = {"ifp_distance"}
if not self.validate_kwargs.get("pocket_coverage_reference"):
raise ValueError(
f"Metric '{selector}' requires pocket_coverage_reference."
)
reproduction = self.validate_kwargs.get("reproduction_metrics")
if reproduction is not None and "pocket_coverage" not in reproduction:
raise ValueError(
f"Metric '{selector}' requires the pocket_coverage reproduction metric."
)
elif metric in {"pocket_coverage_ref", "pocket_coverage_ref_mean"}:
if not self.validate_kwargs.get("reference_path"):
raise ValueError(f"Metric '{selector}' requires reference_path.")
reproduction = self.validate_kwargs.get("reproduction_metrics")
if reproduction is not None and "pocket_coverage" not in reproduction:
raise ValueError(
f"Metric '{selector}' requires the pocket_coverage reproduction metric."
)
elif metric in {
"ligand_rmsd_ref",
"protein_rmsd_ref",
"sucos_shape_ref",
"sucos_feature_ref",
"sucos_ref",
"ligand_pose_overlap_ref",
}:
if not self.validate_kwargs.get("reference_path"):
raise ValueError(f"Metric '{selector}' requires reference_path.")
reproduction_metric = (
"ligand_rmsd"
if metric == "ligand_rmsd_ref"
else "protein_rmsd"
if metric == "protein_rmsd_ref"
else "sucos"
)
reproduction = self.validate_kwargs.get("reproduction_metrics")
if reproduction is not None and reproduction_metric not in reproduction:
raise ValueError(
f"Metric '{selector}' requires the {reproduction_metric} "
"reproduction metric."
)
if required_group is not None and not (required_group & scoring):
raise ValueError(
f"Metric '{selector}' requires one of scoring functions {sorted(required_group)}, "
f"but enabled are {sorted(scoring)}."
)
@report_completion(WorkflowKind.ORACLE)
def run(self) -> float:
try:
return self._run_impl()
except Exception as exc:
raw_candidate = self.input_smiles or (
str(self.input_mol_file)
if self.input_mol_file is not None
else "unresolved"
)
identity = self._public_identity(raw_candidate)
source_exc = exc.__cause__ or exc
stage = (
exc.failures[0].stage
if isinstance(exc, WorkflowExecutionError) and exc.failures
else FailureStage.INPUT_VALIDATION
if isinstance(source_exc, InputValidationError)
else FailureStage.ANALYTICS
)
output_dir = self.wrk_dir / "results"
failure = write_failure_bundle(
output_dir=output_dir,
exc=source_exc,
identity=identity,
stage=stage,
error_code=getattr(source_exc, "error_code", "oracle_score_failed"),
)
raise WorkflowExecutionError(
str(exc), failures=(failure,), output_dir=output_dir
) from exc
def _run_impl(self) -> float:
self._validate_config()
self.wrk_dir.mkdir(parents=True, exist_ok=True)
run_dir = self.wrk_dir / "oracle_run"
run_dir.mkdir(parents=True, exist_ok=True)
smiles = self.input_smiles or self._smiles_from_mol_file(self.input_mol_file)
if not smiles:
raise ValueError("Failed to resolve valid input SMILES for oracle run.")
document = load_yaml_document(self.system_path)
if not isinstance(document.value, dict):
raise SystemInputValidationError(
"System YAML root must be a mapping.", source_path=self.system_path
)
base_system_obj = system.System(system=document.value)
resolve_declared_msa_paths(base_system_obj, base_dir=self.system_path.parent)
runner = get_runner(self.runner)
options_obj = runner.load_options(self.options_path)
validated = runner.validate_system(
base_system_obj,
options_obj,
check_atom_names=False,
source_path=self.system_path,
requirements=WorkflowInputRequirements(
require_protein=True, require_ligand=True
),
)
self.base_system = validated.system.system
target = resolve_ligand_target(validated.system, self.ligand_chain)
self.ligand_chain = target.chain_ids[0]
self.query_ligand_chain = self.ligand_chain
normalized_ligand = validate_smiles(
smiles,
source=LigandSourceIdentity(
entity_id=target.entity_id,
chain_ids=target.chain_ids,
source_record_id="oracle_input",
),
source_path=self.input_mol_file,
)
sys_obj = replace_ligand_smiles(
validated.system, target=target, ligand=normalized_ligand
)
runner.ensure_available()
row_system_path = run_dir / "oracle_system.yaml"
write.write_yaml(sys_obj, path=row_system_path)
Validate(
wrk_dir=str(run_dir),
system_path=str(row_system_path),
options_path=str(self.options_path),
runner=self.runner,
**self.validate_kwargs,
).run()
system_df, chain_df = read_metric_frames(run_dir)
raw_metrics = collect_qualified_metric_values(system_df, chain_df)
context = OracleScoreContext(
query_smiles=smiles,
run_dir=run_dir,
system_metrics=system_df.copy(),
chain_metrics=chain_df.copy(),
aggregated_metrics=MappingProxyType(
self._aggregate_all_metrics(raw_metrics)
),
)
base_value, _component_values, _score_mode = self._calculate_base_score(
context, raw_metrics
)
final_value, _gate_values, _failed_gates, action = self._apply_gates(
base_value, raw_metrics
)
identity = self._public_identity(smiles)
records = [
SuccessRecord(
envelope=make_envelope(RecordKind.SUCCESS, identity, "success"),
),
self._oracle_metric_record(identity, "oracle_raw_score", base_value),
self._oracle_metric_record(identity, "oracle_score", final_value),
]
if action != "none":
records.append(
self._oracle_metric_record(
identity, "oracle_gate_adjusted_score", final_value
)
)
write_public_bundle(
PublicOutputBundle(
manifest=PublicManifest(
schema_version=PUBLIC_SCHEMA_VERSION,
identity=identity,
status="success",
evidence=tuple(self._public_evidence()),
requested_metrics=tuple(
sorted(
{
*self.score_components,
*(gate.metric for gate in self.score_gates),
*([self.output_metric] if self.output_metric else []),
}
)
),
),
records=tuple(records),
),
self.wrk_dir / "results",
)
return final_value
def _public_identity(self, candidate: str) -> OutputIdentity:
normalized = self._normalize_candidate(candidate)
digest = hashlib.sha256(normalized.encode("utf-8")).hexdigest()
return OutputIdentity(
workflow=WorkflowKind.ORACLE,
run_id=self.run_id,
system_id=self.system_path.stem,
compound_id=f"sha256:{digest}",
runner_id=self.runner,
)
@staticmethod
def _normalize_candidate(candidate: str) -> str:
value = str(candidate).strip()
try:
from rdkit import Chem
molecule = Chem.MolFromSmiles(value)
if molecule is not None:
return str(Chem.MolToSmiles(molecule, canonical=True))
except (ImportError, RuntimeError, ValueError):
return value
return value
def _public_evidence(self) -> list[EvidenceSource]:
evidence: list[EvidenceSource] = []
reference_path = self.validate_kwargs.get("reference_path")
if reference_path:
path = Path(reference_path)
evidence.append(
EvidenceSource(
kind="reference_structure",
identifier=path.name,
path=str(path),
)
)
if self.validate_kwargs.get("pocket_coverage_reference"):
evidence.append(
EvidenceSource(
kind="custom_pocket",
identifier="configured_custom_pocket",
)
)
return evidence
@staticmethod
def _oracle_metric_record(
identity: OutputIdentity,
metric_name: str,
value: float,
) -> MetricRecord:
definition = get_metric_definition(metric_name)
return MetricRecord(
envelope=make_envelope(RecordKind.METRIC, identity, metric_name, "value"),
metric_name=metric_name,
metric_group=definition.group,
metric_class=definition.metric_class,
status=RecordStatus.COMPUTED,
value=float(value),
unit=definition.unit,
direction=definition.direction,
evidence_regime=EvidenceRegime.REFERENCE_FREE,
)
def _calculate_base_score(self, context, raw_metrics):
if self.output_metric:
value = self._metric_value(self.output_metric, raw_metrics)
return value, {self.output_metric: value}, "metric"
if self.score_components:
values = {
name: self._metric_value(name, raw_metrics)
for name in self.score_components
}
score = sum(
values[name] * weight for name, weight in self.score_components.items()
)
return self._finite_score(score, "Composite score"), values, "composite"
assert self.scoring_function is not None
return (
self._finite_score(
self.scoring_function(context), "Custom scoring function"
),
{},
"custom",
)
def _apply_gates(self, base_value, raw_metrics):
gate_values: dict[str, float | None] = {}
failed: list[str] = []
for gate in self.score_gates:
try:
value = self._metric_value(gate.metric, raw_metrics)
except _MetricUnavailableError:
value = None
gate_values[gate.metric] = value
if value is None or not self._COMPARISONS[gate.operator](
value, gate.threshold
):
failed.append(gate.metric)
if not failed:
return base_value, gate_values, failed, "none"
assert self.gate_policy is not None
final = (
base_value * self.gate_policy.value
if self.gate_policy.mode == "downweight"
else self.gate_policy.value
)
return (
self._finite_score(final, "Gate-adjusted score"),
gate_values,
failed,
self.gate_policy.mode,
)
def _aggregate_all_metrics(self, raw_metrics):
aggregated: dict[str, float] = {}
for selector, values in raw_metrics.items():
try:
aggregated[selector] = self._aggregate_values(selector, values)
except ValueError:
continue
return aggregated
def _metric_value(self, selector, raw_metrics):
qualified = self._resolve_selector(selector, raw_metrics)
return self._aggregate_values(qualified, raw_metrics[qualified])
def _resolve_selector(self, selector, raw_metrics):
name = str(selector).strip()
if "__" in name:
if name not in raw_metrics:
raise _MetricUnavailableError(
f"Metric '{name}' was not found in Oracle outputs."
)
return name
if self.query_ligand_chain:
ligand_name = f"ligand_{self.query_ligand_chain}__{name}"
if ligand_name in raw_metrics:
return ligand_name
ligand_matches = sorted(
metric
for metric in raw_metrics
if metric.startswith("ligand_") and metric.endswith(f"__{name}")
)
if len(ligand_matches) == 1:
return ligand_matches[0]
system_name = f"system__{name}"
if system_name in raw_metrics:
return system_name
matches = sorted(
metric for metric in raw_metrics if metric.endswith(f"__{name}")
)
if len(matches) == 1:
return matches[0]
if len(matches) > 1:
raise ValueError(
f"Metric '{name}' is ambiguous; use a qualified selector such as '{matches[0]}'."
)
raise _MetricUnavailableError(
f"Metric '{name}' was not found in Oracle outputs."
)
def _aggregate_values(self, selector, values):
numeric = pd.to_numeric(pd.Series(list(values)), errors="coerce").dropna()
finite = [float(value) for value in numeric if math.isfinite(float(value))]
if not finite:
raise _MetricUnavailableError(
f"Metric '{selector}' was not found with finite numeric values in Oracle outputs."
)
if self.aggregate == "first":
return finite[0]
if self.aggregate == "mean":
return float(sum(finite) / len(finite))
if self.aggregate == "max":
return max(finite)
if self.aggregate == "min":
return min(finite)
if self.aggregate == "median":
ordered = sorted(finite)
middle = len(ordered) // 2
return (
ordered[middle]
if len(ordered) % 2
else float((ordered[middle - 1] + ordered[middle]) / 2)
)
raise ValueError(f"Unsupported aggregate: {self.aggregate}")
@staticmethod
def _finite_score(value, source):
if isinstance(value, bool):
raise TypeError(f"{source} must return a finite numeric scalar, not bool.")
try:
score = float(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"{source} must return a finite numeric scalar.") from exc
if not math.isfinite(score):
raise ValueError(f"{source} must return a finite numeric scalar.")
return score
@staticmethod
def _find_query_ligand_chain(system_dict):
path = Oracle._find_first_ligand_smiles_path(system_dict)
if path is None:
return None
ligand = system_dict["sequences"][path[1]]["ligand"]
identifier = ligand.get("id") if isinstance(ligand, dict) else None
if isinstance(identifier, list):
return str(identifier[0]) if len(identifier) == 1 else None
return str(identifier) if identifier is not None else None
@staticmethod
def _find_first_ligand_smiles_path(system_dict):
sequences = system_dict.get("sequences", [])
if not isinstance(sequences, list):
return None
for index, entry in enumerate(sequences):
ligand = entry.get("ligand") if isinstance(entry, dict) else None
if isinstance(ligand, dict) and "smiles" in ligand:
return ["sequences", index, "ligand", "smiles"]
return None
@staticmethod
def _smiles_from_mol_file(path):
if path is None:
return None
try:
from rdkit import Chem
except Exception as exc:
raise ValueError("RDKit is required to parse --input_mol_file.") from exc
suffix = path.suffix.lower()
mol = None
if suffix == ".sdf":
mol = next(
(
item
for item in Chem.SDMolSupplier(str(path), removeHs=False)
if item is not None
),
None,
)
elif suffix in {".mol", ".mol2"}:
mol = Chem.MolFromMolFile(str(path), removeHs=False)
else:
mol = next(
(
item
for item in Chem.SDMolSupplier(str(path), removeHs=False)
if item is not None
),
None,
)
if mol is None:
mol = Chem.MolFromMolFile(str(path), removeHs=False)
return None if mol is None else Chem.MolToSmiles(mol)