Skip to content

Commit

Permalink
Merge pull request #1883 from abhisrkckl/noise_resids
Browse files Browse the repository at this point in the history 
Store correlated noise amplitudes instead of the residuals
  • Loading branch information
@scottransom
scottransom authored Feb 19, 2025
2 parents 2071628 + 23870cf commit 1d54c39
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 27 deletions.
1 change: 1 addition & 0 deletions CHANGELOG-unreleased.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ the released changes.

## Unreleased
### Changed
- In `Residuals`, store correlated noise amplitudes instead of noise residuals. `Residuals.noise_resids` is now a `@property`.
### Added
- Simulate correlated DM noise for wideband TOAs
- Type hints in `pint.models.timing_model`
Expand Down
44 changes: 19 additions & 25 deletions src/pint/fitter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1495,19 +1495,16 @@ def fit_toas(self, maxiter=10, threshold=0, full_cov=False, debug=False, **kwarg
# Compute the noise realizations if possible
if not self.full_cov:
noise_dims = self.model.noise_model_dimensions(self.toas)
noise_resids = {}
noise_ampls = {}
ntmpar = len(self.model.free_params)
for comp in noise_dims:
# The first column of designmatrix is "offset", add 1 to match
# the indices of noise designmatrix
p0 = noise_dims[comp][0] + ntmpar + 1
p1 = p0 + noise_dims[comp][1]
noise_resids[comp] = (
np.dot(
self.current_state.M[:, p0:p1], self.current_state.xhat[p0:p1]
)
* u.s
)
noise_ampls[comp] = (self.current_state.xhat / self.current_state.norm)[
p0:p1
] * u.s
if debug:
setattr(
self.resids,
Expand All @@ -1518,7 +1515,7 @@ def fit_toas(self, maxiter=10, threshold=0, full_cov=False, debug=False, **kwarg
),
)
setattr(self.resids, f"{comp}_M_index", (p0, p1))
self.resids.noise_resids = noise_resids
self.resids.noise_ampls = noise_ampls
if debug:
setattr(self.resids, "norm", self.current_state.norm)

Expand Down Expand Up @@ -1786,19 +1783,16 @@ def fit_toas(
# Compute the noise realizations if possibl
if not self.full_cov:
noise_dims = self.model.noise_model_dimensions(self.toas)
noise_resids = {}
noise_ampls = {}
ntmpar = len(self.model.free_params)
for comp in noise_dims:
# The first column of designmatrix is "offset", add 1 to match
# the indices of noise designmatrix
p0 = noise_dims[comp][0] + ntmpar + 1
p1 = p0 + noise_dims[comp][1]
noise_resids[comp] = (
np.dot(
self.current_state.M[:, p0:p1], self.current_state.xhat[p0:p1]
)
* u.s
)
noise_ampls[comp] = (self.current_state.xhat / self.current_state.norm)[
p0:p1
] * u.s
if debug:
setattr(
self.resids,
Expand All @@ -1809,7 +1803,7 @@ def fit_toas(
),
)
setattr(self.resids, f"{comp}_M_index", (p0, p1))
self.resids.noise_resids = noise_resids
self.resids.noise_ampls = noise_ampls
if debug:
setattr(self.resids, "norm", self.current_state.norm)
return r
Expand Down Expand Up @@ -2185,17 +2179,17 @@ def fit_toas(self, maxiter=1, threshold=0, full_cov=False, debug=False):
# Compute the noise realizations if possible
if not full_cov:
noise_dims = self.model.noise_model_dimensions(self.toas)
noise_resids = {}
noise_ampls = {}
for comp in noise_dims:
# The first column of designmatrix is "offset", add 1 to match
# the indices of noise designmatrix
p0 = noise_dims[comp][0] + ntmpar + 1
p1 = p0 + noise_dims[comp][1]
noise_resids[comp] = np.dot(M[:, p0:p1], xhat[p0:p1]) * u.s
noise_ampls[comp] = (xhat / norm)[p0:p1] * u.s
if debug:
setattr(self.resids, f"{comp}_M", (M[:, p0:p1], xhat[p0:p1]))
setattr(self.resids, f"{comp}_M_index", (p0, p1))
self.resids.noise_resids = noise_resids
self.resids.noise_ampls = noise_ampls
if debug:
setattr(self.resids, "norm", norm)

Expand Down Expand Up @@ -2536,17 +2530,17 @@ def fit_toas(self, maxiter=1, threshold=0, full_cov=False, debug=False):
# Compute the noise realizations if possible
if not full_cov:
noise_dims = self.model.noise_model_dimensions(self.toas)
noise_resids = {}
noise_ampls = {}
for comp in noise_dims:
# The first column of designmatrix is "offset", add 1 to match
# the indices of noise designmatrix
p0 = noise_dims[comp][0] + ntmpar + 1
p1 = p0 + noise_dims[comp][1]
noise_resids[comp] = np.dot(M[:, p0:p1], xhat[p0:p1]) * u.s
noise_ampls[comp] = (xhat / norm)[p0:p1] * u.s
if debug:
setattr(self.resids, f"{comp}_M", (M[:, p0:p1], xhat[p0:p1]))
setattr(self.resids, f"{comp}_M_index", (p0, p1))
self.resids.noise_resids = noise_resids
self.resids.noise_ampls = noise_ampls
if debug:
setattr(self.resids, "norm", norm)

Expand Down Expand Up @@ -2743,19 +2737,19 @@ def update_from_state(self, state, debug=False):
# Compute the noise realizations if possible
if not self.full_cov:
noise_dims = self.model.noise_model_dimensions(self.toas)
noise_resids = {}
noise_ampls = {}
ntmpar = len(self.model.free_params)
for comp in noise_dims:
# The first column of designmatrix is "offset", add 1 to match
# the indices of noise designmatrix
p0 = noise_dims[comp][0] + ntmpar + 1
p1 = p0 + noise_dims[comp][1]
noise_resids[comp] = np.dot(state.M[:, p0:p1], state.xhat[p0:p1]) * u.s
noise_ampls[comp] = (state.xhat / state.norm)[p0:p1] * u.s
if debug:
setattr(
self.resids, f"{comp}_M", (state.M[:, p0:p1], state.xhat[p0:p1])
)
setattr(self.resids, f"{comp}_M_index", (p0, p1))
self.resids.noise_resids = noise_resids
self.resids.noise_ampls = noise_ampls
if debug:
setattr(self.resids, "norm", state.norm)
21 changes: 19 additions & 2 deletions src/pint/residuals.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@

import collections
import copy
from typing import Dict, Literal, Optional, Union
import warnings
from typing import Literal, Optional, Union

import astropy.units as u
import numpy as np
Expand Down Expand Up @@ -158,7 +158,6 @@ def __init__(
# also it's expensive
# only relevant if there are correlated errors
self._chi2 = None
self.noise_resids = {}
# For residual debugging
self.debug_info = {}
# We should be carefully for the other type of residuals
Expand All @@ -167,6 +166,13 @@ def __init__(
# class.
self._is_combined = False

self.noise_ampls: Dict[str, u.Quantity] = {}
for component in model.NoiseComponent_list:
if component.introduces_correlated_errors:
self.noise_ampls[component.category] = (
np.zeros_like(component.get_noise_weights(toas)) << u.s
)

@property
def resids(self) -> u.Quantity:
"""Residuals in time units."""
Expand All @@ -179,6 +185,17 @@ def resids_value(self) -> np.ndarray:
"""Residuals in seconds, with the units stripped."""
return self.resids.to_value(self.unit)

@property
def noise_resids(self) -> Dict[str, u.Quantity]:
return {
component.category: (
component.get_noise_basis(self.toas)
@ self.noise_ampls[component.category]
)
for component in self.model.NoiseComponent_list
if component.introduces_correlated_errors
}

def update(self) -> None:
"""Recalculate everything in residuals class after changing model or TOAs"""
if self.toas is None or self.model is None:
Expand Down

0 comments on commit 1d54c39

Please sign in to comment.