Yambo Community Forum

Code: Select all

#Authors: MN 
"""
Script to compute total crystal angular momentum
This script loads Yambo databases, calculates the angular momentum for a specific exciton,
and generates a reciprocal-space (K-space) visualization.
"""

import numpy as np
import os
import matplotlib.pyplot as plt
from yambopy.dbs.excitondb import YamboExcitonDB
from yambopy.dbs.latticedb import YamboLatticeDB
from yambopy.dbs.wfdb import YamboWFDB

## Inputs
iqpt = 1
path = '.'
gw_bse_dir = 'GW_BSE'
iexe = 2
degen_tol = 1e-4

# ======= Load data ========

# Load the lattice db
lattice = YamboLatticeDB.from_db_file(os.path.join(path, 'SAVE', 'ns.db1'))

# Load exciton db
filename = 'ndb.BS_diago_Q%d' % (iqpt)
excdb = YamboExcitonDB.from_db_file(lattice, filename=filename,
                                    folder=os.path.join(path, gw_bse_dir), neigs = iexe + 101)

# Load the wavefunction database
wfdb = YamboWFDB(path=path, latdb=lattice, bands_range=[np.min(excdb.table[:, 1]) - 1, np.max(excdb.table[:, 2])])

# Symmetry definitions
symm_mat_cart = lattice.sym_car[3]
frac_trans_cart = np.zeros(3)

# Compute Total Crystal Angular Momentum
sbasis = excdb.total_crys_angular_momentum(wfdb, iexe, symm_mat_cart, frac_trans_cart, degen_tol=degen_tol)

jvals = sbasis[0]
Akcv_j = sbasis[1]

print('Total crystal angular momentum : ', jvals)

# Find degenerate states (0-based indexing for internal arrays)
idegen = np.array(excdb.get_degenerate(iexe + 1, eps=degen_tol), dtype=int) - 1

# Update the existing wfs with the new rotated simultaneous eigenbasis
old_eig_states = excdb.Akcv[idegen].copy()
excdb.Akcv[idegen] = Akcv_j

# ======= Plotting in K-Space ========

print("="*80)
for i in range(len(jvals)):
    state_idx = idegen[i] + 1
    j_val = jvals[i]

    print('*********** Plotting K-space wf for state %d with j = %.2f ***************' % (state_idx, j_val))
    print("="*80)

    fig = plt.figure(figsize=(6,6))
    ax  = fig.add_axes([0.15, 0.15, 0.80, 0.80])

    excdb.plot_exciton_2D_ax(ax, [state_idx], mode='hexagon', limfactor=0.8, scale=180)

    ax.set_title(f"Exciton State {state_idx} | j = {j_val:.2f}")

    plt.savefig(f"Kspace_state_{state_idx}_j_{j_val:.2f}.png", dpi=300, bbox_inches='tight')
    plt.show()

print("="*80)

# Restore the original database coefficients to leave everything as it was
excdb.Akcv[idegen] = old_eig_states