Switch to Dask

.gitignore

@@ -1,3 +1,6 @@
+# setuptools_scm
+**/_version.py
+
 .DS_Store
 **/.DS_Store
 

maria/atmosphere/__init__.py

@@ -17,7 +17,7 @@
 here, this_filename = os.path.split(__file__)
 
 SPECTRA_DATA_DIRECTORY = f"{here}/data"
-SPECTRA_DATA_CACHE_DIRECTORY = "/tmp/maria_data_cache/spectra"
+SPECTRA_DATA_CACHE_DIRECTORY = "/tmp/maria/spectra"
 SPECTRA_DATA_URL_BASE = (
     "https://github.com/thomaswmorris/maria-data/raw/master/spectra"  # noqa F401
 )
@@ -93,7 +93,7 @@ def _initialize_atmosphere(self):
         This assume that BaseSimulation.__init__() has been called.
         """
 
-        validate_pointing(self.det_coords.az, self.det_coords.el)
+        validate_pointing(self.coords.az, self.coords.el)
 
         if self.atmosphere_model == "2d":
             self.turbulent_layer_depths = np.linspace(
@@ -147,12 +147,18 @@ def _simulate_2d_turbulence(self):
             tqdm(self.atmosphere.layers) if self.verbose else self.atmosphere.layers
         )
         for layer_index, layer in enumerate(layers):
+            layer.generate()
+            layer_data[layer_index] = sp.interpolate.interp1d(
+                layer.sim_time,
+                layer.sample(),
+                axis=-1,
+                bounds_error=False,
+                fill_value="extrapolate",
+            )(self.boresight.time)
+
             if self.verbose:
                 layers.set_description(f"Generating atmosphere (z={layer.depth:.00f}m)")
 
-            layer.generate()
-            layer_data[layer_index] = layer.sample()
-
         return layer_data
 
     def _simulate_2d_atmospheric_fluctuations(self):
@@ -164,7 +170,7 @@ def _simulate_2d_atmospheric_fluctuations(self):
 
         rel_layer_scaling = np.interp(
             self.site.altitude
-            + self.turbulent_layer_depths[:, None, None] * np.sin(self.det_coords.el),
+            + self.turbulent_layer_depths[:, None, None] * np.sin(self.coords.el),
             self.atmosphere.weather.altitude_levels,
             self.atmosphere.weather.absolute_humidity,
         )
@@ -218,7 +224,7 @@ def _simulate_atmospheric_emission(self, units="K_RJ"):
                     (
                         self.zenith_scaled_pwv[band_index],
                         self.atmosphere.weather.temperature[0],
-                        np.degrees(self.det_coords.el[band_index]),
+                        np.degrees(self.coords.el[band_index]),
                     )
                 )
 
@@ -272,7 +278,7 @@ def _simulate_atmospheric_emission(self, units="K_RJ"):
                     (
                         self.zenith_scaled_pwv[band_index],
                         self.atmosphere.weather.temperature[0],
-                        np.degrees(self.det_coords.el[band_index]),
+                        np.degrees(self.coords.el[band_index]),
                     )
                 )
 

maria/atmosphere/turbulent_layer.py

@@ -29,15 +29,21 @@ def __init__(
         weather: Weather,
         depth: float,
         res: float,
-        turbulent_outer_scale: float,
+        turbulent_outer_scale: float = 500,
         verbose: bool = False,
+        timestep: float = 0.1,
         **kwargs,
     ):
         self.instrument = instrument
-        self.boresight = boresight
+        self.boresight = boresight.downsample(timestep=timestep)
         self.weather = weather
         self.depth = depth
         self.res = res
+        self.timestep = timestep
+
+        self.sim_time = self.boresight.time.compute()
+        self.sim_az = self.boresight.az.compute()
+        self.sim_el = self.boresight.el.compute()
 
         # this is approximately correct
         # self.depth = self.depth / np.sin(np.mean(self.plan.el))
@@ -54,9 +60,7 @@ def __init__(
         if verbose:
             print(f"{self.angular_resolution = }")
 
-        self.layer_altitude = self.weather.altitude + self.depth / np.sin(
-            self.boresight.el
-        )
+        self.layer_altitude = self.weather.altitude + self.depth / np.sin(self.sim_el)
 
         layer_wind_north = sp.interpolate.interp1d(
             self.weather.altitude_levels, self.weather.wind_north, axis=0
@@ -66,39 +70,35 @@ def __init__(
         )(self.layer_altitude)
 
         angular_velocity_x = (
-            +layer_wind_east * np.cos(self.boresight.az)
-            - layer_wind_north * np.sin(self.boresight.az)
+            +layer_wind_east * np.cos(self.sim_az)
+            - layer_wind_north * np.sin(self.sim_az)
         ) / self.depth
 
         angular_velocity_y = (
-            -layer_wind_east * np.sin(self.boresight.az)
-            + layer_wind_north * np.cos(self.boresight.az)
+            -layer_wind_east * np.sin(self.sim_az)
+            + layer_wind_north * np.cos(self.sim_az)
         ) / self.depth
 
         if verbose:
             print(f"{(layer_wind_east, layer_wind_north) = }")
 
         # compute the offset with respect to the center of the scan
-        center_az, center_el = get_center_phi_theta(
-            self.boresight.az, self.boresight.el
-        )
+        center_az, center_el = get_center_phi_theta(self.sim_az, self.sim_el)
         dx, dy = self.boresight.offsets(frame="az_el", center=(center_az, center_el))
 
-        # the angular position of each detector over time WRT the atmosphere
-        # this has dimensions (det index, time index)
+        # the angular position of the boresight over time WRT the atmosphere
+        # this has dimensions (2, time index)
         self.boresight_angular_position = np.c_[
-            dx
-            + np.cumsum(angular_velocity_x * np.gradient(self.boresight.time), axis=-1),
-            dy
-            + np.cumsum(angular_velocity_y * np.gradient(self.boresight.time), axis=-1),
+            dx + np.cumsum(angular_velocity_x * np.gradient(self.sim_time)),
+            dy + np.cumsum(angular_velocity_y * np.gradient(self.sim_time)),
         ]
 
         # find the detector offsets which form a convex hull
         self.detector_offsets = np.radians(
             np.c_[self.instrument.sky_x, self.instrument.sky_y]
         )
 
-        # add a small circle of offsets to account for pesky zeros
+        # add a small circle of offsets to account for the beams
         unit_circle_complex = np.exp(1j * np.linspace(0, 2 * np.pi, 64 + 1)[:-1])
         unit_circle_offsets = np.c_[
             np.real(unit_circle_complex), np.imag(unit_circle_complex)
@@ -112,25 +112,25 @@ def __init__(
                 * unit_circle_offsets[None]
             ).reshape(-1, 2)
         )
-        stare_convex_hull_points = stare_convex_hull.points.reshape(-1, 2)[
+        stare_convex_hull_points = stare_convex_hull.points[
             stare_convex_hull.vertices
-        ]
+        ].reshape(-1, 2)
 
         # convex hull downsample index, to get to 1 second
         chds_index = [
             *np.arange(
                 0,
-                len(self.boresight.time),
-                int(np.maximum(np.gradient(self.boresight.time).mean(), 1)),
+                len(self.sim_time),
+                int(np.maximum(np.median(np.diff(self.sim_time)), 1)),
             ),
             -1,
         ]
 
         # this is a convex hull for the atmosphere
         atmosphere_hull = sp.spatial.ConvexHull(
             (
-                stare_convex_hull_points[None, :, None]
-                + self.boresight_angular_position[None, chds_index]
+                self.boresight_angular_position[chds_index, None]
+                + stare_convex_hull_points[None]
             ).reshape(-1, 2)
         )
         self.atmosphere_hull_points = atmosphere_hull.points.reshape(-1, 2)[

maria/atmosphere/weather.py

@@ -15,7 +15,7 @@
 here, this_filename = os.path.split(__file__)
 
 WEATHER_DATA_DIRECTORY = f"{here}/data"
-WEATHER_DATA_CACHE_DIRECTORY = "/tmp/maria_data_cache/weather"
+WEATHER_DATA_CACHE_DIRECTORY = "/tmp/maria/weather"
 WEATHER_DATA_URL_BASE = (
     "https://github.com/thomaswmorris/maria-data/raw/master/weather"  # noqa F401
 )

maria/cmb/__init__.py

@@ -1,64 +1,77 @@
-class CMB:
-    ...
-
-
-class CMBMixin:
-    ...
+import healpy as hp
+import numpy as np
+from tqdm import tqdm
 
+from maria.constants import T_CMB
+from maria.utils.functions import planck_spectrum
 
-# class CMBSimulation(base.BaseSimulation):
-#     """
-#     This simulates scanning over celestial sources.
-#     """
-#     def __init__(self, instrument, plan, site, map_file, add_cmb=False, **kwargs):
-#         super().__init__(instrument, plan, site)
+from ..utils.io import fetch_cache
 
-#         pass
+PLANCK_URL = """https://irsa.ipac.caltech.edu/data/Planck/release_3/all-sky-maps/maps/
+                component-maps/cmb/COM_CMB_IQU-143-fgsub-sevem_2048_R3.00_full.fits"""
 
-#     def _get_CMBPS(self):
 
-#         import camb
-
-#         pars = camb.CAMBparams()
-#         pars.set_cosmology(H0=67.5, ombh2=0.022, omch2=0.122, mnu=0.06, omk=0, tau=0.06)
-#         pars.InitPower.set_params(As=2e-9, ns=0.965, r=0)
-
-#         # correct mode would l=129600 for 5"
-#         pars.set_for_lmax(5000, lens_potential_accuracy=0)
+class CMB:
+    def __init__(self, source="planck", nu=143):
+        planck_url = PLANCK_URL
+        cache_path = "/tmp/maria/cmb/planck/" + planck_url.split("/")[-1]
+
+        fetch_cache(source_url=planck_url, cache_path=cache_path)
+
+        field_dtypes = {
+            "I": np.float32,
+            "Q": np.float32,
+            "U": np.float32,
+            "I_mask": bool,
+            "P_mask": bool,
+        }
+
+        for i, (field, dtype) in tqdm(
+            enumerate(field_dtypes.items()), desc="Loading CMB"
+        ):
+            setattr(
+                self, field, hp.fitsfunc.read_map(cache_path, field=i).astype(dtype)
+            )
+
+        self.source = source
+        self.nside = 2048
+        self.nu = nu
+
+    def plot(self):
+        vmin, vmax = 1e6 * np.quantile(self.I[self.I_mask], q=[0.001, 0.999])
+        hp.visufunc.mollview(
+            1e6 * np.where(self.I_mask, self.I, np.nan), min=vmin, max=vmax, cmap="cmb"
+        )
 
-#         results = camb.get_results(pars)
-#         powers = results.get_cmb_power_spectra(pars, CMB_unit="K")["total"][:, 0]
 
-#         self.CMB_PS = np.empty((len(self.instrument.ubands), len(powers)))
-#         for i in range(len(self.instrument.ubands)):
-#             self.CMB_PS[i] = powers
+class CMBMixin:
+    def _initialize_cmb(self, source="planck", nu=150):
+        self.cmb = CMB(source=source, nu=nu)
 
+    def _simulate_cmb_emission(self):
+        pixel_index = hp.ang2pix(
+            nside=self.cmb.nside, phi=self.coords.l, theta=np.pi / 2 - self.coords.b
+        ).compute()
+        cmb_temperatures = self.cmb.I[pixel_index]
 
-#     def _cmb_imager(self, bandnumber=0):
+        test_nu = np.linspace(1e9, 1e12, 1024)
 
-#         import pymaster as nmt
+        cmb_temperature_samples_K = T_CMB + np.linspace(
+            self.cmb.I.min(), self.cmb.I.max(), 64
+        )
+        cmb_brightness = planck_spectrum(test_nu, cmb_temperature_samples_K[:, None])
 
-#         nx, ny = self.map_data[0].shape
-#         Lx = nx * np.deg2rad(self.sky_data["incell"])
-#         Ly = ny * np.deg2rad(self.sky_data["incell"])
+        self.data["cmb"] = np.zeros((self.instrument.dets.n, self.plan.n_time))
 
-#         self.CMB_map = nmt.synfast_flat(
-#             nx,
-#             ny,
-#             Lx,
-#             Ly,
-#             np.array([self.CMB_PS[bandnumber]]),
-#             [0],
-#             beam=None,
-#             seed=self.plan.seed,
-#         )[0]
+        for band in self.instrument.bands:
+            band_mask = self.instrument.dets.band_name == band.name
 
-#         self.CMB_map += utils.Tcmb
-#         self.CMB_map *= utils.KcmbToKbright(np.unique(self.instrument.dets.band_center)[bandnumber])
+            band_cmb_power_samples_W = np.trapz(
+                y=cmb_brightness * band.passband(1e-9 * test_nu), x=test_nu
+            )
 
-#     #self._cmb_imager(i)
-#         #         cmb_data = sp.interpolate.RegularGridInterpolator(
-#         #                     (map_x, map_y), self.CMB_map, bounds_error=False, fill_value=0
-#         #                     )((lam_x, lam_y))
-#         #         self.noise    = self.lam.temperature + cmb_data
-#         #         self.combined = self.map_data + self.lam.temperature + cmb_data
+            self.data["cmb"][band_mask] = np.interp(
+                T_CMB + cmb_temperatures[band_mask],
+                cmb_temperature_samples_K,
+                band_cmb_power_samples_W,
+            )