diff --git a/tutorials/2-examples/DTEx253_deeptrack_with_agents.ipynb b/tutorials/2-examples/DTEx253_deeptrack_with_agents.ipynb new file mode 100644 index 000000000..9e661f9f2 --- /dev/null +++ b/tutorials/2-examples/DTEx253_deeptrack_with_agents.ipynb @@ -0,0 +1,593 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0036b202", + "metadata": {}, + "source": [ + "# Using Codex to prepare simulations with DeepTrack2\n", + "\n", + "This notebook is an example of how one could prompt an AI agent to generate code for preparing simulations with DeepTrack2. \n", + "It defines the a spherical particle and an\n", + "optical configuration. You can use your own Codex CLI to\n", + "write one reusable simulation.py, and then execute that simulator on the\n", + "requested particle and on a second particle.\n", + "\n", + "The default request is fluorescence. You can change or edit MODALITY\n", + "to try brightfield, darkfield, iSCAT, or inline\n", + "holography. ARM selects either the baseline request or the\n", + "DeepTrack2 request.\n", + "\n", + "Every run is stored below deeptrack_with_agents/." + ] + }, + { + "cell_type": "markdown", + "id": "173c0f94", + "metadata": {}, + "source": [ + "## Before running\n", + "\n", + "1. Use a Python environment with NumPy and Matplotlib. For ARM =\n", + " \"deeptrack2\", use an environment in which DeepTrack2 is installed.\n", + "2. Install the Codex CLI if needed, then sign in from a terminal with your own\n", + " account. The official Codex CLI guide covers installation, sign-in, and\n", + " codex exec: https://learn.chatgpt.com/docs/codex/cli\n", + "3. Run the setup and configuration cells first. Set RUN_CODEX = True only after reviewing the prompt and output\n", + " directory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a231555", + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import copy\n", + "import importlib.util\n", + "import json\n", + "import os\n", + "import re\n", + "import shutil\n", + "import subprocess\n", + "import sys\n", + "import textwrap\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "\n", + "MODALITIES = (\"fluorescence\", \"brightfield\", \"darkfield\", \"iscat\", \"holography\")\n", + "\n", + "# Edit these three values\n", + "MODALITY = \"fluorescence\"\n", + "ARM = \"deeptrack2\" # choose \"baseline\" or \"deeptrack2\"\n", + "RUN_LABEL = \"run_00\" # use a new label to keep an earlier run\n", + "\n", + "NOTEBOOK_ROOT = Path.cwd().resolve()\n", + "REPRO_ROOT = NOTEBOOK_ROOT / \"deeptrack_with_agents\"\n", + "REPRO_ROOT.mkdir(parents=True, exist_ok=True)\n", + "\n", + "assert MODALITY in MODALITIES\n", + "assert ARM in {\"baseline\", \"deeptrack2\"}\n", + "print(f\"working directory: {NOTEBOOK_ROOT}\")\n", + "print(f\"selected task: {MODALITY} | arm: {ARM} | run: {RUN_LABEL}\")" + ] + }, + { + "cell_type": "markdown", + "id": "c2550c5d", + "metadata": {}, + "source": [ + "## 1. Define the requested particle and optical configuration\n", + "\n", + "Positions are image-pixel coordinates, axial position and radius are in micrometres,\n", + "refractive index is [real, imaginary], and fluorescence intensity is an\n", + "arbitrary simulator intensity parameter.\n", + "\n", + "The second particle is only a reusability check. It is not shown to Codex\n", + "while the simulator is generated." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b329229b", + "metadata": {}, + "outputs": [], + "source": [ + "particle = {\n", + " \"position_xy\": [67.0, 59.0],\n", + " \"axial_position_um\": 1.5,\n", + " \"radius_um\": 0.45,\n", + " \"refractive_index\": [1.45, 0.01],\n", + " \"fluorescence_intensity\": 10.0,\n", + "}\n", + "\n", + "reuse_particle = {\n", + " \"position_xy\": [39.0, 82.0],\n", + " \"axial_position_um\": 2.3,\n", + " \"radius_um\": 0.72,\n", + " \"refractive_index\": [1.51, 0.02],\n", + " \"fluorescence_intensity\": 7.0,\n", + "}\n", + "\n", + "BASE_CONFIG = {\n", + " \"seed\": 123,\n", + " \"image_size\": 128,\n", + " \"camera_resolution_um\": 1.0,\n", + " \"magnification\": 10.0,\n", + " \"numerical_aperture\": 0.8,\n", + " \"medium_refractive_index\": 1.33,\n", + " \"fluorescence_wavelength_nm\": 600.0,\n", + " \"coherent_wavelength_nm\": 532.0,\n", + " \"padding_px\": 32,\n", + " \"iscat_illumination_angle_rad\": float(np.pi),\n", + " \"iscat_amplitude_factor\": 1.0,\n", + "}\n", + "\n", + "\n", + "def config_for(modality: str) -> dict:\n", + " config = copy.deepcopy(BASE_CONFIG)\n", + " config[\"modality\"] = modality\n", + " config[\"particle\"] = copy.deepcopy(particle)\n", + " return config\n", + "\n", + "\n", + "print(json.dumps({\"particle\": particle, \"reuse_particle\": reuse_particle}, indent=2))" + ] + }, + { + "cell_type": "markdown", + "id": "5327494b", + "metadata": {}, + "source": [ + "## 2. Choose or edit the task prompt\n", + "\n", + "The task text below is the prompt that will be sent to Codex. You can edit it to change the request." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cf196174", + "metadata": {}, + "outputs": [], + "source": [ + "MODALITY_TASKS = {\n", + " \"fluorescence\": (\n", + " \"Simulate physically motivated incoherent fluorescence image formation \"\n", + " \"for a uniformly fluorescent spherical particle using the configured \"\n", + " \"fluorescence wavelength and numerical aperture.\"\n", + " ),\n", + " \"brightfield\": (\n", + " \"Simulate brightfield imaging of a dielectric spherical particle using \"\n", + " \"physically motivated coherent sphere scattering and image formation. \"\n", + " \"Do not draw a Gaussian, ring, halo, Mexican-hat, or other hand-made \"\n", + " \"brightfield template.\"\n", + " ),\n", + " \"darkfield\": (\n", + " \"Simulate darkfield imaging of a dielectric spherical particle using \"\n", + " \"physically motivated sphere scattering and image formation in which \"\n", + " \"the unscattered illumination is rejected. Do not draw a modality-like \"\n", + " \"template.\"\n", + " ),\n", + " \"iscat\": (\n", + " \"Simulate interferometric scattering microscopy (iSCAT) of a dielectric \"\n", + " \"spherical particle. The detector image must arise from interference \"\n", + " \"between an appropriate reference field and the coherently scattered \"\n", + " \"field of the sphere. Do not draw rings or halos by hand.\"\n", + " ),\n", + " \"holography\": (\n", + " \"Simulate inline holography of a dielectric spherical particle using \"\n", + " \"coherent sphere scattering, propagation, and interference with the \"\n", + " \"reference illumination. Do not draw a hologram-like ring pattern by \"\n", + " \"hand.\"\n", + " ),\n", + "}\n", + "\n", + "task_description = MODALITY_TASKS[MODALITY]\n", + "print(task_description)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02e7ddc6", + "metadata": {}, + "outputs": [], + "source": [ + "def build_task_text(modality: str, description: str) -> str:\n", + " return textwrap.dedent(\n", + " f\"\"\"\n", + " # Task: reusable {modality} simulation pipeline\n", + "\n", + " Create only simulation.py.\n", + "\n", + " It must define:\n", + "\n", + " ~~~python\n", + " def build_pipeline(config):\n", + " ...\n", + " ~~~\n", + "\n", + " The returned callable must accept a particle dictionary and return one\n", + " two-dimensional NumPy-compatible image with shape\n", + " [image_size, image_size].\n", + "\n", + " The particle dictionary contains:\n", + "\n", + " - position_xy: Cartesian [x, y] coordinates in image pixels;\n", + " - axial_position_um;\n", + " - radius_um;\n", + " - refractive_index: [real, imaginary];\n", + " - fluorescence_intensity.\n", + "\n", + " The optical parameters are provided in task_config.json.\n", + "\n", + " {description}\n", + "\n", + " The pipeline must be reusable: do not hard-code the supplied particle\n", + " values into the implementation. It must work for a second particle\n", + " dictionary with different position, radius, refractive index, axial\n", + " position, and fluorescence intensity.\n", + "\n", + " Do not write plotting code, file I/O, argument parsing, training code,\n", + " or benchmark logic.\n", + " \"\"\"\n", + " ).strip() + \"\\n\"\n", + "\n", + "\n", + "ARM_INSTRUCTIONS = {\n", + " \"baseline\": (\n", + " \"Use standard scientific Python packages such as NumPy and SciPy as \"\n", + " \"needed. Do not import DeepTrack2. Implement the requested reusable \"\n", + " \"simulation pipeline in your own Python code.\"\n", + " ),\n", + " \"deeptrack2\": (\n", + " \"Use the installed DeepTrack2 package and its relevant scientific \"\n", + " \"abstractions where appropriate. Compose the requested optical model \"\n", + " \"with DeepTrack2 rather than reimplementing a library-provided model.\"\n", + " ),\n", + "}\n", + "\n", + "\n", + "def build_codex_prompt(task_text: str, arm: str) -> str:\n", + " return textwrap.dedent(\n", + " f\"\"\"\n", + " Complete one simulation-only scientific coding task in the current\n", + " directory.\n", + "\n", + " Read task.md and task_config.json. Create exactly one file named\n", + " simulation.py implementing the reusable interface requested by the\n", + " task. You may run small local tests if useful, but do not write plotting,\n", + " training, command-line, or result-file code. Do not install packages,\n", + " use the network, or access files outside the current workspace.\n", + "\n", + " The reader will later import simulation.py, call\n", + " build_pipeline(config), and evaluate the returned callable on both the\n", + " requested particle and an unseen particle.\n", + "\n", + " Condition:\n", + " {ARM_INSTRUCTIONS[arm]}\n", + "\n", + " Here is the complete task text:\n", + " ---------------- TASK ----------------\n", + " {task_text}\n", + " -------------- END TASK --------------\n", + " \"\"\"\n", + " ).strip() + \"\\n\"\n", + "\n", + "\n", + "def prepare_run(\n", + " modality: str,\n", + " arm: str,\n", + " run_label: str,\n", + " task_description_override: str | None = None,\n", + ") -> dict:\n", + " description = (\n", + " MODALITY_TASKS[modality]\n", + " if task_description_override is None\n", + " else str(task_description_override)\n", + " )\n", + " config = config_for(modality)\n", + " task_text = build_task_text(modality, description)\n", + " prompt = build_codex_prompt(task_text, arm)\n", + " run_dir = REPRO_ROOT / modality / arm / run_label\n", + " run_dir.mkdir(parents=True, exist_ok=True)\n", + " (run_dir / \"task.md\").write_text(task_text, encoding=\"utf-8\")\n", + " (run_dir / \"task_config.json\").write_text(\n", + " json.dumps(config, indent=2) + \"\\n\", encoding=\"utf-8\"\n", + " )\n", + " (run_dir / \"codex_prompt.txt\").write_text(prompt, encoding=\"utf-8\")\n", + " return {\n", + " \"run_dir\": run_dir,\n", + " \"config\": config,\n", + " \"task_text\": task_text,\n", + " \"prompt\": prompt,\n", + " }\n", + "\n", + "\n", + "selected = prepare_run(\n", + " MODALITY,\n", + " ARM,\n", + " RUN_LABEL,\n", + " task_description_override=task_description,\n", + ")\n", + "RUN_DIR = selected[\"run_dir\"]\n", + "config = selected[\"config\"]\n", + "print(f\"prepared run directory: {RUN_DIR}\")\n", + "print(f\"task prompt: {RUN_DIR / 'codex_prompt.txt'}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Connect to your own Codex CLI\n", + "\n", + "The next cell only defines a local subprocess wrapper. It does not send a\n", + "request until the following cell is run with RUN_CODEX = True.\n", + "\n", + "The wrapper uses codex exec with a workspace-write sandbox and saves the JSONL\n", + "transcript and stderr beside the generated simulator. Authentication is\n", + "handled by the user's existing Codex CLI sign-in." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CODEX_BIN = os.environ.get(\"CODEX_BIN\", \"codex\")\n", + "CODEX_MODEL = os.environ.get(\"CODEX_MODEL\") or None # leave None for your CLI default\n", + "CODEX_TIMEOUT_SECONDS = 1800\n", + "ALLOW_OVERWRITE_EXISTING = False\n", + "\n", + "\n", + "def codex_command() -> list[str]:\n", + " executable = shutil.which(CODEX_BIN)\n", + " if executable is None:\n", + " raise FileNotFoundError(\n", + " \"Codex CLI was not found. Install it, sign in from a terminal, \"\n", + " \"or set CODEX_BIN to its executable path.\"\n", + " )\n", + " command = [\n", + " executable,\n", + " \"exec\",\n", + " \"--sandbox\",\n", + " \"workspace-write\",\n", + " \"--skip-git-repo-check\",\n", + " \"--json\",\n", + " ]\n", + " if CODEX_MODEL:\n", + " command.extend([\"-m\", CODEX_MODEL])\n", + " command.append(\"-\")\n", + " return command\n", + "\n", + "\n", + "def run_codex(prompt: str, run_dir: Path) -> subprocess.CompletedProcess[str]:\n", + " simulation_path = run_dir / \"simulation.py\"\n", + " if simulation_path.exists() and not ALLOW_OVERWRITE_EXISTING:\n", + " raise FileExistsError(\n", + " f\"{simulation_path} already exists. Choose a new RUN_LABEL or set \"\n", + " \"ALLOW_OVERWRITE_EXISTING = True after reviewing the file.\"\n", + " )\n", + " result = subprocess.run(\n", + " codex_command(),\n", + " cwd=run_dir,\n", + " input=prompt,\n", + " text=True,\n", + " capture_output=True,\n", + " timeout=CODEX_TIMEOUT_SECONDS,\n", + " check=False,\n", + " )\n", + " (run_dir / \"codex.stdout.jsonl\").write_text(result.stdout, encoding=\"utf-8\")\n", + " (run_dir / \"codex.stderr.txt\").write_text(result.stderr, encoding=\"utf-8\")\n", + " if result.returncode != 0:\n", + " raise RuntimeError(\n", + " f\"Codex exited with status {result.returncode}. See \"\n", + " f\"{run_dir / 'codex.stderr.txt'}\"\n", + " )\n", + " return result\n", + "\n", + "\n", + "print(\"Codex executable:\", shutil.which(CODEX_BIN) or \"not found\")\n", + "print(\"Codex model:\", CODEX_MODEL or \"CLI default\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "RUN_CODEX = False # set to True to run Codex after signing in from a terminal\n", + "\n", + "if RUN_CODEX:\n", + " codex_result = run_codex(selected[\"prompt\"], RUN_DIR)\n", + " print(f\"Codex completed successfully in {RUN_DIR}\")\n", + " print(f\"transcript: {RUN_DIR / 'codex.stdout.jsonl'}\")\n", + "else:\n", + " print(\"Set RUN_CODEX = True and run this cell after signing in to Codex.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Load and execute the generated simulator\n", + "\n", + "Run the next cell after Codex has created simulation.py. The loader checks\n", + "the benchmark interface and the required image shape before rendering." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "simulation_path = RUN_DIR / \"simulation.py\"\n", + "\n", + "\n", + "def load_pipeline(path: Path, config: dict):\n", + " if not path.is_file():\n", + " raise FileNotFoundError(\n", + " f\"{path} does not exist. Set RUN_CODEX = True in the previous cell \"\n", + " \"or point simulation_path user simulation.py.\"\n", + " )\n", + " module_name = \"user_generated_\" + re.sub(r\"[^a-zA-Z0-9_]\", \"_\", str(path))\n", + " spec = importlib.util.spec_from_file_location(module_name, path)\n", + " if spec is None or spec.loader is None:\n", + " raise ImportError(f\"could not import {path}\")\n", + " module = importlib.util.module_from_spec(spec)\n", + " spec.loader.exec_module(module)\n", + " if not hasattr(module, \"build_pipeline\"):\n", + " raise AttributeError(\"simulation.py must define build_pipeline(config)\")\n", + " pipeline = module.build_pipeline(config)\n", + " if not callable(pipeline):\n", + " raise TypeError(\"build_pipeline(config) must return a callable\")\n", + " return pipeline\n", + "\n", + "\n", + "pipeline = load_pipeline(simulation_path, config)\n", + "print(f\"loaded: {simulation_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def image2d(value, size: int) -> np.ndarray:\n", + " image = np.squeeze(np.asarray(value))\n", + " if image.shape != (size, size):\n", + " raise ValueError(\n", + " f\"simulation returned shape {image.shape}; expected {(size, size)}\"\n", + " )\n", + " image = np.real_if_close(image)\n", + " if np.iscomplexobj(image):\n", + " image = np.real(image)\n", + " return np.asarray(image, dtype=float)\n", + "\n", + "\n", + "def benchmark_gray(image: np.ndarray) -> np.ndarray:\n", + " \"\"\"Match the benchmark's independent 1st-to-99th percentile scaling.\"\"\"\n", + " x = np.squeeze(np.real(np.asarray(image, dtype=float)))\n", + " finite = np.isfinite(x)\n", + " if not finite.any():\n", + " return np.zeros_like(x)\n", + " lo, hi = np.percentile(x[finite], [1, 99])\n", + " if hi <= lo:\n", + " lo, hi = float(x[finite].min()), float(x[finite].max())\n", + " if hi <= lo:\n", + " return np.zeros_like(x)\n", + " return np.clip((x - lo) / (hi - lo), 0, 1)\n", + "\n", + "\n", + "requested_image = image2d(pipeline(dict(particle)), int(config[\"image_size\"]))\n", + "reuse_image = image2d(pipeline(dict(reuse_particle)), int(config[\"image_size\"]))\n", + "\n", + "np.savez_compressed(\n", + " RUN_DIR / \"images.npz\",\n", + " requested_image=requested_image,\n", + " reuse_image=reuse_image,\n", + ")\n", + "(RUN_DIR / \"reuse_particle.json\").write_text(\n", + " json.dumps(reuse_particle, indent=2) + \"\\n\", encoding=\"utf-8\"\n", + ")\n", + "print(\"requested image:\", requested_image.shape, requested_image.dtype)\n", + "print(\"reuse image:\", reuse_image.shape, reuse_image.dtype)\n", + "print(f\"saved: {RUN_DIR / 'images.npz'}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(1, 2, figsize=(7, 3.5), constrained_layout=True)\n", + "for ax, image, item, title in (\n", + " (axes[0], requested_image, particle, \"requested particle\"),\n", + " (axes[1], reuse_image, reuse_particle, \"second particle\"),\n", + "):\n", + " ax.imshow(benchmark_gray(image), cmap=\"gray\", interpolation=\"nearest\")\n", + " ax.plot(\n", + " item[\"position_xy\"][0],\n", + " item[\"position_xy\"][1],\n", + " marker=\"+\",\n", + " markersize=9,\n", + " markeredgewidth=1.25,\n", + " color=\"tab:red\",\n", + " )\n", + " ax.set_title(title)\n", + " ax.axis(\"off\")\n", + "fig.suptitle(f\"{MODALITY} | {ARM}\")\n", + "fig.savefig(RUN_DIR / \"comparison.png\", dpi=200, bbox_inches=\"tight\")\n", + "plt.show()\n", + "print(f\"saved: {RUN_DIR / 'comparison.png'}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Prepare all five modality requests\n", + "\n", + "The following optional cell writes the five task prompts and configuration\n", + "files without calling Codex. To generate all five, review the prepared\n", + "directories and then invoke run_codex once per directory. Use a new\n", + "RUN_LABEL for another independent set of reader runs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prepared_runs = {\n", + " modality: prepare_run(modality, ARM, RUN_LABEL)\n", + " for modality in MODALITIES\n", + "}\n", + "\n", + "for modality, item in prepared_runs.items():\n", + " print(f\"{modality:12s} -> {item['run_dir']}\")\n", + "\n", + "# After reviewing the prompts, an explicit all-five generation loop is:\n", + "#\n", + "# for modality, item in prepared_runs.items():\n", + "# run_codex(item[\"prompt\"], item[\"run_dir\"])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv-deeptrack (3.14.6)", + "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.14.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}