# 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.
"""
Grid Inhomogeneous Model (3D Grid Data)
Base class for 3D inhomogeneous materials defined on a regular grid.
User provides a 3D array of refractive index values.
GPU support is automatic via trilinear interpolation.
"""
from __future__ import annotations
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
[docs]
class GridInhomogeneousModel(MaterialField):
"""
Base class for 3D inhomogeneous materials defined on a grid.
User provides a 3D array of refractive index values. The system provides:
- Trilinear interpolation for arbitrary positions
- Automatic gradient computation via finite differences
- GPU support via 3D array interpolation
Parameters
----------
name : str
Name of the material model
n_grid : ndarray (nx, ny, nz)
Refractive index values on 3D grid
bounds : tuple
((x_min, x_max), (y_min, y_max), (z_min, z_max))
gradient_grid : ndarray (nx, ny, nz, 3), optional
Pre-computed gradients. If None, computed via finite differences.
alpha_grid : ndarray (nx, ny, nz), optional
Absorption coefficient values on 3D grid. If None, defaults to zeros.
Example
-------
>>> n_grid = np.ones((100, 100, 50)) * 1.000293
>>> n_grid += generate_turbulence(100, 100, 50) # Add perturbations
>>> material = GridInhomogeneousModel(
... name="Turbulent Field",
... n_grid=n_grid,
... bounds=((-5000, 5000), (-5000, 5000), (0, 25000))
... )
"""
# =========================================================================
# COMPATIBILITY DECLARATIONS
# =========================================================================
@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.GRID_EULER,
PropagationKernelID.GRID_RK4,
]
cls._default_kernel = PropagationKernelID.GRID_RK4
cls._supported_propagators = [
PropagatorID.GPU_GRADIENT,
PropagatorID.CPU_GRADIENT,
]
cls._default_propagator = PropagatorID.GPU_GRADIENT
[docs]
def __init__(
self,
name: str = "Grid Inhomogeneous",
n_grid: NDArray[np.float32] | None = None,
bounds: tuple[tuple[float, float], ...] | None = None,
gradient_grid: NDArray[np.float32] | None = None,
alpha_grid: NDArray[np.float32] | None = None,
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
if n_grid is None:
raise ValueError("n_grid is required")
if bounds is None:
raise ValueError("bounds is required")
if len(bounds) != 3:
raise ValueError("bounds must have 3 elements: ((x_min, x_max), ...)")
self._n_grid = n_grid.astype(np.float32)
self._bounds = bounds
self._shape = n_grid.shape
# Compute grid spacing
(x_min, x_max), (y_min, y_max), (z_min, z_max) = bounds
nx, ny, nz = self._shape
self._dx = (x_max - x_min) / (nx - 1) if nx > 1 else 1.0
self._dy = (y_max - y_min) / (ny - 1) if ny > 1 else 1.0
self._dz = (z_max - z_min) / (nz - 1) if nz > 1 else 1.0
self._origin = (x_min, y_min, z_min)
# Compute or store gradient grid
if gradient_grid is not None:
self._grad_grid = gradient_grid.astype(np.float32)
else:
self._grad_grid = self._compute_gradient_grid()
# Store or create zero absorption grid
if alpha_grid is not None:
self._alpha_grid = alpha_grid.astype(np.float32)
else:
self._alpha_grid = np.zeros(self._shape, dtype=np.float32)
def _compute_gradient_grid(self) -> NDArray[np.float32]:
"""Compute gradient via finite differences."""
grad_x = np.gradient(self._n_grid, self._dx, axis=0)
grad_y = np.gradient(self._n_grid, self._dy, axis=1)
grad_z = np.gradient(self._n_grid, self._dz, axis=2)
return np.stack([grad_x, grad_y, grad_z], axis=-1).astype(np.float32)
def _trilinear_interpolate(
self,
x: float | NDArray,
y: float | NDArray,
z: float | NDArray,
grid: NDArray,
) -> float | NDArray:
"""Trilinear interpolation in 3D grid."""
x = np.asarray(x)
y = np.asarray(y)
z = np.asarray(z)
# Normalize to grid coordinates
x_norm = (x - self._origin[0]) / self._dx
y_norm = (y - self._origin[1]) / self._dy
z_norm = (z - self._origin[2]) / self._dz
# Clamp to valid range
x_norm = np.clip(x_norm, 0, self._shape[0] - 1.001)
y_norm = np.clip(y_norm, 0, self._shape[1] - 1.001)
z_norm = np.clip(z_norm, 0, self._shape[2] - 1.001)
# Get integer indices and fractions
x0 = np.floor(x_norm).astype(int)
y0 = np.floor(y_norm).astype(int)
z0 = np.floor(z_norm).astype(int)
x1 = np.minimum(x0 + 1, self._shape[0] - 1)
y1 = np.minimum(y0 + 1, self._shape[1] - 1)
z1 = np.minimum(z0 + 1, self._shape[2] - 1)
xd = x_norm - x0
yd = y_norm - y0
zd = z_norm - z0
# Trilinear interpolation
if grid.ndim == 3:
# Scalar field
c000 = grid[x0, y0, z0]
c001 = grid[x0, y0, z1]
c010 = grid[x0, y1, z0]
c011 = grid[x0, y1, z1]
c100 = grid[x1, y0, z0]
c101 = grid[x1, y0, z1]
c110 = grid[x1, y1, z0]
c111 = grid[x1, y1, z1]
c00 = c000 * (1 - xd) + c100 * xd
c01 = c001 * (1 - xd) + c101 * xd
c10 = c010 * (1 - xd) + c110 * xd
c11 = c011 * (1 - xd) + c111 * xd
c0 = c00 * (1 - yd) + c10 * yd
c1 = c01 * (1 - yd) + c11 * yd
return c0 * (1 - zd) + c1 * zd
else:
# Vector field (gradient)
result = []
for i in range(grid.shape[-1]):
result.append(self._trilinear_interpolate(x, y, z, grid[..., i]))
return tuple(result)
[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 via trilinear interpolation."""
return self._trilinear_interpolate(x, y, z, self._n_grid)
[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 via trilinear interpolation in gradient grid."""
return self._trilinear_interpolate(x, y, z, self._grad_grid)
[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
# =========================================================================
@property
def gpu_material_id(self) -> int:
"""Return GPU material ID for kernel dispatch."""
from ...propagation.propagator_protocol import GPUMaterialID
return GPUMaterialID.GRID_INHOMOGENEOUS
[docs]
def get_gpu_kernels(self) -> dict:
"""Return GPU kernels for propagation."""
from ...propagation.kernels.propagation import (
_kernel_grid_inhomogeneous_euler,
_kernel_grid_inhomogeneous_rk4,
)
return {
"euler": _kernel_grid_inhomogeneous_euler,
"rk4": _kernel_grid_inhomogeneous_rk4,
}
[docs]
def get_gpu_parameters(self) -> tuple:
"""Return scalar parameters for GPU kernel."""
(x_min, x_max), (y_min, y_max), (z_min, z_max) = self._bounds
nx, ny, nz = self._shape
return (
float(x_min),
float(y_min),
float(z_min),
float(self._dx),
float(self._dy),
float(self._dz),
int(nx),
int(ny),
int(nz),
)
[docs]
def get_gpu_arrays(self) -> dict:
"""Return device arrays for GPU kernel."""
return {
"n_grid": self._n_grid,
"grad_grid": self._grad_grid,
"alpha_grid": self._alpha_grid,
}