diff --git a/src/toast/ops/CMakeLists.txt b/src/toast/ops/CMakeLists.txt index ddff8d27e..7f6138bf7 100644 --- a/src/toast/ops/CMakeLists.txt +++ b/src/toast/ops/CMakeLists.txt @@ -37,6 +37,7 @@ install(FILES scan_healpix_detector.py scan_wcs.py scan_wcs_detector.py + derivatives_weights.py mapmaker_binning.py mapmaker_solve.py mapmaker_templates.py diff --git a/src/toast/ops/__init__.py b/src/toast/ops/__init__.py index c6b2faf24..48a69010b 100644 --- a/src/toast/ops/__init__.py +++ b/src/toast/ops/__init__.py @@ -76,6 +76,7 @@ from .sss import SimScanSynchronousSignal from .statistics import Statistics from .stokes_weights import StokesWeights +from .derivatives_weights import DerivativesWeights from .time_constant import TimeConstant from .totalconvolve import SimTotalconvolve from .weather_model import WeatherModel diff --git a/src/toast/ops/derivatives_weights.py b/src/toast/ops/derivatives_weights.py new file mode 100644 index 000000000..87907c843 --- /dev/null +++ b/src/toast/ops/derivatives_weights.py @@ -0,0 +1,323 @@ +# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file. +# All rights reserved. Use of this source code is governed by +# a BSD-style license that can be found in the LICENSE file. + +import numpy as np +import traitlets +from astropy import units as u + +from ..accelerator import ImplementationType +from ..observation import default_values as defaults +from ..timing import function_timer +from ..traits import Bool, Instance, Int, Unicode, trait_docs +from ..utils import Environment, Logger +from ..qarray import mult, to_iso_angles +from .operator import Operator + +@trait_docs +class DerivativesWeights(Operator): + """Operator which generates pointing weights for I and derivatives of I with + respect to theta and phi, to order 1 (mode dI) or 2 (mode d2I). + + Given the individual detector pointing, this computes the pointing weights + assuming that the detector is a linear polarizer followed by a total + power measurement. By definition, the detector coordinate frame has the X-axis + aligned with the polarization sensitive direction. An optional dictionary of + beam error factors may be specified for each observation. + + These factors are an overall calibration factor cal, beam centroid error dx/dy, + differential beam fwhm dsigma, and ellipticity dp/dc. Since we are focused + on total intensity, there is no HWP term or detector polarisation efficiency. + + The timestream model without a HWP in COSMO convention is: + + .. math:: + d = cal*I + d_\\theta I \\left[ dx\\sin\\psi - dy\\cos\\psi \\right] + + d_\\phi I \\left[ -dx\\cos\\psi - dy\\sin\\psi + (dp\\sin(2\\psi) - dc\\cos(2\\psi))\frac{\\cos\\theta}{\\sin\\theta} \\right] + + d^2_\\theta I \\left[dsigma + dp\\cos(2\\psi) - dc\\sin(2\\psi)\\right] + + d_\\phi d_\\theta I \\left[-2dp\\sin(2\\psi) + 2dc\\cos(2\\psi) \\right] + + d^2_\\phi I \\left[\\ dsigma + dp\\cos(2\\psi) + dc\\sin(2\\psi) \\right] + + The detector orientation angle "psi" in COSMO convention is measured in a + right-handed sense from the local meridian. + + By default, this operator uses the "COSMO" convention for Q/U. If the "IAU" trait + is set to True, then resulting weights will differ as psi will jump around. + + If the view trait is not specified, then this operator will use the same data + view as the detector pointing operator when computing the pointing matrix pixels + and weights. + + """ + + # Class traits + + API = Int(0, help="Internal interface version for this operator") + + detector_pointing = Instance( + klass=Operator, + allow_none=True, + help="Operator that translates boresight pointing into detector frame", + ) + + mode = Unicode("dI", help="The Stokes weights to generate (dI or d2I)") + + view = Unicode( + None, allow_none=True, help="Use this view of the data in all observations" + ) + + weights = Unicode( + defaults.weights, help="Observation detdata key for output weights" + ) + + single_precision = Bool(False, help="If True, use 32bit float in output") + + IAU = Bool(False, help="If True, use the IAU convention rather than COSMO") + + @traitlets.validate("detector_pointing") + def _check_detector_pointing(self, proposal): + detpointing = proposal["value"] + if detpointing is not None: + if not isinstance(detpointing, Operator): + raise traitlets.TraitError( + "detector_pointing should be an Operator instance" + ) + # Check that this operator has the traits we expect + for trt in [ + "view", + "boresight", + "shared_flags", + "shared_flag_mask", + "det_mask", + "quats", + "coord_in", + "coord_out", + ]: + if not detpointing.has_trait(trt): + msg = f"detector_pointing operator should have a '{trt}' trait" + raise traitlets.TraitError(msg) + return detpointing + + @traitlets.validate("mode") + def _check_mode(self, proposal): + check = proposal["value"] + if check not in ["dI", "d2I"]: + raise traitlets.TraitError("Invalid mode (must be 'dI' or 'd2I')") + return check + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + @function_timer + def _exec(self, data, detectors=None, use_accel=None, **kwargs): + env = Environment.get() + log = Logger.get() + if self.mode == "d2I": + self._nnz = 6 + else: + self._nnz = 3 + + # Kernel selection + implementation, use_accel = self.select_kernels(use_accel=use_accel) + + if self.detector_pointing is None: + raise RuntimeError("The detector_pointing trait must be set") + + # Expand detector pointing + quats_name = self.detector_pointing.quats + + view = self.view + if view is None: + # Use the same data view as detector pointing + view = self.detector_pointing.view + + # Expand detector pointing + self.detector_pointing.apply(data, detectors=detectors, use_accel=use_accel) + + for ob in data.obs: + # Get the detectors we are using for this observation + dets = ob.select_local_detectors( + detectors, flagmask=self.detector_pointing.det_mask + ) + if len(dets) == 0: + # Nothing to do for this observation + continue + + # Check that our view is fully covered by detector pointing. If the + # detector_pointing view is None, then it has all samples. If our own + # view was None, then it would have been set to the detector_pointing + # view above. + if (view is not None) and (self.detector_pointing.view is not None): + if ob.intervals[view] != ob.intervals[self.detector_pointing.view]: + # We need to check intersection + intervals = ob.intervals[self.view] + detector_intervals = ob.intervals[self.detector_pointing.view] + intersection = detector_intervals & intervals + if intersection != intervals: + msg = ( + f"view {self.view} is not fully covered by valid " + "detector pointing" + ) + raise RuntimeError(msg) + + # Create (or re-use) output data for the weights + if self.single_precision: + exists = ob.detdata.ensure( + self.weights, + sample_shape=(self._nnz,), + dtype=np.float32, + detectors=dets, + accel=use_accel, + ) + else: + exists = ob.detdata.ensure( + self.weights, + sample_shape=(self._nnz,), + dtype=np.float64, + detectors=dets, + accel=use_accel, + ) + + # Do we already have pointing for all requested detectors? + if exists: + # Yes + if data.comm.group_rank == 0: + msg = ( + f"Group {data.comm.group}, ob {ob.name}, derivative weights " + f"already computed for {dets}" + ) + log.verbose(msg) + continue + + # FIXME: temporary hack until instrument classes are also pre-staged + # to GPU + focalplane = ob.telescope.focalplane + focalplane.detector_data + #Get the boresight pointing + qbore = ob.shared["boresight_radec"] + nsamp = len(qbore) + ndets = len(dets) + theta = np.empty(nsamp) + psi = np.empty(nsamp) + # Get the per-detector pointing for orientation/sine theta purposes + for idet, d in enumerate(dets): + theta, _, psi = to_iso_angles(mult(qbore, focalplane[d]["quat"])) + psi -= focalplane[d]["pol_angle"].value + wc = np.cos(psi) + wc2 = np.cos(2*psi) + ws = np.sin(psi) + ws2 = np.sin(2*psi) + inv_tan_theta = np.cos(theta)/np.sin(theta) + # Get the per-detector calibration. For now we sidestep with the Nones + cal = focalplane[d].get("cal", 1.0) + fwhm = focalplane[d].get("FWHM", 1.4*u.arcmin).to(u.rad).value + dx = focalplane[d].get("dx", 0.0) + dy = focalplane[d].get("dy", 0.0) + dsigma = focalplane[d].get("dsigma", 0.0) + dp = focalplane[d].get("dp", 0.0) + dc = focalplane[d].get("dc", 0.0) + + b_std = fwhm/np.sqrt(8*np.log(2)) #beam standard deviation + + weights = np.empty((nsamp, self._nnz)) + weights[:,0] = cal # gain error + weights[:,1] = dx * ws - dy * wc #dtheta + weights[:,2] = -dx * wc - dy * ws + b_std**2 * (dp * ws2 - dc * wc2) * inv_tan_theta #dphi + if self.mode == "d2I": + weights[:,3] = b_std * dsigma + 0.5 * b_std * b_std * (dp * wc2 - dc * ws2) #d2theta + weights[:,4] = b_std**2 * (-2.0 * dp * ws2 + 2.0 * dc * wc2) #dphi dtheta + weights[:,5] = b_std * dsigma + b_std**2 * (dp * wc2 + dc * ws2) #dphi2 + ob.detdata[self.weights][d, :] = weights + return + """ + # Get the per-detector calibration + if self.cal is None: + cal = np.array([1.0 for x in dets], np.float64) + else: + cal = np.array([ob[self.cal][x] for x in dets], np.float64) + cal = np.stack([cal for _ in range(nsamp)], axis=1) + # Per-detector pointing error + if self.dx is None: + dx = np.array([0.0 for x in dets], np.float64) + else: + dx = np.array([ob[self.dx][x] for x in dets], np.float64) + if self.dy is None: + dy = np.array([0.0 for x in dets], np.float64) + else: + dy = np.array([ob[self.dy][x] for x in dets], np.float64) + dx = np.stack([dx for _ in range(nsamp)], axis=1) + dy = np.stack([dy for _ in range(nsamp)], axis=1) + #Per-detector fwhm/sigma error + if self.dsigma is None: + dsigma = np.array([0.0 for x in dets], np.float64) + else: + dsigma = np.array([ob[self.dsigma][x] for x in dets], np.float64) + dsigma = np.stack([dsigma for _ in range(nsamp)], axis=1) + #Per-detector ellipticity + if self.dp is None: + dp = np.array([0.0 for x in dets], np.float64) + else: + dp = np.array([ob[self.dp][x] for x in dets], np.float64) + if self.dc is None: + dc = np.array([0.0 for x in dets], np.float64) + else: + dc = np.array([ob[self.dc][x] for x in dets], np.float64) + dp = np.stack([dp for _ in range(nsamp)], axis=1) + dc = np.stack([dc for _ in range(nsamp)], axis=1) + + wc = np.cos(psi) + wc2 = np.cos(2*psi) + ws = np.sin(psi) + ws2 = np.sin(2*psi) + inv_tan_theta = np.cos(theta)/np.sin(theta) + + weights = np.empty((ndets,nsamp,self._nnz)) + weights[:,:,0] = cal # gain error + weights[:,:,1] = dx * ws - dy * wc #dtheta + weights[:,:,2] = -dx * wc - dy * ws + (dp * ws2 - dc * wc2) * inv_tan_theta #dphi + if self.mode == "d2I": + weights[:,:,3] = dsigma + dp * wc2 - dc * ws2 #d2theta + weights[:,:,4] = -2.0 * dp * ws2 + 2.0 * dc * wc2 #dphi dtheta + weights[:,:,5] = dsigma + dp * wc2 + dc * ws2 #dphi2 + + for idet, d in enumerate(dets): + ob.detdata[self.weights][d, :] = weights[idet] + """ + + + def _finalize(self, data, **kwargs): + return + + def _requires(self): + req = self.detector_pointing.requires() + if "detdata" not in req: + req["detdata"] = list() + req["detdata"].append(self.weights) + if self.cal is not None: + req["meta"].append(self.cal) + if self.hwp_angle is not None: + req["shared"].append(self.hwp_angle) + if self.view is not None: + req["intervals"].append(self.view) + return req + + def _provides(self): + prov = self.detector_pointing.provides() + prov["detdata"].append(self.weights) + return prov + + def _implementations(self): + return [ + ImplementationType.DEFAULT, + ImplementationType.COMPILED, + ImplementationType.NUMPY, + ImplementationType.JAX, + ] + + def _supports_accel(self): + if (self.detector_pointing is not None) and ( + self.detector_pointing.supports_accel() + ): + return True + else: + return False diff --git a/src/toast/ops/mapmaker_utils/mapmaker_utils.py b/src/toast/ops/mapmaker_utils/mapmaker_utils.py index 2dd4db71d..ae354af8e 100644 --- a/src/toast/ops/mapmaker_utils/mapmaker_utils.py +++ b/src/toast/ops/mapmaker_utils/mapmaker_utils.py @@ -441,7 +441,6 @@ def _exec(self, data, detectors=None, use_accel=None, **kwargs): if len(dets) == 0: # Nothing to do for this observation continue - # Check that the noise model exists if self.noise_model not in ob: msg = "Noise model {} does not exist in observation {}".format( @@ -721,6 +720,7 @@ def _exec(self, data, detectors=None, use_accel=None, **kwargs): weight_nnz = zmap.n_value else: weight_nnz = 0 + log.info_rank(str(self.weights), data.comm.comm_world) for ob in data.obs: # Get the detectors we are using for this observation dets = ob.select_local_detectors( diff --git a/src/toast/ops/scan_healpix.py b/src/toast/ops/scan_healpix.py index d956ffcd9..8c6f06833 100644 --- a/src/toast/ops/scan_healpix.py +++ b/src/toast/ops/scan_healpix.py @@ -74,6 +74,12 @@ class ScanHealpixMap(Operator): allow_none=True, help="This must be an instance of a Stokes weights operator", ) + + derivatives_weights = Instance( + klass=Operator, + allow_none=True, + help="This must be an instance of a derivatives weights operator", + ) save_map = Bool(False, help="If True, do not delete map during finalize") @@ -119,7 +125,22 @@ def _check_stokes_weights(self, proposal): msg = f"stokes_weights operator should have a '{trt}' trait" raise traitlets.TraitError(msg) return weights - + + @traitlets.validate("derivatives_weights") + def _check_derivatives_weights(self, proposal): + weights = proposal["value"] + if weights is not None: + if not isinstance(weights, Operator): + raise traitlets.TraitError( + "derivatives_weights should be an Operator instance" + ) + # Check that this operator has the traits we expect + for trt in ["weights", "view"]: + if not weights.has_trait(trt): + msg = f"derivatives_weights operator should have a '{trt}' trait" + raise traitlets.TraitError(msg) + return weights + def __init__(self, **kwargs): self.map_names = [] super().__init__(**kwargs) @@ -155,10 +176,21 @@ def _exec(self, data, detectors=None, **kwargs): dist = data[self.pixel_dist] if not isinstance(dist, PixelDistribution): raise RuntimeError("The pixel_dist must be a PixelDistribution instance") + + # Use the pixel distribution and pointing configuration to allocate our # map data and read it in. - nnz = len(self.stokes_weights.mode) + nnz = None + if self.stokes_weights is None or self.stokes_weights.mode == "I": + nnz = 1 + elif self.stokes_weights.mode in ("IQU", "dI"): + nnz = 3 + elif self.stokes_weights.mode == "d2I": + nnz = 6 + else: + msg = f"Unknown Stokes/Derivatives weights mode '{self.stokes_weights.mode}'" + raise RuntimeError(msg) filenames = self.file.split(";") detdata_keys = self.det_data.split(";") @@ -191,7 +223,6 @@ def _exec(self, data, detectors=None, **kwargs): ) # Configure the low-level map scanning operator - scanner = ScanMap( det_data=self.det_data_keys[0], det_data_units=self.det_data_units, diff --git a/src/toast/tests/CMakeLists.txt b/src/toast/tests/CMakeLists.txt index 747b87758..268200121 100644 --- a/src/toast/tests/CMakeLists.txt +++ b/src/toast/tests/CMakeLists.txt @@ -27,6 +27,7 @@ install(FILES ops_sim_ground.py ops_sim_tod_noise.py ops_stokes_weights.py + ops_derivatives_weights.py ops_common_mode_noise.py ops_sim_tod_dipole.py ops_sim_tod_atm.py diff --git a/src/toast/tests/ops_derivatives_weights.py b/src/toast/tests/ops_derivatives_weights.py new file mode 100644 index 000000000..9e0e848e1 --- /dev/null +++ b/src/toast/tests/ops_derivatives_weights.py @@ -0,0 +1,128 @@ +# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file. +# All rights reserved. Use of this source code is governed by +# a BSD-style license that can be found in the LICENSE file. + +import os + +import healpy as hp +import numpy as np +from astropy import units as u + +from .. import ops as ops +from .. import qarray as qa +from ..instrument_coords import quat_to_xieta +from ..instrument_sim import plot_focalplane +from ..observation import default_values as defaults +from ..pixels import PixelData +from ._helpers import close_data, create_healpix_ring_satellite, create_outdir +from .mpi import MPITestCase + + +class DerivativesWeightsTest(MPITestCase): + def setUp(self): + fixture_name = os.path.splitext(os.path.basename(__file__))[0] + self.outdir = create_outdir(self.comm, fixture_name) + self.nside = 64 + + def healpix_derivatives(self, I_sky): + npix = hp.nside2npix(self.nside) + alm = np.array(hp.map2alm(I_sky), dtype=complex) + theta, phi = hp.pix2ang(self.nside, np.arange(npix)) + _, dtheta, dphi = hp.alm2map_der1(alm, self.nside) #1/sin(theta) dphi + _, d2theta, dphidtheta = hp.alm2map_der1(hp.map2alm(dtheta), self.nside) + _, dthetadphi, d2phi = hp.alm2map_der1(hp.map2alm(dphi), self.nside) #1/sin^2(theta) d2phi + return dtheta, dphi, d2theta, dphidtheta, d2phi + + def create_sky(self, data, pixels, I_sky): + + # Build the pixel distribution + build_dist = ops.BuildPixelDistribution( + pixel_pointing=pixels, + ) + build_dist.apply(data) + ops.Delete(detdata=[pixels.pixels]).apply(data) + + # Create a fake sky with fixed I/dI/d2I values at all pixels. Just + # one submap on all processes. + dist = data["pixel_dist"] + npix = 12*self.nside*self.nside + sm, si = dist.global_pixel_to_submap(np.arange(npix)) + pix_data = PixelData(dist, np.float64, n_value=6, units=u.K) + dtheta, dphi, d2theta, dphidtheta, d2phi = self.healpix_derivatives(I_sky) + #Forgive me for the worse distribution job ever + pix_data.data[sm, si, 0] = I_sky + pix_data.data[sm, si, 1] = dtheta + pix_data.data[sm, si, 2] = dphi + pix_data.data[sm, si, 3] = d2theta + pix_data.data[sm, si, 4] = dphidtheta + pix_data.data[sm, si, 5] = d2phi + return pix_data + + def test_generation(self): + I_sky = np.ones(12*self.nside*self.nside) + #Create sky map and its derivatives and distribute it + + # Create a telescope with 2 boresight pixels per process, so that + # each process can compare 4 detector orientations. The boresight + # pointing will be centered on each pixel exactly once. + data = create_healpix_ring_satellite( + self.comm, pix_per_process=2, nside=self.nside + ) + + #Add a detector calibration dictionary + for ob in data.obs: + detcal = dict() + detdx = dict() + detdy = dict() + detdsigma = dict() + detdc = dict() + detdp = dict() + for det in ob.local_detectors: + detcal[det] = 1.0 + detdx[det] = 0.01 + detdy[det] = -0.01 + detdsigma[det] = 0.03 + detdc[det] = 0.995 + detdp[det] = 1.05 + ob["det_cal"] = detcal + ob["det_dx"] = detdx + ob["det_dy"] = detdy + ob["det_dsigma"] = detdsigma + ob["det_dc"] = detdc + ob["det_dp"] = detdp + + # Pointing operator + detpointing = ops.PointingDetectorSimple() + pixels = ops.PixelsHealpix( + nside=self.nside, + detector_pointing=detpointing, + ) + # Create the input sky and distribute it + data["sky"] = self.create_sky(data, pixels, I_sky) + + weights = ops.DerivativesWeights( + mode="d2I", + detector_pointing=detpointing, + fp_gamma="gamma", + cal="det_cal", + dx="det_dx", + dy="det_dy", + dsigma="det_dsigma", + dp="det_dp", + dc="det_dc", + IAU=False, + ) + + scan_map = ops.ScanMap( + pixels=pixels.pixels, + weights=weights.weights, + map_key="sky", + ) + + # Scan map into timestream + pipe = ops.Pipeline(operators=[pixels, weights, scan_map]) + pipe.apply(data) + + ops.Delete(detdata=[pixels.pixels, weights.weights]).apply(data) + close_data(data) + \ No newline at end of file diff --git a/src/toast/tests/runner.py b/src/toast/tests/runner.py index 5701ca52a..8e47a55c6 100644 --- a/src/toast/tests/runner.py +++ b/src/toast/tests/runner.py @@ -79,6 +79,7 @@ from . import ops_sss as test_ops_sss from . import ops_statistics as test_ops_statistics from . import ops_stokes_weights as test_ops_stokes_weights +from . import ops_derivatives_weights as test_ops_derivatives_weights from . import ops_time_constant as test_ops_time_constant from . import ops_yield_cut as test_ops_yield_cut from . import pixels_healpix as test_pixels_healpix @@ -239,6 +240,7 @@ def test(name=None, verbosity=2): suite.addTest(loader.loadTestsFromModule(test_ops_crosslinking)) suite.addTest(loader.loadTestsFromModule(test_ops_sss)) suite.addTest(loader.loadTestsFromModule(test_ops_stokes_weights)) + suite.addTest(loader.loadTestsFromModule(test_ops_derivatives_weights)) suite.addTest(loader.loadTestsFromModule(test_ops_demodulate)) suite.addTest(loader.loadTestsFromModule(test_ops_demod_common_mode)) suite.addTest(loader.loadTestsFromModule(test_ops_perturbhwp))