Metropolis-Hastings Algorithm - Significantly slower than Python

Question

I would like to code a Metropolis Hasting algorithm and thought I would get the general scheme of things going with a toy example. I think I got it running but my Mathematica code runs significantly slower than a quick Python version I sketched, with just basic vectorised operations. I would really like to hear your feedback on what could be done better. I am sure the difference would be minimal for properly written Mathematica code.

For the toy example, I am simulating the motion of a particle in a double well $ H = -x^2 + x^4$. moving by discrete steps of length eps. The particle starts at the left well $x = -1 / \sqrt{2}$.I am interested in the time it takes the particle to get over the position $x=0$. I am simulating two different temperature cases, and I would like to get many different realisations (defined in the variable nrealis below).

The code is as follows. I add some comments, would be glad to add more of course.

(* auxiliary variables *)
nrealis = 100;
ntemp = 2;
tem = List [0.05, 0.06];
eps = 0.1;
(* energy of the particle *)
energy[x_] := -x^2 + x^4;
(*get a random integer in {-1,0,1} to define next candidate step*)
candMov[] := RandomInteger[ {-1, 1}]*eps;
(*function adding current displacement to picked step*)
candPos[x_, y_] := x + y;
(* calculating Metropolis Hasting transition probability, based on the energy of current and candidate particle position *)
alfa[T_, xOld_, xNew_] := 
 Min[1, Exp[-(1/T)*(energy[xNew] - energy[xOld])]]
(* decide if the step is to be made, by comparing MH transition probability to a value picked from a uniform distribution in (0,1) *)
newPos[beta_, alfa_, xNew_, xOld_] := If[beta < alfa, xNew, xOld];

(*I am interested in the number of iterations for the particle to reach $x=0$ for the first time. The following function applies to the current position, and if it exceeds 0 for the first time it sets the function result to the time counter in the main body of the code. *)

timings[pos_, meta_, time_] := 
 If[pos > 0, If[meta != 0, meta, time], meta]

(* initialising vector of current positions. I try to calculate all the realisations, of which there are ntemp * nrealis, in one go *)
xOld = Table[-1 / Sqrt[2], {i, 1, ntemp*nrealis}];
(*initialising iteration counter*)
time = Table[0, {i, 1, ntemp*nrealis}];
(* initialising results vector, i.e. time for first x >0 excursion *)
metaTime = Table[0, {i, 1, ntemp*nrealis}];
(*create a vector of temperatures to use with MapThread, the first nrealis elements get assigned the first temperature, the second half, the second temperature value*)
tempvec = 
  Table[ tem[[Quotient[i, nrealis] + 1]], {i, 0, ntemp *nrealis - 1}];
(* Main Body
 calculate new candidate positions
 calculate MH transition probabilities
calculate uniform probabilites
decide if to perform transition
update current positions vector
update results vector
continue until all particles have hit 0
*)
Timing[While[MemberQ[metaTime, 0], 
   xNew = MapThread[
     candPos, {Table[candMov[], {i, 1, ntemp*nrealis}], xOld}]; 
   alfas =  MapThread [alfa, {tempvec, xOld, xNew}]; 
   betas =  RandomReal[{0, 1}, ntemp*nrealis]; 
   xOld = MapThread[newPos, {betas, alfas, xNew, xOld}]; 
   metaTime =   MapThread[timings, {xOld, metaTime, time}]; time++];]

The Python code (which I would gladly share of course) runs the job with the current temperature values in 1.04 seconds, Mathematica needed 62 seconds.

I would really be grateful for any hints or suggestions, thanks in advance.

ADDENDUM: Python reference code

Please find enclosed the Python code I mentioned and used for the comparison. I did not attempt a one-to-one translation. this code is for sure sub-optimal, that is why I was surprised to be a factor 100 slower in Mathematica, in my first attempt. Thanks to the excellent answer a lot of improvements were made, but I am still not matching it (I do not have a C compiler so I cannot use that option in the Compile function, I guess some improvements are to be reaped there as well).

import numpy as np
import random
from random import randrange

def energy (x):
    x = np.asarray(x)
    energy = -np.power(x,2) + np.power(x,4)
    return energy

## creating the vector named tempvec in Mathematica ##
def tempInit(nres, T):
    temp = np.zeros(nres * len(T))
    for i in range (len(T)):
        for j in range(nres):
            a = j + nres*(i)
            temp [a] = T[i]
    return temp

## Initial position vector 
def PosInit (nres, ntemp):
    initpos = np.zeros(nres*ntemp) - 1/ np.sqrt(2)
    return initpos

## candidate position ##
def newPos (x,eps, nres, ntemp):
    x = x + np.random.randint(-1,2, size=nres*ntemp) * eps
    return x
## equivalent of NewPos function in Mathematica code (with the alfa and beta vectors calculated internally, differently from Mathematica)

def acceptance(xOld, xNew, T, nres, ntemp):
    energy_old = energy(xOld)
    energy_new = energy (xNew)
    alfa = np.minimum(1, np.exp(-(1/T)*(energy_new-energy_old) ) )
    beta = np.random.uniform(0,1, nres * ntemp )
    xOld = np.where(beta < alfa, xNew, xOld)
    return  xOld

## Main body 
def main (T, nrealis, eps):
    ntemp = len(T)
    poszero = np.zeros(nrealis * ntemp)
    time = np.zeros(nrealis * ntemp)
    metatime = np.zeros(nrealis * ntemp)
    actual = PosInit(nrealis, ntemp)
    tempe = tempInit (nrealis, T)
    while np.any(metatime == 0):
                candidate = newPos(actual,eps, nrealis, ntemp)
                actual = acceptance(actual, candidate, tempe, nrealis, ntemp)
                metatime = np.where(actual > poszero , np.where(metatime !=0, metatime, time), metatime )
                time = time + 1

## Temperature Values ##
temparr = [0.05,0.06]

## Launching run ##
res = main(temparr, 100,0.1)

EDIT

The final version by HenrikSchumacher is amazingly fast, but I am not entirely sure it does the same computation my naive Python and Mathematica codes, as results are quite different. Just for completeness I will add a couple of lines to postprocess the results and show the difference. In short, the escape time over an energetic barrier is simulated at two temperatures, for an ensemble of particles. The barrier height equals the energy difference between $x = -1 / \sqrt{2}$ and $x=0$ and equals $\Delta E = (1/ \sqrt{2})^2 -(1/ \sqrt{2})^4 = 1/4$. On theoretical grounds the expected escape time scales as $\exp {-\Delta E / T}$, for sufficiently low temperatures. This can be checked by computing the mean of the logarithm of the result vector metaTime for each temperature, and checking their gradient versus $1/T$. For example with

TLowRes = metaTime[[1 ;; nrealis]];
THighRes = metaTime[[nrealis + 1 ;; 2*nrealis]];
LogTLowRes = Log[TLowRes];
LogTHighRes = Log[THighRes];
(Mean[LogTLowRes] - Mean[LogTHighRes]) / (1/tem[[1]] - 1/tem[[2]])

I get values reasonably close to 0.25 for both the naive code versions in my post, while I get much lower values for the code in the answer. I am not implying it is wrong of course, mine might well be, I just cannot pinpoint where the difference is, as the general logic seems the same. Also, the running time should increase with the exponential of the inverse temperature, while it does not seem the case at all.

Answer 1

Here a bit of inspiration for you. The following code should be approximately as fast as your Python code.

Preparation:

nrealis = 100;
ntemp = 2;
nn = ntemp nrealis;
tem = Developer`ToPackedArray[{0.05, 0.06}];
eps = 0.1;
candmoves = Developer`ToPackedArray[{-eps, 0., eps}];
xInit = ConstantArray[-1./Sqrt[2.], {nn}];
tmp = (-1.)/Developer`ToPackedArray[Table[tem[[Quotient[i, nrealis] + 1]], {i, 0, nn - 1}]];

Running the simulation:

min = 0;
time = 0;
metaTime = ConstantArray[0, {nn}];
xOld = xInit;
enOld = With[{square = xOld xOld}, (square - 1.) square];
AbsoluteTiming[
 While[min == 0, xNew = RandomChoice[candmoves, nn] + xOld;
   enNew = With[{square = xNew xNew}, (square - 1.) square];
   alfas = Exp[tmp Ramp[Subtract[enNew, enOld]]];
   betas = RandomReal[{0., 1.}, nn];
   With[{λ = UnitStep[Subtract[alphas, betas]]},
     xOld = λ xNew + (1 - λ) xOld;
     enOld = λ enNew + (1 - λ) enOld;
     ];
   metaTime = With[{
       λ = UnitStep[-xOld],
       μ = Unitize[metaTime]
       },
      Subtract[1, λ] (μ metaTime + Subtract[1, μ] time) + λ metaTime
      ];
   min = Min[metaTime];
   ++time
   ];
 ]

On my machinge, this takes about 0.3 seconds on average (measured with RepeatedTiming). When I take OP's Python code, store it as string in the symbol code, and call Python from Mathematica with

ExternalEvaluate["Python", code]

it takes 3 seconds on average.

Here some of the crucial steps:

• Use Developer`ToPackedArray to make sure to work with packed arrays.

• Don't use MapThread if you want to exploit vectorizations (as it is built in, e.g., for Plus, Times, Exp,...)

• Compute energies of each state only once (and reuse them in the next iteration).

• Use Clip in order to get a vectorized version of Min[1,#]&. (I found this one quite unobvious.) I found out that Ramp is more suitable here. In contrast to Clip, it is compilable and it probably sets many in input values of Exp to 0., so that Exp can short circuit to return 1..

• Use UnitStep and Unitize to generate 0/1 vectors λ and (1 - λ) instead of If to enable vectorization. (There should actually be a better way to a vectorized If, but I don't know any.)

I have not yet used any parallelism here. So there is certainly room for improvement.

You can get out a bit more speed by wrapping the code in Compile (and by replacing -∞ by a really small number):

cf = Compile[{{xInit, _Real, 1}, {tmp, _Real, 1}, {eps, _Real}},
   Block[{candmoves, enOld, xOld, enNew, xNew, nn, alphas, betas, 
     metaTime, min, time},
    time = 0;
    min = 0;
    candmoves = {-eps, 0., eps};
    nn = Length[xInit];
    metaTime = Table[0, {nn}];
    xOld = xInit;
    enOld = With[{square = xOld xOld}, (square - 1.) square];
    While[min == 0,
     xNew = RandomChoice[candmoves, nn] + xOld;
     enNew = With[{square = xNew xNew}, (square - 1.) square];
     alphas = Exp[tmp Ramp[Subtract[enNew, enOld]]];
     betas = RandomReal[{0., 1.}, nn];
     With[{λ = UnitStep[Subtract[alphas, betas]]},
       xOld = λ xNew + (1 - λ) xOld;
       enOld = λ enNew + (1 - λ) enOld;
       ];
     metaTime = With[{
         λ = UnitStep[-xOld],
         μ = Unitize[metaTime]
         },
        Subtract[1, λ] (μ metaTime + Subtract[1, μ] time) + λ metaTime
        ];
     min = Min[metaTime];
     ++time
    ];
    metaTime
   ],
   CompilationTarget -> "C",
   RuntimeAttributes -> {Listable},
   Parallelization -> True,
   RuntimeOptions -> "Speed"
   ];

Then you can use it as follows:

cf[xInit, tmp, eps]

It is up to 50% faster, but getting the results is a bit more difficult because CompiledFunctions can return only a single tensor. However, now that I have remove Clip, we can use the options RuntimeAttributes -> {Listable} and Parallelization -> True to run several experiments in parallel. On a 4 Core machine as mine, we would just generate 4 datasets like this

corecount = 4;
xInit = ConstantArray[-1./Sqrt[2.], {corecount, nn}];
tmp = ConstantArray[
   (-1.)/Developer`ToPackedArray[
     Table[tem[[Quotient[i, nrealis] + 1]], {i, 0, nn - 1}]],
   {4}];

and then run

 cf[xInit, tmp, eps];

The runtime is essentially the same as above, but now four times the data is processed.

Disclaimer:

I did not check whether the code does exactly the same thing as yours. So please check carefully and see whether the results are plausible.