"""
pycsamt.agents.occam2d_agent
=============================
:class:`Occam2DAgent` — Write a complete Occam2D inversion input file set.
Produces four files in the output directory:
``OccamDataFile.dat``
MT response data in Occam2D format.
Written by :meth:`~pycsamt.models.occam2d.OccamData.write`.
``Occam2DMesh``
2-D FD mesh derived from the data geometry.
Written by :meth:`~pycsamt.models.occam2d.mesh.OccamMesh.write`.
``Occam2DModel``
Mesh-cell → inversion-parameter mapping.
Written by :meth:`~pycsamt.models.occam2d.model.OccamModel.write`.
``OccamStartup``
Starting-model and inversion parameters file.
Written by :meth:`~pycsamt.models.occam2d.startup.OccamStartup.write`.
The agent also validates the written files and returns a brief data summary
so the user can verify the output before running Occam2D.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
from typing import Any
from ._base import AgentResult, BaseAgent
_SYSTEM_PROMPT = """\
You are an expert in 2-D MT inversion setup using Occam2D.
Given the data file statistics and mesh parameters, write 3–4 sentences that:
1. Confirm the data file contains the expected stations and period bands.
2. Comment on the mesh geometry (cell sizes, depth extent, padding).
3. Recommend regularisation parameters (roughness penalty, target RMS).
4. Note any data gaps or stations that should be excluded.
Reply in plain English.
"""
[docs]
class Occam2DAgent(BaseAgent):
"""Generate a complete Occam2D inversion input file set.
Parameters
----------
api_key, model, llm_provider : str
modes : list of str or None
Occam2D mode codes. Common choices:
``["ZXXR","ZXXI","ZXYR","ZXYI","ZYXR","ZYXI","ZYYR","ZYYI"]`` (all Z),
``["RhoZXY","PhsZXY","RhoZYX","PhsZYX"]`` (ρa/φ off-diagonal).
Default: ``None`` → OccamData auto-selects from available data.
error_floor : float
Minimum relative error floor applied to all data (default 0.05 = 5 %).
target_rms : float
Target RMS for the startup file (default 1.0).
Input keys
----------
``sites`` / ``path`` : Sites or str
``output_dir`` : str
``period_range`` : [T_min, T_max], optional
``modes`` : list of str, optional — overrides constructor default
``title`` : str, optional
Output data keys
----------------
``data_path`` Path — OccamDataFile.dat
``mesh_path`` Path — Occam2DMesh
``model_path`` Path — Occam2DModel
``startup_path`` Path — OccamStartup
``n_stations`` int
``n_periods`` int
``n_data`` int
``output_dir`` str
Examples
--------
>>> agent = Occam2DAgent()
>>> result = agent.execute({
... "path": "/data/L22PLT",
... "output_dir": "/out/occam2d",
... })
>>> print(result["data_path"])
/out/occam2d/OccamDataFile.dat
"""
SYSTEM_PROMPT = _SYSTEM_PROMPT
def __init__(
self,
*,
api_key: str | None = None,
model: str | None = None,
llm_provider: str = "claude",
modes: list[str] | None = None,
error_floor: float = 0.05,
target_rms: float = 1.0,
) -> None:
super().__init__(
"Occam2DAgent",
api_key=api_key,
model=model,
llm_provider=llm_provider,
section_preset="inversion",
)
self.modes = modes
self.error_floor = error_floor
self.target_rms = target_rms
[docs]
def execute(self, input_data: dict[str, Any]) -> AgentResult:
self._last_cost = 0.0
t0 = time.time()
warnings: list[str] = []
# ── resolve sites ─────────────────────────────────────────────────────
from ..emtools._core import ensure_sites
sites_raw = input_data.get("sites") or input_data.get("path")
if sites_raw is None:
return AgentResult.failed(
"No 'sites' or 'path'.", elapsed=time.time() - t0
)
try:
sites = ensure_sites(sites_raw, verbose=0)
except Exception as exc:
return AgentResult.failed(str(exc), elapsed=time.time() - t0)
output_dir = input_data.get("output_dir", "pycsamt_occam2d")
input_data.get("period_range")
modes = input_data.get("modes", self.modes)
title = input_data.get(
"title", "pycsamt Occam2D — generated by OccamAgent"
)
os.makedirs(output_dir, exist_ok=True)
# ── import Occam2D model layer ─────────────────────────────────────────
try:
from ..models.occam2d import OccamData
from ..models.occam2d.config import OccamConfig
from ..models.occam2d.mesh import OccamMesh
from ..models.occam2d.model import OccamModel
from ..models.occam2d.startup import OccamStartup
except ImportError as exc:
return AgentResult.failed(
f"pycsamt.models.occam2d not available: {exc}",
elapsed=time.time() - t0,
)
# ── build OccamData ───────────────────────────────────────────────────
data_path: Path | None = None
occ_data: Any = None
n_sites = n_periods = n_data = 0
try:
cfg = OccamConfig(
error_floor_rho=self.error_floor,
error_floor_phase=self.error_floor * 10.0,
target_misfit=self.target_rms,
)
occ_data = OccamData.from_edi(
sites,
modes=modes,
config=cfg,
title=title,
)
n_sites = int(occ_data.n_sites)
n_periods = int(occ_data.n_frequencies)
n_data = int(occ_data.n_data)
data_path = Path(output_dir) / "OccamDataFile.dat"
occ_data.write(data_path)
self._log.info(
"OccamDataFile.dat written: %d sites × %d freqs",
n_sites,
n_periods,
)
except Exception as exc:
return AgentResult.failed(
f"OccamData.from_edi failed: {exc}",
hint="Check that the EDIs contain valid Z data and coordinates.",
elapsed=time.time() - t0,
)
# ── build OccamMesh ───────────────────────────────────────────────────
mesh_path: Path | None = None
occ_mesh: Any = None
try:
occ_mesh = OccamMesh.from_data(occ_data)
mesh_path = Path(output_dir) / "Occam2DMesh"
occ_mesh.write(mesh_path)
self._log.info(
"Occam2DMesh written: %d × %d cells",
getattr(occ_mesh, "n_xcells", "?"),
getattr(occ_mesh, "n_zcells", "?"),
)
except Exception as exc:
warnings.append(f"OccamMesh.from_data failed: {exc}")
# ── build OccamModel (parameterisation) ───────────────────────────────
model_path: Path | None = None
occ_model: Any = None
try:
if occ_mesh is not None:
occ_model = OccamModel.from_mesh(occ_mesh, config=cfg)
model_path = Path(output_dir) / "Occam2DModel"
occ_model.write(model_path)
self._log.info(
"Occam2DModel written: %d params",
getattr(occ_model, "n_params", 0),
)
except Exception as exc:
warnings.append(f"OccamModel.from_mesh failed: {exc}")
occ_model = None
model_path = None
# ── build OccamStartup ────────────────────────────────────────────────
startup_path: Path | None = None
try:
if occ_model is None:
occ_model = OccamModel(description="SMOOTH INVERSION")
occ_startup = OccamStartup.from_model(occ_model, config=cfg)
startup_path = Path(output_dir) / "OccamStartup"
occ_startup.write(startup_path)
self._log.info("OccamStartup written")
except Exception as exc:
warnings.append(f"OccamStartup.from_model failed: {exc}")
# ── validate file sizes ───────────────────────────────────────────────
for label, p in [
("data", data_path),
("mesh", mesh_path),
("model", model_path),
("startup", startup_path),
]:
if p and p.exists():
size_kb = p.stat().st_size // 1024
self._log.info("%s file: %s (%d KB)", label, p.name, size_kb)
elif p:
warnings.append(f"{label} file not found after write: {p}")
# ── LLM interpretation ────────────────────────────────────────────────
interp: str | None = None
if self.api_key and n_sites:
n_x = getattr(occ_mesh, "n_xcells", "?") if occ_mesh else "?"
n_z = getattr(occ_mesh, "n_zcells", "?") if occ_mesh else "?"
prompt = (
f"Occam2D input file summary:\n"
f" Stations: {n_sites}, periods: {n_periods}, data points: {n_data}\n"
f" Mesh: {n_x} × {n_z} cells\n"
f" Error floor: {self.error_floor:.0%}\n"
f" Modes: {modes or 'auto'}\n"
f" Warnings: {warnings[:3] if warnings else 'none'}\n\n"
"Advise on the inversion setup and recommend regularisation."
)
interp = self.query_llm(prompt, max_tokens=200)
elapsed = time.time() - t0
files_written = sum(
1
for p in [data_path, mesh_path, model_path, startup_path]
if p and p.exists()
)
return AgentResult(
status="success"
if data_path and data_path.exists()
else "needs_review",
summary=(
f"Occam2D prep: {n_sites} stations × {n_periods} periods. "
f"{files_written}/4 files written to {output_dir}."
),
data={
"data_path": data_path,
"mesh_path": mesh_path,
"model_path": model_path,
"startup_path": startup_path,
"n_stations": n_sites,
"n_periods": n_periods,
"n_data": n_data,
"output_dir": output_dir,
},
warnings=warnings,
llm_interpretation=interp,
elapsed_seconds=elapsed,
cost_estimate_usd=self._last_cost,
)
__all__ = ["Occam2DAgent"]
