Skip to content

Commit

Permalink
Merge pull request #122 from jerry-sun1/master
Browse files Browse the repository at this point in the history 
add BWM functionality
  • Loading branch information
@stevertaylor
stevertaylor authored Aug 31, 2021
2 parents 0950720 + f6e637d commit 216eb7d
Show file tree
Hide file tree
Showing 4 changed files with 370 additions and 115 deletions.
124 changes: 115 additions & 9 deletions enterprise_extensions/blocks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,18 @@
from enterprise.signals import utils
from enterprise.signals import gp_bases as gpb
from enterprise.signals import gp_priors as gpp
from enterprise.signals import deterministic_signals
from enterprise import constants as const
from . import gp_kernels as gpk
from . import chromatic as chrom
from . import model_orfs

from enterprise_extensions import deterministic as ee_deterministic

__all__ = ['white_noise_block',
'red_noise_block',
'bwm_block',
'bwm_sglpsr_block'
'dm_noise_block',
'scattering_noise_block',
'chromatic_noise_block',
Expand Down Expand Up @@ -101,7 +106,7 @@ def white_noise_block(vary=False, inc_ecorr=False, gp_ecorr=False,
def red_noise_block(psd='powerlaw', prior='log-uniform', Tspan=None,
components=30, gamma_val=None, coefficients=False,
select=None, modes=None, wgts=None,
break_flat=False, break_flat_fq=None):
break_flat=False, break_flat_fq=None, logmin=None, logmax=None):
"""
Returns red noise model:
1. Red noise modeled as a power-law with 30 sampling frequencies
Expand All @@ -124,15 +129,21 @@ def red_noise_block(psd='powerlaw', prior='log-uniform', Tspan=None,
if psd in ['powerlaw', 'powerlaw_genmodes', 'turnover',
'tprocess', 'tprocess_adapt', 'infinitepower']:
# parameters shared by PSD functions
if prior == 'uniform':
log10_A = parameter.LinearExp(-20, -11)
elif prior == 'log-uniform' and gamma_val is not None:
if np.abs(gamma_val - 4.33) < 0.1:
log10_A = parameter.Uniform(-20, -11)
if logmin is not None and logmax is not None:
if prior == 'uniform':
log10_A = parameter.LinearExp(logmin, logmax)
elif prior == 'log-uniform':
log10_A = parameter.Uniform(logmin, logmax)
else:
if prior == 'uniform':
log10_A = parameter.LinearExp(-20, -11)
elif prior == 'log-uniform' and gamma_val is not None:
if np.abs(gamma_val - 4.33) < 0.1:
log10_A = parameter.Uniform(-20, -11)
else:
log10_A = parameter.Uniform(-20, -11)
else:
log10_A = parameter.Uniform(-20, -11)
else:
log10_A = parameter.Uniform(-20, -11)

if gamma_val is not None:
gamma = parameter.Constant(gamma_val)
Expand Down Expand Up @@ -214,6 +225,83 @@ def red_noise_block(psd='powerlaw', prior='log-uniform', Tspan=None,

return rn

def bwm_block(Tmin, Tmax, amp_prior='log-uniform',
skyloc=None, logmin=-18, logmax=-11,
name='bwm'):
"""
Returns deterministic GW burst with memory model:
1. Burst event parameterized by time, sky location,
polarization angle, and amplitude
:param Tmin:
Min time to search, probably first TOA (MJD).
:param Tmax:
Max time to search, probably last TOA (MJD).
:param amp_prior:
Prior on log10_A. Default if "log-uniform". Use "uniform" for
upper limits.
:param skyloc:
Fixed sky location of BWM signal search as [cos(theta), phi].
Search over sky location if ``None`` given.
:param logmin:
log of minimum BWM amplitude for prior (log10)
:param logmax:
log of maximum BWM amplitude for prior (log10)
:param name:
Name of BWM signal.
"""

# BWM parameters
amp_name = '{}_log10_A'.format(name)
if amp_prior == 'uniform':
log10_A_bwm = parameter.LinearExp(logmin, logmax)(amp_name)
elif amp_prior == 'log-uniform':
log10_A_bwm = parameter.Uniform(logmin, logmax)(amp_name)

pol_name = '{}_pol'.format(name)
pol = parameter.Uniform(0, np.pi)(pol_name)

t0_name = '{}_t0'.format(name)
t0 = parameter.Uniform(Tmin, Tmax)(t0_name)

costh_name = '{}_costheta'.format(name)
phi_name = '{}_phi'.format(name)
if skyloc is None:
costh = parameter.Uniform(-1, 1)(costh_name)
phi = parameter.Uniform(0, 2*np.pi)(phi_name)
else:
costh = parameter.Constant(skyloc[0])(costh_name)
phi = parameter.Constant(skyloc[1])(phi_name)


# BWM signal
bwm_wf = ee_deterministic.bwm_delay(log10_h=log10_A_bwm, t0=t0,
cos_gwtheta=costh, gwphi=phi, gwpol=pol)
bwm = deterministic_signals.Deterministic(bwm_wf, name=name)

return bwm

def bwm_sglpsr_block(Tmin, Tmax, amp_prior='log-uniform',
logmin=-17, logmax=-12, name='ramp', fixed_sign=None):

if fixed_sign is None:
sign = parameter.Uniform(-1, 1)("sign")
else:
sign = np.sign(fixed_sign)

amp_name = '{}_log10_A'.format(name)
if amp_prior == 'uniform':
log10_A_ramp = parameter.LinearExp(logmin, logmax)(amp_name)
elif amp_prior == 'log-uniform':
log10_A_ramp = parameter.Uniform(logmin, logmax)(amp_name)

t0_name = '{}_t0'.format(name)
t0 = parameter.Uniform(Tmin, Tmax)(t0_name)


ramp_wf = ee_deterministic.bwm_sglpsr_delay(log10_A=log10_A_ramp, t0=t0, sign = sign)
ramp = deterministic_signals.Deterministic(ramp_wf, name=name)

return ramp

def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic',
prior='log-uniform', dt=15, df=200,
Expand Down Expand Up @@ -493,6 +581,7 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw',
def common_red_noise_block(psd='powerlaw', prior='log-uniform',
Tspan=None, components=30,
log10_A_val = None, gamma_val=None, delta_val=None,
logmin = None, logmax = None,
orf=None, orf_ifreq=0, leg_lmax=5,
name='gw', coefficients=False,
pshift=False, pseed=None):
Expand All @@ -519,7 +608,11 @@ def common_red_noise_block(psd='powerlaw', prior='log-uniform',
models. By default spectral index is varied of range [0,7]
:param delta_val:
Value of spectral index for high frequencies in broken power-law
and turnover models. By default spectral index is varied in range [0,7].
and turnover models. By default spectral index is varied in range [0,7].\
:param logmin:
Specify the lower bound of the prior on the amplitude.
:param logmax:
Specify the upper bound of the prior on the amplitude
:param orf:
String representing which overlap reduction function to use.
By default we do not use any spatial correlations. Permitted
Expand All @@ -539,6 +632,7 @@ def common_red_noise_block(psd='powerlaw', prior='log-uniform',
"""


orfs = {'crn': None, 'hd': model_orfs.hd_orf(),
'gw_monopole': model_orfs.gw_monopole_orf(),
'gw_dipole': model_orfs.gw_dipole_orf(),
Expand All @@ -565,6 +659,18 @@ def common_red_noise_block(psd='powerlaw', prior='log-uniform',
amp_name = '{}_log10_A'.format(name)
if log10_A_val is not None:
log10_Agw = parameter.Constant(log10_A_val)(amp_name)

if logmin is not None and logmax is not None:
if prior == 'uniform':
log10_Agw = parameter.LinearExp(logmin, logmax)(amp_name)
elif prior == 'log-uniform' and gamma_val is not None:
if np.abs(gamma_val - 4.33) < 0.1:
log10_Agw = parameter.Uniform(logmin, logmax)(amp_name)
else:
log10_Agw = parameter.Uniform(logmin, logmax)(amp_name)
else:
log10_Agw = parameter.Uniform(logmin, logmax)(amp_name)

else:
if prior == 'uniform':
log10_Agw = parameter.LinearExp(-18, -11)(amp_name)
Expand Down
122 changes: 68 additions & 54 deletions enterprise_extensions/deterministic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,60 +10,6 @@
from enterprise import constants as const


def bwm_block(Tmin, Tmax, amp_prior='log-uniform',
skyloc=None, logmin=-18, logmax=-11,
name='bwm'):
"""
Returns deterministic GW burst with memory model:
1. Burst event parameterized by time, sky location,
polarization angle, and amplitude
:param Tmin:
Min time to search, probably first TOA (MJD).
:param Tmax:
Max time to search, probably last TOA (MJD).
:param amp_prior:
Prior on log10_A. Default if "log-uniform". Use "uniform" for
upper limits.
:param skyloc:
Fixed sky location of BWM signal search as [cos(theta), phi].
Search over sky location if ``None`` given.
:param logmin:
log of minimum BWM amplitude for prior (log10)
:param logmax:
log of maximum BWM amplitude for prior (log10)
:param name:
Name of BWM signal.
"""

# BWM parameters
amp_name = '{}_log10_A'.format(name)
if amp_prior == 'uniform':
log10_A_bwm = parameter.LinearExp(logmin, logmax)(amp_name)
elif amp_prior == 'log-uniform':
log10_A_bwm = parameter.Uniform(logmin, logmax)(amp_name)

pol_name = '{}_pol'.format(name)
pol = parameter.Uniform(0, np.pi)(pol_name)

t0_name = '{}_t0'.format(name)
t0 = parameter.Uniform(Tmin, Tmax)(t0_name)

costh_name = '{}_costheta'.format(name)
phi_name = '{}_phi'.format(name)
if skyloc is None:
costh = parameter.Uniform(-1, 1)(costh_name)
phi = parameter.Uniform(0, 2*np.pi)(phi_name)
else:
costh = parameter.Constant(skyloc[0])(costh_name)
phi = parameter.Constant(skyloc[1])(phi_name)


# BWM signal
bwm_wf = utils.bwm_delay(log10_h=log10_A_bwm, t0=t0,
cos_gwtheta=costh, gwphi=phi, gwpol=pol)
bwm = deterministic_signals.Deterministic(bwm_wf, name=name)

return bwm


def fdm_block(Tmin, Tmax, amp_prior='log-uniform', name='fdm',
Expand Down Expand Up @@ -461,6 +407,74 @@ def cw_delay(toas, pos, pdist,

return res

@signal_base.function
def bwm_delay(toas, pos, log10_h=-14.0, cos_gwtheta=0.0, gwphi=0.0, gwpol=0.0, t0=55000,
antenna_pattern_fn=None):
"""
Function that calculates the earth-term gravitational-wave
burst-with-memory signal, as described in:
Seto et al, van haasteren and Levin, phsirkov et al, Cordes and Jenet.
This version uses the F+/Fx polarization modes, as verified with the
Continuous Wave and Anisotropy papers.
:param toas: Time-of-arrival measurements [s]
:param pos: Unit vector from Earth to pulsar
:param log10_h: log10 of GW strain
:param cos_gwtheta: Cosine of GW polar angle
:param gwphi: GW azimuthal polar angle [rad]
:param gwpol: GW polarization angle
:param t0: Burst central time [day]
:param antenna_pattern_fn:
User defined function that takes `pos`, `gwtheta`, `gwphi` as
arguments and returns (fplus, fcross)
:return: the waveform as induced timing residuals (seconds)
"""

# convert
h = 10 ** log10_h
gwtheta = np.arccos(cos_gwtheta)
t0 *= const.day

# antenna patterns
if antenna_pattern_fn is None:
apc = create_gw_antenna_pattern(pos, gwtheta, gwphi)
else:
apc = antenna_pattern_fn(pos, gwtheta, gwphi)

# grab fplus, fcross
fp, fc = apc[0], apc[1]

# combined polarization
pol = np.cos(2 * gwpol) * fp + np.sin(2 * gwpol) * fc

# Return the time-series for the pulsar
return pol * h * np.heaviside(toas - t0, 0.5) * (toas - t0)

@signal_base.function
def bwm_sglpsr_delay(toas, sign, log10_A=-15, t0=55000):

"""
Function that calculates the earth-term gravitational-wave
burst-with-memory signal for an optimally oriented source in a single pulsar
:param toas: Time-of-arrival measurements [s]
:param log10_A: log10 of the amplitude of the ramp (delta_f/f)
:param t0: Burst central time [day]
:return: the waveform as induced timing residuals (seconds)
"""


A = 10 ** log10_A
t0 *= const.day
# Return the time-series for the pulsar
heaviside = lambda x: 0.5 * (np.sign(x) + 1)

#return 0 #Fix the return to 0 in order to test what the heck is wrong with red noise detection in bwm
return A * np.sign(sign) * heaviside(toas - t0) * (toas - t0)



@signal_base.function
def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi,
Expand Down
Loading

0 comments on commit 216eb7d

Please sign in to comment.