Homogenize nthreads handling across the framework

CHANGELOG.md

@@ -1,5 +1,7 @@
 # HEAD
 
+-   Homogenize `nthreads` handling across the framework: every function now defaults to `nthreads: int | None = None` and resolves it via `resolve_nthreads()`/`resolve_numba_nthreads()`, removing the previous mix of `0`/`1` literal defaults and duplicated `OMP_NUM_THREADS`-reading code. Thread counts are resolved once at import time (`lbs.NUM_THREADS`, `lbs.NUMBA_NUM_THREADS`, mirroring `lbs.MPI_COMM_WORLD`) instead of on every call. Numba threads can now be tuned independently via `NUMBA_NUM_THREADS`, falling back to `OMP_NUM_THREADS`. **Behavior change**: when neither environment variable is set, threading now defaults to a single thread instead of every hardware thread available to the process, to avoid oversubscription in MPI/multi-process runs; set `OMP_NUM_THREADS` (and, if needed, `NUMBA_NUM_THREADS`) explicitly to use more than one core. Reference PR: [#550](https://github.com/litebird/litebird_sim/pull/550).
+
 -   Fix wrong output units when `bandpass_integration=True` in `input_sky` (CMB and dipole were off by a large factor because `pysm3.bandpass_unit_conversion` wasn't told the input unit), remove the resulting dead code path, and document that `cmb_ps_file` spectra are expected in $\mu K_{CMB}^2$. Fixes [#548](https://github.com/litebird/litebird_sim/issues/548). Reference PR: [#549](https://github.com/litebird/litebird_sim/pull/549).
 
 -   Stop supporting unused MPI processes and drop `MPI_COMM_GRID` class and its attributes. Now throwing an error when `comm_size != n_blocks_det * n_blocks_time`. It fixes [#364](https://github.com/litebird/litebird_sim/issues/364), while incorporating selective changes from [#372](https://github.com/litebird/litebird_sim/pull/372) and [#427](https://github.com/litebird/litebird_sim/pull/427). Reference PR: [#539](https://github.com/litebird/litebird_sim/pull/539).

docs/source/beam_convolution.rst

@@ -392,8 +392,7 @@ For a single-task execution, refer to the following example:
     sim.convolve_sky(sky_alms=alms,
                      beam_alms=blms,
                      convolution_params=Convparams,
-                     pointings_dtype=np.float32,
-                     nthreads = 0)
+                     pointings_dtype=np.float32)
 
 
 API reference

docs/source/maps_and_harmonics.rst

@@ -107,8 +107,13 @@ These functions handle the complexity of spin-0 (Temperature) vs spin-2 (Polariz
 * :func:`~litebird_sim.maps_and_harmonics.compute_dl`: Compute $D_{\ell} = \ell(\ell+1)C_{\ell}/(2\pi)$ spectra
 
 .. tip::
-   All transform functions accept a ``nthreads`` argument. 
-   Setting ``nthreads=0`` (default) uses all available hardware threads, which is optimal for standalone scripts but should be adjusted when running inside an MPI environment.
+   All transform functions accept an ``nthreads`` argument, which defaults to ``None``.
+   When left unset, the number of threads is resolved automatically (see :ref:`nthreads_ducc0`):
+   it honours the ``OMP_NUM_THREADS`` environment variable if set, and otherwise falls back
+   to ``1`` thread. This makes the default safe to use both in standalone scripts and inside
+   an MPI environment, but it also means you need to set ``OMP_NUM_THREADS`` yourself if you
+   want these functions to use more than one core. Pass an explicit ``nthreads`` to override
+   this for a single call.
 
 .. note::
    All transform functions support multi-frequency data. When operating on multi-frequency objects, transforms are applied independently to each frequency, and the output maintains the multi-frequency structure.

docs/source/mpi.rst

@@ -30,20 +30,37 @@ Some parts of the LiteBIRD Simulation Framework are able to
 exploit multiple cores because several of its modules rely on
 the `Numba <https://numba.pydata.org/>`_ library.
 
-If you are running your code on your multi-core laptop, you do not
-have to do anything fancy in order to use all the CPUs on your machine:
-in its default configuration, the Framework should be able to take
-advantage all the available CPU cores.
-
-However, if you want to tune the way the Framework uses the CPUs,
-you can either set the environment variable ``OMP_NUM_THREADS``
-to the number of CPUs to use, or use two parameters in
-the constructor of the class :class:`.Simulation`:
-
-- `numba_num_of_threads`: this is the number of CPUs that Numba will
+To use more than one CPU, you must explicitly tell the Framework how
+many to use, either by setting the environment variable
+``OMP_NUM_THREADS`` to the number of CPUs to use, or by using two
+parameters in the constructor of the class :class:`.Simulation`.
+The Framework does **not** default to using every CPU core available
+on the machine: doing so would be unsafe under MPI, where several
+ranks typically share a node and would otherwise all try to grab every
+core at once. Leaving thread counts unconfigured makes the Framework
+run single-threaded, both for Numba and for ducc0 (see
+:ref:`nthreads_ducc0` below); this is a safe default, but it means
+that on a single-user laptop or workstation you *do* need to set
+``OMP_NUM_THREADS`` (or the equivalent parameters below) yourself if
+you want to take advantage of all the available cores.
+
+- `numba_threads`: this is the number of CPUs that Numba will
   use for parallel calculations. The parameter defaults to ``None``,
-  which means that Numba will check how many CPUs are available and will
-  use all of them.
+  in which case it is resolved once, following this order of precedence:
+
+  1. the value of `numba_threads` itself, if you passed one explicitly
+     (or if it was set in a TOML parameter file, see below);
+  2. the environment variable ``NUMBA_NUM_THREADS``, if set;
+  3. the environment variable ``OMP_NUM_THREADS``, if set;
+  4. ``1``, if none of the above is set.
+
+  This is the same resolution order used for the ``nthreads`` parameters
+  accepted by the low-level `ducc0 <https://gitlab.mpcdf.mpg.de/mtr/ducc>`_-based
+  functions of the framework (spherical harmonic transforms, beam
+  convolution, map scanning, ...) — see :ref:`nthreads_ducc0` below. Using
+  ``NUMBA_NUM_THREADS`` instead of ``OMP_NUM_THREADS`` lets you give Numba a
+  different thread count than ducc0, if you ever need to; leaving
+  ``NUMBA_NUM_THREADS`` unset makes the two agree by default.
 
 - `numba_threading_layer`: this parameter is a string that specifies
   which threading library should be used by Numba. The value depends
@@ -71,7 +88,7 @@ These parameters can be passed through a TOML parameter file (see
    # This is file "my_conf.toml"
    [simulation]
    random_seed = 12345
-   numba_num_of_threads = 32
+   numba_threads = 32
    numba_threading_layer = "tbb"
 
 Both ``tbb`` and ``omp`` require that the relevant library be available on
@@ -87,6 +104,42 @@ of running a command like the following:
    $ module load openmp  # OpenMP
 
 
+.. _nthreads_ducc0:
+
+Threads and ducc0
+~~~~~~~~~~~~~~~~~~
+
+Besides Numba, most of the computationally-heavy, low-level functions of
+the framework (spherical harmonic transforms in
+:mod:`litebird_sim.maps_and_harmonics`, beam convolution, map scanning, the
+map-makers, ...) delegate their parallel work to the
+`ducc0 <https://gitlab.mpcdf.mpg.de/mtr/ducc>`_ library through an
+``nthreads`` parameter. Every such function defaults to ``nthreads=None``,
+which is resolved through :func:`litebird_sim.resolve_nthreads` using the
+same precedence as ``OMP_NUM_THREADS`` for Numba above: explicit value,
+then ``OMP_NUM_THREADS``, then ``1``. The values used by the current
+process, resolved once when ``litebird_sim`` is imported, are exposed as
+:data:`litebird_sim.NUM_THREADS` (for ducc0) and
+:data:`litebird_sim.NUMBA_NUM_THREADS` (for Numba).
+
+This ``1``-thread fallback is deliberately conservative: it keeps both
+serial and MPI runs safe by default (no risk of a rank grabbing every
+core on a shared node), at the cost of not using extra cores unless you
+ask for them. If you want multithreaded ducc0/Numba execution — whether
+on a laptop or per MPI rank on a cluster — set ``OMP_NUM_THREADS`` (and,
+if needed, ``NUMBA_NUM_THREADS``) to the number of cores you want to use
+before launching your job, e.g. for 8 MPI ranks on a 64-core node, using
+8 threads per rank:
+
+.. code-block:: sh
+
+   $ export OMP_NUM_THREADS=8
+   $ mpirun -n 8 python3 my_script.py
+
+You can always override the resolved value for a single call by passing an
+explicit ``nthreads`` argument to the function you are calling.
+
+
 MPI
 ~~~
 

litebird_sim/__init__.py

@@ -42,6 +42,12 @@
     EARTH_L2_DISTANCE_KM,
     NUM_THREADS_ENVVAR,
 )
+from .utilities import (
+    NUM_THREADS,
+    NUMBA_NUM_THREADS,
+    resolve_nthreads,
+    resolve_numba_nthreads,
+)
 from .coordinates import (
     DEFAULT_COORDINATE_SYSTEM,
     DEFAULT_TIME_SCALE,
@@ -253,6 +259,11 @@
     "SOLAR_VELOCITY_GAL_LON_RAD",
     "EARTH_L2_DISTANCE_KM",
     "NUM_THREADS_ENVVAR",
+    # utilities.py
+    "NUM_THREADS",
+    "NUMBA_NUM_THREADS",
+    "resolve_nthreads",
+    "resolve_numba_nthreads",
     # units.py
     "Units",
     "UnitUtils",

litebird_sim/beam_convolution.py

@@ -51,7 +51,7 @@ def add_convolved_sky_to_one_detector(
     mueller_matrix,
     hwp_angle,
     convolution_params: BeamConvolutionParameters | None = None,
-    nthreads: int = 0,
+    nthreads: int | None = None,
 ):
     """
     Convolve given sky alms with a detector beam alms and add the result to the TOD of a single detector.
@@ -74,9 +74,10 @@ def add_convolved_sky_to_one_detector(
     convolution_params : BeamConvolutionParameters, optional
         Parameters controlling the convolution, such as resolution and precision. If None,
         reasonable defaults are chosen based on the sky and beam properties.
-    nthreads : int, default=0
-        Number of threads to use for convolution. If set to 0, all available CPU cores
-        will be used.
+    nthreads : int or None, default=None
+        Number of threads to use for convolution. If None, resolved via
+        :func:`.resolve_nthreads` (``OMP_NUM_THREADS``, or all available
+        threads if unset).
 
     Raises
     ------
@@ -96,6 +97,8 @@ def add_convolved_sky_to_one_detector(
     - The function modifies `tod_det` in place by adding the convolved signal.
     """
 
+    nthreads = resolve_nthreads(nthreads)
+
     if not convolution_params:
         sky_lmax = sky_alms_det.lmax
 
@@ -214,7 +217,7 @@ def add_convolved_sky(
     convolution_params: BeamConvolutionParameters | None = None,
     pointings_dtype=np.float64,
     nside_centering: int | None = None,
-    nthreads: int = 0,
+    nthreads: int | None = None,
 ):
     """
     Convolve a set of sky maps with detector beams and add the resulting signals to the
@@ -253,9 +256,10 @@ def add_convolved_sky(
     nside_centering : int, default=None
         If set, shifts the detector pointings to the centers of the corresponding HEALPix pixels
         at the given NSIDE resolution. If None, no centering is applied.
-    nthreads : int, default=0
+    nthreads : int or None, default=None
         Number of threads to use for convolution and in case for HEALPix operations.
-        If set to 0, all available CPU cores will be used.
+        If None, resolved via :func:`.resolve_nthreads` (``OMP_NUM_THREADS``, or all
+        available threads if unset).
 
     Raises
     ------
@@ -273,6 +277,8 @@ def add_convolved_sky(
     - The function modifies `tod` in place by adding the convolved signals for all detectors.
     """
 
+    nthreads = resolve_nthreads(nthreads)
+
     if mueller_hwp is not None:
         assert tod.shape[0] == mueller_hwp.shape[0]
 
@@ -327,6 +333,7 @@ def add_convolved_sky(
             output_coordinate_system=coordinates,
             nside_centering=nside_centering,
             pointings_dtype=pointings_dtype,
+            nthreads=nthreads,
         )
 
         # FIXME: Fix this at some point, ducc wants phi 0 -> 2pi

litebird_sim/constants.py

@@ -2,9 +2,13 @@
 from astropy.constants import c as c_light
 from astropy.constants import h, k_B
 
-# Name of the environment variable used in the convolution
+# Environment variables used to size ducc0's and Numba's thread pools.
+# NUMBA_NUM_THREADS_ENVVAR takes precedence over NUM_THREADS_ENVVAR when
+# resolving the number of threads for Numba, so the two runtimes can be
+# sized independently if needed (see resolve_nthreads/resolve_numba_nthreads
+# in utilities.py).
 NUM_THREADS_ENVVAR = "OMP_NUM_THREADS"
-NUMBA_NUM_THREADS_ENVVAR = "OMP_NUM_THREADS"
+NUMBA_NUM_THREADS_ENVVAR = "NUMBA_NUM_THREADS"
 
 ARCMIN_TO_RAD = np.pi / 180 / 60
 

litebird_sim/grasp2alm.py

@@ -71,6 +71,7 @@ def to_alm(
         mmax: int,
         epsilon=1e-8,
         max_num_of_iterations=20,
+        nthreads: int | None = None,
     ) -> np.ndarray:
         """Converts the beam map to spherical harmonic coefficients.
 
@@ -79,6 +80,8 @@ def to_alm(
             mmax (`int`): Maximum m value for the spherical harmonic expansion.
             epsilon (`float`): Precision of the result
             max_num_of_iterations (`int`): Maximum number of iterations
+            nthreads (`int` or `None`): Number of threads to use. If None,
+                resolved via :func:`.resolve_nthreads`.
 
         Returns:
             `numpy.ndarray`: The spherical harmonic coefficients, as a (3, N) array.
@@ -88,6 +91,8 @@ def to_alm(
 
         """
 
+        nthreads = resolve_nthreads(nthreads)
+
         if not self.map.shape[0] <= 3:
             raise ValueError(
                 "Error in BeamMap.to_alm: map has more than 3 Stokes parameters"
@@ -105,7 +110,7 @@ def to_alm(
             lmax=lmax,
             mmax=mmax,
             spin=0,
-            nthreads=0,
+            nthreads=nthreads,
             maxiter=max_num_of_iterations,
             epsilon=epsilon,
             **geom,
@@ -123,7 +128,7 @@ def to_alm(
             lmax=lmax,
             mmax=mmax,
             spin=2,
-            nthreads=0,
+            nthreads=nthreads,
             maxiter=max_num_of_iterations,
             epsilon=epsilon,
             **geom,

litebird_sim/hwp_harmonics/hwp_harmonics.py

@@ -1,5 +1,4 @@
 import logging
-import os
 
 import numpy as np
 import numpy.typing as npt
@@ -11,7 +10,6 @@
 from litebird_sim.hwp_jones_parameters import HWPJonesParams
 
 from ..bandpass_template_module import bandpass_profile
-from ..constants import NUM_THREADS_ENVVAR
 from ..coordinates import CoordinateSystem
 from ..hwp_non_ideal import HWPFormalism, NonIdealHWP
 from ..input_sky import SkyInput
@@ -20,6 +18,7 @@
 from ..pointings_in_obs import (
     _get_pointings_array,
 )
+from ..utilities import resolve_nthreads
 from .jones_methods import (
     compute_signal_for_one_detector as compute_signal_for_one_detector_jones,
 )
@@ -228,9 +227,7 @@ def fill_tod_with_hwp_harmonics(
             )
     assert maps is not None, "You need to pass input maps to fill_tod."
 
-    # Set number of threads
-    if nthreads is None:
-        nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+    nthreads = resolve_nthreads(nthreads)
 
     if pointings is None:
         if hwp_angle is not None:
@@ -390,6 +387,7 @@ def fill_tod_with_hwp_harmonics(
                 hwp_angle=cur_hwp_angle,
                 output_coordinate_system=coordinates,
                 pointings_dtype=pointings_dtype,
+                nthreads=nthreads,
             )
 
             tod_det = tod[idet, :]

litebird_sim/input_sky.py

@@ -25,6 +25,7 @@
     synthesize_alm,
 )
 from .units import Units, UnitUtils
+from .utilities import resolve_nthreads
 
 # --- Utility Functions ---
 
@@ -145,7 +146,7 @@ def __init__(
         # Bandpass
         bandpass_integration: bool = False,
         # Parallelism
-        nthreads: int = 0,  # 0 usually means "use all available" in ducc0
+        nthreads: int | None = None,  # None: resolved via resolve_nthreads()
         # Components to generate
         make_cmb: bool = True,
         make_fg: bool = False,
@@ -178,7 +179,7 @@ def __init__(
         self.bandpass_integration = bandpass_integration
         self.maxiter = maxiter
         self.epsilon = epsilon
-        self.nthreads = nthreads
+        self.nthreads = resolve_nthreads(nthreads)
         self.make_cmb = make_cmb
         self.make_fg = make_fg
         self.make_dipole = make_dipole
@@ -635,7 +636,7 @@ def _dipole_map_values(self) -> np.ndarray:
         npix = hpx.npix()
 
         vec = dh.ang2vec(np.array([[lat, lon]]))[0]
-        pix_vecs = hpx.pix2vec(np.arange(npix))
+        pix_vecs = hpx.pix2vec(np.arange(npix), nthreads=self.params.nthreads)
 
         dipole_map_val = np.zeros((3, npix))
         dipole_map_val[0] = np.dot(pix_vecs, vec) * amp