Search code examples
pythonoptimizationodegekko

Why is my basic Gekko ODE solver much slower than Scipy?


In this minimal example I want to solve the basic integrator ODE dT/dt = k[F(t) - T(t)] where F(t) is a forcing vector, which is selected to be a square wave. I have implemented two procedures to solve the ODE: Using Scipy's solve_ivp and using Gekko. The former takes 6 milliseconds to run, the latter takes 8 seconds to run. Am I doing something wrong with Gekko? Are there some additional parameters that I can tune to improve performance?

I have omitted the testing part of the code, but I have compared both solutions to the analytic solution, and they are both accurate within a 3-4 significant digits.

import numpy as np
from scipy.integrate import solve_ivp
from gekko import GEKKO

def step_problem_forcing(time_arr, param):
    T_low = param['T_low']
    T_high = param['T_high']
    t_break_l = param['t_break_l']
    t_break_r = param['t_break_r']

    forcing = np.full(len(time_arr), T_low)
    idxs_pulse = np.logical_and(time_arr >= t_break_l,
                                time_arr <= t_break_r)
    forcing[idxs_pulse] = T_high
    return forcing

def step_problem_generate(tmin, tmax):
    tAvg = (tmin+tmax) / 2
    tDur = tmax - tmin
    return {
        'k' : 10 ** np.random.uniform(-3, 3),
        'T_low' : np.random.uniform(20, 50),
        'T_high' : np.random.uniform(20, 50),
        'T0' : np.random.uniform(20, 50),
        't_break_l' : np.random.uniform(tmin, tAvg - 0.1*tDur),
        't_break_r' : np.random.uniform(tAvg + 0.1*tDur, tmax)
    }

def ode_scipy_step_solver(time_arr, param: dict):
    f_this = lambda t: np.interp(t, time_arr, param['forcing'])
    
    def _dxdt(t, x, param: tuple):
        # if (t < tmin) or (t > tmax):
        #     raise ValueError(f"{t} is out of bounds [{tmin}, {tmax}]")
        k, = param
        return k*(f_this(t) - x)
    
    k = param['k']
    sol = solve_ivp(_dxdt, (time_arr[0], time_arr[-1]),
                    [param['T0'], ], t_eval=time_arr, args=([k, ],),
                    rtol=1.0E-6, method='RK45')
    return sol['y'].T

def  ode_gekko_step_solver(time_arr, param: dict) -> np.ndarray:
    m = GEKKO()  # create GEKKO model

    # Define variables
    m.time = time_arr
    T = m.Var(value=param['T0'])
    F = m.Param(value=param['forcing'])
    k = m.Const(value=param['k'])
    # t = m.Param(value=m.time)

    # equations
    m.Equation(T.dt() == k * (F - T))
    m.options.IMODE = 7  # dynamic simulation
    m.solve(disp=False)  # solve locally (remote=False)

    return T.value

tmin = 0
tmax = 10
time_arr = np.linspace(tmin, tmax, 1000)
param = step_problem_generate(tmin, tmax)
param['forcing'] = step_problem_forcing(time_arr, param)

# Takes 6.8ms to run
rezScipy = ode_scipy_step_solver(time_arr, param)

# Takes 8.5s to run
rezGekko = ode_gekko_step_solver(time_arr, param)

Solution

  • Switch to remote=False to solve locally and avoid the overhead due to network communication to the public server. Switching to IMODE=4 solves all 1000 time steps simultaneously instead of sequentially and can improve solve time. Here are 10 trials with the 1 original and 9 modified versions. I'm on Ubuntu Linux. Timing will likely be different based on the computer.

    # IMODE=7, Original Script
    ODE Scipy Time: 0.014239072799682617
    Gekko Time: 8.92684006690979
    
    # IMODE=4, Remote=False
    ODE Scipy Time: 0.2553427219390869
    Gekko Time: 0.12177252769470215
    
    ODE Scipy Time: 0.0017511844635009766
    Gekko Time: 0.11760997772216797
    
    ODE Scipy Time: 0.0032286643981933594
    Gekko Time: 0.11528849601745605
    
    ODE Scipy Time: 0.06327152252197266
    Gekko Time: 0.12140154838562012
    
    ODE Scipy Time: 0.01037144660949707
    Gekko Time: 0.12077593803405762
    
    ODE Scipy Time: 0.0036330223083496094
    Gekko Time: 0.11592268943786621
    
    ODE Scipy Time: 0.1100163459777832
    Gekko Time: 0.11649894714355469
    
    ODE Scipy Time: 0.0032815933227539062
    Gekko Time: 0.1154778003692627
    
    ODE Scipy Time: 0.0015087127685546875
    Gekko Time: 0.1158149242401123
    

    Here's the modified script:

    import numpy as np
    from scipy.integrate import solve_ivp
    from gekko import GEKKO
    import time
    
    def step_problem_forcing(time_arr, param):
        T_low = param['T_low']
        T_high = param['T_high']
        t_break_l = param['t_break_l']
        t_break_r = param['t_break_r']
    
        forcing = np.full(len(time_arr), T_low)
        idxs_pulse = np.logical_and(time_arr >= t_break_l,
                                    time_arr <= t_break_r)
        forcing[idxs_pulse] = T_high
        return forcing
    
    def step_problem_generate(tmin, tmax):
        tAvg = (tmin+tmax) / 2
        tDur = tmax - tmin
        return {
            'k' : 10 ** np.random.uniform(-3, 3),
            'T_low' : np.random.uniform(20, 50),
            'T_high' : np.random.uniform(20, 50),
            'T0' : np.random.uniform(20, 50),
            't_break_l' : np.random.uniform(tmin, tAvg - 0.1*tDur),
            't_break_r' : np.random.uniform(tAvg + 0.1*tDur, tmax)
        }
    
    def ode_scipy_step_solver(time_arr, param: dict):
        f_this = lambda t: np.interp(t, time_arr, param['forcing'])
        
        def _dxdt(t, x, param: tuple):
            # if (t < tmin) or (t > tmax):
            #     raise ValueError(f"{t} is out of bounds [{tmin}, {tmax}]")
            k, = param
            return k*(f_this(t) - x)
        
        k = param['k']
        sol = solve_ivp(_dxdt, (time_arr[0], time_arr[-1]),
                        [param['T0'], ], t_eval=time_arr, args=([k, ],),
                        rtol=1.0E-6, method='RK45')
        return sol['y'].T
    
    def  ode_gekko_step_solver(time_arr, param: dict) -> np.ndarray:
        m = GEKKO(remote=False)  # create GEKKO model
    
        # Define variables
        m.time = time_arr
        T = m.Var(value=param['T0'])
        F = m.Param(value=param['forcing'])
        k = m.Const(value=param['k'])
        # t = m.Param(value=m.time)
    
        # equations
        m.Equation(T.dt() == k * (F - T))
        m.options.IMODE = 4  # dynamic simulation
        m.solve(disp=False)  # solve locally (remote=False)
    
        return T.value
    
    tmin = 0
    tmax = 10
    time_arr = np.linspace(tmin, tmax, 1000)
    param = step_problem_generate(tmin, tmax)
    param['forcing'] = step_problem_forcing(time_arr, param)
    
    st = time.time()
    rezScipy = ode_scipy_step_solver(time_arr, param)
    print(f'ODE Scipy Time: {time.time()-st}')
    
    st = time.time()
    rezGekko = ode_gekko_step_solver(time_arr, param)
    print(f'Gekko Time: {time.time()-st}')
    

    Scipy solve_ivp is a specialized ODE solver while Gekko solves higher index DAEs, mixed integer optimization, regression, and machine learning models. It uses a simultaneous method and only has basic adaptive step lengths to control error. I recommend solve_ivp if it is an ODE problem.