diff --git a/bin/pycbc_live b/bin/pycbc_live index 763d4ef0f09..3dfdbe93d13 100755 --- a/bin/pycbc_live +++ b/bin/pycbc_live @@ -42,6 +42,7 @@ from pycbc import mchirp_area from pycbc.detector import ppdets from pycbc.filter import resample from pycbc.psd import estimate +from pycbc.live import snr_optimizer # Use cached class-based FFTs in the resample and estimate module resample.USE_CACHING_FOR_LFILTER = True @@ -119,6 +120,7 @@ class LiveEventManager(object): if platform.system() != 'Darwin': available_cores = len(os.sched_getaffinity(0)) self.fu_cores = available_cores - analysis_cores + self.optimizer = args.snr_opt_method if self.fu_cores <= 0: logging.warning( 'Insufficient number of CPU cores (%d) to ' @@ -130,6 +132,8 @@ class LiveEventManager(object): else: # To enable mac testing, this is just set to 1 self.fu_cores = 1 + # Convert SNR optimizer options into a string + self.snr_opt_options = snr_optimizer.args_to_string(args) if args.enable_embright_has_massgap: if args.embright_massgap_max < self.mc_area_args['mass_bdary']['ns_max']: @@ -321,6 +325,11 @@ class LiveEventManager(object): cmd = 'timeout {} '.format(args.snr_opt_timeout) exepath = which('pycbc_optimize_snr') cmd += exepath + ' ' + + # Add SNR optimization options to the command + cmd += self.snr_opt_options + + # Add data files according to the event we are optimizing data_fils_str = '--data-files ' psd_fils_str = '--psd-files ' for ifo in live_ifos: @@ -968,15 +977,13 @@ parser.add_argument('--size-override', type=int, metavar='N', parser.add_argument('--fftw-planning-limit', type=float, help="Time in seconds to allow for a plan to be created") parser.add_argument('--run-snr-optimization', action='store_true', - default=False, help='Run spawned followup processes to maximize SNR for ' 'any trigger uploaded to GraceDB') parser.add_argument('--snr-opt-timeout', type=int, default=400, metavar='SECONDS', help='Maximum allowed duration of followup process to maximize SNR') -parser.add_argument('--snr-opt-label', type=str, default='SNR_OPTIMIZED', +parser.add_argument('--snr-opt-label', default='SNR_OPTIMIZED', help='Label to apply to snr-optimized GraceDB uploads') - parser.add_argument('--enable-embright-has-massgap', action='store_true', default=False, help='Estimate HasMassGap probability for EMBright info. Lower limit ' 'of the mass gap is equal to the maximum NS mass used for ' @@ -994,6 +1001,7 @@ Coincer.insert_args(parser) SingleDetSGChisq.insert_option_group(parser) mchirp_area.insert_args(parser) livepau.insert_live_pastro_option_group(parser) +snr_optimizer.insert_snr_optimizer_options(parser) args = parser.parse_args() diff --git a/bin/pycbc_optimize_snr b/bin/pycbc_optimize_snr index ce38ac7f793..c10e58cda6e 100755 --- a/bin/pycbc_optimize_snr +++ b/bin/pycbc_optimize_snr @@ -4,244 +4,18 @@ import os import argparse, numpy, h5py -import time -import types import logging -from scipy.optimize import differential_evolution, shgo import pycbc from pycbc import ( - DYN_RANGE_FAC, fft, scheme, version, waveform + fft, scheme, version ) -from pycbc.types import MultiDetOptionAction, zeros, load_frequencyseries -from pycbc.filter import matched_filter_core -import pycbc.waveform.bank +from pycbc.types import MultiDetOptionAction, load_frequencyseries import pycbc.conversions as cv from pycbc.io import live from pycbc.io.hdf import load_hdf5_to_dict from pycbc.detector import Detector from pycbc.psd import interpolate - -try: - import pyswarms as ps -except: - ps = None - -# Set a minimum mass for points tried in optimization allowing for -# minimal slop relative to the lightest template -MIN_CPT_MASS = 0.99 - -# Set a large maximum total mass -MAX_MTOTAL = 500. - -# Set a notional maximum and minimum possible eta -MAX_ETA = 0.24999999999 -MIN_ETA = 0.01 - -Nfeval = 0 -start_time = time.time() - -def callback_func(Xi, convergence=0): - global Nfeval - logging.info("Currently at %d %s", Nfeval, convergence) - if (time.time() - start_time) > 360: - return True - Nfeval += 1 - - -def compute_network_snr_core(v, *argv, raise_err=False): - """ - Compute network SNR as a function over mchirp, eta and two aligned - spin components, stored in that order in the sequence v - """ - data = argv[0] - coinc_times = argv[1] - ifos = argv[2] - flen = argv[3] - approximant = argv[4] - flow = argv[5] - f_end = argv[6] - delta_f = argv[7] - sample_rate = argv[8] - distance = 1.0 / DYN_RANGE_FAC - mtotal = cv.mtotal_from_mchirp_eta(v[0], v[1]) - mass1 = cv.mass1_from_mtotal_eta(mtotal, v[1]) - mass2 = cv.mass2_from_mtotal_eta(mtotal, v[1]) - - # enforce broadly accepted search space boundaries - if mass1 < MIN_CPT_MASS or mass2 < MIN_CPT_MASS or mtotal > MAX_MTOTAL: - return -numpy.inf, {} - - try: - htilde = waveform.get_waveform_filter( - zeros(flen, dtype=numpy.complex64), - approximant=approximant, - mass1=mass1, mass2=mass2, spin1z=v[2], spin2z=v[3], - f_lower=flow, f_final=f_end, delta_f=delta_f, - delta_t=1./sample_rate, distance=distance) - except RuntimeError: - if raise_err: - raise - # assume a failure in the waveform approximant - # due to the choice of parameters - else: - return -numpy.inf, {} - - if not hasattr(htilde, 'params'): - htilde.params = dict(mass1=mass1, mass2=mass2, - spin1z=v[2], spin2z=v[3]) - if not hasattr(htilde, 'end_idx'): - htilde.end_idx = int(f_end / htilde.delta_f) - htilde.approximant = approximant - htilde.sigmasq = types.MethodType(pycbc.waveform.bank.sigma_cached, - htilde) - htilde.min_f_lower = flow - htilde.end_frequency = f_end - htilde.f_lower = flow - network_snrsq = 0 - snr_series_dict = {} - for ifo in ifos: - sigmasq = htilde.sigmasq(data[ifo].psd) - snr, _, norm = matched_filter_core(htilde, data[ifo], - h_norm=sigmasq) - duration = 0.095 - half_dur_samples = int(snr.sample_rate * duration / 2) - onsource_idx = float(coinc_times[ifo] - snr.start_time) * snr.sample_rate - onsource_idx = int(round(onsource_idx)) - onsource_slice = slice(onsource_idx - half_dur_samples, - onsource_idx + half_dur_samples + 1) - snr_series = snr[onsource_slice] * norm - snr_series_dict[ifo] = snr * norm - snr_series_dict['sigmasq_' + ifo] = sigmasq - network_snrsq += max(abs(snr_series._data)) ** 2. - - return network_snrsq ** 0.5, snr_series_dict - - -def compute_minus_network_snr(v, *argv): - if len(argv) == 1: - argv = argv[0] - nsnr, _ = compute_network_snr_core(v, *argv) - return -nsnr - - -def compute_minus_network_snr_pso(v, *argv, **kwargs): - argv = kwargs['args'] - nsnr_array = numpy.array([compute_network_snr_core(v_i, *argv)[0] for v_i in v]) - return -nsnr_array - - -def normalize_initial_point(initial_point, bounds): - return (initial_point - bounds[:,0]) / (bounds[:,1] - bounds[:,0]) - -def optimize_di(bounds, cli_args, extra_args, initial_point): - bounds = numpy.array([ - bounds['mchirp'], - bounds['eta'], - bounds['spin1z'], - bounds['spin2z'] - ]) - - - # Currently only implemented for random seed initial array - rng = numpy.random.mtrand._rand - population_shape=(cli_args.di_popsize, 4) - population = rng.uniform(size=population_shape) - if cli_args.include_candidate_in_optimizer: - # Re-normalize the initial point into the correct range - point_init = normalize_initial_point(initial_point, bounds) - # add the initial point to the population - population = numpy.concatenate((population[:-1], point_init)) - - results = differential_evolution( - compute_minus_network_snr, - bounds, - maxiter=cli_args.di_maxiter, - workers=(cli_args.cores or -1), - popsize=cli_args.di_popsize, - mutation=(0.5, 1), - recombination=0.7, - callback=callback_func, - args=extra_args, - init=population - ) - return results.x - - -def optimize_shgo(bounds, cli_args, extra_args, initial_point): # pylint: disable=unused-argument - bounds = [ - bounds['mchirp'], - bounds['eta'], - bounds['spin1z'], - bounds['spin2z'] - ] - results = shgo( - compute_minus_network_snr, - bounds=bounds, - args=extra_args, - iters=args.shgo_iters, - n=args.shgo_samples, - sampling_method="sobol" - ) - return results.x - -def normalize_population(population, min_bounds, max_bounds): - norm_pop = min_bounds + population * (max_bounds - min_bounds) - - return norm_pop - -def optimize_pso(bounds, cli_args, extra_args, initial_point): - options = { - 'c1': cli_args.pso_c1, - 'c2': cli_args.pso_c2, - 'w': cli_args.pso_w - } - min_bounds = numpy.array([ - bounds['mchirp'][0], - bounds['eta'][0], - bounds['spin1z'][0], - bounds['spin2z'][0] - ]) - max_bounds = numpy.array([ - bounds['mchirp'][1], - bounds['eta'][1], - bounds['spin1z'][1], - bounds['spin2z'][1] - ]) - - # Manually generate the initial points, this is the same as the default - # method, but allows us to make some modifications - population = numpy.random.uniform( - low=0.0, high=1.0, size=(cli_args.pso_particles, 4) - ) - population = normalize_population(population, min_bounds, max_bounds) - - if cli_args.include_candidate_in_optimizer: - # add the initial point to the population - population = numpy.concatenate((population[:-1], - initial_point)) - - optimizer = ps.single.GlobalBestPSO( - n_particles=cli_args.pso_particles, - dimensions=4, - options=options, - bounds=(min_bounds, max_bounds), - init_pos=population - ) - _, results = optimizer.optimize( - compute_minus_network_snr_pso, - iters=cli_args.pso_iters, - n_processes=cli_args.cores, - args=extra_args - ) - return results - - -optimize_funcs = { - 'differential_evolution': optimize_di, - 'shgo': optimize_shgo, - 'pso': optimize_pso -} - +from pycbc.live import snr_optimizer parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--version', action='version', @@ -286,60 +60,22 @@ parser.add_argument('--output-path', required=True, help='Path to a directory to store results in') parser.add_argument('--cores', type=int, help='Restrict calculation to given number of CPU cores') -parser.add_argument('--optimizer', type=str, default='differential_evolution', - choices=sorted(optimize_funcs), - help='The optimizer to use, differential_evolution, shgo or pso') -parser.add_argument('--di-maxiter', type=int, default=50, - help='Only relevant for --optimizer differential_evolution: ' - 'The maximum number of generations over which the entire ' - 'population is evolved.') -parser.add_argument('--di-popsize', type=int, default=100, - help='Only relevant for --optimizer differential_evolution: ' - 'A multiplier for setting the total population size.') -parser.add_argument('--shgo-samples', type=int, default=76, - help='Only relevant for --optimizer shgo: ' - 'Number of sampling points used in the construction of the ' - 'simplicial complex.') -parser.add_argument('--shgo-iters', type=int, default=3, - help='Only relevant for --optimizer shgo: ' - 'Number of iterations used in the construction of the simplicial complex.') -parser.add_argument('--pso-iters', type=int, default=5, - help='Only relevant for --optimizer pso: ' - 'Number of iterations used in the particle swarm optimization.') -parser.add_argument('--pso-particles', type=int, default=250, - help='Only relevant for --optimizer pso: ' - 'Number of particles used in the swarm.') -parser.add_argument('--pso-c1', type=float, default=0.5, - help='Only relevant for --optimizer pso: ' - 'The hyperparameter c1: the cognitive parameter.') -parser.add_argument('--pso-c2', type=float, default=2.0, - help='Only relevant for --optimizer pso: ' - 'The hyperparameter c2: the social parameter.') -parser.add_argument('--pso-w', type=float, default=0.01, - help='Only relevant for --optimizer pso: ' - 'The hyperparameter w: the inertia parameter.') -parser.add_argument('--include-candidate-in-optimizer', action='store_true', - help='Include parameters of the candidate event in the ' - 'initialised array for the optimizer. Only relevant for ' - '--optimizer pso or differential_evolution') -parser.add_argument('--seed', type=int, default=42, - help='Seed to supply to the random generation of initial ' - 'array to pass to the optimizer. Only relevant for ' - '--optimizer pso or differential_evolution') - +snr_optimizer.insert_snr_optimizer_options(parser) scheme.insert_processing_option_group(parser) fft.insert_fft_option_group(parser) -# Input checking args = parser.parse_args() -if args.snr_threshold != 4: - parser.error("Sorry, the SNR threshold option doesn't work yet") -if args.optimizer == 'pso' and ps == None: - parser.error('You need to install pyswarms to use the pso optimizer.') + pycbc.init_logging(args.verbose) -if args.seed: - numpy.random.seed(args.seed) +if args.snr_opt_seed is not None and args.snr_opt_seed != 'random': + logging.info('Setting snr optimizer random seed.') + numpy.random.seed(int(args.snr_opt_seed)) + +# Input checking +if args.snr_threshold != 4: + parser.error("Sorry, the SNR threshold option doesn't work yet") +snr_optimizer.check_snr_optimizer_options(args, parser) scheme.verify_processing_options(args, parser) fft.verify_fft_options(args, parser) @@ -347,6 +83,14 @@ fft.verify_fft_options(args, parser) scheme_context = scheme.from_cli(args) fft.from_cli(args) +# Set a minimum mass for points tried in optimization allowing for +# minimal slop relative to the lightest template +MIN_CPT_MASS = snr_optimizer.MIN_CPT_MASS + +# Set a notional maximum and minimum possible eta +MAX_ETA = 0.24999999999 +MIN_ETA = 0.01 + logging.info('Starting optimize SNR') data = {} @@ -420,7 +164,7 @@ bounds = { 'spin2z': (minspin2z, maxspin2z) } -if args.include_candidate_in_optimizer: +if args.snr_opt_include_candidate: # Initial point from found candidate mchirp_init = cv.mchirp_from_mass1_mass2(fp['mass1'][()], fp['mass2'][()]) eta_init = cv.eta_from_mass1_mass2(fp['mass1'][()], fp['mass2'][()]) @@ -439,7 +183,7 @@ else: with scheme_context: logging.info('Starting optimization') - optimize_func = optimize_funcs[args.optimizer] + optimize_func = snr_optimizer.optimize_funcs[args.snr_opt_method] opt_params = optimize_func(bounds, args, extra_args, initial_point) logging.info('Optimization complete') @@ -450,7 +194,7 @@ with scheme_context: extra_args[2] = ifos - _, snr_series_dict = compute_network_snr_core(opt_params, *extra_args) + _, snr_series_dict = snr_optimizer.compute_network_snr_core(opt_params, *extra_args) mtotal = cv.mtotal_from_mchirp_eta(opt_params[0], opt_params[1]) mass1 = cv.mass1_from_mtotal_eta(mtotal, opt_params[1]) diff --git a/examples/live/run.sh b/examples/live/run.sh index 3cda2773664..2e5e81615d4 100755 --- a/examples/live/run.sh +++ b/examples/live/run.sh @@ -195,6 +195,10 @@ python -m mpi4py `which pycbc_live` \ --src-class-eff-to-lum-distance 0.74899 \ --src-class-lum-distance-to-delta -0.51557 -0.32195 \ --run-snr-optimization \ +--snr-opt-di-maxiter 50 \ +--snr-opt-di-popsize 100 \ +--snr-opt-include-candidate \ +--snr-opt-seed 42 \ --sngl-ifar-est-dist conservative \ --single-newsnr-threshold 9 \ --single-duration-threshold 7 \ @@ -202,6 +206,15 @@ python -m mpi4py `which pycbc_live` \ --single-fit-file single_trigger_fits.hdf \ --verbose +# If you would like to use the pso optimizer, change --optimizer to pso +# and include these arguments while removing other optimizer args. +# You will need to install the pyswarms package into your environment. +# --snr-opt-pso-iters 5 \ +# --snr-opt-pso-particles 250 \ +# --snr-opt-pso-c1 0.5 \ +# --snr-opt-pso-c2 2.0 \ +# --snr-opt-pso-w 0.01 \ + # note that, at this point, some SNR optimization processes may still be # running, so the checks below may ignore their results diff --git a/pycbc/live/__init__.py b/pycbc/live/__init__.py new file mode 100644 index 00000000000..c1f3cddf6c8 --- /dev/null +++ b/pycbc/live/__init__.py @@ -0,0 +1,5 @@ +""" +This packages contains modules to help with pycbc live running +""" + +from .snr_optimizer import * diff --git a/pycbc/live/snr_optimizer.py b/pycbc/live/snr_optimizer.py new file mode 100644 index 00000000000..5924d067304 --- /dev/null +++ b/pycbc/live/snr_optimizer.py @@ -0,0 +1,406 @@ +# Copyright (C) 2023 Arthur Tolley, Gareth Cabourn Davies +# +# This program is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 3 of the License, or (at your +# option) any later version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General +# Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + +# +# ============================================================================= +# +# Preamble +# +# ============================================================================= +# +""" This module contains functions for optimizing the signal-to-noise ratio +of triggers produced by PyCBC Live. Also contained within this module are the +command line arguments required and options group for the SNR optimization. +This module is primarily used in the pycbc_optimize_snr program. +""" + +import time +import logging +import types +import numpy +from scipy.optimize import differential_evolution, shgo +from pycbc import ( + DYN_RANGE_FAC, waveform +) +from pycbc.types import zeros +import pycbc.waveform.bank +from pycbc.filter import matched_filter_core +import pycbc.conversions as cv + +try: + import pyswarms as ps +except: + ps = None + +# Set a minimum mass for points tried in optimization allowing for +# minimal slop relative to the lightest template +MIN_CPT_MASS = 0.99 + +# Set a large maximum total mass +MAX_MTOTAL = 500. + +Nfeval = 0 +start_time = time.time() + + +def callback_func(Xi, convergence=0): + global Nfeval + logging.info("Currently at %d %s", Nfeval, convergence) + # Time out if the optimization takes longer than 6 minutes + if (time.time() - start_time) > 360: + return True + Nfeval += 1 + + +def compute_network_snr_core(v, data, coinc_times, ifos, flen, approximant, + flow, f_end, delta_f, sample_rate, raise_err=False): + """ + Compute network SNR as a function over mchirp, eta, and two aligned + spin components, stored in that order in the sequence v. + + Parameters + ---------- + v : list + A list containing the input values for mchirp, eta, and spin + components. + data : dict + A dictionary containing keys of ifos ('H1', 'L1') and + values of the frequency series data for those ifos + coinc_times : dict + A dictionary containing the coincidence times for the network. + ifos : list + A list of the ifos, e.g. ['H1', 'L1'] + flen : float + The length of the data. + approximant : str + The approximant used for the waveform model. + flow : float + The lower frequency bound. + f_end : float + The upper frequency bound. + delta_f : float + The frequency spacing. + sample_rate : float + The sampling rate of the data. + raise_err : bool, optional + A flag indicating whether to raise an error if an exception + occurs during the computation. Defaults to False. + + Returns + ------- + network_snr : float + The computed network SNR (Signal-to-Noise Ratio) value. + snr_series_dict : dict + A dictionary containing the snr timeseries from each ifo. + + """ + distance = 1.0 / DYN_RANGE_FAC + mtotal = cv.mtotal_from_mchirp_eta(v[0], v[1]) + mass1 = cv.mass1_from_mtotal_eta(mtotal, v[1]) + mass2 = cv.mass2_from_mtotal_eta(mtotal, v[1]) + + # enforce broadly accepted search space boundaries + if mass1 < MIN_CPT_MASS or mass2 < MIN_CPT_MASS or mtotal > MAX_MTOTAL: + return -numpy.inf, {} + + try: + htilde = waveform.get_waveform_filter( + zeros(flen, dtype=numpy.complex64), + approximant=approximant, + mass1=mass1, mass2=mass2, spin1z=v[2], spin2z=v[3], + f_lower=flow, f_final=f_end, delta_f=delta_f, + delta_t=1./sample_rate, distance=distance) + except RuntimeError: + if raise_err: + raise + # Assume a failure in the waveform approximant + # due to the choice of parameters and carry on + return -numpy.inf, {} + + if not hasattr(htilde, 'params'): + htilde.params = dict(mass1=mass1, mass2=mass2, + spin1z=v[2], spin2z=v[3]) + if not hasattr(htilde, 'end_idx'): + htilde.end_idx = int(f_end / htilde.delta_f) + htilde.approximant = approximant + htilde.sigmasq = types.MethodType(pycbc.waveform.bank.sigma_cached, + htilde) + htilde.min_f_lower = flow + htilde.end_frequency = f_end + htilde.f_lower = flow + network_snrsq = 0 + snr_series_dict = {} + for ifo in ifos: + sigmasq = htilde.sigmasq(data[ifo].psd) + snr, _, norm = matched_filter_core(htilde, data[ifo], + h_norm=sigmasq) + duration = 0.095 + half_dur_samples = int(snr.sample_rate * duration / 2) + onsource_idx = (float(coinc_times[ifo] - snr.start_time) * + snr.sample_rate) + onsource_idx = int(round(onsource_idx)) + onsource_slice = slice(onsource_idx - half_dur_samples, + onsource_idx + half_dur_samples + 1) + snr_series = snr[onsource_slice] * norm + snr_series_dict[ifo] = snr * norm + snr_series_dict['sigmasq_' + ifo] = sigmasq + network_snrsq += max(abs(snr_series._data)) ** 2. + + return network_snrsq ** 0.5, snr_series_dict + + +def compute_minus_network_snr(v, *argv): + if len(argv) == 1: + argv = argv[0] + nsnr, _ = compute_network_snr_core(v, *argv) + return -nsnr + + +def compute_minus_network_snr_pso(v, *argv, **kwargs): + argv = kwargs['args'] + nsnr_array = numpy.array([ + compute_network_snr_core(v_i, *argv)[0] + for v_i in v]) + return -nsnr_array + + +def normalize_initial_point(initial_point, bounds): + return (initial_point - bounds[:,0]) / (bounds[:,1] - bounds[:,0]) + + +def optimize_di(bounds, cli_args, extra_args, initial_point): + bounds = numpy.array([ + bounds['mchirp'], + bounds['eta'], + bounds['spin1z'], + bounds['spin2z'] + ]) + + # Currently only implemented for random seed initial array + rng = numpy.random.mtrand._rand + population_shape = (int(cli_args.snr_opt_di_popsize), 4) + population = rng.uniform(size=population_shape) + if cli_args.snr_opt_include_candidate: + # Re-normalize the initial point into the correct range + point_init = normalize_initial_point(initial_point, bounds) + # add the initial point to the population + population = numpy.concatenate((population[:-1], point_init)) + + results = differential_evolution( + compute_minus_network_snr, + bounds, + maxiter=int(cli_args.snr_opt_di_maxiter), + workers=(cli_args.cores or -1), + popsize=int(cli_args.snr_opt_di_popsize), + mutation=(0.5, 1), + recombination=0.7, + callback=callback_func, + args=extra_args, + init=population + ) + return results.x + + +def optimize_shgo(bounds, cli_args, extra_args, initial_point): # pylint: disable=unused-argument + bounds = [ + bounds['mchirp'], + bounds['eta'], + bounds['spin1z'], + bounds['spin2z'] + ] + results = shgo( + compute_minus_network_snr, + bounds=bounds, + args=extra_args, + iters=cli_args.snr_opt_shgo_iters, + n=cli_args.snr_opt_shgo_samples, + sampling_method="sobol" + ) + return results.x + + +def normalize_population(population, min_bounds, max_bounds): + norm_pop = min_bounds + population * (max_bounds - min_bounds) + return norm_pop + + +def optimize_pso(bounds, cli_args, extra_args, initial_point): + options = { + 'c1': cli_args.snr_opt_pso_c1, + 'c2': cli_args.snr_opt_pso_c2, + 'w': cli_args.snr_opt_pso_w + } + min_bounds = numpy.array([ + bounds['mchirp'][0], + bounds['eta'][0], + bounds['spin1z'][0], + bounds['spin2z'][0] + ]) + max_bounds = numpy.array([ + bounds['mchirp'][1], + bounds['eta'][1], + bounds['spin1z'][1], + bounds['spin2z'][1] + ]) + + # Manually generate the initial points, this is the same as the default + # method, but allows us to make some modifications + population = numpy.random.uniform( + low=0.0, high=1.0, size=(int(cli_args.snr_opt_pso_particles), 4) + ) + population = normalize_population(population, min_bounds, max_bounds) + + if cli_args.snr_opt_include_candidate: + # add the initial point to the population + population = numpy.concatenate((population[:-1], + initial_point)) + + optimizer = ps.single.GlobalBestPSO( + n_particles=int(cli_args.snr_opt_pso_particles), + dimensions=4, + options=options, + bounds=(min_bounds, max_bounds), + init_pos=population + ) + _, results = optimizer.optimize( + compute_minus_network_snr_pso, + iters=int(cli_args.snr_opt_pso_iters), + n_processes=cli_args.cores, + args=extra_args + ) + return results + + +optimize_funcs = { + 'differential_evolution': optimize_di, + 'shgo': optimize_shgo, + 'pso': optimize_pso +} + +# The following sets the default values of the options, but allows us to check +# if the option has been given on the command line + +# For each optimizer, we have a dictionary of the options, its help +# message and default value + +option_dict = { + 'differential_evolution': { + 'maxiter': ('The maximum number of generations over which the entire ' + 'population is evolved.', 50), + 'popsize': ('A multiplier for setting the total population size.', + 100), + }, + 'shgo': { + 'samples': ('Number of sampling points used in the construction of ' + 'the simplicial complex.', 76), + 'iters': ('Number of iterations used in the construction of the ' + 'simplicial complex.', 3), + }, + 'pso': { + 'iters': ('Number of iterations used in the particle swarm ' + 'optimization.', 5), + 'particles': ('Number of particles used in the swarm.', 250), + 'c1': ('The hyperparameter c1: the cognitive parameter.', 0.5), + 'c2': ('The hyperparameter c2: the social parameter.', 2.0), + 'w': ('The hyperparameter w: the inertia parameter.', 0.01), + } +} + +def insert_snr_optimizer_options(parser): + opt_opt_group = parser.add_argument_group("SNR optimizer configuration " + "options.") + # Option to choose which optimizer to use: + optimizer_choices = sorted(list(option_dict.keys())) + opt_opt_group.add_argument('--snr-opt-method', + default='differential_evolution', + choices=optimizer_choices, + help='SNR Optimizer choices: ' + ', '.join(optimizer_choices)) + + # Add the generic options + opt_opt_group.add_argument('--snr-opt-include-candidate', + action='store_true', + help='Include parameters of the candidate event in the initialized ' + 'array for the optimizer. Only relevant for --optimizer pso or ' + 'differential_evolution') + opt_opt_group.add_argument('--snr-opt-seed', + default='42', + help='Seed to supply to the random generation of initial array to ' + 'pass to the optimizer. Only relevant for --optimizer pso or ' + 'differential_evolution. Set to ''random'' for a random seed') + + # For each optimizer, add the possible options + for optimizer, option_subdict in option_dict.items(): + optimizer_name = optimizer.replace('_', '-') + if optimizer_name == 'differential-evolution': + optimizer_name = 'di' + for opt_name, opt_help_default in option_subdict.items(): + option_name = f"--snr-opt-{optimizer_name}-{opt_name}" + opt_opt_group.add_argument(option_name, + type=float, + help=f'Only relevant for --optimizer {optimizer}: ' + + opt_help_default[0] + + f' Default = {opt_help_default[1]}') + + +def check_snr_optimizer_options(args, parser): + """ + Deal with default options and required parameters given optimizer option + """ + options = {} + options['differential_evolution'] = [args.snr_opt_di_maxiter, + args.snr_opt_di_popsize] + options['shgo'] = [args.snr_opt_shgo_samples, args.snr_opt_shgo_iters] + options['pso'] = [args.snr_opt_pso_iters, args.snr_opt_pso_particles, + args.snr_opt_pso_c1, args.snr_opt_pso_c2, + args.snr_opt_pso_w] + + if args.snr_opt_method == 'pso' and ps is None: + parser.error('You need to install pyswarms to use the pso optimizer.') + + # Check all the options are suitable for the chosen optimizer + for k in options.keys(): + if args.snr_opt_method == k: + continue + if any(options[k]): + parser.error("Argument has been supplied which is not suitable " + + f"for the optimizer given ({args.snr_opt_method})") + + # Give the arguments the default values according to the dictionary + optimizer_name = args.snr_opt_method.replace('_', '-') + if optimizer_name == 'differential-evolution': + optimizer_name = 'di' + for key, value in option_dict[args.snr_opt_method].items(): + key_name = f'snr_opt_{optimizer_name}_{key}' + if not getattr(args, key_name): + setattr(args, key_name, value[1]) + + +def args_to_string(args): + """ + Convert the supplied arguments for SNR optimization config into + a string - this is to be used when running subprocesses + """ + argstr = f'--snr-opt-method {args.snr_opt_method} ' + optimizer_name = args.snr_opt_method.replace('_', '-') + if optimizer_name == 'differential-evolution': + optimizer_name = 'di' + for opt in option_dict[args.snr_opt_method]: + option_fullname = f'--snr-opt-{optimizer_name}-{opt}' + key_name = f'snr_opt_{optimizer_name}_{opt}' + option_value = getattr(args, key_name) + argstr += f'{option_fullname} {option_value} ' + + return argstr