Add agent-first CLI for k-Wave simulations

Structured JSON I/O at every step, deterministic exit codes,
file-backed session model. Commands: session init/status/reset,
phantom generate, sensor define, plan, run.

Acceptance test: new_api_ivp_2D.py runs entirely via CLI with
exact numerical parity to the Python API.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: waltsims

kwave/cli/__init__.py

@@ -0,0 +1,9 @@
+def __getattr__(name):
+    if name == "Session":
+        from kwave.cli.session import Session
+
+        return Session
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["Session"]

kwave/cli/commands/__init__.py

@@ -0,0 +1,108 @@
+"""Phantom generation and loading commands."""
+
+import click
+import numpy as np
+
+from kwave.cli.main import pass_session
+from kwave.cli.schema import CLIError, CLIResponse, ValidationError, json_command
+
+
+@click.group("phantom")
+def phantom():
+    """Define the simulation phantom (medium + initial pressure)."""
+    pass
+
+
+@phantom.command()
+@click.option("--type", "phantom_type", required=True, type=click.Choice(["disc", "spherical", "layered"]))
+@click.option("--grid-size", required=True, help="Grid dimensions, e.g. 128,128")
+@click.option("--spacing", required=True, type=float, help="Grid spacing in meters, e.g. 0.1e-3")
+@click.option("--sound-speed", type=float, default=1500, help="Medium sound speed (m/s)")
+@click.option("--density", type=float, default=1000, help="Medium density (kg/m^3)")
+@click.option("--disc-center", default=None, help="Disc center, e.g. 64,64")
+@click.option("--disc-radius", type=int, default=5, help="Disc radius in grid points")
+@pass_session
+@json_command("phantom.generate")
+def generate(sess, phantom_type, grid_size, spacing, sound_speed, density, disc_center, disc_radius):
+    """Generate an analytical phantom."""
+    sess.load()
+
+    grid_n = tuple(int(x) for x in grid_size.split(","))
+    ndim = len(grid_n)
+    grid_spacing = (spacing,) * ndim
+
+    if phantom_type == "disc":
+        if ndim != 2:
+            raise ValidationError(
+                CLIError(
+                    code="DISC_REQUIRES_2D",
+                    field="grid_size",
+                    value=grid_size,
+                    constraint="disc phantom requires 2D grid",
+                    suggestion="Use --grid-size Nx,Ny (two dimensions)",
+                )
+            )
+        from kwave.data import Vector
+        from kwave.utils.mapgen import make_disc
+
+        if disc_center is None:
+            center = Vector([n // 2 for n in grid_n])
+        else:
+            center = Vector([int(x) for x in disc_center.split(",")])
+
+        p0 = make_disc(Vector(list(grid_n)), center, disc_radius).astype(float)
+
+    elif phantom_type == "spherical":
+        # Simple spherical inclusion centered in grid
+        center = np.array([n // 2 for n in grid_n])
+        coords = np.mgrid[tuple(slice(0, n) for n in grid_n)]
+        dist = np.sqrt(sum((c - cn) ** 2 for c, cn in zip(coords, center)))
+        p0 = (dist <= disc_radius).astype(float)
+
+    elif phantom_type == "layered":
+        p0 = np.zeros(grid_n)
+        layer_pos = grid_n[0] // 4
+        p0[layer_pos, ...] = 1.0
+
+    # Save arrays
+    p0_path = sess.save_array("p0", p0)
+
+    # Update session
+    sess.update(
+        "grid",
+        {
+            "N": list(grid_n),
+            "spacing": list(grid_spacing),
+            "sound_speed_for_time": sound_speed,
+        },
+    )
+    sess.update(
+        "medium",
+        {
+            "sound_speed": sound_speed,
+            "density": density,
+        },
+    )
+    sess.update(
+        "source",
+        {
+            "type": "initial-pressure",
+            "p0_path": p0_path,
+        },
+    )
+
+    return CLIResponse(
+        result={
+            "phantom_type": phantom_type,
+            "grid_size": list(grid_n),
+            "spacing": list(grid_spacing),
+            "p0_shape": list(p0.shape),
+            "p0_max": float(p0.max()),
+            "sound_speed": sound_speed,
+            "density": density,
+        },
+        derived={
+            "ndim": ndim,
+            "grid_points": int(np.prod(grid_n)),
+        },
+    )

kwave/cli/commands/phantom.py

@@ -0,0 +1,108 @@
+"""Plan command: derive full simulation config, validate, estimate cost."""
+
+import click
+import numpy as np
+
+from kwave.cli.main import pass_session
+from kwave.cli.schema import CLIResponse, SessionError, json_command
+
+
+@click.command("plan")
+@pass_session
+@json_command("plan")
+def plan(sess):
+    """Derive full simulation config and validate before running."""
+    sess.load()
+
+    # Check completeness
+    comp = sess._completeness()
+    missing = [k for k, v in comp.items() if not v]
+    if missing:
+        raise SessionError(f"Cannot plan: missing {', '.join(missing)}. Complete setup first.")
+
+    kgrid = sess.make_grid()
+    medium = sess.make_medium()
+
+    g = sess.state["grid"]
+    grid_n = tuple(g["N"])
+    spacing = tuple(g["spacing"])
+    ndim = len(grid_n)
+    grid_points = int(np.prod(grid_n))
+
+    # Time stepping
+    dt = float(kgrid.dt)
+    Nt = int(kgrid.Nt)
+
+    # PPW check
+    c_min = float(np.min(medium.sound_speed)) if hasattr(medium.sound_speed, "__len__") else float(medium.sound_speed)
+    max_spacing = max(spacing)
+    # For IVP problems, use grid-based wavelength estimate
+    min_wavelength = c_min * dt * Nt / 2  # rough estimate
+    ppw = c_min / (max_spacing * (1 / (2 * dt)))  # Nyquist-based PPW
+
+    # CFL
+    c_max = float(np.max(medium.sound_speed)) if hasattr(medium.sound_speed, "__len__") else float(medium.sound_speed)
+    cfl = c_max * dt / min(spacing)
+
+    # Memory estimate: ~(3 + 2*ndim) fields of float64
+    n_fields = 3 + 2 * ndim
+    memory_bytes = grid_points * n_fields * 8
+    memory_mb = memory_bytes / (1024 * 1024)
+
+    # Runtime estimate
+    cost_per_point_step_ns = 50  # ~50ns per grid point per step on CPU
+    estimated_runtime_s = grid_points * Nt * cost_per_point_step_ns / 1e9
+
+    # PML
+    pml_size = 20
+
+    # Warnings
+    warnings = []
+    if ppw < 4:
+        warnings.append(
+            {
+                "code": "LOW_PPW",
+                "detail": f"PPW={ppw:.1f} is below recommended minimum of 4",
+                "suggestion": "Increase grid resolution or reduce maximum frequency",
+            }
+        )
+    if cfl > 0.5:
+        warnings.append(
+            {
+                "code": "HIGH_CFL",
+                "detail": f"CFL={cfl:.3f} exceeds 0.5, simulation may be unstable",
+                "suggestion": "Reduce time step or increase grid spacing",
+            }
+        )
+
+    result = {
+        "grid": {
+            "N": list(grid_n),
+            "spacing": list(spacing),
+            "ndim": ndim,
+            "dt": dt,
+            "Nt": Nt,
+        },
+        "pml": {"size": pml_size},
+        "medium": {
+            "sound_speed": c_min if c_min == c_max else f"{c_min}-{c_max}",
+        },
+        "source": sess.state["source"],
+        "sensor": sess.state["sensor"],
+        "backend": "python",
+        "device": "cpu",
+    }
+
+    derived = {
+        "ppw": round(ppw, 2),
+        "cfl": round(cfl, 4),
+        "grid_points": grid_points,
+        "estimated_memory_mb": round(memory_mb, 1),
+        "estimated_runtime_s": round(estimated_runtime_s, 1),
+    }
+
+    return CLIResponse(
+        result=result,
+        derived=derived,
+        warnings=warnings,
+    )

kwave/cli/commands/plan.py

@@ -0,0 +1,94 @@
+"""Run command: execute simulation with structured JSON progress."""
+
+import json
+import sys
+import time
+
+import click
+import numpy as np
+
+from kwave.cli.main import pass_session
+from kwave.cli.schema import CLIResponse, SessionError, json_command
+
+
+def _emit_event(event: dict):
+    """Write a JSON event to stdout and flush."""
+    click.echo(json.dumps(event, default=str))
+
+
+@click.command("run")
+@click.option("--backend", default="python", type=click.Choice(["python", "cpp"]))
+@click.option("--device", default="cpu", type=click.Choice(["cpu", "gpu"]))
+@pass_session
+@json_command("run")
+def run(sess, backend, device):
+    """Execute the simulation."""
+    sess.load()
+
+    # Check completeness
+    comp = sess._completeness()
+    missing = [k for k, v in comp.items() if not v]
+    if missing:
+        raise SessionError(f"Cannot run: missing {', '.join(missing)}. Complete setup first.")
+
+    kgrid = sess.make_grid()
+    medium = sess.make_medium()
+    source = sess.make_source()
+    sensor = sess.make_sensor()
+
+    Nt = int(kgrid.Nt)
+
+    _emit_event({"event": "started", "backend": backend, "device": device, "Nt": Nt})
+
+    t_start = time.time()
+    last_pct = -5  # emit at most every 5%
+
+    def progress_callback(step, total):
+        nonlocal last_pct
+        pct = round(100 * step / total, 1)
+        if pct - last_pct >= 5 or step == total:
+            last_pct = pct
+            _emit_event(
+                {
+                    "event": "progress",
+                    "step": step,
+                    "total": total,
+                    "pct": pct,
+                    "elapsed_s": round(time.time() - t_start, 2),
+                }
+            )
+
+    from kwave.kspaceFirstOrder import kspaceFirstOrder
+
+    result = kspaceFirstOrder(
+        kgrid,
+        medium,
+        source,
+        sensor,
+        backend=backend,
+        device=device,
+        quiet=True,
+        progress_callback=progress_callback,
+    )
+
+    elapsed = round(time.time() - t_start, 2)
+
+    # Save results
+    result_info = {}
+    for key, val in result.items():
+        if isinstance(val, np.ndarray):
+            path = sess.save_array(f"result_{key}", val)
+            result_info[key] = {"shape": list(val.shape), "path": path}
+        else:
+            result_info[key] = val
+
+    sess.update("result_path", str(sess.data_dir))
+
+    _emit_event({"event": "completed", "elapsed_s": elapsed, "output_keys": list(result.keys())})
+
+    return CLIResponse(
+        result={
+            "elapsed_s": elapsed,
+            "outputs": result_info,
+        },
+    )

kwave/cli/commands/run.py

@@ -0,0 +1,40 @@
+"""Sensor definition command."""
+
+import click
+
+from kwave.cli.main import pass_session
+from kwave.cli.schema import CLIResponse, json_command
+
+
+@click.group("sensor")
+def sensor():
+    """Define sensor configuration."""
+    pass
+
+
+@sensor.command()
+@click.option("--mask", required=True, help="Sensor mask: 'full-grid' or path to .npy file")
+@click.option("--record", default="p,p_final", help="Comma-separated fields to record, e.g. p,p_final,ux")
+@pass_session
+@json_command("sensor.define")
+def define(sess, mask, record):
+    """Define what and where to record."""
+    sess.load()
+
+    record_fields = [r.strip() for r in record.split(",")]
+
+    sensor_config = {"record": record_fields}
+    if mask == "full-grid":
+        sensor_config["mask_type"] = "full-grid"
+    else:
+        sensor_config["mask_type"] = "file"
+        sensor_config["mask_path"] = mask
+
+    sess.update("sensor", sensor_config)
+
+    return CLIResponse(
+        result={
+            "mask_type": sensor_config["mask_type"],
+            "record": record_fields,
+        }
+    )