pycsamt.agents.ai_inversion#

pycsamt.agents.ai_inversion#

AIInversionAgent — End-to-end AI 1-D MT inversion.

Workflow#

  1. Generate synthetic training data — generate_dataset() builds a ForwardDataset of (response, model) pairs covering a realistic resistivity range.

  2. Train — EMInverter1D is fitted on the synthetic data using the selected network architecture.

  3. Predict — the trained inverter predicts a layered resistivity model for every station in the observed dataset.

  4. Evaluate — per-station RMS between observed and re-computed forward response validates the prediction quality.

  5. Visualise — convergence curve + predicted model sections.

Requires PyTorch or TensorFlow (lazy import — no hard dependency at import time). When neither is available the agent returns a clear error message with an installation hint.

Classes

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

Train an AI inverter on synthetic data then predict on observed sites.
class pycsamt.agents.ai_inversion.AIInversionAgent(*, api_key=None, model=None, llm_provider='claude', arch='resnet', n_layers=5, n_train_samples=2000, epochs=30, freqs=None, pretrained=None)[source]#

Bases: BaseAgent

Train an AI inverter on synthetic data then predict on observed sites.

Parameters:

  • api_key (str)

  • model (str)

  • llm_provider (str)

  • arch (optional overrides) – Neural network architecture.

  • n_layers (int) – Number of model layers the inverter will predict (default 5).

  • n_train_samples (optional overrides) – Number of synthetic training samples (default 2 000).

  • epochs (optional overrides) – Training epochs (default 30). Increase for better models.

  • freqs (array-like or None) – Frequencies used for both training synthesis and observed data interpolation. Default: 40 log-spaced 10⁻⁴–10³ Hz.

  • pretrained (str or None) – Path to a pre-trained model checkpoint. When set, skips training.

  • keys (Output data)

  • ----------

  • path (sites /)

  • output_dir (str, optional)

  • arch

  • epochs

  • n_train_samples

  • keys

  • ----------------

:param inverter EMInverter1D: :param predictions dict {station: :type predictions dict {station: ndarray of log₁₀ ρ values} :param best_model dict with “resistivity” and “thickness” for first station: :param rms_per_station dict {station: :type rms_per_station dict {station: float} :param rms_global float: :param train_history dict (loss curves): :param figures dict: :param figure_paths dict:

Examples

>>> agent  = AIInversionAgent(arch="resnet", n_layers=5, epochs=30)
>>> result = agent.execute({
...     "path":       "/data/L22PLT",
...     "output_dir": "/out/ai_inv",
... })
>>> result["rms_global"]
0.24
SYSTEM_PROMPT: str = 'You are an expert in AI-based MT inversion and deep learning for geophysics.\nGiven an AI inversion result, write 4–5 sentences that:\n1. Describe the neural network architecture used and training convergence.\n2. State the prediction quality (RMS, layer count, depth range).\n3. Identify stations where the AI prediction is most / least reliable.\n4. Compare AI results with classical Bostick depth estimates if available.\n5. Recommend next steps (fine-tuning, ensemble, switch to 2-D inversion).\nReply in plain scientific English.\n'#

Override in subclasses to give the LLM its domain expertise.

classmethod from_pretrained(model_name, *, api_key=None, model=None, llm_provider='claude', cache_dir=None, force_download=False)[source]#

Return an AIInversionAgent pre-loaded with a zoo checkpoint.

Parameters:

  • model_name (str) – Registry name — see list_pretrained().

  • cache_dir (str or None) – Override default cache ~/.pycsamt/model_zoo/.

  • force_download (bool) – Re-download even if cached.

  • api_key (str | None)

  • model (str | None)

  • llm_provider (str)

Return type:

AIInversionAgent

Examples

>>> agent = AIInversionAgent.from_pretrained("mt1d-resnet-5layer-v1")
>>> result = agent.execute({"path": "/data/L22PLT"})
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[str, Any])

Return type:

AgentResult