Start adding sqr_pow as part of ctx_tau

src/acb_theta.h

@@ -139,6 +139,7 @@ acb_theta_ctx_z_struct * acb_theta_ctx_z_vec_init(slong nb, slong g);
 void acb_theta_ctx_z_vec_clear(acb_theta_ctx_z_struct * vec, slong nb);
 
 void acb_theta_ctx_tau_set(acb_theta_ctx_tau_t ctx, const acb_mat_t tau, slong prec);
+void acb_theta_ctx_tau_fit_eld(acb_theta_ctx_tau_t ctx, const acb_theta_eld_t E, slong prec);
 void acb_theta_ctx_tau_dupl(acb_theta_ctx_tau_t ctx, slong prec);
 int acb_theta_ctx_tau_overlaps(const acb_theta_ctx_tau_t ctx1, const acb_theta_ctx_tau_t ctx2);
 

src/acb_theta/ctx_tau_clear.c

@@ -20,6 +20,7 @@ acb_theta_ctx_tau_clear(acb_theta_ctx_tau_t ctx)
 {
     slong g = ctx->g;
     slong n = 1 << g;
+    slong j;
 
     arb_mat_clear(&ctx->yinv);
     arb_mat_clear(&ctx->cho);
@@ -30,6 +31,13 @@ acb_theta_ctx_tau_clear(acb_theta_ctx_tau_t ctx)
     acb_mat_clear(ctx->exp_tau_div_2_inv);
     acb_mat_clear(ctx->exp_tau_inv);
 
+    for (j = 0; j < g; j++)
+    {
+        _acb_vec_clear(ctx->sqr_pow[j], ctx->sqr_pow_len[j]);
+    }
+    flint_free(ctx->sqr_pow);
+    flint_free(ctx->sqr_pow_len);
+
     if (ctx->allow_shift)
     {
         _acb_vec_clear(ctx->exp_tau_a, n * g);

src/acb_theta/ctx_tau_dupl.c

@@ -49,6 +49,17 @@ acb_theta_ctx_tau_dupl(acb_theta_ctx_tau_t ctx, slong prec)
         }
     }
 
+    /* Empty sqr_pow */
+    for (j = 0; j < g; j++)
+    {
+        if (ctx->sqr_pow_len[j] > 0)
+        {
+            _acb_vec_clear(ctx->sqr_pow[j], ctx->sqr_pow_len[j]);
+            ctx->sqr_pow[j] = _acb_vec_init(0);
+            ctx->sqr_pow_len[j] = 0;
+        }
+    }
+
     if (ctx->allow_shift)
     {
         /* Update exp_tau_a */

src/acb_theta/ctx_tau_fit_eld.c

@@ -0,0 +1,61 @@
+/*
+    Copyright (C) 2024 Jean Kieffer
+
+    This file is part of FLINT.
+
+    FLINT is free software: you can redistribute it and/or modify it under
+    the terms of the GNU Lesser General Public License (LGPL) as published
+    by the Free Software Foundation; either version 3 of the License, or
+    (at your option) any later version.  See <https://www.gnu.org/licenses/>.
+*/
+
+#include "acb.h"
+#include "acb_mat.h"
+#include "acb_theta.h"
+
+void
+acb_theta_ctx_tau_fit_eld(acb_theta_ctx_tau_t ctx, const acb_theta_eld_t E, slong prec)
+{
+    slong g = ctx->g;
+    acb_ptr cpy;
+    acb_t d, dd;
+    slong j, k, len;
+
+    acb_init(d);
+    acb_init(dd);
+
+    for (j = 0; j < g; j++)
+    {
+        len = ctx->sqr_pow_len[j];
+        if (len < (E->box[j] + 1))
+        {
+            cpy = ctx->sqr_pow[j];
+            ctx->sqr_pow[j] = _acb_vec_init(E->box[j] + 1);
+            ctx->sqr_pow_len[j] = (E->box[j] + 1);
+            _acb_vec_set(ctx->sqr_pow[j], cpy, len);
+            _acb_vec_clear(cpy, len);
+
+            acb_sqr(dd, acb_mat_entry(ctx->exp_tau, j, j), prec);
+            if (len == 0)
+            {
+                acb_one(&(ctx->sqr_pow[j])[len]);
+                acb_set(d, acb_mat_entry(ctx->exp_tau, j, j));
+            }
+            else
+            {
+                acb_pow_ui(d, acb_mat_entry(ctx->exp_tau, j, j), 2 * len - 1, prec);
+                acb_mul(&(ctx->sqr_pow[j])[len], &(ctx->sqr_pow[j])[len - 1], d, prec);
+                acb_mul(d, d, dd, prec);
+            }
+
+            for (k = len + 1; k <= (E->box[j]); k++)
+            {
+                acb_mul(&(ctx->sqr_pow[j])[k], &(ctx->sqr_pow[j])[k - 1], d, prec);
+                acb_mul(d, d, dd, prec);
+            }
+        }
+    }
+
+    acb_clear(d);
+    acb_clear(dd);
+}

src/acb_theta/ctx_tau_init.c

@@ -19,6 +19,8 @@ void
 acb_theta_ctx_tau_init(acb_theta_ctx_tau_t ctx, int allow_shift, slong g)
 {
     slong n = 1 << g;
+    slong j;
+
     FLINT_ASSERT(g >= 1);
 
     ctx->g = g;
@@ -33,6 +35,13 @@ acb_theta_ctx_tau_init(acb_theta_ctx_tau_t ctx, int allow_shift, slong g)
     acb_mat_init(ctx->exp_tau_div_2_inv, g, g);
     acb_mat_init(ctx->exp_tau_inv, g, g);
 
+    ctx->sqr_pow_len = flint_malloc(g * sizeof(slong));
+    ctx->sqr_pow = flint_malloc(g * sizeof(acb_ptr));
+    for (j = 0; j < g; j++)
+    {
+        ctx->sqr_pow[j] = _acb_vec_init(0);
+    }
+
     if (allow_shift)
     {
         ctx->exp_tau_a = _acb_vec_init(g * n);

src/acb_theta/sum.c

@@ -110,6 +110,8 @@ acb_theta_sum_0x(acb_ptr th, const acb_theta_ctx_z_struct * vec, slong nb,
 
     if (b)
     {
+        acb_theta_ctx_tau_fit_eld(ctx_tau, E, prec);
+
         for (j = 0; j < nb; j++)
         {
             acb_theta_sum_work(th + j * n, n, (&vec[j])->exp_2z,

src/acb_theta/sum_jet.c

@@ -421,6 +421,8 @@ acb_theta_sum_jet(acb_ptr th, const acb_theta_ctx_z_struct * vec, slong nb,
 
         if (b)
         {
+            acb_theta_ctx_tau_fit_eld(ctx_tau, E, prec);
+
             /* Sum series, rescale by c factor */
             for (j = 0; j < nb; j++)
             {

src/acb_theta_types.h

@@ -45,6 +45,9 @@ struct acb_theta_ctx_tau_struct
     acb_mat_t exp_tau_div_2_inv;
     acb_mat_t exp_tau_inv;
 
+    slong * sqr_pow_len;
+    acb_ptr * sqr_pow;
+
     /* allow_shift only */
     acb_ptr exp_tau_a;
     acb_ptr exp_tau_a_inv;