{ "cells": [ { "cell_type": "markdown", "id": "2f60e9d4", "metadata": {}, "source": [ "# Closed-Loop Manifold Active Learning with AFL Servers\n", "\n", "This case study runs the manifold active learning campaign in a format amenable with the usage of AFL's microservice concept:\n", "\n", "1. Connect to the services (Tiled, instrument, agent).\n", "2. Specify the problem (campaign ID, components, bounds, number of iterations).\n", "3. Configure data collection: tell the agent how to find each measurement in Tiled.\n", "4. Convert agent suggestions into measurement targets.\n", "5. Run one experiment: the instrument server measures and writes each exposure to Tiled, the agent appends the entry IDs to its dataset.\n", "6. Seed the campaign with the initial composition.\n", "7. Build the pipeline on the agent server.\n", "8. Close the loop: `predict`, measure the suggestion, append, repeat.\n", " \n", "- The virtual instrument `ProteinGlycerolTemperaturePhaseMap` runs as an AFL-automation `APIServer`, exactly like a hardware driver. `AFL.double_agent.VirtualInstrumentDriver` serves whichever virtual instrument the `instrument` key of its `--config` file names (`phase_map_instrument.json` here, which also carries the instrument's `instrument_overrides`), takes that instrument's class name as the server and `driver_name`, and exposes every public method of the instrument as a task: this notebook assigns it a sample with `set_sample` and enqueues one `expose` task per temperature; the server writes the returned dataset to Tiled with the sample identity and campaign in its attributes and reports the Tiled entry ID. \n", "- The agent, a `DoubleAgentDriver` server, reads the entries by ID, fits the pipeline and returns the next composition and temperature profile.\n", "\n", "## Prerequisites\n", "\n", "Three services must be running, all pointed at the same Tiled instance through `~/.afl/config.json` (`tiled_server`, `tiled_api_key` and the `\"ports\"` of the two drivers):\n", "\n", "| Service | Example address | Launch |\n", "| --- | --- | --- |\n", "| Tiled | `http://localhost:8000` | `tiled serve config config.yml` (AFL-automation `tiled/config.yml` with `single_user_api_key` set) |\n", "| `ProteinGlycerolTemperaturePhaseMap` | `localhost:5052` | `python -m AFL.double_agent.VirtualInstrumentDriver --config phase_map_instrument.json` in this directory |\n", "| `DoubleAgentDriver` | `localhost:5050` | `PYTHONPATH=. python /AFL/double_agent/AgentDriver.py` in this directory |\n", "\n", "The agent server is launched from this directory because the pipeline holds lambdas, the SRSF manifold and the simulator, which have no JSON form. \n", "\n", "The settings, the acquisition scores and the `Pipeline` are defined in `helpers.py` in this directory, which this notebook and the companion notebook import. The notebook sends the JSON of the pipeline, in which those objects appear as references, and names the module as the pipeline's `environment`; the agent server imports it to obtain the live objects behind the references.\n", "\n", "The analysis environment needs `AFL.double_agent` with the `manifold` extra, `AFL.automation` and the Tiled client." ] }, { "cell_type": "markdown", "id": "dc9f0ee5", "metadata": {}, "source": [ "## Set Up the Environment\n", "\n", "Import the client of the AFL servers, the Tiled helper and, from `helpers.py`, the settings and functions of the campaign. Everything the pipeline needs lives in that module so that the notebook and the agent server agree on it." ] }, { "cell_type": "code", "execution_count": null, "id": "731d12b5", "metadata": {}, "outputs": [], "source": [ "# Uncomment this cell if you need to install AFL-agent with the ``manifold`` extra\n", "# (torch, gpytorch, pyro-ppl, botorch, geoopt, alphashape, trimesh) and AFL-automation.\n", "# !pip install \"AFL-agent[manifold] @ git+https://github.com/usnistgov/afl-agent\"\n", "# !pip install \"AFL-automation @ git+https://github.com/usnistgov/afl-automation\"" ] }, { "cell_type": "code", "execution_count": 1, "id": "43cb7afc", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/knv2/codebase/AFL-bioformulations/.venv/lib/python3.11/site-packages/torch/cuda/__init__.py:182: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 11040). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:119.)\n", " return torch._C._cuda_getDeviceCount() > 0\n" ] } ], "source": [ "import uuid\n", "\n", "import numpy as np\n", "\n", "from AFL.automation.APIServer.Client import Client\n", "from AFL.automation.shared.tiled import get_tiled_client\n", "\n", "from helpers import BOUNDS_T, LIMITS, XY, pipeline, xarray_to_target # settings, scores and pipeline, see helpers.py in this folder\n", "\n", "np.set_printoptions(precision=3, suppress=True)" ] }, { "cell_type": "code", "execution_count": 2, "id": "8a4a160a", "metadata": {}, "outputs": [], "source": [ "# services\n", "INSTRUMENT_HOST, INSTRUMENT_PORT = \"localhost\", 5052 # ProteinGlycerolTemperaturePhaseMap APIServer\n", "AGENT_HOST, AGENT_PORT = \"localhost\", 5050 # DoubleAgentDriver APIServer\n", "USERNAME = \"afl-agent\"\n", "\n", "# campaign\n", "N_ITERATIONS = 4 # compositions (one temperature profile each) acquired after the seed\n", "INITIAL_COMPOSITION = {\"protein\": 40.0, \"glycerol\": 5.0} # the one composition measured before the campaign\n", "INITIAL_TEMPERATURES = [5.0, 10.0, 20.0, 30.0, 40.0, 45.0] # ... and the temperatures it is measured at\n", "SEED = 1000\n", "\n", "# how the instrument's entries are identified in Tiled\n", "INSTRUMENT_NAME = \"ProteinGlycerolTemperaturePhaseMap\" # driver_name written to every Tiled entry\n", "EXPOSE_TASK = \"expose\" # task_name written to every Tiled entry\n", "SAMPLE_DIM = \"sample\"\n", "INPUT_VARIABLES = [\"design_space\", \"phases\", \"composition\"] # pulled from each entry into the agent dataset\n", "PIPELINE_ENVIRONMENT = \"helpers\" # module importable on the agent server that defines ``pipeline``" ] }, { "cell_type": "markdown", "id": "14f6b3e5", "metadata": {}, "source": [ "## Connect to the Services\n", "\n", "`get_tiled_client` reads the Tiled server and API key from `~/.afl/config.json`, the same configuration the servers use. Create a client for the instrument and the agent and log in to both." ] }, { "cell_type": "code", "execution_count": 3, "id": "82ff4784", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "['Idle', 'Configurations Loaded=3', 'Phases Traced=3', 'Noise Level=0.0']" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tiled = get_tiled_client()\n", "\n", "instrument = Client(INSTRUMENT_HOST, port=INSTRUMENT_PORT)\n", "agent = Client(AGENT_HOST, port=AGENT_PORT)\n", "for client in (instrument, agent):\n", " client.login(USERNAME)\n", "\n", "instrument.enqueue(task_name=\"status\", interactive=True)[\"return_val\"]" ] }, { "cell_type": "markdown", "id": "cac72eae", "metadata": {}, "source": [ "## Specify the Optimization Problem\n", "\n", "Use one campaign ID everywhere. A unique campaign ID prevents measurements from different runs from being combined. The composition components and their bounds are those of the instrument; the temperature range is the one the stage-2 profile is mapped onto." ] }, { "cell_type": "code", "execution_count": 4, "id": "9d16873e", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "('manifold-95dad5a5',\n", " {'protein': (0.5, 75.0), 'glycerol': (0.5, 10.0)},\n", " [0.5, 45.0])" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "campaign_id = f\"manifold-{uuid.uuid4().hex[:8]}\"\n", "components = XY # composition components searched in stage 1: protein, glycerol\n", "bounds = {c: LIMITS[c] for c in components}\n", "np.random.seed(SEED)\n", "campaign_id, bounds, BOUNDS_T" ] }, { "cell_type": "markdown", "id": "d366696b", "metadata": {}, "source": [ "## Configure Data Collection\n", "\n", "Tell the agent how to find and align the measurements. Every Tiled entry written by the instrument server holds one exposure; `driver_name` and `task_name` select the instrument's `expose` entries and `campaign_id_path` points at the campaign attribute that `set_sample` adds to each of them. The three variables are read from the entry's dataset and stacked along `sample`." ] }, { "cell_type": "code", "execution_count": 5, "id": "e07482a2", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'campaign_id': 'manifold-95dad5a5',\n", " 'sources': ['design_space', 'phases', 'composition'],\n", " 'started_at': '2026-10-27T19:33:48.676258+00:00',\n", " 'status': 'success'}" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "input_spec = {\n", " \"campaign_id\": campaign_id,\n", " \"sample_dim\": SAMPLE_DIM,\n", " \"sources\": {\n", " name: {\n", " \"driver_name\": INSTRUMENT_NAME,\n", " \"task_name\": EXPOSE_TASK,\n", " \"campaign_id_path\": \"attrs.AL_campaign_name\",\n", " \"source\": \"dataset\",\n", " \"path\": name,\n", " }\n", " for name in INPUT_VARIABLES\n", " },\n", "}\n", "agent.enqueue(task_name=\"setup_data_collection\", input_spec=input_spec, interactive=True)[\"return_val\"]" ] }, { "cell_type": "markdown", "id": "eb9f65b9", "metadata": {}, "source": [ "## Run One Experiment\n", "\n", "A target is one composition at several temperatures. Each temperature is one exposure of the instrument: assign the sample identity and campaign with `set_sample`, enqueue `expose`, and keep the Tiled entry ID the server reports. The agent then reads the entry by ID; since all three input variables come from the same entry, the same ID is given for each of them.\n", "\n", "The phase label of each exposure is read back from Tiled only for the progress report below." ] }, { "cell_type": "code", "execution_count": 6, "id": "400057c1", "metadata": {}, "outputs": [], "source": [ "def measure_sample(iteration, target):\n", " entry_ids = []\n", " for temperature in target[\"temperatures\"]:\n", " temperature = float(temperature)\n", " instrument.enqueue(\n", " task_name=\"set_sample\",\n", " sample_name=f\"{campaign_id}-sample-{iteration:03d}-T{temperature:.1f}\",\n", " sample_uuid=f\"SAM-{uuid.uuid4()}\",\n", " AL_campaign_name=campaign_id,\n", " interactive=True,\n", " )\n", " exposure = instrument.enqueue(\n", " task_name=EXPOSE_TASK,\n", " protein=target[\"protein\"],\n", " glycerol=target[\"glycerol\"],\n", " temperature=temperature,\n", " interactive=True,\n", " )\n", " entry_id = exposure.get(\"tiled_entry_id\")\n", " if not entry_id:\n", " raise RuntimeError(\n", " f\"the instrument server did not write to Tiled (tiled_status={exposure.get('tiled_status')!r}); \"\n", " \"check tiled_server and tiled_api_key in ~/.afl/config.json\"\n", " )\n", " agent.enqueue(\n", " task_name=\"append_dataset\",\n", " entries={name: entry_id for name in INPUT_VARIABLES},\n", " interactive=True,\n", " )\n", " entry_ids.append(entry_id)\n", " return entry_ids\n", "\n", "\n", "def read_phases(entry_ids):\n", " return [int(tiled[\"run_documents\"][entry_id][\"phases\"].read()) for entry_id in entry_ids]" ] }, { "cell_type": "markdown", "id": "4e5a9ad2", "metadata": {}, "source": [ "## Seed the Campaign\n", "\n", "The campaign starts from a single composition measured at six temperatures." ] }, { "cell_type": "code", "execution_count": 7, "id": "a6b39f9e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "seed: phases=[1, 1, 1, 0, 0, 0] (n_measurements=6)\n" ] } ], "source": [ "initial_target = {**INITIAL_COMPOSITION, \"temperatures\": INITIAL_TEMPERATURES}\n", "initial_entry_ids = measure_sample(0, initial_target)\n", "n_measurements = len(initial_entry_ids)\n", "print(f\"seed: phases={read_phases(initial_entry_ids)} (n_measurements={n_measurements})\")" ] }, { "cell_type": "markdown", "id": "2c8b7961", "metadata": {}, "source": [ "## Build the Pipeline\n", "\n", "`pipeline` is the manifold pipeline of `helpers.py`: the `DirichletGPExtrapolator` phase classifier, the `PhaseMapF1Score` benchmark against the simulator, `BoTorchAcquisition` over the composition with `composition_score` and `SRSFHilbertSphereAcquisitionFunction` over the temperature profile with `profile_score`. Its JSON form records the lambdas, the manifold and the simulator as references; `initialize_pipeline` receives the JSON together with `environment`, the module the agent server imports to obtain the live objects behind those references." ] }, { "cell_type": "code", "execution_count": 8, "id": "c4b24839", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "0 Grid-DesignSpace\n", "1 Grid-Composition\n", "2 DirichletGPExtrapolator-PhaseLabels\n", "3 PhaseMapF1Score\n", "4 BoTorchAcquisition-Composition\n", "5 SRSF-TemperatureProfile\n" ] }, { "data": { "text/plain": [ "'Success!'" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "for i, op in enumerate(pipeline):\n", " print(i, op.name)\n", "\n", "agent.enqueue(\n", " task_name=\"initialize_pipeline\",\n", " pipeline=pipeline.to_dict()[\"ops\"],\n", " name=pipeline.name,\n", " environment=PIPELINE_ENVIRONMENT,\n", " interactive=True,\n", ")[\"exit_state\"]" ] }, { "cell_type": "markdown", "id": "f65cf2a1", "metadata": {}, "source": [ "## Close the Loop\n", "\n", "Each call to `predict` evaluates the accumulated campaign data on the agent server. Retrieve `last_results`, turn its suggestion into a target, measure it and append the new exposures before requesting another prediction.\n", "\n", "The benchmark variables of a result belong to the model fitted *before* the new batch, so the F1 score reported at iteration `i` is that of a model trained on `i - 1` acquired compositions." ] }, { "cell_type": "code", "execution_count": 9, "id": "d242a7ce", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[ 1/4] protein= 72.95 glycerol=10.00 temperatures=[0.5, 3.0, 27.0, 45.0] phases=[1, 1, 1, 0] (n_measurements=10) f1=0.495 accuracy=0.597\n", "[ 2/4] protein= 6.29 glycerol= 0.80 temperatures=[0.5, 24.0, 27.0, 45.0] phases=[1, 0, 0, 0] (n_measurements=14) f1=0.505 accuracy=0.603\n", "[ 3/4] protein= 4.72 glycerol=10.00 temperatures=[0.5, 21.0, 45.0] phases=[1, 1, 0] (n_measurements=17) f1=0.512 accuracy=0.620\n", "[ 4/4] protein= 72.42 glycerol= 0.50 temperatures=[0.5, 18.0, 21.0, 45.0] phases=[2, 1, 1, 0] (n_measurements=21) f1=0.504 accuracy=0.607\n" ] } ], "source": [ "for iteration in range(1, N_ITERATIONS + 1):\n", " agent.enqueue(task_name=\"predict\", AL_campaign_name=campaign_id, interactive=True)\n", " results = agent.get_driver_object(\"last_results\")\n", " target = xarray_to_target(results)\n", "\n", " entry_ids = measure_sample(iteration, target)\n", " n_measurements += len(entry_ids)\n", "\n", " print(\n", " f\"[{iteration:3d}/{N_ITERATIONS}] protein={target['protein']:6.2f} glycerol={target['glycerol']:5.2f} \"\n", " f\"temperatures={[round(float(t), 1) for t in target['temperatures']]} \"\n", " f\"phases={read_phases(entry_ids)} \"\n", " f\"(n_measurements={n_measurements}) \"\n", " f\"f1={results.benchmark_f1.item():.3f} accuracy={results.benchmark_accuracy.item():.3f}\"\n", " )" ] }, { "cell_type": "markdown", "id": "55c80e14", "metadata": {}, "source": [ "After the final iteration, refit on everything measured. `PhaseMapF1Score` scores the classifier against the simulator, and the agent archives the result dataset in Tiled next to the campaign's entries." ] }, { "cell_type": "code", "execution_count": 10, "id": "68939ec7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "score: campaign=manifold-95dad5a5 n_measurements=21 f1=0.519 accuracy=0.627\n" ] } ], "source": [ "agent.enqueue(task_name=\"predict\", AL_campaign_name=campaign_id, interactive=True)\n", "results = agent.get_driver_object(\"last_results\")\n", "print(\n", " f\"score: campaign={campaign_id} n_measurements={n_measurements} \"\n", " f\"f1={results.benchmark_f1.item():.3f} accuracy={results.benchmark_accuracy.item():.3f}\"\n", ")" ] }, { "cell_type": "code", "execution_count": 11, "id": "671c9d2e", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.Dataset> Size: 2kB\n",
       "Dimensions:           (next_sample: 1, component: 2, profile_t: 100,\n",
       "                       profile_dim: 1)\n",
       "Coordinates:\n",
       "  * component         (component) <U8 64B 'protein' 'glycerol'\n",
       "  * profile_t         (profile_t) float64 800B 0.0 0.0101 0.0202 ... 0.9899 1.0\n",
       "Dimensions without coordinates: next_sample, profile_dim\n",
       "Data variables:\n",
       "    composition_next  (next_sample, component) float64 16B 73.83 2.059\n",
       "    profile_function  (profile_t, profile_dim) float64 800B 0.0 4.79e-07 ... 1.0\n",
       "    benchmark_f1      float64 8B 0.5189\n",
       "Attributes:\n",
       "    sample_uuid:       SAM-9bd9e417-7c07-4a54-82de-120fa519b8f1\n",
       "    ag_uuid:           AG-e508f98e-2584-46c0-9fd2-010eb4f3ccd4\n",
       "    AL_campaign_name:  manifold-95dad5a5
" ], "text/plain": [ " Size: 2kB\n", "Dimensions: (next_sample: 1, component: 2, profile_t: 100,\n", " profile_dim: 1)\n", "Coordinates:\n", " * component (component)