Skip to content

feat: add experimental Colab Gemini notebook #45

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Sep 15, 2026
Merged

feat: add experimental Colab Gemini notebook #45

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
291 changes: 291 additions & 0 deletions notebooks/experimental_colab.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
Loading
Loading