Add split-operator method in Python (#332)

* Implement split-operator method in Python.

* Add Python split_op to book.

Author: berquist

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 <Psi|H|Psi>."""
+    # 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()

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 %}
 
 <script>