Merge pull request #33 from christianhbye/v3

change how coordinate transforms are done

.gitignore

@@ -128,4 +128,6 @@ dmypy.json
 # Pyre type checker
 .pyre/
 
-TODO.txt
+*.npy
+notebooks/hera_sim_test.ipynb
+notebooks/ulsa_test.ipynb

croissant/__init__.py

@@ -1,4 +1,9 @@
 __author__ = "Christian Hellum Bye"
-__version__ = "2.1.1"
+__version__ = "3.0.0"
 
-from . import beam, constants, dpss, healpix, rotations, sky, simulator
+from . import constants, dpss, sphtransform
+from .healpix import Alm, HealpixMap
+from .beam import Beam
+from .rotations import Rotator
+from .simulator import Simulator
+from .sky import Sky

croissant/beam.py

@@ -1,113 +1,44 @@
-from healpy import nside2npix
+from healpy import npix2nside, pix2ang
 import numpy as np
-import warnings
 
 from .constants import Y00
-from .healpix import grid2healpix
+from .healpix import Alm
+from .sphtransform import map2alm
 
 
-class Beam:
-    def __init__(
-        self,
-        data,
-        theta=None,
-        phi=None,
-        frequencies=None,
-        alm=False,
-        coords="topocentric",
-    ):
-        """
-        Initialize beam object. The beam must be specified at a rectangular
-        grid, that is, theta and phi must be coordinate axis of a grid.
-
-        Parameters
-        ----------
-        data : array-like
-            The power beam. Must have shape ([freqs,] theta, phi).
-        theta : array-like
-            Zenith angle(s) in radians. Must be in [0, pi].
-        phi : array-like
-            Azimuth angle(s) in radians. Must be in [0, 2*pi).
-        frequencies : array-like (optional)
-            The frequencies in MHz of the beam. Necessary if the beam is
-            specified at more than one frequency.
-
-        """
-        data = np.array(data, copy=True)
-        self.frequencies = np.ravel(frequencies).copy()
-        self.nfreqs = self.frequencies.size
-        if not alm:
-            self.theta = np.ravel(theta).copy()
-            self.phi = np.ravel(phi).copy()
-            if self.theta.min() < 0 or self.theta.max() > np.pi:
-                raise ValueError("Theta must be in the range [0, pi].")
-            if self.phi.min() < 0 or self.phi.max() >= 2 * np.pi:
-                raise ValueError("Phi must be in the range [0, 2pi).")
-            if not (
-                np.allclose(np.diff(self.theta), self.theta[1] - self.theta[0])
-                and np.allclose(np.diff(self.phi), self.phi[1] - self.phi[0])
-            ):
-                raise ValueError(
-                    "The data must be sampled on a rectangular grid."
-                )
-            data.shape = (self.nfreqs, theta.size, phi.size)
-            self.data = data
-            self.alm = None
-        else:
-            data.shape = (self.nfreqs, -1)
-            self.data = None
-            self.alm = data
-        self.total_power = self.compute_total_power()  # before horizon cut
-        self.coords = coords
-
-    def compute_total_power(self, nside=128):
+class Beam(Alm):
+    def compute_total_power(self):
         """
         Compute the total integrated power in the beam at each frequency. This
         is a necessary normalization constant for computing the visibilities.
         It should be computed before applying the horizon cut in order to
         account for ground loss.
         """
-        if self.data is not None:  # from grid
-            healpix_beam = grid2healpix(
-                self.data, nside, theta=self.theta, phi=self.phi
-            )
-            npix = nside2npix(nside)
-            power = healpix_beam.sum(axis=-1) * 4 * np.pi / npix
-        else:  # from alm
-            power = self.alm[0, 0] * Y00 * 4 * np.pi
-        return power
+        if self.alm.ndim == 2:
+            a00 = self[:, 0, 0]
+        else:
+            a00 = self[0, 0]
+        power = a00.real * Y00 * 4 * np.pi
+        self.total_power = power
 
-    def horizon_cut(self, horizon=None):
+    def horizon_cut(self, horizon=None, nside=128):
         """
-        horizon : array-like (optional)
-            An array indicating if a given (frequency/)theta/phi combination is
-            above the horizon or not. Must have shape (theta, phi) or
-            (freqs, theta, phi) if frequency dependent. The array will be
-            multiplied by the antenna beam, hence 0 is intepreted as below the
-            horizon and 1 is interpreted as above it. The elements must be in
-            the range [0, 1].
+        horizon : array-like
+        nside : int
+            The resolution of the healpix map for the intermediate step.
         """
-        # has no effect on alm beam
-        if self.data is None:
-            warnings.warn(
-                "Horizon cut is not implemented for alm beams.",
-                UserWarning,
-            )
-            return
-
-        if horizon is None:
-            horizon = np.ones_like(self.data)
-            horizon[:, self.theta > np.pi / 2] = 0
-        else:
-            horizon = np.array(horizon, copy=True)
-            horizon.shape = (-1, self.theta.size, self.phi.size)
+        if horizon is not None:
             if horizon.min() < 0 or horizon.max() > 1:
                 raise ValueError("Horizon elements must be in [0, 1].")
-        self.data = self.data * horizon
+            npix = horizon.shape[-1]
+            nside = npix2nside(npix)
 
-    @classmethod
-    def from_file(path):
-        raise NotImplementedError
+        hp_beam = self.hp_map(nside=nside)
+        if horizon is None:
+            horizon = np.ones_like(hp_beam)
+            npix = horizon.shape[-1]
+            theta = pix2ang(nside, np.arange(npix))[0]
+            horizon[:, theta > np.pi / 2] = 0
 
-    def to_file(fname):
-        raise NotImplementedError
+        hp_beam *= horizon
+        self.alm = map2alm(hp_beam, lmax=self.lmax)

croissant/constants.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 # nside's for which healpix has computed pixel weights:
-PIX_WEIGHTS_NSIDE = [32, 64, 128, 256, 512, 1024, 2048, 4096]
+PIX_WEIGHTS_NSIDE = (32, 64, 128, 256, 512, 1024, 2048, 4096)
 
 # sidereal days in seconds
 # https://nssdc.gsfc.nasa.gov/planetary/factsheet/earthfact.html