diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md index e5ca814a7..0056fb6c9 100644 --- a/CHANGELOG-unreleased.md +++ b/CHANGELOG-unreleased.md @@ -12,6 +12,8 @@ the released changes. - Command line scripts now automatically do `allow_tcb` and `allow_T2` while reading par files. - Updated the `plot_chains` function in `event_optimize` so that the subplots are a fixed size to prevent the subplots from being condensed in the case of many fit parameters. ### Added +- Time derivatives of NE_SW in `SolarWindDispersion` +- New prefix pattern for `split_prefixed_name` to handle derivatives of NE_SW - Added an option `nbin` to `photonphase` to decide how many phase bins to use for the phaseogram - Added an option `linearize_model` to speed up the photon phases calculation within `event_optimize` through the designmatrix. - Added AIC and BIC calculation to be written in the post fit parfile from `event_optimize` diff --git a/src/pint/models/dispersion_model.py b/src/pint/models/dispersion_model.py index 368392dd9..5dd7c11cf 100644 --- a/src/pint/models/dispersion_model.py +++ b/src/pint/models/dispersion_model.py @@ -21,7 +21,7 @@ taylor_horner, taylor_horner_deriv, ) -from pint import DMconst +from pint import DMconst, dmu from pint.exceptions import MissingParameter, MissingTOAs # This value is cited from Duncan Lorimer, Michael Kramer, Handbook of Pulsar @@ -76,7 +76,7 @@ def dm_value(self, toas): DM values at given TOAs in the unit of DM. """ toas_table = toas if isinstance(toas, Table) else toas.table - dm = np.zeros(len(toas_table)) * self._parent.DM.units + dm = np.zeros(len(toas_table)) * dmu for dm_f in self.dm_value_funcs: dm += dm_f(toas) @@ -183,8 +183,8 @@ def __init__(self): def setup(self): super().setup() - base_dms = list(self.get_prefix_mapping_component("DM").values()) - base_dms += ["DM"] + + base_dms = ["DM"] + list(self.get_prefix_mapping_component("DM").values()) for dm_name in base_dms: self.register_deriv_funcs(self.d_delay_d_dmparam, dm_name) @@ -206,10 +206,10 @@ def validate(self): ) def DM_derivative_unit(self, n): - return "pc cm^-3/yr^%d" % n if n else "pc cm^-3" + return f"pc cm^-3/yr^{n}" if n != 0 else "pc cm^-3" def DM_derivative_description(self, n): - return "%d'th time derivative of the dispersion measure" % n + return f"{n}'th time derivative of the dispersion measure" def get_DM_terms(self): """Return a list of the DM term values in the model: [DM, DM1, ..., DMn]""" @@ -223,7 +223,7 @@ def base_dm(self, toas): if DMEPOCH is None: # Should be ruled out by validate() raise ValueError( - f"DMEPOCH not set but some derivatives are not zero: {dm_terms}" + f"DMEPOCH not set but some DM derivatives are not zero: {dm_terms}" ) else: dt = (toas["tdbld"] - DMEPOCH) * u.day @@ -673,7 +673,7 @@ def dmx_dm(self, toas): condition, tbl["mjd_float"] ) # Get DMX delays - dm = np.zeros(len(tbl)) * self._parent.DM.units + dm = np.zeros(len(tbl)) * dmu for k, v in select_idx.items(): dm[v] += getattr(self, k).quantity return dm diff --git a/src/pint/models/solar_wind_dispersion.py b/src/pint/models/solar_wind_dispersion.py index dddaeb0f2..8a282fee3 100644 --- a/src/pint/models/solar_wind_dispersion.py +++ b/src/pint/models/solar_wind_dispersion.py @@ -10,7 +10,12 @@ from pint import DMconst from pint.models.dispersion_model import Dispersion -from pint.models.parameter import floatParameter, intParameter, prefixParameter +from pint.models.parameter import ( + MJDParameter, + floatParameter, + intParameter, + prefixParameter, +) import pint.utils from pint.models.timing_model import MissingTOAs from pint.toa_select import TOASelect @@ -302,6 +307,26 @@ def __init__(self): tcb2tdb_scale_factor=(const.c * DMconst), ) ) + self.add_param( + prefixParameter( + name="NE_SW1", + units="cm^-3 / yr", + value=0.0, + description="Derivative of Solar Wind density at 1 AU", + unit_template=self.NE_SW_derivative_unit, + description_template=self.NE_SW_derivative_description, + type_match="float", + tcb2tdb_scale_factor=(const.c * DMconst), + ) + ) + self.add_param( + MJDParameter( + name="SWEPOCH", + description="Epoch of NE_SW measurement", + time_scale="tdb", + tcb2tdb_scale_factor=u.Quantity(1), + ) + ) self.add_param( floatParameter( name="SWP", @@ -324,11 +349,24 @@ def __init__(self): def setup(self): super().setup() - self.register_dm_deriv_funcs(self.d_dm_d_ne_sw, "NE_SW") - self.register_deriv_funcs(self.d_delay_d_ne_sw, "NE_SW") + self.register_dm_deriv_funcs(self.d_dm_d_swp, "SWP") self.register_deriv_funcs(self.d_delay_d_swp, "SWP") + base_nesws = ["NE_SW"] + list( + self.get_prefix_mapping_component("NE_SW").values() + ) + + for nesw_name in base_nesws: + self.register_deriv_funcs(self.d_delay_d_ne_sw, nesw_name) + self.register_dm_deriv_funcs(self.d_dm_d_ne_sw, nesw_name) + + def NE_SW_derivative_unit(self, n): + return f"cm^-3/yr^{n}" if n != 0 else "cm^-3" + + def NE_SW_derivative_description(self, n): + return f"{n}'th time derivative of the Solar Wind density at 1 AU" + def solar_wind_geometry(self, toas): """Return the geometry of solar wind dispersion. @@ -375,22 +413,51 @@ def solar_wind_dm(self, toas): SWM==1: Hazboun et al. 2022 """ - if self.NE_SW.value == 0: + ne_sw_terms = self.get_NE_SW_terms() + + if len(ne_sw_terms) == 1: + ne_sw = self.NE_SW.quantity * np.ones(len(toas)) + else: + if any(t.value != 0 for t in ne_sw_terms[1:]): + SWEPOCH = self.SWEPOCH.value + if SWEPOCH is None: + # Should be ruled out by validate() + raise ValueError( + f"SWEPOCH not set but some NE_SW derivatives are not zero: {ne_sw_terms}" + ) + else: + dt = (toas["tdbld"] - SWEPOCH) * u.day + dt_value = dt.to_value(u.yr) + else: + dt_value = np.zeros(len(toas), dtype=np.longdouble) + + ne_sw_terms_value = [d.value for d in ne_sw_terms] + ne_sw = ( + pint.utils.taylor_horner(dt_value, ne_sw_terms_value) * self.NE_SW.units + ) + + if np.all(ne_sw.value == 0): return np.zeros(len(toas)) * u.pc / u.cm**3 + if self.SWM.value not in [0, 1]: raise NotImplementedError( f"Solar Dispersion Delay not implemented for SWM {self.SWM.value}" ) + solar_wind_geometry = self.solar_wind_geometry(toas) - solar_wind_dm = self.NE_SW.quantity * solar_wind_geometry + solar_wind_dm = ne_sw * solar_wind_geometry return solar_wind_dm.to(u.pc / u.cm**3) def solar_wind_delay(self, toas, acc_delay=None): """This is a wrapper function to compute solar wind dispersion delay.""" - if self.NE_SW.value == 0: - return np.zeros(len(toas)) * u.s return self.dispersion_type_delay(toas) + def get_NE_SW_terms(self): + """Return a list of the NE_SW term values in the model: [NE_SW, NE_SW1, ..., NE_SWn]""" + return [self.NE_SW.quantity] + self._parent.get_prefix_list( + "NE_SW", start_index=1 + ) + def d_dm_d_ne_sw(self, toas, param_name, acc_delay=None): """Derivative of of DM wrt the solar wind ne amplitude.""" if self.SWM.value in [0, 1]: @@ -399,17 +466,31 @@ def d_dm_d_ne_sw(self, toas, param_name, acc_delay=None): raise NotImplementedError( "Solar Dispersion Delay not implemented for SWM %d" % self.SWM.value ) - return solar_wind_geometry + + par = getattr(self, param_name) + order = 0 if param_name == "NE_SW" else par.index + + ne_sws = self.get_NE_SW_terms() + + ne_sw_terms = np.longdouble(np.zeros(len(ne_sws))) + ne_sw_terms[order] = np.longdouble(1.0) + + if self.SWEPOCH.value is None: + if any(t.value != 0 for t in ne_sws[1:]): + # Should be ruled out by validate() + raise ValueError(f"SWEPOCH is not set but {param_name} is not zero") + SWEPOCH = 0 + else: + SWEPOCH = self.SWEPOCH.value + + dt = (toas["tdbld"] - SWEPOCH) * u.day + + return (solar_wind_geometry * dt**order / scipy.special.factorial(order)).to( + u.pc / u.cm**3 / par.units + ) def d_delay_d_ne_sw(self, toas, param_name, acc_delay=None): - try: - bfreq = self._parent.barycentric_radio_freq(toas) - except AttributeError: - warn("Using topocentric frequency for dedispersion!") - bfreq = toas.table["freq"] - deriv = self.d_delay_d_dmparam(toas, "NE_SW") - deriv[bfreq < 1.0 * u.MHz] = 0.0 - return deriv + return self.d_delay_d_dmparam(toas, param_name) def d_solar_wind_geometry_d_swp(self, toas, param_name, acc_delay=None): """Derivative of solar_wind_geometry (path length) wrt power-law index p @@ -452,10 +533,12 @@ def d_delay_d_swp(self, toas, param_name, acc_delay=None): def print_par(self, format="pint"): result = "" - result += getattr(self, "NE_SW").as_parfile_line(format=format) - result += getattr(self, "SWM").as_parfile_line(format=format) - if self.SWM.value == 1: - result += getattr(self, "SWP").as_parfile_line(format=format) + for par in self.params: + if par == "SWP" and self.SWM.value != 1: + continue + else: + result += getattr(self, par).as_parfile_line(format=format) + return result def get_max_dm(self): diff --git a/src/pint/utils.py b/src/pint/utils.py index 59f0d3076..ff5fc5317 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -357,7 +357,7 @@ def has_astropy_unit(x: Any) -> bool: re.compile(r"^([a-zA-Z]*\d+[a-zA-Z]+)(\d+)$"), # For the prefix like T2EFAC2 re.compile(r"^([a-zA-Z]+)0*(\d+)$"), # For the prefix like F12 re.compile(r"^([a-zA-Z0-9]+_)(\d+)$"), # For the prefix like DMXR1_3 - # re.compile(r'([a-zA-Z]\d[a-zA-Z]+)(\d+)'), # for prefixes like PLANET_SHAPIRO2? + re.compile(r"([a-zA-Z]+_[a-zA-Z]+)(\d+)$"), # for prefixes like NE_SW2? ] diff --git a/tests/test_solar_wind.py b/tests/test_solar_wind.py index bfa8e0eea..69a79fd52 100644 --- a/tests/test_solar_wind.py +++ b/tests/test_solar_wind.py @@ -359,3 +359,83 @@ def test_overlapping_swx(): dm3 = model3.components["SolarWindDispersionX"].swx_dm(toas) dm3[toas.get_mjds() >= 54180 * u.d] = 0 assert np.allclose(dm2 - dm, dm3) + + +def test_nesw_derivatives(): + par = """ + PSRJ J1744-1134 + ELONG 266.119458498 6.000e-09 + ELAT 11.80517508 3.000e-08 + DM 3.1379 4.000e-04 + PEPOCH 55000 + F0 245.4261196898081 5.000e-13 + F1 -5.38156E-16 3.000e-21 + POSEPOCH 59150 + DMEPOCH 55000 + CLK TT(BIPM2018) + EPHEM DE436 + RM 1.62 9.000e-02 + PX 2.61 9.000e-02 + DM1 -7.0E-5 3.000e-05 + PMELONG 19.11 2.000e-02 + DM2 1.7E-5 4.000e-06 + PMELAT -9.21 1.300e-01 + UNITS TDB + NE_SW 8 1 + NE_SW1 0 1 + NE_SW2 0 1 + SWEPOCH 55000 + """ + m = get_model(StringIO(par)) + t = make_fake_toas_uniform(54500, 55500, 1000, m, add_noise=True) + ftr = Fitter.auto(t, m) + ftr.fit_toas() + + assert ( + m.NE_SW.value - ftr.model.NE_SW.value + ) / ftr.model.NE_SW.uncertainty_value < 3 + assert ( + m.NE_SW1.value - ftr.model.NE_SW1.value + ) / ftr.model.NE_SW1.uncertainty_value < 3 + assert ( + m.NE_SW2.value - ftr.model.NE_SW2.value + ) / ftr.model.NE_SW2.uncertainty_value < 3 + + m.components["SolarWindDispersion"] + + +def test_expression(): + par = """ + PSRJ J1744-1134 + ELONG 266.119458498 6.000e-09 + ELAT 11.80517508 3.000e-08 + DM 3.1379 4.000e-04 + PEPOCH 55000 + F0 245.4261196898081 5.000e-13 + F1 -5.38156E-16 3.000e-21 + POSEPOCH 59150 + DMEPOCH 55000 + CLK TT(BIPM2018) + EPHEM DE436 + PX 2.61 9.000e-02 + PMELONG 19.11 2.000e-02 + PMELAT -9.21 1.300e-01 + UNITS TDB + NE_SW 8 1 + NE_SW1 1 1 + SWEPOCH 55000 + """ + m = get_model(StringIO(par)) + t = make_fake_toas_uniform(54500, 55500, 10, m, add_noise=True) + + t0 = m.SWEPOCH.value * u.day + t1 = t["tdbld"][-1] * u.day + dt = t1 - t0 + + assert np.isclose( + ( + m.components["SolarWindDispersion"].solar_wind_dm(t) + / m.components["SolarWindDispersion"].solar_wind_geometry(t) + )[-1], + (m.NE_SW.quantity + dt * m.NE_SW1.quantity), + )