Metrics#

Metric framework#

Metrics for model evaluation.

Re-exports the base metric framework so from rfproteina.metrics import MetricManager keeps working; the concrete metrics live in submodules.

Bases: ABC

Abstract base class for Foundry metrics.

Defines a framework for computing metrics based on arbitrary keyword arguments. A Metric can specify which batches a MetricManager applies it to, based on tags in the input batch.

Parameters:

required_tags_all – A set of tags that must all be present in the input batch for this metric to be computed.
required_tags_any – A set of tags where at least one must be present in the input batch for this metric to be computed.
prohibited_tags – A set of tags that must not be present in the input batch for this metric to be computed.
aggregation – Reduction(s) applied to this metric’s per-example outputs when summarizing across samples. One of, or a list of, SUPPORTED_REDUCTIONS. Defaults to ["mean"].

To implement a new metric, subclass this class and implement the compute method, at a minimum.

SUPPORTED_REDUCTIONS = ('mean', 'median')#

aggregate(column: str, values: Series) → dict[str, float]#

Aggregate per-example values for one of this metric’s output columns.

Parameters:

column – Bare output-key name (the metric-name prefix already stripped), or f"{save_id}.{output_key}" for a metric an oracle’s pre_oracle_transform fanned out.
values – Per-example values for that column (NaNs already dropped).

Returns:

Mapping of reduction name to value, one entry per configured reduction.

applies_to_tags(tags: set[str] | None, *, explicit: bool = False) → bool#

Return True if this metric would run for a batch with the given tag set.

Mirrors the tag-filtering logic in MetricManager.__call__().

Parameters:

tags – Tags for a single dataset. None means the dataset defines no tag filtering, so this metric always runs for it.
explicit – Ignore dataset-selection tags; retain oracle and refold routing tags.

abstractmethod compute(**kwargs: Any) → dict[str, Any] | list[dict[str, Any]]#

Implement actual metric calculation here

Override e.g. as:

def compute(self, y_true, y_pred, extra_info):
    print(extra_info)
    return lddt(y_true, y_pred, thres=self.custom_thresholds)

compute_from_kwargs(**kwargs: Any) → dict[str, Any]#

Run compute with an arbitrary dictionary of input keys and values.

The ‘kwargs_to_compute_args’ property here will determine where in the kwargs we will look for the values to pass to the compute method.

Parameters marked in ‘optional_kwargs’ will only be passed if present in kwargs.

differentiable = False#

property kwargs_to_compute_args: dict[str, Any]#

Map input keys to a flat dictionary.

If not implemented, we return None, and pass the kwargs directly to the compute method.

Override e.g. as:

@cached_property
def kwargs_to_compute_args(self) -> dict[str, Any]:
    return {
        "y_true": ("network_input", "coords_unnoised"),
        "y_pred": ("network_output", "coords_pred"),
        "extra_info": ("extra_info",),
    }

property optional_kwargs: frozenset[str]#

Set of compute argument names that are optional.

Optional parameters will only be passed to compute() if present in kwargs. The compute() method should have sensible defaults for these parameters.

Override e.g. as:

@property
def optional_kwargs(self) -> frozenset[str]:
    return frozenset(["confidence_indices", "interfaces_to_score"])

property required_compute_args: frozenset[str]#: Required input keys for this metric

property required_kwargs: frozenset[str]#: Required input keys for this metric

setup() → None#: Once-per-epoch Metric setup. Override to load expensive resources.

teardown() → None#: Called at validation epoch end. Override to clean up resources.

exception rfproteina.metrics.MetricInputError#

Bases: Exception

Exception raised when a metric fails to compute.

class rfproteina.metrics.MetricManager(metrics: dict[str, Metric] = {}, *, raise_errors: bool = True)#

Bases: object

Manages and computes a set of Metrics, where each Metric inherits from the Metric class.

For model validation, additional metrics can be added through the Hydra configuration; they will be computed with the __call__ method automatically.

For example, during AF-3, Metrics will receive network_input, network_output, extra_info, ground_truth_atom_array_stack, and predicted_atom_array_stack as input arguments.

Example

>>> class ExampleMetric(Metric):
...     @cached_property
...     def kwargs_to_compute_args(self):
...         return {"x": "x", "y": "y", "extra_info": "extra_info"}
...
...     def compute(self, x, y, extra_info):
...         return {"value": x + y}
>>> metric = ExampleMetric()
>>> manager = MetricManager({"my_metric": metric}, raise_errors=True)
>>> manager(x=1, y=2, extra_info={"example_id": "123"})
{'example_id': '123', 'my_metric.value': 3}

active_metrics(tags: set[str] | None = None, metric_names: list[str] | None = None) → dict[str, Metric]#

Select by explicit names or dataset tags, preserving refold routing constraints.

Explicit names may span several managers; each returns its own matching metrics. An empty list selects none. None uses the normal dataset-tag selection.

classmethod from_metrics(metrics: dict[str, Metric] | list[tuple[str, Metric]], raise_errors: bool = True) → MetricManager#

Create MetricManager from metric objects.

Parameters:

metrics – Either dict mapping names to Metric objects, or list of (name, Metric) tuples.
raise_errors – Whether to raise errors on metric failures. Defaults to True.

classmethod instantiate_from_hydra(metrics_cfg: dict[str, Any] | DictConfig) → MetricManager#

Instantiate a MetricManager from a dictionary of metrics.

Parameters:: metrics_cfg – A dictionary where keys are metric names and values are Hydra configurations for the metrics.

on_validation_epoch_end() → None#: Teardown all metrics after validation.

on_validation_epoch_start(val_dataset_tags: list[set[str] | None] | None = None, metric_names: set[str] | None = None) → None#

Setup metrics for validation.

Parameters:

val_dataset_tags – Per-dataset tag sets for all validation datasets. Each element is the metrics_tags for one dataset, or None if that dataset defines no tag filtering (meaning all metrics run for it). When the outer list itself is None, all metrics are set up unconditionally (backward-compatible default).
metric_names – Additional metrics to set up by name, regardless of dataset tags.

RASA metrics#

rfproteina.metrics.rasa_metrics.get_rasa_metrics(atom_array: AtomArray, probe_radius: float = 1.4, atom_radii: str = 'element', point_number: int = 100) → dict[str, float][source]#

Compute RASA adherence metrics for a given input atom array.

Parameters:

atom_array (AtomArray) – Input AtomArray with predicted structure and rasa annotations.
probe_radius (float) – Van-der-Waals radius of the probe in Angstrom. Defaults to 1.4 (for water).
atom_radii (str) – Atom radii set to use for calculation. Defaults to “element”, matching the radii used for RASA conditioning during training.
point_number (int) – Number of points in the Shrake-Rupley algorithm. Defaults to 100.

Returns:

Dictionary containing RASA adherence metrics

Return type:

dict

Note

RASA is computed with the Shrake-Rupley algorithm. Training uses buried values in [0.0, 0.02), intermediate values in [0.02, 0.1], and exposed values greater than 0.1, after rounding to the SASA sampling grid.

Bases: Metric

Compute adherence to RASA conditioning.

property kwargs_to_compute_args: dict[str, Any][source]#

Map input keys to a flat dictionary.

If not implemented, we return None, and pass the kwargs directly to the compute method.

Override e.g. as:

@cached_property
def kwargs_to_compute_args(self) -> dict[str, Any]:
    return {
        "y_true": ("network_input", "coords_unnoised"),
        "y_pred": ("network_output", "coords_pred"),
        "extra_info": ("extra_info",),
    }

compute(atom_array_or_stack: AtomArray | AtomArrayStack | list[AtomArray]) → dict[str, float][source]#

Implement actual metric calculation here

Override e.g. as:

def compute(self, y_true, y_pred, extra_info):
    print(extra_info)
    return lddt(y_true, y_pred, thres=self.custom_thresholds)

Design metrics#

Bond length, steric clash, and sequence-similarity metrics for generated protein structures.

rfproteina.metrics.design_metrics.compute_bond_length_stats(atom_array: AtomArray, verbose: bool = False) → dict[str, Any][source]#

Compute bond length reasonableness statistics.

Uses atomworks element-dependent distance thresholds to flag bonds that are too long to be physically reasonable.

Parameters:

atom_array – Structure with explicit bonds and coordinates.
verbose – If True, report breakdown by bond category (within/between residue, within/between chain, polymer/non-polymer).

Returns:

Dict with bond length statistics. Empty if no bonds are present.

rfproteina.metrics.design_metrics.compute_clash_stats(atom_array: AtomArray, threshold: float = 1.5) → dict[str, Any][source]#

Fraction of atoms in a steric clash, by clash type.

A clash is a non-bonded pair of atoms closer than threshold A; pairs where both atoms are CRD-conditioned (fixed input) are excluded. See _clash_category_masks() for the 5-category breakdown.

class rfproteina.metrics.design_metrics.BondLengthMetric(verbose: bool = False, **kwargs: Any)[source]#

Bases: Metric

Report unreasonable bond lengths using atomworks element-dependent thresholds.

In verbose mode, unreasonable bonds are broken down by spatial category: within-residue, between-residue within-chain, between-chain polymer-polymer, and the special non-polymer-to-polymer bin (often unreliable in predictions).

property kwargs_to_compute_args: dict[str, Any][source]#

Map input keys to a flat dictionary.

If not implemented, we return None, and pass the kwargs directly to the compute method.

Override e.g. as:

@cached_property
def kwargs_to_compute_args(self) -> dict[str, Any]:
    return {
        "y_true": ("network_input", "coords_unnoised"),
        "y_pred": ("network_output", "coords_pred"),
        "extra_info": ("extra_info",),
    }

compute(atom_array: AtomArray) → dict[str, Any][source]#

Implement actual metric calculation here

Override e.g. as:

def compute(self, y_true, y_pred, extra_info):
    print(extra_info)
    return lddt(y_true, y_pred, thres=self.custom_thresholds)

class rfproteina.metrics.design_metrics.ClashMetric(threshold: float = 1.5, **kwargs: Any)[source]#

Bases: Metric

Fraction of atoms in a steric clash, by clash type, at a configurable threshold.

Emits frac_atoms_clashing plus a 5-category breakdown (see compute_clash_stats()). The threshold is encoded in the metric name (e.g. ..._1p5A); pairs where both atoms are CRD-conditioned are excluded.

property kwargs_to_compute_args: dict[str, Any][source]#

Map input keys to a flat dictionary.

If not implemented, we return None, and pass the kwargs directly to the compute method.

Override e.g. as:

@cached_property
def kwargs_to_compute_args(self) -> dict[str, Any]:
    return {
        "y_true": ("network_input", "coords_unnoised"),
        "y_pred": ("network_output", "coords_pred"),
        "extra_info": ("extra_info",),
    }

compute(atom_array: AtomArray) → dict[str, Any][source]#

Implement actual metric calculation here

Override e.g. as:

def compute(self, y_true, y_pred, extra_info):
    print(extra_info)
    return lddt(y_true, y_pred, thres=self.custom_thresholds)

Bases: Metric

Fraction of residues with identical amino acid type between output and input structure.

Returns nan when the predicted and input protein sequences differ in length. Intended for use in partial diffusion.

property kwargs_to_compute_args: dict[str, Any][source]#

Map input keys to a flat dictionary.

If not implemented, we return None, and pass the kwargs directly to the compute method.

Override e.g. as:

@cached_property
def kwargs_to_compute_args(self) -> dict[str, Any]:
    return {
        "y_true": ("network_input", "coords_unnoised"),
        "y_pred": ("network_output", "coords_pred"),
        "extra_info": ("extra_info",),
    }

compute(predicted_atom_array: AtomArray, ground_truth_atom_array: AtomArray) → dict[str, Any][source]#

Compute sequence identity between generated and input structures.

Parameters:

predicted_atom_array – Generated protein structure.
ground_truth_atom_array – Input structure used to seed partial diffusion.

Returns:

Dict with sequence_identity_to_input in [0, 1], or nan if sequence lengths differ.