cms.equation.equation_newton_raphson のソースコード

import csv
import numpy as np
from numpy import sin, cos, tan, pi, exp
import matplotlib.pyplot as plt
import time
import sys


# self-consistent parameters
x0 = 0.0
dump  = 0.0
eps    = 1.0e-10
nmaxiter = 200
iprintiterval = 1

# graph plot parameters
ngdata = 51
xgmin = -1.0
xgmax =  2.0
tsleep = 1.0


if __name__ == "__main__":
    argv = sys.argv
    n = len(argv)
    if n >= 2: 
        x0 = float(argv[1])
    if n >= 3: 
        dump = float(argv[2])
    if n >= 4: 
        tsleep = float(argv[3])


# first derivative of func(x)


def diff(x):
    return exp(x) - 3.0




def func(x):
    return exp(x) - 3.0 * x





def main():
    global plt, ngdata, xgmin, xgmax
    global x0, dump, eps, nmaxiter, iprintinterval

    print("")
    print("Solution of transcendental equation by Newton-Raphson method")
    print("")
    print("x0 =", x0)
    print("dumping factor =", dump)
    print("")

    xg = []
    yg = []
    yz = [0.0] * ngdata
    xgstep = (xgmax - xgmin) / (ngdata - 1)
    for i in range(ngdata):
        xg.append(xgmin + i * xgstep)
        yg.append(func(xg[i]))

    fig, ax = plt.subplots(1, 1)
    plt.title("Solve equation")
    plt.xlabel("x")
    plt.ylabel("f(x)")
    ax.set_xlim([xgmin, xgmax])
    ax.set_ylim([min(yg), max(yg)])
    data, = ax.plot(xg, yg, color = 'black', linewidth = 0.3)
    yzero, = ax.plot(xg, yz, color = 'red', linewidth = 0.3)
    solve, = ax.plot([], color = 'blue', marker = 'o', linestyle = '')
    plt.pause(0.1)
#    time.sleep(3.0)

    x = x0
    xt = []
    yt = []
    for i in range(nmaxiter):
        f  = func(x)
        f1 = diff(x)
        f1 *= (1.0 + dump)
        xnext = x - f / f1
        dx = xnext - x
        if i % iprintiterval == 0:
            print("Iter {:5d}: x: {:>16.12f} => {:>16.12f}, dx = {:>10.4g}".format(i, x, xnext, dx))
        if abs(dx) < eps:
            print("  Success: Convergence reached: dx = {} < eps = {}".format(dx, eps))
            break

        xt.append(x)
        yt.append(f)
        solve.set_data(xt, yt)
        plt.pause(1.0e-5)
        time.sleep(tsleep)

        x = xnext
    else:
        print("  Failed: Convergence did not reach: dx = {} > eps = {}".format(dx, eps))
        return 0

    print("Press enter to terminate:", end = '')
    ret = input()
    print("")



if __name__ == "__main__":
    main()