def test_500_io(self): # Test save/load HDF5 dets = ['det0', 'det1', 'det2'] n, ofs = 1000, 0 aman = core.AxisManager(core.LabelAxis('dets', dets), core.OffsetAxis('samps', n, ofs), core.IndexAxis('indexaxis', 12)) # Make sure this has axes, scalars, a string array ... aman.wrap_new('test1', ('dets', 'samps'), dtype='float32') aman.wrap_new('flag1', shape=('dets', 'samps'), cls=so3g.proj.RangesMatrix.zeros) aman.wrap('scalar', 8) aman.wrap('test_str', np.array(['a', 'b', 'cd'])) aman.wrap('flags', core.FlagManager.for_tod(aman, 'dets', 'samps')) aman.wrap('a', np.int32(12)) aman.wrap('b', np.float32(12.)) aman.wrap('c', np.str_('twelve')) aman.wrap('d', np.bool_(False)) with tempfile.TemporaryDirectory() as tempdir: filename = os.path.join(tempdir, 'test.h5') aman.save(filename, 'my_axisman') aman2 = aman.load(filename, 'my_axisman') shutil.copy(filename, 'debug.h5') # This is not a very satisfying comparison ... support for == # should be required for all AxisManager members! for k in aman._fields.keys(): self.assertEqual(aman[k].__class__, aman2[k].__class__) if hasattr(aman[k], 'shape'): self.assertEqual(aman[k].shape, aman2[k].shape) else: self.assertEqual(aman[k], aman2[k]) # scalar
def get_pca_model(tod=None, pca=None, n_modes=None, signal=None, wrap=None, wrap_pca=None): """Convert a PCA decomposition into the signal basis, i.e. into time-dependent modes that one might use for cleaning or calibrating. The generalization of "common mode" computation is to convert the eigen-decomposition of the covariance matrix into a limited set of modes that explain a high fraction of the power in the input signals. Here we select the strongest eigenmodes from the PCA and compute the time-dependent modes and the projection of each detector signal onto each mode. An approximation for the input signal may then be computed as signal_model = weights . modes where weights has shape (dets, eigen), modes has shape (eigen, samps), and . indicates matrix multipliaction. The size of the eigen axis can be as large as the number of dets, but is typically smaller and in this routine is set by n_modes. Arguments: tod: AxisManager with dets and samps axes. pca: AxisManager holding a PCA decomposition, as returned by get_pca. If not specified, it will be computed from the data. n_modes: integer specifying the number of modes to compute; the strongest modes are taken and this sets the size of the "eigen" axis in the output. Defaults to len(dets), but beware that the resulting data object will be the same size as the input signal. signal: array of shape (dets, samps) that is used to construct the requested eigen modes. If pca is not passed in, this signal is also used to compute the covariance for PCA. wrap: string; if specified then the returned result is also stored in tod under that name. Returns: An AxisManager with (dets, eigen, samps) axes. The field 'weights' has shape (dets, eigen) and the field 'modes' has shape (eigen, samps). """ if pca is None: pca = get_pca(tod=tod, signal=signal) if n_modes is None: n_modes = pca.eigen.count mode_axis = core.IndexAxis('eigen', n_modes) output = core.AxisManager(tod.dets, mode_axis, tod.samps) if signal is None: signal = tod.signal R = pca.R[:, :n_modes] output.wrap('weights', R, [(0, 'dets'), (1, 'eigen')]) output.wrap('modes', np.dot(R.transpose(), signal), [(0, 'eigen'), (1, 'samps')]) return output
def test_100_index(self): a1 = np.zeros(100) a1[10] = 1. aman = core.AxisManager(core.IndexAxis('samps', len(a1))) aman.wrap('a1', a1, [(0, 'samps')]) aman.restrict('samps', (10, 30)) self.assertNotEqual(aman.a1[0], 0.) self.assertEqual(len(aman.a1), 20)
def test_100_index(self): a1 = np.zeros(100) a1[10] = 1. aman = core.AxisManager(core.IndexAxis('samps', len(a1))) aman.wrap('a1', a1, [(0, 'samps')]) # Don't let people wrap the same field twice with self.assertRaises(ValueError): aman.wrap('a1', 2 * a1, [(0, 'samps')]) aman.restrict('samps', (10, 30)) self.assertNotEqual(aman.a1[0], 0.) self.assertEqual(len(aman.a1), 20)
def test_130_not_inplace(self): a1 = np.zeros(100) a1[10] = 1. aman = core.AxisManager(core.IndexAxis('samps', len(a1))) aman.wrap('a1', a1, [(0, 'samps')]) # This should return a separate thing. rman = aman.restrict('samps', (10, 30), in_place=False) #self.assertNotEqual(aman.a1[0], 0.) self.assertEqual(len(aman.a1), 100) self.assertEqual(len(rman.a1), 20) self.assertNotEqual(aman.a1[10], 0.) self.assertNotEqual(rman.a1[0], 0.)
def get_pca(tod=None, cov=None, signal=None, wrap=None): """Compute a PCA decomposition of the kind useful for signal analysis. A symmetric non-negative matrix cov of shape(n_dets, n_dets) can be decomposed into matrix R (same shape) and vector E (length n_dets) such that cov = R . diag(E) . R^T with . denoting matrix multiplication and T denoting matrix transposition. Arguments: tod: AxisManager with dets and samps axes. cov: covariance matrix to decompose; if None then cov is computed from tod.signal (or signal). signal: array of shape (dets, samps). If cov is not provided, it will be computed from this matrix. Defaults to tod.signal. wrap: string; if set then the returned result is also stored in tod under this name. Returns: AxisManager with axes 'dets' and 'eigen' (of the same length), containing fields 'R' of shape (dets, eigen) and 'E' of shape (eigen). The eigenmodes are sorted from strongest to weakest. """ if cov is None: # Compute it from signal if signal is None: signal = tod.signal cov = np.cov(signal) dets = tod.dets mode_axis = core.IndexAxis('eigen', dets.count) output = core.AxisManager(dets, mode_axis) output.wrap('cov', cov, [(0, dets.name), (1, dets.name)]) E, R = np.linalg.eig(cov) # eigh nans sometimes... E[np.isnan(E)] = 0. idx = np.argsort(-E) output.wrap('E', E[idx], [(0, mode_axis.name)]) output.wrap('R', R[:, idx], [(0, dets.name), (1, mode_axis.name)]) return output
def get_tod(sig_type='trendy'): tod = core.AxisManager(core.LabelAxis('dets', ['a', 'b', 'c']), core.IndexAxis('samps', 1000)) tod.wrap_new('signal', ('dets', 'samps'), dtype='float32') tod.wrap_new('timestamps', ('samps', ))[:] = (np.arange(tod.samps.count) / SAMPLE_FREQ_HZ) if sig_type == 'zero': pass elif sig_type == 'trendy': x = np.linspace(0, 1., tod.samps.count) tod.signal[:] = [(i + 1) + (i + 1)**2 * x for i in range(tod.dets.count)] elif sig_type == 'white': tod.signal = np.random.normal(size=tod.shape) elif sig_type == 'red': tod.signal = np.random.normal(size=tod.shape) tod.signal[:] = np.cumsum(tod.signal, axis=1) else: raise RuntimeError(f'sig_type={sig_type}?') return tod
def get_trends(tod, remove=False, size=1, signal=None): """Computes trends for each detector signal that remove the slope connecting first and last points, as well as the mean of the signal. The returned object can be treated like PCA model (e.g., it can be passed as the model input to add_model). Arguments: tod: AxisManager with dets and samps axes. remove: boolean, if True then the computed trends (and means) are removed from the signal. size: the number of samples on each end of the signal to use for trend level computation. Defaults to 1. signal: array of shape (dets, samps) to compute trends on. Defaults to tod.signal. Returns: An AxisManager with (dets, eigen, samps) axes. The field 'weights' has shape (dets, eigen) and the field 'modes' has shape (eigen, samps). There are two modes, which always have the same form: index 0 is all ones, and index1 is a smooth line from -0.5 to +0.5. """ if signal is None: signal = tod.signal trends = core.AxisManager(tod.dets, core.IndexAxis('eigen', 2), tod.samps) modes = np.ones((trends.eigen.count, trends.samps.count)) modes[1] = np.linspace(-0.5, 0.5, modes.shape[1]) weights = np.empty((trends.dets.count, trends.eigen.count)) weights[:, 0] = signal.mean(axis=1) size = max(1, min(size, signal.shape[1] // 2)) weights[:, 1] = (signal[:, -size:].mean(axis=1) - signal[:, :size].mean(axis=1)) trends.wrap('modes', modes) trends.wrap('weights', weights) if remove: add_model(tod, trends, scale=-1, signal=signal) return trends