Obtain Stationary States via Imaginary Time Propagation

The provided Python code defines a class ImgTimePropagation that calculates stationary states using imaginary time propagation. This process can be described as follows:

1. Initialization: Initialize an empty list to store the stationary states: self.stationary_states = [].

2. Iterate Through Desired States:

   • For each desired stationary state from 0 to (nstates-1):
     • Initialize the wavefunction`$ based on parity:
       • If (n) is even: wavefunction = exp(-x^2)
       • If (n) is odd: wavefunction = x * exp(-x^2)
     • Perform nsteps imaginary time steps:
       • Multiply the wavefunction by img_exp_v [$= (-1)^k \cdot e^{-0.5 \cdot \text{dt} \cdot V}$]to account for the potential energy.
       • Transform the wavefunction to momentum representation using FFT and multiply it by img_exp_k[$= e^{-\text{dt} \cdot K}$] for kinetic energy.
       • Transform the wavefunction back to the coordinate representation using inverse FFT and multiply it by img_exp_v.
     • Normalize the wavefunction: $\text{wavefunction} \div \left( \left\| \text{wavefunction} \right\| \cdot \sqrt{\text{dx}} \right)$
     • Project out previously obtained stationary states:
       • Calculate the projections (projs) onto the stationary states.
       • $\text{projs} = \left[ \langle \psi_i, \text{wavefunction} \rangle \cdot \text{dx} \right] \quad \text{for } i = 1, 2, \ldots, n_{\text{states}}$
       • Subtract the projections of stationary states to orthogonalize.
       • $\text{For each } i = 1, 2, \ldots, n_{\text{states}}: \text{wavefunction} -= \text{proj}_i \cdot \psi_i$
     • Normalize the wavefunction again: $\text{wavefunction} \div \left( \left\| \text{wavefunction} \right\| \cdot \sqrt{\text{dx}} \right)$
     • Save the obtained stationary state in self.stationary_states.

This process repeats for each desired stationary state.

Symbols and Variables:

• nstates: Number of stationary states to obtain.
• nsteps: Number of imaginary time steps.
• wavefunction: The wavefunction.
• img_exp_v: Imaginary time exponent of the potential energy.
• img_exp_k: Imaginary time exponent of the kinetic energy.
• even: Boolean flag for the parity of the wavefunction.
• self.stationary_states: List to store stationary states.

This procedure allows obtaining stationary states of a quantum system through imaginary time propagation.

from split_op_schrodinger1D import SplitOpSchrodinger1D, fftpack, np, linalg

class ImgTimePropagation(SplitOpSchrodinger1D):

    def get_stationary_states(self, nstates, nsteps=10000):
        """
        Obtain stationary states via the imaginary time propagation
        :param nstates: number of states to obtaine.
                If nstates = 1, only the ground state is obtained. If nstates = 2,
                the ground and first exited states are obtained, etc
        :param nsteps: number of the imaginary time steps to take
        :return:self
        """
        # since there is no time dependence (self.t) during the imaginary time propagation
        # pre-calculate imaginary time exponents of the potential and kinetic energy
        # img_expV = ne.evaluate("(-1) ** k * exp(-0.5 * dt * ({}))".format(self.V), local_dict=vars(self))

        img_exp_v = (-1) ** np.arange(self.wavefunction.size) * np.exp(-0.5 * self.dt * self.v(self.x, self.t))

        # img_expK = ne.evaluate("exp(-dt * ({}))".format(self.K), local_dict=vars(self))
        img_exp_k = np.exp(-self.dt * self.k(self.p, self.t))


        # initialize the list where the stationary states will be saved
        self.stationary_states = []

        # boolean flag determining the parity of wavefunction
        even = True

        for n in range(nstates):

            # allocate and initialize the wavefunction depending on the parity.
            # Note that you do not have to be fancy and can choose any initial guess (even random).
            # the more reasonable initial guess, the faster the convergence.
            wavefunction = (np.exp(-self.x ** 2) if even else self.x * np.exp(-self.x ** 2))
            even = not even

            for _ in range(nsteps):
                #################################################################################
                #
                #   Make an imaginary time step
                #
                #################################################################################
                wavefunction *= img_exp_v

                # going to the momentum representation
                wavefunction = fftpack.fft(wavefunction, overwrite_x=True)
                wavefunction *= img_exp_k

                # going back to the coordinate representation
                wavefunction = fftpack.ifft(wavefunction, overwrite_x=True)
                wavefunction *= img_exp_v

                #################################################################################
                #
                #    Project out all previously calculated stationary states
                #
                #################################################################################

                # normalize
                wavefunction /= linalg.norm(wavefunction) * np.sqrt(self.dx)

                # calculate the projections
                projs = [np.vdot(psi, wavefunction) * self.dx for psi in self.stationary_states]

                # project out the stationary states
                for psi, proj in zip(self.stationary_states, projs):
                    wavefunction -= proj * psi
                    # ne.evaluate("wavefunction - proj * psi", out=wavefunction)

                # normalize
                wavefunction /= linalg.norm(wavefunction) * np.sqrt(self.dx)

            # save obtained approximation to the stationary state
            self.stationary_states.append(wavefunction)

        return self