Chern number and Berry curvature
================================

In this tutorial, we demonstrate the usage of :class:`.Berry` class by calculating the Chern number
and Berry curvature of Haldane model. The corresponding script is located at ``examples/prim_cell/berry.py``.
We begin with importing the necessary packages:

.. code-block:: python

    import numpy as np

    import tbplas as tb

and define the function for generating the Haldane model:

.. code-block:: python
    :linenos:

    def make_haldane() -> tb.PrimitiveCell:
        """
        Make Haldane model for testing Berry utilities.
        Reference: examples/haldane.py of pythtb
        """
        lat = 1.0
        delta = 0.2
        t = -1.0
        t2 = 0.15 * np.exp(1j * np.pi / 2.)
        t2c = t2.conjugate()

        vectors = tb.gen_lattice_vectors(a=lat, b=lat, c=1.0, gamma=60)
        cell = tb.PrimitiveCell(vectors, unit=tb.NM)
        cell.add_orbital((1 / 3., 1 / 3.), energy=-delta)
        cell.add_orbital((2 / 3., 2 / 3.), energy=delta)
        cell.add_hopping((0, 0), 0, 1, t)
        cell.add_hopping((1, 0), 1, 0, t)
        cell.add_hopping((0, 1), 1, 0, t)
        cell.add_hopping((1, 0), 0, 0, t2)
        cell.add_hopping((1, -1), 1, 1, t2)
        cell.add_hopping((0, 1), 1, 1, t2)
        cell.add_hopping((1, 0), 1, 1, t2c)
        cell.add_hopping((1, -1), 0, 0, t2c)
        cell.add_hopping((0, 1), 0, 0, t2c)
        return cell

Chern number
------------

We define the following function to demonstrate the calculation of Chern number:

.. code-block:: python
    :linenos:

    def calc_chern_number() -> None:
        """Calculate Chern number of Haldane model."""
        # Create model and solver.
        cell = make_haldane()
        solver = tb.Berry(cell)

        # Set up parameters
        solver.config.k_grid_size = (100, 100, 1)
        solver.config.num_occ = 1

        # Calculate Berry curvature and Chern number using Kubo formula
        solver.config.prefix = "chern_kubo"
        solver.calc_berry_curvature_kubo()

        # Calculate Berry curvature and Chern number using Wilson loop
        solver.config.prefix = "chern_wilson"
        solver.calc_berry_curvature_wilson()

Just as in other tutorials, we firstly create the model and the solver, then set up the calculation
parameters. We need to sample the first Brillouin zone (FBZ) with k-grid specified by ``k_grid_size``,
and the number of occupied states by ``num_occ``. Then we call ``calc_berry_curvature_kubo`` and
``calc_berry_curvature_wilson`` to calculate the Berry curvature and Chern number using Kubo formula
and Wilson loop methods, respectively. The output should look like:

.. code-block:: text
    :linenos:
    
    ... ...
    Chern number for band 0: -1
    Chern number for band 1: 1

    ... ...
    Chern number for num_occ 1: -1

indicating the Chern number for the two bands is -1 and 1, respectively. The total Chern number is -1,
since there is only one occupied band.

Berry curvature
---------------

The ``calc_berry_curvature_kubo`` and ``calc_berry_curvature_wilson`` methods already output the
Berry curvatures in FBZ. For better appearance, we need to extend k-grid to second Brillouin zone.
This is done by the ``bz_size`` parameter, as demonstrated in the following function:

.. code-block:: python
    :linenos:
    
    def calc_berry_curvature() -> None:
        """Calculate and visualize Berry curvature of Haldane model."""
        # Create model and solver.
        cell = make_haldane()
        solver = tb.Berry(cell)

        # Set up parameters
        solver.config.k_grid_size = (100, 100, 1)
        solver.config.bz_size = (2, 2, 1)
        solver.config.num_occ = 1

        # Calculate Berry curvature and Chern number using Kubo formula
        solver.config.prefix = "berry_kubo"
        data_kubo = solver.calc_berry_curvature_kubo()

        # Calculate Berry curvature and Chern number using Wilson loop
        solver.config.prefix = "berry_wilson"
        data_wilson = solver.calc_berry_curvature_wilson()

        # Visualization
        if solver.is_master:
            plot_omega_xy(data_kubo, ib=0)
            plot_omega_xy(data_kubo, ib=1)
            plot_omega_xy(data_wilson)

The functions for visualizing Berry curvature is defined as:

.. code-block:: python
    :linenos:

    def plot_omega_xy(berry_data: tb.BerryData, ib: int = 0) -> None:
        """
        Visualize Berry curvature.
        :param berry_data: Berry curvature to plot
        :param ib: band index of Berry curvature to plot
        :return: None
        """
        kpt_cart = berry_data.kpt_cart
        omega_xy = berry_data.omega_xy[:, ib]
        vis = tb.Visualizer()
        vis.plot_scalar(x=kpt_cart[:, 0], y=kpt_cart[:, 1], z=omega_xy, scatter=True,
                        num_grid=(480, 480), cmap="jet", with_colorbar=True)

The output is shown the figure below:

.. figure:: images/berry/berry.png
    :align: center
    :scale: 30%

    Berry curvatures from Kubo formula and Wilson loop methods.