Source code for pycsamt.ai.domain_gap.survey_fit

# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
"""Fit plausible :class:`~.simulator.CorruptionConfig` ranges from real data.

This module is the bridge between a real field survey (AMT, CSAMT, MT,
or otherwise) and the numpy-only :mod:`pycsamt.ai.domain_gap.simulator`.
It is deliberately split in two:

:func:`survey_data_from_sites`
    Converts EDI/``Sites``/``APISurvey`` input into a canonical
    :class:`~pycsamt.ai.data.contracts.SurveyData`, using only the
    quantities the M1 data contract defines. This path stays numpy-only
    once the bridge itself has run.

:func:`fit_corruption_config`
    Derives plausible noise/error-floor/dropout ranges purely from a
    :class:`~pycsamt.ai.data.contracts.SurveyData`'s own declared errors and
    coverage. No EDI or pandas dependency is required at this stage.

:func:`fit_distortion_priors_from_sites`
    Calls the heavier, pandas-based
    :mod:`pycsamt.emtools.gb`/:mod:`pycsamt.emtools.ss` diagnostics directly
    on real sites to estimate empirical static-shift and galvanic-distortion
    spreads. This is the one part of M3 that genuinely depends on a real
    survey's own QC diagnostics rather than literature defaults.

The frequency grid required by :class:`SurveyData` is *not* silently
interpolated across stations with different sampling: mismatched
frequency grids raise a clear error, per the plan's non-negotiable
principles.
"""

from __future__ import annotations

from typing import Any

import numpy as np

from ...emtools._core import ensure_sites
from ...emtools.gb import groom_bailey_table
from ...emtools.ss import estimate_ss_ama
from ..data.contracts import SurveyData
from .simulator import CorruptionConfig

__all__ = [
    "survey_data_from_sites",
    "fit_corruption_config",
    "fit_distortion_priors_from_sites",
]

_COMPONENT_ORDER = ("xx", "xy", "yx", "yy")
_COMP_IJ = {"xx": (0, 0), "xy": (0, 1), "yx": (1, 0), "yy": (1, 1)}
_DEG_THRESHOLD_M = 1e-5  # ~1 m of latitude/longitude spread

_MV_KM_NT_TO_SI = 4.0e-4 * np.pi
"""Convert EDI-native impedance ([mV/km]/[nT]) to SI (V/A = Ohm).

``pycsamt.site.Site.z`` returns the impedance tensor exactly as
stored in the EDI ``Z`` block, in the SEG/EMTF field convention of
millivolts per kilometre per nanotesla -- not SI. This is the same
factor ``Site.rho``'s own ``rho = 0.2 * |Z|**2 / f`` formula encodes
implicitly (``0.2 == _MV_KM_NT_TO_SI ** 2 / (2 * pi * mu0)`` after
substituting into the SI apparent-resistivity formula
``rho = |Z_SI|**2 / (omega * mu0)``), and the same factor
``pycsamt.forward.maxwell.modem3d`` already documents using when it
converts ModEM's own ``[mV/km]/[nT]`` output back to SI.
"""


def _project_local_metres(
    lats: np.ndarray, lons: np.ndarray, *, station_spacing: float
) -> np.ndarray:
    """Project (lat, lon) degrees to local metres; fall back to a grid.

    Mirrors the equirectangular approximation used by
    :func:`pycsamt.ai.inversion._sites_bridge.sites_to_coords_3d`: adequate
    for the short baselines of a single survey line, not a general-purpose
    geodetic projection.
    """
    n = lats.size
    ok = np.isfinite(lats) & np.isfinite(lons)
    lat_range = float(lats[ok].max() - lats[ok].min()) if ok.any() else 0.0
    lon_range = float(lons[ok].max() - lons[ok].min()) if ok.any() else 0.0
    if ok.sum() >= 2 and (
        lat_range > _DEG_THRESHOLD_M or lon_range > _DEG_THRESHOLD_M
    ):
        lat_ref = float(np.nanmean(lats[ok]))
        lon_ref = float(np.nanmean(lons[ok]))
        cos_lat = float(np.cos(np.radians(lat_ref)))
        x = np.where(ok, (lons - lon_ref) * cos_lat * 111320.0, 0.0)
        y = np.where(ok, (lats - lat_ref) * 111320.0, 0.0)
        return np.column_stack([x, y])
    side = int(np.ceil(np.sqrt(n)))
    xs = np.array([(i % side) * station_spacing for i in range(n)])
    ys = np.array([(i // side) * station_spacing for i in range(n)])
    return np.column_stack([xs, ys])


[docs] def survey_data_from_sites( sites: Any, *, crs: str | None = None, freq_rtol: float = 1e-6, station_spacing: float = 500.0, recursive: bool = True, on_dup: str = "replace", verbose: int = 0, metadata: dict[str, Any] | None = None, ) -> SurveyData: """Bridge EDI/``Sites``/``APISurvey`` input to canonical ``SurveyData``. Parameters ---------- sites : Any Anything accepted by :func:`pycsamt.emtools._core.ensure_sites`: a filesystem path/glob/directory, ``EDIFile``/``EDICollection``, ``Site``/``Sites``, ``APISurvey``, or an iterable of these. crs : str, optional Coordinate reference system identifier to record. Station positions are always projected with a local equirectangular approximation (see :func:`pycsamt.ai.inversion._sites_bridge.sites_to_coords_3d`); pass a CRS string only if it genuinely describes that projection. freq_rtol : float, default=1e-6 Relative tolerance used when checking that every station shares the same frequency grid. station_spacing : float, default=500.0 Forwarded to the coordinate bridge as a uniform-grid fallback spacing, used only when no station reports finite coordinates. recursive, on_dup, verbose Forwarded to ``ensure_sites``. metadata : dict, optional Extra provenance recorded on the returned survey. Returns ------- SurveyData Canonical survey with full ``xx, xy, yx, yy`` components, with impedance and its declared error converted from ``Site``'s EDI-native ``[mV/km]/[nT]`` convention to SI (V/A), matching :class:`~pycsamt.ai.data.contracts.SurveyData`'s default :class:`~pycsamt.ai.data.contracts.ImpedanceConvention`. Raises ------ ValueError If no station has usable impedance data, or stations do not share a common frequency grid within ``freq_rtol``. Notes ----- This function performs no frequency interpolation: a survey whose stations were sampled on different frequency grids must be resolved by an explicit, survey-matched frequency selector (an M1 concern) before reaching this bridge. Examples -------- >>> survey = survey_data_from_sites( ... "data/AMT/WILLY_DATA/L18PLT", recursive=False, verbose=0 ... ) # doctest: +SKIP >>> survey.components # doctest: +SKIP ('xx', 'xy', 'yx', 'yy') """ collection = ensure_sites( sites, recursive=recursive, on_dup=on_dup, verbose=verbose ) stations = list(collection) if not stations: raise ValueError("no stations were found in the given sites input.") reference_freq: np.ndarray | None = None reference_name = None kept: list[Any] = [] for site in stations: freq = getattr(site, "freq", None) z = getattr(site, "z", None) if freq is None or z is None: continue freq = np.asarray(freq, dtype=float) if reference_freq is None: reference_freq = freq reference_name = site.name elif freq.shape != reference_freq.shape or not np.allclose( freq, reference_freq, rtol=freq_rtol, atol=0.0 ): raise ValueError( f"station {getattr(site, 'name', '?')!r} does not share the " f"frequency grid of {reference_name!r}; a survey-matched " "frequency selector must resolve this before bridging to " "SurveyData." ) kept.append(site) if reference_freq is None: raise ValueError("no station exposed both freq and z arrays.") n_station = len(kept) n_frequency = reference_freq.size impedance = np.full((n_station, n_frequency, 4), np.nan, dtype=complex) error = np.full((n_station, n_frequency, 4), np.nan, dtype=float) names: list[str] = [] for row, site in enumerate(kept): names.append(str(site.name)) z = np.asarray(site.z) * _MV_KM_NT_TO_SI z_err = getattr(site, "z_err", None) z_err = ( None if z_err is None else np.asarray(z_err, dtype=float) * _MV_KM_NT_TO_SI ) if z.ndim == 3 and z.shape[1:] == (2, 2): for col, name in enumerate(_COMPONENT_ORDER): i, j = _COMP_IJ[name] impedance[row, :, col] = z[:, i, j] if z_err is not None: error[row, :, col] = z_err[:, i, j] elif z.ndim == 2 and z.shape[1] == 4: impedance[row] = z if z_err is not None: error[row] = z_err else: raise ValueError( f"station {site.name!r} has an unsupported z shape {z.shape}." ) lats = np.full(n_station, np.nan) lons = np.full(n_station, np.nan) elevation = np.full(n_station, np.nan) for row, site in enumerate(kept): coords = getattr(site, "coords", None) if coords is None: continue try: if len(coords) >= 2: lats[row] = float(coords[0]) lons[row] = float(coords[1]) if len(coords) >= 3: elevation[row] = float(coords[2]) except (TypeError, ValueError): pass xy = _project_local_metres(lats, lons, station_spacing=station_spacing) coordinates = np.column_stack([xy, elevation]) has_error = np.any(np.isfinite(error)) return SurveyData( impedance=impedance, frequencies_hz=reference_freq, station_names=tuple(names), components=_COMPONENT_ORDER, coordinates_m=coordinates, impedance_error=error if has_error else None, crs=crs, metadata=metadata or {}, )
[docs] def fit_corruption_config( survey: SurveyData, *, severity_scale: float = 1.0, ) -> CorruptionConfig: """Derive plausible noise/dropout ranges from a real survey's QC. Only quantities already present in the canonical :class:`~pycsamt.ai.data.contracts.SurveyData` contract are used: the ``impedance_error``-to-``|Z|`` ratio for heteroscedastic noise and error floor, and :meth:`~pycsamt.ai.data.contracts.SurveyData.coverage` for dropout rates. Parameters ---------- survey : SurveyData Real (or realistically corrupted) survey to profile. severity_scale : float, default=1.0 Multiplier applied to every fitted range/rate, letting a caller derive a milder or harsher preset from the same empirical fit. Returns ------- CorruptionConfig Configuration whose noise range spans the interquartile range of the observed relative error, whose error floor is the fifth percentile of that ratio, and whose dropout rates equal the observed missing fractions. Distortion and outlier parameters are left at zero; see :func:`fit_distortion_priors_from_sites` for those. Raises ------ ValueError If ``survey`` has no declared ``impedance_error`` to profile. Examples -------- >>> import numpy as np >>> from pycsamt.ai.data.contracts import SurveyData >>> z = np.full((4, 6, 2), 100 + 50j) >>> err = np.full((4, 6, 2), 3.0) >>> survey = SurveyData( ... z, ... np.linspace(1000, 1, 6), ... ["A", "B", "C", "D"], ... ["xy", "yx"], ... np.zeros((4, 2)), ... impedance_error=err, ... ) >>> config = fit_corruption_config(survey) >>> config.noise_level_range[0] >= 0.0 True """ if severity_scale <= 0.0 or not np.isfinite(severity_scale): raise ValueError("severity_scale must be finite and positive.") if survey.impedance_error is None: raise ValueError( "survey has no impedance_error to profile; supply a survey with " "declared errors, or build a CorruptionConfig from literature " "defaults instead." ) valid = survey.valid ratio = np.abs(survey.impedance_error[valid]) / np.maximum( np.abs(survey.impedance[valid]), 1e-24 ) ratio = ratio[np.isfinite(ratio)] if ratio.size == 0: raise ValueError( "no valid observations with finite error ratio were found." ) lo = float(np.percentile(ratio, 25)) * severity_scale hi = float(np.percentile(ratio, 75)) * severity_scale floor = float(np.percentile(ratio, 5)) * severity_scale lo, hi = sorted((lo, hi)) coverage = survey.coverage() station_dropout = float( np.clip((1.0 - coverage.by_station).mean() * severity_scale, 0.0, 1.0) ) frequency_dropout = float( np.clip( (1.0 - coverage.by_frequency).mean() * severity_scale, 0.0, 1.0 ) ) random_dropout = float( np.clip((1.0 - coverage.overall) * severity_scale, 0.0, 1.0) ) return CorruptionConfig( noise_level_range=(lo, hi), error_floor_fraction=floor, station_dropout_rate=station_dropout, frequency_dropout_rate=frequency_dropout, random_dropout_rate=random_dropout, )
[docs] def fit_distortion_priors_from_sites( sites: Any, *, recursive: bool = True, on_dup: str = "replace", verbose: int = 0, **kwargs: Any, ) -> dict[str, float]: """Estimate empirical static-shift and distortion spreads from real EDI. This is the one M3 entry point that genuinely depends on the heavier, pandas-based EM diagnostics in :mod:`pycsamt.emtools.gb` and :mod:`pycsamt.emtools.ss`, run directly on real sites (e.g. a WILLY line) rather than on the numpy-only :class:`SurveyData` contract. Parameters ---------- sites : Any Anything accepted by :func:`pycsamt.emtools._core.ensure_sites`. recursive, on_dup, verbose Forwarded to the underlying diagnostics. **kwargs Forwarded to :func:`pycsamt.emtools.gb.groom_bailey_table`. Returns ------- dict ``static_shift_log10_sigma``, ``distortion_gain_log10_sigma``, ``distortion_twist_deg_sigma``, ``distortion_shear_sigma``, and ``distortion_anisotropy_sigma``, each the population standard deviation of the corresponding per-station fitted parameter across stations with a successful fit. A parameter is ``0.0`` when fewer than two stations produced a usable fit. Examples -------- >>> priors = fit_distortion_priors_from_sites( ... "data/AMT/WILLY_DATA/L18PLT", recursive=False, verbose=0 ... ) # doctest: +SKIP >>> sorted(priors) # doctest: +SKIP ['distortion_anisotropy_sigma', 'distortion_gain_log10_sigma', \ 'distortion_shear_sigma', 'distortion_twist_deg_sigma', \ 'static_shift_log10_sigma'] """ zero = { "static_shift_log10_sigma": 0.0, "distortion_gain_log10_sigma": 0.0, "distortion_twist_deg_sigma": 0.0, "distortion_shear_sigma": 0.0, "distortion_anisotropy_sigma": 0.0, } gb_table = groom_bailey_table( sites, recursive=recursive, on_dup=on_dup, verbose=verbose, api=False, **kwargs, ) ok = gb_table[gb_table["status"] == "ok"] if len(gb_table) else gb_table if len(ok) >= 2: zero["distortion_gain_log10_sigma"] = float( np.std(np.log10(ok["gain"].to_numpy())) ) zero["distortion_twist_deg_sigma"] = float( np.std(ok["twist_deg"].to_numpy()) ) zero["distortion_shear_sigma"] = float(np.std(ok["shear"].to_numpy())) zero["distortion_anisotropy_sigma"] = float( np.std(ok["anisotropy"].to_numpy()) ) ss_table = estimate_ss_ama( sites, recursive=recursive, on_dup=on_dup, verbose=verbose, api=False ) if len(ss_table) >= 2: # delta_log10_rho is a shift on rho_a; Z scales as sqrt(rho_a). zero["static_shift_log10_sigma"] = 0.5 * float( np.std(ss_table["delta_log10_rho"].to_numpy()) ) return zero