stla / EllipticFunctions.jl

Jacobi theta functions and related functions.
https://stla.github.io/EllipticFunctions.jl/dev/
MIT License
16 stars 4 forks source link

Bug: `jthetaN` fails to work when the first argument is an matrix #30

Closed hammerfunctor closed 10 months ago

hammerfunctor commented 11 months ago

promote(z,taufromq(q)) throws the error for Array typed z. We can move all the promotes into _jthetaN_raw

stla commented 11 months ago

Thank you.

Or we could use such a function:

function promote_z_tau(z, tau)
  if length(size(z)) == 0
    return promote(z, tau)
  end
  eltype = Base.promote_eltype(z, tau)
  return (convert(Array{eltype}, z), convert(eltype, tau))
end

I'm not very fluent in Julia, what do you think?

hammerfunctor commented 11 months ago

In my opinion, handling matrix input is not really necessary in such a package. On the one hand, extra overheads for each z input are only qfromtau or taufromq, which is fast enough. Dealing with many zs in one run won't bring practical performance advantage and users could readily broadcast to arrays like jtheta1.(z,q) when they want. On the other hand, when the overhead of GC needs to be considered, in julia we typically create a buffer matrix and assign output of jtheta1(z,q) manually.

simonp0420 commented 11 months ago

I agree with @hammerfunctor's thesis that handling matrix inputs is not necessary. In fact it is unJulian and as we see it makes the code more complicated. The slight performance gain from not repeating taufromq isn't worth it, IMO. However, if you want to retain this, then I would break up promote_z_tau into two methods, one for z::Number and the other for z::AbstractArray, thus eliminating the run-time testing on the size of z.

Another alternative is to add a new method to each of the jthetaN functions, which accepts a single positional argument z::Number and a named argument tau::Number which would not require conversion from q. This could be used with standard Julia dot-broadcasting for array-valued z without repeatedly calling taufromq.

simonp0420 commented 10 months ago

I took a look at adding a method of jtheta1 for Array-valued z arguments:

function jtheta1(z::AbstractArray{Tz}, q::Tq) where {Tz<:Number, Tq<:Number}
    validateq(q)
    T = promote_type(Tz, Tq)
    tau = T(taufromq(q))
    return _jtheta1.(T.(z), tau)
  end

As you can see, it only does the taufromq call a single time. Here is a timing comparison versus calling taufromq for each element of z:

julia> using EllipticFunctions, BenchmarkTools

julia> q = 0.7 + 0.7im
0.7 + 0.7im

julia> zs = rand(ComplexF64, 10_000);

julia> myfun(zs, q) = [jtheta1(z,q) for z in zs] # calls taufromq for each element of z
myfun (generic function with 1 method)

julia> myfun(zs, q) == jtheta1(zs, q) # Test that same result is obtained
true

julia> @btime myfun($zs, $q);
  5.186 ms (2 allocations: 156.30 KiB)

julia> @btime jtheta1($zs, $q);
  4.528 ms (2 allocations: 156.30 KiB)

Eliminating the redundant calls to taufromq for a 10,000 element array saves about 13% of the run time. IMO this isn't enough savings to justify the extra code complexity needed to directly support array-valued arguments.

hammerfunctor commented 10 months ago

This benefit is gone when you want different $\tau$ for different $z$, lmao, If I want critical performance, I would definitely check and regularize input values by hand and use _jtheta1.

stla commented 10 months ago

I've removed all these "vectorizations".