def testDescendants(self):
constants = [ops.Constant(i) for i in range(6)]
# (1 + 2*3**4) / 5 - 6
expression = ops.Sub(
ops.Div(
ops.Add(
constants[0],
ops.Mul(
constants[1],
ops.Pow(
constants[2],
constants[3]))),
constants[4]),
constants[5])
descendants = expression.descendants()
descendants = ops._flatten(descendants)
for constant in constants:
self.assertIn(constant, descendants)
self.assertEqual(descendants.count(constant), 1)
# Also test top-level.
self.assertEqual(constants[0].descendants(), [constants[0]])
# Also general structure.
constant = ops.Constant(3)
expression = ops.Neg(constant)
self.assertEqual(set(expression.descendants()), set([constant, expression]))
def monomial(coefficient, variables, powers):
"""Makes a simple monomial term."""
if not isinstance(variables, (list, tuple)):
variables = [variables]
if not isinstance(powers, (list, tuple, np.ndarray)):
powers = [powers]
terms = []
for variable, power in zip(variables, powers):
if power == 0:
continue
elif power == 1:
terms.append(variable)
else:
terms.append(ops.Pow(variable, power))
if (not terms or isinstance(coefficient, sympy.Symbol)
or abs(coefficient) != 1):
if isinstance(coefficient, sympy.Symbol):
terms.insert(0, coefficient)
else:
terms.insert(0, abs(coefficient))
if len(terms) > 1:
term = ops.Mul(*terms)
else:
term = terms[0]
if not isinstance(coefficient, sympy.Symbol) and coefficient < 0:
term = ops.Neg(term)
return term
def testSub(self):
sub = ops.Sub(2, 3)
self.assertEqual(str(sub), '2 - 3')
self.assertEqual(sub.sympy(), -1)
sub = ops.Sub(ops.Sub(1, 2), 3)
self.assertEqual(str(sub), '1 - 2 - 3')
self.assertEqual(sub.sympy(), -4)
sub = ops.Sub(1, ops.Sub(2, 3))
self.assertEqual(str(sub), '1 - (2 - 3)')
self.assertEqual(sub.sympy(), 2)
sub = ops.Sub(ops.Neg(1), 2)
self.assertEqual(str(sub), '-1 - 2')
self.assertEqual(sub.sympy(), -3)
def _make_equals_zero_split(monomials):
"""Returns an `ops.Eq` containing sum of monomials split on left and right."""
left = []
right = []
for monomial in monomials:
if random.choice([False, True]):
left.append(monomial)
else:
right.append(ops.Neg(monomial))
if not left:
left = [0]
if not right:
right = [0]
left = ops.Add(*left)
right = ops.Add(*right)
return ops.Eq(left, right)
def testNeg(self):
op = ops.Neg(2)
self.assertEqual(str(op), '-2')
self.assertEqual(op.sympy(), -2)
op = ops.Add(ops.Neg(2), 3)
self.assertEqual(str(op), '-2 + 3')
self.assertEqual(op.sympy(), 1)
op = ops.Add(3, ops.Neg(2))
self.assertEqual(str(op), '3 - 2')
self.assertEqual(op.sympy(), 1)
op = ops.Add(ops.Add(ops.Neg(2), 5), 3)
self.assertEqual(str(op), '-2 + 5 + 3')
self.assertEqual(op.sympy(), 6)
op = ops.Add(3, ops.Add(ops.Identity(ops.Neg(2)), 5))
self.assertEqual(str(op), '3 - 2 + 5')
self.assertEqual(op.sympy(), 6)
op = ops.Add(3, ops.Add(2, ops.Neg(5)))
self.assertEqual(str(op), '3 + 2 - 5')
self.assertEqual(op.sympy(), 0)
def testNumberConstants(self): constant = ops.Constant(3) expression = ops.Neg(constant) constants = ops.number_constants([expression]) self.assertEqual(constants, [constant])