Line data Source code
1 : /* Copyright (C) 2015 The PARI group.
2 :
3 : This file is part of the PARI/GP package.
4 :
5 : PARI/GP is free software; you can redistribute it and/or modify it under the
6 : terms of the GNU General Public License as published by the Free Software
7 : Foundation; either version 2 of the License, or (at your option) any later
8 : version. It is distributed in the hope that it will be useful, but WITHOUT
9 : ANY WARRANTY WHATSOEVER.
10 :
11 : Check the License for details. You should have received a copy of it, along
12 : with the package; see the file 'COPYING'. If not, write to the Free Software
13 : Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. */
14 :
15 : /********************************************************************/
16 : /** **/
17 : /** Dirichlet series through Euler product **/
18 : /** **/
19 : /********************************************************************/
20 : #include "pari.h"
21 : #include "paripriv.h"
22 :
23 : static void
24 28 : err_direuler(GEN x)
25 28 : { pari_err_DOMAIN("direuler","constant term","!=", gen_1,x); }
26 :
27 : /* s = t_POL (tolerate t_SER of valuation 0) of constant term = 1
28 : * d = minimal such that p^d > X
29 : * V indexed by 1..X will contain the a_n
30 : * v[1..n] contains the indices nj such that V[nj] != 0 */
31 : static long
32 28700 : dirmuleuler_small(GEN V, GEN v, long n, ulong p, GEN s, long d)
33 : {
34 28700 : long i, j, m = n, D = minss(d+2, lg(s));
35 28700 : ulong q = 1, X = lg(V)-1;
36 :
37 94724 : for (i = 3, q = p; i < D; i++, q *= p) /* q*p does not overflow */
38 : {
39 66024 : GEN aq = gel(s,i);
40 66024 : if (gequal0(aq)) continue;
41 : /* j = 1 */
42 53753 : gel(V,q) = aq;
43 53753 : v[++n] = q;
44 3268013 : for (j = 2; j <= m; j++)
45 : {
46 3214260 : ulong nj = umuluu_le(uel(v,j), q, X);
47 3214260 : if (!nj) continue;
48 192017 : gel(V,nj) = gmul(aq, gel(V,v[j]));
49 192017 : v[++n] = nj;
50 : }
51 : }
52 28700 : return n;
53 : }
54 :
55 : /* ap != 0 for efficiency, p > sqrt(X) */
56 : static void
57 308798 : dirmuleuler_large(GEN V, ulong p, GEN ap)
58 : {
59 308798 : long j, jp, X = lg(V)-1;
60 308798 : gel(V,p) = ap;
61 1506547 : for (j = 2, jp = 2*p; jp <= X; j++, jp += p) gel(V,jp) = gmul(ap, gel(V,j));
62 308798 : }
63 :
64 : static ulong
65 10269 : direulertou(GEN a, GEN fl(GEN))
66 : {
67 10269 : if (typ(a) != t_INT)
68 : {
69 49 : a = fl(a);
70 28 : if (typ(a) != t_INT) pari_err_TYPE("direuler", a);
71 : }
72 10248 : return signe(a)<=0 ? 0: itou(a);
73 : }
74 :
75 : static GEN
76 3724 : direuler_Sbad(GEN V, GEN v, GEN Sbad, ulong *n)
77 : {
78 3724 : long i, l = lg(Sbad);
79 3724 : ulong X = lg(V)-1;
80 3724 : GEN pbad = gen_1;
81 9646 : for (i = 1; i < l; i++)
82 : {
83 5957 : GEN ai = gel(Sbad,i);
84 : ulong q;
85 5957 : if (typ(ai) != t_VEC || lg(ai) != 3)
86 14 : pari_err_TYPE("direuler [bad primes]",ai);
87 5943 : q = gtou(gel(ai,1));
88 5936 : if (q <= X)
89 : {
90 4809 : long d = ulogint(X, q) + 1;
91 4809 : GEN s = direuler_factor(gel(ai,2), d);
92 4795 : *n = dirmuleuler_small(V, v, *n, q, s, d);
93 4795 : pbad = muliu(pbad, q);
94 : }
95 : }
96 3689 : return pbad;
97 : }
98 :
99 : GEN
100 672 : direuler_bad(void *E, GEN (*eval)(void *,GEN,long), GEN a,GEN b,GEN c, GEN Sbad)
101 : {
102 : ulong au, bu, X, sqrtX, n, p;
103 672 : pari_sp av0 = avma;
104 : GEN gp, v, V;
105 : forprime_t T;
106 672 : au = direulertou(a, gceil);
107 665 : bu = direulertou(b, gfloor);
108 658 : X = c ? direulertou(c, gfloor): bu;
109 651 : if (X == 0) return cgetg(1,t_VEC);
110 644 : if (bu > X) bu = X;
111 644 : if (!u_forprime_init(&T, au, bu)) { set_avma(av0); return mkvec(gen_1); }
112 630 : v = vecsmall_ei(X, 1);
113 630 : V = vec_ei(X, 1);
114 630 : n = 1;
115 630 : if (Sbad) Sbad = direuler_Sbad(V, v, Sbad, &n);
116 595 : p = 1; gp = cgetipos(3); sqrtX = usqrt(X);
117 8316 : while (p <= sqrtX && (p = u_forprime_next(&T)))
118 7742 : if (!Sbad || umodiu(Sbad, p))
119 : {
120 7637 : long d = ulogint(X, p) + 1; /* minimal d such that p^d > X */
121 : GEN s;
122 7637 : gp[2] = p; s = eval(E, gp, d);
123 7616 : n = dirmuleuler_small(V, v, n, p, s, d);
124 : }
125 740201 : while ((p = u_forprime_next(&T))) /* sqrt(X) < p <= X */
126 739627 : if (!Sbad || umodiu(Sbad, p))
127 : {
128 : GEN s;
129 739620 : gp[2] = p; s = eval(E, gp, 2); /* s either t_POL or t_SER of val 0 */
130 739620 : if (lg(s) > 3 && !gequal0(gel(s,3)))
131 139153 : dirmuleuler_large(V, p, gel(s,3));
132 : }
133 574 : return gerepilecopy(av0,V);
134 : }
135 :
136 : /* return a t_SER or a truncated t_POL to precision n */
137 : GEN
138 752066 : direuler_factor(GEN s, long n)
139 : {
140 752066 : long t = typ(s);
141 752066 : if (is_scalar_t(t))
142 : {
143 33194 : if (!gequal1(s)) err_direuler(s);
144 33180 : return scalarpol_shallow(s,0);
145 : }
146 718872 : switch(t)
147 : {
148 5712 : case t_POL: break; /* no need to RgXn_red */
149 712845 : case t_RFRAC:
150 : {
151 712845 : GEN p = gel(s,1), q = gel(s,2);
152 712845 : q = RgXn_red_shallow(q,n);
153 712845 : s = RgXn_inv(q, n);
154 712845 : if (typ(p) == t_POL && varn(p) == varn(q))
155 : {
156 28 : p = RgXn_red_shallow(p, n);
157 28 : s = RgXn_mul(s, p, n);
158 : }
159 : else
160 712817 : if (!gequal1(p)) s = RgX_Rg_mul(s, p);
161 712845 : if (!signe(s) || !gequal1(gel(s,2))) err_direuler(s);
162 712831 : break;
163 : }
164 308 : case t_SER:
165 308 : if (!signe(s) || valser(s) || !gequal1(gel(s,2))) err_direuler(s);
166 308 : break;
167 7 : default: pari_err_TYPE("direuler", s);
168 : }
169 718851 : return s;
170 : }
171 :
172 : struct eval_bad
173 : {
174 : void *E;
175 : GEN (*eval)(void *, GEN);
176 : };
177 : static GEN
178 688303 : eval_bad(void *E, GEN p, long n)
179 : {
180 688303 : struct eval_bad *d = (struct eval_bad*) E;
181 688303 : return direuler_factor(d->eval(d->E, p), n);
182 : }
183 : GEN
184 301 : direuler(void *E, GEN (*eval)(void *, GEN), GEN a, GEN b, GEN c)
185 : {
186 : struct eval_bad d;
187 301 : d.E= E; d.eval = eval;
188 301 : return direuler_bad((void*)&d, eval_bad, a, b, c, NULL);
189 : }
190 :
191 : static GEN
192 31157 : primelist(forprime_t *T, GEN Sbad, long n, long *running)
193 : {
194 31157 : GEN P = cgetg(n+1, t_VECSMALL);
195 : long i, j;
196 302631 : for (i = 1, j = 1; i <= n; i++)
197 : {
198 275548 : ulong p = u_forprime_next(T);
199 275548 : if (!p) { *running = 0; break; }
200 271474 : if (Sbad && umodiu(Sbad, p)==0) continue;
201 266791 : uel(P,j++) = p;
202 : }
203 31157 : setlg(P, j);
204 31157 : return P;
205 : }
206 :
207 : GEN
208 4081 : pardireuler(GEN worker, GEN a, GEN b, GEN c, GEN Sbad)
209 : {
210 : ulong au, bu, X, sqrtX, n, snX, nX;
211 4081 : pari_sp av0 = avma;
212 : GEN v, V;
213 : forprime_t T;
214 : struct pari_mt pt;
215 4081 : long running = 1, pending = 0;
216 4081 : au = direulertou(a, gceil);
217 4081 : bu = direulertou(b, gfloor);
218 4081 : X = c ? direulertou(c, gfloor): bu;
219 4081 : if (X == 0) return cgetg(1,t_VEC);
220 4081 : if (bu > X) bu = X;
221 4081 : if (!u_forprime_init(&T, au, bu)) { set_avma(av0); return mkvec(gen_1); }
222 4074 : v = vecsmall_ei(X, 1);
223 4074 : V = vec_ei(X, 1);
224 4074 : n = 1;
225 4074 : if (Sbad) Sbad = direuler_Sbad(V, v, Sbad, &n);
226 4074 : sqrtX = usqrt(X); snX = uprimepi(sqrtX); nX = uprimepi(X);
227 4074 : if (snX)
228 : {
229 4060 : GEN P = primelist(&T, Sbad, snX, &running);
230 4060 : GEN R = gel(closure_callgenvec(worker, mkvec2(P, utoi(X))), 2);
231 4060 : long i, l = lg(P);
232 20349 : for (i = 1; i < l; i++)
233 : {
234 16289 : GEN s = gel(R,i);
235 16289 : n = dirmuleuler_small(V, v, n, uel(P,i), s, lg(s));
236 : }
237 14 : } else snX = 1;
238 4074 : mt_queue_start_lim(&pt, worker, (nX+snX-1)/snX);
239 34212 : while (running || pending)
240 : {
241 : GEN done;
242 30138 : GEN P = running? primelist(&T, Sbad, snX, &running): NULL;
243 30138 : mt_queue_submit(&pt, 0, P ? mkvec2(P, utoi(X)): NULL);
244 30138 : done = mt_queue_get(&pt, NULL, &pending);
245 30138 : if (done)
246 : {
247 27097 : GEN P = gel(done,1), R = gel(done,2);
248 27097 : long j, l = lg(P);
249 277599 : for (j=1; j<l; j++)
250 : {
251 250502 : GEN F = gel(R,j);
252 250502 : if (degpol(F) && !gequal0(gel(F,3)))
253 169645 : dirmuleuler_large(V, uel(P,j), gel(F,3));
254 : }
255 : }
256 : }
257 4074 : mt_queue_end(&pt);
258 4074 : return gerepilecopy(av0,V);
259 : }
260 :
261 : /********************************************************************/
262 : /** **/
263 : /** DIRPOWERS and DIRPOWERSSUM **/
264 : /** **/
265 : /********************************************************************/
266 :
267 : /* [1^B,...,N^B] */
268 : GEN
269 686 : vecpowuu(long N, ulong B)
270 : {
271 : GEN v;
272 : long p, i;
273 : forprime_t T;
274 :
275 686 : if (B <= 8000)
276 : {
277 686 : if (!B) return const_vec(N,gen_1);
278 679 : v = cgetg(N+1, t_VEC); if (N == 0) return v;
279 679 : gel(v,1) = gen_1;
280 679 : if (B == 1)
281 92736 : for (i = 2; i <= N; i++) gel(v,i) = utoipos(i);
282 469 : else if (B == 2)
283 : {
284 : ulong o, s;
285 273 : if (N & HIGHMASK)
286 0 : for (i = 2, o = 3; i <= N; i++, o += 2)
287 0 : gel(v,i) = addiu(gel(v,i-1), o);
288 : else
289 31073 : for (i = 2, s = 1, o = 3; i <= N; i++, s += o, o += 2)
290 30800 : gel(v,i) = utoipos(s + o);
291 : }
292 196 : else if (B == 3)
293 840 : for (i = 2; i <= N; i++) gel(v,i) = powuu(i, B);
294 : else
295 : {
296 182 : long k, Bk, e = expu(N);
297 7553 : for (i = 3; i <= N; i += 2) gel(v,i) = powuu(i, B);
298 1239 : for (k = 1; k <= e; k++)
299 : {
300 1057 : N >>= 1; Bk = B * k;
301 8498 : for (i = 1; i <= N; i += 2) gel(v, i << k) = shifti(gel(v, i), Bk);
302 : }
303 : }
304 679 : return v;
305 : }
306 0 : v = const_vec(N, NULL);
307 0 : u_forprime_init(&T, 3, N);
308 0 : while ((p = u_forprime_next(&T)))
309 : {
310 : long m, pk, oldpk;
311 0 : gel(v,p) = powuu(p, B);
312 0 : for (pk = p, oldpk = p; pk; oldpk = pk, pk = umuluu_le(pk,p,N))
313 : {
314 0 : if (pk != p) gel(v,pk) = mulii(gel(v,oldpk), gel(v,p));
315 0 : for (m = N/pk; m > 1; m--)
316 0 : if (gel(v,m) && m%p) gel(v, m*pk) = mulii(gel(v,m), gel(v,pk));
317 : }
318 : }
319 0 : gel(v,1) = gen_1;
320 0 : for (i = 2; i <= N; i+=2)
321 : {
322 0 : long vi = vals(i);
323 0 : gel(v,i) = shifti(gel(v,i >> vi), B * vi);
324 : }
325 0 : return v;
326 : }
327 :
328 : /* does n^s require log(x) ? */
329 : static long
330 50565 : get_needlog(GEN s)
331 : {
332 50565 : switch(typ(s))
333 : {
334 294 : case t_REAL: return 2; /* yes but not powcx */
335 47143 : case t_COMPLEX: return 1; /* yes using powcx */
336 3128 : default: return 0; /* no */
337 : }
338 : }
339 : /* [1^B,...,N^B] */
340 : GEN
341 12024 : vecpowug(long N, GEN B, long prec)
342 : {
343 12024 : GEN v, logp = NULL;
344 12024 : long gp[] = {evaltyp(t_INT)|_evallg(3), evalsigne(1)|evallgefint(3),0};
345 12024 : long p, precp = 2, prec0, prec1, needlog;
346 : forprime_t T;
347 12024 : if (N == 1) return mkvec(gen_1);
348 12017 : if (typ(B) == t_INT && lgefint(B) <= 3 && signe(B) >= 0)
349 168 : return vecpowuu(N, itou(B));
350 11849 : needlog = get_needlog(B);
351 11849 : prec1 = prec0 = prec;
352 11849 : if (needlog == 1) prec1 = powcx_prec(log2((double)N), B, prec);
353 11849 : u_forprime_init(&T, 2, N);
354 11849 : v = const_vec(N, NULL);
355 11849 : gel(v,1) = gen_1;
356 1564234 : while ((p = u_forprime_next(&T)))
357 : {
358 : long m, pk, oldpk;
359 : GEN u;
360 1552385 : gp[2] = p;
361 1552385 : if (needlog)
362 : {
363 96234 : if (!logp)
364 17486 : logp = logr_abs(utor(p, prec1));
365 : else
366 : { /* Assuming p and precp are odd,
367 : * log p = log(precp) + 2 atanh((p - precp) / (p + precp)) */
368 78748 : ulong a = p >> 1, b = precp >> 1; /* p = 2a + 1, precp = 2b + 1 */
369 78748 : GEN z = atanhuu(a - b, a + b + 1, prec1); /* avoid overflow */
370 78748 : shiftr_inplace(z, 1); logp = addrr(logp, z);
371 : }
372 92710 : u = needlog == 1? powcx(gp, logp, B, prec0)
373 96234 : : mpexp(gmul(B, logp));
374 96234 : if (p == 2) logp = NULL; /* reset: precp must be odd */
375 : }
376 : else
377 1456151 : u = gpow(gp, B, prec0);
378 1552385 : precp = p;
379 1552385 : gel(v,p) = u; /* p^B */
380 1552385 : if (prec0 != prec) gel(v,p) = gprec_wtrunc(gel(v,p), prec);
381 3218353 : for (pk = p, oldpk = p; pk; oldpk = pk, pk = umuluu_le(pk,p,N))
382 : {
383 1665968 : if (pk != p) gel(v,pk) = gmul(gel(v,oldpk), gel(v,p));
384 46314871 : for (m = N/pk; m > 1; m--)
385 44648903 : if (gel(v,m) && m%p) gel(v, m*pk) = gmul(gel(v,m), gel(v,pk));
386 : }
387 : }
388 11849 : return v;
389 : }
390 :
391 : GEN
392 665 : dirpowers(long n, GEN x, long prec)
393 : {
394 : pari_sp av;
395 : GEN v;
396 665 : if (n <= 0) return cgetg(1, t_VEC);
397 651 : av = avma; v = vecpowug(n, x, prec);
398 651 : if (typ(x) == t_INT && lgefint(x) <= 3 && signe(x) >= 0 && cmpiu(x, 2) <= 0)
399 133 : return v;
400 518 : return gerepilecopy(av, v);
401 : }
402 :
403 : static GEN
404 252 : vecmulsqlv(GEN Q, GEN V)
405 : {
406 : long lq, i;
407 : GEN W;
408 252 : if (typ(V) != t_VEC) return RgV_Rg_mul(Q, V);
409 21 : lq = lg(Q); W = cgetg(lq, t_VEC);
410 672 : for (i = 1; i < lq; i++) gel(W, i) = vecmul(gel(Q, i), V);
411 21 : return W;
412 : }
413 :
414 : /* P = prime divisors of (squarefree) n, V[i] = i^s for i <= sq.
415 : * Return NULL if n is not sq-smooth, else f(n)n^s */
416 : static GEN
417 10691182 : smallfact(ulong n, GEN P, ulong sq, GEN V)
418 : {
419 : long i, l;
420 : ulong p, m, o;
421 : GEN c;
422 10691182 : if (n <= sq) return gel(V,n);
423 10656774 : l = lg(P); m = p = uel(P, l-1); if (p > sq) return NULL;
424 2239953 : for (i = l-2; i > 1; i--, m = o) { p = uel(P,i); o = m*p; if (o > sq) break; }
425 2220151 : c = gel(V,m); n /= m; /* m <= sq, o = m * p > sq */
426 2220151 : if (n > sq) { c = vecmul(c, gel(V,p)); n /= p; }
427 2209791 : return vecmul(c, gel(V,n));
428 : }
429 :
430 : static GEN
431 140 : _Qtor(GEN x, long prec)
432 140 : { return typ(x) == t_FRAC? fractor(x, prec): x; }
433 : static GEN
434 315 : Qtor(GEN x, long prec)
435 : {
436 315 : long tx = typ(x);
437 357 : if (tx == t_VEC || tx == t_COL) pari_APPLY_same(_Qtor(gel(x, i), prec));
438 294 : return tx == t_FRAC? fractor(x, prec): x;
439 : }
440 :
441 : /* Here N > 0 is small */
442 : static GEN
443 147 : naivedirpowerssum(long N, GEN s, void *E, GEN (*f)(void *, ulong, long),
444 : long prec)
445 : {
446 147 : GEN V = vecpowug(N, s, prec), S;
447 147 : if (!f) S = RgV_sum(V);
448 : else
449 : {
450 : long n;
451 35 : S = f(E, 1, prec);
452 308 : for (n = 2; n <= N; n++) S = gadd(S, gmul(gel(V, n), f(E, n, prec)));
453 : }
454 147 : return Qtor(S, prec);
455 : }
456 :
457 : static GEN
458 126 : smalldirpowerssum(long N, GEN s, void *E, GEN (*f)(void *, ulong, long),
459 : long both, long prec)
460 : {
461 126 : GEN S = naivedirpowerssum(N, s, E, f, prec), SB, sb;
462 126 : if (!both) return S;
463 42 : sb = gconj(gsubsg(-1, s));
464 42 : SB = both==2 && gequal(s,sb)? S: gconj(naivedirpowerssum(N,sb,E,f,prec));
465 42 : return mkvec2(S, SB);
466 : }
467 :
468 : static GEN
469 63 : dirpowsuminit(GEN s, void *E, GEN (*f)(void *, ulong, long), GEN data,
470 : long both, long prec)
471 : {
472 63 : GEN onef = gel(data, 1), zervec = gel(data, 2), sqlpp = gel(data, 3);
473 63 : long sq = sqlpp[1], needlog = sqlpp[2], prec0 = sqlpp[3], prec1 = sqlpp[4];
474 63 : GEN V = cgetg(sq+1, t_VEC), W = cgetg(sq+1, t_VEC), Q = cgetg(sq+1, t_VEC);
475 63 : GEN VB = NULL, WB = NULL, QB = NULL, c2, c2B = NULL;
476 63 : GEN Q2, Q3, Q6, Q2B = NULL, Q3B = NULL, Q6B = NULL;
477 63 : GEN logp, R, RB = NULL;
478 63 : long gp[] = {evaltyp(t_INT)|_evallg(3), evalsigne(1)|evallgefint(3),0};
479 : long n;
480 63 : if (both == 1 || (both == 2 && !gequal(real_i(s), gneg(ghalf))))
481 21 : { VB = cgetg(sq+1, t_VEC); WB = cgetg(sq+1, t_VEC); QB = cgetg(sq+1, t_VEC);}
482 63 : gel(V, 1) = gel(W, 1) = gel(Q, 1) = onef;
483 63 : if (VB) { gel(VB, 1) = gel(WB, 1) = gel(QB, 1) = onef; }
484 63 : c2 = gpow(gen_2, s, prec0); if (VB) c2B = ginv(gmul2n(gconj(c2), 1));
485 63 : if (f)
486 : {
487 42 : GEN tmp2 = f(E, 2, prec);
488 42 : c2 = gmul(c2, tmp2); if (VB) c2B = gmul(c2B, tmp2);
489 : }
490 63 : gel(V,2) = c2; /* f(2) 2^s */
491 63 : gel(W,2) = Qtor(gadd(c2, onef), prec0);
492 63 : gel(Q,2) = Qtor(gadd(vecsqr(c2), onef), prec0);
493 63 : if (VB)
494 : {
495 21 : gel(VB, 2) = c2B; gel(WB, 2) = Qtor(gadd(c2B, onef), prec0);
496 21 : gel(QB, 2) = Qtor(gadd(vecsqr(c2B), onef), prec0);
497 : }
498 63 : logp = NULL;
499 4074 : for (n = 3; n <= sq; n++)
500 : {
501 4011 : GEN u = NULL, uB = NULL, ks = f ? f(E, n, prec0) : gen_1;
502 4011 : long zks = !gequal0(ks);
503 4011 : if (odd(n))
504 : {
505 2023 : gp[2] = n;
506 2023 : if (needlog)
507 : {
508 476 : if (!logp)
509 42 : logp = logr_abs(utor(n, prec1));
510 : else
511 : { /* log n = log(n-2) + 2 atanh(1 / (n - 1)) */
512 434 : GEN z = atanhuu(1, n - 1, prec1);
513 434 : shiftr_inplace(z, 1); logp = addrr(logp, z);
514 : }
515 476 : if (zks)
516 476 : u = needlog == 1? powcx(gp, logp, s, prec0) : mpexp(gmul(s, logp));
517 : }
518 1547 : else if (zks) u = gpow(gp, s, prec0);
519 2023 : if (zks)
520 : {
521 2009 : if (VB) uB = gmul(ginv(gmulsg(n, gconj(u))), ks);
522 2009 : u = gmul(u, ks); /* f(n) n^s */
523 : }
524 : }
525 : else
526 : {
527 1988 : u = vecmul(c2, gel(V, n >> 1));
528 1988 : if (VB) uB = vecmul(c2B, gel(VB, n >> 1));
529 : }
530 4011 : if (zks)
531 : { /* V[n]=f(n)n^s, W[n]=sum_{i<=n} f(i)i^s, Q[n]=sum_{i<=n} f(i^2)i^2s */
532 3983 : gel(V,n) = u;
533 3983 : gel(W,n) = gadd(gel(W, n-1), gel(V,n));
534 3983 : gel(Q,n) = gadd(gel(Q, n-1), vecsqr(gel(V,n)));
535 3983 : if (VB)
536 : {
537 462 : gel(VB,n) = uB;
538 462 : gel(WB,n) = gadd(gel(WB,n-1), gel(VB,n));
539 462 : gel(QB,n) = gadd(gel(QB,n-1), vecsqr(gel(VB,n)));
540 : }
541 : }
542 : else
543 : {
544 28 : gel(V,n) = zervec; gel(W,n) = gel(W, n-1); gel(Q,n) = gel(Q, n-1);
545 28 : if (VB)
546 : {
547 0 : gel(VB,n) = zervec; gel(WB,n) = gel(WB, n-1);
548 0 : gel(QB,n) = gel(QB, n-1);
549 : }
550 : }
551 : }
552 63 : Q2 = vecmulsqlv(Q, gel(V,2));
553 63 : Q3 = vecmulsqlv(Q, gel(V,3));
554 63 : Q6 = vecmulsqlv(Q, gel(V,6));
555 63 : if (VB)
556 : {
557 21 : Q2B = vecmulsqlv(QB, gel(VB,2));
558 21 : Q3B = vecmulsqlv(QB, gel(VB,3));
559 21 : Q6B = vecmulsqlv(QB, gel(VB,6));
560 : }
561 63 : R = mkvecn(6, V, W, Q, Q2, Q3, Q6);
562 63 : if (VB) RB = mkvecn(6, VB, WB, QB, Q2B, Q3B, Q6B);
563 63 : return VB ? mkvec2(R, RB) : mkvec(R);
564 : }
565 :
566 : static GEN
567 63 : dirpowsumprimeloop(ulong N, GEN s, void *E, GEN (*f)(void *, ulong, long),
568 : GEN data, GEN W, GEN WB)
569 : {
570 : pari_sp av2;
571 63 : GEN zervec = gel(data, 2), S = zervec, SB = zervec, logp = NULL;
572 63 : GEN sqlpp = gel(data, 3);
573 : forprime_t T;
574 63 : long gp[] = {evaltyp(t_INT)|_evallg(3), evalsigne(1)|evallgefint(3),0};
575 63 : long p, precp = 0, sq = sqlpp[1], needlog = sqlpp[2];
576 63 : long prec0 = sqlpp[3], prec1 = sqlpp[4];
577 63 : u_forprime_init(&T, sq + 1, N);
578 63 : av2 = avma;
579 80969 : while ((p = u_forprime_next(&T)))
580 : {
581 80906 : GEN u = NULL, ks = f ? f(E, p, prec1) : gen_1;
582 80906 : long zks = !gequal0(ks);
583 80906 : gp[2] = p;
584 80906 : if (needlog)
585 : {
586 4690 : if (!logp)
587 42 : logp = logr_abs(utor(p, prec1));
588 : else
589 : { /* log p = log(precp) + 2 atanh((p - precp) / (p + precp)) */
590 4648 : ulong a = p >> 1, b = precp >> 1; /* p = 2a + 1, precp = 2b + 1 */
591 4648 : GEN z = atanhuu(a - b, a + b + 1, prec1); /* avoid overflow */
592 4648 : shiftr_inplace(z, 1); logp = addrr(logp, z);
593 : }
594 4690 : if (zks)
595 4690 : u = needlog == 1? powcx(gp, logp, s, prec0) : mpexp(gmul(s, logp));
596 : }
597 76216 : else { if (zks) u = gpow(gp, s, prec0); }
598 80906 : if (zks)
599 : {
600 80906 : S = gadd(S, vecmul(gel(W, N / p), gmul(ks, u)));
601 80906 : if (WB)
602 2345 : SB = gadd(SB, gdiv(vecmul(ks, gel(WB, N / p)), gmulsg(p, gconj(u))));
603 : }
604 80906 : precp = p;
605 80906 : if ((p & 0x1ff) == 1)
606 : {
607 280 : if (!logp) gerepileall(av2, SB? 2: 1, &S, &SB);
608 0 : else gerepileall(av2, SB? 3: 2, &S, &logp, &SB);
609 : }
610 : }
611 63 : return SB ? mkvec2(S, SB) : mkvec(S);
612 : }
613 :
614 : static GEN
615 413 : dirpowsumsqfloop(long N, long x1, long x2, long sq, GEN P, GEN allvecs, GEN S,
616 : GEN allvecsB, GEN SB)
617 : {
618 413 : GEN V = gel(allvecs, 1), Q = gel(allvecs, 2), Q2 = gel(allvecs, 3);
619 413 : GEN Q3 = gel(allvecs, 4), Q6 = gel(allvecs, 5), Z = gel(allvecs, 6);
620 413 : GEN VB = NULL, QB = NULL, Q2B = NULL, Q3B = NULL, Q6B = NULL, ZB = NULL;
621 413 : GEN v = vecfactorsquarefreeu_coprime(x1, x2, P);
622 413 : long lv = lg(v), j;
623 413 : if (allvecsB)
624 : {
625 21 : VB = gel(allvecsB, 1), QB = gel(allvecsB, 2), Q2B = gel(allvecsB, 3);
626 21 : Q3B = gel(allvecsB, 4), Q6B = gel(allvecsB, 5), ZB = gel(allvecsB, 6);
627 : }
628 806813 : for (j = 1; j < lv; j++)
629 806400 : if (gel(v,j))
630 : {
631 245126 : ulong d = x1 - 1 + j; /* squarefree, coprime to 6 */
632 245126 : GEN t = smallfact(d, gel(v,j), sq, V), u;
633 245126 : GEN tB = NULL, uB = NULL; /* = d^s */
634 : long a, b, c, e, q;
635 245126 : if (!t || gequal0(t)) continue;
636 48748 : if (VB) tB = vecinv(gmulsg(d, gconj(t)));
637 : /* warning: gives 1/conj(f(d)) d^(-1-conj(s)), equal to
638 : f(d) d^(-1-conj(s)) only if |f(d)|=1. */
639 : /* S += f(d) * d^s * Z[q] */
640 48748 : q = N / d;
641 48748 : if (q == 1)
642 : {
643 17339 : S = gadd(S, t); if (VB) SB = gadd(SB, tB);
644 17339 : continue;
645 : }
646 31409 : if (q <= sq) { u = gel(Z, q); if (VB) uB = gel(ZB, q); }
647 : else
648 : {
649 1274 : a = usqrt(q); b = usqrt(q / 2); c = usqrt(q / 3); e = usqrt(q / 6);
650 1274 : u = gadd(gadd(gel(Q,a), gel(Q2,b)), gadd(gel(Q3,c), gel(Q6,e)));
651 1274 : if (VB)
652 161 : uB = gadd(gadd(gel(QB,a), gel(Q2B,b)), gadd(gel(Q3B,c), gel(Q6B,e)));
653 : }
654 31409 : S = gadd(S, vecmul(t, u)); if (VB) SB = gadd(SB, vecmul(tB, uB));
655 : }
656 413 : return VB ? mkvec2(S, SB) : mkvec(S);
657 : }
658 :
659 : static GEN
660 63 : dirpowsummakez(GEN V, GEN W, GEN VB, GEN WB, GEN onef, ulong sq)
661 : {
662 63 : GEN Z = cgetg(sq+1, t_VEC), ZB = NULL;
663 : ulong a, b, c, e, q;
664 : /* a,b,c,e = sqrt(q), sqrt(q/2), sqrt(q/3), sqrt(q/6)
665 : * Z[q] = Q[a] + 2^s Q[b] + 3^s Q[c] + 6^s Q[e], with Q[0] = 0 */
666 63 : gel(Z, 1) = onef;
667 63 : gel(Z, 2) = gel(W, 2);
668 63 : gel(Z, 3) = gel(W, 3);
669 63 : gel(Z, 4) = gel(Z, 5) = gel(W, 4);
670 63 : gel(Z, 6) = gel(Z, 7) = gadd(gel(W, 4), gel(V, 6));
671 63 : if (VB)
672 : {
673 21 : ZB = cgetg(sq+1, t_VEC);
674 21 : gel(ZB, 1) = onef;
675 21 : gel(ZB, 2) = gel(WB, 2);
676 21 : gel(ZB, 3) = gel(WB, 3);
677 21 : gel(ZB, 4) = gel(ZB, 5) = gel(WB, 4);
678 21 : gel(ZB, 6) = gel(ZB, 7) = gadd(gel(WB, 4), gel(VB, 6));
679 : }
680 63 : a = 2; b = c = e = 1;
681 3759 : for (q = 8; q <= sq; q++)
682 : { /* Gray code: at most one of a,b,c,d differs (by 1) from previous value */
683 3696 : GEN z = gel(Z, q - 1), zB = NULL;
684 : ulong na, nb, nc, ne, na2, nb2, nc2, ne2;
685 3696 : if (VB) zB = gel(ZB, q - 1);
686 3696 : if ((na = usqrt(q)) != a)
687 280 : { a = na; na2 = na * na; z = gadd(z, gel(V, na2));
688 280 : if (VB) zB = gadd(zB, gel(VB, na2)); }
689 3416 : else if ((nb = usqrt(q / 2)) != b)
690 203 : { b = nb; nb2 = 2 * nb * nb; z = gadd(z, gel(V, nb2));
691 203 : if (VB) zB = gadd(zB, gel(VB, nb2)); }
692 3213 : else if ((nc = usqrt(q / 3)) != c)
693 161 : { c = nc; nc2 = 3 * nc * nc; z = gadd(z, gel(V, nc2));
694 161 : if (VB) zB = gadd(zB, gel(VB, nc2)); }
695 3052 : else if ((ne = usqrt(q / 6)) != e)
696 98 : { e = ne; ne2 = 6 * ne * ne; z = gadd(z, gel(V, ne2));
697 98 : if (VB) zB = gadd(zB, gel(VB, ne2)); }
698 3696 : gel(Z, q) = z; if (VB) gel(ZB, q) = zB;
699 : }
700 63 : return VB ? mkvec2(Z, ZB) : mkvec(Z);
701 : }
702 :
703 : /* both =
704 : * 0: sum_{n<=N}f(n)n^s
705 : * 1: sum for (f,s) and (conj(f),-1-s)
706 : * 2: sum for (f,s) and (f,-1-s), assuming |f(n)| in {0,1} */
707 : GEN
708 203 : dirpowerssumfun(ulong N, GEN s, void *E, GEN (*f)(void *, ulong, long),
709 : long both, long prec)
710 : {
711 203 : const ulong step = 2048;
712 203 : pari_sp av = avma, av2;
713 : GEN P, V, W, Q, Q2, Q3, Q6, S, Z, onef, zervec;
714 203 : GEN VB = NULL, WB = NULL, QB = NULL;
715 203 : GEN Q2B = NULL, Q3B = NULL, Q6B = NULL, SB = NULL, ZB = NULL;
716 203 : GEN R, RS, data, allvecs, allvecsB = NULL;
717 : ulong x1, sq;
718 : long prec0, prec1, needlog;
719 :
720 203 : if ((f && N < 49) || (!f && N < 1000))
721 : {
722 140 : if (!N)
723 : {
724 14 : if (!f) return gen_0;
725 0 : return gerepileupto(av, gmul(gen_0, f(E, 1, prec)));
726 : }
727 126 : return gerepilecopy(av, smalldirpowerssum(N, s, E, f, both, prec));
728 : }
729 63 : onef = f ? f(E, 1, prec) : gen_1;
730 63 : zervec = gmul(gen_0, onef);
731 63 : sq = usqrt(N);
732 63 : prec1 = prec0 = prec + EXTRAPREC64;
733 63 : s = gprec_w(s, prec0);
734 63 : needlog = get_needlog(s);
735 63 : if (needlog == 1) prec1 = powcx_prec(log2((double)N), s, prec);
736 63 : data = mkvec3(onef, zervec, mkvecsmall4(sq, needlog, prec0, prec1));
737 63 : RS = dirpowsuminit(s, E, f, data, both, prec);
738 63 : R = gel(RS, 1); V = gel(R, 1); W = gel(R, 2); Q = gel(R, 3);
739 63 : Q2 = gel(R, 4); Q3 = gel(R, 5); Q6 = gel(R, 6);
740 63 : if (lg(RS) > 2)
741 : {
742 21 : GEN RB = gel(RS, 2);
743 21 : VB = gel(RB, 1); WB = gel(RB, 2); QB = gel(RB, 3);
744 21 : Q2B = gel(RB, 4); Q3B = gel(RB, 5); Q6B = gel(RB, 6);
745 : }
746 63 : RS = dirpowsumprimeloop(N, s, E, f, data, W, WB);
747 63 : S = gel(RS, 1); if (VB) SB = gel(RS, 2);
748 63 : RS = dirpowsummakez(V, W, VB, WB, onef, sq);
749 63 : Z = gel(RS, 1); if (VB) ZB = gel(RS, 2);
750 63 : P = mkvecsmall2(2, 3);
751 63 : allvecs = mkvecn(6, V, Q, Q2, Q3, Q6, Z);
752 63 : if (VB) allvecsB = mkvecn(6, VB, QB, Q2B, Q3B, Q6B, ZB);
753 63 : av2 = avma;
754 63 : for(x1 = 1;; x1 += step)
755 350 : { /* beware overflow, fuse last two bins (avoid a tiny remainder) */
756 413 : ulong x2 = (N >= 2*step && N - 2*step >= x1)? x1-1 + step: N;
757 413 : RS = dirpowsumsqfloop(N, x1, x2, sq, P, allvecs, S, allvecsB, SB);
758 413 : S = gel(RS, 1); if (VB) SB = gel(RS, 2);
759 413 : if (x2 == N) break;
760 350 : gerepileall(av2, SB? 2: 1, &S, &SB);
761 : }
762 63 : if (both == 0) return gerepileupto(av, S);
763 42 : return gerepilecopy(av, mkvec2(S, gconj(VB? SB: S)));
764 : }
765 :
766 : GEN
767 133 : dirpowerssum(ulong N, GEN s, long both, long prec)
768 133 : { return dirpowerssumfun(N, s, NULL, NULL, both, prec); }
769 :
770 : static GEN
771 13818 : gp_callUp(void *E, ulong x, long prec)
772 : {
773 13818 : long court[] = {evaltyp(t_INT)|_evallg(3), evalsigne(1)|evallgefint(3),0};
774 13818 : court[2] = x; return gp_callprec(E, court, prec);
775 : }
776 :
777 : GEN
778 154 : dirpowerssum0(GEN N, GEN s, GEN f, long both, long prec)
779 : {
780 154 : if (typ(N) != t_INT) pari_err_TYPE("dirpowerssum", N);
781 147 : if (signe(N) <= 0) N = gen_0;
782 147 : if (!f) return dirpowerssum(itou(N), s, both, prec);
783 70 : if (typ(f) != t_CLOSURE) pari_err_TYPE("dirpowerssum", f);
784 70 : return dirpowerssumfun(itou(N), s, (void*)f, gp_callUp, both, prec);
785 : }
786 :
787 : /*******************************************************************/
788 : /* Parallel dirpowerssumfun */
789 : /*******************************************************************/
790 : /* f is a totally multiplicative function of modulus 0 or 1
791 : * (essentially a Dirichlet character). Compute simultaneously
792 : * sum_{0 < n <= N} f(n)n^s and sum_{0 < n <= N} f(n)n^{-1-conj(s)}
793 : * Warning: s is conjugated, but not f. Main application for Riemann-Siegel,
794 : * where we need R(chi,s) and conj(R(chi,1-conj(s))). */
795 :
796 : static GEN
797 292513 : mycallvec(GEN f, ulong n, long prec)
798 : {
799 292513 : if (typ(f) == t_CLOSURE) return gp_callUp((void*)f, n, prec);
800 292513 : return gel(f, (n-1) % (lg(f)-1) + 1);
801 : }
802 :
803 : static GEN
804 70 : _smalldirpowerssum(long N, GEN s, long fl, long prec)
805 : {
806 70 : GEN vg = vecpowug(N, s, prec), S1 = _Qtor(RgV_sum(vg), prec);
807 : long n;
808 70 : if (!fl) return mkvec2(S1, S1);
809 616 : for (n = 1; n <= N; n++) gel(vg, n) = ginv(gmulsg(n, gconj(gel(vg, n))));
810 28 : return mkvec2(S1, _Qtor(RgV_sum(vg), prec));
811 : }
812 :
813 : static GEN
814 2226 : gmulvecsqlv(GEN Q, GEN V)
815 : {
816 : long lq, i;
817 : GEN W;
818 2226 : if (typ(V) != t_VEC) return RgV_Rg_mul(Q, V);
819 1128 : lq = lg(Q); W = cgetg(lq, t_VEC);
820 51240 : for (i = 1; i < lq; i++) gel(W, i) = vecmul(gel(Q, i), V);
821 1128 : return W;
822 : }
823 :
824 : GEN
825 17157 : parsqfboth_worker(GEN gk, GEN vZ, GEN vVQ, GEN vV, GEN P, GEN Nsqstep)
826 : {
827 17157 : pari_sp av2 = avma;
828 17157 : GEN Z1 = gel(vZ, 1), Z2 = gel(vZ, 2), V1 = gel(vV, 1), V2 = gel(vV, 2);
829 17157 : GEN VQ1 = gel(vVQ, 1), VQN, v, S1, S2;
830 17157 : GEN Q1 = gel(VQ1, 1), Q2 = gel(VQ1, 2), Q3 = gel(VQ1, 3), Q6 = gel(VQ1, 4);
831 17157 : GEN QN = gen_0, Q2N = gen_0, Q3N = gen_0, Q6N = gen_0;
832 17157 : long k = itos(gk), N = Nsqstep[1], step = Nsqstep[3];
833 17157 : long x1 = 1 + step * k, x2, j, lv, fl = !gcmp0(V2), lvv = 0;
834 17157 : ulong a, b, c, e, q, sq = Nsqstep[2];
835 17157 : if (typ(gel(V1, 1)) == t_VEC)
836 1332 : { lvv = lg(gel(V1, 1)) - 1; S1 = zerovec(lvv); S2 = zerovec(lvv); }
837 15825 : else { S1 = gen_0; S2 = gen_0; }
838 17157 : if (fl)
839 0 : { VQN = gel(vVQ, 2); QN = gel(VQN, 1); Q2N = gel(VQN, 2);
840 0 : Q3N = gel(VQN, 3); Q6N = gel(VQN, 4); }
841 : /* beware overflow, fuse last two bins (avoid a tiny remainder) */
842 17157 : x2 = (N >= 2*step && N - 2*step >= x1)? x1 - 1 + step: N;
843 17157 : v = vecfactorsquarefreeu_coprime(x1, x2, P);
844 17158 : lv = lg(v);
845 34401072 : for (j = 1; j < lv; j++)
846 34384972 : if (gel(v,j))
847 : {
848 10448994 : ulong d = x1 - 1 + j; /* squarefree, coprime to 6 */
849 10448994 : GEN t1 = smallfact(d, gel(v,j), sq, V1), t2 = gen_0, u1, u2 = gen_0;
850 : /* = f(d) d^s */
851 10406888 : if (!t1 || gequal0(t1)) continue;
852 2141174 : if (fl) t2 = vecinv(gmulsg(d, gconj(t1)));
853 : /* warning: gives 1/conj(f(d)) d^(-1-conj(s)), equal to
854 : f(d) d^(-1-conj(s)) only if |f(d)|=1. */
855 : /* S += f(d) d^s * Z[q] */
856 2226087 : q = N / d;
857 2226087 : if (q == 1)
858 : {
859 846159 : S1 = gadd(S1, t1); if (fl) S2 = gadd(S2, t2);
860 848098 : continue;
861 : }
862 1379928 : if (q <= sq) { u1 = gel(Z1, q); if (fl) u2 = gel(Z2, q); }
863 : else
864 : {
865 32458 : a = usqrt(q); b = usqrt(q / 2); c = usqrt(q / 3); e = usqrt(q / 6);
866 32538 : u1 = gadd(gadd(gel(Q1,a), gel(Q2,b)), gadd(gel(Q3,c), gel(Q6,e)));
867 32538 : if (fl)
868 0 : u2 = gadd(gadd(gel(QN,a), gel(Q2N,b)), gadd(gel(Q3N,c), gel(Q6N,e)));
869 : }
870 1380008 : S1 = gadd(S1, vecmul(t1, u1)); if (fl) S2 = gadd(S2, vecmul(t2, u2));
871 : }
872 16100 : return gerepilecopy(av2, mkvec2(S1, S2));
873 : }
874 :
875 : GEN
876 37918 : parsumprimeWfunboth_worker(GEN gk, GEN s, GEN W1, GEN W2, GEN f, GEN Nsqprec)
877 : {
878 : pari_sp av2;
879 37918 : GEN S1, S2 = gen_0, logp;
880 : forprime_t T;
881 37918 : long k = itou(gk), N = Nsqprec[1], sq = Nsqprec[2], precp;
882 37911 : long STEP = Nsqprec[3], prec0 = Nsqprec[4], prec1 = Nsqprec[5], p;
883 37911 : long gp[] = {evaltyp(t_INT)|_evallg(3), evalsigne(1)|evallgefint(3),0};
884 37911 : long needlog = get_needlog(s), fl = !gequal0(W2), lv;
885 :
886 37905 : if (f && gequal0(f)) f = NULL;
887 37895 : if (!f) lv = 0;
888 : else
889 : {
890 15958 : GEN tmp = mycallvec(f, 1, prec1);
891 15958 : lv = typ(tmp) == t_VEC ? lg(tmp) - 1 : 0;
892 : }
893 37895 : precp = 0; logp = NULL;
894 37895 : if (lv) { S1 = const_vec(lv, real_0(prec1)); if (fl) S2 = const_vec(lv, real_0(prec1)); }
895 21938 : else { S1 = real_0(prec1); if (fl) S2 = real_0(prec1); }
896 37880 : u_forprime_init(&T, k * STEP + sq + 1, minss(N, (k + 1) * STEP + sq));
897 37942 : av2 = avma;
898 4170028 : while ((p = u_forprime_next(&T)))
899 : {
900 4134137 : GEN u = gen_0, ks = f ? mycallvec(f, p, prec1) : gen_1;
901 4134131 : long zks = !gequal0(ks);
902 4133990 : gp[2] = p;
903 4133990 : if (needlog)
904 : {
905 4134108 : if (!logp)
906 30123 : logp = logr_abs(utor(p, prec1));
907 : else
908 : { /* log p = log(precp) + 2 atanh((p - precp) / (p + precp)) */
909 4103985 : ulong a = p >> 1, b = precp >> 1; /* p = 2a + 1, precp = 2b + 1 */
910 4103985 : GEN z = atanhuu(a - b, a + b + 1, prec1); /* avoid overflow */
911 4101496 : shiftr_inplace(z, 1); logp = addrr(logp, z);
912 : }
913 4132345 : if (zks)
914 4132187 : u = needlog == 1? powcx(gp, logp, s, prec0) : mpexp(gmul(s, logp));
915 : }
916 0 : else { if (zks) u = gpow(gp, s, prec0); }
917 4134353 : if (zks)
918 : {
919 4134308 : S1 = gadd(S1, vecmul(gel(W1, N / p), gmul(ks, u)));
920 4132041 : if (fl)
921 0 : S2 = gadd(S2, gdiv(vecmul(ks, gel(W2, N / p)), gmulsg(p, gconj(u))));
922 : }
923 4132086 : precp = p;
924 4132086 : if ((p & 0x1ff) == 1)
925 : {
926 15508 : if (!logp) gerepileall(av2, 2, &S1, &S2);
927 15508 : else gerepileall(av2, 3, &S1, &S2, &logp);
928 : }
929 : }
930 37731 : return gcopy(mkvec2(S1, S2));
931 : }
932 :
933 : static GEN
934 14 : smalldirpowerssumfunvec_i(GEN f, ulong N, GEN s, long prec)
935 : {
936 14 : GEN S = real_0(prec);
937 : ulong n;
938 14 : if (f && gequal0(f)) f = NULL;
939 14 : if (f)
940 : {
941 14 : GEN tmp = mycallvec(f, 1, prec);
942 14 : if (typ(tmp) == t_VEC) S = const_vec(lg(tmp) - 1, real_0(prec));
943 : }
944 28 : for (n = 1; n <= N; n++)
945 : {
946 14 : GEN ks = mycallvec(f, n, prec);
947 14 : if (!gequal0(ks)) S = gadd(S, gmul(ks, gpow(utoipos(n), s, prec)));
948 : }
949 14 : return S;
950 : }
951 :
952 : static GEN
953 14 : smalldirpowerssumfunvec(GEN f, ulong N, GEN s, long fl, long prec)
954 : {
955 14 : GEN tmp = smalldirpowerssumfunvec_i(f, N, s, prec);
956 14 : if (fl) return mkvec2(tmp, smalldirpowerssumfunvec_i(f, N, gsubgs(gneg(gconj(s)), 1), prec));
957 14 : else return mkvec2(tmp, tmp);
958 : }
959 :
960 : static GEN
961 1392 : getscal(GEN V, long isscal)
962 : {
963 : long lv;
964 1392 : if (isscal != 1 || typ(V) != t_VEC) return V;
965 336 : lv = lg(V);
966 336 : switch (lv)
967 : {
968 0 : case 2: return gel(V, 1);
969 336 : case 3: return mkvec2(getscal(gel(V, 1), 1), getscal(gel(V, 2), 1));
970 0 : default: pari_err_TYPE("getscal", V);
971 : }
972 : return NULL; /*LCOV_EXCL_LINE*/
973 : }
974 :
975 : GEN
976 854 : pardirpowerssumfun(GEN f, ulong N, GEN s, long both, long prec)
977 : {
978 854 : const ulong step = 2048;
979 854 : pari_sp av = avma;
980 854 : GEN P, V1, W1, Q1, V2 = gen_0, W2 = gen_0, QN = gen_0, c2, c2N;
981 854 : GEN Q2, Q2N = gen_0, Q3, Q3N = gen_0, Q6, Q6N = gen_0;
982 854 : GEN S1, S2, RES, Z1, Z2 = gen_0, logp;
983 854 : long gp[] = {evaltyp(t_INT)|_evallg(3), evalsigne(1)|evallgefint(3),0};
984 : ulong a, b, c, e, q, n, sq, fl;
985 854 : long prec0, prec1, needlog, lv = 0, isscal = 1;
986 854 : GEN unvec = gen_1, zervec = gen_0, re0, re1, tmp2 = NULL;
987 :
988 854 : if (f)
989 : {
990 390 : if (typ(f) == t_CLOSURE) isscal = -1; else isscal = 0;
991 390 : tmp2 = mycallvec(f, 2, prec + EXTRAPRECWORD);
992 390 : if (typ(tmp2) == t_VEC)
993 : {
994 390 : lv = lg(tmp2) - 1;
995 390 : unvec = const_vec(lv, gen_1);
996 390 : zervec = const_vec(lv, gen_0);
997 : }
998 : }
999 854 : if (N <= 0)
1000 28 : { GEN v = mkvec2(gen_0, gen_0); return f ? const_vec(lv, v) : v; }
1001 826 : fl = both && !gequal(real_i(s), gneg(ghalf));
1002 826 : if (f && N < 49)
1003 14 : return gerepilecopy(av, getscal(smalldirpowerssumfunvec(f, N, s, fl, prec), isscal));
1004 812 : if (!f && N < 10000UL)
1005 70 : return gerepilecopy(av, _smalldirpowerssum(N, s, fl, prec));
1006 742 : sq = usqrt(N);
1007 742 : V1 = cgetg(sq+1, t_VEC); W1 = cgetg(sq+1, t_VEC); Q1 = cgetg(sq+1, t_VEC);
1008 742 : if (fl)
1009 : {
1010 0 : V2 = cgetg(sq+1, t_VEC); W2 = cgetg(sq+1, t_VEC); QN = cgetg(sq+1, t_VEC);
1011 : }
1012 742 : prec1 = prec0 = prec + EXTRAPRECWORD;
1013 742 : s = gprec_w(s, prec0);
1014 742 : needlog = get_needlog(s);
1015 742 : if (needlog == 1) prec1 = powcx_prec(log2((double)N), s, prec);
1016 742 : gel(V1,1) = gel(W1,1) = gel(Q1,1) = unvec;
1017 742 : if (fl) { gel(V2,1) = gel(W2,1) = gel(QN,1) = unvec; }
1018 742 : c2 = gpow(gen_2, s, prec0); c2N = ginv(gmul2n(gconj(c2), 1));
1019 742 : re0 = real_0(prec0); re1 = real_1(prec0);
1020 742 : if (f) { c2 = gmul(c2, tmp2); c2N = gmul(c2N, tmp2); }
1021 742 : gel(V1,2) = c2; /* f(2) 2^s */
1022 742 : gel(W1,2) = gmul(re1, gadd(c2, unvec));
1023 742 : gel(Q1,2) = gmul(re1, gadd(vecsqr(c2), unvec));
1024 742 : if (fl)
1025 : {
1026 0 : gel(V2,2) = c2N; gel(W2,2) = gmul(re1, gadd(c2N, unvec));
1027 0 : gel(QN,2) = gmul(re1, gadd(vecsqr(c2N), unvec));
1028 : }
1029 742 : logp = NULL;
1030 107430 : for (n = 3; n <= sq; n++)
1031 : {
1032 106688 : GEN u1 = zervec, u2 = zervec, ks = f ? mycallvec(f, n, prec) : gen_1;
1033 106688 : long zks = !gequal0(ks);
1034 106688 : if (odd(n))
1035 : {
1036 53692 : gp[2] = n;
1037 53692 : if (needlog)
1038 : {
1039 53692 : if (!logp)
1040 742 : logp = logr_abs(utor(n, prec1));
1041 : else
1042 : { /* log n = log(n-2) + 2 atanh(1 / (n - 1)) */
1043 52950 : GEN z = atanhuu(1, n - 1, prec1);
1044 52950 : shiftr_inplace(z, 1); logp = addrr(logp, z);
1045 : }
1046 53692 : if (zks)
1047 52276 : u1 = needlog == 1? powcx(gp, logp, s, prec0) : mpexp(gmul(s, logp));
1048 : }
1049 0 : else { if (zks) u1 = gpow(gp, s, prec0); }
1050 53692 : if (zks)
1051 : {
1052 52276 : if(fl) u2 = gmul(ks, ginv(gmulsg(n, gconj(u1))));
1053 52276 : u1 = gmul(ks, u1); /* f(n) n^s */
1054 : }
1055 : }
1056 : else
1057 : {
1058 52996 : u1 = vecmul(c2, gel(V1, n >> 1));
1059 52996 : if (fl) u2 = vecmul(c2N, gel(V2, n >> 1));
1060 : }
1061 106688 : gel(V1,n) = u1; /* f(n) n^s */
1062 106688 : gel(W1,n) = gadd(gel(W1,n-1), gel(V1,n)); /* = sum_{i<=n} f(i)i^s */
1063 106688 : gel(Q1,n) = gadd(gel(Q1,n-1), vecsqr(gel(V1,n))); /* = sum_{i<=n} f(i^2)i^2s */
1064 106688 : if (fl)
1065 : {
1066 0 : gel(V2,n) = u2;
1067 0 : gel(W2,n) = gadd(gel(W2,n-1), gel(V2,n));
1068 0 : gel(QN,n) = gadd(gel(QN,n-1), vecsqr(gel(V2,n)));
1069 : }
1070 : }
1071 742 : Q2 = gmulvecsqlv(Q1, gel(V1,2));
1072 742 : Q3 = gmulvecsqlv(Q1, gel(V1,3));
1073 742 : Q6 = gmulvecsqlv(Q1, gel(V1,6));
1074 742 : if (fl)
1075 : {
1076 0 : Q2N = gmulvecsqlv(QN, gel(V2,2));
1077 0 : Q3N = gmulvecsqlv(QN, gel(V2,3));
1078 0 : Q6N = gmulvecsqlv(QN, gel(V2,6));
1079 : }
1080 742 : if (typ(zervec) == t_VEC)
1081 376 : { S1 = const_vec(lv, re0); S2 = const_vec(lv, re0); }
1082 366 : else { S1 = re0; S2 = re0; }
1083 742 : RES = mkvec2(S1, S2);
1084 : {
1085 742 : GEN fspec = f ? f : gen_0;
1086 742 : long m = mt_nbthreads();
1087 742 : long STEP = maxss(N / (m * m), 1);
1088 742 : GEN VS = mkvecsmalln(5, N, sq, STEP, prec0, prec1);
1089 742 : GEN FUN = strtoclosure("_parsumprimeWfunboth_worker", 5, s, W1, W2, fspec, VS);
1090 742 : RES = gadd(RES, parsum(gen_0, utoipos((N - 1) / STEP), FUN));
1091 : }
1092 742 : P = mkvecsmall2(2, 3);
1093 742 : Z1 = cgetg(sq+1, t_VEC);
1094 : /* a,b,c,e = sqrt(q), sqrt(q/2), sqrt(q/3), sqrt(q/6)
1095 : * Z[q] = Q[a] + 2^s Q[b] + 3^s Q[c] + 6^s Q[e], with Q[0] = 0 */
1096 742 : gel(Z1, 1) = unvec;
1097 742 : gel(Z1, 2) = gel(W1, 2);
1098 742 : gel(Z1, 3) = gel(W1, 3);
1099 742 : gel(Z1, 4) = gel(Z1, 5) = gel(W1, 4);
1100 742 : gel(Z1, 6) = gel(Z1, 7) = gadd(gel(W1, 4), gel(V1, 6));
1101 742 : if (fl)
1102 : {
1103 0 : Z2 = cgetg(sq+1, t_VEC);
1104 0 : gel(Z2, 1) = unvec;
1105 0 : gel(Z2, 2) = gel(W2, 2);
1106 0 : gel(Z2, 3) = gel(W2, 3);
1107 0 : gel(Z2, 4) = gel(Z2, 5) = gel(W2, 4);
1108 0 : gel(Z2, 6) = gel(Z2, 7) = gadd(gel(W2, 4), gel(V2, 6));
1109 : }
1110 742 : a = 2; b = c = e = 1;
1111 103720 : for (q = 8; q <= sq; q++)
1112 : { /* Gray code: at most one of a,b,c,d differs (by 1) from previous value */
1113 102978 : GEN z1 = gel(Z1, q - 1), z2 = gen_0;
1114 : ulong na, nb, nc, ne, na2, nb2, nc2, ne2;
1115 102978 : if (fl) z2 = gel(Z2, q - 1);
1116 102978 : if ((na = usqrt(q)) != a)
1117 6097 : { a = na; na2 = na * na; z1 = gadd(z1, gel(V1, na2)); if (fl) z2 = gadd(z2, gel(V2, na2)); }
1118 96881 : else if ((nb = usqrt(q / 2)) != b)
1119 4118 : { b = nb; nb2 = 2 * nb * nb; z1 = gadd(z1, gel(V1, nb2)); if (fl) z2 = gadd(z2, gel(V2, nb2)); }
1120 92763 : else if ((nc = usqrt(q / 3)) != c)
1121 3616 : { c = nc; nc2 = 3 * nc * nc; z1 = gadd(z1, gel(V1, nc2)); if (fl) z2 = gadd(z2, gel(V2, nc2)); }
1122 89147 : else if ((ne = usqrt(q / 6)) != e)
1123 2068 : { e = ne; ne2 = 6 * ne * ne; z1 = gadd(z1, gel(V1, ne2)); if (fl) z2 = gadd(z2, gel(V2, ne2)); }
1124 102978 : gel(Z1,q) = z1; if (fl) gel(Z2,q) = z2;
1125 : }
1126 : {
1127 742 : GEN vQ1 = mkvec4(Q1, Q2, Q3, Q6), vQ2 = gen_0, FUN;
1128 742 : if (fl) vQ2 = mkvec4(QN, Q2N, Q3N, Q6N);
1129 742 : FUN = strtoclosure("_parsqfboth_worker", 5, mkvec2(Z1, Z2), mkvec2(vQ1, vQ2), mkvec2(V1, V2), P, mkvecsmall3(N, sq, step));
1130 742 : RES = gadd(RES, parsum(gen_0, utoipos(maxss(0, (N - 1) / step - 1)), FUN));
1131 : }
1132 706 : if (!fl) gel(RES, 2) = gel(RES, 1);
1133 706 : return gerepilecopy(av, getscal(RES, isscal));
1134 : }
1135 :
1136 : GEN
1137 0 : pardirpowerssum0(GEN N, GEN s, GEN f, long both, long prec)
1138 : {
1139 0 : pari_sp av = avma;
1140 : GEN R;
1141 0 : if (typ(N) != t_INT) pari_err_TYPE("pardirpowerssum", N);
1142 0 : R = pardirpowerssumfun(f, itou(N), s, both, prec);
1143 0 : return both ? R : gerepilecopy(av, gel(R, 1));
1144 : }
1145 :
1146 : GEN
1147 0 : pardirpowerssum(ulong N, GEN s, long prec)
1148 : {
1149 0 : pari_sp av = avma;
1150 0 : return gerepilecopy(av, gel(pardirpowerssumfun(NULL, N, s, 0, prec), 1));
1151 : }
|
