mcmc_based_searches.py 77.8 KB
Newer Older
1001
1002
1003
        """ Writes a .par of the best-fit params with an estimated std """
        logging.info('Writing {}/{}.par using the {} method'.format(
            self.outdir, self.label, method))
1004
1005
1006
1007

        median_std_d = self.get_median_stds()
        max_twoF_d, max_twoF = self.get_max_twoF()

Gregory Ashton's avatar
Gregory Ashton committed
1008
        logging.info('Writing par file with max twoF = {}'.format(max_twoF))
1009
1010
1011
        filename = '{}/{}.par'.format(self.outdir, self.label)
        with open(filename, 'w+') as f:
            f.write('MaxtwoF = {}\n'.format(max_twoF))
Gregory Ashton's avatar
Gregory Ashton committed
1012
            f.write('tref = {}\n'.format(self.tref))
1013
1014
            if hasattr(self, 'theta0_index'):
                f.write('theta0_index = {}\n'.format(self.theta0_idx))
1015
            if method == 'med':
1016
1017
                for key, val in median_std_d.iteritems():
                    f.write('{} = {:1.16e}\n'.format(key, val))
1018
            if method == 'twoFmax':
1019
1020
1021
1022
                for key, val in max_twoF_d.iteritems():
                    f.write('{} = {:1.16e}\n'.format(key, val))

    def print_summary(self):
Gregory Ashton's avatar
Gregory Ashton committed
1023
        max_twoFd, max_twoF = self.get_max_twoF()
1024
        median_std_d = self.get_median_stds()
Gregory Ashton's avatar
Gregory Ashton committed
1025
        logging.info('Summary:')
1026
        if hasattr(self, 'theta0_idx'):
Gregory Ashton's avatar
Gregory Ashton committed
1027
1028
            logging.info('theta0 index: {}'.format(self.theta0_idx))
        logging.info('Max twoF: {} with parameters:'.format(max_twoF))
Gregory Ashton's avatar
Gregory Ashton committed
1029
1030
        for k in np.sort(max_twoFd.keys()):
            print('  {:10s} = {:1.9e}'.format(k, max_twoFd[k]))
Gregory Ashton's avatar
Gregory Ashton committed
1031
        logging.info('Median +/- std for production values')
1032
        for k in np.sort(median_std_d.keys()):
1033
            if 'std' not in k:
Gregory Ashton's avatar
Gregory Ashton committed
1034
                logging.info('  {:10s} = {:1.9e} +/- {:1.9e}'.format(
1035
                    k, median_std_d[k], median_std_d[k+'_std']))
Gregory Ashton's avatar
Gregory Ashton committed
1036
        logging.info('\n')
1037

1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
    def CF_twoFmax(self, theta, twoFmax, ntrials):
        Fmax = twoFmax/2.0
        return (np.exp(1j*theta*twoFmax)*ntrials/2.0
                * Fmax*np.exp(-Fmax)*(1-(1+Fmax)*np.exp(-Fmax))**(ntrials-1))

    def pdf_twoFhat(self, twoFhat, nglitch, ntrials, twoFmax=100, dtwoF=0.1):
        if np.ndim(ntrials) == 0:
            ntrials = np.zeros(nglitch+1) + ntrials
        twoFmax_int = np.arange(0, twoFmax, dtwoF)
        theta_int = np.arange(-1/dtwoF, 1./dtwoF, 1./twoFmax)
        CF_twoFmax_theta = np.array(
            [[np.trapz(self.CF_twoFmax(t, twoFmax_int, ntrial), twoFmax_int)
              for t in theta_int]
             for ntrial in ntrials])
        CF_twoFhat_theta = np.prod(CF_twoFmax_theta, axis=0)
        pdf = (1/(2*np.pi)) * np.array(
            [np.trapz(np.exp(-1j*theta_int*twoFhat_val)
             * CF_twoFhat_theta, theta_int) for twoFhat_val in twoFhat])
        return pdf.real

    def p_val_twoFhat(self, twoFhat, ntrials, twoFhatmax=500, Npoints=1000):
1059
        """ Caluculate the p-value for the given twoFhat in Gaussian noise
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078

        Parameters
        ----------
        twoFhat: float
            The observed twoFhat value
        ntrials: int, array of len Nglitch+1
            The number of trials for each glitch+1
        """
        twoFhats = np.linspace(twoFhat, twoFhatmax, Npoints)
        pdf = self.pdf_twoFhat(twoFhats, self.nglitch, ntrials)
        return np.trapz(pdf, twoFhats)

    def get_p_value(self, delta_F0, time_trials=0):
        """ Get's the p-value for the maximum twoFhat value """
        d, max_twoF = self.get_max_twoF()
        if self.nglitch == 1:
            tglitches = [d['tglitch']]
        else:
            tglitches = [d['tglitch_{}'.format(i)] for i in range(self.nglitch)]
1079
        tboundaries = [self.minStartTime] + tglitches + [self.maxStartTime]
1080
        deltaTs = np.diff(tboundaries)
1081
1082
        ntrials = [time_trials + delta_F0 * dT for dT in deltaTs]
        p_val = self.p_val_twoFhat(max_twoF, ntrials)
1083
        print('p-value = {}'.format(p_val))
1084
1085
        return p_val

1086
    def get_evidence(self):
1087
1088
1089
1090
1091
1092
        fburnin = float(self.nsteps[-2])/np.sum(self.nsteps[-2:])
        lnev, lnev_err = self.sampler.thermodynamic_integration_log_evidence(
            fburnin=fburnin)

        log10evidence = lnev/np.log(10)
        log10evidence_err = lnev_err/np.log(10)
1093
1094
1095
1096
        return log10evidence, log10evidence_err

    def compute_evidence_long(self):
        """ Computes the evidence/marginal likelihood for the model """
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
        betas = self.betas
        alllnlikes = self.sampler.lnlikelihood[:, :, self.nsteps[-2]:]
        mean_lnlikes = np.mean(np.mean(alllnlikes, axis=1), axis=1)

        mean_lnlikes = mean_lnlikes[::-1]
        betas = betas[::-1]

        fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(6, 8))

        if any(np.isinf(mean_lnlikes)):
            print("WARNING mean_lnlikes contains inf: recalculating without"
                  " the {} infs".format(len(betas[np.isinf(mean_lnlikes)])))
            idxs = np.isinf(mean_lnlikes)
            mean_lnlikes = mean_lnlikes[~idxs]
            betas = betas[~idxs]
            log10evidence = np.trapz(mean_lnlikes, betas)/np.log(10)
            z1 = np.trapz(mean_lnlikes, betas)
            z2 = np.trapz(mean_lnlikes[::-1][::2][::-1],
                          betas[::-1][::2][::-1])
            log10evidence_err = np.abs(z1 - z2) / np.log(10)

        ax1.semilogx(betas, mean_lnlikes, "-o")
        ax1.set_xlabel(r"$\beta$")
        ax1.set_ylabel(r"$\langle \log(\mathcal{L}) \rangle$")
        print("log10 evidence for {} = {} +/- {}".format(
              self.label, log10evidence, log10evidence_err))
        min_betas = []
        evidence = []
        for i in range(len(betas)/2):
            min_betas.append(betas[i])
            lnZ = np.trapz(mean_lnlikes[i:], betas[i:])
            evidence.append(lnZ/np.log(10))

        ax2.semilogx(min_betas, evidence, "-o")
        ax2.set_ylabel(r"$\int_{\beta_{\textrm{Min}}}^{\beta=1}" +
                       r"\langle \log(\mathcal{L})\rangle d\beta$", size=16)
        ax2.set_xlabel(r"$\beta_{\textrm{min}}$")
        plt.tight_layout()
        fig.savefig("{}/{}_beta_lnl.png".format(self.outdir, self.label))

1137

Gregory Ashton's avatar
Gregory Ashton committed
1138
1139
class MCMCGlitchSearch(MCMCSearch):
    """ MCMC search using the SemiCoherentGlitchSearch """
Gregory Ashton's avatar
Gregory Ashton committed
1140
    @helper_functions.initializer
1141
    def __init__(self, label, outdir, sftfilepath, theta_prior, tref,
Gregory Ashton's avatar
Gregory Ashton committed
1142
                 minStartTime, maxStartTime, nglitch=1, nsteps=[100, 100],
1143
1144
                 nwalkers=100, ntemps=1, log10temperature_min=-5,
                 theta_initial=None, scatter_val=1e-10, dtglitchmin=1*86400,
1145
                 theta0_idx=0, detectors=None, BSGL=False, minCoverFreq=None,
1146
                 maxCoverFreq=None, earth_ephem=None, sun_ephem=None):
Gregory Ashton's avatar
Gregory Ashton committed
1147
1148
        """
        Parameters
Gregory Ashton's avatar
Gregory Ashton committed
1149
        ----------
Gregory Ashton's avatar
Gregory Ashton committed
1150
1151
        label, outdir: str
            A label and directory to read/write data from/to
Gregory Ashton's avatar
Gregory Ashton committed
1152
        sftfilepath: str
1153
            File patern to match SFTs
Gregory Ashton's avatar
Gregory Ashton committed
1154
1155
1156
1157
1158
1159
1160
1161
        theta_prior: dict
            Dictionary of priors and fixed values for the search parameters.
            For each parameters (key of the dict), if it is to be held fixed
            the value should be the constant float, if it is be searched, the
            value should be a dictionary of the prior.
        theta_initial: dict, array, (None)
            Either a dictionary of distribution about which to distribute the
            initial walkers about, an array (from which the walkers will be
Gregory Ashton's avatar
Gregory Ashton committed
1162
            scattered by scatter_val), or None in which case the prior is used.
1163
1164
1165
1166
        scatter_val, float or ndim array
            Size of scatter to use about the initialisation step, if given as
            an array it must be of length ndim and the order is given by
            theta_keys
Gregory Ashton's avatar
Gregory Ashton committed
1167
1168
        nglitch: int
            The number of glitches to allow
1169
        tref, minStartTime, maxStartTime: int
Gregory Ashton's avatar
Gregory Ashton committed
1170
1171
1172
1173
1174
1175
1176
1177
1178
            GPS seconds of the reference time, start time and end time
        nsteps: list (m,)
            List specifying the number of steps to take, the last two entries
            give the nburn and nprod of the 'production' run, all entries
            before are for iterative initialisation steps (usually just one)
            e.g. [1000, 1000, 500].
        dtglitchmin: int
            The minimum duration (in seconds) of a segment between two glitches
            or a glitch and the start/end of the data
1179
1180
1181
1182
1183
1184
        nwalkers, ntemps: int,
            The number of walkers and temperates to use in the parallel
            tempered PTSampler.
        log10temperature_min float < 0
            The  log_10(tmin) value, the set of betas passed to PTSampler are
            generated from np.logspace(0, log10temperature_min, ntemps).
1185
1186
1187
1188
        theta0_idx, int
            Index (zero-based) of which segment the theta refers to - uyseful
            if providing a tight prior on theta to allow the signal to jump
            too theta (and not just from)
1189
        detectors: str
1190
1191
            Two character reference to the data to use, specify None for no
            contraint.
Gregory Ashton's avatar
Gregory Ashton committed
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
        minCoverFreq, maxCoverFreq: float
            Minimum and maximum instantaneous frequency which will be covered
            over the SFT time span as passed to CreateFstatInput
        earth_ephem, sun_ephem: str
            Paths of the two files containing positions of Earth and Sun,
            respectively at evenly spaced times, as passed to CreateFstatInput
            If None defaults defined in BaseSearchClass will be used

        """

Gregory Ashton's avatar
Gregory Ashton committed
1202
1203
        if os.path.isdir(outdir) is False:
            os.mkdir(outdir)
1204
        self.add_log_file()
Gregory Ashton's avatar
Gregory Ashton committed
1205
1206
        logging.info(('Set-up MCMC glitch search with {} glitches for model {}'
                      ' on data {}').format(self.nglitch, self.label,
1207
                                            self.sftfilepath))
Gregory Ashton's avatar
Gregory Ashton committed
1208
1209
1210
        self.pickle_path = '{}/{}_saved_data.p'.format(self.outdir, self.label)
        self.unpack_input_theta()
        self.ndim = len(self.theta_keys)
1211
1212
1213
1214
        if self.log10temperature_min:
            self.betas = np.logspace(0, self.log10temperature_min, self.ntemps)
        else:
            self.betas = None
Gregory Ashton's avatar
Gregory Ashton committed
1215
1216
1217
1218
1219
1220
1221
1222
1223
        if earth_ephem is None:
            self.earth_ephem = self.earth_ephem_default
        if sun_ephem is None:
            self.sun_ephem = self.sun_ephem_default

        if args.clean and os.path.isfile(self.pickle_path):
            os.rename(self.pickle_path, self.pickle_path+".old")

        self.old_data_is_okay_to_use = self.check_old_data_is_okay_to_use()
1224
        self.log_input()
Gregory Ashton's avatar
Gregory Ashton committed
1225
1226
1227
1228

    def inititate_search_object(self):
        logging.info('Setting up search object')
        self.search = SemiCoherentGlitchSearch(
1229
            label=self.label, outdir=self.outdir, sftfilepath=self.sftfilepath,
1230
1231
            tref=self.tref, minStartTime=self.minStartTime,
            maxStartTime=self.maxStartTime, minCoverFreq=self.minCoverFreq,
Gregory Ashton's avatar
Gregory Ashton committed
1232
            maxCoverFreq=self.maxCoverFreq, earth_ephem=self.earth_ephem,
1233
            sun_ephem=self.sun_ephem, detectors=self.detectors, BSGL=self.BSGL,
1234
            nglitch=self.nglitch, theta0_idx=self.theta0_idx)
Gregory Ashton's avatar
Gregory Ashton committed
1235
1236
1237

    def logp(self, theta_vals, theta_prior, theta_keys, search):
        if self.nglitch > 1:
1238
1239
            ts = ([self.minStartTime] + list(theta_vals[-self.nglitch:])
                  + [self.maxStartTime])
Gregory Ashton's avatar
Gregory Ashton committed
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
            if np.array_equal(ts, np.sort(ts)) is False:
                return -np.inf
            if any(np.diff(ts) < self.dtglitchmin):
                return -np.inf

        H = [self.generic_lnprior(**theta_prior[key])(p) for p, key in
             zip(theta_vals, theta_keys)]
        return np.sum(H)

    def logl(self, theta, search):
Gregory Ashton's avatar
Gregory Ashton committed
1250
        if self.nglitch > 1:
1251
1252
            ts = ([self.minStartTime] + list(theta_vals[-self.nglitch:])
                  + [self.maxStartTime])
Gregory Ashton's avatar
Gregory Ashton committed
1253
1254
1255
            if np.array_equal(ts, np.sort(ts)) is False:
                return -np.inf

Gregory Ashton's avatar
Gregory Ashton committed
1256
1257
1258
1259
1260
1261
1262
1263
1264
        for j, theta_i in enumerate(self.theta_idxs):
            self.fixed_theta[theta_i] = theta[j]
        FS = search.compute_nglitch_fstat(*self.fixed_theta)
        return FS

    def unpack_input_theta(self):
        glitch_keys = ['delta_F0', 'delta_F1', 'tglitch']
        full_glitch_keys = list(np.array(
            [[gk]*self.nglitch for gk in glitch_keys]).flatten())
1265
1266
1267
1268

        if 'tglitch_0' in self.theta_prior:
            full_glitch_keys[-self.nglitch:] = [
                'tglitch_{}'.format(i) for i in range(self.nglitch)]
1269
1270
1271
1272
            full_glitch_keys[-2*self.nglitch:-1*self.nglitch] = [
                'delta_F1_{}'.format(i) for i in range(self.nglitch)]
            full_glitch_keys[-4*self.nglitch:-2*self.nglitch] = [
                'delta_F0_{}'.format(i) for i in range(self.nglitch)]
Gregory Ashton's avatar
Gregory Ashton committed
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
        full_theta_keys = ['F0', 'F1', 'F2', 'Alpha', 'Delta']+full_glitch_keys
        full_theta_keys_copy = copy.copy(full_theta_keys)

        glitch_symbols = ['$\delta f$', '$\delta \dot{f}$', r'$t_{glitch}$']
        full_glitch_symbols = list(np.array(
            [[gs]*self.nglitch for gs in glitch_symbols]).flatten())
        full_theta_symbols = (['$f$', '$\dot{f}$', '$\ddot{f}$', r'$\alpha$',
                               r'$\delta$'] + full_glitch_symbols)
        self.theta_keys = []
        fixed_theta_dict = {}
        for key, val in self.theta_prior.iteritems():
            if type(val) is dict:
                fixed_theta_dict[key] = 0
                if key in glitch_keys:
                    for i in range(self.nglitch):
                        self.theta_keys.append(key)
                else:
                    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))
            if key in glitch_keys:
                for i in range(self.nglitch):
                    full_theta_keys_copy.pop(full_theta_keys_copy.index(key))
            else:
                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]

        # Correct for number of glitches in the idxs
        self.theta_idxs = np.array(self.theta_idxs)
        while np.sum(self.theta_idxs[:-1] == self.theta_idxs[1:]) > 0:
            for i, idx in enumerate(self.theta_idxs):
                if idx in self.theta_idxs[:i]:
                    self.theta_idxs[i] += 1

1324
1325
1326
1327
1328
1329
1330
1331
    def get_save_data_dictionary(self):
        d = dict(nsteps=self.nsteps, nwalkers=self.nwalkers,
                 ntemps=self.ntemps, theta_keys=self.theta_keys,
                 theta_prior=self.theta_prior, scatter_val=self.scatter_val,
                 log10temperature_min=self.log10temperature_min,
                 theta0_idx=self.theta0_idx, BSGL=self.BSGL)
        return d

Gregory Ashton's avatar
Gregory Ashton committed
1332
1333
1334
1335
1336
1337
1338
    def apply_corrections_to_p0(self, p0):
        p0 = np.array(p0)
        if self.nglitch > 1:
            p0[:, :, -self.nglitch:] = np.sort(p0[:, :, -self.nglitch:],
                                               axis=2)
        return p0

Gregory Ashton's avatar
Gregory Ashton committed
1339
1340
1341
1342
1343
1344
1345
1346
    def plot_cumulative_max(self):

        fig, ax = plt.subplots()
        d, maxtwoF = self.get_max_twoF()
        for key, val in self.theta_prior.iteritems():
            if key not in d:
                d[key] = val

1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
        if self.nglitch > 1:
            delta_F0s = [d['delta_F0_{}'.format(i)] for i in
                         range(self.nglitch)]
            delta_F0s.insert(self.theta0_idx, 0)
            delta_F0s = np.array(delta_F0s)
            delta_F0s[:self.theta0_idx] *= -1
            tglitches = [d['tglitch_{}'.format(i)] for i in
                         range(self.nglitch)]
        elif self.nglitch == 1:
            delta_F0s = [d['delta_F0']]
            delta_F0s.insert(self.theta0_idx, 0)
            delta_F0s = np.array(delta_F0s)
            delta_F0s[:self.theta0_idx] *= -1
            tglitches = [d['tglitch']]
Gregory Ashton's avatar
Gregory Ashton committed
1361

1362
        tboundaries = [self.minStartTime] + tglitches + [self.maxStartTime]
Gregory Ashton's avatar
Gregory Ashton committed
1363
1364

        for j in range(self.nglitch+1):
1365
1366
            ts = tboundaries[j]
            te = tboundaries[j+1]
Gregory Ashton's avatar
Gregory Ashton committed
1367
1368
1369
1370
1371
1372
1373
1374
            if (te - ts)/86400 < 5:
                logging.info('Period too short to perform cumulative search')
                continue
            if j < self.theta0_idx:
                summed_deltaF0 = np.sum(delta_F0s[j:self.theta0_idx])
                F0_j = d['F0'] - summed_deltaF0
                taus, twoFs = self.search.calculate_twoF_cumulative(
                    F0_j, F1=d['F1'], F2=d['F2'], Alpha=d['Alpha'],
1375
                    Delta=d['Delta'], tstart=ts, tend=te)
Gregory Ashton's avatar
Gregory Ashton committed
1376
1377
1378
1379
1380
1381

            elif j >= self.theta0_idx:
                summed_deltaF0 = np.sum(delta_F0s[self.theta0_idx:j+1])
                F0_j = d['F0'] + summed_deltaF0
                taus, twoFs = self.search.calculate_twoF_cumulative(
                    F0_j, F1=d['F1'], F2=d['F2'], Alpha=d['Alpha'],
1382
                    Delta=d['Delta'], tstart=ts, tend=te)
Gregory Ashton's avatar
Gregory Ashton committed
1383
1384
1385
1386
1387
            ax.plot(ts+taus, twoFs)

        ax.set_xlabel('GPS time')
        fig.savefig('{}/{}_twoFcumulative.png'.format(self.outdir, self.label))

Gregory Ashton's avatar
Gregory Ashton committed
1388

1389
1390
class MCMCSemiCoherentSearch(MCMCSearch):
    """ MCMC search for a signal using the semi-coherent ComputeFstat """
Gregory Ashton's avatar
Gregory Ashton committed
1391
    @helper_functions.initializer
1392
    def __init__(self, label, outdir, theta_prior, tref, sftfilepath=None,
1393
1394
                 nsegs=None, nsteps=[100, 100, 100], nwalkers=100, binary=False,
                 ntemps=1, log10temperature_min=-5, theta_initial=None,
1395
                 scatter_val=1e-10, detectors=None, BSGL=False,
1396
                 minStartTime=None, maxStartTime=None, minCoverFreq=None,
1397
                 maxCoverFreq=None, earth_ephem=None, sun_ephem=None,
1398
                 injectSources=None, assumeSqrtSX=None):
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
        """

        """

        if os.path.isdir(outdir) is False:
            os.mkdir(outdir)
        self.add_log_file()
        logging.info(('Set-up MCMC semi-coherent search for model {} on data'
                      '{}').format(
            self.label, self.sftfilepath))
        self.pickle_path = '{}/{}_saved_data.p'.format(self.outdir, self.label)
        self.unpack_input_theta()
        self.ndim = len(self.theta_keys)
        if self.log10temperature_min:
            self.betas = np.logspace(0, self.log10temperature_min, self.ntemps)
        else:
            self.betas = None
        if earth_ephem is None:
            self.earth_ephem = self.earth_ephem_default
        if sun_ephem is None:
            self.sun_ephem = self.sun_ephem_default

        if args.clean and os.path.isfile(self.pickle_path):
            os.rename(self.pickle_path, self.pickle_path+".old")

        self.log_input()

    def inititate_search_object(self):
        logging.info('Setting up search object')
        self.search = SemiCoherentSearch(
            label=self.label, outdir=self.outdir, tref=self.tref,
            nsegs=self.nsegs, sftfilepath=self.sftfilepath, binary=self.binary,
            BSGL=self.BSGL, minStartTime=self.minStartTime,
            maxStartTime=self.maxStartTime, minCoverFreq=self.minCoverFreq,
1433
            maxCoverFreq=self.maxCoverFreq, detectors=self.detectors,
1434
            earth_ephem=self.earth_ephem, sun_ephem=self.sun_ephem,
1435
            injectSources=self.injectSources, assumeSqrtSX=self.assumeSqrtSX)
1436
1437
1438
1439
1440
1441
1442
1443
1444

    def logp(self, theta_vals, theta_prior, theta_keys, search):
        H = [self.generic_lnprior(**theta_prior[key])(p) for p, key in
             zip(theta_vals, theta_keys)]
        return np.sum(H)

    def logl(self, theta, search):
        for j, theta_i in enumerate(self.theta_idxs):
            self.fixed_theta[theta_i] = theta[j]
Gregory Ashton's avatar
Gregory Ashton committed
1445
1446
        FS = search.run_semi_coherent_computefstatistic_single_point(
            *self.fixed_theta)
1447
1448
1449
        return FS


Gregory Ashton's avatar
Gregory Ashton committed
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
class MCMCFollowUpSearch(MCMCSemiCoherentSearch):
    """ A follow up procudure increasing the coherence time in a zoom """
    def get_save_data_dictionary(self):
        d = dict(nwalkers=self.nwalkers, ntemps=self.ntemps,
                 theta_keys=self.theta_keys, theta_prior=self.theta_prior,
                 scatter_val=self.scatter_val,
                 log10temperature_min=self.log10temperature_min,
                 BSGL=self.BSGL, run_setup=self.run_setup)
        return d

Gregory Ashton's avatar
Gregory Ashton committed
1460
1461
1462
1463
    def update_search_object(self):
        logging.info('Update search object')
        self.search.init_computefstatistic_single_point()

1464
1465
1466
1467
1468
1469
1470
1471
    def get_width_from_prior(self, prior, key):
        if prior[key]['type'] == 'unif':
            return prior[key]['upper'] - prior[key]['lower']

    def get_mid_from_prior(self, prior, key):
        if prior[key]['type'] == 'unif':
            return .5*(prior[key]['upper'] + prior[key]['lower'])

1472
    def init_V_estimate_parameters(self):
1473
1474
1475
1476
        if 'Alpha' in self.theta_keys:
            DeltaAlpha = self.get_width_from_prior(self.theta_prior, 'Alpha')
            DeltaDelta = self.get_width_from_prior(self.theta_prior, 'Delta')
            DeltaMid = self.get_mid_from_prior(self.theta_prior, 'Delta')
1477
            DeltaOmega = np.sin(np.pi/2 - DeltaMid) * DeltaDelta * DeltaAlpha
Gregory Ashton's avatar
Gregory Ashton committed
1478
            logging.info('Search over Alpha and Delta')
1479
        else:
Gregory Ashton's avatar
Gregory Ashton committed
1480
            logging.info('No sky search requested')
1481
            DeltaOmega = 0
1482
1483
1484
        if 'F0' in self.theta_keys:
            DeltaF0 = self.get_width_from_prior(self.theta_prior, 'F0')
        else:
Gregory Ashton's avatar
Gregory Ashton committed
1485
            raise ValueError("You aren't searching over F0?")
1486
        DeltaFs = [DeltaF0]
1487
1488
        if 'F1' in self.theta_keys:
            DeltaF1 = self.get_width_from_prior(self.theta_prior, 'F1')
1489
            DeltaFs.append(DeltaF1)
Gregory Ashton's avatar
Gregory Ashton committed
1490
1491
1492
1493
1494
            if 'F2' in self.theta_keys:
                DeltaF2 = self.get_width_from_prior(self.theta_prior, 'F2')
                DeltaFs.append(DeltaF2)
        logging.info('Searching over Frequency and {} spin-down components'
                     .format(len(DeltaFs)-1))
1495

1496
1497
1498
1499
1500
        if type(self.theta_prior['F0']) == dict:
            fiducial_freq = self.get_mid_from_prior(self.theta_prior, 'F0')
        else:
            fiducial_freq = self.theta_prior['F0']

1501
        return fiducial_freq, DeltaOmega, DeltaFs
1502

1503
1504
1505
1506
1507
    def read_setup_input_file(self, run_setup_input_file):
        with open(run_setup_input_file, 'r+') as f:
            d = pickle.load(f)
        return d

1508
    def write_setup_input_file(self, run_setup_input_file, R, Nsegs0,
1509
                               nsegs_vals, V_vals, DeltaOmega, DeltaFs):
1510
        d = dict(R=R, Nsegs0=Nsegs0, nsegs_vals=nsegs_vals, V_vals=V_vals,
1511
                 DeltaOmega=DeltaOmega, DeltaFs=DeltaFs)
1512
1513
1514
        with open(run_setup_input_file, 'w+') as f:
            pickle.dump(d, f)

1515
1516
    def check_old_run_setup(self, old_setup, **kwargs):
        try:
1517
1518
            truths = [val == old_setup[key] for key, val in kwargs.iteritems()]
            return all(truths)
1519
1520
1521
        except KeyError:
            return False

1522
1523
1524
1525
1526
1527
1528
    def init_run_setup(self, run_setup=None, R=10, Nsegs0=None, log_table=True,
                       gen_tex_table=True):

        if run_setup is None and Nsegs0 is None:
            raise ValueError(
                'You must either specify the run_setup, or Nsegs0 from which '
                'the optimial run_setup given R can be estimated')
1529
        fiducial_freq, DeltaOmega, DeltaFs = self.init_V_estimate_parameters()
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
        if run_setup is None:
            logging.info('No run_setup provided')

            run_setup_input_file = '{}/{}_run_setup.p'.format(
                self.outdir, self.label)

            if os.path.isfile(run_setup_input_file):
                logging.info('Checking old setup input file {}'.format(
                    run_setup_input_file))
                old_setup = self.read_setup_input_file(run_setup_input_file)
1540
1541
                if self.check_old_run_setup(old_setup, R=R,
                                            Nsegs0=Nsegs0,
1542
1543
                                            DeltaOmega=DeltaOmega,
                                            DeltaFs=DeltaFs):
1544
1545
                    logging.info('Using old setup with R={}, Nsegs0={}'.format(
                        R, Nsegs0))
1546
                    nsegs_vals = old_setup['nsegs_vals']
1547
                    V_vals = old_setup['V_vals']
1548
                    generate_setup = False
1549
                else:
1550
                    logging.info('Old setup does not match requested R, Nsegs0')
1551
1552
1553
                    generate_setup = True
            else:
                generate_setup = True
1554

1555
            if generate_setup:
1556
                nsegs_vals, V_vals = get_optimal_setup(
1557
                    R, Nsegs0, self.tref, self.minStartTime,
1558
1559
1560
                    self.maxStartTime, DeltaOmega, DeltaFs, fiducial_freq,
                    self.search.detector_names, self.earth_ephem,
                    self.sun_ephem)
1561
                self.write_setup_input_file(run_setup_input_file, R, Nsegs0,
1562
1563
                                            nsegs_vals, V_vals, DeltaOmega,
                                            DeltaFs)
1564
1565
1566
1567
1568

            run_setup = [((self.nsteps[0], 0),  nsegs, False)
                         for nsegs in nsegs_vals[:-1]]
            run_setup.append(
                ((self.nsteps[0], self.nsteps[1]), nsegs_vals[-1], False))
1569
1570
1571
1572

        else:
            logging.info('Calculating the number of templates for this setup')
            V_vals = []
1573
1574
1575
1576
1577
1578
1579
1580
            for i, rs in enumerate(run_setup):
                rs = list(rs)
                if len(rs) == 2:
                    rs.append(False)
                if np.shape(rs[0]) == ():
                    rs[0] = (rs[0], 0)
                run_setup[i] = rs

1581
1582
1583
                if args.no_template_counting:
                    V_vals.append([1, 1, 1])
                else:
1584
1585
1586
1587
1588
1589
                    V, Vsky, Vpe = get_V_estimate(
                        rs[1], self.tref, self.minStartTime, self.maxStartTime,
                        DeltaOmega, DeltaFs, fiducial_freq,
                        self.search.detector_names, self.earth_ephem,
                        self.sun_ephem)
                    V_vals.append([V, Vsky, Vpe])
1590
1591

        if log_table:
1592
            logging.info('Using run-setup as follows:')
1593
            logging.info('Stage | nburn | nprod | nsegs | Tcoh d | resetp0 |'
1594
                         ' V = Vsky x Vpe')
1595
            for i, rs in enumerate(run_setup):
1596
                Tcoh = (self.maxStartTime - self.minStartTime) / rs[1] / 86400
1597
                if V_vals[i] is None:
1598
                    vtext = 'N/A'
1599
                else:
1600
                    vtext = '{:1.0e} = {:1.0e} x {:1.0e}'.format(
1601
                            V_vals[i][0], V_vals[i][1], V_vals[i][2])
1602
                logging.info('{} | {} | {} | {} | {} | {} | {}'.format(
1603
1604
                    str(i).ljust(5), str(rs[0][0]).ljust(5),
                    str(rs[0][1]).ljust(5), str(rs[1]).ljust(5),
1605
1606
                    '{:6.1f}'.format(Tcoh), str(rs[2]).ljust(7),
                    vtext))
1607
1608
1609

        if gen_tex_table:
            filename = '{}/{}_run_setup.tex'.format(self.outdir, self.label)
1610
1611
1612
1613
1614
1615
1616
            if DeltaOmega > 0:
                with open(filename, 'w+') as f:
                    f.write(r'\begin{tabular}{c|cccccc}' + '\n')
                    f.write(r'Stage & $\Nseg$ & $\Tcoh^{\rm days}$ &'
                            r'$\Nsteps$ & $\V$ & $\Vsky$ & $\Vpe$ \\ \hline'
                            '\n')
                    for i, rs in enumerate(run_setup):
1617
1618
                        Tcoh = float(
                            self.maxStartTime - self.minStartTime)/rs[1]/86400
1619
                        line = r'{} & {} & {} & {} & {} & {} & {} \\' + '\n'
1620
1621
1622
1623
                        if V_vals[i][0] is None:
                            V = Vsky = Vpe = 'N/A'
                        else:
                            V, Vsky, Vpe = V_vals[i]
1624
1625
1626
1627
                        if rs[0][-1] == 0:
                            nsteps = rs[0][0]
                        else:
                            nsteps = '{},{}'.format(*rs[0])
1628
                        line = line.format(i, rs[1], '{:1.1f}'.format(Tcoh),
Gregory Ashton's avatar
Gregory Ashton committed
1629
1630
1631
1632
                                           nsteps,
                                           helper_functions.texify_float(V),
                                           helper_functions.texify_float(Vsky),
                                           helper_functions.texify_float(Vpe))
1633
1634
1635
1636
1637
1638
1639
1640
1641
                        f.write(line)
                    f.write(r'\end{tabular}' + '\n')
            else:
                with open(filename, 'w+') as f:
                    f.write(r'\begin{tabular}{c|cccc}' + '\n')
                    f.write(r'Stage & $\Nseg$ & $\Tcoh^{\rm days}$ &'
                            r'$\Nsteps$ & $\Vpe$ \\ \hline'
                            '\n')
                    for i, rs in enumerate(run_setup):
1642
1643
                        Tcoh = float(
                            self.maxStartTime - self.minStartTime)/rs[1]/86400
1644
                        line = r'{} & {} & {} & {} & {} \\' + '\n'
1645
1646
1647
1648
                        if V_vals[i] is None:
                            V = Vsky = Vpe = 'N/A'
                        else:
                            V, Vsky, Vpe = V_vals[i]
1649
1650
1651
1652
                        if rs[0][-1] == 0:
                            nsteps = rs[0][0]
                        else:
                            nsteps = '{},{}'.format(*rs[0])
1653
                        line = line.format(i, rs[1], '{:1.1f}'.format(Tcoh),
Gregory Ashton's avatar
Gregory Ashton committed
1654
1655
                                           nsteps,
                                           helper_functions.texify_float(Vpe))
1656
1657
                        f.write(line)
                    f.write(r'\end{tabular}' + '\n')
1658
1659
1660
1661

        if args.setup_only:
            logging.info("Exit as requested by setup_only flag")
            sys.exit()
1662
1663
1664
        else:
            return run_setup

1665
    def run(self, run_setup=None, proposal_scale_factor=2, R=10, Nsegs0=None,
1666
1667
            create_plots=True, log_table=True, gen_tex_table=True, fig=None,
            axes=None, return_fig=False, **kwargs):
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
        """ Run the follow-up with the given run_setup

        Parameters
        ----------
        run_setup: list of tuples

        """

        self.nsegs = 1
        self.inititate_search_object()
1678
        run_setup = self.init_run_setup(
1679
            run_setup, R=R, Nsegs0=Nsegs0, log_table=log_table,
1680
            gen_tex_table=gen_tex_table)
1681
        self.run_setup = run_setup
Gregory Ashton's avatar
Gregory Ashton committed
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691

        self.old_data_is_okay_to_use = self.check_old_data_is_okay_to_use()
        if self.old_data_is_okay_to_use is True:
            logging.warning('Using saved data from {}'.format(
                self.pickle_path))
            d = self.get_saved_data()
            self.sampler = d['sampler']
            self.samples = d['samples']
            self.lnprobs = d['lnprobs']
            self.lnlikes = d['lnlikes']
1692
            self.nsegs = run_setup[-1][1]
Gregory Ashton's avatar
Gregory Ashton committed
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
            return

        nsteps_total = 0
        for j, ((nburn, nprod), nseg, reset_p0) in enumerate(run_setup):
            if j == 0:
                p0 = self.generate_initial_p0()
                p0 = self.apply_corrections_to_p0(p0)
            elif reset_p0:
                p0 = self.get_new_p0(sampler)
                p0 = self.apply_corrections_to_p0(p0)
                # self.check_initial_points(p0)
            else:
                p0 = sampler.chain[:, :, -1, :]

            self.nsegs = nseg
1708
            self.search.nsegs = nseg
Gregory Ashton's avatar
Gregory Ashton committed
1709
            self.update_search_object()
1710
            self.search.init_semicoherent_parameters()
Gregory Ashton's avatar
Gregory Ashton committed
1711
1712
1713
1714
1715
1716
            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)

1717
            Tcoh = (self.maxStartTime-self.minStartTime)/nseg/86400.
Gregory Ashton's avatar
Gregory Ashton committed
1718
1719
            logging.info(('Running {}/{} with {} steps and {} nsegs '
                          '(Tcoh={:1.2f} days)').format(
1720
                j+1, len(run_setup), (nburn, nprod), nseg, Tcoh))
1721
            sampler = self.run_sampler(sampler, p0, nburn=nburn, nprod=nprod)
Gregory Ashton's avatar
Gregory Ashton committed
1722
1723
1724
1725
1726
            logging.info("Mean acceptance fraction: {}"
                         .format(np.mean(sampler.acceptance_fraction, axis=1)))
            if self.ntemps > 1:
                logging.info("Tswap acceptance fraction: {}"
                             .format(sampler.tswap_acceptance_fraction))
1727
1728
            logging.info('Max detection statistic of run was {}'.format(
                np.max(sampler.lnlikelihood)))
Gregory Ashton's avatar
Gregory Ashton committed
1729

1730
1731
1732
            if create_plots:
                fig, axes = self.plot_walkers(
                    sampler, symbols=self.theta_symbols, fig=fig, axes=axes,
1733
                    nprod=nprod, xoffset=nsteps_total, **kwargs)
1734
1735
                for ax in axes[:self.ndim]:
                    ax.axvline(nsteps_total, color='k', ls='--', lw=0.25)
Gregory Ashton's avatar
Gregory Ashton committed
1736

1737
            nsteps_total += nburn+nprod
Gregory Ashton's avatar
Gregory Ashton committed
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747

        samples = sampler.chain[0, :, nburn:, :].reshape((-1, self.ndim))
        lnprobs = sampler.lnprobability[0, :, nburn:].reshape((-1))
        lnlikes = sampler.lnlikelihood[0, :, nburn:].reshape((-1))
        self.sampler = sampler
        self.samples = samples
        self.lnprobs = lnprobs
        self.lnlikes = lnlikes
        self.save_data(sampler, samples, lnprobs, lnlikes)

1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
        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(
                    self.outdir, self.label), dpi=200)

Gregory Ashton's avatar
Gregory Ashton committed
1759

Gregory Ashton's avatar
Gregory Ashton committed
1760
1761
1762
1763
1764
1765
1766
1767
1768
class MCMCTransientSearch(MCMCSearch):
    """ MCMC search for a transient signal using the ComputeFstat """

    def inititate_search_object(self):
        logging.info('Setting up search object')
        self.search = ComputeFstat(
            tref=self.tref, sftfilepath=self.sftfilepath,
            minCoverFreq=self.minCoverFreq, maxCoverFreq=self.maxCoverFreq,
            earth_ephem=self.earth_ephem, sun_ephem=self.sun_ephem,
1769
            detectors=self.detectors, transient=True,
Gregory Ashton's avatar
Gregory Ashton committed
1770
            minStartTime=self.minStartTime, maxStartTime=self.maxStartTime,
Gregory Ashton's avatar
Gregory Ashton committed
1771
            BSGL=self.BSGL, binary=self.binary)
Gregory Ashton's avatar
Gregory Ashton committed
1772
1773
1774
1775

    def logl(self, theta, search):
        for j, theta_i in enumerate(self.theta_idxs):
            self.fixed_theta[theta_i] = theta[j]
1776
1777
1778
        in_theta = copy.copy(self.fixed_theta)
        in_theta[1] = in_theta[0] + in_theta[1]
        if in_theta[1] > self.maxStartTime:
1779
            return -np.inf
1780
        FS = search.run_computefstatistic_single_point(*in_theta)
Gregory Ashton's avatar
Gregory Ashton committed
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
        return FS

    def unpack_input_theta(self):
        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]


1828
For faster browsing, not all history is shown. View entire blame