Performance tuning : alternative to the Do loop

I wonder if anyone could help me to improve the following code. I used the Do loop, which is not recommended by MMA. What I need is the final value of 'fB' after looping over arrB. I think maybe I should use Fold, but I cannot think of any good way. Please note that both arrA and arrB are lists containing 6 real numbers. s is also a real number. The code is as follows.

fB=0;
pre=s^(-2);
Do[
argu=arrB[[i]]*s^2;
Which[
argu<.1, fB=fB+arrA[[i]]*arrB[[i]]*(1-.5*argu),
argu>20, fB=fB+arrA[[i]]*pre,
argu>=.1 && argu<=20,fB=fB+arrA[[i]]*(1-Exp[-argu])*pre
]
,{i,6}];
fB;

I'd also like to know is Piecewise faster than Which in general? Thanks.

• Please include the definition of arrA and arrB and s in your question. – xzczd Jul 9 at 7:25

Let's start by assuming arrA, arrB and s are some random real numbers:

s = RandomReal[];
arrA = RandomReal[1, 6];
arrB = RandomReal[1, 6];

Then a function to be used with Fold can be formulated as follows:

Clear[f]

f[fB_, {a_, b_}] := With[
{argu = b*s^2, pre = s^(-2)},
Piecewise[{
{fB + a*b*(1 - .5*argu), argu < .1},
{fB + a*pre, argu > 20},
{fB + a*(1 - Exp[-argu])*pre, argu >= .1 && argu <= 20}
}]]

Note several keys here:

1. initial value fB as the first argument in f;

2. {a, b} are the Fold-ed extra values.

Now we can use Fold to simply your code:

Fold[f, 0, Transpose[{arrA, arrB}]]
• This solution is even slower than the original for large lists, mainly because f is defined based on pattern matching, which stops Fold from auto-compiling. – xzczd Jul 10 at 7:02

Here is my attempt by casting Boolean arithmetic to integer arithmetic and using vectorization:

The lists in the OP are a bit too short in order to obtain relyable timing results, so I enlarged them.

n = 100000;
s = .5;
pre = s^(-2);
A = RandomReal[{-1, 1}, n];
B = RandomReal[{-1, 1}, n];

First@ RepeatedTiming[
fB = 0;
Do[
argu = B[[i]]*s^2;
fB += Which[
argu < .1,
A[[i]]*B[[i]]*(1 - .5*argu),
argu > 20,
A[[i]]*pre,
argu >= .1 && argu <= 20,
A[[i]]*(1 - Exp[-argu])*pre
],
{i, Lenght[A]}];
fB
]

0.33

Here is my approach:

First@ RepeatedTiming[
argu = B s^2;
b1 = N@UnitStep[argu - 0.1];(*argu\[GreaterEqual].1*)
b2 = N@UnitStep[argu - (1. + \$MachineEpsilon) 20];(*argu>20*)
fB2 = A.Plus[
Subtract[1., b1] Subtract[1., .5*argu] B,
b2 pre,
b1 Subtract[1., b2] Subtract[1., Exp[-argu]] pre
]
]

0.0048

So this is almost 70 times faster.

Checking accuracy:

Abs[fB - fB2]

3.41061*10^-13

With your definitions for the values of s and arrA and arrB try this version using Fold

pre=s^-2;
fun[fB_,{aA_,aB_}]:=fB+
Which[
aB*s^2<.1, aA*aB*(1-.5*aB*s^2),
aB*s^2>20, aA*pre,
.1<=aB*s^2<=20,aA*(1-Exp[-aB*s^2])*pre
];
Fold[fun,0,Transpose[{arrA,arrB}]]

Test this carefully to see if it correctly calculates your result with all variations of your input data.

s = RandomReal[];
arrA = RandomReal[1, 6];
arrB = RandomReal[1, 6];
pre = s^(-2);

Module[{arrA, arrB, argu},
func = Compile @@ {{#, _Real, 1} & /@ {arrA, arrB},
Which[argu < .1, arrA arrB (1 - .5 argu), argu > 20, arrA pre, True,
arrA (1 - Exp[-argu]) pre] /. argu -> arrB s^2 // PiecewiseExpand //
Simplify`PWToUnitStep}]

fB = func[arrA, arrB] // Total

Or

cf = Compile[{{arrA, _Real, 1}, {arrB, _Real, 1}, s}, Module[{fB = 0.,
pre = s^(-2), argu},
Do[argu = arrB[[i]]*s^2;
Which[argu < .1, fB = fB + arrA[[i]]*arrB[[i]]*(1 - .5*argu), argu > 20,
fB = fB + arrA[[i]]*pre, argu >= .1 && argu <= 20,
fB = fB + arrA[[i]]*(1 - Exp[-argu])*pre], {i, Length@arrA}];
fB], CompilationTarget -> C, RuntimeOptions -> "Speed"]

cf[arrA, arrB, s] // AbsoluteTiming

Remember to remove CompilationTarget -> C if you don't have a C compiler installed.