Просмотр исходного кода

multipoly: simplify addition code

narodnik 5 лет назад
Родитель
Сommit
4497eb0f83
1 измененных файлов с 18 добавлено и 8 удалено
  1. 18 8
      scripts/halo/multipoly.py

+ 18 - 8
scripts/halo/multipoly.py

@@ -112,16 +112,26 @@ class MultivariatePolynomial:
             assert isinstance(term, MultivariatePolynomial)
             assert isinstance(term, MultivariatePolynomial)
             result = self.copy()
             result = self.copy()
             for other_term in term.terms:
             for other_term in term.terms:
-                found = False
-                for self_term in result.terms:
-                    if self_term.matches(other_term):
-                        self_term.coeff += other_term.coeff
-                        found = True
-                        break
-                if not found:
-                    result.terms.append(other_term)
+                result += other_term
             return result
             return result
 
 
+    def __mul__(self, other):
+        if isinstance(term, Variable):
+            term = term.termify()
+
+        if hasattr(term, "field"):
+            expr = MultiplyExpression()
+            expr.coeff = term
+            term = expr
+
+        if isinstance(term, MultiplyExpression):
+            # Delete ^0 variables
+            term.clean()
+            return None
+        else:
+            assert isinstance(term, MultivariatePolynomial)
+            return None
+
     def find(self, other):
     def find(self, other):
         for term in self.terms:
         for term in self.terms:
             if term.matches(other):
             if term.matches(other):