From d69eef1e2a70f25c5d2cac8a4181e7e1748bbb14 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 15 Sep 2026 08:22:54 +0100 Subject: [PATCH] feat: add experimental Colab Gemini notebook Co-Authored-By: Claude --- README.md | 14 ++ notebooks/experimental_colab.ipynb | 291 +++++++++++++++++++++++++++++ scripts/experimental_colab.py | 202 ++++++++++++++++++++ 3 files changed, 507 insertions(+) create mode 100644 notebooks/experimental_colab.ipynb create mode 100644 scripts/experimental_colab.py diff --git a/README.md b/README.md index bed9f58..984f24a 100644 --- a/README.md +++ b/README.md @@ -466,6 +466,20 @@ the priors, picking the non-linear search, and how to read the posterior. So I come away understanding the workflow, not just the commands. ``` +## Experimental: Google Colab + Gemini + +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PyAutoLabs/autofit_assistant/blob/main/notebooks/experimental_colab.ipynb) + +Try the assistant's bundled Gaussian fit in a fresh CPU runtime, with context to +copy into Colab's Gemini sidebar. The notebook installs PyAutoFit, introduces the +model and likelihood, runs a small fit, and lets you download the results. + +This is experimental: cloning the assistant does not automatically load its +instructions into Gemini. Follow the notebook's explicit context step. Gemini +availability depends on your Google account and Colab settings. Save a notebook +copy and download results before the runtime resets; this is a small browser +experiment, with larger inference projects continuing in the full assistant. + ## License The assistant ships agent instructions and reference material derived from the public diff --git a/notebooks/experimental_colab.ipynb b/notebooks/experimental_colab.ipynb new file mode 100644 index 0000000..a5f6399 --- /dev/null +++ b/notebooks/experimental_colab.ipynb @@ -0,0 +1,291 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Experimental: PyAutoFit assistant in Colab\n", + "\n", + "Try a small inference problem alongside Colab's Gemini sidebar. Use a **fresh CPU\n", + "runtime**, then run the cells in order. Installation can take a few minutes.\n", + "\n", + "This is a notebook experiment, not the full CLI assistant or an HPC service.\n", + "Cloning the repository does **not** make Gemini read its instructions. Step 2\n", + "provides context to copy into the sidebar explicitly. Gemini availability depends\n", + "on your Google account and Colab settings; the fit also works without Gemini.\n", + "\n", + "Save a copy of this notebook in Drive to keep your edits. Runtime files disappear\n", + "when Colab resets: download your results at the end.\n", + "\n", + "__Contents__\n", + "\n", + "1. Setup\n", + "2. Gemini context\n", + "3. Data inspection\n", + "4. Model and likelihood\n", + "5. Fit\n", + "6. Download and continue" + ], + "id": "cell-0" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "from pathlib import Path\n", + "import subprocess\n", + "import sys\n", + "\n", + "# The public release includes the model, likelihood and assistant context used here.\n", + "ASSISTANT_REF = \"2026.9.15.1\"\n", + "STACK_VERSION = \"2026.9.15.1\"\n", + "COLAB_HOME = Path(\"/content\")\n", + "repo = COLAB_HOME / \"autofit_assistant\"\n", + "if not repo.exists():\n", + " subprocess.run(\n", + " [\"git\", \"clone\", \"--depth\", \"1\", \"--branch\", ASSISTANT_REF,\n", + " \"https://github.com/PyAutoLabs/autofit_assistant.git\", str(repo)],\n", + " check=True,\n", + " )\n", + "else:\n", + " actual = subprocess.check_output(\n", + " [\"git\", \"-C\", str(repo), \"describe\", \"--tags\", \"--exact-match\"], text=True\n", + " ).strip()\n", + " if actual != ASSISTANT_REF:\n", + " raise RuntimeError(\"This folder has another assistant version; use a fresh runtime.\")\n", + "\n", + "subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", f\"autofit=={STACK_VERSION}\",\n", + " f\"autonerves=={STACK_VERSION}\"], check=True,\n", + ")\n", + "os.environ.setdefault(\"NUMBA_CACHE_DIR\", \"/tmp/numba_cache\")\n", + "os.environ.setdefault(\"MPLCONFIGDIR\", \"/tmp/matplotlib\")\n", + "os.chdir(repo)\n", + "sys.path.insert(0, str(repo))\n", + "\n", + "import autofit as af\n", + "import autonerves\n", + "\n", + "print(\"Assistant:\", subprocess.check_output(\n", + " [\"git\", \"rev-parse\", \"HEAD\"], text=True).strip())\n", + "print(\"PyAutoFit:\", af.__version__, \"PyAutoNerves:\", autonerves.__version__)" + ], + "execution_count": null, + "outputs": [], + "id": "cell-1" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Give Gemini the assistant context\n", + "\n", + "Run this cell, then **copy its printed text into the Gemini sidebar**. Ask Gemini\n", + "to acknowledge the supplied context before discussing the next cell. The context\n", + "comes from this pinned assistant checkout; no API key or Gemini SDK is needed.\n", + "\n", + "Only the text you supply is guaranteed to be part of that conversation. If Gemini\n", + "needs another skill, open that file in Colab's Files panel and paste the relevant\n", + "section. Sidebar access to notebook content can vary: include the cell or error\n", + "in your message when needed." + ], + "id": "cell-2" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "context_paths = [\n", + " \"AGENTS.md\",\n", + " \"modes/teacher.md\",\n", + " \"skills/af_compose_model.md\",\n", + " \"skills/af_wrap_likelihood.md\",\n", + "]\n", + "bootstrap = (\n", + " \"Help me work through this experimental PyAutoFit notebook in teacher mode. \"\n", + " \"Use only notebook-relevant guidance from the reference material below and \"\n", + " \"the code I share. Skip CLI onboarding, filesystem checks, project-memory \"\n", + " \"creation and maintainer workflows. Ask one question at \"\n", + " \"a time; give hints before solutions. Do not claim to have read runtime files \"\n", + " \"or run code unless you actually have. Preserve the supplied likelihood's \"\n", + " \"numerics. Before adapting this example to real data, help me inspect the \"\n", + " \"data and discuss artefacts, outliers and selection effects. Explain proposed \"\n", + " \"changes before I run them. If an API is uncertain, ask me to inspect it.\\n\"\n", + ")\n", + "context = bootstrap + \"\\nAssistant reference: \" + ASSISTANT_REF + \"\\n\"\n", + "for relative in context_paths:\n", + " context += \"\\n--- \" + relative + \" ---\\n\" + Path(relative).read_text()\n", + "print(context)" + ], + "execution_count": null, + "outputs": [], + "id": "cell-3" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Inspect the bundled simulated data\n", + "\n", + "The assistant's Gaussian is **peak-normalised**: normalization is its height,\n", + "not its integrated area. The simulation used centre 50, height 25 and width 10.\n", + "Its likelihood assumes independent Gaussian measurement errors and omits the\n", + "constant noise-normalisation term. Evidence values therefore use that convention.\n", + "\n", + "Ask Gemini: *What features of this plot should constrain the three parameters?*\n", + "The model and likelihood below are imported unchanged from the assistant's tour." + ], + "id": "cell-4" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "sys.path.insert(0, str(repo / \"scripts\" / \"start_here\"))\n", + "from gaussian import Gaussian\n", + "from analysis import Analysis\n", + "\n", + "data = np.asarray(json.loads(Path(\"dataset/gaussian_x1/data.json\").read_text()))\n", + "noise_map = np.asarray(json.loads(Path(\"dataset/gaussian_x1/noise_map.json\").read_text()))\n", + "x = np.arange(data.size)\n", + "plt.errorbar(x, data, yerr=noise_map, fmt=\".k\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"Profile amplitude\")\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [], + "id": "cell-5" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Compose the model and check the likelihood\n", + "\n", + "Ask Gemini: *Explain these priors. What would happen if I fixed the width?*\n", + "The check below compares the imported likelihood with its written equation." + ], + "id": "cell-6" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = af.Model(\n", + " Gaussian,\n", + " centre=af.UniformPrior(lower_limit=0.0, upper_limit=100.0),\n", + " normalization=af.UniformPrior(lower_limit=0.0, upper_limit=100.0),\n", + " sigma=af.UniformPrior(lower_limit=0.1, upper_limit=30.0),\n", + ")\n", + "analysis = Analysis(data=data, noise_map=noise_map)\n", + "truth = Gaussian(centre=50.0, normalization=25.0, sigma=10.0)\n", + "expected = -0.5 * np.sum(((data - truth.model_data_from(x)) / noise_map) ** 2)\n", + "assert np.isclose(analysis.log_likelihood_function(truth), expected)\n", + "print(model.info)\n", + "print(\"Log likelihood at simulation inputs:\", expected)" + ], + "execution_count": null, + "outputs": [], + "id": "cell-7" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Run a small fit\n", + "\n", + "Run this cell when ready. It performs actual nested sampling with 50 live points,\n", + "chosen for experimentation, not precision results. Runtime varies. The same\n", + "completed run reloads if you execute this cell again. Change `name` for a new\n", + "experiment, especially after changing the data or likelihood.\n", + "\n", + "Ask Gemini: *Before I run this, what do you expect the recovered parameters to be?*" + ], + "id": "cell-8" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.DynestyStatic(\n", + " path_prefix=\"experimental_colab\", name=\"gaussian_x1\", nlive=50,\n", + " number_of_cores=1,\n", + ")\n", + "result = search.fit(model=model, analysis=analysis)\n", + "print(result.info)\n", + "best = result.max_log_likelihood_instance\n", + "plt.errorbar(x, data, yerr=noise_map, fmt=\".k\", label=\"data\")\n", + "plt.plot(x, best.model_data_from(x), label=\"maximum likelihood fit\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"Profile amplitude\")\n", + "plt.legend()\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [], + "id": "cell-9" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Keep the results and continue\n", + "\n", + "Ask Gemini: *Help me explain the parameter uncertainties and one limitation of\n", + "this fit. Then suggest one change I can predict and test.* Share `result.info`\n", + "with your question. The fit and residual figures are also under\n", + "`output/experimental_colab/` in the Files panel.\n", + "\n", + "Run this last cell to download the fit outputs plus the model, likelihood, data,\n", + "configuration and context used here. Save your notebook copy separately: the\n", + "archive does not contain your edited notebook or Gemini conversation. Download\n", + "before disconnecting. For larger inference projects, continue in a local clone\n", + "of [autofit_assistant](https://github.com/PyAutoLabs/autofit_assistant), carrying\n", + "your scientific question, priors, likelihood and results with you." + ], + "id": "cell-10" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import zipfile\n", + "from google.colab import files\n", + "\n", + "archive = repo.parent / \"autofit_colab_results.zip\"\n", + "with zipfile.ZipFile(archive, \"w\", zipfile.ZIP_DEFLATED) as bundle:\n", + " for directory in [\"output/experimental_colab\", \"scripts/start_here\",\n", + " \"dataset/gaussian_x1\", \"config\"]:\n", + " for path in Path(directory).rglob(\"*\"):\n", + " if path.is_file() and \"__pycache__\" not in path.parts:\n", + " bundle.write(path, str(path))\n", + " bundle.writestr(\"gemini_context.txt\", context)\n", + " bundle.writestr(\"environment.txt\", subprocess.check_output(\n", + " [sys.executable, \"-m\", \"pip\", \"freeze\"], text=True))\n", + " bundle.writestr(\"assistant_ref.txt\", ASSISTANT_REF + \"\\n\" +\n", + " subprocess.check_output([\"git\", \"rev-parse\", \"HEAD\"], text=True))\n", + "files.download(str(archive))" + ], + "execution_count": null, + "outputs": [], + "id": "cell-11" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/experimental_colab.py b/scripts/experimental_colab.py new file mode 100644 index 0000000..6c764b1 --- /dev/null +++ b/scripts/experimental_colab.py @@ -0,0 +1,202 @@ +""" +# Experimental: PyAutoFit assistant in Colab + +Try a small inference problem alongside Colab's Gemini sidebar. Use a **fresh CPU +runtime**, then run the cells in order. Installation can take a few minutes. + +This is a notebook experiment, not the full CLI assistant or an HPC service. +Cloning the repository does **not** make Gemini read its instructions. Step 2 +provides context to copy into the sidebar explicitly. Gemini availability depends +on your Google account and Colab settings; the fit also works without Gemini. + +Save a copy of this notebook in Drive to keep your edits. Runtime files disappear +when Colab resets: download your results at the end. + +__Contents__ + +1. Setup +2. Gemini context +3. Data inspection +4. Model and likelihood +5. Fit +6. Download and continue +""" + +import os +from pathlib import Path +import subprocess +import sys + +# The public release includes the model, likelihood and assistant context used here. +ASSISTANT_REF = "2026.9.15.1" +STACK_VERSION = "2026.9.15.1" +COLAB_HOME = Path("/content") +repo = COLAB_HOME / "autofit_assistant" +if not repo.exists(): + subprocess.run( + ["git", "clone", "--depth", "1", "--branch", ASSISTANT_REF, + "https://github.com/PyAutoLabs/autofit_assistant.git", str(repo)], + check=True, + ) +else: + actual = subprocess.check_output( + ["git", "-C", str(repo), "describe", "--tags", "--exact-match"], text=True + ).strip() + if actual != ASSISTANT_REF: + raise RuntimeError("This folder has another assistant version; use a fresh runtime.") + +subprocess.run( + [sys.executable, "-m", "pip", "install", f"autofit=={STACK_VERSION}", + f"autonerves=={STACK_VERSION}"], check=True, +) +os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") +os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") +os.chdir(repo) +sys.path.insert(0, str(repo)) + +import autofit as af +import autonerves + +print("Assistant:", subprocess.check_output( + ["git", "rev-parse", "HEAD"], text=True).strip()) +print("PyAutoFit:", af.__version__, "PyAutoNerves:", autonerves.__version__) + +""" +## 2. Give Gemini the assistant context + +Run this cell, then **copy its printed text into the Gemini sidebar**. Ask Gemini +to acknowledge the supplied context before discussing the next cell. The context +comes from this pinned assistant checkout; no API key or Gemini SDK is needed. + +Only the text you supply is guaranteed to be part of that conversation. If Gemini +needs another skill, open that file in Colab's Files panel and paste the relevant +section. Sidebar access to notebook content can vary: include the cell or error +in your message when needed. +""" + +context_paths = [ + "AGENTS.md", + "modes/teacher.md", + "skills/af_compose_model.md", + "skills/af_wrap_likelihood.md", +] +bootstrap = ( + "Help me work through this experimental PyAutoFit notebook in teacher mode. " + "Use only notebook-relevant guidance from the reference material below and " + "the code I share. Skip CLI onboarding, filesystem checks, project-memory " + "creation and maintainer workflows. Ask one question at " + "a time; give hints before solutions. Do not claim to have read runtime files " + "or run code unless you actually have. Preserve the supplied likelihood's " + "numerics. Before adapting this example to real data, help me inspect the " + "data and discuss artefacts, outliers and selection effects. Explain proposed " + "changes before I run them. If an API is uncertain, ask me to inspect it.\n" +) +context = bootstrap + "\nAssistant reference: " + ASSISTANT_REF + "\n" +for relative in context_paths: + context += "\n--- " + relative + " ---\n" + Path(relative).read_text() +print(context) + +""" +## 3. Inspect the bundled simulated data + +The assistant's Gaussian is **peak-normalised**: normalization is its height, +not its integrated area. The simulation used centre 50, height 25 and width 10. +Its likelihood assumes independent Gaussian measurement errors and omits the +constant noise-normalisation term. Evidence values therefore use that convention. + +Ask Gemini: *What features of this plot should constrain the three parameters?* +The model and likelihood below are imported unchanged from the assistant's tour. +""" + +import json +import numpy as np +import matplotlib.pyplot as plt +sys.path.insert(0, str(repo / "scripts" / "start_here")) +from gaussian import Gaussian +from analysis import Analysis + +data = np.asarray(json.loads(Path("dataset/gaussian_x1/data.json").read_text())) +noise_map = np.asarray(json.loads(Path("dataset/gaussian_x1/noise_map.json").read_text())) +x = np.arange(data.size) +plt.errorbar(x, data, yerr=noise_map, fmt=".k") +plt.xlabel("x") +plt.ylabel("Profile amplitude") +plt.show() + +""" +## 4. Compose the model and check the likelihood + +Ask Gemini: *Explain these priors. What would happen if I fixed the width?* +The check below compares the imported likelihood with its written equation. +""" + +model = af.Model( + Gaussian, + centre=af.UniformPrior(lower_limit=0.0, upper_limit=100.0), + normalization=af.UniformPrior(lower_limit=0.0, upper_limit=100.0), + sigma=af.UniformPrior(lower_limit=0.1, upper_limit=30.0), +) +analysis = Analysis(data=data, noise_map=noise_map) +truth = Gaussian(centre=50.0, normalization=25.0, sigma=10.0) +expected = -0.5 * np.sum(((data - truth.model_data_from(x)) / noise_map) ** 2) +assert np.isclose(analysis.log_likelihood_function(truth), expected) +print(model.info) +print("Log likelihood at simulation inputs:", expected) + +""" +## 5. Run a small fit + +Run this cell when ready. It performs actual nested sampling with 50 live points, +chosen for experimentation, not precision results. Runtime varies. The same +completed run reloads if you execute this cell again. Change `name` for a new +experiment, especially after changing the data or likelihood. + +Ask Gemini: *Before I run this, what do you expect the recovered parameters to be?* +""" + +search = af.DynestyStatic( + path_prefix="experimental_colab", name="gaussian_x1", nlive=50, + number_of_cores=1, +) +result = search.fit(model=model, analysis=analysis) +print(result.info) +best = result.max_log_likelihood_instance +plt.errorbar(x, data, yerr=noise_map, fmt=".k", label="data") +plt.plot(x, best.model_data_from(x), label="maximum likelihood fit") +plt.xlabel("x") +plt.ylabel("Profile amplitude") +plt.legend() +plt.show() + +""" +## 6. Keep the results and continue + +Ask Gemini: *Help me explain the parameter uncertainties and one limitation of +this fit. Then suggest one change I can predict and test.* Share `result.info` +with your question. The fit and residual figures are also under +`output/experimental_colab/` in the Files panel. + +Run this last cell to download the fit outputs plus the model, likelihood, data, +configuration and context used here. Save your notebook copy separately: the +archive does not contain your edited notebook or Gemini conversation. Download +before disconnecting. For larger inference projects, continue in a local clone +of [autofit_assistant](https://github.com/PyAutoLabs/autofit_assistant), carrying +your scientific question, priors, likelihood and results with you. +""" + +import zipfile +from google.colab import files + +archive = repo.parent / "autofit_colab_results.zip" +with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as bundle: + for directory in ["output/experimental_colab", "scripts/start_here", + "dataset/gaussian_x1", "config"]: + for path in Path(directory).rglob("*"): + if path.is_file() and "__pycache__" not in path.parts: + bundle.write(path, str(path)) + bundle.writestr("gemini_context.txt", context) + bundle.writestr("environment.txt", subprocess.check_output( + [sys.executable, "-m", "pip", "freeze"], text=True)) + bundle.writestr("assistant_ref.txt", ASSISTANT_REF + "\n" + + subprocess.check_output(["git", "rev-parse", "HEAD"], text=True)) +files.download(str(archive))