Why the function nroots returns a Matrix{Array{Int64,0}}

import Polynomials: Polynomial, derivative
import StatsBase: sample

function nroots!(p, z; eps=1e-05, maxiter=50)
   dp= derivative(p)
   ddp=derivative(dp)
   prod=ddp*p 
   nr=zeros(Int, size(z)) 
   for _ in 1:maxiter 
       z, zprev=  z-p(z)/(dp(z) - prod(z)/(2*dp(z))), z
       nr[findall(x->x^2>eps, abs(z-zprev))] .+= 1
   end 
   return nr 
end 
x= 0:0.05:1
z= x' .+ x*im
p = Polynomial(sample([1.0, -1.0, 1im, -1im], 5))
nr=nroots!.(p, z)

How shoud it be modified to return a Matrix{Int64}?

1 Like

If you run nroots! without broadcasting on a single element, you get the following result:

julia> nroots!(p, z[1])
0-dimensional Array{Int64, 0}:
2

I do not see any documetation for nroots!, so I’m going to assume you know what you are doing with this function.

Since you used broadcasting on a matrix, you get a matrix of these results.

If you just want to unbox all the values, you can do the following:

julia> nr=first.(nr)
21×21 Matrix{Int64}:
...

For a single line,

nr=first.(nroots!.(p,z))
1 Like

Note that this is a continuation of this thread: How to convert a Matrix{Array{Int64, 0}} to a Matrix{Int64}

You wrote your nroots! function in a classic Matlab-ish “vectorized” style where it assumes that z is an array, and allocates the result nr as an array too. If z is a scalar, then size(z) = () and nr is a zero-dimensional array (which is not the same thing as a scalar in Julia).

In this case, it doesn’t seem like there is any advantage in a vectorized style — even if z is an array, the different z’s don’t seem to share any computation. I would suggest just writing a function that works for scalar z and then broadcasting it as needed. If I understand your code correctly, this should be:

function nroots(p, z; eps=1e-05, maxiter=50)
   dp= derivative(p)
   ddp=derivative(dp)
   prod=ddp*p 
   nr=0
   for _ in 1:maxiter 
       z, zprev=  z-p(z)/(dp(z) - prod(z)/(2*dp(z))), z
       nr += abs(z-zprev)^2 > eps
   end 
   return nr 
end 

(Note that neither this nor your original nroots! function should end with a !, by Julia convention, since neither of them mutate their arguments. Reassigning z is not mutation.)

Then you can broadcast this as nr = nroots.(p, z) over an array z and it will return an array of scalars.

4 Likes