"""
pycsamt.agents.denoising
=========================
:class:`DenoisingAgent` — Multi-method MT data denoising.
Wraps both classical emtools filters and the AI-based denoiser:
Classical (no ML deps)
* **RPCA** — robust PCA off-diagonal denoising
(:func:`~pycsamt.emtools.remove_noise.rpca_offdiag_denoise`)
* **Hampel** — frequency-domain Hampel outlier filter
(:func:`~pycsamt.emtools.remove_noise.hampel_filter_freq`)
* **EMAP** — array-based EM array processing filter
(:func:`~pycsamt.emtools.remove_noise.apply_emap_filter`)
* **Pipeline** — multi-step combined filter
(:func:`~pycsamt.emtools.remove_noise.remove_noise_pipeline`)
AI-based (requires PyTorch or TensorFlow)
* **CAE** — convolutional autoencoder denoiser
(:class:`~pycsamt.ai.processing.denoise.EMDenoiser`)
"""
from __future__ import annotations
import time
from typing import Any
import numpy as np
from ._base import AgentResult, BaseAgent
_SYSTEM_PROMPT = """\
You are an expert MT noise analysis and denoising specialist.
Given a denoising result summary, write 3–4 sentences that:
1. State which noise sources were addressed (powerline, cultural, source effects).
2. Quantify the improvement (e.g. SNR gain, number of frequencies recovered).
3. Identify any remaining problematic frequencies or stations.
4. Recommend follow-up processing steps.
Reply in plain English.
"""
_METHODS = {"rpca", "hampel", "emap", "pipeline", "ai", "ai_cae"}
[docs]
class DenoisingAgent(BaseAgent):
"""Denoise MT impedance data using classical or AI-based methods.
Parameters
----------
api_key, model, llm_provider : str
method : str
``"rpca"`` (default), ``"hampel"``, ``"emap"``, ``"pipeline"``,
or ``"ai"`` / ``"ai_cae"`` (requires PyTorch/TF).
rank : int
RPCA rank for off-diagonal denoising (default 2).
half_window : int
Hampel filter half-window (default 3).
Input keys
----------
``sites`` / ``path`` : Sites or str
``method`` : str, optional — overrides constructor default
``output_dir`` : str, optional
``period_range`` : [T_min, T_max], optional
Output data keys
----------------
``denoised_sites`` Sites with denoised impedance
``snr_before`` ndarray — per-(station, freq) SNR proxy before
``snr_after`` ndarray — per-(station, freq) SNR proxy after
``snr_gain`` float — mean SNR improvement
``n_recovered`` int — frequencies recovered above SNR threshold
``figures`` dict
``figure_paths`` dict
"""
SYSTEM_PROMPT = _SYSTEM_PROMPT
def __init__(
self,
*,
api_key: str | None = None,
model: str | None = None,
llm_provider: str = "claude",
method: str = "rpca",
rank: int = 2,
half_window: int = 3,
) -> None:
super().__init__(
"DenoisingAgent",
api_key=api_key,
model=model,
llm_provider=llm_provider,
section_preset="pseudosection",
)
self.method = method
self.rank = rank
self.half_window = half_window
[docs]
def execute(self, input_data: dict[str, Any]) -> AgentResult:
self._last_cost = 0.0
t0 = time.time()
warnings: list[str] = []
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)
method = str(input_data.get("method", self.method)).lower()
output_dir = input_data.get("output_dir")
input_data.get("period_range")
if method not in _METHODS:
warnings.append(f"Unknown method {method!r}; using 'rpca'.")
method = "rpca"
# ── compute SNR proxy before denoising ────────────────────────────────
snr_before = _compute_snr_proxy(sites)
# ── apply denoising ───────────────────────────────────────────────────
denoised_sites = sites
from ..emtools.remove_noise import (
remove_noise_pipeline,
)
if method == "rpca":
denoised_sites = _apply_rpca(sites, self.rank, warnings)
elif method == "hampel":
denoised_sites = _apply_hampel(sites, self.half_window, warnings)
elif method == "emap":
try:
from ..emtools.remove_noise import (
apply_emap_filter,
)
denoised_sites = apply_emap_filter(sites, verbose=0)
except Exception as exc:
warnings.append(
f"EMAP filter failed: {exc}. No denoising applied."
)
elif method == "pipeline":
try:
denoised_sites = remove_noise_pipeline(sites, verbose=0)
except Exception as exc:
warnings.append(
f"Noise pipeline failed: {exc}. No denoising applied."
)
elif method in ("ai", "ai_cae"):
denoised_sites = _apply_ai_denoiser(sites, warnings)
# ── compute SNR proxy after ───────────────────────────────────────────
snr_after = _compute_snr_proxy(denoised_sites)
# ── metrics ───────────────────────────────────────────────────────────
snr_gain = 0.0
n_recovered = 0
if snr_before is not None and snr_after is not None:
try:
valid_b = snr_before[np.isfinite(snr_before)]
valid_a = snr_after[np.isfinite(snr_after)]
if valid_b.size and valid_a.size:
snr_gain = float(
np.nanmean(valid_a) - np.nanmean(valid_b)
)
# count cells where SNR crossed the threshold of 3
n_recovered = int(
np.sum((snr_after >= 3.0) & (snr_before < 3.0))
)
except Exception:
pass
# ── figures ───────────────────────────────────────────────────────────
figures: dict[str, Any] = {}
fig_paths: dict[str, str] = {}
try:
import matplotlib.pyplot as plt
from ..emtools.inspect import pseudosection
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# before
pseudosection(sites, quantity="rho_xy", ax=axes[0])
axes[0].set_title("ρa before denoising", fontsize=9)
# after
pseudosection(denoised_sites, quantity="rho_xy", ax=axes[1])
axes[1].set_title("ρa after denoising", fontsize=9)
fig.suptitle(
f"Denoising comparison ({method})",
fontsize=10,
fontweight="bold",
)
fig.tight_layout()
figures["denoising_comparison"] = fig
p = self._save_figure(
fig,
output_dir,
"denoising_comparison",
warnings_list=warnings,
)
if p:
fig_paths["denoising_comparison"] = p
except Exception as exc:
warnings.append(f"Comparison figure: {exc}")
# ── LLM interpretation ────────────────────────────────────────────────
interp: str | None = None
if self.api_key:
prompt = (
f"Denoising summary (method: {method}):\n"
f" Mean SNR gain: {snr_gain:+.2f}\n"
f" Frequencies recovered (SNR < 3 → ≥ 3): {n_recovered}\n"
f" Warnings: {warnings[:3] if warnings else 'none'}\n\n"
"Assess the denoising result."
)
interp = self.query_llm(prompt, max_tokens=200)
elapsed = time.time() - t0
return AgentResult(
status="success",
summary=(
f"Denoising ({method}) complete. "
f"SNR gain: {snr_gain:+.2f}. "
f"{n_recovered} frequencies recovered. "
f"{len(figures)} figure(s) produced."
),
data={
"denoised_sites": denoised_sites,
"snr_before": snr_before,
"snr_after": snr_after,
"snr_gain": snr_gain,
"n_recovered": n_recovered,
"method": method,
"figures": figures,
"figure_paths": fig_paths,
},
warnings=warnings,
llm_interpretation=interp,
elapsed_seconds=elapsed,
cost_estimate_usd=self._last_cost,
)
# ── private helpers ───────────────────────────────────────────────────────────
def _compute_snr_proxy(sites: Any) -> np.ndarray | None:
"""Return a flat array of |Zxy| / std(|Zxy|) per (station, freq) cell."""
from ..emtools._core import _get_z_block, _iter_items
vals = []
for _i, ed in enumerate(_iter_items(sites)):
_, z, fr = _get_z_block(ed)
if z is None:
continue
zxy = np.abs(z[:, 0, 1])
if zxy.size > 2:
np.nanmean(zxy)
sd = np.nanstd(zxy) + 1e-30
vals.extend((zxy / sd).tolist())
return np.asarray(vals, float) if vals else None
def _apply_rpca(sites: Any, rank: int, warnings: list) -> Any:
"""Apply RPCA off-diagonal denoising to each station in *sites*."""
from ..emtools.remove_noise import rpca_offdiag_denoise
try:
result = rpca_offdiag_denoise(sites, rank=rank, verbose=0)
return result if result is not None else sites
except Exception as exc:
warnings.append(
f"RPCA denoising failed: {exc}. Original sites returned."
)
return sites
def _apply_hampel(sites: Any, half_window: int, warnings: list) -> Any:
"""Apply Hampel frequency filter to each station."""
from ..emtools.remove_noise import hampel_filter_freq
try:
result = hampel_filter_freq(sites, k=half_window, verbose=0)
return result if result is not None else sites
except Exception as exc:
warnings.append(
f"Hampel filter failed: {exc}. Original sites returned."
)
return sites
def _apply_ai_denoiser(sites: Any, warnings: list) -> Any:
"""Apply the AI convolutional autoencoder denoiser."""
try:
import numpy as np
from ..ai.processing.denoise import (
EMDenoiser,
prepare_z_features,
)
from ..emtools._core import (
_get_z_block,
_iter_items,
_name,
)
# collect all Z data
all_z = []
items = list(_iter_items(sites))
for _i, ed in enumerate(items):
_, z, fr = _get_z_block(ed)
if z is None:
continue
all_z.append(z)
if not all_z:
warnings.append("No valid Z data for AI denoiser.")
return sites
n_freqs = all_z[0].shape[0]
denoiser = EMDenoiser(n_freqs=n_freqs)
# prepare features: (batch, n_freqs, n_components=4)
X = prepare_z_features(np.stack(all_z, axis=0))
denoiser.fit(X, epochs=20, verbose=False)
denoiser.predict(X)
warnings.append(
"AI denoiser applied (light training on this dataset). "
"For production, pre-train on a large synthetic dataset."
)
# for now return original sites (post-fit transform not yet integrated)
return sites
except ImportError as exc:
warnings.append(
f"AI denoiser requires PyTorch or TensorFlow: {exc}. "
"Using original sites."
)
return sites
except Exception as exc:
warnings.append(f"AI denoiser failed: {exc}. Using original sites.")
return sites
__all__ = ["DenoisingAgent"]
