Yambo-py - new script for atomic-shell contribution to exciton

[malwi]

Dear Yambo Team,

I am working on a comparison of the BSE exciton analysis with the Judd-Ofelt theory (JOT) that assumes only the on-site excitations.
JOT usually deals with f-f transitions for the rare earth impurities in crystals.
I have BSE results and mostly charge-transfer excitons.

It would be nice to add to the exciton analysis a Yambo-py tool that takes the BSE amplitudes and sums the weights over bands and k-points with additional factor.
This additional factor is the composition of band (k-point) with respect to the f-shell of the impurity (something what we know from projecting bands on the atomic shells).

I am not an expert on python and tried to ask AI for writing the script.
The script is listed below.
The main problem is with reading QE/tmp/prefix.save for the atomic_orbital.xml - I do not know how to get it.

If someone (an expert on Yambopy) would have some time to prepare such script for the community it would be very useful for the rare-earths and transition metals.
I can share all my inputs for a check/example/trial.

With best regards,
Gosia

===================== AI proposition for the script ===========
import yambopy
from yambopy import *
from qepy import *
import numpy as np
import matplotlib.pyplot as plt


# ==========================================
# CONFIGURATION
# ==========================================
#SAVE_DIR = './SAVE'
BSE_DIR = '/net/scratch/hscra/plgrid/plgmalwi/YAM/Er/s9-2'
QE_SAVE = '/net/scratch/hscra/plgrid/plgmalwi/YAM/Er/tmp'
PREFIX = 'Er'

save_path='/net/scratch/hscra/plgrid/plgmalwi/YAM/Er/'
bse_path ='/net/scratch/hscra/plgrid/plgmalwi/YAM/Er/s9-2'

# Indeksy domieszek w pliku QE (uwaga: w pythonie od 0, w QE od 1)
# Zmień na indeksy atomów, które są Twoją domieszką f-elektronową
DOPANT_ATOM_INDICES = [0] # Np. pierwszy atom w strukturze to lantanowiec/aktynowiec
L_QUANTUM_NUMBER = 3 # l=3 oznacza orbitale f

# Wybrany stan ekscytonu do analizy (np. 1 dla najniższego stanu)
EXCITON_INDEX = 16

# ==========================================
# 1. READ DATABASES
# ==========================================
print("Loading Yambo and QE databases...")
#lattice = YamboLatticeDB.from_db_file(folder=SAVE_DIR)
#yexc = YamboExcitonDB.from_db_file(lattice, folder=BSE_DIR)

# Create "lattice" object by reading the ns.db1 database inside the yambo SAVE
lattice = YamboLatticeDB.from_db_file(filename=save_path+'/SAVE/ns.db1')
# Read exciton data at Q=iQ
yexc = YamboExcitonDB.from_db_file(lattice,filename=bse_path+'/ndb.BS_diago_Q1')
proj = ProjwfcXML(prefix=PREFIX, path=QE_SAVE)

# Pobranie wag ekscytonu dla stanu \lambda
# Struktura słownika yexc.get_exciton_weights(): [{'k':k, 'v':v, 'c':c, 'weight':|A|^2}]
exc_weights = yexc.get_exciton_weights(EXCITON_INDEX)

# ==========================================
# 2. MAP FULL BZ (YAMBO) TO IBZ (QE)
# ==========================================
# Yambo przechowuje wektory w Full BZ. Musimy rzutować punkty k z Yambo
# na punkty k w IBZ z Quantum ESPRESSO przy użyciu tabeli mapowania z lattice
kpts_yambo_to_ibz = lattice.kpts_map # Tablica mapująca indeksy k_full -> k_ibz

# ==========================================
# 3. COMPUTE f -> f CONTRIBUTIONS
# ==========================================
total_f_to_f = 0.0
total_exciton_norm = 0.0

print(f"\nAnalyzing Exciton {EXCITON_INDEX}...")
print(f"Total transitions in BSE vector: {len(exc_weights)}")

for transition in exc_weights:
# Wyciągamy indeksy z Yambo (pamiętaj o konwencji indeksowania 0-based/1-based w Yambopy)
k_full = transition['k'] - 1 # Korekta indeksu do tablic pythona
v_band = transition['v'] - 1
c_band = transition['c'] - 1
weight = transition['weight'] # To jest |A_{vck}|^2

total_exciton_norm += weight

# Znajdujemy odpowiadający punkt k w strefie nieprzywiedlnej dla QE
k_ibz = kpts_yambo_to_ibz[k_full]

# Pobieramy rzutowania orbitalne z QE dla pasma walencyjnego (dziura) i przewodnictwa (elektron)
# proj.get_state_projection przyjmuje indeksy k i pasma oraz filtruje po atomie i liczbie l
# Zwraca wartość zzakresu [0.0, 1.0] określającą udział danego orbitalu w pasmie

p_v_f = 0.0
p_c_f = 0.0

for atom_idx in DOPANT_ATOM_INDICES:
# Sumujemy wkład f-orbitali po wszystkich atomach domieszki (jeśli jest ich więcej niż jeden)
p_v_f += proj.get_state_projection(k_ibz, v_band, atom_index=atom_idx, l=L_QUANTUM_NUMBER)
p_c_f += proj.get_state_projection(k_ibz, c_band, atom_index=atom_idx, l=L_QUANTUM_NUMBER)

# Obliczamy wkład przejścia f -> f ważony amplitudą BSE
f_to_f_contribution = weight * p_v_f * p_c_f
total_f_to_f += f_to_f_contribution

# Normowanie wyniku względem całkowitej wagi ekscytonu w bazie danych
f_to_f_percentage = (total_f_to_f / total_exciton_norm) * 100

# ==========================================
# 4. PRINT RESULTS
# ==========================================
print("\n" + "="*40)
print(f"ANALYSIS RESULTS FOR EXCITON {EXCITON_INDEX}")
print("="*40)
print(f"Exciton Norm Check: {total_exciton_norm:.4f}")
print(f"Absolute f->f weight: {total_f_to_f:.6f}")
print(f"Percentage of f->f character: {f_to_f_percentage:.2f} %")
print("="*40)

------------------ END ------------------------