Chromatic noise updates

.github/workflows/ci_test.yml

@@ -16,8 +16,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest, macos-latest]
-        python-version: [3.7, 3.8, 3.9, '3.10']
+        os: [ubuntu-latest, macos-13]
+        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
 
     steps:
     - name: Checkout repository

enterprise_extensions/chromatic/chromatic.py

@@ -204,7 +204,7 @@ def dmx_delay(toas, freqs, dmx_ids, **kwargs):
     return wf
 
 
-def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp'):
+def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp', vary=True):
     """
     Returns chromatic exponential dip (i.e. TOA advance):
 
@@ -216,15 +216,24 @@ def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp'):
     :param sign:
         set sign of dip: 'positive', 'negative', or 'vary'
     :param name: Name of signal
+    :param vary: Whether to vary the parameters or use constant values.
 
     :return dmexp:
         chromatic exponential dip waveform.
     """
-    t0_dmexp = parameter.Uniform(tmin, tmax)
-    log10_Amp_dmexp = parameter.Uniform(-10, -2)
-    log10_tau_dmexp = parameter.Uniform(0, 2.5)
-    if sign == 'vary':
+    if vary:
+        t0_dmexp = parameter.Uniform(tmin, tmax)
+        log10_Amp_dmexp = parameter.Uniform(-10, -2)
+        log10_tau_dmexp = parameter.Uniform(0, 2.5)
+    else:
+        t0_dmexp = parameter.Constant()
+        log10_Amp_dmexp = parameter.Constant()
+        log10_tau_dmexp = parameter.Constant()
+
+    if sign == 'vary' and vary:
         sign_param = parameter.Uniform(-1.0, 1.0)
+    elif sign == 'vary' and not vary:
+        sign_param = parameter.Constant()
     elif sign == 'positive':
         sign_param = 1.0
     else:
@@ -238,7 +247,7 @@ def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp'):
 
 
 def dm_exponential_cusp(tmin, tmax, idx=2, sign='negative',
-                        symmetric=False, name='dm_cusp'):
+                        symmetric=False, name='dm_cusp', vary=True):
     """
     Returns chromatic exponential cusp (i.e. TOA advance):
 
@@ -250,25 +259,35 @@ def dm_exponential_cusp(tmin, tmax, idx=2, sign='negative',
     :param sign:
         set sign of dip: 'positive', 'negative', or 'vary'
     :param name: Name of signal
+    :param vary: Whether to vary the parameters or use constant values.
 
     :return dmexp:
         chromatic exponential dip waveform.
     """
-    t0_dm_cusp = parameter.Uniform(tmin, tmax)
-    log10_Amp_dm_cusp = parameter.Uniform(-10, -2)
-    log10_tau_dm_cusp_pre = parameter.Uniform(0, 2.5)
+    if vary:
+        t0_dm_cusp = parameter.Uniform(tmin, tmax)
+        log10_Amp_dm_cusp = parameter.Uniform(-10, -2)
+        log10_tau_dm_cusp_pre = parameter.Uniform(0, 2.5)
+    else:
+        t0_dm_cusp = parameter.Constant()
+        log10_Amp_dm_cusp = parameter.Constant()
+        log10_tau_dm_cusp_pre = parameter.Constant()
 
-    if sign == 'vary':
+    if sign == 'vary' and vary:
         sign_param = parameter.Uniform(-1.0, 1.0)
+    elif sign == 'vary' and not vary:
+        sign_param = parameter.Constant()
     elif sign == 'positive':
         sign_param = 1.0
     else:
         sign_param = -1.0
 
     if symmetric:
         log10_tau_dm_cusp_post = 1
-    else:
+    elif vary:
         log10_tau_dm_cusp_post = parameter.Uniform(0, 2.5)
+    else:
+        log10_tau_dm_cusp_post = parameter.Constant()
 
     wf = chrom_exp_cusp(log10_Amp=log10_Amp_dm_cusp, sign_param=sign_param,
                         t0=t0_dm_cusp, log10_tau_pre=log10_tau_dm_cusp_pre,
@@ -280,7 +299,7 @@ def dm_exponential_cusp(tmin, tmax, idx=2, sign='negative',
 
 
 def dm_dual_exp_cusp(tmin, tmax, idx1=2, idx2=4, sign='negative',
-                     symmetric=False, name='dual_dm_cusp'):
+                     symmetric=False, name='dual_dm_cusp', vary=True):
     """
     Returns chromatic exponential cusp (i.e. TOA advance):
 
@@ -292,18 +311,28 @@ def dm_dual_exp_cusp(tmin, tmax, idx1=2, idx2=4, sign='negative',
     :param sign:
         set sign of dip: 'positive', 'negative', or 'vary'
     :param name: Name of signal
+    :param vary: Whether to vary the parameters or use constant values.
 
     :return dmexp:
         chromatic exponential dip waveform.
     """
-    t0_dual_cusp = parameter.Uniform(tmin, tmax)
-    log10_Amp_dual_cusp_1 = parameter.Uniform(-10, -2)
-    log10_Amp_dual_cusp_2 = parameter.Uniform(-10, -2)
-    log10_tau_dual_cusp_pre_1 = parameter.Uniform(0, 2.5)
-    log10_tau_dual_cusp_pre_2 = parameter.Uniform(0, 2.5)
+    if vary:
+        t0_dual_cusp = parameter.Uniform(tmin, tmax)
+        log10_Amp_dual_cusp_1 = parameter.Uniform(-10, -2)
+        log10_Amp_dual_cusp_2 = parameter.Uniform(-10, -2)
+        log10_tau_dual_cusp_pre_1 = parameter.Uniform(0, 2.5)
+        log10_tau_dual_cusp_pre_2 = parameter.Uniform(0, 2.5)
+    else:
+        t0_dual_cusp = parameter.Constant()
+        log10_Amp_dual_cusp_1 = parameter.Constant()
+        log10_Amp_dual_cusp_2 = parameter.Constant()
+        log10_tau_dual_cusp_pre_1 = parameter.Constant()
+        log10_tau_dual_cusp_pre_2 = parameter.Constant()
 
-    if sign == 'vary':
+    if sign == 'vary' and vary:
         sign_param = parameter.Uniform(-1.0, 1.0)
+    elif sign == 'vary' and not vary:
+        sign_param = parameter.Constant()
     elif sign == 'positive':
         sign_param = 1.0
     else:
@@ -312,9 +341,12 @@ def dm_dual_exp_cusp(tmin, tmax, idx1=2, idx2=4, sign='negative',
     if symmetric:
         log10_tau_dual_cusp_post_1 = 1
         log10_tau_dual_cusp_post_2 = 1
-    else:
+    elif vary:
         log10_tau_dual_cusp_post_1 = parameter.Uniform(0, 2.5)
         log10_tau_dual_cusp_post_2 = parameter.Uniform(0, 2.5)
+    else:
+        log10_tau_dual_cusp_post_1 = parameter.Constant()
+        log10_tau_dual_cusp_post_2 = parameter.Constant()
 
     wf = chrom_dual_exp_cusp(t0=t0_dual_cusp, sign_param=sign_param,
                              symmetric=symmetric,
@@ -330,41 +362,53 @@ def dm_dual_exp_cusp(tmin, tmax, idx1=2, idx2=4, sign='negative',
     return dm_cusp
 
 
-def dmx_signal(dmx_data, name='dmx_signal'):
+def dmx_signal(dmx_data, name='dmx_signal', vary=True):
     """
     Returns DMX signal:
 
     :param dmx_data: dictionary of DMX data for each pulsar from parfile.
     :param name: Name of signal.
+    :param vary: Whether to vary the parameters or use constant values.
 
     :return dmx_sig:
         dmx signal waveform.
     """
     dmx = {}
-    for dmx_id in sorted(dmx_data):
-        dmx_data_tmp = dmx_data[dmx_id]
-        dmx.update({dmx_id: parameter.Normal(mu=dmx_data_tmp['DMX_VAL'],
-                                             sigma=dmx_data_tmp['DMX_ERR'])})
+    if vary:
+        for dmx_id in sorted(dmx_data):
+            dmx_data_tmp = dmx_data[dmx_id]
+            dmx.update({dmx_id: parameter.Normal(mu=dmx_data_tmp['DMX_VAL'],
+                                                 sigma=dmx_data_tmp['DMX_ERR'])})
+    else:
+        for dmx_id in sorted(dmx_data):
+            dmx_data_tmp = dmx_data[dmx_id]
+            dmx.update({dmx_id: parameter.Constant()})
+
     wf = dmx_delay(dmx_ids=dmx_data, **dmx)
     dmx_sig = deterministic_signals.Deterministic(wf, name=name)
 
     return dmx_sig
 
 
-def dm_annual_signal(idx=2, name='dm_s1yr'):
+def dm_annual_signal(idx=2, name='dm_s1yr', vary=True):
     """
     Returns chromatic annual signal (i.e. TOA advance):
 
     :param idx:
         index of radio frequency dependence (i.e. DM is 2). If this is set
         to 'vary' then the index will vary from 1 - 6
     :param name: Name of signal
+    :param vary: Whether to vary the parameters or use constant values.
 
     :return dm1yr:
         chromatic annual waveform.
     """
-    log10_Amp_dm1yr = parameter.Uniform(-10, -2)
-    phase_dm1yr = parameter.Uniform(0, 2*np.pi)
+    if vary:
+        log10_Amp_dm1yr = parameter.Uniform(-10, -2)
+        phase_dm1yr = parameter.Uniform(0, 2*np.pi)
+    else:
+        log10_Amp_dm1yr = parameter.Constant()
+        phase_dm1yr = parameter.Constant()
 
     wf = chrom_yearly_sinusoid(log10_Amp=log10_Amp_dm1yr,
                                phase=phase_dm1yr, idx=idx)

enterprise_extensions/chromatic/solar_wind.py

@@ -186,7 +186,7 @@ def solar_wind_block(n_earth=None, ACE_prior=False, include_swgp=True,
         Prior function for solar wind Gaussian process. Default is a power law.
     :param swgp_basis:
         Basis to be used for solar wind Gaussian process.
-        Options includes ['powerlaw'.'periodic','sq_exp']
+        Options includes ['powerlaw','periodic','sq_exp','triangular']
     :param Tspan:
         Sets frequency sampling f_i = i / Tspan. Default will
         use overall time span for individual pulsar. Default is to use 15
@@ -241,6 +241,12 @@ def solar_wind_block(n_earth=None, ACE_prior=False, include_swgp=True,
             sw_prior = gpk.se_dm_kernel(log10_sigma=log10_sigma,
                                         log10_ell=log10_ell)
 
+        elif swgp_basis == 'triangular':
+            # white noise kernel for SW DM, using delta n_earth as coefficients
+            log10_sigma_ne = parameter.Uniform(-4, 2)
+            sw_basis = gpk.sw_dm_triangular_basis()
+            sw_prior = gpk.sw_dm_wn_prior(log10_sigma_ne=log10_sigma_ne)
+
         gp_sw = gp_signals.BasisGP(sw_prior, sw_basis, name='gp_sw')
         sw_model += gp_sw
 

enterprise_extensions/gp_kernels.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 
 import numpy as np
-from enterprise.signals import signal_base, utils
+from enterprise.signals import signal_base, utils, parameter
+import enterprise.constants as const
+from enterprise_extensions import chromatic as chrom
 
 __all__ = ['linear_interp_basis_dm',
            'linear_interp_basis_freq',
@@ -47,9 +49,9 @@ def linear_interp_basis_freq(freqs, df=64):
 
 
 @signal_base.function
-def dmx_ridge_prior(avetoas, log10_sigma=-7):
+def dmx_ridge_prior(avetoas, log10_sigma_ridge=-7):
     """DMX-like signal with Gaussian prior"""
-    sigma = 10**log10_sigma
+    sigma = 10**log10_sigma_ridge
     return sigma**2 * np.ones_like(avetoas)
 
 
@@ -64,7 +66,7 @@ def periodic_kernel(avetoas, log10_sigma=-7, log10_ell=2,
     l = 10**log10_ell * 86400
     p = 10**log10_p * 3.16e7
     gam_p = 10**log10_gam_p
-    d = np.eye(r.shape[0]) * (sigma/500)**2
+    d = np.eye(r.shape[0]) * (sigma/50000.)**2
     K = sigma**2 * np.exp(-r**2/2/l**2 - gam_p*np.sin(np.pi*r/p)**2) + d
     return K
 
@@ -76,7 +78,7 @@ def se_kernel(avefreqs, log10_sigma=-7, log10_lam=3):
 
     lam = 10**log10_lam
     sigma = 10**log10_sigma
-    d = np.eye(tm.shape[0]) * (sigma/500)**2
+    d = np.eye(tm.shape[0]) * (sigma/50000.)**2
     return sigma**2 * np.exp(-tm**2/2/lam) + d
 
 
@@ -88,11 +90,57 @@ def se_dm_kernel(avetoas, log10_sigma=-7, log10_ell=2):
     # Convert everything into seconds
     l = 10**log10_ell * 86400
     sigma = 10**log10_sigma
-    d = np.eye(r.shape[0]) * (sigma/500)**2
+    d = np.eye(r.shape[0]) * (sigma/50000.)**2
     K = sigma**2 * np.exp(-r**2/2/l**2) + d
     return K
 
 
+@signal_base.function
+def sw_dm_triangular_basis(toas, planetssb, sunssb, pos_t, freqs, fref=1400):
+    """
+    Construct SWGP basis from Nitu et al 2024 using triangle functions centered on SW conjunctions
+    :param toas: vector of time series in seconds
+    :param planetssb: solar system barycenter positions
+    :param pos_t: pulsar position as 3-vector
+    :param freqs: radio frequencies of observations [MHz]
+    :param fref: reference frequency [MHz]
+
+    :return: V: Nc x Ntoa design matrix
+    :return: Tc: SW conjunctions
+    """
+
+    # First get SW conjunctions
+    theta, R_earth, _, _ = chrom.solar_wind.theta_impact(planetssb, sunssb, pos_t)
+    # Estimate conjunction times using TOA of closest approach
+    toa_min_theta = toas[np.argmin(theta)]
+    Tc = toa_min_theta + np.arange(100)*const.yr - 50*const.yr  # This might break after the NANOGrav 50 yr dataset
+    Tc = Tc[(Tc > np.min(toas))*(Tc < np.max(toas))]
+
+    # Set up triangular basis matrix functions
+    Nc = len(Tc)
+    Nt = len(toas)
+    Lambda = np.max(np.array([1 - np.abs(toas[:, None] - Tc[None, :])/const.yr, np.zeros((Nt, Nc))]), axis=0)
+
+    # Geometric factor (units of DM)
+    S_theta = chrom.solar_wind.dm_solar(1.0, theta, R_earth)
+
+    # Convert from DM to a time delay
+    S_theta *= 1e-12/const.DM_K/freqs**2
+
+    # return basis and conjunctions
+    V = S_theta[:, None]*Lambda
+    return V, Tc
+
+
+@parameter.function
+def sw_dm_wn_prior(Tc, log10_sigma_ne=-4):
+    """
+    Gaussian prior on variance of n_earth at SW conjunctions
+    """
+    sigma_ne = 10**log10_sigma_ne
+    return sigma_ne**2 * np.ones_like(Tc)
+
+
 @signal_base.function
 def get_tf_quantization_matrix(toas, freqs, dt=30*86400, df=None, dm=False, dm_idx=2):
     """
@@ -164,7 +212,7 @@ def tf_kernel(labels, log10_sigma=-7, log10_ell=2, log10_gam_p=0,
     gam_p = 10**log10_gam_p
     alpha_wgt = 10**log10_alpha_wgt
 
-    d = np.eye(r.shape[0]) * (sigma/500)**2
+    d = np.eye(r.shape[0]) * (sigma/50000.)**2
     Kt = sigma**2 * np.exp(-r**2/2/l**2 - gam_p*np.sin(np.pi*r/p)**2)
     Kv = (1+r2**2/2/alpha_wgt/l2**2)**(-alpha_wgt)
 
@@ -190,7 +238,7 @@ def sf_kernel(labels, log10_sigma=-7, log10_ell=2,
     l2 = 10**log10_ell2
     alpha_wgt = 10**log10_alpha_wgt
 
-    d = np.eye(r.shape[0]) * (sigma/500)**2
+    d = np.eye(r.shape[0]) * (sigma/50000.)**2
     Kt = sigma**2 * np.exp(-r**2/2/l**2)
     Kv = (1+r2**2/2/alpha_wgt/l2**2)**(-alpha_wgt)