ability to fit source position / generate trials

[HansN87]

example code snippet (without gradients for now) below:

class location_trials(object):
    def __init__(self, ana, rss, source, mean_ns, gamma):
        '''
        ana: skyllh analysis object w. custom background gen method
        source: skyllh pointsource object
        rss: skyllh random state service
        '''
        self.ana = ana
        self.max_scan_dist = np.deg2rad(3) # scan for location within 3 degrees
        self.source = source
        self.rss = rss
        self.mean_ns = mean_ns
        self.gamma = gamma
        
        # temp trial storage
        self._n_sig = None
        self._n_events_list = None
        self._events_list = None
        
        
    def _GreatCircleDistance(self, ra_1, dec_1, ra_2, dec_2):
        #Compute the great circle distance between two events#
        #All coordinates must be given in radians#
        delta_dec = np.abs(dec_1 - dec_2)
        delta_ra = np.abs(ra_1 - ra_2)
        x = (np.sin(delta_dec / 2.))**2. + np.cos(dec_1) *\
            np.cos(dec_2) * (np.sin(delta_ra / 2.))**2.
        return 2. * np.arcsin(np.sqrt(x))

    
    def _objective_location(self, x):
        '''
        x : location tuple (ra, dec) in rad
        '''
        current_ra, current_dec = x  
        dist = self._GreatCircleDistance(self.source.ra, self.source.dec, current_ra, current_dec) 
        if dist > self.max_scan_dist:
            return 1.e10
        
        current_source = PointLikeSource(current_ra, current_dec) 
        self.ana.change_source(current_source)
    
        result = self.ana.do_trial_with_given_pseudo_data(rss=self.rss, mean_n_sig=self.mean_ns, n_sig=self._n_sig,
           n_events_list=self._n_events_list, events_list=self._events_list)
        
        return -result['ts']
    
    
    def do_location_trial(self):
        # generate pseudo data on source
        self.ana.change_source(self.source) 
        (n_sig, n_events_list, events_list) = self.ana.generate_pseudo_data(
            rss=self.rss, mean_n_sig=self.mean_ns)
        self._n_sig = n_sig
        self._n_events_list = n_events_list
        self._events_list = events_list
          
        result_on_source = self.ana.do_trial_with_given_pseudo_data(rss=self.rss, mean_n_sig=self.mean_ns, n_sig=self._n_sig,
           n_events_list=self._n_events_list, events_list=self._events_list)
        
        # find best location in vicinity of source
        r = minimize(self._objective_location, x0=[self.source.ra, self.source.dec],  method='Nelder-Mead')
        
        # distance of solution to source
        best_ra, best_dec = r.x
        dist = self._GreatCircleDistance(self.source.ra, self.source.dec, best_ra, best_dec)
        
        best_ts = -r.fun
        
        # get other fit parameters at best-fit
        best_source =  PointLikeSource(best_ra, best_dec)
        self.ana.change_source(best_source)
        result_at_best_fit = self.ana.do_trial_with_given_pseudo_data(rss=self.rss, mean_n_sig=self.mean_ns, n_sig=self._n_sig,
           n_events_list=self._n_events_list, events_list=self._events_list)
        
        # merge result_at_best_fit and result_on_source
        names = []
        vals = []
        for name in result_on_source.dtype.names:
            names.append(name+'_on_source')
            vals.append(result_on_source[name])
            
        names.append('distance_to_source')
        vals.append(dist)
        names.append('best_ra')
        vals.append(best_ra)
        names.append('best_dec')
        vals.append(best_dec)

        combined_result = np.lib.recfunctions.append_fields(result_at_best_fit, names, vals, usemask=False, asrecarray=False)
        return combined_result
    
    def do_location_trials(self, ntrials, outfile=None):
        results = []
        
        for i in range(ntrials):
            results.append(self.do_location_trial())
            
        if outfile:
            np.save(outfile, results)
        
        return np.lib.recfunctions.stack_arrays(results, asrecarray=False, usemask=False)

and

def get_custom_bkg_gen_method():
        bkg_event_rate_field_names = ['astro', 'conv']
        # Custom background generation method without event selection optimizations.
        # Fixes non overlapping events_list distribution in space using the same
        # rss_seed state.
        # Create background generation method.
        # Define the data scrambler with its data scrambling method, which is used
        # for background event generation.
        data_scrambler = DataScrambler(
            UniformRAScramblingMethod())
        def get_bkg_event_prob_func(bkg_event_rate_field_names):
            def bkg_event_prob_func(dataset, data, events):
                p = np.copy(events[bkg_event_rate_field_names[0]])
                for k in bkg_event_rate_field_names[1:]:
                    p += events[k]
                p /= np.sum(p)
                return p
            return bkg_event_prob_func
        def get_bkg_mean_func(bkg_event_rate_field_names):
            def bkg_mean_func(dataset, data, events):
                """Function for calculating the mean number of background events
                for the given data set.
                """
                mean = 0
                for k in bkg_event_rate_field_names:
                    mean += np.sum(events[k])
                mean *= data.livetime*60*60*24
                return mean
            return bkg_mean_func
        bkg_gen_method = MCDataSamplingBkgGenMethod(
            get_bkg_event_prob_func(bkg_event_rate_field_names),
            get_bkg_mean_func(bkg_event_rate_field_names),
            unique_events=False,
            data_scrambler=data_scrambler,
            mc_inplace_scrambling=False,
            keep_mc_data_fields=bkg_event_rate_field_names,
            pre_event_selection_method=None
        )
        return bkg_gen_method

and

ana = create_analysis(
        datasets,
        source,
        bkg_gen_method=get_custom_bkg_gen_method(),
        bkg_event_rate_field_names=['astro', 'conv'],
        refplflux_gamma=gamma,
        fit_gamma=True,
        compress_data=True)