Marimo Flow 🌊

Like marimo algae drifting in crystal waters, your data flows and evolves – each cell a living sphere of computation, gently touching others, creating ripples of reactive change. In this digital ocean, data streams like currents, models grow like organic formations, and insights emerge naturally from the depths. Let your ML experiments flow freely, tracked and nurtured, as nature intended.

https://github.com/user-attachments/assets/3bc24463-ff42-44a7-ae61-5d500d29688c

What is marimo-flow? 🚀

An agentic scientific-computing platform for physics-informed ML. A reactive multi-agent team orchestrates PINA / PINN workflows end-to-end over a pydantic-graph state machine — backed by MLflow for tracing and persistence, and exposed through marimo's chat UI plus optional A2A and AG-UI ASGI servers.

Describe a PDE in natural language; the team composes the Problem, designs the network, wires the solver, trains it, and grades the run — every handoff typed, logged to MLflow, and indexed in a DuckDB provenance store. A classic marimo_flow.core API is still there for when you'd rather drive PINA by hand. Composition-first throughout: no hardcoded PDE factories, plus a Docker deployment story for CPU / NVIDIA / Intel GPUs.

Architecture 🏗️

Natural language goes in; a pydantic-ai multi-agent team turns it into a composed, trained, tracked PINN. PINA, MLflow, and marimo are the three backends the team drives:

   Natural language  ("Solve a 1D Poisson equation on [0,1] with u(0)=u(1)=0")
          │
          ▼
 ┌─────────────────────────────────────────────────────────┐
 │  pydantic-ai  ·  multi-agent team                        │
 │  Triage → Route → Problem · Model · Solver · Training     │
 │         → Validation → MLflow · Notebook                 │
 │  (pydantic-graph state machine · typed spec handoffs)    │
 └─────────────────────────────────────────────────────────┘
          │                    │                    │
          ▼                    ▼                    ▼
       ┌──────┐            ┌────────┐           ┌────────┐
       │ PINA │            │ MLflow │           │ marimo │
       └──────┘            └────────┘           └────────┘
    compose + train      runs · metrics       reactive .py
    PINNs (PyTorch)      registry · tracing    notebooks · chat
    Plotly 2D/3D viz     DuckDB provenance     MCP-wired UI

A classic marimo_flow.core API drives PINA by hand when you'd rather skip the agents — both paths write to the same MLflow backend.

Stack 🧱

🧠 PINA — composition-first (no hardcoded PDE factories)

The core idea: agents emit typed EquationSpec + SubdomainSpec + ConditionSpec and the composer compiles a pina.Problem subclass at runtime via sympy.lambdify + pina.operator.grad/laplacian. There is no ProblemKind enum — any PDE sympy can express is reachable.

Code lives in src/marimo_flow/core/ (PINA solvers + training + viz3d), src/marimo_flow/control/ (MPC), and src/marimo_flow/agents/services/ (composer, mesh-domain, design aggregator).

Features ✨

• 🧑‍🚀 Multi-agent PINA team — free-form intent → typed TaskSpec → composed ProblemSpec / ModelSpec / SolverPlan → trained surrogate → validation verdict, over a 9-node pydantic-graph state machine (see docs/agents.md).
• 🔌 Three transports — in-notebook mo.ui.chat, an A2A ASGI server, and an AG-UI ASGI server, all wrapping the same lead agent.
• 🧩 Provider-agnostic LLMs — per-role model specs resolved through pydantic-ai's infer_model; defaults target Ollama Cloud, any provider in the catalogue works via config.yaml / env vars.
• 🔬 MLflow tracing + DuckDB provenance — every run, model, and metric tracked under data/mlflow/{db,artifacts}/; typed specs, decisions, handoffs, and lineage mirrored into a 16-table DuckDB store. Lightning checkpoints land inside the active run.
• 🧠 Composition-first PINA — agents compose PDEs from primitives; no hardcoded ProblemKind (see below).
• 📓 Reactive notebooks — Git-friendly .py notebooks with automatic dependency tracking.
• 🤖 MCP-powered AI — marimo, mlflow, and context7 MCP servers wired up; live library docs without leaving your notebook.
• 🐳 Multi-platform Docker — CPU, CUDA, Intel XPU images on GHCR.

Quick Start 🏃‍♂️

Talk to the agent team

The fastest path is the in-notebook chat. The lead agent dispatches to the specialist team and streams its answer back:

from marimo_flow.agents import lead_chat, FlowDeps
import marimo as mo

deps = FlowDeps()  # resolves per-role models + MLflow URI from config / .env / env vars
chat = mo.ui.chat(
    lead_chat(deps=deps),
    prompts=["Solve a 1D Poisson equation on [0,1] with u(0)=u(1)=0 using a PINN."],
)
chat

Defaults target Ollama Cloud (:cloud tags via a local Ollama endpoint). Copy config.yaml.example to config.yaml to point any role at OpenAI, Anthropic, Groq, … — auth comes from each provider's standard env var.

Prefer a terminal? The CLI runs the same graph:

uv run marimo-flow solve "Solve the Burgers equation with a small PINN"
uv run marimo-flow solve -m lead=anthropic:claude-sonnet-4-6 "..."   # override a role
uv run marimo-flow config-show        # print the resolved models + URIs
uv run marimo-flow lab                # open examples/lab.py in marimo

Or expose the team as a standalone ASGI server:

uv run python -m marimo_flow.agents.server.a2a     # A2A protocol   → :8000
uv run python -m marimo_flow.agents.server.ag_ui   # AG-UI protocol → :8001

With Docker (Recommended)

git clone https://github.com/synapticore-io/marimo-flow.git
cd marimo-flow
docker compose -f docker/docker-compose.yaml up --build -d

# Marimo:  http://localhost:2718
# MLflow:  http://localhost:5000

Image Variants

GPU compose files: docker-compose.cuda.yaml (requires nvidia-docker) and docker-compose.xpu.yaml (requires Intel GPU drivers).

Local Development

→ See SETUP.md — bare-metal uv path, MLflow + marimo processes, MCP setup table, troubleshooting. Five minutes from clone to running notebook.

Example Notebooks 📚

All notebooks live in examples/ and open with uv run marimo edit examples/<file>.py.

Project Structure 📁

marimo-flow/
├── examples/                     # Demo notebooks
├── src/marimo_flow/
│   ├── core/                     # PINA solvers, training, plotly viz3d
│   ├── control/                  # Rolling-horizon MPC (scipy SLSQP)
│   └── agents/                   # Multi-agent team (pydantic-graph + MLflow)
│       ├── nodes/                # 9 graph nodes
│       ├── schemas/              # Typed Pydantic specs (ProblemSpec, …)
│       ├── toolsets/             # FunctionToolset per role
│       └── services/             # composer, mesh_domain, design,
│                                 #   provenance (DuckDB, 16 tables)
├── tests/                        # 232 passing, 0 xfailed
├── docker/                       # Dockerfiles + compose (CPU/CUDA/XPU)
├── docs/                         # Project documentation (see docs/INDEX.md)
└── data/mlflow/                  # MLflow storage (db + artifacts)

Two Workflows

Both write to the same MLflow backend (data/mlflow/). The two packages do not depend on each other — pick whichever matches the task.

PINA Multi-Agent Team 🧑‍🚀

The team drives PINA workflows end-to-end — it doesn't just answer questions. Free-form intent → typed TaskSpec → composed ProblemSpec → trained surrogate → validation verdict, all logged to MLflow + a DuckDB provenance store.

TriageNode → RouteNode ─┬─ ProblemNode
                        ├─ ModelNode
                        ├─ SolverNode
                        ├─ TrainingNode
                        ├─ ValidationNode → (accept/retry/escalate/reject)
                        ├─ MLflowNode
                        └─ NotebookNode

Nine graph nodes (pydantic-graph), one toolset + skill per role, provider-agnostic LLM config, and three transport options (in-notebook chat, A2A server, AG-UI server). See the Quick Start for the chat / CLI / server entry points.

→ Full architecture, role list, provenance schema, and config docs in docs/agents.md.

MCP Integration 🔌

marimo and AI-assisted IDEs share MCP servers for live documentation and notebook operations. Pre-configured in .marimo.toml (in-notebook) and .vscode/mcp.json (VS Code / Claude Code).

→ Full configuration reference in docs/mcp-setup.md.

Claude Code Integration 🤖

marimo-flow is built to be driven from Claude Code. Two integration points work together:

• marimo MCP server — start marimo with --mcp and Claude Code can introspect the running notebook: list cells, read variables/DataFrames, surface cell errors, and inspect DuckDB/SQL connections.

  marimo edit examples/ --mcp --no-token --headless --port 2718
  

• marimo-pair — the marimo-claude plugin's pair-programming skill. It attaches to that same live kernel so Claude Code can execute code, create cells, and build the notebook as an artifact interactively — not just read it. Ask Claude to "start a marimo notebook" or "work in the active marimo session" and it drives the kernel directly.

Together they turn a marimo notebook into a shared workspace: Claude reads runtime state over MCP and writes cells back through marimo-pair, with MLflow + the DuckDB provenance store recording everything the agent team produces alongside.

Documentation 📚

Contributing 🤝

We welcome contributions! See CONTRIBUTING.md for development setup, code standards, and the PR process.

# Quick path
git checkout -b my-feature
uv run pytest                         # 232 passing
uv run ruff format . && uv run ruff check --fix .
# open a PR

License 📄

MIT — see LICENSE.

Built with ❤️ using marimo, MLflow, and PINA.