tbkit 0.2.0__py3-none-any.whl

tbkit/propagation.py

@@ -0,0 +1,338 @@
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+import scipy.sparse as sparse
+import scipy.linalg as LA
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+from matplotlib.animation import FuncAnimation
+import os
+import tbkit.error_handling as error_handling
+from tbkit.lattice import Lattice
+
+
+
+#################################
+# CLASS PROPAGATION
+#################################
+
+
+class Propagation():
+    '''
+    Get lattice time evolution. Time dependent Schrodinger equation solved by
+    Crank-Nicolson method.
+
+    :param lat: **lattice** class instance.
+    '''
+
+    def __init__(self, lat: Lattice) -> None:
+        error_handling.lat(lat)
+        self.lat = lat
+        self.prop = np.array([], 'c16')
+
+    def get_propagation(
+        self, ham: sparse.spmatrix, psi_init: NDArray, steps: int, dz: float, norm: bool = False,
+    ) -> None:
+        '''
+        Get the time evolution.
+
+        :param ham: sparse.csr_matrix. Tight-Binding Hamilonian.
+        :param psi_init: np.ndarray. Initial state.
+        :param steps: Positive Integer. Number of steps.
+        :param dz: Positive number. Step.
+        :param norm: Boolean. Default value True. Normalize the norm to 1 at each step.
+        '''
+        error_handling.empty_ham(ham)
+        error_handling.ndarray(psi_init, 'psi_init', self.lat.sites)
+        error_handling.positive_int(steps, 'steps')
+        error_handling.positive_real(dz, 'dz')
+        error_handling.boolean(norm, 'norm')
+        self.steps = steps
+        self.dz = dz
+        self.prop = np.empty((self.lat.sites, self.steps), 'c16')
+        self.prop[:, 0] = psi_init
+        diag = 1j*np.ones(self.lat.sites, 'c16')
+        A = (sparse.diags(diag, 0) - 0.5 * self.dz * ham).toarray()
+        B = (sparse.diags(diag, 0) + 0.5 * self.dz * ham).toarray()
+        mat = (np.dot(LA.inv(A), B))
+        for i in range(1, self.steps):
+            self.prop[:, i] = np.dot(mat, self.prop[:, i-1])
+            if norm:
+                self.prop[:, i] /= np.abs(self.prop[:, i]).sum()
+
+    def get_pumping(
+        self, hams: list[sparse.spmatrix], psi_init: NDArray, steps: int, dz: float, norm: bool = True,
+    ) -> None:
+        '''
+        Get the time evolution with adiabatic pumpings.
+
+        :param hams: List of sparse.csr_matrices. Tight-Binding Hamilonians.
+        :param psi_init: np.ndarray. Initial state.
+        :param steps: Positive integer. Number of steps.
+        :param dz: Positive number. Step.
+        :param norm: Boolean. Default value True. Normalize the norm to 1 at each step.
+        '''
+        error_handling.get_pump(hams)
+        error_handling.ndarray(psi_init, 'psi_init', self.lat.sites)
+        error_handling.positive_int(steps, 'steps')
+        error_handling.positive_real(dz, 'dz')
+        error_handling.boolean(norm, 'norm')
+        self.steps = steps
+        self.dz = dz
+        no = len(hams)
+        self.prop = np.empty((self.lat.sites, self.steps), 'c16')
+        self.prop[:, 0] = psi_init
+        diag = 1j * np.ones(self.lat.sites, 'c16')
+        delta = self.steps // (1 + no)
+        A = (sparse.diags(diag, 0) - 0.5 * self.dz * hams[0]).toarray()
+        B = (sparse.diags(diag, 0) + 0.5 * self.dz * hams[0]).toarray()
+        mat = (np.dot(LA.inv(A), B))
+        # before pumping
+        for i in range(1, delta):
+           self.prop[:, i] = np.dot(mat, self.prop[:, i-1])
+           if norm:
+               self.prop[:, i] /= np.abs(self.prop[:, i]).sum()
+      # pumping
+        c = np.linspace(0, 1, delta)
+        for j in range(0, no-1):
+            for i in range(0, delta):
+                ham = (1-c[i])*hams[j]+c[i]*hams[j+1]
+                A = (sparse.diags(diag, 0) - 0.5 * self.dz * ham).toarray()
+                B = (sparse.diags(diag, 0) + 0.5 * self.dz * ham).toarray()
+                mat = (np.dot(LA.inv(A), B))
+                self.prop[:, (j+1)*delta+i] = np.dot(mat, self.prop[:,  (j+1)*delta+i-1])
+                if norm:
+                    self.prop[:,  (j+1)*delta+i] /= \
+                        np.abs(self.prop[:,  (j+1)*delta+i]).sum() 
+      # after pumping
+        j = no
+        for i in range(0, self.steps - no*delta):
+            self.prop[:,  no*delta+i] = np.dot(mat, self.prop[:,  no*delta+i-1])
+            if norm:
+                self.prop[:,  no*delta+i] /= np.abs(self.prop[:,  no*delta+i]).sum()
+
+    def plt_propagation_1d(
+        self, prop_type: str = 'real', fs: float = 20, figsize: tuple[float, float] | None = None,
+    ) -> Figure:
+        '''
+        Plot time evolution for 1D systems. 
+
+        :param fs: Default value 20. Fontsize.
+        '''
+        error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')
+        error_handling.positive_real(fs, 'fs')
+        error_handling.prop_type(prop_type)
+        error_handling.tuple_2elem(figsize, 'figsize')
+        fig, ax = plt.subplots(figsize=figsize)
+        plt.ylabel('n', fontsize=fs)
+        plt.xlabel('z', fontsize=fs)
+        if prop_type == 'real':
+            color = self.prop_smooth_1d(self.prop.real)
+            max_val = max(np.max(color), -np.min(color))
+            ticks = [-max_val, max_val]
+            cmap = 'seismic'
+        elif prop_type == 'imag':
+            color = self.prop_smooth_1d(self.prop.imag)
+            max_val = max(np.max(color), -np.min(color))
+            ticks = [-max_val, max_val]
+            cmap = 'seismic'
+        else:
+            color = self.prop_smooth_1d(np.abs(self.prop) ** 2)
+            ticks = [0., np.max(color[:, -1])]
+            cmap = plt.cm.hot
+        extent = (-0, self.steps*self.dz, self.lat.sites-.5, -.5)
+        aspect = 'auto'
+        interpolation = 'nearest'
+        im = plt.imshow(color, cmap=cmap, aspect=aspect,
+                                  interpolation=interpolation, extent=extent,
+                                  vmin=ticks[0], vmax=ticks[-1])
+        for label in ax.xaxis.get_majorticklabels():
+            label.set_fontsize(fs)
+        for label in ax.yaxis.get_majorticklabels():
+            label.set_fontsize(fs)
+        ax.get_yaxis().set_major_locator(plt.MaxNLocator(integer=True))
+        if prop_type == 'norm':
+            cbar = fig.colorbar(im, ticks=ticks)
+            cbar.ax.set_yticklabels(['0','max'])
+        else:
+            cbar = fig.colorbar(im, ticks=[ticks[0], 0, ticks[1]])
+            cbar.ax.set_yticklabels(['min', '0','max'])
+        cbar.ax.tick_params(labelsize=fs)
+        return fig
+
+    def prop_smooth_1d(self, prop: NDArray, a: float = 10, no: int = 40) -> NDArray:
+        r'''
+        Private function. Used in *plt_propagation_1d*.
+        Smooth propagation for 1D systems.
+        Perform Gaussian interpolation :math:`e^{-a(x-x_i)^2}`,
+
+        :param prop: Propagation.
+        :param a: Default value 15. Gaussian Parameter.
+        :param no: Default value 40. Number of points of each Gaussian.
+
+        :returns:
+           * **smooth** -- Smoothed propagation.
+        '''
+        func = np.exp(- a * np.linspace(-0.5, 0.5, no) ** 2) 
+        smooth = np.empty((self.lat.sites * no, self.steps))
+        for iz in range(0, self.steps):
+            for i in range(self.lat.sites):
+                smooth[i*no: (i+1)*no, iz] = prop[i, iz] * func
+        return smooth
+
+    def get_animation(
+        self, s: float = 300., fs: float = 20., prop_type: str = 'real',
+        figsize: tuple[float, float] | None = None,
+    ) -> FuncAnimation:
+        '''
+        Get time evolution animation.
+
+        :param s: Default value 300. Circle size.
+        :param fs: Default value 20. Fontsize.
+        :param figsize: Tuple. Default value None. Figsize.
+        :param prop_type: Default value None. Figsize.
+
+        :returns:
+          * **ani** -- Animation.
+        '''
+        error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')
+        error_handling.positive_real(s, 's')
+        error_handling.positive_real(fs, 'fs')
+        error_handling.prop_type(prop_type)
+        error_handling.tuple_2elem(figsize, 'figsize')
+        if os.name == 'posix':
+            blit = False
+        else:
+            blit = True
+        if prop_type == 'real':
+            color = self.prop.real
+            max_val = max(np.max(color), -np.min(color))
+            ticks = [-max_val, max_val]
+            cmap = 'seismic'
+        elif prop_type == 'imag':
+            color = self.prop.imag
+            max_val = max(np.max(color), -np.min(color))
+            ticks = [-max_val, max_val]
+            cmap = 'seismic'
+        else:
+            color = np.abs(self.prop) ** 2
+            ticks = [0., np.max(color)]
+            cmap = 'Reds'
+        fig, ax = plt.subplots(figsize=figsize)
+        plt.xlim([self.lat.coor['x'][0]-1., self.lat.coor['x'][-1]+1.])
+        plt.ylim([self.lat.coor['y'][0]-1., self.lat.coor['y'][-1]+1.])
+        scat = plt.scatter(self.lat.coor['x'], self.lat.coor['y'], c=color[:, 0],
+                                   s=s, vmin=ticks[0], vmax=ticks[1],
+                                   cmap=plt.get_cmap(cmap))
+        frame = plt.gca()
+        frame.axes.get_xaxis().set_ticks([])
+        frame.axes.get_yaxis().set_ticks([])
+        ax.set_aspect('equal')
+        if prop_type == 'norm':
+            cbar = fig.colorbar(scat, ticks=ticks)
+            cbar.ax.set_yticklabels(['0','max'])
+        else:
+            cbar = fig.colorbar(scat, ticks=[ticks[0], 0, ticks[1]])
+            cbar.ax.set_yticklabels(['min', '0','max'])
+
+        def update(i, color, scat):
+            scat.set_array(color[:, i])
+            return scat,
+
+        ani = FuncAnimation(fig, update, frames=self.steps,
+                                                  fargs=(color, scat), blit=blit, repeat=False)
+        return ani
+
+    def get_animation_nb(
+        self, s: float = 300., fs: float = 20., prop_type: str = 'real',
+        figsize: tuple[float, float] | None = None,
+    ) -> FuncAnimation:
+        '''
+        Get time evolution animation for iPython notebooks.
+
+        :param s: Default value 300. Circle shape.
+        :param fs: Default value 20. Fontsize.
+
+        :returns:
+           * **ani** -- Animation.
+        '''
+        '''
+        Get time evolution animation.
+
+        :param s: Default value 300. Circle size.
+        :param fs: Default value 20. Fontsize.
+        :param figsize: Tuple. Default value None. Figsize.
+        :param prop_type: Default value None. Figsize.
+
+        :returns:
+          * **ani** -- Animation.
+        '''
+        error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')
+        error_handling.positive_real(s, 's')
+        error_handling.positive_real(fs, 'fs')
+        error_handling.prop_type(prop_type)
+        error_handling.tuple_2elem(figsize, 'figsize')
+        if prop_type == 'real' or prop_type == 'imag':
+            color = self.prop.real
+            max_val = max(np.max(color[:, -1]), -np.min(color[:, -1]))
+            ticks = [-max_val, max_val]
+            cmap = 'seismic'
+        else:
+            color = np.abs(self.prop) ** 2
+            ticks = [0., np.max(color)]
+            cmap = 'Reds'
+        fig = plt.figure()
+        ax = plt.axes(xlim=(np.min(self.lat.coor['x']-.5), np.max(self.lat.coor['x']+.5)), 
+                             ylim=(np.min(self.lat.coor['y']-.5), np.max(self.lat.coor['y']+.5)))
+        ax.set_aspect('equal')
+        frame = plt.gca()
+        frame.axes.get_xaxis().set_ticks([])
+        frame.axes.get_yaxis().set_ticks([])
+        scat = plt.scatter(self.lat.coor['x'], self.lat.coor['y'], c=color[:, 0],
+                                    s=s, vmin=ticks[0], vmax=ticks[1],
+                                    cmap=cmap)
+        if prop_type == 'real' or prop_type == 'imag':
+            cbar = fig.colorbar(scat, ticks=[ticks[0], 0, ticks[1]])
+            cbar.ax.set_yticklabels(['min', '0','max'])
+        else:
+            cbar = fig.colorbar(scat, ticks=[0, ticks[1]])
+            cbar.ax.set_yticklabels(['0','max'])
+
+        def init():
+            scat.set_array(color[:, 0])
+            return scat,
+
+        def animate(i):
+            scat.set_array(color[:, i])    
+            return scat,
+
+        return FuncAnimation(fig, animate, init_func=init,
+                                   frames=self.steps, interval=120, blit=True)
+
+    def plt_prop_dimer(self, lw: float = 5, fs: float = 20) -> Figure:
+        '''
+        Plot time evolution for dimers.
+        
+        :param lw: Default value 5. Linewidth.
+        :param fs: Default value 20. Fontsize.
+
+        :returns:
+           * **fig** -- Figure.
+        '''
+        if not self.prop.any():
+            raise Exception('\n\nRun method get_prop() or get_pump() first.\n')
+        color = ['b', 'r']
+        fig, ax = plt.subplots()
+        z = self.dz * np.arange(self.steps)
+        for i, c in zip([0, 1], color):
+            plt.plot(z, np.abs(self.prop[i, :])**2, c, lw=lw)
+        plt.title('Intensity', fontsize=fs)
+        plt.xlabel('$z$', fontsize=fs)
+        plt.ylabel(r'$|\psi_j|^2$', fontsize=fs)
+        plt.xlim([0, z[-1]])
+        return fig
+
+
+# Backward-compatible lowercase alias (pre-0.2 API).
+propagation = Propagation

tbkit/save.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import os
+from matplotlib.figure import Figure
+from matplotlib.animation import FuncAnimation
+import tbkit.error_handling as error_handling
+
+
+#################################
+# CLASS SAVE
+#################################
+
+
+class Save():
+    '''
+    Create folder and save figures / animations obtained via
+    **plot** or **propagation**.
+
+    :param dir_name: String. Name of the sub-directory the figures are stored in.
+    :param dir_main: String. Default value None. Path of the main directory.
+      If None, figures are stored under ``'figs/'``.
+    :param params: Dictionary. Default value None. Parameters appended to file names.
+    :param file_format: Default value 'png'. Figure format.
+    '''
+    def __init__(
+        self,
+        dir_name: str,
+        dir_main: str | None = None,
+        params: dict | None = None,
+        file_format: str = 'png',
+    ) -> None:
+        error_handling.string(dir_name, 'dir_name')
+        error_handling.string(dir_main, 'dir_main')
+        error_handling.file_format(file_format)
+        self.params = {} if params is None else params
+        self.file_format = file_format
+        if dir_main is None:
+            self.dir_main = 'figs/'
+        else:
+            self.dir_main = dir_main
+        self.dir_name = self.dir_main + dir_name
+        self.create_dir()
+
+    def create_dir(self) -> None:
+        '''
+        Create the directory to store the figures exists.
+        '''
+        if not os.path.exists(self.dir_main):
+            os.makedirs(self.dir_main)
+        if not os.path.exists(self.dir_name):
+            os.makedirs(self.dir_name)
+
+    def file_name(self) -> str:
+        '''
+        Create the file name.
+
+        :returns:
+            * **file_name** -- File name.
+        '''
+        file_name = ''
+        for key, val in self.params.items():
+            file_name += '_' + key + str(complex(val+0)).replace('.', ',')
+        return file_name
+
+    def fig(self, fig: Figure, name: str) -> None:
+        '''
+        Save the figure in the directory defined by the method *dir_name()*.
+
+        :param fig: Matplotlib fig.
+        :param name:  String. Fist part of the file name.
+        '''
+        error_handling.fig(fig)
+        error_handling.string(name, 'name')
+        name_file = self.dir_name + '/' + name + self.file_name() + '.' + self.file_format
+        fig.savefig(name_file, format=self.file_format)
+
+    def fig_lat(self, fig: Figure, name: str) -> None:
+        '''
+        Save the figure in the directory defined by the method *dir_name()*.
+
+        :param fig: Matplotlib fig.
+        :param name:  String. First part of the file name.
+        '''
+        error_handling.fig(fig)
+        error_handling.string(name, 'name')
+        name_file = self.dir_name + '/' + name + '.' + self.file_format
+        fig.savefig(name_file, format=self.file_format)
+
+    def ani(self, ani: FuncAnimation, name: str, fps: int = 10) -> None:
+        error_handling.ani(ani)
+        error_handling.string(name, 'name')
+        error_handling.positive_int(fps, 'fps')
+        name_file = self.dir_name + '/' + name + '.mp4'
+        ani.save(name_file, fps=fps, extra_args=['-vcodec', 'libx264'])
+
+
+# Backward-compatible lowercase alias (pre-0.2 API).
+save = Save