diff --git a/Python/tests/test_filtering_workers.py b/Python/tests/test_filtering_workers.py new file mode 100644 index 00000000..8fafc208 --- /dev/null +++ b/Python/tests/test_filtering_workers.py @@ -0,0 +1,86 @@ +"""Tests for the threaded FDK/FBP filtering step (CPU only, no GPU call).""" +import numpy as np +import pytest + +import tigre +import tigre.utilities.filtering as filtering_module +from tigre.utilities.filtering import default_workers, filter as make_filter, filtering, nextpow2 +from tigre.utilities.filtering import ramp_flat + + +def _geo(nangles, nv, nu, filt=None): + geo = tigre.geometry(mode="cone", nVoxel=np.array([8, 8, 8]), default=True) + geo.nDetector = np.array([nv, nu]) + geo.dDetector = np.array([0.8, 0.8]) + geo.sDetector = geo.nDetector * geo.dDetector + angles = np.linspace(0, 2 * np.pi, nangles, endpoint=False, dtype=np.float32) + geo.check_geo(angles) + geo.filter = filt + return geo, angles + + +def _proj(nangles, nv, nu): + return np.random.default_rng(nangles + nv + nu).random((nangles, nv, nu), dtype=np.float32) + + +def _plain_filter(proj, geo, angles): + """One projection at a time, float64, no packing, no blocks, no threads.""" + nu = int(geo.nDetector[1]) + filt_len = max(64, 2 ** nextpow2(2 * nu)) + filt = make_filter(geo.filter, ramp_flat(filt_len)[0], filt_len, 1).astype(np.float64) + padding = (filt_len - nu) // 2 + scale = (geo.DSD[0] / geo.DSO[0]) * (2 * np.pi / len(angles)) / (4 * geo.dDetector[1]) + out = np.empty(proj.shape, dtype=np.float64) + for i in range(proj.shape[0]): + padded = np.zeros((proj.shape[1], filt_len)) + padded[:, padding:padding + nu] = proj[i] + filtered = np.fft.ifft(np.fft.fft(padded, axis=1) * filt, axis=1).real + out[i] = filtered[:, padding:padding + nu] * scale + return out + + +@pytest.mark.parametrize("shape", [(6, 40, 64), (7, 40, 64), (1, 3, 50), (9, 1, 200), (4, 300, 256)]) +@pytest.mark.parametrize("filt", [None, "ram_lak", "shepp_logan", "cosine", "hamming", "hann"]) +def test_matches_a_plain_per_projection_filter(shape, filt): + geo, angles = _geo(*shape, filt=filt) + proj = _proj(*shape) + expected = _plain_filter(proj, geo, angles) + out = filtering(proj.copy(), geo, angles, parker=False) + assert out.dtype == np.float32 + np.testing.assert_allclose(out, expected, rtol=0, atol=2e-6 * np.abs(expected).max()) + + +@pytest.mark.parametrize("shape", [(6, 40, 64), (7, 13, 33), (9, 1, 200), (5, 300, 256)]) +def test_same_bits_whatever_the_number_of_workers(shape, monkeypatch): + geo, angles = _geo(*shape) + proj = _proj(*shape) + serial = filtering(proj.copy(), geo, angles, parker=False, workers=1) + for workers in (None, 2, 3, 8): + assert np.array_equal(serial, filtering(proj.copy(), geo, angles, parker=False, + workers=workers)) + # the smallest blocks the function will use: many blocks per projection + monkeypatch.setattr(filtering_module, "BLOCK_BYTES", 1) + for workers in (1, 4): + assert np.array_equal(serial, filtering(proj.copy(), geo, angles, parker=False, + workers=workers)) + + +def test_filters_in_place_and_returns_its_argument(): + geo, angles = _geo(4, 20, 32) + proj = _proj(4, 20, 32) + before = proj.copy() + out = filtering(proj, geo, angles, parker=False) + assert out is proj + assert not np.array_equal(before, proj) + + +def test_error_in_a_worker_thread_reaches_the_caller(): + geo, angles = _geo(8, 600, 512) + with pytest.raises(ValueError): + filtering(_proj(8, 600, 500), geo, angles, parker=False, workers=4) + + +def test_default_workers(monkeypatch): + assert 1 <= default_workers() <= filtering_module.MAX_WORKERS + monkeypatch.setattr(filtering_module, "MAX_WORKERS", 1) + assert default_workers() == 1 diff --git a/Python/tigre/utilities/filtering.py b/Python/tigre/utilities/filtering.py index a6cc53ae..447a5374 100644 --- a/Python/tigre/utilities/filtering.py +++ b/Python/tigre/utilities/filtering.py @@ -5,14 +5,22 @@ import numpy as np from scipy.fft import fft, ifft +import os +import threading import warnings +from concurrent.futures import ThreadPoolExecutor import numpy as np from tigre.utilities.parkerweight import parkerweight +# bytes of padded complex data a thread filters at a time, small enough to stay in cache +BLOCK_BYTES = 2**21 +# threads used when the caller does not choose; more than this was not faster where measured +MAX_WORKERS = 16 + # TODO: Fix parker -def filtering(proj, geo, angles, parker, verbose=False): +def filtering(proj, geo, angles, parker, verbose=False, workers=None): if parker: proj=parkerweight(proj.transpose(0,2,1),geo,angles,parker).transpose(0,2,1) @@ -21,36 +29,65 @@ def filtering(proj, geo, angles, parker, verbose=False): d=1 filt=filter(geo.filter,ramp_kernel[0],filt_len,d,verbose=verbose) - filt=np.kron(np.ones((np.int64(geo.nDetector[0]),1),dtype=np.float32),filt) padding = int((filt_len-geo.nDetector[1])//2 ) scale_factor = (geo.DSD[0]/geo.DSO[0]) * (2 * np.pi/ len(angles)) / ( 4 * geo.dDetector[1] ) - #filter 2 projection at a time packing in to complex container - fproj=np.empty((geo.nDetector[0],filt_len),dtype=np.complex64) - for i in range(0,angles.shape[0]-1,2): - fproj.fill(0) - fproj.real[:,padding:padding+geo.nDetector[1]]=proj[i] - fproj.imag[:,padding:padding+geo.nDetector[1]]=proj[i+1] - - fproj=fft(fproj,axis=1) - fproj=fproj*filt - fproj=ifft(fproj,axis=1) - - proj[i]=fproj.real[:,padding:padding+geo.nDetector[1]] * scale_factor - proj[i+1]=fproj.imag[:,padding:padding+geo.nDetector[1]] * scale_factor + nangles=angles.shape[0] + nrows=int(geo.nDetector[0]) + cols=slice(padding,padding+geo.nDetector[1]) + if workers is None: + workers=default_workers() + + #detector rows are filtered independently of each other: cut them in blocks that fit + #in cache and share the blocks among threads (numpy and scipy.fft release the GIL). + #A multiple of 16 rows, so that the FFT packs the same rows in a SIMD vector as it + #does on a whole projection, and rounds the same way + block=max(16,BLOCK_BYTES//(8*filt_len)//16*16) + local=threading.local() + + def filter_blocks(tasks): + if not hasattr(local,"fproj"): + local.fproj=np.empty((block,filt_len),dtype=np.complex64) + + for i,r0,r1 in tasks: + #filter 2 projection at a time packing in to complex container + fproj=local.fproj[:r1-r0] + fproj.fill(0) + fproj.real[:,cols]=proj[i,r0:r1] + #if odd number of projections filter last solo + if i+11: + with ThreadPoolExecutor(workers) as pool: + #list() so that an exception raised in a thread is raised here + list(pool.map(filter_blocks,batches)) + else: + filter_blocks(tasks) - #if odd number of projections filter last solo - if angles.shape[0] % 2: - fproj.fill(0) - fproj.real[:,padding:padding+geo.nDetector[1]]=proj[angles.shape[0]-1] + return proj - fproj=fft(fproj,axis=1) - fproj=fproj*filt - fproj=np.real(ifft(fproj,axis=1)) - proj[angles.shape[0]-1]=fproj[:,padding:padding+geo.nDetector[1]] * scale_factor - return proj +def default_workers(): + try: + ncpu = len(os.sched_getaffinity(0)) # the cores this process is allowed to use + except AttributeError: # not available on Windows and macOS + ncpu = os.cpu_count() or 1 + return min(ncpu, MAX_WORKERS) def ramp_flat(n, verbose=False):