Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNR optimisation options for pycbc_live #4398

Closed
wants to merge 2 commits into from
Closed

SNR optimisation options for pycbc_live #4398

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7d878d1
Move SNR optimisation functions and arguments into their own module
GarethCabournDavies Mar 8, 2023
3d4dae2
typo
GarethCabournDavies May 5, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions bin/pycbc_live
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,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
Expand Down Expand Up @@ -114,12 +115,15 @@ class LiveEventManager(object):
bg_cores = len(tuple(itertools.combinations(ifos, 2)))
analysis_cores = 1 + bg_cores
self.fu_cores = available_cores - analysis_cores
self.optimizer = args.optimizer
if self.fu_cores <= 0:
logging.warning('Insufficient number of CPU cores (%d) to '
'run search and trigger followups. Uploaded '
'triggers will momentarily increase the lag',
available_cores)
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']:
Expand Down Expand Up @@ -313,6 +317,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:
Expand Down Expand Up @@ -894,15 +903,12 @@ 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',
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 '
Expand All @@ -918,6 +924,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()
scheme.verify_processing_options(args, parser)
Expand Down
234 changes: 11 additions & 223 deletions bin/pycbc_optimize_snr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,12 @@

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
from pycbc.conversions import (
mchirp_from_mass1_mass2, mtotal_from_mchirp_eta,
mass1_from_mtotal_eta, mass2_from_mtotal_eta
Expand All @@ -23,183 +18,7 @@ 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


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 = mtotal_from_mchirp_eta(v[0], v[1])
mass1 = mass1_from_mtotal_eta(mtotal, v[1])
mass2 = mass2_from_mtotal_eta(mtotal, v[1])

# enforce broadly accepted search space boundaries
if mass1 < 1 or mass2 < 1 or mtotal > 500:
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.0/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 optimize_di(bounds, cli_args, extra_args):
bounds = [
bounds['mchirp'],
bounds['eta'],
bounds['spin1z'],
bounds['spin2z']
]
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
)
return results.x


def optimize_shgo(bounds, cli_args, extra_args):
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 optimize_pso(bounds, cli_args, extra_args):
options = {
'c1': cli_args.pso_c1,
'c2': cli_args.pso_c2,
'w': cli_args.pso_w
}
min_bounds = [
bounds['mchirp'][0],
bounds['eta'][0],
bounds['spin1z'][0],
bounds['spin2z'][0]
]
max_bounds = [
bounds['mchirp'][1],
bounds['eta'][1],
bounds['spin1z'][1],
bounds['spin2z'][1]
]
optimizer = ps.single.GlobalBestPSO(
n_particles=cli_args.pso_particles,
dimensions=4,
options=options,
bounds=(min_bounds, max_bounds)
)
_, 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__)
Expand Down Expand Up @@ -241,49 +60,18 @@ 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.')

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()

pycbc.init_logging(args.verbose)

# Input checking
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)
snr_optimizer.check_snr_optimizer_options(args, parser)

scheme.verify_processing_options(args, parser)
fft.verify_fft_options(args, parser)
Expand Down Expand Up @@ -347,7 +135,7 @@ bounds = {
with scheme_context:
logging.info('Starting optimization')

optimize_func = optimize_funcs[args.optimizer]
optimize_func = snr_optimizer.optimize_funcs[args.optimizer]
opt_params = optimize_func(bounds, args, extra_args)

logging.info('Optimization complete')
Expand All @@ -358,7 +146,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 = mtotal_from_mchirp_eta(opt_params[0], opt_params[1])
mass1 = mass1_from_mtotal_eta(mtotal, opt_params[1])
Expand Down
5 changes: 5 additions & 0 deletions pycbc/live/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""
This packages contains modules to help with pycbc live running
"""

from .snr_optimizer import *
Loading