Non-linear root-finding: Works in Matlab, but not in Julia

General Usage
,
1

Hey everyone,

I have a small numerical problem here. The main idea is that I try to approximate a univariate function, that needs to satisfy some constraint, with a polynomial. I had this already working with an older version of NLsolve and I run this code under Julia 0.6.4 using the package CompEcon. However, after letting the code rest for a while and updating packages, I stopped finding a solution to my problem. Here is the code:

using CompEcon
# Parameters
global γ = 2;    # Risk aversion
global α = 0.4;    # Capital share
global β = 0.9;   # Discount factor

# First, create a grid using the QuantEcon toolbox
# Set the endpoints of approximation interval:
a =  2.0                           # left endpoint
b =  3.0                           # right endpoint
# Note: They have to be both integer or real. Important in Julia.

# Choose an approximation scheme. In this case, let us use an order 10
n       = 3                         # order of approximation
fspace  = fundefn(:cheb, n, a, b)   # define fspace
nodes   = funnode(fspace)[1]        # safe grid/nodes
Phi     = funbase(fspace,nodes)     # derive basis matrix Φ


using NLsolve

function f!(res, c_guess, fspace, nodes) # we want to find roots for f
  global γ,α,β

  K        = nodes
  con      = funeval(c_guess,fspace,K)[1]
  K_next   = K.^α - con
  con_next = funeval(c_guess,fspace,K_next)[1]
  res      = con.^(-γ) - β*con_next.^(-γ).*α.*K_next.^(α-1)
end

g!(res, c_guess) = f!(res, c_guess, fspace, nodes)
c_guess           = Phi\nodes.^0.3    # some initial guess for parameters c
result           = nlsolve(g!,c_guess)

I have put equivalent code in MATLAB, where it works and delivers a sensible solution. What may be the reason for this not to work in Julia and what can I do about it?

The Matlab code is the following (before running this code, you need to install the CompEcon Toolbox).


function [res] = ExampleProblem(c_guess,fspace,nodes)

gamma = 2;
alpha = 0.4;
beta  = 0.9;

% Unscreened Investment

K        = nodes;
con      = funeval(c_guess,fspace,K);
K_next   = K.^alpha - con;
con_next = funeval(c_guess,fspace,K_next);

% Save the residual
res      = con.^(-gamma) - beta*con_next.^(-gamma).*alpha.*K_next.^(alpha-1);
end


n = 3;
a = 2;
b = 3;
fspace  = fundefn('cheb', n, a, b);   
nodes   = funnode(fspace);       
Phi     = funbas(fspace,nodes);  

c_guess = Phi\nodes.^0.3;

options = optimset('display','iter');

y = fsolve(@(x)ExampleProblem(x,fspace,nodes),c_guess,options);
2

In your f!, the last line isn’t actually updating res. Try this:

res[:]   = con.^(-γ) - β*con_next.^(-γ).*α.*K_next.^(α-1)
1 Like
3

Thanks! That helped. Can you give me an explanation why using res[:] is different? That would be very helpful :slight_smile:

4

res = creates a new variable. It has the same name as the function argument, so it looks confusing.

res[:] is a slicing operation that puts the result from the right-hand side into res.

2 Likes
5

Makes sense, thanks a lot!

6

If the shapes match, in-place assignment with res .= might be better than res[:] =

3 Likes
7

I think I tried that, and I don’t think the shapes matched.

8

By the way, there are a lot of pretty simple things you can do to make this code much more performant (see the performance tips for more).

For reference the original version of your code (with the res[:] fix), runs in:

julia> @btime nlsolve($g!, $c_guess)
  3.965 ms (14859 allocations: 751.09 KiB)
  1. Don’t use global variables inside your inner loop. Instead, you can create a function which takes in all of your parameters (gamma, alpha, nodes, etc.) and returns a new g! that will use those values internally (or, in other words, a closure):
function create_cost_function(fspace, nodes, γ, α, β)
   function f!(res, c_guess, fspace, nodes) # we want to find roots for f
     K        = nodes
     con      = funeval(c_guess,fspace,K)[1]
     K_next   = K.^α - con
     con_next = funeval(c_guess,fspace,K_next)[1]
     res[:]   = con.^(-γ) - β*con_next.^(-γ).*α.*K_next.^(α-1)
   end

   g!(res, c_guess) = f!(res, c_guess, fspace, nodes)
   return g!
end

g! = create_cost_function(fspace, nodes, γ, α, β)

which gives:

julia> @btime nlsolve($g!, $c_guess)
  2.836 ms (12824 allocations: 680.64 KiB)

or about a 30% improvement. Not much, but we’re just getting started.

  1. Take advantage of broadcast fusion to update res completely in-place. For an explanation of what’s going on, check out More Dots: Syntactic Loop Fusion in Julia . In this case, I’m changing res[:] = into res .= and adding dots everywhere on that line so that the entire assignment fuses into a single efficient loop with no extra allocated memory. To do this, I also have to make con and con_next into vectors, since CompEcon makes the very Matlab-like choice to return them as Nx1 matrices instead.
function create_cost_function(fspace, nodes, γ, α, β)
    function f!(res, c_guess, fspace, nodes) # we want to find roots for f
      K        = nodes
      con      = vec(funeval(c_guess,fspace,K)[1])
      K_next   = K.^α - con
      con_next = vec(funeval(c_guess,fspace,K_next)[1])
      res     .= con.^(-γ) .- β .* con_next.^(-γ) .* α .* K_next.^(α-1)
    end

    g!(res, c_guess) = f!(res, c_guess, fspace, nodes)
    return g!
end

g! = create_cost_function(fspace, nodes, γ, α, β)

Performance is now up to:

julia> @btime nlsolve($g!, $c_guess)
  415.725 μs (6004 allocations: 412.52 KiB)

or a solid 10X faster than the original code.

I suspect there’s a lot more that could be done, but it might require updating CompEcon’s internals.

5 Likes
9

Wow! That is super useful. I will have a good look and update my code accordingly. I think the time has come when I switch to Julia for solving my root-finding problems :slight_smile:

10

Great! If you want further performance improvements, it would be useful to try to pre-allocate your con and con_next variables, as in:

function create_cost(....)
  con = zeros(...)
  function f!(...)
    # Update `con` in-place here instead of creating a new vector
  end
end

Unfortunately, I don’t know how to tell CompEcon.funeval to update con in-place rather than creating a brand-new vector (that’s what I was referring to about possibly requiring some internal changes to CompEcon).

2 Likes
11

This was a really helpful illustration. Thanks!