Implementation

import numpy as np

class Eigenstates1D:
    def __init__(self, mass=1.0, grid=None, potential=None):
        self.mass = mass
        self.potential = potential
        self.gridx = None
        self.dim = None
        self.deltax = None
        self.hamiltonian = None
        if grid is not None:
            self.set_grid(grid)

    def set_grid(self, grid):
        self.gridx = grid
        self.dim = len(self.gridx)
        self.deltax = (self.gridx[1] - self.gridx[0])

    def set_potential(self, pot):
        self.potential = pot

    def compute_hamiltonian(self):
        self.hamiltonian = np.zeros((self.dim, self.dim), dtype=complex)
        idx_range = np.arange(self.dim)
        pk = (2.0 * np.pi / (self.dim * self.deltax)) * (idx_range - self.dim / 2.)
        tk = (pk ** 2) / (2.0 * self.mass)
        W = np.exp(2 * np.pi * 1.0j * idx_range / self.dim)
        one_minus_one = np.array([(-1)**k for k in range(self.dim)])

        for i in range(self.dim):
            tmp = tk * (W**i) * ((-1)**i)
            self.hamiltonian[i, :] = one_minus_one * np.fft.fft(tmp) / self.dim

        for i in range(self.dim):
            self.hamiltonian[i, i] += self.potential[i]
        
        self.hamiltonian = np.real(
            0.5 * (self.hamiltonian + self.hamiltonian.conj().T)
        )

    def get_eigenstates(self):
        return np.linalg.eigh(self.hamiltonian)