lisaanalysistools 1.0.0__cp312-cp312-macosx_10_9_x86_64.whl

lisatools/sampling/stopping.py

@@ -0,0 +1,320 @@
+import time 
+
+import numpy as np
+
+from eryn.utils.stopping import Stopping
+from eryn.utils.utility import thermodynamic_integration_log_evidence
+
+
+class SNRStopping(Stopping):
+    def __init__(self, snr_limit=100.0, verbose=False):
+        self.snr_limit = snr_limit
+        self.verbose = verbose
+
+    def __call__(self, iter, sample, sampler):
+
+        ind = sampler.get_log_like().argmax()
+
+        log_best = sampler.get_log_like().max()
+        snr_best = sampler.get_blobs()[:, :, :, 0].flatten()[ind]
+        # d_h_best = sampler.get_blobs()[:, :, :, 1].flatten()[ind]
+        # h_h_best = sampler.get_blobs()[:, :, :, 2].flatten()[ind]
+
+        if self.verbose:
+            print(
+                "snr_best",
+                snr_best,
+                "limit:",
+                self.snr_limit,
+                "loglike:",
+                log_best,
+                # d_h_best,
+                # h_h_best,
+            )
+
+        if snr_best > self.snr_limit:
+            return True
+
+        else:
+            return False
+
+
+class NLeavesSearchStopping:
+    def __init__(self, convergence_iter=5, verbose=False):
+        self.convergence_iter = convergence_iter
+        self.verbose = verbose
+
+    def __call__(self, current_info):
+
+        if not hasattr(self, "st"):
+            self.st = time.perf_counter()
+
+        current_iter = current_info.gb_info["reader"].iteration
+
+        if current_iter > self.convergence_iter:
+
+            nleaves_cc = curr.gb_info["reader"].get_nleaves()["gb"][:, 0]
+
+            # do not include most recent
+            nleaves_cc_max_old = nleaves_cc[:-self.convergence_iter].max()
+            nleaves_cc_max_new = nleaves_cc[-self.convergence_iter:].max()
+
+            if nleaves_cc_max_old > nleaves_cc_max_new:
+                stop = True
+
+            else:
+                stop = False
+
+            if self.verbose:
+                dur = (time.perf_counter() - self.st) / 3600.0  # hours
+                print(
+                    "\nnleaves max old:\n",
+                    nleaves_cc_max_old,
+                    "\nnleaves max new:\n",
+                    nleaves_cc_max_newf,
+                    f"\nTIME TO NOW: {dur} hours"
+                )
+
+        return stop
+
+
+class SearchConvergeStopping(Stopping):
+    def __init__(self, n_iters=30, diff=1.0, verbose=False, start_iteration=0):
+        self.n_iters = n_iters
+        self.iters_consecutive = 0
+        self.past_like_best = -np.inf
+        self.diff = diff
+        self.verbose = verbose
+        self.start_iteration = start_iteration
+
+    def __call__(self, iter, sample, sampler):
+
+        like_best = sampler.get_log_like(discard=self.start_iteration).max()
+
+        if np.abs(like_best - self.past_like_best) < self.diff:
+            self.iters_consecutive += 1
+
+        else:
+            self.iters_consecutive = 0
+            self.past_like_best = like_best
+
+        if self.verbose:
+            print(
+                "\nITERS CONSECUTIVE:\n",
+                self.iters_consecutive,
+                self.past_like_best,
+                like_best,
+            )
+
+        if self.iters_consecutive >= self.n_iters:
+            self.iters_consecutive = 0
+            return True
+
+        else:
+            return False
+
+
+
+class GBBandLogLConvergeStopping(Stopping):
+
+    def __init__(self, fd, band_edges, n_iters=30, diff=1.0, verbose=False, start_iteration=0):
+        self.band_edge_inds = np.searchsorted(fd, band_edges, side="right") - 1
+        self.num_bands = self.band_edge_inds.shape[0] - 1
+        self.converged = np.zeros(self.num_bands, dtype=bool)
+        self.iters_consecutive = np.zeros(self.num_bands, dtype=int)
+        self.past_like_best = np.full(self.num_bands, -np.inf)
+        self.n_iters = n_iters
+        self.diff = diff
+        self.verbose = verbose
+        self.start_iteration = start_iteration
+
+    def add_mgh(self, mgh):
+        self.mgh = mgh
+
+    def __call__(self, i, sample, sampler):
+        
+        ll_per_band = self.mgh.get_ll(band_edge_inds=self.band_edge_inds).max(axis=0)
+        
+        ll_movement = (ll_per_band - self.past_like_best) > self.diff
+
+        self.iters_consecutive[~ll_movement] += 1
+        self.iters_consecutive[ll_movement] = 0
+
+        self.converged = self.iters_consecutive >= self.n_iters
+
+        self.past_like_best[ll_movement] = ll_per_band[ll_movement]
+
+        # for move in sampler.all_moves:
+        #     move.converged_sub_bands = self.converged.copy()
+
+        if self.verbose:
+            print("Num still going:", (~self.converged).sum(), "\nChanged here:", (ll_movement).sum())
+
+        if np.all(self.converged):
+            return True
+        else:
+            return False
+
+        
+
+            
+
+
+class SearchConvergeStopping2(Stopping):
+    def __init__(self, n_iters=30, diff=0.1, verbose=False, start_iteration=0, iter_back_check=-1):
+        self.n_iters = n_iters
+        self.iters_consecutive = 0
+        self.past_like_best = -np.inf
+        self.diff = diff
+        self.verbose = verbose
+        self.start_iteration = start_iteration
+        self.iter_back_check = iter_back_check
+        self.time = 0
+        self.back_check = [None for _ in range(self.iter_back_check)]
+        self.last_sampler_iteration = self.start_iteration
+        self.back_check_ind = 0
+        self.stop_here = True
+
+    def __call__(self, iter, sample, sampler):
+
+        self.time += 1
+
+        if sampler.iteration <= self.start_iteration:
+            return False
+
+        lps = sampler.get_log_like(discard=self.start_iteration)[self.last_sampler_iteration - self.start_iteration:]
+        try:
+           like_best = lps.max()
+        except:
+            breakpoint()
+        self.last_sampler_iteration = sampler.iteration
+        
+        if np.any(np.asarray(self.back_check) == None):
+            for i in range(len(self.back_check)):
+                if self.back_check[i] is None:
+                    self.back_check[i] = like_best
+                    return False
+
+        first_check = like_best - self.past_like_best > self.diff
+        second_check = np.all(like_best >= np.asarray(self.back_check))
+
+        # spread in stored values is below difference
+        third_check = np.asarray(self.back_check).max() - np.asarray(self.back_check).min() < self.diff
+
+        update = (
+            (first_check and second_check and self.past_like_best == -np.inf) 
+            or (self.past_like_best == -np.inf and third_check)
+            or (self.past_like_best > -np.inf and first_check)
+        )
+
+        self.back_check[self.back_check_ind] = like_best
+        self.back_check_ind = (self.back_check_ind + 1) % len(self.back_check)
+        
+        if update:
+            self.past_like_best = like_best
+            self.iters_consecutive = 0
+
+        elif self.past_like_best > -np.inf:
+            self.iters_consecutive += 1
+
+        if self.verbose:
+            print(
+                "\nITERS CONSECUTIVE:\n",
+                self.iters_consecutive,
+                f"previous best: {self.past_like_best}, overall best: {like_best},",
+                "first check:", first_check,
+                "second check:", second_check
+            )
+
+        if self.iters_consecutive >= self.n_iters:
+            self.iters_consecutive = 0
+            return True
+
+        else:
+            return False
+
+
+
+class EvidenceStopping(Stopping):
+    def __init__(self, diff=0.5, verbose=False):
+        self.diff = diff
+        self.verbose = verbose
+
+    def __call__(self, iter, sample, sampler):
+
+        betas = sampler.get_betas()[-1]
+        logls = sampler.get_log_like().mean(axis=(0, 2))
+
+        logZ, dlogZ = thermodynamic_integration_log_evidence(betas, logls)
+        print(logZ, dlogZ)
+        return False
+        
+
+        if self.verbose:
+            print(
+                "snr_best",
+                snr_best,
+                "limit:",
+                self.snr_limit,
+                "loglike:",
+                log_best,
+                # d_h_best,
+                # h_h_best,
+            )
+
+        if snr_best > self.snr_limit:
+            return True
+
+        else:
+            return False
+
+
+class MPICommunicateStopping(Stopping):
+
+    def __init__(self, stopper_rank, other_ranks, stop_fn=None):
+
+        self.stopper_rank = stopper_rank
+        self.other_ranks = other_ranks
+        self.stop_fn = stop_fn
+
+    def add_comm(self, comm):
+        self.comm = comm
+
+    def __call__(self, *args, **kwargs):
+
+        if not hasattr(self, "comm"):
+            raise ValueError("Must add comm via add_comm method before __call__ is used.")
+
+        if not hasattr(self, "rank"):
+            self.rank = self.comm.Get_rank()
+            if not self.rank == self.stopper_rank and not self.rank in self.other_ranks:
+                raise ValueError("Rank is not available in other ranks list. Must be either stopper rank or in other ranks list.")
+
+            if self.stopper_rank == self.rank and self.stop_fn is None:
+                raise ValueError("Rank is equivalent to stopper rank but stop_fn is not provided. It must be provided.")
+
+        if self.rank == self.stopper_rank:
+            stop = self.stop_fn(*args, **kwargs)
+
+            if stop:
+                for rank in self.other_ranks:
+                    tag = int(str(rank) + "1000")
+                    self.comm.isend(True, dest=rank, tag=tag)
+
+        else:
+            tag = int(str(self.rank) + "1000")
+            check_stop = self.comm.irecv(source=self.stopper_rank, tag=tag)
+
+            if check_stop.get_status():
+                stop = check_stop.wait()
+
+            else:
+                check_stop.cancel()
+                stop = False
+        
+        return stop
+
+
+
+
+        

lisatools/sampling/utility.py

@@ -0,0 +1,324 @@
+from multiprocessing.sharedctypes import Value
+import os
+
+import numpy as np
+
+try:
+    import cupy as xp
+
+except (ImportError, ModuleNotFoundError) as e:
+    pass
+
+from eryn.state import State, BranchSupplimental
+from eryn.utils.utility import groups_from_inds
+
+
+class DetermineGBGroups:
+    def __init__(self, gb_wave_generator, transform_fn=None, waveform_kwargs={}):
+        self.gb_wave_generator = gb_wave_generator
+        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"):
+        # TODO: mess with mismatch lim setting
+        # TODO: some time of mismatch annealing may be useful
+        if isinstance(last_sample, State):
+            state = last_sample
+            coords = state.branches_coords[name_here][check_temp]
+            inds = state.branches_inds[name_here][check_temp]
+        elif isinstance(last_sample, dict):
+            coords = last_sample[name_here][check_temp]["coords"]
+            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 
+            if start_term == "max":
+                start_walker_ind = inds[check_temp].sum(axis=-1).argmax()
+            elif start_term == "first":
+                start_walker_ind = 0
+            elif start_term == "random":
+                start_walker_ind = np.random.randint(0, nwalkers)
+            else:
+                raise ValueError("start_term must be 'max', 'first', or 'random'.")
+
+            # get all the good leaves in this walker
+            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([])
+                groups_inds.append([])
+                groups[leaf_i].append(coords[start_walker_ind, leaf].copy())
+                groups_inds[leaf_i].append([start_walker_ind, leaf])
+        else:
+            # allows us to check groups based on groups we already have
+            groups = input_groups
+            groups_inds = input_groups_inds
+
+        if len(groups) == 0:
+            return [], [], []
+        for w in range(coords.shape[0]):
+
+            # we have already loaded this group
+            if input_groups is None and w == start_walker_ind:
+                continue
+
+            # walker has no binaries 
+            if not np.any(inds[w]):
+                continue
+
+            # coords in this walker
+            coords_here = coords[w][inds[w]]
+            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)
+
+                if index_within_group == "first":
+                    test_walker_ind = 0
+                elif index_within_group == "random":
+                    test_walker_ind = np.random.randint(0, group_params.shape[0])
+                else:
+                    raise ValueError("start_term must be 'max', 'first', or 'random'.")
+
+                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)
+            
+            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)]
+
+            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"]
+            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)
+            ])
+
+            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,
+                coords_here_in_full,
+                data_minus_template,
+                psd,
+                start_freq_ind=0,
+                data_index=None,
+                noise_index=None,
+                **waveform_kwargs_fill,
+            )
+
+            numerator = self.gb_wave_generator.add_remove
+            norm_here = self.gb_wave_generator.add_add
+            norm_for_test = self.gb_wave_generator.remove_remove
+
+            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
+            
+            # 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)
+            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)])
+
+
+            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:
+                    groups[best[leaf]].append(coords_here[leaf].copy())
+                    groups_inds[best[leaf]].append([w, inds_for_group_stuff[leaf]])
+
+                elif not fix_group_count:
+                    # this only works for high snr limit
+                    groups.append([coords_here[leaf]].copy())
+                    groups_inds.append([[w, inds_for_group_stuff[leaf]]])
+
+        group_lens = [len(group) for group in groups]
+
+        return groups, groups_inds, group_lens
+
+
+class GetLastGBState:
+    def __init__(self, gb_wave_generator, transform_fn=None, waveform_kwargs={}):
+        self.gb_wave_generator = gb_wave_generator
+        self.xp = self.gb_wave_generator.xp
+        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={}):
+
+        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.")
+
+        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'.")
+
+        #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]:
+        #    lengths.append(len(group))
+        #breakpoint()
+        try:
+            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.")
+
+                    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]
+
+            ntemps, nwalkers, nleaves_max_old, ndim = state.branches["gb"].shape
+            if nleaves_max_in is None:
+                nleaves_max = nleaves_max_old
+            else:
+                nleaves_max = nleaves_max_in
+            if nleaves_max_old <= nleaves_max:
+                coords_tmp = np.zeros((ntemps, nwalkers, nleaves_max, ndim))
+                coords_tmp[:, :, :nleaves_max_old, :] = state.branches["gb"].coords
+
+                inds_tmp = np.zeros((ntemps, nwalkers, nleaves_max), dtype=bool)
+                inds_tmp[:, :, :nleaves_max_old] = state.branches["gb"].inds
+                state.branches["gb"].coords = coords_tmp
+                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)
+
+            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.)
+            
+
+        except KeyError:
+            # no "gb"
+            pass
+
+        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.)
+
+        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)
+
+        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)
+        state.supplimental = supps
+
+        return state
+
+
+class HeterodynedUpdate:
+    def __init__(self, update_kwargs, set_d_d_zero=False):
+        self.update_kwargs = update_kwargs
+        self.set_d_d_zero = set_d_d_zero
+
+    def __call__(self, it, sample_state, sampler, **kwargs):
+
+        samples = sample_state.branches_coords["mbh"].reshape(-1, sampler.ndims[0])
+        lp_max = sample_state.log_like.argmax()
+        best = samples[lp_max]
+
+        lp = sample_state.log_like.flatten()
+        sorted = np.argsort(lp)
+        inds_best = sorted[-1000:]
+        inds_worst = sorted[:1000]
+
+        best_full = sampler.log_like_fn.f.parameter_transforms["mbh"].both_transforms(
+            best, copy=True
+        )
+
+        sampler.log_like_fn.f.template_model.init_heterodyne_info(
+            best_full, **self.update_kwargs
+        )
+
+        if self.set_d_d_zero:
+            sampler.log_like_fn.f.template_model.reference_d_d = 0.0
+
+        # TODO: make this a general update function in Eryn (?)
+        # samples[inds_worst] = samples[inds_best].copy()
+        samples = samples.reshape(sampler.ntemps, sampler.nwalkers, 1, sampler.ndims[0])
+        logp = sampler.compute_log_prior({"mbh": samples})
+        logL, blobs = sampler.compute_log_like({"mbh": samples}, logp=logp)
+
+        sample_state.branches["mbh"].coords = samples
+        sample_state.log_like = logL
+        sample_state.blobs = blobs
+
+        # sampler.backend.save_step(sample_state, np.full_like(lp, True))