diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2e54ca16..4cccec87 100644
--- a/CHANGELOG.md
+++ b/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).
diff --git a/docs/source/beam_convolution.rst b/docs/source/beam_convolution.rst
index fa162a41..bed67d6a 100644
--- a/docs/source/beam_convolution.rst
+++ b/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
diff --git a/docs/source/maps_and_harmonics.rst b/docs/source/maps_and_harmonics.rst
index b84ef7eb..abe4ae65 100644
--- a/docs/source/maps_and_harmonics.rst
+++ b/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.
diff --git a/docs/source/mpi.rst b/docs/source/mpi.rst
index f4627b08..32d26431 100644
--- a/docs/source/mpi.rst
+++ b/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 `_ 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 `_-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 `_ 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
~~~
diff --git a/litebird_sim/__init__.py b/litebird_sim/__init__.py
index a426a2d5..603829ca 100644
--- a/litebird_sim/__init__.py
+++ b/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",
diff --git a/litebird_sim/beam_convolution.py b/litebird_sim/beam_convolution.py
index a02a0f5c..266ef11f 100644
--- a/litebird_sim/beam_convolution.py
+++ b/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
diff --git a/litebird_sim/constants.py b/litebird_sim/constants.py
index 226c9107..c1cc688b 100644
--- a/litebird_sim/constants.py
+++ b/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
diff --git a/litebird_sim/grasp2alm.py b/litebird_sim/grasp2alm.py
index e7911290..407278aa 100644
--- a/litebird_sim/grasp2alm.py
+++ b/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,
diff --git a/litebird_sim/hwp_harmonics/hwp_harmonics.py b/litebird_sim/hwp_harmonics/hwp_harmonics.py
index 51484e97..c6cb9ddb 100644
--- a/litebird_sim/hwp_harmonics/hwp_harmonics.py
+++ b/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, :]
diff --git a/litebird_sim/input_sky.py b/litebird_sim/input_sky.py
index 0efdc467..9866d521 100644
--- a/litebird_sim/input_sky.py
+++ b/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
diff --git a/litebird_sim/mapmaking/common.py b/litebird_sim/mapmaking/common.py
index 23d25bc4..5a95e0ee 100644
--- a/litebird_sim/mapmaking/common.py
+++ b/litebird_sim/mapmaking/common.py
@@ -10,6 +10,7 @@
from litebird_sim.coordinates import CoordinateSystem
from litebird_sim.observations import Observation
from litebird_sim.pointings_in_obs import _get_pointings_array, _get_pol_angle
+from litebird_sim.utilities import resolve_nthreads
# The threshold on the conditioning number used to determine if a pixel
# was really “seen” or not
@@ -161,6 +162,7 @@ def _compute_pixel_indices(
hwp_angle=hwp_angle,
output_coordinate_system=output_coordinate_system,
pointings_dtype=pointings_dtype,
+ nthreads=nthreads,
)
if hmap_generation:
polang_all[idet] = curr_pointings_det[:, 2]
@@ -193,10 +195,13 @@ def _compute_pixel_indices_single_detector(
output_coordinate_system: CoordinateSystem,
pointings_dtype=np.float64,
hmap_generation: bool = False,
+ nthreads: int | None = None,
) -> tuple[npt.NDArray, npt.NDArray]:
"""
Same as _compute_pixel_indices but for a single detector, thus returning only the pixel indices and polarization angles for that detector.
"""
+ nthreads = resolve_nthreads(nthreads)
+
pixidx = np.empty((num_of_samples), dtype=np.int32)
polang = np.empty((num_of_samples), dtype=pointings_dtype)
curr_pointings_det, hwp_angle = _get_pointings_array(
@@ -205,6 +210,7 @@ def _compute_pixel_indices_single_detector(
hwp_angle=hwp_angle,
output_coordinate_system=output_coordinate_system,
pointings_dtype=pointings_dtype,
+ nthreads=nthreads,
)
if hmap_generation:
@@ -215,7 +221,7 @@ def _compute_pixel_indices_single_detector(
hwp_angle=hwp_angle,
pol_angle_detectors=pol_angle_detector,
)
- pixidx = hpx.ang2pix(curr_pointings_det[:, :2])
+ pixidx = hpx.ang2pix(curr_pointings_det[:, :2], nthreads=nthreads)
if output_coordinate_system == CoordinateSystem.Galactic:
# Free curr_pointings_det if the output map is already in Galactic coordinates
diff --git a/litebird_sim/mapmaking/h_maps.py b/litebird_sim/mapmaking/h_maps.py
index 9d079a30..3cc1fa25 100644
--- a/litebird_sim/mapmaking/h_maps.py
+++ b/litebird_sim/mapmaking/h_maps.py
@@ -18,6 +18,7 @@
_get_hwp_angle,
_normalize_observations_and_pointings,
)
+from litebird_sim.utilities import resolve_nthreads
from .common import (
_build_mask_detector_split,
@@ -240,6 +241,7 @@ def make_h_maps(
pointings_dtype=np.float64,
save_to_file: bool = True,
output_directory: str = "./h_n_maps",
+ nthreads: int | None = None,
) -> HMapsResult:
"""Generate complex harmonic maps :math:`h_{n,m}` from observations.
@@ -269,10 +271,15 @@ def make_h_maps(
:type save_to_file: bool
:param output_directory: Output directory for generated HDF5 files.
:type output_directory: str
+ :param nthreads: Number of threads to use for ducc0 calls. If None,
+ resolved via :func:`.resolve_nthreads`.
+ :type nthreads: int | None
:returns: Result container with all computed maps and metadata.
:rtype: HnMapResult
"""
+ nthreads = resolve_nthreads(nthreads)
+
assert (
n_m_couples.shape[1] == 2
), """the n,m couples should be passed in this shape: array([[n1, m1],
@@ -331,6 +338,7 @@ def make_h_maps(
output_coordinate_system=output_coordinate_system,
pointings_dtype=pointings_dtype,
hmap_generation=True,
+ nthreads=nthreads,
)
log.info(
f"Pixel indices and angles for detector {all_dets_list[idet]} computed, now building nobs matrices"
diff --git a/litebird_sim/maps_and_harmonics.py b/litebird_sim/maps_and_harmonics.py
index 6668dfe1..c37decfb 100644
--- a/litebird_sim/maps_and_harmonics.py
+++ b/litebird_sim/maps_and_harmonics.py
@@ -12,6 +12,7 @@
from .coordinates import ECL_TO_GAL_EULER, GAL_TO_ECL_EULER, CoordinateSystem
from .units import Units
+from .utilities import resolve_nthreads
# ======================================================================
# SphericalHarmonics
@@ -2369,7 +2370,7 @@ def interpolate_alm(
locations: np.ndarray,
*,
epsilon: float | None = None,
- nthreads: int = 0,
+ nthreads: int | None = None,
) -> np.ndarray | tuple[np.ndarray, np.ndarray, np.ndarray]:
r"""Interpolate spherical-harmonic coefficients at arbitrary positions
using :func:`ducc0.sht.synthesis_general`.
@@ -2407,9 +2408,9 @@ def interpolate_alm(
* complex64 → ``1e-6``
* complex128 → ``1e-13``
- nthreads : int, optional
- Number of threads for ducc. If 0 (default), ducc uses the
- number of hardware threads.
+ nthreads : int or None, optional
+ Number of threads for ducc. If None, resolved via
+ :func:`.resolve_nthreads`.
Returns
-------
@@ -2426,6 +2427,8 @@ def interpolate_alm(
If the input shapes are inconsistent or ``nstokes`` is not 1 or 3.
"""
+ nthreads = resolve_nthreads(nthreads)
+
alm = np.asarray(alms.values)
loc = np.asarray(locations, dtype=np.float64)
if loc.ndim != 2 or loc.shape[1] != 2:
@@ -2531,7 +2534,7 @@ def pixelize_alm(
nest: bool = False,
lmax: int | None = None,
mmax: int | None = None,
- nthreads: int = 0,
+ nthreads: int | None = None,
) -> "HealpixMap":
r"""
Convert spherical harmonics coefficients to a HEALPix map using
@@ -2584,9 +2587,9 @@ def pixelize_alm(
``lmax``/``mmax`` is requested than what is stored in ``alms``,
the coefficient array is truncated accordingly.
- nthreads : int, optional
- Number of threads passed to ducc. If zero (default), ducc chooses
- the number of threads.
+ nthreads : int or None, optional
+ Number of threads passed to ducc. If None, resolved via
+ :func:`.resolve_nthreads`.
Returns
-------
@@ -2605,6 +2608,8 @@ def pixelize_alm(
or if requested ``lmax``/``mmax`` are inconsistent with the
available coefficients.
"""
+ nthreads = resolve_nthreads(nthreads)
+
# --- basic checks / effective lmax, mmax -------------------------------
if alms.nstokes not in (1, 3):
raise ValueError(
@@ -2732,7 +2737,7 @@ def estimate_alm(
mmax: int | None = None,
maxiter: int | None = None,
epsilon: float | None = None,
- nthreads: int = 0,
+ nthreads: int | None = None,
) -> SphericalHarmonics:
r"""
Estimate spherical harmonic coefficients ($a_{\ell m}$) from a HEALPix map.
@@ -2802,9 +2807,9 @@ def estimate_alm(
* ``complex64`` output -> ``1e-6``
* ``complex128`` output -> ``1e-13``
- nthreads : int, optional
- Number of threads passed to ducc. If zero (default), ducc chooses
- the number of threads.
+ nthreads : int or None, optional
+ Number of threads passed to ducc. If None, resolved via
+ :func:`.resolve_nthreads`.
Returns
-------
@@ -2826,6 +2831,8 @@ def estimate_alm(
If the HealpixMap object has unsupported ``nstokes``,
or if requested ``lmax``/``mmax`` are inconsistent.
"""
+ nthreads = resolve_nthreads(nthreads)
+
# --- validate maxiter ---------------------------------------------------
if maxiter is not None:
if not isinstance(maxiter, int):
@@ -3028,7 +3035,7 @@ def rotate_alm(
phi: float | None = None,
mmax_out: int | None = None,
inplace: bool = False,
- nthreads: int = 0,
+ nthreads: int | None = None,
) -> SphericalHarmonics:
"""
Rotate spherical harmonic coefficients using ducc0.
@@ -3055,8 +3062,9 @@ def rotate_alm(
If True, modifies the input `alms` object in place.
Note: In-place rotation is only possible if `mmax_out` is equal to
`alms.mmax` (output size must match input size).
- nthreads : int, optional, keyword-only
- Number of threads to use for the rotation. Default is 0 (use all available).
+ nthreads : int or None, optional, keyword-only
+ Number of threads to use for the rotation. If None, resolved via
+ :func:`.resolve_nthreads`.
Returns
-------
@@ -3071,6 +3079,8 @@ def rotate_alm(
- If `inplace=True` is requested but `mmax_out` differs from `alms.mmax`.
"""
+ nthreads = resolve_nthreads(nthreads)
+
# 1. Retrieve Geometry from Input
lmax_in = alms.lmax
mmax_in = alms.mmax
diff --git a/litebird_sim/pointings_in_obs.py b/litebird_sim/pointings_in_obs.py
index edef0bc9..35b75e19 100644
--- a/litebird_sim/pointings_in_obs.py
+++ b/litebird_sim/pointings_in_obs.py
@@ -11,6 +11,7 @@
from .hwp import HWP
from .observations import Observation
from .scanning import RotQuaternion
+from .utilities import resolve_nthreads
def prepare_pointings(
@@ -228,7 +229,7 @@ def _get_pointings_array(
output_coordinate_system: CoordinateSystem,
nside_centering: int | None = None,
pointings_dtype=np.float64,
- nthreads: int = 0,
+ nthreads: int | None = None,
) -> tuple[np.ndarray, np.ndarray | None]:
"""Compute the pointings (θ, φ) and HWP angle for a given detector.
@@ -247,9 +248,9 @@ def _get_pointings_array(
If provided, the pointings will be aligned to the center of the HEALPix pixel.
pointings_dtype : np.dtype, optional
Data type for computed pointings and angles. Default is `np.float64`.
- nthreads : int, optional
+ nthreads : int or None, optional
Number of threads to use for HEALPix operations when centering pointings.
- Default is 0 (use all available threads).
+ If None, resolved via :func:`.resolve_nthreads`.
Returns
-------
@@ -259,6 +260,8 @@ def _get_pointings_array(
N_cols is typically 2 ([θ, φ]) or 3 ([θ, φ, ψ]).
- `hwp_angle` is either the provided array or the one computed by the callable.
"""
+ nthreads = resolve_nthreads(nthreads)
+
if isinstance(pointings, np.ndarray):
curr_pointings_det = pointings[detector_idx, :, :]
computed_hwp_angle = None
diff --git a/litebird_sim/scan_map.py b/litebird_sim/scan_map.py
index a5379beb..97573cf5 100644
--- a/litebird_sim/scan_map.py
+++ b/litebird_sim/scan_map.py
@@ -121,7 +121,7 @@ def scan_map(
input_names: str | None = None,
interpolation: str | None = "",
pointings_dtype=np.float64,
- nthreads: int = 0,
+ nthreads: int | None = None,
):
"""
Scan a sky map and fill time-ordered data (TOD) based on detector observations.
@@ -190,9 +190,9 @@ def scan_map(
pointings_dtype : dtype, optional
Data type for pointings generated on the fly.
- 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`.
Raises
------
@@ -210,6 +210,8 @@ def scan_map(
experiments.
"""
+ nthreads = resolve_nthreads(nthreads)
+
n_detectors = tod.shape[0]
if type(pointings) is np.ndarray:
@@ -262,6 +264,7 @@ def scan_map(
hwp_angle=hwp_angle,
output_coordinate_system=coordinates,
pointings_dtype=pointings_dtype,
+ nthreads=nthreads,
)
# ----------------------------------------------------------
diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py
index 10ee60ff..a275b8da 100644
--- a/litebird_sim/simulations.py
+++ b/litebird_sim/simulations.py
@@ -33,7 +33,7 @@
add_convolved_sky_to_observations,
)
from .beam_synthesis import generate_gauss_beam_alms
-from .constants import NUMBA_NUM_THREADS_ENVVAR
+from .constants import NUMBA_NUM_THREADS_ENVVAR # noqa: F401 (re-exported)
from .coordinates import CoordinateSystem
from .detectors import UUID, DetectorInfo, FreqChannelInfo, InstrumentInfo
from .dipole import DipoleType, add_dipole_to_observations
@@ -75,6 +75,7 @@
from .seeding import RNGHierarchy
from .spacecraft import SpacecraftOrbit, spacecraft_pos_and_vel
from .units import Units
+from .utilities import resolve_numba_nthreads
from .version import (
__author__ as litebird_sim_author,
)
@@ -415,9 +416,10 @@ def __init__(
else:
self.imo = Imo()
- if not numba_threads and NUMBA_NUM_THREADS_ENVVAR in os.environ:
- numba_threads = int(os.environ[NUMBA_NUM_THREADS_ENVVAR])
-
+ # Resolved below, after _init_missing_params() has had a chance to
+ # fill this in from a parameter file: explicit argument takes
+ # precedence, then the parameter file, then resolve_numba_nthreads()
+ # (NUMBA_NUM_THREADS/OMP_NUM_THREADS/hardware default).
self.numba_threads = numba_threads
self.numba_threading_layer = numba_threading_layer
@@ -442,8 +444,8 @@ def __init__(
self._init_missing_params()
- if self.numba_threads:
- numba.set_num_threads(self.numba_threads)
+ self.numba_threads = resolve_numba_nthreads(self.numba_threads)
+ numba.set_num_threads(self.numba_threads)
if self.numba_threading_layer:
numba.config.THREADING_LAYER = self.numba_threading_layer # type: ignore[attr-defined]
diff --git a/litebird_sim/utilities.py b/litebird_sim/utilities.py
index 05d2cbdc..c636951f 100644
--- a/litebird_sim/utilities.py
+++ b/litebird_sim/utilities.py
@@ -1,16 +1,58 @@
import os
-from .constants import NUM_THREADS_ENVVAR
+from .constants import NUM_THREADS_ENVVAR, NUMBA_NUM_THREADS_ENVVAR
+
+
+def _compute_nthreads() -> int:
+ if NUM_THREADS_ENVVAR in os.environ:
+ return int(os.environ[NUM_THREADS_ENVVAR])
+ # No explicit thread count was requested. Default to a single thread
+ # rather than every hardware thread available to the process: MPI/OpenMP
+ # jobs are expected to set NUM_THREADS_ENVVAR (OMP_NUM_THREADS) per rank,
+ # so this fallback only matters for un-configured runs (e.g. a laptop),
+ # where using every core by default would otherwise silently starve
+ # other processes on the machine.
+ return 1
+
+
+def _compute_numba_nthreads() -> int:
+ if NUMBA_NUM_THREADS_ENVVAR in os.environ:
+ return int(os.environ[NUMBA_NUM_THREADS_ENVVAR])
+ return _compute_nthreads()
+
+
+# Resolved once at import time (each MPI process reads its own environment),
+# instead of re-reading the environment on every call. Mirrors the
+# lbs.MPI_COMM_WORLD pattern: a single, process-wide value computed once and
+# reused everywhere `nthreads` isn't explicitly overridden.
+NUM_THREADS = _compute_nthreads()
+NUMBA_NUM_THREADS = _compute_numba_nthreads()
def resolve_nthreads(nthreads: int | None) -> int:
- """Resolve thread count from an explicit value or environment.
+ """Resolve the number of threads to use for ducc0/general parallel code.
- If ``nthreads`` is not ``None``, return it unchanged. Otherwise read
- :data:`NUM_THREADS_ENVVAR` and fall back to 0 when the variable is unset.
+ If ``nthreads`` is given explicitly, return it unchanged. Otherwise
+ return :data:`NUM_THREADS`, which was resolved once at import time from
+ :data:`.constants.NUM_THREADS_ENVVAR` (``OMP_NUM_THREADS``), falling back
+ to ``1`` if the environment variable is not set.
"""
if nthreads is not None:
return nthreads
+ return NUM_THREADS
- return int(os.environ.get(NUM_THREADS_ENVVAR, 0))
+
+def resolve_numba_nthreads(nthreads: int | None) -> int:
+ """Resolve the number of threads to use for Numba.
+
+ If ``nthreads`` is given explicitly, return it unchanged. Otherwise
+ return :data:`NUMBA_NUM_THREADS`, resolved once at import time from
+ :data:`.constants.NUMBA_NUM_THREADS_ENVVAR` (``NUMBA_NUM_THREADS``) if
+ set, falling back to :data:`NUM_THREADS` (``OMP_NUM_THREADS``) so Numba
+ and ducc0 agree by default while still allowing independent tuning.
+ """
+
+ if nthreads is not None:
+ return nthreads
+ return NUMBA_NUM_THREADS