tbkit 0.2.0__py3-none-any.whl

tbkit/lattice.py

@@ -0,0 +1,424 @@
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+import tbkit.error_handling as error_handling
+
+
+PI = np.pi
+COOR_DTYPE = [('x', 'f8'), ('y', 'f8'), ('tag', 'U1')]
+
+
+#################################
+# CLASS LATTICE
+#################################
+
+
+class Lattice():
+    r'''
+    Build up 1D or 2D lattice.
+    Lattice is defined by the discrete operation:
+
+    .. math::
+
+        \mathbf{R} = n_1\mathbf{a}_1 + n_2\mathbf{a}_2
+
+    where :math:`\mathbf{a}_1` and :math:`\mathbf{a}_2` are the two primitive
+    vectors and :math:`n_1` and :math:`n_2` are the number of unit cells along
+    :math:`\mathbf{a}_1` and :math:`\mathbf{a}_2`.
+
+    :param unit_cell: List of dictionaries.
+     One dictionary per site within the unit cell. Each dictionary has two keys:
+
+        * 'tag', one-character string. Label of the associated sublattice.
+        * 'r0', Tuple. Position.
+    :param prim_vec: List of tuples.
+     Define the primitive vectors. List of one/two tuples for 1D/2D respectively:
+
+        * Tuple, cartesian coordinate of the primitive vector :math:`\mathbf{a}_1`.
+        * Tuple, cartesian coordinate of the primitive vector :math:`\mathbf{a}_2`.
+
+    Example usage::
+
+        # Line-Centered Square lattice
+        unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}]
+        prim_vec = [(0, 2), (2, 0)]
+        lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
+    '''
+
+    def __init__(self, unit_cell: list[dict], prim_vec: list[tuple[float, float]]) -> None:
+        error_handling.unit_cell(unit_cell)
+        error_handling.prim_vec(prim_vec)
+        self.unit_cell = unit_cell
+        self.prim_vec = prim_vec
+        self.tags = np.unique(np.array([dic['tag'] for dic in self.unit_cell]))
+        self.n1, self.n2 = 0, 0
+        self.coor = np.array([], dtype=COOR_DTYPE)
+        self.sites = 0
+
+    def get_lattice(self, n1: int, n2: int = 1) -> None:
+        r'''
+        Get the lattice positions.
+
+        :param n1: Positive Integer.
+            Number of unit cells along :math:`\mathbf{a}_1`.
+        :param n2: Positive Integer. Default value 1.
+            Number of unit cells along :math:`\mathbf{a}_2`.
+
+        Example usage::
+
+            # Line-Centered Square lattice
+            unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}]
+            prim_vec = [(0, 2), (2, 0)]
+            lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
+            lat.get_lattice(n1=4, n2=5)
+        '''
+        error_handling.get_lattice(self.prim_vec, n1, n2)
+        sites_uc = len(self.unit_cell)
+        sites_tag = n1*n2
+        self.sites = sites_uc * sites_tag
+        self.coor = np.empty(self.sites, dtype=COOR_DTYPE)
+        self.n1, self.n2 = n1, n2
+        x = self.prim_vec[0][0] * np.arange(n1, dtype='f8')
+        y = self.prim_vec[0][1] * np.arange(n1, dtype='f8')
+        xx = np.empty(n1*n2)
+        yy = np.empty(n1*n2)
+        xx[:n1] = x
+        yy[:n1] = y
+        for i in range(1, n2):
+            xx[i*n1: (i+1)*n1] = x + i * self.prim_vec[1][0]
+            yy[i*n1: (i+1)*n1] = y + i * self.prim_vec[1][1]
+        for i, dic in enumerate(self.unit_cell):
+            self.coor['x'][i*sites_tag: (i+1)*sites_tag] = xx + dic['r0'][0]
+            self.coor['y'][i*sites_tag: (i+1)*sites_tag] = yy + dic['r0'][1]
+            self.coor['tag'][i*sites_tag: (i+1)*sites_tag] = dic['tag']
+        self.coor = np.sort(self.coor, order=('y', 'x'))
+
+    def add_sites(self, coor: NDArray) -> None:
+        '''
+        Add sites.
+
+        :param coor: Structured array with keys: {'x', 'y', 'tag'}.
+
+        Example usage::
+
+            # Square lattice
+            unit_cell = [{'tag': 'a', 'r0': (0., 0.)}]
+            prim_vec = [(0, 1), (1, 0)]
+            lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
+            lat.get_lattice(n1=2, n2=2)
+            coor = np.array([(-1., -1, 'b'), (-2., -2, 'c')],
+                                      dtype=[('x', 'f8'), ('y', 'f8'), ('tag', 'U1')])
+            lat.add_sites(coor)
+        '''
+        error_handling.coor(coor)
+        self.coor = np.concatenate([self.coor, coor])
+        self.sites += len(coor)
+        self.tags = np.unique(np.concatenate([self.tags, coor['tag']]))
+        self.coor = np.sort(self.coor, order=('y', 'x'))
+
+    def remove_sites(self, index: list[int]) -> None:
+        '''
+        Remove sites defined by their indices
+        (use method lattice.plot(plt_index=True)
+        to get access to the site indices).
+
+        :param index: List. Site indices to be removed.
+
+        Example usage::
+
+            # Square lattice
+            unit_cell = [{'tag': 'a', 'r0': (0., 0.)}]
+            prim_vec = [(0, 1), (1, 0)]
+            lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
+            lat.get_lattice(n1=2, n2=2)
+            lat.remove_sites([0, 2])
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.remove_sites(index, self.sites)
+        mask = np.ones(self.sites, bool)
+        mask[index] = False
+        self.coor = self.coor[mask]
+        self.sites = self.coor.size
+
+    def remove_dangling(self) -> None:
+        '''
+        Remove dangling sites
+        (sites connected with just another site).
+        '''
+        error_handling.empty_coor(self.coor)
+        while True:
+            dif_x = self.coor['x'] - self.coor['x'].reshape(self.sites, 1)
+            dif_y = self.coor['y'] - self.coor['y'].reshape(self.sites, 1)
+            dis = np.sqrt(dif_x ** 2 + dif_y ** 2)
+            dis_unique = np.unique(dis)
+            len_hop = dis_unique[1]
+            ind = np.argwhere(np.isclose(dis, len_hop))
+            dang = []
+            for i in range(self.sites):
+                if (ind[:, 0] == i).sum() == 1:
+                    dang.append(i)
+            self.coor = np.delete(self.coor, dang, axis=0)
+            self.sites -= len(dang)
+            if dang == []:
+                break
+
+    def shift_x(self, shift: float) -> None:
+        '''
+        Shift the x coordinates.
+
+        :param shift: Real number. Shift value.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.real_number(shift, 'shift')
+        self.coor['x'] += shift
+
+    def shift_y(self, shift: float) -> None:
+        '''
+        Shift by *delta_x* the x coordinates.
+
+        :param shift: Real number. Shift value.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.real_number(shift, 'shift')
+        self.coor['y'] += shift
+
+    def change_sign_x(self) -> None:
+        '''
+        Change x coordinates sign.
+        '''
+        error_handling.empty_coor(self.coor)
+        self.coor['x'] *= -1
+
+    def change_sign_y(self) -> None:
+        '''
+        Change y coordinates sign.
+        '''
+        error_handling.empty_coor(self.coor)
+        self.coor['y'] *= -1
+
+    def boundary_line(self, cx: float, cy: float, co: float) -> None:
+        r'''
+        Select sites according to :math:`c_yy+c_xx > c_0`.
+
+        :param cx: Real number. cx value.
+        :param cy: Real number. cy value.
+        :param co: Real number. co value.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.real_number(cx, 'cx')
+        error_handling.real_number(cy, 'cy')
+        error_handling.real_number(co, 'co')
+        self.coor = self.coor[cy * self.coor['y'] + cx * self.coor['x'] > co]
+        self.sites = len(self.coor)
+
+    def ellipse_in(self, rx: float, ry: float, x0: float, y0: float) -> None:
+        r'''
+        Select sites according to
+
+        .. math::
+
+            (x-x_0)^2/a^2+(y-y_0)^2/b^2 < 1\,  .
+
+        :param list_hop: List of Dictionary (see set_hopping definition).
+        :param rx: Positive Real number. Radius along :math:`x`.
+        :param ry: Positive Real number. Radius along :math:`y`.
+        :param x0: Real number. :math:`x` center.
+        :param y0: Real number. :math:`y` center.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.positive_real(rx, 'rx')
+        error_handling.positive_real(ry, 'ry')
+        error_handling.real_number(x0, 'x0')
+        error_handling.real_number(y0, 'y0')
+        self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 +  \
+                                        (self.coor['y'] -y0) ** 2 / ry ** 2 < 1.]
+        self.sites = len(self.coor)
+
+    def ellipse_out(self, rx: float, ry: float, x0: float, y0: float) -> None:
+        r'''
+        Select sites according to
+
+        .. math::
+
+            (x-x_0)^2/a^2+(y-y_0)^2/b^2 > 1\,  .
+
+
+        :param list_hop: List of Dictionary (see set_hopping definition).
+        :param rx: Positive Real number. Radius along :math:`x`.
+        :param ry: Positive Real number. Radius along :math:`y`.
+        :param x0: Real number. :math:`x` center.
+        :param y0: Real number. :math:`y` center.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.positive_real(rx, 'rx')
+        error_handling.positive_real(ry, 'ry')
+        error_handling.real_number(x0, 'x0')
+        error_handling.real_number(y0, 'y0')
+        self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 +  \
+                                        (self.coor['y'] -y0) ** 2 / ry ** 2 > 1.]
+        self.sites = len(self.coor)
+
+    def center(self) -> None:
+        '''
+        Fix the center of mass of the lattice at (0, 0).
+        '''
+        error_handling.empty_coor(self.coor)
+        self.coor['x'] -= np.mean(self.coor['x'])
+        self.coor['y'] -= np.mean(self.coor['y'])
+
+    def rotation(self, theta: float) -> None:
+        r'''
+        Rotate the lattice structure by the angle :math:`\theta`.
+
+        :param theta: Rotation angle in degrees.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.real_number(theta, 'theta')
+        theta *= PI / 360
+        for dic in self.unit_cell:
+            x  = self.coor['x'] - dic['r0'][0]
+            y  = self.coor['y'] - dic['r0'][1]
+            self.coor['x'] = x * np.cos(theta) - y * np.sin(theta) + dic['r0'][0]
+            self.coor['y'] = y * np.cos(theta) + x* np.sin(theta) + dic['r0'][1]
+
+    def clean_coor(self) -> None:
+        '''
+        Keep only the sites with different coordinates.
+        '''
+        error_handling.empty_coor(self.coor)
+        coor = self.coor[['x', 'y']].copy()
+        coor['x'], coor['y'] = self.coor['x'].round(4), self.coor['y'].round(4)
+        _, idx = np.unique(coor, return_index=True)
+        self.coor = self.coor[idx]
+        self.sites = len(self.coor)
+
+    def __add__(self, other: 'Lattice') -> 'Lattice':
+        '''
+        Overloading operator +.
+        '''
+        error_handling.lat(other)
+        error_handling.empty_coor(self.coor)
+        error_handling.empty_coor(other.coor)
+        coor = np.concatenate([self.coor, other.coor])
+        tags = np.concatenate([self.tags, other.tags])
+        lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec)
+        lat.add_sites(coor)
+        lat.sites = self.sites + other.sites
+        lat.tags = np.unique(tags)
+        return lat
+
+    def __iadd__(self, other: 'Lattice') -> 'Lattice':
+        '''
+        Overloading operator +=.
+        '''
+        error_handling.lat(other)
+        error_handling.empty_coor(self.coor)
+        error_handling.empty_coor(other.coor)
+        self.coor = np.concatenate([self.coor, other.coor])
+        self.sites += other.sites
+        self.tags = np.unique(np.concatenate([self.tags, other.tags]))
+        return self
+
+    def __sub__(self, other: 'Lattice') -> 'Lattice':
+        '''
+        Overloading operator -.
+
+        .. note::
+
+            The tags are not considered in the lattice subtraction.
+        '''
+        error_handling.lat(other)
+        error_handling.empty_coor(self.coor)
+        error_handling.empty_coor(other.coor)
+        boo = np.zeros(self.sites, bool)
+        for c in other.coor:
+            boo += np.isclose(c['x'], self.coor['x']) & np.isclose(c['y'], self.coor['y'])
+        coor = self.coor[np.logical_not(boo)]
+        lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec)
+        lat.add_sites(coor)
+        lat.sites = len(lat.coor)
+        lat.tags = self.tags
+        return lat
+
+    def __isub__(self, other: 'Lattice') -> 'Lattice':
+        '''
+        Overloading operator -=.
+
+        .. note::
+
+            The tags are not considered in the lattice subtraction.
+        '''
+        error_handling.lat(other)
+        error_handling.empty_coor(self.coor)
+        error_handling.empty_coor(other.coor)
+        boo = np.zeros(self.sites, bool)
+        for c in other.coor:
+            boo += np.isclose(c['x'], self.coor['x']) & np.isclose(c['y'], self.coor['y'])
+        self.coor = self.coor[np.logical_not(boo)]
+        self.sites = int(np.sum(np.logical_not(boo)))
+        return self
+
+    def plot(
+        self,
+        ms: float = 20,
+        fs: float = 20,
+        plt_index: bool = False,
+        axis: bool = False,
+        figsize: tuple[float, float] | None = None,
+    ) -> Figure:
+        '''
+        Plot lattice in hopping space.
+
+        :param ms: Positive number. Default value 20. Markersize.
+        :param fs: Positve number. Default value 20. Fontsize.
+        :param plt_index: Boolean. Default value False. Plot site labels.
+        :param axis: Boolean. Default value False. Plot axis.
+        :param figsize: Tuple. Default value None. Figsize.
+
+        :returns:
+            * **fig** -- Figure.
+        '''
+        error_handling.empty_coor(self.coor)
+        error_handling.positive_real(ms, 'ms')
+        error_handling.positive_real(fs, 'fs')
+        error_handling.boolean(plt_index, 'plt_index')
+        error_handling.boolean(axis, 'axis')
+        if figsize is None:
+            figsize = (5, 5)
+        error_handling.list_tuple_2elem(figsize, 'figsize')
+        error_handling.positive_real(figsize[0], 'figsize[0]')
+        error_handling.positive_real(figsize[1], 'figsize[1]')
+        fig, ax = plt.subplots(figsize=figsize)
+        # plot sites
+        colors = ['b', 'r', 'g', 'y', 'm', 'k']
+        for color, tag in zip(colors, self.tags):
+            plt.plot(self.coor['x'][self.coor['tag'] == tag],
+                        self.coor['y'][self.coor['tag'] == tag],
+                       'o', color=color, ms=ms, markeredgecolor='none')
+        ax.set_aspect('equal')
+        ax.set_xlim([np.min(self.coor['x'])-1., np.max(self.coor['x'])+1.])
+        ax.set_ylim([np.min(self.coor['y'])-1., np.max(self.coor['y'])+1.])
+        if not axis:
+            ax.axis('off')
+        # plot indices
+        if plt_index:
+            indices = ['{}'.format(i) for i in range(self.sites)]
+            for l, x, y in zip(indices, self.coor['x'], self.coor['y']):
+                plt.annotate(l, xy=(x, y), xytext=(0, 0),
+                                    textcoords='offset points',
+                                    ha='right', va='bottom', size=fs)
+        plt.draw()
+        return fig
+
+    def show(self) -> None:
+        """
+        Emulate Matplotlib method plt.show().
+        """
+        plt.show()
+
+
+# Backward-compatible lowercase alias (pre-0.2 API).
+lattice = Lattice

tbkit/lattices.py

@@ -0,0 +1,126 @@
+"""
+A small library of common 2D Bravais lattices, ready to feed into
+:class:`tbkit.lattice.Lattice`, :class:`tbkit.system.System`, or
+:class:`tbkit.kspace.KSpace`.
+
+Each function returns a fresh ``Lattice`` instance with *unit_cell* and
+*prim_vec* already set (call ``get_lattice`` yourself to build a finite
+flake, or hand it straight to ``KSpace`` for a periodic/band-structure
+calculation). Nearest-neighbor sites are a distance *a* apart.
+
+Example usage::
+
+    import tbkit.lattices as lattices
+    lat = lattices.kagome()
+    lat.get_lattice(n1=6, n2=6)
+"""
+from __future__ import annotations
+
+from math import sqrt
+
+from tbkit.lattice import Lattice
+import tbkit.error_handling as error_handling
+
+
+def _lat(unit_cell: list[dict], prim_vec: list[tuple[float, float]]) -> Lattice:
+    return Lattice(unit_cell=unit_cell, prim_vec=prim_vec)
+
+
+def chain(a: float = 1.) -> Lattice:
+    '''
+    1D chain: one site per unit cell.
+
+    :param a: Positive real number. Default value 1. Lattice constant.
+    '''
+    error_handling.positive_real(a, 'a')
+    return _lat([{'tag': 'a', 'r0': (0., 0.)}], [(a, 0.)])
+
+
+def square(a: float = 1.) -> Lattice:
+    '''
+    Square lattice: one site per unit cell.
+
+    :param a: Positive real number. Default value 1. Lattice constant.
+    '''
+    error_handling.positive_real(a, 'a')
+    return _lat([{'tag': 'a', 'r0': (0., 0.)}], [(a, 0.), (0., a)])
+
+
+def triangular(a: float = 1.) -> Lattice:
+    '''
+    Triangular lattice: one site per unit cell.
+
+    :param a: Positive real number. Default value 1. Lattice constant.
+    '''
+    error_handling.positive_real(a, 'a')
+    return _lat([{'tag': 'a', 'r0': (0., 0.)}],
+                       [(a, 0.), (0.5*a, 0.5*sqrt(3)*a)])
+
+
+def honeycomb(a: float = 1.) -> Lattice:
+    '''
+    Honeycomb lattice (e.g. graphene): two sites per unit cell, nearest
+    neighbors a distance *a* apart. See also :class:`tbkit.graphene.GrapheneLattice`
+    for ready-made finite flakes of various shapes.
+
+    :param a: Positive real number. Default value 1. Nearest-neighbor distance.
+    '''
+    error_handling.positive_real(a, 'a')
+    dx, dy = 0.5*sqrt(3)*a, 0.5*a
+    unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'b', 'r0': (dx, dy)}]
+    prim_vec = [(2*dx, 0.), (dx, 1.5*a)]
+    return _lat(unit_cell, prim_vec)
+
+
+def kagome(a: float = 1.) -> Lattice:
+    '''
+    Kagome lattice: three sites per unit cell (tags 'a', 'b', 'c'), arranged
+    as corner-sharing triangles on a triangular Bravais lattice. With
+    uniform nearest-neighbor hopping, this lattice famously has an exactly
+    flat band (at E = -2t for hopping amplitude t).
+
+    :param a: Positive real number. Default value 1. Nearest-neighbor distance.
+
+    Example usage (nearest-neighbor hopping, for :class:`tbkit.kspace.KSpace`)::
+
+        lat = lattices.kagome()
+        kag = KSpace(lat)
+        kag.set_hopping([{'i': 0, 'j': 1, 'R': (0, 0), 't': t},
+                                {'i': 0, 'j': 1, 'R': (-1, 0), 't': t},
+                                {'i': 0, 'j': 2, 'R': (0, 0), 't': t},
+                                {'i': 0, 'j': 2, 'R': (0, -1), 't': t},
+                                {'i': 1, 'j': 2, 'R': (0, 0), 't': t},
+                                {'i': 1, 'j': 2, 'R': (1, -1), 't': t}])
+    '''
+    error_handling.positive_real(a, 'a')
+    a1 = (2*a, 0.)
+    a2 = (a, sqrt(3)*a)
+    unit_cell = [{'tag': 'a', 'r0': (0., 0.)},
+                       {'tag': 'b', 'r0': (a, 0.)},
+                       {'tag': 'c', 'r0': (0.5*a, 0.5*sqrt(3)*a)}]
+    return _lat(unit_cell, [a1, a2])
+
+
+def lieb(a: float = 1.) -> Lattice:
+    '''
+    Lieb lattice: three sites per unit cell (tag 'a': corner site; tags
+    'b', 'c': edge-center sites) on a square Bravais lattice. With uniform
+    nearest-neighbor hopping, this lattice famously has an exactly flat
+    band (at E = 0), squeezed between two dispersive bands.
+
+    :param a: Positive real number. Default value 1. Nearest-neighbor distance.
+
+    Example usage (nearest-neighbor hopping, for :class:`tbkit.kspace.KSpace`)::
+
+        lat = lattices.lieb()
+        lb = KSpace(lat)
+        lb.set_hopping([{'i': 0, 'j': 1, 'R': (0, 0), 't': t},
+                               {'i': 0, 'j': 1, 'R': (-1, 0), 't': t},
+                               {'i': 0, 'j': 2, 'R': (0, 0), 't': t},
+                               {'i': 0, 'j': 2, 'R': (0, -1), 't': t}])
+    '''
+    error_handling.positive_real(a, 'a')
+    unit_cell = [{'tag': 'a', 'r0': (0., 0.)},
+                       {'tag': 'b', 'r0': (a, 0.)},
+                       {'tag': 'c', 'r0': (0., a)}]
+    return _lat(unit_cell, [(2*a, 0.), (0., 2*a)])