# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
"""Core data extraction helpers for :mod:`pycsamt.map`."""
from __future__ import annotations
import re
from collections.abc import Iterable, Mapping
from dataclasses import dataclass, field
from typing import Any
import numpy as np
import pandas as pd
from pycsamt.emtools._core import ensure_sites
_NON_ALNUM_RE = re.compile(r"[^a-z0-9]")
[docs]
def normalize_station_id(name: Any) -> str:
"""Loosely-normalized station id for fuzzy matching.
Lowercases and strips everything but letters/digits, so minor
formatting differences between sources — ``23-18-001A`` vs
``23_18_001a`` vs ``23 18 001A`` — still match. Used as a
fallback tier (never the only lookup) wherever a station from
one source (EDI, ModEM, an uploaded elevation file, ...) needs
to be matched against another.
"""
return _NON_ALNUM_RE.sub("", str(name).strip().lower())
[docs]
@dataclass
class StationRecord:
"""A normalized station row used by map renderers."""
id: str
latitude: float | None = None
longitude: float | None = None
elevation: float | None = None
line: str | None = None
index: int = 0
source: Any = None
def __post_init__(self) -> None:
"""Normalize station metadata."""
self.id = str(self.id)
self.latitude = _finite_or_none(self.latitude)
self.longitude = _finite_or_none(self.longitude)
self.elevation = _finite_or_none(self.elevation)
if self.line not in (None, ""):
self.line = str(self.line)
[docs]
@dataclass
class ProfileLine:
"""A named sequence of stations in one profile."""
name: str
stations: tuple[StationRecord, ...] = ()
[docs]
@dataclass
class MapData:
"""Normalized survey data shared by map renderers."""
sites: Any
stations: tuple[StationRecord, ...] = ()
profiles: tuple[ProfileLine, ...] = ()
metadata: dict[str, Any] = field(default_factory=dict)
def __post_init__(self) -> None:
"""Normalize station/profile containers."""
self.stations = tuple(self.stations or ())
if self.profiles:
self.profiles = tuple(self.profiles)
else:
self.profiles = _build_profiles(self.stations)
self.metadata.setdefault(
"n_stations",
len(self.stations),
)
self.metadata.setdefault(
"n_profiles",
len(self.profiles),
)
[docs]
@property
def station_ids(self) -> tuple[str, ...]:
"""Return normalized station IDs."""
return tuple(station.id for station in self.stations)
[docs]
@property
def lines(self) -> tuple[str, ...]:
"""Return available profile/line names."""
return tuple(profile.name for profile in self.profiles)
[docs]
@property
def has_geo(self) -> bool:
"""Whether stations have finite map coordinates."""
if not self.stations:
return False
return all(
station.latitude is not None and station.longitude is not None
for station in self.stations
)
[docs]
def iter_edis(self) -> tuple[Any, ...]:
"""Return EDI-like source objects."""
return _as_edi_tuple(self.sites)
TENSOR_COMPONENTS: dict[str, tuple[int, int]] = {
"xx": (0, 0),
"xy": (0, 1),
"yx": (1, 0),
"yy": (1, 1),
}
COMPONENT_ALIASES: dict[str, str] = {
"zxx": "xx",
"zxy": "xy",
"zyx": "yx",
"zyy": "yy",
"rho_xy": "xy",
"rho_yx": "yx",
"phase_xy": "xy",
"phase_yx": "yx",
"determinant": "det",
"det": "det",
"avg": "avg",
"average": "avg",
"mean": "avg",
}
_LAT_NAMES = (
"latitude",
"lat",
"Latitude",
"LAT",
"LATITUDE",
)
_LON_NAMES = (
"longitude",
"lon",
"long",
"Longitude",
"LON",
"LONG",
"LONGITUDE",
)
_ELEV_NAMES = (
"elevation",
"elev",
"altitude",
"alt",
"Elevation",
"ELEV",
"ELEVATION",
)
_ID_NAMES = ("station", "id", "name", "Station", "ID", "Name")
_LINE_NAMES = (
"line",
"profile",
"survey_line",
"Line",
"Profile",
)
_META_NAMES = (
"Head",
"Header",
"HEAD",
"head",
"header",
"metadata",
"meta",
"info",
"station_info",
)
[docs]
@dataclass(frozen=True)
class ComponentSpec:
"""Parsed impedance component request."""
name: str
indices: tuple[int, int] | None = None
mode: str = "tensor"
[docs]
@dataclass(frozen=True)
class FrequencySelection:
"""Nearest-frequency selection metadata."""
requested: float | None
actual: float
index: int
delta: float
relative_delta: float
within_tolerance: bool = True
[docs]
@dataclass(frozen=True)
class FrequencyValue:
"""Station value and selected-frequency metadata."""
station: str
value: float
selection: FrequencySelection
[docs]
def ensure_map_data(
data: Any,
*,
recursive: bool = True,
line_map: dict[str, Iterable[str]] | None = None,
verbose: int = 0,
) -> MapData:
"""Return normalized map data.
Parameters
----------
data :
EDI path, directory, iterable, or ``Sites``.
recursive :
Passed to :func:`pycsamt.emtools._core.ensure_sites`.
line_map :
Optional ``line -> station names`` mapping used when
line metadata is not embedded in EDI objects.
verbose :
Verbosity passed to :func:`ensure_sites`.
"""
if isinstance(data, MapData):
return data
sites = ensure_sites(
data,
recursive=recursive,
verbose=verbose,
)
edis = _as_edi_tuple(sites)
stations = tuple(
_station_record(edi, index=i, line_map=line_map)
for i, edi in enumerate(edis)
)
profiles = _build_profiles(stations)
metadata = {
"n_stations": len(stations),
"n_profiles": len(profiles),
}
return MapData(
sites=sites,
stations=stations,
profiles=profiles,
metadata=metadata,
)
[docs]
def station_records(
data: Any,
**kwargs: Any,
) -> tuple[StationRecord, ...]:
"""Return normalized station records for *data*."""
return ensure_map_data(data, **kwargs).stations
[docs]
def profile_lines(
data: Any,
**kwargs: Any,
) -> tuple[ProfileLine, ...]:
"""Return normalized profile lines for *data*."""
return ensure_map_data(data, **kwargs).profiles
[docs]
def station_dataframe(data: MapData) -> pd.DataFrame:
"""Return station records as a DataFrame."""
rows = [
{
"ID": s.id,
"Latitude": s.latitude,
"Longitude": s.longitude,
"Elevation": s.elevation,
"Line": s.line or "line",
"Index": s.index,
}
for s in data.stations
]
return pd.DataFrame(rows)
[docs]
def normalize_component(component: str | None) -> str:
"""Return the canonical map component name."""
raw = str(component or "xy").strip().lower()
raw = raw.replace("-", "_")
name = COMPONENT_ALIASES.get(raw, raw)
if name in TENSOR_COMPONENTS or name in {"det", "avg"}:
return name
choices = sorted([*TENSOR_COMPONENTS, "avg", "det"])
msg = (
f"Unknown impedance component {component!r}. "
f"Expected one of {choices}."
)
raise ValueError(msg)
[docs]
def component_spec(component: str | None) -> ComponentSpec:
"""Return a parsed impedance component specification."""
name = normalize_component(component)
if name in TENSOR_COMPONENTS:
return ComponentSpec(
name=name,
indices=TENSOR_COMPONENTS[name],
)
return ComponentSpec(name=name, mode=name)
[docs]
def component_index(component: str) -> tuple[int, int]:
"""Return tensor indices for an impedance component."""
spec = component_spec(component)
if spec.indices is None:
msg = (
f"Component {component!r} is derived and has no "
"single tensor index."
)
raise ValueError(msg)
return spec.indices
[docs]
def component_values(
arr: Any,
component: str | None,
*,
quantity: str = "rho",
) -> np.ndarray:
"""Extract a 1-D value series for one component."""
values = np.asarray(arr, dtype=float)
if values.ndim < 3 or values.shape[-2:] != (2, 2):
msg = "Expected impedance values with shape (nfreq, 2, 2)."
raise ValueError(msg)
spec = component_spec(component)
if spec.indices is not None:
r, c = spec.indices
return values[:, r, c]
xy = values[:, 0, 1]
yx = values[:, 1, 0]
if spec.name == "det":
if quantity.lower() in {"phase", "phi"}:
return np.nanmean(
np.column_stack([xy, yx]),
axis=1,
)
return np.sqrt(np.abs(xy * yx))
return np.nanmean(
np.column_stack([xy, yx]),
axis=1,
)
[docs]
def frequency_axis(data: MapData) -> np.ndarray:
"""Return the first finite frequency axis found."""
for edi in data.iter_edis():
z_obj = getattr(edi, "Z", None)
if z_obj is None:
continue
freq = np.asarray(
getattr(z_obj, "freq", []),
dtype=float,
)
freq = freq[np.isfinite(freq) & (freq > 0)]
if freq.size:
return freq
return np.array([], dtype=float)
[docs]
def select_frequency(
frequencies: Any,
requested: float | None = None,
*,
tolerance: float | None = None,
) -> FrequencySelection | None:
"""Select the closest finite positive frequency."""
freq = np.asarray(frequencies, dtype=float)
good = np.isfinite(freq) & (freq > 0)
if not good.any():
return None
valid_idx = np.flatnonzero(good)
valid = freq[good]
if requested is None:
local_idx = 0
req = None
delta = 0.0
rel = 0.0
else:
req = float(requested)
local_idx = int(np.nanargmin(np.abs(valid - req)))
delta = float(abs(valid[local_idx] - req))
rel = delta / abs(req) if req else delta
actual = float(valid[local_idx])
within = True
if tolerance is not None and requested is not None:
within = delta <= float(tolerance)
return FrequencySelection(
requested=req,
actual=actual,
index=int(valid_idx[local_idx]),
delta=delta,
relative_delta=float(rel),
within_tolerance=within,
)
[docs]
def station_distance_km(data: MapData) -> np.ndarray:
"""Return approximate station distance in km."""
stations = list(data.stations)
if not stations:
return np.array([], dtype=float)
lat = np.array(
[s.latitude if s.latitude is not None else np.nan for s in stations],
dtype=float,
)
lon = np.array(
[
s.longitude if s.longitude is not None else np.nan
for s in stations
],
dtype=float,
)
if not np.isfinite(lat).all() or not np.isfinite(lon).all():
return np.arange(len(stations), dtype=float)
x = lon * 111.320 * np.cos(np.deg2rad(np.nanmean(lat)))
y = lat * 110.574
dist = np.zeros(len(stations), dtype=float)
for i in range(1, len(stations)):
dist[i] = dist[i - 1] + float(
np.hypot(
x[i] - x[i - 1],
y[i] - y[i - 1],
)
)
return dist
[docs]
def value_at_frequency(
data: MapData,
*,
frequency: float,
quantity: str = "rho",
component: str = "xy",
tolerance: float | None = None,
) -> dict[str, float]:
"""Return station values at the closest frequency."""
details = value_at_frequency_details(
data,
frequency=frequency,
quantity=quantity,
component=component,
tolerance=tolerance,
)
return {station: item.value for station, item in details.items()}
[docs]
def value_at_frequency_details(
data: MapData,
*,
frequency: float,
quantity: str = "rho",
component: str = "xy",
tolerance: float | None = None,
) -> dict[str, FrequencyValue]:
"""Return values with selected-frequency metadata."""
out: dict[str, FrequencyValue] = {}
for edi in data.iter_edis():
sid = _station_id_from_edi(edi)
z_obj = getattr(edi, "Z", None)
if z_obj is None:
continue
freq = np.asarray(
getattr(z_obj, "freq", []),
dtype=float,
)
selection = select_frequency(
freq,
frequency,
tolerance=tolerance,
)
if selection is None or not selection.within_tolerance:
continue
q_name = quantity.lower()
if q_name in {"phase", "phi"}:
arr = getattr(z_obj, "phase", None)
else:
arr = getattr(z_obj, "resistivity", None)
if arr is None:
continue
values = component_values(
arr,
component,
quantity=q_name,
)
if selection.index >= values.size:
continue
value = float(values[selection.index])
if np.isfinite(value):
out[sid] = FrequencyValue(
station=sid,
value=value,
selection=selection,
)
return out
[docs]
def skin_depth_at_frequency(
data: MapData,
*,
frequency: float,
component: str = "xy",
tolerance: float | None = None,
) -> dict[str, float]:
"""Return skin depth at the closest frequency."""
details = value_at_frequency_details(
data,
frequency=frequency,
quantity="rho",
component=component,
tolerance=tolerance,
)
out: dict[str, float] = {}
for sid, item in details.items():
freq = item.selection.actual
value = item.value
if value > 0 and freq > 0:
depth = 503.0 * np.sqrt(value / freq)
out[sid] = float(depth)
return out
[docs]
def pseudosection_table(
data: MapData,
*,
quantity: str = "rho",
component: str = "xy",
) -> pd.DataFrame:
"""Return long-form profile data for pseudosections."""
rows: list[dict[str, float | str]] = []
distances = station_distance_km(data)
dist_by_id = {
station.id: float(distances[i])
for i, station in enumerate(data.stations)
}
for edi in data.iter_edis():
sid = _station_id_from_edi(edi)
z_obj = getattr(edi, "Z", None)
if z_obj is None:
continue
freq = np.asarray(
getattr(z_obj, "freq", []),
dtype=float,
)
if freq.size == 0:
continue
arr = getattr(
z_obj,
"resistivity" if quantity == "rho" else "phase",
None,
)
if arr is None:
continue
vals = component_values(
arr,
component,
quantity=quantity,
)
periods = np.where(freq > 0, 1.0 / freq, np.nan)
for period, value in zip(periods, vals):
if not np.isfinite(period) or not np.isfinite(value):
continue
if quantity == "rho" and value <= 0:
continue
rows.append(
{
"station": sid,
"distance": dist_by_id.get(sid, np.nan),
"period": float(period),
"value": float(value),
}
)
return pd.DataFrame(rows)
def _station_record(
edi: Any,
*,
index: int,
line_map: dict[str, Iterable[str]] | None,
) -> StationRecord:
station_id = _station_id_from_edi(edi) or f"S{index:03d}"
return StationRecord(
id=station_id,
latitude=_metadata_float(edi, _LAT_NAMES),
longitude=_metadata_float(edi, _LON_NAMES),
elevation=_metadata_float(edi, _ELEV_NAMES),
line=_resolve_line(edi, station_id, line_map),
index=index,
source=edi,
)
def _metadata_float(
obj: Any,
names: tuple[str, ...],
) -> float | None:
value = _first_float(obj, *names)
if value is not None:
return value
coords = _coords_float(obj, names)
if coords is not None:
return coords
for meta in _metadata_sources(obj):
value = _first_float(meta, *names)
if value is not None:
return value
coords = _coords_float(meta, names)
if coords is not None:
return coords
return None
def _coords_float(
obj: Any,
names: tuple[str, ...],
) -> float | None:
if obj is None:
return None
try:
coords = _get_field(obj, "coords")
if coords is None or len(coords) < 2:
return None
raw = coords[0] if _is_latitude_names(names) else coords[1]
except (TypeError, ValueError, IndexError):
return None
return _finite_or_none(raw)
def _is_latitude_names(names: tuple[str, ...]) -> bool:
return any(str(name).lower().startswith("lat") for name in names)
def _first_float(obj: Any, *names: str) -> float | None:
for name in names:
value = _get_field(obj, name)
value_f = _finite_or_none(value)
if value_f is not None:
return value_f
return None
def _get_field(obj: Any, name: str) -> Any:
if obj is None:
return None
if isinstance(obj, Mapping):
for key in (name, name.lower(), name.upper()):
if key in obj:
return obj[key]
return None
return getattr(obj, name, None)
def _metadata_sources(obj: Any) -> tuple[Any, ...]:
sources: list[Any] = []
for name in _META_NAMES:
value = _get_field(obj, name)
if value is not None:
sources.append(value)
# pyCSAMT EDIFile exposes its header via ``get_section("head")``
# (a method), whose ``.Location`` carries lat/lon/elev. Pull those
# section objects — and their Location — into the metadata search.
getter = getattr(obj, "get_section", None)
if callable(getter):
for section in ("head", "Head", "definemeas", "emeassect"):
try:
sec = getter(section)
except Exception: # noqa: BLE001 - optional section
sec = None
if sec is None:
continue
sources.append(sec)
loc = getattr(sec, "Location", None) or getattr(
sec, "location", None
)
if loc is not None:
sources.append(loc)
direct_loc = getattr(obj, "Location", None) or getattr(
obj, "location", None
)
if direct_loc is not None:
sources.append(direct_loc)
return tuple(sources)
def _finite_or_none(value: Any) -> float | None:
if value is None:
return None
try:
value_f = float(value)
except (TypeError, ValueError):
return None
if np.isfinite(value_f):
return value_f
return None
def _station_id_from_edi(edi: Any) -> str:
value = _first_text(edi, _ID_NAMES)
if value:
return value
for meta in _metadata_sources(edi):
value = _first_text(meta, _ID_NAMES)
if value:
return value
return ""
def _first_text(obj: Any, names: tuple[str, ...]) -> str:
for name in names:
value = _get_field(obj, name)
if value not in (None, ""):
return str(value)
return ""
def _as_edi_tuple(sites: Any) -> tuple[Any, ...]:
if sites is None:
return ()
if hasattr(sites, "as_list"):
return tuple(sites.as_list())
if isinstance(sites, Iterable) and not isinstance(
sites,
(str, bytes),
):
return tuple(sites)
return (sites,)
def _resolve_line(
edi: Any,
station_id: str,
line_map: dict[str, Iterable[str]] | None,
) -> str | None:
value = _first_text(edi, _LINE_NAMES)
if value:
return value
for meta in _metadata_sources(edi):
value = _first_text(meta, _LINE_NAMES)
if value:
return value
if line_map:
for line, stations in line_map.items():
if station_id in {str(s) for s in stations}:
return str(line)
return None
def _build_profiles(
stations: tuple[StationRecord, ...],
) -> tuple[ProfileLine, ...]:
grouped: dict[str, list[StationRecord]] = {}
for station in stations:
line = station.line or "line"
grouped.setdefault(line, []).append(station)
return tuple(
ProfileLine(name=name, stations=tuple(items))
for name, items in grouped.items()
)
