from __future__ import annotations
import sys
import os
import time
import logging
import warnings
import argparse
import re
import tomllib
import textwrap
from dataclasses import dataclass, asdict
from contextlib import nullcontext
import numpy as np
import xmltodict
import termplotlib as tplt
from qiskit.quantum_info import Statevector
import dopyqo
from dopyqo.fci_vector_matrix import statevector_from_fci
from dopyqo import units
from dopyqo.scripts import banners
from dopyqo.colors import *
from dopyqo import DopyqoConfig
import dopyqo.wannier90
import dopyqo.wfc
def has_save_folder(folder: str) -> str | None:
for item in os.listdir(None if folder == "" else folder):
item_path = os.path.join(folder, item)
if os.path.isdir(item_path) and item.endswith(".save"):
return item[:-5]
return None
[docs]
def run(
config: DopyqoConfig | dict,
mats: dopyqo.MatrixElements = dopyqo.MatrixElements(),
return_h: bool = True,
return_wfc: bool = True,
return_mats: bool = True,
show_banner=True,
verbosity: int = 2,
show_mat_calc_progress: bool = False,
) -> tuple[dict, dopyqo.Wfc | None, dopyqo.Hamiltonian | None, dopyqo.MatrixElements | None]:
"""Run a Dopyqo calculation
Args:
config (DopyqoConfig | dict): Dopyqo config or dictionary that can be converted into a DopyqoConfig object
mats (dopyqo.MatrixElements, optional): dopyqo.MatrixElements object. Matrix elements that are not None are not calculate and used from this parameter instead.
Defaults to dopyqo.MatrixElements().
return_h (bool, optional): If the the created dopyqo.Hamiltonian object is returned. Defaults to True.
return_wfc (bool, optional): If the the created dopyqo.Wftc object is returned. Defaults to True.
return_mats (bool, optional): If the the created dopyqo.MatrixElements object is returned. Defaults to True.
show_banner (bool, optional): If the Dopyqo banner is shown. Defaults to True.
verbosity (int, optional): Verbosity of the printed output. Higher is more verbose. 0 is minimal prints, 1 prints the active space,
2 prints summaries of everything. Defaults to 2.
show_mat_calc_progress (bool, optional): Show the progress of the matrix element calculation using spinners and timers. Defaults to False.
Returns:
tuple[dict, dopyqo.Wfc | None, dopyqo.Hamiltonian | None, dopyqo.MatrixElements | None]: Dictionary of energies, and dopyqo.Wfc object,
dopyqo.Hamiltonian object, dopyqo.MatrixElements object
"""
block_color = SOFT_GREEN
yes_color = NO_COLOR
no_color = NO_COLOR
verbosity_summary = 2
verbosity_active_space = 1
if show_mat_calc_progress:
try:
from rich.progress import Progress, TextColumn, TimeElapsedColumn, SpinnerColumn
except ImportError:
show_mat_calc_progress = False
print(f"{ORANGE}Import warning: Could not import rich package. Matrix calculation progress will not be shown.{RESET_COLOR}")
def make_progress():
return Progress(
SpinnerColumn(spinner_name="dots", style="magenta"),
TextColumn("[bold cyan]{task.description}"),
TimeElapsedColumn(),
# transient=True,
)
start_time_all = time.perf_counter()
######################### BANNER #########################
if show_banner:
banner_attributes = vars(banners)
banner_strings = [value for name, value in banner_attributes.items() if isinstance(value, str) and not name.startswith("_")]
rng = np.random.default_rng()
random_banner = rng.choice(banner_strings)
dopyqo.print_banner(random_banner, flush=True)
######################### CHECK CONFIG #########################
if isinstance(config, dict):
config = DopyqoConfig(**config)
if config.run_vqe:
if config.uccsd_reps < 1:
print(f"{RED}Config error: uccsd_reps must be greater or equal to 1 but is {config.uccsd_reps}!{RESET_COLOR}", file=sys.stderr)
sys.exit(1)
######################### SETTINGS SUMMARY #########################
if config.logging_flag:
logging.basicConfig(level=logging.INFO)
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Settings summary", color=block_color, flush=True)
print(f"Active space of ({config.active_electrons}e, {config.active_orbitals}o)")
if config.kpoint_idx is not None:
if isinstance(config.kpoint_idx, int):
print(f"k-point index: {config.kpoint_idx}")
else:
print(f"k-points: {config.kpoint_idx}")
if config.wannier_transform:
msg_tmp = f"U-matrix path: {config.wannier_umat}"
if config.wannier_input_file is not None:
msg_tmp += f", Wannier input file: {config.wannier_input_file}"
print(f"Using Wannier functions: {yes_color}YES{RESET_COLOR} ({msg_tmp})")
else:
print(f"Using Wannier functions: {no_color}NO{RESET_COLOR}")
if config.run_fci:
print(f"Running a FCI calculation: {yes_color}YES{RESET_COLOR}")
else:
print(f"Running a FCI calculation: {no_color}NO{RESET_COLOR}")
if config.run_vqe:
reps_str = f"{config.uccsd_reps} " + ("repetition" if config.uccsd_reps == 1 else "repetitions") + " of the UCCSD ansatz"
vqe_simulator_str = "Qiskit" if config.use_qiskit else "TenCirChem"
occ_str = f"and using occupations {config.occupations}" if config.occupations is not None else ""
if config.vqe_parameters is not None:
print(
f"Running a VQE calculation: {yes_color}YES{RESET_COLOR} by executing the ansatz with {vqe_simulator_str} with the given parameters and {reps_str} {occ_str}"
)
else:
print(
f"Running a VQE calculation: {yes_color}YES{RESET_COLOR} with {vqe_simulator_str} using the {config.vqe_optimizer} optimizer and {reps_str} {occ_str}"
)
else:
print(f"Running a VQE calculation: {no_color}NO{RESET_COLOR}")
if config.logging_flag:
logging.basicConfig(level=logging.INFO)
print(f"Showing logging: {yes_color}YES{RESET_COLOR}")
else:
print(f"Showing logging: {no_color}NO{RESET_COLOR}")
if config.n_threads > 1:
print(f"CPU threads: {config.n_threads}")
sys.stdout.flush()
binary_occupations = True # TODO: Put this in DopyqoConfig?
######################### READING FILES #########################
xml_file = os.path.join(config.base_folder, f"{config.prefix}.save", "data-file-schema.xml")
pp_files = [
os.path.join(config.base_folder, f"{config.prefix}.save", filename)
for filename in os.listdir(os.path.join(config.base_folder, f"{config.prefix}.save"))
if filename.lower().endswith(".upf")
]
pps = [dopyqo.pseudopot.Pseudopot(pp_file) for pp_file in pp_files]
logging.info("Reading data from files...")
wfc_files = [x for x in os.listdir(os.path.join(config.base_folder, f"{config.prefix}.save")) if x.startswith("wfc")]
wfc_files = sorted(wfc_files, key=lambda x: float(x[3:].split(".")[0]))
wfc_files = [os.path.join(config.base_folder, f"{config.prefix}.save", x) for x in wfc_files]
if config.kpoint_idx == "all":
wfc_obj_lst = dopyqo.Wfc.from_file_all_kpoints(wfc_files, xml_file, pseudopots=pps)
else:
if config.kpoint_idx is None:
wfc_file = wfc_files[0]
else:
wfc_file = wfc_files[config.kpoint_idx]
wfc_obj_lst = [dopyqo.Wfc.from_file(wfc_file, xml_file, pseudopots=pps, kpoint_idx=config.kpoint_idx)]
if len(wfc_obj_lst) == 1:
str_tmp = ""
kpoint_tmp = wfc_obj_lst[0].kpoint
if np.allclose(kpoint_tmp, 0.0):
str_tmp += " (Γ-point)"
if verbosity >= verbosity_summary:
print(f"k-point: {kpoint_tmp}{str_tmp}", flush=True)
sys.stdout.flush()
######################### CHANGING ATOM POSITIONS AND LATTICE VECTORS #########################
block_title = ""
if config.atom_positions is not None:
block_title += "atomic positions"
atom_positions_hartree_old = wfc_obj_lst[0].atom_positions_hartree.copy()
for wfc_obj in wfc_obj_lst:
wfc_obj.set_atom_positions(config.atom_positions, config.unit)
if config.lattice_vectors is not None:
block_title += " and lattice vectors"
lattice_vectors_old = wfc_obj_lst[0].a.copy()
for wfc_obj in wfc_obj_lst:
wfc_obj.set_lattice_vectors(config.lattice_vectors, config.unit)
if (config.atom_positions is not None or config.lattice_vectors is not None) and verbosity >= verbosity_summary:
print()
dopyqo.print_block(f"Changing {block_title}", color=block_color, flush=True)
if config.atom_positions is not None:
elements_tmp = [dopyqo.elements_to_atomic_number[x] for x in wfc_obj_lst[0].atomic_numbers]
old_str_tmp = textwrap.indent(np.array2string(atom_positions_hartree_old), "\t")
new_str_tmp = textwrap.indent(np.array2string(wfc_obj_lst[0].atom_positions_hartree), "\t")
old_str_tmp_max_len = max([len(x) for x in old_str_tmp.split("\n")])
old_str_tmp = "\n".join(
[x + (old_str_tmp_max_len - len(x)) * " " + " <- " + elements_tmp[i] for i, x in enumerate(old_str_tmp.split("\n"))]
)
new_str_tmp_max_len = max([len(x) for x in new_str_tmp.split("\n")])
new_str_tmp = "\n".join(
[x + (new_str_tmp_max_len - len(x)) * " " + " <- " + elements_tmp[i] for i, x in enumerate(new_str_tmp.split("\n"))]
)
print(f"Move atomic positions (in Hartree units) from \n{old_str_tmp}\nto\n{new_str_tmp}")
if config.lattice_vectors is not None:
postfix_tmp = ["a1", "a2", "a3"]
old_str_tmp = textwrap.indent(np.array2string(lattice_vectors_old), "\t")
new_str_tmp = textwrap.indent(np.array2string(wfc_obj_lst[0].a), "\t")
old_str_tmp_max_len = max([len(x) for x in old_str_tmp.split("\n")])
old_str_tmp = "\n".join(
[x + (old_str_tmp_max_len - len(x)) * " " + " <- " + postfix_tmp[i] for i, x in enumerate(old_str_tmp.split("\n"))]
)
new_str_tmp_max_len = max([len(x) for x in new_str_tmp.split("\n")])
new_str_tmp = "\n".join(
[x + (new_str_tmp_max_len - len(x)) * " " + " <- " + postfix_tmp[i] for i, x in enumerate(new_str_tmp.split("\n"))]
)
print(f"Change lattice vectors (in Hartree units) from \n{old_str_tmp}\nto\n{new_str_tmp}")
sys.stdout.flush()
######################### RUN CALCULATION FOR EVERY K-POINT #########################
reference_energy = wfc_obj_lst[0].etot
mats_lst = []
h_ks_lst = []
vqe_result_lst = []
hf_energy_lst = []
hf_energy_pyscf_lst = []
fci_energy_lst = []
vqe_energy_lst = []
kpoint_weights = []
for wfc_idx, wfc_obj in enumerate(wfc_obj_lst):
######################### LOADING QE DATA #########################
if len(wfc_obj_lst) > 1 and verbosity >= verbosity_summary:
print()
dopyqo.print_block(
f"k-point {wfc_idx+1}/{len(wfc_obj_lst)}\n{wfc_obj.kpoint}\nweight: {wfc_obj.kpoint_weight}", color=block_color, flush=True
)
overlaps_ncpp = wfc_obj.get_overlaps()
if not np.allclose(overlaps_ncpp, np.eye(overlaps_ncpp.shape[0])):
print(
f"{RED}Wavefunction error: The wavefunctions are not orthonormal!{RESET_COLOR}",
file=sys.stderr,
)
sys.exit(1)
p = wfc_obj.k_plus_G # shape (#waves, 3)
c_ip = wfc_obj.evc # shape (#bands, #waves)
logging.info("Number of plane waves: %i", p.shape[0])
orbital_indices_core, orbital_indices_active = wfc_obj.active_space(
active_electrons=config.active_electrons, active_orbitals=config.active_orbitals
)
occs_tmp = wfc_obj.occupations_binary if binary_occupations else wfc_obj.occupations
occ_str = dopyqo.active_space_string(occs_tmp, orbital_indices_active)
if verbosity >= verbosity_active_space:
print(f"The following active space will be considered:\n{occ_str}")
print(f"Number of plane waves: {p.shape[0]}")
sys.stdout.flush()
occupations_full, c_ip_full = wfc_obj.get_orbitals_by_index(
orbital_indices_core + orbital_indices_active, binary_occupations=binary_occupations
)
occupations_active, c_ip_active = wfc_obj.get_orbitals_by_index(orbital_indices_active, binary_occupations=binary_occupations)
with warnings.catch_warnings(): # Ignore warning that all core orbitals are occupied
warnings.filterwarnings("ignore", category=Warning)
occupations_active_xml, _ = wfc_obj.get_orbitals_by_index(orbital_indices_active, binary_occupations=False)
occupations_core, c_ip_core = wfc_obj.get_orbitals_by_index(orbital_indices_core, binary_occupations=binary_occupations)
# Calculate matrix elements
my_context = nullcontext()
if show_mat_calc_progress:
my_context = make_progress()
with my_context as progress:
if show_mat_calc_progress:
task = progress.add_task("Calculating kinetic energy", total=None)
else:
print("Calculating kinetic energy...", flush=True)
start_time_ham = time.perf_counter()
######################### KINETIC ENERGY #########################
iTj_orbitals = mats.h_pq_kin
if iTj_orbitals is None:
logging.info("Calculate kinetic energy...")
start_time = time.perf_counter()
iTj_orbitals = dopyqo.calc_matrix_elements.iTj(p, c_ip_active)
time_kin = time.perf_counter() - start_time
time_kin_str = f"{time_kin:.2f}s"
logging.info("Kinetic energy calc. took %.2fs", time_kin)
else:
logging.info("Loading kinetic energy...")
time_kin_str = "- Given by " + f"{mats.h_pq_kin=}".split("=")[0] + " parameter"
######################### PSEUDOPOTENTIAL #########################
# Progress bar does not update because of blocking Rust function execution,
# need to put rust function in separate thread and properly catch the return value :$
my_context = nullcontext()
if show_mat_calc_progress:
my_context = make_progress()
with my_context as progress:
if show_mat_calc_progress:
task = progress.add_task("Calculating pseudopotential", total=None)
else:
print("Calculating pseudopotential...", flush=True)
pp_orbitals = mats.h_pq_pp
pp_core = None
if pp_orbitals is None:
# print("Calculating pseudopotential ", end="", flush=True)
n_threads_pp = config.n_threads
start_time = time.perf_counter()
pp_full = dopyqo.calc_pseudo_pot.calc_pps(
p,
c_ip_full,
wfc_obj.cell_volume,
wfc_obj.atom_positions_hartree,
wfc_obj.atomic_numbers,
pps,
n_threads=n_threads_pp,
# os.path.join(config.base_folder, "pp_pw.npy"),
)
time_pp = time.perf_counter() - start_time
time_pp_str = f"{time_pp:.2f}s"
pp_orbitals = pp_full[orbital_indices_active, :][:, orbital_indices_active]
pp_core = pp_full[orbital_indices_core, :][:, orbital_indices_core]
hours = int(time_pp // 3600)
minutes = int(time_pp // 60 - hours * 60)
seconds = int(time_pp - minutes * 60 - hours * 3600)
# print(f"{hours:01d}:{minutes:02d}:{seconds:02d}")
# print(time_pp_str, flush=True)
else:
logging.info("Loading pseudopotentials...")
time_pp_str = "- Given by " + f"{mats.h_pq_pp=}".split("=")[0] + " parameter"
######################### ELECTRON REPULSION INTEGRALS #########################
my_context = nullcontext()
if show_mat_calc_progress:
my_context = make_progress()
with my_context as progress:
if show_mat_calc_progress:
task = progress.add_task("Calculating electron repulsion integrals", total=None)
else:
print("Calculating electron repulsion integrals...", flush=True)
h_pqrs = mats.h_pqrs
if h_pqrs is None:
logging.info("Calculate ERIs...")
start_time = time.perf_counter()
h_pqrs: np.ndarray = (
dopyqo.eri(
c_ip=c_ip_active,
b=wfc_obj.b,
mill=wfc_obj.mill,
fft_grid=wfc_obj.fft_grid,
use_gpu=config.use_gpu,
)
/ wfc_obj.cell_volume
)
time_eri = time.perf_counter() - start_time
time_eri_str = f"{time_eri:.2f}s"
logging.info("ERI calc. took %.2fs", time_eri)
else:
logging.info("Loading ERIs...")
time_eri_str = "- Given by " + f"{mats.h_pqrs=}".split("=")[0] + " parameter"
######################### FROZEN CORE POTENTIAL AND ENERGY #########################
my_context = nullcontext()
if show_mat_calc_progress:
my_context = make_progress()
with my_context as progress:
if show_mat_calc_progress:
task = progress.add_task("Calculating frozen core", total=None)
else:
print("Calculating frozen core...", flush=True)
h_pq_core = np.zeros_like(iTj_orbitals)
energy_frozen_core = 0.0
if len(c_ip_core) > 0:
h_pq_core = mats.h_pq_core
energy_frozen_core = mats.energy_frozen_core
if h_pq_core is None and energy_frozen_core is None:
logging.info("Calculate frozen core potential and frozen core energy...")
start_time = time.perf_counter()
h_pq_core, energy_frozen_core = dopyqo.get_frozen_core_pot_and_energy_given_pp(
p=p,
c_ip_core=c_ip_core,
c_ip_active=c_ip_active,
b=wfc_obj.b,
mill=wfc_obj.mill,
cell_volume=wfc_obj.cell_volume,
pp_core=pp_core,
occupations_core=occupations_core,
fft_grid=wfc_obj.fft_grid,
use_gpu=config.use_gpu,
)
energy_frozen_core = energy_frozen_core.real
time_energy_core = time.perf_counter() - start_time
time_energy_core_str = f"{time_energy_core:.2f}s"
logging.info("Frozen core potential and energy calc. took %.2fs", time_energy_core)
elif h_pq_core is None:
logging.info("Calculate frozen core potential...")
start_time = time.perf_counter()
h_pq_core = dopyqo.get_frozen_core_pot(
c_ip_core=c_ip_core,
c_ip_active=c_ip_active,
b=wfc_obj.b,
mill=wfc_obj.mill,
fft_grid=wfc_obj.fft_grid,
cell_volume=wfc_obj.cell_volume,
use_gpu=config.use_gpu,
)
time_energy_core = time.perf_counter() - start_time
time_energy_core_str = (
f"{time_energy_core:.2f}s (potential) - Energy given by " + f"{mats.energy_frozen_core=}".split("=")[0] + " parameter"
)
logging.info("Frozen core potential and energy calc. took %.2fs", time_energy_core)
elif energy_frozen_core is None:
logging.info("Calculate frozen core energy...")
start_time = time.perf_counter()
energy_frozen_core = dopyqo.get_frozen_core_energy_pp(
p=p,
c_ip_core=c_ip_core,
b=wfc_obj.b,
mill=wfc_obj.mill,
cell_volume=wfc_obj.cell_volume,
atom_positions=wfc_obj.atom_positions_hartree,
atomic_numbers=wfc_obj.atomic_numbers,
occupations_core=occupations_core,
pseudopots=pps,
fft_grid=wfc_obj.fft_grid,
use_gpu=config.use_gpu,
).real
time_energy_core = time.perf_counter() - start_time
time_energy_core_str = (
f"{time_energy_core:.2f}s (energy) - Potential given by " + f"{mats.h_pq_core=}".split("=")[0] + " parameter"
)
else:
logging.info("Loading frozen core...")
time_energy_core_str = (
"- Given by " + f"{mats.h_pq_core=}".split("=")[0] + " and " + f"{mats.energy_frozen_core=}".split("=")[0] + " parameters"
)
######################### NUCLEAR REPULSION #########################
my_context = nullcontext()
if show_mat_calc_progress:
my_context = make_progress()
with my_context as progress:
if show_mat_calc_progress:
task = progress.add_task("Calculating ewald energy", total=None)
else:
print("Calculating ewald energy...", flush=True)
h_pq_ewald = mats.h_pq_ewald
energy_ewald = mats.energy_ewald
if energy_ewald is None:
if wfc_obj.ewald != 0.0 and config.qe_ewald:
logging.info("Using Ewald energy from QE (%s)...", wfc_obj.ewald)
overlap = np.einsum("ij, kj -> ik", c_ip_active.conj(), c_ip_active)
energy_ewald = wfc_obj.ewald
h_pq_ewald = overlap * wfc_obj.ewald
else:
logging.info("Calculate nuclear repulsion...")
start_time = time.perf_counter()
lattice_vectors = np.array([wfc_obj.a1, wfc_obj.a2, wfc_obj.a3])
lattice_vectors_reciprocal = np.array([wfc_obj.b1, wfc_obj.b2, wfc_obj.b3])
energy_ewald = dopyqo.nuclear_repulsion_energy_ewald(
wfc_obj.atom_positions_hartree,
wfc_obj.atomic_numbers_valence,
lattice_vectors,
lattice_vectors_reciprocal,
wfc_obj.cell_volume,
gcut=wfc_obj.gcutrho,
)
time_nucl_rep = time.perf_counter() - start_time
time_nucl_rep_str = f"{time_nucl_rep:.2f}s"
logging.info("Nuclear repulsion calc. took %.2fs", time_nucl_rep)
logging.info(
"Ewald energy: %s (QE: %s, Diff.: %s)",
energy_ewald,
wfc_obj.ewald,
np.abs(energy_ewald - wfc_obj.ewald),
)
if config.atom_positions is None and config.lattice_vectors is None and not np.isclose(energy_ewald, wfc_obj.ewald):
print(
f"{ORANGE}Ewald warning: Calculated Ewald energy is not close to the QE Ewald energy (QE: {wfc_obj.ewald}, Dopyqo: {energy_ewald.real}, Diff. {np.abs(energy_ewald-wfc_obj.ewald)})!{RESET_COLOR}"
)
overlap = np.einsum("ij, kj -> ik", c_ip_active.conj(), c_ip_active)
h_pq_ewald = overlap * energy_ewald
# If nucl. rep. energy (E_nuc) is a matrix it needs to be rescaled
# Imagine E_nuc is a number and we want to build a matrix from that.
# A diagonal matrix with E_nuc on its diagonal will lead to the wrong energy
# contribution. The following rescales it correctly.
if isinstance(h_pq_ewald, np.ndarray):
h_pq_ewald = h_pq_ewald / (2 * np.sum(occupations_active))
else:
overlap = np.einsum("ij, kj -> ik", c_ip_active.conj(), c_ip_active)
h_pq_ewald = overlap * energy_ewald
if config.atom_positions is not None or config.lattice_vectors is not None:
print(
f"{ORANGE}Ewald warning: Ewald matrix elements are given by "
+ f"{mats.h_pq_ewald=}".split("=")[0]
+ f" parameter but atom positions or lattice vectors were changed! Are you sure you want to do this?{RESET_COLOR}"
)
logging.info("Loading nuclear repulsion...")
time_nucl_rep_str = "- Given by " + f"{mats.energy_ewald=}".split("=")[0] + " parameter"
h_pq = iTj_orbitals + pp_orbitals + h_pq_core # + h_pq_ewald
######################### Electron self-energy #########################
# Source: Fraser et al. 1996, "Finite-size effects and Coulomb interactions
# in quantum Monte Carlo calculations for homogeneous
# systems with periodic boundary conditions",
# Equations (19, 28). Basically describing the interaction energy, one electron has
# with itself in neighbouring cells.
# \sum_{T \neq 0} 1/|T|
energy_e_self = mats.energy_e_self
if energy_e_self is None:
print("Calculating electron self-energy...", flush=True)
start_time = time.perf_counter()
energy_e_self = wfc_obj.nelec * dopyqo.nuclear_repulsion_energy_ewald(
atom_positions=np.array([[0.0, 0.0, 0.0]]),
atomic_numbers=np.array([1]),
lattice_vectors=np.array([wfc_obj.a1, wfc_obj.a2, wfc_obj.a3]),
lattice_vectors_reciprocal=np.array([wfc_obj.b1, wfc_obj.b2, wfc_obj.b3]),
cell_volume=wfc_obj.cell_volume,
gcut=wfc_obj.gcutrho,
)
time_e_self = time.perf_counter() - start_time
time_e_self_str = f"{time_e_self:.2f}s"
else:
logging.info("Loading electron self-energy...")
time_e_self_str = "- Given by " + f"{mats.energy_e_self=}".split("=")[0] + " parameter"
time_ham = time.perf_counter() - start_time_ham
logging.info("Calculation of all matrix elements finished. Took %.2fs", time_ham)
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Hamiltonian calculation summary", color=block_color)
print(f"Number of plane waves: {p.shape[0]}")
print(f"Kinetic energy: {time_kin_str}")
print(f"Pseudopotentials: {time_pp_str}")
print(f"ERIs: {time_eri_str}")
if len(c_ip_core) > 0:
print(f"Frozen core potential and energy: {time_energy_core_str}")
if not (wfc_obj.ewald != 0.0 and config.qe_ewald):
print(f"Nuclear repulsion: {time_nucl_rep_str}")
print(f"Electron self-energy: {time_e_self_str}")
print(f"Total: {time_ham:.2f}s")
# print(f"Matrix elements saved to {filename}")
sys.stdout.flush()
######################### WANNIER TRANSFORMATION #########################
if config.wannier_transform:
# Read Wannier90 transform matrix
transform_matrix = dopyqo.read_u_mat(config.wannier_umat)[tuple(wfc_obj.kpoint)]
if not dopyqo.wannier90.check_unitarity_u(transform_matrix):
print(f"{RED}Wannier error: Transformation matrix is not unitary{RESET_COLOR}")
sys.exit(1)
if config.active_orbitals != transform_matrix.shape[0]:
print(
f"{RED}Wannier error: Number of active orbitals ({config.active_orbitals}) "
+ f"does not match number of Wannier orbitals ({transform_matrix.shape[0]})!{RESET_COLOR}"
)
sys.exit(1)
######################### READING WANNIER INPUT FILE #########################
if config.wannier_input_file is not None:
results = {}
keys = ["num_wann", "exclude_bands"]
with open(config.wannier_input_file, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
for key in keys:
if line.startswith(key):
if "=" in line:
_, sep, val = line.partition("=")
elif ":" in line:
_, sep, val = line.partition(":")
else: # whitespace split
parts = line.split(None, 1)
val = parts[1] if len(parts) == 2 else None
results[key] = val.strip()
break
num_wann = int(results["num_wann"])
exclude_bands = []
if "exclude_bands" in results:
for x in results["exclude_bands"].split(","):
if "-" not in x:
exclude_bands.append(int(x))
else:
start_stop = [int(val) for val in x.split("-")]
start = start_stop[0]
stop = start_stop[1] + 1
exclude_bands.extend(list(range(start, stop)))
exclude_bands = [x - 1 for x in exclude_bands] # Wannier90 1-indexing to python 0-indexing
orbital_indices_active_wannier = list(set(exclude_bands) ^ set(range(wfc_obj.nbnd)))
if orbital_indices_active != orbital_indices_active_wannier:
print(
f"{RED}Wannier error: Orbitals in the active space ({orbital_indices_active}) are not the "
+ f"same orbitals transformed with Wannier90 ({orbital_indices_active_wannier})!{RESET_COLOR}"
)
sys.exit(1)
######################### CALCULATE OCCUPATIONS AND TRANSFORM MATRICES #########################
# density_mat = dopyqo.to_density_matrix(config.active_electrons, config.active_orbitals) / 2.0
density_mat = np.diag(occupations_active_xml)
# density_mat = np.diag(occupations_active)
wf_density_matrix = transform_matrix.T.conj() @ density_mat @ transform_matrix
if not np.allclose(wf_density_matrix, wf_density_matrix.real):
print(f"{ORANGE}Wannier warning: Calculated density matrix in the Wannier basis is not real!{RESET_COLOR}")
else:
wf_density_matrix = wf_density_matrix.real
occupations_new = wf_density_matrix.real.diagonal()
make_binary = False
for occ in occupations_new:
if (not np.isclose(occ, 0.0)) and (not np.isclose(occ, 1.0)):
make_binary = True
break
occupations_new_binary = occupations_new.copy().round().astype(np.int64)
if make_binary:
occupations_new_binary = np.concatenate(
(
np.ones((round(np.sum(occupations_new)),), dtype=np.int64),
np.zeros(
(len(occupations_new) - round(np.sum(occupations_new)),),
dtype=np.int64,
),
),
dtype=np.int64,
)
if not np.isclose(sum(occupations_new) * 2, config.active_electrons) or not np.isclose(
sum(occupations_new_binary) * 2, config.active_electrons
):
print(
f"{RED}Wannier error: Sum of Wannier occupations ({sum(occupations_new) * 2} or {sum(occupations_new_binary) * 2}) does not match number of active electrons ({config.active_electrons}). "
+ f"This should not happen. Please, contact a developer!{RESET_COLOR}"
)
sys.exit(1)
overlaps_wannier = np.einsum("mq, np, nm -> pq", transform_matrix, transform_matrix.conj(), wfc_obj.get_overlaps(orbital_indices_active))
if not np.allclose(overlaps_wannier, np.eye(overlaps_wannier.shape[0])):
print(f"{ORANGE}Wannier warning: The wannier wavefunctions are not orthonormal!{RESET_COLOR}")
# Calculate Hamiltonian and eigenenergies in Wannier basis
ks_energies = wfc_obj.ks_energies[orbital_indices_active]
wannier_h_mat = np.einsum(
"mp, mq, m -> pq",
transform_matrix.conj(),
transform_matrix,
ks_energies,
)
wannier_energies, wannier_eigvecs = np.linalg.eig(wannier_h_mat)
if not np.allclose(np.sort(wannier_energies), np.sort(ks_energies)):
print(f"{ORANGE}Wannier warning: The wannier energies do not coincide with the KS-energies!{RESET_COLOR}")
d = np.linalg.inv(wannier_eigvecs) @ wannier_h_mat @ wannier_eigvecs
if not np.allclose(d.diagonal(), wannier_energies):
print(
f"{RED}Wannier error: Eigenvectors of DFT hamiltonian in Wannier basis do not "
+ f"diagonalize DFT hamiltonian in Wannier basis! This should never happen. Please, contact a developer{RESET_COLOR}"
)
sys.exit(1)
h_pq_wannier = dopyqo.transform_one_body_matrix(h_pq, transform_matrix)
h_pqrs_wannier = dopyqo.transform_two_body_matrix(h_pqrs, transform_matrix)
# h_pq_ks = h_pq.copy()
# h_pqrs_ks = h_pqrs.copy()
h_pq = h_pq_wannier
h_pqrs = h_pqrs_wannier
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Wannier transformation summary", color=block_color, flush=True)
print(f"Kohn-Sham active-space occupations: {occupations_active}")
print(f"Kohn-Sham xml occupations: {occupations_active_xml}")
str_tmp = "Kohn-Sham density matrix: "
ks_density_matrix_str = np.array2string(
density_mat, prefix=str_tmp, max_line_width=1000, precision=3, suppress_small=True, separator=" "
)
print(str_tmp + ks_density_matrix_str)
str_tmp = (
""
if not make_binary
else f" (non-binary: {np.array2string(occupations_new, max_line_width=1000, precision=3, suppress_small=True, separator=' ')})"
)
print(f"Wannier occupations: {occupations_new_binary}{str_tmp}")
str_tmp = "Wannier density matrix: "
wf_density_matrix_str = np.array2string(
wf_density_matrix, prefix=str_tmp, max_line_width=1000, precision=3, suppress_small=True, separator=" "
)
print(str_tmp + wf_density_matrix_str)
sys.stdout.flush()
# IDEA: Wannier VQE initialization based on density matrix using approriate single- and double exications?
occupations_active = np.array(occupations_new_binary)
if config.occupations is not None:
print(f"Using custom occupations: {config.occupations}")
occupations_active = config.occupations
######################### SAVE MATRIX ELEMENTS #########################
filename = f"matrix_elements_ks_{config.active_electrons}e_{config.active_orbitals}o_kpoint{wfc_obj.kpoint}.npz"
# TODO: Save MatrixElements and Wfc object together with energy_dict instead of saving indiviual matrices separately. Do it like this:
# filename = f"energydict_wfcobj_mats_{config.active_electrons}e_{config.active_orbitals}o.npz"
# np.savez(
# os.path.join(config.base_folder, filename),
# active_electrons=config.active_electrons,
# active_orbitals=config.active_orbitals,
# energy_dict=energy_dict,
# wfc_obj_ks=wfc_obj,
# mats=mats,
# )
np.savez(
os.path.join(config.base_folder, filename),
active_electrons=config.active_electrons,
active_orbitals=config.active_orbitals,
h_pq_kin=iTj_orbitals,
h_pq_pp=pp_orbitals,
h_pq_core=h_pq_core,
h_pqrs=h_pqrs,
h_pq_ewald=h_pq_ewald,
energy_frozen_core=energy_frozen_core,
energy_ewald=energy_ewald,
energy_e_self=energy_e_self,
transform_matrix=None if not config.wannier_transform else transform_matrix,
mats=dopyqo.MatrixElements(
h_pq_kin=iTj_orbitals,
h_pq_pp=pp_orbitals,
h_pq_core=h_pq_core,
h_pqrs=h_pqrs,
h_pq_ewald=h_pq_ewald,
energy_frozen_core=energy_frozen_core,
energy_ewald=energy_ewald,
energy_e_self=energy_e_self,
transform_matrix=None if not config.wannier_transform else transform_matrix,
),
)
######################### CREATE HAMILTONIAN OBJECT #########################
h_ks = dopyqo.Hamiltonian(
h_pq.copy(),
h_pqrs.copy(),
occupations_active.copy(),
reference_energy=reference_energy,
constants=energy_frozen_core + energy_e_self + energy_ewald,
)
######################### RUN PySCF HF CALCULATION #########################
if config.run_hf:
hf_pyscf_energy = h_ks.solve_hf()
mo_coeff = h_ks.mo_coeff
#
h_ks = dopyqo.Hamiltonian(
dopyqo.transform_one_body_matrix(h_pq.copy(), mo_coeff),
dopyqo.transform_two_body_matrix(h_pqrs.copy(), mo_coeff),
occupations_active.copy(),
reference_energy=reference_energy,
constants=energy_frozen_core + energy_e_self + energy_ewald,
)
######################### RUN FCI CALCULATION #########################
fci_state = Statevector(np.array([]))
if config.run_fci:
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Running FCI", color=block_color, flush=True)
start_time_fci = time.perf_counter()
if wfc_obj.gamma_only:
fci_energy = h_ks.solve_fci_spin1(n_energies=config.n_fci_energies).real
if config.n_fci_energies == 1:
fci_energy = fci_energy[0]
fci_state = h_ks.fci_statevector()
else:
fci_energy = h_ks.solve_fci(n_energies=config.n_fci_energies).real
if config.n_fci_energies == 1:
fci_energy = fci_energy[0]
time_fci = time.perf_counter() - start_time_fci
h_ks.reference_energy = fci_energy
print(f"FCI energy: {fci_energy}")
if verbosity >= verbosity_summary:
print(f"\tTook {time_fci:.2f}s")
# TODO: Generalize fci_statevector to support statevector generated from solve_fci_general
if wfc_obj.gamma_only:
fci_state = h_ks.fci_statevector()
else:
print(
f"{ORANGE}FCI warning: The FCI statevector will not be saved. The FCI statevector cannot be calculated for a non-gamma-only calculation, yet!{RESET_COLOR}"
)
######################### RUN VQE CALCULATION #########################
vqe_state = Statevector(np.array([]))
if config.run_vqe:
######################### VQE #########################
if not config.vqe_adapt:
if config.vqe_parameters is not None:
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Executing VQE circuit with given parameters", color=block_color, flush=True)
start_time_vqe = time.perf_counter()
if config.use_qiskit:
vqe_energy = h_ks.run_qiskit_ansatz(config.vqe_parameters, config.uccsd_reps)
else:
vqe_energy = h_ks.run_tcc_ansatz(
config.vqe_parameters, config.uccsd_reps, excitations=config.vqe_excitations, qiskit_equivalent=False
)
time_vqe = time.perf_counter() - start_time_vqe
if verbosity >= verbosity_summary:
print(f"\nVQE energy: {vqe_energy}")
print(f"\tTook {time_vqe:.2f}s")
else:
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Running VQE", color=block_color, flush=True)
if config.use_qiskit: # Qiskit
if config.vqe_initial_parameters is not None:
print(
f"{RED}Config error: Setting vqe_initial_parameters is not supported when using qiskit!{RESET_COLOR}", file=sys.stderr
)
sys.exit(1)
start_time_vqe = time.perf_counter()
h_ks.solve_vqe_qiskit(UCCSD_reps=config.uccsd_reps, optimizer=config.vqe_optimizer)
time_vqe = time.perf_counter() - start_time_vqe
else: # TenCirChem
optimizer = config.vqe_optimizer
start_time_vqe = time.perf_counter()
h_ks.solve_vqe(
UCCSD_reps=config.uccsd_reps,
optimizer=optimizer,
initial_params=config.vqe_initial_parameters,
maxiter=config.vqe_maxiter,
excitations=config.vqe_excitations,
)
time_vqe = time.perf_counter() - start_time_vqe
vqe_energy = h_ks.vqe_values[-1]
if verbosity >= verbosity_summary:
print(f"\nVQE energy: {vqe_energy}")
print(f"\tTook {time_vqe:.2f}s")
print(f"VQE Result:\n{h_ks.vqe_result}")
vqe_state = h_ks.vqe_statevector_qiskit() if config.use_qiskit else h_ks.vqe_statevector()
if verbosity >= verbosity_summary:
print(f"VQE Statevector: {vqe_state.draw('latex_source')}")
if verbosity >= verbosity_summary:
fig = tplt.figure()
if config.run_fci:
print()
dopyqo.print_block("VQE convergence plot", color=block_color, flush=True)
print("Y-axis shows the absolute difference to the FCI energy and is logarithmic (x -> 10^x).")
abs_diff = np.abs(fci_energy - h_ks.vqe_values)
if np.any(abs_diff < 1e-15):
abs_diff[abs_diff < 1e-15] = 1e-15 # np.log10 raises "divide by zero" error if values are 0.0
fig.plot(
h_ks.vqe_counts,
np.log10(abs_diff),
)
else:
print("VQE convergence plot below. Y-axis shows VQE energy.")
fig.plot(
h_ks.vqe_counts,
h_ks.vqe_values,
)
fig.show()
######################### ADAPT-VQE #########################
else:
if verbosity >= verbosity_summary:
print()
dopyqo.print_block("Running VQE", color=block_color, flush=True)
if config.use_qiskit:
print(
f"{RED}Dopyqo error: Both "
+ f"{config.vqe_adapt=}".split("=")[0]
+ f" and "
+ f"{config.use_qiskit=}".split("=")[0]
+ f" are set to True. ADAPT-VQE is not supported when using Qiskit!{RESET_COLOR}",
file=sys.stderr,
)
sys.exit(1)
else:
optimizer = config.vqe_optimizer
start_time_vqe = time.perf_counter()
h_ks.solve_vqe_adapt(
optimizer=optimizer,
maxiter=config.vqe_maxiter,
excitation_pool=config.vqe_excitations,
drain_pool=config.vqe_adapt_drain_pool,
)
time_vqe = time.perf_counter() - start_time_vqe
vqe_energy = h_ks.vqe_values[-1]
if verbosity >= verbosity_summary:
print(f"\nVQE energy: {vqe_energy}")
print(f"\tTook {time_vqe:.2f}s")
print(f"VQE Result:\n{h_ks.vqe_result}")
if config.use_qiskit:
print(f"{RED}Qiskit cannot be used with ADAPT-VQE, yet. Cannot calculate statevector!{RESET_COLOR}")
sys.exit(1)
# vqe_state = h_ks.vqe_statevector_qiskit() if config.use_qiskit else h_ks.vqe_statevector()
vqe_state = h_ks.vqe_statevector()
if verbosity >= verbosity_summary:
print(f"VQE Statevector: {vqe_state.draw('latex_source')}")
if verbosity >= verbosity_summary:
fig = tplt.figure()
if config.run_fci:
print()
dopyqo.print_block("VQE convergence plot", color=block_color, flush=True)
print("Y-axis shows the absolute difference to the FCI energy and is logarithmic (x -> 10^x).")
abs_diff = np.abs(fci_energy - h_ks.vqe_values)
if np.any(abs_diff < 1e-15):
abs_diff[abs_diff < 1e-15] = 1e-15 # np.log10 raises "divide by zero" error if values are 0.0
fig.plot(
h_ks.vqe_counts,
np.log10(abs_diff),
)
else:
print("VQE convergence plot below. Y-axis shows VQE energy.")
fig.plot(
h_ks.vqe_counts,
h_ks.vqe_values,
)
fig.show()
######################### SAVE STATEVECTORS AND APPENDING RESULTS #########################
if config.run_fci or config.run_vqe:
logging.info("Saving statevectors...")
filename = f"statevecs_{config.active_electrons}e_{config.active_orbitals}o.npz"
np.savez(os.path.join(config.base_folder, filename), fci_state_data=fci_state.data, vqe_state_data=vqe_state.data)
if verbosity >= verbosity_summary:
print(f"Statevectors saved to {filename}")
# IDEA: If error occurs after matrix calculation, do not exit with sys.exit(1) but stop and return
# everything, i.e., wfc objects, matrix elements for every processed k-point
# return success flag to signal if error occured?
hf_energy_lst.append(h_ks.hf_energy())
if config.run_hf:
hf_energy_pyscf_lst.append(hf_pyscf_energy)
if config.run_fci:
fci_energy_lst.append(fci_energy)
if config.run_vqe:
vqe_energy_lst.append(vqe_energy)
if return_h:
h_ks_lst.append(h_ks)
vqe_result_lst.append(h_ks.vqe_result)
if return_mats:
mats_lst.append(
dopyqo.MatrixElements(
h_pq_kin=iTj_orbitals,
h_pq_pp=pp_orbitals,
h_pq_core=h_pq_core,
h_pqrs=h_pqrs,
h_pq_ewald=h_pq_ewald,
energy_frozen_core=energy_frozen_core,
energy_ewald=energy_ewald,
energy_e_self=energy_e_self,
transform_matrix=None if not config.wannier_transform else transform_matrix,
)
)
kpoint_weights.append(wfc_obj.kpoint_weight)
if verbosity >= verbosity_summary:
print()
msg_tmp = "Energy summary"
if len(wfc_obj_lst) > 1:
msg_tmp += f"\n(weighted over {len(wfc_obj_lst)} k-points)"
dopyqo.print_block(msg_tmp, color=block_color, flush=True)
print(f"DFT energy: {wfc_obj_lst[0].etot}")
######################### WEIGHT ENERGIES #########################
# Weighted HF energies
hf_energy = 0.0
for hf_nrg, weight in zip(hf_energy_lst, kpoint_weights):
hf_energy += weight * hf_nrg.real
hf_energy /= sum(kpoint_weights)
hf_energy_pyscf = 0.0
if config.run_hf:
for hf_nrg, weight in zip(hf_energy_pyscf_lst, kpoint_weights):
hf_energy_pyscf += weight * hf_nrg.real
hf_energy_pyscf /= sum(kpoint_weights)
if verbosity >= verbosity_summary:
print(f"HF energy (PySCF): {hf_energy_pyscf}")
if verbosity >= verbosity_summary:
print(f'"HF" energy: {hf_energy}')
energy_dict = {"dft_energy": wfc_obj_lst[0].etot, "hf_energy": hf_energy, "hf_energy_pyscf": hf_energy_pyscf}
if len(wfc_obj_lst) > 1:
energy_dict["hf_energy_per_kpoint"] = hf_energy_lst
energy_dict["hf_energy_pyscf_per_kpoint"] = hf_energy_pyscf_lst
if config.run_fci:
# Weighted FCI energies
fci_energy = 0.0
for fci_nrg, weight in zip(fci_energy_lst, kpoint_weights):
fci_energy += weight * fci_nrg.real
fci_energy /= sum(kpoint_weights)
if verbosity >= verbosity_summary:
print(f"FCI energy: {fci_energy}")
print(f"\tTook {time_fci:.2f}s") # TODO: sum times for multiple kpoints
energy_dict["fci_energy"] = fci_energy
if len(wfc_obj_lst) > 1:
energy_dict["fci_energy_per_kpoint"] = fci_energy_lst
if config.run_vqe:
# Weighted VQE energies
vqe_energy = 0.0
for vqe_nrg, weight in zip(vqe_energy_lst, kpoint_weights):
vqe_energy += weight * vqe_nrg.real
vqe_energy /= sum(kpoint_weights)
energy_dict["vqe_energy"] = vqe_energy
if len(wfc_obj_lst) > 1:
energy_dict["vqe_energy_per_kpoint"] = vqe_energy_lst
if config.vqe_parameters is None:
if len(wfc_obj_lst) > 1:
energy_dict["vqe_result"] = vqe_result_lst
else:
energy_dict["vqe_result"] = vqe_result_lst[0]
if verbosity >= verbosity_summary:
if config.run_fci:
print(f"VQE energy: {vqe_energy} | Diff. to FCI: {np.abs(vqe_energy-fci_energy)}")
else:
print(f"VQE energy: {vqe_energy}")
print(f"\tTook {time_vqe:.2f}s")
sys.stdout.flush()
filename = f"energies_{config.active_electrons}e_{config.active_orbitals}o.npz"
np.savez(os.path.join(config.base_folder, filename), **energy_dict)
if verbosity >= verbosity_summary:
print(f"Energies saved to {filename}")
time_all = time.perf_counter() - start_time_all
print(f"{GREEN}All done in {time_all:.2f}s! Exiting...{RESET_COLOR}", flush=True)
return_tuple = (energy_dict,)
if return_wfc:
return_tuple += (wfc_obj_lst[0] if len(wfc_obj_lst) == 1 else wfc_obj_lst,)
else:
return_tuple += (None,)
if return_h:
return_tuple += (h_ks_lst[0] if len(wfc_obj_lst) == 1 else h_ks_lst,)
else:
return_tuple += (None,)
if return_mats:
return_tuple += (mats_lst[0] if len(wfc_obj_lst) == 1 else mats_lst,)
else:
return_tuple += (None,)
return return_tuple
def main():
######################### BANNER AND ARGUMENT PARSING #########################
banner_attributes = vars(banners)
banner_strings = [value for name, value in banner_attributes.items() if isinstance(value, str) and not name.startswith("_")]
# random_banner = banner_strings[np.random.randint(0, len(banner_strings) + 1)]
rng = np.random.default_rng()
random_banner = rng.choice(banner_strings)
dopyqo.print_banner(random_banner)
parser = argparse.ArgumentParser(
epilog="Either provide an input file or at least the number of electrons and number of active orbitals.",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("-i", "--input_file", type=str, default=argparse.SUPPRESS, help="Input file")
parser.add_argument("-e", "--active_electrons", type=int, default=argparse.SUPPRESS, help="Number of electrons in the active space")
parser.add_argument("-o", "--active_orbitals", type=int, default=argparse.SUPPRESS, help="Number of spatial orbitals in the active space")
parser.add_argument("-fci", required=False, action="store_true", default=argparse.SUPPRESS, help="Running a FCI calculation")
parser.add_argument("-vqe", required=False, action="store_true", default=argparse.SUPPRESS, help="Running a VQE calculation")
parser.add_argument("-uccsd_reps", type=int, default=argparse.SUPPRESS, help="Number of repetitions of the UCCSD in the VQE circuit")
parser.add_argument("-vqe_opt", type=str, default=argparse.SUPPRESS, help="Optimizer used in the VQE")
parser.add_argument("-qiskit", required=False, action="store_true", default=argparse.SUPPRESS, help="Using qiskit instead of TenCirChem")
parser.add_argument("-logging", required=False, action="store_true", default=argparse.SUPPRESS, help="Show logging")
args = parser.parse_args()
n_given_args = len(args.__dict__)
if "input_file" not in args and "i" not in args:
args.input_file = None
if "active_electrons" not in args and "e" not in args:
args.active_electrons = None
if "active_orbitals" not in args and "o" not in args:
args.active_orbitals = None
if "fci" not in args:
args.fci = False
if "vqe" not in args:
args.vqe = False
if "uccsd_reps" not in args:
args.uccsd_reps = 1
if "vqe_opt" not in args:
args.vqe_opt = "L-BFGS-B"
if "qiskit" not in args:
args.qiskit = False
if "logging" not in args:
args.logging = False
#
if args.input_file is None and (args.active_electrons is None or args.active_orbitals is None):
parser.print_help(sys.stderr)
sys.exit(1)
######################### READING INPUT FILE #########################
if args.input_file is not None:
if n_given_args != 1:
argument_s_str = "arguments were" if n_given_args - 1 > 1 else "argument was"
print(
f"{RED}Error in arguments: If input file is given no other arguments are allowed but {n_given_args-1} other {argument_s_str} given!{RESET_COLOR}",
file=sys.stderr,
)
sys.exit(1)
input_filename = args.input_file
config_dict = {}
print(f"Reading {input_filename} input file...")
config_dict = dopyqo.read_input_toml(input_filename)
if config_dict.get("base_folder") is None:
base_folder = os.path.join(*os.path.split(input_filename)[:-1]) # os.getcwd()
base_folder_tmp = "./" if len(base_folder) == 0 else base_folder
print(f"base_folder not given in input file, defaulting to folder of input file: {base_folder_tmp}")
config_dict["base_folder"] = base_folder
if config_dict.get("prefix") is None:
base_folder_tmp = "./" if base_folder is None or len(base_folder) == 0 else base_folder
print(f"prefix not given in input file. Searching for '.save' folder in base_folder ({base_folder_tmp})...")
prefix = has_save_folder(base_folder)
if prefix is None:
print(
f"{RED}Error finding files: Could not find a folder ending with '.save' in base_folder ({base_folder_tmp}){RESET_COLOR}",
file=sys.stderr,
)
sys.exit(1)
print(f"prefix set to: {prefix}")
config_dict["prefix"] = prefix
config = DopyqoConfig(**config_dict)
######################### READING CLI ARGUMENTS #########################
else:
base_folder = os.getcwd()
prefix = has_save_folder(base_folder)
if prefix is None:
print(
f"{RED}Error finding files: Could not find a folder ending with '.save' in current folder ({base_folder}){RESET_COLOR}",
file=sys.stderr,
)
sys.exit(1)
# parser = argparse.ArgumentParser(description="")
# parser.add_argument("-e", "--active_electrons", type=int, required=True, help="Number of electrons in the active space")
# parser.add_argument("-o", "--active_orbitals", type=int, required=True, help="Number of spatial orbitals in the active space")
# parser.add_argument("-fci", required=False, action="store_true", help="Running a FCI calculation")
# parser.add_argument("-vqe", required=False, action="store_true", help="Running a VQE calculation")
# parser.add_argument("-uccsd_reps", type=int, default=1, required=False, help="Number of repetitions of the UCCSD in the VQE circuit")
# parser.add_argument("-vqe_opt", type=str, default="L-BFGS-B", required=False, help="Optimizer used in the VQE")
# parser.add_argument("-qiskit", action="store_true", required=False, help="Using qiskit instead of TenCirChem")
# parser.add_argument("-logging", action="store_true", required=False, help="Show logging")
# args = parser.parse_args()
active_electrons = args.active_electrons
active_orbitals = args.active_orbitals
run_fci_flag = args.fci
run_vqe_flag = args.vqe
uccsd_reps = args.uccsd_reps
vqe_opt = args.vqe_opt
use_qiskit_flag = args.qiskit
logging_flag = args.logging
config = DopyqoConfig(
base_folder,
prefix,
active_electrons,
active_orbitals,
logging_flag,
run_vqe_flag,
run_fci_flag,
use_qiskit_flag,
vqe_optimizer=vqe_opt,
uccsd_reps=uccsd_reps,
)
# print(f"asdict(Config(**config_dict)): {asdict(DopyqoConfig(**config_dict))}")
# print(f"config_dict: {config_dict}")
_ = run(config, show_banner=False)