diff --git a/src/flint/test/test_acb_theta.py b/src/flint/test/test_acb_theta.py index 4ea6237a..c1470d8d 100644 --- a/src/flint/test/test_acb_theta.py +++ b/src/flint/test/test_acb_theta.py @@ -45,3 +45,22 @@ def test_acb_theta_shape_assertions() -> None: assert raises(lambda: acb_theta(object(), tau), TypeError) # type: ignore[arg-type] assert raises(lambda: acb_theta(z, object()), TypeError) # type: ignore[arg-type] + +def test_acb_theta_jets_basic() -> None: + if not _has_acb_theta(): + return + + from flint.types.acb_theta import acb_theta_jets + + z = acb(1 + 1j) + tau = acb(1.25 + 3j) + zmat = acb_mat([[z]]) + taumat = acb_mat([[tau]]) + ord = 2 + + direct = acb_theta_jets(zmat, taumat, ord) + via_method = taumat.theta_jets(zmat, ord) + assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12) + assert direct.nrows() == 4 + assert direct.ncols() == 3 + diff --git a/src/flint/types/acb_mat.pyi b/src/flint/types/acb_mat.pyi index 87a8736b..755071a4 100644 --- a/src/flint/types/acb_mat.pyi +++ b/src/flint/types/acb_mat.pyi @@ -97,6 +97,9 @@ class acb_mat(flint_mat[acb]): def theta(self, z: acb_mat, square: bool = False) -> acb_mat: ... + def theta_jets(self, z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ... + + def str(self, n: int = 0, radius: bool = True, more: bool = False, condense: int = 0) -> _str: ... def repr(self) -> _str: ... def __str__(self) -> _str: ... diff --git a/src/flint/types/acb_mat.pyx b/src/flint/types/acb_mat.pyx index ae65fe71..0b67f2a4 100644 --- a/src/flint/types/acb_mat.pyx +++ b/src/flint/types/acb_mat.pyx @@ -839,3 +839,22 @@ cdef class acb_mat(flint_mat): except ImportError: raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0") return acb_theta(z, tau, square=square) + + def theta_jets(tau, z, ord): + r""" + Computes Taylor approximation for the vector-valued Riemann theta function + `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares, + where `\tau` is given by ``self``. + + This is a wrapper for :func:`.acb_theta_jet.acb_theta_jet`; see the + documentation for that method for details and examples. + This follows the same conventions of the C-function + `acb_theta_jet `_ + for the ordering of the theta characteristics. + + """ + try: + from .acb_theta import acb_theta_jets + except ImportError: + raise NotImplementedError("acb_mat.theta needs Flint >= 3.1.0") + return acb_theta_jets(z, tau, ord) diff --git a/src/flint/types/acb_theta.pyi b/src/flint/types/acb_theta.pyi index 37f1e1e3..5d06b495 100644 --- a/src/flint/types/acb_theta.pyi +++ b/src/flint/types/acb_theta.pyi @@ -4,3 +4,5 @@ from flint.types.acb_mat import acb_mat def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ... + +def acb_theta_jets(z: acb_mat, ord: int, square: bool = False) -> list[acb_mat]: ... diff --git a/src/flint/types/acb_theta.pyx b/src/flint/types/acb_theta.pyx index 80ed5df1..cf3696db 100644 --- a/src/flint/types/acb_theta.pyx +++ b/src/flint/types/acb_theta.pyx @@ -95,3 +95,83 @@ def acb_theta(acb_mat z, acb_mat tau, ulong square=False): res.append(r) _acb_vec_clear(theta, nb) return acb_mat([res]) + + +def acb_theta_jets(acb_mat z, acb_mat tau, slong ord): + r""" + Computes the coefficients of the Taylor expansion of the vector valued Riemann + theta function `(\theta_{a,b}(z, \tau) : a, b \in \{0,1\}^{g})` or its squares. + + This is a wrapper for the C-function + `acb_theta_jet `_ + and it follows the same conventions for the ordering of the theta characteristics. + + This should be used via the method :meth:`.acb_mat.theta_jets`, explicitly ``tau.theta_jets(z, ord)``. + + >>> from flint import acb, acb_mat, showgood, ctx + >>> z = acb(1+1j); tau = acb(1.25+3j) + >>> acb_mat([[tau]]).theta_jets(acb_mat([[z]]),2) # doctest: +SKIP + [[0.969443038779670 +/- 5.67e-16] + [-0.0305569612081680 +/- 5.13e-17]j, + [-0.191993710594950 +/- 4.89e-16] + [0.191993710747776 +/- 7.42e-16]j, + [0.60317023860834 +/- 2.93e-15] + [0.60317023764810 +/- 4.86e-15]j] + [[1.03055696119601 +/- 3.89e-15] + [0.0305569612081680 +/- 5.13e-17]j, + [0.191993710594950 +/- 4.89e-16] + [-0.191993710442123 +/- 4.62e-16]j, + [-0.60317023668787 +/- 5.90e-15] + [-0.60317023764810 +/- 4.86e-15]j] + [[-1.22079026757697 +/- 4.36e-15] + [-1.82705551679115 +/- 5.17e-15]j, + [-5.71849316258739 +/- 7.02e-15] + [3.82088827346268 +/- 5.75e-15]j, + [6.0241074288587 +/- 3.88e-14] + [9.0163253443780 +/- 2.05e-14]j] + [[-1.82023591012499 +/- 2.67e-15] + [1.21625195015448 +/- 4.14e-15]j, + [3.8353056542516 +/- 4.99e-14] + [5.73981078971270 +/- 6.74e-15]j, + [8.9823364151977 +/- 2.44e-14] + [-6.0022138700195 +/- 3.72e-14]j] + """ + g = tau.nrows() + if g == 0: + return [] + + # Calculate the length of the jet for one characteristic + # This is the number of multi-indices (alpha) such that |alpha| < ord + cdef slong nj = acb_theta_jet_nb(g, ord) + + # Total number of characteristics + cdef slong nb = 1 << (2 * g) + + # Total number of acb elements to allocate + # FLINT stores nj coefficients for each of the nb characteristics + cdef slong total_size = nb * nj + + cdef acb_ptr zvec = _acb_vec_init(g) + cdef slong i, j + for i in range(g): + acb_set(zvec + i, acb_mat_entry(z.val, i, 0)) + + # Parameters for characteristics + cdef slong nb_in = 1 # Number of input z vectors + cdef ulong ab = 0 # Base characteristic + cdef ulong all = True # Compute all 2^2g characteristics + cdef ulong square = False # Don't compute the squares of the thetas. + + # Initialize the output buffer + cdef acb_ptr theta = _acb_vec_init(total_size) + + # Call the FLINT C function + # Note: Computes all partial derivatives up to total order 'ord' + acb_theta_jet(theta, zvec, nb_in, tau.val, ord, ab, all, square, getprec()) + + # Copy the output into a structured format + # We return a list of lists: res[char_idx] = [coeff_0, coeff_1, ...] + res = [] + cdef acb r + for i in range(nb): + char_jet = [] + for j in range(nj): + r = acb.__new__(acb) + # Offset: characteristic index * length of one jet + coefficient index + acb_set(r.val, theta + (i * nj + j)) + char_jet.append(r) + res.append(char_jet) + + # Cleanup + _acb_vec_clear(zvec, g) + _acb_vec_clear(theta, total_size) + + return acb_mat(res) \ No newline at end of file