|
- import torchquantum as tq
- import torch
- from torchquantum.vqe_utils import parse_hamiltonian_file
- import random
- import numpy as np
- import argparse
- import torch.optim as optim
-
- from torch.optim.lr_scheduler import CosineAnnealingLR
- from torchquantum.measurement import expval_joint_analytical
-
-
- class QVQEModel(tq.QuantumModule):
- def __init__(self, arch, hamil_info):
- super().__init__()
- self.arch = arch
- self.hamil_info = hamil_info
- self.n_wires = hamil_info["n_wires"]
- self.n_blocks = arch["n_blocks"]
- self.u3_layers = tq.QuantumModuleList()
- self.cu3_layers = tq.QuantumModuleList()
- for _ in range(self.n_blocks):
- self.u3_layers.append(
- tq.Op1QAllLayer(
- op=tq.U3,
- n_wires=self.n_wires,
- has_params=True,
- trainable=True,
- )
- )
- self.cu3_layers.append(
- tq.Op2QAllLayer(
- op=tq.CU3,
- n_wires=self.n_wires,
- has_params=True,
- trainable=True,
- circular=True,
- )
- )
-
- def forward(self):
- qdev = tq.QuantumDevice(
- n_wires=self.n_wires, bsz=1, device=next(self.parameters()).device
- )
-
- for k in range(self.n_blocks):
- self.u3_layers[k](qdev)
- self.cu3_layers[k](qdev)
-
- expval = 0
- for hamil in self.hamil_info["hamil_list"]:
- expval += (
- expval_joint_analytical(qdev, observable=hamil["pauli_string"])
- * hamil["coeff"]
- )
-
- return expval
-
-
- def train(model, optimizer, n_steps=1):
- for _ in range(n_steps):
- loss = model()
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
- print(f"Expectation of energy: {loss.item()}")
-
-
- def valid_test(model):
- with torch.no_grad():
- loss = model()
-
- print(f"validation: expectation of energy: {loss.item()}")
-
-
- def process_hamil_info(hamil_info):
- hamil_list = hamil_info["hamil_list"]
- n_wires = hamil_info["n_wires"]
- all_info = []
-
- for hamil in hamil_list:
- pauli_string = ""
- for i in range(n_wires):
- if i in hamil["wires"]:
- wire = hamil["wires"].index(i)
- pauli_string += hamil["observables"][wire].upper()
- else:
- pauli_string += "I"
- all_info.append({"pauli_string": pauli_string, "coeff": hamil["coefficient"]})
- hamil_info["hamil_list"] = all_info
- return hamil_info
-
-
- def main():
- parser = argparse.ArgumentParser()
- parser.add_argument("--pdb", action="store_true", help="debug with pdb")
- parser.add_argument(
- "--n_blocks",
- type=int,
- default=2,
- help="number of blocks, each contain one layer of "
- "U3 gates and one layer of CU3 with "
- "ring connections",
- )
- parser.add_argument(
- "--steps_per_epoch", type=int, default=10, help="number of training epochs"
- )
- parser.add_argument(
- "--epochs", type=int, default=100, help="number of training epochs"
- )
- parser.add_argument(
- "--hamil_filename",
- type=str,
- default="./h2.txt",
- help="number of training epochs",
- )
-
- args = parser.parse_args()
-
- if args.pdb:
- import pdb
-
- pdb.set_trace()
-
- seed = 0
- random.seed(seed)
- np.random.seed(seed)
- torch.manual_seed(seed)
-
- hamil_info = process_hamil_info(parse_hamiltonian_file(args.hamil_filename))
-
- use_cuda = torch.cuda.is_available()
- device = torch.device("cuda" if use_cuda else "cpu")
- model = QVQEModel(arch={"n_blocks": args.n_blocks}, hamil_info=hamil_info)
-
- model.to(device)
-
- n_epochs = args.epochs
- optimizer = optim.Adam(model.parameters(), lr=5e-3, weight_decay=1e-4)
- scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs)
-
- for epoch in range(1, n_epochs + 1):
- # train
- print(f"Epoch {epoch}, LR: {optimizer.param_groups[0]['lr']}")
- train(model, optimizer, n_steps=args.steps_per_epoch)
-
- scheduler.step()
-
- # final valid
- valid_test(model)
-
-
- if __name__ == "__main__":
- main()
|