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

# Here's a tweet I saw last week:
# 
# <blockquote class="twitter-tweet" data-lang="en"><p lang="en" dir="ltr">Durand–Kerner method: surprisingly simple way of computing all the roots of a polynomial. <a href="https://t.co/Ikv22o5jDO">https://t.co/Ikv22o5jDO</a> <a href="https://t.co/BAoI0ayfwi">pic.twitter.com/BAoI0ayfwi</a></p>&mdash; Gabriel Peyré (@gabrielpeyre) <a href="https://twitter.com/gabrielpeyre/status/923832635193024512?ref_src=twsrc%5Etfw">October 27, 2017</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script> 
# 
# This looks interesting: a method for finding roots in the complex plane. Let's see if we can implement it in Python.

# [Wikipedia](https://en.wikipedia.org/wiki/Durand–Kerner_method) has a nice post on the topic. It gives the following example polynomial:
# 
# $$
# x^3 − 3x^2 + 3 x − 5
# $$

# Let's plot its values in the complex plane:

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


def f(z):
    """Sample function for root finding."""
    return z**3 - 3*z**2 + 3*z - 5


# In[3]:


import domain_coloring as dc


# In[4]:


fig, ax = plt.subplots(dpi=150)
dc.plot_domain(dc.modulus_domaincol, f, re=[-2, 4], im=[-3, 3], daxis=True)


# As one can see above, there are three roots. Let's try to apply the Durand-Kerner scheme described on Wikipedia. Basically, at each step we move a little from the previous estimate. If this sort of equation looks familiar to you, it's because it is Newton's root finding method for polynoms (where the derivate can be computed exactly)! 

# In[5]:


def durand_kerner(p, q, r, func):
    """Performs a step of the Durand-Kerner algorithm for order 3 polynomial."""
    pp = p - func(p) / ((p - q) * (p - r))
    qq = q - func(q) / ((q - p) * (q - r))
    rr = r - func(r) / ((r - p) * (r - q))
    return pp, qq, rr


# Let's see if we find the same values than in Wikipedia.

# In[6]:


p, q, r = 1 + 0*1j, 0.4 + 0.9j, -0.65 + 0.72j 

vals = [(p, q, r)]  
for _ in range(6):
    p, q, r = durand_kerner(p, q, r, f)
    vals.append((p, q, r))


# In[7]:


import pandas as pd


# In[8]:


pd.DataFrame(vals, columns=['p', 'q', 'r'])


# Indeed, these are the values Wikipedia shows.

# Let's plot how the roots evolve as a function of iteration step.

# In[9]:


from moviepy.editor import VideoClip
from moviepy.video.io.bindings import mplfig_to_npimage

p, q, r = 1 + 0*1j, 0.4 + 0.9j, -0.65 + 0.72j 
iteration = 0
duration = 2

fig, ax = plt.subplots(figsize=(5, 5), dpi=150)
def make_frame(t):
    ax.clear()
    dc.plot_domain(dc.modulus_domaincol, f, re=[-2, 4], im=[-3, 3], daxis=True, ax=ax)
    global p, q, r, iteration
    ax.plot(np.c_[[np.real(z) for z in (p, q, r)]],
            np.c_[[np.imag(z) for z in (p, q, r)]], 'ok')
    ax.set_title("iteration {}".format(iteration))
    p, q, r = durand_kerner(p, q, r, f)
    iteration += 1
    return mplfig_to_npimage(fig)

animation = VideoClip(make_frame, duration=duration)
plt.close(fig)
animation.ipython_display(fps=4, loop=True, autoplay=True)


# Can we also plot the trajectory of the iterated values themselves?

# In[10]:


import matplotlib.pyplot as plt
from moviepy.editor import VideoClip
from moviepy.video.io.bindings import mplfig_to_npimage

p, q, r = 1 + 0*1j, 0.4 + 0.9j, -0.65 + 0.72j 
vals = [(p, q, r)]  
    
iteration = 0
duration = 2

fig, ax = plt.subplots(figsize=(5, 5), dpi=150)
def make_frame(t):
    ax.clear()
    dc.plot_domain(dc.modulus_domaincol, f, re=[-2, 4], im=[-3, 3], daxis=True, ax=ax)
    global p, q, r, iteration, vals
    ax.plot(np.c_[[np.real(z) for z in (p, q, r)]],
            np.c_[[np.imag(z) for z in (p, q, r)]], 'ok')
    for ind in range(3):
        ax.plot(np.c_[[np.real(z[ind]) for z in vals]],
                np.c_[[np.imag(z[ind]) for z in vals]], '-r')
    ax.set_title("iteration {}".format(iteration))
    p, q, r = durand_kerner(p, q, r, f)
    vals.append((p, q, r))
    iteration += 1
    return mplfig_to_npimage(fig)

animation = VideoClip(make_frame, duration=duration)
plt.close(fig)
animation.ipython_display(fps=5, loop=True, autoplay=True)


# Not very spectacular, but it works. An interesting point in the above animation is that it highlights the fact that the roots are really following the gradient of the function, as shown by the trajectory that is perpendicular to the level-surfaces.

# *This post was entirely written using the IPython notebook. Its content is BSD-licensed. You can see a static view or download this notebook with the help of nbviewer at [20171110_DurandKerner.ipynb](http://nbviewer.ipython.org/urls/raw.github.com/flothesof/posts/master/20171110_DurandKerner.ipynb).*