# The Clear BSD License
#
# Copyright (c) 2026 Tobias Heibges
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted (subject to the limitations in the disclaimer
# below) provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from this
# software without specific prior written permission.
#
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""
Simple Inhomogeneous Model (Radially Symmetric)
Base class for radially-symmetric inhomogeneous materials where the
refractive index depends only on distance from a center point (altitude).
User implements only `n_at_altitude()` - a simple 1D function.
GPU support is automatic via precomputed lookup tables.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
import numpy as np
from numpy.typing import NDArray
from .material_field import MaterialField
if TYPE_CHECKING:
from ...propagation.kernels.registry import PropagationKernelID, PropagatorID
# Physical constant
EARTH_RADIUS = 6_371_000.0 # meters
[docs]
class SimpleInhomogeneousModel(MaterialField, ABC):
"""
Base class for radially-symmetric inhomogeneous materials.
User implements only `n_at_altitude()` - a simple 1D function that returns
the refractive index at a given altitude. The system automatically provides:
- `get_refractive_index(x, y, z, wavelength)` via altitude computation
- `get_refractive_index_gradient(x, y, z, wavelength)` via radial direction
- GPU support via precomputed lookup tables
Parameters
----------
name : str
Name of the material model
center : tuple
Center of spherical symmetry (x, y, z). Default: (0, 0, 0)
reference_radius : float
Radius at which altitude = 0 (e.g., Earth radius). Default: 6,371,000 m
altitude_range : tuple
(min_altitude, max_altitude) for lookup table. Default: (0, 200,000)
lut_resolution : int
Number of samples in lookup table. Default: 10,000
kernel : PropagationKernelID, optional
Override the default kernel. Supported: SIMPLE_EULER, SIMPLE_RK4
propagator : PropagatorID, optional
Override the default propagator. Supported: GPU_GRADIENT, CPU_GRADIENT
Example
-------
>>> class MyAtmosphere(SimpleInhomogeneousModel):
... def n_at_altitude(self, altitude, wavelength=None):
... return 1.0 + 0.000293 * math.exp(-altitude / 8500)
...
>>> atm = MyAtmosphere(name="My Atmosphere")
>>> propagator = GPUGradientPropagator(atm) # Auto GPU support
"""
# =========================================================================
# COMPATIBILITY DECLARATIONS
# =========================================================================
# Import lazily to avoid circular imports at class definition time
@classmethod
def _init_compatibility(cls):
"""Initialize compatibility declarations (called once on first use)."""
from ...propagation.kernels.registry import PropagationKernelID, PropagatorID
if not cls._supported_kernels:
cls._supported_kernels = [
PropagationKernelID.SIMPLE_EULER,
PropagationKernelID.SIMPLE_RK4,
]
cls._default_kernel = PropagationKernelID.SIMPLE_RK4
cls._supported_propagators = [
PropagatorID.GPU_GRADIENT,
PropagatorID.CPU_GRADIENT,
]
cls._default_propagator = PropagatorID.GPU_GRADIENT
[docs]
def __init__(
self,
name: str = "Simple Inhomogeneous",
center: tuple[float, float, float] = (0.0, 0.0, 0.0),
reference_radius: float = EARTH_RADIUS,
altitude_range: tuple[float, float] = (0.0, 200_000.0),
lut_resolution: int = 10000,
kernel: PropagationKernelID | None = None,
propagator: PropagatorID | None = None,
):
# Initialize compatibility declarations before calling super().__init__
self._init_compatibility()
super().__init__(name, kernel=kernel, propagator=propagator)
self._is_homogeneous = False
self.center = center
self.reference_radius = reference_radius
self.altitude_range = altitude_range
self.lut_resolution = lut_resolution
# LUT arrays initialized lazily on first GPU use
self._lut_initialized = False
self._lut_altitudes: NDArray[np.float32] | None = None
self._lut_n: NDArray[np.float32] | None = None
self._lut_dn_dh: NDArray[np.float32] | None = None
self._lut_alpha: NDArray[np.float32] | None = None # Absorption coefficient LUT
self._lut_delta_h: float = 0.0
# =========================================================================
# USER MUST IMPLEMENT
# =========================================================================
[docs]
@abstractmethod
def n_at_altitude(self, altitude: float, wavelength: float | None = None) -> float:
"""
Return refractive index at given altitude.
This is the only method users must implement.
Parameters
----------
altitude : float
Altitude above reference surface in meters (altitude >= 0)
wavelength : float, optional
Wavelength in meters (for dispersion models)
Returns
-------
float
Refractive index n >= 1.0
"""
pass
# =========================================================================
# USER MAY OVERRIDE
# =========================================================================
[docs]
def dn_dh_at_altitude(
self, altitude: float, wavelength: float | None = None
) -> float:
"""
Return dn/d(altitude) at given altitude.
Default implementation uses numerical differentiation.
Override for analytical derivatives (more accurate).
Parameters
----------
altitude : float
Altitude above reference surface in meters
wavelength : float, optional
Wavelength in meters
Returns
-------
float
Derivative dn/dh in m^-1
"""
eps = 1.0 # 1 meter step for numerical derivative
h_plus = max(altitude + eps, 0.0)
h_minus = max(altitude - eps, 0.0)
if h_plus == h_minus:
return 0.0
return (
self.n_at_altitude(h_plus, wavelength)
- self.n_at_altitude(h_minus, wavelength)
) / (h_plus - h_minus)
[docs]
def alpha_at_altitude(
self, altitude: float, wavelength: float | None = None
) -> float:
"""
Return absorption coefficient at given altitude.
Default implementation returns 0 (no absorption).
Override in subclasses that model absorbing media.
Parameters
----------
altitude : float
Altitude above reference surface in meters
wavelength : float, optional
Wavelength in meters
Returns
-------
float
Absorption coefficient α in m^-1 (for Beer-Lambert: I = I₀·exp(-α·s))
"""
return 0.0
# =========================================================================
# AUTO-GENERATED FROM n_at_altitude()
# =========================================================================
def _compute_altitude(
self,
x: float | NDArray[np.float64],
y: float | NDArray[np.float64],
z: float | NDArray[np.float64],
) -> float | NDArray[np.float64]:
"""Compute altitude from Cartesian coordinates."""
cx, cy, cz = self.center
dx = np.asarray(x) - cx
dy = np.asarray(y) - cy
dz = np.asarray(z) - cz
r = np.sqrt(dx**2 + dy**2 + dz**2)
altitude = r - self.reference_radius
# Clamp to non-negative
if isinstance(altitude, np.ndarray):
return np.maximum(altitude, 0.0)
return max(altitude, 0.0)
def _compute_radial_direction(
self,
x: float | NDArray[np.float64],
y: float | NDArray[np.float64],
z: float | NDArray[np.float64],
) -> tuple[
float | NDArray[np.float64],
float | NDArray[np.float64],
float | NDArray[np.float64],
]:
"""Compute unit radial vector (outward from center)."""
cx, cy, cz = self.center
dx = np.asarray(x) - cx
dy = np.asarray(y) - cy
dz = np.asarray(z) - cz
r = np.sqrt(dx**2 + dy**2 + dz**2)
if isinstance(r, np.ndarray):
r = np.where(r < 1e-10, 1.0, r)
elif r < 1e-10:
return 0.0, 0.0, 1.0 # Default to +z at center
return dx / r, dy / r, dz / r
[docs]
def get_refractive_index(
self,
x: float | NDArray[np.float64],
y: float | NDArray[np.float64],
z: float | NDArray[np.float64],
wavelength: float | NDArray[np.float64],
) -> float | NDArray[np.float64]:
"""Get refractive index at position (auto-generated from n_at_altitude)."""
altitude = self._compute_altitude(x, y, z)
if isinstance(altitude, np.ndarray):
# Vectorized evaluation
return np.array(
[self.n_at_altitude(float(h), wavelength) for h in altitude.flat]
).reshape(altitude.shape)
return self.n_at_altitude(float(altitude), wavelength)
[docs]
def get_refractive_index_gradient(
self,
x: float | NDArray[np.float64],
y: float | NDArray[np.float64],
z: float | NDArray[np.float64],
wavelength: float | NDArray[np.float64],
) -> tuple[
float | NDArray[np.float64],
float | NDArray[np.float64],
float | NDArray[np.float64],
]:
"""Get gradient of n at position (auto-generated)."""
altitude = self._compute_altitude(x, y, z)
rx, ry, rz = self._compute_radial_direction(x, y, z)
if isinstance(altitude, np.ndarray):
dn_dh = np.array(
[self.dn_dh_at_altitude(float(h), wavelength) for h in altitude.flat]
).reshape(altitude.shape)
else:
dn_dh = self.dn_dh_at_altitude(float(altitude), wavelength)
# grad(n) = (dn/dh) * r_hat
return dn_dh * rx, dn_dh * ry, dn_dh * rz
[docs]
def get_absorption_coefficient(
self,
x: float | NDArray[np.float64],
y: float | NDArray[np.float64],
z: float | NDArray[np.float64],
wavelength: float | NDArray[np.float64],
) -> float | NDArray[np.float64]:
"""Return absorption coefficient (default: 0)."""
if isinstance(x, np.ndarray):
return np.zeros_like(x)
return 0.0
[docs]
def get_scattering_coefficient(
self,
x: float | NDArray[np.float64],
y: float | NDArray[np.float64],
z: float | NDArray[np.float64],
wavelength: float | NDArray[np.float64],
) -> float | NDArray[np.float64]:
"""Return scattering coefficient (default: 0)."""
if isinstance(x, np.ndarray):
return np.zeros_like(x)
return 0.0
# =========================================================================
# GPU INTERFACE (Automatic via LUT)
# =========================================================================
def _initialize_lut(self, wavelength: float | None = None) -> None:
"""Build lookup table arrays for GPU."""
if self._lut_initialized:
return
min_alt, max_alt = self.altitude_range
altitudes = np.linspace(min_alt, max_alt, self.lut_resolution)
n_values = np.array(
[self.n_at_altitude(float(h), wavelength) for h in altitudes],
dtype=np.float32,
)
dn_dh_values = np.array(
[self.dn_dh_at_altitude(float(h), wavelength) for h in altitudes],
dtype=np.float32,
)
alpha_values = np.array(
[self.alpha_at_altitude(float(h), wavelength) for h in altitudes],
dtype=np.float32,
)
self._lut_altitudes = altitudes.astype(np.float32)
self._lut_n = n_values
self._lut_dn_dh = dn_dh_values
self._lut_alpha = alpha_values
self._lut_delta_h = float(altitudes[1] - altitudes[0])
self._lut_initialized = True
@property
def gpu_material_id(self) -> int:
"""Return GPU material ID for kernel dispatch."""
from ...propagation.propagator_protocol import GPUMaterialID
return GPUMaterialID.SIMPLE_INHOMOGENEOUS
[docs]
def get_gpu_kernels(self) -> dict:
"""Return GPU kernels for propagation."""
from ...propagation.kernels.propagation import (
_kernel_simple_inhomogeneous_euler,
_kernel_simple_inhomogeneous_rk4,
)
return {
"euler": _kernel_simple_inhomogeneous_euler,
"rk4": _kernel_simple_inhomogeneous_rk4,
}
[docs]
def get_gpu_parameters(self) -> tuple:
"""Return scalar parameters for GPU kernel."""
self._initialize_lut()
return (
float(self.center[0]),
float(self.center[1]),
float(self.center[2]),
float(self.reference_radius),
float(self.altitude_range[0]),
float(self._lut_delta_h),
int(self.lut_resolution),
)
[docs]
def get_gpu_arrays(self) -> dict:
"""Return device arrays for GPU kernel."""
self._initialize_lut()
return {
"lut_n": self._lut_n,
"lut_dn_dh": self._lut_dn_dh,
"lut_alpha": self._lut_alpha,
}