diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md index 139313db4..5ac8acfd0 100644 --- a/CHANGELOG-unreleased.md +++ b/CHANGELOG-unreleased.md @@ -30,6 +30,8 @@ the released changes. - `Residuals.d_lnlikelihood_d_whitenoise_param` will throw a `NotImplementedError` when correlated noise is present. - `DownhillFitter._fit_noise()` doesn't use derivatives when correlated noise is present. - Documentation: Noise fitting example notebook. +- `freeze_params` option in `wavex_setup` and `dmwavex_setup` +- `plrednoise_from_wavex`, `pldmnoise_from_dmwavex`, and `find_optimal_nharms` functions ### Fixed - `MCMC_walkthrough` notebook now runs - Fixed runtime data README diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py new file mode 100644 index 000000000..c8a93e9cb --- /dev/null +++ b/docs/examples/rednoise-fit-example.py @@ -0,0 +1,374 @@ +# %% [markdown] +# # Red noise and DM noise fitting examples +# +# This notebook provides an example on how to fit for red noise +# and DM noise using PINT using simulated datasets. +# +# We will use the `PLRedNoise` and `PLDMNoise` models to generate +# noise realizations (these models provide Fourier Gaussian process +# descriptions of achromatic red noise and DM noise respectively). +# +# We will fit the generated datasets using the `WaveX` and `DMWaveX` models, +# which provide deterministic Fourier representations of achromatic red noise +# and DM noise respectively. +# +# Finally, we will convert the `WaveX`/`DMWaveX` amplitudes into spectral +# parameters and compare them with the injected values. + +# %% +from pint import DMconst +from pint.models import get_model +from pint.simulation import make_fake_toas_uniform +from pint.logging import setup as setup_log +from pint.fitter import WLSFitter +from pint.utils import ( + dmwavex_setup, + find_optimal_nharms, + wavex_setup, + plrednoise_from_wavex, + pldmnoise_from_dmwavex, +) + +from io import StringIO +import numpy as np +import astropy.units as u +from matplotlib import pyplot as plt +from copy import deepcopy + +setup_log(level="WARNING") + +# %% [markdown] +# ## Red noise fitting + +# %% [markdown] +# ### Simulation +# The first step is to generate a simulated dataset for demonstration. +# Note that we are adding PHOFF as a free parameter. This is required +# for the fit to work properly. + +# %% +par_sim = """ + PSR SIM3 + RAJ 05:00:00 1 + DECJ 15:00:00 1 + PEPOCH 55000 + F0 100 1 + F1 -1e-15 1 + PHOFF 0 1 + DM 15 1 + TNREDAMP -13 + TNREDGAM 3.5 + TNREDC 30 + TZRMJD 55000 + TZRFRQ 1400 + TZRSITE gbt + UNITS TDB + EPHEM DE440 + CLOCK TT(BIPM2019) +""" + +m = get_model(StringIO(par_sim)) + +# %% +# Now generate the simulated TOAs. +ntoas = 2000 +toaerrs = np.random.uniform(0.5, 2.0, ntoas) * u.us +freqs = np.linspace(500, 1500, 8) * u.MHz + +t = make_fake_toas_uniform( + startMJD=53001, + endMJD=57001, + ntoas=ntoas, + model=m, + freq=freqs, + obs="gbt", + error=toaerrs, + add_noise=True, + add_correlated_noise=True, + name="fake", + include_bipm=True, + include_gps=True, + multi_freqs_in_epoch=True, +) + +# %% [markdown] +# ### Optimal number of harmonics +# The optimal number of harmonics can be estimated by minimizing the +# Akaike Information Criterion (AIC). This is implemented in the +# `pint.utils.find_optimal_nharms` function. + +# %% +m1 = deepcopy(m) +m1.remove_component("PLRedNoise") + +nharm_opt, d_aics = find_optimal_nharms(m1, t, "WaveX", 30) + +print("Optimum no of harmonics = ", nharm_opt) + +# %% +print(np.argmin(d_aics)) + +# %% +# The Y axis is plotted in log scale only for better visibility. +plt.scatter(list(range(len(d_aics))), d_aics + 1) +plt.axvline(nharm_opt, color="red", label="Optimum number of harmonics") +plt.axvline( + int(m.TNREDC.value), color="black", ls="--", label="Injected number of harmonics" +) +plt.xlabel("Number of harmonics") +plt.ylabel("AIC - AIC$_\\min{} + 1$") +plt.legend() +plt.yscale("log") +# plt.savefig("sim3-aic.pdf") + +# %% +# Now create a new model with the optimum number of harmonics +m2 = deepcopy(m1) +Tspan = t.get_mjds().max() - t.get_mjds().min() +wavex_setup(m2, T_span=Tspan, n_freqs=nharm_opt, freeze_params=False) + +ftr = WLSFitter(t, m2) +ftr.fit_toas(maxiter=10) +m2 = ftr.model + +print(m2) + +# %% [markdown] +# ### Estimating the spectral parameters from the WaveX fit. + +# %% +# Get the Fourier amplitudes and powers and their uncertainties. +idxs = np.array(m2.components["WaveX"].get_indices()) +a = np.array([m2[f"WXSIN_{idx:04d}"].quantity.to_value("s") for idx in idxs]) +da = np.array([m2[f"WXSIN_{idx:04d}"].uncertainty.to_value("s") for idx in idxs]) +b = np.array([m2[f"WXCOS_{idx:04d}"].quantity.to_value("s") for idx in idxs]) +db = np.array([m2[f"WXCOS_{idx:04d}"].uncertainty.to_value("s") for idx in idxs]) +print(len(idxs)) + +P = (a**2 + b**2) / 2 +dP = ((a * da) ** 2 + (b * db) ** 2) ** 0.5 + +f0 = (1 / Tspan).to_value(u.Hz) +fyr = (1 / u.year).to_value(u.Hz) + +# %% +# We can create a `PLRedNoise` model from the `WaveX` model. +# This will estimate the spectral parameters from the `WaveX` +# amplitudes. +m3 = plrednoise_from_wavex(m2) +print(m3) + +# %% +# Now let us plot the estimated spectrum with the injected +# spectrum. +plt.subplot(211) +plt.errorbar( + idxs * f0, + b * 1e6, + db * 1e6, + ls="", + marker="o", + label="$\\hat{a}_j$ (WXCOS)", + color="red", +) +plt.errorbar( + idxs * f0, + a * 1e6, + da * 1e6, + ls="", + marker="o", + label="$\\hat{b}_j$ (WXSIN)", + color="blue", +) +plt.axvline(fyr, color="black", ls="dotted") +plt.axhline(0, color="grey", ls="--") +plt.ylabel("Fourier coeffs ($\mu$s)") +plt.xscale("log") +plt.legend(fontsize=8) + +plt.subplot(212) +plt.errorbar( + idxs * f0, P, dP, ls="", marker="o", label="Spectral power (PINT)", color="k" +) +P_inj = m.components["PLRedNoise"].get_noise_weights(t)[::2][:nharm_opt] +plt.plot(idxs * f0, P_inj, label="Injected Spectrum", color="r") +P_est = m3.components["PLRedNoise"].get_noise_weights(t)[::2][:nharm_opt] +print(len(idxs), len(P_est)) +plt.plot(idxs * f0, P_est, label="Estimated Spectrum", color="b") +plt.xscale("log") +plt.yscale("log") +plt.ylabel("Spectral power (s$^2$)") +plt.xlabel("Frequency (Hz)") +plt.axvline(fyr, color="black", ls="dotted", label="1 yr$^{-1}$") +plt.legend() + +# %% [markdown] +# Note the outlier in the 1 year^-1 bin. This is caused by the covariance with RA and DEC, which introduce a delay with the same frequency. + +# %% [markdown] +# ## DM noise fitting +# Let us now do a similar kind of analysis for DM noise. + +# %% +par_sim = """ + PSR SIM4 + RAJ 05:00:00 1 + DECJ 15:00:00 1 + PEPOCH 55000 + F0 100 1 + F1 -1e-15 1 + PHOFF 0 1 + DM 15 1 + TNDMAMP -13 + TNDMGAM 3.5 + TNDMC 30 + TZRMJD 55000 + TZRFRQ 1400 + TZRSITE gbt + UNITS TDB + EPHEM DE440 + CLOCK TT(BIPM2019) +""" + +m = get_model(StringIO(par_sim)) + +# %% +# Generate the simulated TOAs. +ntoas = 2000 +toaerrs = np.random.uniform(0.5, 2.0, ntoas) * u.us +freqs = np.linspace(500, 1500, 8) * u.MHz + +t = make_fake_toas_uniform( + startMJD=53001, + endMJD=57001, + ntoas=ntoas, + model=m, + freq=freqs, + obs="gbt", + error=toaerrs, + add_noise=True, + add_correlated_noise=True, + name="fake", + include_bipm=True, + include_gps=True, + multi_freqs_in_epoch=True, +) + +# %% +# Find the optimum number of harmonics by minimizing AIC. +m1 = deepcopy(m) +m1.remove_component("PLDMNoise") + +m2 = deepcopy(m1) + +nharm_opt, d_aics = find_optimal_nharms(m2, t, "DMWaveX", 30) +print("Optimum no of harmonics = ", nharm_opt) + +# %% +# The Y axis is plotted in log scale only for better visibility. +plt.scatter(list(range(len(d_aics))), d_aics + 1) +plt.axvline(nharm_opt, color="red", label="Optimum number of harmonics") +plt.axvline( + int(m.TNDMC.value), color="black", ls="--", label="Injected number of harmonics" +) +plt.xlabel("Number of harmonics") +plt.ylabel("AIC - AIC$_\\min{} + 1$") +plt.legend() +plt.yscale("log") +# plt.savefig("sim3-aic.pdf") + +# %% +# Now create a new model with the optimum number of +# harmonics +m2 = deepcopy(m1) + +Tspan = t.get_mjds().max() - t.get_mjds().min() +dmwavex_setup(m2, T_span=Tspan, n_freqs=nharm_opt, freeze_params=False) + +ftr = WLSFitter(t, m2) +ftr.fit_toas(maxiter=10) +m2 = ftr.model + +print(m2) + +# %% [markdown] +# ### Estimating the spectral parameters from the `DMWaveX` fit. + +# %% +# Get the Fourier amplitudes and powers and their uncertainties. +# Note that the `DMWaveX` amplitudes have the units of DM. +# We multiply them by a constant factor to convert them to dimensions +# of time so that they are consistent with `PLDMNoise`. +scale = DMconst / (1400 * u.MHz) ** 2 + +idxs = np.array(m2.components["DMWaveX"].get_indices()) +a = np.array( + [(scale * m2[f"DMWXSIN_{idx:04d}"].quantity).to_value("s") for idx in idxs] +) +da = np.array( + [(scale * m2[f"DMWXSIN_{idx:04d}"].uncertainty).to_value("s") for idx in idxs] +) +b = np.array( + [(scale * m2[f"DMWXCOS_{idx:04d}"].quantity).to_value("s") for idx in idxs] +) +db = np.array( + [(scale * m2[f"DMWXCOS_{idx:04d}"].uncertainty).to_value("s") for idx in idxs] +) +print(len(idxs)) + +P = (a**2 + b**2) / 2 +dP = ((a * da) ** 2 + (b * db) ** 2) ** 0.5 + +f0 = (1 / Tspan).to_value(u.Hz) +fyr = (1 / u.year).to_value(u.Hz) + +# %% +# We can create a `PLDMNoise` model from the `DMWaveX` model. +# This will estimate the spectral parameters from the `DMWaveX` +# amplitudes. +m3 = pldmnoise_from_dmwavex(m2) +print(m3) + +# %% +# Now let us plot the estimated spectrum with the injected +# spectrum. +plt.subplot(211) +plt.errorbar( + idxs * f0, + b * 1e6, + db * 1e6, + ls="", + marker="o", + label="$\\hat{a}_j$ (DMWXCOS)", + color="red", +) +plt.errorbar( + idxs * f0, + a * 1e6, + da * 1e6, + ls="", + marker="o", + label="$\\hat{b}_j$ (DMWXSIN)", + color="blue", +) +plt.axvline(fyr, color="black", ls="dotted") +plt.axhline(0, color="grey", ls="--") +plt.ylabel("Fourier coeffs ($\mu$s)") +plt.xscale("log") +plt.legend(fontsize=8) + +plt.subplot(212) +plt.errorbar( + idxs * f0, P, dP, ls="", marker="o", label="Spectral power (PINT)", color="k" +) +P_inj = m.components["PLDMNoise"].get_noise_weights(t)[::2][:nharm_opt] +plt.plot(idxs * f0, P_inj, label="Injected Spectrum", color="r") +P_est = m3.components["PLDMNoise"].get_noise_weights(t)[::2][:nharm_opt] +print(len(idxs), len(P_est)) +plt.plot(idxs * f0, P_est, label="Estimated Spectrum", color="b") +plt.xscale("log") +plt.yscale("log") +plt.ylabel("Spectral power (s$^2$)") +plt.xlabel("Frequency (Hz)") +plt.axvline(fyr, color="black", ls="dotted", label="1 yr$^{-1}$") +plt.legend() diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 6fca7ab2e..ff7ae8b6c 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -63,6 +63,7 @@ are not included in the default build because they take too long, but you can do examples/How_to_build_a_timing_model_component.ipynb examples/understanding_fitters.ipynb examples/noise-fitting-example.ipynb + examples/rednoise-fit-example.ipynb examples/WorkingWithFlags.ipynb examples/Wideband_TOA_walkthrough.ipynb examples/Simulate_and_make_MassMass.ipynb diff --git a/src/pint/fitter.py b/src/pint/fitter.py index 52ace0245..c34ed3332 100644 --- a/src/pint/fitter.py +++ b/src/pint/fitter.py @@ -76,7 +76,6 @@ AngleParameter, boolParameter, strParameter, - funcParameter, ) from pint.pint_matrix import ( CorrelationMatrix, @@ -1111,6 +1110,7 @@ def fit_toas( max_chi2_increase=1e-2, min_lambda=1e-3, noisefit_method="Newton-CG", + compute_noise_uncertainties=True, debug=False, ): """Carry out a cautious downhill fit. @@ -1182,7 +1182,7 @@ def fit_toas( debug=debug, ) - if ii == noise_fit_niter - 1: + if ii == noise_fit_niter - 1 and compute_noise_uncertainties: values, errors = self._fit_noise( noisefit_method=noisefit_method, uncertainty=True ) diff --git a/src/pint/models/dmwavex.py b/src/pint/models/dmwavex.py index bd51f8465..d61560dd2 100644 --- a/src/pint/models/dmwavex.py +++ b/src/pint/models/dmwavex.py @@ -167,7 +167,7 @@ def add_dmwavex_components( frozens = np.repeat(frozens, len(dmwxfreqs)) if len(frozens) != len(dmwxfreqs): raise ValueError( - f"Number of base frequencies must match number of frozen values" + "Number of base frequencies must match number of frozen values" ) #### If indices is None, increment the current max DMWaveX index by 1. Increment using DMWXFREQ dct = self.get_prefix_mapping_component("DMWXFREQ_") @@ -331,16 +331,15 @@ def validate(self): warn(f"Frequency DMWXFREQ_{index:04d} is negative") wfreqs[j] = getattr(self, f"DMWXFREQ_{index:04d}").value wfreqs.sort() - if np.any(np.diff(wfreqs) <= (1.0 / (2.0 * 364.25))): - warn("Frequency resolution is greater than 1/yr") - if self.DMWXEPOCH.value is None: - if self._parent is not None: - if self._parent.PEPOCH.value is None: - raise MissingParameter( - "DMWXEPOCH or PEPOCH are required if DMWaveX is being used" - ) - else: - self.DMWXEPOCH.quantity = self._parent.PEPOCH.quantity + # if np.any(np.diff(wfreqs) <= (1.0 / (2.0 * 364.25))): + # warn("Frequency resolution is greater than 1/yr") + if self.DMWXEPOCH.value is None and self._parent is not None: + if self._parent.PEPOCH.value is None: + raise MissingParameter( + "DMWXEPOCH or PEPOCH are required if DMWaveX is being used" + ) + else: + self.DMWXEPOCH.quantity = self._parent.PEPOCH.quantity def validate_toas(self, toas): return super().validate_toas(toas) diff --git a/src/pint/models/wavex.py b/src/pint/models/wavex.py index d45b82a7a..a5e9acbbb 100644 --- a/src/pint/models/wavex.py +++ b/src/pint/models/wavex.py @@ -181,7 +181,7 @@ def add_wavex_components( frozens = np.repeat(frozens, len(wxfreqs)) if len(frozens) != len(wxfreqs): raise ValueError( - f"Number of base frequencies must match number of frozen values" + "Number of base frequencies must match number of frozen values" ) #### If indices is None, increment the current max WaveX index by 1. Increment using WXFREQ dct = self.get_prefix_mapping_component("WXFREQ_") @@ -341,16 +341,15 @@ def validate(self): warn(f"Frequency WXFREQ_{index:04d} is negative") wfreqs[j] = getattr(self, f"WXFREQ_{index:04d}").value wfreqs.sort() - if np.any(np.diff(wfreqs) <= (1.0 / (2.0 * 364.25))): - warn("Frequency resolution is greater than 1/yr") - if self.WXEPOCH.value is None: - if self._parent is not None: - if self._parent.PEPOCH.value is None: - raise MissingParameter( - "WXEPOCH or PEPOCH are required if WaveX is being used" - ) - else: - self.WXEPOCH.quantity = self._parent.PEPOCH.quantity + # if np.any(np.diff(wfreqs) <= (1.0 / (2.0 * 364.25))): + # warn("Frequency resolution is greater than 1/yr") + if self.WXEPOCH.value is None and self._parent is not None: + if self._parent.PEPOCH.value is None: + raise MissingParameter( + "WXEPOCH or PEPOCH are required if WaveX is being used" + ) + else: + self.WXEPOCH.quantity = self._parent.PEPOCH.quantity def validate_toas(self, toas): return super().validate_toas(toas) diff --git a/src/pint/pintk/pulsar.py b/src/pint/pintk/pulsar.py index 6ed95e7f0..b38be572d 100644 --- a/src/pint/pintk/pulsar.py +++ b/src/pint/pintk/pulsar.py @@ -20,7 +20,7 @@ ) from pint.residuals import Residuals from pint.toa import get_TOAs, merge_TOAs -from pint.utils import FTest +from pint.utils import FTest, akaike_information_criterion import pint.logging from loguru import logger as log @@ -584,6 +584,10 @@ def fit(self, selected, iters=4, compute_random=False): # adds extra prefix params for fitting self.add_model_params() + print( + f"Akaike information criterion = {akaike_information_criterion(self.fitter.model, self.fitter.toas)}" + ) + def random_models(self, selected): """Compute and plot random models""" log.info("Computing random models based on parameter covariance matrix.") diff --git a/src/pint/utils.py b/src/pint/utils.py index b9dc39bda..1c0bcf21c 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -40,8 +40,8 @@ from contextlib import contextmanager from pathlib import Path from warnings import warn -import scipy -import uncertainties +from scipy.optimize import minimize +from numdifftools import Hessian import astropy.constants as const import astropy.coordinates as coords @@ -183,11 +183,11 @@ class PosVel: def __init__(self, pos, vel, obj=None, origin=None): if len(pos) != 3: - raise ValueError("Position vector has length %d instead of 3" % len(pos)) + raise ValueError(f"Position vector has length {len(pos)} instead of 3") self.pos = pos if isinstance(pos, u.Quantity) else np.asarray(pos) if len(vel) != 3: - raise ValueError("Position vector has length %d instead of 3" % len(pos)) + raise ValueError(f"Position vector has length {len(pos)} instead of 3") self.vel = vel if isinstance(vel, u.Quantity) else np.asarray(vel) if len(self.pos.shape) != len(self.vel.shape): @@ -601,10 +601,8 @@ def dmx_ranges_old( # Round off the dates to 0.1 days and only keep unique values so we ignore closely spaced TOAs loMJDs = np.unique(loMJDs.round(1)) hiMJDs = np.unique(hiMJDs.round(1)) - log.info("There are {} dates with freqs > {} MHz".format(len(hiMJDs), divide_freq)) - log.info( - "There are {} dates with freqs < {} MHz\n".format(len(loMJDs), divide_freq) - ) + log.info(f"There are {len(hiMJDs)} dates with freqs > {divide_freq} MHz") + log.info(f"There are {len(loMJDs)} dates with freqs < {divide_freq} MHz\n") DMXs = [] @@ -691,7 +689,7 @@ def dmx_ranges_old( # Mark TOAs as True if they are in any DMX bin for DMX in DMXs: mask[np.logical_and(MJDs > DMX.min - offset, MJDs < DMX.max + offset)] = True - log.info("{} out of {} TOAs are in a DMX bin".format(mask.sum(), len(mask))) + log.info(f"{mask.sum()} out of {len(mask)} TOAs are in a DMX bin") # Instantiate a DMX component dmx_class = Component.component_types["DispersionDMX"] dmx_comp = dmx_class() @@ -768,12 +766,10 @@ def dmx_ranges(toas, divide_freq=1000.0 * u.MHz, binwidth=15.0 * u.d, verbose=Fa DMXs = [] prevbinR2 = MJDs[0] - 0.001 * u.d - while True: + while np.any(MJDs > prevbinR2): # Consider all TOAs with times after the last bin up through a total span of binwidth # Get indexes that should be in this bin # If there are no more MJDs to process, we are done. - if not np.any(MJDs > prevbinR2): - break startMJD = MJDs[MJDs > prevbinR2][0] binidx = np.logical_and(MJDs > prevbinR2, MJDs <= startMJD + binwidth) if not np.any(binidx): @@ -802,7 +798,7 @@ def dmx_ranges(toas, divide_freq=1000.0 * u.MHz, binwidth=15.0 * u.d, verbose=Fa # Mark TOAs as True if they are in any DMX bin for DMX in DMXs: mask[np.logical_and(MJDs >= DMX.min, MJDs <= DMX.max)] = True - log.info("{} out of {} TOAs are in a DMX bin".format(mask.sum(), len(mask))) + log.info(f"{mask.sum()} out of {len(mask)} TOAs are in a DMX bin") # Instantiate a DMX component dmx_class = Component.component_types["DispersionDMX"] dmx_comp = dmx_class() @@ -984,8 +980,8 @@ def dmxparse(fitter, save=False): # Get number of DMX epochs try: DMX_mapping = fitter.model.get_prefix_mapping("DMX_") - except ValueError: - raise RuntimeError("No DMX values in model!") + except ValueError as e: + raise RuntimeError("No DMX values in model!") from e dmx_epochs = [f"{x:04d}" for x in DMX_mapping.keys()] DMX_keys = list(DMX_mapping.values()) DMXs = np.zeros(len(dmx_epochs)) @@ -1017,7 +1013,7 @@ def dmxparse(fitter, save=False): # access by label name to make sure we get the right values # make sure they are sorted in ascending order cc = fitter.parameter_covariance_matrix.get_label_matrix( - sorted(["DMX_" + x for x in dmx_epochs]) + sorted([f"DMX_{x}" for x in dmx_epochs]) ) n = len(DMX_Errs) - np.sum(mask_idxs) # Find error in mean DM @@ -1049,26 +1045,15 @@ def dmxparse(fitter, save=False): if save is not None and save: if isinstance(save, bool): save = "dmxparse.out" - DMX = "DMX" - lines = [] - lines.append("# Mean %s value = %+.6e \n" % (DMX, DMX_mean)) - lines.append("# Uncertainty in average %s = %.5e \n" % ("DM", DMX_mean_err)) - lines.append( - "# Columns: %sEP %s_value %s_var_err %sR1 %sR2 %s_bin \n" - % (DMX, DMX, DMX, DMX, DMX, DMX) + lines = [ + f"# Mean DMX value = {DMX_mean:+.6e} \n", + f"# Uncertainty in average DM = {DMX_mean_err:.5e} \n", + f"# Columns: DMXEP DMX_value DMX_var_err DMXR1 DMXR2 %s_bin \n", + ] + lines.extend( + f"{DMX_center_MJD[k]:.4f} {DMXs[k] - DMX_mean:+.7e} {DMX_vErrs[k]:.3e} {DMX_R1[k]:.4f} {DMX_R2[k]:.4f} {DMX_keys[k]} \n" + for k in range(len(dmx_epochs)) ) - for k in range(len(dmx_epochs)): - lines.append( - "%.4f %+.7e %.3e %.4f %.4f %s \n" - % ( - DMX_center_MJD[k], - DMXs[k] - DMX_mean, - DMX_vErrs[k], - DMX_R1[k], - DMX_R2[k], - DMX_keys[k], - ) - ) with open_or_use(save, mode="w") as dmxout: dmxout.writelines(lines) if isinstance(save, (str, Path)): @@ -1080,18 +1065,16 @@ def dmxparse(fitter, save=False): DMX_units = getattr(fitter.model, "DMX_{:}".format(dmx_epochs[0])).units DMXR_units = getattr(fitter.model, "DMXR1_{:}".format(dmx_epochs[0])).units - # define the output dictionary - dmx = {} - dmx["dmxs"] = mean_sub_DMXs * DMX_units - dmx["dmx_verrs"] = DMX_vErrs * DMX_units - dmx["dmxeps"] = DMX_center_MJD * DMXR_units - dmx["r1s"] = DMX_R1 * DMXR_units - dmx["r2s"] = DMX_R2 * DMXR_units - dmx["bins"] = DMX_keys - dmx["mean_dmx"] = DMX_mean * DMX_units - dmx["avg_dm_err"] = DMX_mean_err * DMX_units - - return dmx + return { + "dmxs": mean_sub_DMXs * DMX_units, + "dmx_verrs": DMX_vErrs * DMX_units, + "dmxeps": DMX_center_MJD * DMXR_units, + "r1s": DMX_R1 * DMXR_units, + "r2s": DMX_R2 * DMXR_units, + "bins": DMX_keys, + "mean_dmx": DMX_mean * DMX_units, + "avg_dm_err": DMX_mean_err * DMX_units, + } def get_prefix_timerange(model, prefixname): @@ -1140,7 +1123,7 @@ def get_prefix_timeranges(model, prefixname): """ if prefixname.endswith("_"): prefixname = prefixname[:-1] - prefix_mapping = model.get_prefix_mapping(prefixname + "_") + prefix_mapping = model.get_prefix_mapping(f"{prefixname}_") r1 = np.zeros(len(prefix_mapping)) r2 = np.zeros(len(prefix_mapping)) indices = np.zeros(len(prefix_mapping), dtype=np.int32) @@ -1317,7 +1300,7 @@ def split_swx(model, time): return index, newindex -def wavex_setup(model, T_span, freqs=None, n_freqs=None): +def wavex_setup(model, T_span, freqs=None, n_freqs=None, freeze_params=False): """ Set-up a WaveX model based on either an array of user-provided frequencies or the wave number frequency calculation. Sine and Cosine amplitudes are initially set to zero @@ -1395,10 +1378,15 @@ def wavex_setup(model, T_span, freqs=None, n_freqs=None): wave_freqs = wave_numbers / T_span model.WXFREQ_0001.quantity = wave_freqs[0] model.components["WaveX"].add_wavex_components(wave_freqs[1:]) + + for p in model.params: + if p.startswith("WXSIN") or p.startswith("WXCOS"): + model[p].frozen = freeze_params + return model.components["WaveX"].get_indices() -def dmwavex_setup(model, T_span, freqs=None, n_freqs=None): +def dmwavex_setup(model, T_span, freqs=None, n_freqs=None, freeze_params=False): """ Set-up a DMWaveX model based on either an array of user-provided frequencies or the wave number frequency calculation. Sine and Cosine amplitudes are initially set to zero @@ -1476,6 +1464,11 @@ def dmwavex_setup(model, T_span, freqs=None, n_freqs=None): wave_freqs = wave_numbers / T_span model.DMWXFREQ_0001.quantity = wave_freqs[0] model.components["DMWaveX"].add_dmwavex_components(wave_freqs[1:]) + + for p in model.params: + if p.startswith("DMWXSIN") or p.startswith("DMWXCOS"): + model[p].frozen = freeze_params + return model.components["DMWaveX"].get_indices() @@ -1528,15 +1521,12 @@ def _translate_wavex_freqs(wxfreq, k): wxfreq *= u.d**-1 if len(wxfreq) == 1: return (2.0 * np.pi * wxfreq) / (k + 1.0) - else: - wave_om = [ - ((2.0 * np.pi * wxfreq[i]) / (k[i] + 1.0)) for i in range(len(wxfreq)) - ] - if np.allclose(wave_om, wave_om[0], atol=1e-3): - om = sum(wave_om) / len(wave_om) - return om - else: - return False + wave_om = [((2.0 * np.pi * wxfreq[i]) / (k[i] + 1.0)) for i in range(len(wxfreq))] + return ( + sum(wave_om) / len(wave_om) + if np.allclose(wave_om, wave_om[0], atol=1e-3) + else False + ) def translate_wave_to_wavex(model): @@ -1613,10 +1603,10 @@ def get_wavex_freqs(model, index=None, quantity=False): ] elif isinstance(index, (int, float, np.int64)): idx_rf = f"{int(index):04d}" - values = getattr(model.components["WaveX"], "WXFREQ_" + idx_rf) + values = getattr(model.components["WaveX"], f"WXFREQ_{idx_rf}") elif isinstance(index, (list, set, np.ndarray)): idx_rf = [f"{int(idx):04d}" for idx in index] - values = [getattr(model.components["WaveX"], "WXFREQ_" + ind) for ind in idx_rf] + values = [getattr(model.components["WaveX"], f"WXFREQ_{ind}") for ind in idx_rf] else: raise TypeError( f"index most be a float, int, set, list, array, or None - not {type(index)}" @@ -1654,30 +1644,28 @@ def get_wavex_amps(model, index=None, quantity=False): model.components["WaveX"].get_prefix_mapping_component("WXSIN_").keys() ) if len(indices) == 1: - values = ( - getattr(model.components["WaveX"], "WXSIN_" + f"{int(indices):04d}"), - getattr(model.components["WaveX"], "WXCOS_" + f"{int(indices):04d}"), - ) + values = getattr( + model.components["WaveX"], f"WXSIN_{int(indices):04d}" + ), getattr(model.components["WaveX"], f"WXCOS_{int(indices):04d}") else: values = [ ( - getattr(model.components["WaveX"], "WXSIN_" + f"{int(idx):04d}"), - getattr(model.components["WaveX"], "WXCOS_" + f"{int(idx):04d}"), + getattr(model.components["WaveX"], f"WXSIN_{int(idx):04d}"), + getattr(model.components["WaveX"], f"WXCOS_{int(idx):04d}"), ) for idx in indices ] elif isinstance(index, (int, float, np.int64)): idx_rf = f"{int(index):04d}" - values = ( - getattr(model.components["WaveX"], "WXSIN_" + idx_rf), - getattr(model.components["WaveX"], "WXCOS_" + idx_rf), + values = getattr(model.components["WaveX"], f"WXSIN_{idx_rf}"), getattr( + model.components["WaveX"], f"WXCOS_{idx_rf}" ) elif isinstance(index, (list, set, np.ndarray)): idx_rf = [f"{int(idx):04d}" for idx in index] values = [ ( - getattr(model.components["WaveX"], "WXSIN_" + ind), - getattr(model.components["WaveX"], "WXCOS_" + ind), + getattr(model.components["WaveX"], f"WXSIN_{ind}"), + getattr(model.components["WaveX"], f"WXCOS_{ind}"), ) for ind in idx_rf ] @@ -1689,7 +1677,7 @@ def get_wavex_amps(model, index=None, quantity=False): if isinstance(values, tuple): values = tuple(v.quantity for v in values) if isinstance(values, list): - values = [tuple((v[0].quantity, v[1].quantity)) for v in values] + values = [(v[0].quantity, v[1].quantity) for v in values] return values @@ -1776,10 +1764,7 @@ def weighted_mean(arrin, weights_in, inputmean=None, calcerr=False, sdev=False): weights = weights_in wtot = weights.sum() # user has input a mean value - if inputmean is None: - wmean = (weights * arr).sum() / wtot - else: - wmean = float(inputmean) + wmean = (weights * arr).sum() / wtot if inputmean is None else float(inputmean) # how should error be calculated? if calcerr: werr2 = (weights**2 * (arr - wmean) ** 2).sum() @@ -1830,10 +1815,12 @@ def ELL1_check( lhs = A1 / const.c * E**4.0 rhs = TRES / np.sqrt(NTOA) if outstring: - s = "Checking applicability of ELL1 model -- \n" - s += " Condition is asini/c * ecc**4 << timing precision / sqrt(# TOAs) to use ELL1\n" - s += " asini/c * ecc**4 = {:.3g} \n".format(lhs.to(u.us)) - s += " TRES / sqrt(# TOAs) = {:.3g} \n".format(rhs.to(u.us)) + s = ( + f"Checking applicability of ELL1 model -- \n" + f" Condition is asini/c * ecc**4 << timing precision / sqrt(# TOAs) to use ELL1\n" + f" asini/c * ecc**4 = {lhs.to(u.us):.3g} \n" + f" TRES / sqrt(# TOAs) = {rhs.to(u.us):.3g} \n" + ) if lhs * 50.0 < rhs: if outstring: s += " Should be fine.\n" @@ -1888,20 +1875,15 @@ def FTest(chi2_1, dof_1, chi2_2, dof_2): delta_dof = dof_1 - dof_2 new_redchi2 = chi2_2 / dof_2 F = float((delta_chi2 / delta_dof) / new_redchi2) # fdtr doesn't like float128 - ft = fdtrc(delta_dof, dof_2, F) + return fdtrc(delta_dof, dof_2, F) elif dof_1 == dof_2: log.warning("Models have equal degrees of freedom, cannot perform F-test.") - ft = np.nan - elif delta_chi2 <= 0: + return np.nan + else: log.warning( "Chi^2 for Model 2 is larger than Chi^2 for Model 1, cannot perform F-test." ) - ft = 1.0 - else: - raise ValueError( - f"Mystery problem in Ftest - maybe NaN? {chi2_1} {dof_1} {chi2_2} {dof_2}" - ) - return ft + return 1.0 def add_dummy_distance(c, distance=1 * u.kpc): @@ -1922,13 +1904,13 @@ def add_dummy_distance(c, distance=1 * u.kpc): if c.frame.data.differentials == {}: log.warning( - "No proper motions available for %r: returning coordinates unchanged" % c + f"No proper motions available for {c}: returning coordinates unchanged" ) return c if isinstance(c.frame, coords.builtin_frames.icrs.ICRS): - if hasattr(c, "pm_ra_cosdec"): - cnew = coords.SkyCoord( + return ( + coords.SkyCoord( ra=c.ra, dec=c.dec, pm_ra_cosdec=c.pm_ra_cosdec, @@ -1937,11 +1919,8 @@ def add_dummy_distance(c, distance=1 * u.kpc): distance=distance, frame=coords.ICRS, ) - else: - # it seems that after applying proper motions - # it changes the RA pm to pm_ra instead of pm_ra_cosdec - # although the value seems the same - cnew = coords.SkyCoord( + if hasattr(c, "pm_ra_cosdec") + else coords.SkyCoord( ra=c.ra, dec=c.dec, pm_ra_cosdec=c.pm_ra, @@ -1950,10 +1929,9 @@ def add_dummy_distance(c, distance=1 * u.kpc): distance=distance, frame=coords.ICRS, ) - - return cnew + ) elif isinstance(c.frame, coords.builtin_frames.galactic.Galactic): - cnew = coords.SkyCoord( + return coords.SkyCoord( l=c.l, b=c.b, pm_l_cosb=c.pm_l_cosb, @@ -1962,9 +1940,8 @@ def add_dummy_distance(c, distance=1 * u.kpc): distance=distance, frame=coords.Galactic, ) - return cnew elif isinstance(c.frame, pint.pulsar_ecliptic.PulsarEcliptic): - cnew = coords.SkyCoord( + return coords.SkyCoord( lon=c.lon, lat=c.lat, pm_lon_coslat=c.pm_lon_coslat, @@ -1974,10 +1951,9 @@ def add_dummy_distance(c, distance=1 * u.kpc): obliquity=c.obliquity, frame=pint.pulsar_ecliptic.PulsarEcliptic, ) - return cnew else: log.warning( - "Do not know coordinate frame for %r: returning coordinates unchanged" % c + f"Do not know coordinate frame for {c}: returning coordinates unchanged" ) return c @@ -1998,12 +1974,12 @@ def remove_dummy_distance(c): if c.frame.data.differentials == {}: log.warning( - "No proper motions available for %r: returning coordinates unchanged" % c + f"No proper motions available for {c}: returning coordinates unchanged" ) return c if isinstance(c.frame, coords.builtin_frames.icrs.ICRS): - if hasattr(c, "pm_ra_cosdec"): - cnew = coords.SkyCoord( + return ( + coords.SkyCoord( ra=c.ra, dec=c.dec, pm_ra_cosdec=c.pm_ra_cosdec, @@ -2011,11 +1987,8 @@ def remove_dummy_distance(c): obstime=c.obstime, frame=coords.ICRS, ) - else: - # it seems that after applying proper motions - # it changes the RA pm to pm_ra instead of pm_ra_cosdec - # although the value seems the same - cnew = coords.SkyCoord( + if hasattr(c, "pm_ra_cosdec") + else coords.SkyCoord( ra=c.ra, dec=c.dec, pm_ra_cosdec=c.pm_ra, @@ -2023,9 +1996,9 @@ def remove_dummy_distance(c): obstime=c.obstime, frame=coords.ICRS, ) - return cnew + ) elif isinstance(c.frame, coords.builtin_frames.galactic.Galactic): - cnew = coords.SkyCoord( + return coords.SkyCoord( l=c.l, b=c.b, pm_l_cosb=c.pm_l_cosb, @@ -2033,9 +2006,8 @@ def remove_dummy_distance(c): obstime=c.obstime, frame=coords.Galactic, ) - return cnew elif isinstance(c.frame, pint.pulsar_ecliptic.PulsarEcliptic): - cnew = coords.SkyCoord( + return coords.SkyCoord( lon=c.lon, lat=c.lat, pm_lon_coslat=c.pm_lon_coslat, @@ -2044,10 +2016,9 @@ def remove_dummy_distance(c): obliquity=c.obliquity, frame=pint.pulsar_ecliptic.PulsarEcliptic, ) - return cnew else: log.warning( - "Do not know coordinate frame for %r: returning coordinates unchanged" % c + f"Do not know coordinate frame for {c}: returning coordinates unchanged" ) return c @@ -2216,10 +2187,7 @@ def info_string(prefix_string="# ", comment=None, detailed=False): } ) - s = "" - for key, val in info_dict.items(): - s += f"{key}: {val}\n" - + s = "".join(f"{key}: {val}\n" for key, val in info_dict.items()) s = textwrap.dedent(s) # remove blank lines s = os.linesep.join([x for x in s.splitlines() if x]) @@ -2510,8 +2478,7 @@ def divide_times(t, t0, offset=0.5): """ dt = t - t0 values = (dt.to(u.yr).value + offset) // 1 - indices = np.digitize(values, np.unique(values), right=True) - return indices + return np.digitize(values, np.unique(values), right=True) def convert_dispersion_measure(dm, dmconst=None): @@ -2764,3 +2731,244 @@ def woodbury_dot(Ndiag, U, Phidiag, x, y): logdet_C = logdet_N + logdet_Phi + logdet_Sigma return x_Cinv_y, logdet_C + + +def _get_wx2pl_lnlike(model, component_name, ignore_fyr=True): + from pint.models.noise_model import powerlaw + from pint import DMconst + + assert component_name in ["WaveX", "DMWaveX"] + prefix = "WX" if component_name == "WaveX" else "DMWX" + + idxs = np.array(model.components[component_name].get_indices()) + + fs = np.array( + [model[f"{prefix}FREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] + ) + f0 = np.min(fs) + fyr = (1 / u.year).to_value(u.Hz) + + assert np.allclose( + np.diff(np.diff(fs)), 0 + ), "[DM]WaveX frequencies must be uniformly spaced." + + if ignore_fyr: + year_mask = np.abs(((fs - fyr) / f0)) > 0.5 + + idxs = idxs[year_mask] + fs = np.array( + [model[f"{prefix}FREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] + ) + f0 = np.min(fs) + + scaling_factor = 1 if component_name == "WaveX" else DMconst / (1400 * u.MHz) ** 2 + + a = np.array( + [ + (scaling_factor * model[f"{prefix}SIN_{idx:04d}"].quantity).to_value(u.s) + for idx in idxs + ] + ) + da = np.array( + [ + (scaling_factor * model[f"{prefix}SIN_{idx:04d}"].uncertainty).to_value(u.s) + for idx in idxs + ] + ) + b = np.array( + [ + (scaling_factor * model[f"{prefix}COS_{idx:04d}"].quantity).to_value(u.s) + for idx in idxs + ] + ) + db = np.array( + [ + (scaling_factor * model[f"{prefix}COS_{idx:04d}"].uncertainty).to_value(u.s) + for idx in idxs + ] + ) + + def powl_model(params): + """Get the powerlaw spectrum for the WaveX frequencies for a given + set of parameters. This calls the powerlaw function used by `PLRedNoise`/`PLDMNoise`. + """ + gamma, log10_A = params + return (powerlaw(fs, A=10**log10_A, gamma=gamma) * f0) ** 0.5 + + def mlnlike(params): + """Negative of the likelihood function that acts on the + `[DM]WaveX` amplitudes.""" + sigma = powl_model(params) + return 0.5 * np.sum( + (a**2 / (sigma**2 + da**2)) + + (b**2 / (sigma**2 + db**2)) + + np.log(sigma**2 + da**2) + + np.log(sigma**2 + db**2) + ) + + return mlnlike + + +def plrednoise_from_wavex(model, ignore_fyr=True): + """Convert a `WaveX` representation of red noise to a `PLRedNoise` + representation. This is done by minimizing a likelihood function + that acts on the `WaveX` amplitudes over the powerlaw spectral + parameters. + + Parameters + ---------- + model: pint.models.timing_model.TimingModel + The timing model with a `WaveX` component. + ignore_fyr: bool + Whether to ignore the frequency bin containinf 1 yr^-1 + while fitting for the spectral parameters. + + Returns + ------- + pint.models.timing_model.TimingModel + The timing model with a converted `PLRedNoise` component. + """ + from pint.models.noise_model import PLRedNoise + + mlnlike = _get_wx2pl_lnlike(model, "WaveX", ignore_fyr=ignore_fyr) + + result = minimize(mlnlike, [4, -13], method="Nelder-Mead") + if not result.success: + raise ValueError("Log-likelihood maximization failed to converge.") + + gamma_val, log10_A_val = result.x + + hess = Hessian(mlnlike) + gamma_err, log10_A_err = np.sqrt( + np.diag(np.linalg.pinv(hess((gamma_val, log10_A_val)))) + ) + + tnredc = len(model.components["WaveX"].get_indices()) + + model1 = deepcopy(model) + model1.remove_component("WaveX") + model1.add_component(PLRedNoise()) + model1.TNREDAMP.value = log10_A_val + model1.TNREDGAM.value = gamma_val + model1.TNREDC.value = tnredc + model1.TNREDAMP.uncertainty_value = log10_A_err + model1.TNREDGAM.uncertainty_value = gamma_err + + return model1 + + +def pldmnoise_from_dmwavex(model, ignore_fyr=False): + """Convert a `DMWaveX` representation of red noise to a `PLDMNoise` + representation. This is done by minimizing a likelihood function + that acts on the `DMWaveX` amplitudes over the powerlaw spectral + parameters. + + Parameters + ---------- + model: pint.models.timing_model.TimingModel + The timing model with a `DMWaveX` component. + + Returns + ------- + pint.models.timing_model.TimingModel + The timing model with a converted `PLDMNoise` component. + """ + from pint.models.noise_model import PLDMNoise + + mlnlike = _get_wx2pl_lnlike(model, "DMWaveX", ignore_fyr=ignore_fyr) + + result = minimize(mlnlike, [4, -13], method="Nelder-Mead") + if not result.success: + raise ValueError("Log-likelihood maximization failed to converge.") + + gamma_val, log10_A_val = result.x + + hess = Hessian(mlnlike) + + H = hess((gamma_val, log10_A_val)) + assert np.all(np.linalg.eigvals(H) > 0), "The Hessian is not positive definite!" + + Hinv = np.linalg.pinv(H) + assert np.all( + np.linalg.eigvals(Hinv) > 0 + ), "The inverse Hessian is not positive definite!" + + gamma_err, log10_A_err = np.sqrt(np.diag(Hinv)) + + tndmc = len(model.components["DMWaveX"].get_indices()) + + model1 = deepcopy(model) + model1.remove_component("DMWaveX") + model1.add_component(PLDMNoise()) + model1.TNDMAMP.value = log10_A_val + model1.TNDMGAM.value = gamma_val + model1.TNDMC.value = tndmc + model1.TNDMAMP.uncertainty_value = log10_A_err + model1.TNDMGAM.uncertainty_value = gamma_err + + return model1 + + +def find_optimal_nharms(model, toas, component, nharms_max=45): + """Find the optimal number of harmonics for `WaveX`/`DMWaveX` using the Akaike Information + Criterion. + + Parameters + ---------- + model: `pint.models.timing_model.TimingModel` + The timing model. Should not already contain `WaveX`/`DMWaveX` or `PLRedNoise`/`PLDMNoise`. + toas: `pint.toa.TOAs` + Input TOAs + component: str + Component name; "WaveX" or "DMWaveX" + nharms_max: int + Maximum number of harmonics + + Returns + ------- + nharms_opt: int + Optimal number of harmonics + aics: ndarray + Array of normalized AIC values. + """ + from pint.fitter import Fitter + + assert component in ["WaveX", "DMWaveX"] + assert ( + component not in model.components + ), f"{component} is already included in the model." + assert ( + "PLRedNoise" not in model.components and "PLDMNoise" not in model.components + ), "PLRedNoise/PLDMNoise cannot be included in the model." + + model1 = deepcopy(model) + + ftr = Fitter.auto(toas, model1, downhill=False) + ftr.fit_toas(maxiter=5) + aics = [akaike_information_criterion(model1, toas)] + model1 = ftr.model + + T_span = toas.get_mjds().max() - toas.get_mjds().min() + setup_component = wavex_setup if component == "WaveX" else dmwavex_setup + setup_component(model1, T_span, n_freqs=1, freeze_params=False) + + for _ in range(nharms_max): + ftr = Fitter.auto(toas, model1, downhill=False) + ftr.fit_toas(maxiter=5) + aics.append(akaike_information_criterion(ftr.model, toas)) + + model1 = ftr.model + if component == "WaveX": + model1.components[component].add_wavex_component( + (len(model1.components[component].get_indices()) + 1) / T_span, + frozen=False, + ) + else: + model1.components[component].add_dmwavex_component( + (len(model1.components[component].get_indices()) + 1) / T_span, + frozen=False, + ) + + assert all(np.isfinite(aics)), "Infs/NaNs found in AICs!" + + return np.argmin(aics), np.array(aics) - np.min(aics) diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py new file mode 100644 index 000000000..c10ba73dc --- /dev/null +++ b/tests/test_wx2pl.py @@ -0,0 +1,173 @@ +import pytest +from pint.models import get_model +from pint.simulation import make_fake_toas_uniform +from pint.fitter import WLSFitter +from pint.utils import ( + dmwavex_setup, + find_optimal_nharms, + wavex_setup, + plrednoise_from_wavex, + pldmnoise_from_dmwavex, +) + +from io import StringIO +import numpy as np +import astropy.units as u +from copy import deepcopy + + +@pytest.fixture +def data_wx(): + par_sim_wx = """ + PSR SIM3 + RAJ 05:00:00 1 + DECJ 15:00:00 1 + PEPOCH 55000 + F0 100 1 + F1 -1e-15 1 + PHOFF 0 1 + DM 15 1 + TNREDAMP -12.5 + TNREDGAM 3.5 + TNREDC 10 + TZRMJD 55000 + TZRFRQ 1400 + TZRSITE gbt + UNITS TDB + EPHEM DE440 + CLOCK TT(BIPM2019) + """ + m = get_model(StringIO(par_sim_wx)) + + ntoas = 200 + toaerrs = np.random.uniform(0.5, 2.0, ntoas) * u.us + freqs = np.linspace(500, 1500, 2) * u.MHz + + t = make_fake_toas_uniform( + startMJD=54001, + endMJD=56001, + ntoas=ntoas, + model=m, + freq=freqs, + obs="gbt", + error=toaerrs, + add_noise=True, + add_correlated_noise=True, + name="fake", + include_bipm=True, + include_gps=True, + multi_freqs_in_epoch=True, + ) + + return m, t + + +@pytest.fixture +def data_dmwx(): + par_sim_dmwx = """ + PSR SIM3 + RAJ 05:00:00 1 + DECJ 15:00:00 1 + PEPOCH 55000 + F0 100 1 + F1 -1e-15 1 + PHOFF 0 1 + DM 15 1 + TNDMAMP -13 + TNDMGAM 3.5 + TNDMC 10 + TZRMJD 55000 + TZRFRQ 1400 + TZRSITE gbt + UNITS TDB + EPHEM DE440 + CLOCK TT(BIPM2019) + """ + + m = get_model(StringIO(par_sim_dmwx)) + + ntoas = 200 + toaerrs = np.random.uniform(0.5, 2.0, ntoas) * u.us + freqs = np.linspace(500, 1500, 4) * u.MHz + + t = make_fake_toas_uniform( + startMJD=54001, + endMJD=56001, + ntoas=ntoas, + model=m, + freq=freqs, + obs="gbt", + error=toaerrs, + add_noise=True, + add_correlated_noise=True, + name="fake", + include_bipm=True, + include_gps=True, + multi_freqs_in_epoch=True, + ) + + return m, t + + +def test_wx2pl(data_wx): + m, t = data_wx + + m1 = deepcopy(m) + m1.remove_component("PLRedNoise") + + Tspan = t.get_mjds().max() - t.get_mjds().min() + wavex_setup(m1, Tspan, n_freqs=int(m.TNREDC.value), freeze_params=False) + + ftr = WLSFitter(t, m1) + ftr.fit_toas(maxiter=10) + m1 = ftr.model + + m2 = plrednoise_from_wavex(m1) + + assert "PLRedNoise" in m2.components + assert abs(m.TNREDAMP.value - m2.TNREDAMP.value) / m2.TNREDAMP.uncertainty_value < 5 + assert abs(m.TNREDGAM.value - m2.TNREDGAM.value) / m2.TNREDGAM.uncertainty_value < 5 + + +def test_dmwx2pldm(data_dmwx): + m, t = data_dmwx + + m1 = deepcopy(m) + m1.remove_component("PLDMNoise") + + Tspan = t.get_mjds().max() - t.get_mjds().min() + dmwavex_setup(m1, Tspan, n_freqs=int(m.TNDMC.value), freeze_params=False) + + ftr = WLSFitter(t, m1) + ftr.fit_toas(maxiter=10) + m1 = ftr.model + + m2 = pldmnoise_from_dmwavex(m1) + + assert "PLDMNoise" in m2.components + assert abs(m.TNDMAMP.value - m2.TNDMAMP.value) / m2.TNDMAMP.uncertainty_value < 5 + assert abs(m.TNDMGAM.value - m2.TNDMGAM.value) / m2.TNDMGAM.uncertainty_value < 5 + + +def test_find_optimal_nharms_wx(data_wx): + m, t = data_wx + + m1 = deepcopy(m) + m1.remove_component("PLRedNoise") + + nharm, aics = find_optimal_nharms(m1, t, "WaveX", nharms_max=7) + + assert nharm <= 7 + assert np.all(aics >= 0) + + +def test_find_optimal_nharms_dmwx(data_dmwx): + m, t = data_dmwx + + m1 = deepcopy(m) + m1.remove_component("PLDMNoise") + + nharm, aics = find_optimal_nharms(m1, t, "DMWaveX", nharms_max=7) + + assert nharm <= 7 + assert np.all(aics >= 0)