mcmc_based_searches.py 92 KB
Newer Older
2001
                    vtext))
2002
2003
2004

        if gen_tex_table:
            filename = '{}/{}_run_setup.tex'.format(self.outdir, self.label)
2005
            with open(filename, 'w+') as f:
2006
                f.write(r'\begin{tabular}{c|ccc}' + '\n')
2007
2008
                f.write(r'Stage & $N_\mathrm{seg}$ &'
                        r'$T_\mathrm{coh}^{\rm days}$ &'
2009
                        r'$\mathcal{N}^*$ \\ \hline'
2010
2011
2012
2013
                        '\n')
                for i, rs in enumerate(run_setup):
                    Tcoh = float(
                        self.maxStartTime - self.minStartTime)/rs[1]/86400
2014
                    line = r'{} & {} & {} & {} \\' + '\n'
2015
2016
                    if Nstar_vals[i] is None:
                        Nstar = 'N/A'
2017
                    else:
2018
                        Nstar = Nstar_vals[i]
2019
                    line = line.format(i, rs[1], '{:1.1f}'.format(Tcoh),
2020
                                       helper_functions.texify_float(Nstar))
2021
2022
                    f.write(line)
                f.write(r'\end{tabular}' + '\n')
2023
2024
2025
2026

        if args.setup_only:
            logging.info("Exit as requested by setup_only flag")
            sys.exit()
2027
2028
2029
        else:
            return run_setup

Gregory Ashton's avatar
Gregory Ashton committed
2030
2031
2032
    def run(self, run_setup=None, proposal_scale_factor=2, NstarMax=10,
            Nsegs0=None, create_plots=True, log_table=True, gen_tex_table=True,
            fig=None, axes=None, return_fig=False, **kwargs):
2033
2034
2035
2036
        """ Run the follow-up with the given run_setup

        Parameters
        ----------
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
        run_setup: list of tuples, optional
        proposal_scale_factor: float
            The proposal scale factor used by the sampler, see Goodman & Weare
            (2010). If the acceptance fraction is too low, you can raise it by
            decreasing the a parameter; and if it is too high, you can reduce
            it by increasing the a parameter [Foreman-Mackay (2013)].
        create_plots: bool
            If true, save trace plots of the walkers
        c: int
            The minimum number of autocorrelation times needed to trust the
            result when estimating the autocorrelation time (see
            emcee.autocorr.integrated_time for further details. Default is 5
        **kwargs:
            Passed to _plot_walkers to control the figures
2051
2052
2053
2054

        """

        self.nsegs = 1
2055
        self._initiate_search_object()
2056
        run_setup = self.init_run_setup(
Gregory Ashton's avatar
Gregory Ashton committed
2057
            run_setup, NstarMax=NstarMax, Nsegs0=Nsegs0, log_table=log_table,
2058
            gen_tex_table=gen_tex_table)
2059
        self.run_setup = run_setup
Gregory Ashton's avatar
Gregory Ashton committed
2060

2061
        self.old_data_is_okay_to_use = self._check_old_data_is_okay_to_use()
Gregory Ashton's avatar
Gregory Ashton committed
2062
2063
2064
        if self.old_data_is_okay_to_use is True:
            logging.warning('Using saved data from {}'.format(
                self.pickle_path))
2065
            d = self.get_saved_data_dictionary()
Gregory Ashton's avatar
Gregory Ashton committed
2066
2067
2068
            self.samples = d['samples']
            self.lnprobs = d['lnprobs']
            self.lnlikes = d['lnlikes']
2069
            self.all_lnlikelihood = d['all_lnlikelihood']
2070
            self.nsegs = run_setup[-1][1]
Gregory Ashton's avatar
Gregory Ashton committed
2071
2072
2073
2074
2075
            return

        nsteps_total = 0
        for j, ((nburn, nprod), nseg, reset_p0) in enumerate(run_setup):
            if j == 0:
2076
2077
                p0 = self._generate_initial_p0()
                p0 = self._apply_corrections_to_p0(p0)
Gregory Ashton's avatar
Gregory Ashton committed
2078
            elif reset_p0:
2079
2080
2081
                p0 = self._get_new_p0(sampler)
                p0 = self._apply_corrections_to_p0(p0)
                # self._check_initial_points(p0)
Gregory Ashton's avatar
Gregory Ashton committed
2082
2083
2084
2085
            else:
                p0 = sampler.chain[:, :, -1, :]

            self.nsegs = nseg
2086
            self._set_likelihoodcoef()
2087
            self.search.nsegs = nseg
Gregory Ashton's avatar
Gregory Ashton committed
2088
            self.update_search_object()
2089
            self.search.init_semicoherent_parameters()
Gregory Ashton's avatar
Gregory Ashton committed
2090
2091
2092
2093
2094
2095
            sampler = emcee.PTSampler(
                self.ntemps, self.nwalkers, self.ndim, self.logl, self.logp,
                logpargs=(self.theta_prior, self.theta_keys, self.search),
                loglargs=(self.search,), betas=self.betas,
                a=proposal_scale_factor)

2096
            Tcoh = (self.maxStartTime-self.minStartTime)/nseg/86400.
Gregory Ashton's avatar
Gregory Ashton committed
2097
2098
            logging.info(('Running {}/{} with {} steps and {} nsegs '
                          '(Tcoh={:1.2f} days)').format(
2099
                j+1, len(run_setup), (nburn, nprod), nseg, Tcoh))
2100
            sampler = self._run_sampler(sampler, p0, nburn=nburn, nprod=nprod)
2101
2102
            logging.info('Max detection statistic of run was {}'.format(
                np.max(sampler.lnlikelihood)))
Gregory Ashton's avatar
Gregory Ashton committed
2103

2104
            if create_plots:
2105
                fig, axes = self._plot_walkers(
2106
                    sampler, symbols=self.theta_symbols, fig=fig, axes=axes,
2107
                    nprod=nprod, xoffset=nsteps_total, **kwargs)
2108
2109
                for ax in axes[:self.ndim]:
                    ax.axvline(nsteps_total, color='k', ls='--', lw=0.25)
Gregory Ashton's avatar
Gregory Ashton committed
2110

2111
            nsteps_total += nburn+nprod
Gregory Ashton's avatar
Gregory Ashton committed
2112
2113
2114
2115

        samples = sampler.chain[0, :, nburn:, :].reshape((-1, self.ndim))
        lnprobs = sampler.lnprobability[0, :, nburn:].reshape((-1))
        lnlikes = sampler.lnlikelihood[0, :, nburn:].reshape((-1))
2116
        all_lnlikelihood = sampler.lnlikelihood
Gregory Ashton's avatar
Gregory Ashton committed
2117
2118
2119
        self.samples = samples
        self.lnprobs = lnprobs
        self.lnlikes = lnlikes
2120
2121
        self.all_lnlikelihood = all_lnlikelihood
        self._save_data(sampler, samples, lnprobs, lnlikes, all_lnlikelihood)
Gregory Ashton's avatar
Gregory Ashton committed
2122

2123
2124
2125
2126
2127
2128
2129
2130
2131
        if create_plots:
            try:
                fig.tight_layout()
            except (ValueError, RuntimeError) as e:
                logging.warning('Tight layout encountered {}'.format(e))
            if return_fig:
                return fig, axes
            else:
                fig.savefig('{}/{}_walkers.png'.format(
Gregory Ashton's avatar
Gregory Ashton committed
2132
                    self.outdir, self.label))
2133

Gregory Ashton's avatar
Gregory Ashton committed
2134

Gregory Ashton's avatar
Gregory Ashton committed
2135
class MCMCTransientSearch(MCMCSearch):
Gregory Ashton's avatar
Gregory Ashton committed
2136
2137
2138
2139
2140
2141
    """ MCMC search for a transient signal using ComputeFstat

    See parent MCMCSearch for a list of all additional parameters, here we list
    only the additional init parameters of this class.

    """
Gregory Ashton's avatar
Gregory Ashton committed
2142

Gregory Ashton's avatar
Gregory Ashton committed
2143
2144
    symbol_dictionary = dict(
        F0='$f$', F1='$\dot{f}$', F2='$\ddot{f}$',
2145
        Alpha=r'$\alpha$', Delta='$\delta$',
2146
        transient_tstart='$t_\mathrm{start}$', transient_duration='$\Delta T$')
Gregory Ashton's avatar
Gregory Ashton committed
2147
    unit_dictionary = dict(
2148
        F0='Hz', F1='Hz/s', F2='Hz/s$^2$', Alpha=r'rad', Delta='rad',
2149
        transient_tstart='s', transient_duration='s')
Gregory Ashton's avatar
Gregory Ashton committed
2150

2151
    transform_dictionary = dict(
2152
        transient_duration={'multiplier': 1/86400.,
2153
2154
                            'unit': 'day',
                            'symbol': 'Transient duration'},
2155
2156
        transient_tstart={
            'multiplier': 1/86400.,
2157
2158
            'subtractor': 'minStartTime',
            'unit': 'day',
2159
2160
2161
            'label': 'Transient start-time \n days after minStartTime'}
            )

2162
    def _initiate_search_object(self):
Gregory Ashton's avatar
Gregory Ashton committed
2163
        logging.info('Setting up search object')
2164
        self.search = core.ComputeFstat(
2165
            tref=self.tref, sftfilepattern=self.sftfilepattern,
Gregory Ashton's avatar
Gregory Ashton committed
2166
            minCoverFreq=self.minCoverFreq, maxCoverFreq=self.maxCoverFreq,
2167
            detectors=self.detectors, transient=True,
Gregory Ashton's avatar
Gregory Ashton committed
2168
            minStartTime=self.minStartTime, maxStartTime=self.maxStartTime,
Gregory Ashton's avatar
Gregory Ashton committed
2169
2170
            BSGL=self.BSGL, binary=self.binary,
            injectSources=self.injectSources)
2171
2172
2173
2174
        if self.minStartTime is None:
            self.minStartTime = self.search.minStartTime
        if self.maxStartTime is None:
            self.maxStartTime = self.search.maxStartTime
Gregory Ashton's avatar
Gregory Ashton committed
2175
2176
2177
2178

    def logl(self, theta, search):
        for j, theta_i in enumerate(self.theta_idxs):
            self.fixed_theta[theta_i] = theta[j]
2179
2180
2181
        in_theta = copy.copy(self.fixed_theta)
        in_theta[1] = in_theta[0] + in_theta[1]
        if in_theta[1] > self.maxStartTime:
2182
            return -np.inf
2183
2184
        twoF = search.get_fullycoherent_twoF(*in_theta)
        return twoF/2.0 + self.likelihoodcoef
Gregory Ashton's avatar
Gregory Ashton committed
2185

2186
    def _unpack_input_theta(self):
Gregory Ashton's avatar
Gregory Ashton committed
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
        full_theta_keys = ['transient_tstart',
                           'transient_duration', 'F0', 'F1', 'F2', 'Alpha',
                           'Delta']
        if self.binary:
            full_theta_keys += [
                'asini', 'period', 'ecc', 'tp', 'argp']
        full_theta_keys_copy = copy.copy(full_theta_keys)

        full_theta_symbols = [r'$t_{\rm start}$', r'$\Delta T$',
                              '$f$', '$\dot{f}$', '$\ddot{f}$',
                              r'$\alpha$', r'$\delta$']
        if self.binary:
            full_theta_symbols += [
                'asini', 'period', 'period', 'ecc', 'tp', 'argp']

        self.theta_keys = []
        fixed_theta_dict = {}
        for key, val in self.theta_prior.iteritems():
            if type(val) is dict:
                fixed_theta_dict[key] = 0
                self.theta_keys.append(key)
            elif type(val) in [float, int, np.float64]:
                fixed_theta_dict[key] = val
            else:
                raise ValueError(
                    'Type {} of {} in theta not recognised'.format(
                        type(val), key))
            full_theta_keys_copy.pop(full_theta_keys_copy.index(key))

        if len(full_theta_keys_copy) > 0:
            raise ValueError(('Input dictionary `theta` is missing the'
                              'following keys: {}').format(
                                  full_theta_keys_copy))

        self.fixed_theta = [fixed_theta_dict[key] for key in full_theta_keys]
        self.theta_idxs = [full_theta_keys.index(k) for k in self.theta_keys]
        self.theta_symbols = [full_theta_symbols[i] for i in self.theta_idxs]

        idxs = np.argsort(self.theta_idxs)
        self.theta_idxs = [self.theta_idxs[i] for i in idxs]
        self.theta_symbols = [self.theta_symbols[i] for i in idxs]
        self.theta_keys = [self.theta_keys[i] for i in idxs]