Coverage for /usr/lib/python3/dist-packages/sympy/solvers/bivariate.py: 10%
210 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1from sympy.core.add import Add
2from sympy.core.exprtools import factor_terms
3from sympy.core.function import expand_log, _mexpand
4from sympy.core.power import Pow
5from sympy.core.singleton import S
6from sympy.core.sorting import ordered
7from sympy.core.symbol import Dummy
8from sympy.functions.elementary.exponential import (LambertW, exp, log)
9from sympy.functions.elementary.miscellaneous import root
10from sympy.polys.polyroots import roots
11from sympy.polys.polytools import Poly, factor
12from sympy.simplify.simplify import separatevars
13from sympy.simplify.radsimp import collect
14from sympy.simplify.simplify import powsimp
15from sympy.solvers.solvers import solve, _invert
16from sympy.utilities.iterables import uniq
19def _filtered_gens(poly, symbol):
20 """process the generators of ``poly``, returning the set of generators that
21 have ``symbol``. If there are two generators that are inverses of each other,
22 prefer the one that has no denominator.
24 Examples
25 ========
27 >>> from sympy.solvers.bivariate import _filtered_gens
28 >>> from sympy import Poly, exp
29 >>> from sympy.abc import x
30 >>> _filtered_gens(Poly(x + 1/x + exp(x)), x)
31 {x, exp(x)}
33 """
34 # TODO it would be good to pick the smallest divisible power
35 # instead of the base for something like x**4 + x**2 -->
36 # return x**2 not x
37 gens = {g for g in poly.gens if symbol in g.free_symbols}
38 for g in list(gens):
39 ag = 1/g
40 if g in gens and ag in gens:
41 if ag.as_numer_denom()[1] is not S.One:
42 g = ag
43 gens.remove(g)
44 return gens
47def _mostfunc(lhs, func, X=None):
48 """Returns the term in lhs which contains the most of the
49 func-type things e.g. log(log(x)) wins over log(x) if both terms appear.
51 ``func`` can be a function (exp, log, etc...) or any other SymPy object,
52 like Pow.
54 If ``X`` is not ``None``, then the function returns the term composed with the
55 most ``func`` having the specified variable.
57 Examples
58 ========
60 >>> from sympy.solvers.bivariate import _mostfunc
61 >>> from sympy import exp
62 >>> from sympy.abc import x, y
63 >>> _mostfunc(exp(x) + exp(exp(x) + 2), exp)
64 exp(exp(x) + 2)
65 >>> _mostfunc(exp(x) + exp(exp(y) + 2), exp)
66 exp(exp(y) + 2)
67 >>> _mostfunc(exp(x) + exp(exp(y) + 2), exp, x)
68 exp(x)
69 >>> _mostfunc(x, exp, x) is None
70 True
71 >>> _mostfunc(exp(x) + exp(x*y), exp, x)
72 exp(x)
73 """
74 fterms = [tmp for tmp in lhs.atoms(func) if (not X or
75 X.is_Symbol and X in tmp.free_symbols or
76 not X.is_Symbol and tmp.has(X))]
77 if len(fterms) == 1:
78 return fterms[0]
79 elif fterms:
80 return max(list(ordered(fterms)), key=lambda x: x.count(func))
81 return None
84def _linab(arg, symbol):
85 """Return ``a, b, X`` assuming ``arg`` can be written as ``a*X + b``
86 where ``X`` is a symbol-dependent factor and ``a`` and ``b`` are
87 independent of ``symbol``.
89 Examples
90 ========
92 >>> from sympy.solvers.bivariate import _linab
93 >>> from sympy.abc import x, y
94 >>> from sympy import exp, S
95 >>> _linab(S(2), x)
96 (2, 0, 1)
97 >>> _linab(2*x, x)
98 (2, 0, x)
99 >>> _linab(y + y*x + 2*x, x)
100 (y + 2, y, x)
101 >>> _linab(3 + 2*exp(x), x)
102 (2, 3, exp(x))
103 """
104 arg = factor_terms(arg.expand())
105 ind, dep = arg.as_independent(symbol)
106 if arg.is_Mul and dep.is_Add:
107 a, b, x = _linab(dep, symbol)
108 return ind*a, ind*b, x
109 if not arg.is_Add:
110 b = 0
111 a, x = ind, dep
112 else:
113 b = ind
114 a, x = separatevars(dep).as_independent(symbol, as_Add=False)
115 if x.could_extract_minus_sign():
116 a = -a
117 x = -x
118 return a, b, x
121def _lambert(eq, x):
122 """
123 Given an expression assumed to be in the form
124 ``F(X, a..f) = a*log(b*X + c) + d*X + f = 0``
125 where X = g(x) and x = g^-1(X), return the Lambert solution,
126 ``x = g^-1(-c/b + (a/d)*W(d/(a*b)*exp(c*d/a/b)*exp(-f/a)))``.
127 """
128 eq = _mexpand(expand_log(eq))
129 mainlog = _mostfunc(eq, log, x)
130 if not mainlog:
131 return [] # violated assumptions
132 other = eq.subs(mainlog, 0)
133 if isinstance(-other, log):
134 eq = (eq - other).subs(mainlog, mainlog.args[0])
135 mainlog = mainlog.args[0]
136 if not isinstance(mainlog, log):
137 return [] # violated assumptions
138 other = -(-other).args[0]
139 eq += other
140 if x not in other.free_symbols:
141 return [] # violated assumptions
142 d, f, X2 = _linab(other, x)
143 logterm = collect(eq - other, mainlog)
144 a = logterm.as_coefficient(mainlog)
145 if a is None or x in a.free_symbols:
146 return [] # violated assumptions
147 logarg = mainlog.args[0]
148 b, c, X1 = _linab(logarg, x)
149 if X1 != X2:
150 return [] # violated assumptions
152 # invert the generator X1 so we have x(u)
153 u = Dummy('rhs')
154 xusolns = solve(X1 - u, x)
156 # There are infinitely many branches for LambertW
157 # but only branches for k = -1 and 0 might be real. The k = 0
158 # branch is real and the k = -1 branch is real if the LambertW argumen
159 # in in range [-1/e, 0]. Since `solve` does not return infinite
160 # solutions we will only include the -1 branch if it tests as real.
161 # Otherwise, inclusion of any LambertW in the solution indicates to
162 # the user that there are imaginary solutions corresponding to
163 # different k values.
164 lambert_real_branches = [-1, 0]
165 sol = []
167 # solution of the given Lambert equation is like
168 # sol = -c/b + (a/d)*LambertW(arg, k),
169 # where arg = d/(a*b)*exp((c*d-b*f)/a/b) and k in lambert_real_branches.
170 # Instead of considering the single arg, `d/(a*b)*exp((c*d-b*f)/a/b)`,
171 # the individual `p` roots obtained when writing `exp((c*d-b*f)/a/b)`
172 # as `exp(A/p) = exp(A)**(1/p)`, where `p` is an Integer, are used.
174 # calculating args for LambertW
175 num, den = ((c*d-b*f)/a/b).as_numer_denom()
176 p, den = den.as_coeff_Mul()
177 e = exp(num/den)
178 t = Dummy('t')
179 args = [d/(a*b)*t for t in roots(t**p - e, t).keys()]
181 # calculating solutions from args
182 for arg in args:
183 for k in lambert_real_branches:
184 w = LambertW(arg, k)
185 if k and not w.is_real:
186 continue
187 rhs = -c/b + (a/d)*w
189 sol.extend(xu.subs(u, rhs) for xu in xusolns)
190 return sol
193def _solve_lambert(f, symbol, gens):
194 """Return solution to ``f`` if it is a Lambert-type expression
195 else raise NotImplementedError.
197 For ``f(X, a..f) = a*log(b*X + c) + d*X - f = 0`` the solution
198 for ``X`` is ``X = -c/b + (a/d)*W(d/(a*b)*exp(c*d/a/b)*exp(f/a))``.
199 There are a variety of forms for `f(X, a..f)` as enumerated below:
201 1a1)
202 if B**B = R for R not in [0, 1] (since those cases would already
203 be solved before getting here) then log of both sides gives
204 log(B) + log(log(B)) = log(log(R)) and
205 X = log(B), a = 1, b = 1, c = 0, d = 1, f = log(log(R))
206 1a2)
207 if B*(b*log(B) + c)**a = R then log of both sides gives
208 log(B) + a*log(b*log(B) + c) = log(R) and
209 X = log(B), d=1, f=log(R)
210 1b)
211 if a*log(b*B + c) + d*B = R and
212 X = B, f = R
213 2a)
214 if (b*B + c)*exp(d*B + g) = R then log of both sides gives
215 log(b*B + c) + d*B + g = log(R) and
216 X = B, a = 1, f = log(R) - g
217 2b)
218 if g*exp(d*B + h) - b*B = c then the log form is
219 log(g) + d*B + h - log(b*B + c) = 0 and
220 X = B, a = -1, f = -h - log(g)
221 3)
222 if d*p**(a*B + g) - b*B = c then the log form is
223 log(d) + (a*B + g)*log(p) - log(b*B + c) = 0 and
224 X = B, a = -1, d = a*log(p), f = -log(d) - g*log(p)
225 """
227 def _solve_even_degree_expr(expr, t, symbol):
228 """Return the unique solutions of equations derived from
229 ``expr`` by replacing ``t`` with ``+/- symbol``.
231 Parameters
232 ==========
234 expr : Expr
235 The expression which includes a dummy variable t to be
236 replaced with +symbol and -symbol.
238 symbol : Symbol
239 The symbol for which a solution is being sought.
241 Returns
242 =======
244 List of unique solution of the two equations generated by
245 replacing ``t`` with positive and negative ``symbol``.
247 Notes
248 =====
250 If ``expr = 2*log(t) + x/2` then solutions for
251 ``2*log(x) + x/2 = 0`` and ``2*log(-x) + x/2 = 0`` are
252 returned by this function. Though this may seem
253 counter-intuitive, one must note that the ``expr`` being
254 solved here has been derived from a different expression. For
255 an expression like ``eq = x**2*g(x) = 1``, if we take the
256 log of both sides we obtain ``log(x**2) + log(g(x)) = 0``. If
257 x is positive then this simplifies to
258 ``2*log(x) + log(g(x)) = 0``; the Lambert-solving routines will
259 return solutions for this, but we must also consider the
260 solutions for ``2*log(-x) + log(g(x))`` since those must also
261 be a solution of ``eq`` which has the same value when the ``x``
262 in ``x**2`` is negated. If `g(x)` does not have even powers of
263 symbol then we do not want to replace the ``x`` there with
264 ``-x``. So the role of the ``t`` in the expression received by
265 this function is to mark where ``+/-x`` should be inserted
266 before obtaining the Lambert solutions.
268 """
269 nlhs, plhs = [
270 expr.xreplace({t: sgn*symbol}) for sgn in (-1, 1)]
271 sols = _solve_lambert(nlhs, symbol, gens)
272 if plhs != nlhs:
273 sols.extend(_solve_lambert(plhs, symbol, gens))
274 # uniq is needed for a case like
275 # 2*log(t) - log(-z**2) + log(z + log(x) + log(z))
276 # where substituting t with +/-x gives all the same solution;
277 # uniq, rather than list(set()), is used to maintain canonical
278 # order
279 return list(uniq(sols))
281 nrhs, lhs = f.as_independent(symbol, as_Add=True)
282 rhs = -nrhs
284 lamcheck = [tmp for tmp in gens
285 if (tmp.func in [exp, log] or
286 (tmp.is_Pow and symbol in tmp.exp.free_symbols))]
287 if not lamcheck:
288 raise NotImplementedError()
290 if lhs.is_Add or lhs.is_Mul:
291 # replacing all even_degrees of symbol with dummy variable t
292 # since these will need special handling; non-Add/Mul do not
293 # need this handling
294 t = Dummy('t', **symbol.assumptions0)
295 lhs = lhs.replace(
296 lambda i: # find symbol**even
297 i.is_Pow and i.base == symbol and i.exp.is_even,
298 lambda i: # replace t**even
299 t**i.exp)
301 if lhs.is_Add and lhs.has(t):
302 t_indep = lhs.subs(t, 0)
303 t_term = lhs - t_indep
304 _rhs = rhs - t_indep
305 if not t_term.is_Add and _rhs and not (
306 t_term.has(S.ComplexInfinity, S.NaN)):
307 eq = expand_log(log(t_term) - log(_rhs))
308 return _solve_even_degree_expr(eq, t, symbol)
309 elif lhs.is_Mul and rhs:
310 # this needs to happen whether t is present or not
311 lhs = expand_log(log(lhs), force=True)
312 rhs = log(rhs)
313 if lhs.has(t) and lhs.is_Add:
314 # it expanded from Mul to Add
315 eq = lhs - rhs
316 return _solve_even_degree_expr(eq, t, symbol)
318 # restore symbol in lhs
319 lhs = lhs.xreplace({t: symbol})
321 lhs = powsimp(factor(lhs, deep=True))
323 # make sure we have inverted as completely as possible
324 r = Dummy()
325 i, lhs = _invert(lhs - r, symbol)
326 rhs = i.xreplace({r: rhs})
328 # For the first forms:
329 #
330 # 1a1) B**B = R will arrive here as B*log(B) = log(R)
331 # lhs is Mul so take log of both sides:
332 # log(B) + log(log(B)) = log(log(R))
333 # 1a2) B*(b*log(B) + c)**a = R will arrive unchanged so
334 # lhs is Mul, so take log of both sides:
335 # log(B) + a*log(b*log(B) + c) = log(R)
336 # 1b) d*log(a*B + b) + c*B = R will arrive unchanged so
337 # lhs is Add, so isolate c*B and expand log of both sides:
338 # log(c) + log(B) = log(R - d*log(a*B + b))
340 soln = []
341 if not soln:
342 mainlog = _mostfunc(lhs, log, symbol)
343 if mainlog:
344 if lhs.is_Mul and rhs != 0:
345 soln = _lambert(log(lhs) - log(rhs), symbol)
346 elif lhs.is_Add:
347 other = lhs.subs(mainlog, 0)
348 if other and not other.is_Add and [
349 tmp for tmp in other.atoms(Pow)
350 if symbol in tmp.free_symbols]:
351 if not rhs:
352 diff = log(other) - log(other - lhs)
353 else:
354 diff = log(lhs - other) - log(rhs - other)
355 soln = _lambert(expand_log(diff), symbol)
356 else:
357 #it's ready to go
358 soln = _lambert(lhs - rhs, symbol)
360 # For the next forms,
361 #
362 # collect on main exp
363 # 2a) (b*B + c)*exp(d*B + g) = R
364 # lhs is mul, so take log of both sides:
365 # log(b*B + c) + d*B = log(R) - g
366 # 2b) g*exp(d*B + h) - b*B = R
367 # lhs is add, so add b*B to both sides,
368 # take the log of both sides and rearrange to give
369 # log(R + b*B) - d*B = log(g) + h
371 if not soln:
372 mainexp = _mostfunc(lhs, exp, symbol)
373 if mainexp:
374 lhs = collect(lhs, mainexp)
375 if lhs.is_Mul and rhs != 0:
376 soln = _lambert(expand_log(log(lhs) - log(rhs)), symbol)
377 elif lhs.is_Add:
378 # move all but mainexp-containing term to rhs
379 other = lhs.subs(mainexp, 0)
380 mainterm = lhs - other
381 rhs = rhs - other
382 if (mainterm.could_extract_minus_sign() and
383 rhs.could_extract_minus_sign()):
384 mainterm *= -1
385 rhs *= -1
386 diff = log(mainterm) - log(rhs)
387 soln = _lambert(expand_log(diff), symbol)
389 # For the last form:
390 #
391 # 3) d*p**(a*B + g) - b*B = c
392 # collect on main pow, add b*B to both sides,
393 # take log of both sides and rearrange to give
394 # a*B*log(p) - log(b*B + c) = -log(d) - g*log(p)
395 if not soln:
396 mainpow = _mostfunc(lhs, Pow, symbol)
397 if mainpow and symbol in mainpow.exp.free_symbols:
398 lhs = collect(lhs, mainpow)
399 if lhs.is_Mul and rhs != 0:
400 # b*B = 0
401 soln = _lambert(expand_log(log(lhs) - log(rhs)), symbol)
402 elif lhs.is_Add:
403 # move all but mainpow-containing term to rhs
404 other = lhs.subs(mainpow, 0)
405 mainterm = lhs - other
406 rhs = rhs - other
407 diff = log(mainterm) - log(rhs)
408 soln = _lambert(expand_log(diff), symbol)
410 if not soln:
411 raise NotImplementedError('%s does not appear to have a solution in '
412 'terms of LambertW' % f)
414 return list(ordered(soln))
417def bivariate_type(f, x, y, *, first=True):
418 """Given an expression, f, 3 tests will be done to see what type
419 of composite bivariate it might be, options for u(x, y) are::
421 x*y
422 x+y
423 x*y+x
424 x*y+y
426 If it matches one of these types, ``u(x, y)``, ``P(u)`` and dummy
427 variable ``u`` will be returned. Solving ``P(u)`` for ``u`` and
428 equating the solutions to ``u(x, y)`` and then solving for ``x`` or
429 ``y`` is equivalent to solving the original expression for ``x`` or
430 ``y``. If ``x`` and ``y`` represent two functions in the same
431 variable, e.g. ``x = g(t)`` and ``y = h(t)``, then if ``u(x, y) - p``
432 can be solved for ``t`` then these represent the solutions to
433 ``P(u) = 0`` when ``p`` are the solutions of ``P(u) = 0``.
435 Only positive values of ``u`` are considered.
437 Examples
438 ========
440 >>> from sympy import solve
441 >>> from sympy.solvers.bivariate import bivariate_type
442 >>> from sympy.abc import x, y
443 >>> eq = (x**2 - 3).subs(x, x + y)
444 >>> bivariate_type(eq, x, y)
445 (x + y, _u**2 - 3, _u)
446 >>> uxy, pu, u = _
447 >>> usol = solve(pu, u); usol
448 [sqrt(3)]
449 >>> [solve(uxy - s) for s in solve(pu, u)]
450 [[{x: -y + sqrt(3)}]]
451 >>> all(eq.subs(s).equals(0) for sol in _ for s in sol)
452 True
454 """
456 u = Dummy('u', positive=True)
458 if first:
459 p = Poly(f, x, y)
460 f = p.as_expr()
461 _x = Dummy()
462 _y = Dummy()
463 rv = bivariate_type(Poly(f.subs({x: _x, y: _y}), _x, _y), _x, _y, first=False)
464 if rv:
465 reps = {_x: x, _y: y}
466 return rv[0].xreplace(reps), rv[1].xreplace(reps), rv[2]
467 return
469 p = f
470 f = p.as_expr()
472 # f(x*y)
473 args = Add.make_args(p.as_expr())
474 new = []
475 for a in args:
476 a = _mexpand(a.subs(x, u/y))
477 free = a.free_symbols
478 if x in free or y in free:
479 break
480 new.append(a)
481 else:
482 return x*y, Add(*new), u
484 def ok(f, v, c):
485 new = _mexpand(f.subs(v, c))
486 free = new.free_symbols
487 return None if (x in free or y in free) else new
489 # f(a*x + b*y)
490 new = []
491 d = p.degree(x)
492 if p.degree(y) == d:
493 a = root(p.coeff_monomial(x**d), d)
494 b = root(p.coeff_monomial(y**d), d)
495 new = ok(f, x, (u - b*y)/a)
496 if new is not None:
497 return a*x + b*y, new, u
499 # f(a*x*y + b*y)
500 new = []
501 d = p.degree(x)
502 if p.degree(y) == d:
503 for itry in range(2):
504 a = root(p.coeff_monomial(x**d*y**d), d)
505 b = root(p.coeff_monomial(y**d), d)
506 new = ok(f, x, (u - b*y)/a/y)
507 if new is not None:
508 return a*x*y + b*y, new, u
509 x, y = y, x