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..e4fe2c7d3 --- /dev/null +++ b/contents/split-operator_method/code/python/split_op.py @@ -0,0 +1,127 @@ +# 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: + # 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) + 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() 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 %}