Очень низкая производительность в моем сценарии matplotlib

Question

Мой код работает очень плохо. Я едва получаю более 10 кадров в секунду при смене вещей на слайдере. Конечно, я не очень хорошо разбираюсь в matplotlib, но кто-то может указать, что я делаю неправильно и как это исправить?

Примечание: я обрабатываю много данных, около 3 * 100000 точек в наихудший сценарий ... Также не уверен, нужно ли это, но я работаю на бэкенде 'TkAgg'.

Вот мой код (это код для построения и запуска математической модели эпидемиологии SIR):

import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button
import matplotlib.patches as patches

p = 1                                                       #population
i = 0.01*p                                                  #infected
s = p-i                                                     #susceptible
r = 0                                                       #recovered/removed

a = 3.2                                                     #transmission parameter
b = 0.23                                                    #recovery parameter

initialTime = 0
deltaTime = 0.001                                           #smaller the delta, better the approximation to a real derivative
maxTime = 10000                                             #more number of points, better is the curve generated

def sPrime(oldS, oldI, transmissionRate):                   #differential equations being expressed as functions to
    return -1*((transmissionRate*oldS*oldI)/p)              #calculate rate of change between time intervals of the
                                                            #different quantities i.e susceptible, infected and recovered/removed
def iPrime(oldS, oldI, transmissionRate, recoveryRate):             
    return (((transmissionRate*oldS)/p)-recoveryRate)*oldI

def rPrime(oldI, recoveryRate):
    return recoveryRate*oldI

maxTimeInitial = maxTime

def genData(transRate, recovRate, maxT):
    global a, b, maxTimeInitial
    a = transRate
    b = recovRate
    maxTimeInitial = maxT

    sInitial = s
    iInitial = i
    rInitial = r

    time = []
    sVals = []
    iVals = []
    rVals = []

    for t in range(initialTime, maxTimeInitial+1):              #generating the data through a loop
        time.append(t)
        sVals.append(sInitial)
        iVals.append(iInitial)
        rVals.append(rInitial)

        newDeltas = (sPrime(sInitial, iInitial, transmissionRate=a), iPrime(sInitial, iInitial, transmissionRate=a, recoveryRate=b), rPrime(iInitial, recoveryRate=b))
        sInitial += newDeltas[0]*deltaTime
        iInitial += newDeltas[1]*deltaTime
        rInitial += newDeltas[2]*deltaTime

        if sInitial < 0 or iInitial < 0 or rInitial < 0:        #as soon as any of these value become negative, the data generated becomes invalid
            break                                               #according to the SIR model, we assume all values of S, I and R are always positive.

    return (time, sVals, iVals, rVals)

fig, ax = plt.subplots()
plt.subplots_adjust(bottom=0.4, top=0.94)

plt.title('SIR epidemiology curves for a disease')

plt.xlim(0, maxTime+1)
plt.ylim(0, p*1.4)

plt.xlabel('Time (t)')
plt.ylabel('Population (p)')

initialData = genData(a, b, maxTimeInitial)

susceptible, = ax.plot(initialData[0], initialData[1], label='Susceptible', color='b')
infected, = ax.plot(initialData[0], initialData[2], label='Infected', color='r')
recovered, = ax.plot(initialData[0], initialData[3], label='Recovered/Removed', color='g')

plt.legend()

transmissionAxes = plt.axes([0.125, 0.25, 0.775, 0.03], facecolor='white')
recoveryAxes = plt.axes([0.125, 0.2, 0.775, 0.03], facecolor='white')
timeAxes = plt.axes([0.125, 0.15, 0.775, 0.03], facecolor='white')

transmissionSlider = Slider(transmissionAxes, 'Transmission parameter', 0, 10, valinit=a, valstep=0.01)
recoverySlider = Slider(recoveryAxes, 'Recovery parameter', 0, 10, valinit=b, valstep=0.01)
timeSlider = Slider(timeAxes, 'Max time', 0, 100000, valinit=maxTime, valstep=1, valfmt="%i")

def updateTransmission(newVal):
    newData = genData(newVal, b, maxTimeInitial)

    susceptible.set_ydata(newData[1])
    infected.set_ydata(newData[2])
    recovered.set_ydata(newData[3])

    r_o.set_text(r'$R_O$={:.2f}'.format(a/b))

    fig.canvas.draw_idle()

def updateRecovery(newVal):
    newData = genData(a, newVal, maxTimeInitial)

    susceptible.set_ydata(newData[1])
    infected.set_ydata(newData[2])
    recovered.set_ydata(newData[3])

    r_o.set_text(r'$R_O$={:.2f}'.format(a/b))

    fig.canvas.draw_idle()

def updateMaxTime(newVal):
    global susceptible, infected, recovered

    newData = genData(a, b, int(newVal.item()))

    del ax.lines[:3]

    susceptible, = ax.plot(newData[0], newData[1], label='Susceptible', color='b')
    infected, = ax.plot(newData[0], newData[2], label='Infected', color='r')
    recovered, = ax.plot(newData[0], newData[3], label='Recovered/Removed', color='g')

transmissionSlider.on_changed(updateTransmission)
recoverySlider.on_changed(updateRecovery)
timeSlider.on_changed(updateMaxTime)

resetAxes = plt.axes([0.8, 0.025, 0.1, 0.05])
resetButton = Button(resetAxes, 'Reset', color='white')

r_o = plt.text(0.1, 1.5, r'$R_O$={:.2f}'.format(a/b), fontsize=12)

def reset(event):
    transmissionSlider.reset()
    recoverySlider.reset()
    timeSlider.reset()

resetButton.on_clicked(reset)

plt.show()

Answer 1

Используйте соответствующий решатель ODE, например, scipy.integrate.odeint для скорости. Тогда вы можете использовать большие временные шаги для вывода. С неявным решателем, таким как odeint или solve_ivp с method="Radau", координатные плоскости, которые являются границами в точном решении, также будут границами в численном решении, так что значения никогда не станут отрицательными.

Сокращение набора графических данных для соответствия фактическому разрешению изображения графика. Разница от 300 до 1000 баллов все еще может быть видимой, не будет видимой разницы от 1000 до 5000 баллов, вероятно, даже не фактическая разница.

matplotlib dr aws его изображения через дерево сцены в качестве объектов, используя медленную python итерацию. Это делает его очень медленным, если нужно нарисовать более пары 10000 объектов, поэтому лучше ограничить количество деталей этим числом.

Код для решения ODE

для решения ODE Я использовал solve_ivp, но без разницы, если используется odeint,

def SIR_prime(t,SIR,trans, recov): # solver expects t argument, even if not used
    S,I,R = SIR
    dS = (-trans*I/p) * S 
    dI = (trans*S/p-recov) * I
    dR = recov*I
    return [dS, dI, dR]

def genData(transRate, recovRate, maxT):
    SIR = solve_ivp(SIR_prime, [0,maxT], [s,i,r], args=(transRate, recovRate), method="Radau", dense_output=True)
    time = np.linspace(0,SIR.t[-1],1001)
    sVals, iVals, rVals = SIR.sol(time)
    return (time, sVals, iVals, rVals)

Оптимизированный код для процедур обновления графика

Можно удалить большую часть дублированного кода. Я также добавил строку, чтобы ось времени изменялась с помощью переменной maxTime, чтобы можно было действительно увеличивать

def updateTransmission(newVal):
    global trans_rate
    trans_rate = newVal
    updatePlot()

def updateRecovery(newVal):
    global recov_rate
    recov_rate = newVal
    updatePlot()

def updateMaxTime(newVal):
    global maxTime
    maxTime = newVal
    updatePlot()

def updatePlot():
    newData = genData(trans_rate, recov_rate, maxTime)

    susceptible.set_data(newData[0],newData[1])
    infected.set_data(newData[0],newData[2])
    recovered.set_data(newData[0],newData[3])

    ax.set_xlim(0, maxTime+1)

    r_o.set_text(r'$R_O$={:.2f}'.format(trans_rate/recov_rate))

    fig.canvas.draw_idle()

Код между ними и вокруг остается неизменным.