diff --git a/src/sweep_agent/prompts.py b/src/sweep_agent/prompts.py index fcb85e7..5362526 100644 --- a/src/sweep_agent/prompts.py +++ b/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) → diff --git a/src/sweep_agent/tools/__init__.py b/src/sweep_agent/tools/__init__.py index f070fb8..96fe729 100644 --- a/src/sweep_agent/tools/__init__.py +++ b/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"] diff --git a/src/sweep_agent/tools/selection.py b/src/sweep_agent/tools/selection.py index 023ddfd..07a8365 100644 --- a/src/sweep_agent/tools/selection.py +++ b/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", ] diff --git a/src/sweep_agent/tools/visco_forward_sweep.py b/src/sweep_agent/tools/visco_forward_sweep.py new file mode 100644 index 0000000..09e05eb --- /dev/null +++ b/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}.", + } \ No newline at end of file diff --git a/tests/test_visco_forward_sweep.py b/tests/test_visco_forward_sweep.py new file mode 100644 index 0000000..19067fa --- /dev/null +++ b/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"