1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
|
'''schroedinger.py
Collection of method common to solve and tests
'''
import sys
from pathlib import Path
from typing import TextIO, Type
import numpy as np
from numpy.polynomial import Polynomial
from numpy.typing import NDArray
from scipy.interpolate import CubicSpline
from scipy.linalg import eigh_tridiagonal
Interpolator = Type[callable] | Type[Polynomial] | Type[CubicSpline]
# type Interpolator = callable | Polynomial | CubicSpline
class Config:
'''Wrapper for reading, parsing and storing the contents of a configuration
file
'''
def __init__(self, path):
current_line = 0
# Ensure Path object
if path is not Path:
path = Path(path)
def next_parameter(file: TextIO):
'''Read next parameter, ignoring comments or empty lines'''
content = None
nonlocal current_line
while not content:
line = file.readline()
current_line += 1
index = line.find('#')
content = line[0:index].strip()
return content
with open(path, 'r', encoding='utf8') as file:
try:
self.mass = float(next_parameter(file))
start, end, steps = next_parameter(file).split()
self.interval = [float(start), float(end), int(steps)]
self.eig_interval = [int(attr) - 1 for attr in next_parameter(file).split()]
self.interpolation = next_parameter(file)
npoints = int(next_parameter(file))
self.points = np.zeros((npoints, 2))
for i in range(npoints):
line = next_parameter(file)
self.points[i] = np.array([float(comp) for comp in line.split()])
except ValueError:
print(f'Syntax error in \'{path.name}\' line {current_line}')
sys.exit(1)
def potential_interp(
interpolation: str,
points: NDArray[np.float64]
) -> Interpolator:
'''Create an interpolator for a set of predefined potential points
:param interpolation: Kind of interpolation
:param points: Points to interpolate within
:return: Interpolating object
'''
interpolator = None
if interpolation == 'linear':
def line(pos):
return np.interp(pos, points[:, 0], points[:, 1])
interpolator = line
elif interpolation == 'polynomial':
interpolator = Polynomial.fit(points[:, 0], points[:, 1],
points.shape[0] - 1)
elif interpolation == 'cspline':
interpolator = CubicSpline(points[:, 0], points[:, 1])
else:
raise ValueError('Invalid interpolator kind. Use any of: linear, polynomial or cspline')
return interpolator
def build_potential(
config: Config
) -> tuple[NDArray[np.float64], np.float64]:
'''Build a potential based on the options inside a configuration file
:param config: System parameters
:return: Potential and distance between samples
'''
start, end, steps = config.interval
potential = np.zeros((steps, 2))
potential[:, 0] = np.linspace(start, end, steps)
delta = np.abs(potential[1, 0] - potential[0, 0])
interp = potential_interp(config.interpolation, config.points)
potential[:, 1] = interp(potential[:, 0])
return potential, delta
def solve_schroedinger(
mass: float,
potential: NDArray[np.float64],
delta: float,
eig_interval: tuple[int, int]=None
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
'''Solve the one dimensional, time independent Schrödinger's equation for a
particle of given mass inside a discretized potential
:param mass: Mass of the particle
:param potential: Discretized potential
:param delta: Distance between points in the potential
:param eig_interval: Interval of quantum numbers for which the states are to be calculated
:return: Eigenvalues and normalized wave functions
'''
qnumber = potential.shape[0]
coeff = 1 / mass / delta**2
eig, vec = eigh_tridiagonal(coeff + potential,
-coeff * np.ones(qnumber - 1, dtype=np.float64) / 2.0,
select='i', select_range=eig_interval)
# Normalize eigenfunctions
for i in range(eig.shape[0]):
vec[:, i] /= np.sqrt(delta * np.sum(np.abs(vec[:, i])**2))
return eig, vec
|
