Notebook 6 - Fitting disorder linewidth parameters

This short notebooks shows an example of how to fit L and R in a given material

import sys

import h5py
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from ase.io import read
from scipy.interpolate import interp1d

from tqdm import tqdm

from smooth_disorder.structural import obtain_density, THzToCm, THz, Angstrom
from smooth_disorder.vis.interactive import *

from smooth_disorder.disorder_linewidth import lorentzian_numpy, lorentzian_torch
from smooth_disorder.disorder_linewidth import prepare_fitting_inputs
from smooth_disorder.disorder_linewidth import evaluate_linewidth_and_model_prediction
from smooth_disorder.disorder_linewidth import PDCModel

import torch

CRYSTAL_POSCAR    = "./1_graphite/POSCAR"
DISORDERED_POSCAR = "./2_irg_t2/irg_t2_14009.vasp"

WORK_DIR = "./dl_tutorial"

CRYSTAL_VEL_SAVE     = f"{WORK_DIR}/crystal_vdos_group_vel"
DISORDERED_VDOS_SAVE = f"{WORK_DIR}/disordered_vdos"
SHIFTED_SAVE         = f"{WORK_DIR}/reduced_density_crystal_vdos_group_vel"

# read the input data + setup torch arrays

(density_crystal,
density_disordered,
freq_disordered,
vdos_disordered,
interp_shifted_freq_crystal,
interp_shifted_vdos_crystal,
interp_shifted_speed_crystal) = prepare_fitting_inputs(
    CRYSTAL_POSCAR,
    DISORDERED_POSCAR,
    DISORDERED_VDOS_SAVE,
    SHIFTED_SAVE,
)

# X and Y are the normalised crystal and disordered VDOS used in the loss
X = torch.from_numpy(interp_shifted_vdos_crystal / density_crystal)
Y = torch.from_numpy(vdos_disordered / density_disordered)

# Initial parameter values — intentionally away from the optimum
L0, R0 = 3.3, 5.54

# instantiate the model for fitting the linewidths and setup LBFGS search
model = PDCModel(
    L0, R0,
    density_crystal, density_disordered,
    freq_disordered,
    interp_shifted_freq_crystal,
    interp_shifted_vdos_crystal,
    interp_shifted_speed_crystal,
)

optim = torch.optim.LBFGS(
    model.parameters(),
    lr=1.0,
    max_iter=50,
    line_search_fn="strong_wolfe",
)

loss_fn = torch.nn.MSELoss()

losses, model_parameters_history = {}, {}

# MAIN ITERATION LOOP

def closure():
    optim.zero_grad()
    preds = model(X)
    loss = loss_fn(preds, Y)
    loss.backward()
    print(loss, model.model_params.detach().cpu().numpy().copy())
    return loss

loss = optim.step(closure)

tensor(0.0057, dtype=torch.float64, grad_fn=<MseLossBackward0>) [3.3  5.54]
tensor(0.0057, dtype=torch.float64, grad_fn=<MseLossBackward0>) [3.29906297 5.54084807]
tensor(0.0057, dtype=torch.float64, grad_fn=<MseLossBackward0>) [3.29062973 5.54848067]
tensor(0.0055, dtype=torch.float64, grad_fn=<MseLossBackward0>) [3.20629726 5.62480669]
tensor(0.0043, dtype=torch.float64, grad_fn=<MseLossBackward0>) [2.36297261 6.38806689]
tensor(0.5871, dtype=torch.float64, grad_fn=<MseLossBackward0>) [-2.72415394  8.06097748]
tensor(0.0043, dtype=torch.float64, grad_fn=<MseLossBackward0>) [2.35731212 6.38992835]
tensor(0.0039, dtype=torch.float64, grad_fn=<MseLossBackward0>) [1.84916552 6.55703326]
tensor(0.0039, dtype=torch.float64, grad_fn=<MseLossBackward0>) [1.50392337 6.85466844]
tensor(0.0038, dtype=torch.float64, grad_fn=<MseLossBackward0>) [1.65907103 6.72091478]
tensor(0.0038, dtype=torch.float64, grad_fn=<MseLossBackward0>) [1.70118642 6.97404115]
tensor(0.0035, dtype=torch.float64, grad_fn=<MseLossBackward0>) [1.94422296 8.89727369]
tensor(0.0035, dtype=torch.float64, grad_fn=<MseLossBackward0>) [2.01215958 9.34020213]
tensor(0.0035, dtype=torch.float64, grad_fn=<MseLossBackward0>) [2.03352871 9.45896675]
tensor(0.0035, dtype=torch.float64, grad_fn=<MseLossBackward0>) [2.03621991 9.46528819]
tensor(0.0035, dtype=torch.float64, grad_fn=<MseLossBackward0>) [2.03709562 9.46466436]

Hence final estimated values for L and R are L = 20.37 A and R = 9.46e-6 THz cm nm^3