anharmonic.py

anharmonic.py

from __future__ import print_function

import numpy as np
from scipy import linalg

# scipy.linalg.eigh
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.eigh.html


def make_second_deriv(dim, dx):

    # make a diagonal matrix its shifted ones (periodic boundary condition is assumed)
    f0 = np.identity(dim, dtype=int)  # f_{i}
    f1 = np.roll(f0, 1, axis=1)  # f_{i+1}
    f2 = np.roll(f0, 2, axis=1)  # f_{i+2}
    f_1 = f1.transpose()  # f_{i-1}
    f_2 = f2.transpose()  # f_{i-2}

    # O(dx^2): (f_{i+1} - 2f_{i} + f_{i-1}) / dx^2
    deriv2 = (f1 - 2.0 * f0 + f_1) / (dx*dx)

    # O(dx^4): [(-1/12)f_{i+2) + (4/3)f_{i+1} - (5/2)f_{i} + (4/3)f_{i-1} - (1/12)f_{i-2}] / dx^2
    # deriv2 = ((-1./12.)*f2 + (4./3.)*f1 - (5./2.)*f0 + (4./3.)*f_1 - (1./12.)*f_2) / (dx*dx)

    return deriv2


def solve_anharmonic(x_max, nx, anharmonicity, file_val="eigenval.dat", file_vec="eigenvec.dat", n_vec=10):

    print("making Hamiltonian...")

    # arrays containing x, x^2, x^4
    x = np.linspace(-x_max, x_max, 2*nx)
    x2 = x**2  # element-wise square
    x4 = x2**2

    dx = x[1] - x[0]
    print("dx =", dx)

    # potential term
    v_diag = 0.5 * x2 + anharmonicity * x4
    v_matrix = np.diag(v_diag)

    # kinetic term, -(1/2) d^2/dx^2
    kin_matrix = -0.5 * make_second_deriv(2*nx, dx)

    # Hamiltonian
    h_matrix = v_matrix + kin_matrix

    # diagonalize
    print("diagonalizing Hamiltonian...")
    eigval, eigvec = linalg.eigh(h_matrix)

    # normalize eigenvectors so that the result does not depend on dx
    eigvec /= np.sqrt(dx)

    # save result for eigenvalues
    print("saving results...")
    np.savetxt(file_val, eigval)

    # save result for eigenvectors
    temp = np.hstack([x[:, None], eigvec[:, 0:n_vec]])
    np.savetxt(file_vec, temp)

    return eigval, eigvec


def main():
    nx = 100
    x_max = 10
    anharmonicity = 0

    solve_anharmonic(x_max, nx, anharmonicity)


if __name__ == '__main__':
    main()