Adding a Custom Estimator#

cinnabar uses estimator objects to turn relative free energy measurements into per-ligand absolute free energy estimates. The built-in default is the maximum likelihood estimator (MLE), but you can provide your own estimator by subclassing Estimator and passing it to FEMap.generate_absolute_values. This tutorial provides a step-by-step example of how to do this.

Important: This tutorial implements a deliberately simple estimator. It should be treated as a toy example for learning the API, not as a statistically rigorous replacement for MLE.

The custom estimator will:

  1. Choose one ligand as an arbitrary zero point.

  2. Walk through the relative free energy graph using shortest paths.

  3. Assign each ligand a cumulative DG along its path from the anchor ligand.

  4. Store small metadata describing how the estimate was made.

For background on why absolute estimates are needed, see the Absolute Free Energy Estimators concept page.

Building a small FEMap#

First we create a connected toy network of relative binding free energies. Each relative calculation describes the free energy change from labelA to labelB.

from openff.units import unit

from cinnabar import FEMap

femap = FEMap()

femap.add_relative_calculation(
    labelA="ligand-a",
    labelB="ligand-b",
    value=1.2 * unit.kilocalorie_per_mole,
    uncertainty=0.2 * unit.kilocalorie_per_mole,
    source="toy-rbfe",
)
femap.add_relative_calculation(
    labelA="ligand-b",
    labelB="ligand-c",
    value=-0.4 * unit.kilocalorie_per_mole,
    uncertainty=0.3 * unit.kilocalorie_per_mole,
    source="toy-rbfe",
)
femap.add_relative_calculation(
    labelA="ligand-d",
    labelB="ligand-a",
    value=-0.7 * unit.kilocalorie_per_mole,
    uncertainty=0.2 * unit.kilocalorie_per_mole,
    source="toy-rbfe",
)

print(f"FEMap contains {femap.n_ligands} ligands")
print(f"FEMap is weakly connected: {femap.check_weakly_connected()}")

FEMap contains 4 ligands
FEMap is weakly connected: True

Implementing an estimator#

A custom estimator needs two pieces:

  • EstimatorResult: Subclass to store estimator-specific metadata which might be needed for downstream analysis

  • Estimator: Subclass that implements the _estimate method which is the main worker function of an Estimator.

FEMap.generate_absolute_values calls Estimator.estimate for you. That public method handles source grouping, connectivity checks, and storing metadata. Your custom _estimate method receives the measurements for one computational source this includes relative and absolute estimates, plus any experimental measurements in the graph and the source label that should be stamped onto returned measurements.

from dataclasses import dataclass

import networkx as nx

from cinnabar.estimators import Estimator, EstimatorResult
from cinnabar.measurements import Measurement, ReferenceState


@dataclass
class ShortestPathEstimatorResult(EstimatorResult):
    """Metadata produced by the toy shortest-path estimator."""

    anchor: str
    ligand_order: list[str]
    path_lengths: dict[str, int]


class ShortestPathEstimator(Estimator):
    """Toy estimator that accumulates relative free energies along shortest paths."""

    def __init__(self, anchor: str | None = None, source: str = "SP"):
        self.anchor = anchor
        self.source = source

    def _estimate(
        self,
        measurements: list[Measurement],
        source: str,
    ) -> tuple[list[Measurement], ShortestPathEstimatorResult]:
        graph = nx.DiGraph()
        units = {m.DG.u for m in measurements}
        if len(units) != 1:
            raise ValueError("All measurements must use the same free energy units")
        energy_unit = next(iter(units))

        for measurement in measurements:
            # only consider the computational relative measurements that connect ligands
            if not measurement.computational:
                continue
            # only consider relative measurements between ligands, not reference states (absolute calculation measurements)
            if isinstance(measurement.labelA, ReferenceState):
                continue

            # build a graph of edges to find the shortest path between ligands
            graph.add_edge(
                measurement.labelA,
                measurement.labelB,
                dg=measurement.DG.to(energy_unit).m,
                uncertainty=measurement.uncertainty.to(energy_unit).m,
                reverse=False,
            )
            graph.add_edge(
                measurement.labelB,
                measurement.labelA,
                dg=measurement.DG.to(energy_unit).m * -1.0,
                uncertainty=measurement.uncertainty.to(energy_unit).m,
                reverse=True,
            )

        if not graph:
            raise ValueError("No computational relative measurements were provided")

        # if we don't have an anchor ligand, just pick the first one in sorted order
        anchor = self.anchor or sorted(graph.nodes)[0]
        if anchor not in graph:
            raise ValueError(f"Anchor ligand {anchor} is not present in the graph")

        estimator_ref = ReferenceState(label=source)
        out_measurements = []
        path_lengths = {}

        paths = nx.shortest_path(graph, anchor)
        for ligand, path in paths.items():
            dg = 0.0
            variance = 0.0

            # for each edge along the path, accumulate the free energy and uncertainty accounting for the direction of the edge
            for label_a, label_b in zip(path[:-1], path[1:]):
                edge = graph[label_a][label_b]
                dg += edge["dg"]
                variance += edge["uncertainty"] ** 2

            out_measurements.append(
                Measurement(
                    labelA=estimator_ref,
                    labelB=ligand,
                    DG=dg * energy_unit,
                    uncertainty=variance**0.5 * energy_unit,
                    computational=True,
                    source=source,
                )
            )
            path_lengths[ligand] = len(path) - 1

        result = ShortestPathEstimatorResult(
            anchor=anchor,
            ligand_order=sorted(graph.nodes),
            path_lengths=path_lengths,
        )
        return out_measurements, result

The returned measurements are computational absolute free energies. The source argument passed into _estimate is used for both the returned measurements and the metadata key stored on the FEMap.

Running the estimator on an FEMap#

Custom estimators can be applied to an FEMap using the generate_absolute_values method and passing an instance of the estimator to apply, by default this method uses the built-in MLE estimator.

estimator = ShortestPathEstimator(anchor="ligand-a")
femap.generate_absolute_values(estimator=estimator)

absolute_df = femap.get_absolute_dataframe()
absolute_df
label DG (kcal/mol) uncertainty (kcal/mol) source computational
0 ligand-a 0.0 0.000000 SP True
1 ligand-b 1.2 0.200000 SP True
2 ligand-c 0.8 0.360555 SP True
3 ligand-d 0.7 0.200000 SP True

The values are on an arbitrary scale because ligand-a was chosen as the zero point. That is also true for many graph-based absolute reconstructions: adding a constant offset to every ligand does not change any relative free energy difference.

Accessing estimator metadata#

FEMap stores the result object returned by the estimator. For a single computational source, the metadata key is the estimator source name.

metadata = femap.get_estimator_metadata("SP")
metadata

ShortestPathEstimatorResult(estimator='ShortestPathEstimator', source='SP', anchor='ligand-a', ligand_order=['ligand-a', 'ligand-b', 'ligand-c', 'ligand-d'], path_lengths={'ligand-a': 0, 'ligand-b': 1, 'ligand-d': 1, 'ligand-c': 2})

metadata.anchor, metadata.path_lengths

('ligand-a', {'ligand-a': 0, 'ligand-b': 1, 'ligand-d': 1, 'ligand-c': 2})

If the FEMap contains multiple computational sources, cinnabar runs the estimator separately for each source. In that case, output sources are composed from the estimator source and the input source, for example SP(Method A) and SP(Method B).

Recap#

  • Subclass Estimator and implement _estimate(measurements, source).

  • Set self.source to the label users should see in FEMap dataframes.

  • Return absolute computational Measurement objects with labelA=ReferenceState(label=source).

  • Return an EstimatorResult object for any metadata that should be stored on the FEMap.

  • Run the estimator through FEMap.generate_absolute_values(estimator=...) rather than calling _estimate directly.

For production analysis, prefer the built-in MLE estimator unless you have a specific estimator model you want to validate.