DeOwl
96548e60ea
Changes: - added rabbitmq hearbeat, consumer and publisher - fixed vqe to run with rabbitmq - added frontend via fasthtml - added proper .env configuration
169 lines
5.1 KiB
Python
169 lines
5.1 KiB
Python
import os
|
|
from multiprocessing.connection import Connection
|
|
|
|
import jax
|
|
import pennylane as qml
|
|
import pennylane.numpy as np
|
|
from jax import numpy as jnp
|
|
from pennylane import qchem
|
|
from pennylane.devices import Device
|
|
from pennylane.optimize import GradientDescentOptimizer
|
|
|
|
jax.config.update("jax_enable_x64", True)
|
|
|
|
os.environ["OMP_NUM_THREADS"] = "16"
|
|
|
|
|
|
def parse_xyz_from_text(text: str):
|
|
"""Parse XYZ format from text content."""
|
|
lines = text.strip().split("\n")
|
|
|
|
# First line: number of atoms
|
|
num_atoms = int(lines[0].strip())
|
|
|
|
# Second line: Charge/Multiplicity/Electrons/Orbitals (optional)
|
|
# Skip or parse as needed
|
|
|
|
symbols = []
|
|
coordinates = []
|
|
|
|
# Parse atom lines (after the second line)
|
|
for line in lines[2 : 2 + num_atoms]:
|
|
parts = line.strip().split()
|
|
if len(parts) >= 4:
|
|
symbol = parts[0]
|
|
x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
|
|
symbols.append(symbol)
|
|
coordinates.append([x, y, z])
|
|
|
|
return symbols, coordinates
|
|
|
|
|
|
def extract_electron_info(text: str):
|
|
"""Extract electron and orbital counts from the second line."""
|
|
lines = text.strip().split("\n")
|
|
if len(lines) >= 2:
|
|
second_line = lines[1]
|
|
# Parse "Charge=0 Multiplicity=1 Electrons=3 Orbitals=3"
|
|
electrons = 3 # default
|
|
orbitals = 3 # default
|
|
charge = 0 # default
|
|
multiplicity = 1 # default
|
|
|
|
for part in second_line.split():
|
|
if "Electrons=" in part:
|
|
electrons = int(part.split("=")[1])
|
|
elif "Orbitals=" in part:
|
|
orbitals = int(part.split("=")[1])
|
|
elif "Charge=" in part:
|
|
charge = int(part.split("=")[1])
|
|
elif "Multiplicity=" in part:
|
|
multiplicity = int(part.split("=")[1])
|
|
|
|
return electrons, orbitals, charge, multiplicity
|
|
|
|
return 3, 3, 0, 1 # fallback defaults
|
|
|
|
|
|
def prepare_data(data):
|
|
text_content = data.get("text", "")
|
|
|
|
# Parse the molecular data (assuming it's in XYZ format)
|
|
symbols, coordinates = parse_xyz_from_text(text_content)
|
|
|
|
# Extract electron/orbital info (from the Charge/Multiplicity line)
|
|
# "Charge=0 Multiplicity=1 Electrons=3 Orbitals=3"
|
|
electrons, orbitals, charge, multiplicity = extract_electron_info(text_content)
|
|
|
|
return {
|
|
"symbols": symbols,
|
|
"coordinates": coordinates,
|
|
"charge": charge,
|
|
"multiplicity": multiplicity,
|
|
"active_electrons": electrons,
|
|
"active_orbitals": orbitals,
|
|
"max_iterations": data.get("max_iterations", 200),
|
|
"conv_tol": data.get("conv_tol", 1e-6),
|
|
"step_size": data.get("step_size", 0.05),
|
|
}
|
|
|
|
|
|
def run_vqe(conn: Connection, dev1: Device, data: dict):
|
|
coordinates = jnp.array(data.get("coordinates"))
|
|
charge = int(data.get("charge"))
|
|
multiplicity = int(data.get("multiplicity"))
|
|
molecule = qchem.Molecule(
|
|
data.get("symbols"),
|
|
coordinates,
|
|
charge=charge,
|
|
mult=multiplicity,
|
|
)
|
|
|
|
active_electrons = int(data.get("active_electrons"))
|
|
active_orbitals = int(data.get("active_orbitals"))
|
|
|
|
max_iterations = int(data.get("max_iterations", 200))
|
|
step_size = float(data.get("step_size", 0.05))
|
|
conv_tol = float(data.get("conv_tol", 1e-6))
|
|
|
|
H, qubits = qchem.molecular_hamiltonian(
|
|
molecule,
|
|
active_electrons=active_electrons,
|
|
active_orbitals=active_orbitals,
|
|
method="openfermion",
|
|
) # type: ignore
|
|
|
|
singles, doubles = qml.qchem.excitations(active_electrons, qubits)
|
|
|
|
params = np.array(np.zeros(len(singles) + len(doubles)), requires_grad=True)
|
|
|
|
conn.send(
|
|
{
|
|
"iter_num": 0,
|
|
"energy": None,
|
|
"conv": None,
|
|
"params": params.tolist() if hasattr(params, "tolist") else list(params),
|
|
}
|
|
)
|
|
|
|
@qml.qnode(dev1)
|
|
def circuit(param, wires):
|
|
# Map excitations to the wires the UCCSD circuit will act on
|
|
s_wires, d_wires = qml.qchem.excitations_to_wires(singles, doubles)
|
|
qml.UCCSD(
|
|
param,
|
|
wires,
|
|
s_wires=s_wires,
|
|
d_wires=d_wires,
|
|
init_state=qml.qchem.hf_state(active_electrons, qubits),
|
|
)
|
|
return qml.expval(H)
|
|
|
|
def cost_fn(param):
|
|
return circuit(param, wires=range(qubits))
|
|
|
|
opt = GradientDescentOptimizer(stepsize=step_size)
|
|
|
|
for n in range(max_iterations):
|
|
# Take step
|
|
params, prev_energy = opt.step_and_cost(cost_fn, params)
|
|
|
|
energy = cost_fn(params)
|
|
|
|
# Calculate difference between new and old energies
|
|
conv = np.abs(energy - prev_energy)
|
|
|
|
conn.send(
|
|
{
|
|
"iter_num": n,
|
|
"energy": float(energy),
|
|
"conv": float(conv),
|
|
"params": params.tolist()
|
|
if hasattr(params, "tolist")
|
|
else list(params),
|
|
}
|
|
)
|
|
|
|
if conv <= conv_tol:
|
|
break
|