cp2k/tools/rixs_postprocess.py

230 lines
6.7 KiB
Python
Executable file

#!/usr/bin/env python3
# author: BSG
import numpy as np
import matplotlib.pyplot as plt
import argparse
def parse_args():
parser = argparse.ArgumentParser(description="Run RIXS convolution script.")
parser.add_argument(
"--filename", type=str, help="Input data file from cp2k (*.rixs)"
)
parser.add_argument("--v_states", type=int, help="Number of valence excited states")
parser.add_argument(
"--gamma", type=float, default=0.075, help="Broadening width in eV"
)
parser.add_argument(
"--kappa", type=float, default=60.0, help="Scattering angle in degrees"
)
return parser.parse_args()
# hbar = 6.6260689 * 10**-34 / (2 * np.pi)
hbar = 1 # a.u.
dirs = ["x", "y", "z"]
ndirs = len(dirs)
def get_sigma_tensor(f_state, w_in, w_out, w_0i, w_if, mu_0i, mu_if, Gamma):
"""
Calculate absorption cross-section tensor of final valence state f:
sigma_{alpha beta gamma delta}^f(w)
Parameters:
- f_index: int, index of the final valence state
- w_in: float, incident photon frequency
- w_out: float, outgoing photon frequency
- w_0i: array (c_states,)
- w_if: array (v_states,)
- mu_0i: array (3) dipole moments (i->0)
- mu_if: array (c_states, v_states, 3) dipole moments (i->f)
- Gamma: float, broadening parameter
Returns:
- sigma: array (3,3,3,3) tensor
"""
sigma = np.zeros((ndirs, ndirs, ndirs, ndirs), dtype=np.float64)
c_states = len(w_0i)
for i in range(c_states):
denom = (w_in - w_0i[i]) ** 2 + Gamma**2
prefac = (w_if[f] ** 2 * w_0i[i] ** 2) / denom
for alpha in range(ndirs):
for beta in range(ndirs):
for gamma in range(ndirs):
for delta in range(ndirs):
sigma[alpha, beta, gamma, delta] += (
mu_if[i, f_state, alpha]
* mu_if[i, f_state, beta]
* mu_0i[i, gamma]
* mu_0i[i, delta]
)
sigma *= prefac
return sigma
def parse_spectrum_file(filename, v_states):
"""
Parameters:
- filename: str, path to the spectrum file
- v_states: int, number of valence excited states
"""
w_0i = []
mu_0i = []
w_if = []
mu_if = []
w_f0 = []
with open(filename, "r") as f:
lines = f.readlines()
in_core_block = False
in_valence_block = False
count = 0
for line in lines:
line = line.strip()
if "Excitation from ground-state" in line:
in_core_block = True
in_valence_block = False
continue
elif "Emission from core-excited state" in line:
in_core_block = False
in_valence_block = True
continue
elif line.startswith("=") or line.startswith("w_0i") or not line:
continue
tokens = line.split()
if in_core_block and len(tokens) >= 5:
e0i, mu_x, mu_y, mu_z = map(float, tokens[0:4])
w_0i.append(e0i / hbar)
mu_0i.append([mu_x, mu_y, mu_z])
elif in_valence_block and len(tokens) >= 6:
count += 1
ei, ef = map(float, tokens[0:2])
mu_x, mu_y, mu_z = map(float, tokens[2:5])
mu_if.append((mu_x, mu_y, mu_z))
w_if.append((ef - ei) / hbar)
if count <= v_states:
w_f0.append(ef / hbar)
xas_w_0i = np.array(w_0i)
xas_mu_0i = np.array(mu_0i)
rixs_w_if = np.array(w_if)
tddft_w_f0 = np.array(w_f0)
c_states = len(mu_0i)
# sanity check
expected_total = c_states * v_states
if len(mu_if) != expected_total:
raise ValueError(
f"Expected {expected_total} spectrum lines, but found {len(mu_if)}"
)
rixs_mu_if = np.zeros((c_states, v_states, 3))
for idx, (mx, my, mz) in enumerate(mu_if):
i = idx // v_states
f = idx % v_states
if i < c_states and f < v_states:
rixs_mu_if[i, f] = [mx, my, mz]
return xas_w_0i, xas_mu_0i, rixs_w_if, rixs_mu_if, tddft_w_f0
def convolve_sigma_averaged_tensor(
sigma_tensor, v_states, w_f0, w_0i, w_in_window, kappa_rad, Gamma
):
"""
Compute the RIXS spectrum with custom broadening.
Parameters:
- sigma_tensor: shape (v_states, 3, 3, 3, 3)
- v_states: int, number of valence excited states
- w_f0: (v_states,) energy loss values (valence excitation energies)
- w_0i: (N,) array of incident photon energies (XAS axis)
- kappa_rad: float, scattering angle in radians
- Gamma: float, broadening width
Returns:
- 2D array representing the RIXS spectrum
"""
ndirs = 3
A = 3 + np.cos(kappa_rad) ** 2
B = 0.5 * (1 - 3 * np.cos(kappa_rad) ** 2)
w_out_window = np.linspace(0, np.max(w_f0), len(w_in_window))
N_incident = len(w_in_window)
N_loss = len(w_out_window)
spectrum = np.zeros((N_incident, N_loss))
W_IN, W_LOSS = np.meshgrid(w_in_window, w_out_window, indexing="ij")
W_OUT = W_IN - W_LOSS
for f in range(v_states):
denom = (W_IN - w_0i) ** 2 + Gamma**2
prefac = (w_if[f] ** 2 * w_0i**2) / denom
Omega = W_OUT + w_f0[f] - W_IN
Delta = Gamma / np.pi / (Omega**2 + Gamma**2)
val = 0.0
for alpha in range(ndirs):
for beta in range(ndirs):
val += A * sigma_tensor[f, alpha, alpha, beta, beta]
val += B * (
sigma_tensor[f, alpha, beta, alpha, beta]
+ sigma_tensor[f, alpha, beta, beta, alpha]
)
spectrum += prefac * val * Delta / 30.0
return spectrum / np.max(spectrum)
if __name__ == "__main__":
args = parse_args()
filename = args.filename
v_states = args.v_states
kappa_deg = args.kappa
kappa_rad = np.radians(kappa_deg)
Gamma = args.gamma
molecule = filename.split(".")[0]
w_0i, mu_0i, w_if, mu_if, w_f0 = parse_spectrum_file(filename, v_states)
print(np.shape(w_if), np.shape(mu_if), np.shape(w_f0))
w_in = w_0i[0]
w_out = w_if[0]
sigma_tensor = np.zeros((v_states, 3, 3, 3, 3), dtype=np.float64)
for f in range(v_states):
sigma_tensor[f] = get_sigma_tensor(
f, w_in, w_out, w_0i, w_if, mu_0i, mu_if, Gamma
)
w_in = w_0i[0] # incident photon energy = first xas energy
e_window = 0.2
n_points = 50
w_in_window = np.linspace(w_in - e_window, w_in + e_window, n_points)
spectrum_2d = convolve_sigma_averaged_tensor(
sigma_tensor, v_states, w_f0, w_0i[0], w_in_window, kappa_rad, Gamma
)
project_name = filename.split(".")[0]
np.savetxt(f"{project_name}_rixs.dat", spectrum_2d, fmt="%.6f")