Round trip and pickling for arb and acb

bin/rst_to_pxd.py

@@ -62,6 +62,9 @@
 doc_to_header = {'flint/gr_domains': 'flint/gr'}
 void_parameter_list = re.compile(r"\(\s*void\s*\)")
 macro_declarations = {
+    'flint/mag': {
+        'MAG_BITS': 'cdef const int MAG_BITS',
+    },
     'flint/arf': {
         'ARF_PREC_EXACT': 'cdef const slong ARF_PREC_EXACT',
     },
@@ -185,16 +188,16 @@ def get_macro_declarations(macros, h_name):
     Untyped macros (which may also introduce variables or require lvalues)
     cannot in general be represented this way.
     """
-    functions, constants = [], []
+    functions = []
     overrides = macro_declarations.get(h_name, {})
+    # Some constants (e.g. MAG_BITS) are not documented in older RST docs.
+    constants = [decl for decl in overrides.values() if decl.startswith("cdef ")]
     for macro in macros:
         signature = macro.split('::', 1)[1].strip()
         name = signature.split('(', 1)[0].strip()
         declaration = overrides.get(name)
         if declaration is not None:
-            if declaration.startswith('cdef '):
-                constants.append(declaration)
-            else:
+            if not declaration.startswith('cdef '):
                 functions.append(clean_types(declaration))
         elif typed_macro.match(signature):
             functions.append(clean_types(signature))

src/flint/flintlib/functions/mag.pxd

@@ -6,6 +6,7 @@ from flint.flintlib.types.fmpq cimport fmpq_t
 
 
 cdef extern from "flint/mag.h":
+    cdef const int MAG_BITS
     void mag_init(mag_t x)
     void mag_clear(mag_t x)
     void mag_swap(mag_t x, mag_t y)

src/flint/test/test_acb.py

@@ -1,11 +1,57 @@
 """Tests for python-flint's `acb` constructor."""
 
 import math
+import pickle
 
-from flint import acb, arb, dirichlet_char, fmpz_poly
+from flint import acb, arb, ctx, dirichlet_char, fmpz_poly
 from flint.test.helpers import is_close_acb, is_close_arb, raises
 
 
+def _assert_same_acb(a: acb, b: acb) -> None:
+    for original, restored in [(a.real, b.real), (a.imag, b.imag)]:
+        if original.is_nan():
+            assert restored.is_nan()
+        else:
+            assert original.mid() == restored.mid()
+        assert original.rad() == restored.rad()
+
+
+def test_acb_pickle() -> None:
+    """Pickling preserves both balls, including special and large values."""
+    with ctx.workprec(200):
+        components = [
+            arb(0), arb(-3), arb.pi(), arb(69).zeta(69),
+            arb((2**20000 + 1, -123), (2**30 - 1, -1000)),
+            arb((1, 10**30), (1, -10**30)),
+            arb("inf"), arb("-inf"), arb("nan"), arb(0, "inf"),
+            arb("inf", (3, 4)), arb("nan", (3, 4)),
+        ]
+        values = [acb(x, y) for x, y in zip(components, reversed(components))]
+        values.extend([acb(), acb(3), acb(0, 3)])
+    with ctx.workprec(10):
+        for a in values:
+            for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
+                b = pickle.loads(pickle.dumps(a, protocol=protocol))
+                assert type(b) is acb
+                _assert_same_acb(a, b)
+
+
+def test_acb_repr_roundtrip() -> None:
+    """Both components round-trip via repr and integer midpoint/radius data."""
+    with ctx.workprec(53):
+        zeta = arb(69).zeta(69)
+    values = [acb(), acb(zeta), acb(0, zeta), acb(zeta, zeta)]
+    for man in [2**29 - 1, 2**29 + 1, 2**30 - 1]:
+        values.append(acb(arb((3, 1000), (man, -1000)), arb((5, -1000), (man, 1000))))
+    with ctx.workprec(10):
+        for a in values:
+            b = eval(a.repr())
+            _assert_same_acb(a, b)
+            real = arb(a.real.mid().man_exp(), a.real.rad().man_exp())
+            imag = arb(a.imag.mid().man_exp(), a.imag.rad().man_exp())
+            _assert_same_acb(a, acb(real, imag))
+
+
 def test_acb_constructor() -> None:
     """Cover all branches of acb.__init__."""
 

src/flint/test/test_arb.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import math
+import pickle
 
 from flint import arb, ctx
 from flint.test.helpers import is_close_arb, raises
@@ -72,6 +73,64 @@ def test_is_exact() -> None:
     ]:
         assert arb_val.is_exact() == exact
 
+def test_pickle() -> None:
+    """Pickling preserves midpoint and radius at any context precision."""
+    with ctx.workprec(200):
+        values = [
+            arb(0), arb(-3), arb.pi(), arb(69).zeta(69),
+            arb((2**20000 + 1, -123), (2**30 - 1, -1000)),
+            arb((1, 10**30), (1, -10**30)),
+            arb("inf"), arb("-inf"), arb("nan"),
+            arb(0, "inf"), arb(1, "inf"),
+            arb("inf", (3, 4)), arb("nan", (3, 4)),
+        ]
+    with ctx.workprec(10):
+        for a in values:
+            for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
+                b = pickle.loads(pickle.dumps(a, protocol=protocol))
+                assert type(b) is arb
+                if a.is_nan():
+                    assert b.is_nan()
+                else:
+                    assert a.mid() == b.mid()
+                assert a.rad() == b.rad()
+
+def test_tuple_radius_exact() -> None:
+    """Representable tuple radii are preserved, independently of precision."""
+    for prec in [10, 53, 200]:
+        with ctx.workprec(prec):
+            for man in [0, 1, 2**29 - 1, 2**29, 2**29 + 1, 2**30 - 1]:
+                for exp in [-1000, 0, 1000, 10**30]:
+                    for sign in [-1, 1]:
+                        # Extra powers of two must not count towards precision.
+                        a = arb((2**100 + 1, -50), (sign * man * 8, exp - 3))
+                        assert a.mid() == arb((2**100 + 1, -50))
+                        assert a.rad() == arb((man, exp))
+                        b = eval(a.repr())
+                        assert a.mid() == b.mid()
+                        assert a.rad() == b.rad()
+
+def test_tuple_radius_rounds_outward() -> None:
+    """Radii requiring more than 30 bits still round upwards."""
+    for man in [2**30 + 1, 2**31 - 1, 2**100 + 1]:
+        for exp in [-1000, 0, 1000]:
+            exact = arb((man, exp))
+            for sign in [-1, 1]:
+                a = arb(0, (sign * man, exp))
+                assert a.rad() >= exact
+                assert a.contains(exact)
+                assert a.contains(-exact)
+
+def test_zeta_repr_roundtrip() -> None:
+    """Regression for gh-438: serialization preserves a full-width radius."""
+    with ctx.workprec(53):
+        a = arb(69).zeta(69)
+    # Reconstruct at a lower precision to check that no rounding is introduced.
+    with ctx.workprec(10):
+        for b in [eval(a.repr()), arb(a.mid().man_exp(), a.rad().man_exp())]:
+            assert a.mid() == b.mid()
+            assert a.rad() == b.rad()
+
 def test_is_finite() -> None:
     """Tests `arb.is_finite`."""
     assert not (arb(-float("inf")).is_finite())
@@ -756,6 +815,10 @@ def test_no_tests_missing() -> None:
     test_from_float_inf,
     test_from_man_exp,
     test_from_midpoint_radius,
+    test_pickle,
+    test_tuple_radius_exact,
+    test_tuple_radius_rounds_outward,
+    test_zeta_repr_roundtrip,
     test_is_exact,
     test_is_finite,
     test_is_nan,

src/flint/types/acb.pyx

@@ -256,6 +256,9 @@ cdef class acb(flint_scalar):
         arf_set_mag(arb_midref(acb_imagref(u.val)), arb_radref(acb_imagref(self.val)))
         return u
 
+    def __reduce__(self):
+        return (acb, (self.real, self.imag))
+
     def repr(self):
         real = self.real
         imag = self.imag

src/flint/types/arb.pyx

@@ -14,7 +14,7 @@ from flint.types.fmpz cimport fmpz
 from flint.flintlib.types.flint cimport FMPZ_UNKNOWN, FMPZ_TMP, FMPZ_REF
 from flint.flintlib.functions.flint cimport flint_free
 from flint.flintlib.functions.mag cimport *
-from flint.flintlib.functions.fmpz cimport fmpz_init, fmpz_clear
+from flint.flintlib.functions.fmpz cimport fmpz_init, fmpz_clear, fmpz_bits
 from flint.flintlib.functions.arf cimport *
 from flint.flintlib.types.arf cimport ARF_RND_NEAR, ARF_RND_DOWN
 from flint.flintlib.functions.arb cimport *
@@ -34,6 +34,15 @@ cimport cython
 
 ctx = thectx
 
+
+def _arb_from_pickle(mid, rad):
+    if isinstance(mid, tuple):
+        mid = tuple(int(x, 16) for x in mid)
+    if isinstance(rad, tuple):
+        rad = tuple(int(x, 16) for x in rad)
+    return arb(mid, rad)
+
+
 # cdef _str_trunc(s, trunc=0):
 #     if trunc > 0 and len(s) > 3 * trunc:
 #         left = right = trunc
@@ -207,7 +216,13 @@ cdef class arb(flint_scalar):
                 # eval(a.repr()) round trips.
                 _rad = arf(rad)
                 man, exp = _rad.man_exp()
-                mag_set_fmpz_2exp_fmpz(arb_radref(self.val), man.val, exp.val)
+                if fmpz_bits(man.val) <= MAG_BITS:
+                    # The upper-bound conversion can round up even when the
+                    # radius is exactly representable. The lower bound is
+                    # exact here since man_exp() removes trailing zero bits.
+                    mag_set_fmpz_2exp_fmpz_lower(arb_radref(self.val), man.val, exp.val)
+                else:
+                    mag_set_fmpz_2exp_fmpz(arb_radref(self.val), man.val, exp.val)
             else:
                 rad = arb(rad)
                 arb_add_error(self.val, (<arb>rad).val)
@@ -400,6 +415,23 @@ cdef class arb(flint_scalar):
         arf_set_mag(rad.val, arb_radref(self.val))
         return (mid, rad)
 
+    def __reduce__(self):
+        # Hex avoids the decimal digit limit in text pickle protocols.
+        mid, rad = self._to_arfs()
+        if mid.is_finite():
+            mid = tuple(hex(x) for x in mid.man_exp())
+        elif mid.is_nan():
+            mid = "nan"
+        elif mid.is_pos_inf():
+            mid = "inf"
+        else:
+            mid = "-inf"
+        if rad.is_finite():
+            rad = tuple(hex(x) for x in rad.man_exp())
+        else:
+            rad = "inf"
+        return (_arb_from_pickle, (mid, rad))
+
     def repr(self):
         mid, rad = self._to_arfs()
         if rad.is_zero() or mid.is_nan():