Performance and parallelization suggestions for Mathematica codes

Davi C. Rodrigues, 06.2023

Summary

Context: an example  with physical motivation (just something I am working now)

Simple application of using parallel computing to find a table: ParallelTable.

Introduction to kernels, subkernels and subsessions

Memorization techniques

General recommendations on code efficiency

Parallelization and different loops approaches

Introduction to specific parallelization options and related functions

An example to contextualize

The code

w0 = -1; (*w0 is the standard value for w, the DE equation of state parameter*)
k0 = 3; (*k0 is the standard growth exponent for CBH.*)
tdmax0 = 13.0; (*Gyrs. The standard maximum delay time. 
13.5 Gyrs from GWTC-3 2111.03634v4 (population paper), but we adopt here a slightly smaller value*)
tdmin0 = 0.05; (*Gyrs. The standard minimum delay time. From GWTC-3 population paper*)
zMax0 = 10; (* Highest redshift where BBH start to form. In accordance with GWTC-3, arXiv:2111.03634v4 (population paper) *)

Ωm = 0.32;
H0Gy = 70.0 / (kpc 10^3) gyr; (* With these rounded values, the universe age is 13.2 Gyr*)
gyr = 10^9 60 60 24 365; (*Gyr in seconds*)
kpc = 3.08568 (10^16); (*kpc in km*)

hubble[z_, w_] = H0Gy Sqrt[Ωm(((1+z))^3)+(1-Ωm)(((1+z))^(3(1+w)))]; 

Clear[timeZ];
timeZ[z1_?NumberQ, z2_?NumberQ, w_] := NIntegrate[
  ((((1+z) hubble[z, w]))^(-1)), 
  {z, z1, z2}, 
  Method-> {Automatic, "SymbolicProcessing" -> 0}, (*Improves computational time by a factor 5-10, without changes in the answers.*)
  PrecisionGoal -> 5,  (*Roughtly, this implies that the numerical iteration will stop when 5 digits of precision are achieved.*)
  AccuracyGoal -> Infinity,
  MaxRecursion -> 10  (*A few cases require this MaxRecursion*)
];

zFormation[zObs_?NumberQ, delayTime_?NumberQ, zMax_, w_] := If[
  timeZ[zObs, zMax, w] > delayTime, (*Time between zObs and zMax > td*)
  (*then*)
  FindRoot[
    timeZ[zObs, zformation, w] == delayTime, 
    {zformation, zObs, zMax}, 
    PrecisionGoal -> 5, (*Reminder: AccuracyGoal behaves differently in FindRoot. This is just a minimum precision requirement, probably it will be higher.*)
    Method -> "Brent" (*This is the automatic method for this case, but I am stating it explicitly.*)
  ][[1,2]],
  (*else*)
  zMax  
];

zFormation[1,3, zMax0, w0] // EchoTiming

(10^6) 0.0035 / 60.

This seems to be very fast, doesn’t it? 

But actually if we want to do statistical evations that depend on zFormation, it may be necessary to evaluated it at millions of points several times. Which imply  several hours.

Several hours may not seem that much either. Say, one can start the code in the morning, go to the beach, have lunch with friends, play football in the afternoon and by the end of the day have the results ready for publication. But this is a clearly unrealistic picture. 

In practice it is like run the code, see the results, correct the bugs, see the results,  find some result hard to interpret physically, understand the physical meaning, run everything again with more realistic approximations, improve the numerics, new bugs.... Some day it converges, as we hope.

Physical context

Binary black holes (BBH) are expected to be formed mainly from binary star systems.

After the binary stellar pair becomes a BBH system, it will possibly merge. If this happens in a reasonable time, it can in principle be detected by LIGO/Virgo.

The time between the BBH formation and the merger time is called here delay time. And it is expected to happen between the window of 0.05 Gyr (tdmin0) and 13 Gyr (tdmax0).

There is an additional physical restriction, mergers that are observed at large redshifts need to be bounded by a stricter limit than 13 Gyr. This since it since for sure no BBH were formed before the big bang, and more precisely no BBH formed before the first stars. We take this maximum formation redshift as 10 (zMax0).

The function zFormation[zObs, delayTime, zMax, w] will yield the formation redshift  of a BBH merger that happened at redshift zObs, with a given delayTime and with the maximum redshift zMax. 

Not essential for this presentation, but we also consider the possibility of changes in the dark energy equation of state parameter w.

A good approximation for zFormation: its interpolated version

First some very straightforward examples of ParallelTable

Table[PrimeQ[x],{x,10^1000,10^1000+5000}];//EchoTiming
Table[Pause[0.5],{x, 10}];//EchoTiming

ParallelTable[PrimeQ[x],{x,10^1000,10^1000+5000}];//EchoTiming
ParallelTable[Pause[0.5],{x, 10}];//EchoTiming

The code

ResourceFunction["DarkMode"][]

Directory[]

isComputeTableZformation = True; (*Change to True to compute and regenerate the mx file. About 10 min.*)

listDelayTime = Range[tdmin0/10, tdmax0, 0.1]; (*Low resolution tables for this example*)
listzObs = Range[0.0, 1.0, 0.1];
listw = Range[-1.5, -0.5, 0.1]; 


If[isComputeTableZformation, 
  Echo["Computing tableZformation..."];
  tableZformation = ParallelTable[
    {{zObs, delayTime, w1}, zFormation[zObs, delayTime, zMax0, w1]}, (*this table uses zMax0, zFornationI can restrict that*)
    {zObs, listzObs},
    {delayTime, listDelayTime}, (*It is computed up to tdmax0*)
    {w1, listw}
  ]; // EchoTiming;
  tableZformationFlat = Flatten[tableZformation, 2];
  DumpSave["tableZformation.mx", tableZformationFlat],
  (*else*)
  Get["tableZformation.mx"];
  Echo["tableZformation loaded."];
];

Clear[zFormationIraw, zFormationIfast, zFormationI]
zFormationInoZmax[zObs_, delayTime_, w_] = Interpolation[
  tableZformationFlat,
  InterpolationOrder -> 1 
][zObs, delayTime, w]; (*the "I" in the name stands for interpolated. "raw" for no options.*)

(*zFormationIfast[zObs_, delayTime_, opts:OptionsPattern[optionstkw]] = zFormationIraw[zObs, delayTime, OptionValue@w]; *)
(*the "fast" version does not implement the zMax option, thus zMax = zMax0, about 2-3 times faster*)

SetAttributes[zFormationI, Listable]; (*zFormationInoZmax is automatically listable, but the Min function is not Listable*)
zFormationI[zObs_, delayTime_, zMax_, w_] = Min[
  zFormationInoZmax[zObs, delayTime, w],
  zMax
];

tableZformation[[1,1,1]]

zFormationI[0.2, #, 10, -1] & /@ Range[0.1, 12, 0.1]// EchoTiming
zFormation[0.2, #, 10, -1] & /@ Range[0.1, 12, 0.1]// EchoTiming

And zFormationInoZmax is even faster, since this one is simply the interpolation -- Note, InterpolatingFunction is automatically Listable.

zFormationInoZmax[z, td, w]
zFormationInoZmax[0.2, Range[0.1, 12, 0.1], -1]; //EchoTiming

A brief comment about saving and loading intermediate steps of a code: DumpSave

If part of your code takes more than a couple of minutes to compute, it may be worth to export the intermediate results and then load them latter.

This can be done in several ways. Say, you can use Export to generate a CSV, a TSV, an Excell file or several other formats. CSV and TSV are very light files that can be read anywhere. 

If exporting your results to be read by arbitrary users is not relevant, there other better options. In particular, DumpSave was made for that. It is an efficient format to export and import files within Wolfram Language (and within given CPU type). It generates .mx files.

An exercise

Exercíse:

Consider the function

fun[x_] := Sin[x] + 0.1 RandomVariate[NormalDistribution[]];

Export this function values at the following points,

points = Range[0,10, 0.001];

and using two different formats: .csv and .mx

After that, load the results and do a ListPlot with the data.

Answer...

data = {#, fun[#]} & /@ points;

data // ListPlot

Kernels, subkernels and subsessions

* What is the kernel?

* With one license of Mathematica can we run two kernels at the same computer? (Yes)

* How to configure differenet kernels? -- See Evaluation > Kernel Configuration Options  (By the way, do you know the name of the Kernel that I am using?)

* How to set a notebook to a given kernel? -- See Evaluation > Notebook’s Kernel

* Can we do parallel computation with two kernels? (Yes, but it is not in general a good idea to split a table computation into two or more kernels)

* ParallelTable simply calls several kernels? (Actually it calls several subkernels. There are differences between kernels and subkernels.)

* What is a subsession computation? 

* Can I use a subsession it to compute things faster? (No)

* Can I use it to monitor a long computation in the kernel? (Yes)

* Can I use it to do unrelated computations while the kernel is busy? (Yes)

Set, SetDelayed and memorization

a = 2

It is good to recall that common opperations are functions for Mathematica

FullForm[Hold[a=b]]
Equal[c,d]
Plus[c,d]

Nothing appears.

Clear[a,b]
Set[a, Pause[2]; Echo["End"];]
SetDelayed[b, Pause[2]; Echo["End"];]

b

b

f[x_] = NIntegrate[BesselJ[0, y], {y, 0 ,x}];
Clear[f];
Echo["Now trying with SetDelayed"];
f[x_] := NIntegrate[BesselJ[0, y], {y, 0 ,x}];

Nothing appears because the right hand side was not evaluated.

f[x]

f[1]

To prevent to evaluate f with an arbitrary symbol we use

Clear[f];
f[x_?NumberQ] := NIntegrate[BesselJ[0, y], {y, 0 ,x}];

f[x]

f[1.2]

Definition[f]

Now lets do a new definition for f with Set:

f[2] = 2;

Definition[f]

f[1.999999]

f[2]

The Set definition has preference over the SetDelayed one. This is essential for a convenient memorization structure that Mathematica allows.

Memorization

The simplest memorization is to use somehting like:

fE[3] = f[3]

f[3] // EchoTiming
fE[3] // EchoTiming

But note that we would have to introduce a new definition for each evaluated point, and they would be called with fE, instead of f. This is not very convenient.

Mathematica allows for the following convenient structure, which memorizes the value of f into a Set definition of f.

Clear[f];
f[x_?NumberQ]:=f[x] = NIntegrate[BesselJ[0,y],{y,0,x}];

f[2000] // EchoTiming
f[2000] // EchoTiming
f[2000] // EchoTiming

Definition[f]

A “famous” application of memorization

Clear[fib];
fib[1] = 1;
fib[2] = 2;
fib[n_]:=fib[n-1]+fib[n-2];
fib[30] // EchoTiming

Clear[fib];
fib[1] = 1;
fib[2] = 2;
fib[n_]:=fib[n]=fib[n-1]+fib[n-2];
fib[30] // EchoTiming

Other general recommendations

* For numerical computations in which  MachinePrecision is sufficient, use float point numbers as soon as possible. That is, write “10.” instead of “10”, or 0.5 instead of 1/2. It is  faster to compute with approximate (float point) numbers than with exact numbers.

N[Det[Table[1/(1+Abs[i-j]),{i,1,150},{j,1,150}]]]//AbsoluteTiming

Det[Table[1/(1.+Abs[i-j]),{i,1.,150.},{j,1.,150.}]]//AbsoluteTiming

* If there is a specific Mathematica function for your task, use it. 

It will be more efficient than trying to code it within Mathematica. Built-in functions are commonly directly compiled in C. 

You should build functions that do things that a single Mathematica function cannot do.

Other general recommendations (continuation)

* NEVER TRY TO MIMICK Python PROGRAMMING. 

Actually, or any other language. 

Each language has its own programming paradigms that makes it more efficient to compute and easy to read.  

Mathematica (for the good or the bad) is particularly idiossincratic language with a very different approach. 

By they way, it is mostly rooted on the Functional programming and Rule based mannipulation paradigms. Python is far from a pure functional paradigm. 

About rule based manipulations, try to find the equivalient of these examples in Python. For sure it is possible to code them, but everything below works immediately in Matheamtica, and it is very close to its fundamentals.

Clear[f]
f[x_] = 1/(1+Sin[x^2]); 
f[Plot[x, {x,0,1}]]

f[]

Clear[a,b,c]
{a,b,c}/. List->Plus
{a,b,c}/. List->Sequence
  /.{AbsoluteThickness[_] ->  AbsoluteThickness[10], RGBColor[__] -> RGBColor[0.5, 0.5,0.4], List[0.4900439105662921`,0.4900439105662921`] -> List[0.4900439105662921`,1]}

This is not to show that Mathematica is better, it is just to show that it has its own procedures to code. There are several examples of codes in Python that are not trivial to be written within the Mathematica code paradigm.

Here is a cell of Python in a Mathematica notebook that shows that a + 2 is not trivial for Python for arbitrary a:

import numpy as np
print(np.array([1,2,3]))
print("As you can see, numpy works here. But without simpy, a + 2 does not make sense...")
a+2

a = 2

a + 2

It is also possible to use Jupyter (or Jupyter-lab or VS Code) to run Mathematica. This can be done freely, without a Mathematica license. But... The best notebook for Mathematica is clearly its own. Since we have a license, it is better to use it.

Other general recommendations (continuation)

* Compile can be useful in specific situations, particularly for combining a large amount of trivial mathematical operations. But first try to find a specific Mathematica function. Compile is not useful if it includes Mathematica functions that cannot be promptly translated into C functions. Say, don’t try to compile things like FindRoot, NDSolve, Simplify...

Other general recommendations (continuation)

* ALWAYS use DarkMode.  :)

Ok, this is a fully personal opinion that is useful for me.

But the “recent” function DarkMode is really nice. 

Also, the default cells are code cells, instead of the Input ones. I find better to use the code cells, less unsolicited authomatic stuff, lighter and easier to be saved as .wl files.

ResourceFunction["DarkMode"]["SetDefault"]

Comments on different loop efficiencies

Efficiency of Do, Table, ParallelDo and ParallelTable

Efficiency is not easy to predict. Before doing a large time consumming computation, it is important to test with a subset of variables.

$VersionNumber

$KernelCount

t[n_]:=First[(Table[i, {i,n}]; // AbsoluteTiming)]
d[n_]:=First[(Do[i, {i,n}]; // AbsoluteTiming)]

tab = Table[t[n],{n,1000,1000000, 1000}];
dtab = Table[d[n],{n,1000,1000000, 1000}];

tabExt = Table[t[n],{n,(10^6),(10^8), (10^6)}];
dtabExt = Table[d[n],{n,(10^6),(10^8), (10^6)}];

tabExtF = Flatten @ Join[{tab, tabExt}];
dtabExtF = Flatten @ Join[{dtab, dtabExt}];

Use alt+shift+enter (Mac) to evaluate in subssession the current status of n, while the table above is being computed.

Clear[pt, pd]
pt[n_]:=First[(ParallelTable[i, {i,n}]; // AbsoluteTiming)]
pd[n_]:=First[(ParallelDo[i, {i,n}]; // AbsoluteTiming)]

ptab = Table[pt[n],{n,1000,(10^6), 1000}];
pdtab = Table[pd[n],{n,1000,(10^6), 1000}];

ptabExt = Table[pt[n],{n,(10^6),(10^8), (10^6)}];
pdtabExt = Table[pd[n],{n,(10^6),(10^8), (10^6)}];

ListPlot[{tab, dtab, ptab, pdtab}, PlotLegends->{"Table", "Do", "ParallelTable", "ParallelDo"}]

ListPlot[{tabExt, dtabExt, ptabExt, pdtabExt}, PlotLegends->{"Table", "Do", "ParallelTable", "ParallelDo"}]

Another variation, from https://mathematica.stackexchange.com/questions/134609/why-should-i-avoid-the-for-loop-in-mathematica

Conclusions: 
* Efficiency is not something easy to predict. I have not tested this type of loop with For, but probably it would be worse than Do... 
* The communication time between the parallel kernels and the main kernel is relevant if the parallallel kernels computations is too fast.
* The bump above in the ParallelDo was probably trigered by the status display of the parallel computations, it can be turned off.

Attention: This is not to be seen as an arguemnt in favor of using Table in place of Do. This difference in time is commonly negligible, since the computation of each Table or Do element is commonly much more time demanding than each iteration time. 

It is  more important to decide using Table or Do based on the clarity of the code. If the purpose is to build a matrix/table, go for Table. If it is just a loop, Do is preferred. While is also another loop structure, its use depend on the purpose of the computation (it is expected to be slower than Do).

PS : this approach is based on an old notebook of mine. It works, but there are more efficient ways of comparing the efficiency. ;)

Comments on different loop efficiencies (continuation...)

What about For?

What about For in Mathematica? Essentially, forget it. 

Probably the best application of For is for durty and fast translation of a code from another language. 

For breaks the functional paradigm, and it is not easier to read than Do or While.   

Further reading: https://mathematica.stackexchange.com/questions/134609/why-should-i-avoid-the-for-loop-in-mathematica

Clear[k];
Do[Print@k, {k,5}];
Echo[k, "After Do, k = "];
For[k=1, k<=  5, k++, Print[k]];
Echo[k, "After For, k = "];

Comments on different loop efficiencies (continuation...)

Range with Map

Instead of

f[x_] := Print[x];
Do[f[x], {x, 1, 5}];

You may try

f /@ Range@5; (*This is equivalent to Map[f, Range[5]];*)

If necessary, you may also adapt the steps used in Range, for instance,

Range[3, 7, 0.7]

Listable function

SetAttributes[f, Listable];
f[Range@5];

If it is a function of several arguments, either use Listable or use pure functions instead:

g[x_, y_] := Print[x+y];
g[2, #] & /@ Range[5];

SetAttributes[g, Listable];
g[2, Range@5];

If you want to clear Attributes of a function, don’t forget to use ClearAll instead of Clear.

Clear[g];
g[2, Range@5]
ClearAll[g];
g[2, Range@5]

Further discussion: https://mathematica.stackexchange.com/questions/35150/two-ways-to-map-a-function-on-a-list-which-one-is-faster

Warning: if a given function is by default Listable,  use its property directly. Built in Listable functions are faster than user-set Listable functions.

Example: interpolation functions are by default Listable. So, if f[x] is an interpolation function and you want to apply it in several points, f[Range[...]] is faster than other options.

ClearAll[f];
f[y_] := NIntegrate[(Sqrt[(x^2) + x Sin[x]]), {x, 0, y}];
dataf = Table[{x, f[x]}, {x, 1, 20, 1}];
fI = Interpolation[dataf]

f[2.]

EchoTiming[Table[f[i], {i, 1, 20, 0.01}], "f with step 0.01: "];  

EchoTiming[Table[fI[i], {i, 1, 20, 0.00001}];, "fI with step 0.000001: "];
EchoTiming[fI[Range[1, 20, 0.00001]];,  "fI with step 0.000001 and using Listable: "];

More advanced parallelization options and functions

ParallelTable considerations and commonly related functions:
   * DistributeDefinitions  (if not all subkernels are working, probably some definitions were not distributed)
   * ParallelEvaluate
   * SetSharedFunction (avoid using whenever possible)
   * SetSharedVariable (avoid using whenever possible)
   * CoarsestGrained  vs FinestGrained
   * CloseKernels[] -- It is useful to clear the subkernels without quiting the main Kernel. Ater that, probably you will want to execute LaunchKernels[] and DistributeDefinitions.
 
Other functions not directly related to ParallelTable

* ParallelTry: Stops after the first result is found. It can also take the first k results.
  Example of application. There is no way to parallelize a single realization of NMinimize. But one can do several NMinimize in parallel with different meta parameters for the minimization. And them use ParallelTry to pick the first result that converges. The work of the other subkernels is aborted. 
 
ParallelSubmit & WaitNext: Can wait iterativelly for the subkernels results, that is, after the first result is found, the subkernels just stop, they are not aborted. It has different kind of parallelization monitor.

tableSubmitted = Table[ParallelSubmit[{i},Length[FactorInteger[2^i-1]]],{i,181,184}]

WaitAll[tableSubmitted]

{res, id, tableSubmitted} = WaitNext[tableSubmitted]

Parallelize (Parallelize[Table[...]] = ParallelTable, but avoid using Parallelize if a more specific version solves your problem)

ParallelMap
 
An example of nontrivial parallelization issue: Parallelization and InterpolationFunction (https://mathematica.stackexchange.com/questions/13818/using-interpolation-functions-in-parallelmap)