dirichlet: add type stubs and tests

Author: oscarbenjamin

src/flint/test/meson.build

@@ -16,6 +16,7 @@ pyfiles = [
   'test_fmpq_vec.py',
   'test_fmpz_vec.py',
   'test_docstrings.py',
+  'test_dirichlet.py',
 ]
 
 py.install_sources(

src/flint/test/test_dirichlet.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from flint import dirichlet_char, dirichlet_group, fmpz
+from flint.test.helpers import is_close_acb, raises
+
+
+def test_dirichlet_group_basics() -> None:
+    g = dirichlet_group(5)
+    assert g.size() == 4
+    assert g.q == 5
+    assert g.exponent() == 4
+    assert repr(g) == "Dirichlet group mod q = 5"
+    assert str(g) == "Dirichlet group mod q = 5"
+
+    assert raises(lambda: dirichlet_group(0), AssertionError)
+
+
+def test_dirichlet_char_properties() -> None:
+    chi = dirichlet_char(7, 3)
+
+    assert isinstance(chi.group(), dirichlet_group)
+    assert chi.index() == 1
+    assert chi.modulus() == 7
+    assert chi.number() == 3
+    assert chi.order() == 6
+    assert chi.is_real() is False
+    assert chi.is_primitive() is True
+    assert chi.conductor() == 7
+    assert chi.is_principal() is False
+    assert chi.parity() == 1
+
+    assert repr(chi) == "dirichlet_char(7, 3)"
+    assert str(chi) == "dirichlet_char(7, 3)"
+
+    assert raises(lambda: dirichlet_char(8, 2), AssertionError)
+
+
+def test_dirichlet_char_values_and_mul() -> None:
+    chi = dirichlet_char(7, 3)
+    principal = dirichlet_char(7, 1)
+    chi5 = dirichlet_char(7, 5)
+
+    assert is_close_acb(chi(7), 0)
+    assert is_close_acb(chi(2), -0.5 + 0.866025403784439j, tol=1e-12, rel_tol=1e-12)
+    assert is_close_acb(chi(fmpz(2)), -0.5 + 0.866025403784439j, tol=1e-12, rel_tol=1e-12)
+
+    assert chi.chi_exponent(7) is None
+    assert chi.chi_exponent(2) == 2
+    assert principal.chi_exponent(2) == 0
+
+    prod = chi * chi5
+    assert prod.number() == 1
+    for n in [1, 2, 3, 4, 5, 6, 7, 8]:
+        assert is_close_acb(prod(n), chi(n) * chi5(n), tol=1e-12, rel_tol=1e-12)
+
+    assert raises(lambda: chi * dirichlet_char(5, 1), AssertionError)
+    assert raises(lambda: chi(object()), TypeError)  # type: ignore[arg-type]
+    assert raises(lambda: chi.chi_exponent(object()), TypeError)  # type: ignore[arg-type]
+
+
+def test_dirichlet_char_special_functions() -> None:
+    chi = dirichlet_char(7, 3)
+    assert is_close_acb(chi.l_function(2), 0.902247025301257 + 0.232548981277895j, tol=1e-12, rel_tol=1e-12)
+    assert is_close_acb(chi.l(2), chi.l_function(2), tol=1e-12, rel_tol=1e-12)
+
+    zeta_chi = dirichlet_char(1, 1)
+    assert is_close_acb(zeta_chi.hardy_z(1), -0.736305462867318, tol=1e-12, rel_tol=1e-12)

src/flint/types/acb_series.pyi

@@ -8,6 +8,7 @@ from .acb_poly import acb_poly
 from .arb import arb
 from .arb_poly import arb_poly
 from .arb_series import arb_series
+from .dirichlet import dirichlet_char
 from .fmpq import fmpq
 from .fmpq_poly import fmpq_poly
 from .fmpq_series import fmpq_series
@@ -76,7 +77,7 @@ class acb_series(flint_series[acb]):
     def lgamma(self) -> acb_series: ...
     def rising(self, n: int, /) -> acb_series: ...
     def zeta(self, a: iacb = ..., deflate: bool = ...) -> acb_series: ...
-    def dirichlet_l(self, chi: tuple[int, int], deflate: bool = ...) -> acb_series: ...
+    def dirichlet_l(self, chi: dirichlet_char | tuple[int, int], deflate: bool = ...) -> acb_series: ...
 
     @classmethod
     def polylog(cls, s: iacb_series, z: iacb) -> acb_series: ...

src/flint/types/dirichlet.pyi

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from flint.types.acb import acb
+from flint.types.fmpz import fmpz
+
+
+class dirichlet_group:
+    def __init__(self, q: int, /) -> None: ...
+    def size(self) -> int: ...
+    @property
+    def q(self) -> fmpz: ...
+    def exponent(self) -> fmpz: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+
+
+class dirichlet_char:
+    def __init__(self, q: int, l: int = 1, /) -> None: ...
+    def group(self) -> dirichlet_group: ...
+    def index(self) -> int: ...
+    def modulus(self) -> fmpz: ...
+    def number(self) -> int: ...
+    def order(self) -> int: ...
+    def is_real(self) -> bool: ...
+    def is_primitive(self) -> bool: ...
+    def conductor(self) -> int: ...
+    def is_principal(self) -> bool: ...
+    def parity(self) -> int: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+    def __mul__(self, other: dirichlet_char, /) -> dirichlet_char: ...
+    def __call__(self, n: int | fmpz, /) -> acb: ...
+    def chi_exponent(self, n: int | fmpz, /) -> fmpz | None: ...
+    def l_function(self, s: int | float | complex | str | acb | fmpz, /) -> acb: ...
+    def l(self, s: int | float | complex | str | acb | fmpz, /) -> acb: ...
+    def hardy_z(self, s: int | float | complex | str | acb | fmpz, /) -> acb: ...

src/flint/types/meson.build

@@ -43,7 +43,7 @@ pyfiles = [
   'acb_mat.pyi',
   'acb_series.pyi',
   #
-  # 'dirichlet.pyi',
+  'dirichlet.pyi',
   #
   # '_gr.pyi',
 ]