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 

17 

18 

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. 

23 

24 Examples 

25 ======== 

26 

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)} 

32 

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 

45 

46 

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. 

50 

51 ``func`` can be a function (exp, log, etc...) or any other SymPy object, 

52 like Pow. 

53 

54 If ``X`` is not ``None``, then the function returns the term composed with the 

55 most ``func`` having the specified variable. 

56 

57 Examples 

58 ======== 

59 

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 

82 

83 

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``. 

88 

89 Examples 

90 ======== 

91 

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 

119 

120 

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 

151 

152 # invert the generator X1 so we have x(u) 

153 u = Dummy('rhs') 

154 xusolns = solve(X1 - u, x) 

155 

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 = [] 

166 

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. 

173 

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()] 

180 

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 

188 

189 sol.extend(xu.subs(u, rhs) for xu in xusolns) 

190 return sol 

191 

192 

193def _solve_lambert(f, symbol, gens): 

194 """Return solution to ``f`` if it is a Lambert-type expression 

195 else raise NotImplementedError. 

196 

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: 

200 

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 """ 

226 

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``. 

230 

231 Parameters 

232 ========== 

233 

234 expr : Expr 

235 The expression which includes a dummy variable t to be 

236 replaced with +symbol and -symbol. 

237 

238 symbol : Symbol 

239 The symbol for which a solution is being sought. 

240 

241 Returns 

242 ======= 

243 

244 List of unique solution of the two equations generated by 

245 replacing ``t`` with positive and negative ``symbol``. 

246 

247 Notes 

248 ===== 

249 

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. 

267 

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)) 

280 

281 nrhs, lhs = f.as_independent(symbol, as_Add=True) 

282 rhs = -nrhs 

283 

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() 

289 

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) 

300 

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) 

317 

318 # restore symbol in lhs 

319 lhs = lhs.xreplace({t: symbol}) 

320 

321 lhs = powsimp(factor(lhs, deep=True)) 

322 

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}) 

327 

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)) 

339 

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) 

359 

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 

370 

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) 

388 

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) 

409 

410 if not soln: 

411 raise NotImplementedError('%s does not appear to have a solution in ' 

412 'terms of LambertW' % f) 

413 

414 return list(ordered(soln)) 

415 

416 

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:: 

420 

421 x*y 

422 x+y 

423 x*y+x 

424 x*y+y 

425 

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``. 

434 

435 Only positive values of ``u`` are considered. 

436 

437 Examples 

438 ======== 

439 

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 

453 

454 """ 

455 

456 u = Dummy('u', positive=True) 

457 

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 

468 

469 p = f 

470 f = p.as_expr() 

471 

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 

483 

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 

488 

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 

498 

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