# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.6
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %% [markdown]
# ---
# # <center><font color=purple>SciPy -  [integrate](https://docs.scipy.org/doc/scipy/reference/integrate.html)</font> </center>
#
# ### <center>Solve ordinary differential equations (ODEs).</center>
#
#
# ---

# %% [markdown]
#
# ## Standard import

# %%
from scipy import integrate as intg


# %% [markdown]
# Notice that NumPy is not necessary. SciPy will convert native Python arrays to NumPy arrays internally.

# %% [markdown]
# ---
# # <center> <font color=darkgreen>Initial value problems: [`solve_ivp`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html#scipy.integrate.solve_ivp)</font>  </center>
# ---
#
# An *initial value problem* (IVP) consists of an ODE: 
#
# \begin{equation} \frac{d\textbf{x}}{dt} = f(t,\textbf{x}) \end{equation}
#
# and an *intial condition*,
#
# \begin{equation} \textbf{x}(t_0) = \textbf{x}_0 \end{equation}
#
#
#
# + The function $f(t,\textbf{x})$ gives us the rate of change of $\textbf{x}$, as a function of $t$ and the current value of $\textbf{x}$.
# + $\textbf{x}$ is a list containing all of the variables of the ODE.
#
# The steps are,
# 1. Define $f(t,\textbf{x})$.
# 2. Define the time span. 
# 3. Define the initial condition.
# 4. Call `solve_ivp(fun, t_span, x_initial, optional parameters ...)`.
#
# The optional parameters include:
# + `method`: choose the solver. Default is RK45.
# + `args`: Pass in parameters for the ODE.
# + `max_step`: Parameters for controlling the step size. 
# + `t_eval`: Time vector used for reporting.

# %% [markdown]
# ## Example: Lotka-Volterra model of preditor-prey dynamics
#
# + $x_1$ ... number of prey (rabbits)
# + $x_2$ ... number of predators (foxes)
#
# \begin{align}
# &\dot{x}_1 = \alpha x_1 - \beta x_1 x_2 \\
# &\dot{x}_2 = \delta x_1 x_2 -\gamma x_2 
# \end{align}
#
# + The growth rate of the population of prey ($\dot{x}_1$) is 
#     + increased by the size of the population times the birth rate ($\alpha x_1$)
#     + decreased by the rate of predation ($\beta x_1 x_2$)
#
# + The growth rate of the population of predators ($\dot{x}_2$) is 
#     + increased by the availability of food, which is proportional to predation ($\delta x_1 x_2$)
#     + decreased by the rate of death of predators ($\gamma x_2 $)
#
# Use these values for the parameters:
# \begin{align}
# \alpha &= 1.5 \\
# \beta &= 1 \\ 
# \gamma &= 3 \\
# \delta &= 1 
# \end{align}
#
#

# %%
# 1. Define the function f(t,x,parameters)
def f(t,X,a,b,c,d):
    x1, x2 = X
    x1dot = a*x1 - b*x1*x2
    x2dot = d*x1*x2 - c*x2 
    return [x1dot,x2dot]

# 2. Time span
t_span = [0,10]

# 3. Initial condition
Xo = [25,5]

# 4. Parameters
a=1.5
b=1
c=3
d=1

# 5. Call the solver, obtain the solution
sol = intg.solve_ivp(f, t_span, Xo, args=(a,b,c,d))

# %%
sol

# %% [markdown]
# `sol` will contain information about the result.
# + whether or not it succeeded: `sol.success`,
# + the solution: `sol.t`, `sol.y`
#
#
# ### Plot the solution

# %%
import matplotlib.pyplot as plt

plt.figure(figsize=(10,3))
plt.plot(sol.t,sol.y[0],label='prey')
plt.plot(sol.t,sol.y[1],label='predators')
plt.grid()
plt.xlim(t_span)
plt.ylim([0,30])
plt.legend()
plt.xlabel('time')
plt.ylabel('population')


# %% [markdown]
# ### A somewhat nicer plot

# %%
def plotsol(sol):
    plt.figure(figsize=(10,3))
    plt.plot(sol.t,sol.y[0],label='prey')
    plt.fill_between(sol.t,sol.y[0],0,alpha=0.3)
    plt.plot(sol.t,sol.y[1],label='predator')
    plt.fill_between(sol.t,sol.y[1],0,alpha=0.3)
    plt.grid()
    plt.xlim(t_span)
    plt.ylim([0,30])
    plt.legend(fontsize=16)
    plt.xlabel('time',fontsize=16)
    plt.ylabel('population',fontsize=16)
    
plotsol(sol)

# %% [markdown]
# The plot is too coarse. We have three options:
# 1. Reduce the maximum step size,
# 2. Provide the time axis to the solver,
# 3. Request a 'continuous' solution. (we won't cover this).
#
# ## Reduce the maximum step size: `max_step`

# %%
sol = intg.solve_ivp(f, t_span, Xo, args=(a,b,c,d), max_step=.01 )

plotsol(sol)

# %% [markdown]
# ## Provide the time axis to the solver: `t_eval`
# `t_eval` specifies the times where a solution will be returned.

# %%
sol = intg.solve_ivp(f, t_span, Xo, args=(a,b,c,d), t_eval=[0,1,4,6,10])

plotsol(sol)
