calibrax.metrics.plugins.scientific

Scientific domain metrics for molecular and protein modeling. Evaluates molecular geometry generation (chemical validity), binding affinity prediction, and conformational sampling quality. Mostly pure JAX math with no external chemistry dependencies.

No External Dependencies

The [scientific] extra exists for semantic grouping but requires no external packages — all functions use pure JAX math.

Scientific metrics for molecular and protein modeling.

Mostly pure JAX math -- no external chemistry dependencies required. Designed for evaluating molecular geometry generation, binding affinity prediction, and conformational sampling.

Typical use cases: - Molecular dynamics simulations - Drug discovery pipelines - Protein structure prediction

chemical_validity(bond_lengths: jnp.ndarray, bond_angles: jnp.ndarray, *, length_thresholds: tuple[float, float] = (0.8, 2.0), angle_thresholds: tuple[float, float] = (1.5, 3.5)) -> Any

Fraction of bonds and angles within acceptable physical ranges.

Checks if generated molecular geometries are physically plausible by validating bond lengths (in Angstroms) and bond angles (in radians) against threshold ranges.

Parameters:

Name	Type	Description	Default
bond_lengths	ndarray	1D array of bond lengths in Angstroms.	required
bond_angles	ndarray	1D array of bond angles in radians.	required
length_thresholds	tuple[float, float]	(min, max) acceptable bond length range. Defaults to (0.8, 2.0) Angstroms.	(0.8, 2.0)
angle_thresholds	tuple[float, float]	(min, max) acceptable bond angle range. Defaults to (1.5, 3.5) radians (~86-200 degrees).	(1.5, 3.5)

Returns:

Type	Description
Any	Fraction of valid bonds and angles in [0, 1]. Higher is better.

Examples:

>>> lengths = jnp.array([1.0, 1.2, 1.5])
>>> angles = jnp.array([2.0, 2.5, 3.0])
>>> chemical_validity(lengths, angles)
1.0

binding_affinity_metrics(predictions: jnp.ndarray, targets: jnp.ndarray) -> dict[str, Any]

Compute regression metrics for binding affinity predictions.

Evaluates pKd/pKi predictions using standard regression metrics. Delegates to existing regression and statistical functions (DRY).

Parameters:

Name	Type	Description	Default
predictions	ndarray	Predicted binding affinity values (e.g., pKd).	required
targets	ndarray	Ground truth binding affinity values.	required

Returns:

Type	Description
dict[str, Any]	Dictionary with binding_mse, binding_mae, binding_r_squared,
dict[str, Any]	and binding_pearson values.

Examples:

>>> preds = jnp.array([6.5, 7.2, 8.1])
>>> targets = jnp.array([6.8, 7.0, 8.3])
>>> result = binding_affinity_metrics(preds, targets)
>>> result["binding_mse"]  # MSE of predictions

conformational_diversity(coordinates: jnp.ndarray) -> Any

Mean pairwise RMSD across molecular conformations.

Measures diversity of molecular conformations by computing the root mean square deviation (RMSD) between all pairs of structures.

Parameters:

Name	Type	Description	Default
coordinates	ndarray	Array of shape (num_conformations, num_atoms, 3). Each conformation is a set of 3D atomic coordinates.	required

Returns:

Type	Description
Any	Mean pairwise RMSD in the same units as input coordinates
Any	(typically Angstroms). Higher means more diverse conformations.

Examples:

>>> coords = jnp.array([
...     [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
...     [[1.0, 1.0, 1.0], [2.0, 1.0, 1.0]],
... ])
>>> conformational_diversity(coords)  # > 0.0