The Split-Operator Method

The Split-Operator Method (also called the Split-Step Method), was actually the primary method I used to solve the Schrödinger equation during my PhD. It is one of the simplest and fastest methods for this purpose and is widely used throughout modern quantum research in the area, in particular when dealing with the Non-linear Schrödinger Equation (NLSE):

which follows from the notation provided in the quantum systems chapter: is a quantum wave-function with spatial () and time () dependence, is a Laplacian, and is a potential of some sort (like or something). In this case, we also add an interaction term next to a nonlinear term. This is the system I studied for most of my PhD (granted, we played a few tricks with parallelization and such, so it was slightly more complicated).

At its heart, the split-op method is nothing more than a pseudo-spectral differential equation solver... That is to say, it solves the Schrödinger equation with FFTs. In fact, there is a large class of spectral and pseudo-spectral methods used to solve a number of different physical systems, and we'll definitely be covering those in the future. As mentioned in the quantum systems section, we can represent a quantum wavefunction in momentum space, which is parameterized with the wavevector . In the Hamiltonian shown above, we can split our system into position space components, , and momentum space components, . I'll be honest, I didn't know what notation to use for because is used to describe momentum. I settled on for real space, but that is somewhat notationally ambiguous. In addition, will indicate momentum space because it is a sum of all wavevectors, typically notated as . Bad notation aside, let's continue.

If we assume a somewhat general solution to our quantum system:

and assume we are simulating our system by a series of small timesteps (), we can perform similar splitting by using the Baker-Campbell-Housdorff formula:

This accrues a small amount of error () related to the commutation of the real and momentum-space components of the Hamiltonian. This is a relatively large error and that's not okay. In order to change the error to , we can split the system by performing a half-step in position space before doing a full-step in momentum space, through a process called Strang Splitting like so:

We can then address each part of this solution in chunks, first in position space, then in momentum space, then in position space again by using Fourier Transforms. Which looks something like this:

where , , and and indicate forward and inverse Fourier Transforms. Here's a flowchart of what we are looking for every timestep:

For the most part, that's it:

  1. Multiply the wavefunction in real space with the real-space operator.
  2. Flip to momentum space with a Fourier transform.
  3. Multiply the momentum-space wavefunction by the momentum-space operator.
  4. Flip to position space with an inverse Fourier transform.
  5. Repeat 1-4 until satisfied.

If we guess that our initial wavefunction is Gaussian-like and is slightly offset from the center or the trap, this should allow us to see our wavefunction "sloshing" back and forth in our trap, like so:

As a small concession, using this method enforces periodic boundary conditions, where the wavefunction will simply slide from one side of your simulation box to the other, but that's fine for most cases. In fact, for many cases (such as large-scale turbulence models) it's ideal.

That said, there is more to the story. As we mentioned in the quantum systems section, many simulations of quantum systems desire to find the ground state of our system. The split-operator method can be used for that too! If we run this simulation in imaginary time, by simply setting and stepping through instead of , we will no longer see an "real-world" example of how the atoms should behave, but will instead see an exponential decay of higher-energy states. If we run the simulation for long enough with a small enough timestep, all higher energy states will vanish. This means that we can find the ground state of our system by running the simulation in imaginary time, which is an incredibly useful feature! If we run the same simulation as above in imaginary time, we should see our wavefunction smoothly move to the center of our trap (the lowest energy position), like so:

The Algorithm

Luckily, the code in this case is pretty straightforward. As a note before starting, we will be using normalized units in this simulation where . These units are often called natural units. Many of you (cough experimentalists cough) will probably think that these units are completely unphysical, and they are; however, they allow us to output fractions and whole numbers. For example, if we are trying to find the energy of the ground state of atoms in a simple harmonic oscillator, we know it should be , where is the coefficient in front of the term known as the frequency of the trap. If we were to calculate the energy in real units, our simulation would output , which is hard to interpret. By instead using natural units, we get precisely and we know that those are in units of . There is no doubt that it makes the simulation easier to understand (albeit a little misleading in the end).

Regardless, we first need to set all the initial parameters, including the initial grids in real and momentum space:

struct Param
    xmax::Float64
    res::Int64
    dt::Float64
    timesteps::Int64
    dx::Float64
    x::Vector{Float64}
    dk::Float64
    k::Vector{Float64}
    im_time::Bool

    Param() = new(10.0, 512, 0.05, 1000, 2 * 10.0/512,
                  Vector{Float64}(-10.0 + 10.0/512 : 20.0/512 : 10.0),
                  pi / 10.0,
                  Vector{Float64}(vcat(0:512/2 - 1, -512/2 : -1) * pi/10.0),
                  false)
    Param(xmax::Float64, res::Int64, dt::Float64, timesteps::Int64,
          im_val::Bool) = new(
              xmax, res, dt, timesteps,
              2*xmax/res, Vector{Float64}(-xmax+xmax/res:2*xmax/res:xmax),
              pi/xmax, Vector{Float64}(vcat(0:res/2-1, -res/2:-1)*pi/(xmax)),
              im_val
          )
end

struct params {
    double xmax;
    unsigned int res;
    double dt;
    unsigned int timesteps;
    double dx;
    double *x;
    double dk;
    double *k;
    bool im_time;
};

void init_params(struct params *par, double xmax, unsigned int res, double dt,
                 unsigned int timesteps, bool im) {

    par->xmax = xmax;
    par->res = res;
    par->dt = dt;
    par->timesteps = timesteps;
    par->dx = 2.0 * xmax / res;
    par->x = malloc(sizeof(double) * res);
    par->dk = M_PI / xmax;
    par->k = malloc(sizeof(double) * res);
    par->im_time = im;

    for (size_t i = 0; i < res; ++i) {
        par->x[i] = xmax / res - xmax + (double)i * (2.0 * xmax / res);
        if (i < res / 2) {
            par->k[i] = (double)i * M_PI / xmax;
        } else {
            par->k[i] = ((double)i - res) * M_PI / xmax;
        }
    }
}

using complex = std::complex<double>;
using vector_real = std::vector<double>;
using vector_complex = std::vector<complex>;

struct Params {
    Params(double _xmax, unsigned int _res, double _dt, unsigned int _timesteps, bool im) {
        xmax = _xmax;
        res = _res;
        dt = _dt;
        timesteps = _timesteps;
        dx = 2.0 * xmax / res;
        x.reserve(res);
        dk = M_PI / xmax;
        k.reserve(res);
        im_time = im;

        for (size_t i = 0; i < res; ++i) {
            x.emplace_back(xmax / res - xmax + static_cast<double>(i) * (2.0 * xmax / res));
            if (i < res / 2) {
                k.push_back(static_cast<double>(i) * M_PI / xmax);
            } else {
                k.push_back((static_cast<double>(i) - res) * M_PI / xmax);
            }
        }
    }

    double xmax;
    unsigned int res;
    double dt;
    unsigned int timesteps;
    double dx;
    vector_real x;
    double dk;
    vector_real k;
    bool im_time;
};

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

data Parameters = Parameters
  { xmax :: Double
  , res :: Int
  , dt :: Double
  , timesteps :: Int
  , dx :: Double
  , x :: Vector
  , dk :: Double
  , ks :: Vector
  , imTime :: Bool
  }

defaultParameters :: Parameters
defaultParameters = makeParameters 10 512 0.01 1000 True

makeParameters :: Double -> Int -> Double -> Int -> Bool -> Parameters
makeParameters xmax res dt timesteps imTime =
  let fi = fromIntegral
      rng = (0, res - 1)
      ks = [0 .. div res 2 - 1] ++ [-div res 2 .. -1]
   in Parameters
        xmax
        res
        dt
        timesteps
        (2 * xmax / fi res)
        (listArray rng $
         map (\n -> xmax * (-1 + 2 * fi n / fi res) :+ 0) [1 .. res])
        (pi / xmax)
        (listArray rng $ map ((:+ 0) . (pi / xmax *) . fi) ks)
        imTime

    xmax: f64,
    res: usize,
    dt: f64,
    timesteps: usize,
    dx: f64,
    x: Vec<f64>,
    dk: f64,
    k: Vec<f64>,
    im_time: bool,
}

impl Parameters {
    pub fn new(xmax: f64, res: usize, dt: f64, timesteps: usize, im_time: bool) -> Parameters {
        let dx = 2.0_f64 * xmax / (res as f64);
        let mut x: Vec<f64> = Vec::with_capacity(res);
        let dk = PI / xmax;
        let mut k: Vec<f64> = Vec::with_capacity(res);
        for i in 0..res {
            x.push(xmax / (res as f64) - xmax + (i as f64) * dx);
            match i {
                i if (i < res / 2) => k.push((i as f64) * PI / xmax),
                _ => k.push(((i as f64) - (res as f64)) * PI / xmax),
            }
        }
        Parameters {
            xmax,
            res,
            dt,
            timesteps,
            im_time,
            dx,
            x,
            dk,
            k,
        }
    }
}

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. This is simply because the FFT will naturally assume that the 0 in our grid is at the left side of the simulation, so we shift k-space to match this expectation. Also, for this code we will be using notation to what we used above: opr.R will be the real space operators and opr.K will be the momentum space operators. There is another Boolean value here called im_time, which is for imaginary time evolution.

Afterwards, we turn them into operators:

mutable struct Operators
    V::Vector{Complex{Float64}}
    R::Vector{Complex{Float64}}
    K::Vector{Complex{Float64}}
    wfc::Vector{Complex{Float64}}

    Operators(res) = new(zeros(res),
                         zeros(res),
                         zeros(res),
                         zeros(res))
end

# Function to initialize the wfc and potential
function init(par::Param, voffset::Float64, wfcoffset::Float64)
    opr = Operators(length(par.x))
    opr.V = 0.5 * (par.x .- voffset).^2
    opr.wfc = exp.(-(par.x .- wfcoffset).^2/2)
    if (par.im_time)
        opr.K = exp.(-0.5*par.k.^2*par.dt)
        opr.R = exp.(-0.5*opr.V*par.dt)
    else
        opr.K = exp.(-im*0.5*par.k.^2*par.dt)
        opr.R = exp.(-im*0.5*opr.V*par.dt)
    end

    return opr
end

struct operators {
    size_t size;
    double complex *v;
    double complex *pe;
    double complex *ke;
    double complex *wfc;
};

void init_operators(struct operators *opr, struct params par, double voffset,
                    double wfcoffset) {

    opr->size = par.res;
    opr->v = malloc(sizeof(double complex) * par.res);
    opr->pe = malloc(sizeof(double complex) * par.res);
    opr->ke = malloc(sizeof(double complex) * par.res);
    opr->wfc = malloc(sizeof(double complex) * par.res);

    for (size_t i = 0; i < par.res; ++i) {
        opr->v[i] = 0.5 * cpow(par.x[i] - voffset, 2);
        opr->wfc[i] = cexp(-cpow(par.x[i] - wfcoffset, 2) / 2.0);

        if (par.im_time) {
            opr->ke[i] = cexp(-0.5 * par.dt * cpow(par.k[i], 2));
            opr->pe[i] = cexp(-0.5 * par.dt * opr->v[i]);
        } else {
            opr->ke[i] = cexp(-0.5 * par.dt * cpow(par.k[i], 2) * I);
            opr->pe[i] = cexp(-0.5 * par.dt * opr->v[i] * I);
        }
    }
}

struct Operators {
public:
    Operators(Params &par, double voffset,
              double wfcoffset) {
        size = par.res;
        v.reserve(size);
        pe.reserve(size);
        ke.reserve(size);
        wfc.reserve(size);

        for (size_t i = 0; i < size; ++i) {
            v.push_back(0.5 * pow(par.x[i] - voffset, 2));
            wfc.push_back(exp(-pow(par.x[i] - wfcoffset, 2) / 2.0));

            if (par.im_time) {
                ke.push_back(exp(-0.5 * par.dt * pow(par.k[i], 2)));
                pe.push_back(exp(-0.5 * par.dt * v[i]));
            } else {
                ke.push_back(exp(-0.5 * par.dt * pow(par.k[i], 2) * complex(0.0, 1.0)));
                pe.push_back(exp(-0.5 * par.dt * v[i] * complex(0.0, 1.0)));
            }
        }
    }

    size_t size;
    vector_complex v;
    vector_complex pe;
    vector_complex ke;
    vector_complex wfc;
};

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

data Operators = Operators
  { v :: Vector
  , rStep :: Vector
  , kStep :: Vector
  , wfc :: Vector
  }

makeOperators :: Parameters -> Complex Double -> Complex Double -> Operators
makeOperators param v0 wfc0 =
  let rng = (0, res param - 1)
      time
        | imTime param = dt param :+ 0
        | otherwise = 0 :+ dt param
      v = liftArray (\x -> 0.5 * (x - v0) ^ 2) (x param)
      rStep = liftArray (\x -> exp (-0.5 * time * x)) v
      kStep = liftArray (\k -> exp (-0.5 * time * k ^ 2)) (ks param)
      wfc = liftArray (\x -> exp (-(x - wfc0) ^ 2 / 2)) (x param)
   in Operators v rStep kStep (normalize (dx param) wfc)

    v: Vec<Complex<f64>>,
    pe: Vec<Complex<f64>>,
    ke: Vec<Complex<f64>>,
    wfc: Vec<Complex<f64>>,
}

impl Operators {
    pub fn new(par: &Parameters, v_offset: f64, wfc_offset: f64) -> Operators {
        let mut v: Vec<Complex<f64>> = Vec::with_capacity(par.res);
        let mut pe: Vec<Complex<f64>> = Vec::with_capacity(par.res);
        let mut ke: Vec<Complex<f64>> = Vec::with_capacity(par.res);
        let mut wfc: Vec<Complex<f64>> = Vec::with_capacity(par.res);

        for i in 0..par.res {
            v.push(Complex::new(
                0.5_f64 * (par.x[i] - v_offset).powi(2),
                0.0_f64,
            ));
            wfc.push(Complex::new(
                (-((par.x[i] - wfc_offset).powi(2)) / 2.0_f64).exp(),
                0.0_f64,
            ));
            if par.im_time {
                ke.push(Complex::new(
                    (-0.5_f64 * par.dt * par.k[i].powi(2)).exp(),
                    0.0_f64,
                ));
                pe.push(Complex::new((-0.5_f64 * par.dt * v[i].re).exp(), 0.0_f64));
            } else {
                ke.push(Complex::new(
                    0.0_f64,
                    (-0.5_f64 * par.dt * par.k[i].powi(2)).exp(),
                ));
                pe.push(Complex::new(0.0_f64, (-0.5_f64 * par.dt * v[i].re).exp()));
            }
        }
        Operators { v, pe, ke, wfc }
    }
}

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. If we give either the trap or the atoms a slight offset (so the Gaussian distribution of atoms does not quite rest at the bottom of the potential, we can see the atoms moving back and forth in the potential as we move the simulation forward in time. This means that we can easily see the dynamics of our quantum system! If we run the simulation in imaginary time, we will see the Gaussian distribution of atoms move towards the center of the potential, which is the location with the lowest energy. Both of these have been shown in the figures above.

The final step is to do the iteration, itself.

function split_op!(par::Param, opr::Operators)

    for i = 1:par.timesteps
        # Half-step in real space
        opr.wfc = opr.wfc .* opr.R

        # fft to momentum space
        opr.wfc = fft(opr.wfc)

        # Full step in momentum space
        opr.wfc = opr.wfc .* opr.K

        # ifft back
        opr.wfc = ifft(opr.wfc)

        # final half-step in real space
        opr.wfc = opr.wfc .* opr.R

        # density for plotting and potential
        density = abs2.(opr.wfc)

        # renormalizing for imaginary time
        if (par.im_time)
            renorm_factor = sum(density) * par.dx

            for j = 1:length(opr.wfc)
                opr.wfc[j] /= sqrt(renorm_factor)
            end
        end

        # 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
        if ((i-1) % div(par.timesteps, 100) == 0)
            outfile = open("output" * string(lpad(string(i-1), 5, string(0)))
                                    * ".dat","w")

            # Outputting for gnuplot. Any plotter will do.
            for j = 1:length(density)
                write(outfile, string(par.x[j]) * "\t"
                               * string(density[j]) * "\t"
                               * string(real(opr.V[j])) * "\n")
            end

            close(outfile)
            println("Outputting step: ", i)
        end
    end
end

void split_op(struct params par, struct operators opr) {
    double density[opr.size];

    for (size_t i = 0; i < par.timesteps; ++i) {
        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        fft(opr.wfc, opr.size, false);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.ke[j];
        }

        fft(opr.wfc, opr.size, true);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        for (size_t j = 0; j < opr.size; ++j) {
            density[j] = pow(cabs(opr.wfc[j]), 2);
        }

        if (par.im_time) {
            double sum = 0;

            for (size_t j = 0; j < opr.size; ++j) {
                sum += density[j];
            }

            sum *= par.dx;

            for (size_t j = 0; j < opr.size; ++j) {
                opr.wfc[j] /= sqrt(sum);
            }
        }

        // Writing data into a file in the format of:
        // index, density, real potential.
        char filename[256];
        sprintf(filename, "output%lu.dat", i);
        FILE *fp = fopen(filename, "w");

        for (size_t i = 0; i < opr.size; ++i) {
            fprintf(fp, "%ld\t%f\t%f\n", i, density[i], creal(opr.v[i]));
        }

        fclose(fp);
    }
}

void split_op(Params &par, Operators &opr) {
    auto density = std::vector<double>(opr.size, 0);

    for (size_t i = 0; i < par.timesteps; ++i) {
        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        fft(opr.wfc, false);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.ke[j];
        }

        fft(opr.wfc, true);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        for (size_t j = 0; j < opr.size; ++j) {
            density[j] = pow(abs(opr.wfc[j]), 2);
        }

        if (par.im_time) {
            double sum = 0;

            for (size_t j = 0; j < opr.size; ++j) {
                sum += density[j];
            }

            sum *= par.dx;

            for (size_t j = 0; j < opr.size; ++j) {
                opr.wfc[j] /= sqrt(sum);
            }
        }

        // Writing data into a file in the format of:
        // index, density, real potential.
        std::stringstream filename_stream;
        filename_stream << "output" << i << ".dat";

        std::ofstream fstream = std::ofstream(filename_stream.str());

        if (fstream) {
            for (std::size_t i = 0; i < opr.size; ++i) {
                std::stringstream data_stream;

                data_stream << i << "\t" << density[i] << "\t" << real(opr.v[i]) << "\n";

                fstream.write(data_stream.str().c_str(), data_stream.str().length());
            }
        }

        fstream.close();
    }
}

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)

evolve :: Parameters -> Operators -> [Operators]
evolve param op@(Operators _ rStep kStep _) = iterate splitop op
  where
    splitop op = op {wfc = wfc' op}
    wfc' = norm . (rStep .*) . idft . (kStep .*) . dft . (rStep .*) . wfc
    norm = if imTime param then normalize (dx param) else id

    let mut density: Vec<f64>;

    for i in 0..par.timesteps {
        for j in 0..par.res {
            opr.wfc[j] *= opr.pe[j];
        }

        fft(&mut opr.wfc, false);

        for j in 0..par.res {
            opr.wfc[j] *= opr.ke[j];
        }

        fft(&mut opr.wfc, true);

        for j in 0..par.res {
            opr.wfc[j] *= opr.pe[j];
        }

        density = opr.wfc.iter().map(|x| x.norm().powi(2)).collect();

        if par.im_time {
            let sum = density.iter().sum::<f64>() * par.dx;

            for j in 0..par.res {
                opr.wfc[j] /= sum.sqrt();
            }
        }

        // Writing data into a file in the format of:
        // index, density, real potential.
        let path_name = format!("output{}.dat", i);
        let path = Path::new(&path_name);
        let display = path.display();

        let mut file = match File::create(&path) {
            Err(why) => panic!("Couldn't create {}: {}", display, why),
            Ok(good) => good,
        };

        for j in 0..par.res {
            if let Err(why) = writeln!(file, "{}\t{}\t{}", j, density[j], opr.v[j].re) {
                panic!("Couldn't write to {}: {}", display, why)
            }
            if let Err(why) = file.flush() {
                panic!("Couldn't flush {}: {}", display, why)
            }
        }
    }
}

And that's it.

There is something a bit odd about the simulation in imaginary time, though. Basically, in imaginary time, we see an exponential decay of all the higher energy states, which means we are technically losing a large amount of our wavefunction density every timestep! To solve this issue, we renormalize by enforcing that . As you can see from the code, this involves summing the density, multiplying that sum by dx, and then dividing each element in the wavefunction by the sqrt() of that value.

The Split-Operator method is one of the most commonly used quantum simulation algorithms because of how straightforward it is to code and how quickly you can start really digging into the physics of the simulation results!

Video Explanation

Here is a video describing the split-operator method:

Example Code

This example code is a simulation of a Gaussian distribution of atoms slightly offset in a harmonic trap in imaginary time. So long as the code is written appropriately, this means that the atoms should move towards the center of the trap and the energy should decay to , which will be simply in this simulation. Checking to make sure your code can output the correct energy for a harmonic trap is a good test to make sure it is all working under-the-hood before simulating systems with more complicated Hamiltonians.

#------------split_op.jl-------------------------------------------------------#
#
# Plotting: to plot individual timesteps, use gnuplot like so:
#               p "output00000.dat" u 1:2 w l
#               rep "output00000.dat" u 1:3 w l
#
#------------------------------------------------------------------------------#

using FFTW

struct Param
    xmax::Float64
    res::Int64
    dt::Float64
    timesteps::Int64
    dx::Float64
    x::Vector{Float64}
    dk::Float64
    k::Vector{Float64}
    im_time::Bool

    Param() = new(10.0, 512, 0.05, 1000, 2 * 10.0/512,
                  Vector{Float64}(-10.0 + 10.0/512 : 20.0/512 : 10.0),
                  pi / 10.0,
                  Vector{Float64}(vcat(0:512/2 - 1, -512/2 : -1) * pi/10.0),
                  false)
    Param(xmax::Float64, res::Int64, dt::Float64, timesteps::Int64,
          im_val::Bool) = new(
              xmax, res, dt, timesteps,
              2*xmax/res, Vector{Float64}(-xmax+xmax/res:2*xmax/res:xmax),
              pi/xmax, Vector{Float64}(vcat(0:res/2-1, -res/2:-1)*pi/(xmax)),
              im_val
          )
end

mutable struct Operators
    V::Vector{Complex{Float64}}
    R::Vector{Complex{Float64}}
    K::Vector{Complex{Float64}}
    wfc::Vector{Complex{Float64}}

    Operators(res) = new(zeros(res),
                         zeros(res),
                         zeros(res),
                         zeros(res))
end

# Function to initialize the wfc and potential
function init(par::Param, voffset::Float64, wfcoffset::Float64)
    opr = Operators(length(par.x))
    opr.V = 0.5 * (par.x .- voffset).^2
    opr.wfc = exp.(-(par.x .- wfcoffset).^2/2)
    if (par.im_time)
        opr.K = exp.(-0.5*par.k.^2*par.dt)
        opr.R = exp.(-0.5*opr.V*par.dt)
    else
        opr.K = exp.(-im*0.5*par.k.^2*par.dt)
        opr.R = exp.(-im*0.5*opr.V*par.dt)
    end

    return opr
end

# Function for the split-operator loop
function split_op!(par::Param, opr::Operators)

    for i = 1:par.timesteps
        # Half-step in real space
        opr.wfc = opr.wfc .* opr.R

        # fft to momentum space
        opr.wfc = fft(opr.wfc)

        # Full step in momentum space
        opr.wfc = opr.wfc .* opr.K

        # ifft back
        opr.wfc = ifft(opr.wfc)

        # final half-step in real space
        opr.wfc = opr.wfc .* opr.R

        # density for plotting and potential
        density = abs2.(opr.wfc)

        # renormalizing for imaginary time
        if (par.im_time)
            renorm_factor = sum(density) * par.dx

            for j = 1:length(opr.wfc)
                opr.wfc[j] /= sqrt(renorm_factor)
            end
        end

        # 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
        if ((i-1) % div(par.timesteps, 100) == 0)
            outfile = open("output" * string(lpad(string(i-1), 5, string(0)))
                                    * ".dat","w")

            # Outputting for gnuplot. Any plotter will do.
            for j = 1:length(density)
                write(outfile, string(par.x[j]) * "\t"
                               * string(density[j]) * "\t"
                               * string(real(opr.V[j])) * "\n")
            end

            close(outfile)
            println("Outputting step: ", i)
        end
    end
end

# We are calculating the energy to check <Psi|H|Psi>
function calculate_energy(par, opr)
    # Creating real, momentum, and conjugate wavefunctions
    wfc_r = opr.wfc
    wfc_k = fft(wfc_r)
    wfc_c = conj(wfc_r)

    # Finding the momentum and real-space energy terms
    energy_k = 0.5*wfc_c.*ifft((par.k.^2) .* wfc_k)
    energy_r = wfc_c.*opr.V .* wfc_r

    # Integrating over all space
    energy_final = 0
    for i = 1:length(energy_k)
        energy_final += real(energy_k[i] + energy_r[i])
    end

    return energy_final*par.dx
end

# main function
function main()
    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)
    println("Energy is: ", energy)
end

main()

#include <complex.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

// Using fftw3 library.
#include <fftw3.h>

struct params {
    double xmax;
    unsigned int res;
    double dt;
    unsigned int timesteps;
    double dx;
    double *x;
    double dk;
    double *k;
    bool im_time;
};

struct operators {
    size_t size;
    double complex *v;
    double complex *pe;
    double complex *ke;
    double complex *wfc;
};

void fft(double complex *x, size_t n, bool inverse) {
    double complex y[n];
    memset(y, 0, sizeof(y));
    fftw_plan p;

    if (inverse) {
        p = fftw_plan_dft_1d((int)n, (fftw_complex*)x, (fftw_complex*)y,
                             FFTW_BACKWARD, FFTW_ESTIMATE);
    } else {
        p = fftw_plan_dft_1d((int)n, (fftw_complex*)x, (fftw_complex*)y,
                             FFTW_FORWARD, FFTW_ESTIMATE);
    }

    fftw_execute(p);
    fftw_destroy_plan(p);

    for (size_t i = 0; i < n; ++i) {
        x[i] = y[i] / sqrt((double)n);
    }
}

void init_params(struct params *par, double xmax, unsigned int res, double dt,
                 unsigned int timesteps, bool im) {

    par->xmax = xmax;
    par->res = res;
    par->dt = dt;
    par->timesteps = timesteps;
    par->dx = 2.0 * xmax / res;
    par->x = malloc(sizeof(double) * res);
    par->dk = M_PI / xmax;
    par->k = malloc(sizeof(double) * res);
    par->im_time = im;

    for (size_t i = 0; i < res; ++i) {
        par->x[i] = xmax / res - xmax + (double)i * (2.0 * xmax / res);
        if (i < res / 2) {
            par->k[i] = (double)i * M_PI / xmax;
        } else {
            par->k[i] = ((double)i - res) * M_PI / xmax;
        }
    }
}

void init_operators(struct operators *opr, struct params par, double voffset,
                    double wfcoffset) {

    opr->size = par.res;
    opr->v = malloc(sizeof(double complex) * par.res);
    opr->pe = malloc(sizeof(double complex) * par.res);
    opr->ke = malloc(sizeof(double complex) * par.res);
    opr->wfc = malloc(sizeof(double complex) * par.res);

    for (size_t i = 0; i < par.res; ++i) {
        opr->v[i] = 0.5 * cpow(par.x[i] - voffset, 2);
        opr->wfc[i] = cexp(-cpow(par.x[i] - wfcoffset, 2) / 2.0);

        if (par.im_time) {
            opr->ke[i] = cexp(-0.5 * par.dt * cpow(par.k[i], 2));
            opr->pe[i] = cexp(-0.5 * par.dt * opr->v[i]);
        } else {
            opr->ke[i] = cexp(-0.5 * par.dt * cpow(par.k[i], 2) * I);
            opr->pe[i] = cexp(-0.5 * par.dt * opr->v[i] * I);
        }
    }
}

void split_op(struct params par, struct operators opr) {
    double density[opr.size];

    for (size_t i = 0; i < par.timesteps; ++i) {
        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        fft(opr.wfc, opr.size, false);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.ke[j];
        }

        fft(opr.wfc, opr.size, true);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        for (size_t j = 0; j < opr.size; ++j) {
            density[j] = pow(cabs(opr.wfc[j]), 2);
        }

        if (par.im_time) {
            double sum = 0;

            for (size_t j = 0; j < opr.size; ++j) {
                sum += density[j];
            }

            sum *= par.dx;

            for (size_t j = 0; j < opr.size; ++j) {
                opr.wfc[j] /= sqrt(sum);
            }
        }

        // Writing data into a file in the format of:
        // index, density, real potential.
        char filename[256];
        sprintf(filename, "output%lu.dat", i);
        FILE *fp = fopen(filename, "w");

        for (size_t i = 0; i < opr.size; ++i) {
            fprintf(fp, "%ld\t%f\t%f\n", i, density[i], creal(opr.v[i]));
        }

        fclose(fp);
    }
}

double calculate_energy(struct params par, struct operators opr) {
    double complex wfc_r[opr.size];
    double complex wfc_k[opr.size];
    double complex wfc_c[opr.size];
    memcpy(wfc_r, opr.wfc, sizeof(wfc_r));

    memcpy(wfc_k, opr.wfc, sizeof(wfc_k));
    fft(wfc_k, opr.size, false);

    for (size_t i = 0; i < opr.size; ++i) {
        wfc_c[i] = conj(wfc_r[i]);
    }

    double complex energy_k[opr.size];
    double complex energy_r[opr.size];

    for (size_t i = 0; i < opr.size; ++i) {
        energy_k[i] = wfc_k[i] * cpow(par.k[i] + 0.0*I, 2);
    }

    fft(energy_k, opr.size, true);

    for (size_t i = 0; i < opr.size; ++i) {
        energy_k[i] *= 0.5 * wfc_c[i];
        energy_r[i] = wfc_c[i] * opr.v[i] * wfc_r[i];
    }

    double energy_final = 0;

    for (size_t i = 0; i < opr.size; ++i) {
        energy_final += creal(energy_k[i] + energy_r[i]);
    }

    return energy_final * par.dx;
}

void free_params(struct params par) {
    free(par.x);
    free(par.k);
}

void free_operators(struct operators opr) {
    free(opr.v);
    free(opr.pe);
    free(opr.ke);
    free(opr.wfc);
}

int main() {
    struct params par;
    struct operators opr;

    init_params(&par, 5.0, 256, 0.05, 100, true);
    init_operators(&opr, par, 0.0, -1.0);

    split_op(par, opr);

    printf("the energy is %f\n", calculate_energy(par, opr));

    free_params(par);
    free_operators(opr);

    return 0;
}

#include <complex>
#include <vector>
#include <iostream>
#include <cstring>
#include <fstream>

// Using fftw3 library.
#include <fftw3.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

using complex = std::complex<double>;
using vector_real = std::vector<double>;
using vector_complex = std::vector<complex>;

struct Params {
    Params(double _xmax, unsigned int _res, double _dt, unsigned int _timesteps, bool im) {
        xmax = _xmax;
        res = _res;
        dt = _dt;
        timesteps = _timesteps;
        dx = 2.0 * xmax / res;
        x.reserve(res);
        dk = M_PI / xmax;
        k.reserve(res);
        im_time = im;

        for (size_t i = 0; i < res; ++i) {
            x.emplace_back(xmax / res - xmax + static_cast<double>(i) * (2.0 * xmax / res));
            if (i < res / 2) {
                k.push_back(static_cast<double>(i) * M_PI / xmax);
            } else {
                k.push_back((static_cast<double>(i) - res) * M_PI / xmax);
            }
        }
    }

    double xmax;
    unsigned int res;
    double dt;
    unsigned int timesteps;
    double dx;
    vector_real x;
    double dk;
    vector_real k;
    bool im_time;
};

struct Operators {
public:
    Operators(Params &par, double voffset,
              double wfcoffset) {
        size = par.res;
        v.reserve(size);
        pe.reserve(size);
        ke.reserve(size);
        wfc.reserve(size);

        for (size_t i = 0; i < size; ++i) {
            v.push_back(0.5 * pow(par.x[i] - voffset, 2));
            wfc.push_back(exp(-pow(par.x[i] - wfcoffset, 2) / 2.0));

            if (par.im_time) {
                ke.push_back(exp(-0.5 * par.dt * pow(par.k[i], 2)));
                pe.push_back(exp(-0.5 * par.dt * v[i]));
            } else {
                ke.push_back(exp(-0.5 * par.dt * pow(par.k[i], 2) * complex(0.0, 1.0)));
                pe.push_back(exp(-0.5 * par.dt * v[i] * complex(0.0, 1.0)));
            }
        }
    }

    size_t size;
    vector_complex v;
    vector_complex pe;
    vector_complex ke;
    vector_complex wfc;
};

void fft(vector_complex &x, bool inverse) {
    std::vector<std::complex<double>> y(x.size(), std::complex<double>(0.0, 0.0));
    fftw_plan p;

    fftw_complex *in = reinterpret_cast<fftw_complex*>(x.data());
    fftw_complex *out = reinterpret_cast<fftw_complex*>(y.data());
    p = fftw_plan_dft_1d(static_cast<int>(x.size()), in, out,
                         (inverse ? FFTW_BACKWARD : FFTW_FORWARD), FFTW_ESTIMATE);

    fftw_execute(p);
    fftw_destroy_plan(p);

    for (size_t i = 0; i < x.size(); ++i) {
        x[i] = y[i] / sqrt(static_cast<double>(x.size()));
    }
}

void split_op(Params &par, Operators &opr) {
    auto density = std::vector<double>(opr.size, 0);

    for (size_t i = 0; i < par.timesteps; ++i) {
        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        fft(opr.wfc, false);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.ke[j];
        }

        fft(opr.wfc, true);

        for (size_t j = 0; j < opr.size; ++j) {
            opr.wfc[j] *= opr.pe[j];
        }

        for (size_t j = 0; j < opr.size; ++j) {
            density[j] = pow(abs(opr.wfc[j]), 2);
        }

        if (par.im_time) {
            double sum = 0;

            for (size_t j = 0; j < opr.size; ++j) {
                sum += density[j];
            }

            sum *= par.dx;

            for (size_t j = 0; j < opr.size; ++j) {
                opr.wfc[j] /= sqrt(sum);
            }
        }

        // Writing data into a file in the format of:
        // index, density, real potential.
        std::stringstream filename_stream;
        filename_stream << "output" << i << ".dat";

        std::ofstream fstream = std::ofstream(filename_stream.str());

        if (fstream) {
            for (std::size_t i = 0; i < opr.size; ++i) {
                std::stringstream data_stream;

                data_stream << i << "\t" << density[i] << "\t" << real(opr.v[i]) << "\n";

                fstream.write(data_stream.str().c_str(), data_stream.str().length());
            }
        }

        fstream.close();
    }
}

double calculate_energy(Params &par, Operators &opr) {
    vector_complex wfc_r(opr.wfc);
    vector_complex wfc_k(opr.wfc);
    vector_complex wfc_c(opr.size);
    fft(wfc_k, false);

    for (size_t i = 0; i < opr.size; ++i) {
        wfc_c[i] = conj(wfc_r[i]);
    }

    vector_complex energy_k(opr.size);
    vector_complex energy_r(opr.size);

    for (size_t i = 0; i < opr.size; ++i) {
        energy_k[i] = wfc_k[i] * pow(complex(par.k[i], 0.0), 2);
    }

    fft(energy_k, true);

    for (size_t i = 0; i < opr.size; ++i) {
        energy_k[i] *= 0.5 * wfc_c[i];
        energy_r[i] = wfc_c[i] * opr.v[i] * wfc_r[i];
    }

    double energy_final = 0;

    for (size_t i = 0; i < opr.size; ++i) {
        energy_final += real(energy_k[i] + energy_r[i]);
    }

    return energy_final * par.dx;
}

int main() {
    Params par = Params(5.0, 256, 0.05, 100, true);
    Operators opr = Operators(par, 0.0, -1.0);

    split_op(par, opr);

    std::cout << "The energy is " << calculate_energy(par, opr) << "\n";

    return 0;
}

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()

import Data.Array.CArray
import Data.Complex
import Data.List (intercalate, transpose)
import Math.FFT (dft, idft)

type Vector = CArray Int (Complex Double)

(.*), (.+) :: Vector -> Vector -> Vector
a .* b = liftArray2 (*) a b
a .+ b = liftArray2 (+) a b

normalize :: Double -> Vector -> Vector
normalize dx v =
  let factor = 1 / sqrt dx / norm2 v :+ 0
   in liftArray (factor *) v

data Parameters = Parameters
  { xmax :: Double
  , res :: Int
  , dt :: Double
  , timesteps :: Int
  , dx :: Double
  , x :: Vector
  , dk :: Double
  , ks :: Vector
  , imTime :: Bool
  }

defaultParameters :: Parameters
defaultParameters = makeParameters 10 512 0.01 1000 True

makeParameters :: Double -> Int -> Double -> Int -> Bool -> Parameters
makeParameters xmax res dt timesteps imTime =
  let fi = fromIntegral
      rng = (0, res - 1)
      ks = [0 .. div res 2 - 1] ++ [-div res 2 .. -1]
   in Parameters
        xmax
        res
        dt
        timesteps
        (2 * xmax / fi res)
        (listArray rng $
         map (\n -> xmax * (-1 + 2 * fi n / fi res) :+ 0) [1 .. res])
        (pi / xmax)
        (listArray rng $ map ((:+ 0) . (pi / xmax *) . fi) ks)
        imTime

data Operators = Operators
  { v :: Vector
  , rStep :: Vector
  , kStep :: Vector
  , wfc :: Vector
  }

makeOperators :: Parameters -> Complex Double -> Complex Double -> Operators
makeOperators param v0 wfc0 =
  let rng = (0, res param - 1)
      time
        | imTime param = dt param :+ 0
        | otherwise = 0 :+ dt param
      v = liftArray (\x -> 0.5 * (x - v0) ^ 2) (x param)
      rStep = liftArray (\x -> exp (-0.5 * time * x)) v
      kStep = liftArray (\k -> exp (-0.5 * time * k ^ 2)) (ks param)
      wfc = liftArray (\x -> exp (-(x - wfc0) ^ 2 / 2)) (x param)
   in Operators v rStep kStep (normalize (dx param) wfc)

evolve :: Parameters -> Operators -> [Operators]
evolve param op@(Operators _ rStep kStep _) = iterate splitop op
  where
    splitop op = op {wfc = wfc' op}
    wfc' = norm . (rStep .*) . idft . (kStep .*) . dft . (rStep .*) . wfc
    norm = if imTime param then normalize (dx param) else id

calculateEnergy :: Parameters -> Operators -> Double
calculateEnergy param ops = (* dx param) . sum . map realPart $ elems totalE
  where
    totalE = potentialE .+ kineticE
    potentialE = wfcConj .* v ops .* wfc ops
    kineticOp = liftArray ((/ 2) . (^ 2)) (ks param)
    kineticE = wfcConj .* idft (kineticOp .* dft (wfc ops))
    wfcConj = liftArray conjugate $ wfc ops

-- Use gnuplot to make an animated  GIF using ../gnuplot/plot_output.plt
-- $ gnuplot -e "folder='../haskell'" plot_output.plt
printEvolution :: Parameters -> [Operators] -> IO ()
printEvolution param =
  mapM_ (export . (format <$>)) . zip [0 ..] . take 100 . skip
  where
    skip (x:xs) = x : skip (drop (div (timesteps param) 100 - 1) xs)
    format (Operators v _ _ wfc) =
      let density = liftArray ((^ 2) . abs) wfc
          values = map (map (show . realPart) . elems) [x param, density, v]
       in intercalate "\n" $ map (intercalate "\t") $ transpose values
    export (i, f) = writeFile ("output" ++ pad (show i) ++ ".dat") f
    pad n = replicate (5 - length n) '0' ++ n

main :: IO ()
main = do
  let p = defaultParameters
      o = makeOperators p 0 4
      evol = evolve p o
  print $ calculateEnergy p (evol !! timesteps p)
  printEvolution p evol

extern crate rustfft;

use rustfft::num_complex::Complex;
use rustfft::FFTplanner;
use std::f64::consts::PI;
use std::fs::File;
use std::io::Write;
use std::path::Path;

// This implementation is based on the C and C++ implementations.

#[derive(Clone)]
struct Parameters {
    xmax: f64,
    res: usize,
    dt: f64,
    timesteps: usize,
    dx: f64,
    x: Vec<f64>,
    dk: f64,
    k: Vec<f64>,
    im_time: bool,
}

impl Parameters {
    pub fn new(xmax: f64, res: usize, dt: f64, timesteps: usize, im_time: bool) -> Parameters {
        let dx = 2.0_f64 * xmax / (res as f64);
        let mut x: Vec<f64> = Vec::with_capacity(res);
        let dk = PI / xmax;
        let mut k: Vec<f64> = Vec::with_capacity(res);
        for i in 0..res {
            x.push(xmax / (res as f64) - xmax + (i as f64) * dx);
            match i {
                i if (i < res / 2) => k.push((i as f64) * PI / xmax),
                _ => k.push(((i as f64) - (res as f64)) * PI / xmax),
            }
        }
        Parameters {
            xmax,
            res,
            dt,
            timesteps,
            im_time,
            dx,
            x,
            dk,
            k,
        }
    }
}

struct Operators {
    v: Vec<Complex<f64>>,
    pe: Vec<Complex<f64>>,
    ke: Vec<Complex<f64>>,
    wfc: Vec<Complex<f64>>,
}

impl Operators {
    pub fn new(par: &Parameters, v_offset: f64, wfc_offset: f64) -> Operators {
        let mut v: Vec<Complex<f64>> = Vec::with_capacity(par.res);
        let mut pe: Vec<Complex<f64>> = Vec::with_capacity(par.res);
        let mut ke: Vec<Complex<f64>> = Vec::with_capacity(par.res);
        let mut wfc: Vec<Complex<f64>> = Vec::with_capacity(par.res);

        for i in 0..par.res {
            v.push(Complex::new(
                0.5_f64 * (par.x[i] - v_offset).powi(2),
                0.0_f64,
            ));
            wfc.push(Complex::new(
                (-((par.x[i] - wfc_offset).powi(2)) / 2.0_f64).exp(),
                0.0_f64,
            ));
            if par.im_time {
                ke.push(Complex::new(
                    (-0.5_f64 * par.dt * par.k[i].powi(2)).exp(),
                    0.0_f64,
                ));
                pe.push(Complex::new((-0.5_f64 * par.dt * v[i].re).exp(), 0.0_f64));
            } else {
                ke.push(Complex::new(
                    0.0_f64,
                    (-0.5_f64 * par.dt * par.k[i].powi(2)).exp(),
                ));
                pe.push(Complex::new(0.0_f64, (-0.5_f64 * par.dt * v[i].re).exp()));
            }
        }
        Operators { v, pe, ke, wfc }
    }
}

fn fft(x: &mut Vec<Complex<f64>>, inverse: bool) {
    let mut y = vec![Complex::new(0.0_f64, 0.0_f64); x.len()];
    let mut p = FFTplanner::new(inverse);
    let fft = p.plan_fft(x.len());
    fft.process(x.as_mut_slice(), y.as_mut_slice());

    for i in 0..x.len() {
        x[i] = y[i] / (x.len() as f64).sqrt();
    }
}

fn split_op(par: &Parameters, opr: &mut Operators) {
    let mut density: Vec<f64>;

    for i in 0..par.timesteps {
        for j in 0..par.res {
            opr.wfc[j] *= opr.pe[j];
        }

        fft(&mut opr.wfc, false);

        for j in 0..par.res {
            opr.wfc[j] *= opr.ke[j];
        }

        fft(&mut opr.wfc, true);

        for j in 0..par.res {
            opr.wfc[j] *= opr.pe[j];
        }

        density = opr.wfc.iter().map(|x| x.norm().powi(2)).collect();

        if par.im_time {
            let sum = density.iter().sum::<f64>() * par.dx;

            for j in 0..par.res {
                opr.wfc[j] /= sum.sqrt();
            }
        }

        // Writing data into a file in the format of:
        // index, density, real potential.
        let path_name = format!("output{}.dat", i);
        let path = Path::new(&path_name);
        let display = path.display();

        let mut file = match File::create(&path) {
            Err(why) => panic!("Couldn't create {}: {}", display, why),
            Ok(good) => good,
        };

        for j in 0..par.res {
            if let Err(why) = writeln!(file, "{}\t{}\t{}", j, density[j], opr.v[j].re) {
                panic!("Couldn't write to {}: {}", display, why)
            }
            if let Err(why) = file.flush() {
                panic!("Couldn't flush {}: {}", display, why)
            }
        }
    }
}

fn calculate_energy(par: &Parameters, opr: &Operators) -> f64 {
    let wfc_r = opr.wfc.clone();
    let mut wfc_k = opr.wfc.clone();
    let mut wfc_c = vec![Complex::new(0.0_f64, 0.0_f64); par.res];

    fft(&mut wfc_k, false);

    for i in 0..par.res {
        wfc_c[i] = wfc_r[i].conj();
    }

    let mut energy_k = vec![Complex::new(0.0_f64, 0.0_f64); par.res];
    let mut energy_r = vec![Complex::new(0.0_f64, 0.0_f64); par.res];

    for i in 0..par.res {
        energy_k[i] = wfc_k[i] * Complex::new(par.k[i], 0.0_f64).powi(2);
    }

    fft(&mut energy_k, true);

    for i in 0..par.res {
        energy_k[i] *= wfc_c[i].scale(0.5_f64);
        energy_r[i] = wfc_c[i] * opr.v[i] * wfc_r[i];
    }

    let energy_final = energy_k
        .into_iter()
        .zip(energy_r.into_iter())
        .fold(0.0_f64, |acc, x| acc + (x.0 + x.1).re);

    energy_final * par.dx
}

fn main() {
    let par = Parameters::new(5.0, 256, 0.05, 100, true);
    let mut opr = Operators::new(&par, 0.0, -1.0);

    split_op(&par, &mut opr);

    println!("The energy is {}", calculate_energy(&par, &opr));
}

License

Code Examples

The code examples are licensed under the MIT license (found in LICENSE.md).

Text

The text of this chapter was written by James Schloss and is licensed under the Creative Commons Attribution-ShareAlike 4.0 International License.

Images/Graphics
Pull Requests

After initial licensing (#560), the following pull requests have modified the text or graphics of this chapter:

  • none

results matching ""

    No results matching ""