{ "cells": [ { "cell_type": "markdown", "id": "6528c710a249ee9f", "metadata": {}, "source": [ "# Adding a Custom Estimator\n", "\n", "``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.\n", "\n", "
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.
\n", "\n", "\n", "The custom estimator will:\n", "\n", "1. Choose one ligand as an arbitrary zero point.\n", "2. Walk through the relative free energy graph using shortest paths.\n", "3. Assign each ligand a cumulative ``DG`` along its path from the anchor ligand.\n", "4. Store small metadata describing how the estimate was made.\n", "\n", "For background on why absolute estimates are needed, see the {ref}`estimators` concept page." ] }, { "cell_type": "markdown", "id": "5a50d90617bcea41", "metadata": {}, "source": [ "## Building a small FEMap\n", "\n", "First we create a connected toy network of relative binding free energies. Each relative calculation describes the free energy change from ``labelA`` to ``labelB``." ] }, { "cell_type": "code", "execution_count": 1, "id": "28c29c527e6c07ac", "metadata": { "ExecuteTime": { "end_time": "2026-07-08T11:48:16.983181Z", "start_time": "2026-07-08T11:48:13.111261Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "FEMap contains 4 ligands\n", "FEMap is weakly connected: True\n" ] } ], "source": [ "from openff.units import unit\n", "\n", "from cinnabar import FEMap\n", "\n", "femap = FEMap()\n", "\n", "femap.add_relative_calculation(\n", " labelA=\"ligand-a\",\n", " labelB=\"ligand-b\",\n", " value=1.2 * unit.kilocalorie_per_mole,\n", " uncertainty=0.2 * unit.kilocalorie_per_mole,\n", " source=\"toy-rbfe\",\n", ")\n", "femap.add_relative_calculation(\n", " labelA=\"ligand-b\",\n", " labelB=\"ligand-c\",\n", " value=-0.4 * unit.kilocalorie_per_mole,\n", " uncertainty=0.3 * unit.kilocalorie_per_mole,\n", " source=\"toy-rbfe\",\n", ")\n", "femap.add_relative_calculation(\n", " labelA=\"ligand-d\",\n", " labelB=\"ligand-a\",\n", " value=-0.7 * unit.kilocalorie_per_mole,\n", " uncertainty=0.2 * unit.kilocalorie_per_mole,\n", " source=\"toy-rbfe\",\n", ")\n", "\n", "print(f\"FEMap contains {femap.n_ligands} ligands\")\n", "print(f\"FEMap is weakly connected: {femap.check_weakly_connected()}\")" ] }, { "cell_type": "markdown", "id": "fe5e963fb482e2b", "metadata": {}, "source": [ "## Implementing an estimator\n", "\n", "A custom estimator needs two pieces:\n", "\n", "- ``EstimatorResult``: Subclass to store estimator-specific metadata which might be needed for downstream analysis\n", "- ``Estimator``: Subclass that implements the ``_estimate`` method which is the main worker function of an ``Estimator``.\n", "\n", "``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." ] }, { "cell_type": "code", "execution_count": 2, "id": "c27f1ef4efa13101", "metadata": { "ExecuteTime": { "end_time": "2026-07-08T11:48:17.004409Z", "start_time": "2026-07-08T11:48:16.993068Z" } }, "outputs": [], "source": [ "from dataclasses import dataclass\n", "\n", "import networkx as nx\n", "\n", "from cinnabar.estimators import Estimator, EstimatorResult\n", "from cinnabar.measurements import Measurement, ReferenceState\n", "\n", "\n", "@dataclass\n", "class ShortestPathEstimatorResult(EstimatorResult):\n", " \"\"\"Metadata produced by the toy shortest-path estimator.\"\"\"\n", "\n", " anchor: str\n", " ligand_order: list[str]\n", " path_lengths: dict[str, int]\n", "\n", "\n", "class ShortestPathEstimator(Estimator):\n", " \"\"\"Toy estimator that accumulates relative free energies along shortest paths.\"\"\"\n", "\n", " def __init__(self, anchor: str | None = None, source: str = \"SP\"):\n", " self.anchor = anchor\n", " self.source = source\n", "\n", " def _estimate(\n", " self,\n", " measurements: list[Measurement],\n", " source: str,\n", " ) -> tuple[list[Measurement], ShortestPathEstimatorResult]:\n", " graph = nx.DiGraph()\n", " units = {m.DG.u for m in measurements}\n", " if len(units) != 1:\n", " raise ValueError(\"All measurements must use the same free energy units\")\n", " energy_unit = next(iter(units))\n", "\n", " for measurement in measurements:\n", " # only consider the computational relative measurements that connect ligands\n", " if not measurement.computational:\n", " continue\n", " # only consider relative measurements between ligands, not reference states (absolute calculation measurements)\n", " if isinstance(measurement.labelA, ReferenceState):\n", " continue\n", "\n", " # build a graph of edges to find the shortest path between ligands\n", " graph.add_edge(\n", " measurement.labelA,\n", " measurement.labelB,\n", " dg=measurement.DG.to(energy_unit).m,\n", " uncertainty=measurement.uncertainty.to(energy_unit).m,\n", " reverse=False,\n", " )\n", " graph.add_edge(\n", " measurement.labelB,\n", " measurement.labelA,\n", " dg=measurement.DG.to(energy_unit).m * -1.0,\n", " uncertainty=measurement.uncertainty.to(energy_unit).m,\n", " reverse=True,\n", " )\n", "\n", " if not graph:\n", " raise ValueError(\"No computational relative measurements were provided\")\n", "\n", " # if we don't have an anchor ligand, just pick the first one in sorted order\n", " anchor = self.anchor or sorted(graph.nodes)[0]\n", " if anchor not in graph:\n", " raise ValueError(f\"Anchor ligand {anchor} is not present in the graph\")\n", "\n", " estimator_ref = ReferenceState(label=source)\n", " out_measurements = []\n", " path_lengths = {}\n", "\n", " paths = nx.shortest_path(graph, anchor)\n", " for ligand, path in paths.items():\n", " dg = 0.0\n", " variance = 0.0\n", "\n", " # for each edge along the path, accumulate the free energy and uncertainty accounting for the direction of the edge\n", " for label_a, label_b in zip(path[:-1], path[1:]):\n", " edge = graph[label_a][label_b]\n", " dg += edge[\"dg\"]\n", " variance += edge[\"uncertainty\"] ** 2\n", "\n", " out_measurements.append(\n", " Measurement(\n", " labelA=estimator_ref,\n", " labelB=ligand,\n", " DG=dg * energy_unit,\n", " uncertainty=variance**0.5 * energy_unit,\n", " computational=True,\n", " source=source,\n", " )\n", " )\n", " path_lengths[ligand] = len(path) - 1\n", "\n", " result = ShortestPathEstimatorResult(\n", " anchor=anchor,\n", " ligand_order=sorted(graph.nodes),\n", " path_lengths=path_lengths,\n", " )\n", " return out_measurements, result" ] }, { "cell_type": "markdown", "id": "fd06dbbe5be7817a", "metadata": {}, "source": [ "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``." ] }, { "cell_type": "markdown", "id": "5ff39dcc7183b699", "metadata": {}, "source": [ "## Running the estimator on an ``FEMap``\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": 3, "id": "268b44d41a6ac43e", "metadata": { "ExecuteTime": { "end_time": "2026-07-08T11:48:17.030815Z", "start_time": "2026-07-08T11:48:17.010021Z" } }, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
labelDG (kcal/mol)uncertainty (kcal/mol)sourcecomputational
0ligand-a0.00.000000SPTrue
1ligand-b1.20.200000SPTrue
2ligand-c0.80.360555SPTrue
3ligand-d0.70.200000SPTrue
\n", "
" ], "text/plain": [ " label DG (kcal/mol) uncertainty (kcal/mol) source computational\n", "0 ligand-a 0.0 0.000000 SP True\n", "1 ligand-b 1.2 0.200000 SP True\n", "2 ligand-c 0.8 0.360555 SP True\n", "3 ligand-d 0.7 0.200000 SP True" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "estimator = ShortestPathEstimator(anchor=\"ligand-a\")\n", "femap.generate_absolute_values(estimator=estimator)\n", "\n", "absolute_df = femap.get_absolute_dataframe()\n", "absolute_df" ] }, { "cell_type": "markdown", "id": "d71f359026555fd2", "metadata": {}, "source": [ "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." ] }, { "cell_type": "markdown", "id": "7a16d92b821aa541", "metadata": {}, "source": [ "## Accessing estimator metadata\n", "\n", "``FEMap`` stores the result object returned by the estimator. For a single computational source, the metadata key is the estimator source name." ] }, { "cell_type": "code", "execution_count": 4, "id": "1451ba901f7b8f48", "metadata": { "ExecuteTime": { "end_time": "2026-07-08T11:48:17.054081Z", "start_time": "2026-07-08T11:48:17.032966Z" } }, "outputs": [ { "data": { "text/plain": [ "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})" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "metadata = femap.get_estimator_metadata(\"SP\")\n", "metadata" ] }, { "cell_type": "code", "execution_count": 5, "id": "5299ec843711a4bb", "metadata": { "ExecuteTime": { "end_time": "2026-07-08T11:48:17.067221Z", "start_time": "2026-07-08T11:48:17.055362Z" } }, "outputs": [ { "data": { "text/plain": [ "('ligand-a', {'ligand-a': 0, 'ligand-b': 1, 'ligand-d': 1, 'ligand-c': 2})" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "metadata.anchor, metadata.path_lengths" ] }, { "cell_type": "markdown", "id": "5c41165775e2aceb", "metadata": {}, "source": [ "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)``." ] }, { "cell_type": "markdown", "id": "2a77fc8f85091d4f", "metadata": {}, "source": [ "## Recap\n", "\n", "- Subclass ``Estimator`` and implement ``_estimate(measurements, source)``.\n", "- Set ``self.source`` to the label users should see in ``FEMap`` dataframes.\n", "- Return absolute computational ``Measurement`` objects with ``labelA=ReferenceState(label=source)``.\n", "- Return an ``EstimatorResult`` object for any metadata that should be stored on the ``FEMap``.\n", "- Run the estimator through ``FEMap.generate_absolute_values(estimator=...)`` rather than calling ``_estimate`` directly.\n", "\n", "For production analysis, prefer the built-in MLE estimator unless you have a specific estimator model you want to validate." ] } ], "metadata": { "kernelspec": { "display_name": "openfe", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.2" } }, "nbformat": 4, "nbformat_minor": 5 }