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
+

meson.build

@@ -73,9 +73,6 @@ endif
 # flint.pc was missing -lflint until Flint 3.1.0
 if flint_dep.version().version_compare('<3.1')
   flint_dep = cc.find_library('flint')
-  have_acb_theta = false
-else
-  have_acb_theta = true
 endif
 
 pyflint_deps = [gmp_dep, mpfr_dep, flint_dep]

setup.py

@@ -119,6 +119,7 @@
     ("flint.types.acb_poly", ["src/flint/types/acb_poly.pyx"]),
     ("flint.types.acb_mat", ["src/flint/types/acb_mat.pyx"]),
     ("flint.types.acb_series", ["src/flint/types/acb_series.pyx"]),
+    ("flint.types.acb_theta", ["src/flint/types/acb_theta.pyx"]),
 
     ("flint.types.dirichlet", ["src/flint/types/dirichlet.pyx"]),
 

src/flint/__init__.py

@@ -51,22 +51,6 @@
 __version__ = "0.8.0"
 
 
-def _flint_version_at_least(major: int, minor: int) -> bool:
-    version_parts = __FLINT_VERSION__.split(".")
-    if len(version_parts) < 2:
-        return False
-    try:
-        current_major = int(version_parts[0])
-        current_minor = int(version_parts[1])
-    except ValueError:
-        return False
-    return (current_major, current_minor) >= (major, minor)
-
-
-def _has_acb_theta() -> bool:
-    return _flint_version_at_least(3, 1)
-
-
 __all__ = [
     "ctx",
     "fmpz",

src/flint/test/test_acb_mat.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
-import sys
-from unittest.mock import patch
-
-from flint import _has_acb_theta, acb, acb_mat, arb_mat, ctx, fmpq_mat, fmpz_mat
-from flint.test.helpers import is_close_acb, is_close_acb_mat as is_close, is_close_arb_mat, raises
+from flint import acb, acb_mat, arb_mat, ctx, fmpq_mat, fmpz_mat
+from flint.test.helpers import (
+    is_close_acb,
+    is_close_acb_mat as is_close,
+    is_close_arb_mat,
+    raises,
+)
+from flint.types.acb_theta import acb_theta
 
 
 class _DummyMatrix:
@@ -304,17 +307,14 @@ def test_acb_mat_eig_theta_and_helper() -> None:
     assert acb_mat(0, 0).eig() == []
 
     tau = acb_mat([[1j]])
-    if _has_acb_theta():
+    if acb_theta is not None:
         theta_vals = acb_mat.theta(tau, acb_mat([[0]]))
         assert isinstance(theta_vals, acb_mat)
         assert theta_vals.nrows() == 1
         assert theta_vals.ncols() == 4
     else:
         assert raises(lambda: acb_mat.theta(tau, acb_mat([[0]])), NotImplementedError)
 
-    with patch.dict(sys.modules, {"flint.types.acb_theta": None}):
-        assert raises(lambda: acb_mat.theta(tau, acb_mat([[0]])), NotImplementedError)
-
     assert is_close(a, [[1, 0], [0, 2]]) is True
     assert is_close(a, acb_mat([[1, 0], [0, 2]])) is True
     assert is_close(a, [[1, 0, 0], [0, 2, 0]]) is False

src/flint/test/test_acb_theta.py

@@ -1,37 +1,35 @@
 from __future__ import annotations
 
-from flint import _has_acb_theta, acb, acb_mat
+from flint import acb, acb_mat
 from flint.test.helpers import is_close_acb_mat as is_close, raises
+from flint.types.acb_theta import acb_theta, acb_theta_jets
 
 
 def test_acb_theta_basic() -> None:
-    if not _has_acb_theta():
+    if acb_theta is None:
         return
-
-    from flint.types.acb_theta import acb_theta
+    theta = acb_theta
 
     z = acb(1 + 1j)
     tau = acb(1.25 + 3j)
     zmat = acb_mat([[z]])
     taumat = acb_mat([[tau]])
 
-    direct = acb_theta(zmat, taumat)
+    direct = theta(zmat, taumat)
     via_method = taumat.theta(zmat)
     assert is_close(direct, via_method, tol=1e-12, rel_tol=1e-12, max_width=1e-12)
     assert direct.nrows() == 1
     assert direct.ncols() == 4
 
-    squared = acb_theta(zmat, taumat, square=True)
+    squared = theta(zmat, taumat, square=True)
     assert squared.nrows() == 1
     assert squared.ncols() == 4
 
 
 def test_acb_theta_shape_assertions() -> None:
-    if not _has_acb_theta():
+    if acb_theta is None:
         return
 
-    from flint.types.acb_theta import acb_theta
-
     z = acb_mat([[0]])
     tau_bad = acb_mat([[1j, 0]])
     assert raises(lambda: acb_theta(z, tau_bad), AssertionError)
@@ -45,3 +43,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 acb_theta_jets is None:
+        return
+
+    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/test/test_docstrings.py

@@ -10,7 +10,7 @@
 test_flint_at_least = {
     "flint.types._gr.gr_ctx.gens": 30100,
     "flint.types._gr.gr_ctx.neg": 30100,
-    "flint.types.acb_theta.acb_theta": 30300,
+    "flint.types.acb_theta.acb_theta": 30100,
 }
 
 

src/flint/types/acb_mat.pyi

@@ -96,6 +96,7 @@ class acb_mat(flint_mat[acb]):
     ) -> Any: ...
 
     def theta(self, z: acb_mat, square: bool = False) -> acb_mat: ...
+    def theta_jets(self, z: acb_mat, ord: int) -> acb_mat: ...
 
     def str(self, n: int = 0, radius: bool = True, more: bool = False, condense: int = 0) -> _str: ...
     def repr(self) -> _str: ...

src/flint/types/acb_mat.pyx

@@ -5,9 +5,8 @@ from flint.types.arb_mat cimport arb_mat
 from flint.types.fmpz_mat cimport fmpz_mat
 from flint.types.fmpq_mat cimport fmpq_mat
 from flint.types.arb cimport arb
-from flint.types.acb cimport acb
+from flint.types.acb cimport acb, any_as_acb
 from flint.types.acb_poly cimport acb_poly
-from flint.types.acb cimport any_as_acb
 from flint.types.fmpz cimport fmpz
 from flint.types.fmpq cimport fmpq
 
@@ -38,6 +37,8 @@ from flint.flintlib.types.acb cimport (
 
 cimport cython
 
+from flint.types.acb_theta import acb_theta, acb_theta_jets
+
 cdef acb_mat_coerce_operands(x, y):
     if isinstance(y, (fmpz_mat, fmpq_mat, arb_mat)):
         return x, acb_mat(y)
@@ -834,8 +835,25 @@ cdef class acb_mat(flint_mat):
         for the ordering of the theta characteristics.
 
         """
-        try:
-            from .acb_theta import acb_theta
-        except ImportError:
+        if acb_theta is None:
             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.
+        The result is an ``acb_mat`` with one row for each theta characteristic and
+        one column for each jet coefficient.
+
+        """
+        if acb_theta_jets is None:
+            raise NotImplementedError("acb_mat.theta_jets needs Flint >= 3.3.0")
+        return acb_theta_jets(z, tau, ord)

src/flint/types/acb_theta.pyi

@@ -1,6 +1,18 @@
 from __future__ import annotations
 
+from typing import Protocol
+
 from flint.types.acb_mat import acb_mat
 
 
-def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
+class _AcbThetaFunc(Protocol):
+    def __call__(self, z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...
+
+
+class _AcbThetaJetsFunc(Protocol):
+    def __call__(self, z: acb_mat, tau: acb_mat, ord: int) -> acb_mat: ...
+
+
+acb_theta: _AcbThetaFunc | None
+
+acb_theta_jets: _AcbThetaJetsFunc | None