nu-waves 1.0.0__py3-none-any.whl

nu_waves/__init__.py

@@ -0,0 +1,4 @@
+from .version import __version__  # noqa: F401
+
+# Optionally expose top-level API symbols
+# from .propagation import probability, ...

nu_waves/backend.py

@@ -0,0 +1,15 @@
+def pick_backend(prefer: str | None = None):
+    if prefer == "torch":
+        try:
+            import torch as xp  # type: ignore
+            return xp, "torch"
+        except Exception:
+            pass
+    if prefer == "cupy":
+        try:
+            import cupy as xp  # type: ignore
+            return xp, "cupy"
+        except Exception:
+            pass
+    import numpy as xp
+    return xp, "numpy"

nu_waves/backends/__init__.py

@@ -0,0 +1,3 @@
+from .numpy_backend import make_numpy_backend
+from .torch_backend import make_torch_mps_backend  # new
+__all__ = ["make_numpy_backend", "make_torch_mps_backend"]

nu_waves/backends/interface.py

@@ -0,0 +1,21 @@
+# Minimal array-backend interface used by nu_waves numerics.
+
+from __future__ import annotations
+from dataclasses import dataclass
+
+@dataclass
+class ArrayBackend:
+    """
+    Thin wrapper around a numeric library (NumPy for now; Torch/JAX later).
+    """
+    xp: any                 # array namespace (np or torch)
+    linalg: any             # linalg namespace (np.linalg or torch.linalg)
+    rng: any                # random generator (np.random.Generator-like)
+    dtype_real: any         # default real dtype
+    dtype_complex: any      # default complex dtype
+
+    def to_device(self, x):  # no-op for NumPy; overridden in GPU backends
+        return x
+
+    def from_device(self, x):  # no-op for NumPy; overridden in GPU backends
+        return x

nu_waves/backends/numpy_backend.py

@@ -0,0 +1,104 @@
+import numpy as np
+
+class _NumpyXP:
+    def __init__(self):
+        # allow x[..., xp.newaxis] in user code
+        self.newaxis = np.newaxis
+
+    def __getattr__(self, name):
+        # fallback to numpy for anything we didn't wrap explicitly
+        if hasattr(np, name):
+            return getattr(np, name)
+        raise AttributeError(f"_NumpyXP has no attribute {name}")
+
+    # common helpers used by the codebase
+    def normal(self, loc=0.0, scale=1.0, size=None):
+        return np.random.normal(loc=loc, scale=scale, size=size)
+
+    def uniform(self, low=0.0, high=1.0, size=None):
+        return np.random.uniform(low=low, high=high, size=size)
+
+    def random(self, size=None):
+        return np.random.random(size=size)
+
+    def asarray(self, x, dtype=None):
+        return np.asarray(x, dtype=dtype)
+
+    def repeat_last(self, x, n):
+        return np.repeat(x[..., np.newaxis], n, axis=-1)
+
+    def isscalar(self, x):
+        return np.isscalar(x)
+
+    def stack(self, arrays, axis=0):
+        return np.stack(arrays, axis=axis)
+
+    def argmax(self, x, axis=None):
+        return np.argmax(x, axis=axis)
+
+    def angle(self, z):
+        return np.angle(z)
+
+    def eye(self, n, dtype=None):
+        return np.eye(n, dtype=dtype)
+
+    def abs(self, x):
+        return np.abs(x)
+
+    def conj(self, x):
+        return np.conjugate(x)
+
+    def conjugate(self, x):
+        return np.conjugate(x)
+
+    def exp(self, x):
+        return np.exp(x)
+
+    def einsum(self, subs, *ops):
+        return np.einsum(subs, *ops)
+
+
+class _NumpyLinalg:
+    def eigh(self, A):
+        # works for (N,N) or (...,N,N)
+        return np.linalg.eigh(A)
+
+    def matrix_exp(self, A):
+        """
+        Stable numpy implementation without SciPy: eig-based exp.
+        Supports (N,N) or (...,N,N).
+        """
+        A = np.asarray(A)
+        if A.ndim == 2:
+            w, V = np.linalg.eig(A)
+            Vinv = np.linalg.inv(V)
+            return (V * np.exp(w)[np.newaxis, :]) @ Vinv
+
+        # batched
+        *b, n, _ = A.shape
+        A2 = A.reshape((-1, n, n))
+        out = np.empty_like(A2, dtype=A.dtype)
+        for i in range(A2.shape[0]):
+            w, V = np.linalg.eig(A2[i])
+            Vinv = np.linalg.inv(V)
+            out[i] = (V * np.exp(w)[np.newaxis, :]) @ Vinv
+        return out.reshape((*b, n, n))
+
+
+def make_numpy_backend(seed=0):
+    class Backend: ...
+    backend = Backend()
+    backend.device = "cpu"
+    backend.dtype_real = np.float64
+    backend.dtype_complex = np.complex128
+    backend.xp = _NumpyXP()
+    backend.linalg = _NumpyLinalg()
+    backend.from_device = lambda x: np.asarray(x)
+
+    def to_device(x): return x
+    def from_device(x): return x
+
+    backend.to_device = to_device
+    backend.from_device = from_device
+
+    return backend

nu_waves/backends/torch_backend.py

@@ -0,0 +1,224 @@
+# torch_backend.py
+import torch
+import numpy as np
+
+def _map_dtype_torch(requested, dtype_real, dtype_complex):
+    if requested is None:
+        return dtype_real
+    # ---- integers (FIX) ----
+    if requested in (int, np.int64, torch.int64):
+        return torch.int64
+    if requested in (np.int32, torch.int32):
+        return torch.int32
+    # ---- floats ----
+    if requested in (float, np.float32, torch.float32):
+        return torch.float32
+    if requested in (np.float64, torch.float64):
+        return torch.float64
+    # ---- complex ----
+    if requested in (complex, np.complex64, torch.complex64):
+        return torch.complex64
+    if requested in (np.complex128, torch.complex128):
+        return torch.complex128
+    # fallback to backend defaults
+    return dtype_real
+
+
+class _TorchXP:
+    def __init__(self, device, dtype_real, dtype_complex):
+        self.device = device
+        self.dtype_real = dtype_real
+        self.dtype_complex = dtype_complex
+        # allow x[..., xp.newaxis] indexing
+        self.newaxis = None
+
+    def __getattr__(self, name):
+        # delegate missing ops to torch when possible
+        if hasattr(torch, name):
+            return getattr(torch, name)
+        raise AttributeError(f"_TorchXP has no attribute {name}")
+
+    def maximum(self, x, y):
+        # torch.maximum only works on tensors — make sure both are tensors
+        x_t = x if isinstance(x, torch.Tensor) else torch.as_tensor(x, device=self.device)
+        y_t = y if isinstance(y, torch.Tensor) else torch.as_tensor(y, device=self.device)
+        return torch.maximum(x_t, y_t)
+
+    def minimum(self, x, y):
+        x_t = x if isinstance(x, torch.Tensor) else torch.as_tensor(x, device=self.device)
+        y_t = y if isinstance(y, torch.Tensor) else torch.as_tensor(y, device=self.device)
+        return torch.minimum(x_t, y_t)
+
+    def normal(self, loc=0.0, scale=1.0, size=None):
+        """
+        Torch equivalent of np.random.normal, supporting tensor broadcasting
+        and explicit size expansion (unlike torch.normal).
+        """
+        # Convert numpy arrays to tensors if needed
+        if not isinstance(loc, torch.Tensor):
+            loc = torch.as_tensor(loc, device=self.device)
+        if not isinstance(scale, torch.Tensor):
+            scale = torch.as_tensor(scale, device=self.device)
+
+        # Case 1: both tensors
+        if isinstance(loc, torch.Tensor) and isinstance(scale, torch.Tensor):
+            if size is None or tuple(size) == tuple(loc.shape):
+                # Default: same shape as loc
+                return torch.normal(mean=loc, std=scale)
+            else:
+                # Manually expand via broadcasting to requested size
+                loc_b = loc.expand(size)
+                scale_b = scale.expand(size)
+                return torch.normal(mean=loc_b, std=scale_b)
+
+        # Case 2: at least one scalar → use size directly
+        mean_val = float(loc) if not isinstance(loc, torch.Tensor) else float(loc)
+        std_val  = float(scale) if not isinstance(scale, torch.Tensor) else float(scale)
+        return torch.normal(mean=mean_val, std=std_val, size=size, device=self.device)
+
+    def uniform(self, low=0.0, high=1.0, size=None):
+        if size is None:
+            return (high - low) * torch.rand((), device=self.device) + low
+        return (high - low) * torch.rand(size, device=self.device) + low
+
+    def random(self, size=None):
+        return torch.rand(size, device=self.device)
+
+    def broadcast_arrays(self, *xs):
+        # torch.broadcast_tensors requires tensors, not python scalars
+        ts = [x if isinstance(x, torch.Tensor) else torch.as_tensor(x, device=self.device) for x in xs]
+        return torch.broadcast_tensors(*ts)
+
+    def asarray(self, x, dtype=None):
+        dt = _map_dtype_torch(dtype, self.dtype_real, self.dtype_complex)
+        return torch.as_tensor(x, dtype=dt, device=self.device)
+
+    def repeat_last(self, x, n):
+        # x[..., None].repeat(..., n) would also work
+        return x.unsqueeze(-1).repeat_interleave(n, dim=-1)
+
+    def isscalar(self, x):
+        return not isinstance(x, torch.Tensor) or x.ndim == 0
+
+    def stack(self, arrays, axis=0):
+        return torch.stack(arrays, dim=axis)
+
+    def argmax(self, x, axis=None):
+        return torch.argmax(x, dim=axis)
+
+    def angle(self, z):
+        return torch.atan2(torch.imag(z), torch.real(z))
+
+    def eye(self, n, dtype=None):
+        dt = _map_dtype_torch(dtype, self.dtype_real, self.dtype_complex)
+        return torch.eye(n, dtype=dt, device=self.device)
+
+    def abs(self, x):
+        return torch.abs(x)
+
+    def conj(self, x):
+        return torch.conj(x)
+
+    def conjugate(self, x):
+        return torch.conj(x)
+
+    def exp(self, x):
+        return torch.exp(x)
+
+    def einsum(self, subs, *ops):
+        return torch.einsum(subs, *ops)
+
+    def zeros(self, shape, dtype=None):
+        dt = _map_dtype_torch(dtype, self.dtype_real, self.dtype_complex)
+        return torch.zeros(shape, dtype=dt, device=self.device)
+
+    def ones(self, shape, dtype=None):
+        dt = _map_dtype_torch(dtype, self.dtype_real, self.dtype_complex)
+        return torch.ones(shape, dtype=dt, device=self.device)
+
+    def full(self, shape, fill_value, dtype=None):
+        dt = _map_dtype_torch(dtype, self.dtype_real, self.dtype_complex)
+        return torch.full(shape, fill_value, dtype=dt, device=self.device)
+
+    def zeros_like(self, x, dtype=None):
+        dt = _map_dtype_torch(dtype, x.dtype, self.dtype_complex if x.is_complex() else self.dtype_real)
+        return torch.zeros_like(x, dtype=dt, device=x.device)
+
+    def ones_like(self, x, dtype=None):
+        dt = _map_dtype_torch(dtype, x.dtype, self.dtype_complex if x.is_complex() else self.dtype_real)
+        return torch.ones_like(x, dtype=dt, device=x.device)
+
+    def diag(self, v, diagonal=0):
+        # ensure v is on device first
+        vt = v if isinstance(v, torch.Tensor) else torch.as_tensor(v, device=self.device)
+        return torch.diag(vt, diagonal=diagonal)
+
+
+class _TorchLinalg:
+    def __init__(self, device, dtype_real, dtype_complex):
+        self.device = device
+        self.dtype_real = dtype_real
+        self.dtype_complex = dtype_complex
+
+    def eigh(self, A):
+        """
+        Deterministic & stable Hermitian eigendecomposition:
+        run on CPU in float64/complex128, then cast back.
+        Works for (N,N) or (...,N,N).
+        """
+        A_cpu = A.to("cpu")
+        if A_cpu.dtype.is_complex:
+            A_cpu = A_cpu.to(torch.complex128)
+        else:
+            A_cpu = A_cpu.to(torch.float64)
+
+        w_cpu, V_cpu = torch.linalg.eigh(A_cpu)  # batched OK
+        w = w_cpu.to(self.device, dtype=self.dtype_real)
+        V = V_cpu.to(self.device, dtype=self.dtype_complex)
+        return w, V
+
+    def matrix_exp(self, A):
+        """
+        Use CPU fallback in float64/complex128 for numerical stability
+        and MPS compatibility; cast result back to original device/dtype.
+        """
+        A_cpu = A.to("cpu")
+        if A_cpu.dtype.is_complex:
+            A_cpu = A_cpu.to(torch.complex128)
+        else:
+            A_cpu = A_cpu.to(torch.float64)
+        Y_cpu = torch.linalg.matrix_exp(A_cpu)
+        return Y_cpu.to(self.device, dtype=A.dtype)
+
+
+def make_torch_mps_backend(seed=0, use_complex64=True):
+    torch.manual_seed(seed)
+
+    if torch.backends.mps.is_available():
+        device = torch.device("mps")
+    elif torch.backends.cudnn.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    dtype_real    = torch.float32 if use_complex64 else torch.float64
+    dtype_complex = torch.complex64 if use_complex64 else torch.complex128
+
+    class Backend: ...
+    backend = Backend()
+    backend.device = device
+    backend.dtype_real = dtype_real
+    backend.dtype_complex = dtype_complex
+    backend.xp = _TorchXP(device, dtype_real, dtype_complex)
+    backend.linalg = _TorchLinalg(device, dtype_real, dtype_complex)
+
+    def to_device(x):
+        return torch.as_tensor(x, device=device)
+
+    def from_device(x):
+        # detach is safe for tensors; passthrough for numpy scalars
+        return x.detach().cpu().numpy() if isinstance(x, torch.Tensor) else np.asarray(x)
+
+    backend.to_device = to_device
+    backend.from_device = from_device
+    return backend

nu_waves/hamiltonian/base.py

@@ -0,0 +1,91 @@
+from nu_waves.backends import make_numpy_backend
+from nu_waves.utils.units import GEV_TO_EV
+from nu_waves.utils.units import VCOEFF_EV
+
+
+class Hamiltonian:
+    def __init__(self, mixing_matrix, m2_diag, backend=None):
+        """
+        U: (N,N) complex PMNS (ou 3+N)
+        m2_diag: (N,N) diag(m_i^2) [eV^2]
+        """
+        self.backend = backend or make_numpy_backend()
+        xp = self.backend.xp
+
+        self.U = xp.asarray(mixing_matrix, dtype=self.backend.dtype_complex)
+        m2 = xp.asarray(m2_diag, dtype=self.backend.dtype_real)
+        if m2.ndim == 2:
+            # accept diagonal matrix but store vector
+            m2 = xp.diag(m2)
+        self.m2_diag = m2.reshape(-1)
+        # self.U = mixing_matrix
+        # self.m2_diag = m2_diag
+
+    def vacuum(self, E_GeV, antineutrino: bool = False):
+        """
+        Return the flavor-basis Hamiltonian in vacuum.
+
+        Parameters
+        ----------
+        E_GeV : float or array
+            Neutrino energy in GeV.
+        antineutrino : bool, optional
+            If True, uses complex-conjugated mixing matrix (U*).
+        """
+        xp = self.backend.xp
+        E = xp.asarray(E_GeV, dtype=self.backend.dtype_real)
+        E = E.reshape(()) if E.ndim == 0 else E.reshape(-1)  # () or (nE,)
+
+        EV_PER_GEV = self.backend.xp.asarray(GEV_TO_EV, dtype=self.backend.dtype_real)
+        E_eV = E * EV_PER_GEV
+        # E_eV = E * GEV_TO_EV
+
+        U = xp.conjugate(self.U) if antineutrino else self.U
+        # (U * m2) @ U^†  (scale columns by m2)
+        # H0 = (U * self.m2_diag[xp.newaxis, :]) @ xp.conjugate(U).T  # (N,N)
+        # H0 = (U * self.m2_diag[xp.newaxis, :]) @ xp.swapaxes(xp.conj(U), -1, -2)
+        H0 = (U * self.m2_diag[self.backend.xp.newaxis, :]) @ self.backend.xp.swapaxes(self.backend.xp.conj(U), -1, -2)
+        H = H0 / (2.0 * E_eV[..., xp.newaxis, xp.newaxis])  # ()→(N,N) or (nE,N,N)
+        return H
+
+        # E_eV = np.asarray(E_GeV, dtype=float) * GEV_TO_EV
+        # U = np.conjugate(self.U) if antineutrino else self.U
+        # D = np.diag(self.m2_diag)
+        # H = U @ D @ U.conj().T
+        # H = H / (2.0 * E_eV[..., None, None])  # broadcast over E
+        # return H
+
+    def matter_constant(self,
+                        E_GeV,
+                        rho_gcm3: float,
+                        Ye: float = 0.5,
+                        antineutrino: bool = False):
+        """
+        H_matter(E) = U diag(m^2)/(2E) U^† + diag(Ve, 0, 0, ...),
+        with Ve = + 7.632e-14 * rho[g/cm^3] * Ye [eV] for neutrinos,
+                = - Ve for antineutrinos.
+        Works for scalar or vector E and arbitrary dimension >= 1.
+        """
+        xp = self.backend.xp
+
+        # Energies on backend dtype/device
+        E = xp.asarray(E_GeV, dtype=self.backend.dtype_real)
+        E = E.reshape(()) if E.ndim == 0 else E.reshape(-1)  # () or (nE,)
+        EV_PER_GEV = xp.asarray(GEV_TO_EV, dtype=self.backend.dtype_real)
+
+        # Vacuum term (respect anti-ν via U*)
+        U = xp.conj(self.U) if antineutrino else self.U
+        H0 = (U * self.m2_diag[xp.newaxis, :]) @ xp.swapaxes(xp.conj(U), -1, -2)  # (N,N)
+        H_vac = H0 / (2.0 * (E * EV_PER_GEV)[..., xp.newaxis, xp.newaxis])  # (..,N,N)
+
+        # Matter potential in flavor basis (complex dtype for uniformity)
+        N = self.U.shape[0]
+        V_f = xp.zeros((N, N), dtype=self.backend.dtype_complex)
+        Ve = xp.asarray(VCOEFF_EV * rho_gcm3 * Ye, dtype=self.backend.dtype_real)
+        sign = -1.0 if antineutrino else +1.0
+        V_f = V_f + 0  # no-op; safe to remove if you prefer
+        V_f[..., 0, 0] = sign * Ve  # only e-flavor gets the CC potential
+
+        # Broadcast across energies if vector E
+        return H_vac + (V_f if E.ndim == 0 else V_f[xp.newaxis, :, :])
+

nu_waves/matter/prem.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+from dataclasses import dataclass
+import numpy as np
+
+from .profile import MatterLayer, MatterProfile
+
+@dataclass
+class PREMModel:
+    """
+    Minimal PREM density model (Dziewonski & Anderson 1981) with two discretization modes:
+      - 'prem_layers': cut the chord at PREM radial boundaries and use the local density.
+      - 'hist_density': sample along the chord, bin by density (nbins), and merge contiguous bins.
+
+    Densities are in g/cm^3. Electron fraction Ye defaults: mantle≈0.495, core≈0.467.
+    """
+    R_earth_km: float = 6371.0
+    # PREM region boundaries in km (increasing radii)
+    # 0–1221.5 (inner core), 1221.5–3480 (outer core), 3480–... mantle shells, crustal layers
+    prem_boundaries_km: tuple = (0.0, 1221.5, 3480.0, 5701.0, 5771.0, 5971.0,
+                                 6151.0, 6346.6, 6356.0, 6368.0, 6371.0)
+    # Ye defaults
+    Ye_mantle: float = 0.495
+    Ye_core:   float = 0.467
+
+    # --- PREM density polynomials ρ(r) in g/cm^3; x = r/R_earth ---
+    def rho(self, r_km: np.ndarray) -> np.ndarray:
+        r = np.asarray(r_km, float)
+        x = r / self.R_earth_km
+        rho = np.empty_like(x)
+
+        # regions defined by r
+        b = self.prem_boundaries_km
+
+        # 0–1221.5 km (inner core)
+        m = (r >= b[0]) & (r < b[1])
+        rho[m] = 13.0885 - 8.8381 * x[m]**2
+
+        # 1221.5–3480 km (outer core)
+        m = (r >= b[1]) & (r < b[2])
+        rho[m] = 12.5815 - 1.2638 * x[m] - 3.6426 * x[m]**2 - 5.5281 * x[m]**3
+
+        # 3480–5701 km (lower mantle)
+        m = (r >= b[2]) & (r < b[3])
+        rho[m] = 7.9565 - 6.4761 * x[m] + 5.5283 * x[m]**2 - 3.0807 * x[m]**3
+
+        # 5701–5771 km (TZ 1)
+        m = (r >= b[3]) & (r < b[4])
+        rho[m] = 5.3197 - 1.4836 * x[m]
+
+        # 5771–5971 km (TZ 2)
+        m = (r >= b[4]) & (r < b[5])
+        rho[m] = 11.2494 - 8.0298 * x[m]
+
+        # 5971–6151 km (upper mantle low)
+        m = (r >= b[5]) & (r < b[6])
+        rho[m] = 7.1089 - 3.8045 * x[m]
+
+        # 6151–6346.6 km (upper mantle high)
+        m = (r >= b[6]) & (r < b[7])
+        rho[m] = 2.6910 + 0.6924 * x[m]
+
+        # 6346.6–6356 km (crust 1)
+        m = (r >= b[7]) & (r < b[8])
+        rho[m] = 2.900
+
+        # 6356–6368 km (crust 2)
+        m = (r >= b[8]) & (r < b[9])
+        rho[m] = 2.600
+
+        # 6368–6371 km (ocean)
+        m = (r >= b[9]) & (r <= b[10])
+        rho[m] = 1.020
+
+        return rho
+
+    def Ye(self, r_km: np.ndarray) -> np.ndarray:
+        """Simple two-zone Ye: core vs mantle/crust."""
+        r = np.asarray(r_km, float)
+        return np.where(r <= 3480.0, self.Ye_core, self.Ye_mantle)
+
+    # --- chord geometry helpers ---
+    def _chord_length_km(self, cosz: float) -> float:
+        cz = float(cosz)
+        return 0.0 if cz >= 0.0 else -2.0 * self.R_earth_km * cz
+
+    def _impact_parameter_km(self, cosz: float) -> float:
+        cz = float(cosz)
+        return self.R_earth_km * np.sqrt(max(0.0, 1.0 - cz*cz))
+
+    # atmosphere thickness
+    def _atm_path_km(self, cosz: float, h_km: float) -> float:
+        if h_km <= 0: return 0.0
+        R, Rp = self.R_earth_km, self.R_earth_km + h_km
+        b = self._impact_parameter_km(cosz)
+        if b >= Rp:  # ray misses the atmosphere
+            return 0.0
+        seg = lambda r: np.sqrt(max(r * r - b * b, 0.0))
+        # ONE segment from production altitude Rp down to the surface R, irrespective of sign(cosz)
+        return seg(Rp) - seg(R)
+
+    # --- layer builders ---
+    def profile_from_coszen(self, cosz: float,
+                            scheme: str = "prem_layers",
+                            n_bins: int = 800,
+                            nbins_density: int = 24,
+                            merge_tol: float = 0.0,
+                            h_atm_km: float = 15.0) -> MatterProfile:
+        cz = float(cosz)
+        Ltot = self._chord_length_km(cz)
+        b = self._impact_parameter_km(cz)
+
+        layers: list[MatterLayer] = []
+
+        # --- Earth chord (only if upgoing) ---
+        if Ltot > 0.0:
+            half = 0.5 * Ltot
+
+            def r_of_t(t):
+                return np.sqrt(b * b + t * t)
+
+            if scheme == "prem_layers":
+                # cut at PREM boundaries intersected by the chord
+                t_pts = [-half, +half]
+                for rb in self.prem_boundaries_km[1:-1]:
+                    if rb > b:
+                        dt = np.sqrt(rb * rb - b * b)
+                        t_pts += [-dt, +dt]
+                t_pts = np.array(sorted(tp for tp in t_pts if -half <= tp <= half))
+                for t0, t1 in zip(t_pts[:-1], t_pts[1:]):
+                    dL = float(t1 - t0)
+                    r_mid = r_of_t(0.5 * (t0 + t1))
+                    rho_mid = float(self.rho(r_mid))
+                    Ye_mid = float(self.Ye(r_mid))
+                    layers.append(MatterLayer(rho_mid, Ye_mid, dL, "absolute"))
+
+            elif scheme == "hist_density":
+                t_edges = np.linspace(-half, +half, n_bins + 1)
+                t_mid = 0.5 * (t_edges[:-1] + t_edges[1:])
+                dL = np.diff(t_edges)
+                r_mid = r_of_t(t_mid)
+                rho_mid = self.rho(r_mid)
+                Ye_mid = self.Ye(r_mid)
+
+                rmin, rmax = rho_mid.min(), rho_mid.max()
+                edges = np.linspace(rmin, rmax, nbins_density + 1)
+                idx = np.minimum(np.searchsorted(edges, rho_mid, side="right") - 1, nbins_density - 1)
+
+                accL = accR = accY = 0.0
+                prev = idx[0]
+
+                def flush():
+                    nonlocal accL, accR, accY
+                    if accL > 0.0:
+                        layers.append(MatterLayer(accR / accL, accY / accL, accL, "absolute"))
+                        accL = accR = accY = 0.0
+
+                for i in range(len(dL)):
+                    if idx[i] != prev:
+                        flush()
+                        prev = idx[i]
+                    accL += float(dL[i])
+                    accR += float(dL[i] * rho_mid[i])
+                    accY += float(dL[i] * Ye_mid[i])
+                flush()
+
+                if merge_tol > 0 and len(layers) > 1:
+                    merged = [layers[0]]
+                    for nxt in layers[1:]:
+                        cur = merged[-1]
+                        if abs(nxt.rho_gcm3 - cur.rho_gcm3) <= merge_tol:
+                            Lsum = cur.weight + nxt.weight
+                            merged[-1] = MatterLayer(
+                                (cur.rho_gcm3 * cur.weight + nxt.rho_gcm3 * nxt.weight) / Lsum,
+                                (cur.Ye * cur.weight + nxt.Ye * nxt.weight) / Lsum,
+                                Lsum, "absolute"
+                            )
+                        else:
+                            merged.append(nxt)
+                    layers = merged
+            else:
+                raise ValueError(f"Unknown scheme '{scheme}'")
+        # else: no Earth layers for cz>=0
+
+        # --- Atmosphere around the Earth chord (both signs of cosz) ---
+        L_atm = self._atm_path_km(cosz, h_atm_km)
+        if L_atm > 0.0:
+            atm = MatterLayer(0.0, self.Ye_mantle, L_atm, "absolute")
+            if Ltot > 0.0:  # upgoing: atmosphere is on the FAR side, before Earth
+                layers = [atm] + layers
+            else:  # downgoing: atmosphere is on the NEAR side, straight to detector
+                layers = [atm]
+
+        # Ensure non-empty profile to keep the engine happy
+        if not layers:
+            layers = [MatterLayer(0.0, self.Ye_mantle, 0.0, "absolute")]
+
+        return MatterProfile(layers)