import torch
import numpy as np
from typing import Union
from typing import Tuple, Union, Optional, List, Set, Dict, Any
def _nice_float(x,just,rnd):
return str(round(x,rnd)).rjust(just)
[docs]
class kpoints_generator:
"""
Used to generate K point path
"""
def __init__(self, dim_k: int=3, lat: Union[np.array, list]=None, per: Union[List, Tuple] = None):
self._dim_k = dim_k
self._lat = lat
# choose which self._dim_k out of self._dim_r dimensions are
# to be considered periodic.
if per==None:
# by default first _dim_k dimensions are periodic
self._per=list(range(self._dim_k))
else:
if len(per)!=self._dim_k:
raise Exception("\n\nWrong choice of periodic/infinite direction!")
# store which directions are the periodic ones
self._per=per
[docs]
def k_path(self,kpts,nk,report=False):
r"""
Interpolates a path in reciprocal space between specified
k-points. In 2D or 3D the k-path can consist of several
straight segments connecting high-symmetry points ("nodes"),
and the results can be used to plot the bands along this path.
The interpolated path that is returned contains as
equidistant k-points as possible.
:param kpts: Array of k-vectors in reciprocal space between
which interpolated path should be constructed. These
k-vectors must be given in reduced coordinates. As a
special case, in 1D k-space kpts may be a string:
* *"full"* -- Implies *[ 0.0, 0.5, 1.0]* (full BZ)
* *"fullc"* -- Implies *[-0.5, 0.0, 0.5]* (full BZ, centered)
* *"half"* -- Implies *[ 0.0, 0.5]* (half BZ)
:param nk: Total number of k-points to be used in making the plot.
:param report: Optional parameter specifying whether printout
is desired (default is True).
:returns:
* **k_vec** -- Array of (nearly) equidistant interpolated
k-points. The distance between the points is calculated in
the Cartesian frame, however coordinates themselves are
given in dimensionless reduced coordinates! This is done
so that this array can be directly passed to function
:func:`pythtb.tb_model.solve_all`.
* **k_dist** -- Array giving accumulated k-distance to each
k-point in the path. Unlike array *k_vec* this one has
dimensions! (Units are defined here so that for an
one-dimensional crystal with lattice constant equal to for
example *10* the length of the Brillouin zone would equal
*1/10=0.1*. In other words factors of :math:`2\pi` are
absorbed into *k*.) This array can be used to plot path in
the k-space so that the distances between the k-points in
the plot are exact.
* **k_node** -- Array giving accumulated k-distance to each
node on the path in Cartesian coordinates. This array is
typically used to plot nodes (typically special points) on
the path in k-space.
Example usage::
# Construct a path connecting four nodal points in k-space
# Path will contain 401 k-points, roughly equally spaced
path = [[0.0, 0.0], [0.0, 0.5], [0.5, 0.5], [0.0, 0.0]]
(k_vec,k_dist,k_node) = my_model.k_path(path,401)
# solve for eigenvalues on that path
evals = tb.solve_all(k_vec)
# then use evals, k_dist, and k_node to plot bandstructure
# (see examples)
"""
# processing of special cases for kpts
if kpts=='full':
# full Brillouin zone for 1D case
k_list=np.array([[0.],[0.5],[1.]])
elif kpts=='fullc':
# centered full Brillouin zone for 1D case
k_list=np.array([[-0.5],[0.],[0.5]])
elif kpts=='half':
# half Brillouin zone for 1D case
k_list=np.array([[0.],[0.5]])
else:
k_list=np.array(kpts)
# in 1D case if path is specified as a vector, convert it to an (n,1) array
if len(k_list.shape)==1 and self._dim_k==1:
k_list=np.array([k_list]).T
# make sure that k-points in the path have correct dimension
if k_list.shape[1]!=self._dim_k:
print('input k-space dimension is',k_list.shape[1])
print('k-space dimension taken from model is',self._dim_k)
raise Exception("\n\nk-space dimensions do not match")
# must have more k-points in the path than number of nodes
if nk<k_list.shape[0]:
raise Exception("\n\nMust have more points in the path than number of nodes.")
# number of nodes
n_nodes=k_list.shape[0]
# extract the lattice vectors from the TB model
lat_per=np.copy(self._lat)
# choose only those that correspond to periodic directions
lat_per=lat_per[self._per]
# compute k_space metric tensor
k_metric = np.linalg.inv(np.dot(lat_per,lat_per.T))
# Find distances between nodes and set k_node, which is
# accumulated distance since the start of the path
# initialize array k_node
k_node=np.zeros(n_nodes,dtype=float)
for n in range(1,n_nodes):
dk = k_list[n]-k_list[n-1]
dklen = np.sqrt(np.dot(dk,np.dot(k_metric,dk)))
k_node[n]=k_node[n-1]+dklen
# Find indices of nodes in interpolated list
node_index=[0]
for n in range(1,n_nodes-1):
frac=k_node[n]/k_node[-1]
node_index.append(int(round(frac*(nk-1))))
node_index.append(nk-1)
# initialize two arrays temporarily with zeros
# array giving accumulated k-distance to each k-point
k_dist=np.zeros(nk,dtype=float)
# array listing the interpolated k-points
k_vec=np.zeros((nk,self._dim_k),dtype=float)
# go over all kpoints
k_vec[0]=k_list[0]
for n in range(1,n_nodes):
n_i=node_index[n-1]
n_f=node_index[n]
kd_i=k_node[n-1]
kd_f=k_node[n]
k_i=k_list[n-1]
k_f=k_list[n]
for j in range(n_i,n_f+1):
frac=float(j-n_i)/float(n_f-n_i)
k_dist[j]=kd_i+frac*(kd_f-kd_i)
k_vec[j]=k_i+frac*(k_f-k_i)
np.set_printoptions(precision=5)
if (lat_per.shape[0]==lat_per.shape[1]):
# lat_per is invertible
lat_per_inv=np.linalg.inv(lat_per).T
return (k_vec,k_dist,k_node,lat_per_inv)
[docs]
def compute_reciprocal_lattice_vectors(lattice_vectors):
"""
Calculates the reciprocal lattice vectors for a given set of lattice vectors using PyTorch.
Parameters:
lattice_vectors (torch.Tensor): A 2D tensor containing three lattice vectors
as rows.
Returns:
torch.Tensor: A 2D tensor containing the reciprocal lattice vectors as rows.
"""
# Extract lattice vectors
a1, a2, a3 = lattice_vectors
# Calculate the volume (V) of the unit cell
volume = torch.dot(a1, torch.cross(a2, a3))
# Compute the reciprocal lattice vectors using the volume
b1 = (2 * torch.pi * torch.cross(a2, a3)) / volume
b2 = (2 * torch.pi * torch.cross(a3, a1)) / volume
b3 = (2 * torch.pi * torch.cross(a1, a2)) / volume
# Stack the vectors into a 2D tensor
return torch.stack([b1, b2, b3], dim=0)
[docs]
def generate_kpoint_grid(lattice_vectors, kpoint_density=0.04):
"""
Generate a k-point grid with density proportional to reciprocal lattice vectors.
Parameters:
lattice_vectors (torch.Tensor): A 2D tensor with shape (3,3) containing lattice vectors as rows
kpoint_density (float): Target k-point density in Å^-1 (typical values: 0.03-0.05)
Returns:
torch.Tensor: A 1D tensor containing the number of k-points along each reciprocal lattice vector
"""
# Calculate reciprocal lattice vectors
recip_lattice = compute_reciprocal_lattice_vectors(lattice_vectors)
# Calculate lengths of reciprocal lattice vectors
recip_lengths = torch.norm(recip_lattice, dim=1)
# Calculate number of k-points along each direction
# Formula: Ni = max(1, ceiling(|bi|/(2π*kpoint_density)))
kpoint_numbers = torch.maximum(
torch.ones(3, dtype=torch.int64),
torch.ceil(recip_lengths / (2 * torch.pi * kpoint_density)).to(torch.int64)
)
return kpoint_numbers