# ground_state_search.py
import argparse
import itertools
import os
from pathlib import Path
import numpy as np
import pandas as pd
import yaml
from dotmap import DotMap
from ttnopt.hamiltonian import hamiltonian
from ttnopt.src import GroundStateSearch, GroundStateSearchSparse, TreeTensorNetwork
[docs]def ground_state_search():
parser = argparse.ArgumentParser(description="Ground state search simulation")
parser.add_argument(
"config_file", type=str, help="Path to the YAML configuration file"
)
args = parser.parse_args()
with open(args.config_file, "r") as file:
config = yaml.safe_load(file)
config = DotMap(config)
psi = TreeTensorNetwork.mps(config.system.N)
if config.numerics.init_tree == 1:
if config.system.N > 0 and (config.system.N & (config.system.N - 1)) == 0:
psi = TreeTensorNetwork.tree(config.system.N)
else:
print("=" * 50)
print("⚠️ Note: N is not a power of 2.")
print(" Using MPS structure as the default.")
print("=" * 50)
ham = hamiltonian(config.system)
numerics = config.numerics
path = (
Path(config.output.dir)
if not isinstance(config.output.dir, DotMap)
else Path("./")
)
os.makedirs(path, exist_ok=True)
save_onesite_expval = (
config.output.single_site
if not isinstance(config.output.single_site, DotMap)
else False
)
save_twosite_expval = (
config.output.two_site
if not isinstance(config.output.two_site, DotMap)
else False
)
u1_symmetry = True if not isinstance(numerics.U1_symmetry, DotMap) else False
if u1_symmetry and config.model.type == "XYZ":
print("=" * 50)
print("⚠️ Error: U1 symmetry is not supported for the XYZ model.")
print(" Please set the U1 symmetry to False or change the model to XXZ.")
print("=" * 50)
exit()
if u1_symmetry and not (isinstance(config.system.MF_X, DotMap)):
print("=" * 50)
print("⚠️ Error: U1 symmetry is not allowed for the X magnetic field.")
print("=" * 50)
exit()
if u1_symmetry and not (isinstance(config.system.MF_Y, DotMap)):
print("=" * 50)
print("⚠️ Error: U1 symmetry is not allowed for the Y magnetic field.")
print("=" * 50)
exit()
if u1_symmetry and not isinstance(config.system.DM_X, DotMap):
print("=" * 50)
print(
"⚠️ Error: U1 symmetry is not allowed for the X term of Dzyaloshinskii-Moriya interaction."
)
print("=" * 50)
exit()
if u1_symmetry and not isinstance(config.system.DM_Y, DotMap):
print("=" * 50)
print(
"⚠️ Error: U1 symmetry is not allowed for the Y term of Dzyaloshinskii-Moriya interaction."
)
print("=" * 50)
exit()
if u1_symmetry and not isinstance(config.system.SOD_X, DotMap):
print("=" * 50)
print(
"⚠️ Error: U1 symmetry is not allowed for the X term of symmetric exchange anisotropy interaction."
)
print("=" * 50)
exit()
if u1_symmetry and not isinstance(config.system.SOD_Y, DotMap):
print("=" * 50)
print(
"⚠️ Error: U1 symmetry is not allowed for the Y term of symmetric exchange anisotropy interaction."
)
print("=" * 50)
exit()
if u1_symmetry and not isinstance(config.system.SOD_Z, DotMap):
print("=" * 50)
print(
"⚠️ Error: U1 symmetry is not allowed for the Z term of symmetric exchange anisotropy interaction."
)
print("=" * 50)
exit()
opt_structure = (
numerics.opt_structure.type
if (isinstance(numerics.opt_structure.type, int))
else 0
)
temperature = (
float(numerics.opt_structure.temperature)
if (
isinstance(numerics.opt_structure.temperature, float)
or isinstance(numerics.opt_structure.temperature, int)
)
else 0.0
)
seed = (
numerics.opt_structure.seed
if isinstance(numerics.opt_structure.seed, int)
else 0
)
if isinstance(numerics.energy_degeneracy_threshold, DotMap):
energy_degeneracy_threshold = 1.0e-8
else:
energy_degeneracy_threshold = float(numerics.energy_degeneracy_threshold)
if isinstance(numerics.entanglement_degeneracy_threshold, DotMap):
entanglement_degeneracy_threshold = 1.0e-8
else:
entanglement_degeneracy_threshold = float(
numerics.entanglement_degeneracy_threshold
)
if isinstance(numerics.energy_convergence_threshold, DotMap):
energy_convergence_threshold = 1.0e-8
else:
energy_convergence_threshold = float(numerics.energy_convergence_threshold)
if isinstance(numerics.entanglement_convergence_threshold, DotMap):
entanglement_convergence_threshold = 1.0e-8
else:
entanglement_convergence_threshold = float(
numerics.entanglement_convergence_threshold
)
np.random.seed(seed)
tau = (
numerics.opt_structure.tau
if isinstance(numerics.opt_structure.tau, int)
else numerics.max_bond_dimensions[0] // 2
)
if isinstance(numerics.energy_degeneracy_threshold, DotMap):
energy_degeneracy_threshold = 1.0e-8
else:
energy_degeneracy_threshold = float(numerics.energy_degeneracy_threshold)
if isinstance(numerics.entanglement_degeneracy_threshold, DotMap):
entanglement_degeneracy_threshold = 1.0e-8
else:
entanglement_degeneracy_threshold = float(
numerics.entanglement_degeneracy_threshold
)
if isinstance(numerics.energy_convergence_threshold, DotMap):
energy_convergence_threshold = 1.0e-8
else:
energy_convergence_threshold = float(numerics.energy_convergence_threshold)
if isinstance(numerics.entanglement_convergence_threshold, DotMap):
entanglement_convergence_threshold = 1.0e-8
else:
entanglement_convergence_threshold = float(
numerics.entanglement_convergence_threshold
)
if u1_symmetry:
gss = GroundStateSearchSparse(
psi,
ham,
numerics.U1_symmetry,
init_bond_dim=numerics.initial_bond_dimension,
max_bond_dim=numerics.max_bond_dimensions[0],
energy_degeneracy_threshold=energy_degeneracy_threshold,
entanglement_degeneracy_threshold=entanglement_degeneracy_threshold,
)
if gss.init_u1_num != gss.u1_num:
sz_sign = np.sign(gss.u1_num - gss.init_u1_num)
sz_diff = abs(gss.u1_num - gss.init_u1_num) // 2
gss.run(
opt_structure=0,
max_num_sweep=sz_diff,
energy_convergence_threshold=0.0,
entanglement_convergence_threshold=0.0,
sz_sign=sz_sign,
)
print(f"Complete the warmup sweeps for M={gss.u1_num}/2.")
else:
gss = GroundStateSearch(
psi,
ham,
init_bond_dim=numerics.initial_bond_dimension,
max_bond_dim=numerics.initial_bond_dimension,
energy_degeneracy_threshold=energy_degeneracy_threshold,
entanglement_degeneracy_threshold=entanglement_degeneracy_threshold,
)
for i, (max_bond_dim, max_num_sweep) in enumerate(
zip(numerics.max_bond_dimensions, numerics.max_num_sweeps)
):
gss.max_bond_dim = max_bond_dim
if i == 0:
gss.run(
opt_structure=opt_structure,
energy_convergence_threshold=energy_convergence_threshold,
entanglement_convergence_threshold=entanglement_convergence_threshold,
max_num_sweep=max_num_sweep,
temperature=temperature,
tau=tau,
)
print("Calculating the expectation values for the initial structure")
# re-run the first iteration to save the expectation values
gss.run(
opt_structure=0,
max_num_sweep=1,
eval_onesite_expval=save_onesite_expval,
eval_twosite_expval=save_twosite_expval,
)
else:
gss.run(
opt_structure=0,
energy_convergence_threshold=energy_convergence_threshold,
entanglement_convergence_threshold=entanglement_convergence_threshold,
max_num_sweep=max_num_sweep,
eval_onesite_expval=save_onesite_expval,
eval_twosite_expval=save_twosite_expval,
)
nodes_list = {}
for edge_id in gss.energy.keys():
tmp = []
for node_id, edges in enumerate(psi.edges):
node_id += config.system.N
if edge_id in edges:
tmp.append(node_id)
nodes_list[edge_id] = tmp
for edge_id in psi.physical_edges:
tmp = []
tmp.append(edge_id)
for node_id, edges in enumerate(psi.edges):
node_id += config.system.N
if edge_id in edges:
tmp.append(node_id)
nodes_list[edge_id] = tmp
gss.energy[edge_id] = 0.0
gss.error[edge_id] = 0.0
all_keys = set(nodes_list.keys())
df = pd.DataFrame(
[nodes_list[k] for k in all_keys],
columns=["node1", "node2"],
index=None,
)
df["energy"] = [gss.energy[k] for k in all_keys]
df["entanglement"] = [gss.entanglement[k] for k in all_keys]
df["error"] = [gss.error[k] for k in all_keys]
path_ = path / f"run{i + 1}"
os.makedirs(path_, exist_ok=True)
df.to_csv(path_ / "basic.csv", header=True, index=None)
np.savetxt(path_ / "graph.dat", gss.psi.edges, fmt="%d", delimiter=",")
if save_onesite_expval:
df = pd.DataFrame(psi.physical_edges, columns=["site"], index=None)
sp = np.zeros(len(psi.physical_edges))
sm = np.zeros(len(psi.physical_edges))
if not u1_symmetry:
sp = np.array(
[
gss.one_site_expval[edge_id]["S+"]
for edge_id in psi.physical_edges
]
)
sm = np.array(
[
gss.one_site_expval[edge_id]["S-"]
for edge_id in psi.physical_edges
]
)
df["Sx"] = np.real((sp + sm) / 2.0)
df["Sy"] = np.real((sp - sm) / 2.0j)
df["Sz"] = np.real(
[gss.one_site_expval[edge_id]["Sz"] for edge_id in psi.physical_edges]
)
path_ = path / f"run{i + 1}"
os.makedirs(path_, exist_ok=True)
df.to_csv(path_ / "single_site.csv", header=True, index=None)
if save_twosite_expval:
pairs = [(i, j) for i, j in itertools.combinations(psi.physical_edges, 2)]
df = pd.DataFrame(pairs, columns=["site1", "site2"], index=None)
spp = np.zeros(len(pairs))
smm = np.zeros(len(pairs))
szp = np.zeros(len(pairs))
spz = np.zeros(len(pairs))
szm = np.zeros(len(pairs))
smz = np.zeros(len(pairs))
if not u1_symmetry:
spp = np.array([gss.two_site_expval[pair]["S+S+"] for pair in pairs])
smm = np.array([gss.two_site_expval[pair]["S-S-"] for pair in pairs])
szp = np.array([gss.two_site_expval[pair]["SzS+"] for pair in pairs])
spz = np.array([gss.two_site_expval[pair]["S+Sz"] for pair in pairs])
szm = np.array([gss.two_site_expval[pair]["SzS-"] for pair in pairs])
smz = np.array([gss.two_site_expval[pair]["S-Sz"] for pair in pairs])
szz = np.array([gss.two_site_expval[pair]["SzSz"] for pair in pairs])
spm = np.array([gss.two_site_expval[pair]["S+S-"] for pair in pairs])
smp = np.array([gss.two_site_expval[pair]["S-S+"] for pair in pairs])
df["SxSx"] = np.real((spp + spm + smp + smm) / 4.0)
df["SySy"] = np.real((-spp + spm + smp - smm) / 4.0)
df["SzSz"] = np.real(szz)
df["SySz"] = np.real((spz - smz) / 2.0j)
df["SzSy"] = np.real((szp - szm) / 2.0j)
df["SzSx"] = np.real((szp + szm) / 2.0)
df["SxSz"] = np.real((spz + smz) / 2.0)
df["SxSy"] = np.real((spp - spm + smp - smm) / 4.0j)
df["SySx"] = np.real((spp + spm - smp - smm) / 4.0j)
path_ = path / f"run{i + 1}"
os.makedirs(path_, exist_ok=True)
df.to_csv(path_ / "two_site.csv", header=True, index=None)
return 0