lisaanalysistools 1.0.10__cp312-cp312-macosx_10_9_x86_64.whl → 1.0.11__cp312-cp312-macosx_10_9_x86_64.whl

{lisaanalysistools-1.0.10.dist-info → lisaanalysistools-1.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lisaanalysistools
-Version: 1.0.10
+Version: 1.0.11
 Home-page: https://github.com/mikekatz04/LISAanalysistools
 Author: Michael Katz
 Author-email: mikekatz04@gmail.com
@@ -66,7 +66,7 @@ Please read [CONTRIBUTING.md](CONTRIBUTING.md) for details on our code of conduc
 
 We use [SemVer](http://semver.org/) for versioning. For the versions available, see the [tags on this repository](https://github.com/mikekatz04/LISAanalysistools/tags).
 
-Current Version: 1.0.10
+Current Version: 1.0.11
 
 ## Authors/Developers
 

{lisaanalysistools-1.0.10.dist-info → lisaanalysistools-1.0.11.dist-info}/RECORD

@@ -1,13 +1,13 @@
 lisatools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lisatools/_version.py,sha256=rrQg9RuwmpsEKf63sR6w0kWC6dP3jibZZD8wvMlcM8E,124
-lisatools/analysiscontainer.py,sha256=kvTP0KSzfAlLBkeF3RNn0BcZs_K1Z0ntQA4xglfxqgI,15345
+lisatools/_version.py,sha256=kORDw-aHXqe534RIAm9NvG4SSRug8egjhH2hTnHhut0,124
+lisatools/analysiscontainer.py,sha256=UQft6SvyDueDtuL1H1auiEzG5IEsFwRiiXf_0DHtu5U,16289
 lisatools/datacontainer.py,sha256=QVz0twD46Fl_J-ueGQUCXOJkkXUEum7yrwp9LaVeohU,9853
-lisatools/detector.py,sha256=RhWIj5SaMpOGOc5_i7GkYf8vkoEqzxUU_or1eSEwb8Q,21007
+lisatools/detector.py,sha256=1ATZToFoy14U3ACTyver2x9DKavJ3l6ToU6nUSGSTmo,21161
 lisatools/diagnostic.py,sha256=o9vtfXbY3yMDk4cGNeOsTrlP7gYG_BRROw7gyuThDaM,34192
-lisatools/sensitivity.py,sha256=cUgYvfsfd-MrKCAUk2rwb-oQm6IfHI-IMndbRLOKqVs,27380
+lisatools/sensitivity.py,sha256=3BEVnRjFsQ7HZX6PRjd6B7K2FyIJRL3lpQpGaHjNgMA,27393
 lisatools/stochastic.py,sha256=CysAQ5BIVD349fLjAMpZAXaT0X0dFjrd4QFi0vfqops,9458
 lisatools/cutils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lisatools/cutils/detector_cpu.cpython-312-darwin.so,sha256=sYWmMV8OJKuqkbrnr9jfPomur9WR5pVVU2eWG_XumJk,154720
+lisatools/cutils/detector_cpu.cpython-312-darwin.so,sha256=xhTNwTCXaSHLmHFOWOVAl3sP4tDy2o7TMExj_G5MORE,154720
 lisatools/cutils/include/Detector.hpp,sha256=Ic37OgP-gvEg8qouhv9aFYf7vQP98tJ_i6PGu8RI1YQ,2050
 lisatools/cutils/include/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lisatools/cutils/include/global.hpp,sha256=3VPiqglTMRrIBXlEvDUJO2-CjKy_SLUXZt-9A1fH6lQ,572
@@ -16,10 +16,10 @@ lisatools/cutils/src/Detector.cu,sha256=JvAK8UCoYxgTACBhx3hi0EK9hT-HRYOxX2XozHuH
 lisatools/cutils/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lisatools/cutils/src/pycppdetector.pyx,sha256=e0tz79VNS7Sxxwyx9RQwsFVdpuClBsUf7OArSkyRBa0,7613
 lisatools/sampling/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lisatools/sampling/likelihood.py,sha256=G2kAQ43qlhAzIFWvYsrSmHXd7WKJAzcCN2o07vRE8vc,29585
+lisatools/sampling/likelihood.py,sha256=A9a1McgjqJre8c3Bzoyxakl3onjxaE6c7zVwY2-12iI,29585
 lisatools/sampling/prior.py,sha256=1K1PMStpwO9WT0qG0aotKSyoNjuehXNbzTDtlk8Q15M,21407
 lisatools/sampling/stopping.py,sha256=Q8q7nM0wnJervhRduf2tBXZZHUVza5kJiAUAMUQXP5o,9682
-lisatools/sampling/utility.py,sha256=rOGotS0Aj8-DAWqsTVy2xWNsxsoz74BVrHEnG2mOkwU,14340
+lisatools/sampling/utility.py,sha256=R5LKNK0A36L1PzLjCENQjDWcDjsR1oUBQGrR9PqasSA,15482
 lisatools/sampling/moves/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lisatools/sampling/moves/skymodehop.py,sha256=0nf72eFhFMGwi0dLJci6XZz-bIMGqco2B2_J72hQvf8,3348
 lisatools/sources/__init__.py,sha256=Fm085xHQ3VpRjqaSlws0bdVefFofAxNzVZyGQQYrQic,140
@@ -36,8 +36,8 @@ lisatools/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lisatools/utils/constants.py,sha256=r1kVwkpbZS13JTOxj2iRxT5sMgTYX30y-S0JdVmD5Oo,1354
 lisatools/utils/pointeradjust.py,sha256=2sT-7qeYWr1pd_sHk9leVHUTSJ7jJgYIRoWQOtYqguE,2995
 lisatools/utils/utility.py,sha256=MKqRsG8qRI1xCsj51mt6QRa80-deUUgedxibyMV3KD4,6776
-lisaanalysistools-1.0.10.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-lisaanalysistools-1.0.10.dist-info/METADATA,sha256=hOqVWPDpsoy9kyzDpSwLA9dgJxF9rY1SnSSWJxiUyMY,4202
-lisaanalysistools-1.0.10.dist-info/WHEEL,sha256=Vn5rrdwmXMBVNlKjgSWDNoqXXbGpTeTFNnBfj813LFw,110
-lisaanalysistools-1.0.10.dist-info/top_level.txt,sha256=qXUn8Xi8yvK0vr3QH0vvT5cmoccjSj-UPeKGLAxdYGo,10
-lisaanalysistools-1.0.10.dist-info/RECORD,,
+lisaanalysistools-1.0.11.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+lisaanalysistools-1.0.11.dist-info/METADATA,sha256=xjT_alb_5mDWKzhKyK9KldcQjlndzlQ8-eyrZ9Fp8fk,4202
+lisaanalysistools-1.0.11.dist-info/WHEEL,sha256=zL8OKpwMaCM_hEdCwF-bHbg5JyuZtrW4vMYNMNxkj9k,110
+lisaanalysistools-1.0.11.dist-info/top_level.txt,sha256=qXUn8Xi8yvK0vr3QH0vvT5cmoccjSj-UPeKGLAxdYGo,10
+lisaanalysistools-1.0.11.dist-info/RECORD,,

{lisaanalysistools-1.0.10.dist-info → lisaanalysistools-1.0.11.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (75.3.0)
 Root-Is-Purelib: false
 Tag: cp312-cp312-macosx_10_9_x86_64
 

lisatools/_version.py

@@ -1,4 +1,4 @@
-__version__ = '1.0.10'
+__version__ = '1.0.11'
 __copyright__ = "Michael L. Katz 2024"
 __name__ = "lisaanalysistools"
 __author__ = "Michael L. Katz"

lisatools/analysiscontainer.py

@@ -11,6 +11,8 @@ import numpy as np
 from scipy import interpolate
 import matplotlib.pyplot as plt
 
+from eryn.utils import TransformContainer
+
 
 try:
     import cupy as cp
@@ -286,6 +288,7 @@ class AnalysisContainer:
         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.
@@ -308,8 +311,14 @@ class AnalysisContainer:
         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, **waveform_kwargs), **data_res_arr_kwargs
+            self.signal_gen(*args_in, **waveform_kwargs), **data_res_arr_kwargs
         )
 
         args_2 = (template,)
@@ -426,7 +435,30 @@ class AnalysisContainer:
             **kwargs,
         )
 
-    def eryn_likelihood_function(self, x, *args, **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)

lisatools/detector.py

@@ -9,6 +9,7 @@ import h5py
 from scipy import interpolate
 
 from .utils.constants import *
+from .utils.utility import get_array_module
 
 import numpy as np
 
@@ -453,9 +454,10 @@ class Orbits(ABC):
             squeeze = False
             t = self.xp.asarray(t)
             sc = self.xp.asarray(sc).astype(np.int32)
+
         else:
             raise ValueError(
-                "(t, sc) can be (float, int), (np.ndarray, int), (np.ndarray, np.ndarray)."
+                "(t, sc) can be (float, int), (np.ndarray, int), (np.ndarray, np.ndarray). If the inputs follow this, make sure the orbits class GPU setting matches the arrays coming in (GPU or CPU)."
             )
 
         # buffer arrays for input into c code

lisatools/sampling/likelihood.py

@@ -1,5 +1,5 @@
 import warnings
-from eryn.state import Branch, BranchSupplimental
+from eryn.state import Branch, BranchSupplemental
 
 import numpy as np
 

lisatools/sampling/utility.py

@@ -9,7 +9,7 @@ try:
 except (ImportError, ModuleNotFoundError) as e:
     pass
 
-from eryn.state import State, BranchSupplimental
+from eryn.state import State, BranchSupplemental
 from eryn.utils.utility import groups_from_inds
 
 
@@ -19,8 +19,21 @@ class DetermineGBGroups:
         self.xp = self.gb_wave_generator.xp
         self.transform_fn = transform_fn
         self.waveform_kwargs = waveform_kwargs
-            
-    def __call__(self, last_sample, name_here, check_temp=0, input_groups=None, input_groups_inds=None, fix_group_count=False, mismatch_lim=0.2, double_check_lim=0.2, start_term="random", waveform_kwargs={}, index_within_group="random"):
+
+    def __call__(
+        self,
+        last_sample,
+        name_here,
+        check_temp=0,
+        input_groups=None,
+        input_groups_inds=None,
+        fix_group_count=False,
+        mismatch_lim=0.2,
+        double_check_lim=0.2,
+        start_term="random",
+        waveform_kwargs={},
+        index_within_group="random",
+    ):
         # TODO: mess with mismatch lim setting
         # TODO: some time of mismatch annealing may be useful
         if isinstance(last_sample, State):
@@ -32,13 +45,13 @@ class DetermineGBGroups:
             inds = last_sample[name_here][check_temp]["inds"]
 
         waveform_kwargs = {**self.waveform_kwargs, **waveform_kwargs}
-        
+
         # get coordinates and inds of the temperature you are considering.
-        
+
         nwalkers, nleaves_max, ndim = coords.shape
         if input_groups is None:
 
-            # figure our which walker to start with 
+            # figure our which walker to start with
             if start_term == "max":
                 start_walker_ind = inds[check_temp].sum(axis=-1).argmax()
             elif start_term == "first":
@@ -52,7 +65,7 @@ class DetermineGBGroups:
             inds_good = np.where(inds[start_walker_ind])[0]
             groups = []
             groups_inds = []
-            
+
             # set up this information to load the information into the group lists
             for leaf_i, leaf in enumerate(inds_good):
                 groups.append([])
@@ -72,7 +85,7 @@ class DetermineGBGroups:
             if input_groups is None and w == start_walker_ind:
                 continue
 
-            # walker has no binaries 
+            # walker has no binaries
             if not np.any(inds[w]):
                 continue
 
@@ -81,7 +94,6 @@ class DetermineGBGroups:
             inds_for_group_stuff = np.arange(len(inds[w]))[inds[w]]
             nleaves, ndim = coords_here.shape
 
-            
             params_for_test = []
             for group in groups:
                 group_params = np.asarray(group)
@@ -95,39 +107,49 @@ class DetermineGBGroups:
 
                 params_for_test.append(group_params[test_walker_ind])
             params_for_test = np.asarray(params_for_test)
-            
+
             # transform coords
             if self.transform_fn is not None:
-                params_for_test_in = self.transform_fn[name_here].both_transforms(params_for_test, return_transpose=False)
-                coords_here_in = self.transform_fn[name_here].both_transforms(coords_here, return_transpose=False)
-            
+                params_for_test_in = self.transform_fn[name_here].both_transforms(
+                    params_for_test, return_transpose=False
+                )
+                coords_here_in = self.transform_fn[name_here].both_transforms(
+                    coords_here, return_transpose=False
+                )
+
             else:
                 params_for_test_in = params_for_test.copy()
                 coords_here_in = coords_here.copy()
 
             inds_tmp_test = np.arange(len(params_for_test_in))
             inds_tmp_here = np.arange(len(coords_here_in))
-            inds_tmp_test, inds_tmp_here = [tmp.ravel() for tmp in np.meshgrid(inds_tmp_test, inds_tmp_here)]
+            inds_tmp_test, inds_tmp_here = [
+                tmp.ravel() for tmp in np.meshgrid(inds_tmp_test, inds_tmp_here)
+            ]
 
             params_for_test_in_full = params_for_test_in[inds_tmp_test]
             coords_here_in_full = coords_here_in[inds_tmp_here]
             # build the waveforms at the same time
 
-            df = 1. / waveform_kwargs["T"]
-            max_f = 1. / 2 * 1/waveform_kwargs["dt"]
+            df = 1.0 / waveform_kwargs["T"]
+            max_f = 1.0 / 2 * 1 / waveform_kwargs["dt"]
             frqs = self.xp.arange(0.0, max_f, df)
-            data_minus_template = self.xp.asarray([
-                self.xp.ones_like(frqs, dtype=complex),
-                self.xp.ones_like(frqs, dtype=complex)
-            ])[None, :, :]
-            psd = self.xp.asarray([
-                self.xp.ones_like(frqs, dtype=np.float64),
-                self.xp.ones_like(frqs, dtype=np.float64)
-            ])
+            data_minus_template = self.xp.asarray(
+                [
+                    self.xp.ones_like(frqs, dtype=complex),
+                    self.xp.ones_like(frqs, dtype=complex),
+                ]
+            )[None, :, :]
+            psd = self.xp.asarray(
+                [
+                    self.xp.ones_like(frqs, dtype=np.float64),
+                    self.xp.ones_like(frqs, dtype=np.float64),
+                ]
+            )
 
             waveform_kwargs_fill = waveform_kwargs.copy()
             waveform_kwargs_fill.pop("start_freq_ind")
- 
+
             # TODO: could use real data and get observed snr for each if needed
             check = self.gb_wave_generator.swap_likelihood_difference(
                 params_for_test_in_full,
@@ -147,26 +169,36 @@ class DetermineGBGroups:
             normalized_autocorr = numerator / np.sqrt(norm_here * norm_for_test)
             normalized_against_test = numerator / norm_for_test
 
-            normalized_autocorr = normalized_autocorr.reshape(coords_here_in.shape[0], params_for_test_in.shape[0]).real
-            normalized_against_test = normalized_against_test.reshape(coords_here_in.shape[0], params_for_test_in.shape[0]).real
-            
+            normalized_autocorr = normalized_autocorr.reshape(
+                coords_here_in.shape[0], params_for_test_in.shape[0]
+            ).real
+            normalized_against_test = normalized_against_test.reshape(
+                coords_here_in.shape[0], params_for_test_in.shape[0]
+            ).real
+
             # TODO: do based on Likelihood? make sure on same posterior
             # TODO: add check based on amplitude
-            test1 = np.abs(1.0 - normalized_autocorr.real)  # (numerator / norm_for_test[None, :]).real)
+            test1 = np.abs(
+                1.0 - normalized_autocorr.real
+            )  # (numerator / norm_for_test[None, :]).real)
             best = test1.argmin(axis=1)
             try:
                 best = best.get()
             except AttributeError:
                 pass
             best_mismatch = test1[(np.arange(test1.shape[0]), best)]
-            check_normalized_against_test = np.abs(1.0 - normalized_against_test[(np.arange(test1.shape[0]), best)])
-
+            check_normalized_against_test = np.abs(
+                1.0 - normalized_against_test[(np.arange(test1.shape[0]), best)]
+            )
 
             f0_here = coords_here[:, 1]
             f0_test = params_for_test[best, 1]
 
             for leaf in range(nleaves):
-                if best_mismatch[leaf] < mismatch_lim and check_normalized_against_test[leaf] < double_check_lim:
+                if (
+                    best_mismatch[leaf] < mismatch_lim
+                    and check_normalized_against_test[leaf] < double_check_lim
+                ):
                     groups[best[leaf]].append(coords_here[leaf].copy())
                     groups_inds[best[leaf]].append([w, inds_for_group_stuff[leaf]])
 
@@ -187,41 +219,65 @@ class GetLastGBState:
         self.transform_fn = transform_fn
         self.waveform_kwargs = waveform_kwargs
 
-    def __call__(self, mgh, reader, df, supps_base_shape, fix_temp_initial_ind:int=None, fix_temp_inds:list=None, nleaves_max_in=None, waveform_kwargs={}):
+    def __call__(
+        self,
+        mgh,
+        reader,
+        df,
+        supps_base_shape,
+        fix_temp_initial_ind: int = None,
+        fix_temp_inds: list = None,
+        nleaves_max_in=None,
+        waveform_kwargs={},
+    ):
 
         xp.cuda.runtime.setDevice(mgh.gpus[0])
 
         if fix_temp_initial_ind is not None or fix_temp_inds is not None:
             if fix_temp_initial_ind is None or fix_temp_inds is None:
-                raise ValueError("If giving fix_temp_initial_ind or fix_temp_inds, must give both.")
+                raise ValueError(
+                    "If giving fix_temp_initial_ind or fix_temp_inds, must give both."
+                )
 
         state = reader.get_last_sample()
 
         waveform_kwargs = {**self.waveform_kwargs, **waveform_kwargs}
         if "start_freq_ind" not in waveform_kwargs:
-            raise ValueError("In get_last_gb_state, waveform_kwargs must include 'start_freq_ind'.")
+            raise ValueError(
+                "In get_last_gb_state, waveform_kwargs must include 'start_freq_ind'."
+            )
 
-        #check = reader.get_last_sample()
+        # check = reader.get_last_sample()
         ntemps, nwalkers, nleaves_max_old, ndim = state.branches["gb"].shape
-        
-        #out = get_groups_for_remixing(check, check_temp=0, input_groups=None, input_groups_inds=None, fix_group_count=False, name_here="gb")
 
-        #lengths = []
-        #for group in out[0]:
+        # out = get_groups_for_remixing(check, check_temp=0, input_groups=None, input_groups_inds=None, fix_group_count=False, name_here="gb")
+
+        # lengths = []
+        # for group in out[0]:
         #    lengths.append(len(group))
-        #breakpoint()
+        # breakpoint()
         try:
-            if fix_temp_initial_ind is not None: 
+            if fix_temp_initial_ind is not None:
                 for i in fix_temp_inds:
                     if i < fix_temp_initial_ind:
-                        raise ValueError("If providing fix_temp_initial_ind and fix_temp_inds, all values in fix_temp_inds must be greater than fix_temp_initial_ind.")
+                        raise ValueError(
+                            "If providing fix_temp_initial_ind and fix_temp_inds, all values in fix_temp_inds must be greater than fix_temp_initial_ind."
+                        )
 
                     state.log_like[i] = state.log_like[fix_temp_initial_ind]
                     state.log_prior[i] = state.log_prior[fix_temp_initial_ind]
-                    state.branches_coords["gb"][i] = state.branches_coords["gb"][fix_temp_initial_ind]
-                    state.branches_coords["gb"][i] = state.branches_coords["gb"][fix_temp_initial_ind]
-                    state.branches_inds["gb"][i] = state.branches_inds["gb"][fix_temp_initial_ind]
-                    state.branches_inds["gb"][i] = state.branches_inds["gb"][fix_temp_initial_ind]
+                    state.branches_coords["gb"][i] = state.branches_coords["gb"][
+                        fix_temp_initial_ind
+                    ]
+                    state.branches_coords["gb"][i] = state.branches_coords["gb"][
+                        fix_temp_initial_ind
+                    ]
+                    state.branches_inds["gb"][i] = state.branches_inds["gb"][
+                        fix_temp_initial_ind
+                    ]
+                    state.branches_inds["gb"][i] = state.branches_inds["gb"][
+                        fix_temp_initial_ind
+                    ]
 
             ntemps, nwalkers, nleaves_max_old, ndim = state.branches["gb"].shape
             if nleaves_max_in is None:
@@ -238,23 +294,34 @@ class GetLastGBState:
                 state.branches["gb"].inds = inds_tmp
                 state.branches["gb"].nleaves_max = nleaves_max
                 state.branches["gb"].shape = (ntemps, nwalkers, nleaves_max, ndim)
-                
+
             else:
                 raise ValueError("new nleaves_max is less than nleaves_max_old.")
 
             # add "gb" if there are any
             data_index_in = groups_from_inds({"gb": state.branches_inds["gb"]})["gb"]
 
-            data_index = xp.asarray(mgh.get_mapped_indices(data_index_in)).astype(xp.int32)
+            data_index = xp.asarray(mgh.get_mapped_indices(data_index_in)).astype(
+                xp.int32
+            )
+
+            params_add_in = self.transform_fn["gb"].both_transforms(
+                state.branches_coords["gb"][state.branches_inds["gb"]]
+            )
 
-            params_add_in = self.transform_fn["gb"].both_transforms(state.branches_coords["gb"][state.branches_inds["gb"]])
-            
             # batch_size is ignored if waveform_kwargs["use_c_implementation"] is True
-            #  -1 is to do -(-d + h) = d - h  
-            mgh.multiply_data(-1.)
-            self.gb_wave_generator.generate_global_template(params_add_in, data_index, mgh.data_list, data_length=mgh.data_length, data_splits=mgh.gpu_splits, batch_size=1000, **waveform_kwargs)
-            mgh.multiply_data(-1.)
-            
+            #  -1 is to do -(-d + h) = d - h
+            mgh.multiply_data(-1.0)
+            self.gb_wave_generator.generate_global_template(
+                params_add_in,
+                data_index,
+                mgh.data_list,
+                data_length=mgh.data_length,
+                data_splits=mgh.gpu_splits,
+                batch_size=1000,
+                **waveform_kwargs,
+            )
+            mgh.multiply_data(-1.0)
 
         except KeyError:
             # no "gb"
@@ -263,22 +330,42 @@ class GetLastGBState:
         data_index_in = groups_from_inds({"gb": state.branches_inds["gb"]})["gb"]
         data_index = xp.asarray(mgh.get_mapped_indices(data_index_in)).astype(xp.int32)
 
-        params_add_in = self.transform_fn["gb"].both_transforms(state.branches_coords["gb"][state.branches_inds["gb"]])
-        
-        #  -1 is to do -(-d + h) = d - h  
-        mgh.multiply_data(-1.)
-        self.gb_wave_generator.generate_global_template(params_add_in, data_index, mgh.data_list, data_length=mgh.data_length, data_splits=mgh.gpu_splits, batch_size=1000, **waveform_kwargs)
-        mgh.multiply_data(-1.)
+        params_add_in = self.transform_fn["gb"].both_transforms(
+            state.branches_coords["gb"][state.branches_inds["gb"]]
+        )
+
+        #  -1 is to do -(-d + h) = d - h
+        mgh.multiply_data(-1.0)
+        self.gb_wave_generator.generate_global_template(
+            params_add_in,
+            data_index,
+            mgh.data_list,
+            data_length=mgh.data_length,
+            data_splits=mgh.gpu_splits,
+            batch_size=1000,
+            **waveform_kwargs,
+        )
+        mgh.multiply_data(-1.0)
 
         self.gb_wave_generator.d_d = np.asarray(mgh.get_inner_product(use_cpu=True))
-        
-        state.log_like = -1/2 * self.gb_wave_generator.d_d.real.reshape(ntemps, nwalkers)
+
+        state.log_like = (
+            -1 / 2 * self.gb_wave_generator.d_d.real.reshape(ntemps, nwalkers)
+        )
 
         temp_inds = mgh.temp_indices.copy()
         walker_inds = mgh.walker_indices.copy()
         overall_inds = mgh.overall_indices.copy()
-            
-        supps = BranchSupplimental({ "temp_inds": temp_inds, "walker_inds": walker_inds, "overall_inds": overall_inds,}, obj_contained_shape=supps_base_shape, copy=True)
+
+        supps = BranchSupplemental(
+            {
+                "temp_inds": temp_inds,
+                "walker_inds": walker_inds,
+                "overall_inds": overall_inds,
+            },
+            obj_contained_shape=supps_base_shape,
+            copy=True,
+        )
         state.supplimental = supps
 
         return state

lisatools/sensitivity.py

@@ -162,11 +162,11 @@ class Sensitivity(ABC):
         ):
             if stochastic_function is None:
                 stochastic_function = FittedHyperbolicTangentGalacticForeground
+                assert len(stochastic_params) == 1
 
-                check = stochastic_function.get_Sh(
-                    f, *stochastic_params, **stochastic_kwargs
-                )
-            sgal[:] = check
+            sgal[:] = stochastic_function.get_Sh(
+                f, *stochastic_params, **stochastic_kwargs
+            )
 
         if squeeze:
             sgal = sgal.squeeze()