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.