#!/usr/bin/env python
# coding: utf-8

# # Stability and Stability Functions
# 
# Copyright (C) 2026 Andreas Kloeckner
# 
# <details>
# <summary>MIT License</summary>
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# </details>

# In[1]:


import numpy as np
import numpy.linalg as la
import matplotlib.pyplot as pt


# In[2]:


A = np.array([[0, 1], [-1, 0]])
def harmonic_rhs(yvec):
    return A@yvec


# Write out the step functions for forward and backward Euler as functions `step_XXX(y, h, f)`:

# In[3]:


def step_fw(y, h, f):
    return y + h*f(y)


def step_bw(y, h, f):
    return la.solve(np.eye(2) - h*A, y)


# Now write Heun and RK4:
# $$
# \begin{array}{c|cc}
# 0 & & \\
# 1 & 1 & \\ \hline
#   & \tfrac{1}{2} & \tfrac{1}{2}
# \end{array}
# \qquad
# \begin{array}{c|cccc}
# 0          &             &             &             & \\
# \tfrac{1}{2} & \tfrac{1}{2} &             &             & \\
# \tfrac{1}{2} & 0           & \tfrac{1}{2} &             & \\
# 1          & 0           & 0           & 1           & \\ \hline
#            & \tfrac{1}{6} & \tfrac{1}{3} & \tfrac{1}{3} & \tfrac{1}{6}
# \end{array}
# $$

# In[4]:


def step_heun(y, h, f):
    k1 = f(y)
    k2 = f(y + h*k1)
    return y + h/2*(k1 + k2)

def step_rk4(y, h, f):
    k1 = f(y)
    k2 = f(y + h/2*k1)
    k3 = f(y + h/2*k2)
    k4 = f(y + h*k3)
    return y + h/6*(k1 + 2*k2 + 2*k3 + k4)


# Now write out the stability functions: (see below for Heun, RK4 and trapezoidal):

# In[5]:


def stab_fw(z):
    return 1 + z

def stab_bw(z):
    return 1/(1 - z)

def stab_heun(z):
    return (1/2)*z**2 + z + 1

def stab_rk4(z):
    return (1/24)*z**4 + (1/6)*z**3 + (1/2)*z**2 + z + 1

def stab_trapezoidal(z):
    return -z**2/(2*z - 4) + (1/2)*z - z/(z - 2) + 1


# In[26]:


if 0:
    step = step_fw
    stab = stab_fw
elif 0:
    step = step_bw
    stab = stab_bw
elif 0:
    step = step_heun
    stab = stab_heun
elif 1:
    step = step_rk4
    stab = stab_rk4
elif 0:
    step = None
    stab = stab_trapezoidal


# In[27]:


times = [0]
ys = [np.array([0.1, 0.9])]

dt = 0.01

while times[-1] < 6*np.pi:
    y = ys[-1]
    ynext = step(y, dt, harmonic_rhs)
    ys.append(ynext)
    times.append(times[-1] + dt)

ys = np.array(ys)
times = np.array(times)


# In[28]:


pt.plot(times, ys[:, 0])


# In[29]:


ext = 3
re, im = np.mgrid[-ext:ext:400j, -ext:ext:400j]
z = re + 1j*im

pt.imshow(np.log10(np.abs(stab(z))[..., ::-1].T), extent=(-ext, ext, -ext, ext))
pt.colorbar()
pt.contour(np.log10(np.abs(stab(z))[..., ::-1].T), extent=(-ext, ext, -ext, ext), levels=[0], colors=["r"])
pt.gca().set_aspect("equal")


# Why are these symmetric about the $x$ axis?

# ### Order stars
# 
# Plot the order star for each method:

# In[30]:


order_star = stab(z)/np.exp(z)

pt.imshow(np.log10(np.abs(order_star)[..., ::-1].T), extent=(-ext, ext, -ext, ext))
pt.colorbar()
pt.contour(np.log10(np.abs(order_star)[..., ::-1].T), extent=(-ext, ext, -ext, ext), levels=[0], colors=["r"])
pt.gca().set_aspect("equal")


# ## Deriving stability functions

# In[118]:


import sympy as sp
from sympy.printing.numpy import NumPyPrinter

h = sp.Symbol("h")
y = sp.Symbol("y")
ynext = sp.Symbol("y_n")
lam = sp.Symbol("lamda")
k1, k2 = sp.symbols("k1, k2")
z = sp.Symbol("z")

gen_code = NumPyPrinter().doprint

def rhs(y):
    return lam*y


# ### Explicit
# 
# Use `.expand()` and `.subs(h*lam, z)` to find the stability function. Make sure to divide by `y`.
# 
# Do this for both Heun and RK4.

# In[119]:


heun(y, h, rhs).expand().subs(h*lam, z)


# In[120]:


print(gen_code((step_heun(y, h, rhs)/y).expand().subs(h*lam, z)))


# In[121]:


print(gen_code((step_rk4(y, h, rhs)/y).expand().subs(h*lam, z)))


# Do you notice something about these stability functions?
# 
# <details>
#     They're truncations of $e^x$.
# </details>

# ### Implicit

# Derive the stability function for the trapezoidal rule:
# 
# $$
# \begin{array}{c|cc}
# 0 & 0           & 0           \\
# 1 & \tfrac{1}{2} & \tfrac{1}{2} \\ \hline
#   & \tfrac{1}{2} & \tfrac{1}{2}
# \end{array}
# $$
# 
# Set up the equations for the stage values:

# In[124]:


equations = [
    k1 - rhs(y),
    k2 - rhs(y + h/2*(k1+k2))
]
equations


# Solve for `ynext`:

# In[126]:


soln = sp.solve(equations, [k1, k2])
soln


# Assemble the overall scheme:

# In[128]:


trapezoidal = y + h/2*(soln[k1] + soln[k2])
trapezoidal


# And find the stability function as above:

# In[129]:


print(gen_code((trapezoidal/y).expand().subs(h*lam, z)))


# In[ ]: