diff --git a/enterprise/signals/signal_base.py b/enterprise/signals/signal_base.py index e0fd3d32..adc3125e 100644 --- a/enterprise/signals/signal_base.py +++ b/enterprise/signals/signal_base.py @@ -25,6 +25,7 @@ from enterprise.signals.parameter import function # noqa: F401 from enterprise.signals.parameter import ConstantParameter from enterprise.signals.utils import KernelMatrix +from enterprise.signals.utils import indices_from_slice from enterprise import __version__ from sys import version @@ -1118,6 +1119,7 @@ class BlockMatrix(object): def __init__(self, blocks, slices, nvec=0): self._blocks = blocks self._slices = slices + self._idxs = [indices_from_slice(slc) for slc in slices] self._nvec = nvec if np.any(nvec != 0): @@ -1152,15 +1154,15 @@ def _solve_ZNX(self, X, Z): ZNXr = np.dot(Z[self._idx, :].T, X[self._idx, :] / self._nvec[self._idx, None]) else: ZNXr = 0 - for slc, block in zip(self._slices, self._blocks): - Zblock = Z[slc, :] - Xblock = X[slc, :] + for idx, block in zip(self._idxs, self._blocks): + Zblock = Z[idx, :] + Xblock = X[idx, :] - if slc.stop - slc.start > 1: - cf = sl.cho_factor(block + np.diag(self._nvec[slc])) + if len(idx) > 1: + cf = sl.cho_factor(block + np.diag(self._nvec[idx])) bx = sl.cho_solve(cf, Xblock) else: - bx = Xblock / self._nvec[slc][:, None] + bx = Xblock / self._nvec[idx][:, None] ZNX += np.dot(Zblock.T, bx) ZNX += ZNXr return ZNX.squeeze() if len(ZNX) > 1 else float(ZNX) @@ -1173,11 +1175,11 @@ def _solve_NX(self, X): X = X.reshape(X.shape[0], 1) NX = X / self._nvec[:, None] - for slc, block in zip(self._slices, self._blocks): - Xblock = X[slc, :] - if slc.stop - slc.start > 1: - cf = sl.cho_factor(block + np.diag(self._nvec[slc])) - NX[slc] = sl.cho_solve(cf, Xblock) + for idx, block in zip(self._idxs, self._blocks): + Xblock = X[idx, :] + if len(idx) > 1: + cf = sl.cho_factor(block + np.diag(self._nvec[idx])) + NX[idx] = sl.cho_solve(cf, Xblock) return NX.squeeze() def _get_logdet(self): @@ -1188,12 +1190,12 @@ def _get_logdet(self): logdet = np.sum(np.log(self._nvec[self._idx])) else: logdet = 0 - for slc, block in zip(self._slices, self._blocks): - if slc.stop - slc.start > 1: - cf = sl.cho_factor(block + np.diag(self._nvec[slc])) + for idx, block in zip(self._idxs, self._blocks): + if len(idx) > 1: + cf = sl.cho_factor(block + np.diag(self._nvec[idx])) logdet += np.sum(2 * np.log(np.diag(cf[0]))) else: - logdet += np.sum(np.log(self._nvec[slc])) + logdet += np.sum(np.log(self._nvec[idx])) return logdet def solve(self, other, left_array=None, logdet=False): @@ -1218,6 +1220,7 @@ class ShermanMorrison(object): def __init__(self, jvec, slices, nvec=0.0): self._jvec = jvec self._slices = slices + self._idxs = [indices_from_slice(slc) for slc in slices] self._nvec = nvec def __add__(self, other): @@ -1235,12 +1238,12 @@ def _solve_D1(self, x): """Solves :math:`N^{-1}x` where :math:`x` is a vector.""" Nx = x / self._nvec - for slc, jv in zip(self._slices, self._jvec): - if slc.stop - slc.start > 1: - rblock = x[slc] - niblock = 1 / self._nvec[slc] + for idx, jv in zip(self._idxs, self._jvec): + if len(idx) > 1: + rblock = x[idx] + niblock = 1 / self._nvec[idx] beta = 1.0 / (np.einsum("i->", niblock) + 1.0 / jv) - Nx[slc] -= beta * np.dot(niblock, rblock) * niblock + Nx[idx] -= beta * np.dot(niblock, rblock) * niblock return Nx def _solve_1D1(self, x, y): @@ -1250,11 +1253,11 @@ def _solve_1D1(self, x, y): Nx = x / self._nvec yNx = np.dot(y, Nx) - for slc, jv in zip(self._slices, self._jvec): - if slc.stop - slc.start > 1: - xblock = x[slc] - yblock = y[slc] - niblock = 1 / self._nvec[slc] + for idx, jv in zip(self._idxs, self._jvec): + if len(idx) > 1: + xblock = x[idx] + yblock = y[idx] + niblock = 1 / self._nvec[idx] beta = 1.0 / (np.einsum("i->", niblock) + 1.0 / jv) yNx -= beta * np.dot(niblock, xblock) * np.dot(niblock, yblock) return yNx @@ -1265,11 +1268,11 @@ def _solve_2D2(self, X, Z): """ ZNX = np.dot(Z.T / self._nvec, X) - for slc, jv in zip(self._slices, self._jvec): - if slc.stop - slc.start > 1: - Zblock = Z[slc, :] - Xblock = X[slc, :] - niblock = 1 / self._nvec[slc] + for idx, jv in zip(self._idxs, self._jvec): + if len(idx) > 1: + Zblock = Z[idx, :] + Xblock = X[idx, :] + niblock = 1 / self._nvec[idx] beta = 1.0 / (np.einsum("i->", niblock) + 1.0 / jv) zn = np.dot(niblock, Zblock) xn = np.dot(niblock, Xblock) @@ -1281,9 +1284,9 @@ def _get_logdet(self): is a quantization matrix. """ logdet = np.einsum("i->", np.log(self._nvec)) - for slc, jv in zip(self._slices, self._jvec): - if slc.stop - slc.start > 1: - niblock = 1 / self._nvec[slc] + for idx, jv in zip(self._idxs, self._jvec): + if len(idx) > 1: + niblock = 1 / self._nvec[idx] beta = 1.0 / (np.einsum("i->", niblock) + 1.0 / jv) logdet += np.log(jv) - np.log(beta) return logdet diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py index df949af0..eff440e1 100644 --- a/enterprise/signals/utils.py +++ b/enterprise/signals/utils.py @@ -767,26 +767,37 @@ def create_quantization_matrix(toas, dt=1, nmin=2): return U, weights -def quant2ind(U): +def quant2ind(U, as_slice=False): """ - Use quantization matrix to return slices of non-zero elements. + Use quantization matrix to return indices of non-zero elements. :param U: quantization matrix + :param as_slice: whether to return a slice object - :return: list of `slice`s for non-zero elements of U + :return: list of `slice`s or indices for non-zero elements of U - .. note:: This function assumes that the pulsar TOAs were sorted by time. + .. note:: For slice objects the TOAs need to be sorted by time """ inds = [] for cc, col in enumerate(U.T): epinds = np.flatnonzero(col) - if epinds[-1] - epinds[0] + 1 != len(epinds): - raise ValueError("ERROR: TOAs not sorted properly!") - inds.append(slice(epinds[0], epinds[-1] + 1)) + if epinds[-1] - epinds[0] + 1 != len(epinds) or not as_slice: + inds.append(epinds) + else: + inds.append(slice(epinds[0], epinds[-1] + 1)) return inds +def indices_from_slice(slc): + """Given a slice object, return an index arrays""" + + if isinstance(slc, np.ndarray): + return slc + else: + return np.arange(*slc.indices(slc.stop)) + + def linear_interp_basis(toas, dt=30 * 86400): """Provides a basis for linear interpolation. diff --git a/enterprise/signals/white_signals.py b/enterprise/signals/white_signals.py index a1df6f86..9c318b8e 100644 --- a/enterprise/signals/white_signals.py +++ b/enterprise/signals/white_signals.py @@ -11,6 +11,7 @@ from enterprise.signals import parameter, selections, signal_base, utils from enterprise.signals.parameter import function from enterprise.signals.selections import Selection +from enterprise.signals.utils import indices_from_slice try: import fastshermanmorrison.fastshermanmorrison as fastshermanmorrison @@ -217,6 +218,7 @@ def __init__(self, psr): nepoch = sum(U.shape[1] for U in Umats) U = np.zeros((len(psr.toas), nepoch)) self._slices = {} + self._idxs = {} netot = 0 for ct, (key, mask) in enumerate(zip(keys, masks)): nn = Umats[ct].shape[1] @@ -224,6 +226,10 @@ def __init__(self, psr): self._slices.update({key: utils.quant2ind(U[:, netot : nn + netot])}) netot += nn + self._idxs.update( + {key: [indices_from_slice(slc) for slc in slices] for (key, slices) in self._slices.items()} + ) + # initialize sparse matrix self._setup(psr) @@ -252,17 +258,17 @@ def _setup(self, psr): def _setup_sparse(self, psr): Ns = scipy.sparse.csc_matrix((len(psr.toas), len(psr.toas))) - for key, slices in self._slices.items(): - for slc in slices: - if slc.stop - slc.start > 1: - Ns[slc, slc] = 1.0 + for key, idxs in self._idxs.items(): + for idx in idxs: + if len(idx) > 1: + Ns[np.ix_(idx, idx)] = 1.0 self._Ns = signal_base.csc_matrix_alt(Ns) def _get_ndiag_sparse(self, params): for p in self._params: - for slc in self._slices[p]: - if slc.stop - slc.start > 1: - self._Ns[slc, slc] = 10 ** (2 * self.get(p, params)) + for idx in self._idxs[p]: + if len(idx) > 1: + self._Ns[np.ix_(idx, idx)] = 10 ** (2 * self.get(p, params)) return self._Ns def _get_ndiag_sherman_morrison(self, params): @@ -274,21 +280,18 @@ def _get_ndiag_fast_sherman_morrison(self, params): return fastshermanmorrison.FastShermanMorrison(jvec, slices) def _get_ndiag_block(self, params): - slices, jvec = self._get_jvecs(params) + idxs, jvec = self._get_jvecs(params) blocks = [] - for jv, slc in zip(jvec, slices): - nb = slc.stop - slc.start + for jv, idx in zip(jvec, idxs): + nb = len(idx) blocks.append(np.ones((nb, nb)) * jv) - return signal_base.BlockMatrix(blocks, slices) + return signal_base.BlockMatrix(blocks, idxs) def _get_jvecs(self, params): - slices = sum([self._slices[key] for key in sorted(self._slices.keys())], []) + idxs = sum([self._idxs[key] for key in sorted(self._idxs.keys())], []) jvec = np.concatenate( - [ - np.ones(len(self._slices[key])) * 10 ** (2 * self.get(key, params)) - for key in sorted(self._slices.keys()) - ] + [np.ones(len(self._idxs[key])) * 10 ** (2 * self.get(key, params)) for key in sorted(self._idxs.keys())] ) - return (slices, jvec) + return (idxs, jvec) return EcorrKernelNoise diff --git a/tests/test_utils.py b/tests/test_utils.py index f8268cf3..a07af86e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -128,6 +128,27 @@ def test_quantization_matrix(self): assert U.shape == (4005, 235), msg1 assert all(np.sum(U, axis=0) > 1), msg2 + inds = utils.quant2ind(U, as_slice=False) + slcs = utils.quant2ind(U, as_slice=True) + inds_check = [utils.indices_from_slice(slc) for slc in slcs] + + msg3 = "Quantization Matrix slice not equal to quantization indices" + for ind, ind_c in zip(inds, inds_check): + assert np.all(ind == ind_c), msg3 + + def test_indices_from_slice(self): + """Test conversion of slices to numpy indices""" + ind_np = np.array([2, 4, 6, 8]) + ind_np_check = utils.indices_from_slice(ind_np) + + msg1 = "Numpy indices not left as-is by indices_from_slice" + assert np.all(ind_np == ind_np_check), msg1 + + slc = slice(2, 10, 2) + ind_np_check = utils.indices_from_slice(slc) + msg2 = "Slice not converted properly by indices_from_slice" + assert np.all(ind_np == ind_np_check), msg2 + def test_psd(self): """Test PSD functions.""" Tmax = self.psr.toas.max() - self.psr.toas.min() diff --git a/tests/test_white_signals.py b/tests/test_white_signals.py index 7f67d272..5ef2c33b 100644 --- a/tests/test_white_signals.py +++ b/tests/test_white_signals.py @@ -58,7 +58,14 @@ def setUpClass(cls): cls.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim") # IPTA-like pulsar - cls.ipsr = Pulsar(datadir + "/1713.Sep.T2.par", datadir + "/1713.Sep.T2.tim") + cls.ipsr = Pulsar(datadir + "/1713.Sep.T2.par", datadir + "/1713.Sep.T2.tim", sort=True) + + # Same pulsar, but with TOAs shuffled + cls.ipsr_shuffled = Pulsar(datadir + "/1713.Sep.T2.par", datadir + "/1713.Sep.T2.tim", sort=True) + rng = np.random.default_rng(seed=123) + rng.shuffle(cls.ipsr_shuffled._isort) + for ii, p in enumerate(cls.ipsr_shuffled._isort): + cls.ipsr_shuffled._iisort[p] = ii def test_efac(self): """Test that efac signal returns correct covariance.""" @@ -384,8 +391,13 @@ def _ecorr_test(self, method="sparse"): msg = "EFAC/ECORR {} 2D2 solve incorrect.".format(method) assert np.allclose(N.solve(T, left_array=T), np.dot(T.T, wd.solve(T)), rtol=1e-10), msg - def _ecorr_test_ipta(self, method="sparse"): + def _ecorr_test_ipta(self, method="sparse", shuffled=False): """Test of sparse/sherman-morrison ecorr signal and solve methods.""" + if shuffled: + ipsr = self.ipsr_shuffled + else: + ipsr = self.ipsr + selection = Selection(selections.nanograv_backends) efac = parameter.Uniform(0.1, 5) @@ -394,7 +406,7 @@ def _ecorr_test_ipta(self, method="sparse"): ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection, method=method) tm = gp_signals.TimingModel() s = ef + ec + tm - m = s(self.ipsr) + m = s(ipsr) # set parameters efacs = [1.3] @@ -412,18 +424,18 @@ def _ecorr_test_ipta(self, method="sparse"): } # get EFAC Nvec - nvec0 = efacs[0] ** 2 * self.ipsr.toaerrs**2 + nvec0 = efacs[0] ** 2 * ipsr.toaerrs**2 # get the basis flags = ["ASP-L", "ASP-S", "GASP-8", "GASP-L", "GUPPI-8", "GUPPI-L", "PUPPI-L", "PUPPI-S"] - bflags = self.ipsr.backend_flags + bflags = ipsr.backend_flags Umats = [] for flag in np.unique(bflags): if flag in flags: mask = bflags == flag - Umats.append(utils.create_quantization_matrix(self.ipsr.toas[mask], nmin=2)[0]) + Umats.append(utils.create_quantization_matrix(ipsr.toas[mask], nmin=2)[0]) nepoch = sum(U.shape[1] for U in Umats) - U = np.zeros((len(self.ipsr.toas), nepoch)) + U = np.zeros((len(ipsr.toas), nepoch)) jvec = np.zeros(nepoch) netot, ct = 0, 0 for flag in np.unique(bflags): @@ -441,23 +453,21 @@ def _ecorr_test_ipta(self, method="sparse"): # test msg = "EFAC/ECORR {} logdet incorrect.".format(method) N = m.get_ndiag(params) - assert np.allclose(N.solve(self.ipsr.residuals, logdet=True)[1], wd.logdet(), rtol=1e-8), msg + assert np.allclose(N.solve(ipsr.residuals, logdet=True)[1], wd.logdet(), rtol=1e-8), msg msg = "EFAC/ECORR {} D1 solve incorrect.".format(method) - assert np.allclose(N.solve(self.ipsr.residuals), wd.solve(self.ipsr.residuals), rtol=1e-8), msg + assert np.allclose(N.solve(ipsr.residuals), wd.solve(ipsr.residuals), rtol=1e-8), msg msg = "EFAC/ECORR {} 1D1 solve incorrect.".format(method) assert np.allclose( - N.solve(self.ipsr.residuals, left_array=self.ipsr.residuals), - np.dot(self.ipsr.residuals, wd.solve(self.ipsr.residuals)), + N.solve(ipsr.residuals, left_array=ipsr.residuals), + np.dot(ipsr.residuals, wd.solve(ipsr.residuals)), rtol=1e-8, ), msg msg = "EFAC/ECORR {} 2D1 solve incorrect.".format(method) T = m.get_basis() - assert np.allclose( - N.solve(self.ipsr.residuals, left_array=T), np.dot(T.T, wd.solve(self.ipsr.residuals)), rtol=1e-8 - ), msg + assert np.allclose(N.solve(ipsr.residuals, left_array=T), np.dot(T.T, wd.solve(ipsr.residuals)), rtol=1e-8), msg msg = "EFAC/ECORR {} 2D2 solve incorrect.".format(method) assert np.allclose(N.solve(T, left_array=T), np.dot(T.T, wd.solve(T)), rtol=1e-8), msg @@ -480,15 +490,18 @@ def test_ecorr_block(self): def test_ecorr_sparse_ipta(self): """Test of sparse ecorr signal and solve methods.""" - self._ecorr_test_ipta(method="sparse") + self._ecorr_test_ipta(method="sparse", shuffled=False) + self._ecorr_test_ipta(method="sparse", shuffled=True) def test_ecorr_sherman_morrison_ipta(self): """Test of sherman-morrison ecorr signal and solve methods.""" - self._ecorr_test_ipta(method="sherman-morrison") + self._ecorr_test_ipta(method="sherman-morrison", shuffled=False) + self._ecorr_test_ipta(method="sherman-morrison", shuffled=True) def test_ecorr_block_ipta(self): """Test of block matrix ecorr signal and solve methods.""" - self._ecorr_test_ipta(method="block") + self._ecorr_test_ipta(method="block", shuffled=False) + self._ecorr_test_ipta(method="block", shuffled=True) class TestWhiteSignalsPint(TestWhiteSignals): @@ -506,5 +519,14 @@ def setUpClass(cls): # IPTA-like pulsar cls.ipsr = Pulsar( - datadir + "/1713.Sep.T2.par", datadir + "/1713.Sep.T2.tim", ephem="DE421", timint_package="pint" + datadir + "/1713.Sep.T2.par", datadir + "/1713.Sep.T2.tim", ephem="DE421", timint_package="pint", sort=True + ) + + # Same pulsar, but with TOAs shuffled + cls.ipsr_shuffled = Pulsar( + datadir + "/1713.Sep.T2.par", datadir + "/1713.Sep.T2.tim", ephem="DE421", timint_package="pint", sort=True ) + rng = np.random.default_rng(seed=123) + rng.shuffle(cls.ipsr_shuffled._isort) + for ii, p in enumerate(cls.ipsr_shuffled._isort): + cls.ipsr_shuffled._iisort[p] = ii