How do I improve the style, efficiency, and versatility of this code?

Question

I have a math problem that I have solved analytically. I now want to generate a number of graphs based on different parameter values. My problem is threefold:

1. It takes a couple of minutes to produce a single graph.

2. I am wasting time waiting until one graph is finished so that I can start the next one for different parameter values, or one for e.g. different domain and range of the plot.

3. My code seems generally badly designed.

What I'd like is:

1. Make the computation much more efficient, and maybe less messy in terms of code.

2. More importantly: Make it so that I can generate an arbitrary amount of graphs for different settings very easily, and wait until the computer finishes them all. An make it so that after the computation is done, I can easily change the range and domain of the graph, and its design, without the computer having to redo the underlying computations.

3. Generally improve the design of the code.

---

Here is the mathematical problem (the code is given below):

The heaviest part of the computation seems to be in finding $t$.

I have taken as a first example, $$A(i)= \begin{cases} 0 &\text { if } i<a\\ tech\cdot(i-a) &\text { if } i\geq a\end{cases}$$ For some real number $tech$.

I use the following code to compute this problem:

ClearAll[x, A, i, L, K, phi, t, intA]

A[i_, tech_] := If[i > a, tech*(i - a), 0];
intA[t_?NumericQ, tech_] := 
  NIntegrate[A[i, tech]^(phi/(1 - phi)), {i, Max[a, t], 1}];

t[LK_, tech_] := 
 t[LK, tech] = 
  x /. FindRoot[A[x, tech] - ((LK/x)*intA[x, tech])^(1 - phi),
    {x, (a + 1)/2}]


L = 5;
K = 5;
phi = -0.5;
a = 0.5;
tech = 20;
cL[L_, K_, tech_] := L/t[L/K, tech];
cK[L_, K_, tech_] := K/intA[t[L/K, tech], tech];
Lab[i_, L_, K_, tech_] := If[i > t[L/K, tech], 0, cL[L, K, tech]];
Kap[i_, L_, K_, tech_] := 
  If[i > t[L/K, tech], (cK[L, K, tech])*(A[i, tech])^(phi/(1 - phi)), 
   0];

Plot[{Lab[i, L, K, tech], Kap[i, L, K, tech]*A[i, tech], 
  A[i, tech]}, {i, 0, 1}, PlotRange -> All]

tildeF[L_, K_, tech_] := 
  tildeF[L, K, tech] = 
   NIntegrate[(Lab[i, L, K, tech] + A[i, tech]*Kap[i, L, K, tech])^
     phi, {i, 0, 1}];
F[L_, K_, tech_] := tildeF[L, K, tech]^(1/phi);

F[5, 5, 30]
Plot3D[{F[L, K, tech]}, {L, 0, 10}, {K, 0, 10}]

It is especially the Plot3D that is taking most of the computation time.

Answer 1

I don't have enough time, but here are some suggestions for Step 1:

Speeding up $\tilde A$

Many of the integrals are computed with overlapping domains, so there is a lot of redundant integration going on.

A first strategy to speed this up is to replace $\tilde A$ by a piecewise linear approximation that is generated by computing the integral on small intervals once; I use a three-point Gauss quadrature for that, which is faster than letting NIntegrate figure out again and again how to do it.

Fixing some constants

phi = -0.5;
a = 0.5;
tech = 20.;
A[i_] := Evaluate[If[i > a, Evaluate[tech (i - a)], 0.]];
intA[t_?NumericQ] := NIntegrate[A[i]^(phi/(1 - phi)), {i, Max[a, t],1}];

This is where integration takes place:

n = 1000000;
Δi = (1 - a)/n;
ilist = Table[i, {i, a, 1., Δi}];
Block[{x, ω},
   {x, ω} = Most[NIntegrate`GaussRuleData[3,$MachinePrecision]];
   sample = Transpose@D[α x + (1 - x) β, {{α, β}, 1}];
   cA = With[{code = A[i]},
     Compile[{{i, _Real}},
      code,
      CompilationTarget -> "C",
      RuntimeAttributes -> {Listable},
      Parallelization -> True
      ]
     ];
   integrant = With[{code = A[i]^(phi/(1. - phi))},
     Compile[{{i, _Real}},
      code,
      CompilationTarget -> "C",
      RuntimeAttributes -> {Listable},
      Parallelization -> True
      ]
     ];
   intervals = Transpose[{Most[ilist], Rest[ilist]}];
   intervalintegrals = 
    integrant[intervals.sample].(Δi ω);
   ];
f0 = NIntegrate[A[i]^(phi/(1 - phi)), {i, a, 1.}];
flist = Join[{f0}, Reverse@Accumulate[Reverse[intervalintegrals]]];
tildeA = Interpolation[Transpose[{ilist, flist}], InterpolationOrder -> 1];

An error plot:

Speeding up the root finding

Finding the roots of a piecewise linear function is straightforward. Here is some handwritten code that does it for solving $A(t) = \left(\frac{L}{K} \frac{1}{t} \tilde A(t) \right)^{1-\varphi}$. I use a suitable reformulation of this problem along with nearest to detect a point whose neighboring intervals contain the intersection point. Note that we do not use tildeA here.

cfindroot = 
  Compile[{{t, _Real, 1}, {y, _Real, 1}, {y0, _Real}, {index, _Integer}},
   Block[{rootinleft, α, β, yα, yβ, root = 0.1},
    rootinleft = index > 1;
    If[rootinleft,
     yα = y[[index - 1]] - y0;
     yβ = y[[index]] - y0;
     α = t[[index - 1]];
     β = t[[index]];
     If [Abs[yα - yβ] < 1. 10^-8,
      root = 0.5 (α + β),
      root = α (1. - yα/(yα - yβ)) + β yα/(yα - yβ);
      rootinleft = t[[1]] - 1. 10^-8 <= root <= t[[-1]] + 1. 10^-8;
      ]
     ];
    If[! rootinleft,
     yα = y[[index]] - y0;
     yβ = y[[index + 1]] - y0;
     α = t[[index]];
     β = t[[index + 1]];
     If [Abs[yα - yβ] < 1. 10^-8,
      root = 0.5 (α + β),
      root = α (1. - yα/(yα - yβ)) + β yα/(yα - yβ);
      If[! (t[[1]] - 1. 10^-8 <= root <= t[[-1]] + 1. 10^-8), root = -1.];
      ]
     ];
    root
    ],
   CompilationTarget -> "C",
   RuntimeAttributes -> {Listable},
   Parallelization -> True
  ];
buffer = ilist/flist cA[ilist]^(1/(1 - phi));
nearestfun = Nearest[buffer -> Automatic];
findroot = {k, l} \[Function] With[{
  root = Flatten[cfindroot[ilist, buffer, l/k, nearestfun[l/k]]]
 },
   If[Length[root] == 1, root[[1]], root]
];

In order to get fair and reliable timing results with RepeatedTiming, I have to remove memoization of t (it wasn't too helpfull anyways). Moreover, I make the root finding more robust by enforcing the secant-based Brent method:

ClearAll[t]
t[LK_] := x /. FindRoot[A[x] - ((LK/x) intA[x])^(1 - phi), {x, a, 1.},  Method -> "Brent"]

So, here is the test:

SeedRandom[666];
k = RandomReal[{0, 1}];
l = RandomReal[{0, 1}];
G[x_, k_, l_] := cA[x] - ((l/k/x) tildeA[x])^(1 - phi);
Gtrue[x_, k_, l_] := A[x] - ((l/k/x) intA[x])^(1 - phi);
Plot[{G[x, k, l], Gtrue[x, k, l]}, {x, a, 1}]
root = findroot[k, l]; // RepeatedTiming // First
trueroot = t[l/k]; // RepeatedTiming // First
G[root, k, l]
Gtrue[root, k, l]
trueroot - root

0.000026

0.055

1.66567*10^-12

1.2359*10^-6

-5.25296*10^-8

This is roughly a 2000-fold speedup with an accuracy that should suffice for most applications. In order to increase the accuracy, you may consider cranking up $n$ even further.

In the way the code is written, it also threads over lists with automatic parallelization:

m = 100000;
l = RandomReal[{0.1, 10.}, m];
k = ConstantArray[1., m];
(root = findroot[k, l]; // RepeatedTiming // First)/m

1.1*10^-7

This is roughly 200 times faster per root than applying findroot to pairs of single real numbers k and l, leading of a total speedup for Step 1 of something over 400000.

Answer 2

I will just address one part of your computation. For your particular choice of A and $\phi$, you can solve for t in terms of L/K (I use lk = L/K). First, here is your equation for t:

eqn = A[t] == (lk/t Integrate[A[i]^(\[Phi]/(1-\[Phi])),{i,t,1}])^(1-\[Phi]);
eqn //TeXForm

$A(t)=\left(\frac{\operatorname{lk} \int_t^1 A(i)^{\frac{\phi }{1-\phi }} \, di}{t}\right)^{1-\phi }$

Now, let's assume that t > 1/2 and simplify the equation:

eqn2 = Assuming[t>1/2, Refine[eqn/.A->(Ramp[tech (#-a)]&)/.{tech->30,a->1/2,\[Phi]->-1/2}]];
eqn2 //TeXForm

$30 \left(t-\frac{1}{2}\right)=\frac{\left(\operatorname{lk} \left( \begin{array}{cc} \{ & \begin{array}{cc} -\frac{3^{2/3} \left((2 t-1)^{2/3}-1\right)}{4 \sqrt[3]{5}} & t>1 \\ \left((2 t-1)^{2/3}-1\right) \operatorname{Root}\left[320 \#^3+9\&,1\right] & t<1 \\ \end{array} \\ \end{array} \right)\right)^{3/2}}{t^{3/2}}$

Finally, let's choose the 1/2 < t < 1 branch, and solve for t:

Refine[
    Solve[
        (eqn2 /. pw_Piecewise :> Simplify[pw, 1/2<t<1]) && lk>0,
        t,
        Reals
    ],
    lk>0
]

{{t -> Root[ 6075 lk^6 Root[9 + 320 #1^3 &, 1]^6 + 64 lk^9 Root[9 + 320 #1^3 &, 1]^9 + (-48600 lk^6 Root[9 + 320 #1^3 &, 1]^6 - 192 lk^9 Root[9 + 320 #1^3 &, 1]^9) #1 + (156600 lk^6 Root[ 9 + 320 #1^3 &, 1]^6 + 192 lk^9 Root[9 + 320 #1^3 &, 1]^9) #1^2 + (-259200 lk^6 Root[ 9 + 320 #1^3 &, 1]^6 - 64 lk^9 Root[9 + 320 #1^3 &, 1]^9) #1^3 + (607500 lk^3 Root[ 9 + 320 #1^3 &, 1]^3 + 237600 lk^6 Root[9 + 320 #1^3 &, 1]^6) #1^4 + (-5467500 lk^3 Root[ 9 + 320 #1^3 &, 1]^3 - 129600 lk^6 Root[9 + 320 #1^3 &, 1]^6) #1^5 + (11390625 + 19440000 lk^3 Root[9 + 320 #1^3 &, 1]^3 + 43200 lk^6 Root[9 + 320 #1^3 &, 1]^6) #1^6 + (-136687500 - 34020000 lk^3 Root[9 + 320 #1^3 &, 1]^3) #1^7 + (683437500 + 29160000 lk^3 Root[9 + 320 #1^3 &, 1]^3) #1^8 + (-1822500000 - 9720000 lk^3 Root[9 + 320 #1^3 &, 1]^3) #1^9 + 2733750000 #1^10 - 2187000000 #1^11 + 729000000 #1^12 &, 2]}}

So, for any values of L and K, one can use the above solution to compute t.