Solving a system of matrix differential equations in Python

Question

I'm interesting in solving this system of differential equations( this is a simplification of my real problem), the goal is: write two matrices ,with the solutions that are "Jsol" and "Cmatrix", and plot a table of SS, that is the multiplication of this matrices (np.abs(Jsol @ Cmatrix)), for each value of a*Hubble/k.

In wolfram mathematica i can make this easily:

ClearAll["Global*"]
    n = 2;
    Ntot = 1; 
    V = \[Lambda]*(\[Phi]^(2*n)/(2*n)); 
    dV = D[V, \[Phi]]; 
    ddV = D[dV, \[Phi]]; 
    H[t_] = Sqrt[(1/3/Mp^2)*(d\[Phi][t]^2/2 + (V /. {\[Phi] -> \[Phi][t]}) + \[Rho][t])]; 
    Mp = 1; 
    Cr = gstar*(Pi^2/30); 
    V0 = \[Lambda]; 
    gstar = 12.5; 
    (* Initial conditions *)`

    \[Lambda] = 4.0394888902589096*^-15; 
    Cupsilon = 0.014985474358746776; 
    \[Phi]0 = 12.327368461463733; 
    d\[Phi]0 = -7.956663634476879*\[Lambda]^(1/2); 
    \[Rho]0 = 36.962219515053384*\[Lambda]; 
    H0 = Sqrt[(1/3/Mp^2)*(d\[Phi]0^2/2 + (V /. {\[Phi] -> \[Phi]0}) + \[Rho]0)]; 
    T[t_] = (\[Rho][t]/Cr)^(1/4); 
    T0 = T[t] /. {\[Rho][t] -> \[Rho]0}; 
    \[Alpha] = 0; 
    \[Beta] = 1; 
    Q[t_] = Cupsilon*\[Phi][t]^\[Alpha]*(T[t]^\[Beta]/(3*H[t])); 
    \[CapitalGamma][t_] = 
    Cupsilon*Mp*(\[Phi][t]/Mp)^\[Alpha]*(T[t]/Mp)^\[Beta]; 
    \[CapitalGamma]T[t_] = \[Beta]*
    Cupsilon*(T[t]/Mp)^(-1 + \[Beta])*(\[Phi][t]/Mp)^\[Alpha]; 
    (\[CapitalGamma]\[Phi][t_] = 0; ); 

    k = 100*H0; 

    f\[Rho][t_] = 1/(6*Mp^2*H[t]^2); 
    g\[Rho][t_] = 
    4 - \[CapitalGamma]T[t]*H[t]*
    T[t]*((d\[Phi][t]/H[t])^2/(4*\[Rho][t])) - k^2/(3*a[t]^2*H[t]^2); 
    G\[Rho][t_] = g\[Rho][t] + k^2/(3*a[t]^2*H[t]^2); 
    A[t_] = {{G\[Rho][t] + 4*\[Rho][t]*
    f\[Rho][t], (-H[t])*(k^2/(a[t]^2*H[t]^2))}, {1/(3*H[t]), 3}}; 
    B[t_] = {{(-(d\[Phi][t]/H[t]))*
    Sqrt[2*\[CapitalGamma][t]*T[t]*(H[t]/a[t]^3)]}, {0}}; 
    J[t_] = {{j11[t], j12[t]}, {j21[t], j22[t]}}; 

    (*System*)

    sys = {Derivative[1][\[Phi]][t] == d\[Phi][t]/H[t], Derivative[1][d\[Phi]][
     t] == -3*(1 + Q[t])*d\[Phi][t] - (dV /. {\[Phi] -> \[Phi][t]})/
      H[t], Derivative[1][\[Rho]][t] == -4*\[Rho][t] + 
     3*Q[t]*d\[Phi][t]^2, 
       Derivative[1][a][t] == a[t], 
     Derivative[1][J][t] == -A[t] . J[t] - J[t] . Transpose[A[t]] + 
     B[t] . Transpose[B[t]]}; 

     ICs = {\[Phi][0] == \[Phi]0, 
     d\[Phi][0] == d\[Phi]0, \[Rho][0] == \[Rho]0, a[0] == 1, 
     J[0] == (1/(2*k))*{{0, 0}, {0, 0}}}; 

     soln = NDSolve[{sys, ICs}, 
     Flatten[{\[Phi], d\[Phi], \[Rho], J[t], a}], {t, 0, Ntot}, 
     Method -> {"StiffnessSwitching"}, AccuracyGoal -> 16, 
     PrecisionGoal -> 16]; 

     (*Take the solutions*)

    phi[t_] = Re[First[\[Phi][t] /. soln]]; 
    dphi[t_] = Re[First[d\[Phi][t] /. soln]]; 
    rad[t_] = Re[First[\[Rho][t] /. soln]]; 
    Hubble[t_] = Re[Sqrt[(1/3/Mp^2)*(dphi[t]^2/2 + (V /. {\[Phi] -> phi[t]}) + 
       rad[t])]]; 
    scale[t_] = Re[First[a[t] /. soln]]; 
    epsilon0[t_] = -D[H[t], t]/
     H[t] /. {Derivative[1][\[Phi]][t] -> d\[Phi][t]/H[t], 
     Derivative[1][d\[Phi]][
       t] -> -3*(1 + \[CapitalGamma][t]/(3*H[t]))*
        d\[Phi][t] - (dV /. {\[Phi] -> \[Phi][t]})/H[t], 
           
     Derivative[1][\[Rho]][
       t] -> -4*\[Rho][t] + (\[CapitalGamma][t]/H[t])*
        d\[Phi][t]^2} /. {\[Phi][t] -> phi[t], 
    d\[Phi][t] -> dphi[t], \[Rho][t] -> rad[t]}; 
    Qdiss0[t_] = 
    Re[Q[t] /. {\[Phi][t] -> phi[t], 
     d\[Phi][t] -> dphi[t], \[Rho][t] -> rad[t]}]; 


    J11[t_] = First[j11[t] /. soln]; 
    J12[t_] = First[j12[t] /. soln]; 
    J21[t_] = First[j21[t] /. soln]; 
    J22[t_] = First[j22[t] /. soln];

    (*create the matrix with solutions*)
 
    Jsol[t_] = {{J11[t], J12[t]}, {J21[t], J22[t]}};
 
    Cmatrix[t_] = (1/(3*dphi[t]^2 + 4*rad[t]))*{{0}, {3*Hubble[t]}};
 
    Spectrum[t_] = (k^3/(2*Pi^2))*
    Transpose[Cmatrix[t]] . Jsol[t] . Cmatrix[t]; 

     (*create a table*)

    tabela = Table[{scale[t]*(Hubble[t]/k), Abs[Spectrum[t]]}, {t, 0, 
    Ntot, 0.1}]; 
    Data = tabela; 

    simplifiedData = Replace[Data, {a_, {{b_}}} :> {a, b}, {1}]; 

    TableForm[Data, TableHeadings -> {Style, {"aH/k", "Spectrum"}}];

     (*Plot*)
 
    ListLinePlot[Abs[simplifiedData], Mesh -> All, 
    MeshStyle -> {Opacity[0.5, Darker[Blue]], PointSize[0.00015]}, 
    MeshShading -> {Automatic}, Frame -> True, 
     FrameLabel -> {"aH/k", "|<\[ScriptCapitalP]>|"}, 
    PlotLabel -> "|<\[ScriptCapitalP]>| vs. aH/k", 
    GridLines -> Automatic, GridLinesStyle -> Directive[Dotted, Gray], 
    PlotStyle -> {Opacity[0.9, Pink], Thickness -> 0.007}, 
    ImageSize -> 500]

But in Python i have many difficulties.

Python code:

    import numpy as np
    import matplotlib.pyplot as plt
    from scipy.integrate import solve_ivp
    import numpy as np
    from scipy.integrate import odeint
    import sympy as sp

Numerical constants:

    Mp=1
    n=2
    Ntotal=10
    Lambda= 4.0394888902589096*10**(-15)
    Cupsilon= 0.014985474358746776

Initial conditions:

    phi0=12.327368461463733
    dphi0=-7.95666363447687*Lambda**(1/2)
    rad0=36.962219515053384*Lambda
    a0=1

    J11_0= 0
    J12_0= 0

    J21_0= 0
    J22_0= 0

Building a system of differential equations:

     def system_matricial_m(w, t):
        phi, dphi, rad, a,J11, J12,J21, J22= w

        pot= Lambda*phi**(2*n)/(2*n)
        dpot= Lambda*phi**(2*n-1)
        ddpot = Lambda*(2*n-1)*phi**(2*n-2)
        dpot0= Lambda*phi0**(2*n-1)
 
        H = np.sqrt(Mp**2/2*(dphi**2/2+dpot+rad))
        H0 = np.sqrt(Mp**2/2*(dphi0**2/2+dpot0+rad0))
        gstar=12.5
        Cr = gstar*np.pi**2/30
        T=(rad/Cr)**(1/4);
        k=100*H0

        Alpha=0
        Beta=1

        Q=(Cupsilon*phi**(Alpha)*T**Beta)/(3*H)
        gamma= Cupsilon*phi**(Alpha)*T**Beta
        gammaT=Beta*Cupsilon*T**(-1+Beta)*(phi/Mp)**Alpha
        gammaPhi=0

        frho=1/(6*Mp**2*H**2)
        grho=4 - gammaT*H*T*((dphi/H))**2/(4*rad) - k**2/(3*a**2*H**2)
        hrho=T*gammaT/(4*rad*H)*(dphi/H)
        Grho=grho + k**2/(3*a**2*H**2)

        A = np.array([[Grho+4*rad*frho,-H*k**2/(a**2*H**2)],
                     [1/(3*H),3]])

        B=np.array([[-(dphi/H)*np.sqrt(2*gamma*T*H/a**3)],[0]])
        J = np.array([[J11, J12], [J21, J22]])

        dphidt = dphi/H
        ddphidt = -3*(1+Q)*dphi-dpot/H
        draddt = -4*rad+3*Q*dphi**2
        dadt=a
        dJdt =- (np.multiply(A, J) -np.multiply( np.transpose(A) , J))+np.multiply(B,np.transpose(B))

        dwdt = [dphidt, ddphidt, draddt,dadt,
                dJdt[0, 0], dJdt[0, 1],  
                dJdt[1, 0], dJdt[1, 1]]

        return dwdt

Initial conditions:

 w0 = [phi0, dphi0, rad0, a0, J11_0, J12_0,J21_0, J22_0]

time interval:

    t=np.arange(0,60)

Solving the system:

    sol = odeint(system_matricial_m, w0, t)

Take the solutions:

    PHI = sol[:, 0]
    DPHI = sol[:, 1]
    RAD= sol[:, 2]
    scale = sol[:, 3]
    J11 = sol[:, 4]
    J12 = sol[:, 5]
    J21 = sol[:, 6]
    J22 = sol[:, 7]

    k=100
    gstar=12.5
    Cr = gstar*np.pi**2/30
    TEMP=(RAD/Cr)**(1/4)

    DPOT=Lambda*PHI**(2*n-1)
    GAMMA= Cupsilon*PHI**(0)*TEMP**(1)
    HUBBLE=np.real(np.sqrt(Mp**2/2*(DPHI**2/2+DPOT+RAD)))
    Q=GAMMA/(3*HUBBLE)
    epsilon0=-(DPHI**2*GAMMA/HUBBLE-4*RAD+(-3*DPHI*(1+Q)-DPOT/HUBBLE)*DPHI+(4.03949*10**(-15)*DPHI*PHI**3/HUBBLE))/(2*(DPHI**2/2+RAD+1.00987222*10**(-15)*PHI**4))

Writing the matrices of solutions:

    print("Dimension of phi:", PHI.shape)
    print("Dimension of dphi:", DPHI.shape)
    print("Dimension of Jsol:", RAD.shape)
    print("Dimension of  rad:", scale.shape)
    print("Dimension of  J11:", J11.shape)

    Jsol = np.array([[J11, J12], [J21, J22] ])

    Cmatrix = (1 / (3 * DPHI**2 + 4 * RAD)) * np.array([[0], [3 * HUBBLE]])

    SS= np.abs(Jsol@Cmatrix)

if i understand correctly, the problem is that the shape of jsol and Cmatrix are different, but i don't understand why this occurs. In Mathematica, solving the system and take the solutions are a directly process. In python for some reason that i cant see, the system dont return the correct dimension. How can i solve this?

In resume: I tried to solve a system of differential equations with a matrix ODE in python, in wolfram mathematica I managed to do this, I took the solutions, created two matrices, Jsol and Cmatrix, and then plotted a table of the values of the multiplication module of Jsol and Cmatrix for each value of aH/k (which are also solutions), my idea was to do the same procedure in Python but for some reason the Jsol and Cmatrix matrices, when multiplied, return an error saying that their formats are different. I don't understand why this happens if in mathematica the code runs well.

Answer 1

Jsol should have dimension 2x2xN, while you should get an error of incompatible lengths in the construction of Cmatrix. If you correct this, for instance via np.array([[0 * HUBBLE], [3 * HUBBLE]]), then you get an array of shape 2x1xN. In the matrix-vector product you want to multiply and sum up over the middle index and the first index.

The default however is to try to join over the last and second-to-last index, here N and 1, as the documentation says. The idea is that you have a collection/stack/array or other structure of matrices in both operands.

You can now try to rotate around the dimensions, or you can use the np.tensordot function where you can specify which axis to use.