acb_theta_jet

README.md

@@ -173,6 +173,8 @@ Contributors (0.9.0):
 - Oscar Benjamin (OB)
 - Daniel Simmons-Marengo (DSM)
 - ForeverHaibara (FH)
+- Jeroen Hanselman (JH)
+- Sam Schiavone (SS)
 
 Changes (0.9.0):
 
@@ -198,6 +200,7 @@ Changes (0.9.0):
 - [gh-359](https://github.com/flintlib/python-flint/pull/359),
   Sort factorisations of all mpoly types. (OB)
 - [gh-374](https://github.com/flintlib/python-flint/pull/374), Fixed a bug in `nmod.__hash__`. (FH)
+- [gh-392](https://github.com/flintlib/python-flint/pull/392), Add interface to `acb_theta_jet` (JH, SS)
 
 0.8.0
 -----

doc/source/acb_theta.rst

@@ -3,3 +3,5 @@
 
 .. autofunction :: flint.types.acb_theta.acb_theta
 
+.. autofunction :: flint.types.acb_theta.acb_theta_jets
+

src/flint/__init__.py

@@ -66,6 +66,9 @@ def _flint_version_at_least(major: int, minor: int) -> bool:
 def _has_acb_theta() -> bool:
     return _flint_version_at_least(3, 1)
 
+def _has_acb_theta_jet() -> bool:
+    return _flint_version_at_least(3, 3)
+
 
 __all__ = [
     "ctx",

src/flint/test/test_acb_theta.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from flint import _has_acb_theta, acb, acb_mat
+from flint import _has_acb_theta, _has_acb_theta_jet acb, acb_mat
 from flint.test.helpers import is_close_acb_mat as is_close, raises
 
 
@@ -45,3 +45,20 @@ 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_jet():
+        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,76 @@ 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 acb_mat(0, 0)
+
+    # 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
+    res_mat = acb_mat(nb, nj)
+    for i in range(nb):
+        for j in range(nj):
+            acb_set(acb_mat_entry(res_mat.val, i, j), theta + (i * nj + j))
+
+    # Cleanup
+    _acb_vec_clear(zvec, g)
+    _acb_vec_clear(theta, total_size)
+
+    return res_mat