Colab-first dependencies, vendored plotting helpers, Colab smoke tests

.github/workflows/colab-compat.yml

@@ -0,0 +1,44 @@
+name: Colab compatibility
+
+# Runs the tutorials inside Google's published Colab runtime image
+# (which lags the live Colab runtime by a few weeks), checking that
+# - colab_bootstrap.sh installs cleanly,
+# - it modifies no preinstalled package (=> no "restart runtime" prompt),
+# - both notebooks execute end-to-end.
+# See tests/colab/ and the dependency policy in pyproject.toml.
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 8 1 * *" # 1st of every month: catches Colab's upgrades
+  push:
+    branches: [master]
+    paths:
+      - requirements-colab.txt
+      - colab_bootstrap.sh
+      - notebooks/**
+      - tests/colab/**
+      - .github/workflows/colab-compat.yml
+  pull_request:
+    paths:
+      - requirements-colab.txt
+      - colab_bootstrap.sh
+      - notebooks/**
+      - tests/colab/**
+      - .github/workflows/colab-compat.yml
+
+jobs:
+  colab:
+    runs-on: ubuntu-latest
+    timeout-minutes: 90
+    container:
+      image: us-docker.pkg.dev/colab-images/public/runtime:latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Smoke test (bootstrap, freeze diff, execute notebooks)
+        env:
+          REPO_DIR: ${{ github.workspace }}
+          SMOKE_MODE: local
+        run: bash tests/colab/smoke_inner.sh

.pre-commit-config.yaml

@@ -35,7 +35,7 @@ repos:
           - tags
 
   - repo: https://github.com/mwouts/jupytext
-    rev: v1.15.1 # ensure equal to main venv
+    rev: v1.15.2 # keep equal to the `jupytext` pin in pyproject.toml
     hooks:
       - id: jupytext
         args: [--sync]

README.md

@@ -18,24 +18,20 @@ Use this option for development, or if you simply want faster computations
 *You could instead download & unzip, but then you will
 have to manually download any later updates.*
 
-**Install** (Python >= 3.12) with either
+**Install** (Python >= 3.12) with
+[uv](https://docs.astral.sh/uv/getting-started/installation/),
+a single, cross-platform tool that also fetches Python for you:
 
-- [uv](https://docs.astral.sh/uv/getting-started/installation/) -- recommended:
-  a single, cross-platform tool that also fetches Python for you.
+```bash
+uv sync
+uv run jupyter notebook
+```
 
-  ```bash
-  uv sync
-  uv run jupyter notebook
-  ```
-
-- Or any other installer, e.g. `pip` inside an active environment
-  created by [venv](https://docs.python.org/3/library/venv.html),
-  [conda](https://www.anaconda.com/download), ... :
-
-  ```bash
-  pip install .
-  jupyter notebook
-  ```
+*Without uv:* this project is not a package (nothing is imported from it),
+so there is nothing to `pip install .`. Instead, in an active environment
+(venv, conda, ...) do `pip install uv && uv pip install -r pyproject.toml`.
+Or, if you already have the scientific stack, just add what
+`requirements-colab.txt` lists.
 
 The `jupyter notebook` command opens a file navigator in your web browser.
 Click on `notebooks/HistoryMatch.ipynb`.
@@ -44,8 +40,7 @@ Click on `notebooks/HistoryMatch.ipynb`.
 
 The dev tooling (jupytext, ruff, pre-commit, ...) lives in the `dev`
 [dependency group](https://peps.python.org/pep-0735/),
-which `uv sync` installs by default
-(with `pip`: `pip install --group dev`).
+which `uv sync` installs by default.
 
 I prefer to develop mostly in the format of standard python script,
 which is why each notebook corresponds to a `.py` file synced via [jupytext](https://jupytext.readthedocs.io/en/latest/).
@@ -56,6 +51,36 @@ then the notebooks will get synced with the `.py` files before committing.
 Linting (which is, as of now, just a suggestion) can be run with
 `ruff check --output-format=grouped`.
 
+### Dependencies and Colab
+
+The single-click Colab experience is the design constraint for dependencies:
+Colab's base environment is not ours to control and gets upgraded regularly,
+and re-installing anything it has already imported (numpy, matplotlib, ...)
+forces a runtime restart.
+Therefore `colab_bootstrap.sh` installs `requirements-colab.txt` with `--no-deps`,
+i.e. only what Colab lacks, while `pyproject.toml` leaves Colab-preinstalled
+packages unconstrained and `uv.lock` provides reproducibility locally.
+The full reasoning is documented in `pyproject.toml`.
+The two files must be kept in sync by hand.
+
+### Smoke tests
+
+Google publishes the Colab runtime image, so the Colab experience can be tested
+without Colab. This runs monthly (and on relevant pushes) on GitHub Actions,
+see `.github/workflows/colab-compat.yml`, which is the primary way to catch
+breakage from Colab's upgrades. The same test can be run locally
+(requires `podman` or `docker`; the image is large, see the script):
+
+```bash
+tests/colab/smoke.sh            # local checkout
+tests/colab/smoke.sh --remote   # exactly what students get (GitHub master)
+```
+
+It checks that the bootstrap installs cleanly, that it modifies no
+preinstalled package, and that both notebooks execute end-to-end.
+The notebook runner also works locally, without the image:
+`uv run tests/colab/run_nb.py notebooks/HistoryMatch.ipynb --skip-bootstrap`.
+
 ## Contributors
 
 This work has been developed by *Patrick N. Raanes*, researcher at *NORCE*.

colab_bootstrap.sh

@@ -1,42 +1,60 @@
 #!/usr/bin/env bash
 
-# Colab doesn't provide
-# - Auto-installing dependencies (neither pyproject.toml nor requirements.txt)
-# - Pre-loading data/modules (aside from the notebook itself)
-# This script takes care of the above by cloning the full (shallow) repo.
+# Bootstrap for Google Colab, invoked from the first code cell of each notebook:
+#
+#     !wget -qO- https://raw.githubusercontent.com/patnr/HistoryMatching/master/colab_bootstrap.sh | bash -s
+#
+# Colab does not auto-install a project's dependencies (no pyproject/requirements
+# support), nor does it fetch anything but the notebook itself. So this script
+# 1. shallow-clones the repo (or copies it from `--repo=PATH`, for tests),
+# 2. installs ONLY the packages Colab lacks, with `--no-deps`, so that nothing
+#    preinstalled (and already imported) gets touched, which would otherwise
+#    require a runtime restart. See requirements-colab.txt and pyproject.toml.
+# 3. copies the notebooks dir (incl. `tools/`) into the working directory.
+#
+# Options: `--debug`/`-v` for verbose output; `--branch=NAME` (default: master).
 
 main () {
     set -e
 
-    # Clear any existing REPO for a fresh git clone
+    # Clear any existing REPO for a fresh copy
     rm -rf REPO
 
-    # Download repo
-    URL=https://github.com/patnr/HistoryMatching.git
-    if [[ ! -d REPO ]]; then git clone --depth=1 $URL REPO; fi
-
-    # https://pythonspeed.com/articles/upgrade-pip/
-    pip install --upgrade pip
+    if [[ -n "$SRC" ]]; then
+        # Local checkout (used by tests/colab/smoke.sh)
+        mkdir REPO
+        tar -C "$SRC" --exclude=.git -cf - . | tar -C REPO -xf -
+    else
+        git clone --depth=1 --branch "$BRANCH" https://github.com/patnr/HistoryMatching.git REPO
+    fi
 
-    # Install dependencies (declared in REPO/pyproject.toml)
-    pip install ./REPO
+    # Install what Colab lacks. NB: no pip upgrade, no dependency resolution.
+    pip install --no-deps -r REPO/requirements-colab.txt
 
-    # Put repo contents (including hidden files) in PWD
+    # Put notebooks (including hidden files) in PWD
     shopt -s dotglob
     cp -r REPO/notebooks/* ./
 }
 
-# Only run if we're on colab
-if python -c "import google.colab" 2>/dev/null; then
-
-    # Use `bash -s -- --debug` to get verbose output
-    if echo $@ | grep -E -- '(--debug|-v)' > /dev/null ; then
+# Parse args
+SRC=""
+BRANCH="master"
+DEBUG=""
+for arg in "$@"; do
+    case "$arg" in
+        --repo=*)   SRC="${arg#--repo=}" ;;
+        --branch=*) BRANCH="${arg#--branch=}" ;;
+        --debug|-v) DEBUG=1 ;;
+    esac
+done
+
+# Only run if we're on Colab (the env var is also set in Colab's runtime image)
+if [[ -n "$COLAB_RELEASE_TAG" ]] || python -c "import google.colab" 2>/dev/null; then
+    if [[ -n "$DEBUG" ]]; then
         main
     else
-        # Quiet
         main > /dev/null 2>&1
     fi
-
     echo "Initialization for Colab done."
 else
     echo "Not running on Colab => Didn't do anything."

notebooks/tools/geostat.py

@@ -3,9 +3,10 @@
 import numpy as np
 import scipy.linalg as sla
 from matplotlib import pyplot as plt
-from mpl_tools.misc import nRowCol
 from numpy.random import randn
 
+from tools.plotting import nRowCol
+
 
 def variogram_gauss(xx, r, n=0, a=1 / 3):
     """Compute the Gaussian variogram for the 1D points xx.