Working calibration example

CALIB/Dockerfile

@@ -0,0 +1,23 @@
+# Use Rocky Linux as base
+FROM rockylinux:9
+
+# Install Python 3.12 and dependencies
+RUN python3 -m ensurepip && \
+    python3 -m pip install --upgrade pip
+
+# Set the working directory inside the container
+WORKDIR /app
+
+# Copy the requirements file and install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy your script into the container
+COPY laser.py .
+COPY objective.py .
+COPY worker.py .
+COPY mysql_storage.py .
+
+# Run the script when the container starts
+#CMD ["python3", "laser.py", "-p", "shared/params.json", "-o", "shared/simulation_results.csv"]
+CMD ["python3", "worker.py"]

CALIB/README.rst

@@ -0,0 +1,122 @@
+LASER Calibration with Optuna
+=============================
+
+This repository demonstrates how to calibrate a spatial **Susceptible-Infected-Recovered (SIR) model** using **Optuna**, an optimization framework for hyperparameter tuning.
+
+Our goal is to identify **transmission rate** and **migration rate** values that produce a **second prevalence peak** in **node 0**, with a noticeable **trough between the peaks**.
+
+Model Overview
+--------------
+
+The core model (``laser.py``) implements a spatial **SIR model** with the following characteristics:
+
+- **500,000 agents** spread **heterogeneously** across **20 nodes**.
+- **Node 0** is the **largest**, while **node 19** is the **smallest**.
+- The outbreak is **seeded in 1%** of the **population of node 0**.
+- **Gravity-based migration** determines agent movement between nodes.
+- Infections last **5 to 15 days** (uniformly distributed).
+- **Configurable** transmission and migration rates.
+
+Calibration Goal
+----------------
+
+We use **Optuna** to optimize the **transmission rate** and **migration rate** to achieve:
+
+✅ A **second prevalence peak** in node 0.
+✅ A **clear trough** between the two peaks.
+
+The calibration process runs multiple simulations, adjusting parameters until the desired epidemic curve is achieved.
+
+Files and Structure
+-------------------
+
+- ``laser.py`` – Implements the SIR model.
+- ``run.py`` – Main calibration script (starts the Optuna process).
+- ``objective.py`` – Defines the **objective function**, evaluating how well each trial matches the target epidemic curve.
+- ``impatient.py`` – Allows you to inspect the **current best parameters** while calibration is still running.
+- ``optuna_review_contour.py`` – Generates **Optuna visualizations** to analyze the search process.
+
+Installation
+------------
+
+Before running the calibration, install the required dependencies:
+
+.. code-block:: bash
+
+    pip install -r requirements.txt
+
+Running Calibration
+-------------------
+
+To start calibration, run:
+
+.. code-block:: bash
+
+    python run.py
+
+By default, this will:
+
+- Run **100 trials** with **4 replicates** per trial.
+- Simulate **500 timesteps per run** (each taking ~10 seconds).
+- Identify the best parameter set and run **a final simulation** with those values.
+
+Checking Calibration Progress
+-----------------------------
+
+To monitor the best parameters found so far, run:
+
+.. code-block:: bash
+
+    python impatient.py
+
+To visualize the parameter search space explored by Optuna, run:
+
+.. code-block:: bash
+
+    python optuna_review_contour.py
+
+Expected Results
+----------------
+
+If calibration is successful, the final prevalence plot for **node 0** should display:
+
+✅ A **clear second peak** in infections.
+✅ A **noticeable trough** between the two peaks.
+
+You can modify parameters in the scripts to explore different calibration behaviors.
+
+Next Steps
+----------
+
+- Try adjusting the search space or evaluation criteria in ``objective.py``.
+- Increase the number of trials to improve calibration accuracy.
+- Experiment with different outbreak seeding strategies.
+
+Summary
+-------
+
+This project demonstrates **LASER and Optuna-based epidemic model calibration** and is designed for researchers interested in large scale spatial disease modeling and parameter estimation.
+
+Dockerized
+----------
+
+Network Start
+^^^^^^^^^^^^^
+
+docker network create optuna-network
+
+DB Start
+^^^^^^^^
+
+docker run -d --name optuna-mysql --network optuna-network -p 3306:3306 -e MYSQL_ALLOW_EMPTY_PASSWORD=yes -e MYSQL_DATABASE=optuna_db mysql:latest
+
+Optuna Workers
+^^^^^^^^^^^^^^
+
+docker build -t docker.io/library/calib-worker:latest
+docker run --rm --name calib-worker --network optuna-network -e STORAGE_URL="mysql+pymysql://root@optuna-mysql/mysql" docker.io/library/calib-worker:latest&
+
+View Results
+^^^^^^^^^^^^
+
+python3 impatient_mysqldocker.py

CALIB/calib_arch_roles_steps2.dot

@@ -0,0 +1,69 @@
+digraph OptunaCalibrationWorkflow {
+    rankdir=TB;
+    node [shape=box, fontname="Arial"];
+
+    # Define roles (placed outside of steps)
+    modeler [label="Modeler (User)", style=filled, fillcolor=lightblue, shape=ellipse];
+    developer [label="Developer", style=filled, fillcolor=lightgreen, shape=ellipse];
+    it_admin [label="IT Admin", style=filled, fillcolor=lightgray, shape=ellipse];
+
+    # Step 1: Model Definition (Only Modeler)
+    subgraph cluster_step1 {
+        label="1️⃣ Model Definition (Python Only)";
+        style=dashed;
+
+        model_definition [label="Define Model Code (Python)", style=filled, fillcolor=lightblue];
+        model_file [label="model.py", shape=parallelogram, style=filled, fillcolor=white];
+
+        model_definition -> model_file [label="Writes"];
+    }
+    modeler -> model_definition [label="Defines Model"];
+
+    # Step 2: Local Calibration (Modeler + Developer)
+    subgraph cluster_step2 {
+        label="2️⃣ Local Calibration (Docker)";
+        style=dashed;
+
+        objective_function [label="Define Calibration Inputs/Outputs\n(objective.py)", style=filled, fillcolor=lightblue];
+        local_run [label="Run Local Calibration (`run.py`)", style=filled, fillcolor=lightblue];
+
+        local_workers [label="Local Worker Containers\n(laser.py + Optuna)", style=filled, fillcolor=lightgreen];
+        local_db [label="Local Database\n(SQLite/MySQL)", shape=cylinder, style=filled, fillcolor=lightgray];
+        local_storage [label="Local Persistent Storage\n(Mounted Volume)", shape=cylinder, style=filled, fillcolor=gold];
+
+        objective_function -> local_run [label="Uses"];
+        local_run -> local_workers [label="Spawns Workers"];
+        local_workers -> local_db [label="Store Results"];
+        local_db -> local_storage [label="File I/O", dir=both];
+    }
+    modeler -> objective_function [label="Defines Inputs/Outputs"];
+    modeler -> local_run [label="Runs Calibration Locally"];
+    developer -> local_workers [label="Supports Debugging & Docker"];
+
+    # Step 3: Remote Calibration (Modeler + Developer + IT Admin)
+    subgraph cluster_step3 {
+        label="3️⃣ Remote Calibration (K8s)";
+        style=dashed;
+
+        remote_run [label="Run Full Calibration (`run.py`)\n(No Extra Config)", style=filled, fillcolor=lightblue];
+        monitor_results [label="Monitor Progress & Plots", style=filled, fillcolor=lightblue];
+
+        cloud_controller [label="Controller (K8s Pod) + Optuna", style=filled, fillcolor=lightblue];
+        cloud_workers [label="Worker Pods (Auto-Scaling)", style=filled, fillcolor=lightgreen];
+        cloud_db [label="Cloud Database\n(PostgreSQL/MySQL)", shape=cylinder, style=filled, fillcolor=lightgray];
+        cloud_storage [label="Cloud Storage\n(Azure Blob / S3 / GCS)", shape=cylinder, style=filled, fillcolor=gold];
+
+        remote_run -> cloud_controller [label="Runs on Cluster"];
+        cloud_controller -> cloud_workers [label="Manages Workers"];
+        cloud_workers -> cloud_db [label="Store Results"];
+        cloud_db -> cloud_storage [label="Persistent Data", dir=both];
+
+        remote_run -> monitor_results [label="Check Progress"];
+    }
+    modeler -> remote_run [label="Runs Calibration on Cluster"];
+    modeler -> monitor_results [label="Monitors Progress"];
+
+    developer -> cloud_workers [label="Supports Debugging"];
+    it_admin -> cloud_controller [label="Configures K8s & Cloud"];
+    it_admin -> cloud_storage [label="Configures"];
+}

CALIB/docker-compose.yml

@@ -0,0 +1,28 @@
+version: '3.8'
+
+services:
+  mysql:
+    image: mysql:latest
+    container_name: optuna-mysql
+    restart: always
+    environment:
+      MYSQL_ROOT_PASSWORD: root
+      MYSQL_DATABASE: optuna_db
+      MYSQL_USER: user
+      MYSQL_PASSWORD: password
+    ports:
+      - "3306:3306"
+    networks:
+      - optuna_network
+
+  worker:
+    build: .
+    depends_on:
+      - mysql
+    environment:
+      - STORAGE_URL=mysql+pymysql://user:password@optuna-mysql/optuna_db
+    networks:
+      - optuna_network
+
+networks:
+  optuna_network:

CALIB/howto_local_docker.txt

@@ -0,0 +1,12 @@
+# cleanup maybe
+docker stop optuna-mysql && docker rm optuna-mysql
+
+# network start
+docker network create optuna-network
+# db start
+docker run -d --name optuna-mysql --network optuna-network -p 3306:3306 -e MYSQL_ALLOW_EMPTY_PASSWORD=yes -e MYSQL_DATABASE=optuna_db mysql:latest
+# optuna workers
+docker run --rm --name calib-worker --network optuna-network -e STORAGE_URL="mysql+pymysql://root@optuna-mysql/mysql" docker.io/library/calib-worker:latest&
+
+# peak results
+python3 impatient_mysqldocker.py

CALIB/impatient.py

@@ -0,0 +1,26 @@
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import optuna
+from mysql_storage import get_storage_url
+
+# Define the storage URL
+# storage_url = "sqlite:///optuna_study.db"
+storage_url = get_storage_url()
+study = optuna.create_study(direction="minimize", storage=storage_url, study_name="spatial_demo_calibr8n", load_if_exists=True)
+
+# Get the best parameters found so far
+best_params = study.best_params
+print("Best parameters so far:", best_params)
+
+# Add fixed parameters
+best_params.update({"population_size": 500000, "nodes": 20, "timesteps": 500, "initial_infected_fraction": 0.01})
+
+# Save best parameters
+Path("params_test.json").write_text(json.dumps(best_params, indent=4))
+print("Saved best parameters to params_test.json.")
+
+# Run the model
+subprocess.run([sys.executable, "laser.py", "--plot", "-o", "simulation_result_test.csv", "-p", "params_test.json"], check=True)