def test_collect_func():
f = ((x + a + 1)**3).expand()
assert collect(f, x) == a**3 + 3*a**2 + 3*a + x**3 + x**2*(3*a + 3) + x*(3*a**2 + 6*a + 3) + 1
assert collect(f, x, factor) == x**3 + 3*x**2*(a + 1) + 3*x*(a + 1)**2 + (a + 1)**3
assert collect(f, x, evaluate=False) == {S.One: a**3 + 3*a**2 + 3*a + 1, x: 3*a**2 + 6*a + 3, x**2: 3*a + 3, x**3: 1}
def test_collect_4():
"""Collect with respect to a power"""
a, b, c, x = symbols('a,b,c,x')
assert collect(a*x**c + b*x**c, x**c) == x**c*(a + b)
# issue 6096: 2 stays with c (unless c is integer or x is positive0
assert collect(a*x**(2*c) + b*x**(2*c), x**c) == x**(2*c)*(a + b)
def test_collect_5():
"""Collect with respect to a tuple"""
a, x, y, z, n = symbols('axyzn')
assert collect(x**2*y**4 + z*(x*y**2)**2 + z + a*z, [x*y**2, z]) in [
z*(1 + a + x**2*y**4) + x**2*y**4,
z*(1 + a) + x**2*y**4*(1 + z) ]
assert collect((1+ (x+y) + (x+y)**2).expand(), [x,y]) == 1 + y + x*(1 + 2*y) + x**2 + y**2
def __init__(self, m, **kwargs):
(a, b), (c, d) = m
a = cancel(a)
b = collect(b, ohm)
c = collect(c, 1/ohm)
d = cancel(d)
super(Transmission, self).__init__([[a, b], [c, d]], **kwargs)
if self.shape != (2, 2):
raise ValueError("Transmission Matrix must be 2x2")
def test_collect_D():
D = Derivative
f = Function('f')
x,a,b = symbols('xab')
fx = D(f(x), x)
fxx = D(f(x), x,x)
assert collect(a*fx + b*fx, fx) == (a + b)*fx
assert collect(a*D(fx,x) + b*D(fx,x), fx) == (a + b)*D(fx, x)
assert collect(a*fxx + b*fxx , fx) == (a + b)*D(fx, x)
def calcTF(A, B, C, n):
s = sympy.Symbol('s')
I = sympy.eye(n) * s
symA = makeSymMat(A)
#sysMat = I - symA
#sysMat = sysMat.inv()
#print 'sysMat'
#print sysMat
charEqn = (I - symA).det()
#print 'determinant'
#print sympy.collect(sympy.expand(sympy.simplify(charEqn)), s, evaluate=False)
myCharEqn = charEqn
myTFNum = C * (I - symA).adjugate() * B
#print 'My characteristic eqn'
#print myCharEqn
#print 'numerator'
#print myTFNum
alphaDicts = []
betaDicts = []
#print 'rows'
for elem in myTFNum:
#print elem
#newEqn = sympy.simplify(sympy.expand(elem))
#newEqn = sympy.ratsimp(newEqn)
#tfFrac = sympy.fraction(newEqn)
elem = sympy.expand(sympy.simplify(elem))
#alphas are the denominator
currDictAlpha = sympy.collect(myCharEqn, s, evaluate=False)
if len(currDictAlpha) > 0:
alphaDicts.append(currDictAlpha)
#betas are the numerator
currDictBeta = sympy.collect(elem, s, evaluate=False)
if len(currDictBeta) > 0:
betaDicts.append(currDictBeta)
#print 'Transfer Function'
#print myTF
#print 'betas'
#print betaDicts
myBetas = betaDicts
myAlphas = alphaDicts
#print 'alphas'
#print alphaDicts
#print 'done'
return(myTFNum, myCharEqn, myAlphas, myBetas)
def test_issue_13143():
fx = f(x).diff(x)
e = f(x) + fx + f(x)*fx
# collect function before derivative
assert collect(e, Wild('w')) == f(x)*(fx + 1) + fx
e = f(x) + f(x)*fx + x*fx*f(x)
assert collect(e, fx) == (x*f(x) + f(x))*fx + f(x)
assert collect(e, f(x)) == (x*fx + fx + 1)*f(x)
e = f(x) + fx + f(x)*fx
assert collect(e, [f(x), fx]) == f(x)*(1 + fx) + fx
assert collect(e, [fx, f(x)]) == fx*(1 + f(x)) + f(x)
def op(self, op_name, *args):
print "RUNNING %s"% op_name, len(self.terms)
new_terms = dict()
for key in self.terms:
X = self.terms[key]
r_val = getattr(X, op_name)(*args)
new_key = str(X)
new_terms[new_key] = X.copy()
old_val = sympy.Symbol(key)
new_val = sympy.Symbol(new_key)
if(op_name=="delete" and r_val):
new_coeff = sympy.Symbol(r_val)
self.val = self.val.subs(old_val,old_val*new_coeff)
self.val = self.val.subs(old_val, new_val)
self.terms = new_terms
# Collect like terms
for X in self.terms:
self.val = sympy.collect(self.val, sympy.Symbol(X))
def exponential_euler_code(self):
'''
Generates Python code for an exponential Euler step.
Not efficient for the moment!
'''
all_variables = self._eq_names + self._diffeq_names + self._alias.keys() + ['t']
vars_tmp = [name + '__tmp' for name in self._diffeq_names]
lines = ','.join(self._diffeq_names) + '=P._S\n'
lines += ','.join(vars_tmp) + '=P._dS\n'
for name in self._diffeq_names:
# Freeze
namespace = self._namespace[name]
expr = optimiser.freeze(self._string[name], all_variables, namespace)
# Find a and b in dx/dt=a*x+b
sym_expr = symbolic_eval(expr)
if isinstance(sym_expr, float):
lines += name + '__tmp[:]=' + name + '+(' + str(expr) + ')*dt\n'
else:
sym_expr = sym_expr.expand()
sname = sympy.Symbol(name)
terms = sympy.collect(sym_expr, name, evaluate=False)
if sname ** 0 in terms:
b = terms[sname ** 0]
else:
b = 0
if sname in terms:
a = terms[sname]
else:
a = 0
lines += name + '__tmp[:]=' + str(-b / a + (sname + b / a) * sympy.exp(a * sympy.Symbol('dt'))) + '\n'
lines += 'P._S[:]=P._dS'
#print lines
return compile(lines, 'Exponential Euler update code', 'exec')
def test_collect_D_0():
D = Derivative
f = Function('f')
x, a, b = symbols('x,a,b')
fxx = D(f(x), x, x)
assert collect(a*fxx + b*fxx, fxx) == (a + b)*fxx
def __init__(self):
h, g, h0, g0, a, d = sympy.symbols('h g h0 g0 a d')
#g1 = sympy.Max(-d, g0 + h0 - 1/(4*a))
g1 = g0 + h0 - 1/(4*a)
h1 = h0 - 1/(2*a)
parabola = d - a*h*h # =g on boundary
slope = g1 + h - h1 # =g on line with slope=1 through z1
h2 = sympy.solve(parabola-slope,h)[0] # First solution is above z1
g2 = h2 - h1 + g1 # Line has slope of 1
g3 = g1
h3 = sympy.solve((parabola-g).subs(g, g1),h)[1] # Second is on right
r_a = sympy.Rational(1,24) # a=1/24 always
self.h = tuple(x.subs(a,r_a) for x in (h0, h1, h2, h3))
self.g = tuple(x.subs(a,r_a) for x in (g0, g1, g2, g3))
def integrate(f):
ia = sympy.integrate(
sympy.integrate(f,(g, g1, g1+h-h1)),
(h, h1, h2)) # Integral of f over right triangle
ib = sympy.integrate(
sympy.integrate(f,(g, g1, parabola)),
(h, h2, h3)) # Integral of f over region against parabola
return (ia+ib).subs(a, r_a)
i0 = integrate(1) # Area = integral of pie slice
E = lambda f:(integrate(f)/i0) # Expected value wrt Lebesgue measure
sigma = lambda f,g:E(f*g) - E(f)*E(g)
self.d = d
self.Eh = collect(E(h),sympy.sqrt(d-g0-h0+6))
self.Eg = E(g)
self.Sigmahh = sigma(h,h)
self.Sigmahg = sigma(h,g)
self.Sigmagg = sigma(g,g)
return
def sympy_polynomial_to_function(expr, symbs,
func_name='sympy_polynomial',
return_string=False):
poly = sp.poly(sp.collect(sp.expand(expr), symbs), *symbs)
input_names = map(attrgetter('name'), poly.gens)
num_names = len(input_names)
def term_string((powers, coeff)):
indices = filter(lambda i: powers[i] != 0,
range(num_names))
if indices:
terms = map(lambda i: '(%s ** %s)' % (input_names[i],
powers[i]),
indices)
return '%s * %s' % (float(coeff), ' * '.join(terms))
else:
return '%s' % float(coeff)
poly_str = ' + '.join(map(term_string, poly.terms()))
function_str = 'def %s(%s):\n return %s' % (
func_name, ', '.join(input_names), poly_str)
if return_string:
return function_str
exec_env = {}
exec function_str in exec_env
return exec_env[func_name]
def coeff(expr, term):
expr = sp.collect(expr, term)
symbols = list(term.atoms(sp.Symbol))
w = sp.Wild("coeff", exclude=symbols)
m = expr.match(w * term + sp.Wild("rest"))
if m:
return m[w]
def test_gcd_terms():
f = 2*(x + 1)*(x + 4)/(5*x**2 + 5) + (2*x + 2)*(x + 5)/(x**2 + 1)/5 + (2*x + 2)*(x + 6)/(5*x**2 + 5)
assert _gcd_terms(f) == ((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)
assert _gcd_terms(Add.make_args(f)) == ((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)
assert gcd_terms(f) == (S(6)/5)*((1 + x)*(5 + x)/(1 + x**2))
assert gcd_terms(Add.make_args(f)) == (S(6)/5)*((1 + x)*(5 + x)/(1 + x**2))
assert gcd_terms((2*x + 2)**3 + (2*x + 2)**2) == 4*(x + 1)**2*(2*x + 3)
assert gcd_terms(0) == 0
assert gcd_terms(1) == 1
assert gcd_terms(x) == x
assert gcd_terms(2 + 2*x) == Mul(2, 1 + x, evaluate=False)
arg = x*(2*x + 4*y)
garg = 2*x*(x + 2*y)
assert gcd_terms(arg) == garg
assert gcd_terms(sin(arg)) == sin(garg)
# issue 3040-like
alpha, alpha1, alpha2, alpha3 = symbols('alpha:4')
a = alpha**2 - alpha*x**2 + alpha + x**3 - x*(alpha + 1)
rep = (alpha, (1 + sqrt(5))/2 + alpha1*x + alpha2*x**2 + alpha3*x**3)
s = (a/(x - alpha)).subs(*rep).series(x, 0, 1)
assert simplify(collect(s, x)) == -sqrt(5)/2 - S(3)/2 + O(x)
# issue 2818
assert _gcd_terms([S.Zero, S.Zero]) == (0, 0, 1)
assert _gcd_terms([2*x + 4]) == (2, x + 2, 1)
def collect_into_dict_include_zero_and_constant_terms(expr, syms):
expr_terms_dict = sympy.collect(expr,syms,exact=True,evaluate=False)
for sym in syms:
if sym not in expr_terms_dict.keys(): expr_terms_dict[sym] = 0
if 1 not in expr_terms_dict: expr_terms_dict[1] = 0
return expr_terms_dict
def run(expr):
# Return semantic (rebuilt expression, factorization candidates)
if expr.is_Number or expr.is_Symbol:
return expr, [expr]
elif expr.is_Indexed or expr.is_Atom:
return expr, []
elif expr.is_Add:
rebuilt, candidates = zip(*[run(arg) for arg in expr.args])
w_numbers = [i for i in rebuilt if any(j.is_Number for j in i.args)]
wo_numbers = [i for i in rebuilt if i not in w_numbers]
w_numbers = collect_const(expr.func(*w_numbers))
wo_numbers = expr.func(*wo_numbers)
if aggressive is True and wo_numbers:
for i in flatten(candidates):
wo_numbers = collect(wo_numbers, i)
rebuilt = expr.func(w_numbers, wo_numbers)
return rebuilt, []
elif expr.is_Mul:
rebuilt, candidates = zip(*[run(arg) for arg in expr.args])
rebuilt = collect_const(expr.func(*rebuilt))
return rebuilt, flatten(candidates)
elif expr.is_Equality:
rebuilt, candidates = zip(*[run(expr.lhs), run(expr.rhs)])
return expr.func(*rebuilt, evaluate=False), flatten(candidates)
else:
rebuilt, candidates = zip(*[run(arg) for arg in expr.args])
return expr.func(*rebuilt), flatten(candidates)
def residue(expr, x, x0):
"""
Finds the residue of ``expr`` at the point x=x0.
The residue is defined as the coefficient of 1/(x-x0) in the power series
expansion about x=x0.
Examples:
>>> from sympy import Symbol, residue, sin
>>> x = Symbol("x")
>>> residue(1/x, x, 0)
1
>>> residue(1/x**2, x, 0)
0
>>> residue(2/sin(x), x, 0)
2
This function is essential for the Residue Theorem [1].
The current implementation uses series expansion to calculate it. A more
general implementation is explained in the section 5.6 of the Bronstein's
book [2]. For purely rational functions, the algorithm is much easier. See
sections 2.4, 2.5, and 2.7 (this section actually gives an algorithm for
computing any Laurent series coefficient for a rational function). The
theory in section 2.4 will help to understand why the resultant works in
the general algorithm. For the definition of a resultant, see section 1.4
(and any previous sections for more review).
[1] http://en.wikipedia.org/wiki/Residue_theorem
[2] M. Bronstein: Symbolic Integration I, Springer Verlag (2005)
"""
from sympy import collect
expr = sympify(expr)
if x0 != 0:
expr = expr.subs(x, x + x0)
s = expr.series(x, 0, 0).removeO()
# TODO: this sometimes helps, but the series expansion should rather be
# fixed, see #1627:
if s == 0:
s = expr.series(x, 0, 1).removeO()
if s == 0:
s = expr.series(x, 0, 6).removeO()
s = collect(s, x)
if x0 != 0:
s = s.subs(x, x - x0)
a = Wild("r", exclude=[x])
c = Wild("c", exclude=[x])
r = s.match(a / (x - x0) + c)
if r:
return r[a]
elif isinstance(s, Add):
# TODO: this is to overcome a bug in match (#1626)
for t in s.args:
r = t.match(a / (x - x0) + c)
if r:
return r[a]
return Integer(0)
def get_conditionally_linear_system(eqs):
'''
Convert equations into a linear system using sympy.
Parameters
----------
eqs : `Equations`
The model equations.
Returns
-------
coefficients : dict of (sympy expression, sympy expression) tuples
For every variable x, a tuple (M, B) containing the coefficients M and
B (as sympy expressions) for M * x + B
Raises
------
ValueError
If one of the equations cannot be converted into a M * x + B form.
Examples
--------
>>> from brian2 import Equations
>>> eqs = Equations("""
... dv/dt = (-v + w**2) / tau : 1
... dw/dt = -w / tau : 1
... """)
>>> system = get_conditionally_linear_system(eqs)
>>> print(system['v'])
(-1/tau, w**2.0/tau)
>>> print(system['w'])
(-1/tau, 0)
'''
diff_eqs = eqs.substituted_expressions
coefficients = {}
for name, expr in diff_eqs:
var = sp.Symbol(name, real=True)
# Coefficients
wildcard = sp.Wild('_A', exclude=[var])
#Additive constant
constant_wildcard = sp.Wild('_B', exclude=[var])
pattern = wildcard*var + constant_wildcard
# Factor out the variable
s_expr = sp.collect(expr.sympy_expr.expand(), var)
matches = s_expr.match(pattern)
if matches is None:
raise ValueError(('The expression "%s", defining the variable %s, '
'could not be separated into linear components') %
(s_expr, name))
coefficients[name] = (matches[wildcard].simplify(),
matches[constant_wildcard].simplify())
return coefficients
def list_collect_v1(row, listin):
rowout = copy.copy(row)
for c, item in enumerate(rowout):
temp = item
for var in listin:
temp = sympy.collect(temp, var)
rowout[c] = temp
return rowout
def test_collect_func_xfail():
# XXX: this test will pass when automatic constant distribution is removed (#1497)
assert collect(f, x, factor, evaluate=False) == {
S.One: (a + 1) ** 3,
x: 3 * (a + 1) ** 2,
x ** 2: 3 * (a + 1),
x ** 3: 1,
}
def test_collect_D_0():
D = Derivative
f = Function('f')
x, a, b = symbols('x,a,b')
fxx = D(f(x), x, x)
# collect does not distinguish nested derivatives, so it returns
# -- (a + b)*D(D(f, x), x)
assert collect(a*fxx + b*fxx, fxx) == (a + b)*fxx
def test_gcd_terms():
f = 2*(x + 1)*(x + 4)/(5*x**2 + 5) + (2*x + 2)*(x + 5)/(x**2 + 1)/5 + \
(2*x + 2)*(x + 6)/(5*x**2 + 5)
assert _gcd_terms(f) == ((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)
assert _gcd_terms(Add.make_args(f)) == \
((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)
newf = (S(6)/5)*((1 + x)*(5 + x)/(1 + x**2))
assert gcd_terms(f) == newf
args = Add.make_args(f)
# non-Basic sequences of terms treated as terms of Add
assert gcd_terms(list(args)) == newf
assert gcd_terms(tuple(args)) == newf
assert gcd_terms(set(args)) == newf
# but a Basic sequence is treated as a container
assert gcd_terms(Tuple(*args)) != newf
assert gcd_terms(Basic(Tuple(1, 3*y + 3*x*y), Tuple(1, 3))) == \
Basic((1, 3*y*(x + 1)), (1, 3))
# but we shouldn't change keys of a dictionary or some may be lost
assert gcd_terms(Dict((x*(1 + y), 2), (x + x*y, y + x*y))) == \
Dict({x*(y + 1): 2, x + x*y: y*(1 + x)})
assert gcd_terms((2*x + 2)**3 + (2*x + 2)**2) == 4*(x + 1)**2*(2*x + 3)
assert gcd_terms(0) == 0
assert gcd_terms(1) == 1
assert gcd_terms(x) == x
assert gcd_terms(2 + 2*x) == Mul(2, 1 + x, evaluate=False)
arg = x*(2*x + 4*y)
garg = 2*x*(x + 2*y)
assert gcd_terms(arg) == garg
assert gcd_terms(sin(arg)) == sin(garg)
# issue 6139-like
alpha, alpha1, alpha2, alpha3 = symbols('alpha:4')
a = alpha**2 - alpha*x**2 + alpha + x**3 - x*(alpha + 1)
rep = (alpha, (1 + sqrt(5))/2 + alpha1*x + alpha2*x**2 + alpha3*x**3)
s = (a/(x - alpha)).subs(*rep).series(x, 0, 1)
assert simplify(collect(s, x)) == -sqrt(5)/2 - S(3)/2 + O(x)
# issue 5917
assert _gcd_terms([S.Zero, S.Zero]) == (0, 0, 1)
assert _gcd_terms([2*x + 4]) == (2, x + 2, 1)
eq = x/(x + 1/x)
assert gcd_terms(eq, fraction=False) == eq
eq = x/2/y + 1/x/y
assert gcd_terms(eq, fraction=True, clear=True) == \
(x**2 + 2)/(2*x*y)
assert gcd_terms(eq, fraction=True, clear=False) == \
(x**2/2 + 1)/(x*y)
assert gcd_terms(eq, fraction=False, clear=True) == \
(x + 2/x)/(2*y)
assert gcd_terms(eq, fraction=False, clear=False) == \
(x/2 + 1/x)/y
def test_collect_Wild():
"""Collect with respect to functions with Wild argument"""
a, b, x, y = symbols('a b x y')
f = Function('f')
w1 = Wild('.1')
w2 = Wild('.2')
assert collect(f(x) + a*f(x), f(w1)) == (1 + a)*f(x)
assert collect(f(x, y) + a*f(x, y), f(w1)) == f(x, y) + a*f(x, y)
assert collect(f(x, y) + a*f(x, y), f(w1, w2)) == (1 + a)*f(x, y)
assert collect(f(x, y) + a*f(x, y), f(w1, w1)) == f(x, y) + a*f(x, y)
assert collect(f(x, x) + a*f(x, x), f(w1, w1)) == (1 + a)*f(x, x)
assert collect(a*(x + 1)**y + (x + 1)**y, w1**y) == (1 + a)*(x + 1)**y
assert collect(a*(x + 1)**y + (x + 1)**y, w1**b) == a*(x + 1)**y + (x + 1)**y
assert collect(a*(x + 1)**y + (x + 1)**y, (x + 1)**w2) == (1 + a)*(x + 1)**y
assert collect(a*(x + 1)**y + (x + 1)**y, w1**w2) == (1 + a)*(x + 1)**y
def mycollectsimp(expr):
from sympy import collect
counts = {}
for subtree in preorder_traversal(expr):
if isinstance(subtree, Symbol):
counts[subtree] = counts.get(subtree, 0)+1
counts = [(count, symbol) for symbol, count in counts.iteritems()]
counts.sort()
for count, symbol in counts:
expr = collect(expr, symbol)
return expr
def test_collect_Wild():
"""Collect with respect to functions with Wild argument"""
a, b, x, y = symbols("a b x y")
f = Function("f")
w1 = Wild(".1")
w2 = Wild(".2")
assert collect(f(x) + a * f(x), f(w1)) == (1 + a) * f(x)
assert collect(f(x, y) + a * f(x, y), f(w1)) == f(x, y) + a * f(x, y)
assert collect(f(x, y) + a * f(x, y), f(w1, w2)) == (1 + a) * f(x, y)
assert collect(f(x, y) + a * f(x, y), f(w1, w1)) == f(x, y) + a * f(x, y)
assert collect(f(x, x) + a * f(x, x), f(w1, w1)) == (1 + a) * f(x, x)
assert collect(a * (x + 1)**y + (x + 1)**y, w1**y) == (1 + a) * (x + 1)**y
assert (collect(a * (x + 1)**y + (x + 1)**y,
w1**b) == a * (x + 1)**y + (x + 1)**y)
assert (collect(a * (x + 1)**y + (x + 1)**y,
(x + 1)**w2) == (1 + a) * (x + 1)**y)
assert collect(a * (x + 1)**y + (x + 1)**y, w1**w2) == (1 + a) * (x + 1)**y
def test_collect_1():
"""Collect with respect to a Symbol"""
x, y, z, n = symbols('x,y,z,n')
assert collect(1, x) == 1
assert collect( x + y*x, x ) == x * (1 + y)
assert collect( x + x**2, x ) == x + x**2
assert collect( x**2 + y*x**2, x ) == (x**2)*(1 + y)
assert collect( x**2 + y*x, x ) == x*y + x**2
assert collect( 2*x**2 + y*x**2 + 3*x*y, [x] ) == x**2*(2 + y) + 3*x*y
assert collect( 2*x**2 + y*x**2 + 3*x*y, [y] ) == 2*x**2 + y*(x**2 + 3*x)
assert collect( ((1 + y + x)**4).expand(), x) == ((1 + y)**4).expand() + \
x*(4*(1 + y)**3).expand() + x**2*(6*(1 + y)**2).expand() + \
x**3*(4*(1 + y)).expand() + x**4
# symbols can be given as any iterable
expr = x + y
assert collect(expr, expr.free_symbols) == expr
def test_collect_Wild():
"""Collect with respect to functions with Wild argument"""
a, b, x, y = symbols('a b x y')
f = Function('f')
w1 = Wild('.1')
w2 = Wild('.2')
assert collect(f(x) + a * f(x), f(w1)) == (1 + a) * f(x)
assert collect(f(x, y) + a * f(x, y), f(w1)) == f(x, y) + a * f(x, y)
assert collect(f(x, y) + a * f(x, y), f(w1, w2)) == (1 + a) * f(x, y)
assert collect(f(x, y) + a * f(x, y), f(w1, w1)) == f(x, y) + a * f(x, y)
assert collect(f(x, x) + a * f(x, x), f(w1, w1)) == (1 + a) * f(x, x)
assert collect(a * (x + 1)**y + (x + 1)**y, w1**y) == (1 + a) * (x + 1)**y
assert collect(a*(x + 1)**y + (x + 1)**y, w1**b) == \
a*(x + 1)**y + (x + 1)**y
assert collect(a*(x + 1)**y + (x + 1)**y, (x + 1)**w2) == \
(1 + a)*(x + 1)**y
assert collect(a * (x + 1)**y + (x + 1)**y, w1**w2) == (1 + a) * (x + 1)**y
def test_collect_1():
"""Collect with respect to a Symbol"""
x, y, z, n = symbols('x,y,z,n')
assert collect(1, x) == 1
assert collect( x + y*x, x ) == x * (1 + y)
assert collect( x + x**2, x ) == x + x**2
assert collect( x**2 + y*x**2, x ) == (x**2)*(1 + y)
assert collect( x**2 + y*x, x ) == x*y + x**2
assert collect( 2*x**2 + y*x**2 + 3*x*y, [x] ) == x**2*(2 + y) + 3*x*y
assert collect( 2*x**2 + y*x**2 + 3*x*y, [y] ) == 2*x**2 + y*(x**2 + 3*x)
assert collect( ((1 + y + x)**4).expand(), x) == ((1 + y)**4).expand() + \
x*(4*(1 + y)**3).expand() + x**2*(6*(1 + y)**2).expand() + \
x**3*(4*(1 + y)).expand() + x**4
# symbols can be given as any iterable
expr = x + y
assert collect(expr, expr.free_symbols) == expr
def __sub__(self, LT):
if self.Ga.name != LT.Ga.name:
raise ValueError("Attempting subtraction of Lt's from different geometric algebras")
self_add_LT = copy(self.lt_dict)
for key in LT.lt_dict.keys():
if key in self_add_LT:
self_add_LT[key] = collect(self_add_LT[key] - LT.lt_dict[key], self.Ga.basis)
else:
self_add_LT[key] = -LT.lt_dict[key]
return(Lt(self_add_LT, ga=self.Ga))
def test_collect_func():
f = ((x + a + 1)**3).expand()
assert collect(f, x) == a**3 + 3*a**2 + 3*a + x**3 + x**2*(3*a + 3) + \
x*(3*a**2 + 6*a + 3) + 1
assert collect(f, x, factor) == x**3 + 3*x**2*(a + 1) + 3*x*(a + 1)**2 + \
(a + 1)**3
assert collect(f, x, evaluate=False) == {
S.One: a**3 + 3 * a**2 + 3 * a + 1,
x: 3 * a**2 + 6 * a + 3,
x**2: 3 * a + 3,
x**3: 1
}
assert collect(f, x, factor, evaluate=False) == {
S.One: (a + 1)**3,
x: 3 * (a + 1)**2,
x**2: umul(S(3), a + 1),
x**3: 1
}
def collect_square(f, term, square):
'''collect_square(x**5+x**2, x, XX) == x*XX**2+XX'''
g = collect(f, term)
h = Poly(g, term)
s = zero
for (e, ), coeff in h.all_terms():
t = square**(e // 2)
if e % 2:
t *= term
s += coeff * t
return s
def __sub__(self, LT):
if self.Ga.name != LT.Ga.name:
raise ValueError("Attempting subtraction of Lt's from different geometric algebras")
self_add_LT = copy(self.lt_dict)
for key in list(LT.lt_dict.keys()):
if key in self_add_LT:
self_add_LT[key] = collect(self_add_LT[key] - LT.lt_dict[key], self.Ga.basis)
else:
self_add_LT[key] = -LT.lt_dict[key]
return(Lt(self_add_LT, ga=self.Ga))
def collects(expression):
'''
This function will rearrange the
expression according to the variables
powers and factored out the coefficient.
The expression MUST be expanded before hand.
Collects CAN be used before coeff method of
an expression
'''
return collect(expression)
def get_conditionally_linear_system(eqs, variables=None):
'''
Convert equations into a linear system using sympy.
Parameters
----------
eqs : `Equations`
The model equations.
Returns
-------
coefficients : dict of (sympy expression, sympy expression) tuples
For every variable x, a tuple (M, B) containing the coefficients M and
B (as sympy expressions) for M * x + B
Raises
------
ValueError
If one of the equations cannot be converted into a M * x + B form.
Examples
--------
>>> from brian2 import Equations
>>> eqs = Equations("""
... dv/dt = (-v + w**2) / tau : 1
... dw/dt = -w / tau : 1
... """)
>>> system = get_conditionally_linear_system(eqs)
>>> print(system['v'])
(-1/tau, w**2.0/tau)
>>> print(system['w'])
(-1/tau, 0)
'''
diff_eqs = eqs.get_substituted_expressions(variables)
coefficients = {}
for name, expr in diff_eqs:
var = sp.Symbol(name, real=True)
# Factor out the variable
s_expr = sp.collect(str_to_sympy(expr.code, variables).expand(),
var, evaluate=False)
if len(s_expr) > 2 or var not in s_expr:
raise ValueError(('The expression "%s", defining the variable %s, '
'could not be separated into linear components') %
(expr, name))
coefficients[name] = (s_expr[var], s_expr.get(1, 0))
return coefficients
def derive_euler_equation_functions(M):
'''Output Python functions dphi/dt and dzeta/dt for M-order
Euler eqs. these functions' input arguments are:
w0, ... , w(M-1), gradphi^2, gradzeta^2, zeta, gradzeta*gradphi,
time, nonlinear_damping_coefficient
taken from west and the zakharov fromultation.'''
# define all symbols
eps = sy.Symbol('eps') # smallness paramater
zeta = sy.Symbol('zeta')
gradphi2 = sy.Symbol('gradphi2')
gradzeta2 = sy.Symbol('gradzeta2')
gradzeta_gradphi = sy.Symbol('gradzeta_gradphi')
time = sy.Symbol('t')
damping_factor = sy.Symbol('ta')
damping= 1-sy.exp(-(time/damping_factor)**2)
# gives order to the list of w(M), w(M) has order M+1
w_list = [None]*M
eps_list = [None]*M
for m in range(M):
w_list[m] = sy.Symbol('w_' + str(m))
eps_list[m] = eps**(m+1)
# make these arrays to allow array multiplication
# w is the sum of all w_m with right order assigned
w_array = np.array(w_list)
eps_array = np.array(eps_list)
w_eps = sum(eps_array*w_array)
# symbolic euler eqs with order assigned
dzeta_dt = -(gradzeta_gradphi*(eps**2)) + w_eps*(1+(gradzeta2*eps**2))
dphi_dt = (-0.5*(gradphi2*eps**2) - (grav*zeta*eps) +
0.5*(w_eps**2)*(1 + (gradzeta2*eps**2)) )
args = tuple(w_list + [gradphi2, gradzeta2, zeta, gradzeta_gradphi, time, damping_factor])
expressions = [dphi_dt, dzeta_dt]
functions = [None]*2
final_eqs = [None]*2
for eqn in expressions:
coeffs = sy.Poly(eqn, eps).all_coeffs() #gives equation as a list with coefficient of eps
coeffs.reverse()
coeffs_array = np.array(coeffs) # allows array multiplication
if M > 1:
coeffs_array[2:M+1] *= damping #applying damping factor on non_linear terms
eqn_truncated = np.sum(coeffs_array[1:M+1]) # sums the needed M coefficients, discards rest
s = expressions.index(eqn)
final_eqs[s] = sy.collect(eqn_truncated, damping) #factorizes to reduce number of computations
functions[s] = sy.lambdify(args, final_eqs[s], modules='numpy') # now it is a numpy python function
return functions[0], functions[1], final_eqs[0], final_eqs[1]
def get_conditionally_linear_system(eqs):
'''
Convert equations into a linear system using sympy.
Parameters
----------
eqs : `Equations`
The model equations.
Returns
-------
coefficients : dict of (sympy expression, sympy expression) tuples
For every variable x, a tuple (M, B) containing the coefficients M and
B (as sympy expressions) for M * x + B
Raises
------
ValueError
If one of the equations cannot be converted into a M * x + B form.
Examples
--------
>>> from brian2 import Equations
>>> eqs = Equations("""
... dv/dt = (-v + w**2) / tau : 1
... dw/dt = -w / tau : 1
... """)
>>> system = get_conditionally_linear_system(eqs)
>>> print(system['v'])
(-1/tau, w**2.0/tau)
>>> print(system['w'])
(-1/tau, 0)
'''
diff_eqs = eqs.substituted_expressions
coefficients = {}
for name, expr in diff_eqs:
var = sp.Symbol(name, real=True)
# Factor out the variable
s_expr = sp.collect(expr.sympy_expr.expand(), var, evaluate=False)
if len(s_expr) > 2 or var not in s_expr:
raise ValueError(
('The expression "%s", defining the variable %s, '
'could not be separated into linear components') %
(expr, name))
coefficients[name] = (s_expr[var], s_expr.get(1, 0))
return coefficients
def print_as_array(m, mname, sufix=None, use_cse=False, header=None,
print_file=True, collect_for=None, pow_by_mul=True, order='C',
op='+='):
ls = []
if use_cse:
subs, m_list = sympy.cse(m)
for i, v in enumerate(m_list):
m[i] = v
if sufix is None:
namesufix = '{0}'.format(mname)
else:
namesufix = '{0}_{1}'.format(mname, sufix)
filename = 'print_{0}.txt'.format(namesufix)
if header:
ls.append(header)
if use_cse:
ls.append('# cdefs')
num = 9
for i, sub in enumerate(subs[::num]):
ls.append('cdef double ' + ', '.join(
map(str, [j[0] for j in subs[num*i:num*(i+1)]])))
ls.append('# subs')
for sub in subs:
ls.append('{0} = {1}'.format(*sub))
ls.append('# {0}'.format(namesufix))
num = len([i for i in list(m) if i])
ls.append('# {0}_num={1}'.format(namesufix, num))
if order == 'C':
miter = enumerate(np.ravel(m))
elif order == 'F':
miter = enumerate(np.ravel(m.T))
miter = list(miter)
for i, v in miter:
if v:
if collect_for is not None:
v = collect(v, collect_for, evaluate=False)
ls.append('{0}[pos+{1}] +='.format(mname, i))
for k, expr in v.items():
ls.append('# collected for {k}'.format(k=k))
ls.append(' {expr}'.format(expr=k*expr))
else:
if pow_by_mul:
v = str(v)
for p in re.findall(r'\w+\*\*\d+', v):
var, exp = p.split('**')
v = v.replace(p, '(' + '*'.join([var]*int(exp)) + ')')
ls.append('{0}[pos+{1}] {2} {3}'.format(mname, i, op, v))
string = '\n'.join(ls)
if print_file:
with open(filename, 'w') as f:
f.write(string)
return string
def as_coeff_exponent(self, x):
""" c*x**e -> c,e where x can be any symbolic expression.
"""
x = sympify(x)
wc = Wild('wc')
we = Wild('we')
p = wc*x**we
from sympy import collect
s = collect(self, x)
d = s.match(p)
if d is not None and we in d:
return d[wc], d[we]
return s, S.Zero
def as_coeff_exponent(self, x):
""" c*x**e -> c,e where x can be any symbolic expression.
"""
x = sympify(x)
wc = Wild('wc')
we = Wild('we')
p = wc*x**we
from sympy import collect
self = collect(self, x)
d = self.match(p)
if d is not None and we in d:
return d[wc], d[we]
return self, S.Zero
def transform(degree):
x = Symbol('x')
A = Symbol('A')
B = Symbol('B')
a0 = Symbol('a0')
xprime = A*(x-B)
expr = a0
for k in xrange(1,degree+1):
a = Symbol('a{}'.format(k))
expr = expr + a*xprime**k
d = collect(expand(expr),x,evaluate=False)
for k in xrange(degree+1):
print '{}:\t{}'.format(x**k,d[x**k])
def separate_odd_pow_terms(f, term):
'''separate_odd_pow_terms(x**3+x**2+x+1, x) = (x**2+1, x**3+x)'''
g = collect(f, term)
h = Poly(g, term)
odd = even = zero
for (e, ), coeff in h.all_terms():
t = term**e * coeff
if e % 2:
odd += t
else:
even += t
return even, odd
def run():
t,x0,x1,x2,x3 = map(Symbol, ['t','x0','x1','x2','x3'])
a01 = x0 + (x1 - x0)*t
a11 = x1 + (x2 - x1)*t
a21 = x2 + (x3 - x2)*t
a02 = a01 + (a11 - a01)*t
a12 = a11 + (a21 - a11)*t
x = a02 + (a12 - a02)*t
a = expand(a11 - a01)
a = collect(a,t)
print(a)
def Q2():
x, y, yp, yc = sp.symbols('x y yp yc')
z, z1, z2 = sp.symbols('z z1 z2')
c0, c1, c2, c3 = sp.symbols('c0 c1 c2 c3')
A, B = sp.symbols('A B')
expr = pow(z, 2) + 4*z + 4
sol = sp.solve(expr, z)
pprint.pprint(sol)
if len(sol) == 1:
yp = (A + B*x)*sp.exp(z1*x)
if len(sol) == 2:
yp = (A * sp.exp(z1)) + (B * sp.exp(z2*x))
yc = c0 + c1*x + c2*(x**2) + c3*(x**3)
dyc = sp.diff(yc, x)
ddyc = sp.diff(dyc, x)
ydict = sp.collect(yc + 4*dyc + 4*ddyc, x, evaluate=False)
c3 = 1
exp0 = ydict[x**0]
exp1 = ydict[x**1]
exp2 = ydict[x**2]
exp3 = ydict[x**3]
sc3 = sp.solve(sp.Eq(exp3, 1))
exp2 = exp2.subs('c3', sc3[0])
sc2 = sp.solve(sp.Eq(exp2, 0))
exp1 = exp1.subs('c2', sc2[0])
exp1 = exp1.subs('c3', sc3[0])
sc1 = sp.solve(sp.Eq(exp1, 0))
exp0 = exp0.subs('c1', sc1[0])
exp0 = exp0.subs('c2', sc2[0])
exp0 = exp0.subs('c3', sc3[0])
sc0 = sp.solve(sp.Eq(exp0, 0))
yc = sc3[0]*x**3 + sc2[0]*x**2 + sc1[0]*x + sc0[0]
y = yp + yc
pprint.pprint("yp = " + str(yp))
pprint.pprint("yc = " + str(yc))
pprint.pprint("y = " + str(y))
return
def compile_expansions(self): #compile expansion term as sympy expr
collect_cartesian_terms = None
for func in self.funcs:
if isinstance(func,Trig):
trigfunc = func
if trigfunc.cartesian == True:
if len(trigfunc.args) > 1:
collect_cartesian_terms = 2
else:
collect_cartesian_terms = 1
free_parameters = []
compiled_expansions = []
for term in self.expansions:
term.coeff.compute()
term.coeff.compile()
for i,func in enumerate(term.funcs):
func.compute()
func.compile()
compiled_term = 1
for part in term.form:
if hasattr(part,'sympy_form') == False:
compiled_term = compiled_term*part
else:
compiled_term = compiled_term*part.sympy_form
#if two e modes:
if collect_cartesian_terms == 2:
compiled_term = sym.collect(compiled_term,['x_alpha','y_alpha','x_beta','y_beta'],exact=True)
#else if 1 e mode:
elif collect_cartesian_terms == 1:
compiled_term = sym.collect(compiled_term,['x','y'],exact=True)
compiled_expansions.append(compiled_term)
if term.coeff.sympy_form not in free_parameters and compiled_term != 0:
free_parameters.append(term.coeff.sympy_form)
self.expansions = compiled_expansions
self.parameters = free_parameters
def get_transfer_functions(t, x, theta, f, x_dot_dot, theta_dot_dot, s, X,
Theta, F):
laplace_tf = {
x: X,
sp.diff(x, t): s * X,
sp.diff(x, (t, 2)): s**2 * X,
theta: Theta,
sp.diff(theta, t): s * Theta,
sp.diff(theta, (t, 2)): s**2 * Theta,
f: F
}
eq1 = sp.Eq(sp.diff(x, (t, 2)), x_dot_dot).subs(laplace_tf)
eq2 = sp.Eq(sp.diff(theta, (t, 2)), theta_dot_dot).subs(laplace_tf)
result = sp.solve([eq1, eq2], [X, Theta])
X_tf = result[X] / F
Theta_tf = result[Theta] / F
# attempt partial fraction decomposition
X_tf = sp.collect(sp.expand(X_tf), s)
Theta_tf = sp.collect(sp.expand(Theta_tf), s)
return X_tf, Theta_tf
def test_gcd_terms():
f = 2*(x + 1)*(x + 4)/(5*x**2 + 5) + (2*x + 2)*(x + 5)/(x**2 + 1)/5 + \
(2*x + 2)*(x + 6)/(5*x**2 + 5)
assert _gcd_terms(f) == ((S(6) / 5) * ((1 + x) / (1 + x**2)), 5 + x, 1)
assert _gcd_terms(Add.make_args(f)) == \
((S(6)/5)*((1 + x)/(1 + x**2)), 5 + x, 1)
newf = (S(6) / 5) * ((1 + x) * (5 + x) / (1 + x**2))
assert gcd_terms(f) == newf
args = Add.make_args(f)
# non-Basic sequences of terms treated as terms of Add
assert gcd_terms(list(args)) == newf
assert gcd_terms(tuple(args)) == newf
assert gcd_terms(set(args)) == newf
# but a Basic sequence is treated as a container
assert gcd_terms(Tuple(*args)) != newf
assert gcd_terms(Basic(Tuple(1, 3*y + 3*x*y), Tuple(1, 3))) == \
Basic((1, 3*y*(x + 1)), (1, 3))
# but we shouldn't change keys of a dictionary or some may be lost
assert gcd_terms(Dict((x*(1 + y), 2), (x + x*y, y + x*y))) == \
Dict({x*(y + 1): 2, x + x*y: y*(1 + x)})
assert gcd_terms((2 * x + 2)**3 +
(2 * x + 2)**2) == 4 * (x + 1)**2 * (2 * x + 3)
assert gcd_terms(0) == 0
assert gcd_terms(1) == 1
assert gcd_terms(x) == x
assert gcd_terms(2 + 2 * x) == Mul(2, 1 + x, evaluate=False)
arg = x * (2 * x + 4 * y)
garg = 2 * x * (x + 2 * y)
assert gcd_terms(arg) == garg
assert gcd_terms(sin(arg)) == sin(garg)
# issue 3040-like
alpha, alpha1, alpha2, alpha3 = symbols('alpha:4')
a = alpha**2 - alpha * x**2 + alpha + x**3 - x * (alpha + 1)
rep = (alpha,
(1 + sqrt(5)) / 2 + alpha1 * x + alpha2 * x**2 + alpha3 * x**3)
s = (a / (x - alpha)).subs(*rep).series(x, 0, 1)
assert simplify(collect(s, x)) == -sqrt(5) / 2 - S(3) / 2 + O(x)
# issue 2818
assert _gcd_terms([S.Zero, S.Zero]) == (0, 0, 1)
assert _gcd_terms([2 * x + 4]) == (2, x + 2, 1)
eq = x / (x + 1 / x)
assert gcd_terms(eq, fraction=False) == eq
def test_collect_1():
"""Collect with respect to a Symbol"""
x, y, z, n = symbols('xyzn')
assert collect( x + y*x, x ) == x * (1 + y)
assert collect( x + x**2, x ) == x + x**2
assert collect( x**2 + y*x**2, x ) == (x**2)*(1+y)
assert collect( x**2 + y*x, x ) == x*y + x**2
assert collect( 2*x**2 + y*x**2 + 3*x*y, [x] ) == x**2*(2+y) + 3*x*y
assert collect( 2*x**2 + y*x**2 + 3*x*y, [y] ) == 2*x**2 + y*(x**2+3*x)
assert collect( ((1 + y + x)**4).expand(), x) == ((1 + y)**4).expand() + \
x*(4*(1 + y)**3).expand() + x**2*(6*(1 + y)**2).expand() + \
x**3*(4*(1 + y)).expand() + x**4
def get_x_sol(Eq_UOU, x_ul, subs_yz):
Eq_UOU_x = Eq_UOU.subs(subs_yz)
Eq_UOU_x_rearr = sp.Eq(
-Eq_UOU_x.args[1].args[1], -Eq_UOU_x.args[0] +
Eq_UOU_x.args[1].args[0] + Eq_UOU_x.args[1].args[2])
Eq_UOU_x_rhs = Eq_UOU_x_rearr.args[1]**2 - Eq_UOU_x_rearr.args[0]**2
Eq_UOU_x_rhs_collect = sp.collect(sp.expand(Eq_UOU_x_rhs), x_ul)
A_ = Eq_UOU_x_rhs_collect.coeff(x_ul, 2)
B_ = Eq_UOU_x_rhs_collect.coeff(x_ul, 1)
C_ = Eq_UOU_x_rhs_collect.coeff(x_ul, 0)
A, B, C = sp.symbols('A, B, C')
x_ul_sol1, x_ul_sol2 = sp.solve(A * x_ul**2 + B * x_ul + C, x_ul)
return x_ul_sol1, x_ul_sol2, A_, B_, C_
def test_action_verbs():
assert nsimplify((1/(exp(3*pi*x/5)+1))) == (1/(exp(3*pi*x/5)+1)).nsimplify()
assert ratsimp(1/x + 1/y) == (1/x + 1/y).ratsimp()
assert trigsimp(log(x), deep=True) == (log(x)).trigsimp(deep = True)
assert radsimp(1/(2+sqrt(2))) == (1/(2+sqrt(2))).radsimp()
assert powsimp(x**y*x**z*y**z, combine='all') == (x**y*x**z*y**z).powsimp(combine='all')
assert simplify(x**y*x**z*y**z) == (x**y*x**z*y**z).simplify()
assert together(1/x + 1/y) == (1/x + 1/y).together()
assert separate((x*(y*z)**3)**2) == ((x*(y*z)**3)**2).separate()
assert collect(a*x**2 + b*x**2 + a*x - b*x + c, x) == (a*x**2 + b*x**2 + a*x - b*x + c).collect(x)
assert apart(y/(y+2)/(y+1), y) == (y/(y+2)/(y+1)).apart(y)
assert combsimp(y/(x+2)/(x+1)) == (y/(x+2)/(x+1)).combsimp()
assert factor(x**2+5*x+6) == (x**2+5*x+6).factor()
assert refine(sqrt(x**2)) == sqrt(x**2).refine()
assert cancel((x**2+5*x+6)/(x+2)) == ((x**2+5*x+6)/(x+2)).cancel()
def PartialTrace(expr, n): ##calculate the partial trace of the n_th photon
dictadd = collect(expr, [HH[l1, l2, l3, l4, l5, l6]], evaluate=False)
TermsCoeff = list(dictadd.items())
ParticleOne = []
ParticleTwo = []
## get the size of the matrix
for ii in range(len(TermsCoeff)):
HHList = TermsCoeff[ii][0]
if HHList.indices[n - 1] == HHList.indices[n + 2]:
ll = [
HHList.indices[0], HHList.indices[1], HHList.indices[2],
HHList.indices[3], HHList.indices[4], HHList.indices[5]
]
del (
ll[n - 1], ll[n + 1]
) ## because cannot del all at the same time, thus do it one by one, the index is not n+2
ParticleOne.append(ll[0])
ParticleTwo.append(ll[1])
# start from 0
Upperone = max(ParticleOne) + 1
Lowerone = min(min(ParticleOne), 0)
Uppertwo = max(ParticleTwo) + 1
Lowertwo = min(min(ParticleTwo), 0)
rangeP1 = Upperone - Lowerone
rangeP2 = Uppertwo - Lowertwo
Msize = (rangeP1 * rangeP2)
SMatrix = SparseMatrix(Msize, Msize, {(0, 0): 0})
for ii in range(len(TermsCoeff)):
HHList = TermsCoeff[ii][0]
if HHList.indices[n - 1] == HHList.indices[n + 2]:
ll = [
HHList.indices[0], HHList.indices[1], HHList.indices[2],
HHList.indices[3], HHList.indices[4], HHList.indices[5]
]
del (
ll[n - 1], ll[n + 1]
) ## because cannot del all at the same time, thus do it one by one, the index is not n+2
# print('rest: ',ll)
# print('rest: ',ll[0]-Lowerone,'',ll[1]-Lowertwo, '',ll[2]-Lowerone,'',ll[3]-Lowertwo)
Dimrow = (ll[0] - Lowerone) * rangeP2 + (ll[1] - Lowertwo)
Dimcol = (ll[2] - Lowerone) * rangeP2 + (ll[3] - Lowertwo)
SMatrix = SparseMatrix(
Msize, Msize, {(Dimrow, Dimcol): TermsCoeff[ii][1]}) + SMatrix
return SMatrix.rank()
def construct(k):
"""Construct Hermite shape functions.
The result can be used for interpolating a function and its first k derivatives on [0,1]."""
order = 2 * k + 1 # minimum polynomial order that has enough degrees of freedom
*A, x = sy.symbols("a0:%d,x" % (order + 1))
# Create the polynomial and its derivatives
#
w = sum(a * x**i for i, a in enumerate(A))
λw = lambda x0: w.subs({x: x0})
wp = [sy.diff(w, x, i) for i in range(1, 1 + k)]
# see my lecture notes, sec. 5.8
λwp = [(lambda expr: lambda x0: expr.subs({x: x0}))(expr) for expr in wp]
# Formulate the interpolation conditions
#
zero, one = sy.S.Zero, sy.S.One
w0, w1 = sy.symbols("w0, w1")
eqs = [
λw(zero) - w0, # see sy.solve()
λw(one) - w1
]
dofs = [w0, w1]
for i, f in enumerate(λwp):
d0_name = "w%s0" % ((i + 1) * "p")
d1_name = "w%s1" % ((i + 1) * "p")
d0, d1 = sy.symbols("%s, %s" % (d0_name, d1_name))
eqs.extend([f(zero) - d0, f(one) - d1])
dofs.extend([d0, d1])
# Solve the interpolation conditions for the polynomial coefficients aj
#
coeffs = sy.solve(eqs, A)
# Substitute solution into w, collect w.r.t. the DOFs
#
solution = w.subs(coeffs)
solution = sy.collect(sy.expand(solution), dofs)
# Extract the shape functions
#
N = [solution.coeff(dof) for dof in dofs]
return (dofs, N)
def coeff(self, x, expand=True):
"""
Returns the coefficient of the term "x" or None if there is no "x".
Optional expand keyword argument allows one to control whether the
expression is expanded before terms are collected, which can be useful
if the term "x" isn't nested inside of terms and you don't want the
returned coefficient to be expanded.
Example:
>>> from sympy import symbols
>>> from sympy.abc import x, y, z
>>> (3+2*x+4*x**2).coeff(1)
>>> (3+2*x+4*x**2).coeff(x)
2
>>> (3+2*x+4*x**2).coeff(x**2)
4
>>> (3+2*x+4*x**2).coeff(x**3)
>>> (z*(x+y)**2).coeff(z)
2*x*y + x**2 + y**2
>>> (z*(x+y)**2).coeff(z, expand=False)
(x + y)**2
>>>
"""
from sympy import collect
x = sympify(x)
const = x.as_coeff_mul()[0] # constant multiplying x
if const != S.One: # get rid of constants
result = self.coeff(x / const)
if result is not None:
return (result / const)
else:
return None
if x.is_Integer:
return
result = self
if expand:
result = result.expand(
) # collect expects its arguments in expanded form
result = collect(result, x, evaluate=False, exact=True)
if x in result:
return result[x]
else:
return None
def test_collect_3():
"""Collect with respect to a product"""
a, b, c = symbols('a,b,c')
f = Function('f')
x, y, z, n = symbols('x,y,z,n')
assert collect(-x/8 + x*y, -x) == x*(y - S(1)/8)
assert collect( 1 + x*(y**2), x*y ) == 1 + x*(y**2)
assert collect( x*y + a*x*y, x*y) == x*y*(1 + a)
assert collect( 1 + x*y + a*x*y, x*y) == 1 + x*y*(1 + a)
assert collect(a*x*f(x) + b*(x*f(x)), x*f(x)) == x*(a + b)*f(x)
assert collect(a*x*log(x) + b*(x*log(x)), x*log(x)) == x*(a + b)*log(x)
assert collect(a*x**2*log(x)**2 + b*(x*log(x))**2, x*log(x)) == x**2*log(x)**2*(a + b)
# with respect to a product of three symbols
assert collect(y*x*z+a*x*y*z, x*y*z) == (1 + a)*x*y*z
def implicit(self, expression):
"Full implicit method, linearising for example (V - E)^2, but this is not desired."
# print('Expression', expression)
# Transform the gradient into a difference TODO: more robust...
new_expression = expression.replace('d' + self.name, '_t_gradient_')
new_expression = re.sub(r'([^\w]+)' + self.name + r'([^\w]+)',
r'\1_' + self.name + r'\2', new_expression)
new_expression = new_expression.replace(
'_t_gradient_', '(_' + self.name + ' - ' + self.name + ')')
# print('New Expression', new_expression)
# Add a sympbol for the next value of the variable
new_var = sp.Symbol('_' + self.name)
self.local_dict['_' + self.name] = new_var
# Parse the string
analysed = self.parse_expression(new_expression,
local_dict=self.local_dict)
self.analysed = analysed
# print('Analysed', analysed)
# Solve the equation for delta_mp
solved = sp.solve(analysed, new_var, check=False, rational=False)
# print('Solved', solved)
if len(solved) > 1:
Global._print(self.expression)
Global._error(
'Parser: the ODE is not linear, can not use the implicit method.'
)
else:
solved = solved[0]
equation = sp.simplify(sp.collect(solved, self.local_dict['dt']))
# Obtain C code
variable_name = self.c_code(self.local_dict[self.name])
explicit_code = Global.config['precision'] + ' _' + self.name + ' = '\
+ self.c_code(equation) + ';'
switch = variable_name + ' = _' + self.name + ' ;'
# Return result
return [{}, explicit_code, switch]
def generate_operator_entries_from_equation(equation, operators: list, fields, grid):
row_of_operators = []
indices = []
def recursive_descent(expr, field_index):
if expr.is_Number:
identity = base.Identity(grid[field_index])
if not expr == sympy.sympify(1):
return base.Scaling(float(expr.evalf()), identity)
else:
return identity
elif expr.is_Symbol:
op_symbol = expr
j = next(k for k, op_info in enumerate(operators) if op_symbol.name == op_info.name)
operator = base.Operator(op_symbol.name, grid[field_index], ConstantStencilGenerator(operators[j].stencil))
return operator
elif expr.is_Mul:
tmp = recursive_descent(expr.args[-1], field_index)
for arg in expr.args[-2::-1]:
if arg.is_Number:
tmp = base.Scaling(float(arg.evalf()), tmp)
else:
lhs = recursive_descent(arg, field_index)
tmp = base.Multiplication(lhs, tmp)
elif expr.is_Add:
tmp = recursive_descent(expr.args[0], field_index)
for arg in expr.args[1:]:
tmp = base.Addition(recursive_descent(arg, field_index), tmp)
else:
raise RuntimeError("Invalid Expression")
return tmp
expanded_expression = sympy.expand(equation.sympy_expr)
for i, field in enumerate(fields):
if field in expanded_expression.free_symbols:
collected_terms = sympy.collect(expanded_expression, field, evaluate=False)
term = collected_terms[field]
entry = recursive_descent(term, i)
row_of_operators.append(entry)
indices.append(i)
for i in range(len(grid)):
if i not in indices:
row_of_operators.append(base.ZeroOperator(grid[i]))
indices.append(i)
result = [operator for (index, operator) in sorted(zip(indices, row_of_operators), key=lambda p: p[0])]
return result
def update_kernel(expr, kernel, variables):
expr = _update_kernel(expr, kernel(*variables))
cls_name = kernel(evaluate=False)
if isinstance(cls_name, RBF):
expr = expr.subs(RBF(*variables), Symbol('RBF'))
expr = collect(expr, Symbol('RBF'))
# ...
a = expr / Symbol('RBF')
a = simplify(a)
expr = a * Symbol('RBF')
# ...
expr = expr.subs(Symbol('RBF'), RBF(*variables))
return expr
def crossing():
t, d = map(Symbol, ['t', 'd'])
a3, a2, a1, a0 = map(Symbol, ['a3', 'a2', 'a1', 'a0'])
# parametric equation
x = a3 * t**3 + a2 * t**2 + a1 * t + a0
x1 = x
x2 = x.subs(t, t + d)
# 'forward difference'
dX = x2 - x1
dX = expand(dX)
dX = collect(dX, t)
print(dX)
def test_collect_D():
D = Derivative
f = Function('f')
x, a, b = symbols('x,a,b')
fx = D(f(x), x)
fxx = D(f(x), x, x)
assert collect(a * fx + b * fx, fx) == (a + b) * fx
assert collect(a * D(fx, x) + b * D(fx, x), fx) == (a + b) * D(fx, x)
assert collect(a * fxx + b * fxx, fx) == (a + b) * D(fx, x)
# issue 4784
assert collect(5 * f(x) + 3 * fx, fx) == 5 * f(x) + 3 * fx
assert collect(f(x) + f(x)*diff(f(x), x) + x*diff(f(x), x)*f(x), f(x).diff(x)) == \
(x*f(x) + f(x))*D(f(x), x) + f(x)
assert collect(f(x) + f(x)*diff(f(x), x) + x*diff(f(x), x)*f(x), f(x).diff(x), exact=True) == \
(x*f(x) + f(x))*D(f(x), x) + f(x)
assert collect(1/f(x) + 1/f(x)*diff(f(x), x) + x*diff(f(x), x)/f(x), f(x).diff(x), exact=True) == \
(1/f(x) + x/f(x))*D(f(x), x) + 1/f(x)
e = (1 + x * fx + fx) / f(x)
assert collect(e.expand(), fx) == fx * (x / f(x) + 1 / f(x)) + 1 / f(x)
def test_helper_functions():
a, b, c = symbols('a b c')
tsteqn = Equation(a, b / c)
raises(ValueError, lambda: integrate(tsteqn, c))
raises(AttributeError, lambda: integrate(tsteqn, c, side='right'))
assert tsteqn.evalf(4, {b: 2.0, c: 4}) == Equation(a, 0.5000)
assert diff(tsteqn, c) == Equation(diff(a, c, evaluate=False), -b / c**2)
tsteqn = Equation(a * c, b / c)
assert diff(tsteqn, c) == Equation(a, -b / c**2)
assert integrate(tsteqn, c, side='rhs') == integrate(tsteqn.rhs, c)
assert integrate(tsteqn, c, side='lhs') == integrate(tsteqn.lhs, c)
def adsq(eqn):
# Arbitrary python function
return eqn + eqn**2
assert adsq(Equation(a * c, b / c)) == Equation(a**2 * c**2 + a * c,
b**2 / c**2 + b / c)
assert Equation(
(a - 1) * (a + 1),
(2 * b + c)**2).expand() == Equation(a**2 - 1,
4 * b**2 + 4 * b * c + c**2)
assert expand(Equation(
(a - 1) * (a + 1),
(2 * b + c)**2)) == Equation(a**2 - 1, 4 * b**2 + 4 * b * c + c**2)
assert Equation(a**2 - 1,
4 * b**2 + 4 * b * c + c**2).factor() == Equation(
(a - 1) * (a + 1), (2 * b + c)**2)
assert factor(Equation(a**2 - 1, 4 * b**2 + 4 * b * c + c**2)) == Equation(
(a - 1) * (a + 1), (2 * b + c)**2)
assert Equation(a**2 - 1,
4 * b**2 + 4 * b * c + c * a).collect(c) == Equation(
a**2 - 1, 4 * b**2 + c * (a + 4 * b))
assert collect(Equation(a**2 - 1, 4 * b**2 + 4 * b * c + c * a),
c) == Equation(a**2 - 1, 4 * b**2 + c * (a + 4 * b))
assert Equation((a + 1)**2 / (a + 1),
exp(log(c))).simplify() == Equation(a + 1, c)
assert simplify(Equation((a + 1)**2 / (a + 1),
exp(log(c)))) == Equation(a + 1, c)
# Check matrix exponentiation is not overridden.
tsteqn5 = Equation(a, Matrix([[1, 1], [1, 1]]))
assert (exp(tsteqn5).lhs == exp(a))
assert (exp(tsteqn5).rhs == exp(Matrix([[1, 1], [1, 1]])))
