Symbolic force computation

force.py

# Compute force on a particle with a pair potential
# Uses Lennard-Jones potential

import sympy

# Assigment looks like a relational operator, but probably shouldn't be
from sympy.core.relational import Assignment

import sympy.printing
import numpy as np
from sympy.strategies.traverse import top_down


# Base class for modeling procedures (sequence of expressions)
class procbase:
    def subs(self, old_name, e2):
        var = self.vars[old_name]
        self.stmts = [s.subs(var,e2) for s in self.stmts]
        self.args = [s.subs(var,e2) for s in self.args]
        return self

    def output(self):
        print 'Function ',self.name
        print '  args = ',self.args
        print '  body:'
        for st in self.stmts:
            print '   ',st

    def latex(self):
        print 'Function ',self.name
        for st in self.stmts:
            print '   ',sympy.latex(st)


# Replace a function name with a concrete function body,
#  evaluting the function by substituting parameters
#ftype - function name to replace
#conc -  concrete function object
#          expected members:
#            stmts - list of statments in the function body
#            args - list of function arguments
          
def evalfunctype(ftype, conc):
    def efunc(expr):
        if isinstance(expr, ftype):
            cexpr = None

            sdict = zip(conc.args, expr.args)
            for s in conc.stmts:
                if isinstance(s, Assignment):
                    cexpr = s.rhs.subs(sdict).doit()
                    sdict.append( (s.lhs, cexpr) )
                else:
                    try:
                        cexpr = s.subs(sdict).evalf()
                    except AttributeError:
                        cexpr = s.subs(sdict).doit()
            return cexpr
        return expr
    return efunc



# fix up expression tree - workaround for a bug
def replacematadd(expr):
    #print 'expr = ',type(expr)
    if isinstance(expr, sympy.Add):
        args_are_mat = True
        for a in expr.args:
            args_are_mat =  args_are_mat and  isinstance(a, sympy.MatrixExpr)
        #print 'args are mat',args_are_mat
        if args_are_mat:
            return sympy.MatAdd(*expr.args)
    if isinstance(expr, sympy.Mul):
        e = sympy.Mul(*expr.args)

    return expr


def lj(r):
    return 4*(1/r**12 - 1/r**6)

class lj_force_sym(procbase):
    def __init__(self):
        r = sympy.Symbol('r') # scalar
        s = - sympy.diff(lj(r), r)
        self.args = [r]
        self.stmts = [s]
        self.vars = {'r':r}
        self.name = 'LJ force'


class pair_pot_sym(procbase):
    def __init__(self):
        n = sympy.Symbol('n_d')
        r2 = sympy.MatrixSymbol('r2',n,1)  #vector
        r1 = sympy.MatrixSymbol('r1',n,1)  # vector
        d = sympy.MatrixSymbol('d',n,1)     # vector

        s1 = sympy.Assign(d, r2 - r1)   # can we type this?
        r = sympy.Symbol('r') # scalar
        s2 = sympy.Assign(r, abs(d))
        V2 = sympy.Symbol('V2')
        ret = V2(r)
        self.args = [r1, r2]
        self.stmts = [s1, s2, ret]
        self.vars = {'r1':r1,'r2':r2,'n_d':n,'V2':V2}
        self.name = 'pair_potential'


class pair_force_sym(procbase):
    def __init__(self):
        n = sympy.Symbol('n')
        r2 = sympy.MatrixSymbol('r2',n,1)  #vector
        r1 = sympy.MatrixSymbol('r1',n,1)  # vector
        d = sympy.MatrixSymbol('d',n,1)     # vector

        s1 = sympy.Assign(d, r2 - r1)   # can we type this?
        r = sympy.Symbol('r') # scalar
        i = sympy.Symbol('i')
        s2 = sympy.Assign(r, abs(d))
        F2 = sympy.Symbol('F2')
        ret = d * F2(r)/r

        self.args = [r1, r2]
        self.stmts = [s1, s2, ret]
        self.vars = {'r1':r1,'r2':r2,'n_d':n,'F2':F2}
        self.name = 'Pair force'


class total_force_sym(procbase):
    def __init__(self):
        n = sympy.Symbol('n')
        F = sympy.Symbol('F')
        i = sympy.Symbol('i')
        j = sympy.Symbol('j')
        r = sympy.IndexedBase('r')
        s = sympy.Sum(F(r[i], r[j]), (j,1,n), constraint=sympy.Ne(i,j))
        self.args = [r, i, n]
        self.stmts = [s]
        self.vars = {'n':n, 'F':F, 'i': i, 'r':r}
        self.name = 'Force on particle i'

# Particle positions

r1 = sympy.Matrix([1.1, 2.3, 4.4])
r2 = sympy.Matrix([0.1, 2.1, 1.4])
r3 = sympy.Matrix([0.2, 1.5, 1.4])
pos = [r1,r2,r3]



# Force
total_force = total_force_sym()
tvars = total_force.vars

total_force.output()

# Specialize to number of of particles

s2 = total_force.subs('n', len(pos)).stmts[0]

# Compute force on particle 1
target_i = 1
s2 = s2.subs('i', target_i)

# evaluate the sum
s2 = s2.doit_direct()
print ''
print 'Specialized force to 3 particles, on particle',target_i
print s2
print ''

s3 = s2.subs({tvars['r'][1]:pos[0], tvars['r'][2]:pos[1], tvars['r'][3]:pos[2]})

pp = pair_force_sym()
pp.output()
pp = pp.subs('n_d', 3)

pp_eval = evalfunctype(sympy.Function('F'), pp)
pp2  = top_down(pp_eval)(s3)


# workaround for a bug - some tree nodes are in an unexpected state after replacement
#print ''
#print 'replacing matadd'

pp2 = top_down(replacematadd)(pp2)

print ''
print 'After replacing pair force'
print pp2
print ''


ljp = lj_force_sym()
print ''
ljp.output()
lj_eval = evalfunctype(sympy.Function('F2'), ljp)
lj2 = top_down(lj_eval)(pp2)

print ''
print 'After replacing LJ force:'
print lj2
print ''

print 'Final answer:'
print lj2.doit()

output.txt

Function  Force on particle i
  args =  [r, i, n]
  body:
    Sum(F(r[i], r[j]), (j, 1, n))

Specialized force to 3 particles, on particle 1
F(r[1], r[2]) + F(r[1], r[3])

Function  Pair force
  args =  [r1, r2]
  body:
    d = (-1)*r1 + r2
    r = MatrixNorm(d)
    (F2(r)/r)*d

After replacing pair force
(0.315597201548902*F2(3.16859590355097))*Matrix([
[-1.0],
[-0.2],
[-3.0]]) + (0.309344112444873*F2(3.23264597504892))*Matrix([
[-0.9],
[-0.8],
[-3.0]])


Function  LJ force
  args =  [r]
  body:
    -24/r**7 + 48/r**13

After replacing LJ force:
((-0.00746937302658855)*0.315597201548902)*Matrix([
[-1.0],
[-0.2],
[-3.0]]) + ((-0.00649445057204655)*0.309344112444873)*Matrix([
[-0.9],
[-0.8],
[-3.0]])

Final answer:
Matrix([
[0.00416543126774035],
[0.00207867868332471],
[ 0.0130989998176291]])