From 795e837e8aeab39d3b5299466c2e02d6b18f1abb Mon Sep 17 00:00:00 2001 From: Metin San Date: Thu, 21 Mar 2024 09:23:04 +0900 Subject: [PATCH 1/2] Remove `parallel` keyword, and make parallelization be determined by `n_proc` which now defaults to 1, in which case we have no parallelization. --- docs/examples/get_emission_ang.py | 1 + docs/examples/get_parallel_emission.py | 6 +++--- docs/usage.md | 4 ++-- pyproject.toml | 2 +- tests/_strategies.py | 2 +- tests/test_get_emission.py | 4 ++-- zodipy/zodipy.py | 26 +++++++++++--------------- 7 files changed, 21 insertions(+), 24 deletions(-) diff --git a/docs/examples/get_emission_ang.py b/docs/examples/get_emission_ang.py index 5608af2..146271a 100644 --- a/docs/examples/get_emission_ang.py +++ b/docs/examples/get_emission_ang.py @@ -2,6 +2,7 @@ import matplotlib.pyplot as plt import numpy as np from astropy.time import Time + from zodipy import Zodipy model = Zodipy("dirbe") diff --git a/docs/examples/get_parallel_emission.py b/docs/examples/get_parallel_emission.py index d708896..e5be445 100644 --- a/docs/examples/get_parallel_emission.py +++ b/docs/examples/get_parallel_emission.py @@ -1,4 +1,3 @@ -import multiprocessing import time import astropy.units as u @@ -11,9 +10,10 @@ nside = 256 pixels = np.arange(hp.nside2npix(nside)) obs_time = Time("2020-01-01") +n_proc = 8 model = Zodipy() -model_parallel = Zodipy(parallel=True) +model_parallel = Zodipy(n_proc=n_proc) start = time.perf_counter() emission = model.get_binned_emission_pix( @@ -33,7 +33,7 @@ obs_time=obs_time, ) print( - f"Time spent on {multiprocessing.cpu_count()} CPUs:", + f"Time spent on {n_proc} CPUs:", round(time.perf_counter() - start, 2), "seconds", ) diff --git a/docs/usage.md b/docs/usage.md index b6bc5e5..76ccc9e 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -83,7 +83,7 @@ read [Cosmoglobe: Simulating Zodiacal Emission with ZodiPy](https://arxiv.org/ab ## Parallelization -If you are not using ZodiPy in an already parallelized environment **and** are working with large pointing sequences, setting `parallel=True` when initializing `Zodipy` will improve the performance. ZodiPy will then automatically distribute the pointing to all available CPU's, given by `multiprocessing.cpu_count()` or to `n_proc` if this argument is provided. +If you are not using ZodiPy in an already parallelized environment, you may specify the number of cores used by ZodiPy through the `n_proc` keyword. By default `n_proc` is set to 1. For values of `n_proc` > 1, the line-of-sight calculations are parallelized using the `multiprocessing` module. ```python hl_lines="15 16" {!examples/get_parallel_emission.py!} @@ -102,7 +102,7 @@ If you are not using ZodiPy in an already parallelized environment **and** are w !!! warning "Using ZodiPy in parallelized environments" If ZodiPy is used in a parallelized environment one may have to specifically set the environment variable `OMP_NUM_THREADS=1` to avoid oversubscription. This is due automatic parallelization in third party libraries such as `healpy` where for instance the `hp.Rotator` object automatically parallelizes rotation of unit vectors. - This means that when using ZodiPy with pointing in a coordinate system other than ecliptic, even if `Zodipy` is initialized with `parallel=False`, `healpy` will under the hood automatically distribute the pointing to available CPU's. + This means that when using ZodiPy with pointing in a coordinate system other than ecliptic, even if `Zodipy` is initialized with `n_proc`=1, `healpy` will under the hood automatically distribute the pointing to available CPU's. ## Visualizing the interplanetary dust distribution of a model diff --git a/pyproject.toml b/pyproject.toml index e0b54f7..4c0bdbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,7 @@ ignore = [ "PLR0913", "ISC001" ] -exclude = ["tests/*"] +exclude = ["tests/*", "docs/*"] [tool.ruff.pydocstyle] convention = "google" \ No newline at end of file diff --git a/tests/_strategies.py b/tests/_strategies.py index 4a9b7b7..477f7e6 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -236,5 +236,5 @@ def model(draw: DrawFn, **static_params: dict[str, Any]) -> zodipy.Zodipy: strategies.pop(key) return draw( - builds(partial(zodipy.Zodipy, parallel=False, **static_params), **strategies) + builds(partial(zodipy.Zodipy, **static_params), **strategies) ) diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index 6e276e4..cdee19e 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -252,8 +252,8 @@ def test_multiprocessing() -> None: without multiprocessing. """ - model = Zodipy(parallel=False) - model_parallel = Zodipy(parallel=True) + model = Zodipy() + model_parallel = Zodipy(n_proc=4) observer = "earth" time = Time("2020-01-01") diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 0fbbab7..f7ed6b5 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -62,10 +62,8 @@ class Zodipy: `solar_cutoff` are masked. Defaults to `None`. solar_cut_fill_value (float): Fill value for pixels masked with `solar_cut`. Defaults to `np.nan`. - parallel (bool): If `True`, input pointing will be split among several cores, and - the emission will be computed in parallel. Default is `False`. - n_proc (int): Number of cores to use when `parallel` is `True`. Defaults is `None`, - which uses all available cores. + n_proc (int): Number of cores to use. If `n_proc` is greater than 1, the line-of-sight + integrals are parallelized using the `multiprocessing` module. Defaults to 1. """ @@ -78,8 +76,7 @@ def __init__( ephemeris: str = "de432s", solar_cut: u.Quantity[u.deg] | None = None, solar_cut_fill_value: float = np.nan, - parallel: bool = False, - n_proc: int | None = None, + n_proc: int = 1, ) -> None: self.model = model self.gauss_quad_degree = gauss_quad_degree @@ -88,7 +85,6 @@ def __init__( self.ephemeris = ephemeris self.solar_cut = solar_cut.to(u.rad) if solar_cut is not None else solar_cut self.solar_cut_fill_value = solar_cut_fill_value - self.parallel = parallel self.n_proc = n_proc self._interpolator = partial( @@ -450,18 +446,18 @@ def _compute_emission( **source_parameters["common"], ) - if self.parallel: - n_proc = multiprocessing.cpu_count() if self.n_proc is None else self.n_proc - - unit_vector_chunks = np.array_split(unit_vectors, n_proc, axis=-1) + if self.n_proc > 1: + unit_vector_chunks = np.array_split(unit_vectors, self.n_proc, axis=-1) integrated_comp_emission = np.zeros((len(self._ipd_model.comps), unit_vectors.shape[1])) - with multiprocessing.get_context(SYS_PROC_START_METHOD).Pool(processes=n_proc) as pool: + with multiprocessing.get_context(SYS_PROC_START_METHOD).Pool( + processes=self.n_proc + ) as pool: for idx, comp_label in enumerate(self._ipd_model.comps.keys()): - stop_chunks = np.array_split(stop[comp_label], n_proc, axis=-1) + stop_chunks = np.array_split(stop[comp_label], self.n_proc, axis=-1) if start[comp_label].size == 1: - start_chunks = [start[comp_label]] * n_proc + start_chunks = [start[comp_label]] * self.n_proc else: - start_chunks = np.array_split(start[comp_label], n_proc, axis=-1) + start_chunks = np.array_split(start[comp_label], self.n_proc, axis=-1) comp_integrands = [ partial( common_integrand, From 9e8c211189d0e0e07dbde0a585bc4a68f8bb751b Mon Sep 17 00:00:00 2001 From: Metin San Date: Thu, 21 Mar 2024 09:48:22 +0900 Subject: [PATCH 2/2] Add test for multiprocessing of model with components featuring inner radial cutoffs --- tests/test_get_emission.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index cdee19e..df83a6a 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -245,7 +245,6 @@ def test_invalid_pixel( obs=observer, ) - def test_multiprocessing() -> None: """ Testing that model with multiprocessing enabled returns the same value as @@ -331,6 +330,35 @@ def test_multiprocessing() -> None: assert np.allclose(emission_binned_ang.value, emission_binned_ang_parallel.value) +def test_inner_radial_cutoff_multiprocessing() -> None: + """ + Testing that model with inner radial cutoffs can be parallelized. + """ + + model = Zodipy("RRM-experimental") + model_parallel = Zodipy("RRM-experimental", n_proc=4) + + observer = "earth" + time = Time("2020-01-01") + frequency = 78 * u.micron + nside = 32 + pix = np.random.randint(0, hp.nside2npix(nside), size=1000) + + emission_pix = model.get_emission_pix( + frequency, + pixels=pix, + nside=nside, + obs_time=time, + obs=observer, + ) + emission_pix_parallel = model_parallel.get_emission_pix( + frequency, + pixels=pix, + nside=nside, + obs_time=time, + obs=observer, + ) + assert np.allclose(emission_pix, emission_pix_parallel) @given( model(extrapolate=True),