diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md index 5487b9482..8200caa75 100644 --- a/CHANGELOG-unreleased.md +++ b/CHANGELOG-unreleased.md @@ -12,6 +12,8 @@ the released changes. - `WAVE` parameters can be added to a `Wave` model with `add_wave_component()` in `wave.py` - Moved design matrix normalization code from `pint.fitter` to the new `pint.utils.normalize_designmatrix()` function. - Made `Residuals` independent of `GLSFitter` (GLS chi2 is now computed using the new function `Residuals._calc_gls_chi2()`). +- `ssb_to_psb_ICRS` implementation is now a lot faster (uses erfa functions directly) +- `ssb_to_psb_ECL` implementation within `AstrometryEcliptic` model is now a lot faster (uses erfa functions directly) - Upgraded versioneer for compatibility with Python 3.12 - Creation of `Fitter` objects will fail if there are free unfittable parameters in the timing model. - Only fittable parameters will be listed as check boxes in the `plk` interface. diff --git a/src/pint/models/astrometry.py b/src/pint/models/astrometry.py index 1d737849c..2d739a13d 100644 --- a/src/pint/models/astrometry.py +++ b/src/pint/models/astrometry.py @@ -11,7 +11,7 @@ from loguru import logger as log -from erfa import ErfaWarning +from erfa import ErfaWarning, pmsafe from pint import ls from pint.models.parameter import ( AngleParameter, @@ -23,7 +23,6 @@ from pint.pulsar_ecliptic import OBL, PulsarEcliptic from pint.utils import add_dummy_distance, remove_dummy_distance - astropy_version = sys.modules["astropy"].__version__ mas_yr = u.mas / u.yr @@ -61,8 +60,20 @@ def ssb_to_psb_xyz_ICRS(self, epoch=None): """Returns unit vector(s) from SSB to pulsar system barycenter under ICRS. If epochs (MJD) are given, proper motion is included in the calculation. + + Parameters + ---------- + epoch : float or astropy.time.Time, optional + + Returns + ------- + np.ndarray : + (len(epoch), 3) array of unit vectors """ # TODO: would it be better for this to return a 6-vector (pos, vel)? + + # this is somewhat slow, since it repeatedly created different SkyCoord Objects + # but for consistency only change the method in the subclasses below return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose() def ssb_to_psb_xyz_ECL(self, epoch=None, ecl=None): @@ -75,6 +86,11 @@ def ssb_to_psb_xyz_ECL(self, epoch=None, ecl=None): epoch : float, optional ecl : str, optional Obliquity (IERS2010 by default) + + Returns + ------- + np.ndarray : + (len(epoch), 3) array of unit vectors """ # TODO: would it be better for this to return a 6-vector (pos, vel)? return self.coords_as_ECL(epoch=epoch, ecl=ecl).cartesian.xyz.transpose() @@ -371,7 +387,7 @@ def coords_as_ECL(self, epoch=None, ecl=None): If `ecl` is left unspecified, the global default IERS2010 will be used. """ if ecl is None: - log.info("ECL not specified; using IERS2010.") + log.debug("ECL not specified; using IERS2010.") ecl = "IERS2010" pos_icrs = self.get_psr_coords(epoch=epoch) @@ -390,6 +406,63 @@ def get_params_as_ICRS(self): "PMDEC": self.PMDEC.quantity, } + def ssb_to_psb_xyz_ICRS(self, epoch=None): + """Returns unit vector(s) from SSB to pulsar system barycenter under ICRS. + + If epochs (MJD) are given, proper motion is included in the calculation. + + Parameters + ---------- + epoch : float or astropy.time.Time, optional + + Returns + ------- + np.ndarray : + (len(epoch), 3) array of unit vectors + """ + # TODO: would it be better for this to return a 6-vector (pos, vel)? + + # this was somewhat slow, since it repeatedly created different SkyCoord Objects + # return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose() + + # Instead look at what https://docs.astropy.org/en/stable/_modules/astropy/coordinates/sky_coordinate.html#SkyCoord.apply_space_motion + # does, which is to use https://github.com/liberfa/erfa/blob/master/src/starpm.c + # and then just use the relevant pieces of that + if epoch is None or (self.PMRA.quantity == 0 and self.PMDEC.quantity == 0): + return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose() + + if isinstance(epoch, Time): + jd1 = epoch.jd1 + jd2 = epoch.jd2 + else: + # assume MJD + jd1 = 2400000.5 + jd2 = epoch + # compared to the general case above we can assume that the coordinates are ICRS + # so just access those components + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ErfaWarning) + # note that starpm wants mu_alpha not mu_alpha * cos(delta) + starpmout = pmsafe( + self.RAJ.quantity.to_value(u.radian), + self.DECJ.quantity.to_value(u.radian), + self.PMRA.quantity.to_value(u.radian / u.yr) + / np.cos(self.DECJ.quantity).value, + self.PMDEC.quantity.to_value(u.radian / u.yr), + self.PX.quantity.to_value(u.arcsec), + 0.0, + self.POSEPOCH.quantity.jd1, + self.POSEPOCH.quantity.jd2, + jd1, + jd2, + ) + # ra,dec now in radians + ra, dec = starpmout[0], starpmout[1] + x = np.cos(ra) * np.cos(dec) + y = np.sin(ra) * np.cos(dec) + z = np.sin(dec) + return u.Quantity([x, y, z]).T + def d_delay_astrometry_d_RAJ(self, toas, param="", acc_delay=None): """Calculate the derivative wrt RAJ @@ -754,6 +827,78 @@ def coords_as_ECL(self, epoch=None, ecl=None): pos_ecl = pos_ecl.transform_to(PulsarEcliptic(ecl=ecl)) return pos_ecl + def ssb_to_psb_xyz_ECL(self, epoch=None, ecl=None): + """Returns unit vector(s) from SSB to pulsar system barycenter under ECL. + + If epochs (MJD) are given, proper motion is included in the calculation. + + Parameters + ---------- + epoch : float or astropy.time.Time, optional + ecl : str, optional + Obliquity (IERS2010 by default) + + Returns + ------- + np.ndarray : + (len(epoch), 3) array of unit vectors + """ + # TODO: would it be better for this to return a 6-vector (pos, vel)? + + # this was somewhat slow, since it repeatedly created different SkyCoord Objects + # return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose() + + # Instead look at what https://docs.astropy.org/en/stable/_modules/astropy/coordinates/sky_coordinate.html#SkyCoord.apply_space_motion + # does, which is to use https://github.com/liberfa/erfa/blob/master/src/starpm.c + # and then just use the relevant pieces of that + + # but we need to check that the obliquity is the same + if ecl is not None and ecl != self.ECL.quantity: + return super().ssb_to_psb_xyz_ECL(epoch=epoch, ecl=ecl) + + if ecl is None: + log.debug("ECL not specified; using IERS2010.") + ecl = "IERS2010" + if epoch is None or (self.PMELONG.value == 0 and self.PMELAT.value == 0): + return self.coords_as_ECL(epoch=epoch, ecl=ecl).cartesian.xyz.transpose() + if isinstance(epoch, Time): + jd1 = epoch.jd1 + jd2 = epoch.jd2 + else: + jd1 = 2400000.5 + jd2 = epoch + # compared to the general case above we can assume that the coordinates are ECL + # so just access those components + lon = self.ELONG.quantity.to_value(u.radian) + lat = self.ELAT.quantity.to_value(u.radian) + pm_lon = ( + self.PMELONG.quantity.to_value(u.radian / u.yr) + / np.cos(self.ELAT.quantity).value + ) + pm_lat = self.PMELAT.quantity.to_value(u.radian / u.yr) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ErfaWarning) + # note that pmsafe wants mu_lon not mu_lon * cos(lat) + starpmout = pmsafe( + lon, + lat, + pm_lon, + pm_lat, + self.PX.quantity.to_value(u.arcsec), + 0.0, + self.POSEPOCH.quantity.jd1, + self.POSEPOCH.quantity.jd2, + jd1, + jd2, + ) + # lon,lat now in radians + lon, lat = starpmout[0], starpmout[1] + x = np.cos(lon) * np.cos(lat) + y = np.sin(lon) * np.cos(lat) + z = np.sin(lat) + return u.Quantity([x, y, z]).T + def get_d_delay_quantities_ecliptical(self, toas): """Calculate values needed for many d_delay_d_param functions.""" # TODO: Move all these calculations in a separate class for elegance diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 5c0df4b5d..0b6350548 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -1898,7 +1898,7 @@ def d_phase_d_param_num(self, toas, param, step=1e-2): res_f = -phase_f[:, 0] + phase_f[:, 1] result = (res_i + res_f) / (2.0 * h * unit) # shift value back to the original value - par.quantity = ori_value + par.value = ori_value return result def d_delay_d_param_num(self, toas, param, step=1e-2): diff --git a/tests/test_B1855.py b/tests/test_B1855.py index e1248e755..605b6953a 100644 --- a/tests/test_B1855.py +++ b/tests/test_B1855.py @@ -1,7 +1,6 @@ """Various tests to assess the performance of the B1855+09.""" import logging import os -import pytest import astropy.units as u import numpy as np diff --git a/tests/test_astrometry.py b/tests/test_astrometry.py index 152cd7c27..73ff00ecb 100644 --- a/tests/test_astrometry.py +++ b/tests/test_astrometry.py @@ -81,3 +81,123 @@ def test_pm_acquires_posepoch_equatorial(): def test_pm_one_set_other_not(): m = get_model(StringIO("\n".join([par_basic_equatorial, "PMRA 7"]))) assert m.POSEPOCH.quantity == m.PEPOCH.quantity + + +parICRS = """PSR 1748-2021E +RAJ 17:48:52.75 1 +DECJ -20:21:29.0 1 +PMRA 1000 +PMDEC -500 +F0 61.485476554 1 +F1 -1.181D-15 1 +PEPOCH 53750.000000 +POSEPOCH 53750.000000 +DM 223.9 1 +SOLARN0 0.00 +EPHEM DE421 +CLK UTC(NIST) +UNITS TDB +TIMEEPH FB90 +T2CMETHOD TEMPO +CORRECT_TROPOSPHERE N +PLANET_SHAPIRO N +DILATEFREQ N +TZRMJD 53801.38605118223 +TZRFRQ 1949.609 +TZRSITE 1 +""" + + +ntests = 200 + + +@pytest.mark.parametrize( + ("ra,dec,pmra,pmdec"), + list( + zip( + np.random.uniform(0, 360, ntests), + np.random.uniform(-90, 90, ntests), + np.random.normal(0, scale=10, size=ntests), + np.random.normal(0, scale=10, size=ntests), + ) + ), +) +def test_ssb_accuracy_ICRS(ra, dec, pmra, pmdec): + m = get_model(StringIO(parICRS), RAJ=ra, DECJ=dec, PMRA=pmra, PMDEC=pmdec) + t0 = m.POSEPOCH.quantity + t0.format = "pulsar_mjd" + t = t0 + np.linspace(0, 20) * u.yr + xyzastropy = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose() + xyznew = m.ssb_to_psb_xyz_ICRS(epoch=t) + assert np.allclose(xyznew, xyzastropy) + + +parECL = """PSR 1748-2021E +ELONG 300.0 +ELAT -55.0 +PMELONG 1000 +PMELAT -500 +F0 61.485476554 1 +F1 -1.181D-15 1 +PEPOCH 53750.000000 +POSEPOCH 53750.000000 +DM 223.9 1 +SOLARN0 0.00 +EPHEM DE421 +CLK UTC(NIST) +UNITS TDB +TIMEEPH FB90 +T2CMETHOD TEMPO +CORRECT_TROPOSPHERE N +PLANET_SHAPIRO N +DILATEFREQ N +TZRMJD 53801.38605118223 +TZRFRQ 1949.609 +TZRSITE 1 +""" + + +@pytest.mark.parametrize( + ("elong,elat,pmelong,pmelat"), + list( + zip( + np.random.uniform(0, 360, ntests), + np.random.uniform(-90, 90, ntests), + np.random.normal(0, scale=10, size=ntests), + np.random.normal(0, scale=10, size=ntests), + ) + ), +) +def test_ssb_accuracy_ICRS_ECLmodel(elong, elat, pmelong, pmelat): + m = get_model( + StringIO(parECL), ELONG=elong, ELAT=elat, PMELONG=pmelong, PMELAT=pmelat + ) + t0 = m.POSEPOCH.quantity + t0.format = "pulsar_mjd" + t = t0 + np.linspace(0, 20) * u.yr + xyzastropy = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose() + xyznew = m.ssb_to_psb_xyz_ICRS(epoch=t) + assert np.allclose(xyznew, xyzastropy) + + +@pytest.mark.parametrize( + ("elong,elat,pmelong,pmelat"), + list( + zip( + np.random.uniform(0, 360, ntests), + np.random.uniform(-90, 90, ntests), + np.random.normal(0, scale=10, size=ntests), + np.random.normal(0, scale=10, size=ntests), + ) + ), +) +def test_ssb_accuracy_ECL_ECLmodel(elong, elat, pmelong, pmelat): + m = get_model( + StringIO(parECL), ELONG=elong, ELAT=elat, PMELONG=pmelong, PMELAT=pmelat + ) + t0 = m.POSEPOCH.quantity + t0.format = "pulsar_mjd" + t = t0 + np.linspace(0, 20) * u.yr + xyzastropy = m.coords_as_ECL(epoch=t).cartesian.xyz.transpose() + xyznew = m.ssb_to_psb_xyz_ECL(epoch=t) + assert np.allclose(xyznew, xyzastropy) diff --git a/tests/test_derivative_utils.py b/tests/test_derivative_utils.py index ca35028ae..650caa37c 100644 --- a/tests/test_derivative_utils.py +++ b/tests/test_derivative_utils.py @@ -68,10 +68,8 @@ def get_derivative_params(model): h = 2e-1 elif p in ["FD2"]: h = 3e-2 - elif p in ["F1", "M2", "PMELONG", "PMELAT"]: + elif p in ["F1", "M2", "PMELONG", "PMELAT", "PMRA", "PMDEC"]: h = 2 - elif p in ["PMDEC"]: - h = 3e-2 elif p in ["PBDOT"]: h = 200 elif p in ["FB1"]: diff --git a/tests/test_modelconversions.py b/tests/test_modelconversions.py index 99b328e63..34a65f0d1 100644 --- a/tests/test_modelconversions.py +++ b/tests/test_modelconversions.py @@ -155,7 +155,10 @@ def test_ICRS_to_ECL_uncertainties(): m2 = fit_ECL.model.as_ICRS() for p in ("RAJ", "DECJ", "PMRA", "PMDEC"): - assert np.isclose(m1.__getitem__(p).value, m2.__getitem__(p).value) + tol = 1e-12 if p in ["RAJ", "DECJ"] else 1e-3 + assert np.isclose( + m1.__getitem__(p).value, m2.__getitem__(p).value, atol=tol + ), f"Paramter {p} with values {m1.__getitem__(p).value}, {m2.__getitem__(p).value} was too far apart (difference={m1.__getitem__(p).value-m2.__getitem__(p).value})" # do a generous test here since the uncertainties could change assert np.isclose( m1.__getitem__(p).uncertainty, m2.__getitem__(p).uncertainty, rtol=0.5