cDKA.py

import numpy as np 
import matplotlib.pyplot as plt 
import sympy as sb


class DKA():
    def __init__(self):
        pass

    def p(self,a,x):
        n = len(a)-1
        y = 0.
        for i in range(n+1):
            y += a[i] * x**(n-i)
        return y

    def S(self,a,x):
        n = len(a)-1
        
        y = x**n
        for i in range(2,n+1):
            y -= np.abs(a[i]) * x**(n-i)
        return y

    def Newton(self,f, x0, MAX=1000, ε=10**-6):
        df  = f.diff(x)
        dff = sb.lambdify(x, df)
        ff  = sb.lambdify(x, f)

        xo = x0
        for n in range(MAX):
            xn = xo - ff(xo)/dff(xo)
            if np.abs(xn-xo) < ε:
                print('FINISH')
                break
            xo = xn
            print(n)
        return xn

    def dfx(self,i,x):
        y = 1
        n = len(x)
        for j in range(n):
            if j is not i:
                y *= (x[i] - x[j])
        return y

    def get_coeff(self,f):
        n = sb.degree(f)
        a = [f.coeff(x**(n-i)) for i in range(n)]+[f.subs(x,0)]
        a = np.array(a,dtype=np.complex64)

        return a

    def calc(self,f):
        p = self.p 
        S = self.S 
        dfx = self.dfx
        Newton = self.Newton
        
        ff = sb.lambdify(x,f)
        # 係数
        a = self.get_coeff(f)
        # 次数
        n = len(a)-1
        # モニック
        b = a/a[0]
        # 重心
        zc = -b[1]/n
        

        g = p(b,x).subs(x,x+zc).expand()
        c = self.get_coeff(g)
        # 初期値
        x0 = 10
        # 半径
        r = Newton(S(c,x),x0)

        # 初期値
        k = np.arange(1,n+1)
        z0 = zc + r * np.exp( 1j*( (2*(k-1)*π)/n + π/(2*n) ))

        # DKA法
        MAX = 200
        ε = 10**-6
        z = np.zeros([n,MAX],dtype=complex)
        z[:,0] = z0

        K = MAX
        for k in range(MAX-1):
            for i in range(n):
                z[i,k+1] = z[i,k] - p(b,z[i,k]) / dfx(i,z[:,k])
            if np.max(np.abs(z[:,k+1] - z[:,k])) < ε and all(np.abs(ff(z[:,k+1])) < ε):
            # if np.max(np.abs(z[:,k+1] - z[:,k])) < ε:
                K = k+1
                print('FINISH')
                break
            print(k)

        # 零点
        α = z[:,k+1]

        print('|f(α)| < %.1e : ' %ε +str(all(np.abs(ff(α)) < ε)))

        self.f  = f
        self.ff =ff 
        self.a  = a
        self.b  = b
        self.c  = c
        self.n  = n
        self.α  = α
        self.z  = z
        self.K  = K
        self.zc = zc
        self.r  = r
        self.z0 = z0
        self.ε  = ε

plt.rcParams["font.size"] = 8
π = np.pi

x = sb.Symbol('x')

f =  (x**4*(x-1)**4*(x+1+1j)).expand() 

dka = DKA()
dka.calc(f)

lm = dka.r+0.5
X = np.linspace(-lm+dka.zc.real,lm+dka.zc.real,1000)
Y = np.linspace(-lm+dka.zc.imag,lm+dka.zc.imag,1000)
XX,YY = np.meshgrid(X,Y)

ZZ = XX + 1j * YY
WW = dka.ff(ZZ)

absWW = np.abs(WW)
θ = np.arange(0,2*np.pi,0.01)

plt.close('all')
################
# Result
################
plt.figure(figsize=(12,4))
plt.subplot(1,3,1)
plt.pcolormesh(XX,YY,np.log(absWW),cmap='inferno')
plt.title('$|f(z)|$')

plt.subplot(1,3,2)
plt.title('All zeros')
plt.pcolormesh(XX,YY,np.log(absWW),cmap='inferno')
plt.plot(dka.α.real,dka.α.imag,'C3o',ms=5,label='Zeros')

plt.plot(dka.r*np.cos(θ)+dka.zc.real,dka.r*np.sin(θ)+dka.zc.imag,'k--')
plt.plot(dka.z0.real,dka.z0.imag,'C0o',ms=5,label='Initial guess')

plt.legend(loc=1)

plt.subplot(1,3,3)
plt.title('Trajectory')
plt.pcolormesh(XX,YY,np.log(absWW),cmap='inferno')

for i in range(dka.n):
    plt.plot(dka.z[i,:dka.K].real,dka.z[i,:dka.K].imag,'o-',lw=1,ms=3)

plt.plot(dka.r*np.cos(θ)+dka.zc.real,dka.r*np.sin(θ)+dka.zc.imag,'k--')
plt.plot(dka.α.real,dka.α.imag,'C3o',ms=5)
plt.plot(dka.z0.real,dka.z0.imag,'C0o',ms=5)

axes = plt.gcf().get_axes()
for axis in axes: 
    plt.sca(axis)   
    plt.gca().set_aspect(1/plt.gca().get_data_ratio())
    plt.xlabel('$Re$')
    plt.ylabel('$Im$')

plt.tight_layout()
# plt.savefig('DKA1.png')

################
# Radius
################
plt.figure(figsize=(8,4))
WWW = dka.S(dka.c,ZZ)
absWWW= np.abs(WWW)
plt.subplot(1,2,1)
plt.pcolormesh(XX,YY,np.log(absWWW),cmap='inferno')
plt.subplot(1,2,2)
plt.pcolormesh(XX,YY,np.log(absWWW),cmap='inferno')
plt.plot(dka.r*np.cos(θ),dka.r*np.sin(θ),'k--')
plt.plot(dka.r.real,dka.r.imag,'C3o')

axes = plt.gcf().get_axes()
for axis in axes: 
    plt.sca(axis)   

    plt.gca().set_aspect(1/plt.gca().get_data_ratio())
    plt.xlabel('$Re$')
    plt.ylabel('$Im$')

plt.tight_layout()

################
# |f(α)|, |f'(α)|
################
plt.figure(figsize=(8,4))
plt.subplot(1,2,1)
for i in range(dka.n):
    plt.plot(i*np.ones(dka.K),np.abs(dka.ff(dka.z[i,:dka.K])),'o-',lw=1,ms=3)

plt.plot(np.abs(dka.ff(dka.z[:,0])),'C0o',label='Initial guess')
plt.plot(np.abs(dka.ff(dka.α)),'C3o',label='Finail result')

plt.axhline(dka.ε,color='k',linestyle='--',label=r'$\varepsilon$')
plt.title('$|f(α)|$')

plt.subplot(1,2,2)
df = sb.lambdify(x,dka.f.diff(x))

for i in range(dka.n):
    plt.plot(i*np.ones(dka.K),np.abs(df(dka.z[i,:dka.K])),'o-',lw=1,ms=3)

plt.plot(np.abs(df(dka.z[:,0])),'C0o',label='Initial guess')
plt.plot(np.abs(df(dka.α)),'C3o',label='Finail result')

plt.title(r'$|f\prime(α)|$')

axes = plt.gcf().get_axes()
for axis in axes: 
    plt.sca(axis)   
    plt.gca().set_yscale("log")
    plt.xlabel('zeros')
    plt.legend(loc=4)
    plt.xlim([-0.1,dka.n-0.9])
plt.tight_layout()
# plt.savefig('DKA2.png')