acb_theta_jet

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
+

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: ...

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 <https://flintlib.org/doc/acb_theta.html#c.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)

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]: ...

src/flint/types/acb_theta.pyx

@@ -95,3 +95,84 @@ 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 <https://flintlib.org/doc/acb_theta.html#c.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)
+