(wf_array_v2) =

wf_array to WFArray in v2.0

This tutorial demonstrates the new WFArray class in PythTB v2.0, which is an extension to the original wf_array class in previous versions. For a more comprehensive overview of all changes, please refer to the release notes.

The first thing to be aware of is that the class names have changed

• < v1.8: tb_model, wf_array, w90

• 2.0: TBModel, WFArray, W90

from pythtb import Lattice, Mesh, TBModel, WFArray
import numpy as np

If you haven’t already, we recommend you read through the TBModel v2.0 tutorial and Mesh tutorial first, as this tutorial will build on that knowledge.

The WFArray class in v2.0 is aimed at addressing a few limitations of the previous wf_array class:

• Mesh awareness: The new WFArray class is designed to be aware of the k-point mesh structure, allowing for more efficient storage and manipulation of wavefunction data on combined \((k, \lambda)\)-meshes.

• Faster Calculations: The new class includes optimized methods for common operations such as overlaps, projections, etc., which are faster than the previous implementation.

• Flexibility: The WFArray class is decoupled from the TBModel class, allowing for more flexible usage.

Let’s look at how to use the new WFArray class in practice.

Initializing and populating a WFArray

In previous versions of PythTB, the wf_array class was tightly coupled to the tb_model class. Upon initializing a wf_array, wf_array would read off the lattice information from the associated tb_model. In v2.0, like in TBModel, we will instead explicitly pass a Lattice object to the WFArray constructor. Aside from increased flexibility, this also allows us to create WFArray objects that are not tied to any specific TBModel. For instance, one may want to store states that are derived from TBModels on an adiabatic cycle, which there is no single TBModel that describes all the states.

Compare to versions 1.x:

# v1.x
from pythtb import tb_model, wf_array
model = tb_model(1, 1, lat=[[1.0, 0.0], [0.0, 1.0]], orb=[[0, 1/3, 2/3]])
wfa = wf_array(model, [20])
wfa.solve_on_grid(0.0)

In solve_on_grid, the wf_array will generate the k-point mesh and populate itself with the wavefunctions from model. Afterwords the mesh is discarded and the wf_array only retains the wavefunction data. This loses the information about the k-point structure, and requires the user to manually keep track of the k-point mesh if needed.

In v2.0, one itializes a WFArray as follows:

1. Create a Lattice object that describes the lattice structure of the system.

lattice = Lattice(
    lat_vecs=[[1.0]], orb_vecs=[[0.0], [1 / 3], [2 / 3]], periodic_dirs=[0]
)

2. Create a Mesh object that describes the k-point and parameter mesh.

mesh = Mesh(["k"])

3. Populate the Mesh with the desired k-points and parameter values.

mesh.build_grid([20])

4. Pass these objects to the WFArray constructor.

wfa = WFArray(lattice=lattice, mesh=mesh)

Now if we want to populate the WFArray with energy eigenstates of a given TBModel, we can do so by calling the WFArray.solve_model() function.

t = -1.3
delta = 2.0

model = TBModel(lattice=lattice)

# nearest-neighbour hoppings (last hop wraps to the next cell)
model.set_hop(t, 0, 1, [0])
model.set_hop(t, 1, 2, [0])
model.set_hop(t, 2, 0, [1])

# Per-orbital onsite as lambdas of lmbda
onsite = [
    lambda lmbda: delta * -np.cos(2 * np.pi * (lmbda - 0 / 3)),
    lambda lmbda: delta * -np.cos(2 * np.pi * (lmbda - 1 / 3)),
    lambda lmbda: delta * -np.cos(2 * np.pi * (lmbda - 2 / 3)),
]

model.set_onsite(onsite)

model_fixed = model.with_parameters(lmbda=0.25)
wfa.solve_model(model_fixed)

We can access the individual wavefunctions stored in the WFArray using indexing, similar to how it was done in previous versions. For example, wfa[k_index] will return the wavefunctions at the specified mesh index.

wfa[2]

array([[ 0.48119032+0.j        ,  0.80505946-0.06370792j,
         0.33915892-0.0353223j ],
       [ 0.77558258+0.j        , -0.54837962-0.05902583j,
         0.21322754+0.22091043j],
       [-0.40856766+0.j        , -0.09282883-0.1870805j ,
         0.80421332+0.3777527j ]])

As before, internally periodic boundary conditions are automatically applied across the BZ boundary in Berry phase calculations and other operations.

WFArray on an adiabatic cycle

We will reuse the same lattice as above, except now we will create a Mesh that includes an adiabatic parameter axis lambda, which varies from \(0\) to \(1\) in 11 steps. This parameter could represent, for example, a distortion of the lattice or a change in an external field.

Recall from the Mesh tutorial that we need to name the \(\lambda\) axis according to the parameter name we are varying in the TBModel function (see function above if you missed it).

mesh = Mesh(["k", "l"], axis_names=["kx", "lmbda"])

Now we will build the grid, specifying the range of the adiabatic parameter from \(0\) to \(1\) in 11 steps.

mesh.build_grid(
    shape=(31, 21),
    gamma_centered=True,
    lambda_start=0.0,
    lambda_stop=1.0,
    lambda_endpoints=True,
)

If we want to impose the \(\lambda\) axis as an adiabatic loop, we need to set loop for the axis index corresponding to \(\lambda\). The component index is 1 (Python counting) since the combined \((k, \lambda)\)-space is two dimensional, and the \(\lambda\) dimensions come after the k-dimensions (of which there are 1). This indicates that traversing axis 1 winds the second (1 in Python counting) component of the mesh vector in a loop, such that the Hamiltonian at the end of the loop connects back to the Hamiltonian at the beginning of the loop.

We also mark the axis as closed, since the endpoint at \(\lambda=1\) is equivalent to the starting point at \(\lambda=0\). We can do this by setting the closed parameter to True.

mesh.loop(
    axis_idx=1, component_idx=1, closed=True
)  # form the lambda axis into a closed loop
print(mesh)

Mesh Summary
========================================
Type: grid
Dimensionality: 1 k-dim(s) + 1 λ-dim(s)
Number of mesh points: 651
Full shape: (31, 21, 2)
k-axes: [Axis(type=k, name=kx, size=31)]
λ-axes: [Axis(type=l, name=lmbda, size=21)]
Is a torus in k-space (all k-axes wind BZ): yes
Loops: (axis 0, comp 0, winds_bz=yes, closed=no), (axis 1, comp 1, winds_bz=no, closed=yes)

We initialize the WFArray

wfa = WFArray(lattice, mesh)

Finally we call the solve_model function to populate the WFArray with the eigenstates of the TBModel at each point in the combined \((k, \lambda)\)-mesh.

wfa.solve_model(model=model)

We can verify that the \(\lambda\) axis was correctly imposed as a loop by comparing the wavefunctions at the start and end of the adiabatic cycle.

np.allclose(wfa[0, 0], wfa[0, -1])  # first k-point, first and last lambda points

True

Compare to versions 1.x:

# v1.x

path_steps = 11
all_lambda = np.linspace(0, 1, path_steps, endpoint=True)
num_kpt = 31

wf_kpt_lambda = wf_array(my_model,[num_kpt, path_steps])
for i_lambda in range(path_steps):
    lmbd = all_lambda[i_lambda]
    model = set_model(t, delta, lmbd)

    (k_vec, k_dist, k_node) = my_model.k_path([[0],[1]],num_kpt, report=False)
    (eval, evec) = my_model.solve_all(k_vec,eig_vectors=True)
    for i_kpt in range(num_kpt):
        # Manually populate the wf_array
        wf_kpt_lambda[i_kpt, i_lambda]=evec[:,i_kpt,:]

# Need to manually impose PBCs since the wf_array does not know about the mesh structure
wf_kpt_lambda.impose_pbc(0, 0)
# Need to manually impose the loop in lambda
wf_kpt_lambda.impose_loop(1, 1)

New methods in WFArray v2.0

Explore the variety of new methods available in the WFArray class in v2.0 by checking out the API documentation