How to speed up the plot of solution from NSolve?

Question

I'm trying to plot the solution of a non-linear equation, but it's very, very slow. From what I could gather, it's not finding the solution that is time consuming, so I cannot quite understand the reason why the plot takes too long. Anyone has an idea about what may be the problem here?

Clear["Global`*"]
pbase = .5;
sbase = .5;
nbase = .35;
abase = .1;
tbase = .7;
lbase = 1.6;
kbase = .2;
ibase = .01;

y[p_, t_, k_, z_, a_] := p/(1 - t)*k*z*(1 - z/(2*a));

m[p_, t_, k_, z_, a_] := Max[k*z, k*z*(1 - z/(2*a))*(1 + p*t/(1 - t))];

u[p_, t_, k_, z_, a_, n_, l_] := 
  p*(l*y[p, t, k, z, a])^(1 - n)/(1 - n) - (1 - p)*
    y[p, t, k, z, a]^(1 - n)/(1 - n);

eq[p_, t_, k_, z_, a_, n_, l_, s_, i_] := -i*m[p, t, k, z, a] + 
   s*(u[p, t, k, z, a, n, l] + (1 - p)*k*z*(1 - z/a) - 
      m[p, t, k, z, a]);

zsol[p_, t_, k_, a_, n_, l_, s_, i_] := 
  z /. NSolve[eq[p, t, k, z, a, n, l, s, i] == 0 && z > 0, z];

ysol[p_, t_, k_, a_, n_, l_, s_, i_] := 
  y[p, t, k, zsol[p, t, k, a, n, l, s, i], a];

yprime[p_, t_, k_, z_, a_] := p/(1 - t)*k*(1 - z/a);

uprime[p_, t_, k_, z_, a_, n_, l_] := (1 + p*(l^(1 - n) - 1))*
   y[p, t, k, z, a]^(-n);

eq2[p_, t_, k_, z_, a_, n_, l_] := 
  uprime[p, t, k, z, a, n, l]*yprime[p, t, k, z, a] - 
   1 - (1 - p)*k*z/a;

zstar[p_, t_, k_, a_, n_, l_] := 
  z /. NSolve[eq2[p, t, k, z, a, n, l] == 0 && z > 0, z];

ystar[p_, t_, k_, a_, n_, l_] := 
  y[p, t, k, Evaluate[zstar[p, t, k, a, n, l]], a];

istar[p_, t_, k_, a_, n_, l_, s_] := 
  s*((u[p, t, k, zstar[p, t, k, a, n, l], a, n, l] + (1 - p)*k*
         zstar[p, t, k, a, n, l]*(1 - zstar[p, t, k, a, n, l]/a))/
      m[p, t, k, zstar[p, t, k, a, n, l], a] - 1);
    
Plot[Evaluate[
  ysol[pbase, tbase, kbase, abase, nbase, lbase, sbase, i]], {i, 
  istar[pbase, tbase, kbase, abase, nbase, lbase, sbase][[1]], 
  istar[pbase, tbase, kbase, abase, nbase, lbase, sbase][[1]] + .3}, 
 PerformanceGoal -> "Speed"]

Answer 1

The following is OPs code with two differences:

• I only keep z and i as function parameters and install all others globally. If OP uses this code, it would be good to check that I did not introduce a mistake while making this change.
• I construct zsol via interpolation. I introduced isol which was obtained by solving eq[z,i]==0 for i which is a linear equation, then generate a table of {i,z} values and interpolate.

Here is the code:

install[data_,nsteps_:10000]:=With[{
  p=data["p"],t=data["t"],k=data["k"],
  a=data["a"],n=data["n"],l=data["l"],s=data["s"]},
    y[z_]:=p/(1-t)*k*z*(1-z/(2*a));
    m[z_]:=Max[k*z,k*z*(1-z/(2*a))*(1+p*t/(1-t))];
    u[z_]:=p*(l*y[z])^(1-n)/(1-n)-(1-p)*y[z]^(1-n)/(1-n);
    eq[z_,i_]:=-i*m[z]+s*(u[z]+(1-p)*k*z*(1-z/a)-m[z]);
    isol[z_]:=s*(u[z]+(1-p)*k*z*(1-z/a)-m[z])/m[z]; (* new *)
    zsol=With[{izvalues=Table[{isol[z],z},{z,Range[0,2*a,2*a/nsteps][[2;;-2]]}]},
            Interpolation[izvalues]];
    ysol[i_]:=y[zsol[i]];
    yprime[z_]:=p/(1-t)*k*(1-z/a);
    uprime[z_]:=(1+p*(l^(1-n)-1))*y[z]^(-n);
    eq2[z_]:=uprime[z]*yprime[z]-1-(1-p)*k*z/a;
    zstar=Module[{z},z/.First[NSolve[eq2[z]==0&&z>0,z]]];
    ystar=y[zstar];
    istar=s*((u[zstar]+(1-p)*k*zstar*(1-zstar/a))/m[zstar]-1);
];

This can now be used as follows:

(* the values are OPs base values *)
install[<|"p"->.5,"t"->.7,"k"->.2,"a"->.1,"n"->.35,"l"->1.6,"s"->.5|>];
Plot[ysol[i],{i,istar,istar+.3},AxesLabel->{"i","y"}]

which produces

---

Edit. Here is a variant of the code that can be used to plot for different values of p:

getsol[data_,nsteps_:10000]:=With[{
  p=data["p"],t=data["t"],k=data["k"],
  a=data["a"],n=data["n"],l=data["l"],s=data["s"]},
    Module[{y,m,u,eq,isol,zsol,ysol,yprime,uprime,eq2,zstar,ystar,istar},
      y[z_]:=p/(1-t)*k*z*(1-z/(2*a));
      m[z_]:=Max[k*z,k*z*(1-z/(2*a))*(1+p*t/(1-t))];
      u[z_]:=p*(l*y[z])^(1-n)/(1-n)-(1-p)*y[z]^(1-n)/(1-n);
      eq[z_,i_]:=-i*m[z]+s*(u[z]+(1-p)*k*z*(1-z/a)-m[z]);
      isol[z_]:=s*(u[z]+(1-p)*k*z*(1-z/a)-m[z])/m[z];
      zsol=With[{izvalues=Table[{isol[z],z},{z,Range[0,2*a,2*a/nsteps][[2;;-2]]}]},
             Interpolation[izvalues]];
      ysol[i_]:=y[zsol[i]];
      yprime[z_]:=p/(1-t)*k*(1-z/a);
      uprime[z_]:=(1+p*(l^(1-n)-1))*y[z]^(-n);
      eq2[z_]:=uprime[z]*yprime[z]-1-(1-p)*k*z/a;
      zstar=Module[{z},z/.First[NSolve[eq2[z]==0&&z>0,z]]];
      ystar=y[zstar];
      istar=s*((u[zstar]+(1-p)*k*zstar*(1-zstar/a))/m[zstar]-1);
      ysol["istar"]=istar;
      ysol
]];

Using

With[{ps=Range[.49,.51,.005]},Show[Table[
  With[{ysol=getsol[<|"p"->ps[[j]],"t"->.7,"k"->.2,"a"->.1,"n"->.35,"l"->1.6,"s"->.5|>]},
     Plot[ysol[i],{i,ysol["istar"],ysol["istar"]+.3},AxesLabel->{"i","y"},
          PlotStyle->ColorData[97,"ColorList"][[j]],
          PlotRange->{{0,0.48},{0.002,0.013}}]],{j,Length[ps]}]]]

one obtains

Answer 2

TL;DR : Use FindRoot instead of NSolve but you can use NSolve for an initial condition for FindRoot before plotting. This is a fix rather than a full answer because I do not know why NSolve creates the issues mentioned below.

Discussion on what seems to be a conflict between NSolve and Plot:

There seems to be a problem with plotting in this case because it seems to take a lot more time to use Plot with options MaxRecursion->0 and PlotPoints->20 then to just use Table.

Maybe this is because Plot is trying to use symbolic processing and is encountering an issue with NSolve. But this seems unlikely because I tried recreating this behavior with simpler cases with FindRoot and Nsolve and I did not encounter any major issues.

(* What I tried

To prevent Plot from manipulating ysol symbolically, I tried imposing that all arguments of ysol should be numerical using NumericQ (use ysol[p_?NumericQ,t_?NumericQ] or add a condition to the definition of the function like ysol[p,t] /; VectorQ[{p,t},NumericQ] (add as many arguments as you need)). But that caused Mathematica to crash each time I tried. Using FindRoot instead of NSolve (discussed below) I was able to get a Plot quickly (without any options for Plot). So there seems to be an issue between NSolve and Plot in this case.

*)

You could maybe consider just using Table and ListPlot if you want a quick solution and the quality of the Plot does not matter much.

In the following I will focus on how using FindRoot makes the computation of ysol faster than using NSolve. If you have further plans that are more CPU time intensive the discussion below might be interesting. If not then maybe a table and ListPlot is enough.

---

NDSolve and FindRoot

According to the documentation for NSolve in the section Details and Options:

NSolve deals primarily with linear and polynomial equations

I guess the algorithms are suited for finding all polynomial roots at once.

Now consider the definition of zsol evaluated with the parameters in the question:

start=istar[pbase, tbase, kbase, abase, nbase, lbase, sbase][[1]];
z /. NSolve[
  eq[pbase, tbase, kbase, z, abase, nbase, lbase, sbase, start] == 0 && 
   z > 0, z]

Looking at eq for these parameters I found :

eq[pbase, tbase, kbase, z, abase, nbase, lbase, sbase, start]

(* 0.5 (0.1 (1 - 10. z) z + 0.134576 ((1 - 5. z) z)^0.65 - Max[0.2 z, 0.433333 (1 - 5. z) z]) - 0.0871035 Max[0.2 z, 0.433333 (1 - 5. z) z] *)

The exponent 0.65 implies that the equation is not polynomial. Nsolve seems to be able to find a solution still but it takes some time:

z /. NSolve[
      eq[pbase, tbase, kbase, z, abase, nbase, lbase, sbase, start] == 0 && 
       z > 0, z] // AbsoluteTiming

output: (* {0.291923, {0.0443099}} ) ( timing= 0.291923*)

Let us compare with FindRoot which is the default function for non linear systems. First, there is the constraint z>0. To remove this constraint we can use the change of variables z=Exp[x] and find a root in terms of x (! Be careful if the name x was already used in your notebook ).

 Exp[x] /. 
  FindRoot[
   Evaluate[
    eq[pbase, tbase, kbase, Exp[x], abase, nbase, lbase, sbase, 
     start]], {x, Log[0.044]}] // AbsoluteTiming

output: {0.000611, 0.0443099} (* timing=0.0006 *)

Why did I choose Log[0.044] for the initial condition ? Because it is the solution that NSolve gave. If you want a quick solution for this problem then you could maybe use Nsolve only for the starting point of the interval considered and then use the same initial condition for FindRoot throughout the interval:

Table[ Exp[x] /. 
   FindRoot[
    Evaluate[
     eq[pbase, tbase, kbase, Exp[x], abase, nbase, lbase, sbase, 
      i]], {x, Log[0.04430987744940078`]}], {i, 
   Subdivide[start, start + 0.3, 20]}] // AbsoluteTiming

output: {0.007084, {0.0443099, 0.0405668,..., 0.0111852, 0.0105612}}

If we used NSolve instead that would take 5.83 s.

However, if the functions you are considering varies quickly over the interval you look at, using the same initial condition for the whole interval might not work in which case see the link(read the parenthesis before the link if you are tempted to click now) in the last paragraph for the method "by continuity".

Making a program that first tries with NSolve and then uses FindRoot can be a hassle depending on what you want to do later. Here is a way to only use FindRoot but it takes some extra work figuring out constraints and change of variables.

If I took a guess instead of using Nsolve, I would have to consider the constraint that (1 - 5. z) z)^0.65 should be real so there is actually another constraint that z should not be too large. You could remove this constraint by using the change of variables $z=(1/(2*5))*(1+\textrm{Tanh}[x])$ but that only works for the choice of parameters p=pbase,t=tbase, etc. If you want to try this method you would need to find the right change of variables for the general case. You can use Reduce and Refine to find out what the constraints in some of your equations are.

Then you could probably just take x=0 as an initial condition, that worked for the parameter choice above and it took 0.008 seconds.

You probably will not need to do more than that but if you are interested in more advanced methods you could look at the method by continuity (by differentiability could be possible but with the Max that you have you would have to first find the conditions for which one of the arguments is larger than the other to remove the Max and only use one of the arguments, you would have to be careful in that case) in this answer that I finished writing shortly before this one.