From a30ac391ec0b905b16371203dfea8212752f5b91 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 8 Sep 2023 14:03:14 -0500 Subject: [PATCH 01/11] dont use decompose --- .../models/stand_alone_psr_binaries/DD_model.py | 6 +++--- .../models/stand_alone_psr_binaries/ELL1_model.py | 8 +++++--- .../models/stand_alone_psr_binaries/ELL1k_model.py | 2 +- .../stand_alone_psr_binaries/binary_generic.py | 2 +- .../stand_alone_psr_binaries/binary_orbits.py | 11 +++++------ src/pint/residuals.py | 13 +++++++------ 6 files changed, 22 insertions(+), 20 deletions(-) diff --git a/src/pint/models/stand_alone_psr_binaries/DD_model.py b/src/pint/models/stand_alone_psr_binaries/DD_model.py index e9b82bdc3..b4b40b32c 100644 --- a/src/pint/models/stand_alone_psr_binaries/DD_model.py +++ b/src/pint/models/stand_alone_psr_binaries/DD_model.py @@ -370,7 +370,7 @@ def d_beta_d_ECC(self): de/dECC = 1 """ eTheta = self.eTheta() - a1 = (self.a1()).decompose() + a1 = (self.a1()).to(u.m) sinOmg = np.sin(self.omega()) cosOmg = np.cos(self.omega()) with u.set_enabled_equivalencies(u.dimensionless_angles()): @@ -392,7 +392,7 @@ def d_beta_d_EDOT(self): de/dEDOT = tt0 """ eTheta = self.eTheta() - a1 = (self.a1()).decompose() + a1 = (self.a1()).to(u.m) sinOmg = np.sin(self.omega()) cosOmg = np.cos(self.omega()) with u.set_enabled_equivalencies(u.dimensionless_angles()): @@ -642,7 +642,7 @@ def delayInverse(self): + 1.0 / 2 * nHat**2 * Dre * Drepp - 1.0 / 2 * e * sinE / (1 - e * cosE) * nHat**2 * Dre * Drep ) - ).decompose() + ).to(u.s) def d_delayI_d_par(self, par): """Derivative on delay inverse.""" diff --git a/src/pint/models/stand_alone_psr_binaries/ELL1_model.py b/src/pint/models/stand_alone_psr_binaries/ELL1_model.py index b3fd8133d..4ccc79ef3 100644 --- a/src/pint/models/stand_alone_psr_binaries/ELL1_model.py +++ b/src/pint/models/stand_alone_psr_binaries/ELL1_model.py @@ -106,7 +106,9 @@ def orbits_ELL1(self): PB = (self.pb()).to("second") PBDOT = self.pbdot() ttasc = self.ttasc() - return (ttasc / PB - 0.5 * PBDOT * (ttasc / PB) ** 2).decompose() + return (ttasc / PB - 0.5 * PBDOT * (ttasc / PB) ** 2).to( + u.dimensionless_unscaled + ) def d_Phi_d_TASC(self): """dPhi/dTASC""" @@ -162,7 +164,7 @@ def delayI(self): return ( Dre * (1 - nhat * Drep + (nhat * Drep) ** 2 + 1.0 / 2 * nhat**2 * Dre * Drepp) - ).decompose() + ).to(u.s) def nhat(self): return 2 * np.pi / self.pb() @@ -319,7 +321,7 @@ def delayR(self): Zhu et al. (2019), Eqn. 1 Fiore et al. (2023), Eqn. 4 """ - return ((self.a1() / c.c) * self.d_delayR_da1()).decompose() + return ((self.a1() / c.c) * self.d_delayR_da1()).to(u.s) def d_Dre_d_par(self, par): """Derivative computation. diff --git a/src/pint/models/stand_alone_psr_binaries/ELL1k_model.py b/src/pint/models/stand_alone_psr_binaries/ELL1k_model.py index 0fe785988..9863b1aaa 100644 --- a/src/pint/models/stand_alone_psr_binaries/ELL1k_model.py +++ b/src/pint/models/stand_alone_psr_binaries/ELL1k_model.py @@ -130,7 +130,7 @@ def delayR(self): + 0.5 * (self.eps2() * np.sin(2 * Phi) - self.eps1() * (np.cos(2 * Phi) + 3)) ) - ).decompose() + ).to(u.s) def d_Dre_d_par(self, par): """Derivative computation. diff --git a/src/pint/models/stand_alone_psr_binaries/binary_generic.py b/src/pint/models/stand_alone_psr_binaries/binary_generic.py index c62f174a3..5d5accf27 100644 --- a/src/pint/models/stand_alone_psr_binaries/binary_generic.py +++ b/src/pint/models/stand_alone_psr_binaries/binary_generic.py @@ -381,7 +381,7 @@ def get_tt0(self, barycentricTOA): def ecc(self): """Calculate eccentricity with EDOT""" if hasattr(self, "_tt0"): - return self.ECC + (self.tt0 * self.EDOT).decompose() + return self.ECC + (self.tt0 * self.EDOT).to(u.dimensionless_unscaled) return self.ECC def d_ecc_d_T0(self): diff --git a/src/pint/models/stand_alone_psr_binaries/binary_orbits.py b/src/pint/models/stand_alone_psr_binaries/binary_orbits.py index 3a1d60bbb..31b2b310d 100644 --- a/src/pint/models/stand_alone_psr_binaries/binary_orbits.py +++ b/src/pint/models/stand_alone_psr_binaries/binary_orbits.py @@ -100,9 +100,9 @@ def orbits(self): PB = self.PB.to("second") PBDOT = self.PBDOT XPBDOT = self.XPBDOT - return ( - self.tt0 / PB - 0.5 * (PBDOT + XPBDOT) * (self.tt0 / PB) ** 2 - ).decompose() + return (self.tt0 / PB - 0.5 * (PBDOT + XPBDOT) * (self.tt0 / PB) ** 2).to( + u.dimensionless_unscaled + ) def pbprime(self): """Instantaneous binary period as a function of time.""" @@ -183,7 +183,7 @@ def _FBXs(self): def orbits(self): """Orbital phase (number of orbits since T0).""" orbits = taylor_horner(self.tt0, self._FBXs()) - return orbits.decompose() + return orbits.to(u.dimensionless_unscaled) def pbprime(self): """Instantaneous binary period as a function of time.""" @@ -213,8 +213,7 @@ def d_orbits_d_FBX(self, FBX): FBXs.append(1.0 * getattr(self, f"FB{ii}").unit) break ii += 1 - d_orbits = taylor_horner(self.tt0, FBXs) / par.unit - return d_orbits.decompose() * 2 * np.pi * u.rad + return taylor_horner(self.tt0, FBXs) * 2 * np.pi * (u.rad / par.unit) def d_pbprime_d_FBX(self, FBX): par = getattr(self, FBX) diff --git a/src/pint/residuals.py b/src/pint/residuals.py index 736c5798b..68c9e9a35 100644 --- a/src/pint/residuals.py +++ b/src/pint/residuals.py @@ -506,7 +506,7 @@ def _calc_gls_chi2(self): If the system is singular, it uses singular value decomposition instead.""" self.update() - residuals = self.time_resids.to(u.s).value + residuals = self.time_resids.to_value(u.s) M = self.model.noise_model_designmatrix(self.toas) phi = self.model.noise_model_basis_weight(self.toas) @@ -520,7 +520,7 @@ def _calc_gls_chi2(self): M, norm = normalize_designmatrix(M, None) phiinv /= norm**2 - Nvec = self.model.scaled_toa_uncertainty(self.toas).to(u.s).value ** 2 + Nvec = self.model.scaled_toa_uncertainty(self.toas).to_value(u.s) ** 2 cinv = 1 / Nvec mtcm = np.dot(M.T, cinv[:, None] * M) mtcm += np.diag(phiinv) @@ -763,10 +763,11 @@ def calc_chi2(self): data_errors = self.get_data_error() if (data_errors == 0.0).any(): return np.inf - try: - return ((self.resids / data_errors) ** 2.0).sum().decompose().value - except ValueError: - return ((self.resids / data_errors) ** 2.0).sum().decompose() + return ( + ((self.resids / data_errors) ** 2.0) + .sum() + .to_value(u.dimensionless_unscaled) + ) def rms_weighted(self): """Compute weighted RMS of the residuals in time.""" From 996fd2f247de76ad61e35e4702b391786229e585 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 8 Sep 2023 14:52:45 -0500 Subject: [PATCH 02/11] -- --- src/pint/models/astrometry.py | 24 ++++++++++-------------- src/pint/residuals.py | 9 +++++---- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/src/pint/models/astrometry.py b/src/pint/models/astrometry.py index 1d737849c..ba700ff9d 100644 --- a/src/pint/models/astrometry.py +++ b/src/pint/models/astrometry.py @@ -203,10 +203,10 @@ def d_delay_astrometry_d_PX(self, toas, param="", acc_delay=None): rd = self.get_d_delay_quantities(toas) px_r = np.sqrt(rd["ssb_obs_r"] ** 2 - rd["in_psr_obs"] ** 2) - dd_dpx = 0.5 * (px_r**2 / (u.AU * const.c)) * (u.mas / u.radian) + dd_dpx = 0.5 * px_r**2 * ((u.mas / u.radian) / (u.AU * const.c)) # We want to return sec / mas - return dd_dpx.decompose(u.si.bases) / u.mas + return dd_dpx.to(u.s) / u.mas def d_delay_astrometry_d_POSEPOCH(self, toas, param="", acc_delay=None): """Calculate the derivative wrt POSEPOCH""" @@ -415,7 +415,7 @@ def d_delay_astrometry_d_RAJ(self, toas, param="", acc_delay=None): ) dd_draj = rd["ssb_obs_r"] * geom / (const.c * u.radian) - return dd_draj.decompose(u.si.bases) + return dd_draj.to(u.s / u.rad) def d_delay_astrometry_d_DECJ(self, toas, param="", acc_delay=None): """Calculate the derivative wrt DECJ @@ -432,7 +432,7 @@ def d_delay_astrometry_d_DECJ(self, toas, param="", acc_delay=None): ) - np.sin(rd["earth_dec"]) * np.cos(psr_dec) dd_ddecj = rd["ssb_obs_r"] * geom / (const.c * u.radian) - return dd_ddecj.decompose(u.si.bases) + return dd_ddecj.to(u.s / u.rad) def d_delay_astrometry_d_PMRA(self, toas, param="", acc_delay=None): """Calculate the derivative wrt PMRA @@ -448,10 +448,9 @@ def d_delay_astrometry_d_PMRA(self, toas, param="", acc_delay=None): geom = np.cos(rd["earth_dec"]) * np.sin(psr_ra - rd["earth_ra"]) deriv = rd["ssb_obs_r"] * geom * te / (const.c * u.radian) - dd_dpmra = deriv * u.mas / u.year # We want to return sec / (mas / yr) - return dd_dpmra.decompose(u.si.bases) / (u.mas / u.year) + return deriv.to(u.s * u.year / u.mas) def d_delay_astrometry_d_PMDEC(self, toas, param="", acc_delay=None): """Calculate the derivative wrt PMDEC @@ -470,10 +469,9 @@ def d_delay_astrometry_d_PMDEC(self, toas, param="", acc_delay=None): ) - np.cos(psr_dec) * np.sin(rd["earth_dec"]) deriv = rd["ssb_obs_r"] * geom * te / (const.c * u.radian) - dd_dpmdec = deriv * u.mas / u.year # We want to return sec / (mas / yr) - return dd_dpmdec.decompose(u.si.bases) / (u.mas / u.year) + return deriv.to(u.s * u.year / u.mas) def change_posepoch(self, new_epoch): """Change POSEPOCH to a new value and update the position accordingly. @@ -823,7 +821,7 @@ def d_delay_astrometry_d_ELONG(self, toas, param="", acc_delay=None): ) dd_delong = rd["ssb_obs_r"] * geom / (const.c * u.radian) - return dd_delong.decompose(u.si.bases) + return dd_delong.to(u.s / u.rad) def d_delay_astrometry_d_ELAT(self, toas, param="", acc_delay=None): """Calculate the derivative wrt DECJ @@ -840,7 +838,7 @@ def d_delay_astrometry_d_ELAT(self, toas, param="", acc_delay=None): ) - np.sin(rd["earth_elat"]) * np.cos(psr_elat) dd_delat = rd["ssb_obs_r"] * geom / (const.c * u.radian) - return dd_delat.decompose(u.si.bases) + return dd_delat.to(u.s / u.rad) def d_delay_astrometry_d_PMELONG(self, toas, param="", acc_delay=None): """Calculate the derivative wrt PMRA @@ -857,10 +855,9 @@ def d_delay_astrometry_d_PMELONG(self, toas, param="", acc_delay=None): geom = np.cos(rd["earth_elat"]) * np.sin(psr_elong - rd["earth_elong"]) deriv = rd["ssb_obs_r"] * geom * te / (const.c * u.radian) - dd_dpmelong = deriv * u.mas / u.year # We want to return sec / (mas / yr) - return dd_dpmelong.decompose(u.si.bases) / (u.mas / u.year) + return deriv.to(u.s * u.year / u.mas) def d_delay_astrometry_d_PMELAT(self, toas, param="", acc_delay=None): """Calculate the derivative wrt PMDEC @@ -879,10 +876,9 @@ def d_delay_astrometry_d_PMELAT(self, toas, param="", acc_delay=None): ) - np.cos(psr_elat) * np.sin(rd["earth_elat"]) deriv = rd["ssb_obs_r"] * geom * te / (const.c * u.radian) - dd_dpmelat = deriv * u.mas / u.year # We want to return sec / (mas / yr) - return dd_dpmelat.decompose(u.si.bases) / (u.mas / u.year) + return deriv.to(u.s * u.year / u.mas) def print_par(self, format="pint"): result = "" diff --git a/src/pint/residuals.py b/src/pint/residuals.py index 68c9e9a35..214fd656a 100644 --- a/src/pint/residuals.py +++ b/src/pint/residuals.py @@ -582,10 +582,11 @@ def calc_chi2(self): # This the fastest way, but highly depend on the assumption of time_resids and # error units. Ensure only a pure number is returned. - try: - return ((self.time_resids / toa_errors.to(u.s)) ** 2.0).sum().value - except ValueError: - return ((self.time_resids / toa_errors.to(u.s)) ** 2.0).sum() + return ( + ((self.time_resids / toa_errors) ** 2.0) + .sum() + .to_value(u.dimensionless_unscaled) + ) def ecorr_average(self, use_noise_model=True): """Uses the ECORR noise model time-binning to compute "epoch-averaged" residuals. From 22796a4f9b8c92b638d464508c81ba517bc76dd1 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 8 Sep 2023 15:27:19 -0500 Subject: [PATCH 03/11] cache parameter masks --- src/pint/models/parameter.py | 10 ++++++++++ src/pint/models/timing_model.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/src/pint/models/parameter.py b/src/pint/models/parameter.py index 016687563..6bd90d41d 100644 --- a/src/pint/models/parameter.py +++ b/src/pint/models/parameter.py @@ -1739,6 +1739,9 @@ def __init__( self.is_prefix = True self._parfile_name = self.origin_name + self.__cache_mask = None + self.__cache_toas = None + def __repr__(self): out = f"{self.__class__.__name__}({self.name}" if self.key is not None: @@ -1931,6 +1934,9 @@ def select_toa_mask(self, toas): array An array of TOA indices selected by the mask. """ + if self.__cache_mask is not None and self.__cache_toas is toas: + return self.__cache_mask + if len(self.key_value) == 1: if not hasattr(self, "toa_selector"): self.toa_selector = TOASelect(is_range=False, use_hash=True) @@ -1966,6 +1972,10 @@ def select_toa_mask(self, toas): select_idx = self.toa_selector.get_select_index(condition, col) return select_idx[self.name] + def cache_toa_mask(self, toas): + self.__cache_mask = self.select_toa_mask(toas) + self.__cache_toas = toas + def compare_key_value(self, other_param): """Compare if the key and value are the same with the other parameter. diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 1f02b8c75..36188fba2 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -349,6 +349,9 @@ def __init__(self, name="", components=[]): for cp in components: self.add_component(cp, setup=False, validate=False) + self.__locked = False + self.__lock_toas = None + def __repr__(self): return "{}(\n {}\n)".format( self.__class__.__name__, @@ -485,6 +488,19 @@ def num_components_of_type(type): # result += str(getattr(cp, pp)) + "\n" # return result + def lock(self, toas): + log.warning( + "Locking the TimingModel object for the given TOAs. Mutating this TOAs object or using a " + "different TOAs object with this model will lead to undefined behavior." + ) + self.__locked = True + self.__lock_toas = toas + + for mp in self.get_params_of_type_top("maskParameter"): + mpar = getattr(self, mp) + if mpar.key is not None: + mpar.cache_toa_mask(toas) + def __getattr__(self, name): if name in ["components", "component_types", "search_cmp_attr"]: raise AttributeError From cc49c731178b66784f489a63c233b6fdc355b350 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Fri, 8 Sep 2023 15:58:44 -0500 Subject: [PATCH 04/11] docs and test --- src/pint/models/parameter.py | 3 +++ src/pint/models/timing_model.py | 10 ++++++++-- tests/test_mask_cache.py | 19 +++++++++++++++++++ 3 files changed, 30 insertions(+), 2 deletions(-) create mode 100644 tests/test_mask_cache.py diff --git a/src/pint/models/parameter.py b/src/pint/models/parameter.py index 6bd90d41d..a62685731 100644 --- a/src/pint/models/parameter.py +++ b/src/pint/models/parameter.py @@ -1973,6 +1973,9 @@ def select_toa_mask(self, toas): return select_idx[self.name] def cache_toa_mask(self, toas): + """Cache the TOA mask for this parameter. This will lead to exceptions and/or + wrong results if the TOAs object is mutated after this. + """ self.__cache_mask = self.select_toa_mask(toas) self.__cache_toas = toas diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 36188fba2..797716138 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -489,10 +489,16 @@ def num_components_of_type(type): # return result def lock(self, toas): + """Lock a timing model object for a given TOAs object by caching parameter masks. This can + provide significant performance gains when mask parameters are present. However, mutating the + TOAs object after calling this function can lead to exceptions or wrong results down the line. + """ + log.warning( - "Locking the TimingModel object for the given TOAs. Mutating this TOAs object or using a " - "different TOAs object with this model will lead to undefined behavior." + "Locking the TimingModel object for the given TOAs object. This is will lead to " + "undefined behavior if the TOAs object is mutated after this." ) + self.__locked = True self.__lock_toas = toas diff --git a/tests/test_mask_cache.py b/tests/test_mask_cache.py new file mode 100644 index 000000000..d28839254 --- /dev/null +++ b/tests/test_mask_cache.py @@ -0,0 +1,19 @@ +from pinttestdata import datadir +from pint.models import get_model_and_toas +from pint.residuals import Residuals + + +def test_mask_cache(): + model, toas = get_model_and_toas( + datadir / "J1713+0747_small.gls.par", datadir / "J1713+0747_small.tim" + ) + + res1 = Residuals(toas, model) + chi2_1 = res1.calc_chi2() + + model.lock(toas) + + res2 = Residuals(toas, model) + chi2_2 = res2.calc_chi2() + + assert chi2_1 == chi2_2 From 92a409a086a6e9239896c0eab731b410c91934a0 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 06:50:12 -0500 Subject: [PATCH 05/11] calc_chi2 --- src/pint/residuals.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/pint/residuals.py b/src/pint/residuals.py index 214fd656a..e0bada80c 100644 --- a/src/pint/residuals.py +++ b/src/pint/residuals.py @@ -571,8 +571,11 @@ def calc_chi2(self): else: # Residual units are in seconds. Error units are in microseconds. toa_errors = self.get_data_error() - if (toa_errors == 0.0).any(): + + # This check is expensive. Skip it if the model is locked. + if (not self.model.locked) and (toa_errors == 0.0).any(): return np.inf + # The self.time_resids is in the unit of "s", the error "us". # This is more correct way, but it is the slowest. # return (((self.time_resids / self.toas.get_errors()).decompose()**2.0).sum()).value @@ -583,10 +586,10 @@ def calc_chi2(self): # This the fastest way, but highly depend on the assumption of time_resids and # error units. Ensure only a pure number is returned. return ( - ((self.time_resids / toa_errors) ** 2.0) - .sum() - .to_value(u.dimensionless_unscaled) - ) + ( + (self.time_resids.to_value(u.s) / toa_errors.to_value(u.s)) ** 2.0 + ).sum() + ) << u.dimensionless_unscaled def ecorr_average(self, use_noise_model=True): """Uses the ECORR noise model time-binning to compute "epoch-averaged" residuals. From 70df575c040c460c17a0a839ee1946230a137757 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 06:51:17 -0500 Subject: [PATCH 06/11] scale_toa_sigma --- src/pint/models/noise_model.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/pint/models/noise_model.py b/src/pint/models/noise_model.py index 02a5d291a..b735ce076 100644 --- a/src/pint/models/noise_model.py +++ b/src/pint/models/noise_model.py @@ -146,24 +146,29 @@ def validate(self): raise ValueError(f"'{el}' have duplicated keys and key values.") def scale_toa_sigma(self, toas): - sigma_scaled = toas.table["error"].quantity.copy() + # @abhisrkckl : This function turned out to be a performance bottleneck. + # So I have changed this code so as to minimize Quantity operations. + + sigma_scaled = toas.table["error"].quantity.to_value(u.s) for equad_name in self.EQUADs: equad = getattr(self, equad_name) if equad.quantity is None: continue mask = equad.select_toa_mask(toas) if np.any(mask): - sigma_scaled[mask] = np.hypot(sigma_scaled[mask], equad.quantity) + sigma_scaled[mask] = np.hypot( + sigma_scaled[mask], equad.quantity.to_value(u.s) + ) else: warnings.warn(f"EQUAD {equad} has no TOAs") for efac_name in self.EFACs: efac = getattr(self, efac_name) mask = efac.select_toa_mask(toas) if np.any(mask): - sigma_scaled[mask] *= efac.quantity + sigma_scaled[mask] *= efac.quantity.to_value(u.dimensionless_unscaled) else: warnings.warn(f"EFAC {efac} has no TOAs") - return sigma_scaled + return sigma_scaled << u.s def sigma_scaled_cov_matrix(self, toas): scaled_sigma = self.scale_toa_sigma(toas).to(u.s).value ** 2 From 10e490a59e63128cbcd3d75a65b8ef2a03324b2c Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 07:01:25 -0500 Subject: [PATCH 07/11] locked, unlock --- src/pint/models/parameter.py | 8 +++++++- src/pint/models/timing_model.py | 28 +++++++++++++++++++++++----- tests/test_mask_cache.py | 11 +++++++++-- 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/src/pint/models/parameter.py b/src/pint/models/parameter.py index a62685731..059ac21dd 100644 --- a/src/pint/models/parameter.py +++ b/src/pint/models/parameter.py @@ -1974,11 +1974,17 @@ def select_toa_mask(self, toas): def cache_toa_mask(self, toas): """Cache the TOA mask for this parameter. This will lead to exceptions and/or - wrong results if the TOAs object is mutated after this. + wrong results if the TOAs object is mutated after this. This can be undone using + the `clear_toa_mask_cache()` method. """ self.__cache_mask = self.select_toa_mask(toas) self.__cache_toas = toas + def clear_toa_mask_cache(self): + """Clear the TOA mask cache. This will undo a `cache_toa_mask()` call.""" + self.__cache_mask = None + self.__cache_toas = None + def compare_key_value(self, other_param): """Compare if the key and value are the same with the other parameter. diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 797716138..28d4db032 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -499,6 +499,11 @@ def lock(self, toas): "undefined behavior if the TOAs object is mutated after this." ) + if (toas.get_errors() == 0).any(): + raise ValueError( + "Model cannot be locked since some TOAs have zero uncertainties." + ) + self.__locked = True self.__lock_toas = toas @@ -507,6 +512,19 @@ def lock(self, toas): if mpar.key is not None: mpar.cache_toa_mask(toas) + @property + def locked(self): + return self.__locked + + def unlock(self): + self.__locked = False + self.__lock_toas = None + + for mp in self.get_params_of_type_top("maskParameter"): + mpar = getattr(self, mp) + if mpar.key is not None: + mpar.clear_toa_mask_cache() + def __getattr__(self, name): if name in ["components", "component_types", "search_cmp_attr"]: raise AttributeError @@ -1451,15 +1469,15 @@ def scaled_toa_uncertainty(self, toas): """ ntoa = toas.ntoas tbl = toas.table - result = np.zeros(ntoa) * u.us + result = np.zeros(ntoa) + # When there is no noise model. if len(self.scaled_toa_uncertainty_funcs) == 0: - result += tbl["error"].quantity - return result + return tbl["error"].quantity.to(u.us) for nf in self.scaled_toa_uncertainty_funcs: - result += nf(toas) - return result + result += nf(toas).to_value(u.us) + return result << u.us def scaled_dm_uncertainty(self, toas): """Get the scaled DM data uncertainties noise models. diff --git a/tests/test_mask_cache.py b/tests/test_mask_cache.py index d28839254..c09f372a4 100644 --- a/tests/test_mask_cache.py +++ b/tests/test_mask_cache.py @@ -1,6 +1,7 @@ from pinttestdata import datadir from pint.models import get_model_and_toas from pint.residuals import Residuals +import pytest def test_mask_cache(): @@ -12,8 +13,14 @@ def test_mask_cache(): chi2_1 = res1.calc_chi2() model.lock(toas) - res2 = Residuals(toas, model) chi2_2 = res2.calc_chi2() - assert chi2_1 == chi2_2 + + model.unlock() + chi2_3 = res1.calc_chi2() + assert chi2_1 == chi2_3 + + toas.table["error"][0] = 0 + with pytest.raises(ValueError): + model.lock(toas) From e7b7c094665ac775c98ea391473b95f8902180ad Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 07:08:57 -0500 Subject: [PATCH 08/11] docs --- src/pint/models/timing_model.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 28d4db032..26a4b1cda 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -489,9 +489,10 @@ def num_components_of_type(type): # return result def lock(self, toas): - """Lock a timing model object for a given TOAs object by caching parameter masks. This can - provide significant performance gains when mask parameters are present. However, mutating the - TOAs object after calling this function can lead to exceptions or wrong results down the line. + """Lock a timing model object for a given TOAs object by caching parameter masks. This also + allows us to do other optimizations such as skipping TOA error checks. This can provide significant + performance gains when mask parameters are present. However, mutating the TOAs object after calling + this function can lead to exceptions or wrong results down the line. """ log.warning( @@ -517,6 +518,7 @@ def locked(self): return self.__locked def unlock(self): + """Undo the lock() call.""" self.__locked = False self.__lock_toas = None From 2207488401db73d698bda0459f377828b747a155 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 07:58:40 -0500 Subject: [PATCH 09/11] scale_toa_sigma --- src/pint/models/noise_model.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/pint/models/noise_model.py b/src/pint/models/noise_model.py index b735ce076..567ae4012 100644 --- a/src/pint/models/noise_model.py +++ b/src/pint/models/noise_model.py @@ -149,13 +149,18 @@ def scale_toa_sigma(self, toas): # @abhisrkckl : This function turned out to be a performance bottleneck. # So I have changed this code so as to minimize Quantity operations. - sigma_scaled = toas.table["error"].quantity.to_value(u.s) + sigma_scaled = ( + self._parent._toa_errors_cache + if self._parent.locked + else toas.table["error"].quantity.to_value(u.s) + ) + for equad_name in self.EQUADs: equad = getattr(self, equad_name) if equad.quantity is None: continue mask = equad.select_toa_mask(toas) - if np.any(mask): + if self._parent.locked or np.any(mask): sigma_scaled[mask] = np.hypot( sigma_scaled[mask], equad.quantity.to_value(u.s) ) @@ -164,7 +169,7 @@ def scale_toa_sigma(self, toas): for efac_name in self.EFACs: efac = getattr(self, efac_name) mask = efac.select_toa_mask(toas) - if np.any(mask): + if self._parent.locked or np.any(mask): sigma_scaled[mask] *= efac.quantity.to_value(u.dimensionless_unscaled) else: warnings.warn(f"EFAC {efac} has no TOAs") From 4616b0ac03e02bd4ce4b1a9f5454eabd209928a6 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 08:00:00 -0500 Subject: [PATCH 10/11] cache_toa_mask --- src/pint/models/parameter.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/pint/models/parameter.py b/src/pint/models/parameter.py index 059ac21dd..b455b4437 100644 --- a/src/pint/models/parameter.py +++ b/src/pint/models/parameter.py @@ -1977,6 +1977,9 @@ def cache_toa_mask(self, toas): wrong results if the TOAs object is mutated after this. This can be undone using the `clear_toa_mask_cache()` method. """ + assert self.select_toa_mask( + toas + ).any(), f"Empty TOA mask found for {self.name}. Will not cache this." self.__cache_mask = self.select_toa_mask(toas) self.__cache_toas = toas From b02420ee45e122eb5dcccf5e5d4038368b41a518 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Sun, 10 Sep 2023 08:00:35 -0500 Subject: [PATCH 11/11] init --- src/pint/models/timing_model.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py index 26a4b1cda..cb7589874 100644 --- a/src/pint/models/timing_model.py +++ b/src/pint/models/timing_model.py @@ -351,6 +351,7 @@ def __init__(self, name="", components=[]): self.__locked = False self.__lock_toas = None + self.__toa_errors_cache = None def __repr__(self): return "{}(\n {}\n)".format( @@ -507,6 +508,7 @@ def lock(self, toas): self.__locked = True self.__lock_toas = toas + self._toa_errors_cache = toas.get_errors().to_value(u.s) for mp in self.get_params_of_type_top("maskParameter"): mpar = getattr(self, mp) @@ -521,6 +523,7 @@ def unlock(self): """Undo the lock() call.""" self.__locked = False self.__lock_toas = None + self._toa_errors_cache = None for mp in self.get_params_of_type_top("maskParameter"): mpar = getattr(self, mp) @@ -1469,17 +1472,21 @@ def scaled_toa_uncertainty(self, toas): toas: pint.toa.TOAs The input data object for TOAs uncertainty. """ - ntoa = toas.ntoas tbl = toas.table - result = np.zeros(ntoa) # When there is no noise model. if len(self.scaled_toa_uncertainty_funcs) == 0: return tbl["error"].quantity.to(u.us) - for nf in self.scaled_toa_uncertainty_funcs: - result += nf(toas).to_value(u.us) - return result << u.us + return ( + sum( + map( + lambda nf: nf(toas).to_value(u.us), + self.scaled_toa_uncertainty_funcs, + ) + ) + << u.us + ) def scaled_dm_uncertainty(self, toas): """Get the scaled DM data uncertainties noise models.