Jacobi theta functions with JavaScript - depends on math.js

jacobi.js

/*
Jacobi theta functions.
Author: Stéphane Laurent.
The functions jtheta1, jtheta2, jtheta3 and jtheta4 are the Jacobi theta 
functions. The function jtheta1 is 1-antiperiodic, and the functions 
jtheta2, jtheta3 and jtheta4 are 1-periodic.
The function jthetaAB is the Jacobi theta function with characteristics.
The function jtheta1Prime0 is the derivative at 0 of jtheta1.
*/

var jacobi = (function() {

  let Epsilon = 1 / Math.pow(2, 51);

  let close = function(z1, z2) {
    const mod_z2 = math.abs(z2);
    const h = (mod_z2 < Epsilon) ? 1 : Math.max(math.abs(z1), mod_z2);
    return math.abs(math.subtract(z1, z2)) < Epsilon * h;
  };

  let calctheta3 = function(z, tau) {
    let out = math.complex(1, 0);
    let n = 0;
    while(true) {
      n++;
      let nipi = math.multiply(math.multiply(n, math.complex(0, 1)), Math.PI);
      let ntau = math.multiply(n, tau);
      let dblz = math.multiply(2, z);
      let lweight1 = math.multiply(nipi, math.add(ntau, dblz)); 
      let lweight2 = math.multiply(nipi, math.subtract(ntau, dblz));
      let qweight = math.add(math.exp(lweight1), math.exp(lweight2));
      out = math.add(out, qweight);
      let modulus = math.abs(out);
      if(!Number.isFinite(modulus)) {
        throw "NaN or infinity has occured in the summation.";
      } else if(n >= 3 && close(math.add(out, qweight), out)) {
        break;
      }
    }
    return math.log(out);
  };

  let argtheta3 = function(z, tau, passes, maxiter) {
    passes++;
    if(passes > maxiter) {
      throw "Reached maximal iteration (argtheta3).";
    }
    let z_img = math.im(z);
    let tau_img = math.im(tau);
    let h = tau_img / 2;
    let zuse = math.complex(math.re(z) % 1, z_img);
    let out;
    if(z_img < -h) {
      out = argtheta3(math.unaryMinus(zuse), tau, passes, maxiter);
    } else if(z_img >= h) {
      let quotient = Math.floor(z_img / tau_img + 0.5);
      let zmin = math.subtract(zuse, math.multiply(quotient, tau));
      out = math.chain(zmin)
                .multiply(math.complex(0, 1))
                .multiply(-2 * Math.PI * quotient)
                .add(argtheta3(zmin, tau, passes, maxiter))
                .subtract(
                  math.multiply(
                    Math.PI * quotient * quotient,
                    math.multiply(
                      math.complex(0, 1), tau
                    )
                  )
                )
                .done();
    } else {
      out = calctheta3(zuse, tau);
    }
    return out;
  };

  let dologtheta4 = function(z, tau, passes, maxiter) {
    return dologtheta3(math.add(z, 0.5), tau, passes + 1, maxiter);
  };

  let dologtheta3 = function(z, tau, passes, maxiter) {
    passes++;
    let tau2;
    let rl = math.re(tau);
    if(rl >= 0) {
      tau2 = math.complex((rl + 1) % 2 - 1, math.im(tau))
    } else {
      tau2 = math.complex((rl - 1) % 2 + 1, math.im(tau))
    }
    rl = math.re(tau2);
    let out;
    if(math.abs(tau2) < 0.98 && math.im(tau2) < 0.98) { 
      let tauprime = math.divide(-1, tau2);
      out = math.chain(tauprime)
                .multiply(math.complex(0, 1))
                .multiply(math.square(z))
                .multiply(Math.PI)
                .add(dologtheta3(
                  math.multiply(z, tauprime), tauprime, passes, maxiter
                ))
                .subtract(
                  math.log(
                    math.divide(
                      math.sqrt(tau2), math.sqrt(math.complex(0, 1))
                    )
                  )
                )
                .done();
    } else if(rl > 0.6) {
      out = dologtheta4(z, math.subtract(tau2, 1), passes, maxiter);
    } else if(rl <= -0.6) {
      out = dologtheta4(z, math.add(tau2, 1), passes, maxiter);
    } else {
      out = argtheta3(z, tau2, 0, maxiter);
    }
    return out;
  };

  let M = function(z, tau) {
    return math.multiply(
      math.complex(0, 1),
      math.multiply(
        Math.PI,
        math.add(
          z, math.divide(tau, 4)
        )
      )
    );
  };

  let ljtheta2 = function(z, tau) {
    return math.add(
      M(z, tau),
      dologtheta3(
        math.add(z, math.divide(tau, 2)), tau, 0, 1000
      )
    );
  };

  let ljtheta1 = function(z, tau) {
    return ljtheta2(math.subtract(z, 0.5), tau);
  };

  let ljtheta3 = function(z, tau) {
    return dologtheta3(z, tau, 0, 1000);
  };

  let ljtheta4 = function(z, tau) {
    return dologtheta4(z, tau, 0, 1000);
  };

  let jtheta1 = function(z, tau) {
    if(math.im(tau) <= 0) {
      throw "The imaginary part of `tau` must be nonnegative.";
    }
    return math.exp(ljtheta1(z, tau))
  };

  let jtheta2 = function(z, tau) {
    if(math.im(tau) <= 0) {
      throw "The imaginary part of `tau` must be nonnegative.";
    }
    return math.exp(ljtheta2(z, tau))
  };

  let jtheta3 = function(z, tau) {
    if(math.im(tau) <= 0) {
      throw "The imaginary part of `tau` must be nonnegative.";
    }
    return math.exp(ljtheta3(z, tau))
  };

  let jtheta4 = function(z, tau) {
    if(math.im(tau) <= 0) {
      throw "The imaginary part of `tau` must be nonnegative.";
    }
    return math.exp(ljtheta4(z, tau))
  };

  let jthetaAB = function(a, b, z, tau) {
    if(math.im(tau) <= 0) {
      throw "The imaginary part of `tau` must be nonnegative.";
    }
    let alpha = math.multiply(a, tau);
    let beta = math.add(z, b);
    let C = math.exp(
      math.multiply(
        math.complex(0, Math.PI),
        math.multiply(
          a, 
          math.add(alpha, math.multiply(2, beta))
        )
      )
    );
    return math.multiply(C, jtheta3(math.add(alpha, beta), tau));
  };

  let jtheta1Prime0 = function(tau) {
    let j = jthetaAB(1/6, 0.5, 0, math.multiply(3, tau));
    return math.multiply(
      math.complex(0, -2), math.multiply(j, math.square(j))
    );
  };

  return {
    jtheta1: jtheta1,
    jtheta2: jtheta2,
    jtheta3: jtheta3,
    jtheta4: jtheta4,
    jthetaAB: jthetaAB,
    jtheta1Prime0: jtheta1Prime0
  };

})();