stla / EllipticFunctions.jl

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

WIP: altjtheta1 #4

Closed simonp0420 closed 2 years ago

simonp0420 commented 2 years ago

As the name indicates, this is a work in progress. I'm soliciting your comments and direction as to whether I should push ahead with these modifications.

TLDR

I've implemented an alternative evaluation scheme for jtheta1(z,tau) which is often much faster (as much as 10x faster for real z and pure imaginary tau) than the presently implemented algorithm. Because my proposed changes are so disruptive, I'd like your take on this and whether I should continue to work on the remaining three functions. Also, I would like to suggest using the nome q = exp(im * pi * tau) rather than tau as the second argument to the functions jtheta1, ..., jtheta4 . I make arguments for this below.

The Details

When I saw your Discourse posts and package it stirred up some memories from more than 30 years ago. I remember seeing some alternative series for the Jacobi theta functions, but couldn't find the source (it was probably a book from the library at my then workplace, which is long since gone). But I played around and reconstructed the series for the $\theta_1$ function using the Poisson summation formula (notes attached below). The alternative series has the nice property that it converges fastest where the imaginary part of tau is small, so it nicely complements the original NIST series, which converges fastest where this quantity is large. So I coded up a function named altjtheta1 which uses these two series and made some timing and accuracy comparisons to jtheta1. The reference values for accuracy are from Wolfram Alpha:

julia> using EllipticFunctions # I have the altthetas branch checked out

julia> using BenchmarkTools 

The first example has the most advantage for this PR: z is pure real and tau is pure imaginary and small:

julia> @btime EllipticFunctions.jtheta1(20, 0.01*im)
  307.236 ns (0 allocations: 0 bytes)
0.03608696080199044 - 2.6516268622268154e-17im

julia> @btime EllipticFunctions.altjtheta1(20, 0.01*im)
  31.570 ns (0 allocations: 0 bytes)
0.03608696080206553 + 0.0im

julia> # Wolfram Alpha: 0.0360869608020636316

altjtheta1 provides several more accurate digits, returns a zero imaginary part, and is almost ten times faster than jtheta1. There are several reasons for the speed difference:

  1. The fact that both z and q = xcispi(tau) are pure real is exploited so that all the calculations are done on Float64 reals. This also accounts for the pure real result.
  2. The Poisson transformed series used in this case converges faster as imag(tau) decreases, and this is a quite small value.

I attribute the greater accuracy to the fact that small errors are introduced in jtheta1 from the multiple modular transformations required.

The choice as to whether the original or transformed series is used is according to whether imag(tau) is greater than (original) or less than or equal to (transformed) 1.3. Here is an example where we are slightly above this decision point:

julia> @btime EllipticFunctions.jtheta1(-2, 1.4*im)
  230.073 ns (0 allocations: 0 bytes)
-0.6056537639933428 + 1.9479310714792054e-16im

julia> @btime EllipticFunctions.altjtheta1(-2, 1.4*im)
  200.333 ns (0 allocations: 0 bytes)
-0.6056537639933426 + 0.0im

julia> # Wolfram Alpha: -0.605653763993342615

Here the speed is comparable and altjtheta1 is only very slightly more accurate.

As the imaginary part of tau increases, alttheta1 starts to regain an advantage:

julia> @btime EllipticFunctions.jtheta1(-2, 2.2*im)
  195.487 ns (0 allocations: 0 bytes)
-0.323094150693883 + 1.0391500434177393e-16im

julia> @btime EllipticFunctions.altjtheta1(-2, 2.2*im)
  110.673 ns (0 allocations: 0 bytes)
-0.32309415069388303 + 0.0im

julia> # Wolfram Alpha: -0.323094150693883056

Here altjtheta1 is significantly faster and perhaps more accurate by one digit.

Here is an example where both arguments are have nonzero real and imaginary parts:

julia> @btime EllipticFunctions.jtheta1(1.3 + 0.7im, 0.3 + 2.2*im)
  187.297 ns (0 allocations: 0 bytes)
0.40103039942452695 + 0.17043051360017036im

julia> @btime EllipticFunctions.altjtheta1(1.3 + 0.7im, 0.3 + 2.2*im)
  123.098 ns (0 allocations: 0 bytes)
0.40103039942452723 + 0.17043051360017042im

julia> # Wolfram Alpha: 0.401030399424527305 + 0.170430513600170463 i

Although this is by no means an exhaustive test, it looks like altjtheta1 is at least as accurate and is usually faster, sometimes much faster, than jtheta.

Additional Discussion

I realize that you have a very elegant code structure in place and this PR, if fully implemented, would mean a drastic change of direction. I would like to make one more suggestion that would be a breaking change (but since you're not yet at version 1.0 I think this is acceptable): I would like you to consider using the nome q = cispi(tau) as the second argument to the theta functions. This would allow Julia's type system to eliminate the need for the shenanigans I implemented to exploit real-valued z and/or q values. If one had jtheta(z, q), then when a user invoked the function with Float64 or other pure real-typed values, the compiler would specialize for these types, and the calculations would automatically be done in that type, returning a Float64 value rather than a ComplexF64 value. This is how virtually all other Julia math functions behave. It would simplify the user experience and result in faster code in the most commonly occurring case.

Thanks in advance for looking at this PR.

ThetaFunctions.pdf

oscardssmith commented 2 years ago

note that this could be faster by replacing the q^... with multiplication since floating point powers are much slower than multiplication.

simonp0420 commented 2 years ago

@oscardssmith Do you mean this: q^(n*(n+1))? Something multiplicative that recurses on a previous value?

oscardssmith commented 2 years ago

yeah. each term of that increases by q^2n, so you can compute q^2n with a multiplication by q^2 and then use that to compute the new value of q^(n+1)*n

simonp0420 commented 2 years ago

Done. It had a very dramatic effect on the runtime for that branch of code for the real result case:

julia> @btime EllipticFunctions.altjtheta1(-2, 1.4*im)
  38.101 ns (0 allocations: 0 bytes)
-0.6056537639933426 + 0.0im

julia> @btime EllipticFunctions.altjtheta1(-2, 2.2*im)
  28.003 ns (0 allocations: 0 bytes)
-0.32309415069388303 + 0.0im

but only a slight improvement for the general case where both z and q are complex:

julia> @btime EllipticFunctions.altjtheta1(1.3 + 0.7im, 0.3 + 2.2*im)
  118.831 ns (0 allocations: 0 bytes)
0.40103039942452723 + 0.17043051360017042im

I guess the time here is dominated by the complex sin call?

stla commented 2 years ago

Awesome! Thanks. I will look more carefully later.

stla commented 2 years ago

I translated your code in C++ and I call it from R. I get an issue:

> jacobi::jtheta1(1+1i, 1i)
[1] 1.181613+0.595897i
> jacobi:::altjtheta1(1+1i, 1i)
Error in jacobi:::altjtheta1(1 + (0+1i), 0+1i) : Reached 3000 terms.

But I possible did a mistake in the code.

stla commented 2 years ago

Ok I did a mistake, this works now. I found a problematic case:

> jacobi:::altjtheta1(1-1i, 0.0000000000001i)
Error in jacobi:::altjtheta1(1 - (0+1i), 0+1e-13i) : Reached 3000 terms.
stla commented 2 years ago

One can define altjtheta2 and altjtheta3 as follows:

function altjtheta2(z, tau)
  return altjtheta1(z + pi/2, tau)
end

function expM(z, tau) 
  return exp(im * (z + tau * pi/4))
end

function altjtheta3(z, tau)
  return altjtheta2(z - pi/2 * tau, tau) * expM(-z, tau)
end

I don't get the correct result for this case:

tau = -0.5 + 0.8660253im
altjtheta3(pi*(tau+1)/2, 3*tau)

Could you try, please?

stla commented 2 years ago

I get the correct result with the help of the tau -> -1/tau transformation, namely:

alpha = sqrt(-im * tau) * exp(im / tau * z * z / pi)
altjtheta1(z, tau)
# same as:
im * altjtheta1(z / tau, -1 / tau) / alpha 
simonp0420 commented 2 years ago

It'll be a couple of hours before I can look at this in detail, but clearly the transformed sum has a limited region of applicability in the complex z plane (which should have been obvious to me before this 😳).

stla commented 2 years ago

By the way the Poisson summation is used in my code: to calculate theta1dash (derivative of theta1). This function can be improved by considering the real case and by using @oscardssmith 's trick.

stla commented 2 years ago

I manually did the changes. I incorporated your modifications, but with my modifications for the case t or q not real (tau -> -1/tau). I don't understand why this is needed by the way. Without my modifications, there are failures in the unit tests.

I close because the pull request is not needed anymore, but we can continue to discuss here.

simonp0420 commented 2 years ago

I'm gratified that you responded positively, but I'm a little concerned that maybe more testing would be advisable before incorporating the new algorithms. What I've been doing is looking at the problematic case you found: altjtheta1(1-1i, 0.0000000000001im) (still using my altthetas branch). The problem is with the n=0 term in the transformed sum. For pure imaginary tau, the condition that the real part of the argument to the exponential be positive is $$ |\Im{z}| > |2\Re{z} + (n-1/2)\pi| $$ When nminus is zero, we have

julia> abs(imag(z)), abs(2*real(z) + (nminus-1/2)*π)
(1, 0.42920367320510344)

so that this condition is satisfied. When the division by t/pi is included, the real part of the argument of the exponential function is about 2e12, which of course results in infinity after exponentiation.

My original thoughts were to proceed with transforming the sums for the other theta functions if things were looking good for theta1 and you concurred. However, since you're able to compute them in terms of theta1, it doesn't seem like a good use of time. I guess I'll let this go for the time being, but would still ask you to consider some way to call the theta functions with the nome, which is real in many common applications, so that a Float64 result can be returned for pure Float64 inputs.

P.S. Your logo graphic is very cool!

stla commented 2 years ago

All the unit tests pass with the new implementation, that's not bad.

You're right for the nome, in addition this would be consistent with Wolfram. I'll do the changes whenever I'll be motivated.

My logo is obtained by targeting a light with the camera of a mobile phone, and taking the photo while shaking the phone.

simonp0420 commented 2 years ago

My logo is obtained by targeting a light with the camera of a mobile phone, and taking the photo while shaking the phone

I don't really understand this. I assumed it was a pseudo color plot of one of the elliptic functions.

stla commented 2 years ago

I talked about my avatar. Maybe you talked about the picture in the README of the repo? Yes this is an elliptic function. I'm currently doing an animation of a Jacobi elliptic function mapped from a square to a circle and transformed with a Möbius transformation :-)

simonp0420 commented 2 years ago

Thanks. I finally figured it out from looking at the example code in the docs and then the name of the file gave away as the cn function. :-). I can be a little slow on the uptake. Must be where my nickname "lightning" came from ;-).

stla commented 2 years ago

I took cn on a "fundamental" square. By "fundamental" I mean that cn is periodic and this square covers a period. In this way I can map the picture to a ball. See the repo of my R package for other such pictures. I also uploaded such pictures to my youtube channel.