Facilitate matrix simulations

The backend allows for n x m number of values and components to vary.
The frontend utilises this to make a 1D grid to vary single component
over a linear range.

Author: lars-petter-hauge

backend/alembic/versions/a7f3c9d21b84_add_sweeps.py

@@ -0,0 +1,35 @@
+"""add grid simulations
+
+Revision ID: a7f3c9d21b84
+Revises: 7c2e1f4b8a90
+Create Date: 2026-06-12 00:00:00.000000
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+
+revision: str = "a7f3c9d21b84"
+down_revision: Union[str, Sequence[str], None] = "7c2e1f4b8a90"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "grid_simulations",
+        sa.Column("owner_id", sa.Uuid(), nullable=True),
+        sa.Column("axes", sa.JSON(), nullable=False),
+        sa.Column("simulation_ids", sa.JSON(), nullable=False),
+        sa.Column("id", sa.Uuid(), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("grid_simulations")

backend/src/acidwatch_api/database.py

@@ -49,6 +49,14 @@ class Simulation(Base):
     model_inputs: Mapped[list[ModelInput]] = relationship(back_populates="simulation")
 
 
+class GridSimulation(Base):
+    __tablename__ = "grid_simulations"
+
+    owner_id: Mapped[UUID | None] = mapped_column(Uuid)
+    axes: Mapped[list[dict]] = mapped_column(JSON)
+    simulation_ids: Mapped[list[str]] = mapped_column(JSON)
+
+
 class ModelInput(Base):
     __tablename__ = "model_inputs"
 

backend/src/acidwatch_api/models/datamodel.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 from typing import Any, Literal, Optional, Dict, TypeAlias
+from uuid import UUID
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 from pydantic.alias_generators import to_camel
 
 
@@ -58,6 +59,60 @@ class SimulationResult(_BaseModel):
     results: list[ModelResult]
 
 
+class AxisRange(_BaseModel):
+    """A linear, inclusive range that is sampled at ``steps`` points."""
+
+    min: float
+    max: float
+    steps: int = Field(default=10, ge=2, le=25)
+
+    @model_validator(mode="after")
+    def _check_bounds(self) -> "AxisRange":
+        if self.max <= self.min:
+            raise ValueError("max must be greater than min")
+        return self
+
+    def values(self) -> list[float]:
+        step = (self.max - self.min) / (self.steps - 1)
+        return [self.min + step * i for i in range(self.steps)]
+
+
+class Axis(_BaseModel):
+    substance: str
+    range: AxisRange
+
+
+class CreateGridSimulation(_BaseModel):
+    """Request body for starting a grid simulation.
+
+    Runs a model chain once for each point in the cartesian product of the
+    axes, substituting each axis's substance in ``concentrations`` with the
+    corresponding value.
+    """
+
+    axes: list[Axis] = Field(min_length=1, max_length=2)
+    concentrations: dict[str, int | float]
+    conditions: Conditions = Field(default_factory=Conditions)
+    models: list[ModelInput] = Field(min_length=1)
+
+
+class GridPoint(_BaseModel):
+    coordinates: list[float]
+    simulation_id: UUID
+    status: Literal["done", "pending", "error"]
+    error: str | None = None
+    concentrations: dict[str, int | float]
+
+
+class GridSimulationResult(_BaseModel):
+    status: Literal["done", "pending"]
+    axes: list[Axis]
+    concentrations: dict[str, int | float]
+    conditions: Conditions
+    models: list[ModelInput]
+    points: list[GridPoint]
+
+
 class JsonResult(BaseModel):
     type: Literal["json"] = "json"
     label: str | None = None

backend/src/acidwatch_api/routes/__init__.py

@@ -4,10 +4,12 @@
 
 from . import models
 from . import oasis
+from . import grid_simulations
 
 router = APIRouter()
 router.include_router(models.router)
 router.include_router(oasis.router)
+router.include_router(grid_simulations.router)
 
 
 __all__ = ["router"]

backend/src/acidwatch_api/routes/grid_simulations.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import Annotated
+from uuid import UUID
+
+from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request
+from sqlalchemy import select
+
+import acidwatch_api.database as db
+from acidwatch_api.authentication import OptionalCurrentUser
+from acidwatch_api.database import GetDB
+from acidwatch_api.models import InputError
+from acidwatch_api.models.datamodel import (
+    Axis,
+    Conditions,
+    CreateGridSimulation,
+    GridPoint,
+    GridSimulationResult,
+    ModelInput,
+    Phase,
+)
+from acidwatch_api.routes.models import (
+    AdapterSet,
+    _phases_to_concentrations,
+    build_adapters,
+    build_model_input_rows,
+    get_adapters,
+    order_chain,
+    run_adapters,
+)
+
+router = APIRouter()
+
+logger = logging.getLogger(__name__)
+
+
+def _cartesian_values(axes: list[Axis]) -> list[list[float]]:
+    ranges = [axis.range.values() for axis in axes]
+    return [list(point) for point in itertools.product(*ranges)]
+
+
+def _summarize_point(
+    ordered: list[tuple[db.ModelInput, db.ModelResult | None]],
+) -> tuple[str, dict[str, int | float], str | None]:
+    if not ordered:
+        return "pending", {}, None
+
+    for _, result in ordered:
+        if result is not None and result.error is not None:
+            return "error", {}, result.error
+
+    if any(result is None for _, result in ordered):
+        return "pending", {}, None
+
+    final_result = ordered[-1][1]
+    assert final_result is not None
+    concentrations = _phases_to_concentrations(
+        [Phase(**p) for p in final_result.phases]
+    )
+    return "done", concentrations, None
+
+
+@router.post("/grid-simulations")
+async def run_grid_simulation(
+    create: CreateGridSimulation,
+    user: OptionalCurrentUser,
+    request: Request,
+    session: GetDB,
+    background_tasks: BackgroundTasks,
+    all_adapters: Annotated[AdapterSet, Depends(get_adapters)],
+) -> UUID:
+    jwt_token = user.jwt_token if user else None
+
+    adapters = build_adapters(
+        create.models, create.conditions, all_adapters, jwt_token
+    )
+
+    for axis in create.axes:
+        if axis.substance not in adapters[0].valid_substances:
+            raise HTTPException(
+                status_code=422,
+                detail={
+                    "axes": [
+                        f"'{axis.substance}' is not supported by the selected model"
+                    ]
+                },
+            )
+
+    test_concentrations = {
+        **create.concentrations,
+        **{axis.substance: 0 for axis in create.axes},
+    }
+    try:
+        adapters[0].validate_concentrations(test_concentrations)
+    except InputError as exc:
+        raise HTTPException(status_code=422, detail=exc.detail)
+
+    grid_points = _cartesian_values(create.axes)
+
+    scheduled: list[tuple[dict[str, int | float], list, list[UUID]]] = []
+    simulation_ids: list[str] = []
+
+    for coordinates in grid_points:
+        point_concentrations = {
+            **create.concentrations,
+            **{
+                axis.substance: value
+                for axis, value in zip(create.axes, coordinates)
+            },
+        }
+        model_input_rows = build_model_input_rows(create.models)
+        simulation = db.Simulation(
+            owner_id=UUID(user.id) if user else None,
+            phases=[{"kind": "co2-rich", "fraction": 1.0, "concentrations": point_concentrations}],
+            conditions=create.conditions.model_dump(),
+            model_inputs=model_input_rows,
+        )
+        session.add(simulation)
+        session.flush()
+        simulation_ids.append(str(simulation.id))
+
+        point_adapters = build_adapters(
+            create.models, create.conditions, all_adapters, jwt_token
+        )
+        scheduled.append(
+            (
+                point_concentrations,
+                point_adapters,
+                [row.id for row in model_input_rows],
+            )
+        )
+
+    grid = db.GridSimulation(
+        owner_id=UUID(user.id) if user else None,
+        axes=[axis.model_dump() for axis in create.axes],
+        simulation_ids=simulation_ids,
+    )
+    session.add(grid)
+    session.commit()
+
+    for point_concentrations, point_adapters, model_input_ids in scheduled:
+        background_tasks.add_task(
+            run_adapters,
+            request.state.session,
+            point_concentrations,
+            point_adapters,
+            model_input_ids,
+        )
+
+    return grid.id
+
+
+@router.get("/grid-simulations/{grid_id}/result")
+def get_grid_simulation_result(
+    grid_id: UUID,
+    session: GetDB,
+) -> GridSimulationResult:
+    grid = session.get_one(db.GridSimulation, grid_id)
+
+    axes = [Axis(**a) for a in grid.axes]
+    sim_uuids = [UUID(sid) for sid in grid.simulation_ids]
+
+    simulations_by_id: dict[UUID, db.Simulation] = {}
+    if sim_uuids:
+        rows = (
+            session.execute(
+                select(db.Simulation).where(db.Simulation.id.in_(sim_uuids))
+            )
+            .scalars()
+            .all()
+        )
+        simulations_by_id = {sim.id: sim for sim in rows}
+
+    rows_by_simulation: dict[UUID, list[tuple[db.ModelInput, db.ModelResult | None]]] = {}
+    if sim_uuids:
+        q = (
+            select(db.ModelInput, db.ModelResult)
+            .where(db.ModelInput.simulation_id.in_(sim_uuids))
+            .outerjoin(db.ModelResult)
+        )
+        for model_input, result in session.execute(q):
+            rows_by_simulation.setdefault(model_input.simulation_id, []).append(
+                (model_input, result)
+            )
+
+    grid_points = _cartesian_values(axes)
+    points: list[GridPoint] = []
+    overall_pending = False
+
+    for index, coordinates in enumerate(grid_points):
+        sim_id = sim_uuids[index] if index < len(sim_uuids) else None
+        if sim_id is None or sim_id not in simulations_by_id:
+            overall_pending = True
+            continue
+
+        ordered = order_chain(rows_by_simulation.get(sim_id, []))
+        status, concentrations, error = _summarize_point(ordered)
+        if status == "pending":
+            overall_pending = True
+
+        points.append(
+            GridPoint(
+                coordinates=coordinates,
+                simulation_id=sim_id,
+                status=status,  # type: ignore[arg-type]
+                error=error,
+                concentrations=concentrations,
+            )
+        )
+
+    first_sim = simulations_by_id.get(sim_uuids[0]) if sim_uuids else None
+    if first_sim is not None:
+        models = [
+            ModelInput(model_id=mi.model_id, parameters=mi.parameters)
+            for mi, _ in order_chain(rows_by_simulation.get(first_sim.id, []))
+        ]
+        base_concentrations = _phases_to_concentrations(
+            [Phase(**p) for p in (first_sim.phases or [])]
+        )
+        conditions = Conditions(**(first_sim.conditions or {}))
+    else:
+        models = []
+        base_concentrations = {}
+        conditions = Conditions()
+
+    return GridSimulationResult(
+        status="pending" if overall_pending else "done",
+        axes=axes,
+        concentrations=base_concentrations,
+        conditions=conditions,
+        models=models,
+        points=points,
+    )