diff --git a/enterprise_extensions/blocks.py b/enterprise_extensions/blocks.py index e1a4b16e..a3c6c76f 100644 --- a/enterprise_extensions/blocks.py +++ b/enterprise_extensions/blocks.py @@ -36,8 +36,8 @@ def channelized_backends(backend_flags): return {flagval: backend_flags == flagval for flagval in flagvals} -def white_noise_block(vary=False, inc_ecorr=False, gp_ecorr=False, - efac1=False, select='backend', tnequad=False, name=None): +def white_noise_block(vary:bool=False, inc_ecorr:bool=False, gp_ecorr:bool=False, + efac1:bool=False, select:str='backend', tnequad:bool=False, name:str=None): """ Returns the white noise block of the model: diff --git a/enterprise_extensions/deterministic.py b/enterprise_extensions/deterministic.py index dae2b6fa..a78a4595 100644 --- a/enterprise_extensions/deterministic.py +++ b/enterprise_extensions/deterministic.py @@ -48,9 +48,9 @@ def fdm_block(Tmin, Tmax, amp_prior='log-uniform', name='fdm', log10_f_fdm = parameter.Uniform(freq_lower, freq_upper)(freq_name) phase_e_name = '{}_phase_e'.format(name) - phase_e_fdm = parameter.Uniform(0, 2*np.pi)(phase_e_name) + phase_e_fdm = parameter.Uniform(0, 2 * np.pi)(phase_e_name) - phase_p = parameter.Uniform(0, 2*np.pi) + phase_p = parameter.Uniform(0, 2 * np.pi) fdm_wf = fdm_delay(log10_A=log10_A_fdm, log10_f=log10_f_fdm, phase_e=phase_e_fdm, phase_p=phase_p) @@ -113,7 +113,7 @@ def cw_block_circ(amp_prior='log-uniform', dist_prior=None, # orbital inclination angle [radians] cosinc = parameter.Uniform(-1.0, 1.0)('{}_cosinc'.format(name)) # initial GW phase [radians] - phase0 = parameter.Uniform(0.0, 2*np.pi)('{}_phase0'.format(name)) + phase0 = parameter.Uniform(0.0, 2 * np.pi)('{}_phase0'.format(name)) # polarization psi_name = '{}_psi'.format(name) @@ -124,7 +124,7 @@ def cw_block_circ(amp_prior='log-uniform', dist_prior=None, 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) + phi = parameter.Uniform(0, 2 * np.pi)(phi_name) else: costh = parameter.Constant(skyloc[0])(costh_name) phi = parameter.Constant(skyloc[1])(phi_name) @@ -201,7 +201,7 @@ def cw_block_ecc(amp_prior='log-uniform', skyloc=None, log10_F=None, e_0 = parameter.Constant(ecc)('{}_e0'.format(name)) # initial mean anomaly [radians] - l_0 = parameter.Uniform(0.0, 2.0*np.pi)('{}_l0'.format(name)) + l_0 = parameter.Uniform(0.0, 2.0 * np.pi)('{}_l0'.format(name)) # mass ratio = M_2/M_1 q = parameter.Constant(1.0)('{}_q'.format(name)) @@ -214,7 +214,7 @@ def cw_block_ecc(amp_prior='log-uniform', skyloc=None, log10_F=None, 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) + phi = parameter.Uniform(0, 2 * np.pi)(phi_name) else: costh = parameter.Constant(skyloc[0])(costh_name) phi = parameter.Constant(skyloc[1])(phi_name) @@ -234,12 +234,20 @@ def cw_block_ecc(amp_prior='log-uniform', skyloc=None, log10_F=None, @signal_base.function def cw_delay(toas, pos, pdist, - cos_gwtheta=0, gwphi=0, cos_inc=0, - log10_mc=9, log10_fgw=-8, log10_dist=None, log10_h=None, - phase0=0, psi=0, - psrTerm=False, p_dist=1, p_phase=None, - evolve=False, phase_approx=False, check=False, - tref=0): + cos_gwtheta: float = 0, gwphi: float = 0, cos_inc: float = 0, + log10_mc: float = 9, log10_fgw: float = -8, log10_dist: float = None, log10_h: float = None, + phase0: float = 0, psi: float = 0, + psrTerm: bool = False, p_dist: float = 1, p_phase: float = None, + evolve: bool = False, phase_approx: bool = False, check: bool = False, + tref: float = 0): + # cos_gwtheta: parameter.Parameter = 0, gwphi:parameter.Parameter = 0, + # cos_inc: parameter.Parameter = 0, + # log10_mc: parameter.Parameter = 9, log10_fgw:parameter.Parameter = -8, + # log10_dist: parameter.Parameter = None, log10_h:parameter.Parameter = None, + # phase0: parameter.Parameter = 0, psi:parameter.Parameter = 0, + # psrTerm: bool = False, p_dist:parameter.Parameter = 1, p_phase:parameter.Parameter = None, + # evolve: bool = False, phase_approx:bool = False, check:bool = False, + # tref: float = 0): """ Function to create GW incuced residuals from a SMBMB as defined in Ellis et. al 2012,2013. @@ -291,30 +299,30 @@ def cw_delay(toas, pos, pdist, """ # convert units to time - mc = 10**log10_mc * const.Tsun - fgw = 10**log10_fgw + mc = 10 ** log10_mc * const.Tsun + fgw = 10 ** log10_fgw gwtheta = np.arccos(cos_gwtheta) inc = np.arccos(cos_inc) - p_dist = (pdist[0] + pdist[1]*p_dist)*const.kpc/const.c + p_dist = (pdist[0] + pdist[1] * p_dist) * const.kpc / const.c if log10_h is None and log10_dist is None: raise ValueError("one of log10_dist or log10_h must be non-None") elif log10_h is not None and log10_dist is not None: raise ValueError("only one of log10_dist or log10_h can be non-None") elif log10_h is None: - dist = 10**log10_dist * const.Mpc / const.c + dist = 10 ** log10_dist * const.Mpc / const.c else: - dist = 2 * mc**(5/3) * (np.pi*fgw)**(2/3) / 10**log10_h + dist = 2 * mc ** (5 / 3) * (np.pi * fgw) ** (2 / 3) / 10 ** log10_h if check: # check that frequency is not evolving significantly over obs. time - fstart = fgw * (1 - 256/5 * mc**(5/3) * fgw**(8/3) * toas[0])**(-3/8) - fend = fgw * (1 - 256/5 * mc**(5/3) * fgw**(8/3) * toas[-1])**(-3/8) + fstart = fgw * (1 - 256 / 5 * mc ** (5 / 3) * fgw ** (8 / 3) * toas[0]) ** (-3 / 8) + fend = fgw * (1 - 256 / 5 * mc ** (5 / 3) * fgw ** (8 / 3) * toas[-1]) ** (-3 / 8) df = fend - fstart # observation time - Tobs = toas.max()-toas.min() - fbin = 1/Tobs + Tobs = toas.max() - toas.min() + fbin = 1 / Tobs if np.abs(df) > fbin: print('WARNING: Frequency is evolving over more than one ' @@ -329,7 +337,7 @@ def cw_delay(toas, pos, pdist, # get pulsar time toas -= tref if p_dist > 0: - tp = toas-p_dist*(1-cosMu) + tp = toas - p_dist * (1 - cosMu) else: tp = toas @@ -341,45 +349,45 @@ def cw_delay(toas, pos, pdist, # evolution if evolve: # calculate time dependent frequency at earth and pulsar - omega = w0 * (1 - 256/5 * mc**(5/3) * w0**(8/3) * toas)**(-3/8) - omega_p = w0 * (1 - 256/5 * mc**(5/3) * w0**(8/3) * tp)**(-3/8) + omega = w0 * (1 - 256 / 5 * mc ** (5 / 3) * w0 ** (8 / 3) * toas) ** (-3 / 8) + omega_p = w0 * (1 - 256 / 5 * mc ** (5 / 3) * w0 ** (8 / 3) * tp) ** (-3 / 8) if p_dist > 0: - omega_p0 = w0 * (1 + 256/5 - * mc**(5/3) * w0**(8/3) * p_dist*(1-cosMu))**(-3/8) + omega_p0 = w0 * (1 + 256 / 5 + * mc ** (5 / 3) * w0 ** (8 / 3) * p_dist * (1 - cosMu)) ** (-3 / 8) else: omega_p0 = w0 # calculate time dependent phase - phase = phase0 + 1/32/mc**(5/3) * (w0**(-5/3) - omega**(-5/3)) + phase = phase0 + 1 / 32 / mc ** (5 / 3) * (w0 ** (-5 / 3) - omega ** (-5 / 3)) if p_phase is None: - phase_p = phase0 + 1/32/mc**(5/3) * (w0**(-5/3) - omega_p**(-5/3)) + phase_p = phase0 + 1 / 32 / mc ** (5 / 3) * (w0 ** (-5 / 3) - omega_p ** (-5 / 3)) else: phase_p = (phase0 + p_phase - + 1/32*mc**(-5/3) * (omega_p0**(-5/3) - omega_p**(-5/3))) + + 1 / 32 * mc ** (-5 / 3) * (omega_p0 ** (-5 / 3) - omega_p ** (-5 / 3))) elif phase_approx: # monochromatic omega = w0 if p_dist > 0: - omega_p = w0 * (1 + 256/5 - * mc**(5/3) * w0**(8/3) * p_dist*(1-cosMu))**(-3/8) + omega_p = w0 * (1 + 256 / 5 + * mc ** (5 / 3) * w0 ** (8 / 3) * p_dist * (1 - cosMu)) ** (-3 / 8) else: omega_p = w0 # phases phase = phase0 + omega * toas if p_phase is not None: - phase_p = phase0 + p_phase + omega_p*toas + phase_p = phase0 + p_phase + omega_p * toas else: - phase_p = (phase0 + omega_p*toas - + 1/32/mc**(5/3) * (w0**(-5/3) - omega_p**(-5/3))) + phase_p = (phase0 + omega_p * toas + + 1 / 32 / mc ** (5 / 3) * (w0 ** (-5 / 3) - omega_p ** (-5 / 3))) # no evolution else: # monochromatic - omega = np.pi*fgw + omega = np.pi * fgw omega_p = omega # phases @@ -387,26 +395,26 @@ def cw_delay(toas, pos, pdist, phase_p = phase0 + omega * tp # define time dependent coefficients - At = -0.5*np.sin(2*phase)*(3+np.cos(2*inc)) - Bt = 2*np.cos(2*phase)*np.cos(inc) - At_p = -0.5*np.sin(2*phase_p)*(3+np.cos(2*inc)) - Bt_p = 2*np.cos(2*phase_p)*np.cos(inc) + At = -0.5 * np.sin(2 * phase) * (3 + np.cos(2 * inc)) + Bt = 2 * np.cos(2 * phase) * np.cos(inc) + At_p = -0.5 * np.sin(2 * phase_p) * (3 + np.cos(2 * inc)) + Bt_p = 2 * np.cos(2 * phase_p) * np.cos(inc) # now define time dependent amplitudes - alpha = mc**(5./3.)/(dist*omega**(1./3.)) - alpha_p = mc**(5./3.)/(dist*omega_p**(1./3.)) + alpha = mc ** (5. / 3.) / (dist * omega ** (1. / 3.)) + alpha_p = mc ** (5. / 3.) / (dist * omega_p ** (1. / 3.)) # define rplus and rcross - rplus = alpha*(-At*np.cos(2*psi)+Bt*np.sin(2*psi)) - rcross = alpha*(At*np.sin(2*psi)+Bt*np.cos(2*psi)) - rplus_p = alpha_p*(-At_p*np.cos(2*psi)+Bt_p*np.sin(2*psi)) - rcross_p = alpha_p*(At_p*np.sin(2*psi)+Bt_p*np.cos(2*psi)) + rplus = alpha * (-At * np.cos(2 * psi) + Bt * np.sin(2 * psi)) + rcross = alpha * (At * np.sin(2 * psi) + Bt * np.cos(2 * psi)) + rplus_p = alpha_p * (-At_p * np.cos(2 * psi) + Bt_p * np.sin(2 * psi)) + rcross_p = alpha_p * (At_p * np.sin(2 * psi) + Bt_p * np.cos(2 * psi)) # residuals if psrTerm: - res = fplus*(rplus_p-rplus)+fcross*(rcross_p-rcross) + res = fplus * (rplus_p - rplus) + fcross * (rcross_p - rcross) else: - res = -fplus*rplus - fcross*rcross + res = -fplus * rplus - fcross * rcross return res @@ -520,11 +528,11 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi, """ # convert from sampling - F = 10.0**log10_F - mc = 10.0**log10_mc - dist = 10.0**log10_dist + F = 10.0 ** log10_F + mc = 10.0 ** log10_mc + dist = 10.0 ** log10_dist if log10_h is not None: - h0 = 10.0**log10_h + h0 = 10.0 ** log10_h else: h0 = None inc = np.arccos(cos_inc) @@ -533,19 +541,19 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi, # define variable for later use cosgwtheta, cosgwphi = np.cos(gwtheta), np.cos(gwphi) singwtheta, singwphi = np.sin(gwtheta), np.sin(gwphi) - sin2psi, cos2psi = np.sin(2*psi), np.cos(2*psi) + sin2psi, cos2psi = np.sin(2 * psi), np.cos(2 * psi) # unit vectors to GW source m = np.array([singwphi, -cosgwphi, 0.0]) - n = np.array([-cosgwtheta*cosgwphi, -cosgwtheta*singwphi, singwtheta]) - omhat = np.array([-singwtheta*cosgwphi, -singwtheta*singwphi, -cosgwtheta]) + n = np.array([-cosgwtheta * cosgwphi, -cosgwtheta * singwphi, singwtheta]) + omhat = np.array([-singwtheta * cosgwphi, -singwtheta * singwphi, -cosgwtheta]) # pulsar position vector - phat = np.array([np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi), + phat = np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]) - fplus = 0.5 * (np.dot(m, phat)**2 - np.dot(n, phat)**2) / (1+np.dot(omhat, phat)) - fcross = (np.dot(m, phat)*np.dot(n, phat)) / (1 + np.dot(omhat, phat)) + fplus = 0.5 * (np.dot(m, phat) ** 2 - np.dot(n, phat) ** 2) / (1 + np.dot(omhat, phat)) + fcross = (np.dot(m, phat) * np.dot(n, phat)) / (1 + np.dot(omhat, phat)) cosMu = -np.dot(omhat, phat) # get values from pulsar object @@ -561,11 +569,11 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi, Fc1, ec1, gc1, phic1 = y[-1, :] # observation time - Tobs = 1/(toas.max()-toas.min()) + Tobs = 1 / (toas.max() - toas.min()) - if np.abs(Fc0-Fc1) > 1/Tobs: + if np.abs(Fc0 - Fc1) > 1 / Tobs: print('WARNING: Frequency is evolving over more than one frequency bin.') - print('F0 = {0}, F1 = {1}, delta f = {2}'.format(Fc0, Fc1, 1/Tobs)) + print('F0 = {0}, F1 = {1}, delta f = {2}'.format(Fc0, Fc1, 1 / Tobs)) return np.ones(len(toas)) * np.nan # get gammadot for earth term @@ -600,7 +608,7 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi, pd *= const.kpc / const.c # get pulsar time - tp = toas.copy() - pd * (1-cosMu) + tp = toas.copy() - pd * (1 - cosMu) # solve coupled system of equations to get pulsar term values y = utils.solve_coupled_ecc_solution(F, e0, gamma0, l0, mc, @@ -646,21 +654,20 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi, gammadot=gammadotp, inc=inc) - rr = (fplus*cos2psi - fcross*sin2psi) * (splusp - splus) + \ - (fplus*sin2psi + fcross*cos2psi) * (scrossp - scross) + rr = (fplus * cos2psi - fcross * sin2psi) * (splusp - splus) + \ + (fplus * sin2psi + fcross * cos2psi) * (scrossp - scross) else: rr = np.ones(len(toas)) * np.nan else: - rr = - (fplus*cos2psi - fcross*sin2psi) * splus - \ - (fplus*sin2psi + fcross*cos2psi) * scross + rr = - (fplus * cos2psi - fcross * sin2psi) * splus - \ + (fplus * sin2psi + fcross * cos2psi) * scross return rr -def CWSignal(cw_wf, ecc=False, psrTerm=False, name='cw'): - +def CWSignal(cw_wf, ecc: bool = False, psrTerm: bool = False, name: str = 'cw'): BaseClass = deterministic_signals.Deterministic(cw_wf, name=name) class CWSignal(BaseClass): @@ -669,9 +676,9 @@ def __init__(self, psr): super(CWSignal, self).__init__(psr) self._wf[''].add_kwarg(psrTerm=psrTerm) if ecc: - pgam = parameter.Uniform(0, 2*np.pi)('_'.join([psr.name, - 'pgam', - name])) + pgam = parameter.Uniform(0, 2 * np.pi)('_'.join([psr.name, + 'pgam', + name])) self._params['pgam'] = pgam self._wf['']._params['pgam'] = pgam @@ -681,7 +688,7 @@ def __init__(self, psr): @signal_base.function def generalized_gwpol_psd(f, log10_A_tt=-15, log10_A_st=-15, log10_A_vl=-15, log10_A_sl=-15, - kappa=10/3, p_dist=1.0): + kappa=10 / 3, p_dist=1.0): """ PSD for a generalized mixture of scalar+vector dipole radiation and tensorial quadrupole radiation from SMBHBs. @@ -691,24 +698,24 @@ def generalized_gwpol_psd(f, log10_A_tt=-15, log10_A_st=-15, euler_e = 0.5772156649 pdist = p_dist * const.kpc / const.c - orf_aa_tt = (2/3) * np.ones(len(f)) - orf_aa_st = (2/3) * np.ones(len(f)) - orf_aa_vl = 2*np.log(4*np.pi*f*pdist) - 14/3 + 2*euler_e - orf_aa_sl = np.pi**2*f*pdist/4 - \ - np.log(4*np.pi*f*pdist) + 37/24 - euler_e - - prefactor = (1 + kappa**2) / (1 + kappa**2 * (f / const.fyr)**(-2/3)) - gwpol_amps = 10**(2*np.array([log10_A_tt, log10_A_st, - log10_A_vl, log10_A_sl])) - gwpol_factors = np.array([orf_aa_tt*gwpol_amps[0], - orf_aa_st*gwpol_amps[1], - orf_aa_vl*gwpol_amps[2], - orf_aa_sl*gwpol_amps[3]]) - - S_psd = prefactor * (gwpol_factors[0, :] * (f / const.fyr)**(-4/3) + + orf_aa_tt = (2 / 3) * np.ones(len(f)) + orf_aa_st = (2 / 3) * np.ones(len(f)) + orf_aa_vl = 2 * np.log(4 * np.pi * f * pdist) - 14 / 3 + 2 * euler_e + orf_aa_sl = np.pi ** 2 * f * pdist / 4 - \ + np.log(4 * np.pi * f * pdist) + 37 / 24 - euler_e + + prefactor = (1 + kappa ** 2) / (1 + kappa ** 2 * (f / const.fyr) ** (-2 / 3)) + gwpol_amps = 10 ** (2 * np.array([log10_A_tt, log10_A_st, + log10_A_vl, log10_A_sl])) + gwpol_factors = np.array([orf_aa_tt * gwpol_amps[0], + orf_aa_st * gwpol_amps[1], + orf_aa_vl * gwpol_amps[2], + orf_aa_sl * gwpol_amps[3]]) + + S_psd = prefactor * (gwpol_factors[0, :] * (f / const.fyr) ** (-4 / 3) + np.sum(gwpol_factors[1:, :], axis=0) * - (f / const.fyr)**(-2)) / \ - (8*np.pi**2*f**3) + (f / const.fyr) ** (-2)) / \ + (8 * np.pi ** 2 * f ** 3) return S_psd * np.repeat(df, 2) diff --git a/enterprise_extensions/models.py b/enterprise_extensions/models.py index 6f5d0c63..e458ff21 100644 --- a/enterprise_extensions/models.py +++ b/enterprise_extensions/models.py @@ -349,10 +349,10 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, return pta -def model_1(psrs, psd='powerlaw', noisedict=None, white_vary=False, - components=30, upper_limit=False, bayesephem=False, tnequad=False, - be_type='orbel', is_wideband=False, use_dmdata=False, - select='backend', tm_marg=False, dense_like=False, tm_svd=False): +def model_1(psrs:list, psd:str='powerlaw', noisedict:dict=None, white_vary:bool=False, + components:int=30, upper_limit:bool=False, bayesephem:bool=False, tnequad:bool=False, + be_type:str='orbel', is_wideband:bool=False, use_dmdata:bool=False, + select:str='backend', tm_marg:bool=False, dense_like:bool=False, tm_svd:bool=False): """ Reads in list of enterprise Pulsar instance and returns a PTA instantiated with only white and red noise: @@ -461,13 +461,14 @@ def model_1(psrs, psd='powerlaw', noisedict=None, white_vary=False, return pta -def model_2a(psrs, psd='powerlaw', noisedict=None, components=30, - n_rnfreqs=None, n_gwbfreqs=None, gamma_common=None, - delta_common=None, upper_limit=False, bayesephem=False, - be_type='setIII', white_vary=False, is_wideband=False, - use_dmdata=False, select='backend', tnequad=False, - pshift=False, pseed=None, psr_models=False, - tm_marg=False, dense_like=False, tm_svd=False): +def model_2a(psrs: list, psd: str='powerlaw', noisedict: dict=None, components: int=30, + n_rnfreqs: int=None, n_gwbfreqs: int=None, gamma_common: float=None, + delta_common: float=None, upper_limit: int=False, bayesephem: int=False, + be_type: str='setIII', white_vary: bool=False, is_wideband: bool=False, + use_dmdata: bool=False, select: str='backend', tnequad: bool=False, + pshift: bool=False, pseed: int=None, psr_models: bool=False, + tm_marg: bool=False, dense_like: bool=False, tm_svd: bool=False, + gp_ecorr: bool=False, individual_Tspan_for_red_noise: bool=False): """ Reads in list of enterprise Pulsar instance and returns a PTA instantiated with model 2A from the analysis paper: @@ -526,12 +527,20 @@ def model_2a(psrs, psd='powerlaw', noisedict=None, components=30, up the likelihood calculation significantly. :param dense_like: Use dense or sparse functions to evalute lnlikelihood :param tm_svd: boolean for svd-stabilised timing model design matrix + :param gp_ecorr: + whether to use the Gaussian process model for ECORR + :param individual_Tspan_for_red_noise: + each pulsar has frequencies defined by 1/Tspan_pulsar for the red noise """ amp_prior = 'uniform' if upper_limit else 'log-uniform' # find the maximum time span to set GW frequency sampling Tspan = model_utils.get_tspan(psrs) + if individual_Tspan_for_red_noise: + Tspan_red_noise = None + else: + Tspan_red_noise = Tspan if n_gwbfreqs is None: n_gwbfreqs = components @@ -563,7 +572,7 @@ def model_2a(psrs, psd='powerlaw', noisedict=None, components=30, s = gp_signals.TimingModel(use_svd=tm_svd) # red noise - s += red_noise_block(prior=amp_prior, Tspan=Tspan, components=n_rnfreqs) + s += red_noise_block(prior=amp_prior, Tspan=Tspan_red_noise, components=n_rnfreqs) # common red noise block s += common_red_noise_block(psd=psd, prior=amp_prior, Tspan=Tspan, @@ -580,11 +589,11 @@ def model_2a(psrs, psd='powerlaw', noisedict=None, components=30, models = [] for p in psrs: if 'NANOGrav' in p.flags['pta'] and not is_wideband: - s2 = s + white_noise_block(vary=white_vary, inc_ecorr=True, + s2 = s + white_noise_block(vary=white_vary, inc_ecorr=True, gp_ecorr=gp_ecorr, tnequad=tnequad, select=select) models.append(s2(p)) else: - s3 = s + white_noise_block(vary=white_vary, inc_ecorr=False, + s3 = s + white_noise_block(vary=white_vary, inc_ecorr=False, gp_ecorr=gp_ecorr, tnequad=tnequad, select=select) models.append(s3(p)) @@ -610,7 +619,7 @@ def model_2a(psrs, psd='powerlaw', noisedict=None, components=30, noisedict = noisedict pta.set_default_params(noisedict) - return pta + return pta def model_general(psrs, tm_var=False, tm_linear=False, tmparam_list=None, @@ -1312,14 +1321,14 @@ def model_2d(psrs, psd='powerlaw', noisedict=None, white_vary=False, return pta -def model_3a(psrs, psd='powerlaw', noisedict=None, white_vary=False, - components=30, n_rnfreqs=None, n_gwbfreqs=None, - gamma_common=None, delta_common=None, upper_limit=False, - bayesephem=False, be_type='setIII', is_wideband=False, - use_dmdata=False, select='backend', - tnequad=False, - pshift=False, pseed=None, psr_models=False, - tm_marg=False, dense_like=False, tm_svd=False): +def model_3a(psrs: list, psd: str='powerlaw', noisedict: dict=None, white_vary: bool=False, + components: int=30, n_rnfreqs: int=None, n_gwbfreqs: int=None, + gamma_common: float=None, delta_common: float=None, upper_limit: bool=False, + bayesephem: bool=False, be_type: str='setIII', is_wideband: bool=False, + use_dmdata: bool=False, select: str='backend', tnequad: bool=False, + pshift: bool=False, pseed: int=None, psr_models: bool=False, + tm_marg: bool=False, dense_like: bool=False, tm_svd: bool=False, + gp_ecorr: bool=False, individual_Tspan_for_red_noise: bool=False): """ Reads in list of enterprise Pulsar instance and returns a PTA instantiated with model 3A from the analysis paper: @@ -1378,12 +1387,20 @@ def model_3a(psrs, psd='powerlaw', noisedict=None, white_vary=False, up the likelihood calculation significantly. :param dense_like: Use dense or sparse functions to evalute lnlikelihood :param tm_svd: boolean for svd-stabilised timing model design matrix + :param gp_ecorr: + whether to use the Gaussian process model for ECORR + :param individual_Tspan_for_red_noise: + each pulsar has frequencies defined by 1/Tspan_pulsar for the red noise """ amp_prior = 'uniform' if upper_limit else 'log-uniform' # find the maximum time span to set GW frequency sampling Tspan = model_utils.get_tspan(psrs) + if individual_Tspan_for_red_noise: + Tspan_red_noise = None + else: + Tspan_red_noise = Tspan if n_gwbfreqs is None: n_gwbfreqs = components @@ -1415,7 +1432,7 @@ def model_3a(psrs, psd='powerlaw', noisedict=None, white_vary=False, s = gp_signals.TimingModel(use_svd=tm_svd) # red noise - s += red_noise_block(prior=amp_prior, Tspan=Tspan, components=n_rnfreqs) + s += red_noise_block(prior=amp_prior, Tspan=Tspan_red_noise, components=n_rnfreqs) # common red noise block s += common_red_noise_block(psd=psd, prior=amp_prior, Tspan=Tspan, @@ -1432,11 +1449,13 @@ def model_3a(psrs, psd='powerlaw', noisedict=None, white_vary=False, models = [] for p in psrs: if 'NANOGrav' in p.flags['pta'] and not is_wideband: - s2 = s + white_noise_block(vary=white_vary, inc_ecorr=True, + s2 = s + white_noise_block(vary=white_vary, + gp_ecorr=gp_ecorr, inc_ecorr=True, tnequad=tnequad, select=select) models.append(s2(p)) else: - s3 = s + white_noise_block(vary=white_vary, inc_ecorr=False, + s3 = s + white_noise_block(vary=white_vary, + gp_ecorr=gp_ecorr, inc_ecorr=False, tnequad=tnequad, select=select) models.append(s3(p)) diff --git a/enterprise_extensions/pipe/__init__.py b/enterprise_extensions/pipe/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/enterprise_extensions/pipe/data_class.py b/enterprise_extensions/pipe/data_class.py new file mode 100644 index 00000000..a155d412 --- /dev/null +++ b/enterprise_extensions/pipe/data_class.py @@ -0,0 +1,375 @@ +import collections.abc +import configparser +import importlib +import inspect +import json +import pickle +import re +from dataclasses import dataclass, field + +import numpy as np +from enterprise.signals import signal_base +import enterprise.signals.parameter # this is used but only implicitly +import enterprise_extensions.models + + +def get_default_args_types_from_function(func): + """ + Given function, returns two dictionaries with default values and types + code modified from: https://stackoverflow.com/questions/12627118/get-a-function-arguments-default-value + """ + signature = inspect.signature(func) + defaults = {} + types = {} + for key, value in signature.parameters.items(): + # get default kwarg value from function + if value.default is not inspect.Parameter.empty: + defaults[key] = value.default + + # get type annotation from function + if value.annotation is inspect.Parameter.empty: + print(f"Warning! in {func} {value} does not have an associated type annotation") + else: + types[key] = value.annotation + return defaults, types + + +def update_dictionary_with_subdictionary(d, u): + """ + Updates dictionary d with preference for contents of dictionary u + code taken from: https://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth + """ + if d is None: + return u + for k, v in u.items(): + if isinstance(v, collections.abc.Mapping): + d[k] = update_dictionary_with_subdictionary(d.get(k, {}), v) + else: + d[k] = v + return d + + +def load_module_globally(package_dict): + # import modules globablly from dictionary import key as item + for import_as, package_name in package_dict.items(): + mod = importlib.import_module(package_name) + globals()[import_as] = mod + return + + +@dataclass() +class RunSettings: + """ + Class for keeping track of enterprise model run settings + TODO: link to examples of how to use + """ + pulsar_pickle: str = None + noise_dict_json: str = None + + # dictionary of functions that create signals + signal_creating_function_keys: list = field(default_factory=list) + + # dictionary of functions that create signals depending on parameters of each pulsar + per_pulsar_signal_creating_function_keys: dict = field(default_factory=dict) + + # dictionary of functions that create pta objects + pta_creating_function_keys: list = field(default_factory=list) + + custom_classes: dict = field(default_factory=dict) + custom_class_parameters: dict = field(default_factory=dict) + + function_parameters: dict = field(default_factory=dict) + functions: dict = field(default_factory=dict) + custom_return: dict = field(default_factory=dict) + + psrs: list = field(default_factory=list) + noise_dict: dict = field(default_factory=dict) + # stores config file as dictionary + config_file_items: dict = field(default_factory=dict) + + def update_from_file(self, config_file: str) -> None: + """ + Set defaults for functions from file + + [modules]: example: np=numpy will load numpy as np globally + """ + config = configparser.ConfigParser(comment_prefixes=';', + interpolation=configparser.ExtendedInterpolation()) + config.optionxform = str + config.read(config_file) + exclude_keys = ['function', 'module', 'class', 'signal_return', 'pta_return', + 'custom_return', 'per_pulsar_signal'] + for section in config.sections(): + config_file_items = dict(config.items(section)) + self.config_file_items[section] = config_file_items + if section == 'modules': + load_module_globally(config_file_items) + elif 'class' in config_file_items.keys(): + """ + Initialize a class given in a config file + """ + if 'module' in config_file_items.keys(): + # get a class defined from a module + module = importlib.import_module(config_file_items['module']) + custom_class = getattr(module, config_file_items['class']) + else: + # or if class module has already been imported, do this! + custom_class = eval(config_file_items['class']) + + default_class_parameters, types = get_default_args_types_from_function(custom_class.__init__) + class_parameters_from_file = self.apply_types(config_file_items, types, + exclude_keys=exclude_keys) + class_parameters = update_dictionary_with_subdictionary(default_class_parameters, + class_parameters_from_file) + self.custom_classes[section] = custom_class(**class_parameters) + self.custom_class_parameters[section] = class_parameters + + elif 'function' in config_file_items.keys(): + if 'module' in config_file_items.keys(): + # get a class defined from a module + module = importlib.import_module(config_file_items['module']) + custom_function = getattr(module, config_file_items['function']) + else: + # or if class module has already been imported, do this! + custom_function = eval(config_file_items['function']) + + try: + default_function_parameters, types = get_default_args_types_from_function(custom_function) + except ValueError: + # builtin functions like dictionary and list don't have default types + default_function_parameters, types = {}, {} + function_parameters_from_file = self.apply_types(config_file_items, types, + exclude_keys=exclude_keys) + self.functions[section] = custom_function + self.function_parameters[section] = update_dictionary_with_subdictionary( + default_function_parameters, + function_parameters_from_file) + + if 'custom_return' in config_file_items.keys(): + # custom_return means to store the return value of this function in self.custom_function_return + self.custom_return[config_file_items['custom_return']] = \ + self.functions[section](**self.function_parameters[section]) + elif 'signal_return' in config_file_items.keys(): + # label this function as something that returns signal models + self.signal_creating_function_keys.append(section) + elif 'per_pulsar_signal' in config_file_items.keys(): + # Per pulsar can either be used as a function applied to every pulsar, or specify one by name + if (config_file_items['per_pulsar_signal'] == 'EVERY_PULSAR') \ + or (config_file_items['per_pulsar_signal'] == 'True'): + try: + self.per_pulsar_signal_creating_function_keys['EVERY_PULSAR'] + except KeyError: + self.per_pulsar_signal_creating_function_keys['EVERY_PULSAR'] = [] + self.per_pulsar_signal_creating_function_keys['EVERY_PULSAR'].append(section) + else: + pulsar_name = config_file_items['per_pulsar_signal'] + try: + self.per_pulsar_signal_creating_function_keys[pulsar_name] + except KeyError: + self.per_pulsar_signal_creating_function_keys[pulsar_name] = [] + self.per_pulsar_signal_creating_function_keys[pulsar_name].append(section) + + elif 'pta_return' in config_file_items.keys(): + # label this function as something that returns ptas + self.pta_creating_function_keys.append(section) + if 'singular_pulsar' in config_file_items.keys(): + # Only apply to given pulsar + pulsar_name = config_file_items['singular_pulsar'] + try: + self.singular_pulsar_function_keys[pulsar_name] + except KeyError: + self.singular_pulsar_function_keys[pulsar_name] = [] + self.singular_pulsar_function_keys[pulsar_name].append(section) + + else: + # If not a class or function or module + # it must be something specified in the RunSettings class + # now read those in from the file + for item in config_file_items.copy(): + if not config_file_items[item]: + config_file_items.pop(item) + self.update_from_dict(**config_file_items) + + def apply_types(self, dictionary, type_dictionary, exclude_keys=None): + """ + Given dictionary (usually created from config_file) and dictionary containing types + apply type to dictionary + + if CUSTOM_CLASS:your_class is in dictionary[key], + instead of applying type it assigns from self.custom_classes + if CUSTOM_FUNCTION_RETURN:whatever is in dictionary[key] + instead of applying type it assigns from self.custom_returns[whatever] + if EVAL:whatever is in dictionary[key] + will call eval("whatever") and assign that + """ + if exclude_keys is None: + exclude_keys = [] + out_dictionary = {} + for key, value in dictionary.items(): + if key in exclude_keys: + continue + if 'CUSTOM_FUNCTION_RETURN:' in value or 'CUSTOM_RETURN' in value: + if 'CUSTOM_FUNCTION_RETURN:' in value: + value = value.replace('CUSTOM_FUNCTION_RETURN', 'CUSTOM_RETURN') + print("CUSTOM_FUNCTION_RETURN has been renamed CUSTOM_RETURN, please use that instead") + out_dictionary[key] = self.custom_return[value.replace('CUSTOM_RETURN:', '')] + continue + # Apply custom class instance stored in custom_classes + if 'CUSTOM_CLASS:' in value: + # Apply custom class instance stored in custom_classes + out_dictionary[key] = self.custom_classes[value.replace('CUSTOM_CLASS:', '')] + continue + if 'EVAL:' in value: + function_call = value.replace('EVAL:', '') + out_dictionary[key] = eval(function_call) + continue + if key not in type_dictionary.keys(): + print(f"WARNING! apply_types: {key} is not within type dictionary!") + print(f"\tObject value is {value} and type is {type(value)}") + print("\tIgnoring value and continuing") + continue + # special comprehension for (1d) numpy arrays + if type_dictionary[key] == np.ndarray: + out_dictionary[key] = np.array([np.float(x) for x in value.split(',')]) + # Special comprehension for bool because otherwise bool('False') == True + elif type_dictionary[key] == bool: + out_dictionary[key] = dictionary[key].lower().capitalize() == "True" + else: + out_dictionary[key] = type_dictionary[key](value) + + return out_dictionary + + def update_from_dict(self, **kwargs): + """ + Given a dictionary, RunSettings attempts to update itself + Will apply dtypes taken from its own definition to those kwargs + """ + ann = getattr(self, "__annotations__", {}) + for name, dtype in ann.items(): + if name in kwargs: + try: + kwargs[name] = dtype(kwargs[name]) + except TypeError: + pass + setattr(self, name, kwargs[name]) + print(f'WARNING: {set(kwargs.keys()) - set(ann.keys())} arguments are getting ignored!') + + def load_pickled_pulsars(self): + """ + Set self.psrs and self.noise_dict + """ + + try: + self.psrs = self.get_pulsars() + self.noise_dict = self.get_noise_dict() + except FileNotFoundError as e: + print(e) + exit(1) + + for par in list(self.noise_dict.keys()): + if 'log10_ecorr' in par and 'basis_ecorr' not in par: + ecorr = par.split('_')[0] + '_basis_ecorr_' + '_'.join(par.split('_')[1:]) + self.noise_dict[ecorr] = self.noise_dict[par] + + # assign noisedict to all enterprise models + for key in self.pta_creating_function_keys: + if 'noisedict' in self.function_parameters[key].keys(): + self.function_parameters[key]['noisedict'] = self.noise_dict + + def get_pulsars(self): + if len(self.psrs) == 0: + self.psrs = pickle.load(open(self.pulsar_pickle, 'rb')) + return self.psrs + + def get_noise_dict(self): + if len(self.noise_dict.keys()) == 0: + self.noise_dict = json.load(open(self.noise_dict_json)) + + for par in list(self.noise_dict.keys()): + if 'log10_equad' in par: + efac = re.sub('log10_equad', 'efac', par) + equad = re.sub('log10_equad', 'log10_t2equad', par) + self.noise_dict[equad] = np.log10(10 ** self.noise_dict[par] / self.noise_dict[efac]) + elif 'log10_ecorr' in par and 'basis_ecorr' not in par: + ecorr = par.split('_')[0] + '_basis_ecorr_' + '_'.join(par.split('_')[1:]) + self.noise_dict[ecorr] = self.noise_dict[par] + return self.noise_dict + + def create_pta_object_from_signals(self): + """ + Using both signals from pta objects and signals from self.signal_creating_functions + Create a pta object + """ + self.load_pickled_pulsars() + + pta_list = self.get_pta_objects() + signal_collections = [self.get_signal_collection_from_pta_object(pta) for pta in pta_list] + for key in self.signal_creating_function_keys: + func = self.functions[key] + signal_collections.append(func(**self.function_parameters[key])) + + signal_collection = sum(signal_collections[1:], signal_collections[0]) + + model_list = [] + # get list of functions to apply to each pulsar + try: + function_keys_for_every_pulsar = self.per_pulsar_signal_creating_function_keys['EVERY_PULSAR'] + except KeyError: + function_keys_for_every_pulsar = [] + + for psr in self.psrs: + # get list of functions to only apply to this pulsar + try: + keys_for_this_pulsar = self.per_pulsar_signal_creating_function_keys[psr.name] + except KeyError: + keys_for_this_pulsar = [] + keys_for_this_pulsar.extend(function_keys_for_every_pulsar) + + per_pulsar_signal = [] + for key in keys_for_this_pulsar: + # TODO this will only work if the parameter is named psr or pulsar, + # Todo maybe it should look for all caps ENTERPRISE_PULSAR string so we can make it more general? + if 'psr' in self.function_parameters[key]: + self.function_parameters[key]['psr'] = psr + if 'pulsar' in self.function_parameters[key]: + self.function_parameters[key]['pulsar'] = psr + if 'enterprise_pulsar' in self.function_parameters[key]: + self.function_parameters[key]['enterprise_pulsar'] = psr + try: + # this allows each pulsar to have signals applied to them + per_pulsar_signal.append(self.functions[key](**self.function_parameters[key])) + except TypeError as e: + print(e) + raise TypeError("ERROR to use per_pulsar_signal to call functions depending on enterprise pulsar," + "\n\t that function must have arguments named one of 'psr', 'pulsar', or 'enterprise_pulsar'") + + # just sums to signal_collection if additional_models is empty + model_list.append(sum(per_pulsar_signal, signal_collection)(psr)) + + pta = signal_base.PTA(model_list) + + # apply noise dictionary to pta + pta.set_default_params(self.noise_dict) + # return pta object + return pta + + def get_signal_collection_from_pta_object(self, pta): + """ + Under assumption that same model has been applied to ALL pulsars + and that there are pulsars inside of this pta, + get signal collection from this pta object + """ + return type(pta.pulsarmodels[0]) + + def get_pta_objects(self): + """ + Using pta creating functions specified in config, get list of pta objects + """ + pta_list = [] + if len(self.psrs) == 0: + print("Loading pulsars") + self.load_pickled_pulsars() + + for key in self.pta_creating_function_keys: + pta_list.append(self.functions[key](psrs=self.psrs, **self.function_parameters[key])) + return pta_list diff --git a/enterprise_extensions/pipe/example.ini b/enterprise_extensions/pipe/example.ini new file mode 100644 index 00000000..c3778c04 --- /dev/null +++ b/enterprise_extensions/pipe/example.ini @@ -0,0 +1,45 @@ +[input] +DATA_DIR = /home/sophie.hourihane/data/PTA/15yr/15yr_v1.1 +pulsar_pickle = ${DATA_DIR}/v1p1_de440_pint_bipm2019.pkl +noise_dict_json = ${DATA_DIR}/v1p1_wn_dict.json +parameter=enterprise.signals.parameter + +[output] +output_directory=test + +[cw_delay] +module=enterprise_extensions.deterministic +function=cw_delay +custom_return=cw_delay_residual +evolve=True +psrTerm=True +cos_gwtheta=EVAL:enterprise.signals.parameter.Uniform(-1,1)('0_cos_gwtheta') +gwphi=EVAL:enterprise.signals.parameter.Uniform(0,2*np.pi)('0_gwphi') +log10_h =EVAL:enterprise.signals.parameter.Uniform(-18, -11)('0_log10_h') +log10_mc =EVAL:enterprise.signals.parameter.Uniform(7,10)('0_log10_mc') +log10_fgw=EVAL:enterprise.signals.parameter.Uniform(np.log10(1e-9), np.log10(10**-7.6))('0_log10_fgw') +phase0 =EVAL:enterprise.signals.parameter.Uniform(0, 2*np.pi)('0_phase0') +psi =EVAL:enterprise.signals.parameter.Uniform(0, np.pi)('0_psi') +cos_inc =EVAL:enterprise.signals.parameter.Uniform(-1, 1)('0_cos_inc') +p_phase =EVAL:enterprise.signals.parameter.Uniform(0, 2*np.pi) +p_dist =EVAL:enterprise.signals.parameter.Normal(0, 1) +tref=EVAL:53000*86400 + +[CWSignal] +module=enterprise_extensions.deterministic +function=CWSignal +signal_return=True +; note cw_wf=CUSTOM_RETURN:cw_delay_residual is equivalent +cw_wf=CUSTOM_FUNCTION_RETURN:${cw_delay:custom_return} +psrTerm=True +name=cw0 + +[model_2a] +module=enterprise_extensions.models +function=model_2a +pta_return=True +tm_marg=False +n_rnfreqs=30 +n_gwbfreqs=14 +gp_ecorr=True +individual_Tspan_for_red_noise=True \ No newline at end of file diff --git a/enterprise_extensions/pipe/example_dataclass.py b/enterprise_extensions/pipe/example_dataclass.py new file mode 100644 index 00000000..3a67a6cf --- /dev/null +++ b/enterprise_extensions/pipe/example_dataclass.py @@ -0,0 +1,16 @@ +import data_class +#from enterprise_extensions.pipe import data_class +import h5py +import numpy as np +import pandas as pd + +params = data_class.RunSettings() +#params.update_from_file('m2a.ini') +params.update_from_file('m2a_CW.ini') + +pta = params.create_pta_object_from_signals() + +print(pta) + + +print("done") \ No newline at end of file diff --git a/setup.py b/setup.py index a984ee80..e69b27e1 100644 --- a/setup.py +++ b/setup.py @@ -23,6 +23,7 @@ "scikit-learn>=0.24", "emcee", "ptmcmcsampler", + "importlib" ] test_requirements = [] @@ -60,6 +61,7 @@ def get_version(): "enterprise_extensions", "enterprise_extensions.frequentist", "enterprise_extensions.chromatic", + "enterprise_extensions.pipe" ], package_data={ "enterprise_extensions.chromatic": [