pycsamt.agents.package_qa#

PackageQAAgent — answer free-form questions about the pycsamt v2 package.

This agent solves the “LLM does not know pycsamt” problem by:

Extracting the full pycsamt v2 API reference (docstrings + signatures + parameter tables) from the live package at import time.
Using query-adaptive injection: instead of always dumping the entire reference into the prompt, it selects the most relevant tiers based on keywords in the user’s question.

Example:
“How do I access impedance Z values?” -> injects TIER_SITES + TIER_EXAMPLES only

(not the 15-agent class dump)

“What does StaticShiftAgent do?”
-> injects TIER_AGENTS + TIER_EXAMPLES

“What workflows are supported?”
-> injects TIER_CORE only

This keeps the LLM’s signal-to-noise ratio high — the relevant information is at the top, not buried in unrelated content.

Offline mode (no api_key)#

Uses keyword-matching over extracted docstrings to return a best-effort answer. No LLM call, no cost, CI-safe.

Online mode (api_key provided)#

Sends question + selected context tiers to the LLM provider and returns the grounded answer.

Functions

inspect_full_doc(obj)

Re-export: return full docstring.

Classes

PackageQAAgent(*[, api_key, model, ...])

Answer free-form questions about pycsamt v2.

pycsamt.agents.package_qa.inspect_full_doc(obj)[source]#

Re-export: return full docstring.

Parameters:: obj (Any)
Return type:: str

class pycsamt.agents.package_qa.PackageQAAgent(*, api_key=None, model=None, llm_provider='claude', use_rag=True)[source]#

Bases: BaseAgent

Answer free-form questions about pycsamt v2.

Parameters:

api_key (str or None) – LLM key. None activates offline docstring-lookup mode (no cost, no network).
model (str) – Passed to BaseAgent.
llm_provider (str) – Passed to BaseAgent.
use_rag (bool)

Notes

Uses query-adaptive context injection: selects the most relevant documentation tiers (agents, Sites data model, usage examples) based on keywords in the question, rather than always dumping the full reference. This keeps the LLM focused on what’s relevant.

Examples

Offline (no LLM):

agent = PackageQAAgent()
r = agent.execute({
    "question": (
        "What does StaticShiftAgent do?"
    )
})
print(r["answer"])

Online (Claude):

agent = PackageQAAgent(
    api_key="sk-ant-...",
    llm_provider="claude",
)
r = agent.execute({
    "question": (
        "How do I access impedance Z?"
    )
})
print(r["answer"])

SYSTEM_PROMPT: str = 'You are a helpful expert on the pycsamt v2 Python library for magnetotelluric data processing.\n\nThe pycsamt v2 API reference extracted from the live package is provided below. Use it as your *primary source of truth*. If an answer is not in the reference, say so — do not invent class names, parameters, or behaviours that are not listed.\n\n---\nPYCSAMT v2 — PACKAGE OVERVIEW\n==============================\npycsamt is a Python library for magnetotelluric (MT/AMT/CSAMT) data processing, correction, and AI-assisted inversion. Agents are self-contained workflow components; Sites is the main data container for loaded EDI data.\n\nWorkflow keywords (pass as config["workflow"])\n----------------------------------------------\nqc Data quality control + station flagging\nstatic_shift Galvanic static-shift detection + correction\nphase_analysis Phase tensor, dimensionality, strike analysis\nforward Forward model computation\npre_inversion Pre-processing before inversion\ninversion_eval Evaluate inversion results\ninterpretation Geological interpretation of resistivity models\nreport Generate PDF/HTML processing report\nfull Full pipeline: load -> qc -> correct -> invert -> report\nai_inversion 1-D neural network / CNN inversion (EMInverter1D)\ninv1d Alias for ai_inversion\ninv2d 2-D U-Net profile inversion\ninv3d 3-D GCN graph-convolutional inversion\nensemble_inversion Ensemble / uncertainty-quantified inversion\npinn_inversion Physics-Informed Neural Network inversion\nhybrid_inversion PINN + AI inverter hybrid (requires checkpoint)\norchestrated_code_gen Generate reproducible Python script for a workflow\ntipper Tipper / induction arrow analysis\nmodem ModEM 3-D inversion interface\noccam2d Occam2D regularised 2-D inversion interface\n\n\nAgent class reference\n=====================\nWorkflowOrchestratorAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', default_workflow: \'str\' = \n Intelligently route an NL request to the correct agent chain.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n default_workflow : str\n Fallback when NL classification is ambiguous (default ``"qc"``).\n Input keys:\n ----------\n ``request`` : str — natural-language processing request\n ``config`` : dict, optional — pre-built config (skips NL parsing)\n ``dry_run`` : bool — preview without executing (default False)\n ``output_dir`` : str\n Output data keys:\n ----------------\n ``workflow_type`` str\n ``reasoning`` str\n ``coordinator`` AgentCoordinator instance\n ``result`` AgentResult from the coordinator\n\nContextInputAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\') -> \'None\'\n Parse a natural-language MT workflow request into a structured config.\n Parameters:\n ----------\n api_key : str or None\n LLM API key. When ``None`` the regex fallback is used exclusively.\n model, llm_provider : str\n Passed to :class:`~pycsamt.agents._base.BaseAgent`.\n\nMTLoaderAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', recursive: \'bool\' = True, \n Load MT data from any pycsamt-supported format and assess quality.\n Parameters:\n ----------\n api_key : str or None\n model, llm_provider : str\n recursive : bool\n When loading a directory, recurse into sub-directories.\n on_dup : str\n Duplicate-station handling: ``"replace"`` (default) or ``"skip"``.\n\nDataQCAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', method: \'str\' = \'composite\n Run data quality control on a MT/AMT dataset.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n LLM configuration (optional).\n method : str\n Confidence scoring method: ``"composite"`` (default), ``"presence"``,\n ``"snr"``, or ``"spatial"``.\n min_frac_ok : float\n Minimum fraction of OK frequencies for a station to pass (0–1).\n min_snr_med : float\n Minimum median SNR for a station to pass.\n max_skew_med : float\n Maximum median |β| skewness for a station to pass.\n Input keys:\n ----------\n ``sites`` : Sites or ``path`` : str\n EDI data to assess.\n ``output_dir`` : str, optional\n Where to save QC figures.\n ``period_range`` : [T_min, T_max], optional\n Restrict QC to this period window.\n Output data keys:\n ----------------\n ``qc_table`` pandas DataFrame — per-station metrics\n ``flags`` pandas DataFrame — pass / fail per station\n ``confidence_table`` pandas DataFrame — per-station confidence scores\n ``freq_conf_table`` pandas DataFrame — per-frequency confidence\n ``n_flagged`` int\n ``flagged_stations`` list[str]\n ``figures`` dict — matplotlib Figure objects\n\nStaticShiftAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', method: \'str\' = \'ama\', hal\n Detect and correct galvanic static shift in MT/AMT data.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n method : {"ama", "loess", "refmedian", "bilateral"}\n Correction algorithm. Default ``"ama"`` (adaptive moving average).\n half_window : int\n Spatial half-window for AMA / LOESS smoothing.\n pband : (T_min, T_max) or None\n Period band used to estimate shift factors.\n inplace : bool\n Modify the input Sites in-place. Default ``False`` (returns a copy).\n Input keys:\n ----------\n ``sites`` / ``path`` : Sites or str\n ``method`` : str, optional — overrides constructor default\n Output data keys:\n ----------------\n ``corrected_sites`` Sites with static shift removed\n ``shift_factors`` dict {station: factor}\n ``rho_before`` ndarray (n_freq × n_sta) — log₁₀ ρa before\n ``rho_after`` ndarray — log₁₀ ρa after\n ``delta_stats`` dict — min/max/mean shift magnitude\n ``figures`` dict — matplotlib Figure objects\n\nDenoisingAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', method: \'str\' = \'rpca\', ra\n Denoise MT impedance data using classical or AI-based methods.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n method : str\n ``"rpca"`` (default), ``"hampel"``, ``"emap"``, ``"pipeline"``,\n or ``"ai"`` / ``"ai_cae"`` (requires PyTorch/TF).\n rank : int\n RPCA rank for off-diagonal denoising (default 2).\n half_window : int\n Hampel filter half-window (default 3).\n Input keys:\n ----------\n ``sites`` / ``path`` : Sites or str\n ``method`` : str, optional — overrides constructor default\n ``output_dir`` : str, optional\n Output data keys:\n ----------------\n ``denoised_sites`` Sites with denoised impedance\n ``snr_before`` ndarray — per-(station, freq) SNR proxy before\n ``snr_after`` ndarray — per-(station, freq) SNR proxy after\n ``snr_gain`` float — mean SNR improvement\n ``n_recovered`` int — frequencies recovered above SNR threshold\n ``figures`` dict\n ``figure_paths`` dict\n\nPhaseAnalysisAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', skew_th: \'float\' = 5.0, el\n Run a full phase tensor, strike, and dimensionality survey analysis.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n skew_th : float\n Skewness |β| threshold for 3-D classification (°).\n ellipt_th : float\n Ellipticity λ threshold for 2-D classification.\n band : (T_min, T_max) or None\n Period band for strike estimation.\n Input keys:\n ----------\n ``sites`` / ``path`` : Sites or str\n ``period_range`` : [T_min, T_max], optional\n ``output_dir`` : str, optional\n ``run_mohr`` : bool, optional — also produce Mohr circles (default False)\n Output data keys:\n ----------------\n ``pt_table`` pandas DataFrame — full PT metrics per (station, period)\n ``strike_consensus`` float — consensus strike angle (°)\n ``strike_iqr`` float — IQR of strike across all stations\n ``dim_table`` pandas DataFrame — per-(station, period) classification\n ``n_1d``, ``n_2d``, ``n_3d`` int — count of observations per class\n ``figures`` dict — matplotlib Figure objects\n\nForwardModelAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', dim: \'int\' = 1, freqs: \'An\n Run a 1-D, 2-D, or 3-D MT forward model.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n dim : {1, 2, 3}\n Forward solver dimensionality.\n freqs : array-like or None\n Frequencies (Hz). Defaults to 40 log-spaced points 10⁻⁴–10³ Hz.\n Input keys:\n ----------\n ``model`` : dict or LayeredModel or None\n **1-D / 2-D from 1-D layers:**\n ``{"resistivities": [...], "thicknesses": [...]}``.\n \n **2-D grid type override:**\n add ``"type": "halfspace" | "anomaly"`` and grid parameters such as\n ``"bg_rho"``, ``"anomaly_rho"``, ``"anomaly_bounds"``.\n Output data keys:\n ----------------\n ``dim`` int\n ``layered_model`` LayeredModel (1-D / 2-D from 1-D)\n ``grid`` Grid2D or Grid3D (2-D / 3-D)\n ``response`` ForwardResponse / ForwardResponse2D / ForwardResponse3D\n ``rho_a`` ndarray — 1-D ρa\n ``phase`` ndarray — 1-D phase (°)\n ``rho_a_te`` ndarray (n_freqs, n_stations) — 2-D TE\n\nInterpretationAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', context: \'str\' = \'\') -> \'N\n Interpret a resistivity model in terms of geological formations.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n context : str\n Optional geological context passed to the LLM\n (e.g. ``"Semi-arid Precambrian terrain, looking for aquifers"``).\n Input keys:\n ----------\n ``model`` : dict or LayeredModel\n ``{"resistivities": [...], "thicknesses": [...]}``\n ``rms`` : float, optional\n ``context`` : str, optional — overrides constructor default\n ``sites`` / ``path`` : optional — for period range context\n Output data keys:\n ----------------\n ``layer_interpretations`` list of dict\n ``summary_text`` str\n ``dominant_lithology`` str\n ``formation_depths_m`` list[float]\n\nReportAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', report_title: \'str\' = \'MT/\n Generate a structured MT survey report from agent results.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n report_title : str\n Title for the report.\n formats : list of {"md", "html", "pdf"}\n Output formats. Default ``["md", "html"]``.\n Input keys:\n ----------\n ``results`` : dict\n Keyed by agent step name → :class:`AgentResult`.\n Expected keys: ``"load"``, ``"qc"``, ``"static_shift"``,\n ``"phase_analysis"``, ``"forward"`` (all optional).\n ``output_dir`` : str\n Output data keys:\n ----------------\n ``report_md`` str — full markdown text\n ``report_html`` str or None\n ``report_path_md`` str — path to .md file\n ``report_path_html`` str or None\n\nAIInversionAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', arch: \'str\' = \'resnet\', n_\n Train an AI inverter on synthetic data then predict on observed sites.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n arch : {"resnet", "cnn1d", "fcn"}\n Neural network architecture.\n n_layers : int\n Number of model layers the inverter will predict (default 5).\n n_train_samples : int\n Number of synthetic training samples (default 2 000).\n epochs : int\n Training epochs (default 30). Increase for better models.\n freqs : array-like or None\n Frequencies used for both training synthesis and observed data\n Input keys:\n ----------\n ``sites`` / ``path`` : Sites or str — observed data\n ``output_dir`` : str, optional\n Output data keys:\n ----------------\n ``inverter`` :class:`~pycsamt.ai.inversion.inv1d.EMInverter1D`\n ``predictions`` dict {station: ndarray of log₁₀ ρ values}\n ``best_model`` dict with "resistivity" and "thickness" for first station\n ``rms_per_station`` dict {station: float}\n ``rms_global`` float\n ``train_history`` dict (loss curves)\n ``figures`` dict\n\nPINNInversionAgent(*, dim: \'int\' = 1, n_layers: \'int\' = 10, depth_max: \'float\' = 2000.0, smoothness_weight: \'float\' = 0.01, lateral_weight\n PINN-based MT inversion without labelled data.\n Parameters:\n ----------\n dim : {1, 2, 3}\n Dimensionality. Default ``1``.\n n_layers : int\n Number of layers including the halfspace.\n Default ``10``.\n depth_max : float\n Maximum investigation depth in metres.\n Default ``2000.0``.\n smoothness_weight : float\n Vertical regularisation weight.\n Default ``0.01``.\n Input keys:\n ----------\n ``sites`` / ``path`` observed data\n ``output_dir`` optional figure/save dir\n Output data keys:\n ----------------\n ``inverter`` fitted inverter object\n ``section`` ndarray (n_layers, n_stations)\n log10-rho section matrix\n ``models`` list of LayeredModel (1-D)\n ``n_stations`` int\n ``rms_per_station`` dict {station: float}\n ``rms_global`` float\n\nHybridInversionAgent(*, dim: \'int\' = 1, max_iter: \'int\' = 200, smoothness_weight: \'float\' = 0.005, lateral_weight: \'float\' = 0.005, graph_we\n Two-stage AI + physics MT inversion.\n Parameters:\n ----------\n dim : {1, 2, 3}\n Dimensionality. Default ``1``.\n max_iter : int\n Physics refinement iterations (Stage 2).\n Default ``200``.\n smoothness_weight : float\n Vertical regularisation weight.\n Default ``0.005``.\n lateral_weight : float\n Lateral smoothness weight (2-D only).\n Default ``0.005``.\n Input keys:\n ----------\n ``sites`` / ``path`` observed data\n ``ai_inverter`` fitted AI inverter object\n or path to checkpoint\n ``checkpoint`` alias for ``ai_inverter``\n ``output_dir`` optional save directory\n ``dim``, ``max_iter``,\n ``smoothness_weight``,\n Output data keys:\n ----------------\n ``inverter`` fitted HybridInverterXD\n ``section`` ndarray (n_layers, n_stations)\n Stage-2 log10-rho section\n ``stage1_section`` ndarray — Stage-1 section\n ``models`` list of LayeredModel (1-D)\n ``stage1_models`` list of LayeredModel (1-D)\n ``n_stations`` int\n\nInv2DAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', n_depth: \'int\' = 40, n_fre\n 2-D MT profile inversion using a U-Net convolutional architecture.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n n_depth : int\n Number of depth cells in the output section (default 40).\n n_freqs : int\n Number of input frequencies (default 32).\n n_components : int\n Number of impedance components in input (default 4: Re/Im × xy/yx).\n arch : str\n U-Net variant (default ``"unet"``).\n n_train_profiles : int\n Number of synthetic 2-D profiles for training (default 200).\n Input keys:\n ----------\n ``sites`` / ``path`` : Sites or str\n ``output_dir`` : str, optional\n Output data keys:\n ----------------\n ``pred_section`` ndarray (n_depth × n_stations) — log₁₀ ρ\n ``depths_km`` ndarray — depth axis (km)\n ``station_names`` list[str]\n ``rms_global`` float\n ``inverter`` EMInverter2D\n ``figures`` dict\n ``figure_paths`` dict\n\nInv3DAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', n_layers: \'int\' = 5, n_fre\n 3-D MT profile inversion using a graph-convolutional network (GCN).\n Parameters:\n ----------\n api_key, model, llm_provider : str\n n_layers : int\n Number of depth layers per station (default 5).\n n_freqs : int\n Number of frequencies used for feature extraction (default 32).\n n_train_profiles : int\n Number of synthetic 3-D training profiles (default 150).\n epochs : int\n Training epochs (default 30).\n radius : float\n Maximum inter-station edge distance in metres for the adjacency graph\n Input keys:\n ----------\n ``sites`` / ``path`` : Sites or str — observed MT dataset\n ``coords`` : ndarray (n_stations, 2), optional — station (x, y) in metres.\n Auto-extracted from EDI lat/lon when absent.\n ``adjacency`` : ndarray (n_stations, n_stations), optional — pre-computed\n normalised adjacency; overrides *radius* when supplied.\n ``output_dir`` : str, optional\n Output data keys:\n ----------------\n ``pred_rho`` ndarray (n_sta, n_layers) — log₁₀ρ\n ``pred_thick`` ndarray (n_sta, n_layers-1) — log₁₀h (metres)\n ``pred_uncertainty`` ndarray (n_sta, n_layers) or None — MC-dropout std\n ``depths_km`` ndarray — depth axis at station midpoints (km)\n ``station_names`` list[str]\n ``station_coords`` ndarray (n_sta, 2) — metres\n ``adjacency`` ndarray (n_sta, n_sta)\n\nCodeGenerationAgent(*, api_key: \'str | None\' = None, model: \'str | None\' = None, llm_provider: \'str\' = \'claude\', script_title: \'str\' = \'pyc\n Generate a reproducible Python script from a completed workflow.\n Parameters:\n ----------\n api_key, model, llm_provider : str\n When an API key is provided the LLM refines and annotates the\n generated code. Otherwise the agent uses static templates.\n Input keys:\n ----------\n ``workflow_config`` : dict\n The config dict produced by :class:`ContextInputAgent`.\n ``results`` : dict\n The agent results dict from :class:`AgentCoordinator`.\n Output data keys:\n ----------------\n ``code`` str — Python source code\n\n\nSites / EDI data model reference\n================================\nSites(edic: \'EDICollection | Sequence[EDIFile]\') -> \'None\'\n Container for multiple :class:`~pycsamt.site.base.Site` objects with convenient indexing, selection, and bulk edit operations.\n Constructor parameters:\n ----------\n edic : pycsamt.seg.collection.EDICollection or sequence of\n pycsamt.seg.edi.EDIFile\n Parsed EDI collection or any sequence of EDI objects.\n Items are wrapped into :class:`Site` instances in the\n order provided.\n Attributes:\n ----------\n _items : list of Site\n Internal sequence of sites. This is considered private.\n Iterate over ``Sites`` instead of accessing it directly.\n Key methods:\n .from_any(source: \'Any\', topo_src: \'Any | None\' = None) -> \'Sites\'\n Construct a container from heterogeneous inputs by using a normalized loading session.\n .write(outdir: \'str | Path\', *, template: \'str\' = \'{station}.edi\', exist_ok: \'bool\' = \n Write one EDI file per site to a directory.\n .select(names: \'Sequence[str] | None\' = None, predicate: \'Callable[[Site], bool] | None\n Filter sites by explicit names or by a boolean predicate.\n .edit_all(*, rename: \'Callable[[str], str] | None\' = None, freq_slice: \'slice | None\' = N\n Bulk-edit all sites with optional rename, frequency slicing, and tensor masking.\n .closest(lat: \'float\', lon: \'float\', tol: \'float | None\' = None) -> \'Site | None\'\n Find the closest site to a target coordinate using geodetic distance.\n .to_profile(origin: \'tuple[float, float]\', azimuth: \'float\', *, crs: \'int | None\' = None) -\n Convert sites to a 1D profile aligned with a specified azimuth, returning either a rich Profile object or a lightweight fallback.\n .to_edicollection(*, copy: \'bool\' = False, progress: \'bool | str\' = False, verbose: \'int\' = 0) ->\n Return the underlying EDI objects as an ``EDICollection``.\n .to_edis(*, copy: \'bool\' = False, progress: \'bool | str\' = False, verbose: \'int\' = 0) ->\n Return the underlying EDI objects as a list.\n .as_list() -> \'list[EDIFile]\'\n Return the underlying list of EDI objects.\n .get(name: \'str\') -> \'Site | None\'\n Safe lookup by case-insensitive station name.\n .by_index(i: \'int\') -> \'Site\'\n Retrieve a site by zero-based index.\n .map(fn: \'Callable[[Site], Any]\') -> \'list[Any]\'\n Apply a function to every site and collect the results.\n\nKey data attributes (set after construction)\n--------------------------------------------\nsites[i] Site object for station i\nsites[i].Z Impedance tensor: shape (n_freq, 2, 2), complex\nsites[i].freq Frequencies array (Hz)\nsites[i].rho Apparent resistivity (Ohm.m), shape (n_freq, 2, 2)\nsites[i].phase Phase (degrees), shape (n_freq, 2, 2)\nsites[i].id Station name / identifier\nlen(sites) Number of stations loaded\n\n\nensure_sites(source) — universal EDI loader\n-------------------------------------------\nAccepts: a Sites object, a directory path (str/Path),\n a list of EDI file paths, or an EDICollection.\nReturns: Sites object.\nImport: from pycsamt.ai.inversion._sites_bridge import ensure_sites\nUse case: any agent or function that needs to accept\n flexible EDI input without caring about the format.\n\n\nwrite_sites(sites, dest, exist_ok=True) — EDI exporter\n-------------------------------------------------------\nWrites one corrected EDI file per site to dest/.\nImport: from pycsamt.site.export import write_sites\nUse case: export corrected Sites after static shift or QC.\n\n\nUsage examples (copy-paste ready)\n=================================\n# 1. Load EDI files and run quality control\nfrom pycsamt.agents import MTLoaderAgent, DataQCAgent\nloader = MTLoaderAgent()\nsites = loader.execute({"path": "/data/L22PLT"})["sites"]\nqc = DataQCAgent()\nreport = qc.execute({"sites": sites})\n\n# 2. Correct static shift\nfrom pycsamt.agents import StaticShiftAgent\nagent = StaticShiftAgent(method="ama")\nresult = agent.execute({"sites": sites})\ncorrected = result["corrected_sites"]\nfactors = result["shift_factors"] # {station: factor}\n\n# 3. Full orchestrated workflow (with LLM router)\nfrom pycsamt.agents import WorkflowOrchestratorAgent\norch = WorkflowOrchestratorAgent(api_key="sk-...")\nresult = orch.execute({\n "config": {"workflow": "qc"},\n "data_path": "/data/L22PLT",\n "output_dir": "/out/qc/",\n})\n\n# 4. Load EDI collection -> Sites\nfrom pycsamt.edi import EDICollection\ncol = EDICollection("/data/L22PLT")\nsites = col.to_sites() # -> Sites object\n\n# 5. Universal loader (accepts dir, list, or Sites)\nfrom pycsamt.ai.inversion._sites_bridge import ensure_sites\nsites = ensure_sites("/data/L22PLT") # dir\nsites = ensure_sites(["a.edi", "b.edi"]) # list\n\n# 6. Access impedance tensor data\nsite0 = sites[0] # first station\nZ = site0.Z # (n_freq, 2, 2) complex\nfreq = site0.freq # frequencies (Hz)\nrho = site0.rho # apparent resistivity\nphase = site0.phase # phase (degrees)\n\n# 7. Export corrected EDI files\nfrom pycsamt.site.export import write_sites\nwrite_sites(corrected, "/out/corrected_edis/")\n\n# 8. PINN inversion\nfrom pycsamt.agents import PINNInversionAgent\nagent = PINNInversionAgent(epochs=200, n_layers=8)\nresult = agent.execute({"sites": sites})\nmodel = result["model"] # layered resistivity model\n\n# 9. Generate workflow script\nfrom pycsamt.agents import CodeGenerationAgent\ncg = CodeGenerationAgent()\nresult = cg.execute({\n "sites": sites,\n "workflow": "qc",\n "output_dir": "/out/",\n})\nprint(result["code"]) # Python source\n\n\nemtools.ss / Static-shift reference\n===================================\nread_edis -- load many EDI files\n---------------------------------\nfrom pycsamt.api import read_edis\n\nSignature:\n read_edis(sources, *, recursive=True,\n strict=False, on_dup="replace",\n progress="auto", leave=False,\n verbose=0) -> APISurvey\n\nsources : str, Path, list, or glob pattern\nReturns APISurvey\n .collection -> Sites / EDICollection\n .name -> survey name string\n\nTypical use:\n survey = read_edis("L22PLT/")\n sites = survey.collection\n\n\nestimate_ss_ama -- AMA static-shift factors\n--------------------------------------------\nfrom pycsamt.emtools.ss import estimate_ss_ama\n\nSignature:\n estimate_ss_ama(sites, *, sort_by="lon",\n half_window=3, weights="tri",\n pband=None, max_skew=6.0,\n robust_freq="median",\n robust_overall="median",\n recursive=True, on_dup="replace",\n strict=False, verbose=0,\n api=None) -> DataFrame\n\nsites : Sites, str, Path, list, EDICollection\nsort_by : \'lon\'|\'lat\'|\'name\' (station order)\nhalf_window : k neighbours each side (default 3)\nweights : \'tri\'|\'gauss\'|\'uniform\'\npband : (p_min_s, p_max_s) period band\nmax_skew : |beta| threshold (default 6.0)\napi : wrap result in APIFrame when True\n\nReturns DataFrame (one row per station):\n station -- station name\n delta_log10_rho -- log10 shift (pos=above trend)\n fac_rho -- 10^(-delta) rho factor\n fac_z -- 10^(-0.5*delta) Z factor\n n_used -- frequencies used\n\nTypical use:\n tbl = estimate_ss_ama(\n sites, half_window=3, sort_by="lon"\n )\n print(\n tbl[["station","delta_log10_rho","fac_z"]]\n )\n\n\ncorrect_ss_ama -- estimate + apply AMA correction\n--------------------------------------------------\nfrom pycsamt.emtools.ss import correct_ss_ama\n\nSignature:\n correct_ss_ama(sites, *, sort_by="lon",\n half_window=3, weights="tri",\n pband=None, max_skew=6.0,\n robust_freq="median",\n robust_overall="median",\n inplace=False, recursive=True,\n on_dup="replace", strict=False,\n verbose=0) -> Sites\n\nCalls estimate_ss_ama then\napply_ss_factors(key="fac_z").\ninplace=False returns a corrected copy.\n\nTypical use:\n sites_corr = correct_ss_ama(\n sites, half_window=3, sort_by="lon"\n )\n\n\nplot_ss_summary -- four-panel correction figure\n------------------------------------------------\nfrom pycsamt.emtools.ss import plot_ss_summary\n\nSignature:\n plot_ss_summary(\n logRho_before, logRho_after, *,\n freqs, station_labels=None, ...\n ) -> matplotlib.Figure\n\nPanels: (a) Before, (b) After pseudosections,\n (c) Delta section, (d) Per-station bars.\n\nlogRho_before, logRho_after : ndarray (n_st, n_f)\n log10 apparent-resistivity arrays.\nfreqs : ndarray (n_f,) Hz.\nstation_labels : list of str or None.\n\nplot_ss_1d_curves -- per-station sounding grid\n-----------------------------------------------\nfrom pycsamt.emtools.ss import plot_ss_1d_curves\n\nSignature:\n plot_ss_1d_curves(\n logRho_before, logRho_after, *,\n freqs, station_labels=None,\n n_cols=4, max_stations=16, ...\n ) -> matplotlib.Figure\n\nGrid of subplots (one per station) showing\nbefore/after log10-rho sounding curves.\n\nBuild logRho arrays from Sites:\n from pycsamt.emtools._core import (\n _get_z_block, _name, _iter_items,\n )\n\n def collect_logRho(S):\n rows, freqs = [], None\n for i, ed in enumerate(_iter_items(S)):\n Z, z, fr = _get_z_block(ed)\n if Z is None:\n continue\n rxy = (\n 0.2*np.abs(z[:,0,1])**2\n /(fr+1e-24)\n )\n ryx = (\n 0.2*np.abs(z[:,1,0])**2\n /(fr+1e-24)\n )\n rows.append(\n np.log10(\n np.sqrt(rxy*ryx)+1e-24\n )\n )\n freqs = fr\n return np.array(rows), freqs\n\n logRho_b, freqs = collect_logRho(sites)\n logRho_a, _ = collect_logRho(sites_corr)\n labels = [\n _name(ed, i)\n for i, ed in enumerate(_iter_items(sites))\n ]\n fig = plot_ss_summary(\n logRho_b, logRho_a,\n freqs=freqs, station_labels=labels,\n )\n fig.savefig("ss_summary.png", dpi=150)\n\n\n# Static-shift correction full workflow\nfrom pycsamt.api import read_edis\nfrom pycsamt.emtools.ss import (\n estimate_ss_ama,\n correct_ss_ama,\n plot_ss_summary,\n plot_ss_1d_curves,\n)\nfrom pycsamt.agents import StaticShiftAgent\nfrom pycsamt.emtools._core import (\n _get_z_block, _name, _iter_items,\n)\nimport numpy as np\n\n# 1. Load EDI files\nsurvey = read_edis("L22PLT/")\nsites = survey.collection\n\n# 2. Inspect shift factors (optional)\nss_table = estimate_ss_ama(\n sites,\n half_window=3,\n sort_by="lon",\n weights="tri",\n max_skew=6.0,\n)\nprint(\n ss_table[[\n "station","delta_log10_rho","fac_z"\n ]]\n)\n\n# 3. Correct Z tensor in place\nsites_corr = correct_ss_ama(\n sites,\n half_window=3,\n sort_by="lon",\n)\n\n# 4. (Alternative) via StaticShiftAgent\nagent = StaticShiftAgent(method="ama")\nresult = agent.execute({"sites": sites})\nsites_corr2 = result["corrected_sites"]\nfactors = result["shift_factors"]\n\n# 5. Build log10-rho arrays for plots\ndef collect_logRho(S):\n rows, freqs = [], None\n for i, ed in enumerate(_iter_items(S)):\n Z, z, fr = _get_z_block(ed)\n if Z is None:\n continue\n rxy = (\n 0.2*np.abs(z[:,0,1])**2/(fr+1e-24)\n )\n ryx = (\n 0.2*np.abs(z[:,1,0])**2/(fr+1e-24)\n )\n rows.append(\n np.log10(np.sqrt(rxy*ryx)+1e-24)\n )\n freqs = fr\n return np.array(rows), freqs\n\nlogRho_b, freqs = collect_logRho(sites)\nlogRho_a, _ = collect_logRho(sites_corr)\nlabels = [\n _name(ed, i)\n for i, ed in enumerate(_iter_items(sites))\n]\n\n# 6. Summary figure\nfig_sum = plot_ss_summary(\n logRho_b, logRho_a,\n freqs=freqs,\n station_labels=labels,\n)\nfig_sum.savefig(\n "ss_summary.png", dpi=150,\n bbox_inches="tight",\n)\n\n# 7. Per-station 1-D sounding curves\nfig_1d = plot_ss_1d_curves(\n logRho_b, logRho_a,\n freqs=freqs,\n station_labels=labels,\n)\nfig_1d.savefig(\n "ss_1d_curves.png", dpi=150,\n bbox_inches="tight",\n)\n\n---\n\nAnswer in clear technical English. Include short code examples when helpful.'#: Override in subclasses to give the LLM its domain expertise.

execute(input_data)[source]#

Run this agent on input_data and return an AgentResult.

Subclasses must implement this method. The contract:

Reset self._last_cost = 0.0 at the top.
Record wall-clock time with t0 = time.time().
Return AgentResult(elapsed_seconds=time.time()-t0, cost_estimate_usd=self._last_cost, ...).

Parameters:: input_data (dict)
Return type:: AgentResult