From c34911994a63871deac808598925336a1fcb27eb Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 15 Dec 2023 11:50:05 -0600 Subject: [PATCH 01/36] plrednoise_from_wavex --- src/pint/utils.py | 73 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 71 insertions(+), 2 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index a06ba5e23..129816749 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 @@ -2630,3 +2630,72 @@ def akaike_information_criterion(model, toas): return 2 * (k - lnL) else: raise NotImplementedError + + +def plrednoise_from_wavex(model): + """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. + + Returns + ------- + pint.models.timing_model.TimingModel + The timing model with a converted `PLRedNoise` component. + """ + from pint.models.noise_model import powerlaw, PLRedNoise + + idxs = np.array(model.components["WaveX"].get_indices()) + a = np.array([model[f"WXSIN_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) + da = np.array([model[f"WXSIN_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) + b = np.array([model[f"WXCOS_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) + db = np.array([model[f"WXCOS_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) + + fs = np.array([model[f"WXFREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs]) + f0 = np.min(fs) + + assert np.allclose( + np.diff(np.diff(fs)), 0 + ), "`get_powerlaw_from_wavex` requires the WaveX frequencies to be uniformly spaced." + + 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`.""" + 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 + `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) + ) + + gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x + + hess = Hessian(mlnlike) + gamma_err, log10_A_err = np.sqrt( + np.diag(np.linalg.pinv(hess((gamma_val, log10_A_val)))) + ) + + # return list(zip(maxlike_params, maxlike_errs)) + + 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 = len(idxs) + model1.TNREDAMP.uncertainty_value = log10_A_err + model1.TNREDGAM.uncertainty_value = gamma_err + + return model1 From 6b026865f0ff9d3c2d65287f1fd38da9087caf7d Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 15 Dec 2023 12:30:49 -0600 Subject: [PATCH 02/36] rednoise-fit-example --- docs/examples/rednoise-fit-example.py | 249 ++++++++++++++++++++++++++ src/pint/models/wavex.py | 17 +- 2 files changed, 257 insertions(+), 9 deletions(-) create mode 100644 docs/examples/rednoise-fit-example.py diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py new file mode 100644 index 000000000..7e9521144 --- /dev/null +++ b/docs/examples/rednoise-fit-example.py @@ -0,0 +1,249 @@ +# %% [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.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 akaike_information_criterion, wavex_setup, plrednoise_from_wavex + +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, +) + +# %% +# Write the model and toas to disk. +m.write_parfile("sim3.par") +t.write_TOA_file("sim3.tim") + +# %% +# We also need the WaveX version of the par file. +m1 = deepcopy(m) +m1.remove_component("PLRedNoise") + +Tspan = t.get_mjds().max() - t.get_mjds().min() +wavex_setup(m1, Tspan, n_freqs=45) + +for p in m1.params: + if p.startswith("WXSIN") or p.startswith("WXCOS"): + m1[p].frozen = False + +m1.write_parfile("sim3.wx.par") + +# %% [markdown] +# ### Initial fitting + +# %% +ftr = WLSFitter(t, m1) + +# %% +ftr.fit_toas(maxiter=15) + +# %% +m1 = ftr.model +print(m1) + +# %% [markdown] +# ### Optimal number of harmonics +# The optimal number of harmonics can be estimated +# using the Akaike Information Criterion (AIC). + +# %% +m2 = deepcopy(m1) + +aics = [] +idxs = m2.components["WaveX"].get_indices() + +ftr = WLSFitter(t, m2) +ftr.fit_toas(maxiter=3) +aic = akaike_information_criterion(ftr.model, t) +aics += [aic] +print(f"{len(idxs)}\t{aic}\t{ftr.resids.chi2_reduced}") + +for idx in reversed(idxs): + if idx == 1: + m2.remove_component("WaveX") + else: + m2.components["WaveX"].remove_wavex_component(idx) + + ftr = WLSFitter(t, m2) + ftr.fit_toas(maxiter=3) + aic = akaike_information_criterion(ftr.model, t) + aics += [aic] + print(f"{idx-1}\t{aic}\t{ftr.resids.chi2_reduced}") + +# %% +# Find the optimum number of harmonics by minimizing AIC. +d_aics = np.array(aics) - np.min(aics) +nharm_opt = len(d_aics) - 1 - np.argmin(d_aics) +print("Optimum no of harmonics = ", nharm_opt) + +# %% +# The Y axis is plotted in log scale only for better visibility. +plt.scatter(list(reversed(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) + +idxs = m2.components["WaveX"].get_indices() +for idx in reversed(idxs): + if idx > nharm_opt: + m2.components["WaveX"].remove_wavex_component(idx) + +ftr = WLSFitter(t, m2) +ftr.fit_toas(maxiter=5) +m2 = ftr.model + +print(m2) + +# %% [markdown] +# ### Estimate 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.legend() diff --git a/src/pint/models/wavex.py b/src/pint/models/wavex.py index 12286169f..c43249617 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_") @@ -343,14 +343,13 @@ def validate(self): 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 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) From d85c6f54435c96a8ce24762e12f6f25f645076ec Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 15 Dec 2023 13:44:51 -0600 Subject: [PATCH 03/36] pldmnoise_from_dmwavex --- src/pint/utils.py | 90 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 87 insertions(+), 3 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 129816749..34224576d 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2661,7 +2661,7 @@ def plrednoise_from_wavex(model): assert np.allclose( np.diff(np.diff(fs)), 0 - ), "`get_powerlaw_from_wavex` requires the WaveX frequencies to be uniformly spaced." + ), "`plrednoise_from_wavex` requires the WaveX frequencies to be uniformly spaced." def powl_model(params): """Get the powerlaw spectrum for the WaveX frequencies for a given @@ -2687,8 +2687,6 @@ def mlnlike(params): np.diag(np.linalg.pinv(hess((gamma_val, log10_A_val)))) ) - # return list(zip(maxlike_params, maxlike_errs)) - model1 = deepcopy(model) model1.remove_component("WaveX") model1.add_component(PLRedNoise()) @@ -2699,3 +2697,89 @@ def mlnlike(params): model1.TNREDGAM.uncertainty_value = gamma_err return model1 + + +def pldmnoise_from_dmwavex(model): + """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 powerlaw, PLDMNoise + from pint import DMconst + + scale = DMconst / (1400 * u.MHz) ** 2 + + idxs = np.array(model.components["DMWaveX"].get_indices()) + a = np.array( + [(scale * model[f"DMWXSIN_{idx:04d}"].quantity).to_value(u.s) for idx in idxs] + ) + da = np.array( + [ + (scale * model[f"DMWXSIN_{idx:04d}"].uncertainty).to_value(u.s) + for idx in idxs + ] + ) + b = np.array( + [(scale * model[f"DMWXCOS_{idx:04d}"].quantity).to_value(u.s) for idx in idxs] + ) + db = np.array( + [ + (scale * model[f"DMWXCOS_{idx:04d}"].uncertainty).to_value(u.s) + for idx in idxs + ] + ) + + fs = np.array( + [model[f"DMWXFREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] + ) + f0 = np.min(fs) + + assert np.allclose( + np.diff(np.diff(fs)), 0 + ), "`pldmnoise_from_dmwavex` requires the DMWaveX frequencies to be uniformly spaced." + + def powl_model(params): + """Get the powerlaw spectrum for the DMWaveX frequencies for a given + set of parameters. This calls the powerlaw function used by `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 + `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) + ) + + gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x + + hess = Hessian(mlnlike) + gamma_err, log10_A_err = np.sqrt( + np.diag(np.linalg.pinv(hess((gamma_val, log10_A_val)))) + ) + + 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 = len(idxs) + model1.TNDMAMP.uncertainty_value = log10_A_err + model1.TNDMGAM.uncertainty_value = gamma_err + + return model1 From 8d2246a485855c5f0dbb23c0b7369767f4caf24f Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Mon, 18 Dec 2023 12:44:47 -0600 Subject: [PATCH 04/36] ignore_fyr --- src/pint/utils.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 34224576d..6487b7e3a 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2632,7 +2632,7 @@ def akaike_information_criterion(model, toas): raise NotImplementedError -def plrednoise_from_wavex(model): +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 @@ -2642,6 +2642,9 @@ def plrednoise_from_wavex(model): ---------- 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 ------- @@ -2651,18 +2654,29 @@ def plrednoise_from_wavex(model): from pint.models.noise_model import powerlaw, PLRedNoise idxs = np.array(model.components["WaveX"].get_indices()) - a = np.array([model[f"WXSIN_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) - da = np.array([model[f"WXSIN_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) - b = np.array([model[f"WXCOS_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) - db = np.array([model[f"WXCOS_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) fs = np.array([model[f"WXFREQ_{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 ), "`plrednoise_from_wavex` requires the WaveX frequencies to be uniformly spaced." + if ignore_fyr: + year_mask = np.abs(((fs - fyr) / f0)) > 0.5 + + idxs = idxs[year_mask] + fs = np.array( + [model[f"WXFREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] + ) + f0 = np.min(fs) + + a = np.array([model[f"WXSIN_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) + da = np.array([model[f"WXSIN_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) + b = np.array([model[f"WXCOS_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) + db = np.array([model[f"WXCOS_{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`.""" @@ -2692,7 +2706,7 @@ def mlnlike(params): model1.add_component(PLRedNoise()) model1.TNREDAMP.value = log10_A_val model1.TNREDGAM.value = gamma_val - model1.TNREDC.value = len(idxs) + model1.TNREDC.value = len(idxs) + 1 if ignore_fyr else len(idxs) model1.TNREDAMP.uncertainty_value = log10_A_err model1.TNREDGAM.uncertainty_value = gamma_err From 8771dc204d54c01d29fa785051c335d548235659 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Mon, 18 Dec 2023 12:45:07 -0600 Subject: [PATCH 05/36] example --- docs/examples/rednoise-fit-example.py | 194 ++++++++++++++++++++++++-- 1 file changed, 184 insertions(+), 10 deletions(-) diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py index 7e9521144..b768a16ff 100644 --- a/docs/examples/rednoise-fit-example.py +++ b/docs/examples/rednoise-fit-example.py @@ -16,11 +16,18 @@ # 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 akaike_information_criterion, wavex_setup, plrednoise_from_wavex +from pint.utils import ( + akaike_information_criterion, + dmwavex_setup, + wavex_setup, + plrednoise_from_wavex, + pldmnoise_from_dmwavex, +) from io import StringIO import numpy as np @@ -84,11 +91,6 @@ multi_freqs_in_epoch=True, ) -# %% -# Write the model and toas to disk. -m.write_parfile("sim3.par") -t.write_TOA_file("sim3.tim") - # %% # We also need the WaveX version of the par file. m1 = deepcopy(m) @@ -101,18 +103,14 @@ if p.startswith("WXSIN") or p.startswith("WXCOS"): m1[p].frozen = False -m1.write_parfile("sim3.wx.par") - # %% [markdown] # ### Initial fitting # %% ftr = WLSFitter(t, m1) -# %% ftr.fit_toas(maxiter=15) -# %% m1 = ftr.model print(m1) @@ -247,3 +245,179 @@ plt.ylabel("Spectral power (s$^2$)") plt.xlabel("Frequency (Hz)") plt.legend() + +# %% [markdown] +# ## DM noise fitting +# Let us now do a similar kind of analysis for DM noise. + +# %% +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 + 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, +) + +# %% +# Create the DMWaveX version of the par file. +m1 = deepcopy(m) +m1.remove_component("PLDMNoise") + +Tspan = t.get_mjds().max() - t.get_mjds().min() +dmwavex_setup(m1, Tspan, n_freqs=45) + +for p in m1.params: + if p.startswith("DMWXSIN") or p.startswith("DMWXCOS"): + m1[p].frozen = False + +# %% [markdown] +# ### Initial fitting + +# %% +ftr = WLSFitter(t, m1) + +ftr.fit_toas(maxiter=15) + +m1 = ftr.model +print(m1) + +# %% [markdown] +# ### Optimal number of harmonics +# The optimal number of harmonics can be estimated +# using the Akaike Information Criterion (AIC). + +# %% +m2 = deepcopy(m1) + +aics = [] +idxs = m2.components["DMWaveX"].get_indices() + +ftr = WLSFitter(t, m2) +ftr.fit_toas(maxiter=3) +aic = akaike_information_criterion(ftr.model, t) +aics += [aic] +print(f"{len(idxs)}\t{aic}\t{ftr.resids.chi2_reduced}") + +for idx in reversed(idxs): + if idx == 1: + m2.remove_component("DMWaveX") + else: + m2.components["DMWaveX"].remove_dmwavex_component(idx) + + ftr = WLSFitter(t, m2) + ftr.fit_toas(maxiter=3) + aic = akaike_information_criterion(ftr.model, t) + aics += [aic] + print(f"{idx-1}\t{aic}\t{ftr.resids.chi2_reduced}") + +# %% +# Find the optimum number of harmonics by minimizing AIC. +d_aics = np.array(aics) - np.min(aics) +nharm_opt = len(d_aics) - 1 - np.argmin(d_aics) +print("Optimum no of harmonics = ", nharm_opt) + +# %% +# The Y axis is plotted in log scale only for better visibility. +plt.scatter(list(reversed(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) + +idxs = m2.components["DMWaveX"].get_indices() +for idx in reversed(idxs): + if idx > nharm_opt: + m2.components["DMWaveX"].remove_dmwavex_component(idx) + +ftr = WLSFitter(t, m2) +ftr.fit_toas(maxiter=5) +m2 = ftr.model + +print(m2) + +# %% [markdown] +# ### Estimate 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) + +# %% From 94254a78d372b518d77f9c1a35ace7d387e2eb6f Mon Sep 17 00:00:00 2001 From: David Kaplan Date: Sun, 17 Dec 2023 15:05:14 -0600 Subject: [PATCH 06/36] fix bug --- docs/tutorials.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 107c8e1e2..99e876477 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -62,6 +62,7 @@ are not included in the default build because they take too long, but you can do examples/build_model_from_scratch.ipynb examples/How_to_build_a_timing_model_component.ipynb examples/understanding_fitters.ipynb + examples/rednoise-fit-example.ipynb examples/WorkingWithFlags.ipynb examples/Wideband_TOA_walkthrough.ipynb examples/Simulate_and_make_MassMass.ipynb From b9339653dc4f0b9a33abe6a5d8940af87d994ddf Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 6 Dec 2023 15:26:13 -0600 Subject: [PATCH 07/36] relax test --- tests/test_ddgr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_ddgr.py b/tests/test_ddgr.py index 8925ece05..a0ec92ae4 100644 --- a/tests/test_ddgr.py +++ b/tests/test_ddgr.py @@ -155,7 +155,7 @@ def test_ddgrfit_noMTOT(self): # chi^2 values don't have to be super close assert ( np.fabs(fDD.model.M2.quantity - fDDGR.model.M2.quantity) - < 4 * fDD.model.M2.uncertainty + < 5 * fDD.model.M2.uncertainty ) # perturn M2 and make sure chi^2 gets worse fDDGR.model.M2.quantity += 3 * fDDGR.model.M2.uncertainty From 51d5ca6f9ac31bd10e03661e20de2ea7eceed766 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 6 Dec 2023 15:26:47 -0600 Subject: [PATCH 08/36] -- --- tests/test_ddgr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_ddgr.py b/tests/test_ddgr.py index a0ec92ae4..8925ece05 100644 --- a/tests/test_ddgr.py +++ b/tests/test_ddgr.py @@ -155,7 +155,7 @@ def test_ddgrfit_noMTOT(self): # chi^2 values don't have to be super close assert ( np.fabs(fDD.model.M2.quantity - fDDGR.model.M2.quantity) - < 5 * fDD.model.M2.uncertainty + < 4 * fDD.model.M2.uncertainty ) # perturn M2 and make sure chi^2 gets worse fDDGR.model.M2.quantity += 3 * fDDGR.model.M2.uncertainty From 271ebc1c977b2d6bb6403206f5ef97dcd9d12376 Mon Sep 17 00:00:00 2001 From: David Kaplan Date: Sun, 10 Dec 2023 16:03:28 -0600 Subject: [PATCH 09/36] moved get_derived_params --- src/pint/fitter.py | 18 +----------------- src/pint/models/timing_model.py | 27 ++++++--------------------- 2 files changed, 7 insertions(+), 38 deletions(-) diff --git a/src/pint/fitter.py b/src/pint/fitter.py index eef71776c..18cb22f36 100644 --- a/src/pint/fitter.py +++ b/src/pint/fitter.py @@ -470,26 +470,10 @@ def get_summary(self, nodmx=False): par.units, ) s += "\n" + self.model.get_derived_params() + s += "\n" + self.model.get_derived_params() return s def get_derived_params(self, returndict=False): - """Return a string with various derived parameters from the fitted model - - Parameters - ---------- - returndict : bool, optional - Whether to only return the string of results or also a dictionary - - Returns - ------- - results : str - parameters : dict, optional - - See Also - -------- - :func:`pint.models.timing_model.TimingModel.get_derived_params` - """ - return self.model.get_derived_params( rms=self.resids.toa.rms_weighted() if self.is_wideband diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 79876de27..fcae9425b 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -32,6 +32,7 @@ import copy import inspect import contextlib +import contextlib from collections import OrderedDict, defaultdict from functools import wraps from warnings import warn @@ -39,6 +40,7 @@ import astropy.time as time from astropy import units as u, constants as c +from astropy import units as u, constants as c import numpy as np from astropy.utils.decorators import lazyproperty import astropy.coordinates as coords @@ -2949,22 +2951,7 @@ def as_ICRS(self, epoch=None): return new_model def get_derived_params(self, rms=None, ntoas=None, returndict=False): - """Return a string with various derived parameters from the fitted model - - Parameters - ---------- - rms : astropy.units.Quantity, optional - RMS of fit for checking ELL1 validity - ntoas : int, optional - Number of TOAs for checking ELL1 validity - returndict : bool, optional - Whether to only return the string of results or also a dictionary - - Returns - ------- - results : str - parameters : dict, optional - """ + """Return a string with various derived parameters from the fitted model""" import uncertainties.umath as um from uncertainties import ufloat @@ -3009,7 +2996,7 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False): s += "\n" px = self.PX.as_ufloat(u.arcsec) s += f"Parallax distance = {1.0/px:.3uP} pc\n" - outdict["Dist (pc)"] = 1.0 / px + outdict["Dist (pc):P"] = 1.0 / px # Now binary system derived parameters if self.is_binary: for x in self.components: @@ -3091,7 +3078,7 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False): # This is the mass function, done explicitly so that we get # uncertainty propagation automatically. # TODO: derived quantities funcs should take uncertainties - fm = 4.0 * np.pi**2 * a1**3 / (4.925490947e-6 * (pb * 86400) ** 2) + fm = 4.0 * np.pi**2 * a1**3 / (4.925490947e-6 * pb**2) s += f"Mass function = {fm:SP} Msun\n" outdict["Mass Function (Msun)"] = fm mcmed = pint.derived_quantities.companion_mass( @@ -3152,9 +3139,7 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False): ) s += f"Pulsar mass (Shapiro Delay) = {psrmass}" outdict["Mp (Msun)"] = psrmass - if not returndict: - return s - return s, outdict + return s if not returndict else s, outdict class ModelMeta(abc.ABCMeta): From 857164e669fbd688a501c1f0c4a408c53517b4a2 Mon Sep 17 00:00:00 2001 From: David Kaplan Date: Sun, 10 Dec 2023 16:23:14 -0600 Subject: [PATCH 10/36] added docstrings --- src/pint/fitter.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/pint/fitter.py b/src/pint/fitter.py index 18cb22f36..0a042b26e 100644 --- a/src/pint/fitter.py +++ b/src/pint/fitter.py @@ -474,6 +474,23 @@ def get_summary(self, nodmx=False): return s def get_derived_params(self, returndict=False): + """Return a string with various derived parameters from the fitted model + + Parameters + ---------- + returndict : bool, optional + Whether to only return the string of results or also a dictionary + + Returns + ------- + results : str + parameters : dict, optional + + See Also + -------- + :func:`pint.models.timing_model.TimingModel.get_derived_params` + """ + return self.model.get_derived_params( rms=self.resids.toa.rms_weighted() if self.is_wideband From dcc6d243b74ae03eea27cfdca5a35a9f9965eeb6 Mon Sep 17 00:00:00 2001 From: David Kaplan Date: Sun, 10 Dec 2023 17:35:48 -0600 Subject: [PATCH 11/36] fix bug --- src/pint/models/timing_model.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index fcae9425b..a95cb4600 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -2951,7 +2951,22 @@ def as_ICRS(self, epoch=None): return new_model def get_derived_params(self, rms=None, ntoas=None, returndict=False): - """Return a string with various derived parameters from the fitted model""" + """Return a string with various derived parameters from the fitted model + + Parameters + ---------- + rms : astropy.units.Quantity, optional + RMS of fit for checking ELL1 validity + ntoas : int, optional + Number of TOAs for checking ELL1 validity + returndict : bool, optional + Whether to only return the string of results or also a dictionary + + Returns + ------- + results : str + parameters : dict, optional + """ import uncertainties.umath as um from uncertainties import ufloat @@ -3139,7 +3154,9 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False): ) s += f"Pulsar mass (Shapiro Delay) = {psrmass}" outdict["Mp (Msun)"] = psrmass - return s if not returndict else s, outdict + if not returndict: + return s + return s, outdict class ModelMeta(abc.ABCMeta): From cccb700a4eaa912af128efe1558d9177b84a216e Mon Sep 17 00:00:00 2001 From: David Kaplan Date: Sun, 10 Dec 2023 19:07:49 -0600 Subject: [PATCH 12/36] fixed units --- src/pint/models/timing_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index a95cb4600..2f85d506f 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -3093,7 +3093,7 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False): # This is the mass function, done explicitly so that we get # uncertainty propagation automatically. # TODO: derived quantities funcs should take uncertainties - fm = 4.0 * np.pi**2 * a1**3 / (4.925490947e-6 * pb**2) + fm = 4.0 * np.pi**2 * a1**3 / (4.925490947e-6 * (pb * 86400) ** 2) s += f"Mass function = {fm:SP} Msun\n" outdict["Mass Function (Msun)"] = fm mcmed = pint.derived_quantities.companion_mass( From af954dae68ea239a5b383e7848c08e2490e2b410 Mon Sep 17 00:00:00 2001 From: David Kaplan Date: Wed, 13 Dec 2023 09:15:04 -0600 Subject: [PATCH 13/36] fixed format of output --- src/pint/models/timing_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 2f85d506f..6aa0e1921 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -3011,7 +3011,7 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False): s += "\n" px = self.PX.as_ufloat(u.arcsec) s += f"Parallax distance = {1.0/px:.3uP} pc\n" - outdict["Dist (pc):P"] = 1.0 / px + outdict["Dist (pc)"] = 1.0 / px # Now binary system derived parameters if self.is_binary: for x in self.components: From ed2bf0e29ca8ca54808c256aee8850f985df1530 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 2 Jan 2024 23:49:37 -0600 Subject: [PATCH 14/36] test_wx2pl --- src/pint/models/timing_model.py | 2 - tests/test_wx2pl.py | 77 +++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 2 deletions(-) create mode 100644 tests/test_wx2pl.py diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index d313f8434..d1bfdafce 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -32,7 +32,6 @@ import copy import inspect import contextlib -import contextlib from collections import OrderedDict, defaultdict from functools import wraps from warnings import warn @@ -40,7 +39,6 @@ import astropy.time as time from astropy import units as u, constants as c -from astropy import units as u, constants as c import numpy as np from astropy.utils.decorators import lazyproperty import astropy.coordinates as coords diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py new file mode 100644 index 000000000..7184b205c --- /dev/null +++ b/tests/test_wx2pl.py @@ -0,0 +1,77 @@ +from pint import DMconst +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, + 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 + + +def test_wx2pl(): + 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 5 + TZRMJD 55000 + TZRFRQ 1400 + TZRSITE gbt + UNITS TDB + EPHEM DE440 + CLOCK TT(BIPM2019) + """ + + m = get_model(StringIO(par_sim)) + + 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, + ) + + 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)) + + for p in m1.params: + if p.startswith("WXSIN") or p.startswith("WXCOS"): + m1[p].frozen = False + + ftr = WLSFitter(t, m1) + ftr.fit_toas(maxiter=3) + m1 = ftr.model + + m2 = plrednoise_from_wavex(m1) + + assert "PLRedNoise" in m2.components From 6ba316986853492ce2f44654611dd7ad9c65d808 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 2 Jan 2024 23:56:22 -0600 Subject: [PATCH 15/36] test_dmwx2pldm --- docs/examples/rednoise-fit-example.py | 2 +- tests/test_wx2pl.py | 62 +++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py index b768a16ff..5655133ad 100644 --- a/docs/examples/rednoise-fit-example.py +++ b/docs/examples/rednoise-fit-example.py @@ -252,7 +252,7 @@ # %% par_sim = """ - PSR SIM3 + PSR SIM4 RAJ 05:00:00 1 DECJ 15:00:00 1 PEPOCH 55000 diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py index 7184b205c..6f84a60bb 100644 --- a/tests/test_wx2pl.py +++ b/tests/test_wx2pl.py @@ -75,3 +75,65 @@ def test_wx2pl(): m2 = plrednoise_from_wavex(m1) assert "PLRedNoise" in m2.components + + +def test_dmwx2pldm(): + 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 + TNDMAMP -13 + TNDMGAM 3.5 + TNDMC 5 + TZRMJD 55000 + TZRFRQ 1400 + TZRSITE gbt + UNITS TDB + EPHEM DE440 + CLOCK TT(BIPM2019) + """ + + m = get_model(StringIO(par_sim)) + + 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, + ) + + 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)) + + for p in m1.params: + if p.startswith("DMWXSIN") or p.startswith("DMWXCOS"): + m1[p].frozen = False + + ftr = WLSFitter(t, m1) + ftr.fit_toas(maxiter=3) + m1 = ftr.model + + m2 = pldmnoise_from_dmwavex(m1) + + assert "PLDMNoise" in m2.components From 4df81c468dcb2a8a62f0592938b7ea884def3bfb Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 3 Jan 2024 10:39:08 -0600 Subject: [PATCH 16/36] freeze_params --- docs/examples/rednoise-fit-example.py | 2 +- src/pint/utils.py | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py index 5655133ad..16c075ce2 100644 --- a/docs/examples/rednoise-fit-example.py +++ b/docs/examples/rednoise-fit-example.py @@ -58,7 +58,7 @@ DM 15 1 TNREDAMP -13 TNREDGAM 3.5 - TNREDC 30 + TNREDC 20 TZRMJD 55000 TZRFRQ 1400 TZRSITE gbt diff --git a/src/pint/utils.py b/src/pint/utils.py index 5e9ab6631..c416d3752 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -1317,7 +1317,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 +1395,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 +1481,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() From 136a0d1d96058d1cb55a055ef12beb808250f86d Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 3 Jan 2024 10:41:24 -0600 Subject: [PATCH 17/36] CHANGELOG --- CHANGELOG-unreleased.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md index 6b26bde74..f37971671 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` and `pldmnoise_from_dmwavex` functions ### Fixed - `MCMC_walkthrough` notebook now runs - Fixed runtime data README From c33b49c8d2ca6fdbe83878a82bfb9c951c423f79 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 3 Jan 2024 11:01:24 -0600 Subject: [PATCH 18/36] rednoise-fit-example --- docs/examples/rednoise-fit-example.py | 62 +++++++++++++++++++++------ 1 file changed, 50 insertions(+), 12 deletions(-) diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py index 16c075ce2..8144520c8 100644 --- a/docs/examples/rednoise-fit-example.py +++ b/docs/examples/rednoise-fit-example.py @@ -7,7 +7,7 @@ # 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. @@ -58,7 +58,7 @@ DM 15 1 TNREDAMP -13 TNREDGAM 3.5 - TNREDC 20 + TNREDC 30 TZRMJD 55000 TZRFRQ 1400 TZRSITE gbt @@ -97,11 +97,7 @@ m1.remove_component("PLRedNoise") Tspan = t.get_mjds().max() - t.get_mjds().min() -wavex_setup(m1, Tspan, n_freqs=45) - -for p in m1.params: - if p.startswith("WXSIN") or p.startswith("WXCOS"): - m1[p].frozen = False +wavex_setup(m1, Tspan, n_freqs=30, freeze_params=False) # %% [markdown] # ### Initial fitting @@ -244,8 +240,12 @@ 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. @@ -301,11 +301,7 @@ m1.remove_component("PLDMNoise") Tspan = t.get_mjds().max() - t.get_mjds().min() -dmwavex_setup(m1, Tspan, n_freqs=45) - -for p in m1.params: - if p.startswith("DMWXSIN") or p.startswith("DMWXCOS"): - m1[p].frozen = False +dmwavex_setup(m1, Tspan, n_freqs=30, freeze_params=False) # %% [markdown] # ### Initial fitting @@ -421,3 +417,45 @@ 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() From f9db977c2eef01664129410e1037089ac545845d Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 3 Jan 2024 11:02:18 -0600 Subject: [PATCH 19/36] test_wx2pl --- tests/test_wx2pl.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py index 6f84a60bb..4201b72c2 100644 --- a/tests/test_wx2pl.py +++ b/tests/test_wx2pl.py @@ -62,11 +62,7 @@ def test_wx2pl(): m1.remove_component("PLRedNoise") Tspan = t.get_mjds().max() - t.get_mjds().min() - wavex_setup(m1, Tspan, n_freqs=int(m.TNREDC.value)) - - for p in m1.params: - if p.startswith("WXSIN") or p.startswith("WXCOS"): - m1[p].frozen = False + wavex_setup(m1, Tspan, n_freqs=int(m.TNREDC.value), freeze_params=False) ftr = WLSFitter(t, m1) ftr.fit_toas(maxiter=3) @@ -124,11 +120,7 @@ def test_dmwx2pldm(): m1.remove_component("PLDMNoise") Tspan = t.get_mjds().max() - t.get_mjds().min() - dmwavex_setup(m1, Tspan, n_freqs=int(m.TNDMC.value)) - - for p in m1.params: - if p.startswith("DMWXSIN") or p.startswith("DMWXCOS"): - m1[p].frozen = False + dmwavex_setup(m1, Tspan, n_freqs=int(m.TNDMC.value), freeze_params=False) ftr = WLSFitter(t, m1) ftr.fit_toas(maxiter=3) From d5a8c8953478c475600bc6cc949ede12ab376c87 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 3 Jan 2024 16:39:59 -0600 Subject: [PATCH 20/36] print aic --- src/pint/pintk/pulsar.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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.") From a1e2ab9e5525e575a99da64650ffed4a5de7e428 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 3 Jan 2024 16:40:07 -0600 Subject: [PATCH 21/36] compute_noise_uncertainties --- src/pint/fitter.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pint/fitter.py b/src/pint/fitter.py index 07dfb0ad7..23848dba7 100644 --- a/src/pint/fitter.py +++ b/src/pint/fitter.py @@ -1112,6 +1112,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. @@ -1183,7 +1184,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 ) From 25a8724a51ac25c10de85f3551668b9dfd8d6dbf Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 12 Jan 2024 10:33:52 -0600 Subject: [PATCH 22/36] refactor repeated code --- src/pint/utils.py | 156 ++++++++++++++++++++++------------------------ 1 file changed, 75 insertions(+), 81 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index c416d3752..c413ed94b 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2776,60 +2776,71 @@ def woodbury_dot(Ndiag, U, Phidiag, x, y): return x_Cinv_y, logdet_C -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. +def _get_wx2pl_lnlike(model, component_name, ignore_fyr=True): + from pint.models.noise_model import powerlaw + from pint import DMconst - Returns - ------- - pint.models.timing_model.TimingModel - The timing model with a converted `PLRedNoise` component. - """ - from pint.models.noise_model import powerlaw, PLRedNoise + assert component_name in ["WaveX", "DMWaveX"] + prefix = "WX" if component_name == "WaveX" else "DMWX" - idxs = np.array(model.components["WaveX"].get_indices()) + idxs = np.array(model.components[component_name].get_indices()) - fs = np.array([model[f"WXFREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs]) + 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 - ), "`plrednoise_from_wavex` requires the WaveX frequencies to be uniformly spaced." + ), "[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"WXFREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] + [model[f"{prefix}FREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] ) f0 = np.min(fs) - a = np.array([model[f"WXSIN_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) - da = np.array([model[f"WXSIN_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) - b = np.array([model[f"WXCOS_{idx:04d}"].quantity.to_value(u.s) for idx in idxs]) - db = np.array([model[f"WXCOS_{idx:04d}"].uncertainty.to_value(u.s) for idx in idxs]) + 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`.""" + 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 - `WaveX` amplitudes.""" + `[DM]WaveX` amplitudes.""" sigma = powl_model(params) return 0.5 * np.sum( (a**2 / (sigma**2 + da**2)) @@ -2838,6 +2849,32 @@ def mlnlike(params): + 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) + gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x hess = Hessian(mlnlike) @@ -2845,19 +2882,21 @@ def mlnlike(params): 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 = len(idxs) + 1 if ignore_fyr else len(idxs) + model1.TNREDC.value = tnredc model1.TNREDAMP.uncertainty_value = log10_A_err model1.TNREDGAM.uncertainty_value = gamma_err return model1 -def pldmnoise_from_dmwavex(model): +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 @@ -2873,56 +2912,9 @@ def pldmnoise_from_dmwavex(model): pint.models.timing_model.TimingModel The timing model with a converted `PLDMNoise` component. """ - from pint.models.noise_model import powerlaw, PLDMNoise - from pint import DMconst - - scale = DMconst / (1400 * u.MHz) ** 2 - - idxs = np.array(model.components["DMWaveX"].get_indices()) - a = np.array( - [(scale * model[f"DMWXSIN_{idx:04d}"].quantity).to_value(u.s) for idx in idxs] - ) - da = np.array( - [ - (scale * model[f"DMWXSIN_{idx:04d}"].uncertainty).to_value(u.s) - for idx in idxs - ] - ) - b = np.array( - [(scale * model[f"DMWXCOS_{idx:04d}"].quantity).to_value(u.s) for idx in idxs] - ) - db = np.array( - [ - (scale * model[f"DMWXCOS_{idx:04d}"].uncertainty).to_value(u.s) - for idx in idxs - ] - ) - - fs = np.array( - [model[f"DMWXFREQ_{idx:04d}"].quantity.to_value(u.Hz) for idx in idxs] - ) - f0 = np.min(fs) - - assert np.allclose( - np.diff(np.diff(fs)), 0 - ), "`pldmnoise_from_dmwavex` requires the DMWaveX frequencies to be uniformly spaced." - - def powl_model(params): - """Get the powerlaw spectrum for the DMWaveX frequencies for a given - set of parameters. This calls the powerlaw function used by `PLDMNoise`.""" - gamma, log10_A = params - return (powerlaw(fs, A=10**log10_A, gamma=gamma) * f0) ** 0.5 + from pint.models.noise_model import PLDMNoise - def mlnlike(params): - """Negative of the likelihood function that acts on the - `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) - ) + mlnlike = _get_wx2pl_lnlike(model, "DMWaveX", ignore_fyr=ignore_fyr) gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x @@ -2931,12 +2923,14 @@ def mlnlike(params): np.diag(np.linalg.pinv(hess((gamma_val, log10_A_val)))) ) + 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 = len(idxs) + model1.TNDMC.value = tndmc model1.TNDMAMP.uncertainty_value = log10_A_err model1.TNDMGAM.uncertainty_value = gamma_err From 6e48c3868a1ecc2418c3ad76ea3a0dddb2c920b0 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 12 Jan 2024 10:41:06 -0600 Subject: [PATCH 23/36] positive definite H --- src/pint/utils.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index c413ed94b..10292852a 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2919,9 +2919,16 @@ def pldmnoise_from_dmwavex(model, ignore_fyr=False): gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x hess = Hessian(mlnlike) - gamma_err, log10_A_err = np.sqrt( - np.diag(np.linalg.pinv(hess((gamma_val, log10_A_val)))) - ) + + 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()) From f88812ced7b8b7005af2dc182353ea121020e29c Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 11:31:21 -0600 Subject: [PATCH 24/36] find_optimal_nharms --- src/pint/utils.py | 56 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/src/pint/utils.py b/src/pint/utils.py index 10292852a..6f26b2bf5 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2942,3 +2942,59 @@ def pldmnoise_from_dmwavex(model, ignore_fyr=False): 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 + model1.components[component].add_wavex_component( + (len(model1.components[component].get_indices()) + 1) / T_span, frozen=False + ) + + return np.argmin(aics), np.array(aics) - np.min(aics) From 75d7b80fefe1f3fb7cbfb4abe00e2dd208ae2e4d Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 11:44:17 -0600 Subject: [PATCH 25/36] test_find_optimal_nharms --- src/pint/utils.py | 13 +++- tests/test_wx2pl.py | 157 +++++++++++++++++++++++++++----------------- 2 files changed, 108 insertions(+), 62 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 6f26b2bf5..01d42fe27 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2993,8 +2993,15 @@ def find_optimal_nharms(model, toas, component, nharms_max=45): aics.append(akaike_information_criterion(ftr.model, toas)) model1 = ftr.model - model1.components[component].add_wavex_component( - (len(model1.components[component].get_indices()) + 1) / T_span, frozen=False - ) + 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, + ) return np.argmin(aics), np.array(aics) - np.min(aics) diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py index 4201b72c2..943ee9a56 100644 --- a/tests/test_wx2pl.py +++ b/tests/test_wx2pl.py @@ -1,9 +1,11 @@ +import pytest from pint import DMconst 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, @@ -15,28 +17,28 @@ from copy import deepcopy -def test_wx2pl(): - 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 5 - TZRMJD 55000 - TZRFRQ 1400 - TZRSITE gbt - UNITS TDB - EPHEM DE440 - CLOCK TT(BIPM2019) - """ - - m = get_model(StringIO(par_sim)) +@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 -13 + TNREDGAM 3.5 + TNREDC 5 + 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 @@ -58,43 +60,32 @@ def test_wx2pl(): multi_freqs_in_epoch=True, ) - 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=3) - m1 = ftr.model - - m2 = plrednoise_from_wavex(m1) - - assert "PLRedNoise" in m2.components - - -def test_dmwx2pldm(): - 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 - TNDMAMP -13 - TNDMGAM 3.5 - TNDMC 5 - TZRMJD 55000 - TZRFRQ 1400 - TZRSITE gbt - UNITS TDB - EPHEM DE440 - CLOCK TT(BIPM2019) - """ - - m = get_model(StringIO(par_sim)) + 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 5 + 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 @@ -116,6 +107,30 @@ def test_dmwx2pldm(): 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=3) + m1 = ftr.model + + m2 = plrednoise_from_wavex(m1) + + assert "PLRedNoise" in m2.components + + +def test_dmwx2pldm(data_dmwx): + m, t = data_dmwx + m1 = deepcopy(m) m1.remove_component("PLDMNoise") @@ -129,3 +144,27 @@ def test_dmwx2pldm(): m2 = pldmnoise_from_dmwavex(m1) assert "PLDMNoise" in m2.components + + +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) From e8c55d377a3220ccf25488abcc676557c6f4e60f Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 12:05:59 -0600 Subject: [PATCH 26/36] CHANGELOG --- CHANGELOG-unreleased.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md index f37971671..e20f497b4 100644 --- a/CHANGELOG-unreleased.md +++ b/CHANGELOG-unreleased.md @@ -31,7 +31,7 @@ the released changes. - `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` and `pldmnoise_from_dmwavex` functions +- `plrednoise_from_wavex`, `pldmnoise_from_dmwavex`, and `find_optimal_nharms` functions ### Fixed - `MCMC_walkthrough` notebook now runs - Fixed runtime data README From 28fdfaa44c3b1c700db6e9c0b7f414a3ac8109f0 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 12:28:52 -0600 Subject: [PATCH 27/36] rednoise-fit-example --- docs/examples/rednoise-fit-example.py | 130 +++++--------------------- 1 file changed, 22 insertions(+), 108 deletions(-) diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py index 8144520c8..988717ece 100644 --- a/docs/examples/rednoise-fit-example.py +++ b/docs/examples/rednoise-fit-example.py @@ -24,6 +24,7 @@ from pint.utils import ( akaike_information_criterion, dmwavex_setup, + find_optimal_nharms, wavex_setup, plrednoise_from_wavex, pldmnoise_from_dmwavex, @@ -91,63 +92,26 @@ multi_freqs_in_epoch=True, ) -# %% -# We also need the WaveX version of the par file. -m1 = deepcopy(m) -m1.remove_component("PLRedNoise") - -Tspan = t.get_mjds().max() - t.get_mjds().min() -wavex_setup(m1, Tspan, n_freqs=30, freeze_params=False) - -# %% [markdown] -# ### Initial fitting - -# %% -ftr = WLSFitter(t, m1) - -ftr.fit_toas(maxiter=15) - -m1 = ftr.model -print(m1) - # %% [markdown] # ### Optimal number of harmonics -# The optimal number of harmonics can be estimated -# using the Akaike Information Criterion (AIC). +# 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. # %% -m2 = deepcopy(m1) +m1 = deepcopy(m) +m1.remove_component("PLRedNoise") -aics = [] -idxs = m2.components["WaveX"].get_indices() +nharm_opt, d_aics = find_optimal_nharms(m1, t, "WaveX", 30) -ftr = WLSFitter(t, m2) -ftr.fit_toas(maxiter=3) -aic = akaike_information_criterion(ftr.model, t) -aics += [aic] -print(f"{len(idxs)}\t{aic}\t{ftr.resids.chi2_reduced}") - -for idx in reversed(idxs): - if idx == 1: - m2.remove_component("WaveX") - else: - m2.components["WaveX"].remove_wavex_component(idx) - - ftr = WLSFitter(t, m2) - ftr.fit_toas(maxiter=3) - aic = akaike_information_criterion(ftr.model, t) - aics += [aic] - print(f"{idx-1}\t{aic}\t{ftr.resids.chi2_reduced}") +print("Optimum no of harmonics = ", nharm_opt) # %% -# Find the optimum number of harmonics by minimizing AIC. -d_aics = np.array(aics) - np.min(aics) -nharm_opt = len(d_aics) - 1 - np.argmin(d_aics) -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(reversed(range(len(d_aics)))), d_aics + 1) +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" @@ -159,23 +123,19 @@ # plt.savefig("sim3-aic.pdf") # %% -# Now create a new model with the optimum number of -# harmonics +# Now create a new model with the optimum number of harmonics m2 = deepcopy(m1) - -idxs = m2.components["WaveX"].get_indices() -for idx in reversed(idxs): - if idx > nharm_opt: - m2.components["WaveX"].remove_wavex_component(idx) +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=5) +ftr.fit_toas(maxiter=10) m2 = ftr.model print(m2) # %% [markdown] -# ### Estimate the spectral parameters from the WaveX fit. +# ### Estimating the spectral parameters from the WaveX fit. # %% # Get the Fourier amplitudes and powers and their uncertainties. @@ -296,62 +256,18 @@ ) # %% -# Create the DMWaveX version of the par file. +# Find the optimum number of harmonics by minimizing AIC. m1 = deepcopy(m) m1.remove_component("PLDMNoise") -Tspan = t.get_mjds().max() - t.get_mjds().min() -dmwavex_setup(m1, Tspan, n_freqs=30, freeze_params=False) - -# %% [markdown] -# ### Initial fitting - -# %% -ftr = WLSFitter(t, m1) - -ftr.fit_toas(maxiter=15) - -m1 = ftr.model -print(m1) - -# %% [markdown] -# ### Optimal number of harmonics -# The optimal number of harmonics can be estimated -# using the Akaike Information Criterion (AIC). - -# %% m2 = deepcopy(m1) -aics = [] -idxs = m2.components["DMWaveX"].get_indices() - -ftr = WLSFitter(t, m2) -ftr.fit_toas(maxiter=3) -aic = akaike_information_criterion(ftr.model, t) -aics += [aic] -print(f"{len(idxs)}\t{aic}\t{ftr.resids.chi2_reduced}") - -for idx in reversed(idxs): - if idx == 1: - m2.remove_component("DMWaveX") - else: - m2.components["DMWaveX"].remove_dmwavex_component(idx) - - ftr = WLSFitter(t, m2) - ftr.fit_toas(maxiter=3) - aic = akaike_information_criterion(ftr.model, t) - aics += [aic] - print(f"{idx-1}\t{aic}\t{ftr.resids.chi2_reduced}") - -# %% -# Find the optimum number of harmonics by minimizing AIC. -d_aics = np.array(aics) - np.min(aics) -nharm_opt = len(d_aics) - 1 - np.argmin(d_aics) +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(reversed(range(len(d_aics)))), d_aics + 1) +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" @@ -367,19 +283,17 @@ # harmonics m2 = deepcopy(m1) -idxs = m2.components["DMWaveX"].get_indices() -for idx in reversed(idxs): - if idx > nharm_opt: - m2.components["DMWaveX"].remove_dmwavex_component(idx) +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=5) +ftr.fit_toas(maxiter=10) m2 = ftr.model print(m2) # %% [markdown] -# ### Estimate the spectral parameters from the `DMWaveX` fit. +# ### Estimating the spectral parameters from the `DMWaveX` fit. # %% # Get the Fourier amplitudes and powers and their uncertainties. From e14030dbe3e498cff250b11426f11d984dab0660 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 12:33:18 -0600 Subject: [PATCH 28/36] sourcery --- src/pint/utils.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 01d42fe27..82d1be9b0 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -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 = [] @@ -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() From 720282df0919693c07ac7bc498100f5cea2fc861 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 12:44:24 -0600 Subject: [PATCH 29/36] sourcery --- src/pint/utils.py | 181 +++++++++++++++++++--------------------------- 1 file changed, 76 insertions(+), 105 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 82d1be9b0..011f28d06 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -689,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() @@ -980,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)) @@ -1013,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 @@ -1046,25 +1046,24 @@ def dmxparse(fitter, save=False): 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( + lines = [ + "# Mean %s value = %+.6e \n" % (DMX, DMX_mean), + "# Uncertainty in average %s = %.5e \n" % ("DM", DMX_mean_err), "# Columns: %sEP %s_value %s_var_err %sR1 %sR2 %s_bin \n" - % (DMX, DMX, DMX, DMX, DMX, DMX) - ) - 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], - ) + % (DMX, DMX, DMX, DMX, DMX, DMX), + ] + lines.extend( + "%.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], ) + for k in range(len(dmx_epochs)) + ) with open_or_use(save, mode="w") as dmxout: dmxout.writelines(lines) if isinstance(save, (str, Path)): @@ -1076,18 +1075,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): @@ -1136,7 +1133,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) @@ -1534,15 +1531,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): @@ -1619,10 +1613,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)}" @@ -1660,30 +1654,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 ] @@ -1695,7 +1687,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 @@ -1782,10 +1774,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() @@ -1836,10 +1825,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" @@ -1894,20 +1885,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): @@ -1933,8 +1919,8 @@ def add_dummy_distance(c, distance=1 * u.kpc): 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, @@ -1943,11 +1929,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, @@ -1956,10 +1939,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, @@ -1968,9 +1950,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, @@ -1980,7 +1961,6 @@ 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 @@ -2008,8 +1988,8 @@ def remove_dummy_distance(c): ) 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, @@ -2017,11 +1997,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, @@ -2029,9 +2006,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, @@ -2039,9 +2016,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, @@ -2050,7 +2026,6 @@ 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 @@ -2222,10 +2197,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]) @@ -2516,8 +2488,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): From 00eac1edd3bf75403dccce85d09a3d0d06c56c29 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 12:45:53 -0600 Subject: [PATCH 30/36] -- --- src/pint/fitter.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pint/fitter.py b/src/pint/fitter.py index 23848dba7..676e97ae7 100644 --- a/src/pint/fitter.py +++ b/src/pint/fitter.py @@ -470,7 +470,6 @@ def get_summary(self, nodmx=False): par.units, ) s += "\n" + self.model.get_derived_params() - s += "\n" + self.model.get_derived_params() return s def get_derived_params(self, returndict=False): From 898a33f4c44baabe8abdd7303b88c71531583567 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Tue, 16 Jan 2024 12:47:30 -0600 Subject: [PATCH 31/36] cleanup --- docs/examples/rednoise-fit-example.py | 1 - tests/test_wx2pl.py | 1 - 2 files changed, 2 deletions(-) diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py index 988717ece..c8a93e9cb 100644 --- a/docs/examples/rednoise-fit-example.py +++ b/docs/examples/rednoise-fit-example.py @@ -22,7 +22,6 @@ from pint.logging import setup as setup_log from pint.fitter import WLSFitter from pint.utils import ( - akaike_information_criterion, dmwavex_setup, find_optimal_nharms, wavex_setup, diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py index 943ee9a56..88200bc0c 100644 --- a/tests/test_wx2pl.py +++ b/tests/test_wx2pl.py @@ -1,5 +1,4 @@ import pytest -from pint import DMconst from pint.models import get_model from pint.simulation import make_fake_toas_uniform from pint.fitter import WLSFitter From a677c82363d7425f4bbae81b8b1b659478a55b06 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Wed, 17 Jan 2024 11:06:50 -0600 Subject: [PATCH 32/36] remove warn --- src/pint/fitter.py | 1 - src/pint/models/dmwavex.py | 21 ++++++++++----------- src/pint/models/wavex.py | 4 ++-- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/src/pint/fitter.py b/src/pint/fitter.py index 676e97ae7..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, 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 bf5bae2af..a5e9acbbb 100644 --- a/src/pint/models/wavex.py +++ b/src/pint/models/wavex.py @@ -341,8 +341,8 @@ 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 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( From c39e00fe450452a853dda1eeb200a69390601227 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 19 Jan 2024 09:23:52 -0600 Subject: [PATCH 33/36] more quantitative tests --- tests/test_wx2pl.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py index 88200bc0c..c10ba73dc 100644 --- a/tests/test_wx2pl.py +++ b/tests/test_wx2pl.py @@ -27,9 +27,9 @@ def data_wx(): F1 -1e-15 1 PHOFF 0 1 DM 15 1 - TNREDAMP -13 + TNREDAMP -12.5 TNREDGAM 3.5 - TNREDC 5 + TNREDC 10 TZRMJD 55000 TZRFRQ 1400 TZRSITE gbt @@ -75,7 +75,7 @@ def data_dmwx(): DM 15 1 TNDMAMP -13 TNDMGAM 3.5 - TNDMC 5 + TNDMC 10 TZRMJD 55000 TZRFRQ 1400 TZRSITE gbt @@ -119,12 +119,14 @@ def test_wx2pl(data_wx): wavex_setup(m1, Tspan, n_freqs=int(m.TNREDC.value), freeze_params=False) ftr = WLSFitter(t, m1) - ftr.fit_toas(maxiter=3) + 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): @@ -137,12 +139,14 @@ def test_dmwx2pldm(data_dmwx): dmwavex_setup(m1, Tspan, n_freqs=int(m.TNDMC.value), freeze_params=False) ftr = WLSFitter(t, m1) - ftr.fit_toas(maxiter=3) + 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): From 6907455e09be150c8d9b193d251d0b6de8491504 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 19 Jan 2024 13:35:30 -0600 Subject: [PATCH 34/36] fstring --- src/pint/utils.py | 30 ++++++++++-------------------- 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 011f28d06..7fde56829 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -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): @@ -1045,23 +1045,13 @@ def dmxparse(fitter, save=False): if save is not None and save: if isinstance(save, bool): save = "dmxparse.out" - DMX = "DMX" lines = [ - "# Mean %s value = %+.6e \n" % (DMX, DMX_mean), - "# Uncertainty in average %s = %.5e \n" % ("DM", DMX_mean_err), - "# Columns: %sEP %s_value %s_var_err %sR1 %sR2 %s_bin \n" - % (DMX, DMX, DMX, DMX, DMX, DMX), + 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( - "%.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], - ) + 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)) ) with open_or_use(save, mode="w") as dmxout: @@ -1914,7 +1904,7 @@ 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 @@ -1963,7 +1953,7 @@ def add_dummy_distance(c, distance=1 * u.kpc): ) 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 @@ -1984,7 +1974,7 @@ 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): @@ -2028,7 +2018,7 @@ def remove_dummy_distance(c): ) 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 From fdaf29224e72a9c88bec08d229b4c2debdfee2ba Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 19 Jan 2024 13:38:49 -0600 Subject: [PATCH 35/36] check convergence --- src/pint/utils.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index 7fde56829..6d565c1bd 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2832,7 +2832,11 @@ def plrednoise_from_wavex(model, ignore_fyr=True): mlnlike = _get_wx2pl_lnlike(model, "WaveX", ignore_fyr=ignore_fyr) - gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x + 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( @@ -2873,7 +2877,11 @@ def pldmnoise_from_dmwavex(model, ignore_fyr=False): mlnlike = _get_wx2pl_lnlike(model, "DMWaveX", ignore_fyr=ignore_fyr) - gamma_val, log10_A_val = minimize(mlnlike, [4, -13], method="Nelder-Mead").x + 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) From bf7df7cfc09ffc1cee8f7fe2762e2d75a0207a4c Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 19 Jan 2024 13:41:58 -0600 Subject: [PATCH 36/36] nan check --- src/pint/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pint/utils.py b/src/pint/utils.py index 6d565c1bd..1c0bcf21c 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -2969,4 +2969,6 @@ def find_optimal_nharms(model, toas, component, nharms_max=45): frozen=False, ) + assert all(np.isfinite(aics)), "Infs/NaNs found in AICs!" + return np.argmin(aics), np.array(aics) - np.min(aics)