qmlkit 0.1.0__py3-none-any.whl

qmlkit/algorithms/rl.py

@@ -0,0 +1,198 @@
+"""Quantum policy gradient — a variational circuit as an RL policy.
+
+REINFORCE, with the policy replaced by a circuit: encode the observation, measure a
+few observables, softmax them into action probabilities, and push up the log-
+probability of whatever earned reward.
+
+    policy = QuantumPolicy(n_observations=2, n_actions=2)
+    result = train_reinforce(policy, ContextualBandit(seed=0), n_episodes=200, seed=0)
+
+The environment is an argument with a three-method protocol (``reset``, ``step``,
+``n_observations``/``n_actions``), so a Gym environment wraps in a few lines and
+nothing here needs Gym installed. The ansatz and feature map are arguments too, as
+everywhere else.
+
+The gradient is the ordinary policy-gradient one — ``grad log pi(a|s) * return`` —
+computed through :func:`qmlkit.grad`, so it is exact rather than finite-differenced.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+import numpy as np
+import numpy.typing as npt
+
+from qmlkit.ansatz.library import Ansatz, hardware_efficient
+from qmlkit.core.execute import BackendLike, expval
+from qmlkit.core.observables import Observable, Z
+from qmlkit.encoding.feature_maps import AngleFeatureMap, FeatureMap
+from qmlkit.gradients.dispatch import grad
+
+__all__ = ["QuantumPolicy", "ContextualBandit", "train_reinforce", "ReinforceResult", "Environment"]
+
+
+class Environment(Protocol):
+    """The three things an RL loop needs. Deliberately not Gym."""
+
+    n_observations: int
+    n_actions: int
+
+    def reset(self) -> npt.NDArray[Any]: ...
+
+    def step(self, action: int) -> tuple[npt.NDArray[Any], float, bool]: ...
+
+
+class ContextualBandit:
+    """A tiny environment with a known optimal policy, so training is checkable.
+
+    The observation is a random vector; the correct action is the sign of its first
+    coordinate. One step per episode, reward 1 for right and 0 for wrong — so the
+    optimal return is exactly 1.0 and "did it learn" has an unambiguous answer.
+    """
+
+    n_observations = 2
+    n_actions = 2
+
+    def __init__(self, seed: int | None = None) -> None:
+        self._rng = np.random.default_rng(seed)
+        self._state: npt.NDArray[Any] = np.zeros(self.n_observations)
+
+    def reset(self) -> npt.NDArray[Any]:
+        self._state = self._rng.uniform(-1.0, 1.0, self.n_observations)
+        return self._state
+
+    def step(self, action: int) -> tuple[npt.NDArray[Any], float, bool]:
+        correct = int(self._state[0] > 0)
+        return self._state, float(action == correct), True
+
+    def optimal_return(self) -> float:
+        return 1.0
+
+
+class QuantumPolicy:
+    """A circuit policy: observation in, action probabilities out."""
+
+    def __init__(
+        self,
+        n_observations: int,
+        n_actions: int,
+        feature_map: FeatureMap | None = None,
+        ansatz: Ansatz | None = None,
+        observables: Sequence[Observable] | None = None,
+        n_layers: int = 2,
+        beta: float = 2.0,
+        backend: BackendLike = None,
+        seed: int | None = None,
+    ) -> None:
+        n_qubits = max(n_observations, n_actions)
+        self.n_observations = n_observations
+        self.n_actions = n_actions
+        self.n_qubits = n_qubits
+        self.feature_map = feature_map or AngleFeatureMap(n_qubits, entangle=n_qubits > 1)
+        self.ansatz = ansatz or hardware_efficient(n_qubits, n_layers)
+        self.observables = list(observables) if observables else [Z(a) for a in range(n_actions)]
+        self.beta = beta  # softmax inverse temperature
+        self.backend = backend
+        # Seeded from the argument, not from a constant: with seed baked in, every
+        # QuantumPolicy started from the same point and "average over seeds" was
+        # impossible to express.
+        self.theta = self.ansatz.init("small", seed=seed)
+        self._spec = self.feature_map.build_parametric().compose(self.ansatz.build())
+        self._n_inputs = self.feature_map.n_angles
+
+    def _full(self, observation: npt.NDArray[Any], theta: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        padded = np.zeros(self.n_qubits)
+        padded[: self.n_observations] = np.asarray(observation, dtype=float).ravel()
+        return np.concatenate([self.feature_map.angles(padded), theta])
+
+    def logits(
+        self, observation: npt.NDArray[Any], theta: npt.NDArray[Any] | None = None
+    ) -> npt.NDArray[Any]:
+        full = self._full(observation, self.theta if theta is None else theta)
+        return np.array(
+            [expval(self._spec, o, theta=full, backend=self.backend) for o in self.observables]
+        )
+
+    def probabilities(
+        self, observation: npt.NDArray[Any], theta: npt.NDArray[Any] | None = None
+    ) -> npt.NDArray[Any]:
+        z = self.beta * self.logits(observation, theta)
+        z = z - z.max()
+        e = np.exp(z)
+        return e / e.sum()
+
+    def sample(self, observation: npt.NDArray[Any], rng: np.random.Generator) -> int:
+        return int(rng.choice(self.n_actions, p=self.probabilities(observation)))
+
+    def grad_log_prob(self, observation: npt.NDArray[Any], action: int) -> npt.NDArray[Any]:
+        r"""``d/dtheta log pi(a|s)``.
+
+        For a softmax over measured observables this is
+        ``beta * (dO_a/dtheta - sum_b pi_b dO_b/dtheta)`` — one exact circuit gradient
+        per action, no finite differences anywhere.
+        """
+        full = self._full(observation, self.theta)
+        probs = self.probabilities(observation)
+        jac = np.stack(
+            [
+                grad(self._spec, full, o, backend=self.backend)[self._n_inputs :]
+                for o in self.observables
+            ]
+        )
+        return self.beta * (jac[action] - probs @ jac)
+
+    def __repr__(self) -> str:
+        return (
+            f"QuantumPolicy(obs={self.n_observations}, actions={self.n_actions}, "
+            f"P={self.theta.size}, ansatz={self.ansatz.name!r})"
+        )
+
+
+@dataclass
+class ReinforceResult:
+    theta: npt.NDArray[Any]
+    returns: list[float] = field(default_factory=list)
+
+    def mean_return(self, last: int = 50) -> float:
+        return float(np.mean(self.returns[-last:])) if self.returns else 0.0
+
+    def __repr__(self) -> str:
+        return (
+            f"ReinforceResult(episodes={len(self.returns)}, "
+            f"final_mean_return={self.mean_return():.4f})"
+        )
+
+
+def train_reinforce(
+    policy: QuantumPolicy,
+    env: Environment,
+    n_episodes: int = 200,
+    lr: float = 0.2,
+    baseline: bool = True,
+    seed: int | None = None,
+) -> ReinforceResult:
+    """REINFORCE with an optional moving-average baseline.
+
+    The baseline subtracts a running mean return before scaling the gradient. It does
+    not change what is being optimised, only the variance of the estimate — which for
+    a policy sampled one episode at a time is the thing that decides whether it
+    learns at all.
+    """
+    rng = np.random.default_rng(seed)
+    returns: list[float] = []
+    running = 0.0
+
+    for episode in range(n_episodes):
+        observation = env.reset()
+        action = policy.sample(observation, rng)
+        _, reward, _ = env.step(action)
+        returns.append(float(reward))
+
+        advantage = reward - running if baseline else reward
+        policy.theta = policy.theta + lr * advantage * policy.grad_log_prob(observation, action)
+        running += (reward - running) / (episode + 1)
+
+    return ReinforceResult(theta=policy.theta, returns=returns)

qmlkit/algorithms/vqe.py

@@ -0,0 +1,198 @@
+"""The variational quantum eigensolver.
+
+VQE is a loop, not an architecture: prepare a parameterised state, measure an energy,
+step downhill, repeat. Everything that makes one VQE differ from another — the
+ansatz, the optimiser, the gradient method, the shot budget — is therefore an
+argument, and the class itself is thin on purpose.
+
+    from qmlkit.algorithms import VQE, ising_hamiltonian
+
+    H = ising_hamiltonian(4, j=1.0, h=0.5)
+    result = VQE(H, n_qubits=4).run(seed=0)
+    print(result.energy, result.error_vs_exact)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from qmlkit.algorithms.hamiltonians import exact_ground_energy
+from qmlkit.ansatz.library import Ansatz, hardware_efficient
+from qmlkit.core.execute import BackendLike, expectation
+from qmlkit.core.observables import Observable, observable_support
+
+__all__ = ["VQE", "VQEResult", "OPTIMIZERS"]
+
+#: ``fn(loss, theta0, **kwargs) -> (theta, history)`` — the shape every optimiser here
+#: already has, so a custom one is a function rather than an adapter class.
+Optimizer = Callable[..., "tuple[npt.NDArray[Any], list[float]]"]
+
+
+def _rotosolve(loss, theta0, **kw):  # type: ignore[no-untyped-def]
+    from qmlkit.optim import minimize_rotosolve
+
+    return minimize_rotosolve(loss, theta0, **kw)
+
+
+def _spsa(loss, theta0, **kw):  # type: ignore[no-untyped-def]
+    from qmlkit.gradients.spsa import minimize_spsa
+
+    return minimize_spsa(loss, theta0, **kw)
+
+
+def _gradient_descent(loss, theta0, *, grad, n_steps=100, lr=0.1, **kw):  # type: ignore[no-untyped-def]
+    theta = np.asarray(theta0, dtype=float).copy()
+    history = [float(loss(theta))]
+    for _ in range(n_steps):
+        theta = theta - lr * grad(theta)
+        history.append(float(loss(theta)))
+    return theta, history
+
+
+OPTIMIZERS: dict[str, Optimizer] = {
+    "rotosolve": _rotosolve,
+    "spsa": _spsa,
+    "gradient-descent": _gradient_descent,
+}
+
+
+@dataclass
+class VQEResult:
+    """What a run produced, and how good it actually is."""
+
+    energy: float
+    theta: npt.NDArray[Any]
+    history: list[float] = field(default_factory=list)
+    exact: float | None = None
+    n_evaluations: int = 0
+
+    @property
+    def error_vs_exact(self) -> float | None:
+        """Absolute error against dense diagonalisation, when that was computed."""
+        return None if self.exact is None else abs(self.energy - self.exact)
+
+    def __repr__(self) -> str:
+        tail = (
+            ""
+            if self.exact is None
+            else f", exact={self.exact:.8f}, error={self.error_vs_exact:.2e}"
+        )
+        return f"VQEResult(energy={self.energy:.8f}{tail}, steps={len(self.history) - 1})"
+
+
+class VQE:
+    """Minimise ``<H>`` over a parameterised state.
+
+    Parameters
+    ----------
+    hamiltonian
+        Any observable. :mod:`qmlkit.algorithms.hamiltonians` has constructors.
+    ansatz
+        The trial state. Defaults to a hardware-efficient circuit wide enough for
+        the Hamiltonian's support — replaceable with anything, including one you
+        invented, because an ``Ansatz`` is the only contract.
+    optimizer
+        A name from :data:`OPTIMIZERS` or any ``fn(loss, theta0, **kw)`` returning
+        ``(theta, history)``.
+    gradient
+        Passed through to :func:`qmlkit.grad`. Only consulted by gradient-based
+        optimisers; Rotosolve and SPSA never ask for one.
+    """
+
+    def __init__(
+        self,
+        hamiltonian: Observable,
+        ansatz: Ansatz | None = None,
+        n_qubits: int | None = None,
+        optimizer: str | Optimizer = "rotosolve",
+        gradient: str = "auto",
+        backend: BackendLike = None,
+        shots: int | None = None,
+    ) -> None:
+        support = observable_support(hamiltonian)
+        width = n_qubits or (max(support) + 1 if support else 1)
+        if ansatz is not None and ansatz.n_qubits < width:
+            raise ValueError(
+                f"the ansatz has {ansatz.n_qubits} qubits but the Hamiltonian acts on {width}"
+            )
+        self.hamiltonian = hamiltonian
+        self.ansatz = ansatz or hardware_efficient(width, n_layers=2)
+        self.n_qubits = self.ansatz.n_qubits
+        self.optimizer = optimizer
+        self.gradient = gradient
+        self.backend = backend
+        self.shots = shots
+        self._spec = self.ansatz.build()
+        self.n_evaluations = 0
+
+    # ------------------------------------------------------------------ energy --
+    def energy(self, theta: Sequence[float]) -> float:
+        """``<H>`` at these parameters."""
+        self.n_evaluations += 1
+        value = expectation(
+            self._spec,
+            self.hamiltonian,
+            theta=np.asarray(theta, dtype=float),
+            shots=self.shots,
+            backend=self.backend,
+        )
+        return float(value)
+
+    def gradient_of_energy(self, theta: Sequence[float]) -> npt.NDArray[Any]:
+        from qmlkit.gradients.dispatch import grad
+
+        return grad(
+            self._spec,
+            np.asarray(theta, dtype=float),
+            self.hamiltonian,
+            method=self.gradient,
+            backend=self.backend,
+            shots=self.shots,
+        )
+
+    # --------------------------------------------------------------------- run --
+    def run(
+        self,
+        theta0: Sequence[float] | None = None,
+        seed: int | None = None,
+        compare_exact: bool | None = None,
+        **optimizer_kwargs: Any,
+    ) -> VQEResult:
+        """Optimise, and report against the exact answer when that is affordable."""
+        start = (
+            np.asarray(theta0, dtype=float)
+            if theta0 is not None
+            else self.ansatz.init("small", seed=seed)
+        )
+        fn = OPTIMIZERS[self.optimizer] if isinstance(self.optimizer, str) else self.optimizer
+        if fn is _gradient_descent:
+            optimizer_kwargs.setdefault("grad", self.gradient_of_energy)
+        if fn is _spsa:
+            optimizer_kwargs.setdefault("seed", seed)
+
+        self.n_evaluations = 0
+        theta, history = fn(self.energy, start, **optimizer_kwargs)
+
+        # dense diagonalisation is exponential, so only offer it where it is cheap
+        if compare_exact is None:
+            compare_exact = self.n_qubits <= 12
+        exact = exact_ground_energy(self.hamiltonian, self.n_qubits) if compare_exact else None
+
+        return VQEResult(
+            energy=float(history[-1]),
+            theta=theta,
+            history=list(history),
+            exact=exact,
+            n_evaluations=self.n_evaluations,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"VQE(n_qubits={self.n_qubits}, ansatz={self.ansatz.name!r}, "
+            f"P={self.ansatz.n_weights}, optimizer={self.optimizer!r})"
+        )

qmlkit/ansatz/__init__.py

@@ -0,0 +1,68 @@
+"""Ansatz: the trainable block, and the vocabulary for inventing new ones."""
+
+from qmlkit.ansatz.blocks import (
+    Block,
+    BuildContext,
+    Custom,
+    EncodingLayer,
+    EntanglerLayer,
+    ParametricEntangler,
+    PoolLayer,
+    RotationLayer,
+    Sequential,
+    repeat,
+    share,
+)
+from qmlkit.ansatz.library import (
+    Ansatz,
+    basic_entangler,
+    conv_block,
+    get_ansatz,
+    hardware_efficient,
+    list_ansatze,
+    list_conv_filters,
+    mps_ansatz,
+    qaoa_ansatz,
+    qcnn_ansatz,
+    random_layers,
+    register_ansatz,
+    register_conv_filter,
+    simplified_two_design,
+    strongly_entangling,
+    tree_tensor_network,
+    two_local,
+)
+from qmlkit.ansatz.reupload import ReuploadModel, reupload
+
+__all__ = [
+    "Ansatz",
+    "register_ansatz",
+    "get_ansatz",
+    "list_ansatze",
+    "Block",
+    "BuildContext",
+    "Sequential",
+    "RotationLayer",
+    "EntanglerLayer",
+    "ParametricEntangler",
+    "PoolLayer",
+    "EncodingLayer",
+    "Custom",
+    "repeat",
+    "share",
+    "hardware_efficient",
+    "strongly_entangling",
+    "simplified_two_design",
+    "tree_tensor_network",
+    "mps_ansatz",
+    "qcnn_ansatz",
+    "qaoa_ansatz",
+    "conv_block",
+    "list_conv_filters",
+    "register_conv_filter",
+    "reupload",
+    "ReuploadModel",
+    "basic_entangler",
+    "two_local",
+    "random_layers",
+]