Merge pull request #777 from HEXRD/parallelize-snip1d

Parallelize rows in snip1d

Author: saransh13

hexrd/fitting/fitpeak.py

@@ -119,7 +119,7 @@ def estimate_pk_parms_1d(x, f, pktype='pvoigt'):
 
     # handle background
     # ??? make kernel width a kwarg?
-    bkg = snip1d(np.atleast_2d(f), w=int(2*npts/3.)).flatten()
+    bkg = snip1d(np.atleast_2d(f), w=int(2*npts/3.), max_workers=1).flatten()
 
     # fit linear bg and grab params
     bp, _ = optimize.curve_fit(lin_fit_obj, x, bkg, jac=lin_fit_jac)
@@ -372,8 +372,11 @@ def estimate_mpk_parms_1d(
     min_val = np.min(f)
 
     # estimate background with SNIP1d
-    bkg = snip1d(np.atleast_2d(f),
-                 w=int(np.floor(0.25*len(f)))).flatten()
+    bkg = snip1d(
+        np.atleast_2d(f),
+        w=int(np.floor(0.25*len(f))),
+        max_workers=1,
+    ).flatten()
 
     # fit linear bg and grab params
     bp, _ = optimize.curve_fit(lin_fit_obj, x, bkg, jac=lin_fit_jac)

hexrd/fitting/spectrum.py

@@ -204,8 +204,11 @@ def _initial_guess(peak_positions, x, f,
 
     # estimate background with snip1d
     # !!! using a window size based on abcissa
-    bkg = snip1d(np.atleast_2d(f),
-                 w=int(np.floor(len(f)/num_pks/2.))).flatten()
+    bkg = snip1d(
+        np.atleast_2d(f),
+        w=int(np.floor(len(f)/num_pks/2.)),
+        max_workers=1,
+    ).flatten()
 
     bkg_mod = chebyshev.Chebyshev(
         [0., 0.], domain=(min(x), max(x))

hexrd/imageutil.py

@@ -1,3 +1,7 @@
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+import os
+
 import numpy as np
 from scipy import signal, ndimage
 
@@ -60,7 +64,7 @@ def fast_snip1d(y, w=4, numiter=2):
     return bkg
 
 
-def snip1d(y, w=4, numiter=2, threshold=None):
+def snip1d(y, w=4, numiter=2, threshold=None, max_workers=os.cpu_count()):
     """
     Return SNIP-estimated baseline-background for given spectrum y.
 
@@ -79,29 +83,47 @@ def snip1d(y, w=4, numiter=2, threshold=None):
         mask = np.zeros_like(y, dtype=bool)
 
     # step through rows
-    for k, z in enumerate(zfull):
-        if np.all(mask[k]):
-            bkg[k, :] = np.nan
-        else:
-            b = z
-            for i in range(numiter):
-                for p in range(w, 0, -1):
-                    kernel = np.zeros(p*2 + 1)
-                    kernel[0] = kernel[-1] = 1./2.
-                    b = np.minimum(
-                        b,
-                        convolution.convolve(
-                            z, kernel, boundary='extend', mask=mask[k],
-                            nan_treatment='interpolate', preserve_nan=True
-                        )
-                    )
-                z = b
-            bkg[k, :] = _scale_image_snip(b, min_val, invert=True)
+    tasks = enumerate(zip(zfull, mask))
+    f = partial(_run_snip1d_row, numiter=numiter, w=w, min_val=min_val)
+
+    if max_workers > 1:
+        # Parallelize over tasks
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            for k, result in executor.map(f, tasks):
+                bkg[k, :] = result
+    else:
+        # Run the tasks in this process
+        for task in tasks:
+            k, result = f(task)
+            bkg[k, :] = result
+
     nan_idx = np.isnan(bkg)
     bkg[nan_idx] = threshold
     return bkg
 
 
+def _run_snip1d_row(task, numiter, w, min_val):
+    k, (z, mask) = task
+
+    if np.all(mask):
+        return k, np.nan
+
+    b = z
+    for i in range(numiter):
+        for p in range(w, 0, -1):
+            kernel = np.zeros(p*2 + 1)
+            kernel[0] = kernel[-1] = 1./2.
+            b = np.minimum(
+                b,
+                convolution.convolve(
+                    z, kernel, boundary='extend', mask=mask,
+                    nan_treatment='interpolate', preserve_nan=True
+                )
+            )
+        z = b
+    return k, _scale_image_snip(b, min_val, invert=True)
+
+
 def snip1d_quad(y, w=4, numiter=2):
     """Return SNIP-estimated baseline-background for given spectrum y.
 

tests/data/expected_snip1d_results.npy

@@ -0,0 +1,69 @@
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+import pytest
+
+from hexrd import imageutil
+from hexrd.instrument.hedm_instrument import HEDMInstrument
+from hexrd.projections.polar import PolarView
+
+
+@pytest.fixture
+def simulated_tardis_path(example_repo_path: Path) -> Path:
+    return example_repo_path / 'tardis' / 'simulated'
+
+
+@pytest.fixture
+def simulated_tardis_images(
+    simulated_tardis_path: Path,
+) -> dict[str, np.ndarray]:
+    path = simulated_tardis_path / 'tardis_images.npz'
+    npz = np.load(path)
+    return {k: v for k, v in npz.items()}
+
+
+@pytest.fixture
+def tardis_instrument(simulated_tardis_path: Path) -> HEDMInstrument:
+    path = simulated_tardis_path / 'ideal_tardis.yml'
+    with open(path, 'r') as rf:
+        conf = yaml.safe_load(rf)
+
+    return HEDMInstrument(conf)
+
+
+@pytest.fixture
+def expected_snip1d_results(test_data_dir: Path) -> np.ndarray:
+    path = test_data_dir / 'expected_snip1d_results.npy'
+    return np.load(path)
+
+
+def test_snip1d(
+    tardis_instrument: HEDMInstrument,
+    simulated_tardis_images: dict[str, np.ndarray],
+    expected_snip1d_results: np.ndarray,
+):
+    instr = tardis_instrument
+    img_dict = simulated_tardis_images
+    ref = expected_snip1d_results
+
+    # Create the PolarView
+    tth_range = [10, 120]
+    eta_min = -180.0
+    eta_max = 180.0
+    pixel_size = (0.1, 1.0)
+
+    pv = PolarView(tth_range, instr, eta_min, eta_max, pixel_size)
+    img = pv.warp_image(img_dict, pad_with_nans=True,
+                        do_interpolation=True)
+
+    snip_width = 100
+    numiter = 2
+    output = imageutil.snip1d(
+        img,
+        snip_width,
+        numiter,
+    )
+
+    assert np.allclose(output.filled(np.nan), ref, equal_nan=True)