Data commit

Task/P-Adic-square-roots/Nim/p-adic-square-roots.nim

@@ -0,0 +1,251 @@
+import strformat
+
+const
+  Emx = 64        # Exponent maximum.
+  Amx = 6000      # Argument maximum.
+  PMax = 32749    # Prime maximum.
+
+type
+
+  Ratio = tuple[a, b: int]
+
+  Padic = object
+    p: int                        # Prime.
+    k: int                        # Precision.
+    v: int
+    d: array[-Emx..(Emx-1), int]
+
+  PadicError = object of ValueError
+
+
+proc sqrt(pa: var Padic; q: Ratio; sw: bool) =
+  ## Return the p-adic square root of q = a/b. Set sw to print.
+
+  var (a, b) = q
+  var i, x: int
+
+  if b == 0:
+    raise newException(PadicError, &"Wrong rational: {a}/{b}" )
+  if b < 0:
+    b = -b
+    a = -a
+  if pa.p  < 2:
+    raise newException(PadicError, &"Wrong value for p: {pa.p}")
+  if pa.k < 1:
+    raise newException(PadicError, &"Wrong value for k: {pa.k}")
+  pa.p = min(pa.p, PMax)      # Maximum short prime.
+
+  if sw: echo &"{a}/{b} + 0({pa.p}^{pa.k})"
+
+  # Initialize.
+  pa.v = 0
+  pa.d.reset()
+  if a == 0: return
+
+  # Valuation.
+  while b mod pa.p == 0:
+    b = b div pa.p
+    dec pa.v
+  while a mod pa.p == 0:
+    a = a div pa.p
+    inc pa.v
+  if (pa.v and 1) != 0:
+    # Odd valuation.
+    raise newException(PadicError, &"Non-residue mod {pa.p}.")
+
+  # Maximum array length.
+  pa.k = min(pa.k + pa.v, Emx - 1) - pa.v
+  pa.v = pa.v shr 1
+
+  if abs(a) > Amx or b > Amx:
+    raise newException(PadicError, &"Rational exceeding limits: {a}/{b}.")
+
+  if pa.p == 2:
+    # 1 / b = b (mod 8); a / b = 1 (mod 8).
+    if (a * b and 7) - 1 != 0:
+      raise newException(PadicError, "Non-residue mod 8.")
+
+    # Initialize.
+    x = 1
+    pa.d[pa.v] = 1
+    pa.d[pa.v + 1] = 0
+    var pk = 4
+    i = pa.v + 2
+    while i < pa.k + pa.v:
+      pk *= 2
+      let f = b * x * x - a
+      let q = f div pk
+      if f != q * pk: break   # Overflow.
+      # Next digit.
+      pa.d[i] = if (q and 1) != 0: 1 else: 0
+      # Lift "x".
+      x += pa.d[i] * (pk shr 1)
+      inc i
+
+  else:
+    # Find root for small "p".
+    var r = 1
+    while r < pa.p:
+      if (b * r * r - a) mod pa.p == 0: break
+      inc r
+    if r == pa.p:
+      raise newException(PadicError, &"Non-residue mod {pa.p}.")
+    let t = (b * r shl 1) mod pa.p
+    var s = 0
+
+    # Modular inverse for small "p".
+    var f1 = 1
+    while f1 < pa.p:
+      inc s, t
+      if s >= pa.p: dec s, pa.p
+      if s == 1: break
+      inc f1
+    if f1 == pa.p:
+      raise newException(PadicError, "Impossible to compute inverse modulo")
+
+    f1 = pa.p - f1
+    x = r
+    pa.d[pa.v] = x
+
+    var pk = 1
+    i = pa.v + 1
+    while i < pa.k + pa.v:
+      pk *= pa.p
+      let f = b * x * x - a
+      let q = f div pk
+      if f != q * pk: break   # Overflow.
+      pa.d[i] = q * f1 mod pa.p
+      if pa.d[i] < 0: pa.d[i] += pa.p
+      x += pa.d[i] * pk
+      inc i
+
+  pa.k = i - pa.v
+  if sw: echo &"lift: {x} mod {pa.p}^{pa.k}"
+
+
+proc crat(pa: Padic; sw: bool): Ratio =
+  ## Rational reconstruction.
+
+  # Weighted digit sum.
+  var
+    s = 0
+    pk = 1
+  for i in min(pa.v, 0)..<(pa.k + pa.v):
+    let pm = pk
+    pk *= pa.p
+    if pk div pm - pa.p != 0:
+      # Overflow.
+      pk = pm
+      break
+    s += pa.d[i] * pm
+
+  # Lattice basis reduction.
+  var
+    m = [pk, s]
+    n = [0, 1]
+    i = 0
+    j = 1
+  s = s * s + 1
+  # Lagrange's algorithm.
+  while true:
+    # Euclidean step.
+    var q = ((m[i] * m[j] + n[i] * n[j]) / s).toInt
+    m[i] -= q * m[j]
+    n[i] -= q * n[j]
+    q = s
+    s = m[i] * m[i] + n[i] * n[i]
+    # Compare norms.
+    if s < q: swap i, j   # Interchange vectors.
+    else: break
+
+  var x = m[j]
+  var y = n[j]
+  if y < 0:
+    y = -y
+    x = -x
+
+  # Check determinant.
+  if abs(m[i] * y - x * n[i]) != pk:
+    raise newException(PadicError, "Rational reconstruction failed.")
+
+  # Negative powers.
+  for i in pa.v..(-1): y *= pa.p
+
+  if sw: echo x, if y > 1: '/' & $y else: ""
+  result = (x, y)
+
+
+func cmpt(pa: Padic): Padic =
+  ## Return the complement.
+  result = Padic(p: pa.p, k: pa.k, v: pa.v)
+  var c = 1
+  for i in pa.v..(pa.k + pa.v):
+    inc c, pa.p - 1 - pa.d[i]
+    if c >= pa.p:
+      result.d[i] = c - pa.p
+      c = 1
+    else:
+      result.d[i] = c
+      c = 0
+
+
+func sqr(pa: Padic): Padic =
+  ## Return the square of a P-adic number.
+  result = Padic(p: pa.p, k: pa.k, v: pa.v * 2)
+  var c = 0
+  for i in 0..pa.k:
+    for j in 0..i:
+      c += pa.d[pa.v + j] * pa.d[pa.v + i - j]
+    # Euclidean step.
+    let q = c div pa.p
+    result.d[result.v + i] = c - q * pa.p
+    c = q
+
+
+func `$`(pa: Padic): string =
+  ## String representation.
+  let t = min(pa.v, 0)
+  for i in countdown(pa.k - 1 + t, t):
+    result.add $pa.d[i]
+    if i == 0 and pa.v < 0: result.add "."
+    result.add " "
+
+
+when isMainModule:
+
+  const Data = [[-7, 1, 2, 7],
+                [9, 1, 2, 8],
+                [17, 1, 2, 9],
+                [497, 10496, 2, 18],
+                [10496, 497, 2, 19],
+                [3141, 5926, 3, 15],
+                [2718,  281, 3, 13],
+                [-1,  1,  5, 8],
+                [86, 25,  5, 8],
+                [2150, 1,  5, 8],
+                [2,1, 7, 8],
+                [-2645, 28518, 7, 9],
+                [3029, 4821, 7, 9],
+                [379, 449, 7, 8],
+                [717, 8, 11, 7],
+                [1414, 213, 41, 5],
+                [-255, 256, 257, 3]]
+
+  for d in Data:
+    try:
+      let q: Ratio = (d[0], d[1])
+      var a = Padic(p: d[2], k: d[3])
+      a.sqrt(q, true)
+      echo "sqrt +/-"
+      echo "...", a
+      a = a.cmpt()
+      echo "...", a
+      let c = sqr(a)
+      echo "sqrt^2"
+      echo "   ", c
+      let r = c.crat(true)
+      if q.a * r.b - r.a * q.b != 0:
+        echo "fail: sqrt^2"
+      echo ""
+    except PadicError:
+      echo getCurrentExceptionMsg()