gmpy.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342
  1. from __future__ import annotations
  2. import os
  3. from ctypes import c_long, sizeof
  4. from functools import reduce
  5. from typing import Type
  6. from warnings import warn
  7. from sympy.external import import_module
  8. from .pythonmpq import PythonMPQ
  9. from .ntheory import (
  10. bit_scan1 as python_bit_scan1,
  11. bit_scan0 as python_bit_scan0,
  12. remove as python_remove,
  13. factorial as python_factorial,
  14. sqrt as python_sqrt,
  15. sqrtrem as python_sqrtrem,
  16. gcd as python_gcd,
  17. lcm as python_lcm,
  18. gcdext as python_gcdext,
  19. is_square as python_is_square,
  20. invert as python_invert,
  21. legendre as python_legendre,
  22. jacobi as python_jacobi,
  23. kronecker as python_kronecker,
  24. iroot as python_iroot,
  25. is_fermat_prp as python_is_fermat_prp,
  26. is_euler_prp as python_is_euler_prp,
  27. is_strong_prp as python_is_strong_prp,
  28. is_fibonacci_prp as python_is_fibonacci_prp,
  29. is_lucas_prp as python_is_lucas_prp,
  30. is_selfridge_prp as python_is_selfridge_prp,
  31. is_strong_lucas_prp as python_is_strong_lucas_prp,
  32. is_strong_selfridge_prp as python_is_strong_selfridge_prp,
  33. is_bpsw_prp as python_is_bpsw_prp,
  34. is_strong_bpsw_prp as python_is_strong_bpsw_prp,
  35. )
  36. __all__ = [
  37. # GROUND_TYPES is either 'gmpy' or 'python' depending on which is used. If
  38. # gmpy is installed then it will be used unless the environment variable
  39. # SYMPY_GROUND_TYPES is set to something other than 'auto', 'gmpy', or
  40. # 'gmpy2'.
  41. 'GROUND_TYPES',
  42. # If HAS_GMPY is 0, no supported version of gmpy is available. Otherwise,
  43. # HAS_GMPY will be 2 for gmpy2 if GROUND_TYPES is 'gmpy'. It used to be
  44. # possible for HAS_GMPY to be 1 for gmpy but gmpy is no longer supported.
  45. 'HAS_GMPY',
  46. # SYMPY_INTS is a tuple containing the base types for valid integer types.
  47. # This is either (int,) or (int, type(mpz(0))) depending on GROUND_TYPES.
  48. 'SYMPY_INTS',
  49. # MPQ is either gmpy.mpq or the Python equivalent from
  50. # sympy.external.pythonmpq
  51. 'MPQ',
  52. # MPZ is either gmpy.mpz or int.
  53. 'MPZ',
  54. 'bit_scan1',
  55. 'bit_scan0',
  56. 'remove',
  57. 'factorial',
  58. 'sqrt',
  59. 'is_square',
  60. 'sqrtrem',
  61. 'gcd',
  62. 'lcm',
  63. 'gcdext',
  64. 'invert',
  65. 'legendre',
  66. 'jacobi',
  67. 'kronecker',
  68. 'iroot',
  69. 'is_fermat_prp',
  70. 'is_euler_prp',
  71. 'is_strong_prp',
  72. 'is_fibonacci_prp',
  73. 'is_lucas_prp',
  74. 'is_selfridge_prp',
  75. 'is_strong_lucas_prp',
  76. 'is_strong_selfridge_prp',
  77. 'is_bpsw_prp',
  78. 'is_strong_bpsw_prp',
  79. ]
  80. #
  81. # Tested python-flint version. Future versions might work but we will only use
  82. # them if explicitly requested by SYMPY_GROUND_TYPES=flint.
  83. #
  84. _PYTHON_FLINT_VERSION_NEEDED = ["0.6", "0.7", "0.8", "0.9", "0.10"]
  85. def _flint_version_okay(flint_version):
  86. major, minor = flint_version.split('.')[:2]
  87. flint_ver = f'{major}.{minor}'
  88. return flint_ver in _PYTHON_FLINT_VERSION_NEEDED
  89. #
  90. # We will only use gmpy2 >= 2.0.0
  91. #
  92. _GMPY2_MIN_VERSION = '2.0.0'
  93. def _get_flint(sympy_ground_types):
  94. if sympy_ground_types not in ('auto', 'flint'):
  95. return None
  96. try:
  97. import flint
  98. # Earlier versions of python-flint may not have __version__.
  99. from flint import __version__ as _flint_version
  100. except ImportError:
  101. if sympy_ground_types == 'flint':
  102. warn("SYMPY_GROUND_TYPES was set to flint but python-flint is not "
  103. "installed. Falling back to other ground types.")
  104. return None
  105. if _flint_version_okay(_flint_version):
  106. return flint
  107. elif sympy_ground_types == 'auto':
  108. return None
  109. else:
  110. warn(f"Using python-flint {_flint_version} because SYMPY_GROUND_TYPES "
  111. f"is set to flint but this version of SymPy is only tested "
  112. f"with python-flint versions {_PYTHON_FLINT_VERSION_NEEDED}.")
  113. return flint
  114. def _get_gmpy2(sympy_ground_types):
  115. if sympy_ground_types not in ('auto', 'gmpy', 'gmpy2'):
  116. return None
  117. gmpy = import_module('gmpy2', min_module_version=_GMPY2_MIN_VERSION,
  118. module_version_attr='version', module_version_attr_call_args=())
  119. if sympy_ground_types != 'auto' and gmpy is None:
  120. warn("gmpy2 library is not installed, switching to 'python' ground types")
  121. return gmpy
  122. #
  123. # SYMPY_GROUND_TYPES can be flint, gmpy, gmpy2, python or auto (default)
  124. #
  125. _SYMPY_GROUND_TYPES = os.environ.get('SYMPY_GROUND_TYPES', 'auto').lower()
  126. _flint = None
  127. _gmpy = None
  128. #
  129. # First handle auto-detection of flint/gmpy2. We will prefer flint if available
  130. # or otherwise gmpy2 if available and then lastly the python types.
  131. #
  132. if _SYMPY_GROUND_TYPES in ('auto', 'flint'):
  133. _flint = _get_flint(_SYMPY_GROUND_TYPES)
  134. if _flint is not None:
  135. _SYMPY_GROUND_TYPES = 'flint'
  136. else:
  137. _SYMPY_GROUND_TYPES = 'auto'
  138. if _SYMPY_GROUND_TYPES in ('auto', 'gmpy', 'gmpy2'):
  139. _gmpy = _get_gmpy2(_SYMPY_GROUND_TYPES)
  140. if _gmpy is not None:
  141. _SYMPY_GROUND_TYPES = 'gmpy'
  142. else:
  143. _SYMPY_GROUND_TYPES = 'python'
  144. if _SYMPY_GROUND_TYPES not in ('flint', 'gmpy', 'python'):
  145. warn("SYMPY_GROUND_TYPES environment variable unrecognised. "
  146. "Should be 'auto', 'flint', 'gmpy', 'gmpy2' or 'python'.")
  147. _SYMPY_GROUND_TYPES = 'python'
  148. #
  149. # At this point _SYMPY_GROUND_TYPES is either flint, gmpy or python. The blocks
  150. # below define the values exported by this module in each case.
  151. #
  152. #
  153. # In gmpy2 and flint, there are functions that take a long (or unsigned long)
  154. # argument. That is, it is not possible to input a value larger than that.
  155. #
  156. LONG_MAX = (1 << (8*sizeof(c_long) - 1)) - 1
  157. #
  158. # Type checkers are confused by what SYMPY_INTS is. There may be a better type
  159. # hint for this like Type[Integral] or something.
  160. #
  161. SYMPY_INTS: tuple[Type, ...]
  162. if _SYMPY_GROUND_TYPES == 'gmpy':
  163. assert _gmpy is not None
  164. flint = None
  165. gmpy = _gmpy
  166. HAS_GMPY = 2
  167. GROUND_TYPES = 'gmpy'
  168. SYMPY_INTS = (int, type(gmpy.mpz(0)))
  169. MPZ = gmpy.mpz
  170. MPQ = gmpy.mpq
  171. bit_scan1 = gmpy.bit_scan1
  172. bit_scan0 = gmpy.bit_scan0
  173. remove = gmpy.remove
  174. factorial = gmpy.fac
  175. sqrt = gmpy.isqrt
  176. is_square = gmpy.is_square
  177. sqrtrem = gmpy.isqrt_rem
  178. gcd = gmpy.gcd
  179. lcm = gmpy.lcm
  180. gcdext = gmpy.gcdext
  181. invert = gmpy.invert
  182. legendre = gmpy.legendre
  183. jacobi = gmpy.jacobi
  184. kronecker = gmpy.kronecker
  185. def iroot(x, n):
  186. # In the latest gmpy2, the threshold for n is ULONG_MAX,
  187. # but adjust to the older one.
  188. if n <= LONG_MAX:
  189. return gmpy.iroot(x, n)
  190. return python_iroot(x, n)
  191. is_fermat_prp = gmpy.is_fermat_prp
  192. is_euler_prp = gmpy.is_euler_prp
  193. is_strong_prp = gmpy.is_strong_prp
  194. is_fibonacci_prp = gmpy.is_fibonacci_prp
  195. is_lucas_prp = gmpy.is_lucas_prp
  196. is_selfridge_prp = gmpy.is_selfridge_prp
  197. is_strong_lucas_prp = gmpy.is_strong_lucas_prp
  198. is_strong_selfridge_prp = gmpy.is_strong_selfridge_prp
  199. is_bpsw_prp = gmpy.is_bpsw_prp
  200. is_strong_bpsw_prp = gmpy.is_strong_bpsw_prp
  201. elif _SYMPY_GROUND_TYPES == 'flint':
  202. assert _flint is not None
  203. flint = _flint
  204. gmpy = None
  205. HAS_GMPY = 0
  206. GROUND_TYPES = 'flint'
  207. SYMPY_INTS = (int, flint.fmpz) # type: ignore
  208. MPZ = flint.fmpz # type: ignore
  209. MPQ = flint.fmpq # type: ignore
  210. bit_scan1 = python_bit_scan1
  211. bit_scan0 = python_bit_scan0
  212. remove = python_remove
  213. factorial = python_factorial
  214. def sqrt(x):
  215. return flint.fmpz(x).isqrt()
  216. def is_square(x):
  217. if x < 0:
  218. return False
  219. return flint.fmpz(x).sqrtrem()[1] == 0
  220. def sqrtrem(x):
  221. return flint.fmpz(x).sqrtrem()
  222. def gcd(*args):
  223. return reduce(flint.fmpz.gcd, args, flint.fmpz(0))
  224. def lcm(*args):
  225. return reduce(flint.fmpz.lcm, args, flint.fmpz(1))
  226. gcdext = python_gcdext
  227. invert = python_invert
  228. legendre = python_legendre
  229. def jacobi(x, y):
  230. if y <= 0 or not y % 2:
  231. raise ValueError("y should be an odd positive integer")
  232. return flint.fmpz(x).jacobi(y)
  233. kronecker = python_kronecker
  234. def iroot(x, n):
  235. if n <= LONG_MAX:
  236. y = flint.fmpz(x).root(n)
  237. return y, y**n == x
  238. return python_iroot(x, n)
  239. is_fermat_prp = python_is_fermat_prp
  240. is_euler_prp = python_is_euler_prp
  241. is_strong_prp = python_is_strong_prp
  242. is_fibonacci_prp = python_is_fibonacci_prp
  243. is_lucas_prp = python_is_lucas_prp
  244. is_selfridge_prp = python_is_selfridge_prp
  245. is_strong_lucas_prp = python_is_strong_lucas_prp
  246. is_strong_selfridge_prp = python_is_strong_selfridge_prp
  247. is_bpsw_prp = python_is_bpsw_prp
  248. is_strong_bpsw_prp = python_is_strong_bpsw_prp
  249. elif _SYMPY_GROUND_TYPES == 'python':
  250. flint = None
  251. gmpy = None
  252. HAS_GMPY = 0
  253. GROUND_TYPES = 'python'
  254. SYMPY_INTS = (int,)
  255. MPZ = int
  256. MPQ = PythonMPQ
  257. bit_scan1 = python_bit_scan1
  258. bit_scan0 = python_bit_scan0
  259. remove = python_remove
  260. factorial = python_factorial
  261. sqrt = python_sqrt
  262. is_square = python_is_square
  263. sqrtrem = python_sqrtrem
  264. gcd = python_gcd
  265. lcm = python_lcm
  266. gcdext = python_gcdext
  267. invert = python_invert
  268. legendre = python_legendre
  269. jacobi = python_jacobi
  270. kronecker = python_kronecker
  271. iroot = python_iroot
  272. is_fermat_prp = python_is_fermat_prp
  273. is_euler_prp = python_is_euler_prp
  274. is_strong_prp = python_is_strong_prp
  275. is_fibonacci_prp = python_is_fibonacci_prp
  276. is_lucas_prp = python_is_lucas_prp
  277. is_selfridge_prp = python_is_selfridge_prp
  278. is_strong_lucas_prp = python_is_strong_lucas_prp
  279. is_strong_selfridge_prp = python_is_strong_selfridge_prp
  280. is_bpsw_prp = python_is_bpsw_prp
  281. is_strong_bpsw_prp = python_is_strong_bpsw_prp
  282. else:
  283. assert False