Merge EPTA version part 2

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

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

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"

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)