From b56c6ea17691e9dc06e9b4a662199e2a752ee08d Mon Sep 17 00:00:00 2001 From: Martin Reinecke Date: Thu, 7 Nov 2024 08:49:05 +0100 Subject: [PATCH] add MuellerConvolver --- litebird_sim/__init__.py | 3 + litebird_sim/mueller_convolver.py | 316 ++++++++++++++++++++++++++++++ 2 files changed, 319 insertions(+) create mode 100644 litebird_sim/mueller_convolver.py diff --git a/litebird_sim/__init__.py b/litebird_sim/__init__.py index d3c3bb3f..69f79b92 100644 --- a/litebird_sim/__init__.py +++ b/litebird_sim/__init__.py @@ -74,6 +74,7 @@ from .madam import save_simulation_for_madam from .mbs.mbs import Mbs, MbsParameters, MbsSavedMapInfo from .mpi import MPI_COMM_WORLD, MPI_ENABLED, MPI_CONFIGURATION +from .mueller_convolver import MuellerConvolver from .noise import ( add_white_noise, add_one_over_f_noise, @@ -221,6 +222,8 @@ def destripe_with_toast2(*args, **kwargs): "MPI_COMM_WORLD", "MPI_ENABLED", "MPI_CONFIGURATION", + # mueller_convolver.py + "MuellerConvolver", # observations.py "Observation", "TodDescription", diff --git a/litebird_sim/mueller_convolver.py b/litebird_sim/mueller_convolver.py new file mode 100644 index 00000000..1101f1cd --- /dev/null +++ b/litebird_sim/mueller_convolver.py @@ -0,0 +1,316 @@ +# fmt: off +# MR: This code is maintained externally, so I'm switching auto-formatting off +# to make updating from outside sources easier. +import ducc0 +import numpy as np + + +__all__ = ["MuellerConvolver"] + + +def nalm(lmax, mmax): + return ((mmax+1)*(mmax+2))//2 + (mmax+1)*(lmax-mmax) + + +# Adri 2020 A25/A35 +def mueller_to_C(mueller): + T = np.zeros((4, 4), dtype=np.complex128) + T[0, 0] = T[3, 3] = 1.0 + T[1, 1] = T[2, 1] = 1.0 / np.sqrt(2.0) + T[1, 2] = 1j / np.sqrt(2.0) + T[2, 2] = -1j / np.sqrt(2.0) + return T.dot(mueller.dot(np.conj(T.T))) + + +def truncate_blm(inp, lmax, kmax, epsilon=1e-10): + limit = epsilon * np.max(np.abs(inp)) + out = [] + for i in range(len(inp)): + maxk = -1 + idx = 0 + for k in range(kmax+1): + if np.max(np.abs(inp[i, :, idx:idx+lmax+1-k])) > limit: + maxk = k + idx += lmax+1-k +# print("component",i,"maxk=",maxk) + if maxk == -1: + out.append(None) + else: + out.append((inp[i, :, : nalm(lmax, maxk)].copy(), maxk)) + return out + + +class MuellerConvolver: + """Class for computing convolutions between arbitrary beams and skies in the + presence of an optical element with arbitrary Mueller matrix in front of the + detector. + + Most of the expressions in this code are derived from + Duivenvoorden et al. 2021, MNRAS 502, 4526 + (https://arxiv.org/abs/2012.10437) + + Parameters + ---------- + lmax : int + maximum l moment of the provided sky and beam a_lm + kmax : int + maximum m moment of the provided beam a_lm + slm : numpy.ndarray((n_comp, n_slm), dtype=complex) + input sky a_lm + ncomp can be 1, 3, or 4, for T, TEB, TEBV components, respectively. + The components have the a_lm format used by healpy + blm : numpy.ndarray((n_comp, n_blm), dtype=complex) + input beam a_lm + ncomp can be 1, 3, or 4, for T, TEB, TEBV components, respectively. + The components have the a_lm format used by healpy + mueller : np.ndarray((4,4), dtype=np.float64) + Mueller matrix of the optical element in front of the detector + single_precision : bool + if True, store internal data in single precision, else double precision + epsilon : float + desired accuracy for the interpolation; a typical value is 1e-4 + npoints : int + total number of irregularly spaced points you want to use this object for + (only used for performance fine-tuning) + sigma_min, sigma_max: float + minimum and maximum allowed oversampling factors + 1.2 <= sigma_min < sigma_max <= 2.5 + nthreads : int + the number of threads to use for computation + """ + + # Very simple class to store a_lm that allow negative m values + class AlmPM: + def __init__(self, lmax, mmax): + if lmax < 0 or mmax < 0 or lmax < mmax: + raise ValueError("bad parameters") + self._lmax, self._mmax = lmax, mmax + self._data = np.zeros((2*mmax+1, lmax+1), dtype=np.complex128) + + def __getitem__(self, lm): + l, m = lm # noqa + + if isinstance(l, slice): + if l.step is not None or l.start < 0 or l.stop-1 > self._lmax: + print(l, m) + raise ValueError("out of bounds read access") + else: + if l < 0 or l > self._lmax: # or abs(m) > l: + print(l, m) + raise ValueError("out of bounds read access") + # if we are asked for elements outside our m range, return 0 + if m < -self._mmax or m > self._mmax: + return 0.0 + 0j + return self._data[m+self._mmax, l] + + def __setitem__(self, lm, val): + l, m = lm # noqa + if isinstance(l, slice): + if l.step is not None or l.start < 0 or l.stop-1 > self._lmax: + print(l, m) + raise ValueError("out of bounds write access") + else: + if l < 0 or l > self._lmax or abs(m) > l or abs(m) > self._mmax: + print(l, m) + raise ValueError("out of bounds write access") + self._data[m+self._mmax, l] = val + + def mueller_tc_prep(self, blm, mueller, lmax, mmax): + ncomp = blm.shape[0] + + # convert input blm to T/P/P*/V blm + blm2 = [self.AlmPM(lmax, mmax+4) for _ in range(4)] + idx = 0 + for m in range(mmax+1): + sign = (-1)**m + lrange = slice(m, lmax+1) + idxrange = slice(idx, idx+lmax+1-m) + # T component + blm2[0][lrange, m] = blm[0, idxrange] + blm2[0][lrange, -m] = np.conj(blm[0, idxrange]) * sign + # V component + if ncomp > 3: + blm2[3][lrange, m] = blm[3, idxrange] + blm2[3][lrange, -m] = np.conj(blm[3, idxrange]) * sign + # E/B components + if ncomp > 2: + # Adri's notes [10], spin +2 + blm2[1][lrange, m] = -blm[1, idxrange] - 1j*blm[2, idxrange] + # Adri's notes [9], spin -2 + blm2[2][lrange, m] = -blm[1, idxrange] + 1j*blm[2, idxrange] + # negative m + # Adri's notes [2] + blm2[1][lrange, -m] = np.conj(blm2[2][lrange, m]) * sign + blm2[2][lrange, -m] = np.conj(blm2[1][lrange, m]) * sign + idx += lmax+1-m + + C = mueller_to_C(mueller) + + # compute the blm for the full beam+Mueller matrix system at angles + # n*pi/5 for n in [0; 5[ + sqrt2 = np.sqrt(2.0) + nbeam = 5 + inc = 4 + res = np.zeros((nbeam, ncomp, nalm(lmax, mmax+inc)), dtype=self._ctype) + blm_eff = [self.AlmPM(lmax, mmax+4) for _ in range(4)] + + for ibeam in range(nbeam): + alpha = ibeam*np.pi/nbeam + e2ia = np.exp(2j*alpha) + e2iac = np.exp(-2j*alpha) + e4ia = np.exp(4j*alpha) + e4iac = np.exp(-4j*alpha) +# FIXME: do I need to calculate anything for negative m? + for m in range(-mmax-4, mmax+4+1): + lrange = slice(abs(m), lmax+1) + # T component, Marta notes [4a] + blm_eff[0][lrange, m] = ( + C[0, 0] * blm2[0][lrange, m] + + C[3, 0] * blm2[3][lrange, m] + + 1.0/sqrt2 * ((C[1, 0]*e2ia ) * blm2[2][lrange, m+2] + +(C[2, 0]*e2iac) * blm2[1][lrange, m-2])) + # E/B components, Marta notes [4b,c] + blm_eff[1][lrange, m] = ( + (sqrt2*e2iac) * (C[0, 1] * blm2[0][lrange, m+2] + + C[3, 1] * blm2[3][lrange, m+2]) + + (C[2, 1]*e4iac) * blm2[2][lrange, m+4] + + C[1, 1] * blm2[1][lrange, m]) + blm_eff[2][lrange, m] = ( + (sqrt2*e2ia) * (C[0, 2] * blm2[0][lrange, m-2] + + C[3, 2] * blm2[3][lrange, m-2]) + + (C[1, 2]*e4ia) * blm2[1][lrange, m-4] + + C[2, 2] * blm2[2][lrange, m]) + # V component, Marta notes [4d] + blm_eff[3][lrange, m] = ( + C[0, 3] * blm2[0][lrange, m] + + C[3, 3] * blm2[3][lrange, m] + + 1.0/sqrt2 * ((C[1, 3]*e2ia ) * blm2[2][lrange, m+2] + +(C[2, 3]*e2iac) * blm2[1][lrange, m-2])) + + # TEMPORARY sanity check ... + def check(xp, xm, lrange, m): + sign = (-1)**m + diff = xp[lrange, m] - sign*np.conj(xm[lrange, -m]) + return np.max(np.abs(diff)) <= 1e-4 + + for m in range(0, mmax+4+1): + sign = (-1)**m + lrange = slice(abs(m), lmax+1) + if not check(blm_eff[0], blm_eff[0], lrange, m): + raise RuntimeError("error T") + if not check(blm_eff[1], blm_eff[2], lrange, m): + raise RuntimeError("error 12") + if not check(blm_eff[2], blm_eff[1], lrange, m): + raise RuntimeError("error 21") + if not check(blm_eff[3], blm_eff[3], lrange, m): + raise RuntimeError("error V") + # ... up to here + + # back to original TEBV b_lm format + idx = 0 + for m in range(mmax+inc+1): + lrange = slice(m, lmax+1) + idxrange = slice(idx, idx+lmax+1-m) + # T component + res[ibeam, 0, idxrange] = blm_eff[0][lrange, m] + # P/P* components + if ncomp > 2: + # Adri's notes [10] + res[ibeam, 1, idxrange] = -0.5 * ( + blm_eff[1][lrange, m] + blm_eff[2][lrange, m] + ) + res[ibeam, 2, idxrange] = 0.5j * ( + blm_eff[1][lrange, m] - blm_eff[2][lrange, m] + ) + # V component + if ncomp > 3: + res[ibeam, 3, idxrange] = blm_eff[3][lrange, m] + idx += lmax+1-m + + return res + + # "Fourier transform" the blm at different alpha to obtain + # blm(alpha) = out[0] + cos(2*alpha)*out[1] + sin(2*alpha)*out[2] + # + cos(4*alpha)*out[3] + sin(4*alpha)*out[4] + def pseudo_fft(self, inp): + out = np.zeros((5, inp.shape[1], inp.shape[2]), dtype=self._ctype) + out[0] = 0.2 * (inp[0]+inp[1]+inp[2]+inp[3]+inp[4]) + # FIXME: I'm not absolutely sure about the sign of the angles yet + c1, s1 = np.cos(2*np.pi/5), np.sin(2*np.pi/5) + c2, s2 = np.cos(4*np.pi/5), np.sin(4*np.pi/5) + out[1] = 0.4 * (inp[0] + c1*(inp[1]+inp[4]) + c2*(inp[2]+inp[3])) + out[2] = 0.4 * ( s1*(inp[1]-inp[4]) + s2*(inp[2]-inp[3])) + out[3] = 0.4 * (inp[0] + c2*(inp[1]+inp[4]) + c1*(inp[2]+inp[3])) + out[4] = 0.4 * ( s2*(inp[1]-inp[4]) - s1*(inp[2]-inp[3])) + # Alternative way via real FFT + # out2 = inp.copy() + # out2 = out2.view(np.float64) + # out2 = ducc0.fft.r2r_fftpack( + # out2, real2hermitian=True, forward=False, axes=(0,), out=out2 + # ) + # out2[0] *=0.2 + # out2[1:] *=0.4 + # out2 = out2.view(np.complex128) + # print(np.max(np.abs(out-out2))) + return out + + def __init__(self, *, lmax, kmax, slm, blm, mueller, single_precision=True, + epsilon=1e-4, npoints=1000000000,sigma_min=1.2, sigma_max=2.5, + nthreads=1): + self._ftype = np.float32 if single_precision else np.float64 + self._ctype = np.complex64 if single_precision else np.complex128 + self._slm = slm.astype(self._ctype) + self._lmax = lmax + self._kmax = kmax + tmp = self.mueller_tc_prep(blm, mueller, lmax, kmax) + tmp = self.pseudo_fft(tmp) + + # construct the five interpolators for the individual components + # All sets of blm are checked up to which kmax they contain significant + # coefficients, and the interpolator is chosen accordingly + tmp = truncate_blm(tmp, self._lmax, self._kmax + 4) + + self._inter = [] + intertype = (ducc0.totalconvolve.Interpolator_f + if self._ctype == np.complex64 + else ducc0.totalconvolve.Interpolator) + for i in range(5): + if tmp[i] is not None: # component is not zero + self._inter.append( + intertype(self._slm, tmp[i][0], False, self._lmax, + tmp[i][1], epsilon=epsilon, npoints=npoints, + sigma_min=sigma_min, sigma_max=sigma_max, + nthreads=nthreads)) + else: # we can ignore this component entirely + self._inter.append(None) + + def signal(self, *, ptg, alpha): + """Computes the convolved signal for a set of pointings and HWP angles. + + Parameters + ---------- + ptg : numpy.ndarray((nptg, 3), dtype=numpy.float32/64) + the input pointings in radians, in (theta, phi, psi) order + alpha : numpy.ndarray((nptg,), dtype=numpy.float32/64) + the HWP angles in radians + + Returns + ------- + signal : numpy.ndarray((nptg,), dtype=numpy.float32/64) + the signal measured by the detector + """ + ptg = ptg.astype(self._ftype) + alpha = alpha.astype(self._ftype) + if self._inter[0] is not None: + res = self._inter[0].interpol(ptg)[0] + else: + res = np.zeros(ptg.shape[0], dtype=self._ftype) + if self._inter[1] is not None: + res += np.cos(2*alpha) * self._inter[1].interpol(ptg)[0] + if self._inter[2] is not None: + res += np.sin(2*alpha) * self._inter[2].interpol(ptg)[0] + if self._inter[3] is not None: + res += np.cos(4*alpha) * self._inter[3].interpol(ptg)[0] + if self._inter[4] is not None: + res += np.sin(4*alpha) * self._inter[4].interpol(ptg)[0] + return res