lisaanalysistools 1.1.6__cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

lisatools/analysiscontainer.py

@@ -0,0 +1,474 @@
+from __future__ import annotations
+
+
+import warnings
+from abc import ABC
+from typing import Any, Tuple, Optional, List
+
+import math
+import numpy as np
+
+from scipy import interpolate
+import matplotlib.pyplot as plt
+
+from eryn.utils import TransformContainer
+
+
+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,
+)
+from .datacontainer import DataResidualArray
+from .sensitivity import SensitivityMatrix
+from .diagnostic import (
+    noise_likelihood_term,
+    residual_full_source_and_noise_likelihood,
+    residual_source_likelihood_term,
+    inner_product,
+    data_signal_source_likelihood_term,
+    data_signal_full_source_and_noise_likelihood,
+)
+
+
+class AnalysisContainer:
+    """Combinatorial container that combines sensitivity and data information.
+
+    Args:
+        data_res_arr: Data / Residual / Signal array.
+        sens_mat: Sensitivity information.
+        signal_gen: Callable object that takes information through ``*args`` and ``**kwargs`` and
+            generates a signal in the proper channel setup employed in ``data_res_arr`` and ``sens_mat``.
+
+    """
+
+    def __init__(
+        self,
+        data_res_arr: DataResidualArray,
+        sens_mat: SensitivityMatrix,
+        signal_gen: Optional[callable] = None,
+    ) -> None:
+        self.data_res_arr = data_res_arr
+        self.sens_mat = sens_mat
+
+        if signal_gen is not None:
+            self.signal_gen = signal_gen
+
+    @property
+    def data_res_arr(self) -> DataResidualArray:
+        """Data information."""
+        return self._data_res_arr
+
+    @data_res_arr.setter
+    def data_res_arr(self, data_res_arr: DataResidualArray) -> None:
+        """Set data."""
+        assert isinstance(data_res_arr, DataResidualArray)
+        self._data_res_arr = data_res_arr
+
+    @property
+    def sens_mat(self) -> SensitivityMatrix:
+        """Sensitivity information."""
+        return self._sens_mat
+
+    @sens_mat.setter
+    def sens_mat(self, sens_mat: SensitivityMatrix) -> None:
+        "Set sensitivity information."
+        assert isinstance(sens_mat, SensitivityMatrix)
+        self._sens_mat = sens_mat
+
+    @property
+    def signal_gen(self) -> callable:
+        """Signal generator."""
+        if not hasattr(self, "_signal_gen"):
+            raise ValueError(
+                "User must input signal_gen kwarg to use the signal generator."
+            )
+        return self._signal_gen
+
+    @signal_gen.setter
+    def signal_gen(self, signal_gen: callable):
+        """Set signal generator."""
+        assert hasattr(signal_gen, "__call__")
+        self._signal_gen = signal_gen
+
+    def loglog(self) -> Tuple[plt.Figure, plt.Axes]:
+        """Produce loglog plot of both source and sensitivity information.
+
+        Returns:
+            Matplotlib figure and axes object in a 2-tuple.
+
+        """
+        fig, ax = self.sens_mat.loglog(char_strain=True)
+        if self.sens_mat.ndim == 3:
+            # 3x3 most likely
+            for i in range(self.sens_mat.shape[0]):
+                for j in range(i, self.sens_mat.shape[1]):
+                    # char strain
+                    ax[i * self.sens_mat.shape[1] + j].loglog(
+                        self.data_res_arr.f_arr,
+                        self.data_res_arr.f_arr * np.abs(self.data_res_arr[i]),
+                    )
+                    ax[i * self.sens_mat.shape[1] + j].loglog(
+                        self.data_res_arr.f_arr,
+                        self.data_res_arr.f_arr * np.abs(self.data_res_arr[j]),
+                    )
+        else:
+            for i in range(self.sens_mat.shape[0]):
+                ax[i].loglog(
+                    self.data_res_arr.f_arr,
+                    self.data_res_arr.f_arr * np.abs(self.data_res_arr[i]),
+                )
+        return (fig, ax)
+
+    def inner_product(self, **kwargs: dict) -> float | complex:
+        """Return the inner product of the current set of information
+
+        Args:
+            **kwargs: Inner product keyword arguments.
+
+        Returns:
+            Inner product value.
+
+        """
+        if "psd" in kwargs:
+            kwargs.pop("psd")
+
+        return inner_product(self.data_res_arr, self.data_res_arr, psd=self.sens_mat)
+
+    def snr(self, **kwargs: dict) -> float:
+        """Return the SNR of the current set of information
+
+        Args:
+            **kwargs: Inner product keyword arguments.
+
+        Returns:
+            SNR value.
+
+        """
+        return self.inner_product(**kwargs).real ** (1 / 2)
+
+    def template_inner_product(
+        self, template: DataResidualArray, **kwargs: dict
+    ) -> float | complex:
+        """Calculate the inner product of a template with the data.
+
+        Args:
+            template: Template signal.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`.
+
+        Returns:
+            Inner product value.
+
+        """
+        if "psd" in kwargs:
+            kwargs.pop("psd")
+
+        if "include_psd_info" in kwargs:
+            kwargs.pop("include_psd_info")
+
+        ip_val = inner_product(self.data_res_arr, template, psd=self.sens_mat, **kwargs)
+        return ip_val
+
+    def template_snr(
+        self, template: DataResidualArray, phase_maximize: bool = False, **kwargs: dict
+    ) -> Tuple[float, float]:
+        """Calculate the SNR of a template, both optimal and detected.
+
+        Args:
+            template: Template signal.
+            phase_maximize: If ``True``, maximize over an overall phase.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`.
+
+        Returns:
+            ``(optimal snr, detected snr)``.
+
+        """
+        kwargs_in = kwargs.copy()
+        if "psd" in kwargs:
+            kwargs_in.pop("psd")
+
+        if "complex" in kwargs_in:
+            kwargs_in.pop("complex")
+
+        # TODO: should we cache?
+        h_h = inner_product(template, template, psd=self.sens_mat, **kwargs_in)
+        non_marg_d_h = inner_product(
+            self.data_res_arr, template, psd=self.sens_mat, complex=True, **kwargs_in
+        )
+        d_h = np.abs(non_marg_d_h) if phase_maximize else non_marg_d_h.copy()
+        self.non_marg_d_h = non_marg_d_h
+
+        opt_snr = np.sqrt(h_h.real)
+        det_snr = d_h.real / opt_snr
+        return (opt_snr, det_snr)
+
+    def template_likelihood(
+        self,
+        template: DataResidualArray,
+        include_psd_info: bool = False,
+        phase_maximize: bool = False,
+        **kwargs: dict,
+    ) -> float:
+        """Calculate the Likelihood of a template against the data.
+
+        Args:
+            template: Template signal.
+            include_psd_info: If ``True``, add the PSD term to the Likelihood value.
+            phase_maximize: If ``True``, maximize over an overall phase.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`.
+
+        Returns:
+            Likelihood value.
+
+        """
+        kwargs_in = kwargs.copy()
+        if "psd" in kwargs_in:
+            kwargs_in.pop("psd")
+
+        if "complex" in kwargs_in:
+            kwargs_in.pop("complex")
+
+        # TODO: should we cache?
+        d_d = inner_product(
+            self.data_res_arr, self.data_res_arr, psd=self.sens_mat, **kwargs_in
+        )
+        h_h = inner_product(template, template, psd=self.sens_mat, **kwargs_in)
+        non_marg_d_h = inner_product(
+            self.data_res_arr, template, psd=self.sens_mat, complex=True, **kwargs_in
+        )
+        d_h = np.abs(non_marg_d_h) if phase_maximize else non_marg_d_h.copy()
+        self.non_marg_d_h = non_marg_d_h
+        like_out = -1 / 2 * (d_d + h_h - 2 * d_h).real
+
+        if include_psd_info:
+            # add noise term if requested
+            like_out += self.likelihood(noise_only=True)
+
+        return like_out
+
+    def likelihood(
+        self, source_only: bool = False, noise_only: bool = False, **kwargs: dict
+    ) -> float | complex:
+        """Return the likelihood of the current arangement.
+
+        Args:
+            source_only: If ``True`` return the source-only Likelihood.
+            noise_only: If ``True``, return the noise part of the Likelihood alone.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`.
+
+        Returns:
+            Likelihood value.
+
+        """
+        if noise_only and source_only:
+            raise ValueError("noise_only and source only cannot both be True.")
+        elif noise_only:
+            return noise_likelihood_term(self.sens_mat)
+        elif source_only:
+            return residual_source_likelihood_term(
+                self.data_res_arr, psd=self.sens_mat, **kwargs
+            )
+        else:
+            return residual_full_source_and_noise_likelihood(
+                self.data_res_arr, self.sens_mat, **kwargs
+            )
+
+    def _calculate_signal_operation(
+        self,
+        calc: str,
+        *args: Any,
+        source_only: bool = False,
+        waveform_kwargs: Optional[dict] = {},
+        data_res_arr_kwargs: Optional[dict] = {},
+        transform_fn: Optional[TransformContainer] = None,
+        **kwargs: dict,
+    ) -> float | complex:
+        """Return the likelihood of a generated signal with the data.
+
+        Args:
+            calc: Type of calculation to do. Options are ``"likelihood"``, ``"inner_product"``, or ``"snr"``.
+            *args: Arguments to waveform generating function. Must include parameters.
+            source_only: If ``True`` return the source-only Likelihood (leave out noise part).
+            waveform_kwargs: Keyword arguments to pass to waveform generator.
+            data_res_arr_kwargs: Keyword arguments for instantiation of :class:`DataResidualArray`.
+                This can be used if any transforms are desired prior to the Likelihood computation. If it is not input,
+                the kwargs are taken to be the same as those used to initalize ``self.data_res_arr``.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`
+
+        Returns:
+            Likelihood value.
+
+        """
+
+        if data_res_arr_kwargs == {}:
+            data_res_arr_kwargs = self.data_res_arr.init_kwargs
+
+        if transform_fn is not None:
+            args_tmp = np.asarray(args)
+            args_in = tuple(transform_fn.both_transforms(args_tmp))
+        else:
+            args_in = args
+
+        template = DataResidualArray(
+            self.signal_gen(*args_in, **waveform_kwargs), **data_res_arr_kwargs
+        )
+
+        args_2 = (template,)
+
+        if "include_psd_info" in kwargs:
+            assert kwargs["include_psd_info"] == (not source_only)
+            kwargs.pop("include_psd_info")
+
+        kwargs = dict(psd=self.sens_mat, **kwargs)
+
+        if calc == "likelihood":
+            kwargs["include_psd_info"] = not source_only
+            return self.template_likelihood(*args_2, **kwargs)
+        elif calc == "inner_product":
+            return self.template_inner_product(*args_2, **kwargs)
+        elif calc == "snr":
+            return self.template_snr(*args_2, **kwargs)
+        else:
+            raise ValueError("`calc` must be 'likelihood', 'inner_product', or 'snr'.")
+
+    def calculate_signal_likelihood(
+        self,
+        *args: Any,
+        source_only: bool = False,
+        waveform_kwargs: Optional[dict] = {},
+        data_res_arr_kwargs: Optional[dict] = {},
+        **kwargs: dict,
+    ) -> float | complex:
+        """Return the likelihood of a generated signal with the data.
+
+        Args:
+            params: Arguments to waveform generating function. Must include parameters.
+            source_only: If ``True`` return the source-only Likelihood (leave out noise part).
+            waveform_kwargs: Keyword arguments to pass to waveform generator.
+            data_res_arr_kwargs: Keyword arguments for instantiation of :class:`DataResidualArray`.
+                This can be used if any transforms are desired prior to the Likelihood computation.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`
+
+        Returns:
+            Likelihood value.
+
+        """
+
+        return self._calculate_signal_operation(
+            "likelihood",
+            *args,
+            source_only=source_only,
+            waveform_kwargs=waveform_kwargs,
+            data_res_arr_kwargs=data_res_arr_kwargs,
+            **kwargs,
+        )
+
+    def calculate_signal_inner_product(
+        self,
+        *args: Any,
+        source_only: bool = False,
+        waveform_kwargs: Optional[dict] = {},
+        data_res_arr_kwargs: Optional[dict] = {},
+        **kwargs: dict,
+    ) -> float | complex:
+        """Return the inner product of a generated signal with the data.
+
+        Args:
+            *args: Arguments to waveform generating function. Must include parameters.
+            source_only: If ``True`` return the source-only Likelihood (leave out noise part).
+            waveform_kwargs: Keyword arguments to pass to waveform generator.
+            data_res_arr_kwargs: Keyword arguments for instantiation of :class:`DataResidualArray`.
+                This can be used if any transforms are desired prior to the Likelihood computation.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`
+
+        Returns:
+            Inner product value.
+
+        """
+
+        return self._calculate_signal_operation(
+            "inner_product",
+            *args,
+            source_only=source_only,
+            waveform_kwargs=waveform_kwargs,
+            data_res_arr_kwargs=data_res_arr_kwargs,
+            **kwargs,
+        )
+
+    def calculate_signal_snr(
+        self,
+        *args: Any,
+        source_only: bool = False,
+        waveform_kwargs: Optional[dict] = {},
+        data_res_arr_kwargs: Optional[dict] = {},
+        **kwargs: dict,
+    ) -> Tuple[float, float]:
+        """Return the SNR of a generated signal with the data.
+
+        Args:
+            *args: Arguments to waveform generating function. Must include parameters.
+            source_only: If ``True`` return the source-only Likelihood (leave out noise part).
+            waveform_kwargs: Keyword arguments to pass to waveform generator.
+            data_res_arr_kwargs: Keyword arguments for instantiation of :class:`DataResidualArray`.
+                This can be used if any transforms are desired prior to the Likelihood computation.
+            **kwargs: Keyword arguments to pass to :func:`lisatools.diagnostic.inner_product`
+
+        Returns:
+            Snr values (optimal, detected).
+
+        """
+
+        return self._calculate_signal_operation(
+            "snr",
+            *args,
+            source_only=source_only,
+            waveform_kwargs=waveform_kwargs,
+            data_res_arr_kwargs=data_res_arr_kwargs,
+            **kwargs,
+        )
+
+    def eryn_likelihood_function(
+        self, x: np.ndarray | list | tuple, *args: Any, **kwargs: Any
+    ) -> np.ndarray | float:
+        """Likelihood function for Eryn sampler.
+
+        This function is not vectorized.
+
+        ``signal_gen`` must be set to use this function.
+
+        Args:
+            x: Parameters. Can be 1D list, tuple, array or 2D array.
+                If a 2D array is input, the computation is done serially.
+            *args: Likelihood args.
+            **kwargs: Likelihood kwargs.
+
+        Returns:
+            Likelihood value(s).
+
+        """
+        assert self.signal_gen is not None
+
+        if isinstance(x, list) or isinstance(x, tuple):
+            x = np.asarray(x)
+
+        if x.ndim == 1:
+            input_vals = tuple(x) + tuple(args)
+            return self.calculate_signal_likelihood(*input_vals, **kwargs)
+        elif x.ndim == 2:
+            likelihood_out = np.zeros(x.shape[0])
+            for i in range(x.shape[0]):
+                input_vals = tuple(x[i]) + tuple(args)
+                likelihood_out[i] = self.calculate_signal_likelihood(
+                    *input_vals, **kwargs
+                )
+
+        else:
+            raise ValueError("x must be a 1D or 2D array.")

lisatools/cutils/__init__.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+import dataclasses
+import enum
+import types
+import typing
+import abc
+from typing import Optional, Sequence, TypeVar, Union
+from ..utils.exceptions import *
+
+from gpubackendtools.gpubackendtools import BackendMethods, CpuBackend, Cuda11xBackend, Cuda12xBackend
+from gpubackendtools.exceptions import *
+
+@dataclasses.dataclass
+class LISAToolsBackendMethods(BackendMethods):
+    pycppDetector: object
+
+class LISAToolsBackend:
+    # TODO: not ClassVar?
+    pycppDetector: object
+
+    def __init__(self, lisatools_backend_methods):
+
+        # set direct lisatools methods
+        # pass rest to general backend
+        assert isinstance(lisatools_backend_methods, LISAToolsBackendMethods)
+        self.pycppDetector = lisatools_backend_methods.pycppDetector
+    
+
+class LISAToolsCpuBackend(CpuBackend, LISAToolsBackend):
+    """Implementation of the CPU backend"""
+    
+    _backend_name = "lisatools_backend_cpu"
+    _name = "lisatools_cpu"
+    def __init__(self, *args, **kwargs):
+        CpuBackend.__init__(self, *args, **kwargs)
+        LISAToolsBackend.__init__(self, self.cpu_methods_loader())
+
+    @staticmethod
+    def cpu_methods_loader() -> LISAToolsBackendMethods:
+        try:
+            import lisatools_backend_cpu.pycppdetector
+
+        except (ModuleNotFoundError, ImportError) as e:
+            raise BackendUnavailableException(
+                "'cpu' backend could not be imported."
+            ) from e
+
+        numpy = LISAToolsCpuBackend.check_numpy()
+
+        return LISAToolsBackendMethods(
+            pycppDetector=lisatools_backend_cpu.pycppdetector.pycppDetector,
+            xp=numpy,
+        )
+
+
+class LISAToolsCuda11xBackend(Cuda11xBackend, LISAToolsBackend):
+
+    """Implementation of CUDA 11.x backend"""
+    _backend_name : str = "lisatools_backend_cuda11x"
+    _name = "lisatools_cuda11x"
+
+    def __init__(self, *args, **kwargs):
+        Cuda11xBackend.__init__(self, *args, **kwargs)
+        LISAToolsBackend.__init__(self, self.cuda11x_module_loader())
+        
+    @staticmethod
+    def cuda11x_module_loader():
+        try:
+            import lisatools_backend_cuda11x.utils
+
+        except (ModuleNotFoundError, ImportError) as e:
+            raise BackendUnavailableException(
+                "'cuda11x' backend could not be imported."
+            ) from e
+
+        try:
+            import cupy
+        except (ModuleNotFoundError, ImportError) as e:
+            raise MissingDependencies(
+                "'cuda11x' backend requires cupy", pip_deps=["cupy-cuda11x"]
+            ) from e
+
+        return LISAToolsBackendMethods(
+            pycppDetector=lisatools_backend_cuda11x.pycppdetector.pycppDetector,
+            xp=cupy,
+        )
+
+class LISAToolsCuda12xBackend(Cuda12xBackend, LISAToolsBackend):
+    """Implementation of CUDA 12.x backend"""
+    _backend_name : str = "lisatools_backend_cuda12x"
+    _name = "lisatools_cuda12x"
+    
+    def __init__(self, *args, **kwargs):
+        Cuda12xBackend.__init__(self, *args, **kwargs)
+        LISAToolsBackend.__init__(self, self.cuda12x_module_loader())
+        
+    @staticmethod
+    def cuda12x_module_loader():
+        try:
+            import lisatools_backend_cuda12x.utils
+
+        except (ModuleNotFoundError, ImportError) as e:
+            raise BackendUnavailableException(
+                "'cuda12x' backend could not be imported."
+            ) from e
+
+        try:
+            import cupy
+        except (ModuleNotFoundError, ImportError) as e:
+            raise MissingDependencies(
+                "'cuda12x' backend requires cupy", pip_deps=["cupy-cuda12x"]
+            ) from e
+
+        return LISAToolsBackendMethods(
+            pycppDetector=lisatools_backend_cuda12x.pycppdetector.pycppDetector,
+            xp=cupy,
+        )
+
+
+KNOWN_BACKENDS = {
+    "cuda12x": LISAToolsCuda12xBackend,
+    "cuda11x": LISAToolsCuda11xBackend,
+    "cpu": LISAToolsCpuBackend,
+}
+
+"""List of existing backends, per default order of preference."""