Optimization of the following code

Question

Consider the function $$ h:[-1,1]\times I_{\sigma}\to I_{\sigma} $$ $$ (\omega, x)\mapsto \sqrt[3]{x + \sigma \omega} $$ where $ \sigma > \frac{2}{3\sqrt{3}}, I_{\sigma} = [x_-(\sigma), x_+(\sigma)] $ and $ x_\pm(\sigma) $ is the unique real root for $ x^3 = x \pm \sigma $.

Now, for a fixed $ x\in I_{\sigma} $, define funtion (Markov Kernel)

$$ P_x: \mathcal{B}(I_\sigma)\to \mathbb{R},\\ I \mapsto \frac{1}{2}m(\left\{\omega\in [-1,1];\ h(\omega,x)\in I\right\}), $$

where $ m $ is the Lebesgue measure and $ I $ is a sub-interval of $ I_{\sigma} $. So, $ P_x(\cdot) $ is a function that associates each sub-interval $ I\subset I_\sigma $ with half of the length of the set $ \left\{\omega\in [-1,1]; h(\omega,x)\in I\right\} $. It is easy to see that

$$ P_x([a,b]) := \frac{1}{2}m\left( \left[\frac{a^3 - x}{\sigma},\frac{b^3 - x}{\sigma}\right]\cap [-1,1]\right). $$

Finally, given a natural number $ n > 0 $, divide $ I_{\sigma} $ into $ n $ subintervals of length $ (x_+(\sigma)-x_-(\sigma))/n $, so we can write $ I = I_1 \cup I_2 \cup\ldots\cup I_n $.

My objetive is to write a program that given as inputs the natural number $ n $ and the real number $ \sigma\in\left(\frac{2}{3\sqrt 3},\infty\right) $, the output is the stochastic matrix $$ Q = \left(\frac{1}{m(I_i)}\int_{I_i} P_x(I_j)\mathrm dx\right)_{(i,j)\in\{1,\ldots,n\}\times\{1,\ldots,n\}}. $$

To solve this problem I have written the following Mathematica code,

F[a_, b_, x_, y_] := 
 1/2 If[a <= b && b <= x && x <= y, 0, 
   If[a <= x && x <= b && b <= y, b - x, 
    If[a <= x && x <= y && y <= b, y - x, 
     If[x <= a && a <= b && b <= y, b - a, 
      If[x <= a && a <= y && y <= b, y - a, 0]]]]]; (*This a function that associates (a,b,x,y) to the length of [a,b]∩[x,y]*) 
A[x_, a_, σ_] := (1/σ) *(a^3 - x);
B[x_, b_, σ_] := (1/σ) *(b^3 - x);
G[x_, a_, b_, σ_] := 
  F[A[x, a, σ], B[x, b, σ], -1, 1]; (*Associates (x,a,b,σ) to P_x([a,b])*)
P[a_, b_, σ_, n_] := 
  P[a, b, σ, n] = 
   ParallelTable[
    NIntegrate[
     n/(a - b)*
      G[x, b + (j - 1) (a - b)/n, b + j (a - b)/n , σ], {x, 
      b + (i - 1) (a - b)/n, b + i (a - b)/n }], {i, 1., n}, {j, 1., 
     n}]; (*when a=x_-(\sigma), b=x_+(\sigma), then P[a,b,\sigma,n] is the stochastic matrix that I want*)

My code works "fine" the problem is that when I try to calculate

n = 10000.;
σ = 1.2;
a = x /. NSolve[x^3 == x + σ, x, Reals][[1]];
b = -a;
Q = P[a, b, σ, n];

I left my program running for 10 hours and it has not terminated.

Can someone help me improve my code in order to decrease the running time? Or propose another way of writing this code?

Answer 1

The function $x \mapsto P_x([a,b])$ is continuous and piecewise linear, so it can be easily and exactly integrated with the trapezoidal rule. Then the resulting quadrature rule can be compiled for better performance. The only problem is to place the quadrature points onto the discontinuities of the first derivative.

The following CompiledFunction getRowofQ does precisely that and computes a single row of the matrix $Q$ at once. Computation of the whole matrix $Q$ can than performed in parallel by exploiting that getRowofQ has attribute Listable.

Block[{x, k, aj, bj, σ},

 P = Function @@ {
    {x, σ, aj, bj},
    Simplify[
     PiecewiseExpand[
      1/2 Ramp[ Min[(bj^3 - x)/σ, 1] - Max[(aj^3 - x)/σ, -1]]
      ],
     σ > 0
     ]
    };

 getRowofQ = 
  With[{Pcode = N@P[Compile`GetElement[x, k], σ, aj, bj]}, 
   Compile[{{ai, _Real}, {bi, _Real}, {a, _Real, 1}, {b, _Real, 1}, {σ, _Real}}, 
    Block[{aj, bj, x, ω, λ, factor},
     x = Table[0., {6}];
     x[[1]] = ai;
     x[[6]] = bi;
     ω = Table[0., {6}];
     factor = 1./(bi - ai);
     Table[
      aj = Compile`GetElement[a, j];
      bj = Compile`GetElement[b, j];
      (*Finding the quadrature points.*)
      x[[2 ;; 5]] = Sort[{
         Max[Min[aj^3 - σ, bi], ai],
         Max[Min[aj^3 + σ, bi], ai],
         Max[Min[bj^3 - σ, bi], ai],
         Max[Min[bj^3 + σ, bi], ai]
         }];

      (*Computing the quadrature weights.*)
      ω[[1]] = 0.;
      Do[
       λ = 0.5 (Compile`GetElement[x, k + 1] - Compile`GetElement[x, k]);
       ω[[k]] += λ;
       ω[[k + 1]] = λ;, {k, 1, 5}];
      (*Computing the integral.*)
      factor Sum[
        If[Compile`GetElement[ω, k] > 0.,
         Compile`GetElement[ω, k] Pcode,
         0.
         ], {k, 1, 6}]
      , {j, 1, Length[a]}]
     ],
    CompilationTarget -> "C",
    RuntimeAttributes -> {Listable},
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
   ];
 ]

So, let's try it on a test problem:

n = 100;
σ = RandomReal[{2/3/Sqrt[3], 5}];
xminus = Root[σ - #1 + #1^3 &, 1];
xplus = Root[-σ - #1 + #1^3 &, 1];
xlist = Subdivide[xminus, xplus, n];
m = Differences[xlist];
a = Most[xlist];
b = Rest[xlist];

Q = getRowofQ[a, b, a, b, σ]; // AbsoluteTiming // First

0.001608

Checking the row sums:

MinMax[Total[Q, {2}]] == {1., 1.}

True

Comparison with NIntegrate:

int[ai_, bi_, aj_, bj_, σ_] := NIntegrate[P[x, σ, aj, bj], {x, ai, bi}]/(bi - ai);
Qref = ParallelTable[
      int[a[[i]], b[[i]], a[[j]], b[[j]], σ], 
      {i, 1, Length[a]}, {j, 1, Length[a]}
      ]; // AbsoluteTiming // First
Max[Abs[Q - Qref]]

20.9859

1.22125*10^-14

This is about $12000$ times faster with negligible error.

For $n = 10000$, the matrix Q is computed by getRowofQ in less than 12 seconds. It is not as fast as the function in the previous version of my post (it can certainly be optimized by unfolding the sort code), but it is hopefully more correct.

As always, please check whether this code really produces the desired result in your use cases.

Matrix-vector multiplication for $Q^\top$

An $O(n^2)$-implementation for the matrix-vector multiplication $u.Q$ is given below. Just us it like this getQTDotU[a, b, a, b, σ, u].

getQTDotU = 
  With[{Pcode = N@P[Compile`GetElement[x, k], σ, aj, bj]}, 
   Compile[{{aj, _Real}, {bj, _Real}, {a, _Real, 1}, {b, _Real, 1}, {σ, _Real}, {u, _Real, 1}},
    Block[{ai, bi, x, ω, λ, factor, sum},
     x = Table[0., {6}];
     ω = Table[0., {6}];
     sum = 0.;
     Do[
      ai = Compile`GetElement[a, i];
      bi = Compile`GetElement[b, i];
      factor = 1./(bi - ai);
      x[[1]] = ai;
      x[[6]] = bi;

      (*Finding the quadrature points.*)
      x[[2 ;; 5]] = Sort[{
         Max[Min[aj^3 - σ, bi], ai],
         Max[Min[aj^3 + σ, bi], ai],
         Max[Min[bj^3 - σ, bi], ai],
         Max[Min[bj^3 + σ, bi], ai]
         }];
      (*Computing the quadrature weights.*)
      ω[[1]] = 0.;
      Do[λ = 0.5 (Compile`GetElement[x, k + 1] - Compile`GetElement[x, k]);
       ω[[k]] += λ;
       ω[[k + 1]] = λ;, {k, 1, 5}];
      (*Computing the integral.*)
      sum += Times[
        Compile`GetElement[u, i],
        factor,
        Sum[
         If[Compile`GetElement[ω, k] > 0., Compile`GetElement[ω, k] Pcode, 0.], 
         {k, 1, 6}]
        ],
      {i, 1, Length[a]}];
     sum
     ],
    CompilationTarget -> "C",
    RuntimeAttributes -> {Listable},
    Parallelization -> True,
    RuntimeOptions -> "Speed"]
   ];

Power Iteration

Setting up the power iteration is straight-forward:

n = 1000;
σ = RandomReal[{2/3/Sqrt[3], 5}];
xminus = Root[σ - #1 + #1^3 &, 1];
xplus = Root[-σ - #1 + #1^3 &, 1];
xlist = Subdivide[xminus, xplus, n];
m = Differences[xlist];
a = Most[xlist];
b = Rest[xlist];
Q = getRowofQ[a, b, a, b, σ];

Of course, I computed Q only for comparison.

u = ConstantArray[1., n];
v = FixedPointList[Normalize[getQTDotU[a, b, a, b, σ, #]] &, u,
    SameTest -> (Norm[#1 - #2]/Sqrt[n] < 1*^-15 &)];
Length[v]
Norm[v[[-1]] - v[[-1]].Q]

15

1.67708*10^-15

As we can see, 15 iterations where enough to produce a dominant eigenvector (with eigenvalue 1).

Of course, this takes roughly 15 times as long as assembling the matrix Q. But as it turns out, the matrix-vector multiplication u.Q can be implemented much faster (see below).

Further optimization potential

It is always worthwhile to plot the matrices that one works with. ArrayPlot[Q] reveals that the structure of Q is pretty boring: The columns of Q are basically step functions with at most 2 steps and essentially two values: 0 and another nonzero value. (There are also 2 or three entries per column around the jump position that take other values.) It should be pretty straight-forward to determine where exaxtly the jump occurs in each column, what the dominant nonzero function value is (and where and what the other 2--3 entries are).

Let's forget for a moment about these 2--3 other values (the should not influence the total outcome of matrix-vector products significantly). Then the i-th column of Q can be fully described by (i) two integers α[[i]] and β[[i]] that mark the jump positions and (ii) a single nonzero value c[[i]]. Then the i-th entry of the matrix-vector product u.Q can be be computed by just summing the entries of u that correspond to the nonzero positions in the i-th column of Q and by multiplying the sum with the corresponding nonzero value in the i-th column of Q, i.e., by

Total[u[[α[[i]];;β[[i]]] c[[i]] 

In fact, it is even more efficient to set

U = Prepend[Accumulate[u],0.]

and to compute the i-th entry of u.Q by

(U[[α[[i]]]]-U[[β[[i]]-1]]) c[[i]]

So the matrix-vector product u.Q can be computed in $O(n)$ time which is considerably less than computing Q once in $O(n^2)$ time and computing u.Q naively also in $O(n^2)$ time.

Unfortunately, I do not have the time to set up the details. But if your work heavily relies on this, you should investigate this further.