modp.py

@@ -1,5 +1,5 @@
 class ModP(int):
-    'Integers mod p.'
+    'Integers mod p, p a prime power.'
     def __new__(cls, p, num):
         self.p = p
         return int.__new__(cls, num % p)

padic.py

@@ -2,15 +2,9 @@ from fractions import Fraction
 from sys import maxint
 from modp import *
 
-def padic(prime):
-    'Return p-adic class with given prime.'
-    class Res(PAdic):
-        p = prime
-    return Res
-
 class PAdic:
-    p = 2 # default
-    def __init__(self):
+    def __init__(self, p):
+        self.p = p
         self.val  = '' # current known value
         self.prec = 0 # current known precision
         self.order = 0 # order/valuation of number
@@ -23,7 +17,7 @@ class PAdic:
         
         while self.prec < prec:
             # update val based on value
-            self._nextdigit()
+            self.val = self._nextdigit() + self.val
             self.prec += 1
         return self.val # TODO add decimal point or trailing zeros
     
@@ -31,6 +25,10 @@ class PAdic:
         'Calculate next digit of p-adic number.'
         raise NotImplementedError
     
+    def getdigit(self, index):
+        'Return digit at given index.'
+        return int(self.get(index+1)[0])
+    
     # return value with precision up to 32 bits
     def __int__(self):
         return int(self.get(32), p)
@@ -41,17 +39,17 @@ class PAdic:
     
     # arithmetic operations
     def __add__(self, other):
-        return PAdicAdd(self, other)
+        return PAdicAdd(self.p, self, other)
     def __radd__(self, other):
-        return PAdicAdd(other, self)
+        return PAdicAdd(self.p, other, self)
     def __sub__(self, other):
-        return PAdicAdd(self, PAdicInv(other))
+        return PAdicAdd(self.p, self, PAdicNeg(self.p, other))
     def __rsub__(self, other):
-        return PAdicAdd(other, PAdicInv(self))
+        return PAdicAdd(self.p, other, PAdicNeg(self.p, self))
     def __smul__(self, other):
-        return PAdicMul(self, other)
+        return PAdicMul(self.p, self, other)
     def __rsub__(self, other):
-        return PAdicMul(other, self)
+        return PAdicMul(self.p, other, self)
     
     # p-adic norm
     def __abs__(self):
@@ -59,14 +57,14 @@ class PAdic:
             return 0
         norm = Fraction(1)
         if self.order > 0:
-            norm.numerator = p ** self.order
+            norm.numerator = self.p ** self.order
         if self.order < 0:
-            norm.denominator = p ** self.order
+            norm.denominator = self.p ** self.order
         return norm
 
 class PAdicConst(PAdic):
-    def __init__(self, value):
-        PAdic.__init__(self)
+    def __init__(self, p, value):
+        PAdic.__init__(self, p)
         value = Fraction(value)
         
         # calculate valuation
@@ -77,21 +75,56 @@ class PAdicConst(PAdic):
         
         self.norm = Fraction(1)
         self.order = 0
-        while not value.numerator % p:
-            self.norm /= p
+        while not value.numerator % self.p:
+            self.norm /= self.p
             self.order += 1
-            value.numerator /= p
-        while not value.denominator % p:
-            self.norm *= p
+            value.numerator /= self.p
+        while not value.denominator % self.p:
+            self.norm *= self.p
             valuation -= 1
-            value.denominator /= p
+            value.denominator /= self.p
         self.value = value
     
     def _nextdigit(self):
         'Calculate next digit of p-adic number.'
-        rem = ModP(p, self.value.numerator) / ModP(p, self.value.denominator)
-        self.val = rem + self.val
+        rem = ModP(self.p, self.value.numerator) / ModP(self.p, self.value.denominator)
         self.value -= rem
-        self.value /= p
+        self.value /= self.p
+        return rem
 
-     
+class PAdicAdd(PAdic):
+    'Sum of two p-adic numbers.'
+    def __init__(self, p, arg1, arg2):
+        PAdic.__init__(self, p)
+        self.carry = 0
+        self.arg1 = arg1
+        self.arg2 = arg2
+    
+    def _nextdigit(self):
+        s = self.arg1.getdigit(self.prec) + self.arg2.getdigit(self.prec) + self.carry
+        self.carry = s // self.p
+        return s % self.p 
+
+class PAdicNeg(PAdic):
+    'Negation of a p-adic number.'
+    def __init__(self, p, arg):
+        PAdic.__init__(self, p)
+        self.arg = arg
+    
+    def _nextdigit(self):
+        if self.prec == 0:
+            return self.p - self.arg.getdigit(0)
+        return self.p - 1 - self.arg.getdigit(0)
+
+class PAdicMul(PAdic):
+    'Product of two p-adic numbers.'
+    def __init__(self, p, arg1, arg2):
+        PAdic.__init__(self, p)
+        self.carry = 0
+        self.arg1 = arg1
+        self.arg2 = arg2
+    
+    def _nextdigit(self):
+        s = sum(self.arg1.getdigit(i) * self.arg2.getdigit(self.prec - i) for i in xrange(self.prec + 1)) + self.carry
+        self.carry = s // self.p
+        return s % self.p