Wavefunctions.py

"""
Provides a DVRWavefunction class that inherits from the base Psience wavefunction
"""
from Psience.Wavefun import Wavefunction, Wavefunctions


class DVRWavefunction(Wavefunction):
    def plot(self, figure=None, grid=None, **opts):
        import numpy as np

        if grid is None:
            grid = self.opts['grid']

        dim = len(grid.shape)
        if dim > 1 and grid.shape[-1] == dim-1:  # check whether we have a mesh of points that we need to reshape
            unroll = np.roll(np.arange(len(grid.shape)), 1)
            grid = grid.transpose(unroll)

        if dim == 1:
            if figure is None:
                from McUtils.Plots import Plot
                return Plot(grid, self.data, **opts)
            else:
                return figure.plot(grid, self.data, **opts)
        else:
            if figure is None:
                from McUtils.Plots import Plot3D
                return Plot3D(*grid, self.data.reshape(grid[0].shape), **opts)
            else:
                return figure.plot(*grid, self.data.reshape(grid[0].shape), **opts)

    def expectation(self, op, other):
        """Computes the expectation value of operator op over the wavefunction other and self

        :param other:
        :type other: Wavefunction | np.ndarray
        :param op:
        :type op:
        :return:
        :rtype:
        """
        import numpy as np

        wf = op(self.data)
        if not isinstance(other, np.ndarray):
            other = other.data
        return np.dot(wf, other)

    def probability_density(self):
        """Computes the probability density of the current wavefunction

        :return:
        :rtype:
        """
        import numpy as np

        return np.power(self.data, 2)


class DVRWavefunctions(Wavefunctions):
    # most evaluations are most efficient done in batch for DVR wavefunctions so we focus on the batch object

    def __init__(self, energies=None, wavefunctions=None,
                 wavefunction_class=None,
                 rephase=True,
                 **opts
                 ):
        import numpy as np

        if rephase:
            phase_gs = np.sign(wavefunctions[:, 0])
            wavefunctions = wavefunctions*phase_gs[:, np.newaxis]
        super().__init__(wavefunctions=wavefunctions, energies=energies, wavefunction_class=DVRWavefunction, **opts)

    def __getitem__(self, item):
        """Returns a single Wavefunction object"""
        # iter comes for free with this
        if isinstance(item, slice):
            return type(self)(
                energies=self.energies[item],
                wavefunctions=self.wavefunctions[:, item],
                wavefunction_class=self.wavefunction_class,
                **self.opts
            )
        else:
            return self.wavefunction_class(self.energies[item], self.wavefunctions[:, item], parent=self, **self.opts)

    def plot(self, figure=None, graphics_class=None, plot_class=None, plot_style=None, **opts):
        import numpy as np
        from McUtils.Plots import Plot, Plot3D

        grid = self.opts['grid']

        dim = len(grid.shape)
        if dim > 1 and grid.shape[-1] == dim-1:  # check whether we have a mesh of points that we need to reshape
            unroll = np.roll(np.arange(len(grid.shape)), 1)
            grid = grid.transpose(unroll)

        # if plot_class is None:
        #     if dim == 1:
        #         plot_class = Plot
        #     elif dim == 2:
        #         plot_class = Plot3D
        #     else:
        #         print("I don't know how to plot that")
        #         # raise DVRException("{}.{}: don't know how to plot {} dimensional potential".format(
        #         #     type(self).__name__,
        #         #     'plot',
        #         #     dim
        #         # ))

        super().plot(
            figure=figure,
            graphics_class=graphics_class,
            plot_style=plot_style,
            **opts
        )

    def expectation(self, op, other):
        """Computes the expectation value of operator op over the wavefunction other and self

        :param other:
        :type other: DVRWavefunctions | np.ndarray
        :param op:
        :type op:
        :return:
        :rtype:
        """
        import numpy as np

        wfs = op(self.wavefunctions)
        if not isinstance(other, np.ndarray):
            other = other.wavefunctions
        return np.dot(other.T, wfs)

    def probability_density(self):
        """Computes the probability density of the set of wavefunctions

        :return:
        :rtype:
        """
        import numpy as np

        return np.power(self.wavefunctions, 2)