Yes, this is exactly what I wanted, Aaron; many thanks for pointing me in the right direction. As shown in the final example code below, I've done as you've suggested and used atoms(Indexed). I've also taken a poet's license with the code and called expand() to be able to distribute through the negative term, but the results are the same. The output from the updated example code below is:

(A + B)*p[1, 2, 1 + n] + (B - A)*p[1, 1, 1 + n] == -A*p[1, 2, n] - B*p[1, 1, n]



Here's a short comment: using Idx(n+1) seems to work a little better for me, since I've noticed that if I use Idx(n) with a very complicated expression, there is a possibility of terms with (n+1) indices ending up on the same side as the (n) indices. So here is the updated output, and the full modified test program once again, but with the Idx(n+1) change. The output is now:

(A - B)*p[1, 1, 1 + n] + (-A - B)*p[1, 2, 1 + n] == A*p[1, 2, n] + B*p[1, 1, n]

Nicholas


#begin sample code
from sympy import *
from sympy.tensor import Indexed, Idx, IndexedBase
from sympy.matrices import *

p = IndexedBase('p')
var('deltax deltay delta deltat A B')
i,j,n = symbols('i j n', integer = True)

#
# function to generate the FDTD scheme
#
# All p[i,j,n+1] indices should be collected and on the LHS of the expression
def generate_fdtd(expr, nsolve):
    expr0 = expr.lhs - expr.rhs
    expr1 = expr0.expand()
    w1 = Wild('w1',integer=True)
    w2 = Wild('w2',integer=True)
    expr2 = collect(expr1, p[w1, w2, nsolve])
expr3 = expr2.as_independent(*filter(lambda t: t.args[-1] == Idx(n+1), expr2.atoms(Indexed) ) )
    left = expr3[0]
    right = expr3[1]
    left = -1 * expr3[0]
    left = left.expand()
    left = collect(left, p[w1,w2,nsolve])
    right = expr3[1]
    return Eq(right, left)

def main():
expr = Eq( A*p[1,1,n+1] - B*p[1,1,n], B*p[1,1,n+1] + A*p[1,2,n] + B*p[1,2,n+1] + A*p[1,2,n+1])
    print generate_fdtd(expr,n+1)

if __name__ == "__main__":
    main()
# end sample code




--
You received this message because you are subscribed to the Google Groups 
"sympy" group.
To post to this group, send email to [email protected].
To unsubscribe from this group, send email to 
[email protected].
For more options, visit this group at 
http://groups.google.com/group/sympy?hl=en.

Reply via email to