From 07ff89d0d8e20c5578dff5596b2283a0cb272c21 Mon Sep 17 00:00:00 2001 From: Eric Berquist Date: Sat, 4 Aug 2018 21:08:37 -0400 Subject: [PATCH 1/2] Implement split-operator method in Python. --- .../code/python/split_op.py | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 contents/split-operator_method/code/python/split_op.py diff --git a/contents/split-operator_method/code/python/split_op.py b/contents/split-operator_method/code/python/split_op.py new file mode 100644 index 000000000..fac7c5c9d --- /dev/null +++ b/contents/split-operator_method/code/python/split_op.py @@ -0,0 +1,128 @@ +# pylint: disable=invalid-name +# pylint: disable=too-many-instance-attributes +# pylint: disable=too-few-public-methods +# pylint: disable=too-many-arguments +from math import pi +from math import sqrt + +import numpy as np + + +class Param: + """Container for holding all simulation parameters.""" + + def __init__(self, + xmax: float, + res: int, + dt: float, + timesteps: int, + im_time: bool) -> None: + + self.xmax = xmax + self.res = res + self.dt = dt + self.timesteps = timesteps + self.im_time = im_time + + self.dx = 2 * xmax / res + self.x = np.arange(-xmax + xmax / res, xmax, self.dx) + self.dk = pi / xmax + self.k = np.concatenate((np.arange(0, res / 2), + np.arange(-res / 2, 0))) * self.dk + + +class Operators: + """Container for holding operators and wavefunction coefficients.""" + + def __init__(self, res: int) -> None: + + self.V = np.empty(res, dtype=complex) + self.R = np.empty(res, dtype=complex) + self.K = np.empty(res, dtype=complex) + self.wfc = np.empty(res, dtype=complex) + + +def init(par: Param, voffset: float, wfcoffset: float) -> Operators: + """Initialize the wavefunction coefficients and the potential.""" + opr = Operators(len(par.x)) + opr.V = 0.5 * (par.x - voffset) ** 2 + opr.wfc = np.exp(-((par.x - wfcoffset) ** 2) / 2, dtype=complex) + if par.im_time: + opr.K = np.exp(-0.5 * (par.k ** 2) * par.dt) + opr.R = np.exp(-0.5 * opr.V * par.dt) + else: + opr.K = np.exp(-0.5 * (par.k ** 2) * par.dt * 1j) + opr.R = np.exp(-0.5 * opr.V * par.dt * 1j) + return opr + + +def split_op(par: Param, opr: Operators) -> None: + + for i in range(par.timesteps): + + # Half-step in real space + opr.wfc *= opr.R + + # FFT to momentum space + opr.wfc = np.fft.fft(opr.wfc) + + # Full step in momentum space + opr.wfc *= opr.K + + # iFFT back + opr.wfc = np.fft.ifft(opr.wfc) + + # Final half-step in real space + opr.wfc *= opr.R + + # Density for plotting and potential + density = np.abs(opr.wfc) ** 2 + + # Renormalizing for imaginary time + if par.im_time: + renorm_factor = sum(density) * par.dx + opr.wfc /= sqrt(renorm_factor) + + # Outputting data to file. Plotting can also be done in a + # similar way. This is set to output exactly 100 files, no + # matter how many timesteps were specified. + if i % (par.timesteps // 100) == 0: + filename = "output{}.dat".format(str(i).rjust(5, str(0))) + with open(filename, "w") as outfile: + for j in range(len(density)): + template = "{}\t{}\t{}\n".format + line = template(par.x[j], density[j].real, opr.V[j].real) + outfile.write(line) + print("Outputting step: ", i + 1) + + +def calculate_energy(par: Param, opr: Operators) -> float: + """Calculate the energy .""" + # Creating real, momentum, and conjugate wavefunctions. + wfc_r = opr.wfc + wfc_k = np.fft.fft(wfc_r) + wfc_c = np.conj(wfc_r) + + # Finding the momentum and real-space energy terms + energy_k = 0.5 * wfc_c * np.fft.ifft((par.k ** 2) * wfc_k) + energy_r = wfc_c * opr.V * wfc_r + + # Integrating over all space + energy_final = sum(energy_k + energy_r).real + + return energy_final * par.dx + + +def main() -> None: + par = Param(5.0, 256, 0.05, 100, True) + + # Starting wavefunction slightly offset so we can see it change + opr = init(par, 0.0, -1.00) + split_op(par, opr) + + energy = calculate_energy(par, opr) + print("Energy is: ", energy) + + +if __name__ == "__main__": + main() From 6044c2639498b5d83e99cfd31ca17f20f96435a3 Mon Sep 17 00:00:00 2001 From: Eric Berquist Date: Sat, 4 Aug 2018 22:18:00 -0400 Subject: [PATCH 2/2] Add Python split_op to book. --- contents/split-operator_method/code/python/split_op.py | 3 +-- contents/split-operator_method/split-operator_method.md | 8 ++++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/contents/split-operator_method/code/python/split_op.py b/contents/split-operator_method/code/python/split_op.py index fac7c5c9d..e4fe2c7d3 100644 --- a/contents/split-operator_method/code/python/split_op.py +++ b/contents/split-operator_method/code/python/split_op.py @@ -10,7 +10,6 @@ class Param: """Container for holding all simulation parameters.""" - def __init__(self, xmax: float, res: int, @@ -33,7 +32,6 @@ def __init__(self, class Operators: """Container for holding operators and wavefunction coefficients.""" - def __init__(self, res: int) -> None: self.V = np.empty(res, dtype=complex) @@ -89,6 +87,7 @@ def split_op(par: Param, opr: Operators) -> None: if i % (par.timesteps // 100) == 0: filename = "output{}.dat".format(str(i).rjust(5, str(0))) with open(filename, "w") as outfile: + # Outputting for gnuplot. Any plotter will do. for j in range(len(density)): template = "{}\t{}\t{}\n".format line = template(par.x[j], density[j].real, opr.V[j].real) diff --git a/contents/split-operator_method/split-operator_method.md b/contents/split-operator_method/split-operator_method.md index feaa33661..9f1a4e6e2 100644 --- a/contents/split-operator_method/split-operator_method.md +++ b/contents/split-operator_method/split-operator_method.md @@ -99,6 +99,8 @@ Regardless, we first need to set all the initial parameters, including the initi {% method %} {% sample lang="jl" %} [import:9-32, lang:"julia"](code/julia/split_op.jl) +{% sample lang="py" %} +[import:11-30, lang:"python"](code/python/split_op.py) {% endmethod %} As a note, when we generate our grid in momentum space `k`, we need to split the grid into two lines, one that is going from `0` to `-kmax` and is then discontinuous and goes from `kmax` to `0`. @@ -111,6 +113,8 @@ Afterwards, we turn them into operators: {% method %} {% sample lang="jl" %} [import:34-60, lang:"julia"](code/julia/split_op.jl) +{% sample lang="py" %} +[import:33-54, lang:"python"](code/python/split_op.py) {% endmethod %} Here, we use a standard harmonic potential for the atoms to sit in and a gaussian distribution for an initial guess for the probability distribution. @@ -124,6 +128,8 @@ The final step is to do the iteration, itself. {% method %} {% sample lang="jl" %} [import:63-109, lang:"julia"](code/julia/split_op.jl) +{% sample lang="py" %} +[import:57-95, lang:"python"](code/python/split_op.py) {% endmethod %} And that's it. @@ -143,6 +149,8 @@ Checking to make sure your code can output the correct energy for a harmonic tra {% method %} {% sample lang="jl" %} [import, lang:"julia"](code/julia/split_op.jl) +{% sample lang="py" %} +[import:5-127, lang:"python"](code/python/split_op.py) {% endmethod %}