python 方程式の解をdfに格納していく

Question

実現したいこと

三次方程式の解をdfに格納したい

発生している問題・分からないこと

①三次方程式
X3 + A*X2 + B = 0
のA,Bそれぞれにdf["A"],df["B"]を与えた時の解Xを、df["X1"],df["X2"],df["X3"]のように格納したい
解の公式を用いるしかないでしょうか

②またその場合　連立方程式
u3 + v3 + 2*q = 0 , u * v + p = 0
のp,qそれぞれにdf["p"],df["q"]を与えた時の解u,vを、df["u"],df["v"]のように格納したい

普通に連立方程式を解く分にはsympy.solveでできるそうですが
for文を使って１行ずつ計算ではなく、以下のように簡潔にできないかと考えています
df["C"] = df["A"] + df["B"]

該当のソースコード

特になし

試したこと・調べたこと

• teratailやGoogle等で検索した
• ソースコードを自分なりに変更した
• 知人に聞いた
• その他

上記の詳細・結果

似た内容にたどり着くことができませんでした

補足

特になし

Answer 1

変数のまま sympy で解いてしまって、それを lambdify で関数化するといいのではないでしょうか。

import pandas as pd
import sympy

df = pd.DataFrame(...)

a, b, x = sympy.symbols('a b x')
x1, x2, x3 = sympy.solve(x**3 + a * x**2 + b, x)

func_x1 = sympy.lambdify([a, b], x1)
func_x2 = sympy.lambdify([a, b], x2)
func_x3 = sympy.lambdify([a, b], x3)

df['X1'] = func_x1(df['A'], df['B'])
df['X2'] = func_x2(df['A'], df['B'])
df['X3'] = func_x3(df['A'], df['B'])

for を使って書くほうがすっきり見えるかもしれません。

a, b, x = sympy.symbols('a b x')
ans = sympy.solve(x**3 + a * x**2 + b, x)
funcs = [sympy.lambdify([a, b], ans_x) for ans_x in ans]

for i, func in enumerate(funcs, 1):
    df[f'X{i}'] = func(df['A'], df['B'])

2 もだいたい同じです。(試しにやったとき nan になったので complex に変換しています)

u, v, p, q = sympy.symbols('u v p q')
ans = sympy.solve([u**3 + v**3 + 2 * q, u * v + p], [u, v])
funcs = [(sympy.lambdify([p, q], ans_u), sympy.lambdify([p, q], ans_v))
         for (ans_u, ans_v) in ans]

for i, (func_u, func_v) in enumerate(funcs, 1):
    df[f'u{i}'] = func_u(df['p'].astype(complex), df['q'].astype(complex))
    df[f'v{i}'] = func_v(df['p'].astype(complex), df['q'].astype(complex))

Comments

[hiroki]

皆様回答ありがとうございました。
そのままでは解がNanとなりましたが、a,bともにpositive=Trueとすることで解決しました！
私の扱っているデータが5000行くらいあるせいか、1行ずつ方程式を解いてたら時間がかなりかかってしまったこともあり、「先に変数のまま解いてしまって後から代入」を提案してくださった方をベストアンサーに選ばせていただきました。

Answer 2

それぞれこんな感じですかね

import pandas as pd
import sympy

x = sympy.Symbol('x')

df = pd.DataFrame(columns=['A', 'B'], data =[[1, 1], [2, 1], [3, 1], [1, 2], [1, 5], [2, 3]])
display(df)

ans = df.apply(lambda row: sympy.solve(x**3 + row['A'] * x**2 + row['B']), axis=1)
df['X1'] = [a[0] for a in ans]
df['X2'] = [a[1] for a in ans]
df['X3'] = [a[2] for a in ans]

display(df)

import pandas as pd
import sympy

u = sympy.Symbol('u')
v = sympy.Symbol('v')

df = pd.DataFrame(columns=['p', 'q'], data =[[1, 1], [2, 1], [3, 1], [1, 2], [1, 5], [2, 3]])
display(df)

ans = df.apply(lambda row: sympy.solve([u**3 + v**3 + 2 * row['q'], u * v + row['p']]), axis=1)
df['u'] = [a[0] for a in ans]
df['v'] = [a[1] for a in ans]

display(df)

Comments

[hiroki]

回答ありがとうございました。

Answer 3

記述例を下記に示します。

①（3次元方程式）の解を求めるには numpy.roots() を用い，②（3次元連立方程式）の解を求めるには sympy.solve() を用いました。

import pandas as pd
import numpy as np
import sympy as sy

df = pd.DataFrame({'c0': [1, 1, 1],
                   'A': [1, 0, 1],
                   'c2': [0, 0, 0],
                   'B': [0, 1, 1]}, dtype=float)

df[['X1', 'X2', 'X3']] = [np.roots(c) for c in df.to_numpy()]

print(df)
#     c0    A   c2    B                  X1                  X2                  X3
# 0  1.0  1.0  0.0  0.0 -1.000000+0.000000j  0.000000+0.000000j  0.000000+0.000000j
# 1  1.0  0.0  0.0  1.0 -1.000000+0.000000j  0.500000+0.866025j  0.500000-0.866025j
# 2  1.0  1.0  0.0  1.0 -1.465571+0.000000j  0.232786+0.792552j  0.232786-0.792552j


df2 = pd.DataFrame({'p': [1, 0, 1],
                    'q': [0, 1, 1]}, dtype=float)

sy.var('u v');
result = [sy.solve([u**3 + v**3 + 2 * q, u * v + p])
          for p, q in df2.to_numpy()]
df2[['u', 'v']] = [[[complex(d[u]) for d in rs], [complex(d[v]) for d in rs]]
                   for rs in result]

for id in df2.index:
    print("p =", df2.loc[id, 'p'], " q =", df2.loc[id, 'q'])
    print("u:", df2.loc[id, 'u'])
    print("v:", df2.loc[id, 'v'])
# p = 1.0  q = 0.0
# u: [(-1+0j), (1+0j), (-0.5-0.8660254037844386j), (-0.5+0.8660254037844386j), (0.5-0.8660254037844386j), (0.5+0.8660254037844386j)]
# v: [(1+0j), (-1+0j), (0.5-0.8660254037844386j), (0.5+0.8660254037844386j), (-0.5-0.8660254037844386j), (-0.5+0.8660254037844386j)]
# p = 0.0  q = 1.0
# u: [(-1.2599210498948732+0j), 0j, 0j, 0j, (0.6299605249474366-1.0911236359717214j), (0.6299605249474366+1.0911236359717214j)]
# v: [0j, (-1.2599210498948732+0j), (0.6299605249474366-1.0911236359717214j), (0.6299605249474366+1.0911236359717214j), 0j, 0j]
# p = 1.0  q = 1.0
# u: [(-1.3415037626305777+0j), (0.7454321246472562+0j), (-0.3727160623236281-0.6455631567415321j), (-0.3727160623236281+0.6455631567415321j), (0.6707518813152888-1.1617763377104897j), (0.6707518813152888+1.1617763377104897j)]
# v: [(0.7454321246472562+0j), (-1.3415037626305777+0j), (0.6707518813152888-1.1617763377104897j), (0.6707518813152888+1.1617763377104897j), (-0.3727160623236281-0.6455631567415321j), (-0.3727160623236281+0.6455631567415321j)]

Comments

[hiroki]

回答ありがとうございました。