Skip to content

Commit 52449f3

Browse files
authored
Add apmath functions power, hypergeometric, and hypergeometric_minus_one (#107)
1 parent d876e8d commit 52449f3

File tree

2 files changed

+250
-0
lines changed

2 files changed

+250
-0
lines changed

functional_algorithms/apmath.py

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,3 +348,140 @@ def rsqrt(ctx, seq, functional=False, size=None, niter=None):
348348
def sqrt(ctx, seq, functional=False, size=None):
349349
"""Square root of an FP expansion."""
350350
return multiply(ctx, seq, rsqrt(ctx, seq, functional=functional, size=size), functional=functional, size=size)
351+
352+
353+
def power(ctx, seq, n, functional=False, size=None):
354+
"""
355+
n-th power of an FP expansion.
356+
"""
357+
if n == 0:
358+
return [ctx.constant(1, seq[0])]
359+
if n == 1:
360+
assert size is None or len(seq) == size # not impl
361+
return seq
362+
if n == 2:
363+
return square(ctx, seq, functional=functional, size=size)
364+
assert isinstance(n, int) and n > 0, n # not impl
365+
r = power(ctx, seq, n // 2, functional=functional, size=size)
366+
sq = square(ctx, r, functional=functional, size=size)
367+
if n % 2 == 0:
368+
return sq
369+
return multiply(ctx, sq, seq, functional=functional, size=size)
370+
371+
372+
def hypergeometric(ctx, a, b, seq, niter, functional=False, size=None):
373+
"""
374+
Generalized hypergeometic series on an FP expansion:
375+
376+
sum(prod((a[i])_n, i=0..p) / prod((b[i])_n), i=0..q * z**n / n!, n=0,1,...,niter-1)
377+
378+
where
379+
p = len(a) - 1
380+
p = len(b) - 1
381+
(k)_n = 1 if n == 0 else (k)_{n-1} * (k + n - 1)
382+
n! = 1 if n == 0 else (n-1)! * n
383+
a and b are lists of integers or Fraction instances.
384+
"""
385+
import numpy
386+
387+
if isinstance(seq[0], (numpy.float64, numpy.float32, numpy.float16)):
388+
return hypergeometric_impl(ctx, type(seq[0]), a, b, seq, niter, functional=functional, size=size)
389+
390+
largest = fpa.get_largest(ctx, seq[0])
391+
r_fp64 = hypergeometric_impl(ctx, numpy.float64, a, b, seq, niter, functional=functional, size=size)
392+
r_fp32 = hypergeometric_impl(ctx, numpy.float32, a, b, seq, niter, functional=functional, size=size)
393+
r_fp16 = hypergeometric_impl(ctx, numpy.float16, a, b, seq, niter, functional=functional, size=size)
394+
395+
return ctx.select(largest > 1e308, r_fp64, ctx.select(largest > 1e38, r_fp32, r_fp16))
396+
397+
398+
def hypergeometric_minus_one(ctx, a, b, seq, niter, functional=False, size=None):
399+
"""
400+
Generalized hypergeometic series on an FP expansion minus one:
401+
402+
sum(prod((a[i])_n, i=0..p) / prod((b[i])_n), i=0..q * z**n / n!, n=1,...,niter-1)
403+
404+
where
405+
p = len(a) - 1
406+
p = len(b) - 1
407+
(k)_n = 1 if n == 0 else (k)_{n-1} * (k + n - 1)
408+
n! = 1 if n == 0 else (n-1)! * n
409+
a and b are lists of integers or Fraction instances.
410+
"""
411+
import numpy
412+
413+
if isinstance(seq[0], (numpy.float64, numpy.float32, numpy.float16)):
414+
return hypergeometric_minus_one_impl(ctx, type(seq[0]), a, b, seq, niter, functional=functional, size=size)
415+
416+
largest = fpa.get_largest(ctx, seq[0])
417+
r_fp64 = hypergeometric_minus_one_impl(ctx, numpy.float64, a, b, seq, niter, functional=functional, size=size)
418+
r_fp32 = hypergeometric_minus_one_impl(ctx, numpy.float32, a, b, seq, niter, functional=functional, size=size)
419+
r_fp16 = hypergeometric_minus_one_impl(ctx, numpy.float16, a, b, seq, niter, functional=functional, size=size)
420+
421+
return ctx.select(largest > 1e308, r_fp64, ctx.select(largest > 1e38, r_fp32, r_fp16))
422+
423+
424+
def hypergeometric_impl(ctx, dtype, a, b, seq, niter, functional=False, size=None):
425+
r = hypergeometric_minus_one_impl(ctx, dtype, a, b, seq, niter, functional=functional, size=size)
426+
return add(ctx, [dtype(1)], r, functional=functional, size=size)
427+
428+
429+
def hypergeometric_minus_one_impl(ctx, dtype, a, b, seq, niter, functional=False, size=None):
430+
import fractions
431+
import functional_algorithms as fa
432+
433+
rcoeffs = []
434+
for n in range(1, niter):
435+
numer_ = 1
436+
denom_ = 1
437+
for a_ in a:
438+
numer_ *= a_ + n - 1
439+
for b_ in b:
440+
denom_ *= b_ + n - 1
441+
denom_ *= n
442+
443+
if not numer_:
444+
break
445+
446+
rc = fa.utils.fraction2expansion(dtype, fractions.Fraction(numer_, denom_))
447+
if not rc:
448+
break
449+
rcoeffs.append(renormalize(ctx, rc, functional=False))
450+
451+
# hypergeometric series evaluation using Horner' scheme as
452+
#
453+
# sum(c[n] * z ** n, n=0..niter-1)
454+
# = c[0] + c[1] * z + c[2] * z ** 2 + c[3] * z ** 3 + ...
455+
# = c[0] + (c[1] + (c[2] + (c[3] + ...) * z) * z) * z
456+
# = fma(fma(fma(fma(..., z, c[3]), z, c[2]), z, c[1]), z, c[0])
457+
#
458+
# is inaccurate because c[n] is a rapidly decreasing sequence and
459+
# the given dtype may not have enough exponent range to represent
460+
# very small coefficients.
461+
#
462+
# In the following, we'll use the property of geometric series that
463+
# the ratio of neighboring coefficients is a rational number so
464+
# that we have
465+
#
466+
# c[n] * z ** n == (c[n-1] * z ** (n-1)) * z * R(n)
467+
# c[0] = 1
468+
#
469+
# Hence
470+
#
471+
# sum(c[n] * z ** n, n=0..niter-1)
472+
# = c[0] + c[0] * z * R(1) + c[0] * z * R(1) * z * R(2) + ...
473+
# = c[0] * (1 + z * R(1) * (1 + z * R(2) * (1 + z * R(3) * (1 + ...))))
474+
# = 1 + z * R(1) * (1 + z * R(2) * (1 + z * R(3) * (1 + ...)))
475+
#
476+
# where R(n) is a slowly varying sequence in n. For instance, for
477+
# float16 hypergeometric([], [], z), R(n) = 1 / n is nonzero for n
478+
# over 2 ** 24, that is, the maximal value for user-specified
479+
# niter is practically unlimited.
480+
def rhorner(rcoeffs, z):
481+
r = multiply(ctx, rcoeffs[0], z, functional=functional, size=size) or [dtype(0)]
482+
if len(rcoeffs) > 1:
483+
h = add(ctx, [dtype(1)], rhorner(rcoeffs[1:], z), functional=functional, size=size)
484+
r = multiply(ctx, r, h, functional=functional, size=size)
485+
return r
486+
487+
return rhorner(rcoeffs, seq)

functional_algorithms/tests/test_apmath.py

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -438,3 +438,116 @@ def test_add(dtype, functional, overlapping):
438438
assert prec > (-fi.negep) * length
439439

440440
fa.utils.show_prec(precs)
441+
442+
443+
@pytest.mark.parametrize(
444+
"function,functional",
445+
[
446+
("exp", "non-functional"),
447+
("exp", "functional"),
448+
("cos", "non-functional"),
449+
("j0", "non-functional"),
450+
("j1", "non-functional"),
451+
],
452+
)
453+
def test_hypergeometric(dtype, function, functional):
454+
if dtype == numpy.longdouble:
455+
pytest.skip(f"test not implemented")
456+
import mpmath
457+
import fractions
458+
459+
size = 100
460+
functional = {"non-functional": False, "functional": True}[functional]
461+
ctx = fa.utils.NumpyContext()
462+
mp_ctx = mpmath.mp
463+
fi = numpy.finfo(dtype)
464+
min_prec = -fi.negep
465+
max_prec = fi.maxexp - numpy.frexp(fi.smallest_subnormal)[1] + 20
466+
467+
min_value = 1
468+
max_value = 20
469+
niter = 20
470+
length = 2
471+
if function == "exp":
472+
a, b = [], [] # exp(z)
473+
sample_func = lambda z: z
474+
length = 2
475+
min_value = 0
476+
max_value = 9
477+
niter = {
478+
numpy.float16: 22, # max z == 9
479+
numpy.float32: 30, # max z == 9
480+
numpy.float64: 45, # max z == 9
481+
}.get(dtype, niter)
482+
elif function == "cos":
483+
a, b = [], [fractions.Fraction(1, 2)] # cos(-4 * sqrt(z))
484+
sample_func = lambda z: -z
485+
min_value = 0
486+
max_value = 10
487+
length = 2
488+
niter = {
489+
numpy.float16: 11, # max z == 5.516
490+
numpy.float32: 28, # max z == 49.94
491+
numpy.float64: 66, # max z == 326.0
492+
}.get(dtype, niter)
493+
niter = {
494+
numpy.float16: 14, # max z == 15
495+
numpy.float32: 19, # max z == 22
496+
numpy.float64: 25, # max z == 21
497+
}.get(dtype, niter)
498+
niter = {
499+
numpy.float16: 14, # max z == 10
500+
numpy.float32: 15, # max z == 10
501+
numpy.float64: 21, # max z == 11
502+
}.get(dtype, niter)
503+
elif function == "j0":
504+
a, b = [], [1] # J0(-4 * sqrt(z)) * 2 / z
505+
sample_func = lambda z: -z
506+
min_value = 0
507+
max_value = 10
508+
length = 2
509+
niter = {
510+
numpy.float16: 13, # max z == 10
511+
numpy.float32: 15, # max z == 10
512+
numpy.float64: 21, # max z == 11
513+
}.get(dtype, niter)
514+
elif function == "j1":
515+
a, b = [], [2] # J1(-4 * sqrt(z)) * (2 / z) ** 2 * 2
516+
sample_func = lambda z: -z
517+
min_value = 0
518+
max_value = 10
519+
length = 3
520+
niter = {
521+
# numpy.float16: 13, # max z == 3.68 [length==2]
522+
numpy.float16: 11, # max z == 10 [length==3]
523+
numpy.float32: 14, # max z == 10 [length==2 or 3]
524+
numpy.float64: 20, # max z == 11
525+
}.get(dtype, niter)
526+
else:
527+
assert 0, function # not implemented
528+
529+
with mp_ctx.workprec(max_prec):
530+
for z in map(sample_func, fa.utils.real_samples(size=size, dtype=dtype, min_value=min_value, max_value=max_value)):
531+
e = fa.apmath.hypergeometric(ctx, a, b, [z], niter=niter, functional=functional, size=length)
532+
533+
# skip samples that result overflow to infinity
534+
s = sum(e[:-1], e[-1])
535+
if not numpy.isfinite(s):
536+
print(f"{z=} {e=}")
537+
break
538+
539+
result_mp = fa.utils.expansion2mpf(mp_ctx, e)
540+
541+
z_mp = fa.utils.float2mpf(mp_ctx, z)
542+
543+
expected_mp = mp_ctx.hyper(a, b, z_mp)
544+
545+
prec = fa.utils.diff_prec(result_mp, expected_mp)
546+
547+
# print(f'{prec=} {z=} {e=}')
548+
549+
if size <= 1000:
550+
assert prec > min_prec
551+
552+
if prec <= min_prec:
553+
break

0 commit comments

Comments
 (0)