molecule_2state_wigner_moyal.py 文件源码

def single_step_propagation(self):
        """
        Perform single step propagation. The final Wigner functions are not normalized.
        """
        ################ p x -> theta x ################
        self.wigner_ge = fftpack.fft(self.wigner_ge, axis=0, overwrite_x=True)
        self.wigner_g = fftpack.fft(self.wigner_g, axis=0, overwrite_x=True)
        self.wigner_e = fftpack.fft(self.wigner_e, axis=0, overwrite_x=True)

        # Construct T matricies
        TgL, TgeL, TeL = self.get_T_left(self.t)
        TgR, TgeR, TeR = self.get_T_right(self.t)

        # Save previous version of the Wigner function
        Wg, Wge, We = self.wigner_g, self.wigner_ge, self.wigner_e

        # First update the complex valued off diagonal wigner function
        self.wigner_ge = (TgL*Wg + TgeL*Wge.conj())*TgeR + (TgL*Wge + TgeL*We)*TeR

        # Slice arrays to employ the symmetry (savings in speed)
        TgL, TgeL, TeL = self.theta_slice(TgL, TgeL, TeL)
        TgR, TgeR, TeR = self.theta_slice(TgR, TgeR, TeR)
        Wg, Wge, We = self.theta_slice(Wg, Wge, We)

        # Calculate the remaning real valued Wigner functions
        self.wigner_g = (TgL*Wg + TgeL*Wge.conj())*TgR + (TgL*Wge + TgeL*We)*TgeR
        self.wigner_e = (TgeL*Wg + TeL*Wge.conj())*TgeR + (TgeL*Wge + TeL*We)*TeR

        ################ Apply the phase factor ################
        self.wigner_ge *= self.expV
        self.wigner_g *= self.expV[:(1 + self.P_gridDIM//2), :]
        self.wigner_e *= self.expV[:(1 + self.P_gridDIM//2), :]

        ################ theta x -> p x ################
        self.wigner_ge = fftpack.ifft(self.wigner_ge, axis=0, overwrite_x=True)
        self.wigner_g = fft.irfft(self.wigner_g, axis=0)
        self.wigner_e = fft.irfft(self.wigner_e, axis=0)

        ################ p x  ->  p lambda ################
        self.wigner_ge = fftpack.fft(self.wigner_ge, axis=1, overwrite_x=True)
        self.wigner_g = fft.rfft(self.wigner_g, axis=1)
        self.wigner_e = fft.rfft(self.wigner_e, axis=1)

        ################ Apply the phase factor ################
        self.wigner_ge *= self.expK
        self.wigner_g *= self.expK[:, :(1 + self.X_gridDIM//2)]
        self.wigner_e *= self.expK[:, :(1 + self.X_gridDIM//2)]

        ################ p lambda  ->  p x ################
        self.wigner_ge = fftpack.ifft(self.wigner_ge, axis=1, overwrite_x=True)
        self.wigner_g = fft.irfft(self.wigner_g, axis=1)
        self.wigner_e = fft.irfft(self.wigner_e, axis=1)

        #self.normalize_wigner_matrix()