add run_visco_forward_sweep: viscoacoustic forward tool with attenuation switches, wired into selector + prompt

src/sweep_agent/prompts.py

@@ -33,6 +33,12 @@
       solver directly and is the ONLY way to run a shot without the sweep_tasks
       tier — use it as the fallback whenever build_forward_spec /
       run_forward_and_plot say "not importable".
+      - ATTENUATING / VISCO-ACOUSTIC forward (Q / attenuation / lossy / 衰减) →
+      run_visco_forward_sweep(vp_path, dh, dt, nt|record_length_s, fm, q,
+      phase_shift, amplitude_damping). Use THIS (not run_forward_sweep) when
+      attenuation is involved — ViscoAcoustic needs three models. Switches:
+      phase_shift (dispersion) + amplitude_damping (dissipation); both off =
+      acoustic, both on = full visco-acoustic.
 - WAVEFIELD — see the wave / snapshots / how it propagates / a movie or gif /
   波场 / 快照 / 看波怎么传播 / 动画
       • To ANIMATE one wavefield (a GIF/movie of the wave propagating) →

src/sweep_agent/tools/__init__.py

@@ -120,5 +120,5 @@ def _decorate(fn: Callable[[BaseModel], Any]) -> Tool:
 import sweep_agent.tools.synth            # noqa: E402,F401
 import sweep_agent.tools.analysis         # noqa: E402,F401
 import sweep_agent.tools.forward_sweep    # noqa: E402,F401
-
+import sweep_agent.tools.visco_forward_sweep
 __all__ = ["Tool", "ToolResult", "Registry", "registry", "register"]

src/sweep_agent/tools/selection.py

@@ -13,13 +13,15 @@
 # Always available — the inspect / run / discover backbone.
 CORE = [
     "inspect_file", "run_task", "read_status", "list_artifacts",
-    "list_equations", "describe_task_schema", "build_spec", "run_forward_sweep",
+    "list_equations", "describe_task_schema", "build_spec", "run_forward_sweep", "run_visco_forward_sweep",
 ]
 
 # (trigger keywords, tool names) — a query adds every group whose keywords hit.
 GROUPS: list[tuple[tuple[str, ...], list[str]]] = [
     (("forward", "正演", "shot gather", "shot-gather", "炮记录", "合成记录", "record", "gather", "synthetic record"),
-     ["build_forward_spec", "run_forward_and_plot", "plot_shot_gather", "run_forward_sweep"]),
+     ["build_forward_spec", "run_forward_and_plot", "plot_shot_gather", "run_forward_sweep", "run_visco_forward_sweep"]),
+         (("visco", "attenuat", "quality factor", "q factor", "lossy", "衰减"),
+     ["run_visco_forward_sweep", "run_forward_sweep"]),
     (("wavefield", "波场", "snapshot", "快照", "animate", "animation", "动画", "movie", "gif", "propagat", "传播", "p wave", "s wave", "p波", "s波", "p-wave", "s-wave"),
      ["build_wavefield_spec", "animate_wavefield", "make_wavefield_gif", "plot_wavefield"]),
     (("anisotrop", "各向异性", "vti", "tti", "wavefront", "波前", "compare equation", "对比方程", "elastic", "弹性"),
@@ -40,7 +42,7 @@
 _FALLBACK = [
     "build_forward_spec", "run_forward_and_plot", "plot_shot_gather",
     "build_wavefield_spec", "animate_wavefield", "make_synthetic_model",
-    "plot_velocity_model", "run_fwi", "check_parameters", "run_forward_sweep",
+    "plot_velocity_model", "run_fwi", "check_parameters", "run_forward_sweep", "run_visco_forward_sweep",
 ]
 
 

src/sweep_agent/tools/visco_forward_sweep.py

@@ -0,0 +1,206 @@
+"""``run_visco_forward_sweep`` — forward modelling straight through the `sweep` solver.
+
+Every other modelling tool here routes through ``sweep_tasks``: build a spec,
+write YAML, hand it to ``TaskRunner``. That tier is unreleased, so on a plain
+``pip install -e .`` plus ``sweep`` the agent can inspect files and list
+equations but cannot actually run a shot.
+
+This tool closes that gap by driving the solver directly —
+``sweep.equations`` + ``sweep.propagator.torch.PropTorch`` — with no
+``sweep_tasks``, ``sweep_io`` or ``sweep_loss`` anywhere in the path. It needs
+only ``sweep`` + torch + numpy + matplotlib, so it runs on a laptop, CPU or
+Apple-Silicon MPS.
+
+STATUS: scaffold. The parameter surface, registration and dependency guard are
+in place; the modelling body is the exercise — see the TODO in
+``run_visco_forward_sweep`` and ``tests/test_forward_sweep.py``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from pydantic import BaseModel, Field
+
+from sweep_agent.tools import register
+
+
+def _require_sweep():
+    """Return ``(modules, None)`` or ``(None, error_dict)`` when sweep is absent.
+
+    Same contract as the rest of the package: a missing layer is reported as
+    data, never raised, so the agent stays up on a base install.
+    """
+    try:
+        import torch  # noqa: F401
+        from sweep.equations import _equation_classes  # noqa: F401
+        from sweep.propagator.torch import PropTorch  # noqa: F401
+        from sweep.signal import ricker  # noqa: F401
+    except ImportError as exc:
+        return None, {
+            "error": (
+                f"the sweep solver (or torch) is not importable: {exc}. "
+                "Install a PyTorch environment, then `pip install .` from a clone of "
+                "https://github.com/DeepWave-KAUST/sweep"
+            )
+        }
+    import torch
+    import sweep.equations as eq_mod
+    from sweep.propagator.torch import PropTorch
+    from sweep.signal import ricker
+
+    return {"torch": torch, "eq_mod": eq_mod, "PropTorch": PropTorch, "ricker": ricker}, None
+
+
+class RunViscoForwardSweepParams(BaseModel):
+    vp_path: str = Field(..., description="Velocity model .npy, shape (nz, nx) in m/s. Call inspect_file first to learn its shape.")
+    dh: float = Field(..., gt=0, description="Grid spacing in metres (isotropic).")
+    dt: float = Field(..., gt=0, description="Time step in seconds. Call check_parameters first — a too-large dt is unstable.")
+    nt: int | None = Field(None, gt=0, description="Number of time samples. Give either nt or record_length_s.")
+    record_length_s: float | None = Field(None, gt=0, description="Record length in seconds; nt = round(record_length_s / dt).")
+    fm: float = Field(8.0, gt=0, description="Ricker peak frequency in Hz.")
+    q: float = Field(30.0, gt=0, description="Homogeneous quality factor Q (lower = stronger attenuation).")
+    phase_shift: bool = Field(True, description="Apply the dispersion (phase shift) term.")
+    amplitude_damping: bool = Field(True, description="Apply the dissipative amplitude damping term. Both off = Acoustic; both on = full visco acoustic.")
+    source_x: int | None = Field(None, ge=0, description="Source column index; defaults to the middle of the model.")
+    source_depth: int = Field(2, ge=0, description="Source row index, in grid cells from the top.")
+    receiver_depth: int = Field(1, ge=0, description="Row index of the receiver line, in grid cells from the top.")
+    receiver_step: int = Field(1, ge=1, description="Place a receiver every Nth column.")
+    spatial_order: int = Field(4, description="Finite-difference spatial order.")
+    abcn: int = Field(40, ge=0, description="PML half-width in grid cells.")
+    free_surface: bool = Field(False, description="Apply a free surface at the top boundary.")
+    device: str = Field("cpu", description="'cpu', 'cuda' (NVIDIA) or 'mps' (Apple Silicon). The runtime-environment note lists what THIS machine has.")
+    out_dir: str = Field("./sweep_runs", description="Directory for the record .npy and the shot-gather PNG.")
+    plot: bool = Field(True, description="Also render the shot gather to a PNG.")
+
+
+@register(
+    name="run_visco_forward_sweep",
+    description=(
+        "Run 2-D acoustic forward modelling directly on the `sweep` solver and return the shot "
+        "gather — no sweep_tasks spec, no TaskRunner. This is the tool to use when only the core "
+        "solver is installed. Give vp_path, dh, dt, and either nt or record_length_s; the source "
+        "defaults to the middle of the model and receivers to a full-width surface line. Call "
+        "inspect_file on the model and check_parameters on (dh, dt, fm) first. Returns the record "
+        ".npy path, its (nt, nrec) shape, the device actually used, and a shot-gather PNG."
+    ),
+    params_model=RunViscoForwardSweepParams,
+)
+def run_visco_forward_sweep(args: RunViscoForwardSweepParams) -> dict[str, Any]:
+    mods, err = _require_sweep()
+    if err is not None:
+        return err
+    from sweep_agent.tools.build_forward import _device_unavailable_reason
+    dev_err = _device_unavailable_reason(args.device)
+    if dev_err is not None:
+        return {"error": dev_err}
+
+
+
+    # ------------------------------------------------------------------
+    # TODO — implement the forward run. Worked reference (a complete, runnable
+    # script using exactly this API):
+    #     sweep/examples/wavefields/topography/acoustic2d_hill_demo.py
+    torch = mods["torch"]
+    ricker = mods["ricker"]
+    # 1. Resolve nt: use args.nt, else round(args.record_length_s / args.dt).
+    #    Return {"error": ...} if neither was given.
+    nt = args.nt
+    if nt is None:
+        if args.record_length_s is None:
+            return {"error": "give either nt or record_length_s."}
+        nt = round(args.record_length_s / args.dt)
+    # 2. np.load the model, sanity-check it is 2-D, move it to args.device.
+    if not Path(args.vp_path).exists():
+        return {"error": f"velocity model not found: {args.vp_path}"}
+    vp_np = np.load(args.vp_path)
+    if vp_np.ndim != 2:
+        return {"error": f"vp must be 2-D (nz, nx); got shape {vp_np.shape}."}
+    vp = torch.from_numpy(vp_np.astype(np.float32)).to(args.device)
+    nz, nx = vp_np.shape
+    Q = torch.full((nz, nx), float(args.q), dtype=torch.float32, device=args.device)
+    omega = torch.full((nz, nx), float(2.0 * np.pi * args.fm), dtype=torch.float32, device=args.device)
+    # 3. Build the Ricker wavelet with mods["ricker"] on a t axis of nt samples
+    #    (give it a delay so the wavelet is causal).
+    delay = 1.0 / args.fm
+    t = np.arange(nt, dtype=np.float32) * args.dt - delay
+    wavelet = torch.tensor((1.0e3 * ricker(t, f=args.fm)).astype(np.float32)).to(args.device)
+    # 4. Build source and receiver index tensors. sources is (nshot, 2) as
+    #    (x, z); receivers is (nshot, nrec, 2). Default the source to the middle
+    #    column, receivers to a line across the model at receiver_depth.
+    src_x = args.source_x if args.source_x is not None else nx // 2
+    sources = torch.from_numpy(
+        np.array([[src_x, args.source_depth]], dtype=np.int64)
+    ).to(args.device)
+    rec_x = np.arange(0, nx, args.receiver_step, dtype=np.int64)
+    rec_z = np.full_like(rec_x, args.receiver_depth)
+    receivers = torch.from_numpy(
+        np.stack([rec_x, rec_z], axis=-1)[None, ...]
+    ).to(args.device)
+    from sweep.equations.visco_acoustic import ViscoAcoustic
+    equation = ViscoAcoustic(
+        spatial_order=args.spatial_order, device=args.device, backend="torch", phase_shift=args.phase_shift, amplitude_damping=args.amplitude_damping,
+    )
+    # 6. prop = mods["PropTorch"](equation, shape=vp.shape, dh=args.dh,
+    #        dt=args.dt, abcn=args.abcn, free_surface=args.free_surface,
+    #        use_ckpt=False, impl="eager")
+    #    record = prop(wavelet, sources, receivers, models=[vp, Q, omega])  under no_grad.
+    prop = mods["PropTorch"](
+        equation, shape=(nz, nx), dh=args.dh, dt=args.dt,
+        abcn=args.abcn, free_surface=args.free_surface, use_ckpt=False, impl="eager",
+    )
+    with torch.no_grad():
+        record = prop(wavelet, sources, receivers, models=[vp, Q, omega])
+    # 7. record comes back as (nshot, nt, nrec, nfield) — squeeze to (nt, nrec),
+    #    save it as .npy under out_dir.
+    rec_np = np.squeeze(record.detach().cpu().numpy())
+    if rec_np.ndim != 2:
+        return {"error": f"unexpected record shape {tuple(record.shape)} -> {rec_np.shape}."}
+    if not np.isfinite(rec_np).all():
+        return {"error": "diverged (non-finite record) - reduce dt; see check_parameters"}
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    record_path = out_dir / "record.npy"
+    np.save(record_path, rec_np)
+
+    # 8. If args.plot: draw the gather with matplotlib (see
+    #    visualize.py::_require_matplotlib for the guarded-import pattern and
+    #    plot_observed_data for a gather-plotting example) and save a PNG.
+    image_path = None
+    if args.plot:
+        try:
+            import matplotlib
+            matplotlib.use("Agg")
+            import matplotlib.pyplot as plt
+        except ImportError:
+            plt = None
+        if plt is not None:
+            pct = np.percentile(np.abs(rec_np), 99.0) or 1.0
+            fig, ax = plt.subplots(figsize=(8, 5), constrained_layout=True)
+            ax.imshow(rec_np, cmap="seismic", vmin=-pct, vmax=pct,
+                      aspect="auto", extent=[rec_x.min(), rec_x.max(), nt * args.dt, 0])
+            ax.set_xlabel("receiver x (cells)"); ax.set_ylabel("time (s)")
+            ax.set_title(f"ViscoAcoustic gather (phase={args.phase_shift}, amp={args.amplitude_damping})")
+            png_path = out_dir / "shot_gather.png"
+            fig.savefig(png_path, dpi=140); plt.close(fig)
+            image_path = str(png_path)
+
+
+
+
+    # 9. Return {"record_path", "image_path", "shape", "device", "nt", "summary"}.
+    # Keep every import inside this function or inside _require_sweep, so the
+    # module still imports on a base install with no solver present.
+    # ------------------------------------------------------------------
+    return {
+        "record_path": str(record_path),
+        "image_path": image_path,
+        "shape": list(rec_np.shape),
+        "device": args.device,
+        "nt": nt,
+        "summary": f"Ran ViscoAcoustic forward (Q={args.q}, phase_shift={args.phase_shift}, "
+                   f"amplitude_damping={args.amplitude_damping}): {rec_np.shape[0]} samples x "
+                   f"{rec_np.shape[1]} receivers on {args.device}.",
+    }

tests/test_visco_forward_sweep.py

@@ -0,0 +1,71 @@
+"""Tests for ``run_visco_forward_sweep`` — the visco-acoustic solver-only path.
+
+The defining behaviour: the two attenuation switches must actually change the
+gather. phase-only and both must differ (the physics the equation demonstrates).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from sweep_agent.tools import registry
+
+
+def test_registered_with_the_attenuation_switches():
+    assert "run_visco_forward_sweep" in registry.tools
+    fields = registry.tools["run_visco_forward_sweep"].params_model.model_fields
+    for required in ("vp_path", "dh", "dt", "q", "phase_shift", "amplitude_damping", "device"):
+        assert required in fields, f"missing parameter: {required}"
+
+
+def test_missing_solver_is_reported_as_data(monkeypatch):
+    import sweep_agent.tools.visco_forward_sweep as vf
+
+    monkeypatch.setattr(vf, "_require_sweep", lambda: (None, {"error": "no solver"}))
+    out = vf.run_visco_forward_sweep.fn(
+        vf.RunViscoForwardSweepParams(vp_path="/nonexistent.npy", dh=10.0, dt=1e-3, nt=100)
+    )
+    assert "error" in out
+
+
+def test_selector_includes_run_visco_forward_sweep():
+    from sweep_agent.tools.selection import select_tool_names
+
+    assert "run_visco_forward_sweep" in select_tool_names("run an attenuating forward shot")
+
+
+def _run(vf, tmp_path, vp_path, phase_shift, amplitude_damping):
+    out = vf.run_visco_forward_sweep.fn(
+        vf.RunViscoForwardSweepParams(
+            vp_path=str(vp_path), dh=10.0, dt=1.0e-3, nt=300, fm=8.0, q=20.0,
+            phase_shift=phase_shift, amplitude_damping=amplitude_damping,
+            abcn=20, out_dir=str(tmp_path / f"{phase_shift}_{amplitude_damping}"),
+            plot=False,
+        )
+    )
+    assert "error" not in out, out
+    return np.load(out["record_path"])
+
+
+def test_toggles_change_the_gather(tmp_path):
+    """phase-only and both must differ — the attenuation switches are live."""
+    pytest.importorskip("sweep")
+    pytest.importorskip("torch")
+
+    import sweep_agent.tools.visco_forward_sweep as vf
+
+    nz, nx = 80, 120
+    vp = np.full((nz, nx), 2000.0, dtype=np.float32)
+    vp[40:, :] = 2600.0
+    vp_path = tmp_path / "vp.npy"
+    np.save(vp_path, vp)
+
+    phase_only = _run(vf, tmp_path, vp_path, phase_shift=True, amplitude_damping=False)
+    both = _run(vf, tmp_path, vp_path, phase_shift=True, amplitude_damping=True)
+    acoustic = _run(vf, tmp_path, vp_path, phase_shift=False, amplitude_damping=False)
+
+    assert phase_only.shape == both.shape
+    assert np.isfinite(both).all() and np.abs(both).max() > 0
+    assert not np.allclose(phase_only, both), "amplitude_damping had no effect"
+    assert not np.allclose(acoustic, phase_only), "phase_shift had no effect"