pytem 0.1.0__py3-none-any.whl

pytem/__init__.py

@@ -0,0 +1,219 @@
+"""
+pyTEM: 1-D layered-earth TEM modelling.
+
+Supports circular and square loop geometries, central and offset receivers,
+DLF and Euler transforms, with NumPy / Numba / CuPy backends.
+
+================================================================================
+Package map (table of contents)
+================================================================================
+
+The package is organised in layers.  Lower layers hold pure data and physics
+with no project-internal dependencies; higher layers compose them into the
+public forward, inversion, and plotting API.  Arrows below read "imports from".
+
+Layer 0 - Data and constants
+    transform_weights.py   MU0 and the digital-filter coefficient tables:
+                           Key (2009) 201/101-pt Hankel J0/J1 filters, Key
+                           81/101-pt Fourier sine/cosine filters, and the
+                           Euler weights.  Exposes the HANKEL_FILTERS,
+                           FOURIER_FILTERS, EULER_PARAMS registries.  No
+                           internal dependencies; everything else builds on it.
+
+Layer 1 - Core physics and backend detection
+    recursion.py           te_reflection_coeff / te_reflection_coeff_grad:
+                           the TE reflection coefficient of the
+                           layer stack and its log-resistivity gradient (NumPy,
+                           complex).  This is the physical heart of the model.
+                               -> transform_weights
+    backends.py            CuPy/CUDA detection (HAS_CUDA) and transfer of the
+                           filter tables to device memory (GPU_HANKEL,
+                           GPU_FOURIER).
+                               -> transform_weights
+
+Layer 2 - Compute kernels (one physics, three backends)
+    kernels_numba.py       Numba JIT forward kernels (circular/square, DLF and
+                           Euler).  Sets HAS_NUMBA.
+    kernels_gpu.py         CuPy/CUDA forward kernels (mirror of kernels_numba).
+                               -> transform_weights, backends
+    kernels_jacobian.py    Adjoint upward recursion: forward+backward pass giving
+                           d(r_TE)/d(ln rho_j) for all layers at once, in both
+                           Numba and CuPy variants.  Backs the analytical
+                           Jacobian.
+                               -> transform_weights, backends
+
+Layer 3 - Forward modelling
+    forward.py             Public fwd_circle_* / fwd_square_* functions, the
+                           analytical half-space references, and the geometry
+                           builders + backend dispatchers (CUDA > Numba >
+                           pure Python).  Also exports the filter-precompute
+                           helpers reused by the inversion module.
+                               -> transform_weights, backends, recursion,
+                                  kernels_numba, kernels_gpu
+
+Layer 4 - Inversion
+    inversion.py           Analytical (getJ_ana) and finite-difference
+                           (getJ_fd) Jacobians, regularisation helpers (getR,
+                           getAlpha[s], getRMS), apparent-resistivity
+                           conversion, and the Gauss-Newton invert() loop.
+                               -> forward (+ its filter helpers),
+                                  kernels_jacobian, kernels_numba, backends,
+                                  transform_weights
+
+Utilities and add-ons (composable, mostly standalone)
+    waveform.py            Convolution of the step response with a piecewise-
+                           linear transmitter waveform (setup_waveform,
+                           convolve_waveform).      -> kernels_numba
+    system_filter.py       Butterworth and WalkTEM cascade H(omega) transfer
+                           functions fed in as system_filter=.   (standalone)
+    ip_models.py           Complex-resistivity IP models (Pelton, Cole-Cole,
+                           double-Pelton, MPA) and tem_forward_ip.
+                               -> transform_weights, recursion
+    euler.py               Standalone Euler inverse-Laplace transform, used to
+                           verify the production Euler path.      (standalone)
+    plotter.py             Matplotlib helpers: plot_sounding, plot_model,
+                           plot_inversion.                        (standalone)
+    data_io.py             Field-data import: read_tem_xyz / read_kenbec_xyz /
+                           read_xyz parse TEM Data Manager and TEMImage-Beta
+                           `.xyz` exports into TEMData / KenbecTEMData, with
+                           shared dbdt/snr/to_pytem and line/UTM/distance
+                           helpers.                               (standalone)
+    survey.py              Survey: quick-look sounding, transect and station-
+                           map plots built on data_io containers. LM/HM are
+                           always plotted separately (different waveforms).
+                               -> data_io
+
+    __init__.py            Re-exports the public API listed in __all__.
+
+--------------------------------------------------------------------------------
+Dependency flow
+--------------------------------------------------------------------------------
+
+    transform_weights
+        |-> backends
+        |-> recursion
+        |-> kernels_numba / kernels_gpu / kernels_jacobian
+                |
+                v
+            forward  ----> inversion
+                ^               ^
+                |               |
+            (waveform, system_filter, ip_models feed in here as options)
+
+--------------------------------------------------------------------------------
+Typical call chains
+--------------------------------------------------------------------------------
+
+Forward:
+    fwd_circle_central()                              (forward.py)
+        _resolve_filters()        -> transform_weights tables
+        _filter_weights()         -> optional system_filter samples
+        _build_circular_geometry()-> per-wavenumber weights
+        _run_circular()           -> CUDA | Numba | pure-Python kernel
+            te_reflection_coeff() -> upward recursion (recursion.py / kernels)
+        _apply_signal_scaling()   -> current and step-off/on/impulse sign
+
+Inversion:
+    invert()                                          (inversion.py)
+        loop:
+            forward call (one of fwd_*)               (forward.py)
+            getJ_ana()  -> adjoint kernels            (kernels_jacobian.py)
+                        or getJ_fd() -> repeated forward calls
+            _alpha_search() / _gn_solve() / _backtrack()  Gauss-Newton step
+        optional: convolve_waveform(), system_filter=H, IP rho(omega)
+"""
+
+__version__ = "0.1.0"
+
+from .transform_weights import MU0, HANKEL_FILTERS, FOURIER_FILTERS, EULER_PARAMS
+from .backends import HAS_CUDA
+from .kernels_numba import HAS_NUMBA
+from .recursion import te_reflection_coeff, te_reflection_coeff_grad
+
+from .forward import (
+    fwd_circle_central,
+    fwd_circle_offset,
+    fwd_square_central,
+    fwd_square_offset,
+    fwd_analytical_central,
+    fwd_analytical_offset,
+)
+
+from .waveform import convolve_waveform, setup_waveform, setup_waveform_matrix, setup_shared_gate_matrices
+from .system_filter import butterworth_filter, cascade_filter
+from .euler import euler_invert
+from .inversion import (getJ_ana, getJ_fd, getR, dbdt_to_apprho, getRMS,
+                        getAlpha, getAlphas, invert, invert_joint, invert_stations,
+                        compute_doi)
+from .plotter import (plot_sounding, plot_model, plot_inversion,
+                      plot_survey_models, plot_survey_responses)
+from .benchmarks import benchmark_step_grids
+from .data_io import (TEMData, KenbecTEMData, TunoeTEMData, read_tem_xyz, read_kenbec_xyz,
+                      read_xyz, read_usf)
+from .survey import Survey
+
+from .ip_models import (
+    pelton_res_rho,
+    cole_cole_rho,
+    double_pelton_rho,
+    get_m_taur_MPA,
+    mpa_rho,
+    tem_forward_ip,
+)
+
+__all__ = [
+    # Constants & flags
+    'MU0', 'HAS_CUDA', 'HAS_NUMBA',
+    'HANKEL_FILTERS', 'FOURIER_FILTERS', 'EULER_PARAMS',
+    # Core
+    'te_reflection_coeff',
+    'te_reflection_coeff_grad',
+    # Forward models
+    'fwd_circle_central',
+    'fwd_circle_offset',
+    'fwd_square_central',
+    'fwd_square_offset',
+    # Analytical
+    'fwd_analytical_central',
+    'fwd_analytical_offset',
+    # Waveform & system filter
+    'convolve_waveform',
+    'setup_waveform',
+    'setup_waveform_matrix',
+    'setup_shared_gate_matrices',
+    'butterworth_filter',
+    'cascade_filter',
+    # IP models
+    'pelton_res_rho',
+    'cole_cole_rho',
+    'double_pelton_rho',
+    'get_m_taur_MPA',
+    'mpa_rho',
+    'tem_forward_ip',
+    # Euler (verification)
+    'euler_invert',
+    # Inversion
+    'getJ_ana',
+    'getJ_fd',
+    'getR',
+    'dbdt_to_apprho',
+    'getRMS',
+    'getAlpha',
+    'getAlphas',
+    'invert',
+    'invert_joint',
+    'compute_doi',
+    # Plotting
+    'plot_sounding',
+    'plot_model',
+    'plot_inversion',
+    # Field-data import & survey plotting
+    'TEMData',
+    'KenbecTEMData',
+    'TunoeTEMData',
+    'read_tem_xyz',
+    'read_kenbec_xyz',
+    'read_xyz',
+    'read_usf',
+    'Survey',
+]

pytem/backends.py

@@ -0,0 +1,94 @@
+"""
+backends.py - CuPy/CUDA detection and GPU filter array transfer.
+
+Provides:
+  - HAS_CUDA flag
+  - GPU_HANKEL, GPU_FOURIER registries (CuPy arrays on device)
+"""
+
+import os
+import glob
+import sys
+
+import numpy as np
+
+from .transform_weights import (
+    _HANKEL_BASE_201, _HANKEL_J0_201, _HANKEL_J1_201,
+    _HANKEL_BASE_101, _HANKEL_J0_101, _HANKEL_J1_101,
+    _FOURIER_BASE_81, _FOURIER_SIN_81, _FOURIER_COS_81,
+    _FOURIER_BASE_101, _FOURIER_SIN_101, _FOURIER_COS_101,
+)
+
+
+def _register_cuda_dll_dirs():
+    """Make CuPy find the CUDA runtime shipped as ``nvidia-*-cu12`` pip wheels.
+
+    On Windows the wheels install their DLLs under ``site-packages/nvidia/*/bin``.
+    CuPy's bundled loader does not search there, and nvrtc loads its builtins
+    DLL via the process PATH, so we add those directories to both PATH and the
+    DLL search path before CuPy is imported. No-op on non-Windows or when the
+    wheels are absent.
+    """
+    if os.name != "nt":
+        return
+    bins = []
+    for site_dir in {os.path.dirname(os.path.dirname(__file__)), *sys.path}:
+        base = os.path.join(site_dir, "nvidia")
+        if os.path.isdir(base):
+            bins.extend(glob.glob(os.path.join(base, "*", "bin")))
+    bins = sorted(set(p for p in bins if os.path.isdir(p)))
+    if not bins:
+        return
+    os.environ["PATH"] = os.pathsep.join(bins) + os.pathsep + os.environ.get("PATH", "")
+    for p in bins:
+        try:
+            os.add_dll_directory(p)
+        except (OSError, AttributeError):
+            pass
+
+
+_register_cuda_dll_dirs()
+
+try:
+    import cupy as cp
+    HAS_CUDA = cp.cuda.is_available()
+except ImportError:
+    HAS_CUDA = False
+
+GPU_HANKEL = {}
+GPU_FOURIER = {}
+
+if HAS_CUDA:
+    import cupy as cp
+
+    _GPU_DEVICE = cp.cuda.Device(0)
+    _GPU_DEVICE.use()
+
+    # 201-pt Hankel
+    _d_HANKEL_BASE_201 = cp.asarray(_HANKEL_BASE_201)
+    _d_HANKEL_J0_201   = cp.asarray(_HANKEL_J0_201, dtype=cp.complex128)
+    _d_HANKEL_J1_201   = cp.asarray(_HANKEL_J1_201, dtype=cp.complex128)
+
+    # 101-pt Hankel
+    _d_HANKEL_BASE_101 = cp.asarray(_HANKEL_BASE_101)
+    _d_HANKEL_J0_101   = cp.asarray(_HANKEL_J0_101, dtype=cp.complex128)
+    _d_HANKEL_J1_101   = cp.asarray(_HANKEL_J1_101, dtype=cp.complex128)
+
+    # 81-pt Fourier
+    _d_FOURIER_BASE_81 = cp.asarray(_FOURIER_BASE_81)
+    _d_FOURIER_SIN_81  = cp.asarray(_FOURIER_SIN_81)
+    _d_FOURIER_COS_81  = cp.asarray(_FOURIER_COS_81)
+
+    # 101-pt Fourier
+    _d_FOURIER_BASE_101 = cp.asarray(_FOURIER_BASE_101)
+    _d_FOURIER_SIN_101  = cp.asarray(_FOURIER_SIN_101)
+    _d_FOURIER_COS_101  = cp.asarray(_FOURIER_COS_101)
+
+    GPU_HANKEL = {
+        'key_201': (_d_HANKEL_BASE_201, _d_HANKEL_J0_201, _d_HANKEL_J1_201),
+        'key_101': (_d_HANKEL_BASE_101, _d_HANKEL_J0_101, _d_HANKEL_J1_101),
+    }
+    GPU_FOURIER = {
+        'key_81':  (_d_FOURIER_BASE_81,  _d_FOURIER_SIN_81,  _d_FOURIER_COS_81),
+        'key_101': (_d_FOURIER_BASE_101, _d_FOURIER_SIN_101, _d_FOURIER_COS_101),
+    }

pytem/benchmarks.py

@@ -0,0 +1,108 @@
+"""Numerical benchmarks for pyTEM forward-model workflows."""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+import pandas as pd
+
+from .forward import fwd_circle_offset, fwd_square_offset
+from .waveform import setup_shared_gate_matrices
+
+
+def benchmark_step_grids(
+    systems,
+    thicknesses,
+    models,
+    geometry,
+    tx_size,
+    rx_x=0.0,
+    rx_y=0.0,
+    grid_points=(20, 70, 120, 170, 220, 270, 320, 370, 400),
+    reference_points=1200,
+    repeats=3,
+    tolerance=1.0,
+    system_filter=None,
+    n_quad=5,
+    transform='dlf',
+    use_numba=True,
+    use_cuda=False,
+):
+    """Benchmark shared step-time grids against a high-resolution reference."""
+    if geometry not in ('circle_offset', 'square_offset'):
+        raise ValueError("geometry must be 'circle_offset' or 'square_offset'.")
+    if repeats < 1 or reference_points < 1:
+        raise ValueError('repeats and reference_points must be positive.')
+
+    thicknesses = np.asarray(thicknesses, dtype=float)
+    models = {str(name): np.asarray(model, dtype=float)
+              for name, model in models.items()}
+    points = sorted(set(int(value) for value in grid_points))
+    if not points or min(points) < 1:
+        raise ValueError('grid_points must contain positive integers.')
+
+    def predict(resistivities, step_times, gate_matrices):
+        options = dict(current=1.0, signal=-1, use_numba=use_numba,
+                       use_cuda=use_cuda, transform=transform,
+                       system_filter=system_filter)
+        if geometry == 'square_offset':
+            response = -fwd_square_offset(
+                thicknesses, resistivities, tx_size, rx_x, rx_y, step_times,
+                n_quad=n_quad, **options)
+        else:
+            response = -fwd_circle_offset(
+                thicknesses, resistivities, tx_size, rx_x, step_times, **options)
+        return {moment: gate_matrices[moment] @ response
+                for moment in systems}
+
+    reference_times, reference_matrices = setup_shared_gate_matrices(
+        systems, n_step=reference_points)
+    check_times, check_matrices = setup_shared_gate_matrices(
+        systems, n_step=2 * reference_points)
+    references = {name: predict(model, reference_times, reference_matrices)
+                  for name, model in models.items()}
+    reference_errors = []
+    for name in models:
+        check = predict(models[name], check_times, check_matrices)
+        for moment in systems:
+            scale = np.maximum(np.abs(check[moment]),
+                               max(np.max(np.abs(check[moment])) * 1e-10, 1e-300))
+            reference_errors.append(
+                np.max(100 * np.abs(references[name][moment] - check[moment]) / scale))
+
+    rows = []
+    for points_count in points:
+        started = time.perf_counter()
+        step_times, gate_matrices = setup_shared_gate_matrices(
+            systems, n_step=points_count)
+        setup_seconds = time.perf_counter() - started
+        for name, model in models.items():
+            predict(model, step_times, gate_matrices)
+            durations = []
+            prediction = None
+            for _ in range(repeats):
+                started = time.perf_counter()
+                prediction = predict(model, step_times, gate_matrices)
+                durations.append(time.perf_counter() - started)
+            for moment in systems:
+                reference = references[name][moment]
+                scale = np.maximum(np.abs(reference),
+                                   max(np.max(np.abs(reference)) * 1e-10, 1e-300))
+                error = 100 * np.abs(prediction[moment] - reference) / scale
+                rows.append({
+                    'Points': points_count,
+                    'Model': name,
+                    'Moment': moment,
+                    'RMS error [%]': np.sqrt(np.mean(error ** 2)),
+                    'Max error [%]': error.max(),
+                    'Prediction [ms]': 1000 * np.median(durations),
+                    'Setup [s]': setup_seconds,
+                    'Below tolerance': bool(error.max() <= tolerance),
+                })
+
+    result = pd.DataFrame(rows)
+    result.attrs['reference_max_error_percent'] = max(reference_errors, default=0.0)
+    result.attrs['reference_points'] = reference_points
+    result.attrs['tolerance'] = tolerance
+    return result