# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
"""Extract elevation and chainage arrays from Sites / EDI collections.
These utilities are the bridge between pycsamt's data model and the
topography rendering pipeline. They operate on any object that contains
EDI-like items with HEAD lat/lon/elev attributes — including
:class:`~pycsamt.site.base.Sites`, plain lists of
:class:`~pycsamt.seg.edi.EDIFile`, and
:class:`~pycsamt.stratagem.io.EDIBatch` objects.
Examples
--------
>>> from pycsamt.topo.extract import extract_elevation, extract_chainage
>>> elev = extract_elevation(sites) # (n_stations,) m a.s.l.
>>> chain = extract_chainage(sites) # (n_stations,) km
"""
from __future__ import annotations
import warnings
from typing import Any
import numpy as np
__all__ = [
"extract_elevation",
"extract_chainage",
"has_elevation",
"extract_station_names",
]
# ---------------------------------------------------------------------------
# Public functions
# ---------------------------------------------------------------------------
[docs]
def has_elevation(sites: Any) -> bool:
"""Return True if any station carries a meaningful non-zero elevation.
Parameters
----------
sites : any station container
Returns
-------
bool
"""
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning)
arr = extract_elevation(sites)
return bool(np.any(arr != 0.0))
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _iter_edis(sites: Any):
"""Yield individual EDI-like objects from any container type."""
# pycsamt Sites — has ._items list of Site objects with .edi attribute
if hasattr(sites, "_items"):
for site in sites._items:
ed = getattr(site, "edi", site)
yield ed
return
# EDIBatch / EDICollection — has .edi_objects_ or ._edis or similar
for attr in ("edi_objects_", "_edis", "edis", "edi_files"):
container = getattr(sites, attr, None)
if container is not None:
try:
for item in container:
yield item
return
except TypeError:
pass
# Generic iterable (list / tuple of EDIFile / Site)
try:
items = list(sites)
for item in items:
yield getattr(item, "edi", item)
return
except TypeError:
pass
# Single EDI-like object
yield sites
def _get_head(ed: Any) -> Any | None:
"""Return the head section object from any EDI-like object.
Handles three HEAD storage patterns:
- attribute access ``ed.Head`` / ``ed.head`` (old pycsamt / MTpy style)
- ``sections`` dict ``ed.sections['head']`` (``seg.edi.EDIFile``)
- ``DEFINEMEAS`` blocks (some legacy parsers)
"""
for attr in ("Head", "head", "DEFINEMEAS", "definemeas"):
h = getattr(ed, attr, None)
if h is not None:
return h
# seg.edi.EDIFile stores sections in a dict
sections = getattr(ed, "sections", None)
if isinstance(sections, dict):
for key in ("head", "Head", "MTSECT", "mtsect"):
h = sections.get(key)
if h is not None:
return h
return None
def _read_elev(ed: Any) -> float | None:
"""Read elevation in metres from an EDI-like object's HEAD."""
h = _get_head(ed)
if h is not None:
# seg.edi.EDIFile exposes a Location sub-object
loc = getattr(h, "Location", None)
if loc is not None:
for name in ("elevation", "elev", "alt"):
v = getattr(loc, name, None)
if v is not None:
try:
f = float(v)
return f if np.isfinite(f) else None
except (TypeError, ValueError):
pass
# Direct attributes on head (old style / MTpy)
for name in ("elev", "elevation", "alt", "ALT", "z", "Z"):
v = getattr(h, name, None)
if v is not None:
try:
f = float(v)
return f if np.isfinite(f) else None
except (TypeError, ValueError):
pass
# Also try direct attributes on the EDI object itself
for name in ("elev", "elevation", "alt"):
v = getattr(ed, name, None)
if v is not None:
try:
f = float(v)
return f if np.isfinite(f) else None
except (TypeError, ValueError):
pass
return None
def _read_latlon(sites: Any) -> list[tuple[float, float]]:
"""Extract (lat, lon) pairs in collection order."""
coords = []
for ed in _iter_edis(sites):
lat = lon = None
h = _get_head(ed)
if h is not None:
# seg.edi.EDIFile — Location sub-object
loc = getattr(h, "Location", None)
if loc is not None:
lat = getattr(loc, "latitude", None) or getattr(
loc, "lat", None
)
lon = (
getattr(loc, "longitude", None)
or getattr(loc, "lon", None)
or getattr(loc, "long", None)
)
# Fall back to direct attributes on head
if lat is None or lon is None:
lat = getattr(h, "lat", None)
lon = getattr(h, "lon", None) or getattr(h, "long", None)
if lat is not None and lon is not None:
try:
coords.append((float(lat), float(lon)))
except (TypeError, ValueError):
pass
return coords
def _read_name(ed: Any) -> str | None:
"""Read station name / dataid from an EDI-like object."""
h = _get_head(ed)
if h is not None:
for name_attr in ("station", "dataid", "sitename", "name"):
v = getattr(h, name_attr, None)
if v is not None:
return str(v).strip()
for name_attr in ("station", "station_id", "name", "id"):
v = getattr(ed, name_attr, None)
if v is not None:
return str(v).strip()
return None
