vqe-pennylane 0.2.2__py3-none-any.whl

vqe/io_utils.py

@@ -0,0 +1,265 @@
+"""
+vqe.io_utils
+------------
+Utility functions for reproducible VQE runs:
+
+- Run configuration construction & hashing
+- JSON-safe serialization
+- File/directory management for results & images
+
+This module is the single source of truth for:
+- Where new package results are stored
+- How run configurations are represented and hashed
+"""
+
+from __future__ import annotations
+
+import json, os
+import hashlib
+from pathlib import Path
+from typing import Any, Dict
+
+# ================================================================
+# BASE PATHS
+# ================================================================
+
+# Root of the repository/package (…/ <repo_root> / vqe / io_utils.py)
+BASE_DIR: Path = Path(__file__).resolve().parent.parent
+
+# New, package-scoped locations for results and images
+RESULTS_DIR = BASE_DIR / "results" / "vqe"
+IMG_DIR = BASE_DIR / "images" / "vqe"
+
+
+# ================================================================
+# INTERNAL HELPERS
+# ================================================================
+
+def _round_floats(x: Any, ndigits: int = 8) -> Any:
+    """
+    Recursively round floats / numpy scalars / arrays for stable hashing.
+
+    Parameters
+    ----------
+    x:
+        Arbitrary Python or numpy object.
+    ndigits:
+        Number of decimal places to round to.
+
+    Returns
+    -------
+    Any
+        Object of the same container structure with floats rounded.
+    """
+    # Plain Python float
+    if isinstance(x, float):
+        return round(x, ndigits)
+
+    # Numpy scalar-like: has .item() that is a float
+    try:
+        if hasattr(x, "item"):
+            scalar = x.item()
+            if isinstance(scalar, float):
+                return round(float(scalar), ndigits)
+    except Exception:
+        # Fall through if .item() misbehaves
+        pass
+
+    # Numpy arrays / array-like -> convert to list and recurse
+    if hasattr(x, "tolist"):
+        return _round_floats(x.tolist(), ndigits)
+
+    # Python containers
+    if isinstance(x, (list, tuple)):
+        return type(x)(_round_floats(v, ndigits) for v in x)
+
+    # Everything else left untouched
+    return x
+
+
+def _to_serializable(obj: Any) -> Any:
+    """
+    Convert tensors / numpy arrays / complex containers into JSON-safe types.
+
+    Rules of thumb:
+    - numpy / tensor-like with .tolist() → list / nested lists
+    - numpy scalar with .item() → float (when possible)
+    - dict / list / tuple → recurse
+    - everything else returned unchanged
+    """
+    # Numpy / tensor scalar
+    if hasattr(obj, "item"):
+        try:
+            return obj.item()
+        except Exception:
+            pass
+
+    # Numpy arrays and friends
+    if hasattr(obj, "tolist"):
+        return obj.tolist()
+
+    # Dicts
+    if isinstance(obj, dict):
+        return {k: _to_serializable(v) for k, v in obj.items()}
+
+    # Lists / tuples
+    if isinstance(obj, (list, tuple)):
+        return [_to_serializable(v) for v in obj]
+
+    # Primitive or unknown: hope json can handle it
+    return obj
+
+
+# ================================================================
+# RUN CONFIGURATION & HASHING
+# ================================================================
+
+def make_run_config_dict(
+    symbols,
+    coordinates,
+    basis: str,
+    ansatz_desc: str,
+    optimizer_name: str,
+    stepsize: float,
+    max_iterations: int,
+    seed: int,
+    mapping: str,
+    noisy: bool = False,
+    depolarizing_prob: float = 0.0,
+    amplitude_damping_prob: float = 0.0,
+    molecule_label: str | None = None,
+) -> Dict[str, Any]:
+    """
+    Construct a canonical dictionary describing a VQE/SSVQE run configuration.
+
+    This dictionary is used to generate a stable hash signature
+    (see :func:`run_signature`) so that identical configurations
+    always map to the same cache key / filename.
+    """
+    cfg: Dict[str, Any] = {
+        "symbols": list(symbols),
+        "geometry": _round_floats(coordinates, 8),
+        "basis": basis,
+        "ansatz": ansatz_desc,
+        "optimizer": {
+            "name": optimizer_name,
+            "stepsize": float(stepsize),
+            "iterations_planned": int(max_iterations),
+        },
+        "optimizer_name": optimizer_name,  # convenient flat copy
+        "seed": int(seed),
+        "noisy": bool(noisy),
+        "depolarizing_prob": float(depolarizing_prob),
+        "amplitude_damping_prob": float(amplitude_damping_prob),
+        "mapping": mapping.lower(),
+    }
+
+    if molecule_label is not None:
+        cfg["molecule"] = str(molecule_label)
+
+    return cfg
+
+
+def run_signature(cfg: Dict[str, Any]) -> str:
+    """
+    Generate a short, stable hash identifier from a run configuration.
+
+    The configuration is JSON-serialized with sorted keys and compact
+    separators, then hashed with SHA-256. The first 12 hex characters
+    are used as the signature.
+
+    Parameters
+    ----------
+    cfg:
+        Configuration dictionary as returned by :func:`make_run_config_dict`.
+
+    Returns
+    -------
+    str
+        12-character hexadecimal signature.
+    """
+    payload = json.dumps(cfg, sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:12]
+
+
+# ================================================================
+# FILESYSTEM UTILITIES
+# ================================================================
+
+def ensure_dirs() -> None:
+    """
+    Ensure that the standard result and image directories exist.
+
+    - RESULTS_DIR: where JSON run records are written (package_results/)
+    - IMG_DIR:     where plots and figures are saved (vqe/images/)
+    """
+    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
+    IMG_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _result_path_from_prefix(prefix: str) -> Path:
+    """
+    Build the full JSON path from a filename prefix (without extension).
+
+    Parameters
+    ----------
+    prefix:
+        Filename prefix, e.g. "H2_Adam_s0__abc123def456".
+
+    Returns
+    -------
+    Path
+        Path to the JSON file within RESULTS_DIR.
+    """
+    return RESULTS_DIR / f"{prefix}.json"
+
+
+def save_run_record(prefix: str, record: Dict[str, Any]) -> str:
+    """
+    Save a run record (config + results) as JSON in RESULTS_DIR.
+
+    The final path is:
+        RESULTS_DIR / f"{prefix}.json"
+
+    Parameters
+    ----------
+    prefix:
+        Unique filename prefix, typically including molecule, optimizer,
+        seed, and the run signature.
+    record:
+        Dictionary containing configuration, results, and any metadata.
+
+    Returns
+    -------
+    str
+        String path to the saved JSON file.
+    """
+    ensure_dirs()
+    path = _result_path_from_prefix(prefix)
+    serializable_record = _to_serializable(record)
+
+    with path.open("w") as f:
+        json.dump(serializable_record, f, indent=2)
+
+    return str(path)
+
+
+def make_filename_prefix(cfg: dict, *, noisy: bool, seed: int, hash_str: str, ssvqe: bool = False):
+    """Return unified Option-C filename prefix."""
+    # Molecule lives in config only indirectly; infer from symbols
+    # OR require callers to inject molecule into cfg beforehand.
+    mol = cfg.get("molecule", "MOL")  # We will fix cfg in core/ssvqe.
+    
+    # Ansatz string taken directly
+    ans = cfg.get("ansatz", "ANSATZ")
+
+    # Optimizer name
+    if "optimizer" in cfg and "name" in cfg["optimizer"]:
+        opt = cfg["optimizer"]["name"]
+    else:
+        opt = "OPT"
+
+    noise_tag = "noisy" if noisy else "noiseless"
+    algo_tag  = "SSVQE" if ssvqe else "VQE"
+
+    return f"{mol}__{ans}__{opt}__{algo_tag}__{noise_tag}__s{seed}__{hash_str}"

vqe/optimizer.py

@@ -0,0 +1,58 @@
+"""
+vqe.optimizer
+-------------
+Lightweight wrapper over PennyLane optimizers with a unified interface.
+
+Provides:
+    - get_optimizer(name, stepsize)
+"""
+
+from __future__ import annotations
+import pennylane as qml
+
+
+# ================================================================
+# AVAILABLE OPTIMIZERS
+# ================================================================
+_OPTIMIZERS = {
+    "Adam": qml.AdamOptimizer,
+    "adam": qml.AdamOptimizer,  # alias
+
+    "GradientDescent": qml.GradientDescentOptimizer,
+    "gd": qml.GradientDescentOptimizer,  # alias
+
+    "Momentum": qml.MomentumOptimizer,
+    "Nesterov": qml.NesterovMomentumOptimizer,
+
+    "RMSProp": qml.RMSPropOptimizer,
+    "Adagrad": qml.AdagradOptimizer,
+
+    # ---- NEW ADDITIONS (SPSA SUPPORT) ----
+    "SPSA": qml.SPSAOptimizer,
+    "spsa": qml.SPSAOptimizer,  # alias
+}
+
+
+# ================================================================
+# MAIN FACTORY
+# ================================================================
+def get_optimizer(name: str = "Adam", stepsize: float = 0.2):
+    """
+    Return a PennyLane optimizer instance by name.
+
+    Args:
+        name: Optimizer identifier (case-insensitive).
+        stepsize: Learning rate.
+
+    Returns:
+        An instantiated optimizer.
+    """
+    key = name.lower()
+    for k, cls in _OPTIMIZERS.items():
+        if k.lower() == key:
+            return cls(stepsize)
+
+    raise ValueError(
+        f"Unknown optimizer '{name}'. "
+        f"Available: {', '.join(_OPTIMIZERS.keys())}"
+    )

vqe/ssvqe.py

@@ -0,0 +1,271 @@
+"""
+vqe.ssvqe
+---------
+Subspace-Search Variational Quantum Eigensolver (SSVQE).
+
+Features:
+    - Ground and excited states via a shared variational ansatz
+    - Orthogonality enforced with |⟨ψ_i|ψ_j⟩|² penalties
+    - Any ansatz and optimizer defined in the vqe package
+    - Optional depolarizing / amplitude-damping noise
+    - Caching and reproducibility via vqe.io_utils
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import itertools
+import pennylane as qml
+from pennylane import numpy as np
+
+from .hamiltonian import build_hamiltonian
+from .engine import (
+    make_device,
+    build_ansatz,
+    build_optimizer,
+    make_energy_qnode,
+    make_overlap00_fn,
+)
+from .io_utils import (
+    ensure_dirs,
+    make_run_config_dict,
+    run_signature,
+    save_run_record,
+    make_filename_prefix,
+    RESULTS_DIR,
+)
+from .visualize import plot_ssvqe_convergence_multi
+
+
+# ================================================================
+# MAIN ENTRYPOINT
+# ================================================================
+def run_ssvqe(
+    molecule: str = "H3+",
+    *,
+    num_states: int = 2,
+    penalty_weight: float = 10.0,
+    ansatz_name: str = "UCCSD",
+    optimizer_name: str = "Adam",
+    steps: int = 100,
+    stepsize: float = 0.4,
+    seed: int = 0,
+    noisy: bool = False,
+    depolarizing_prob: float = 0.0,
+    amplitude_damping_prob: float = 0.0,
+    symbols=None,
+    coordinates=None,
+    basis: str = "sto-3g",
+    plot: bool = True,
+    force: bool = False,
+):
+    """
+    Run a Subspace-Search VQE (SSVQE) optimization to obtain ground and
+    excited states of a molecular Hamiltonian.
+
+    Args:
+        molecule: Molecule label (e.g. "H2", "LiH", "H3+").
+        num_states: Number of eigenstates to target (>= 2).
+        penalty_weight: Weight on the orthogonality penalty term.
+        ansatz_name: Name of the ansatz from `vqe.ansatz`.
+        optimizer_name: Name of optimizer from `vqe.optimizer`.
+        steps: Number of optimization iterations.
+        stepsize: Optimizer step size.
+        seed: Random seed for reproducibility.
+        noisy: If True, use a mixed-state simulator and insert noise.
+        depolarizing_prob: Depolarizing noise probability per wire.
+        amplitude_damping_prob: Amplitude damping probability per wire.
+        symbols, coordinates, basis: Optional explicit molecular data.
+        plot: Whether to plot E0 / E1 convergence (if num_states >= 2).
+        force: If True, ignore cached results and recompute.
+
+    Returns:
+        dict with keys:
+            - "energies_per_state": list[list[float]]
+            - "final_params":       list[list[float]]
+            - "config":             dict
+    """
+    assert num_states >= 2, "SSVQE requires at least two target states."
+
+    np.random.seed(seed)
+    ensure_dirs()
+
+    # ============================================================
+    # 1. Build Hamiltonian and molecular data
+    # ============================================================
+    if symbols is None or coordinates is None:
+        H, num_wires, symbols, coordinates, basis = build_hamiltonian(molecule)
+    else:
+        # Default charge heuristic (matches your prior H3+ usage)
+        charge = +1 if molecule.upper() == "H3+" else 0
+        H, num_wires = qml.qchem.molecular_hamiltonian(
+            symbols, coordinates, charge=charge, basis=basis, unit="angstrom"
+        )
+
+    # ============================================================
+    # 2. Ansatz and initial parameters (shared structure)
+    # ============================================================
+    ansatz_fn, init_params = build_ansatz(
+        ansatz_name,
+        num_wires,
+        seed=seed,
+        symbols=symbols,
+        coordinates=coordinates,
+        basis=basis,
+    )
+    # One parameter set per state, all with same shape
+    param_sets = [np.array(init_params, requires_grad=True) for _ in range(num_states)]
+
+    # ============================================================
+    # 3. Device and QNodes
+    # ============================================================
+    dev = make_device(num_wires, noisy=noisy)
+
+    energy_qnode = make_energy_qnode(
+        H,
+        dev,
+        ansatz_fn,
+        num_wires,
+        noisy=noisy,
+        depolarizing_prob=depolarizing_prob,
+        amplitude_damping_prob=amplitude_damping_prob,
+        symbols=symbols,
+        coordinates=coordinates,
+        basis=basis,
+    )
+
+    overlap00 = make_overlap00_fn(
+        dev,
+        ansatz_fn,
+        num_wires,
+        noisy=noisy,
+        depolarizing_prob=depolarizing_prob,
+        amplitude_damping_prob=amplitude_damping_prob,
+        symbols=symbols,
+        coordinates=coordinates,
+        basis=basis,
+    )
+
+    # ============================================================
+    # 4. Config + caching
+    # ============================================================
+    cfg = make_run_config_dict(
+        symbols=symbols,
+        coordinates=coordinates,
+        basis=basis,
+        ansatz_desc=f"SSVQE({ansatz_name})_{num_states}states",
+        optimizer_name=optimizer_name,
+        stepsize=stepsize,
+        max_iterations=steps,
+        seed=seed,
+        mapping="jordan_wigner",
+        noisy=noisy,
+        depolarizing_prob=depolarizing_prob,
+        amplitude_damping_prob=amplitude_damping_prob,
+        molecule_label=molecule,   # <<< add this
+    )
+    cfg["penalty_weight"] = float(penalty_weight)
+    cfg["num_states"] = int(num_states)
+
+    sig = run_signature(cfg)
+    safe_molecule = molecule.replace("+", "plus")
+    prefix = make_filename_prefix(
+        cfg,
+        noisy=noisy,
+        seed=seed,
+        hash_str=sig,
+        ssvqe=True,
+    )
+    result_path = RESULTS_DIR / f"{prefix}.json"
+
+    if not force and result_path.exists():
+        print(f"📂 Using cached SSVQE result: {result_path}")
+        with open(result_path, "r") as f:
+            record = json.load(f)
+        return record["result"]
+
+    # ============================================================
+    # 5. Cost function: total energy + orthogonality penalties
+    # ============================================================
+    opt = build_optimizer(optimizer_name, stepsize=stepsize)
+
+    def _unpack_flat(flat, templates):
+        """Unpack a flat parameter vector into a list of arrays matching templates."""
+        arrays = []
+        idx = 0
+        for tmpl in templates:
+            size = int(np.prod(tmpl.shape))
+            vec = flat[idx:idx + size]
+            arrays.append(np.reshape(vec, tmpl.shape))
+            idx += size
+        return arrays
+
+    def cost(flat_params):
+        # Split into per-state parameter arrays
+        unpacked = _unpack_flat(flat_params, param_sets)
+
+        # Sum of individual energies
+        total = sum(energy_qnode(p) for p in unpacked)
+
+        # Add orthogonality penalties between all distinct pairs
+        for i, j in itertools.combinations(range(len(unpacked)), 2):
+            total = total + penalty_weight * overlap00(unpacked[i], unpacked[j])
+
+        return total
+
+    # Flatten initial parameters for joint optimization
+    flat = np.concatenate([p.ravel() for p in param_sets])
+    flat = np.array(flat, requires_grad=True)
+
+    energies_per_state = [[] for _ in range(num_states)]
+
+    # ============================================================
+    # 6. Optimization loop
+    # ============================================================
+    for step in range(steps):
+        try:
+            flat, _ = opt.step_and_cost(cost, flat)
+        except AttributeError:
+            flat = opt.step(cost, flat)
+
+        # Unpack back into per-state parameter arrays
+        unpacked = _unpack_flat(flat, param_sets)
+
+        # Record energies for each state
+        for k in range(num_states):
+            energies_per_state[k].append(float(energy_qnode(unpacked[k])))
+
+        # Update param_sets for the next iteration (keep grad enabled)
+        param_sets = [np.array(u, requires_grad=True) for u in unpacked]
+
+    # ============================================================
+    # 7. Package and save results
+    # ============================================================
+    result = {
+        "energies_per_state": energies_per_state,
+        "final_params": [u.tolist() for u in param_sets],
+        "config": cfg,
+    }
+
+    save_run_record(prefix, {"config": cfg, "result": result})
+    print(f"💾 Saved SSVQE run to {result_path}")
+
+    # ============================================================
+    # 8. Optional plotting (E0 / E1 only, if available)
+    # ============================================================
+    if plot and num_states >= 2:
+        try:
+            E0 = energies_per_state[0]
+            E1 = energies_per_state[1]
+            plot_ssvqe_convergence_multi(
+                molecule,
+                E0_list=E0,
+                E1_list=E1,
+                optimizer_name=optimizer_name,
+                ansatz=ansatz_name,
+            )
+        except Exception as exc:  # plotting is non-fatal
+            print(f"⚠️ SSVQE plotting failed (non-fatal): {exc}")
+
+    return result