lisaanalysistools 1.1.20__cp39-cp39-macosx_15_0_arm64.whl

lisatools/sensitivity.py

@@ -0,0 +1,1554 @@
+from __future__ import annotations
+import warnings
+from abc import ABC
+from typing import Any, Tuple, Optional, List
+from copy import deepcopy
+import os
+
+import math
+import numpy as np
+from scipy import interpolate
+import matplotlib.pyplot as plt
+
+try:
+    import cupy as cp
+
+except (ModuleNotFoundError, ImportError):
+    import numpy as cp
+
+from . import detector as lisa_models
+from .utils.utility import AET, get_array_module
+from .utils.constants import *
+from .stochastic import (
+    StochasticContribution,
+    FittedHyperbolicTangentGalacticForeground,
+    check_stochastic,
+)
+
+"""
+The sensitivity code is heavily based on an original code by Stas Babak, Antoine Petiteau for the LDC team.
+"""
+
+
+class Sensitivity(ABC):
+    """Base Class for PSD information.
+
+    The initialization function is only needed if using a file input.
+
+    """
+
+    channel: str = None
+
+    @staticmethod
+    def get_xp(array: np.ndarray) -> object:
+        """Numpy or Cupy (or float)"""
+        try:
+            return get_array_module(array)
+        except ValueError:
+            if isinstance(array, float):
+                return np
+            raise ValueError(
+                "array must be a numpy or cupy array (it can be a float as well)."
+            )
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        """Transform from the base sensitivity functions to the TDI PSDs.
+
+        Args:
+            f: Frequency array.
+            noise_levels: Current noise levels at frequency ``f``.
+            **kwargs: For interoperability.
+
+        Returns:
+            Transformed TDI PSD values.
+
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_Sn(
+        cls,
+        f: float | np.ndarray,
+        model: Optional[lisa_models.LISAModel | str] = lisa_models.sangria,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        """Calculate the PSD
+
+        Args:
+            f: Frequency array.
+            model: Noise model. Object of type :class:`lisa_models.LISAModel`.
+                It can also be a string corresponding to one of the stock models.
+                The model object must include attributes for ``Soms_d`` (shot noise)
+                and ``Sa_a`` (acceleration noise) or a spline as attribute ``Sn_spl``.
+                In the case of a spline, this must be a dictionary with
+                channel names as keys and callable PSD splines. For example,
+                if using ``scipy.interpolate.CubicSpline``, an input option
+                can be:
+
+                ```
+                noise_model.Sn_spl = {
+                    "A": CubicSpline(f, Sn_A)),
+                    "E": CubicSpline(f, Sn_E)),
+                    "T": CubicSpline(f, Sn_T))
+                }
+                ```
+            **kwargs: For interoperability.
+
+        Returns:
+            PSD values.
+
+        """
+        # spline or stock computation
+        if hasattr(model, "Sn_spl") and model.Sn_spl is not None:
+            spl = model.Sn_spl
+            if cls.channel not in spl:
+                raise ValueError("Calling a channel that is not available.")
+
+            Sout = spl[cls.channel](f)
+
+        else:
+            model = lisa_models.check_lisa_model(model)
+            # assert hasattr(model, "Soms_d") and hasattr(model, "Sa_a")
+
+            # get noise values
+            noise_levels = model.lisanoises(f)
+
+            # transform as desired for TDI combination
+            Sout = cls.transform(f, noise_levels, **kwargs)
+
+        # will add zero if ignored
+        stochastic_contribution = cls.stochastic_transform(
+            f, cls.get_stochastic_contribution(f, **kwargs), **kwargs
+        )
+
+        try:
+            Sout += stochastic_contribution
+        except:
+            breakpoint()
+        return Sout
+
+    @classmethod
+    def get_stochastic_contribution(
+        cls,
+        f: float | np.ndarray,
+        stochastic_params: Optional[tuple] = (),
+        stochastic_kwargs: Optional[dict] = {},
+        stochastic_function: Optional[StochasticContribution | str] = None,
+    ) -> float | np.ndarray:
+        """Calculate contribution from stochastic signal.
+
+        This function directs and wraps the calculation of and returns
+        the stochastic signal. The ``stochastic_function`` calculates the
+        sensitivity contribution. The ``transform_factor`` can transform that
+        output to the correct TDI contribution.
+
+        Args:
+            f: Frequency array.
+            stochastic_params: Parameters (arguments) to feed to ``stochastic_function``.
+            stochastic_kwargs: Keyword arguments to feeed to ``stochastic_function``.
+            stochastic_function: Stochastic class or string name of stochastic class. Takes ``stochastic_args`` and ``stochastic_kwargs``.
+                If ``None``, it uses :class:`FittedHyperbolicTangentGalacticForeground`.
+
+        Returns:
+            Contribution from stochastic signal.
+
+
+        """
+        xp = cls.get_xp(f)
+        if isinstance(f, float):
+            f = xp.ndarray([f])
+            squeeze = True
+        else:
+            squeeze = False
+
+        sgal = xp.zeros_like(f)
+
+        if (
+            (stochastic_params != () and stochastic_params is not None)
+            or (stochastic_kwargs != {} and stochastic_kwargs is not None)
+            or stochastic_function is not None
+        ):
+            if stochastic_function is None:
+                stochastic_function = FittedHyperbolicTangentGalacticForeground
+
+            stochastic_function = check_stochastic(stochastic_function)
+
+            sgal[:] = stochastic_function.get_Sh(
+                f, *stochastic_params, **stochastic_kwargs
+            )
+
+        if squeeze:
+            sgal = sgal.squeeze()
+        return sgal
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        """Transform from the base stochastic functions to the TDI PSDs.
+
+        **Note**: If not implemented, the transform will return the input.
+
+        Args:
+            f: Frequency array.
+            Sh: Power spectral density in stochastic term.
+            **kwargs: For interoperability.
+
+        Returns:
+            Transformed TDI PSD values.
+
+        """
+        return Sh
+
+
+class X1TDISens(Sensitivity):
+    channel: str = "X"
+
+    @staticmethod
+    def Cxx(f: float | np.ndarray) -> float | np.ndarray:
+        """Common TDI transform factor.
+        
+        Args:
+            f: Frequencyies to evaluate.
+
+        Returns:
+            Cxx: Transform factor.
+        
+        """
+        x = 2 * np.pi * f * L_SI / C_SI
+        return 16.0 * np.sin(x) ** 2
+                                 
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the XYZ TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        assert noise_levels.units == "relative_frequency"
+        Cxx = X1TDISens.Cxx(f)
+
+        x = 2 * np.pi * f * L_SI / C_SI
+        # TODO: need to check these
+        isi_rfi_readout_transfer = Cxx
+        tmi_readout_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        tm_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        rfi_backlink_transfer = Cxx
+        tmi_backlink_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        t = 4.0 * x**2 * np.sin(x) ** 2
+        return Sh * t
+
+class Y1TDISens(X1TDISens):
+    channel: str = "Y"
+    __doc__ = X1TDISens.__doc__
+    pass
+
+class Z1TDISens(X1TDISens):
+    channel: str = "Z"
+    __doc__ = X1TDISens.__doc__
+    pass
+
+
+class XY1TDISens(Sensitivity):
+    channel: str = "XY"
+
+    @staticmethod
+    def Cxy(f: float | np.ndarray) -> float | np.ndarray:
+        """Common TDI transform factor for CSD.
+        
+        Args:
+            f: Frequencyies to evaluate.
+
+        Returns:
+            Cxy: Transform factor.
+        
+        """
+        x = 2 * np.pi * f * L_SI / C_SI
+        return -4.0 * np.sin(2 * x) * np.sin(x)
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the XYZ TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        assert noise_levels.units == "relative_frequency"
+        Cxy = XY1TDISens.Cxy(f)
+
+        isi_rfi_readout_transfer = Cxy
+        tmi_readout_transfer = 4 * Cxy
+        tm_transfer = 4 * Cxy
+        rfi_backlink_transfer = Cxy
+        tmi_backlink_transfer = 4 * Cxy
+
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        # TODO: check these functions
+        # GB = -0.5 of X
+        t = -0.5 * (4.0 * x**2 * np.sin(x) ** 2)
+        return Sh * t
+
+
+class ZX1TDISens(XY1TDISens):
+    channel: str = "ZX"
+    __doc__ = XY1TDISens.__doc__
+    pass
+
+
+class YZ1TDISens(XY1TDISens):
+    channel: str = "YZ"
+    __doc__ = XY1TDISens.__doc__
+    pass
+
+
+class X2TDISens(Sensitivity):
+    channel: str = "X"
+
+    @staticmethod
+    def Cxx(f: float | np.ndarray) -> float | np.ndarray:
+        """Common TDI transform factor.
+
+        `arXiv:2211.02539 <https://arxiv.org/pdf/2211.02539>`_. 
+        
+        Args:
+            f: Frequencyies to evaluate.
+
+        Returns:
+            Cxx: Transform factor.
+        
+        """
+        x = 2 * np.pi * f * L_SI / C_SI
+        return 16. * np.sin(x) ** 2 * np.sin(2 * x) ** 2  # np.abs(1. - np.exp(-2j * np.pi * f * L_SI / C_SI) ** 2) ** 2
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the XYZ TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        assert noise_levels.units == "relative_frequency"
+        Cxx = X2TDISens.Cxx(f)
+
+        x = 2 * np.pi * f * L_SI / C_SI
+
+        isi_rfi_readout_transfer = 4. * Cxx
+        tmi_readout_transfer = Cxx * (3 + np.cos(2 * x)) 
+        tm_transfer = 4 * Cxx * (3 + np.cos(2 * x)) 
+        rfi_backlink_transfer = 4 * Cxx
+        tmi_backlink_transfer = Cxx * (3 + np.cos(2 * x)) 
+ 
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        # TODO: check these functions for TDI2
+        t = 4.0 * x**2 * np.sin(x) ** 2
+        return Sh * t
+
+
+class Y2TDISens(X2TDISens):
+    channel: str = "Y"
+    __doc__ = X2TDISens.__doc__
+    pass
+
+
+class Z2TDISens(X2TDISens):
+    channel: str = "Z"
+    __doc__ = X2TDISens.__doc__
+    pass
+
+class XY2TDISens(Sensitivity):
+    """
+    Cross-spectral density (CSD) between X and Y channels for TDI2.
+
+    From Table II of Nam et al. (2023) for uncorrelated noises:
+    - Common factor: C_XY(ω) = -16 sin(ωL) sin³(2ωL)
+    - Acceleration contribution: 4 * C_XY * S_pm
+    - Optical path contribution (ISI/RFI): C_XY * S_op
+
+    Total CSD: C_XY * (4*S_pm + S_op)
+
+    Notes:
+        - By circular symmetry, YZ and ZX CSDs have identical transfer functions
+        - For equal armlengths, the CSD is real-valued
+        - This implements the uncorrelated noise case
+    """
+
+    channel: str = "XY"
+
+    @staticmethod
+    def Cxy(f: float | np.ndarray) -> float | np.ndarray:
+        """Common TDI transform factor for CSD.
+
+        `arXiv:2211.02539 <https://arxiv.org/pdf/2211.02539>`_. 
+        
+        Args:
+            f: Frequencyies to evaluate.
+
+        Returns:
+            Cxy: Transform factor.
+        
+        """
+        x = 2 * np.pi * f * L_SI / C_SI
+
+        return -16.0 * np.sin(x) * np.sin(2.0 * x) ** 3
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        """
+        Transform from base sensitivity functions (S_pm, S_op) to TDI2 XY CSD.
+
+        Args:
+            f: Frequency array [Hz].
+            noise_levels: Current noise levels at frequency ``f``.
+            **kwargs: For interoperability.
+
+        Returns:
+            Cross-spectral density between X and Y channels.
+
+        Mathematical form:
+            x = 2π(L/c)f  [dimensionless frequency]
+            C_XY = -16 sin(x) sin³(2x)
+            CSD_XY = C_XY * (4*S_pm + S_op)
+        """
+        assert noise_levels.units == "relative_frequency"
+        Cxy = XY2TDISens.Cxy(f)
+
+        isi_rfi_readout_transfer = Cxy
+        tmi_readout_transfer = Cxy
+        tm_transfer = 4 * Cxy
+        rfi_backlink_transfer = Cxy
+        tmi_backlink_transfer = Cxy
+ 
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        """
+        Transform stochastic background to TDI2 XY CSD.
+
+        Note: For now, using same transform as TDI1 (placeholder).
+        TODO: Verify correct stochastic transform for TDI2 CSDs.
+
+        Args:
+            f: Frequency array [Hz].
+            Sh: Stochastic signal PSD.
+            **kwargs: For interoperability.
+
+        Returns:
+            Stochastic contribution to CSD.
+        """
+        x = 2.0 * np.pi * lisaLT * f
+        # Placeholder - using TDI1 form scaled by -0.5
+        t = -0.5 * (4.0 * x**2 * np.sin(x) ** 2)
+        return Sh * t
+
+
+class YZ2TDISens(XY2TDISens):
+    """
+    Cross-spectral density (CSD) between Y and Z channels for TDI2.
+
+    By circular symmetry of the LISA constellation (for equal armlengths),
+    this has the same transfer function as XY2TDISens.
+    """
+
+    channel: str = "YZ"
+    __doc__ = XY2TDISens.__doc__
+    pass
+
+
+class ZX2TDISens(XY2TDISens):
+    """
+    Cross-spectral density (CSD) between Z and X channels for TDI2.
+
+    By circular symmetry of the LISA constellation (for equal armlengths),
+    this has the same transfer function as XY2TDISens.
+    """
+
+    channel: str = "ZX"
+    __doc__ = XY2TDISens.__doc__
+    pass
+
+class A1TDISens(X1TDISens, Sensitivity):
+    channel: str = "A"
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the A,E TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        # these are WRONG
+        if np.any(np.asarray([
+            noise_levels.rfi_backlink_noise,
+            noise_levels.tmi_backlink_noise,
+            noise_levels.rfi_oms_noise,
+            noise_levels.tmi_oms_noise
+        ]) != 0.0):
+            raise NotImplementedError("ExtendedLISAModel has not been implemented yet for A1/E1/T1.")
+
+        assert noise_levels.units == "relative_frequency"
+        Cxx = X1TDISens.Cxx(f)
+
+        x = 2 * np.pi * f * L_SI / C_SI
+
+        # these are WRONG
+        tmi_readout_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        rfi_backlink_transfer = Cxx
+        tmi_backlink_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        
+        # these are right and were changed accordingly
+        # Need to find a citation for these 1st gen stuff
+        # all that is needed for old model type
+        isi_rfi_readout_transfer = 1/2 * Cxx * (2.0 + np.cos(x))
+        tm_transfer = Cxx * (3.0 + 2.0 * np.cos(x) + np.cos(2 * x))
+
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        t = 4.0 * x**2 * np.sin(x) ** 2
+        return 1.5 * (Sh * t)
+
+
+class E1TDISens(A1TDISens):
+    channel: str = "E"
+    __doc__ = A1TDISens.__doc__
+    pass
+
+
+class T1TDISens(Sensitivity):
+    channel: str = "T"
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the T TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        assert noise_levels.units == "relative_frequency"
+        
+        Cxx = X1TDISens.Cxx(f)
+
+        x = 2 * np.pi * f * L_SI / C_SI
+
+        # these are WRONG
+        if np.any(np.asarray([
+            noise_levels.rfi_backlink_noise,
+            noise_levels.tmi_backlink_noise,
+            noise_levels.rfi_oms_noise,
+            noise_levels.tmi_oms_noise
+        ]) != 0.0):
+            raise NotImplementedError("ExtendedLISAModel has not been implemented yet for A1/E1/T1.")
+        tmi_readout_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        rfi_backlink_transfer = Cxx
+        tmi_backlink_transfer = Cxx * (2.0 * (1.0 + np.cos(x) ** 2))
+        
+        # these are right and were changed accordingly
+        # Need to find a citation for these 1st gen stuff
+        # all that is needed for old model type
+        isi_rfi_readout_transfer = Cxx * (1 - np.cos(x))
+        tm_transfer = 8.0 * Cxx * np.sin(x / 2.) ** 4
+
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+    
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        t = 4.0 * x**2 * np.sin(x) ** 2
+        return 0.0 * (Sh * t)
+
+
+
+class A2TDISens(X2TDISens, Sensitivity):
+    channel: str = "A"
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the XYZ TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        assert noise_levels.units == "relative_frequency"
+        Cxx = X2TDISens.Cxx(f)
+
+        x = 2 * np.pi * f * L_SI / C_SI
+        
+        isi_rfi_readout_transfer = 2. * Cxx * (2 * np.cos(x))
+        tmi_readout_transfer = Cxx * (3 + 2 * np.cos(x) + np.cos(2 * x)) 
+        tm_transfer = 4 * Cxx * (3 + 2 * np.cos(x) + np.cos(2 * x)) 
+        rfi_backlink_transfer = 2 * Cxx * (2 * np.cos(x))
+        tmi_backlink_transfer = Cxx * (3 + 2 * np.cos(x) + np.cos(2 * x)) 
+ 
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        # TODO: check these functions for TDI2
+        t = 4.0 * x**2 * np.sin(x) ** 2
+        return Sh * t
+
+
+class E2TDISens(A2TDISens):
+    channel: str = "E"
+    __doc__ = A2TDISens.__doc__
+    pass
+
+
+class T2TDISens(X2TDISens, Sensitivity):
+    channel: str = "T"
+
+    @staticmethod
+    def transform(
+        f: float | np.ndarray,
+        noise_levels: lisa_models.CurrentNoises,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base sensitivity functions to the XYZ TDI PSDs.\n\n"
+            + Sensitivity.transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+
+        assert noise_levels.units == "relative_frequency"
+        Cxx = X2TDISens.Cxx(f)
+
+        x = 2 * np.pi * f * L_SI / C_SI
+        
+        isi_rfi_readout_transfer = 4. * Cxx * (1 - np.cos(x))
+        tmi_readout_transfer = 8 * Cxx * np.sin(x / 2.) ** 4
+        tm_transfer = 32 * Cxx * np.sin(x / 2.) ** 4
+        rfi_backlink_transfer = 4. * Cxx * (1 - np.cos(x))
+        tmi_backlink_transfer = 8 * Cxx * np.sin(x / 2.) ** 4
+ 
+        isi_oms_ffd = isi_rfi_readout_transfer * noise_levels.isi_oms_noise
+        rfi_oms_ffd = isi_rfi_readout_transfer * noise_levels.rfi_oms_noise
+        tmi_oms_ffd = tmi_readout_transfer * noise_levels.tmi_oms_noise
+        tm_noise_ffd = tm_transfer * noise_levels.tm_noise
+        
+        rfi_backlink_ffd = rfi_backlink_transfer * noise_levels.rfi_backlink_noise
+        tmi_backlink_ffd = tmi_backlink_transfer * noise_levels.tmi_backlink_noise
+
+        total_noise = tm_noise_ffd + isi_oms_ffd + rfi_oms_ffd + tmi_oms_ffd + rfi_backlink_ffd + tmi_backlink_ffd
+        return total_noise
+
+    @staticmethod
+    def stochastic_transform(
+        f: float | np.ndarray, Sh: float | np.ndarray, **kwargs: dict
+    ) -> float | np.ndarray:
+        __doc__ = (
+            "Transform from the base stochastic functions to the XYZ stochastic TDI information.\n\n"
+            + Sensitivity.stochastic_transform.__doc__.split("PSDs.\n\n")[-1]
+        )
+        x = 2.0 * np.pi * lisaLT * f
+        # TODO: check these functions for TDI2
+        t = 4.0 * x**2 * np.sin(x) ** 2
+        return Sh * t
+
+
+class LISASens(Sensitivity):
+    @classmethod
+    def get_Sn(
+        cls,
+        f: float | np.ndarray,
+        model: Optional[lisa_models.LISAModel | str] = lisa_models.sangria,
+        average: bool = True,
+        **kwargs: dict,
+    ) -> float | np.ndarray:
+        """Compute the base LISA sensitivity function.
+
+        Args:
+            f: Frequency array.
+            model: Noise model. Object of type :class:`lisa_models.LISAModel`. It can also be a string corresponding to one of the stock models.
+            average: Whether to apply averaging factors to sensitivity curve.
+                Antenna response: ``av_resp = np.sqrt(5) if average else 1.0``
+                Projection effect: ``Proj = 2.0 / np.sqrt(3) if average else 1.0``
+            **kwargs: Keyword arguments to pass to :func:`get_stochastic_contribution`. # TODO: fix
+
+        Returns:
+            Sensitivity array.
+
+        """
+        model = lisa_models.check_lisa_model(model)
+        
+        if not isinstance(model, lisa_models.LISAModel):
+            raise NotImplementedError("This function has not been implemented for ExtendedLISAModel yet.")
+
+        # get noise values
+        noise_values = model.lisanoises(f, unit="displacement")
+
+        Sa_d = noise_values.tm_noise
+        Sop = noise_values.isi_oms_noise
+
+        all_m = np.sqrt(4.0 * Sa_d + Sop)
+        ## Average the antenna response
+        av_resp = np.sqrt(5) if average else 1.0
+
+        ## Projection effect
+        Proj = 2.0 / np.sqrt(3) if average else 1.0
+
+        ## Approximative transfer function
+        f0 = 1.0 / (2.0 * lisaLT)
+        a = 0.41
+        T = np.sqrt(1 + (f / (a * f0)) ** 2)
+        sens = (av_resp * Proj * T * all_m / lisaL) ** 2
+
+        # will add zero if ignored
+        sens += cls.get_stochastic_contribution(f, **kwargs)
+        return sens
+
+
+class CornishLISASens(LISASens):
+    """PSD from https://arxiv.org/pdf/1803.01944.pdf
+
+    Power Spectral Density for the LISA detector assuming it has been active for a year.
+    I found an analytic version in one of Niel Cornish's paper which he submitted to the arXiv in
+    2018. I evaluate the PSD at the frequency bins found in the signal FFT.
+
+    PSD obtained from: https://arxiv.org/pdf/1803.01944.pdf
+
+    """
+
+    @staticmethod
+    def get_Sn(
+        f: float | np.ndarray, average: bool = True, **kwargs: dict
+    ) -> float | np.ndarray:
+        # TODO: documentation here
+
+        sky_averaging_constant = 20.0 / 3.0 if average else 1.0
+
+        L = 2.5 * 10**9  # Length of LISA arm
+        f0 = 19.09 * 10 ** (-3)  # transfer frequency
+
+        # Optical Metrology Sensor
+        Poms = ((1.5e-11) * (1.5e-11)) * (1 + np.power((2e-3) / f, 4))
+
+        # Acceleration Noise
+        Pacc = (
+            (3e-15)
+            * (3e-15)
+            * (1 + (4e-4 / f) * (4e-4 / f))
+            * (1 + np.power(f / (8e-3), 4))
+        )
+
+        # constants for Galactic background after 1 year of observation
+        alpha = 0.171
+        beta = 292
+        k = 1020
+        gamma = 1680
+        f_k = 0.00215
+
+        # Galactic background contribution
+        Sc = (
+            9e-45
+            * np.power(f, -7 / 3)
+            * np.exp(-np.power(f, alpha) + beta * f * np.sin(k * f))
+            * (1 + np.tanh(gamma * (f_k - f)))
+        )
+
+        # PSD
+        PSD = (sky_averaging_constant) * (
+            (10 / (3 * L * L))
+            * (Poms + (4 * Pacc) / (np.power(2 * np.pi * f, 4)))
+            * (1 + 0.6 * (f / f0) * (f / f0))
+            + Sc
+        )
+
+        return PSD
+
+
+class FlatPSDFunction(LISASens):
+    """White Noise PSD function."""
+
+    @classmethod
+    def get_Sn(
+        cls, f: float | np.ndarray, val: float, **kwargs: dict
+    ) -> float | np.ndarray:
+        # TODO: documentation here
+        xp = cls.get_xp(f)
+        out = xp.full_like(f, val)
+        if isinstance(f, float):
+            out = out.item()
+        return out
+
+
+class SensitivityMatrix:
+    """Container to hold sensitivity information.
+
+    Args:
+        f: Frequency array.
+        sens_mat: Input sensitivity list. The shape of the nested lists should represent the shape of the
+            desired matrix. Each entry in the list must be an array, :class:`Sensitivity`-derived object,
+            or a string corresponding to the :class:`Sensitivity` object.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(
+        self,
+        f: np.ndarray,
+        sens_mat: (
+            List[List[np.ndarray | Sensitivity]]
+            | List[np.ndarray | Sensitivity]
+            | np.ndarray
+            | Sensitivity
+        ),
+        *sens_args: tuple,
+        sens_kwargs_mat = None,
+        **sens_kwargs: dict,
+    ) -> None:
+        self.frequency_arr = f
+        self.data_length = len(self.frequency_arr)
+        self.sens_args = sens_args
+        if sens_kwargs_mat is None:
+            self.sens_kwargs = sens_kwargs
+        else:
+            self.sens_kwargs = sens_kwargs_mat
+
+        self.sens_mat = sens_mat
+
+    @property
+    def frequency_arr(self) -> np.ndarray:
+        return self._frequency_arr
+
+    @frequency_arr.setter
+    def frequency_arr(self, frequency_arr: np.ndarray) -> None:
+        assert frequency_arr.dtype == np.float64 or frequency_arr.dtype == float
+        assert frequency_arr.ndim == 1
+        self._frequency_arr = frequency_arr
+
+    def check_update(self):
+        if not self.can_redo:
+            raise ValueError("Cannot update sensitivities because original input was arrays rather than functions.")
+
+    def update_frequency_arr(self, frequency_arr: np.ndarray) -> None:
+        """Update class with new frequency array.
+
+        Args:
+            frequency_arr: Frequency array.
+
+        """
+        self.check_update()
+        self.frequency_arr = frequency_arr
+        self.sens_mat = self.sens_mat_input
+
+    def update_model(self, model: lisa_models.LISAModel | list | np.ndarray) -> None:
+        """Update class with new sensitivity model.
+
+        Args:
+            model: Noise model. Object of type :class:`lisa_models.LISAModel`. It can also be a string corresponding to one of the stock models.
+
+        """
+        self.check_update()
+        for tmp_kwargs in self.sens_kwargs.flatten():
+            tmp_kwargs["model"] = model
+        self.sens_mat = self.sens_mat_input
+
+    def update_stochastic(self, **kwargs: dict) -> None:
+        """Update class with new stochastic function.
+
+        Args:
+            **kwargs: Keyword arguments update for :func:`lisatools.sensitivity.Sensitivity.get_stochastic_contribution`.
+                This operation will combine the new and old kwarg dictionaries, updating any
+                old information with any added corresponding new information. **Note**: any old information
+                that is not updated will remain in place.
+
+        """
+        self.check_update()
+        tmptmp = self.sens_kwargs.flatten()
+        for i, tmp_kwargs in tmptmp:
+            tmptmp[i] = {**tmp_kwargs, **kwargs}
+        self.sens_kwargs = tmptmp.reshape(self.sens_kwargs.shape)
+        self.sens_mat = self.sens_mat_input
+
+    @property
+    def sens_mat(self) -> np.ndarray:
+        """Get sensitivity matrix."""
+        return self._sens_mat
+
+    @sens_mat.setter
+    def sens_mat(
+        self,
+        sens_mat: (
+            List[List[np.ndarray | Sensitivity]]
+            | List[np.ndarray | Sensitivity]
+            | np.ndarray
+            | Sensitivity
+        ),
+    ) -> None:
+        """Set sensitivity matrix."""
+        
+        if (isinstance(sens_mat, np.ndarray) or isinstance(
+            sens_mat, cp.ndarray)
+        ) and sens_mat.dtype != object:
+            self._sens_mat = sens_mat
+            if not hasattr(self, "sens_mat_input"):
+                self.can_redo = False
+            else:
+                self.can_redo = True
+
+        elif isinstance(sens_mat, list) or (isinstance(sens_mat, np.ndarray) and sens_mat.dtype == object):
+            self.sens_mat_input = deepcopy(sens_mat)
+            _run = True
+            _layer = self.sens_mat_input
+            outer_shape = [len(_layer)]
+            while _run:
+                _test_length = None
+                _type_1 = None
+                for tmp in _layer:
+                    # check each entry is the same type
+                    if _type_1 is None:
+                        _type_1 = type(tmp)
+                    else:
+                        if _type_1 != type(tmp):
+                            raise ValueError("List inputs must be all of the same type.")
+                        
+                    if isinstance(tmp, list):
+                        if _test_length is None:
+                            _test_length = len(tmp)
+                        else:
+                            if len(tmp) != _test_length:
+                                raise ValueError("Input list structure is not Rectangular.")
+                    elif isinstance(tmp, np.ndarray) or isinstance(tmp, cp.ndarray):
+                        if tmp.ndim > 1:
+                            raise ValueError("If entering a list of arrays, arrays must be 1D on the last dimension of the list structure.")
+                        if _test_length is None:
+                            _test_length = len(tmp)
+                        else:
+                            if len(tmp) != _test_length:
+                                raise ValueError("Input list/array structure is not Rectangular.")
+
+                if isinstance(_layer[0], list):
+                    outer_shape.append(len(_layer[0]))
+                    _layer = _layer[0]
+                    continue
+                        
+                elif isinstance(_layer[0], np.ndarray) or isinstance(_layer[0], cp.ndarray):
+                    # hit the array, must be last layer
+                    _run = False
+                    self.can_redo = False
+                    self.is_array_base = True
+                    continue
+
+                # TODO: better way to do this?
+                elif hasattr(_layer[0], "get_Sn"):
+                    _run = False
+                    self.can_redo = True
+                    self.is_array_base = False
+                    continue
+
+                elif isinstance(_layer[0], str):
+                    _run = False
+                    self.can_redo = True
+                    self.is_array_base = False
+                    sensitivity = check_sensitivity(_layer[0])
+                    assert hasattr(sensitivity, "get_Sn")
+                    continue
+
+                else:
+                    raise ValueError("Matrix element must be Sensitivity object, string representing a sensitivity object, or an array with values.")
+                
+        
+            if isinstance(self.sens_kwargs, np.ndarray) or isinstance(self.sens_kwargs, list):
+                tmp_kwargs = np.asarray(self.sens_kwargs, dtype=object)
+                assert tmp_kwargs.shape == tuple(outer_shape)
+
+            elif isinstance(self.sens_kwargs, dict):
+                tmp_kwargs = np.full(outer_shape, self.sens_kwargs, dtype=object)
+            else:
+                raise ValueError("sens_kwargs Must be numpy object array, list, or dict.")
+            
+            # TODO: sens_kwargs property setup
+            self.sens_kwargs = tmp_kwargs
+            
+            num_components = np.prod(outer_shape).item()
+            xp = get_array_module(self.frequency_arr)
+            if self.is_array_base:
+                _sens_mat = xp.asarray(sens_mat)
+            
+            else:
+                _flattened_arr = np.asarray(sens_mat, dtype=object).flatten()
+                _sens_mat = xp.zeros((num_components, len(self.frequency_arr)))
+                for i, matrix_member in enumerate(_flattened_arr):
+                    # calculate it
+                    if hasattr(matrix_member, "get_Sn") or isinstance(matrix_member, str):
+                        _sens_mat[i, :] = get_sensitivity(
+                            self.frequency_arr,
+                            *self.sens_args,
+                            sens_fn=matrix_member,
+                            **self.sens_kwargs.flatten()[i],
+                        )
+
+                    else:
+                        raise ValueError
+
+            # setup in array form
+            self._sens_mat = _sens_mat.reshape(tuple(outer_shape) + (len(self.frequency_arr),))
+            
+        else:
+            raise ValueError("Must input array or list.")
+        
+        self._setup_det_and_inv()
+
+    def _setup_det_and_inv(self):
+        # setup detC
+        """Determinant of TDI matrix."""
+        if self.sens_mat.ndim < 3:
+            self.detC = self.sens_mat
+            self.invC = 1/self.sens_mat
+
+        else:
+            xp = get_array_module(self.sens_mat)
+            self.detC = xp.linalg.det(self.sens_mat.transpose(2, 0, 1))
+            invC = xp.zeros_like(self.sens_mat.transpose(2, 0, 1))
+            if xp.all(self.detC == 0.0):
+                raise ValueError("All determinants are zero.")
+            
+            invC[self.detC != 0.0] = xp.linalg.inv(self.sens_mat.transpose(2, 0, 1)[self.detC != 0.0])
+            invC[self.detC == 0.0] = 1e-100
+            self.invC = invC.transpose(1, 2, 0)
+            
+        xp = get_array_module(self.sens_mat)
+
+        # setup detC
+        """Determinant and inverse of TDI matrix."""
+        if self.sens_mat.ndim < 3:
+            self.detC = xp.prod(self.sens_mat, axis=0)
+            self.invC = 1 / self.sens_mat
+
+        else:
+            self.detC = xp.linalg.det(self.sens_mat.transpose(2, 0, 1))
+            invC = xp.zeros_like(self.sens_mat.transpose(2, 0, 1))
+            invC[self.detC != 0.0] = xp.linalg.inv(
+                self.sens_mat.transpose(2, 0, 1)[self.detC != 0.0]
+            )
+            invC[self.detC == 0.0] = 1e-100
+            self.invC = invC.transpose(1, 2, 0)
+
+    def __getitem__(self, index: Any) -> np.ndarray:
+        """Indexing the class indexes the array."""
+        return self.sens_mat[index]
+
+    def __setitem__(self, index: Any, value: np.ndarray) -> np.ndarray:
+        """Indexing the class indexes the array."""
+        self.sens_mat[index] = value
+        self._setup_det_and_inv()
+
+    @property
+    def ndim(self) -> int:
+        """Dimensionality of sens mat array."""
+        return self.sens_mat.ndim
+
+    def flatten(self) -> np.ndarray:
+        """Flatten sens mat array."""
+        return self.sens_mat.reshape(-1, self.sens_mat.shape[-1])
+
+    @property
+    def shape(self) -> tuple:
+        """Shape of sens mat array."""
+        return self.sens_mat.shape
+
+    def loglog(
+        self,
+        ax: Optional[plt.Axes] = None,
+        fig: Optional[plt.Figure] = None,
+        inds: Optional[int | tuple] = None,
+        char_strain: Optional[bool] = False,
+        **kwargs: dict,
+    ) -> Tuple[plt.Figure, plt.Axes]:
+        """Produce a log-log plot of the sensitivity.
+
+        Args:
+            ax: Matplotlib Axes objects to add plots. Either a list of Axes objects or a single Axes object.
+            fig: Matplotlib figure object.
+            inds: Integer index to select out which data to add to a single access.
+                A list can be provided if ax is a list. They must be the same length.
+            char_strain: If ``True``, plot in characteristic strain representation. **Note**: assumes the sensitivity
+                is input as power spectral density.
+            **kwargs: Keyword arguments to be passed to ``loglog`` function in matplotlib.
+
+        Returns:
+            Matplotlib figure and axes objects in a 2-tuple.
+
+
+        """
+        if (ax is None and fig is None) or (
+            ax is not None and (isinstance(ax, list) or isinstance(ax, np.ndarray))
+        ):
+            if ax is None and fig is None:
+                outer_shape = self.shape[:-1]
+                if len(outer_shape) == 2:
+                    nrows = outer_shape[0]
+                    ncols = outer_shape[1]
+                elif len(outer_shape) == 1:
+                    nrows = 1
+                    ncols = outer_shape[0]
+
+                fig, ax = plt.subplots(nrows, ncols, sharex=True, sharey=True)
+                try:
+                    ax = ax.ravel()
+                except AttributeError:
+                    ax = [ax]  # just one axis object, no list
+
+            else:
+                assert len(ax) == np.prod(self.shape[:-1])
+
+            for i in range(np.prod(self.shape[:-1])):
+                plot_in = self.flatten()[i]
+                if char_strain:
+                    plot_in = np.sqrt(self.frequency_arr * plot_in)
+                ax[i].loglog(self.frequency_arr, plot_in, **kwargs)
+
+        elif fig is not None:
+            raise NotImplementedError
+
+        elif isinstance(ax, plt.axes):
+            if inds is None:
+                raise ValueError(
+                    "When passing a single axes object for `ax`, but also pass `inds` kwarg."
+                )
+            plot_in = self.sens_mat[inds]
+            if char_strain:
+                plot_in = np.sqrt(self.frequency_arr * plot_in)
+            ax.loglog(self.frequency_arr, plot_in, **kwargs)
+
+        else:
+            raise ValueError(
+                "ax must be a list of axes objects or a single axes object."
+            )
+
+        return (fig, ax)
+
+
+class XYZ1SensitivityMatrix(SensitivityMatrix):
+    """Default sensitivity matrix for XYZ (TDI 1)
+
+    This is 3x3 symmetric matrix.
+
+    Args:
+        f: Frequency array.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(self, f: np.ndarray, **sens_kwargs: dict) -> None:
+        sens_mat = [
+            [X1TDISens, XY1TDISens, ZX1TDISens],
+            [XY1TDISens, Y1TDISens, YZ1TDISens],
+            [ZX1TDISens, YZ1TDISens, Z1TDISens],
+        ]
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+class XYZ2SensitivityMatrix(SensitivityMatrix):
+    """
+    Default sensitivity matrix for XYZ channels using TDI2 transfer functions.
+
+    This creates a 3×3 Hermitian covariance matrix accounting for correlations
+    between the X, Y, and Z TDI channels due to shared noise sources (S_pm and S_op).
+
+    Matrix structure:
+        Σ(f) = [ Σ_XX   Σ_XY   Σ_XZ ]
+               [ Σ_YX   Σ_YY   Σ_YZ ]  at each frequency
+               [ Σ_ZX   Σ_ZY   Σ_ZZ ]
+
+    Args:
+        f: Frequency array [Hz].
+        **sens_kwargs: Keyword arguments to pass to Sensitivity.get_Sn()
+            (e.g., model=lisa_models.sangria).
+
+    Notes:
+        - Inherits matrix inversion and determinant computation from SensitivityMatrix
+        - The invC attribute provides Σ⁻¹(f) for likelihood computations
+        - The detC attribute provides det[Σ(f)] for normalization
+    """
+
+    def __init__(self, f: np.ndarray, **sens_kwargs: dict) -> None:
+        """
+        Initialize TDI2 sensitivity matrix.
+
+        Args:
+            f: Frequency array [Hz].
+            **sens_kwargs: Keyword arguments for Sensitivity.get_Sn()
+                Common kwargs:
+                    - model: LISA noise model (e.g., sangria, sangria)
+                    - stochastic_params: Parameters for galactic foreground
+                    - stochastic_function: Custom stochastic function
+        """
+        # Define 3×3 matrix structure
+        # Diagonal: X2, Y2, Z2 PSDs 
+        # Off-diagonal: XY2, YZ2, ZX2 CSDs 
+        sens_mat = [
+            [X2TDISens,   XY2TDISens,  ZX2TDISens],
+            [XY2TDISens,  Y2TDISens,   YZ2TDISens],
+            [ZX2TDISens,  YZ2TDISens,  Z2TDISens],
+        ]
+
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+class AET1SensitivityMatrix(SensitivityMatrix):
+    """Default sensitivity matrix for AET (TDI 1)
+
+    This is just an array because no cross-terms.
+
+    Args:
+        f: Frequency array.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(self, f: np.ndarray, **sens_kwargs: dict) -> None:
+        sens_mat = [A1TDISens, E1TDISens, T1TDISens]
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+
+
+class AET2SensitivityMatrix(SensitivityMatrix):
+    """Default sensitivity matrix for AET (TDI 2)
+
+    This is just an array because no cross-terms.
+
+    Args:
+        f: Frequency array.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(self, f: np.ndarray, **sens_kwargs: dict) -> None:
+        sens_mat = [A2TDISens, E2TDISens, T2TDISens]
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+
+class AE1SensitivityMatrix(SensitivityMatrix):
+    """Default sensitivity matrix for AE (no T) (TDI 1)
+
+    Args:
+        f: Frequency array.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(self, f: np.ndarray, **sens_kwargs: dict) -> None:
+        sens_mat = [A1TDISens, E1TDISens]
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+
+class AE2SensitivityMatrix(SensitivityMatrix):
+    """Default sensitivity matrix for AE (no T) (TDI 1)
+
+    Args:
+        f: Frequency array.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(self, f: np.ndarray, **sens_kwargs: dict) -> None:
+        sens_mat = [A2TDISens, E2TDISens]
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+
+class LISASensSensitivityMatrix(SensitivityMatrix):
+    """Default sensitivity matrix adding :class:`LISASens` for the specified number of channels.
+
+    Args:
+        f: Frequency array.
+        nchannels: Number of channels.
+        **sens_kwargs: Keyword arguments to pass to :func:`Sensitivity.get_Sn`.
+
+    """
+
+    def __init__(self, f: np.ndarray, nchannels: int, **sens_kwargs: dict) -> None:
+        sens_mat = [LISASens for _ in range(nchannels)]
+        super().__init__(f, sens_mat, **sens_kwargs)
+
+
+def get_sensitivity(
+    f: float | np.ndarray,
+    *args: tuple,
+    sens_fn: Optional[Sensitivity | str] = LISASens,
+    return_type="PSD",
+    fill_nans: float = 1e10,
+    **kwargs,
+) -> float | np.ndarray:
+    """Generic sensitivity generator
+
+    Same interface to many sensitivity curves.
+
+    Args:
+        f: Frequency array.
+        *args: Any additional arguments for the sensitivity function ``get_Sn`` method.
+        sens_fn: String or class that represents the name of the desired PSD function.
+        return_type: Described the desired output. Choices are ASD,
+            PSD, or char_strain (characteristic strain). Default is ASD.
+        fill_nans: Value to fill nans in sensitivity (at 0 frequency).
+            If ``None``, thens nans will be left in the array.
+        **kwargs: Keyword arguments to pass to sensitivity function ``get_Sn`` method.
+
+    Return:
+        Sensitivity values.
+
+    """
+
+    if isinstance(sens_fn, str):
+        sensitivity = check_sensitivity(sens_fn)
+
+    elif hasattr(sens_fn, "get_Sn"):
+        sensitivity = sens_fn
+
+    else:
+        raise ValueError(
+            "sens_fn must be a string for a stock option or a class with a get_Sn method."
+        )
+
+    PSD = sensitivity.get_Sn(f, *args, **kwargs)
+
+    if fill_nans is not None:
+        assert isinstance(fill_nans, float)
+        PSD[np.isnan(PSD)] = fill_nans
+
+    if return_type == "PSD":
+        return PSD
+
+    elif return_type == "ASD":
+        return PSD ** (1 / 2)
+
+    elif return_type == "char_strain":
+        return (f * PSD) ** (1 / 2)
+
+    else:
+        raise ValueError("return_type must be PSD, ASD, or char_strain.")
+
+
+__stock_sens_options__ = [
+    "X1TDISens",
+    "Y1TDISens",
+    "Z1TDISens",
+    "XY1TDISens",
+    "YZ1TDISens",
+    "ZX1TDISens",
+    "A1TDISens",
+    "E1TDISens",
+    "T1TDISens",
+    "X2TDISens",
+    "Y2TDISens",
+    "Z2TDISens",
+    "XY2TDISens",
+    "YZ2TDISens",   
+    "ZX2TDISens",
+    "LISASens",
+    "CornishLISASens",
+    "FlatPSDFunction",
+]
+
+
+def get_stock_sensitivity_options() -> List[Sensitivity]:
+    """Get stock options for sensitivity curves.
+
+    Returns:
+        List of stock sensitivity options.
+
+    """
+    return __stock_sens_options__
+
+
+__stock_sensitivity_mat_options__ = [
+    "XYZ1SensitivityMatrix",
+    "XYZ2SensitivityMatrix",
+    "AET1SensitivityMatrix",
+    "AE1SensitivityMatrix",
+]
+
+
+def get_stock_sensitivity_matrix_options() -> List[SensitivityMatrix]:
+    """Get stock options for sensitivity matrix.
+
+    Returns:
+        List of stock sensitivity matrix options.
+
+    """
+    return __stock_sensitivity_mat_options__
+
+
+def get_stock_sensitivity_from_str(sensitivity: str) -> Sensitivity:
+    """Return a LISA sensitivity from a ``str`` input.
+
+    Args:
+        sensitivity: Sensitivity indicated with a ``str``.
+
+    Returns:
+        Sensitivity associated to that ``str``.
+
+    """
+    if sensitivity not in __stock_sens_options__:
+        raise ValueError(
+            "Requested string sensitivity is not available. See lisatools.sensitivity documentation."
+        )
+    return globals()[sensitivity]
+
+
+def check_sensitivity(sensitivity: Any) -> Sensitivity:
+    """Check input sensitivity.
+
+    Args:
+        sensitivity: Sensitivity to check.
+
+    Returns:
+        Sensitivity checked. Adjusted from ``str`` if ``str`` input.
+
+    """
+    if isinstance(sensitivity, str):
+        sensitivity = get_stock_sensitivity_from_str(sensitivity)
+
+    if not issubclass(sensitivity, Sensitivity):
+        raise ValueError("sensitivity argument not given correctly.")
+
+    return sensitivity