Wrapping acb_theta_all

Author: edgarcosta

src/flint/flintlib/acb_theta.pxd

@@ -2,15 +2,30 @@ from flint.flintlib.acb cimport acb_t, acb_srcptr, acb_ptr
 from flint.flintlib.acb_poly cimport acb_poly_struct, acb_poly_t
 from flint.flintlib.arb cimport arb_t, arb_ptr, arb_srcptr
 from flint.flintlib.flint cimport ulong, slong, flint_rand_t
-from flint.flintlib.fmpz_mat cimport fmpz_mat_t
+from flint.flintlib.fmpz_mat cimport fmpz_mat_t, fmpz_mat_struct
 from flint.flintlib.acb_mat cimport acb_mat_t
 from flint.flintlib.arb_mat cimport arb_mat_t
 from flint.flintlib.arf cimport arf_t
 
 
-# unimported types  {'acb_theta_eld_t', 'acb_theta_naive_worker_t', 'acb_theta_ql_worker_t'}
-
 cdef extern from "flint/acb_theta.h":
+    ctypedef struct acb_theta_eld_struct:
+        slong dim
+        slong ambient_dim
+        slong * last_coords
+        slong min
+        slong mid
+        slong max
+        slong nr
+        slong nl
+        acb_theta_eld_struct * rchildren
+        acb_theta_eld_struct * lchildren
+        slong nb_pts, nb_border
+        slong * box
+    ctypedef acb_theta_eld_struct acb_theta_eld_t[1]
+    ctypedef void (*acb_theta_naive_worker_t)(acb_ptr, acb_srcptr, acb_srcptr, const slong *, slong, const acb_t, const slong *, slong, slong, slong, slong)
+    ctypedef int (*acb_theta_ql_worker_t)(acb_ptr, acb_srcptr, acb_srcptr, arb_srcptr, arb_srcptr, const acb_mat_t, slong, slong)
+
     void acb_theta_all(acb_ptr th, acb_srcptr z, const acb_mat_t tau, int sqr, slong prec)
     void acb_theta_naive_fixed_ab(acb_ptr th, ulong ab, acb_srcptr zs, slong nb, const acb_mat_t tau, slong prec)
     void acb_theta_naive_all(acb_ptr th, acb_srcptr zs, slong nb, const acb_mat_t tau, slong prec)

src/flint/types/acb_mat.pyx

@@ -19,6 +19,7 @@ from flint.flintlib.arb_mat cimport *
 from flint.flintlib.arf cimport *
 from flint.flintlib.acb cimport *
 from flint.flintlib.acb_mat cimport *
+from flint.flintlib.acb_theta cimport *
 
 cimport cython
 
@@ -814,3 +815,35 @@ cdef class acb_mat(flint_mat):
         if left:
             return Elist, L
         return Elist, R
+
+    def theta(tau, z, square=False):
+        r"""
+        FIXME
+        """
+        g = tau.nrows()
+        assert g > 1
+        assert tau.ncols() == g
+        assert z.nrows() == g
+        assert z.ncols() == 1
+
+        # convert input
+        cdef acb_ptr zvec
+        zvec = _acb_vec_init(g)
+        cdef long i
+        for i in range(g):
+            acb_set(zvec + i, acb_mat_entry((<acb_mat>z).val, i, 0))
+
+        # initialize the output
+        cdef slong nb = 1 << (2 * g)
+        cdef acb_ptr theta = _acb_vec_init(nb)
+
+        acb_theta_all(theta, zvec, tau.val, square, getprec())
+        _acb_vec_clear(zvec, g)
+        # copy the output
+        res = tuple()
+        for i in range(nb):
+            r = acb.__new__(acb)
+            acb_set((<acb>r).val, theta + i)
+            res += (r,)
+        _acb_vec_clear(theta, nb)
+        return res