Implemented manual step-wise trapezoidal integration.

Drastically improved memory usage for high nsides by moving away from using np.trapz. Also renamed some files.

zodipy/_integration.py

@@ -0,0 +1,31 @@
+from typing import Callable
+
+import numpy as np
+
+
+EmissionCallable = Callable[
+    [float, np.ndarray, np.ndarray, np.ndarray, float], np.ndarray
+]
+
+
+def trapezoidal(
+    emission_func: EmissionCallable,
+    freq: float,
+    x_obs: np.ndarray,
+    x_earth: np.ndarray,
+    x_unit: np.ndarray,
+    R: np.ndarray,
+    npix: int,
+    pixels: np.ndarray,
+) -> np.ndarray:
+    """Integrates the emission for a component using trapezoidal."""
+
+    comp_emission = np.zeros(npix)[pixels]
+    emission_prev = emission_func(freq, x_obs, x_earth, x_unit, R[0])
+    for i in range(1, len(R)):
+        dR = R[i] - R[i - 1]
+        emission_cur = emission_func(freq, x_obs, x_earth, x_unit, R[i])
+        comp_emission += (emission_prev + emission_cur) * dR / 2
+        emission_prev = emission_cur
+
+    return comp_emission

zodipy/_los_config.py

@@ -0,0 +1,50 @@
+from typing import Tuple, Dict
+
+import numpy as np
+
+
+class LOSFactory:
+    """Factory responsible for registring and book-keeping a line-of-sight (LOS)."""
+
+    def __init__(self) -> None:
+        self._configs = {}
+
+    def register_config(
+        self, name: str, components: Dict[str, Tuple[float, int]]
+    ) -> None:
+        """Initializes and stores a LOS."""
+
+        config = {}
+        for key, value in components.items():
+            if isinstance(value, np.ndarray):
+                config[key] = value
+            elif isinstance(value, (tuple, list)):
+                try:
+                    start, stop, n, geom = value
+                except ValueError:
+                    raise ValueError(
+                        "Line-of-sight config must either be an array, or "
+                        "a tuple with the format (start, stop, n, geom)"
+                        "where geom is either 'linear' or 'log'"
+                    )
+                if geom.lower() == "linear":
+                    geom = np.linspace
+                elif geom.lower() == "log":
+                    geom = np.geomspace
+                else:
+                    raise ValueError("geom must be either 'linear' or 'log'")
+                config[key] = geom(start, stop, n)
+
+        self._configs[name] = config
+
+    def get_config(self, name: str) -> np.ndarray:
+        """Returns a registered config."""
+
+        config = self._configs.get(name)
+        if config is None:
+            raise ValueError(
+                f"Config {name} is not registered. Available configs are "
+                f"{list(self._configs.keys())}"
+            )
+
+        return config

zodipy/_simulation.py

@@ -1,11 +1,12 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from typing import Dict
 import warnings
 
 import healpy as hp
 import numpy as np
 
-from zodipy._integration_config import IntegrationConfig
+from zodipy._integration import trapezoidal
 from zodipy._model import InterplanetaryDustModel
 
 
@@ -18,9 +19,8 @@ class SimulationStrategy(ABC):
     model
         Interplanetary dust model with initialized componentents and
         corresponding emissivities.
-    integration_config
-        Configuration object that determines how a component is integrated
-        along a line-of-sight.
+    line_of_sight_config
+        Dictionary mapping a line_of_sight_config to each component.
     observer_locations
         The location(s) of the observer.
     earth_location
@@ -30,7 +30,7 @@ class SimulationStrategy(ABC):
     """
 
     model: InterplanetaryDustModel
-    integration_config: IntegrationConfig
+    line_of_sight_config: Dict[str, np.ndarray]
     observer_locations: np.ndarray
     earth_locations: np.ndarray
     hit_counts: np.ndarray
@@ -83,18 +83,20 @@ def simulate(self, nside: int, freq: float) -> np.ndarray:
         emission = np.zeros((len(components), npix)) + np.NAN
 
         for comp_idx, (comp_name, comp) in enumerate(components.items()):
-            integration_config = self.integration_config[comp_name]
-            R = integration_config.R
-
-            comp_emission = comp.get_emission(freq, X_observer, X_earth, X_unit, R)
-            integrated_comp_emission = integration_config.integrator(
-                comp_emission, R, dx=integration_config.dR, axis=0
-            )
-
             comp_emissivity = emissivities.get_emissivity(comp_name, freq)
-            integrated_comp_emission *= comp_emissivity
+            line_of_sight = self.line_of_sight_config[comp_name]
+            integrated_comp_emission = trapezoidal(
+                comp.get_emission,
+                freq,
+                X_observer,
+                X_earth,
+                X_unit,
+                line_of_sight,
+                npix,
+                pixels,
+            )
 
-            emission[comp_idx, pixels] = integrated_comp_emission
+            emission[comp_idx, pixels] = comp_emissivity * integrated_comp_emission
 
         return emission * 1e20
 
@@ -131,20 +133,22 @@ def simulate(self, nside: int, freq: float) -> np.ndarray:
             unit_vectors = X_unit[:, pixels]
 
             for comp_idx, (comp_name, comp) in enumerate(components.items()):
-                integration_config = self.integration_config[comp_name]
-                R = integration_config.R
-
-                comp_emission = comp.get_emission(
-                    freq, observer_pos, earth_pos, unit_vectors, R
-                )
-                integrated_comp_emission = integration_config.integrator(
-                    comp_emission, R, dx=integration_config.dR, axis=0
+                comp_emissivity = emissivities.get_emissivity(comp_name, freq)
+                line_of_sight = self.line_of_sight_config[comp_name]
+
+                integrated_comp_emission = trapezoidal(
+                    comp.get_emission,
+                    freq,
+                    observer_pos,
+                    earth_pos,
+                    unit_vectors,
+                    line_of_sight,
+                    npix,
+                    pixels,
                 )
 
-                comp_emissivity = emissivities.get_emissivity(comp_name, freq)
-                integrated_comp_emission *= comp_emissivity
                 emission[comp_idx, pixels] += (
-                    integrated_comp_emission * hit_count[pixels]
+                    integrated_comp_emission * comp_emissivity * hit_count[pixels]
                 )
 
         with warnings.catch_warnings():

zodipy/los_configs.py

@@ -0,0 +1,43 @@
+from zodipy._los_config import LOSFactory
+
+import numpy as np
+
+
+EPS = np.finfo(float).eps
+RADIAL_CUTOFF = 6
+
+LOS_configs = LOSFactory()
+
+LOS_configs.register_config(
+    name="default",
+    components={
+        "cloud": (EPS, RADIAL_CUTOFF, 250, "linear"),
+        "band1": (EPS, RADIAL_CUTOFF, 50, "linear"),
+        "band2": (EPS, RADIAL_CUTOFF, 50, "linear"),
+        "band3": (EPS, RADIAL_CUTOFF, 50, "linear"),
+        "ring": (EPS, 2.25, 50, "linear"),
+        "feature": (EPS, 1, 50, "linear"),
+    },
+)
+LOS_configs.register_config(
+    name="high",
+    components={
+        "cloud": (EPS, RADIAL_CUTOFF, 500, "linear"),
+        "band1": (EPS, RADIAL_CUTOFF, 500, "linear"),
+        "band2": (EPS, RADIAL_CUTOFF, 500, "linear"),
+        "band3": (EPS, RADIAL_CUTOFF, 500, "linear"),
+        "ring": (EPS, 2.25, 200, "linear"),
+        "feature": (EPS, 1, 200, "linear"),
+    },
+)
+LOS_configs.register_config(
+    name="fast",
+    components={
+        "cloud": (EPS, RADIAL_CUTOFF, 25, "linear"),
+        "band1": (EPS, RADIAL_CUTOFF, 25, "linear"),
+        "band2": (EPS, RADIAL_CUTOFF, 25, "linear"),
+        "band3": (EPS, RADIAL_CUTOFF, 25, "linear"),
+        "ring": (EPS, 2.25, 25, "linear"),
+        "feature": (EPS, 1, 25, "linear"),
+    },
+)

zodipy/zodi.py

@@ -1,12 +1,11 @@
 from typing import Optional, Union, Iterable, Dict
-import warnings
 
 import astropy.units as u
 import numpy as np
 
 from zodipy._coordinates import get_target_coordinates, to_frame
 from zodipy._simulation import InstantaneousStrategy, TimeOrderedStrategy
-from zodipy.integration_configs import integration_configs
+from zodipy.los_configs import LOS_configs
 from zodipy.models import models
 
 _EpochsType = Optional[Union[float, Iterable[float], Dict[str, str]]]
@@ -43,8 +42,8 @@ class Zodi:
         The Interplanetary dust model used in the simulation. Available
         options are 'planck 2013', 'planck 2015', and 'planck 2018'.
         Defaults to 'planck 2018'.
-    integration_config
-        The integration_config object used when calling `get_emission`.
+    line_of_sight_config
+        The line-of-sight configuration used in the simulation.
         Available options are: 'default', and 'high', and 'fast'. Defaults
         to 'default'.
     """
@@ -55,11 +54,11 @@ def __init__(
         epochs: Optional[_EpochsType] = None,
         hit_counts: Optional[Iterable[np.ndarray]] = None,
         model: str = "planck 2018",
-        integration_config: str = "default",
+        line_of_sight_config: str = "default",
     ) -> None:
 
         model = models.get_model(model)
-        integration_config = integration_configs.get_config(integration_config)
+        line_of_sight_config = LOS_configs.get_config(line_of_sight_config)
 
         observer_locations = get_target_coordinates(observer, epochs)
         earth_locations = get_target_coordinates("earth", epochs)
@@ -82,7 +81,7 @@ def __init__(
         else:
             simulation_strategy = TimeOrderedStrategy
         self._simulation_strategy = simulation_strategy(
-            model, integration_config, observer_locations, earth_locations, hit_counts
+            model, line_of_sight_config, observer_locations, earth_locations, hit_counts
         )
 
     def get_emission(
@@ -117,16 +116,6 @@ def get_emission(
             Simulated Zodiacal emission in units of MJy/sr.
         """
 
-        if nside > 256:
-            warnings.warn(
-                "Memory usage may be very high for the requested simulation. "
-                "In the current implementation of zodipy, intermediate "
-                "results are stored in memory during line-of-sight integration. "
-                "We recommend to simulate at a lower nside, e.g 256, and use "
-                "healpy.ud_grade afterwards instead (Zodiacal emission is very "
-                "smooth, so this should be fine)."
-            )
-
         if isinstance(freq, u.Quantity):
             freq = freq.to("GHz").value