optimal_setup_functions.py 8.55 KB
Newer Older
Gregory Ashton's avatar
Gregory Ashton committed
1
2
"""

3
Provides functions to aid in calculating the optimal setup for zoom follow up
Gregory Ashton's avatar
Gregory Ashton committed
4
5

"""
6

Gregory Ashton's avatar
Gregory Ashton committed
7
8
9
10
11
12

import logging
import numpy as np
import scipy.optimize
import lal
import lalpulsar
13
import pyfstat.helper_functions as helper_functions
Gregory Ashton's avatar
Gregory Ashton committed
14
15
16


def get_optimal_setup(
17
18
    NstarMax, Nsegs0, tref, minStartTime, maxStartTime, prior, detector_names
):
19
20
    """ Using an optimisation step, calculate the optimal setup ladder

21
22
23
24
    The details of the methods are described in Sec Va of arXiv:1802.05450.
    Here we provide implementation details. All equation numbers refer to
    arXiv:1802.05450.

25
26
    Parameters
    ----------
Gregory Ashton's avatar
Gregory Ashton committed
27
    NstarMax : float
28
29
30
31
32
33
        The ratio of the size at the old coherence time to the new coherence
        time for each step, see Eq. (31). Larger values allow a more rapid
        "zoom" of the search space at the cost of convergence. Smaller values
        are more conservative at the cost of additional computing time. The
        exact choice should be optimized for the problem in hand, but values
        of 100-1000 are typically okay.
34
35
    Nsegs0 : int
        The number of segments for the initial step of the ladder
36
37
    tref, minStartTime, maxStartTime : int
        GPS times of the reference, start, and end time.
38
39
40
    prior : dict
        Prior dictionary, each item must either be a fixed scalar value, or
        a uniform prior.
41
42
    detector_names : list or str
        Names of the detectors to use
43
44
45
46
47
48
49
50

    Returns
    -------
    nsegs, Nstar : list
        Ladder of segment numbers and the corresponding Nstar

    """

51
52
53
    logging.info(
        "Calculating optimal setup for NstarMax={}, Nsegs0={}".format(NstarMax, Nsegs0)
    )
Gregory Ashton's avatar
Gregory Ashton committed
54

55
    Nstar_0 = get_Nstar_estimate(
56
57
58
        Nsegs0, tref, minStartTime, maxStartTime, prior, detector_names
    )
    logging.info("Stage {}, nsegs={}, Nstar={}".format(0, Nsegs0, int(Nstar_0)))
Gregory Ashton's avatar
Gregory Ashton committed
59
60

    nsegs_vals = [Nsegs0]
61
    Nstar_vals = [Nstar_0]
Gregory Ashton's avatar
Gregory Ashton committed
62
63
64
65

    i = 0
    nsegs_i = Nsegs0
    while nsegs_i > 1:
66
        nsegs_i, Nstar_i = _get_nsegs_ip1(
67
68
            nsegs_i, NstarMax, tref, minStartTime, maxStartTime, prior, detector_names
        )
Gregory Ashton's avatar
Gregory Ashton committed
69
        nsegs_vals.append(nsegs_i)
70
        Nstar_vals.append(Nstar_i)
Gregory Ashton's avatar
Gregory Ashton committed
71
        i += 1
72
        logging.info("Stage {}, nsegs={}, Nstar={}".format(i, nsegs_i, int(Nstar_i)))
Gregory Ashton's avatar
Gregory Ashton committed
73

74
    return nsegs_vals, Nstar_vals
Gregory Ashton's avatar
Gregory Ashton committed
75
76


77
78
79
def _get_nsegs_ip1(
    nsegs_i, NstarMax, tref, minStartTime, maxStartTime, prior, detector_names
):
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
    """ Calculate Nsegs_{i+1} given Nsegs_{i}

    Perform the optimization step to calculate nsegs and i+1 given the setup
    and i. The "Powell" minimiization method from scipy is used. Below, we give
    help for parameters unique to _get_nsegs_ip1, for help with other
    parameters see get_optimal_setup

    Parameters
    ----------
    nsegs_i: int
        The number of segments at step i

    Returns
    -------
    nsegs_ip1: int
        The number of segments at i + 1

    Raises
    ------
    ValueError: Optimisation unsuccesful
        A value error is raised if the optimization step fails

    """
Gregory Ashton's avatar
Gregory Ashton committed
103

Gregory Ashton's avatar
Gregory Ashton committed
104
    log10NstarMax = np.log10(NstarMax)
105
106
107
108
109
    log10Nstari = np.log10(
        get_Nstar_estimate(
            nsegs_i, tref, minStartTime, maxStartTime, prior, detector_names
        )
    )
Gregory Ashton's avatar
Gregory Ashton committed
110
111
112
113
114
115
116
117
118

    def f(nsegs_ip1):
        if nsegs_ip1[0] > nsegs_i:
            return 1e6
        if nsegs_ip1[0] < 0:
            return 1e6
        nsegs_ip1 = int(nsegs_ip1[0])
        if nsegs_ip1 == 0:
            nsegs_ip1 = 1
119
        Nstarip1 = get_Nstar_estimate(
120
121
            nsegs_ip1, tref, minStartTime, maxStartTime, prior, detector_names
        )
122
        if Nstarip1 is None:
Gregory Ashton's avatar
Gregory Ashton committed
123
124
            return 1e6
        else:
125
            log10Nstarip1 = np.log10(Nstarip1)
Gregory Ashton's avatar
Gregory Ashton committed
126
            return np.abs(log10Nstari + log10NstarMax - log10Nstarip1)
127
128
129
130
131

    res = scipy.optimize.minimize(
        f, 0.4 * nsegs_i, method="Powell", tol=1, options={"maxiter": 10}
    )
    logging.info("{} with {} evaluations".format(res["message"], res["nfev"]))
Gregory Ashton's avatar
Gregory Ashton committed
132
133
134
135
    nsegs_ip1 = int(res.x)
    if nsegs_ip1 == 0:
        nsegs_ip1 = 1
    if res.success:
136
137
138
139
140
141
        return (
            nsegs_ip1,
            get_Nstar_estimate(
                nsegs_ip1, tref, minStartTime, maxStartTime, prior, detector_names
            ),
        )
Gregory Ashton's avatar
Gregory Ashton committed
142
    else:
143
        raise ValueError("Optimisation unsuccesful")
Gregory Ashton's avatar
Gregory Ashton committed
144
145


146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def _extract_data_from_prior(prior):
    """ Calculate the input data from the prior

    Parameters
    ----------
    prior: dict

    Returns
    -------
    p : ndarray
        Matrix with columns being the edges of the uniform bounding box
    spindowns : int
        The number of spindowns
    sky : bool
        If true, search includes the sky position
    fiducial_freq : float
        Fidicual frequency

    """
165
    keys = ["Alpha", "Delta", "F0", "F1", "F2"]
166
167
168
169
170
171
172
    spindown_keys = keys[3:]
    sky_keys = keys[:2]
    lims = []
    lims_keys = []
    lims_idxs = []
    for i, key in enumerate(keys):
        if type(prior[key]) == dict:
173
174
            if prior[key]["type"] == "unif":
                lims.append([prior[key]["lower"], prior[key]["upper"]])
175
176
177
178
                lims_keys.append(key)
                lims_idxs.append(i)
            else:
                raise ValueError(
179
180
                    "Prior type {} not yet supported".format(prior[key]["type"])
                )
181
182
183
184
185
186
187
188
189
190
191
192
        elif key not in spindown_keys:
            lims.append([prior[key], 0])
    lims = np.array(lims)
    lims_keys = np.array(lims_keys)
    base = lims[:, 0]
    p = [base]
    for i in lims_idxs:
        basex = base.copy()
        basex[i] = lims[i, 1]
        p.append(basex)
    spindowns = np.sum([np.sum(lims_keys == k) for k in spindown_keys])
    sky = any([key in lims_keys for key in sky_keys])
193
194
    if type(prior["F0"]) == dict:
        fiducial_freq = prior["F0"]["upper"]
195
    else:
196
        fiducial_freq = prior["F0"]
197
198

    return np.array(p).T, spindowns, sky, fiducial_freq
199
200


201
def get_Nstar_estimate(nsegs, tref, minStartTime, maxStartTime, prior, detector_names):
202
    """ Returns N* estimated from the super-sky metric
Gregory Ashton's avatar
Gregory Ashton committed
203
204
205

    Parameters
    ----------
206
    nsegs : int
Gregory Ashton's avatar
Gregory Ashton committed
207
        Number of semi-coherent segments
208
    tref : int
Gregory Ashton's avatar
Gregory Ashton committed
209
        Reference time in GPS seconds
210
    minStartTime, maxStartTime : int
Gregory Ashton's avatar
Gregory Ashton committed
211
        Minimum and maximum SFT timestamps
212
    prior : dict
213
        The prior dictionary
214
    detector_names : array
Gregory Ashton's avatar
Gregory Ashton committed
215
216
        Array of detectors to average over

217
218
219
220
221
222
223
    Returns
    -------
    Nstar: int
        The estimated approximate number of templates to cover the prior
        parameter space at a mismatch of unity, assuming the normalised
        thickness is unity.

Gregory Ashton's avatar
Gregory Ashton committed
224
    """
225
    earth_ephem, sun_ephem = helper_functions.get_ephemeris_files()
226
    in_phys, spindowns, sky, fiducial_freq = _extract_data_from_prior(prior)
227
228
229
230
231
    out_rssky = np.zeros(in_phys.shape)

    in_phys = helper_functions.convert_array_to_gsl_matrix(in_phys)
    out_rssky = helper_functions.convert_array_to_gsl_matrix(out_rssky)

232
    tboundaries = np.linspace(minStartTime, maxStartTime, nsegs + 1)
Gregory Ashton's avatar
Gregory Ashton committed
233
234
235

    ref_time = lal.LIGOTimeGPS(tref)
    segments = lal.SegListCreate()
236
237
238
239
    for j in range(len(tboundaries) - 1):
        seg = lal.SegCreate(
            lal.LIGOTimeGPS(tboundaries[j]), lal.LIGOTimeGPS(tboundaries[j + 1]), j
        )
Gregory Ashton's avatar
Gregory Ashton committed
240
241
242
243
244
        lal.SegListAppend(segments, seg)
    detNames = lal.CreateStringVector(*detector_names)
    detectors = lalpulsar.MultiLALDetector()
    lalpulsar.ParseMultiLALDetector(detectors, detNames)
    detector_weights = None
245
    detector_motion = lalpulsar.DETMOTION_SPIN + lalpulsar.DETMOTION_ORBIT
Gregory Ashton's avatar
Gregory Ashton committed
246
247
248
    ephemeris = lalpulsar.InitBarycenter(earth_ephem, sun_ephem)
    try:
        SSkyMetric = lalpulsar.ComputeSuperskyMetrics(
249
250
251
252
253
254
255
256
257
258
            lalpulsar.SUPERSKY_METRIC_TYPE,
            spindowns,
            ref_time,
            segments,
            fiducial_freq,
            detectors,
            detector_weights,
            detector_motion,
            ephemeris,
        )
Gregory Ashton's avatar
Gregory Ashton committed
259
    except RuntimeError as e:
260
        logging.warning("Encountered run-time error {}".format(e))
261
        raise RuntimeError("Calculation of the SSkyMetric failed")
Gregory Ashton's avatar
Gregory Ashton committed
262

263
264
265
266
267
268
    if sky:
        i = 0
    else:
        i = 2

    lalpulsar.ConvertPhysicalToSuperskyPoints(
269
270
        out_rssky, in_phys, SSkyMetric.semi_rssky_transf
    )
271

272
    d = out_rssky.data
273

274
    g = SSkyMetric.semi_rssky_metric.data
275

276
    g = g[i:, i:]  # Remove sky if required
Gregory Ashton's avatar
Gregory Ashton committed
277
    parallelepiped = (d[i:, 1:].T - d[i:, 0]).T
278

279
    Nstars = []
280
    for j in range(1, len(g) + 1):
281
282
283
        dV = np.abs(np.linalg.det(parallelepiped[:j, :j]))
        sqrtdetG = np.sqrt(np.abs(np.linalg.det(g[:j, :j])))
        Nstars.append(sqrtdetG * dV)
284
285
286
287
288
    logging.debug(
        "Nstar for each dimension = {}".format(
            ", ".join(["{:1.1e}".format(n) for n in Nstars])
        )
    )
289
    return np.max(Nstars)