class Bias:
"""High-level orchestrator for standalone bias workflow."""
def __init__(
self,
wrk_dir: str,
system_path: str,
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,
):
self.wrk_dir = Path(wrk_dir)
self.system_path = Path(system_path)
self.run_id = str(uuid4())
self.protein_training_data_path = (
Path(protein_training_data_path) if protein_training_data_path else None
)
self.ligand_training_data_path = (
Path(ligand_training_data_path) if ligand_training_data_path else None
)
self.bias_training_data_protein_path = (
Path(bias_training_data_protein_path)
if bias_training_data_protein_path
else None
)
self.bias_training_data_ligand_path = (
Path(bias_training_data_ligand_path)
if bias_training_data_ligand_path
else None
)
self.bias_query_cache_path = (
Path(bias_query_cache_path) if bias_query_cache_path else None
)
self.custom_bias_reference_path = (
Path(custom_bias_reference_path) if custom_bias_reference_path else None
)
self.custom_protein_reference_path = (
Path(custom_protein_reference_path)
if custom_protein_reference_path
else None
)
self.custom_ligand_reference_path = (
Path(custom_ligand_reference_path) if custom_ligand_reference_path else None
)
self.bias_release_cutoff = str(bias_release_cutoff)
self.bias_protein_similarity_threshold = float(
bias_protein_similarity_threshold
)
self.bias_ligand_similarity_threshold = float(bias_ligand_similarity_threshold)
self.bias_chains = normalize_bias_chains(bias_chains)
self.build_bias_training_data = bool(build_bias_training_data)
self.bias_training_components_cif = (
Path(bias_training_components_cif) if bias_training_components_cif else None
)
self.logger = logging.getLogger(__name__)
self._system = read.read_yaml(path=self.system_path)
self.sys = system.System(system=self._system)
self.timings = DebugTimingCollector(logger=self.logger)
def _debug_timer(self, label: str):
return self.timings.measure(label, logger=self.logger)
def _log_timing_summary(self) -> None:
self.timings.log_summary(logger=self.logger)
def preflight(self):
"""Validate the standalone bias request without searches or writes."""
from cofolder.modules.analytics.bias_database import (
bias_query_cache_readiness,
DEFAULT_LIGAND_DATABASE,
DEFAULT_PROTEIN_DATABASE,
_selected_queries,
parse_bias_release_policy,
validate_custom_bias_reference_bundle,
validate_bias_database_bundle,
)
from cofolder.modules.input.system import iter_system_chains
from cofolder.recipes.preflight import PreflightReport
try:
self._validate_source_mode()
policy = parse_bias_release_policy(self.bias_release_cutoff)
for name, value in (
("protein", self.bias_protein_similarity_threshold),
("ligand", self.bias_ligand_similarity_threshold),
):
if not 0.0 <= value <= 1.0:
raise ValueError(
f"Bias {name} similarity threshold must be between 0 and 1."
)
sources: list[str] = []
protein_bundle = None
ligand_bundle = None
database_mode = not (
self.protein_training_data_path
or self.ligand_training_data_path
or self.custom_protein_reference_path
or self.custom_ligand_reference_path
or self.custom_bias_reference_path
or self.build_bias_training_data
) or bool(
self.bias_training_data_protein_path
or self.bias_training_data_ligand_path
)
entity_types = {
chain.entity_type
for chain in iter_system_chains(self.sys)
if not self.bias_chains
or chain.chain_id.strip().upper() in self.bias_chains
}
if database_mode and "protein" in entity_types:
protein_bundle = validate_bias_database_bundle(
self.bias_training_data_protein_path or DEFAULT_PROTEIN_DATABASE,
"protein",
)
sources.append(
f"protein:{protein_bundle.root}:{protein_bundle.fingerprint}"
)
if database_mode and "ligand" in entity_types:
ligand_bundle = validate_bias_database_bundle(
self.bias_training_data_ligand_path or DEFAULT_LIGAND_DATABASE,
"ligand",
)
sources.append(
f"ligand:{ligand_bundle.root}:{ligand_bundle.fingerprint}"
)
if self.custom_bias_reference_path:
custom = validate_custom_bias_reference_bundle(self.custom_bias_reference_path)
sources.append(f"custom:{custom.root}:{custom.fingerprint}")
cache_root, cache_ready, cache_state = bias_query_cache_readiness(
self.bias_query_cache_path
)
ligand_smiles_by_id: dict[str, str] = {}
if ligand_bundle is not None:
source = pd.read_csv(ligand_bundle.data_path)
if {"ligand_id", "smiles"} <= set(source.columns):
ligand_smiles_by_id = {
str(row["ligand_id"]).upper(): str(row["smiles"])
for _, row in source.dropna(
subset=["ligand_id", "smiles"]
).iterrows()
}
proteins, ligands = _selected_queries(
self.sys, self.bias_chains, ligand_smiles_by_id
)
import hashlib
from rdkit import Chem
query_hashes = [
f"protein:{chain}:{hashlib.sha256(''.join(sequence.split()).upper().encode()).hexdigest()}"
for chain, sequence in sorted(proteins.items())
]
for chain, smiles in sorted(ligands.items()):
molecule = Chem.MolFromSmiles(smiles)
canonical = (
Chem.MolToSmiles(molecule, isomericSmiles=True)
if molecule
else smiles
)
query_hashes.append(
f"ligand:{chain}:{hashlib.sha256(canonical.encode()).hexdigest()}"
)
return PreflightReport(
workflow="bias",
ready=True,
selected_chains=tuple(
sorted(chain.chain_id for chain in iter_system_chains(self.sys))
),
database_sources=tuple(sources),
release_cutoff=self.bias_release_cutoff,
release_policy=policy.mode,
query_hashes=tuple(query_hashes),
bias_query_cache=f"{cache_root} ({cache_state})",
bias_cache_ready=cache_ready,
expected_bias_backfill=bool(
protein_bundle is not None
and protein_bundle.sequence_index_path is not None
and proteins
and ligands
),
protein_similarity_threshold=self.bias_protein_similarity_threshold,
ligand_similarity_threshold=self.bias_ligand_similarity_threshold,
output_dir=str(self.wrk_dir / "results"),
)
except Exception as exc:
return PreflightReport(
workflow="bias",
ready=False,
release_cutoff=self.bias_release_cutoff,
output_dir=str(self.wrk_dir / "results"),
messages=(str(exc),),
)
def _validate_source_mode(self) -> None:
if self.custom_bias_reference_path and (
self.custom_protein_reference_path or self.custom_ligand_reference_path
):
raise ValueError(
"--custom_bias_reference_path cannot be combined with paired custom reference paths."
)
if (
self.bias_training_data_protein_path
or self.bias_training_data_ligand_path
) and (
self.protein_training_data_path
or self.ligand_training_data_path
or self.build_bias_training_data
):
raise ValueError(
"Database-backed bias paths cannot be combined with legacy public-reference paths/build mode."
)
@report_completion(WorkflowKind.BIAS)
def run(self) -> tuple[pd.DataFrame, pd.DataFrame]:
try:
return self._run_impl()
except WorkflowExecutionError:
raise
except Exception as exc:
identity = OutputIdentity(
workflow=WorkflowKind.BIAS,
run_id=self.run_id,
system_id=self.system_path.stem,
)
output_dir = self.wrk_dir / "results"
failure = write_failure_bundle(
output_dir=output_dir,
exc=exc,
identity=identity,
stage=(
FailureStage.INPUT_VALIDATION
if isinstance(exc, InputValidationError)
or "database" in str(exc).lower()
else FailureStage.ANALYTICS
),
error_code="bias_analytics_failed",
)
raise WorkflowExecutionError(
str(exc), failures=(failure,), output_dir=output_dir
) from exc
def _run_impl(self) -> tuple[pd.DataFrame, pd.DataFrame]:
with self._debug_timer("bias.total"):
self._validate_source_mode()
self.wrk_dir.mkdir(parents=True, exist_ok=True)
output_dir = self.wrk_dir / "results" / "bias_train"
output_dir.mkdir(parents=True, exist_ok=True)
for legacy_name in ("system_metrics.csv", "chain_metrics.csv"):
(output_dir / legacy_name).unlink(missing_ok=True)
system_df, chain_df = build_bias_dataframes(
self.sys,
system_name=self.system_path.stem,
)
if chain_df.empty:
raise ValueError(
"The system definition must include at least one protein or ligand "
"component with an explicit chain ID for bias assessment."
)
workflow = BiasAssessmentWorkflow(
wrk_dir=self.wrk_dir,
system_path=self.system_path,
sys_obj=self.sys,
protein_training_data_path=self.protein_training_data_path,
ligand_training_data_path=self.ligand_training_data_path,
bias_training_data_protein_path=self.bias_training_data_protein_path,
bias_training_data_ligand_path=self.bias_training_data_ligand_path,
bias_query_cache_path=self.bias_query_cache_path,
custom_bias_reference_path=self.custom_bias_reference_path,
custom_protein_reference_path=self.custom_protein_reference_path,
custom_ligand_reference_path=self.custom_ligand_reference_path,
bias_release_cutoff=self.bias_release_cutoff,
bias_protein_similarity_threshold=self.bias_protein_similarity_threshold,
bias_ligand_similarity_threshold=self.bias_ligand_similarity_threshold,
bias_chains=self.bias_chains,
build_bias_training_data=self.build_bias_training_data,
bias_training_components_cif=self.bias_training_components_cif,
logger=self.logger,
timings=self.timings,
strict_training_data=True,
)
system_df, chain_df = workflow.apply(
system_df=system_df,
chain_df=chain_df,
boltz_cache_path="~/.boltz",
)
evidence = [
EvidenceSource(
kind="training_dataset", identifier=path.name, path=str(path)
)
for path in (
self.protein_training_data_path,
self.ligand_training_data_path,
self.bias_training_data_protein_path,
self.bias_training_data_ligand_path,
self.custom_bias_reference_path,
self.custom_protein_reference_path,
self.custom_ligand_reference_path,
)
if path is not None
]
artifacts = (
ArtifactReference("bias_supporting_outputs", "bias_train", "directory"),
)
write_frame_bundle(
system_df,
chain_df,
output_dir=self.wrk_dir / "results",
identity=OutputIdentity(
workflow=WorkflowKind.BIAS,
run_id=self.run_id,
system_id=self.system_path.stem,
),
evidence=evidence,
requested_metrics={"bias_metrics"},
artifacts=artifacts,
)
self._log_timing_summary()
return system_df, chain_df