p-adic numbers implemented in Python

hensel.py

# Finding roots of polynomials in p-adic integers using Hensel's lemma

from padic import *
from poly import *

def roots(p, poly):
    'Yield all roots of polynomial in the given p-adic integers.'
    for root in xrange(p):
        try:
            yield PAdicPoly(p, poly, root)
        except ValueError:
            pass

class PAdicPoly(PAdic):
    'Result of lifting a root of a polynomial in the integers mod p to the p-adic integers.'
    def __init__(self, p, poly, root):
        PAdic.__init__(self, p)
        self.root = root
        self.poly = poly
        self.deriv = derivative(poly)
        
        # argument checks for the algorithm to work
        if poly(root) % p:
            raise ValueError("%d is not a root of %s modulo %d" % (root, poly, p))
        if self.deriv(root) % p == 0:
            raise ValueError("Polynomial %s is not separable modulo %d" % (poly, p))
        
        # take care of trailing zeros
        digit = self.root
        self.val = str(digit)
        self.exp = self.p
        while digit == 0:
            self.order += 1
            digit = self._nextdigit()
            # self.prec += 1

    def _nextdigit(self):
        self.root = ModP(self.exp * self.p, self.root)
        self.root = self.root - self.poly(self.root) / self.deriv(self.root) # coercions automatically taken care of
        digit = self.root // self.exp
        self.exp *= self.p
        return digit

modp.py

class ModP(int):
    'Integers mod p, p a prime power.'
    def __new__(cls, p, num):
        self = int.__new__(cls, int(num) % p)
        self.p = p
        return self
    
    def __str__(self):
        return "%d (mod %d)" % (self, self.p)
    def __repr__(self):
        return "%d %% %d" % (self, self.p)
    
    # arithmetic
    def __neg__(self):
        return ModP(self.p, self.p - int(self))
    def __add__(self, other):
        return ModP(self.p, int(self) + int(other))
    def __radd__(self, other):
        return ModP(self.p, int(other) + int(self))
    def __sub__(self, other):
        return ModP(self.p, int(self) - int(other))
    def __rsub__(self, other):
        return ModP(self.p, int(other) - int(self))
    def __mul__(self, other):
        return ModP(self.p, int(self) * int(other))
    def __rmul__(self, other):
        return ModP(self.p, int(other) * int(self))
    def __div__(self, other):
        if not isinstance(other, ModP):
            other = ModP(self.p, other)
        return self * other._inv()
    def __rdiv__(self, other):
        return other * self._inv()
    
    def _inv(self):
        'Find multiplicative inverse of self in Z mod p.'
        # extended Euclidean algorithm
        rcurr = self.p
        rnext = int(self)
        tcurr = 0
        tnext = 1
        
        while rnext:
            q = rcurr // rnext
            rcurr, rnext = rnext, rcurr - q * rnext
            tcurr, tnext = tnext, tcurr - q * tnext
        
        if rcurr != 1:
            raise ValueError("%d not a unit modulo %d" % (self, self.p))
        return ModP(self.p, tcurr)

padic.py

from fractions import Fraction
from sys import maxint
from modp import *

class PAdic:
    def __init__(self, p):
        self.p = p
        self.val  = '' # current known value
        self.prec = 0 # current known precision, not containing trailing zeros
        self.order = 0 # order/valuation of number
        pass # initialize object as subclass perhaps
    
    def get(self, prec, decimal = True):
        'Return value of number with given precision.'
        while self.prec < prec:
            # update val based on value
            self.val = str(int(self._nextdigit())) + self.val
            self.prec += 1
        if self.order < 0:
            return (self.val[:self.order] + ('.' if decimal else '') + self.val[self.order:])[-prec-1:]
        return (self.val + self.order * '0')[-prec:]
    
    def _nextdigit(self):
        'Calculate next digit of p-adic number.'
        raise NotImplementedError
    
    def getdigit(self, index):
        'Return digit at given index.'
        return int(self.get(index + 1, False)[0]) #int(self.get(index+1+int(index < -self.order))[-int(index < -self.order)])
    
    # return value with precision up to 32 bits
    def __int__(self):
        return int(self.get(32), self.p)
    def __str__(self):
        return self.get(32)
    
    # arithmetic operations
    def __neg__(self):
        return PAdicNeg(self.p, self)
    def __add__(self, other):
        return PAdicAdd(self.p, self, other)
    def __radd__(self, other):
        return PAdicAdd(self.p, other, self)
    def __sub__(self, other):
        return PAdicAdd(self.p, self, PAdicNeg(self.p, other))
    def __rsub__(self, other):
        return PAdicAdd(self.p, other, PAdicNeg(self.p, self))
    def __mul__(self, other):
        return PAdicMul(self.p, self, other)
    def __rmul__(self, other):
        return PAdicMul(self.p, other, self)
    
    # p-adic norm
    def __abs__(self):
        if self.order == maxint:
            return 0
        numer = denom = 1
        if self.order > 0:
            numer = self.p ** self.order
        if self.order < 0:
            denom = self.p ** self.order
        return Fraction(numer, denom)

class PAdicConst(PAdic):
    def __init__(self, p, value):
        PAdic.__init__(self, p)
        value = Fraction(value)
        
        # calculate valuation
        if value == 0:
            self.value = value
            self.val = '0'
            self.order = maxint
            return
        
        self.order = 0
        while not value.numerator % self.p:
            self.order += 1
            value /= self.p
        while not value.denominator % self.p:
            self.order -= 1
            value *= self.p
        self.value = value
        
        self.zero = not value
    
    def get(self, prec, decimal = True):
        if self.zero:
            return '0' * prec
        return PAdic.get(self, prec, decimal)
    
    def _nextdigit(self):
        'Calculate next digit of p-adic number.'
        rem = ModP(self.p, self.value.numerator) / ModP(self.p, self.value.denominator)
        self.value -= int(rem)
        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
        self.order = self.prec = min(arg1.order, arg2.order) # might be larger than this
        arg1.order -= self.order
        arg2.order -= self.order
        # loop until first nonzero digit is found
        self.index = 0
        digit = self._nextdigit()
        while digit == 0:
            self.index += 1
            self.order += 1
            digit = self._nextdigit()
        self.val += str(int(digit))
        self.prec = 1
    
    def _nextdigit(self):
        s = self.arg1.getdigit(self.index) + self.arg2.getdigit(self.index) + self.carry
        self.carry = s // self.p
        self.index += 1
        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
        self.order = arg.order
    
    def _nextdigit(self):
        if self.prec == 0:
            return self.p - self.arg.getdigit(0) # cannot be p, 0th digit of arg must be nonzero
        return self.p - 1 - self.arg.getdigit(self.prec)

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
        self.order = arg1.order + arg2.order
        self.arg1.order = self.arg2.order = 0 # TODO requires copy
        self.index = 0
    
    def _nextdigit(self):
        s = sum(self.arg1.getdigit(i) * self.arg2.getdigit(self.index - i) for i in xrange(self.index + 1)) + self.carry
        self.carry = s // self.p
        self.index += 1
        return s % self.p 

poly.py

# Polynomial class on Z or Z_p

from collections import defaultdict

class Poly:
    'Polynomial class.'
    def __init__(self, coeffs = None):
        self.coeffs = defaultdict(int, isinstance(coeffs, int) and {0:coeffs} or coeffs or {})
        self.deg = int(len(self.coeffs) and max(self.coeffs.keys()))
    
    def __call__(self, val):
        'Evaluate polynomial for a given value.'
        res = 0
        for i in xrange(self.deg, -1, -1):
            res = res * val + self.coeffs[i]
        return res
    
    def __str__(self):
        def term(coeff, expt):
            if coeff == 1 and expt == 0:
                return '1'
            return ' * '.join(([] if coeff == 1 else [str(coeff)]) + \
                              ([] if expt == 0 else ['X'] if expt == 1 else ['X ** %d' % expt]))
            
        return ' + '.join(term(self.coeffs[i], i) for i in self.coeffs if self.coeffs[i] != 0)
    def __repr__(self):
        return str(self)
    
    # arithmetic
    def __neg__(self):
        return Poly({(i, -self.coeffs[i]) for i in self.coeffs})
    def __add__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        res = Poly()
        res.deg = max(self.deg, other.deg)
        for i in xrange(res.deg+1):
            res.coeffs[i] = self.coeffs[i] + other.coeffs[i]
        return res
    def __radd__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return other.__add__(self)
    def __sub__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return self.__add__(other.__neg__())
    def __rsub__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return other.__add__(self.__neg__())
    
    def __mul__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        res = Poly()
        res.deg = self.deg + other.deg # consider case where other is 0
        for i in xrange(res.deg+1):
            for j in xrange(i+1):
                res.coeffs[i] += self.coeffs[j] * other.coeffs[i - j]
        return res
    def __rmul__(self, other):
        if not isinstance(other, Poly):
            other = Poly(other)
        return other.__mul__(self)
        
    def __pow__(self, other):
        if not isinstance(other, int) or other < 0:
            raise ValueError("Exponent %d is not a natural number" % other)
        res = Poly(1)
        while other:
            res *= self
            other -= 1
        return res

X = Poly({1:1})

def derivative(p):
    'Return derivative of polynomial.'
    return Poly({(i - 1, i * p.coeffs[i]) for i in p.coeffs if i != 0})