-
Notifications
You must be signed in to change notification settings - Fork 1
/
solver.py
216 lines (190 loc) · 8.23 KB
/
solver.py
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
import numpy as _np
import numba as _numba
import potential as _potential
import scipy.sparse as _sparse
from scipy.sparse.linalg import spsolve as _spsolve
class Solver():
_psi = None
x = None
dx = None
dt = None
n_points = None
potential = None
stationary = False
delta_depth = 0.0
def __init__(self, x_max, dx, dt, potential=None, stationary=False):
"""
Creates a solver for the interval [-x_max, x_max] with a spatial step `dx` and a time step `dt`.
:param x_max:
:param dx:
:param dt:
:param potential:
:param stationary: if True, the potential is considered constant in time
"""
self.dx = dx
self.dt = dt
n1 = int(x_max / dx)
if n1 * dx < x_max - 0.01 * dx:
n1 += 1
self.n_points = 2 * n1 + 1
self.x = _np.linspace(- n1 * dx, n1 * dx, self.n_points)
self.stationary = stationary
if potential is None:
self.potential = _numba.vectorize(nopython=True)(lambda x: 0.0j)
self.stationary = True
elif isinstance(potential, _potential.Potential):
if not isinstance(potential, _potential.Potential1D):
raise ValueError("Only 1D potentials are allowed")
self.potential = potential.get_potential()
self.stationary = potential.is_stationary()
self.delta_depth = potential.get_delta_depth()
else:
self.potential = potential
self._psi = _np.zeros(self.x.shape, dtype=_np.complex)
def set_psi(self, psi):
"""
Sets the psi-function.
:param psi: an array, a function, or a tuple of functions.
If `psi` is a tuple of functions, they define the real and imaginary parts, respectively.
"""
if isinstance(psi, _np.ndarray):
if psi.shape != self.x.shape:
raise ValueError(f"Shape of array of psi {psi.shape} is not the same as shape of x {x.shape}")
self._psi = psi
elif isinstance(psi, tuple):
if len(psi) != 2:
raise ValueError(f"Tuple of length 2 expected, got {len(psi0)}")
self._psi = psi[0](self.x) + 1j * psi[1](self.x)
else:
self._psi = psi(self.x) + 1j * 0.0
def iterate(self, t):
"""
Iterates the system over a single time step from the starting time `t`.
:param t:
:return:
"""
raise NotImplementedError()
def execute(self, t_max, output_dt=None, psi0=None):
"""
Calculates the system up to `t_max` for the initial function `psi0`.
Returns a list of times and solutions with the interval `output_dt`
:param t_max:
:param output_dt:
:param psi0: if None, current psi will be used.
:return: A tuple of times and psi-functions.
"""
dt = self.dt
if psi0 is not None:
self.set_psi(psi0)
if output_dt is not None and output_dt < dt:
raise ValueError(f"output_dt = {output_dt} is smaller than dt = {dt}")
iterations = int(t_max / dt)
output_i = 0
ts = []
psis = []
for i in range(iterations):
if (output_dt is not None) and (i * dt >= output_dt * output_i):
ts.append(i * dt)
psis.append(self._psi.copy())
output_i += 1
self.iterate(i * dt)
return (_np.array(ts), psis) if output_dt is not None else self._psi.copy()
class EulerSolver(Solver):
__dpsi = None
__potential_x = None
def __init__(self, x_max, dx, dt, potential=None, stationary=False):
Solver.__init__(self, x_max, dx, dt, potential, stationary)
if self.stationary:
self.__potential_x = self.potential(self.x)
self.__potential_x[self.n_points // 2] -= self.delta_depth / self.dx
self.__dpsi = _np.zeros(self.x.shape, dtype=_np.complex)
@staticmethod
@_numba.jit(nopython=True)
def __iterate(psi, dpsi, dt, dx, n_points, potential, delta_depth):
dpsi[0] = 0.5 * 1j * (psi[1] + psi[-1] - 2 * psi[0]) / dx ** 2
dpsi[-1] = 0.5 * 1j * (psi[0] + psi[-2] - 2 * psi[-1]) / dx ** 2
for i in range(1, n_points - 1):
dpsi[i] = 0.5 * 1j * (psi[i + 1] + psi[i - 1] - 2 * psi[i]) / dx ** 2 - 1j * potential[i] * psi[i]
for i in range(n_points):
psi[i] += dpsi[i] * dt
def iterate(self, t):
potential = self.__potential_x if self.stationary else self.potential(t, self.x)
EulerSolver.__iterate(self._psi, self.__dpsi, self.dt, self.dx, self.n_points, potential, self.delta_depth)
class CrankNicolsonSolver(Solver):
__potential_x = None
__matrix = None
__diagonal = 0.0j
def __init__(self, x_max, dx, dt, potential=None, stationary=False, vector_potential=None):
"""
Creates a solver for the interval [-x_max, x_max] with a spatial step `dx` and a time step `dt`.
:param x_max:
:param dx:
:param dt:
:param potential:
:param stationary: if True, the potential is considered constant in time
:param vector_potential: if a function A(t), then A(t) p term will be added to the Hamiltonian.
"""
Solver.__init__(self, x_max, dx, dt, potential, stationary)
dim = len(self.x)
A = _sparse.dok_matrix((dim, dim), dtype=_np.complex)
dt = self.dt
dx = self.dx
self.__diagonal = 1.0 / dt + 0.5j / dx ** 2
for i in range(dim):
A[i, i] = self.__diagonal
A[i, i - 1] = -0.25j / dx ** 2
for i in range(dim - 1):
A[i, i + 1] = -0.25j / dx ** 2
A[dim - 1, 0] = -0.25j / dx ** 2
self.__matrix = _sparse.csc_matrix(A)
self.__vector_potential = vector_potential
if self.stationary:
self.__potential_x = self.potential(self.x)
self.__potential_x[self.n_points // 2] -= self.delta_depth / self.dx
self.__matrix[_np.arange(dim), _np.arange(dim)] = self.__diagonal + 0.5j * self.__potential_x
def iterate(self, t):
dt = self.dt
dx = self.dx
psi = self._psi
if self.stationary:
v = self.__potential_x
else:
v = self.potential(t, self.x)
v[self.n_points // 2] -= self.delta_depth / self.dx
v2 = self.potential(t + self.dt, self.x)
v2[self.n_points // 2] -= self.delta_depth / self.dx
dim = len(self.x)
self.__matrix[_np.arange(dim), _np.arange(dim)] = self.__diagonal + 0.5j * v2
b = psi / dt + 0.25j * (_np.roll(psi, 1) + _np.roll(psi, -1) - 2 * psi) / dx ** 2 \
- 0.5j * v * psi
if self.__vector_potential is not None:
vp_current = self.__vector_potential(t)
vp_next = self.__vector_potential(t + self.dt)
dim = len(self.x)
self.__matrix[_np.arange(dim), _np.roll(_np.arange(dim), -1)] = -0.25j / dx ** 2 + 0.25 * vp_next / dx
self.__matrix[_np.arange(dim), _np.roll(_np.arange(dim), 1)] = -0.25j / dx ** 2 - 0.25 * vp_next / dx
b -= 0.25 * (_np.roll(psi, -1) - _np.roll(psi, 1)) / dx * vp_current
self._psi = _spsolve(self.__matrix, b)
class SplitOperatorHalfSpectralSolver(Solver):
__potential_exponent = None
__momentum_exponent = None
def __init__(self, x_max, dx, dt, potential=None, stationary=False):
from scipy.fftpack import fftfreq
Solver.__init__(self, x_max, dx, dt, potential, stationary)
if self.stationary:
v = self.potential(self.x)
v[self.n_points // 2] -= self.delta_depth / self.dx
self.__potential_exponent = _np.exp(-0.5j * v * dt)
p = 2.0 * _np.pi * fftfreq(len(self.x)) / self.dx
self.__momentum_exponent = _np.exp(-0.5j * p ** 2 * dt)
def iterate(self, t):
from scipy.fftpack import fft, ifft
if self.stationary:
pot_exp = self.__potential_exponent
else:
v = self.potential(t + 0.5 * self.dt, self.x)
v[self.n_points // 2] -= self.delta_depth / self.dx
pot_exp = _np.exp(-0.5j * v * self.dt)
psi = pot_exp * self._psi
psi = ifft(self.__momentum_exponent * fft(psi))
self._psi = pot_exp * psi