181 lines
7.6 KiB
Python
181 lines
7.6 KiB
Python
import numpy as np
|
|
import numpy.typing as npt
|
|
from numpy.fft import fft, ifft, fftfreq
|
|
from scipy import constants as const
|
|
from scipy import sparse
|
|
from scipy import linalg
|
|
|
|
|
|
def discrete_hamiltonian(
|
|
dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float
|
|
) -> npt.NDArray:
|
|
"""Derive the hamiltonian for an arbitrary potential using the
|
|
discretized laplace operator. The function is implemented for 1 to
|
|
3 dimensions.
|
|
|
|
:param dx: Step size used for the coordinates. If the potential is
|
|
one-dimensional, this can simply be the step size of the
|
|
coordinate used in :param:`potential`.
|
|
For larger dimensions, this parameter must be a tuple with the
|
|
same size as the dimension.
|
|
:param potential: A (multidimensional) potential.
|
|
:param mass: The mass of the particle, used for the kinetic term.
|
|
"""
|
|
ndim = potential.ndim
|
|
if np.shape(dx) != (ndim,):
|
|
if ndim != 1 or np.shape(dx) != ():
|
|
raise ValueError(f"Expected list of {ndim} values for 'dx'")
|
|
else:
|
|
dx = np.array([dx])
|
|
grid_points = potential.shape
|
|
laplace_d2_1 = sparse.diags(
|
|
[1, -2, 1], [-1, 0, 1], shape=(grid_points[0], grid_points[0])
|
|
) / (dx[0] ** 2)
|
|
|
|
if ndim == 1:
|
|
laplace = laplace_d2_1
|
|
elif ndim >= 2:
|
|
laplace_d2_2 = sparse.diags(
|
|
[1, -2, 1], [-1, 0, 1], shape=(grid_points[1], grid_points[1])
|
|
) / (dx[1] ** 2)
|
|
if ndim == 2:
|
|
# laplace = sparse.kron(laplace_d2_1, eye_2) + sparse.kron(
|
|
# eye_1, laplace_d2_2
|
|
# )
|
|
laplace = sparse.kronsum(laplace_d2_1, laplace_d2_2)
|
|
if ndim == 3:
|
|
laplace_d2_3 = sparse.diags(
|
|
[1, -2, 1], [-1, 0, 1], shape=(grid_points[2], grid_points[2])
|
|
) / (dx[2] ** 2)
|
|
laplace = sparse.kronsum(
|
|
sparse.kronsum(laplace_d2_1, laplace_d2_2), laplace_d2_3
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"Number of dimensions not implemented: '{ndim}'")
|
|
t = -(const.hbar**2) / (2 * mass) * laplace
|
|
v = sparse.diags([potential.flatten()], [0])
|
|
return t + v
|
|
|
|
def functional_hamiltonian(
|
|
dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float
|
|
) -> npt.NDArray:
|
|
"""Returns a scipy.sparse.linalg LinearOperator for the hamiltonian of
|
|
an arbitrary potential using the discretized laplace operator.
|
|
The function is implemented for 1 to 3 dimensions.
|
|
|
|
:param dx: Step size used for the coordinates. If the potential is
|
|
one-dimensional, this can simply be the step size of the
|
|
coordinate used in :param:`potential`.
|
|
For larger dimensions, this parameter must be a tuple with the
|
|
same size as the dimension.
|
|
:param potential: A (multidimensional) potential.
|
|
:param mass: The mass of the particle, used for the kinetic term.
|
|
"""
|
|
ndim = potential.ndim
|
|
if np.shape(dx) != (ndim,):
|
|
if ndim != 1 or np.shape(dx) != ():
|
|
raise ValueError(f"Expected list of {ndim} values for 'dx'")
|
|
else:
|
|
dx = np.array([dx])
|
|
grid_points = potential.shape
|
|
|
|
if ndim != 3:
|
|
raise Exception(f"Number of dimensions ({ndim}) is not implemented.")
|
|
|
|
kx = 2*np.pi*fftfreq(grid_points[0], dx[0])
|
|
ky = 2*np.pi*fftfreq(grid_points[1], dx[1])
|
|
kz = 2*np.pi*fftfreq(grid_points[2], dx[2])
|
|
KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing='ij')
|
|
|
|
def mv(psi):
|
|
"""Function for action of hamiltonian on wavefunction"""
|
|
psi_nonflat = psi.reshape(grid_points)
|
|
#calculate the laplacian with a fourier transform
|
|
laplace_psi = (ifft(-KX**2*fft(psi_nonflat, axis = 0), axis = 0) +
|
|
ifft(-KY**2*fft(psi_nonflat, axis = 1), axis = 1) +
|
|
ifft(-KZ**2*fft(psi_nonflat, axis = 2), axis = 2))
|
|
v = sparse.diags([potential.flatten()], [0])
|
|
return -(const.hbar**2) / (2 * mass) *laplace_psi.flatten() + v*psi
|
|
|
|
len_psi = np.prod(grid_points)
|
|
return sparse.linalg.LinearOperator((len_psi,len_psi), matvec=mv)
|
|
|
|
def nstationary_solution(
|
|
dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float, k: int = 1,
|
|
method="sparse") -> tuple[npt.NDArray, npt.NDArray]:
|
|
"""Derive the stationary solution for a discrete potential of 1 to
|
|
3 dimensions numerically.
|
|
|
|
:param dx: Step size used for the coordinates. If the potential is
|
|
one-dimensional, this can simply be the step size of the
|
|
coordinate used in :param:`potential`.
|
|
For larger dimensions, this parameter must be a tuple with the
|
|
same size as the dimension.
|
|
:param potential: A (multidimensional) potential.
|
|
:param mass: The mass of the particle, used for the kinetic term.
|
|
:param k: Number of eigenvalues and eigenstates to return.
|
|
:returns: Tuple of eigenvalues and eigenvectors. Note that the
|
|
eigenvectors are transposed so that they are of shape
|
|
``(k, np.size(potential))``
|
|
"""
|
|
e_min = np.min(potential)
|
|
if method == "matrix_free":
|
|
eigenvals, eigenvects = sparse.linalg.eigsh(
|
|
functional_hamiltonian(dx, potential, mass), k=k, which="LM", sigma=e_min)
|
|
elif method == "sparse":
|
|
eigenvals, eigenvects = sparse.linalg.eigsh(
|
|
discrete_hamiltonian(dx, potential, mass), k=k, which="LM", sigma=e_min)
|
|
else:
|
|
raise ValueError(f'Method "{method}" does not exist.')
|
|
|
|
|
|
return eigenvals, eigenvects.T
|
|
|
|
def nstationary_solution_export(path,
|
|
dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float, k: int = 1
|
|
) -> tuple[npt.NDArray, npt.NDArray]:
|
|
"""Derive the stationary solution for a discrete potential of 1 to
|
|
3 dimensions numerically.
|
|
|
|
:param dx: Step size used for the coordinates. If the potential is
|
|
one-dimensional, this can simply be the step size of the
|
|
coordinate used in :param:`potential`.
|
|
For larger dimensions, this parameter must be a tuple with the
|
|
same size as the dimension.
|
|
:param potential: A (multidimensional) potential.
|
|
:param mass: The mass of the particle, used for the kinetic term.
|
|
:param k: Number of eigenvalues and eigenstates to return.
|
|
:returns: Tuple of eigenvalues and eigenvectors. Note that the
|
|
eigenvectors are transposed so that they are of shape
|
|
``(k, np.size(potential))``
|
|
"""
|
|
e_min = np.min(potential)
|
|
sparse.save_npz(path + r"\hamiltonian.npz",discrete_hamiltonian(dx, potential, mass))
|
|
return e_min
|
|
|
|
def nstationary_solution_mod(
|
|
dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float, k: int = 1
|
|
) -> tuple[npt.NDArray, npt.NDArray]:
|
|
"""
|
|
Modified function to only compute solutions in the local minimum.
|
|
|
|
Derive the stationary solution for a discrete potential of 1 to
|
|
3 dimensions numerically.
|
|
|
|
:param dx: Step size used for the coordinates. If the potential is
|
|
one-dimensional, this can simply be the step size of the
|
|
coordinate used in :param:`potential`.
|
|
For larger dimensions, this parameter must be a tuple with the
|
|
same size as the dimension.
|
|
:param potential: A (multidimensional) potential.
|
|
:param mass: The mass of the particle, used for the kinetic term.
|
|
:param k: Number of eigenvalues and eigenstates to return.
|
|
:returns: Tuple of eigenvalues and eigenvectors. Note that the
|
|
eigenvectors are transposed so that they are of shape
|
|
``(k, np.size(potential))``
|
|
"""
|
|
e_min = np.min(potential)
|
|
eigenvals, eigenvects = linalg.eigh(
|
|
discrete_hamiltonian(dx, potential, mass),subset_by_index=(0,k-1),subset_by_value=(e_min,np.inf)
|
|
)
|
|
return eigenvals, eigenvects.T |