augmentrum 0.0.1__py3-none-any.whl

augmentrum/augmentation/eddy_current.py

@@ -0,0 +1,278 @@
+####################################################################################################
+#                                      eddy_current.py                                             #
+####################################################################################################
+#                                                                                                  #
+# Authors: K. C. Igwe (kci2104@columbia.edu)                                                       #
+#          J. P. Merkofer (j.p.merkofer@tue.nl)                                                    #
+#                                                                                                  #
+# Created: 2026-02-07                                                                              #
+#                                                                                                  #
+# Purpose: Implements water-derived and synthetic eddy current phase distortions for MRS data.     #
+#                                                                                                  #
+####################################################################################################
+
+import numpy as np
+from typing import Optional, List
+from scipy.signal import butter, filtfilt
+from augmentrum.core.base_module import BaseModule
+from augmentrum.core.nifti_mrs_plus import Backend
+from augmentrum.core.nifti_mrs_plus import Backend, NIfTI_MRS_Plus
+from augmentrum.utils.tensor_ops import to_numpy, match_backend
+
+
+class EddyCurrent(BaseModule):
+    """
+    Add eddy current phase distortions to MRS data.
+
+    Supports two modes:
+    - 'synthetic': Generate synthetic eddy current phase trajectory
+    - 'water': Derive eddy current phase from water reference
+
+    Parameters
+    ----------
+    mode : str
+        Eddy current mode: 'synthetic' or 'water' (default: 'synthetic')
+
+    **Synthetic Mode Parameters:**
+    std_rad : float
+        Standard deviation of phase noise in radians (default: 0.6)
+    lp_cut_hz : float
+        Low-pass filter cutoff frequency in Hz (default: 25.0)
+    seed : int, optional
+        Random seed for reproducibility
+
+    **Common Parameters:**
+    strength : float
+        Strength multiplier for eddy current effect (default: 1.0)
+    remove_linear : bool
+        Remove linear trend from phase (default: True)
+
+    **Water Mode Parameters:**
+    Uses water reference from water_list parameter in __call__
+    lp_cut_hz : float
+        Low-pass filter cutoff for water-derived phase (default: 20.0)
+
+    Examples
+    --------
+    >>> # Synthetic eddy current
+    >>> ec = EddyCurrent(mode='synthetic', std_rad=0.8, lp_cut_hz=30.0)
+    >>> result_data, _ = ec(nifti_plus, None)
+    >>>
+    >>> # Water-derived eddy current
+    >>> ec = EddyCurrent(mode='water', lp_cut_hz=20.0, strength=1.0)
+    >>> result_data, result_water = ec(nifti_plus, water_plus)
+    """
+
+    SUPPORTED_BACKENDS = []  # Supports all backends via process_tensor
+
+    def __init__(self, mode: str = 'synthetic',
+                 std_rad: float = 0.6, lp_cut_hz: float = 25.0,
+                 strength: float = 1.0, remove_linear: bool = True, seed: Optional[int] = None):
+        """Initialize eddy current module."""
+        super().__init__(mode=mode, std_rad=std_rad, lp_cut_hz=lp_cut_hz,
+                        strength=strength, remove_linear=remove_linear, seed=seed)
+
+        self.mode = mode.lower()
+        self.std_rad = std_rad
+        self.lp_cut_hz = lp_cut_hz
+        self.strength = strength
+        self.remove_linear = remove_linear
+        self.seed = seed
+
+        if self.mode not in ['synthetic', 'water']:
+            raise ValueError(f"mode must be 'synthetic' or 'water', got '{mode}'")
+
+    def process_nifti_list(self, data_list: List, water_list: Optional[List] = None, **kwargs):
+        """
+        Add eddy current distortion to list of NIFTI_MRS objects.
+
+        Args:
+            data_list: List of NIFTI_MRS objects
+            water_list: Optional list of water reference NIFTI_MRS objects (required for 'water' mode)
+            **kwargs: Additional arguments
+
+        Returns:
+            Tuple of (processed_data_list, processed_water_list)
+        """
+        if self.mode == 'water' and water_list is None:
+            raise ValueError("Water reference required for 'water' mode eddy current")
+
+        processed_data = []
+
+        for idx, nifti in enumerate(data_list):
+            # Get FID data
+            fid = nifti[:]
+
+            # Get spectral width
+            sw_hz = 1.0 / nifti.dwelltime
+
+            # Get water reference if available
+            water_fid = None
+            if water_list is not None and idx < len(water_list):
+                water_fid = water_list[idx][:]
+
+            # Apply eddy current
+            fid_with_ec = self._add_eddy_current(fid, sw_hz, water_fid)
+
+            # Update NIFTI_MRS data
+            nifti[:] = fid_with_ec
+            processed_data.append(nifti)
+
+        return processed_data, water_list
+
+    def process_tensor(self, data_array, water_array=None, backend=None, **kwargs):
+        """
+        Add eddy current distortion to tensor/array data (**any backend**).
+
+        The phase trajectory is generated in NumPy (uses SciPy butter filter),
+        then applied as a complex phasor multiplication which stays in the
+        native backend — ``data * match_backend(phasor, data)``.
+
+        Args:
+            data_array: Input tensor of shape ``(batch, ..., n_points)``
+            water_array: Optional water reference tensor
+            backend: Backend enum (unused)
+            **kwargs: Must contain ``'sw_hz'`` (spectral width in Hz)
+
+        Returns:
+            Tuple of (processed_data, water_array)
+        """
+        sw_hz = kwargs.get('sw_hz')
+        if sw_hz is None:
+            raise ValueError("EddyCurrent.process_tensor requires 'sw_hz' in kwargs")
+
+        if self.mode == 'water' and water_array is None:
+            raise ValueError("Water reference required for 'water' mode eddy current")
+
+        original_shape = data_array.shape
+        N = original_shape[-1]
+        n_batch = int(np.prod(original_shape[:-1]))
+
+        # Build per-FID complex phasors in NumPy (scipy-backed, data-independent for synthetic)
+        phasors_np = np.ones((n_batch, N), dtype=np.complex128)
+
+        if self.mode == 'water':
+            water_np = to_numpy(water_array).reshape(-1, N)
+            for i in range(n_batch):
+                w1d = water_np[min(i, water_np.shape[0] - 1)]
+                phi = self._ec_phase_from_water(w1d, sw_hz)
+                phasors_np[i] = np.exp(1j * self.strength * phi)
+        else:  # synthetic — no data dependency
+            for i in range(n_batch):
+                phi = self._synth_ec_phase(N, sw_hz)
+                phasors_np[i] = np.exp(1j * self.strength * phi)
+
+        phasors_np = phasors_np.reshape(original_shape)
+
+        # Apply phasor: backend-native multiply (preserves gradients for data)
+        return data_array * match_backend(phasors_np, data_array), water_array
+
+    def _add_eddy_current(self, fid: np.ndarray, sw_hz: float, water_fid: Optional[np.ndarray] = None) -> np.ndarray:
+        """
+        Add eddy current to FID data.
+
+        Args:
+            fid: Input FID data
+            sw_hz: Spectral width in Hz
+            water_fid: Optional water reference FID (for 'water' mode)
+
+        Returns:
+            FID with eddy current
+        """
+        # Work with last dimension (FID points)
+        original_shape = fid.shape
+        N = original_shape[-1]
+
+        # Reshape to 2D for processing
+        fid_2d = fid.reshape(-1, N)
+        result = np.zeros_like(fid_2d)
+
+        # If water mode, also reshape water FID
+        if self.mode == 'water' and water_fid is not None:
+            water_2d = water_fid.reshape(-1, N)
+        else:
+            water_2d = None
+
+        # Process each FID
+        for i in range(fid_2d.shape[0]):
+            fid_1d = fid_2d[i]
+
+            # Generate phase distortion based on mode
+            if self.mode == 'water' and water_2d is not None:
+                # Extract phase from water reference
+                water_1d = water_2d[min(i, water_2d.shape[0] - 1)]  # Handle different sizes
+                phi_ec = self._ec_phase_from_water(water_1d, sw_hz)
+            else:
+                # Generate synthetic phase
+                phi_ec = self._synth_ec_phase(N, sw_hz)
+
+            # Apply phase distortion
+            fid_with_ec = fid_1d * np.exp(1j * self.strength * phi_ec)
+            result[i] = fid_with_ec
+
+        # Reshape back to original shape
+        return result.reshape(original_shape)
+
+    def _ec_phase_from_water(self, fid_water: np.ndarray, sw_hz: float) -> np.ndarray:
+        """
+        Extract eddy current phase trajectory from water reference.
+
+        Args:
+            fid_water: Water reference FID
+            sw_hz: Spectral width in Hz
+
+        Returns:
+            Phase values in radians
+        """
+        fid_water = np.asarray(fid_water)
+        N = fid_water.size
+        t = np.arange(N, dtype=float) / float(sw_hz)
+
+        # Unwrap phase
+        phi = np.unwrap(np.angle(fid_water))
+
+        # Apply low-pass filter
+        if self.lp_cut_hz is not None and self.lp_cut_hz > 0:
+            nyq = 0.5 * float(sw_hz)
+            Wn = min(max(self.lp_cut_hz / nyq, 1e-6), 0.999999)
+            b, a = butter(2, Wn, btype='low')
+            phi = filtfilt(b, a, phi)
+
+        # Remove linear trend if requested
+        if self.remove_linear:
+            A = np.c_[np.ones(N), t]
+            k0, k1 = np.linalg.lstsq(A, phi, rcond=None)[0]
+            phi = phi - (k0 + k1 * t)
+
+        return phi
+
+    def _synth_ec_phase(self, N: int, sw_hz: float) -> np.ndarray:
+        """
+        Synthesize eddy current phase trajectory.
+
+        Args:
+            N: Number of points
+            sw_hz: Spectral width in Hz
+
+        Returns:
+            Phase values in radians
+        """
+        rng = np.random.default_rng(self.seed)
+        t = np.arange(N, dtype=float) / float(sw_hz)
+
+        # Generate white noise
+        w = rng.normal(scale=self.std_rad, size=N)
+
+        # Apply low-pass filter
+        nyq = 0.5 * float(sw_hz)
+        Wn = min(max(self.lp_cut_hz / nyq, 1e-6), 0.999999)
+        b, a = butter(2, Wn, btype='low')
+        phi = filtfilt(b, a, w)
+
+        # Remove linear trend if requested
+        if self.remove_linear:
+            A = np.c_[np.ones(N), t]
+            k0, k1 = np.linalg.lstsq(A, phi, rcond=None)[0]
+            phi = phi - (k0 + k1 * t)
+
+        return phi

augmentrum/augmentation/gaussian_noise.py

@@ -0,0 +1,211 @@
+####################################################################################################
+#                                      gaussian_noise.py                                           #
+####################################################################################################
+#                                                                                                  #
+# Authors: K. C. Igwe (kci2104@columbia.edu)                                                       #
+#          J. P. Merkofer (j.p.merkofer@tue.nl)                                                    #
+#                                                                                                  #
+# Created: 2026-02-07                                                                              #
+#                                                                                                  #
+# Purpose: Implements uncorrelated complex Gaussian noise (AWGN) for MRS data augmentation.        #
+#                                                                                                  #
+####################################################################################################
+
+import numpy as np
+from typing import Optional, List
+
+from augmentrum.core.base_module import BaseModule
+from augmentrum.utils.tensor_ops import to_numpy, match_backend, _is_torch, _is_jax, _is_tf
+
+
+class GaussianNoise(BaseModule):
+    """
+    Add uncorrelated complex Gaussian noise (AWGN) to MRS data.
+
+    **Backend-agnostic** (zea pattern): ``process_tensor`` works transparently
+    with NumPy, PyTorch, JAX, or TensorFlow tensors.
+
+    Parameters
+    ----------
+    snr: float, optional
+        Signal-to-noise ratio (unitless, signal_power / noise_power)
+    snr_db : float, optional
+        Signal-to-noise ratio in dB (signal_power / noise_power)
+    sigma : float, optional
+        Standard deviation per dimension (real and imaginary)
+    sigma_frac : float, optional
+        Sigma as fraction of max|FID| (alternative to sigma)
+    seed : int, optional
+        Random seed for reproducibility
+
+    Notes
+    -----
+    Provide EITHER snr_db, sigma, OR sigma_frac (not multiple).
+    If snr_db is used, sigma is computed from mean(|fid|^2).
+
+    Examples
+    --------
+    >>> noise = GaussianNoise(snr=20.0)
+    >>> noise = GaussianNoise(snr_db=10.0)
+    >>> noise = GaussianNoise(sigma=0.01)
+    >>> noise = GaussianNoise(sigma_frac=0.02)
+    """
+
+    SUPPORTED_BACKENDS = []  # Supports all backends
+
+    def __init__(self,
+                 snr: Optional[float] = None,
+                 snr_db: Optional[float] = None,
+                 sigma: Optional[float] = None,
+                 sigma_frac: Optional[float] = None,
+                 seed: Optional[int] = None):
+        super().__init__(snr=snr, snr_db=snr_db, sigma=sigma, sigma_frac=sigma_frac, seed=seed)
+
+        self.snr = snr
+        self.snr_db = snr_db
+        self.sigma = sigma
+        self.sigma_frac = sigma_frac
+        self.seed = seed
+
+        # Validate parameters
+        params_provided = sum([snr is not None, snr_db is not None,
+                               sigma is not None, sigma_frac is not None])
+        if params_provided == 0:
+            raise ValueError("Must provide one of: snr, snr_db, sigma, or sigma_frac")
+        if params_provided > 1:
+            raise ValueError("Provide only ONE of: snr, snr_db, sigma, or sigma_frac")
+
+    def process_nifti_list(self, data_list: List, water_list: Optional[List] = None, **kwargs):
+        """Add Gaussian noise to list of NIFTI_MRS objects."""
+        processed_data = []
+        for nifti in data_list:
+            fid = nifti[:]
+            fid_noisy = self._add_noise_numpy(fid)
+            nifti[:] = fid_noisy
+            processed_data.append(nifti)
+        return processed_data, water_list
+
+    def process_tensor(self, data_array, water_array=None, backend=None, **kwargs):
+        """
+        Add Gaussian noise to tensor/array data (**any backend, natively**).
+
+        Scale statistics (sigma_frac / snr modes) are measured with a tiny
+        ``to_numpy`` scalar reduction.  Noise is then **generated in the
+        native framework's own RNG** — ``torch.randn`` for PyTorch,
+        ``jax.random.normal`` for JAX, ``tf.random.normal`` for TF/Keras,
+        or ``numpy`` as the universal fallback.  This keeps gradients intact
+        and avoids the old "numpy noise → match_backend" workaround that
+        broke under JIT tracing.
+
+        Args:
+            data_array: Input tensor of shape ``(batch, ..., n_points)``
+            water_array: Optional water reference tensor (unchanged)
+            backend: Backend enum (unused — ops dispatch automatically)
+
+        Returns:
+            Tuple of (noisy_data, water_array)
+        """
+        original_shape = data_array.shape
+        N = original_shape[-1]
+        n_batch = int(np.prod(original_shape[:-1]))
+
+        # ── Step 1: compute noise scale (small scalar reduction → numpy OK) ──
+        if self.sigma is not None:
+            # Fixed sigma — no data stats needed at all
+            scale_np = np.full((n_batch, 1), float(self.sigma), dtype=np.float64)
+        else:
+            flat_np = to_numpy(data_array).reshape(n_batch, N)
+            if self.sigma_frac is not None:
+                peak = np.max(np.abs(flat_np), axis=-1, keepdims=True)
+                peak = np.where(peak > 0, peak, 1.0)
+                scale_np = self.sigma_frac * peak
+            else:
+                sig_pow = np.mean(np.abs(flat_np) ** 2, axis=-1, keepdims=True) + 1e-16
+                if self.snr is not None:
+                    noise_pow = sig_pow / float(self.snr)
+                else:
+                    noise_pow = sig_pow / (10.0 ** (float(self.snr_db) / 10.0))
+                scale_np = np.sqrt(noise_pow / 2.0)  # per-channel std
+
+        # scale_np shape: (n_batch, 1) — broadcast over last N axis
+        # Reshape to broadcast correctly over all dims
+        scale_np = scale_np.reshape([-1] + [1] * (len(original_shape) - 1))
+
+        # ── Step 2: generate noise in the native framework ────────────────────
+        if _is_torch(data_array):
+            import torch
+            gen = torch.Generator(device=data_array.device)
+            if self.seed is not None:
+                gen.manual_seed(self.seed)
+            # Determine real dtype for noise
+            if data_array.is_complex():
+                rdtype = data_array.real.dtype
+            else:
+                rdtype = data_array.dtype
+            scale_t = torch.as_tensor(scale_np, device=data_array.device, dtype=rdtype)
+            r = torch.randn(*original_shape, generator=gen, device=data_array.device, dtype=rdtype)
+            i_ = torch.randn(*original_shape, generator=gen, device=data_array.device, dtype=rdtype)
+            noise = torch.complex(r, i_) * scale_t
+            return data_array + noise, water_array
+
+        elif _is_jax(data_array):
+            import jax
+            import jax.numpy as jnp
+            seed = self.seed if self.seed is not None else 0
+            key = jax.random.PRNGKey(seed)
+            k1, k2 = jax.random.split(key)
+            r = jax.random.normal(k1, shape=original_shape, dtype=jnp.float32)
+            i_ = jax.random.normal(k2, shape=original_shape, dtype=jnp.float32)
+            scale_j = jnp.array(scale_np, dtype=jnp.float32)
+            noise = (r + 1j * i_).astype(jnp.complex64) * scale_j
+            return data_array + noise, water_array
+
+        elif _is_tf(data_array):
+            import tensorflow as tf
+            if self.seed is not None:
+                tf.random.set_seed(self.seed)
+            r = tf.random.normal(original_shape, dtype=tf.float32)
+            i_ = tf.random.normal(original_shape, dtype=tf.float32)
+            scale_t = tf.constant(scale_np, dtype=tf.float32)
+            noise = tf.complex(r * scale_t, i_ * scale_t)
+            return data_array + noise, water_array
+
+        else:
+            # NumPy and any other framework (e.g. Keras-native): numpy noise + match_backend
+            rng = np.random.default_rng(self.seed)
+            r = rng.normal(0.0, 1.0, size=original_shape)
+            i_ = rng.normal(0.0, 1.0, size=original_shape)
+            noise = (r + 1j * i_) * scale_np
+            return data_array + match_backend(noise, data_array), water_array
+
+    # ------------------------------------------------------------------
+    # Pure-numpy path for process_nifti_list (input is always numpy)
+    # ------------------------------------------------------------------
+
+    def _add_noise_numpy(self, fid: np.ndarray) -> np.ndarray:
+        """Add Gaussian noise to numpy FID data."""
+        rng = np.random.default_rng(self.seed)
+
+        original_shape = fid.shape
+        fid_flat = fid.reshape(-1, original_shape[-1])
+        out_dtype = np.result_type(fid.dtype, np.complex64)
+        fid_flat = fid_flat.astype(out_dtype, copy=False)
+
+        if self.sigma is not None:
+            scale = float(self.sigma)
+        elif self.sigma_frac is not None:
+            peak = np.max(np.abs(fid_flat), axis=-1)
+            peak = np.where(peak > 0, peak, 1.0)
+            scale = (self.sigma_frac * peak)[:, None]
+        else:
+            sig_pow = np.mean(np.abs(fid_flat) ** 2, axis=-1) + 1e-16
+            if self.snr is not None:
+                noise_pow = sig_pow / float(self.snr)
+            else:
+                noise_pow = sig_pow / (10.0 ** (float(self.snr_db) / 10.0))
+            scale = np.sqrt(noise_pow / 2.0)[:, None]
+
+        noise = (rng.normal(0.0, 1.0, size=fid_flat.shape)
+                 + 1j * rng.normal(0.0, 1.0, size=fid_flat.shape)) * scale
+
+        return (fid_flat + noise).reshape(original_shape)

augmentrum/augmentation/line_broadening.py

@@ -0,0 +1,179 @@
+####################################################################################################
+#                                     line_broadening.py                                           #
+####################################################################################################
+#                                                                                                  #
+# Authors: K. C. Igwe (kci2104@columbia.edu)                                                       #
+#          J. P. Merkofer (j.p.merkofer@tue.nl)                                                    #
+#                                                                                                  #
+# Created: 2026-02-07                                                                              #
+#                                                                                                  #
+# Purpose: Implements Lorentzian, Gaussian, and Voigt line broadening for MRS data.                #
+#                                                                                                  #
+####################################################################################################
+
+import math
+import numpy as np
+from typing import Optional, List
+
+from augmentrum.core.base_module import BaseModule
+from augmentrum.utils.tensor_ops import match_backend
+
+
+class LineBroadening(BaseModule):
+    """
+    Apply line broadening to MRS data.
+
+    Supports Lorentzian (exponential), Gaussian, and Voigt (combined) broadening.
+
+    **Backend-agnostic** (zea pattern): all ``process_tensor`` operations use
+    ``keras.ops`` so they work transparently with NumPy arrays, PyTorch tensors,
+    JAX arrays, and TensorFlow tensors without losing gradients.
+
+    Parameters
+    ----------
+    lb_hz : float
+        Lorentzian broadening Full Width at Half Maximum (FWHM) in Hz
+    gb_hz : float
+        Gaussian broadening FWHM in Hz
+    mode : str
+        Broadening type: 'lorentzian', 'gaussian', or 'voigt' (default)
+
+    Examples
+    --------
+    >>> # Lorentzian broadening only
+    >>> broadening = LineBroadening(lb_hz=5.0, mode='lorentzian')
+    >>>
+    >>> # Gaussian broadening only
+    >>> broadening = LineBroadening(gb_hz=3.0, mode='gaussian')
+    >>>
+    >>> # Voigt (combined) broadening
+    >>> broadening = LineBroadening(lb_hz=5.0, gb_hz=3.0, mode='voigt')
+    """
+
+    SUPPORTED_BACKENDS = []  # Supports all backends
+
+    def __init__(self, lb_hz: float = 0.0, gb_hz: float = 0.0, mode: str = 'voigt'):
+        """Initialize line broadening module."""
+        super().__init__(lb_hz=lb_hz, gb_hz=gb_hz, mode=mode)
+
+        self.lb_hz = lb_hz
+        self.gb_hz = gb_hz
+        self.mode = mode.lower()
+
+        if self.mode not in ['lorentzian', 'gaussian', 'voigt']:
+            raise ValueError(f"mode must be 'lorentzian', 'gaussian', or 'voigt', got '{mode}'")
+
+    def process_nifti_list(self, data_list: List, water_list: Optional[List] = None, **kwargs):
+        """
+        Apply line broadening to list of NIFTI_MRS objects.
+
+        Args:
+            data_list: List of NIFTI_MRS objects
+            water_list: Optional list of water reference NIFTI_MRS objects
+            **kwargs: Additional arguments
+
+        Returns:
+            Tuple of (processed_data_list, processed_water_list)
+        """
+        processed_data = []
+
+        for nifti in data_list:
+            # Get spectral width from NIFTI_MRS
+            sw_hz = 1.0 / nifti.dwelltime
+
+            # Get FID data
+            fid = nifti[:]
+
+            # Apply broadening based on mode
+            if self.mode == 'lorentzian':
+                fid_broadened = self._apply_lorentzian(fid, sw_hz, self.lb_hz)
+            elif self.mode == 'gaussian':
+                fid_broadened = self._apply_gaussian(fid, sw_hz, self.gb_hz)
+            else:  # voigt
+                fid_broadened = self._apply_voigt(fid, sw_hz, self.lb_hz, self.gb_hz)
+
+            # Update NIFTI_MRS data
+            nifti[:] = fid_broadened
+            processed_data.append(nifti)
+
+        return processed_data, water_list
+
+    def process_tensor(self, data_array, water_array=None, backend=None, **kwargs):
+        """
+        Apply line broadening to tensor/array data (**any backend**).
+
+        Uses ``keras.ops`` — works identically with NumPy, PyTorch, JAX, or
+        TensorFlow tensors.  Gradients are preserved for differentiable
+        training pipelines.
+
+        Args:
+            data_array: Input tensor of shape ``(batch, ..., n_points)``
+            water_array: Optional water reference tensor (unchanged)
+            backend: Backend enum (unused — ops dispatch automatically)
+            **kwargs: Must contain ``'sw_hz'`` (spectral width in Hz)
+
+        Returns:
+            Tuple of (processed_data, water_array)
+        """
+        sw_hz = kwargs.get('sw_hz')
+        if sw_hz is None:
+            raise ValueError("LineBroadening.process_tensor requires 'sw_hz' in kwargs")
+
+        if self.mode == 'lorentzian':
+            result = self._apply_lorentzian(data_array, sw_hz, self.lb_hz)
+        elif self.mode == 'gaussian':
+            result = self._apply_gaussian(data_array, sw_hz, self.gb_hz)
+        else:  # voigt
+            result = self._apply_voigt(data_array, sw_hz, self.lb_hz, self.gb_hz)
+
+        return result, water_array
+
+    # ------------------------------------------------------------------
+    # Envelope helpers — fully backend-agnostic via keras.ops
+    #
+    # The time axis `t` is a plain numpy array (no gradients needed for
+    # coordinates).  The envelope is also numpy.  The final multiplication
+    # `fid * envelope` is handled by the framework's own __mul__, which
+    # auto-promotes the numpy envelope to the correct backend and keeps
+    # the fid tensor on its original device / graph.
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _make_time_envelope(fid, sw_hz, lb_hz, gb_hz):
+        """Build the broadening envelope (numpy) shaped for broadcasting."""
+        N = fid.shape[-1]
+        t = np.arange(N, dtype=np.float64) / float(sw_hz)
+
+        # Reshape t for broadcasting: (1, 1, ..., N)
+        t = t.reshape([1] * (len(fid.shape) - 1) + [N])
+
+        env = np.ones_like(t)
+        if lb_hz > 0:
+            env *= np.exp(-math.pi * lb_hz * t)
+        if gb_hz > 0:
+            env *= np.exp(-((math.pi * gb_hz * t) ** 2) / (4 * math.log(2)))
+        return env
+
+    @staticmethod
+    def _apply_lorentzian(fid, sw_hz, lb_hz):
+        """Apply Lorentzian (exponential) line broadening."""
+        if lb_hz <= 0:
+            return fid
+        envelope = LineBroadening._make_time_envelope(fid, sw_hz, lb_hz, 0.0)
+        return fid * match_backend(envelope, fid)
+
+    @staticmethod
+    def _apply_gaussian(fid, sw_hz, gb_hz):
+        """Apply Gaussian line broadening."""
+        if gb_hz <= 0:
+            return fid
+        envelope = LineBroadening._make_time_envelope(fid, sw_hz, 0.0, gb_hz)
+        return fid * match_backend(envelope, fid)
+
+    @staticmethod
+    def _apply_voigt(fid, sw_hz, lb_hz, gb_hz):
+        """Apply Voigt (Lorentzian × Gaussian) line broadening."""
+        if lb_hz <= 0 and gb_hz <= 0:
+            return fid
+        envelope = LineBroadening._make_time_envelope(fid, sw_hz, lb_hz, gb_hz)
+        return fid * match_backend(envelope, fid)