diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py index 9b621f4e..f924be12 100644 --- a/enterprise/signals/gp_bases.py +++ b/enterprise/signals/gp_bases.py @@ -13,9 +13,12 @@ __all__ = [ "createfourierdesignmatrix_red", "createfourierdesignmatrix_dm", + "createfourierdesignmatrix_dm_tn", "createfourierdesignmatrix_env", "createfourierdesignmatrix_ephem", "createfourierdesignmatrix_eph", + "createfourierdesignmatrix_chromatic", + "createfourierdesignmatrix_general", ] @@ -124,6 +127,44 @@ def createfourierdesignmatrix_dm( return F * Dm[:, None], Ffreqs +@function +def createfourierdesignmatrix_dm_tn( + toas, freqs, nmodes=30, Tspan=None, pshift=False, fref=1400, logf=False, fmin=None, fmax=None, idx=2, modes=None +): + """ + Construct DM-variation fourier design matrix. Current + normalization expresses DM signal as a deviation [seconds] + at fref [MHz] + + :param toas: vector of time series in seconds + :param freqs: radio frequencies of observations [MHz] + :param nmodes: number of fourier coefficients to use + :param Tspan: option to some other Tspan + :param pshift: option to add random phase shift + :param fref: reference frequency [MHz] + :param logf: use log frequency spacing + :param fmin: lower sampling frequency + :param fmax: upper sampling frequency + :param idx: index of the radio frequency dependence + :param modes: option to provide explicit list or array of + sampling frequencies + + :return: F: DM-variation fourier design matrix + :return: f: Sampling frequencies + """ + + # get base fourier design matrix and frequencies + F, Ffreqs = createfourierdesignmatrix_red( + toas, nmodes=nmodes, Tspan=Tspan, logf=logf, fmin=fmin, fmax=fmax, pshift=pshift, modes=modes + ) + + # compute the DM-variation vectors in the temponest normalization + # amplitude normalization: sqrt(12)*pi, scaling to 1 MHz from 1400 MHz, DM constant: 2.41e-4 + Dm = (fref / freqs) ** idx * np.sqrt(12) * np.pi / 1400 / 1400 / 2.41e-4 + + return F * Dm[:, None], Ffreqs + + @function def createfourierdesignmatrix_env( toas, @@ -218,7 +259,9 @@ def createfourierdesignmatrix_eph( @function -def createfourierdesignmatrix_chromatic(toas, freqs, nmodes=30, Tspan=None, logf=False, fmin=None, fmax=None, idx=4): +def createfourierdesignmatrix_chromatic( + toas, freqs, nmodes=30, Tspan=None, logf=False, fmin=None, fmax=None, idx=4, modes=None +): """ Construct Scattering-variation fourier design matrix. @@ -231,15 +274,82 @@ def createfourierdesignmatrix_chromatic(toas, freqs, nmodes=30, Tspan=None, logf :param fmin: lower sampling frequency :param fmax: upper sampling frequency :param idx: Index of chromatic effects + :param modes: option to provide explicit list or array of + sampling frequencies :return: F: Chromatic-variation fourier design matrix :return: f: Sampling frequencies """ # get base fourier design matrix and frequencies - F, Ffreqs = createfourierdesignmatrix_red(toas, nmodes=nmodes, Tspan=Tspan, logf=logf, fmin=fmin, fmax=fmax) + F, Ffreqs = createfourierdesignmatrix_red( + toas, nmodes=nmodes, Tspan=Tspan, logf=logf, fmin=fmin, fmax=fmax, modes=modes + ) # compute the DM-variation vectors Dm = (1400 / freqs) ** idx return F * Dm[:, None], Ffreqs + + +@function +def createfourierdesignmatrix_general( + toas, + freqs, + flags, + flagname="group", + flagval=None, + idx=None, + tndm=False, + nmodes=30, + Tspan=None, + psrTspan=True, + logf=False, + fmin=None, + fmax=None, + modes=None, + pshift=None, + pseed=None, +): + """ + Construct fourier design matrix with possibility of adding selection and/or chromatic index envelope. + + :param toas: vector of time series in seconds + :param freqs: radio frequencies of observations [MHz] + :param flags: Flags from timfiles + :param nmodes: number of fourier coefficients to use + :param Tspan: option to some other Tspan + :param psrTspan: option to use pulsar time span. Used only if sub-group of ToAs is chosen + :param logf: use log frequency spacing + :param fmin: lower sampling frequency + :param fmax: upper sampling frequency + :param log10_Amp: log10 of the Amplitude [s] + :param idx: Index of chromatic effects + :param modes: option to provide explicit list or array of + sampling frequencies + + :return: F: fourier design matrix + :return: f: Sampling frequencies + """ + if flagval and not psrTspan: + sel_toas = toas[np.where(flags[flagname] == flagval)] + Tspan = sel_toas.max() - sel_toas.min() + + # get base fourier design matrix and frequencies + F, Ffreqs = createfourierdesignmatrix_red( + toas, nmodes=nmodes, Tspan=Tspan, logf=logf, fmin=fmin, fmax=fmax, modes=modes, pshift=pshift, pseed=pseed + ) + + # compute the chromatic-variation vectors + if idx: + if tndm: + chrom_fac = (1400 / freqs) ** idx * np.sqrt(12) * np.pi / 1400 / 1400 / 2.41e-4 + else: + chrom_fac = (1400 / freqs) ** idx + F *= chrom_fac[:, None] + + # compute the mask for the selection + if flagval: + F *= np.array([flags[flagname] == flagval] * F.shape[1]).T + + return F, Ffreqs diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py index c6f03126..21ad177d 100644 --- a/enterprise/signals/gp_signals.py +++ b/enterprise/signals/gp_signals.py @@ -192,6 +192,9 @@ def FourierBasisGP( components=20, selection=Selection(selections.no_selection), Tspan=None, + logf=False, + fmin=None, + fmax=None, modes=None, name="red_noise", pshift=False, @@ -200,7 +203,9 @@ def FourierBasisGP( """Convenience function to return a BasisGP class with a fourier basis.""" - basis = utils.createfourierdesignmatrix_red(nmodes=components, Tspan=Tspan, modes=modes, pshift=pshift, pseed=pseed) + basis = utils.createfourierdesignmatrix_red( + nmodes=components, Tspan=Tspan, logf=logf, fmin=fmin, fmax=fmax, modes=modes, pshift=pshift, pseed=pseed + ) BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name) class FourierBasisGP(BaseClass): @@ -211,24 +216,24 @@ class FourierBasisGP(BaseClass): return FourierBasisGP -def get_timing_model_basis(use_svd=False, normed=True): +def get_timing_model_basis(use_svd=False, normed=True, idx_exclude=None): if use_svd: if normed is not True: raise ValueError("use_svd == True requires normed == True") - return utils.svd_tm_basis() + return utils.svd_tm_basis(idx_exclude=idx_exclude) elif normed is True: - return utils.normed_tm_basis() + return utils.normed_tm_basis(idx_exclude=idx_exclude) elif normed is not False: - return utils.normed_tm_basis(norm=normed) + return utils.normed_tm_basis(norm=normed, idx_exclude=idx_exclude) else: - return utils.unnormed_tm_basis() + return utils.unnormed_tm_basis(idx_exclude=idx_exclude) -def TimingModel(coefficients=False, name="linear_timing_model", use_svd=False, normed=True): +def TimingModel(coefficients=False, name="linear_timing_model", use_svd=False, normed=True, idx_exclude=None): """Class factory for marginalized linear timing model signals.""" - basis = get_timing_model_basis(use_svd, normed) + basis = get_timing_model_basis(use_svd, normed, idx_exclude) prior = utils.tm_prior() BaseClass = BasisGP(prior, basis, coefficients=coefficients, name=name) @@ -413,6 +418,9 @@ def FourierBasisCommonGP( combine=True, components=20, Tspan=None, + logf=False, + fmin=None, + fmax=None, modes=None, name="common_fourier", pshift=False, @@ -424,7 +432,9 @@ def FourierBasisCommonGP( "With coefficients=True, FourierBasisCommonGP " + "requires that you specify Tspan explicitly." ) - basis = utils.createfourierdesignmatrix_red(nmodes=components, Tspan=Tspan, modes=modes, pshift=pshift, pseed=pseed) + basis = utils.createfourierdesignmatrix_red( + nmodes=components, Tspan=Tspan, logf=logf, fmin=fmin, fmax=fmax, modes=modes, pshift=pshift, pseed=pseed + ) BaseClass = BasisCommonGP(spectrum, basis, orf, coefficients=coefficients, combine=combine, name=name) class FourierBasisCommonGP(BaseClass): diff --git a/enterprise/signals/signal_base.py b/enterprise/signals/signal_base.py index adc3125e..f092f05a 100644 --- a/enterprise/signals/signal_base.py +++ b/enterprise/signals/signal_base.py @@ -735,6 +735,9 @@ def summary(self, include_params=True, to_stdout=False): cpcount += 1 row = [sig.name, sig.__class__.__name__, len(sig.param_names)] summary += "{: <40} {: <30} {: <20}\n".format(*row) + if "BasisGP" in sig.__class__.__name__: + summary += "\nBasis shape (Ntoas x N basis functions): {}".format(str(sig.get_basis().shape)) + summary += "\nN selected toas: {}\n".format(str(len([i for i in sig._masks[0] if i]))) if include_params: summary += "\n" summary += "params:\n" diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py index e2986f55..d896de37 100644 --- a/enterprise/signals/utils.py +++ b/enterprise/signals/utils.py @@ -19,11 +19,14 @@ from enterprise import constants as const from enterprise import signals as sigs # noqa: F401 from enterprise.signals.gp_bases import ( # noqa: F401 + createfourierdesignmatrix_red, createfourierdesignmatrix_dm, + createfourierdesignmatrix_dm_tn, createfourierdesignmatrix_env, - createfourierdesignmatrix_eph, createfourierdesignmatrix_ephem, - createfourierdesignmatrix_red, + createfourierdesignmatrix_eph, + createfourierdesignmatrix_chromatic, + createfourierdesignmatrix_general, ) from enterprise.signals.gp_priors import powerlaw, turnover # noqa: F401 from enterprise.signals.parameter import function @@ -872,12 +875,19 @@ def anis_orf(pos1, pos2, params, **kwargs): @function -def unnormed_tm_basis(Mmat): +def unnormed_tm_basis(Mmat, idx_exclude=None): + if idx_exclude: + idxs = np.array([i for i in range(Mmat.shape[1]) if i not in idx_exclude]) + Mmat = Mmat[:, idxs] return Mmat, np.ones_like(Mmat.shape[1]) @function -def normed_tm_basis(Mmat, norm=None): +def normed_tm_basis(Mmat, norm=None, idx_exclude=None): + if idx_exclude: + idxs = np.array([i for i in range(Mmat.shape[1]) if i not in idx_exclude]) + Mmat = Mmat[:, idxs] + if norm is None: norm = np.sqrt(np.sum(Mmat**2, axis=0)) @@ -888,7 +898,11 @@ def normed_tm_basis(Mmat, norm=None): @function -def svd_tm_basis(Mmat): +def svd_tm_basis(Mmat, idx_exclude=None): + if idx_exclude: + idxs = np.array([i for i in range(Mmat.shape[1]) if i not in idx_exclude]) + Mmat = Mmat[:, idxs] + u, s, v = np.linalg.svd(Mmat, full_matrices=False) return u, np.ones_like(s)