process_ternary.py

#!/usr/bin/env python3
"""Main optimisation driver for ternary alloy band-gap fitting.

This script proposes new Hubbard-U indices with Bayesian optimisation,
writes corresponding QE input files, evaluates losses, and stores the
result history in an HDF5 file.
"""

from re import U

import numpy as np
import json
import os
import sys
import subprocess 
import h5py
from local.U_standard_main import U_it_to_st,U_st_to_it
from local.make_U import write_U_to_files
from local.bandgaps_alloy_bands import bandgaps_xml_csv
from local.do_loss import find_differences,loss_fn
from local.bayes_class import BayesianOpt
from local.meff import get_xml_info,meff_kpoints_eigenvalues

orbitals={"GaAs":['Ga-4p','As-4p'],"InAs":['In-5p','As-4p'],"GaSb":['Ga-4p','Sb-5p'],"InSb":['In-5p','Sb-5p']}

def iteration_old(Us):
    """Legacy single-iteration flow kept for compatibility."""
    os.chdir('intermediate_SR')
    out0=subprocess.run(['rm -r ./*/'],shell=True)
    out1=subprocess.run(['tar','-xf','base_SR.tar'])
    U_new=U_it_to_st(Us)
    with open('U_int.json','w') as uo:
        json.dump(U_new,uo)
    write_U_to_files(U_new)
    out2=subprocess.run(['../../local/run_all.sh'],shell=True)
    data1=bandgaps_xml_csv()
    delta_E=find_differences(data1,data_don)
    sorted_data_don = data_don[np.lexsort((data_don[:, 0], data_don[:, 1]))]
    delta_E_dimless=delta_E/sorted_data_don[:,2]
    loss=loss_fn(delta_E_dimless)
    os.chdir('../')
    return loss

def iteration(Us,ternary):
    """Run one full calculation iteration and return squared loss."""
    os.chdir(f'intermediate/{ternary}')
    out0=subprocess.run(['rm -r ./*/'],shell=True)
    out1=subprocess.run(['tar','-xf',f'{ternary}.tar'])
    out1=subprocess.run(['echo here'],shell=True)
    U_new=Us
    with open('U_int.json','w') as uo:
        json.dump(U_new,uo)
    write_U_to_files(U_new)
    #Please insert your run script here if you would prefer to arrange the cores yourself, 
    # otherwise the default is to run with all cores available
    out2=subprocess.run(['../../local/run_all.sh >script_out.txt'],shell=True)
    data1=bandgaps_xml_csv()
    delta_E=find_differences(data1,data_don)
    sorted_data_don = data_don[np.lexsort((data_don[:, 0], data_don[:, 1]))]
    delta_E_dimless=delta_E/sorted_data_don[:,2]
    loss=loss_fn(delta_E_dimless)**2
    os.chdir('../../')
    return loss

def iteration_test(Us,ternary):
    """Test-mode iteration that skips launching QE and reuses existing outputs."""
    os.chdir(f'intermediate/{ternary}')
    #out0=subprocess.run(['rm -r ./*/'],shell=True)
    #out1=subprocess.run(['tar','-xf',f'{ternary}.tar'])
    #out1=subprocess.run(['echo here'],shell=True)
    print(Us)
    U_new=Us
    #U_new=U_it_to_st(Us)
    with open('U_int.json','w') as uo:
        json.dump(U_new,uo)
    write_U_to_files(U_new)
    #exit(0)
    #Please insert your run script here if you would prefer to arrange the cores yourself, 
    # otherwise the default is to run with all cores available
    #out2=subprocess.run(['../../local/run_all.sh >script_out.txt'],shell=True)
    data1=bandgaps_xml_csv()
    delta_E=find_differences(data1,data_don)
    sorted_data_don = data_don[np.lexsort((data_don[:, 0], data_don[:, 1]))]
    delta_E_dimless=delta_E/sorted_data_don[:,2]
    loss1=loss_fn(delta_E_dimless)
    print(delta_E_dimless)
    print(loss1)
    os.chdir('../../')
    return loss1

def find_delta(eigs_U,binary):
    """Compute RMSE between selected QE and HSE bands near the gap region."""
    # Align both spectra to a common reference before comparing CB/VB edges.
    eigs_U-=eigs_U[:,13].max()
    with h5py.File("../../binary.hse.data.hdf5","r") as hf:
        eigenvalues_hse=hf[binary]['eigenvalues'][:]
    eigenvalues_hse[:,14:]+=eigs_U[:,14:].min()-eigenvalues_hse[:,14:].min()
    delta_energy=eigs_U-eigenvalues_hse
    delta_total= np.sqrt(np.sum(np.square(delta_energy[:,13:15]))/np.prod(delta_energy[:,13:15].shape))
    return delta_total

def extract_data(Us,loss1,dir):
    """Collect Eg, effective-mass, and band-delta metrics for one iteration."""
    os.chdir(dir)
    Eg_array=np.loadtxt('data_Eg.csv',dtype=str,delimiter=',',skiprows=1)
    Eg_data={}
    for i in range(len(Eg_array)):
        Eg_data[Eg_array[i,0]]=Eg_array[i,1]
    #print(Eg_data)
    deltas_total={}
    meff_data={}
    for key in Eg_data.keys():
        if len(key)==4:
            list1=os.listdir("%s/meff"%key)
            xml_file=''
            for item1 in list1:
                if '.xml' in item1:
                    xml_file=item1
            k_points,eigenvalues,a,occupations=get_xml_info('%s/meff/%s'%(key,xml_file))
            homo_i=int(int(np.sum(occupations[0]))-1)
            homo_i=13
            meff_data["meff_%s"%key]=meff_kpoints_eigenvalues(k_points,eigenvalues,a,homo_i+1)
            k_points,eigenvalues,a,occupations=get_xml_info('%s/short_band/%s'%(key,xml_file))
            delta_total=find_delta(eigenvalues,key)
            deltas_total[key]=delta_total
    data_all={}
    for key in Us.keys():
        data_all["lambda_%s"%key]=Us[key]
    for key in Eg_data.keys():
        data_all["Eg_%s"%key]=Eg_data[key]
    for key in meff_data.keys():
        data_all[key]=meff_data[key]
    loss0=0
    for key in deltas_total.keys():
        data_all["delta_%s"%key]=deltas_total[key]
        loss0+=deltas_total[key]
    data_all['loss1']=loss1
    data_all['loss0']=loss0
    os.chdir('../../')
    return data_all

def current_data(direc):
    """Evaluate current folder outputs against reference cube data."""
    os.chdir(direc)
    data1=bandgaps_xml_csv()
    delta_E=find_differences(data1,data_don)
    sorted_data_don = data_don[np.lexsort((data_don[:, 0], data_don[:, 1]))]
    delta_E_dimless=delta_E/sorted_data_don[:,2]
    loss=loss_fn(delta_E_dimless)
    os.chdir('../')
    return loss


def init_keys(binaries,alloys):
    """Build ordered output keys used in historical data records."""
    data_keys=["lambda_%s"%bina for bina in binaries]
    data_keys.append("Eg_%s"%binaries[0])
    for key in alloys:
        data_keys.append('Eg_%s'%key)
    data_keys.append("Eg_%s"%binaries[1])
    for key in binaries:
        data_keys.append('meff_%s'%key)
    data_keys.append('loss0')
    data_keys.append('loss1')
    return data_keys


def ternary_proc(ternary):
    """Run Bayesian optimisation loop for one ternary system."""
    if ternary=='InGaAs':
        binaries=['GaAs','InAs']
        alloys=['In1Ga3As','In2Ga2As','In3Ga1As']
    elif ternary=='InGaSb':
        binaries=['GaSb','InSb']
        alloys=['In1Ga3Sb','In2Ga2Sb','In3Ga1Sb']
    elif ternary=='InAsSb':
        binaries=['InAs','InSb']
        alloys=['InAs1Sb3','InAs2Sb2','InAs3Sb1']
    elif ternary=='GaAsSb':
        binaries=['GaAs','GaSb']
        alloys=['GaAs1Sb3','GaAs2Sb2','GaAs3Sb1']
    else:
        print('not a valid ternary')
        exit(0)
    data_files={}
    pbounds={}
    with h5py.File('data_int.hdf5','r') as hf:
        for i,bina in enumerate(binaries):
            data_files[bina]=hf[bina]
            # Search space is the discrete index range of precomputed U grids.
            pbounds[bina]=(0,len(data_files[bina]['Eg(eV)'])-1)
            
    data_file = f"data_{ternary}.hdf5"
    data_keys = init_keys(binaries, alloys)
    print(data_keys)

    # Initialize HDF5 file with extendable datasets
    with h5py.File(data_file, "w") as hf:

        # --- lambda group ---
        lambda_grp = hf.create_group("lambda")

        for bina in binaries:
            g = lambda_grp.create_group(bina)
            g.create_dataset(
                "index",
                shape=(0,),
                maxshape=(None,),
                dtype="i4"
            )
            g.create_dataset(
                "delta",
                shape=(0,),
                maxshape=(None,),
                dtype="f8"
            )

        # --- Eg group ---
        Eg_grp = hf.create_group("Eg")

        Eg_materials = [binaries[0]] + alloys + [binaries[1]]

        # Sort experimental data same way you use elsewhere
        sorted_data_don = data_don[np.lexsort((data_don[:, 0], data_don[:, 1]))]

        for i, mat in enumerate(Eg_materials):

            ds = Eg_grp.create_dataset(
                mat,
                shape=(0,),
                maxshape=(None,),
                dtype="f8"
            )

            # Attach experimental metadata
            ds.attrs["Eg_exp"] = float(sorted_data_don[i,2])
            ds.attrs["x"] = float(sorted_data_don[i,0])
            ds.attrs["y"] = float(sorted_data_don[i,1])

        # --- alpha group ---
        ell_grp = hf.create_group("ell")
        ell_grp.create_dataset(
            "ell1",
            shape=(0,),
            maxshape=(None,),
            dtype="f8"
        )
        ell_grp.create_dataset(
            "ell0",
            shape=(0,),
            maxshape=(None,),
            dtype="f8"
        )

    print(pbounds)
    BO=BayesianOpt(pbounds)
    #exit(0)
    iterations=1
    #U_register={'GaAs':0,'InAs':0,'GaSb':0,'InSb':0}
    for i in range(iterations):
        length_along=BO.suggest()
        for key in length_along:
            # Bayesian suggestion is continuous; project to valid integer grid index.
            length_along[key]=np.round(length_along[key],0)
        print(i,length_along)
        Us_new={"InAs":{'In-5p':0.0001,'As-4p':0.0001},"GaAs":{'Ga-4p':0.0001,'As-4p':0.0001},
                "InSb":{'In-5p':0.0001,'Sb-5p':0.0001},"GaSb":{'Ga-4p':0.0001,'Sb-5p':0.0001}}
        with h5py.File('data_int.hdf5','r') as hf:
            for bina in binaries:
                # Map selected discrete index to orbital-resolved U parameters.
                Us_new[bina]={orbitals[bina][0]:hf[f"{bina}/U_{orbitals[bina][0]}"][int(length_along[bina])],
                              orbitals[bina][1]:hf[f"{bina}/U_{orbitals[bina][1]}"][int(length_along[bina])]}
        print(Us_new)
        U_register={}
        for binary in binaries:
            U_register[binary]=int(np.round(length_along[binary],0))
        #print(i,U_register)
        loss=iteration_test(Us_new,ternary)
        BO.register_param(U_register,loss)
        data_all=extract_data(U_register,loss,'intermediate/%s'%ternary)
        with h5py.File(data_file, "a") as hf:

            # ----- lambda -----
            for binary in binaries:
                ds = hf["lambda"][binary]["index"]
                ds.resize((ds.shape[0] + 1,))
                ds[-1] = U_register[binary]

                ds1 = hf["lambda"][binary]["delta"]
                ds1.resize((ds1.shape[0] + 1,))
                ds1[-1] = data_all[f"delta_{binary}"]

            # ----- Eg -----
            for key, val in data_all.items():
                if key.startswith("Eg_"):
                    mat = key.replace("Eg_", "")
                    ds = hf["Eg"][mat]
                    ds.resize((ds.shape[0] + 1,))
                    ds[-1] = float(val)

            # ----- alpha -----
            ds = hf["ell"]["ell1"]
            ds.resize((ds.shape[0] + 1,))
            ds[-1] = float(data_all["loss1"])
            ds = hf["ell"]["ell0"]
            ds.resize((ds.shape[0] + 1,))
            ds[-1] = float(data_all["loss0"])

import sys
if __name__=="__main__":
    # Usage: python process_ternary.py [InGaAs|InGaSb|InAsSb|GaAsSb]
    if len(sys.argv)>1:
        ternary=sys.argv[1]
    else:
        ternary='InGaAs'
    data_don=np.loadtxt("cube_donati_{}.csv".format(ternary),delimiter=',',skiprows=1)
    ternary_proc(ternary)