# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
from __future__ import annotations
import copy
import math
import re
import warnings
from collections.abc import Mapping, Sequence
from os import PathLike
from pathlib import Path
from typing import (
Any,
Callable,
)
import numpy as np
import pandas as pd
from ..api.view import iter_progress, maybe_wrap_frame
from ..core.base import CoreObject
from ..gis.utils import (
assert_lat_value,
assert_lon_value,
)
from ..seg.collection import EDICollection
from ..seg.edi import EDIFile
from ..session import normalize_session
from .utils import _ensure_head, _get_head, as_edicollection
from .utils import get_coords as _get_coords
from .utils import maybe_copy as _maybe_copy
from .utils import set_coords as _set_coords
__all__ = ["SiteMixin", "Site", "Sites", "to_sites", "to_edis"]
[docs]
class SiteMixin(CoreObject):
r"""
Lightweight wrapper exposing station-centric accessors
and utilities for a single :class:`~pycsamt.seg.edi.EDIFile`.
The mixin normalizes common site operations around the
underlying EDI content, including coordinate handling,
impedance arrays, tipper, derived resistivity/phase, and
convenient export to :class:`pandas.DataFrame`.
Notes
-----
The station name is resolved from EDI HEAD fields in the
following order: ``dataid``, ``station``, ``sitename``,
``name``, ``STATION``. If none is present, the file stem
is used. See :func:`_station_name`.
The Z tensor is treated as an array of shape ``(n, 2, 2)``
or flattened to ``(n, 4)`` as needed. The order is:
``Zxx, Zxy, Zyx, Zyy``.
See Also
--------
pycsamt.seg.edi.EDIFile
Parsed SEG-EDI container.
pycsamt.site.base.Site
Concrete site wrapper that enforces a stable ID.
pycsamt.site.base.Sites
Collection helper for multiple sites.
References
----------
.. [SiteMixin-1] SEG EDI Format. Society of Exploration Geophysicists.
Commonly used magnetotelluric exchange format.
"""
def __init__(self, source: Any, *, on_loss: str = "warn") -> None:
self._on_loss = on_loss
if _is_xml_like(source):
self._native = "xml"
self._tf_obj: Any = source
self._edi_obj: EDIFile | None = None
else:
self._native = "edi"
self._edi_obj = source
self._tf_obj = None
[docs]
@property
def backend(self) -> str:
r"""
Which backend natively constructed this site.
Returns
-------
str
``"edi"`` or ``"xml"``. The other representation is
always available too (see :attr:`edi` and :attr:`tf`)
but is lazily materialized and cached on first access.
"""
return self._native
@property
def edi(self) -> EDIFile:
r"""
The underlying :class:`~pycsamt.seg.edi.EDIFile`.
Returns
-------
pycsamt.seg.edi.EDIFile
When this site was constructed from an EDI object, the
same object every time. When constructed from an
:class:`~pycsamt.emtf.document.EMTF` XML document, an
EDI view is materialized on first access via
:func:`pycsamt.emtf.converters.edi.emtf_to_edi` and
cached, so this always returns the same object on
repeated access (safe for in-place mutation and
identity-keyed caches elsewhere in the codebase).
Notes
-----
Materializing from XML may emit a
:class:`~pycsamt.emtf.converters.edi.DataLossWarning` (or
raise, depending on ``on_loss`` passed at construction) if
the source document holds metadata the historical EDI
format cannot represent.
"""
if self._edi_obj is None:
from ..emtf.converters.edi import emtf_to_edi
self._edi_obj = emtf_to_edi(self._tf_obj, on_loss=self._on_loss)
return self._edi_obj
[docs]
@edi.setter
def edi(self, value: EDIFile) -> None:
self._edi_obj = value
[docs]
@property
def tf(self) -> Any:
r"""
The underlying :class:`~pycsamt.emtf.document.EMTF` document.
Returns
-------
pycsamt.emtf.document.EMTF
When this site was constructed from an XML document,
the same object every time. When constructed from an
EDI object, a format-neutral view is materialized on
first access via
:func:`pycsamt.emtf.converters.edi.edi_to_emtf` and
cached (lossless -- XML is a strict superset of EDI).
See Also
--------
site_meta, site_layout, provenance, processing, copyright,
quality_meta
Typed :mod:`pycsamt.metadata` accessors built on top of
this document.
"""
if self._tf_obj is None:
from ..emtf.converters.edi import edi_to_emtf
self._tf_obj = edi_to_emtf(self._edi_obj)
return self._tf_obj
[docs]
@property
def name(self) -> str:
r"""
Station identifier resolved from the EDI header or file
stem.
Returns
-------
str
Station name used for lookups and display.
Notes
-----
The resolution order is defined in :func:`_station_name`.
XML-native sites resolve the name directly from the
:class:`~pycsamt.emtf.document.EMTF` document without
forcing EDI materialization.
"""
if self._native == "xml" and self._edi_obj is None:
nm = _tf_station_name(self._tf_obj)
if nm:
return nm
return _station_name(self.edi)
[docs]
@property
def coords(self) -> tuple[float, float, float]:
r"""
Geographic coordinates of the site.
Returns
-------
tuple of float
A 3-tuple ``(lat, lon, elev)`` in decimal degrees and
meters.
Notes
-----
This accessor relies on :func:`_get_coords` which parses
EDI HEAD latitude, longitude and elevation. XML-native
sites read directly from the :class:`EMTF` document without
forcing EDI materialization.
"""
if self._native == "xml" and self._edi_obj is None:
tf = self._tf_obj
lat = float(tf.lat) if tf.lat is not None else float("nan")
lon = float(tf.lon) if tf.lon is not None else float("nan")
elev = float(tf.elev) if tf.elev is not None else float("nan")
return (lat, lon, elev)
c = _get_coords(self.edi)
return (float(c.lat), float(c.lon), float(c.elev))
[docs]
@property
def freq(self) -> Any:
r"""
Frequency vector extracted from the Z section.
Returns
-------
array-like or None
One-dimensional array of frequencies in Hz, or
``None`` if missing.
See Also
--------
to_dataframe : Tabular export of arrays at each frequency.
"""
return _extract_z_arrays(self.edi)["freq"]
[docs]
@property
def z(self) -> Any:
r"""
Complex impedance tensor across frequencies.
Returns
-------
array-like or None
Array with shape ``(n, 2, 2)`` or flattened to
``(n, 4)`` depending on context, or ``None`` if
absent.
Notes
-----
Flattening order is ``Zxx, Zxy, Zyx, Zyy``.
"""
return _extract_z_arrays(self.edi)["z"]
[docs]
@property
def z_err(self) -> Any:
r"""
Uncertainty associated with the impedance tensor.
Returns
-------
array-like or None
Error array aligned with :attr:`SiteMixin.z`, or ``None`` if
absent.
"""
return _extract_z_arrays(self.edi)["z_err"]
[docs]
@property
def rho(self) -> Any:
r"""
Apparent resistivity derived from the impedance tensor.
Returns
-------
array-like or None
Resistivity values per component and frequency, or
``None`` if not computed or absent.
Notes
-----
Derived fields may be recomputed downstream when the Z
tensor or frequency vector is updated.
"""
return _extract_z_arrays(self.edi)["rho"]
[docs]
@property
def phase(self) -> Any:
r"""
Impedance phase (degrees) derived from the tensor.
Returns
-------
array-like or None
Phase values per component and frequency, or
``None`` if not computed or absent.
"""
return _extract_z_arrays(self.edi)["phase"]
[docs]
@property
def tipper(self) -> Any:
r"""
Vertical magnetic transfer function (tipper).
Returns
-------
array-like or None
Two columns ``(Tx, Ty)`` per frequency, or ``None``
if missing.
"""
return _extract_z_arrays(self.edi)["tipper"]
[docs]
def to_dataframe(self, kind: str = "z", *, api: bool | None = None) -> Any:
r"""
Export core arrays to a tidy :class:`pandas.DataFrame`.
Parameters
----------
kind : {"z", "imp", "impedance", "resphase", "rp",
"rho_phase", "tip", "tipper", "t"}, optional
Selects which quantity to export. The default is
``"z"``.
Returns
-------
pandas.DataFrame
Frame indexed by frequency (name ``"f"``). Columns
depend on ``kind``:
- ``"z"``: ``Zxx, Zxy, Zyx, Zyy`` (complex values).
- ``"resphase"``: pairs of columns per component
``rho_*`` and ``phi_*``.
- ``"tipper"``: ``Tx, Ty``.
Raises
------
ValueError
If ``kind`` is not recognized.
Notes
-----
Missing arrays yield empty frames with the correct index.
Examples
--------
>>> df = site.to_dataframe("z")
>>> df.columns
Index(["Zxx","Zxy","Zyx","Zyy"], dtype="object")
See Also
--------
quality_flags : Quick presence checks of available arrays.
"""
arrs = _extract_z_arrays(self.edi)
f = (
np.asarray(arrs["freq"], float)
if arrs["freq"] is not None
else np.asarray([], float)
)
idx = pd.Index(f, name="f")
def _out(df: pd.DataFrame) -> Any:
return maybe_wrap_frame(
df,
api=api,
name=f"site_{kind.strip().lower()}",
kind=f"site.{kind.strip().lower()}",
source=getattr(self, "name", None),
)
k = kind.strip().lower()
if k in ("z", "imp", "impedance"):
z = arrs["z"]
z = None if z is None else _z_to_2d(z)
if z is None or z.size == 0:
return _out(pd.DataFrame(index=idx))
cols = _component_names()
data: dict[str, Any] = {}
for i, c in enumerate(cols):
try:
data[c] = z[:, i]
except Exception:
data[c] = np.full(idx.size, np.nan, float)
return _out(pd.DataFrame(data, index=idx))
if k in ("resphase", "rp", "rho_phase"):
rho = arrs["rho"]
ph = arrs["phase"]
rho = None if rho is None else _z_to_2d(rho)
ph = None if ph is None else _z_to_2d(ph)
cols = _component_names()
data: dict[str, Any] = {}
for i, c in enumerate(cols):
rr = np.full(idx.size, np.nan, float)
pp = np.full(idx.size, np.nan, float)
try:
if rho is not None:
rr = rho[:, i]
except Exception:
pass
try:
if ph is not None:
pp = ph[:, i]
except Exception:
pass
data[f"rho_{c.lower()}"] = rr
data[f"phi_{c.lower()}"] = pp
return _out(pd.DataFrame(data, index=idx))
if k in ("tip", "tipper", "t"):
tip = arrs["tipper"]
if tip is None:
return _out(pd.DataFrame(index=idx))
tip = np.asarray(tip)
data: dict[str, Any] = {}
for i, c in enumerate(("Tx", "Ty")):
try:
data[c] = tip[:, i]
except Exception:
data[c] = np.full(idx.size, np.nan, float)
return _out(pd.DataFrame(data, index=idx))
raise ValueError(f"Unknown kind: {kind!r}")
[docs]
def quality_flags(self) -> dict[str, bool]:
r"""
Report presence and basic validity of key arrays.
Returns
-------
dict
Flags: ``has_freq``, ``has_z``, ``has_z_err``,
``has_rho``, ``has_phase``, ``has_tipper``. A flag
is ``True`` if the array exists, has non-zero size,
and all values are finite.
Examples
--------
>>> site.quality_flags()["has_z"]
True
"""
a = _extract_z_arrays(self.edi)
def _ok(x: Any) -> bool:
if x is None:
return False
arr = np.asarray(x)
return arr.size > 0 and np.all(np.isfinite(arr))
return {
"has_freq": a["freq"] is not None,
"has_z": _ok(a["z"]),
"has_z_err": _ok(a["z_err"]),
"has_rho": _ok(a["rho"]),
"has_phase": _ok(a["phase"]),
"has_tipper": _ok(a["tipper"]),
}
[docs]
def has_component(self, comp: str) -> bool:
r"""
Check if a given component contains any finite value.
Parameters
----------
comp : str
One of ``"Zxx"``, ``"Zxy"``, ``"Zyx"``, ``"Zyy"`` for
impedance, or ``"tip"``, ``"tx"``, ``"ty"``,
``"tipper"`` for tipper.
Returns
-------
bool
``True`` if the component exists and has at least
one finite value.
Notes
-----
Component names are case-insensitive.
Examples
--------
>>> site.has_component("Zxy")
True
"""
c = comp.strip().lower()
if c in ("tip", "tx", "ty", "tipper"):
tip = _extract_z_arrays(self.edi)["tipper"]
if tip is None:
return False
try:
a = np.asarray(tip, dtype=np.complex128)
except (TypeError, ValueError):
return False
return a.size > 0 and np.any(np.isfinite(a))
z = _extract_z_arrays(self.edi)["z"]
if z is None:
return False
z = _z_to_2d(z)
names = [n.lower() for n in _component_names()]
try:
i = names.index(c)
except ValueError:
return False
try:
ncols = z.shape[1]
except Exception:
ncols = 0
if i >= ncols:
return False
a = z[:, i]
return a.size > 0 and np.any(np.isfinite(a))
[docs]
def summary(self) -> dict[str, Any]:
r"""
Summarize site identity, geometry, and data coverage.
Returns
-------
dict
Keys include: ``name``, ``nfreq``, ``lat``, ``lon``,
``elev``, ``components`` (present Z components), and
``tipper`` (boolean).
Examples
--------
>>> s = site.summary()
>>> s["name"], s["nfreq"]
('E01', 37)
"""
lat, lon, elev = self.coords
arrs = _extract_z_arrays(self.edi)
f = (
np.asarray(arrs["freq"])
if arrs["freq"] is not None
else np.asarray([], float)
)
nfreq = int(f.size)
present = [c for c in _component_names() if self.has_component(c)]
return {
"name": self.name,
"nfreq": nfreq,
"lat": lat,
"lon": lon,
"elev": elev,
"components": present,
"tipper": bool(self.has_component("tipper")),
}
[docs]
def rename(
self,
new: str,
*,
inplace: bool = False,
) -> Site:
r"""
Set a new station identifier across common header
fields.
Parameters
----------
new : str
New station name or ID.
inplace : bool, optional
If ``True``, modify this instance. If ``False``,
return a new :class:`Site`. The default is ``False``.
Returns
-------
Site
The modified site (new instance unless ``inplace``).
Notes
-----
The method writes to EDI HEAD fields ``dataid``,
``station``, ``sitename``, ``name`` (best effort) and
also updates ``edi.name``. Downstream containers that
derive station IDs from the file stem may be configured
to prefer ``edi.name`` to preserve the rename.
Examples
--------
>>> s2 = site.rename("X_E01")
>>> s2.name
'X_E01'
"""
if inplace:
_set_station_name(self.edi, new)
return self # type: ignore[return-value]
ed = _clone_edi(self.edi)
_set_station_name(ed, new)
return Site(ed)
[docs]
def set_coords(
self,
lat: float,
lon: float,
elev: float | None = None,
*,
inplace: bool = False,
) -> Site:
r"""
Update the site coordinates in the EDI HEAD section.
Parameters
----------
lat : float
Latitude in decimal degrees.
lon : float
Longitude in decimal degrees.
elev : float, optional
Elevation in meters. If ``None``, elevation is left
unchanged. The default is ``None``.
inplace : bool, optional
If ``True``, modify this instance. If ``False``,
return a new :class:`Site`. The default is ``False``.
Returns
-------
Site
The modified site (new instance unless ``inplace``).
Notes
-----
Latitude and longitude are validated by utility functions
to ensure plausible values.
Examples
--------
>>> site = site.set_coords(10.0, 20.0, 100.0)
>>> site.coords
(10.0, 20.0, 100.0)
"""
if inplace:
_set_coords(
self.edi,
lat=float(lat),
lon=float(lon),
elev=(None if elev is None else float(elev)),
inplace=True,
)
return self # type: ignore[return-value]
ed = _clone_edi(self.edi)
_set_coords(
ed,
lat=float(lat),
lon=float(lon),
elev=(None if elev is None else float(elev)),
inplace=True,
)
return Site(ed)
[docs]
def set_empty(self, *, inplace: bool = False) -> Site:
r"""
Clear Z-related arrays to an empty dataset.
Parameters
----------
inplace : bool, optional
If ``True``, modify this instance. If ``False``,
return a new :class:`Site`. The default is ``False``.
Returns
-------
Site
The modified site (new instance unless ``inplace``).
Notes
-----
The following arrays are replaced with empty arrays:
``freq``, ``z``, ``z_error``, ``rho``, and ``phase``.
Use this to initialize a skeleton record without data.
Examples
--------
>>> s2 = site.set_empty()
>>> s2.to_dataframe("z").empty
True
"""
tgt = self.edi if inplace else _clone_edi(self.edi)
Z = _safe_get(tgt, "Z")
try:
if Z is None:
tgt.Z = type("Z", (), {})() # type: ignore
Z = tgt.Z # type: ignore[attr-defined]
Z.freq = np.asarray([], float)
Z.z = np.asarray([], float)
Z.z_error = np.asarray([], float)
Z.rho = np.asarray([], float)
Z.phase = np.asarray([], float)
except Exception:
pass
return self if inplace else Site(tgt)
[docs]
class Site(SiteMixin):
r"""
High-level wrapper for a single MT/CSAMT site backed by an
:class:`~pycsamt.seg.edi.EDIFile`. The class enforces a
stable, file-stem-based station identifier and exposes
convenient accessors for coordinates, impedance Z, tipper,
and derived quantities.
``Site`` also accepts an :class:`~pycsamt.emtf.document.EMTF` XML
document instead of an EDI object -- see :meth:`from_xml`. Either
way, :attr:`edi` and :attr:`tf` are both always available: the
representation not natively supplied is lazily materialized and
cached from the other via
:func:`pycsamt.emtf.converters.edi.edi_to_emtf`/``emtf_to_edi``,
so every existing EDI-shaped accessor keeps working unchanged
regardless of which format the site was built from.
When constructed from EDI, the constructor normalizes EDI
``HEAD`` fields so that ``dataid`` (and, if absent, ``station``)
matches the site stem resolved by :func:`_stem_from_edi`. If an
in-memory ``edi.name`` exists, the stem may prefer it so that a
prior rename is preserved. This normalization improves name-based
indexing, deterministic selection, and downstream joins in
collections. XML-native sites keep the station identity already
recorded in the document instead.
Parameters
----------
source : pycsamt.seg.edi.EDIFile or pycsamt.emtf.document.EMTF
Parsed SEG-EDI container, or an EMTF XML document, holding
one station. The EDI file may be constructed from disk or
synthesized in memory.
on_loss : {"warn", "raise", "ignore"}, default "warn"
Policy applied if :attr:`edi` is later materialized from an
XML-native site and the document holds metadata the
historical EDI format cannot represent.
Attributes
----------
edi : pycsamt.seg.edi.EDIFile
Underlying (possibly lazily materialized) EDI object. Use
with care; prefer the typed accessors of :class:`SiteMixin`
(e.g., ``freq``, ``z``).
tf : pycsamt.emtf.document.EMTF
Underlying (possibly lazily materialized) EMTF XML document.
See also :attr:`site_meta`, :attr:`provenance`,
:attr:`processing`, :attr:`site_layout`, :attr:`copyright`,
:attr:`quality_meta`.
backend : str
``"edi"`` or ``"xml"`` -- which representation is native.
name : str
Normalized station identifier. By default this equals
the file stem, unless a previous explicit rename is in
effect.
coords : tuple of float
Geographic location as ``(lat, lon, elev)``. Latitude
and longitude are in degrees; elevation is in meters.
Notes
-----
Identity policy
The site identity is derived from header fields in the
following order: ``dataid``, ``station``, ``sitename``,
``name``, ``STATION``. If none are present, the file
stem is used. The constructor ensures that ``dataid`` is
set to the resolved stem to stabilize lookups.
Array conventions
The impedance tensor :math:`Z` may be represented as a
3D array with shape ``(n, 2, 2)`` or flattened to
``(n, 4)`` in the component order
``Zxx, Zxy, Zyx, Zyy``. Frequencies are 1D of length
``n``. The tipper has two columns ``Tx`` and ``Ty``.
Derived fields
Apparent resistivity and phase are derived from
:math:`Z` and the frequency vector. When Z or frequency
slices are applied, recomputation is triggered after
arrays are aligned to avoid inconsistent shapes.
Robustness
Accessors are defensive. Missing arrays produce empty
frames or ``None``. Coordinates are validated for range
using utility checks.
Examples
--------
Basic construction and inspection
>>> from pycsamt.seg.edi import EDIFile
>>> from pycsamt.site.base import Site
>>> e = EDIFile("E01.edi") # parse from disk
>>> s = Site(e)
>>> s.name
'E01'
>>> s.coords # doctest: +ELLIPSIS
(..., ..., ...)
>>> s.summary()["nfreq"] >= 0
True
Tabular export
>>> from pycsamt.site.base import Site
>>> dfz = s.to_dataframe("z")
>>> list(dfz.columns)
['Zxx', 'Zxy', 'Zyx', 'Zyy']
>>> dfrp = s.to_dataframe("resphase")
>>> sorted([c for c in dfrp.columns if c.startswith("rho_")])[:2]
['rho_zxx', 'rho_zxy']
Rename without mutating the original instance
>>> s2 = s.rename("X_E01") # returns a new Site
>>> s2.name
'X_E01'
>>> s.name # original unchanged
'E01'
Coordinate update
>>> s3 = s.set_coords(10.0, 20.0, 100.0)
>>> s3.coords
(10.0, 20.0, 100.0)
With a collection
>>> from pycsamt.site.base import Sites
>>> e2 = EDIFile("E02.edi")
>>> col = Sites([s.edi, Site(e2).edi])
>>> [t.name for t in col]
['E01', 'E02']
>>> col["E02"].summary()["name"]
'E02'
See Also
--------
pycsamt.site.base.SiteMixin
Mixin providing typed accessors and utilities.
pycsamt.site.base.Sites
Collection helper for selection, slicing, and bulk edits.
pycsamt.seg.edi.EDIFile
Low-level SEG-EDI container.
References
----------
.. [Site-1] SEG EDI Format Specification. Society of Exploration
Geophysicists. Exchange format for magnetotelluric
and related EM data.
.. [Site-2] Chave, A. D., and Jones, A. G. (Eds.) (2012).
The Magnetotelluric Method. Cambridge University
Press.
.. [Site-3] Simpson, F., and Bahr, K. (2005). Practical
Magnetotellurics. Cambridge University Press.
"""
def __init__(self, source: Any, *, on_loss: str = "warn") -> None:
super().__init__(source, on_loss=on_loss)
if self._native == "edi":
try:
h = _ensure_head(self.edi)
nm = _stem_from_edi(self.edi)
if nm:
h.dataid = nm
if not getattr(h, "station", None):
h.station = nm
except Exception:
pass
[docs]
@classmethod
def from_edi(cls, source: Any, *, on_loss: str = "warn") -> Site:
r"""
Construct a :class:`Site` from an EDI object or a path to one.
Parameters
----------
source : pycsamt.seg.edi.EDIFile or str or pathlib.Path
Parsed EDI object, or a path that will be read via
``EDIFile(source)``.
on_loss : {"warn", "raise", "ignore"}, default "warn"
Policy used if this site's :attr:`tf` is later
materialized and something needs to be converted back.
Returns
-------
Site
A site backed natively by EDI.
See Also
--------
from_xml : Construct from an EMTF XML document instead.
"""
if isinstance(source, (str, bytes, Path, PathLike)):
source = EDIFile(source)
return cls(source, on_loss=on_loss)
[docs]
@classmethod
def from_xml(cls, source: Any, *, on_loss: str = "warn") -> Site:
r"""
Construct a :class:`Site` from an EMTF XML document or path.
Parameters
----------
source : pycsamt.emtf.document.EMTF or str or pathlib.Path
Parsed EMTF document, or a path that will be read via
:meth:`pycsamt.emtf.document.EMTF.from_xml`.
on_loss : {"warn", "raise", "ignore"}, default "warn"
Policy applied when this site's :attr:`edi` view is
later materialized and the document holds metadata the
historical EDI format cannot represent.
Returns
-------
Site
A site backed natively by the richer EMTF XML model;
:attr:`z`, :attr:`freq`, :attr:`tipper`, and friends
keep working exactly as for an EDI-backed site (via a
cached, lazily materialized EDI view), while
:attr:`tf`, :attr:`site_meta`, :attr:`provenance`, and
related properties expose the full richer metadata.
See Also
--------
from_edi : Construct from a SEG EDI object instead.
Site.to_xml : Serialize a site back to EMTF XML.
"""
from ..emtf.document import EMTF
if not isinstance(source, EMTF):
source = EMTF.from_xml(source)
return cls(source, on_loss=on_loss)
[docs]
def to_xml(self, target: Any = None, **kwargs) -> Any:
r"""
Serialize this site to EMTF XML.
Parameters
----------
target : str or pathlib.Path, optional
Destination path. If ``None``, the XML document is
returned as a string instead of being written.
**kwargs
Forwarded to
:meth:`pycsamt.emtf.document.EMTF.write_xml` (when
``target`` is given) or
:meth:`pycsamt.emtf.document.EMTF.to_xml` (otherwise).
Returns
-------
str or Any
The XML document as a string when ``target`` is
``None``; otherwise the return value of
:meth:`~pycsamt.emtf.document.EMTF.write_xml`.
See Also
--------
from_xml : The inverse constructor.
tf : The underlying :class:`EMTF` document used here.
"""
if target is None:
return self.tf.to_xml(**kwargs)
return self.tf.write_xml(target, **kwargs)
[docs]
@property
def site_layout(self) -> Any:
r"""``pycsamt.metadata.SiteLayout`` channel geometry (via :attr:`tf`)."""
return self.tf.site_layout
[docs]
@property
def provenance(self) -> Any:
r"""``pycsamt.metadata.ProvenanceMeta`` creator/submitter info."""
return self.tf.provenance
[docs]
@property
def processing(self) -> Any:
r"""``pycsamt.metadata.ProcessingMeta`` processing/software info."""
return self.tf.processing
[docs]
@property
def copyright(self) -> Any:
r"""``pycsamt.metadata.CopyrightInfo`` release/citation info."""
return self.tf.copyright
[docs]
def to_edi(self, *, copy: bool = False) -> EDIFile:
r"""
Return the underlying EDI object.
Parameters
----------
copy : bool, default False
If ``True``, return a best-effort deep copy of the
wrapped EDI object. If copying fails, the original
object is returned.
Returns
-------
pycsamt.seg.edi.EDIFile
EDI object wrapped by this ``Site``.
See Also
--------
pycsamt.site.base.to_edis
General unwrapping helper for ``Site``/``Sites`` and
mixed inputs.
"""
return _maybe_copy(self.edi) if copy else self.edi
def __repr__(self) -> str:
r"""
Debug-friendly, one-line summary of the site.
Returns
-------
str
A compact string with site name, number of
frequencies, and coordinates.
Examples
--------
>>> repr(site) # doctest: +ELLIPSIS
"Site(name='E01', nfreq=..., coords=(...,...,...))"
"""
s = self.summary()
return (
f"Site(name={s['name']!r}, nfreq={s['nfreq']}, "
f"coords=({s['lat']:.5f},{s['lon']:.5f},"
f"{s['elev']:.1f}))"
)
[docs]
class Sites(CoreObject):
r"""
Container for multiple :class:`~pycsamt.site.base.Site`
objects with convenient indexing, selection, and bulk edit
operations.
``Sites`` wraps each provided :class:`~pycsamt.seg.edi.EDIFile`
into a :class:`~pycsamt.site.base.Site`, ensuring that station
identity is normalized consistently. You can iterate, index by
integer, or look up by case-insensitive station name. Bulk
operations such as renaming, frequency slicing, and masking
are provided via :meth:`edit_all`.
Parameters
----------
edic : pycsamt.seg.collection.EDICollection or sequence of
pycsamt.seg.edi.EDIFile, pycsamt.emtf.document.EMTF, or
:class:`Site`
Parsed EDI collection, or any sequence mixing EDI objects,
EMTF XML documents, and already-constructed :class:`Site`
instances. Items are wrapped into :class:`Site` instances
(existing ``Site`` items are used as-is) in the order
provided.
Attributes
----------
_items : list of Site
Internal sequence of sites. This is considered private.
Iterate over ``Sites`` instead of accessing it directly.
Notes
-----
Identity and lookup
Each site inside the container uses the same naming
policy as :class:`Site`. Name-based lookups with
``sites["E01"]`` are case-insensitive and match the
normalized station name. If duplicates exist, the first
match is returned.
Order preservation
Ordering of input items is preserved. Integer indexing
with ``sites[i]`` retrieves the i-th :class:`Site`.
Bulk edits
:meth:`edit_all` supports three common operations:
- ``rename``: compute a new name from the old name.
- ``freq_slice``: apply a frequency slice (``slice``)
consistently across Z, freq, errors, and derived fields.
- ``mask``: apply a boolean mask to the Z tensor rows.
Use ``inplace=True`` to modify the container, otherwise a
new ``Sites`` is returned.
Geospatial helpers
:meth:`closest` uses geodetic distance to find the nearest
site to a given latitude and longitude. Sites with missing
or non-finite coordinates are skipped.
Topography integration
:meth:`with_topography` aligns site coordinates and
elevation with a user-provided frame, returning a new
container unless ``inplace=True`` is requested.
Profile conversion
:meth:`to_profile` attempts to build a
:class:`~pycsamt.site.profile.Profile` when that optional
dependency is available. Otherwise, a lightweight dict
describing a chainage-ordered sequence is returned.
Persistence
:meth:`write` emits one EDI file per site into a target
directory. If an EDI object provides ``to_file``, it is
used; otherwise a placeholder is written.
Robustness and errors
- ``__getitem__`` raises ``KeyError`` when a name is not
found. Prefer :meth:`get` to obtain ``None`` instead.
- Bulk operations ignore missing arrays on a best-effort
basis so that other arrays can still be processed.
Examples
--------
Build from a few EDI files
>>> from pycsamt.seg.edi import EDIFile
>>> from pycsamt.site.base import Sites
>>> e1, e2 = EDIFile("E01.edi"), EDIFile("E02.edi")
>>> sites = Sites([e1, e2])
>>> len(sites)
2
>>> [s.name for s in sites]
['E01', 'E02']
Indexing and lookup
>>> sites[0].name
'E01'
>>> sites["e02"].name
'E02'
>>> sites.get("missing") is None
True
Bulk rename without mutating the original container
>>> def rnm(n):
... return "X_" + n
>>> out = sites.edit_all(rename=rnm, inplace=False)
>>> [s.name for s in out]
['X_E01', 'X_E02']
>>> [s.name for s in sites]
['E01', 'E02']
Frequency slice across all arrays
>>> sl = slice(1, None) # drop the first frequency
>>> out2 = sites.edit_all(freq_slice=sl, inplace=False)
>>> f0 = sites["E01"].freq
>>> f1 = out2["E01"].freq
>>> len(f1) == len(f0) - 1
True
Selection by names or predicate
>>> subset = sites.select(names=["E02"])
>>> [s.name for s in subset]
['E02']
>>> # predicate: keep sites with tipper available
>>> subset2 = sites.select(
... predicate=lambda s: s.has_component("tipper")
... )
>>> isinstance(subset2, Sites)
True
Nearest station to a target location
>>> near = sites.closest(lat=10.0, lon=20.0, tol=None)
>>> near is None or hasattr(near, "name")
True
Persist to a directory
>>> import tempfile, pathlib
>>> tmp = pathlib.Path(tempfile.mkdtemp())
>>> out_paths = sites.write(
... tmp, template="{station}.edi", exist_ok=True
... )
>>> all(p.exists() for p in out_paths)
True
See Also
--------
pycsamt.site.base.Site
Site wrapper used for each element in the container.
pycsamt.seg.collection.EDICollection
Parsed collection produced by the core parser.
pycsamt.site.profile.Profile
Optional profile object produced by :meth:`to_profile`.
References
----------
.. [Sites-1] SEG EDI Format Specification. Society of Exploration
Geophysicists.
.. [Sites-2] Chave, A. D., and Jones, A. G. (2012). The
Magnetotelluric Method. Cambridge University Press.
.. [Sites-3] Simpson, F., and Bahr, K. (2005). Practical
Magnetotellurics. Cambridge University Press.
"""
def __init__(
self,
edic: EDICollection | Sequence[EDIFile | Any | Site],
*,
on_loss: str = "warn",
) -> None:
if isinstance(edic, EDIFile) or isinstance(edic, Site) or _is_xml_like(edic):
edic = [edic]
items = list(edic)
self._items: list[Site] = [
it if isinstance(it, Site) else Site(it, on_loss=on_loss)
for it in items
]
[docs]
@property
def edic(self) -> list[EDIFile]:
# Back-compat: expose the iterable of EDIFile expected by
# selection utilities (iter_edifiles works on sequences).
return [s.edi for s in self._items]
def __len__(self) -> int:
r"""
Number of sites in the container.
Returns
-------
int
Count of contained :class:`~pycsamt.site.base.Site`
objects.
Examples
--------
>>> len(sites) >= 0
True
"""
return len(self._items)
def __iter__(self):
r"""
Iterate over contained :class:`~pycsamt.site.base.Site`
objects.
Returns
-------
iterator of Site
Iterator yielding sites in input order.
Examples
--------
>>> names = [s.name for s in sites]
>>> isinstance(names, list)
True
"""
return iter(self._items)
def __getitem__(self, key: int | str) -> Site:
r"""
Retrieve a site by zero-based index or by case-insensitive
station name.
Parameters
----------
key : int or str
Integer index or station name. Name comparison is
case-insensitive.
Returns
-------
Site
The matching site.
Raises
------
KeyError
If ``key`` is a name and no site matches.
Examples
--------
>>> sites[0].name # by index
'E01'
>>> sites["e02"].name # by name, case-insensitive
'E02'
See Also
--------
get : Safe lookup returning ``None`` on failure.
by_index : Explicit index-based accessor.
"""
if isinstance(key, int):
return self._items[key]
nm = str(key).lower()
for s in self._items:
if s.name.lower() == nm:
return s
raise KeyError(key)
[docs]
def by_index(self, i: int) -> Site:
r"""
Retrieve a site by zero-based index.
Parameters
----------
i : int
Position in the container.
Returns
-------
Site
The site at the requested index.
Examples
--------
>>> sites.by_index(0).name == sites[0].name
True
"""
return self._items[i]
[docs]
def get(self, name: str) -> Site | None:
r"""
Safe lookup by case-insensitive station name.
Parameters
----------
name : str
Station identifier to find.
Returns
-------
Site or None
Matching site, or ``None`` if not present.
Examples
--------
>>> sites.get("missing") is None
True
>>> sites.get("E01").name
'E01'
See Also
--------
__getitem__ : Raises on missing names.
"""
try:
return self[name]
except Exception:
return None
[docs]
def as_list(self) -> list[EDIFile]:
r"""
Return the underlying list of EDI objects.
Returns
-------
list of pycsamt.seg.edi.EDIFile
The EDI objects corresponding to each site.
Notes
-----
This is useful when passing the dataset to utilities that
operate on EDI-level structures rather than on sites.
Examples
--------
>>> edis = sites.as_list()
>>> hasattr(edis[0], "get_section")
True
"""
return [s.edi for s in self._items]
[docs]
@property
def ordering(self) -> dict[str, Any]:
"""Describe the most recent ordering decision for this container."""
return dict(
getattr(
self,
"_ordering",
{"requested": "input", "applied": "input"},
)
)
[docs]
def ordered(
self,
by: str | None = None,
*,
inplace: bool = False,
min_linearity: float | None = None,
max_cross_track_ratio: float | None = None,
min_coordinate_fraction: float | None = None,
) -> Sites:
"""Return sites in a deterministic spatial or identity order.
``by='auto'`` applies chainage ordering only when the finite
coordinates describe one credible, approximately straight survey
line. Otherwise it preserves input order. ``by='chainage'`` forces
projection onto the PCA profile axis; sites without usable
coordinates are retained at the end in their original order.
Other modes are ``'input'``, ``'station'`` (natural numeric station
order), ``'latitude'``/``'lat'``, and ``'longitude'``/``'lon'``.
Station names remain identifiers and are never changed.
"""
aliases = {
"name": "station",
"natural": "station",
"lat": "latitude",
"lon": "longitude",
"profile": "chainage",
"spatial": "chainage",
"none": "input",
"preserve": "input",
}
from ..api.ordering import PYCSAMT_ORDERING
requested = (
str(PYCSAMT_ORDERING.mode if by is None else by).strip().lower()
)
min_linearity = float(
PYCSAMT_ORDERING.min_linearity
if min_linearity is None
else min_linearity
)
max_cross_track_ratio = float(
PYCSAMT_ORDERING.max_cross_track_ratio
if max_cross_track_ratio is None
else max_cross_track_ratio
)
min_coordinate_fraction = float(
PYCSAMT_ORDERING.min_coordinate_fraction
if min_coordinate_fraction is None
else min_coordinate_fraction
)
mode = aliases.get(requested, requested)
allowed = {
"auto",
"input",
"station",
"latitude",
"longitude",
"chainage",
}
if mode not in allowed:
raise ValueError(
f"by must be one of {sorted(allowed)}, got {by!r}"
)
# Reorder via the existing Site objects, not `self.as_list()`
# (which forces every site through its materialized `.edi`
# view) -- that would silently downgrade XML-native sites to
# EDI-backed on every non-inplace `.ordered()` call, including
# the one `ensure_sites` always performs.
target = self if inplace else Sites(self._items)
n = len(target._items)
report: dict[str, Any] = {
"requested": mode,
"applied": mode,
"n_sites": n,
"n_coordinates": 0,
}
if n < 2 or mode == "input":
report["applied"] = "input"
target._ordering = report
return target
def natural_key(site: Site) -> tuple[Any, ...]:
parts = re.split(r"(\d+)", site.name)
return tuple(
int(p) if p.isdigit() else p.casefold() for p in parts
)
if mode == "station":
target._items.sort(key=natural_key)
target._ordering = report
return target
coords = np.asarray([s.coords[:2] for s in target._items], dtype=float)
finite = np.isfinite(coords).all(axis=1)
valid_idx = np.flatnonzero(finite)
report["n_coordinates"] = int(valid_idx.size)
if valid_idx.size < 2:
report.update(
applied="input", reason="fewer than two finite coordinates"
)
target._ordering = report
return target
if mode in {"latitude", "longitude"}:
column = 0 if mode == "latitude" else 1
target._items = [
target._items[i]
for i in sorted(
range(n),
key=lambda i: (
not finite[i],
coords[i, column] if finite[i] else np.inf,
natural_key(target._items[i]),
),
)
]
target._ordering = report
return target
# Work in local metres before PCA; raw degree-space PCA biases the
# east component by latitude and is unsuitable for chainage.
ll = coords[finite]
lat0, lon0 = np.mean(ll, axis=0)
east = (ll[:, 1] - lon0) * 111_000.0 * math.cos(math.radians(lat0))
north = (ll[:, 0] - lat0) * 111_000.0
xy = np.column_stack((east, north))
centred = xy - xy.mean(axis=0)
_, singular, vh = np.linalg.svd(centred, full_matrices=False)
axis = vh[0]
along = centred @ axis
across = centred @ np.array([-axis[1], axis[0]])
variance = singular * singular
linearity = (
float(variance[0] / variance.sum()) if variance.sum() else 0.0
)
span = float(np.ptp(along))
cross_ratio = (
float(np.ptp(across) / span) if span > 0.0 else float("inf")
)
# A pair of parallel profiles can still have excellent global PCA
# linearity. Detect a well-populated gap across the fitted axis so
# such lines are ordered independently instead of interleaved.
cross_order = np.argsort(across)
cross_gaps = np.diff(across[cross_order])
profile_groups: list[np.ndarray] = [np.arange(valid_idx.size)]
if cross_gaps.size and np.ptp(across) > 0.0:
split = int(np.argmax(cross_gaps)) + 1
min_group = max(2, int(math.ceil(0.15 * valid_idx.size)))
gap_ratio = float(cross_gaps[split - 1] / np.ptp(across))
if (
split >= min_group
and (valid_idx.size - split) >= min_group
and gap_ratio >= 0.25
and cross_ratio >= 0.05
):
profile_groups = [cross_order[:split], cross_order[split:]]
report["cross_track_gap_ratio"] = gap_ratio
coordinate_fraction = float(valid_idx.size / n)
report.update(
linearity=linearity,
cross_track_ratio=cross_ratio,
coordinate_fraction=coordinate_fraction,
n_profiles=len(profile_groups),
)
if mode == "auto" and (
coordinate_fraction < float(min_coordinate_fraction)
or linearity < float(min_linearity)
or cross_ratio > float(max_cross_track_ratio)
or span <= 0.0
):
report.update(
applied="input",
reason="coordinates do not define one credible straight line",
)
target._ordering = report
return target
if len(profile_groups) > 1:
# Refit and order each detected line independently. A single
# global axis is close for parallel lines but can reverse tiny
# within-line offsets or ties compared with each line's own fit.
groups = [valid_idx[group] for group in profile_groups]
groups.sort(key=lambda group: int(np.min(group)))
ordered_items: list[Site] = []
for group in groups:
subset = Sites([target._items[int(i)].edi for i in group])
ordered_items.extend(subset.ordered("chainage")._items)
ordered_items.extend(
target._items[i] for i in range(n) if not finite[i]
)
target._items = ordered_items
report.update(
applied="chainage_by_line",
azimuth=math.degrees(math.atan2(axis[0], axis[1])) % 360.0,
span_m=span,
)
target._ordering = report
return target
# Resolve PCA's sign ambiguity with natural station numbers when
# they carry a useful monotonic signal; otherwise use a deterministic
# north/east-positive orientation.
numbers = []
for idx in valid_idx:
found = re.findall(r"\d+", target._items[int(idx)].name)
numbers.append(float(found[-1]) if found else float("nan"))
numbers_a = np.asarray(numbers, dtype=float)
if np.isfinite(numbers_a).sum() >= 2 and np.nanstd(numbers_a) > 0:
ok = np.isfinite(numbers_a)
corr = float(np.corrcoef(numbers_a[ok], along[ok])[0, 1])
if np.isfinite(corr) and corr < 0:
along *= -1.0
axis *= -1.0
report["station_chainage_correlation"] = (
abs(corr) if np.isfinite(corr) else None
)
elif axis[int(np.argmax(np.abs(axis)))] < 0:
along *= -1.0
axis *= -1.0
azimuth = math.degrees(math.atan2(axis[0], axis[1])) % 360.0
chainage = {
int(idx): float(value) for idx, value in zip(valid_idx, along)
}
ordered_idx = sorted(
range(n),
key=lambda i: (
i not in chainage,
chainage.get(i, np.inf),
i,
),
)
target._items = [target._items[i] for i in ordered_idx]
report.update(applied="chainage", azimuth=azimuth, span_m=span)
target._ordering = report
return target
[docs]
def to_edis(
self,
*,
copy: bool = False,
progress: bool | str = False,
verbose: int = 0,
) -> list[EDIFile]:
r"""
Return the underlying EDI objects as a list.
Parameters
----------
copy : bool, default False
If ``True``, return best-effort deep copies of the EDI
objects.
progress : bool or {'auto'}, default False
Enable progress display while unwrapping.
verbose : int, default 0
Verbosity forwarded to progress/reporting helpers.
Returns
-------
list of pycsamt.seg.edi.EDIFile
EDI objects in site order.
See Also
--------
to_edicollection : Return an ``EDICollection`` instead.
pycsamt.site.base.to_edis : General unwrapping helper.
"""
out = to_edis(
self,
copy=copy,
progress=progress,
verbose=verbose,
)
return out if isinstance(out, list) else [out]
[docs]
def to_edicollection(
self,
*,
copy: bool = False,
progress: bool | str = False,
verbose: int = 0,
) -> EDICollection:
r"""
Return the underlying EDI objects as an ``EDICollection``.
Parameters
----------
copy : bool, default False
If ``True``, return best-effort deep copies of the EDI
objects.
progress : bool or {'auto'}, default False
Enable progress display while unwrapping.
verbose : int, default 0
Verbosity assigned to the returned collection.
Returns
-------
pycsamt.seg.collection.EDICollection
Collection built from the underlying EDI objects.
"""
return to_edis(
self,
as_collection=True,
copy=copy,
progress=progress,
verbose=verbose,
)
[docs]
def to_emtf_list(self) -> list[Any]:
r"""
Return the underlying EMTF XML documents as a list.
Returns
-------
list of pycsamt.emtf.document.EMTF
One document per site, in site order. EDI-native sites
are lazily converted and cached via
:attr:`Site.tf`; XML-native sites are returned as-is,
with no information loss.
See Also
--------
write_xml : Persist these documents to a directory.
to_edis : The EDI-side equivalent.
"""
return [s.tf for s in self._items]
[docs]
def write_xml(self, outdir: str | Path, **kwargs) -> list[Path]:
r"""
Write one EMTF XML file per site into a directory.
Parameters
----------
outdir : str or pathlib.Path
Destination directory. It is created if missing.
**kwargs
Forwarded to :meth:`Site.to_xml` for each site.
Returns
-------
list of pathlib.Path
Paths to the written files, named ``"{station}.xml"``.
See Also
--------
to_emtf_list : In-memory equivalent without writing to disk.
Sites.write : The EDI-side equivalent.
"""
outdir = Path(outdir)
outdir.mkdir(parents=True, exist_ok=True)
paths: list[Path] = []
for s in self._items:
p = outdir / f"{s.name}.xml"
s.to_xml(p, **kwargs)
paths.append(p)
return paths
[docs]
def closest(
self,
lat: float,
lon: float,
tol: float | None = None,
) -> Site | None:
r"""
Find the closest site to a target coordinate using geodetic
distance.
Parameters
----------
lat : float
Target latitude in decimal degrees.
lon : float
Target longitude in decimal degrees.
tol : float, optional
Maximum allowed distance in meters. If provided and the
nearest site is farther than ``tol``, return ``None``.
The default is ``None``.
Returns
-------
Site or None
Nearest site or ``None`` if all sites are too far or
lack valid coordinates.
Notes
-----
Coordinates are validated. Sites with non-finite values are
skipped. Distance is computed in meters using a geodetic
model.
Examples
--------
>>> near = sites.closest(10.0, 20.0)
>>> near is None or hasattr(near, "name")
True
>>> sites.closest(0.0, 0.0, tol=1.0) is None
True
"""
from .location import distance
la = float(assert_lat_value(lat)) # type: ignore[arg-type]
lo = float(assert_lon_value(lon)) # type: ignore[arg-type]
best: tuple[Site | None, float]
best = (None, float("inf"))
for s in self._items:
sla, slo, _ = s.coords
if not np.isfinite(sla) or not np.isfinite(slo):
continue
d = distance((la, lo), (float(sla), float(slo)), mode="geodetic")
if d < best[1]:
best = (s, d)
if best[0] is None:
return None
if (tol is not None) and (best[1] > tol):
return None
return best[0]
[docs]
def map(self, fn: Callable[[Site], Any]) -> list[Any]:
r"""
Apply a function to every site and collect the results.
Parameters
----------
fn : callable
Function of signature ``fn(site) -> Any``.
Returns
-------
list
Results collected in order.
Examples
--------
>>> sites.map(lambda s: s.name)[:2]
['E01', 'E02']
"""
return [fn(s) for s in self._items]
[docs]
def rename(
self,
names: Any,
*,
inplace: bool = False,
missing: str = "raise",
allow_duplicates: bool = False,
) -> Sites:
"""Rename stations from a mapping, sequence, or callable."""
from .metadata import rename_sites
return rename_sites(
self,
names,
inplace=inplace,
missing=missing,
allow_duplicates=allow_duplicates,
)
[docs]
def edit_all(
self,
*,
rename: Callable[[str], str] | None = None,
freq_slice: slice | None = None,
mask: Callable[[pd.DataFrame], pd.Series] | None = None,
inplace: bool = False,
) -> Sites:
r"""
Bulk-edit all sites with optional rename, frequency slicing,
and tensor masking.
Parameters
----------
rename : callable, optional
Function ``rename(old_name) -> new_name``. If provided,
each site is renamed accordingly.
freq_slice : slice, optional
Row-wise slice applied consistently to frequency,
impedance Z, errors, and derived arrays.
mask : callable, optional
Function ``mask(df) -> bool_series`` where ``df`` is the
output of ``site.to_dataframe("z")``. Rows where the
mask is ``False`` are set to ``NaN`` in Z.
inplace : bool, optional
If ``True``, modify this container and return it. If
``False``, return a new :class:`Sites`. Default is
``False``.
Returns
-------
Sites
The edited container (possibly the same instance).
Notes
-----
- When ``inplace=False``, sites are shallow-cloned so that
edits do not affect the original container.
- Frequency slicing is applied atomically to avoid temporary
shape mismatches between ``freq`` and Z-derived arrays.
- Missing arrays are tolerated on a best-effort basis.
Examples
--------
Rename with a prefix
>>> def rnm(n):
... return "X_" + n
>>> out = sites.edit_all(rename=rnm)
>>> [s.name for s in out][:2]
['X_E01', 'X_E02']
Slice away the first frequency
>>> sl = slice(1, None)
>>> out2 = sites.edit_all(freq_slice=sl)
>>> len(out2["E01"].freq) == len(sites["E01"].freq) - 1
True
Mask the top half of rows in Z
>>> def top_half(frame):
... m = np.ones(len(frame), dtype=bool)
... m[: len(frame) // 2] = False
... return m
>>> out3 = sites.edit_all(mask=top_half)
>>> df = out3["E01"].to_dataframe("z")
>>> np.isnan(df.iloc[0].values).all()
True
See Also
--------
Site.rename : Per-site rename helper.
Site.to_dataframe : Source for building masks.
"""
items: list[EDIFile] = []
for s in self._items:
t = s if inplace else Site(_clone_edi(s.edi))
if rename:
newn = rename(t.name)
t = t.rename(newn, inplace=True)
if freq_slice is not None:
Z = _safe_get(t.edi, "Z", default=None)
if Z is not None:
_slice_fields(Z, freq_slice)
if mask is not None:
try:
df = t.to_dataframe("z")
m = np.asarray(mask(df))
Z = _safe_get(t.edi, "Z", default=None)
if Z is not None and hasattr(Z, "z"):
z = np.asarray(Z.z)
z = z.copy()
z[~m] = np.nan
Z.z = z
except Exception:
pass
items.append(t.edi)
return self if inplace else Sites(items)
[docs]
def with_topography(
self,
frame: Any,
*,
inplace: bool = False,
) -> Sites:
r"""
Align site coordinates and elevation from a tabular frame.
Parameters
----------
frame : Any
A table-like object (e.g., :class:`pandas.DataFrame`)
with station identifiers and columns for latitude,
longitude, and elevation. Column names are resolved by
the topography utility.
inplace : bool, optional
If ``True``, modify this container and return it. If
``False``, return a new :class:`Sites`. Default is
``False``.
Returns
-------
Sites
Container with updated coordinates.
Notes
-----
Sites are matched by normalized station identifiers. The
operation is performed on a best-effort basis; unmatched
stations are left unchanged.
Examples
--------
>>> import pandas as pd
>>> df = pd.DataFrame(
... {
... "station": ["E01", "E02"],
... "latitude": [10.0, 11.0],
... "longitude": [20.0, 21.0],
... "elevation": [100.0, 200.0],
... }
... )
>>> out = sites.with_topography(df, inplace=False)
>>> tuple(round(v, 3) for v in out["E01"].coords)
(10.0, 20.0, 100.0)
"""
from .location import apply_topography
edis = self.as_list()
out = apply_topography(edis, frame, inplace=inplace)
# apply_topography returns list when given list
return self if inplace else Sites(out)
[docs]
def select(
self,
names: Sequence[str] | None = None,
predicate: Callable[[Site], bool] | None = None,
) -> Sites:
r"""
Filter sites by explicit names or by a boolean predicate.
Parameters
----------
names : sequence of str, optional
Case-insensitive station names to retain. If provided,
this takes precedence over ``predicate``.
predicate : callable, optional
Function ``predicate(site) -> bool``. Sites for which
the function returns ``True`` are retained.
Returns
-------
Sites
New container with the selected sites.
Notes
-----
If neither ``names`` nor ``predicate`` is provided, the
method returns a shallow copy of the current container.
Examples
--------
>>> subset = sites.select(names=["E02"])
>>> [s.name for s in subset]
['E02']
>>> subset2 = sites.select(predicate=lambda s: s.has_component("Zxy"))
>>> isinstance(subset2, Sites)
True
"""
# Filter via the existing Site objects, not `.edi`, so an
# XML-native site stays XML-native through a select() (the
# same class of bug fixed in `ordered()`).
out: list[Site]
if names:
wanted = {str(n).lower() for n in names}
out = [s for s in self._items if s.name.lower() in wanted]
elif predicate:
out = [s for s in self._items if predicate(s)]
else:
out = list(self._items)
return Sites(out)
[docs]
@classmethod
def from_any(
cls,
source: Any,
topo_src: Any | None = None,
) -> Sites:
r"""
Construct a container from heterogeneous inputs by using a
normalized loading session.
Parameters
----------
source : Any
Input that can be understood by the loader. Supported
cases include:
- an :class:`EDICollection`,
- a list of :class:`EDIFile`,
- a single :class:`EDIFile`,
- or an iterable that yields EDIs.
topo_src : Any, optional
Optional topography source passed to the session for
use during loading. The default is ``None``.
Returns
-------
Sites
Parsed container. Returns an empty container if the
input cannot be interpreted.
Notes
-----
The method uses :func:`~pycsamt.session.normalize_session`
to handle parsing, discovery, and optional topography
alignment in a consistent way.
Examples
--------
>>> from pycsamt.seg.collection import EDICollection
>>> # edicol = ... # suppose we already parsed a folder
>>> # sites = Sites.from_any(edicol)
>>> # isinstance(sites, Sites)
True
"""
with normalize_session(".tmp", topo_src=topo_src) as nz:
obj = nz.load(source)
if isinstance(obj, EDICollection):
return cls(obj)
if isinstance(obj, list) and all(isinstance(x, EDIFile) for x in obj):
return cls(obj) # type: ignore[arg-type]
if isinstance(obj, EDIFile):
return cls([obj])
try:
return cls(list(obj)) # type: ignore[arg-type]
except Exception:
return cls([])
[docs]
def write(
self,
outdir: str | Path,
*,
template: str = "{station}.edi",
exist_ok: bool = False,
) -> list[Path]:
r"""
Write one EDI file per site to a directory.
Parameters
----------
outdir : str or pathlib.Path
Destination directory. It is created if missing.
template : str, optional
Filename template. The token ``{station}`` is replaced
by the normalized station name. Default is
``"{station}.edi"``.
exist_ok : bool, optional
If ``False`` and a file already exists, raise
``FileExistsError``. If ``True``, overwrite. Default is
``False``.
Returns
-------
list of pathlib.Path
Paths to the written files.
Raises
------
FileExistsError
If a target file exists and ``exist_ok`` is ``False``.
Notes
-----
Serialization goes through
:func:`pycsamt.emtf.converters.edi.write_edi`, which writes
the underlying :class:`~pycsamt.seg.edi.EDIFile` to an exact
target path (working around
:meth:`~pycsamt.seg.edi.EDIFile.write`'s own default of
deriving a ``new_``-prefixed name under a local ``edi_out/``
directory when not given an explicit ``new_edifn``/
``savepath``).
Examples
--------
>>> import tempfile, pathlib
>>> tmp = pathlib.Path(tempfile.mkdtemp())
>>> paths = sites.write(tmp, exist_ok=True)
>>> all(p.exists() for p in paths)
True
"""
from ..emtf.converters.edi import write_edi
outp = Path(outdir)
outp.mkdir(parents=True, exist_ok=True)
paths: list[Path] = []
for s in self._items:
name = s.name or "site"
fn = template.format(station=name)
p = outp / fn
if p.exists() and not exist_ok:
raise FileExistsError(str(p))
write_edi(s.edi, p)
paths.append(p)
return paths
[docs]
def to_profile(
self,
origin: tuple[float, float],
azimuth: float,
*,
crs: int | None = None,
) -> Any:
r"""
Convert sites to a 1D profile aligned with a specified
azimuth, returning either a rich Profile object or a
lightweight fallback.
Parameters
----------
origin : tuple of float
``(lat, lon)`` in decimal degrees defining the profile
origin.
azimuth : float
Profile azimuth in degrees. North is ``0`` and angles
increase clockwise.
crs : int, optional
Optional CRS code for libraries that require it. Not
used in the lightweight fallback. Default is ``None``.
Returns
-------
Profile or dict
If :class:`~pycsamt.site.profile.Profile` is available,
a Profile is returned. Otherwise a dict is returned with
keys ``"origin"``, ``"azimuth"``, and ``"sites"`` in
chainage order.
Notes
-----
The fallback computes local chainage by a flat approximation
around ``origin``:
.. math::
ch = dx * \\sin(az) + dy * \\cos(az)
where ``dx`` and ``dy`` are metric offsets relative to the
origin. Sites lacking valid coordinates are skipped.
Examples
--------
>>> prof = sites.to_profile(origin=(0.0, 0.0), azimuth=90.0)
>>> hasattr(prof, "chainages") or isinstance(prof, dict)
True
See Also
--------
pycsamt.site.profile.Profile
Rich profile object when available.
"""
# Try to return a real Profile first
try:
from .location import Coord
from .profile import Profile # type: ignore
prof = Profile.from_sites(
self._items,
origin=Coord(float(origin[0]), float(origin[1]), 0.0),
azimuth=float(azimuth),
)
return prof
except Exception:
pass
# Lightweight fallback: sort by chainage using local flat metric
la0, lo0 = origin
la0 = float(assert_lat_value(la0)) # type: ignore[arg-type]
lo0 = float(assert_lon_value(lo0)) # type: ignore[arg-type]
az = math.radians(float(azimuth))
items: list[tuple[float, Site]] = []
for s in self._items:
la, lo, _ = s.coords
if not np.isfinite(la) or not np.isfinite(lo):
continue
dy = (la - la0) * 111_000.0
dx = (lo - lo0) * 111_000.0 * math.cos(math.radians(la0))
# correct projection: north=0 -> dx*sin + dy*cos
ch = dx * math.sin(az) + dy * math.cos(az)
items.append((ch, s))
items.sort(key=lambda t: t[0])
return {
"origin": origin,
"azimuth": azimuth,
"sites": [s for _, s in items],
}
[docs]
def to_sites(
x: Any,
*,
recursive: bool = True,
on_dup: str = "replace",
strict: bool = False,
verbose: int = 0,
):
r"""
Coerce an arbitrary EDI-like input into a ``Sites`` wrapper.
This helper normalizes many inputs to a uniform
:class:`~pycsamt.site.base.Sites` interface:
* If ``x`` is already a ``Sites`` instance, it is returned
unchanged.
* If ``x`` is an ``EDICollection`` or a sequence of
``EDIFile`` objects, a new ``Sites`` wrapper is created.
* If ``x`` is an iterable yielding EDI-like items, they are
collected and wrapped.
The operation is light-weight and does not deep-copy the
underlying EDI objects. The returned ``Sites`` simply holds
references to the same items.
Parameters
----------
x : Any
A ``Sites`` instance, an ``EDICollection``, a sequence of
``EDIFile`` objects, or an iterable yielding EDI-like
items.
Returns
-------
pycsamt.site.base.Sites
A ``Sites`` wrapper over the provided items.
Notes
-----
Use this utility at API boundaries to conveniently accept
multiple input forms while providing a consistent downstream
interface. If you need independent copies of the underlying
data, perform your own cloning before calling ``to_sites``.
Examples
--------
Wrap a list of EDIFile objects:
>>> from pycsamt.site.selection import to_sites
>>> s = to_sites([e1, e2, e3])
>>> len(s)
3
Wrap an existing Sites (no-op):
>>> s2 = to_sites(s)
>>> s2 is s
True
Wrap an EDICollection:
>>> s3 = to_sites(coll) # coll is an EDICollection
>>> [t.name for t in s3]
['A01', 'A02']
See Also
--------
pycsamt.site.base.Sites :
Wrapper providing per-site convenience methods.
pycsamt.site.base.Sites.from_any :
Alternate constructor with session normalization.
"""
return _to_sites(
x, recursive=recursive, on_dup=on_dup, strict=strict, verbose=verbose
)
[docs]
def to_edis(
x: Any,
*,
as_collection: bool = False,
copy: bool = False,
recursive: bool = True,
on_dup: str = "replace",
strict: bool = False,
verbose: int = 0,
progress: bool | str = False,
):
r"""
Unwrap site-like inputs to raw EDI objects.
This is the inverse boundary of :func:`to_sites`. It accepts a
single :class:`Site`, a :class:`Sites` collection, an
``EDICollection``, raw EDI objects, path-like inputs, or mixed
iterables containing those forms. The returned objects are the
underlying EDI containers used by low-level writers, exporters,
and EM processing functions.
Parameters
----------
x : Any
Site-like input to unwrap. Supported values include ``Site``,
``Sites``, ``EDIFile``, ``EDICollection``, path-like inputs,
or iterables containing site/EDI-like objects.
as_collection : bool, default False
If ``True``, return an :class:`~pycsamt.seg.collection.EDICollection`.
Otherwise a single input returns one EDI object and multi-item
inputs return a list.
copy : bool, default False
If ``True``, return best-effort deep copies of the EDI objects.
If copying fails for an item, that item is returned unchanged.
recursive : bool, default True
Forwarded to path-like discovery through ``EDICollection``.
on_dup : {'replace', 'keep', 'keep_first', 'keep_last', 'raise'}, default 'replace'
Duplicate station policy. ``replace`` and ``keep`` are forwarded
to path loading. ``keep_first``, ``keep_last``, and ``raise`` are
enforced after collection construction.
strict : bool, default False
If ``True``, raise when an object cannot be unwrapped to EDI.
If ``False``, invalid items are skipped.
verbose : int, default 0
Verbosity forwarded to collection construction and duplicate
policy diagnostics.
progress : bool or {'auto'}, default False
Enable progress display while unwrapping iterable inputs.
Returns
-------
EDIFile, list of EDIFile, or EDICollection
Raw EDI object(s), depending on the input shape and
``as_collection``.
Notes
-----
The operation is shallow by default. It returns the same EDI objects
wrapped by ``Site`` or ``Sites``. Pass ``copy=True`` when the caller
should be able to mutate the returned objects independently.
Examples
--------
>>> from pycsamt.site.base import Site, Sites, to_edis
>>> site = Site(edi)
>>> raw = to_edis(site)
>>> raw is edi
True
>>> raws = to_edis(Sites([edi]))
>>> len(raws)
1
>>> coll = to_edis(Sites([edi]), as_collection=True)
>>> len(coll)
1
See Also
--------
to_sites : Wrap raw EDI-like inputs into ``Sites``.
Site.to_edi : Convenience method for one ``Site``.
Sites.to_edis : Convenience method for a ``Sites`` collection.
"""
single = _is_single_edi_input(x)
items = _collect_edis(
x,
recursive=recursive,
on_dup=on_dup,
strict=strict,
verbose=verbose,
progress=progress,
)
if copy:
items = [_maybe_copy(ed) for ed in items]
coll = _edis_collection_from_items(
items,
on_dup=on_dup,
verbose=verbose,
)
if as_collection:
return coll
if single:
if items:
return items[0]
if strict:
raise ValueError(
"to_edis(strict=True): no EDI-like item could be "
"unwrapped from the provided input."
)
return list(coll)
def _is_pathlike(obj: Any) -> bool:
return isinstance(obj, (str, bytes, Path, PathLike))
def _is_seq_of_pathlike(x: Any) -> bool:
if isinstance(x, (str, bytes, Path, PathLike)):
return False # single path-like handled elsewhere
try:
it = iter(x) # noqa: F841
except Exception:
return False
# Heuristic: non-empty and all items path-like
try:
xs = list(x)
except Exception:
return False
return len(xs) > 0 and all(_is_pathlike(t) for t in xs)
def _is_edi_like(obj: Any) -> bool:
return (
obj is not None and hasattr(obj, "get_section") and hasattr(obj, "Z")
)
def _is_xml_like(obj: Any) -> bool:
"""Duck-type-safe check for an in-memory EMTF XML document.
Uses ``isinstance`` (unlike the EDI checks above) because
:class:`~pycsamt.emtf.document.EMTF` is a concrete, stable class
rather than a historically duck-typed format.
"""
if obj is None:
return False
from ..emtf.document import EMTF
return isinstance(obj, EMTF)
def _is_xml_path(obj: Any) -> bool:
"""True for a path-like object whose suffix is ``.xml``.
Works for plain file paths and simple ``*.xml``-suffixed glob
patterns alike, since :attr:`pathlib.PurePath.suffix` is purely
lexical (it does not touch the filesystem).
"""
if not _is_pathlike(obj):
return False
try:
return Path(str(obj)).suffix.lower() == ".xml"
except Exception:
return False
def _load_xml(obj: Any) -> Any:
"""Return an :class:`EMTF` document from *obj* (object or path)."""
from ..emtf.document import EMTF
if isinstance(obj, EMTF):
return obj
return EMTF.from_xml(obj)
def _tf_station_name(tf: Any) -> str:
for attr in ("station_id", "station"):
v = _safe_get(tf, attr, default=None)
if v:
return str(v)
site = _safe_get(tf, "site", default=None)
for attr in ("site_id", "name"):
v = _safe_get(site, attr, default=None)
if v:
return str(v)
return "site"
def _tf_meta(tf: Any) -> dict[str, Any]:
out: dict[str, Any] = {}
name = _tf_station_name(tf)
if name:
out["station"] = name
out["dataid"] = name
site = _safe_get(tf, "site", default=None)
nm = _safe_get(site, "name", default=None)
if nm:
out["name"] = nm
out["sitename"] = nm
for key in ("lat", "lon", "elev"):
v = _safe_get(tf, key, default=None)
if v is not None:
out[key] = v
return out
def _discover_xml_sites(
dir_path: Any, *, recursive: bool = True, on_loss: str = "warn"
) -> list[Site]:
"""Glob ``*.xml`` under *dir_path* and wrap each as a :class:`Site`.
Malformed/unrelated XML files are skipped rather than raised,
matching the tolerant, best-effort spirit of the EDI-side
discovery this runs alongside.
"""
p = Path(dir_path)
if not p.is_dir():
return []
it = p.rglob("*.xml") if recursive else p.glob("*.xml")
out: list[Site] = []
for m in sorted(it):
if not m.is_file():
continue
try:
out.append(Site(_load_xml(m), on_loss=on_loss))
except Exception:
continue
return out
def _partition_xml_sources(
x: Any, *, recursive: bool = True, on_loss: str = "warn"
) -> tuple[list[Site], Any, bool]:
"""Split XML-native content out of *x* for :func:`_to_sites`.
Returns ``(xml_sites, edi_remainder, touched)``. When ``touched``
is ``False``, *x* contained nothing XML-native and
``edi_remainder`` is *x* itself, unmodified -- callers should
fall through to the original, unchanged EDI-only coercion in
that case. ``edi_remainder`` is what should still be handed to
that EDI-only coercion (possibly *x* itself, a filtered list, or
an empty list).
"""
if _is_xml_like(x):
return [Site(x, on_loss=on_loss)], [], True
if _is_xml_path(x):
return [Site(_load_xml(x), on_loss=on_loss)], [], True
if isinstance(x, Site):
if x.backend == "xml":
return [x], [], True
return [], x, False
# EDI-like objects (and anything that merely looks EDI-like) must
# never be iterated into -- short-circuit before the generic
# iterable fallback below.
if _is_edi_like(x) or isinstance(x, (EDIFile, Sites, EDICollection)):
return [], x, False
if _is_pathlike(x):
p = Path(x)
if p.is_dir():
found = _discover_xml_sites(p, recursive=recursive, on_loss=on_loss)
if found:
return found, x, True
return [], x, False
if hasattr(x, "__iter__"):
try:
items = list(x)
except Exception:
return [], x, False
xml_sites: list[Site] = []
edi_remainder: list[Any] = []
touched = False
for item in items:
if _is_xml_like(item):
xml_sites.append(Site(item, on_loss=on_loss))
touched = True
elif _is_xml_path(item):
xml_sites.append(Site(_load_xml(item), on_loss=on_loss))
touched = True
elif isinstance(item, Site) and item.backend == "xml":
xml_sites.append(item)
touched = True
elif (
_is_pathlike(item)
and not _is_edi_like(item)
and Path(item).is_dir()
):
found = _discover_xml_sites(
item, recursive=recursive, on_loss=on_loss
)
if found:
xml_sites.extend(found)
touched = True
edi_remainder.append(item)
else:
edi_remainder.append(item)
return xml_sites, edi_remainder, touched
return [], x, False
def _dedup_sites(
sites: list[Site], *, policy: str = "replace", verbose: int = 0
) -> list[Site]:
"""List-level counterpart of :func:`_dedup_collection_names`.
Used to apply the same ``on_dup`` vocabulary uniformly across a
combined EDI+XML site list (which may include XML-native
:class:`Site` objects that never pass through an
:class:`EDICollection`).
Notes
-----
``"replace"``/``"keep"`` are *not* no-ops here, unlike a bare
EDI-only collection (where same-format duplicates are already
resolved while the collection is being built). This list combines
two independently-built halves -- the EDI-only sites and the
XML-native sites discovered alongside them -- so a station present
in *both* formats (e.g. ``18-001.edi`` and ``18-001.xml`` in the
same folder) would otherwise survive as two separate entries
regardless of ``on_dup``. ``"replace"`` keeps the last occurrence
(matching :class:`~pycsamt.seg.collection.EDICollection`'s
"new entry wins" semantics -- the XML entry, since it is appended
after the EDI ones) and ``"keep"`` keeps the first.
"""
key = (policy or "replace").strip().lower()
if len(sites) < 2:
return sites
if key == "replace":
key = "keep_last"
elif key == "keep":
key = "keep_first"
name_to_idx: dict[str, int] = {}
names: list[str] = []
for i, s in enumerate(sites):
n = s.name
names.append(n)
if key == "keep_first":
name_to_idx.setdefault(n, i)
elif key == "keep_last":
name_to_idx[n] = i
elif key == "raise":
if n in name_to_idx:
raise ValueError(
f"Duplicate site '{n}' encountered with on_dup='raise'."
)
name_to_idx[n] = i
else:
name_to_idx[i] = i
if verbose > 0:
dups = sorted({n for n in names if names.count(n) > 1})
if dups:
warnings.warn(
f"to_sites: duplicates {dups} handled with policy='{key}'.",
RuntimeWarning,
stacklevel=2,
)
return [sites[i] for i in sorted(name_to_idx.values())]
def _is_single_edi_input(x: Any) -> bool:
if isinstance(x, Site):
return True
if _is_edi_like(x):
return True
if _is_xml_like(x) or _is_xml_path(x):
return True
if _is_pathlike(x):
return False
return False
def _unwrap_one_edi(x: Any, *, strict: bool = False) -> EDIFile | None:
if isinstance(x, Site):
return x.edi
if _is_xml_like(x):
return Site(x).edi
if _is_xml_path(x):
return Site(_load_xml(x)).edi
edi = getattr(x, "edi", None)
if _is_edi_like(edi):
return edi
if _is_edi_like(x):
return x
if strict:
raise TypeError(
"Object cannot be unwrapped to an EDI-like item: "
f"{type(x).__name__}."
)
return None
def _edis_collection_from_items(
items: Sequence[Any],
*,
on_dup: str = "replace",
verbose: int = 0,
) -> EDICollection:
try:
coll = EDICollection(items=items, verbose=verbose)
except TypeError:
coll = EDICollection(items, verbose=verbose) # type: ignore
if on_dup.strip().lower() in {"keep_first", "keep_last", "raise"}:
coll = _dedup_collection_names(
coll,
policy=on_dup,
verbose=verbose,
)
return coll
def _collect_edis(
x: Any,
*,
recursive: bool = True,
on_dup: str = "replace",
strict: bool = False,
verbose: int = 0,
progress: bool | str = False,
) -> list[EDIFile]:
if isinstance(x, Sites):
seq = list(x)
elif isinstance(x, Site) or _is_edi_like(x):
seq = [x]
elif _is_pathlike(x) or _is_seq_of_pathlike(x):
coll_on_dup = _map_on_dup_for_collection(on_dup)
coll = EDICollection.from_sources(
x,
recursive=recursive,
strict=strict,
on_dup=coll_on_dup,
verbose=verbose,
)
if on_dup.strip().lower() == "raise":
coll = _dedup_collection_names(
coll,
policy="raise",
verbose=verbose,
)
seq = list(coll)
elif isinstance(x, EDICollection):
seq = list(x)
else:
try:
seq = list(x)
except Exception:
seq = [x]
out: list[EDIFile] = []
iterator = iter_progress(
seq,
enabled=progress,
desc="Unwrapping EDI",
unit="site",
total=len(seq),
)
for item in iterator:
if isinstance(item, (Sites, EDICollection)) or _is_pathlike(item):
out.extend(
_collect_edis(
item,
recursive=recursive,
on_dup=on_dup,
strict=strict,
verbose=verbose,
progress=False,
)
)
continue
if _is_seq_of_pathlike(item):
out.extend(
_collect_edis(
item,
recursive=recursive,
on_dup=on_dup,
strict=strict,
verbose=verbose,
progress=False,
)
)
continue
ed = _unwrap_one_edi(item, strict=strict)
if ed is not None:
out.append(ed)
if strict and not out:
raise ValueError(
"to_edis(strict=True): no EDI-like items could be "
"unwrapped from the provided input."
)
return out
def _map_on_dup_for_collection(on_dup: str) -> str:
"""
EDICollection.from_sources supports {'replace', 'keep'}.
Map broader API to that set.
"""
key = (on_dup or "replace").strip().lower()
if key in {"replace", "keep"}:
return key
if key in {"keep_first"}:
return "keep"
if key in {"keep_last"}:
return "replace"
if key in {"raise"}:
# no native support; we'll handle after construction
return "replace"
raise ValueError(
"Invalid on_dup policy. Allowed: "
"'replace', 'keep', 'keep_first', 'keep_last', 'raise'."
)
def _dedup_collection_names(coll, *, policy: str, verbose: int = 0):
"""
Apply post-hoc duplicate policy ('keep_first', 'keep_last', 'raise')
when the collection was built from non-path inputs.
"""
key = (policy or "replace").strip().lower()
if key in {"replace", "keep"}:
return coll # already handled at load time
# Build index by item name
try:
items = list(coll.items.values())
except Exception:
# Fallback: assume EDICollection-like iterability
items = list(coll)
def _name(it, i):
for attr in ("station", "name", "site", "id"):
if hasattr(it, attr):
v = getattr(it, attr)
if isinstance(v, str) and v:
return v
return f"site_{i}"
name_to_idx: dict[str, int] = {}
names: list[str] = []
for i, it in enumerate(items):
n = _name(it, i)
names.append(n)
if key == "keep_first":
name_to_idx.setdefault(n, i)
elif key == "keep_last":
name_to_idx[n] = i
elif key == "raise":
if n in name_to_idx:
raise ValueError(
f"Duplicate site '{n}' encountered with on_dup='raise'."
)
name_to_idx[n] = i
if verbose > 0:
dups = sorted({n for n in names if names.count(n) > 1})
if dups:
warnings.warn(
f"to_sites: duplicates {dups} handled with policy='{key}'.",
RuntimeWarning,
stacklevel=2,
)
kept = [items[i] for i in sorted(name_to_idx.values())]
# Rebuild a new collection of the same class with filtered items
CollCls = coll.__class__
try:
return CollCls(items=kept, verbose=getattr(coll, "verbose", 0))
except TypeError:
return CollCls(kept, verbose=getattr(coll, "verbose", 0))
def _to_sites(
x: Any,
*,
recursive: bool = True,
on_dup: str = "replace",
strict: bool = False,
verbose: int = 0,
_skip_xml: bool = False,
):
r"""
Internal helper that implements :func:`to_sites`.
See :func:`to_sites` for parameters and behavior.
This version additionally recognizes path-like inputs and
uses EDICollection.from_sources(...) to honor discovery
controls (recursive/strict/on_dup/verbose).
"""
# Fast path
if isinstance(x, Sites):
return x
# 0) XML-aware pre-pass: pull out anything XML-native (an EMTF
# object, a .xml path, an XML-backed Site, or *.xml files found
# inside a directory) without disturbing the EDI-only logic below.
# When nothing XML-native is found, `touched` is False and
# `edi_remainder` is `x` itself, unmodified -- the original,
# unchanged coercion runs exactly as before. `_skip_xml` guards
# the recursive call below: for a single directory, `edi_remainder`
# is that same directory (its *.edi half is handled by the
# unmodified logic further down), so re-running this pre-pass on
# it would re-discover the same *.xml files and recurse forever.
if not _skip_xml:
xml_sites, edi_remainder, touched = _partition_xml_sources(
x, recursive=recursive, on_loss="warn"
)
else:
xml_sites, edi_remainder, touched = [], x, False
if touched:
# `_partition_xml_sources` already extracted every XML-native
# item in one pass (including *.xml files inside any directory
# it walked), so the recursive call below only ever needs to
# run the original, unmodified EDI-only logic -- never another
# XML-discovery pass, which would otherwise re-find (and
# duplicate) the same *.xml files.
edi_only = _to_sites(
edi_remainder,
recursive=recursive,
on_dup=on_dup,
strict=False,
verbose=verbose,
_skip_xml=True,
)
combined = list(edi_only) + xml_sites
combined = _dedup_sites(combined, policy=on_dup, verbose=verbose)
if strict and not combined:
raise ValueError(
"to_sites(strict=True): no EDI- or EMTF-like items "
"could be coerced from the provided input."
)
return Sites(combined)
# 1) Path-like (single or sequence) → from_sources(...)
if _is_pathlike(x) or _is_seq_of_pathlike(x):
coll_on_dup = _map_on_dup_for_collection(on_dup)
coll = EDICollection.from_sources(
x,
recursive=recursive,
strict=strict,
on_dup=coll_on_dup,
verbose=verbose,
)
# If caller asked for 'raise', enforce it now.
if on_dup.strip().lower() == "raise":
coll = _dedup_collection_names(
coll, policy="raise", verbose=verbose
)
# Wrap in Sites
try:
return Sites(coll)
except TypeError:
return Sites(edic=coll)
# 2a) Single Site object — wrap its EDIFile directly
if isinstance(x, Site):
return Sites([x.edi])
# 2b) Non-path inputs → coerce to collection
try:
coll = as_edicollection(
x,
strict=strict,
verbose=verbose,
)
except TypeError:
coll = as_edicollection(x)
if coll is None:
if strict:
raise ValueError(
"to_sites(strict=True): no EDI-like items could be "
"coerced from the provided input."
)
coll = EDICollection(items=[], verbose=verbose)
# Enforce any post-hoc duplicate policy not covered by the collection
if on_dup.strip().lower() in {"keep_first", "keep_last", "raise"}:
coll = _dedup_collection_names(coll, policy=on_dup, verbose=verbose)
# Sites wrapper
try:
return Sites(coll)
except TypeError:
return Sites(edic=coll)
def _station_name(edi: EDIFile) -> str:
h = _get_head(edi)
if h is not None:
for k in ("dataid", "station", "sitename", "name", "STATION"):
v = _safe_get(h, k, default=None)
if v:
return str(v)
nm = _stem_from_edi(edi)
return nm or "site"
def _z_to_2d(a: Any) -> np.ndarray:
z = np.asarray(a)
if z.ndim == 3 and z.shape[-2:] == (2, 2):
n = z.shape[0]
# row-major flatten: xx, xy, yx, yy
return z.reshape(n, 4)
return z
def _set_station_name(edi: EDIFile, name: str) -> None:
head = _ensure_head(edi)
# best effort: write all common identifiers
for k in ("dataid", "station", "sitename", "name", "STATION"):
try:
setattr(head, k, name)
except Exception:
pass
try:
edi.name = name
except Exception:
pass
def _safe_get(obj: Any, *names: str, default: Any = None) -> Any:
for n in names:
try:
return getattr(obj, n)
except:
pass
try:
return obj[n] # type: ignore[index]
except:
pass
return default
def _clone_edi(ed: EDIFile) -> EDIFile:
try:
return _maybe_copy(ed)
except:
c = EDIFile()
for k, v in getattr(ed, "__dict__", {}).items():
try:
setattr(c, k, copy.copy(v))
except:
setattr(c, k, v)
return c
def _component_names() -> list[str]:
return ["Zxx", "Zxy", "Zyx", "Zyy"]
def _stem_from_edi(edi: EDIFile) -> str:
# Prefer an explicit in-memory name set by rename()
v = _safe_get(edi, "name", default=None)
if v:
return str(v)
# Fall back to on-disk/file-like attributes
for attr in ("path", "filepath", "filename", "file", "source"):
p = _safe_get(edi, attr, default=None)
if p:
try:
return Path(str(p)).stem
except Exception:
pass
return ""
def _extract_z_arrays(ed: EDIFile) -> dict[str, Any]:
Z = _safe_get(ed, "Z", default=None)
out: dict[str, Any] = {}
out["freq"] = _safe_get(Z, "freq", "frequency")
out["z"] = _safe_get(Z, "z", "impedance")
out["z_err"] = _safe_get(Z, "z_error", "z_err", "impedance_err")
out["rho"] = _safe_get(Z, "rho", "res", "resistivity")
out["phase"] = _safe_get(Z, "phase", "phi")
tip = _safe_get(ed, "T", "TIP", "Tip", "tipper", "Tipper")
if tip is None:
tip = _safe_get(Z, "tipper", "tip")
if tip is not None and not isinstance(tip, np.ndarray):
# ``tip`` may be a Tipper wrapper object rather than the raw
# array; unwrap it so downstream consumers see a plain array
# (or None when no tipper data is stored).
tip = _safe_get(tip, "tipper", "tip", default=tip)
out["tipper"] = tip
return out
def _slice_fields(Z: Any, sl: slice) -> None:
"""
Slice Z fields atomically to keep freq/z aligned.
Use private attrs to avoid partial recomputations.
"""
# Fetch (prefer privates if present)
def _ga(name: str, alt: str | None = None):
if hasattr(Z, f"_{name}"):
return getattr(Z, f"_{name}")
if hasattr(Z, name):
return getattr(Z, name)
if alt and hasattr(Z, alt):
return getattr(Z, alt)
return None
f = _ga("freq")
z = _ga("z")
ze = _ga("z_err", "z_error")
rho = _ga("rho", "resistivity")
ph = _ga("phase", "phi")
rot = getattr(Z, "rotation_angle", None)
# Build sliced copies (guard Nones)
fs = np.asarray(f)[sl] if f is not None else None
zs = np.asarray(z)[sl] if z is not None else None
zes = np.asarray(ze)[sl] if ze is not None else None
rhos = np.asarray(rho)[sl] if rho is not None else None
phs = np.asarray(ph)[sl] if ph is not None else None
# Assign to private storage to prevent immediate recompute
if fs is not None and hasattr(Z, "_freq"):
Z._freq = fs
if zs is not None and hasattr(Z, "_z"):
Z._z = zs
if zes is not None and hasattr(Z, "_z_err"):
Z._z_err = zes
if rhos is not None and hasattr(Z, "_rho"):
Z._rho = rhos
if phs is not None and hasattr(Z, "_phase"):
Z._phase = phs
# Keep rotation vector aligned if it’s an array
try:
if rot is not None and np.ndim(rot) == 1 and fs is not None:
ra = np.asarray(rot)
if ra.size == np.asarray(f).size:
Z.rotation_angle = ra[sl]
except Exception:
pass
# Finalize: compute rho/phi once arrays are consistent
try:
if (
getattr(Z, "_z", None) is not None
and getattr(Z, "_freq", None) is not None
):
Z.compute_resistivity_phase()
except Exception:
# Be tolerant; tests only require freq slicing to succeed
pass
