Search code examples
pythonmatrixsympy

How to collect the coefficients of sympy matrix?

I want to collect the coefficients of the variables denoted by {r11,r12,r21,r22} in the below matrix denoted eqs, and find the coefficient matrix. I am pasting the code below to get eqs. I tried using coeff(). I couldn't get what I wanted. In Mathematica for example, this can be achieved by the command "CoefficientArrays"

import sympy as sp
r11,r12,r21,r22,rx,ry,k,s,dt,omega,gamma = sp.symbols('r11,r12,r21,r22,rx,ry,k,s,dt,omega,gamma')
Ad=sp.Matrix([[0,1],[0,0]]); Aa= sp.Matrix([[0,0],[1,0]]);H0= omega*sp.Matrix([[1,0],[0,0]])
rho= sp.Matrix([[r11,r12],[r21,r22]]);sx= sp.Matrix([[0,1],[1,0]])
eqs=(sp.simplify((dt*(-1j*(H0*rho-rho*H0) + k*sp.exp(-s)*Aa*rho*Ad- k/2 *(Ad*Aa*rho +rho*Ad*Aa)) -gamma*(sx*(sx*rho- rho*sx) -(sx*rho -rho*sx)*sx)*dt)/dt))
Solution

  • You can rewrite a system of equations in the A * x = b form using the linear_eq_to_matrix function:

    import sympy as sp
sp.init_printing()
r11,r12,r21,r22,rx,ry,k,s,dt,omega,gamma = sp.symbols('r11,r12,r21,r22,rx,ry,k,s,dt,omega,gamma')
Ad=sp.Matrix([[0,1],[0,0]]); Aa= sp.Matrix([[0,0],[1,0]]);H0= omega*sp.Matrix([[1,0],[0,0]])
rho= sp.Matrix([[r11,r12],[r21,r22]]);sx= sp.Matrix([[0,1],[1,0]])
eqs=(sp.simplify((dt*(-1j*(H0*rho-rho*H0) + k*sp.exp(-s)*Aa*rho*Ad- k/2 *(Ad*Aa*rho +rho*Ad*Aa)) -gamma*(sx*(sx*rho- rho*sx) -(sx*rho -rho*sx)*sx)*dt)/dt))

x = [r11,r12,r21,r22]
A, b = sp.linear_eq_to_matrix(eqs.nsimplify(), x)

    results of linear_eq_to_matrix

    Note that I used nsimplify() to convert the 1.0i number into i. Ideally, when using sympy you should write the imaginary part as sp.I instead of 1j.