acb_theta: add type stubs

Author: oscarbenjamin

src/flint/test/meson.build

@@ -9,6 +9,7 @@ pyfiles = [
   'test_acb_mat.py',
   'test_acb_poly.py',
   'test_acb_series.py',
+  'test_acb_theta.py',
   'test_arb.py',
   'test_arb_mat.py',
   'test_arb_series.py',

src/flint/test/test_acb_theta.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from flint import acb, acb_mat
+from flint.test.helpers import is_close_acb_mat as is_close, raises
+
+
+def _acb_theta_importable() -> bool:
+    try:
+        from flint.types.acb_theta import acb_theta as _  # noqa: F401
+        return True
+    except ImportError:
+        return False
+
+
+def test_acb_theta_basic() -> None:
+    if not _acb_theta_importable():
+        return
+
+    from flint.types.acb_theta import acb_theta
+
+    z = acb(1 + 1j)
+    tau = acb(1.25 + 3j)
+    zmat = acb_mat([[z]])
+    taumat = acb_mat([[tau]])
+
+    direct = acb_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)
+    assert squared.nrows() == 1
+    assert squared.ncols() == 4
+
+
+def test_acb_theta_shape_assertions() -> None:
+    if not _acb_theta_importable():
+        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)
+
+    tau = acb_mat([[1j]])
+    z_bad_rows = acb_mat([[0], [0]])
+    assert raises(lambda: acb_theta(z_bad_rows, tau), AssertionError)
+
+    z_bad_cols = acb_mat([[0, 0]])
+    assert raises(lambda: acb_theta(z_bad_cols, tau), AssertionError)
+
+    assert raises(lambda: acb_theta(object(), tau), TypeError)  # type: ignore[arg-type]
+    assert raises(lambda: acb_theta(z, object()), TypeError)  # type: ignore[arg-type]

src/flint/types/acb_mat.pyi

@@ -95,8 +95,7 @@ class acb_mat(flint_mat[acb]):
         nonstop: bool = False,
     ) -> Any: ...
 
-    @staticmethod
-    def theta(tau: acb_mat, z: object, square: bool = False) -> Any: ...
+    def theta(self, z: acb_mat, square: bool = False) -> 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_theta.pyi

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from flint.types.acb_mat import acb_mat
+
+
+def acb_theta(z: acb_mat, tau: acb_mat, square: bool | int = False) -> acb_mat: ...

src/flint/types/meson.build

@@ -41,6 +41,7 @@ pyfiles = [
   #
   'acb_poly.pyi',
   'acb_mat.pyi',
+  'acb_theta.pyi',
   'acb_series.pyi',
   #
   'dirichlet.pyi',