flake

Author: jeremy-baier

enterprise_extensions/blocks.py

@@ -170,10 +170,10 @@ def white_noise_block(
         else:
             if name is None:
                 ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr,
-                                                    selection=backend_ng, name=name)
+                                                    selection=selection_ecorr)
             else:
                 ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr,
-                                                    selection=backend_ng)
+                                                    selection=selection_ecorr, name=name)
     # combine signals
     if inc_ecorr:
         s = efeq + ec
@@ -670,7 +670,7 @@ def dm_noise_block(
                 elif prior == "gaussian":
                     log10_A_dm = parameter.Normal(logmin, logmax)
             elif not vary:
-              log10_A_dm = parameter.Constant()
+                log10_A_dm = parameter.Constant()
             else:
                 if prior == "uniform":
                     log10_A_dm = parameter.LinearExp(-20, -10)
@@ -783,10 +783,10 @@ def dm_noise_block(
                     log10_rho_dm = parameter.Uniform(-10, -4, size=components)
                 else:
                     log10_rho_dm = parameter.Uniform(-9, -4, size=components)
-                if not vary: # here just overwrite the prior to constant if not varying it
+                if not vary:  # here just overwrite the prior to constant if not varying it
                     log10_rho_dm = parameter.Constant()
 
-            dm_prior = gpp.free_spectrum(log10_rho=log10_rho)
+            dm_prior = gpp.free_spectrum(log10_rho=log10_rho_dm)
 
         if tndm:
             dm_basis = utils.createfourierdesignmatrix_dm_tn(
@@ -893,7 +893,7 @@ def dm_noise_block(
                 log10_sigma_ridge = parameter.Constant()
 
             dm_basis = gpk.linear_interp_basis_dm(dt=dt * const.day)
-            dm_prior = gpk.dmx_ridge_prior(log10_sigma=log10_sigma)
+            dm_prior = gpk.dmx_ridge_prior(log10_sigma=log10_sigma_ridge)
 
     if select is None:
         dmgp = gp_signals.BasisGP(
@@ -1110,7 +1110,7 @@ def chromatic_noise_block(
                 if vary:
                     log10_B = parameter.Uniform(-10, -4)
                 else:
-                    log10B = parameter.Constant()
+                    log10_B = parameter.Constant()
                 chm_prior = gpp.flat_powerlaw(
                     log10_A=log10_A, gamma=gamma, log10_B=log10_B
                 )
@@ -1130,7 +1130,7 @@ def chromatic_noise_block(
                     log10_rho = parameter.Uniform(-9, -4, size=components)
             if vary:
                 log10_rho = parameter.Constant()
- 
+
             chm_prior = gpp.free_spectrum(log10_rho=log10_rho)
 
         if tndm:

enterprise_extensions/chromatic/chromatic.py

@@ -3,7 +3,6 @@
 import numpy as np
 from enterprise import constants as const
 from enterprise.signals import deterministic_signals, parameter, signal_base, gp_bases
-from .solar_wind import solar_wind
 
 __all__ = [
     "chrom_exp_decay",
@@ -429,11 +428,11 @@ def dmx_signal(dmx_data, name="dmx_signal", vary=True):
         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"]
-                )
-              }
+                {
+                    dmx_id: parameter.Normal(
+                        mu=dmx_data_tmp["DMX_VAL"], sigma=dmx_data_tmp["DMX_ERR"]
+                    )
+                }
             )
     else:
         for dmx_id in sorted(dmx_data):
@@ -446,7 +445,6 @@ def dmx_signal(dmx_data, name="dmx_signal", vary=True):
 
 
 def dm_annual_signal(idx=2, tmin=None, tmax=None, name="dm_s1yr", vary=True):
-def dm_annual_signal(idx=2, name="dm_s1yr"):
     """
     Returns chromatic annual signal (i.e. TOA advance):
 
@@ -487,7 +485,7 @@ def construct_chromatic_cached_parts(
     fmax=None,
     modes=None,
     fref=1400,
-    ):
+):
     """
     Using this function alongside `createfourierdesignmatrix_chromatic_with_additional_caching()`
     enables caching of the achromatic portion of the chromatic Fourier designmatrix as well as caching

enterprise_extensions/chromatic/solar_wind.py

@@ -248,9 +248,9 @@ def solar_wind_block(
         n_earth = ACE_SWEPAM_Parameter(size=n_earth_bins.size-1)("n_earth")
     elif n_earth is None and (isinstance(n_earth_bins, list) or
                               isinstance(n_earth_bins, np.ndarray)) and not ACE_prior:
-        n_earth = parameter.Uniform(0,30,size=n_earth_bins.size-1)("n_earth")
+        n_earth = parameter.Uniform(0, 30, size=n_earth_bins.size-1)("n_earth")
     else:
-        pass # set n_earth to the provided value(s) below
+        pass  # set n_earth to the provided value(s) below
 
     deter_sw = solar_wind(n_earth=n_earth, n_earth_bins=n_earth_bins, t_init=t_init, t_final=t_final)
     mean_sw = deterministic_signals.Deterministic(deter_sw, name=det_name)
@@ -262,13 +262,13 @@ def solar_wind_block(
                 sw_basis = createfourierdesignmatrix_solar_dm(modes=modes)
                 nmodes = len(modes)
             elif Tspan is not None:
-                    sw_basis = createfourierdesignmatrix_solar_dm(nmodes=nmodes,
-                                                                  Tspan=Tspan)
+                sw_basis = createfourierdesignmatrix_solar_dm(nmodes=nmodes,
+                                                              Tspan=Tspan)
             if swgp_prior == "powerlaw":
                 if vary_swgp:
                     # sometimes amplitudes larger than 1 break the likelihood
-                    log10_A_sw = parameter.Uniform(-12, 0) # sometimes positive amplitudes break this
-                    gamma_sw = parameter.Uniform(-6, 5) # priors from susurla et al. 2024
+                    log10_A_sw = parameter.Uniform(-12, 0)  # sometimes positive amplitudes break this
+                    gamma_sw = parameter.Uniform(-6, 5)  # priors from susurla et al. 2024
                 else:
                     log10_A_sw = parameter.Constant()
                     gamma_sw = parameter.Constant()
@@ -292,9 +292,9 @@ def solar_wind_block(
             if swgp_prior == "periodic":
                 # Periodic GP kernel for DM
                 if vary_swgp:
-                    log10_sigma = parameter.Uniform(-10, -4) # units are log10(seconds)
-                    log10_ell = parameter.Uniform(1, 4) # units are log10(days)
-                    log10_p = parameter.Uniform(-4, 1.5) # units are log10(years)
+                    log10_sigma = parameter.Uniform(-10, -4)  # units are log10(seconds)
+                    log10_ell = parameter.Uniform(1, 4)  # units are log10(days)
+                    log10_p = parameter.Uniform(-4, 1.5)  # units are log10(years)
                     log10_gam_p = parameter.Uniform(-3, 2)
                 else:
                     log10_sigma = parameter.Constant()
@@ -303,9 +303,9 @@ def solar_wind_block(
                     log10_gam_p = parameter.Constant()
 
                 sw_prior = gpk.periodic_kernel(log10_sigma=log10_sigma,
-                                            log10_ell=log10_ell,
-                                            log10_gam_p=log10_gam_p,
-                                            log10_p=log10_p)
+                                               log10_ell=log10_ell,
+                                               log10_gam_p=log10_gam_p,
+                                               log10_p=log10_p)
             elif swgp_prior == "sq_exp":
                 # squared-exponential GP kernel for DM
                 if vary_swgp:
@@ -338,7 +338,7 @@ def solar_wind_block(
                 sw_prior = gpk.sw_dm_wn_prior(log10_sigma_ne=log10_sigma_ne)
             else:
                 raise ValueError("Invalid triangular-basis SWGP prior specified.")
-        
+
         else:
             raise ValueError("Invalid SWGP basis specified.")
 

enterprise_extensions/models.py

@@ -31,7 +31,6 @@
 from enterprise_extensions.chromatic.solar_wind import solar_wind_block
 from enterprise_extensions.timing import timing_block
 
-# from enterprise.signals.signal_base import LookupLikelihood
 
 def model_singlepsr_noise(
     psr,
@@ -45,6 +44,7 @@ def model_singlepsr_noise(
     tm_svd=False,
     tm_norm=True,
     white_vary=True,
+    gp_ecorr=False,
     components=30,
     upper_limit=False,
     is_wideband=False,
@@ -53,6 +53,7 @@ def model_singlepsr_noise(
     dmjump_var=False,
     gamma_val=None,
     dm_var=False,
+    vary_dm=True,
     dm_type="gp",
     dmgp_kernel="diag",
     dm_psd="powerlaw",
@@ -64,6 +65,7 @@ def model_singlepsr_noise(
     dm_df=200,
     dm_Nfreqs=100,
     chrom_gp=False,
+    vary_chrom=True,
     chrom_gp_kernel="nondiag",
     chrom_psd="powerlaw",
     chrom_idx=4,
@@ -72,6 +74,7 @@ def model_singlepsr_noise(
     chrom_dt=15,
     chrom_df=200,
     chrom_Nfreqs=100,
+    vary_dm_dips=True,
     dm_expdip=False,
     dmexp_sign="negative",
     dm_expdip_idx=2,
@@ -138,6 +141,7 @@ def model_singlepsr_noise(
            is_wideband
     :param gamma_val: red noise spectral index to fix
     :param dm_var: whether to explicitly model DM-variations
+    :param vary_dm: whether to vary the DM model GP hyperparams or use constant values
     :param dm_type: gaussian process ('gp') or dmx ('dmx')
     :param dmgp_kernel: diagonal in frequency or non-diagonal
     :param dm_psd: power-spectral density of DM variations
@@ -149,6 +153,7 @@ def model_singlepsr_noise(
     :param dm_df: frequency-scale for DM linear interpolation basis (MHz)
     :param dm_Nfreqs: Number of Fourier modes to use for the dm_gp model.
     :param chrom_gp: include general chromatic noise
+    :param vary_chrom: whether to vary the chromatic GP hyperparams or use constant values
     :param chrom_gp_kernel: GP kernel type to use in chrom ['diag','nondiag']
     :param chrom_psd: power-spectral density of chromatic noise
         ['powerlaw','tprocess','free_spectrum']
@@ -161,6 +166,7 @@ def model_singlepsr_noise(
     :param chrom_dt: time-scale for chromatic linear interpolation basis (days)
     :param chrom_df: frequency-scale for chromatic linear interpolation basis (MHz)
     :param chrom_Nfreqs: Number of Fourier modes to use for the chromatic GP.
+    :param vary_dm_dips: whether to vary the DM dip parameters or keep them fixed
     :param dm_expdip: inclue a DM exponential dip
     :param dmexp_sign: set the sign parameter for dip
     :param dm_expdip_idx: chromatic index of exponential dip
@@ -424,7 +430,7 @@ def model_singlepsr_noise(
                     sign=dm_dual_cusp_sign,
                     symmetric=dm_dual_cusp_sym,
                     name=dual_cusp_name_base + str(dd),
-                    vary=vary_dm)
+                    vary=vary_dm,
                 )
         if dm_sw_deter:
             Tspan = psr.toas.max() - psr.toas.min()
@@ -445,8 +451,7 @@ def model_singlepsr_noise(
                                    include_quadratic=chrom_quad,
                                    coefficients=coefficients,
                                    Tspan=Tspan,
-                                   vary=vary_chrom,
-                                   idx_prior_upper_bound=chrom_gp_idx_prior_upper_bound)
+                                   vary=vary_chrom,)
     if extra_sigs is not None:
         s += extra_sigs
 

enterprise_extensions/sampler.py

@@ -366,11 +366,11 @@ def __init__(
                         # extend empirical_distr here:
                         print(f'extending {key}\'s empirical distributions to priors...\n')
                         self.empirical_distr[key] = extend_emp_dists(
-                          pta,
-                          self.empirical_distr[key],
-                          npoints=100_000,
-                          save_ext_dists=save_ext_dists,
-                          outdir=outdir
+                            pta,
+                            self.empirical_distr[key],
+                            npoints=100_000,
+                            save_ext_dists=save_ext_dists,
+                            outdir=outdir
                         )
             else:
                 mask = []

tests/test_models.py

@@ -6,8 +6,6 @@
 import json
 import logging
 import os
-
-import pickle
 import numpy as np
 
 import pytest