Solve a system of non-linear equations in Python (scipy.optimize.fsolve)

[scidam]

In case of system of ordinary differential equations you will faced with necessity to solve
algebraic system of size m*s , where m -- the number of differential equations, s -- the number of stages in rk-method.

I slightly modified the code above to be able to handle systems of ODEs, but it still includes hardcoded
components (e.g. the number of stages, s=2), so improvements definitely needed:

import numpy as np
from scipy.optimize import fsolve , newton_krylov
import matplotlib.pyplot as plt



class ImpRK4 :
    def __init__(self, functions, t0, tf, dt, y0):
        self.functions = functions
        self.t0 = t0
        self.tf = tf
        self.dt = dt
        self.u0 = y0
        self.n = round((tf - t0) / dt)
        self.time  = np.linspace(self.t0, self.tf, self.n+1)
        self.u = np.array([self.u0  for i in range(self.n+1)])
        self.m = len(functions) # the number of equations (you can estimate it)
        
    def ode(self, t, u):
        return  np.array([func(t, u) for func in self.functions]) # u is an array of size `m`!
    
    def implicit_equations(self, variables, ti, ui):
        k1, k2 = variables[:self.m], variables[self.m:]
        f1 = -k1 + self.ode(ti + (0.5 + np.sqrt(3) / 6) * self.dt, ui + 0.25 * self.dt * k1 + (0.25 + np.sqrt(3) / 6) * self.dt * k2)
        f2 = -k2 + self.ode(ti + (0.5 - np.sqrt(3) / 6) * self.dt, ui + (0.25 - np.sqrt(3) / 6) * self.dt * k1 + 0.25 * self.dt * k2)
        return np.hstack([f1,f2]).ravel()
 
    def solve(self): 
        s = 2
        for i in range(len(self.time) - 1):
            res = fsolve(self.implicit_equations, np.random.rand(self.m * s), args=(self.time[i], self.u[i])) # s =2, the number of stadies of the num method
            k1, k2 = res[:self.m], res[self.m:]
            self.u[i + 1] = self.u[i] + self.dt / 2 * (k1 + k2)
        print('The system is solved.')



def main():
    func00 = lambda t,u : -10*(t-1)*u[0]
    func01 = lambda t,u : u[1] 
    func02 = lambda t,u : (1-u[0]**2)*u[1] - u[0]
    func0x = np.array([func00])
    func0 = np.array([func01,func02])
    t0 = 0. 
    tf = 2.      
    y01 = np.array([1., 1.])
    dt = 0.008 
    diffeq = ImpRK4(func0,t0,tf,dt,y01)    
    diffeq.solve()
    plt.plot(diffeq.u[:,1], diffeq.u[:,0])
    plt.show()
 
 
if __name__ == '__main__': 
    main()