Bugfix in Subexprs to work with indexed subbexpression containing ano…

…ther subexpressions.
cdsousa · Jun 12, 2014 · 86b8ab1 · 86b8ab1
1 parent 12d8100
commit 86b8ab1
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 18 deletions.
diff --git a/sympybotics/dynamics/rne.py b/sympybotics/dynamics/rne.py
@@ -82,7 +82,7 @@ def coriolismatrix(rbtdef, geom, ifunc=None):
         rbtdeftmp.dq[i] = 1
         geomtmp = Geometry(rbtdeftmp)
         fw_results = rne_forward(rbtdeftmp, geomtmp, ifunc)
-        c = rne_backward(rbtdeftmp, geomtmp, fw_results, ifunc=ifunc)
+        c = rne_backward(rbtdeftmp, geomtmp, fw_results, ifunc)
         for k in range(rbtdef.dof):
             A[k][i, i] = c[k]
 
@@ -91,8 +91,8 @@ def coriolismatrix(rbtdef, geom, ifunc=None):
             rbtdeftmp.dq = zeros((rbtdeftmp.dof, 1))
             rbtdeftmp.dq[i] = rbtdeftmp.dq[j] = 1
             geomtmp = Geometry(rbtdeftmp)
-            fw_results = rne_forward(rbtdeftmp, geomtmp)
-            c = rne_backward(rbtdeftmp, geomtmp, fw_results, ifunc=ifunc)
+            fw_results = rne_forward(rbtdeftmp, geomtmp, ifunc)
+            c = rne_backward(rbtdeftmp, geomtmp, fw_results, ifunc)
             for k in range(rbtdef.dof):
                 A[k][i, j] = ifunc(c[k] - A[k][i, i] - A[k][j, j])
 

diff --git a/sympybotics/symcode/subexprs.py b/sympybotics/symcode/subexprs.py
@@ -220,22 +220,29 @@ def _find_repeated_subexprs(subexpr):
         tmpivs_ivs = dict()
         ordered_iv_se = collections.OrderedDict()
 
-        def _get_subexprs(symb):
-            if symb in ivar_se:
-                if symb in tmpivs_ivs:
-                    return tmpivs_ivs[symb]
-                else:
-                    subexpr = ivar_se[symb]
-                    args = list(map(_get_subexprs, subexpr.args))
-                    subexpr = type(subexpr)(*args)
-                    if symb in repeated:
-                        ivar = next(symbols)
-                        ordered_iv_se[ivar] = subexpr
-                        tmpivs_ivs[symb] = ivar
-                        return ivar
+        def _get_subexprs(subexpr):
+            if subexpr.is_Atom:
+                symb = subexpr
+                if symb in ivar_se:
+                    if symb in tmpivs_ivs:
+                        return tmpivs_ivs[symb]
                     else:
-                        return subexpr
-            return symb
+                        subexpr = ivar_se[symb]
+                        args = list(map(_get_subexprs, subexpr.args))
+                        subexpr = type(subexpr)(*args)
+                        if symb in repeated:
+                            ivar = next(symbols)
+                            ordered_iv_se[ivar] = subexpr
+                            tmpivs_ivs[symb] = ivar
+                            return ivar
+                        else:
+                            return subexpr
+                else:
+                    return symb
+            else:
+                args = list(map(_get_subexprs, subexpr.args))
+                subexpr = type(subexpr)(*args)
+                return subexpr
 
         out_exprs = []
         for expr in exprs: