Parallelize for-loops using numba.prange

[arunoruto]

Amazing work with this framework, and I am really happy to see that you used numba to speed up the computation. I was wondering if there is a reason why the prange function of numba has not been used to parallelize some for-loops. They have a small doc page about prange and its use.

By simply skimming the code, I found some possible places where prange could be used:

• _mie_scalar
• mie
• _mie_S1_S2
• mie_mu_with_uniform_cdf
  and possibly here too:
• _mie_An_Bn

If there is an interest in implementing some prange functionality, I could try to make some changes. It would be nice also to have some standard tests, so the framework can be benchmarked and tested so that the output hasn't altered by the modifications made. Do you have any data for that currently?

[scottprahl]

I am very interested in speeding things up with prange! I did not know it existed.

Some of the loops may not be amenable to parallelizing because earlier array elements affect later elements.

I would not worry about mie_mu_with_uniform_cdf() since I would be surprised if anyone is actually using it.

There are tests for accuracy in the tests folder. However there are not any tests for performance. The closest is the Jupyter notebook https://miepython.readthedocs.io/en/latest/11_performance.html . It should be straightforward to adapt some of that code to a test harness.

Looking forward to seeing a pull request.

[andrej5elin]

Hi Scott,

nice work! I am glad I found your code for mie calculation, very useful.

I have actually done the parallelization for my own needs. Please see the code in the test branch of my fork. I am willing to request a pull, if you find my implementation useful.

The way to do it is to create universal functions using @guvectorize decorator. As a bonus, the code gets vectorized automatically and run on multiple cores. For example, I can do

mu = np.linspace(0,1,1000)
x = np.linspace(1,2,500)
S1, S2 = S1_S2(2.,x,mu)

The resulting arrays will have shapes of (500,1000) and the computation is done in parallel if utilizing numba's target = "parallel" option.

Cheers,

Andrej

[scottprahl]

This is very interesting. Thanks

One day, I will incorporate your idea. However, I want to get the E-field stuff released first before making everything faster.

[scottprahl]

Sorry for the long delay on this. The near-field work I wanted to finish first is
released now (3.3.0), so I went back and actually measured what prange would buy.
Short version: there is a real speedup available, but not in the loops listed above,
and the change as written would silently produce wrong answers. Details below,
because I think they're worth having on the record.

The code moved. Those line references point into miepython/miepython.py, which
no longer exists — the kernels now live in mie_jit.py (numba) and mie_nojit.py
(pure NumPy), with core.py on top and _backend.py choosing between them at import.
Mapping the five suggestions onto the current layout:

Suggested	Now	Can it take prange?
_mie_An_Bn	_an_bn_nb	No — it's a recurrence
_mie_scalar	_single_sphere_nb	Only a ~10–30 term reduction; too small to matter
_mie_S1_S2	_S1_S2_nb	Yes, over angles — but see the race below
mie	core.efficiencies_mx	Loop is in Python, not numba, so prange doesn't apply
mie_mu_with_uniform_cdf	monte_carlo.mu_with_uniform_cdf	Pure NumPy, never compiled

_an_bn is the important one to rule out: it carries chi_nm1 = chi_n from one
iteration to the next, so iteration n cannot start until n−1 finishes. Same for
the ψ and D_n recurrences it depends on. That's not a limitation of the
implementation, it's the algorithm — those loops can't be parallelized at all.

The race. _S1_S2_nb allocates its pi and tau buffers once outside the angle
loop and has _pi_tau_nb overwrite them on every iteration. Replacing range with
prange there lets every thread write those same two arrays at once. I tried it
(66 terms, 200,000 angles):

• wrong answers on 5 of 5 runs
• worst relative error ~590 — not round-off, completely wrong numbers
• not reproducible from one run to the next

Nondeterministic wrong answers are the failure mode I least want in this package, and
tolerance-based tests would pass on a lucky run. The fix is to move the buffers inside
the loop so each iteration gets its own, which is what I did before timing anything.

What the speedup actually is, done correctly, on 12 cores, with results bitwise
identical to serial:

Workload	Serial	prange
100,000 spheres	141 ms	20.6 ms	6.8×
S1_S2, 100,000 angles, x=10	14.9 ms	3.1 ms	4.8×
S1_S2, 100,000 angles, x=200	189 ms	23.6 ms	8.0×

So you were right that there's something here — I'd assumed less.

The catch is small inputs. There's a fixed ~76 µs cost per parallel call, and the
crossover is around 100 elements:

n=1      0.9 µs serial  →  77 µs prange   0.01×
n=10    14.5 µs         →  77 µs          0.19×
n=100  146 µs           → 132 µs          1.11×
n=500  720 µs           → 213 µs          3.38×

Scalar calls are the common case here, and they'd get ~100× slower. So this needs a
length threshold rather than being switched on unconditionally.

Where I'd want it, if it goes in. Not prange inside the existing kernels, but a
new array-at-a-time driver reached through _backend.py: a plain loop in
mie_nojit.py and a prange loop in mie_jit.py, with core.py picking it up for
long inputs and staying on the scalar path otherwise. That keeps core.py
backend-agnostic, which matters because those two kernel files have to stay
numerically identical and there's a parity test suite enforcing it. (For what it's
worth, cache=True and parallel=True do coexist fine — that was my main worry going
in, and it turned out not to be a problem.)

On your question about tests and benchmark data — yes, that all exists now, and
it didn't when you asked:

• make speed runs tests/benchmark_efficiencies.py, which times exactly the
  100,000-particle ensemble you'd want for this, once per backend and then both
  against each other in one process.
• tests/test_backend_parity.py compares the two kernel sets head-on, signatures and
  values, so any change that makes them disagree fails loudly.
• make test runs the suite once per backend (545 tests, 100% branch coverage),
  including Wiscombe's MIEV0 reference cases and Bohren & Huffman values.

That's enough to check both "still correct" and "actually faster" for a change like
this, so the tooling is no longer the blocker it was in March.

If you'd still like to take a run at it, the array driver is the piece I'd welcome —
with the length threshold, and with the parity and MIEV0 tests passing on both
backends. I'd rather not merge prange into the recurrences or into _S1_S2 as it
stands.

@andrej5elin — your @guvectorize suggestion is a different and in some ways more
interesting proposal, since broadcasting over m, x and mu together would be a
real feature and not just a speedup (right now n_pole has to be nonzero when m or
x is an array). That changes the shape of what the public functions return, though,
so I'd want to treat it as its own discussion rather than fold it in here. If you're
still willing to open that PR against current main, I'd like to look at it.