"""
pycsamt.agents.code_gen
========================
:class:`CodeGenerationAgent` — Emit a reproducible Python script from a
completed workflow configuration and agent results.
The generated script imports pycsamt, replicates every processing step that
was executed, and is self-contained (no agent infrastructure needed to run
it). It is the "save your workflow as code" button.
"""
from __future__ import annotations
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from ._base import AgentResult, BaseAgent
_SYSTEM_PROMPT = """\
You are an expert Python developer specialising in geophysics scripting.
Given a pycsamt workflow configuration and execution log, generate a clean,
well-commented Python script that reproduces the workflow step by step.
Use pycsamt v2 public API only. Add a one-line comment above each major
block. Do not use the agents/ subpackage — call emtools, forward, and
site functions directly. Output only valid Python code.
"""
# ── static script template fragments ─────────────────────────────────────────
_HEADER = """\
#!/usr/bin/env python
# -*- coding: utf-8 -*-
\"\"\"
{title}
Generated by pycsamt CodeGenerationAgent — {date}
This script reproduces the processing workflow automatically.
\"\"\"
from __future__ import annotations
import numpy as np
import matplotlib.pyplot as plt
"""
_LOAD_BLOCK = """\
# ── load MT data ──────────────────────────────────────────────
from pycsamt.emtools._core import ensure_sites
sites = ensure_sites({path!r}, verbose=0)
print(f"Loaded {{len(list(sites))}} stations")
"""
_QC_BLOCK = """\
# ── quality control ───────────────────────────────────────────
from pycsamt.emtools.qc import (
build_qc_table,
qc_flags,
plot_frequency_confidence_psection,
)
qc_table = build_qc_table(sites)
flags = qc_flags(sites, min_frac_ok=0.6, min_snr_med=2.0)
fig_qc = (
plot_frequency_confidence_psection(sites).get_figure()
)
fig_qc.savefig(
{out!r} + "/qc_confidence.png",
dpi=150, bbox_inches="tight",
)
"""
_SS_BLOCK = """\
# ── static-shift correction ───────────────────────────────────
from pycsamt.emtools.ss import (
estimate_ss_ama,
correct_ss_ama,
plot_ss_summary,
plot_ss_1d_curves,
)
from pycsamt.emtools._core import (
_get_z_block, _name, _iter_items,
)
ss_table = estimate_ss_ama(
sites,
half_window={hw},
sort_by="lon",
weights="tri",
max_skew=6.0,
)
print(
ss_table[
["station", "delta_log10_rho", "fac_z"]
]
)
sites_corr = correct_ss_ama(
sites,
half_window={hw},
sort_by="lon",
)
def _collect_logRho(S):
rows, freqs = [], None
for i, ed in enumerate(_iter_items(S)):
Z, z, fr = _get_z_block(ed)
if Z is None:
continue
rxy = (
0.2 * np.abs(z[:, 0, 1]) ** 2
/ (fr + 1e-24)
)
ryx = (
0.2 * np.abs(z[:, 1, 0]) ** 2
/ (fr + 1e-24)
)
rows.append(
np.log10(np.sqrt(rxy * ryx) + 1e-24)
)
freqs = fr
return np.array(rows), freqs
logRho_b, ss_freqs = _collect_logRho(sites)
logRho_a, _ = _collect_logRho(sites_corr)
ss_labels = [
_name(ed, i)
for i, ed in enumerate(_iter_items(sites))
]
fig_sum = plot_ss_summary(
logRho_b, logRho_a,
freqs=ss_freqs,
station_labels=ss_labels,
)
fig_sum.savefig(
{out!r} + "/ss_summary.png",
dpi=150, bbox_inches="tight",
)
fig_1d = plot_ss_1d_curves(
logRho_b, logRho_a,
freqs=ss_freqs,
station_labels=ss_labels,
)
fig_1d.savefig(
{out!r} + "/ss_1d_curves.png",
dpi=150, bbox_inches="tight",
)
"""
_PT_BLOCK = """\
# ── phase tensor analysis ─────────────────────────────────────
from pycsamt.emtools.tensor import (
build_phase_tensor_table,
plot_phase_tensor_psection,
)
from pycsamt.emtools.strike import (
estimate_strike_consensus,
plot_strike_analysis,
)
from pycsamt.emtools.dimensionality import classify_dimensionality
pt_table = build_phase_tensor_table(sites_corr)
dim_table = classify_dimensionality(
sites_corr, skew_th={skew_th}, ellipt_th={ellipt_th},
)
st_result = estimate_strike_consensus(sites_corr)
strike_angle = st_result.get("angle_deg")
try:
strike_text = f"{{float(strike_angle):.1f}} deg"
except (TypeError, ValueError):
strike_text = "n/a"
print(f"Consensus strike: {{strike_text}}")
fig_pt = plot_phase_tensor_psection(sites_corr).get_figure()
fig_pt.savefig(
{out!r} + "/pt_psection.png",
dpi=150, bbox_inches="tight",
)
fig_strike = plot_strike_analysis(sites_corr)
fig_strike.savefig(
{out!r} + "/strike_analysis.png",
dpi=150, bbox_inches="tight",
)
"""
_FWD_BLOCK = """\
# ── 1-D MT forward model ──────────────────────────────────────
from pycsamt.forward import (
MT1DForward,
LayeredModel,
plot_response_and_model_1d,
)
layered = LayeredModel(
resistivities={rhos},
thicknesses={ths},
)
freqs = np.logspace(-4, 3, 60)
fwd = MT1DForward(freqs=freqs)
resp = fwd.run(layered)
fig_fwd = plot_response_and_model_1d(resp, layered)
fig_fwd.savefig(
{out!r} + "/forward_response.png",
dpi=150, bbox_inches="tight",
)
"""
_PRE_INV_BLOCK = """\
# Occam2D pre-inversion export
from pycsamt.models.occam2d import InputBuilder, OccamConfig
occam_workdir = os.path.join({out!r}, "occam2d_inputs")
os.makedirs(occam_workdir, exist_ok=True)
occam_config = OccamConfig()
occam_config.to_template(os.path.join(occam_workdir, "occam2d.yml"))
occam_builder = InputBuilder(
sites_corr,
workdir=occam_workdir,
config=occam_config,
)
# Expected Occam2D artifacts: OccamDataFile.dat, Occam2DMesh,
# Occam2DModel, and OccamStartup. The builder records the public API
# used for pre-inversion preparation.
print(f"Occam2D pre-inversion configuration written to {{occam_workdir}}")
"""
_TIPPER_BLOCK = """\
# Tipper and induction-arrow analysis
from pycsamt.emtools.tf import (
plot_induction_arrows,
plot_tipper_hodograms,
plot_tipper_polar,
)
try:
ax_tip = plot_induction_arrows(
sites_corr,
periods=({period},),
convention={convention!r},
)
ax_tip.get_figure().savefig(
{out!r} + "/induction_arrows.png",
dpi=150, bbox_inches="tight",
)
except Exception as exc:
print(f"No plottable induction arrows: {{exc}}")
try:
ax_hodo = plot_tipper_hodograms(sites_corr)
ax_hodo.get_figure().savefig(
{out!r} + "/tipper_hodograms.png",
dpi=150, bbox_inches="tight",
)
except Exception as exc:
print(f"Tipper hodogram skipped: {{exc}}")
try:
ax_polar = plot_tipper_polar(sites_corr, component={component!r})
ax_polar.get_figure().savefig(
{out!r} + "/tipper_polar.png",
dpi=150, bbox_inches="tight",
)
except Exception as exc:
print(f"Tipper polar plot skipped: {{exc}}")
"""
_SENSITIVITY_BLOCK = """\
# Depth-of-investigation and sensitivity analysis
from pycsamt.emtools.csumt import (
depth_coverage_table,
plot_depth_section,
)
doi_table = depth_coverage_table(sites_corr)
doi_table.to_csv(
{out!r} + "/depth_of_investigation.csv",
index=False,
)
ax_doi = plot_depth_section(sites_corr)
ax_doi.get_figure().savefig(
{out!r} + "/depth_of_investigation.png",
dpi=150, bbox_inches="tight",
)
print("Depth-of-investigation and sensitivity outputs written.")
"""
_AI_INV_BLOCK = """\
# ── AI 1-D inversion ──────────────────────────────────────────
from pycsamt.ai.inversion import EMInverter1D
from pycsamt.forward.batch import generate_dataset
freqs_ai = {freqs!r}
n_layers = {n_layers}
n_samples = {n_samples}
epochs = {epochs}
ds = generate_dataset(
solver="mt1d",
n_samples=n_samples,
freqs=freqs_ai,
n_layers=n_layers,
noise_level=0.03,
seed=42,
n_jobs=1,
verbose=False,
)
inv1d = EMInverter1D(
arch={arch!r},
n_layers=n_layers,
solver="mt1d",
)
inv1d.fit(ds.X, ds.y, epochs=epochs, batch_size=32, verbose=False)
import numpy as np
obs_features = {{
f"synthetic_{{i:02d}}": feat
for i, feat in enumerate(ds.X[: min(3, len(ds.X))])
}}
predictions = {{}}
for site_name, feat in obs_features.items():
pred = inv1d.predict(feat.reshape(1, -1))
predictions[site_name] = pred[0]
print(f"AI 1-D inversion done: {{len(predictions)}} stations")
"""
_ENSEMBLE_BLOCK = """\
# ── ensemble 1-D inversion with uncertainty ───────────────────
from pycsamt.ai.inversion import EnsembleInverter
ens = EnsembleInverter(
n_members={n_members},
arch={arch!r},
n_layers={n_layers},
solver="mt1d",
)
ens.fit(ds.X, ds.y, epochs={epochs}, verbose=False)
ensemble_preds = {{}}
ensemble_unc = {{}}
for site_name, feat in obs_features.items():
mean, std = ens.predict_with_uncertainty(
feat.reshape(1, -1)
)
ensemble_preds[site_name] = mean[0]
ensemble_unc[site_name] = std[0]
print(
f"Ensemble inversion done: "
f"{{len(ensemble_preds)}} stations"
)
"""
_INV2D_BLOCK = """\
# ── U-Net 2-D profile inversion ───────────────────────────────
from pycsamt.ai.inversion import EMInverter2D
from scipy.ndimage import zoom
n_depth = {n_depth}
n_freqs_2d = {n_freqs}
n_components = {n_components}
n_sta = len(obs_features)
X_obs_list = list(obs_features.values())
X_obs = np.stack(X_obs_list, axis=2) # (n_freqs, n_comp, n_sta)
X_obs = X_obs.transpose(1, 0, 2) # (n_comp, n_freqs, n_sta)
X_obs_4d = X_obs[None, ...]
inv2d = EMInverter2D(
n_components=n_components,
n_depth=n_depth,
n_stations=n_sta,
n_freqs=n_freqs_2d,
arch={arch!r},
)
inv2d.fit(
X_train_2d, y_train_2d,
epochs={epochs},
batch_size=8,
verbose=False,
)
pred_section = inv2d.predict(X_obs_4d)[0] # (n_depth, n_sta)
print(
f"2-D inversion section shape: "
f"{{pred_section.shape}}"
)
"""
_FULL_AI_BLOCK = """\
# ── full AI workflow summary ───────────────────────────────────
import json, hashlib, platform, sys
from importlib.metadata import version as _v
provenance = {{
"workflow": "full_ai_workflow",
"data_path": {path!r},
"output_dir": {out!r},
"python": sys.version,
"platform": platform.platform(),
"packages": {{
"pycsamt": _v("pycsamt"),
"numpy": _v("numpy"),
}},
}}
prov_path = {out!r} + "/provenance.json"
with open(prov_path, "w") as fh:
json.dump(provenance, fh, indent=2)
print(f"Provenance written to {{prov_path}}")
"""
_FOOTER = """\
plt.show()
print("Workflow complete.")
"""
[docs]
class CodeGenerationAgent(BaseAgent):
"""Generate a reproducible Python script from a completed workflow.
Parameters
----------
api_key, model, llm_provider : str
When an API key is provided the LLM refines and annotates the
generated code. Otherwise the agent uses static templates.
script_title : str
Input keys
----------
``workflow_config`` : dict
The config dict produced by :class:`ContextInputAgent`.
``results`` : dict
The agent results dict from :class:`AgentCoordinator`.
``output_dir`` : str, optional
Output data keys
----------------
``code`` str — Python source code
``script_path`` str or None — path to saved .py file
Examples
--------
>>> agent = CodeGenerationAgent()
>>> result = agent.execute({
... "workflow_config": cfg,
... "results": coord_results,
... "output_dir": "/out",
... })
>>> print(result["script_path"])
/out/workflow_script.py
"""
SYSTEM_PROMPT = _SYSTEM_PROMPT
def __init__(
self,
*,
api_key: str | None = None,
model: str | None = None,
llm_provider: str = "claude",
script_title: str = "pycsamt MT Processing Workflow",
) -> None:
super().__init__(
"CodeGenerationAgent",
api_key=api_key,
model=model,
llm_provider=llm_provider,
)
self.script_title = script_title
[docs]
def execute(self, input_data: dict[str, Any]) -> AgentResult:
self._last_cost = 0.0
t0 = time.time()
warnings: list[str] = []
cfg = input_data.get("workflow_config") or {}
results = input_data.get("results") or {}
output_dir = input_data.get("output_dir", ".")
title = input_data.get("title", self.script_title)
# Optional RAG context (real symbols / recipe) to ground the LLM
# refinement pass; ignored offline.
rag_context = input_data.get("rag_context") or ""
os.makedirs(output_dir, exist_ok=True)
data_path = cfg.get("data_path", "/path/to/EDIs")
workflow = cfg.get("workflow", "qc")
out_dir = cfg.get("output_dir", output_dir)
# ── build script from templates ───────────────────────────────────────
code = _HEADER.format(
title=title,
date=datetime.now().strftime("%Y-%m-%d"),
)
code += f"import os\nos.makedirs({out_dir!r}, exist_ok=True)\n\n"
code += _LOAD_BLOCK.format(path=data_path)
# add blocks based on workflow + available results
if workflow in ("qc", "full") or "qc" in results:
code += _QC_BLOCK.format(out=out_dir)
if workflow in ("static_shift", "full") or "static_shift" in results:
hw = 3
results.get("static_shift")
code += _SS_BLOCK.format(hw=hw, out=out_dir)
else:
code += "sites_corr = sites # no static-shift correction\n\n"
if (
workflow in ("phase_analysis", "full")
or "phase_analysis" in results
):
results.get("phase_analysis")
skew_th = 5.0
ellipt_th = 0.1
code += _PT_BLOCK.format(
skew_th=skew_th,
ellipt_th=ellipt_th,
out=out_dir,
)
if "forward" in results or workflow == "forward":
fwd_r = results.get("forward")
rhos = [100, 10, 1000, 100]
ths = [500, 1000, 2000]
if fwd_r is not None:
lm = fwd_r.get("layered_model")
if lm is not None:
rhos = list(
getattr(
lm,
"resistivity",
getattr(lm, "resistivities", rhos),
)
)
ths = list(
getattr(
lm,
"thickness",
getattr(lm, "thicknesses", ths),
)
)
code += _FWD_BLOCK.format(rhos=rhos, ths=ths, out=out_dir)
if workflow in ("pre_inversion", "occam2d") or "pre_inversion" in results:
code += _PRE_INV_BLOCK.format(out=out_dir)
if workflow in ("tipper", "tipper_plot") or "tipper" in results:
code += _TIPPER_BLOCK.format(
out=out_dir,
period=float(cfg.get("period", 1.0)),
component=str(cfg.get("tipper_component", "real")),
convention=str(cfg.get("tipper_convention", "park")),
)
if workflow == "sensitivity" or "sensitivity" in results:
code += _SENSITIVITY_BLOCK.format(out=out_dir)
# ── AI inversion blocks ───────────────────────────────────
ai_wf = {
"ai_inversion",
"inv1d",
"ensemble_inversion",
"inv2d",
"full_ai_workflow",
"full",
}
if workflow in ai_wf or "ai_inv" in results:
inv_r = results.get("ai_inv") or {}
n_layers = int(inv_r.get("n_layers", cfg.get("n_layers", 5)))
n_samples = int(cfg.get("n_train_samples", 2000))
epochs = int(cfg.get("epochs", 50))
arch = str(cfg.get("arch", "cnn1d"))
import numpy as np
freqs_ai = list(np.logspace(-4, 3, 32).round(6).tolist())
code += _AI_INV_BLOCK.format(
freqs=freqs_ai,
n_layers=n_layers,
n_samples=n_samples,
epochs=epochs,
arch=arch,
)
if workflow in ("ensemble_inversion",) or "ensemble" in results:
results.get("ensemble") or {}
n_layers = int(cfg.get("n_layers", 5))
arch = str(cfg.get("arch", "cnn1d"))
code += _ENSEMBLE_BLOCK.format(
n_members=int(cfg.get("n_members", 5)),
arch=arch,
n_layers=n_layers,
epochs=int(cfg.get("epochs", 50)),
)
if workflow in ("inv2d", "full_ai_workflow"):
code += _INV2D_BLOCK.format(
n_depth=int(cfg.get("n_depth", 40)),
n_freqs=int(cfg.get("n_freqs", 32)),
n_components=4,
arch=str(cfg.get("arch", "unet")),
epochs=int(cfg.get("epochs", 30)),
)
if workflow in ("full_ai_workflow",):
code += _FULL_AI_BLOCK.format(
path=data_path,
out=out_dir,
)
code += _FOOTER
# ── optional LLM refinement ───────────────────────────────────────────
if self.api_key:
grounding = ""
if rag_context:
grounding = (
"Use ONLY the real pyCSAMT symbols and patterns from "
"this retrieved context (do not invent functions):\n"
f"{rag_context}\n\n"
)
prompt = (
f"Refine and annotate the following pycsamt Python script. "
f"Fix any issues, add helpful inline comments, ensure "
f"all imports are correct. Return only valid Python code.\n\n"
f"{grounding}"
f"```python\n{code}\n```"
)
llm_code = self.query_llm(prompt, max_tokens=2000)
if llm_code:
# strip markdown fences if present
if "```python" in llm_code:
llm_code = llm_code.split("```python", 1)[1]
if "```" in llm_code:
llm_code = llm_code.rsplit("```", 1)[0]
code = llm_code.strip()
# ── write file ────────────────────────────────────────────────────────
script_path: str | None = None
try:
script_path = os.path.join(output_dir, "workflow_script.py")
Path(script_path).write_text(code, encoding="utf-8")
except Exception as exc:
warnings.append(f"Could not write script: {exc}")
script_path = None
elapsed = time.time() - t0
n_lines = code.count("\n")
return AgentResult(
status="success",
summary=(
f"Generated {n_lines}-line Python script for "
f"workflow={workflow!r}. "
+ (f"Saved to {script_path}" if script_path else "")
),
data={"code": code, "script_path": script_path},
warnings=warnings,
elapsed_seconds=elapsed,
cost_estimate_usd=self._last_cost,
)
__all__ = ["CodeGenerationAgent"]
