# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
from __future__ import annotations
import copy
import math
import warnings
from collections.abc import 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
----------
.. [1] SEG EDI Format. Society of Exploration Geophysicists.
Commonly used magnetotelluric exchange format.
"""
def __init__(self, edi: EDIFile) -> None:
self.edi = edi
[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`.
"""
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.
"""
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 :pyattr:`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.
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.
Parameters
----------
edi : pycsamt.seg.edi.EDIFile
Parsed SEG-EDI container holding one station. The file
may be constructed from disk or synthesized in memory.
Attributes
----------
edi : pycsamt.seg.edi.EDIFile
Underlying EDI object. Use with care; prefer the typed
accessors of :class:`SiteMixin` (e.g., ``freq``, ``z``).
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
----------
.. [1] SEG EDI Format Specification. Society of Exploration
Geophysicists. Exchange format for magnetotelluric
and related EM data.
.. [2] Chave, A. D., and Jones, A. G. (Eds.) (2012).
The Magnetotelluric Method. Cambridge University
Press.
.. [3] Simpson, F., and Bahr, K. (2005). Practical
Magnetotellurics. Cambridge University Press.
"""
def __init__(self, edi: EDIFile) -> None:
super().__init__(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]
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
Parsed EDI collection or any sequence of EDI objects.
Items are wrapped into :class:`Site` instances 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
----------
.. [1] SEG EDI Format Specification. Society of Exploration
Geophysicists.
.. [2] Chave, A. D., and Jones, A. G. (2012). The
Magnetotelluric Method. Cambridge University Press.
.. [3] Simpson, F., and Bahr, K. (2005). Practical
Magnetotellurics. Cambridge University Press.
"""
def __init__(
self,
edic: EDICollection | Sequence[EDIFile],
) -> None:
if isinstance(edic, EDIFile):
edic = [edic]
items = list(edic)
self._items: list[Site] = [Site(e) for e 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]
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 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 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
"""
out: list[EDIFile] = []
if names:
wanted = {str(n).lower() for n in names}
out.extend(s.edi for s in self._items if s.name.lower() in wanted)
elif predicate:
out.extend(s.edi for s in self._items if predicate(s))
else:
out = self.as_list()
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
-----
If an EDI object implements ``to_file``, it is used to
serialize. Otherwise a small placeholder file is written.
Examples
--------
>>> import tempfile, pathlib
>>> tmp = pathlib.Path(tempfile.mkdtemp())
>>> paths = sites.write(tmp, exist_ok=True)
>>> all(p.exists() for p in paths)
True
"""
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))
edi = s.edi
tf = getattr(edi, "to_file", None)
if callable(tf):
try:
tf(str(p))
paths.append(p)
continue
except Exception:
pass
with open(p, "w", encoding="utf-8") as f:
f.write(f"# EDI placeholder for {name}\n")
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_single_edi_input(x: Any) -> bool:
if isinstance(x, Site):
return True
if _is_edi_like(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
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,
):
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
# 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")
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
