qenvs 0.0.1__py3-none-any.whl

qenvs/__init__.py

@@ -0,0 +1,25 @@
+"""QEnvS: physically validated quantum control environments for RL and optimal control."""
+
+from qenvs import envs, models, observables, tasks, validation
+from qenvs.envs.base import Environment, EnvState, TimeStep
+from qenvs.registry import make, register, registry
+
+try:
+    from qenvs._version import __version__
+except ImportError:  # pragma: no cover - editable install without build hook
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "EnvState",
+    "Environment",
+    "TimeStep",
+    "__version__",
+    "envs",
+    "make",
+    "models",
+    "observables",
+    "register",
+    "registry",
+    "tasks",
+    "validation",
+]

qenvs/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.0.1'
+__version_tuple__ = version_tuple = (0, 0, 1)
+
+__commit_id__ = commit_id = None

qenvs/cli.py

@@ -0,0 +1,167 @@
+"""``qenvs`` command line interface (no extra dependencies; ``--json`` on every command).
+
+Designed to be read by humans and by coding agents alike::
+
+    qenvs list --json
+    qenvs info qubit/StateTransfer-v0 --json
+    qenvs run qubit/StateTransfer-v0 --policy pi-pulse --json
+    qenvs validate --json
+    qenvs doctor --json
+
+Errors are reported as ``{"error": {"code": ..., "message": ...}}`` with exit code 1
+when ``--json`` is given.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections.abc import Sequence
+
+from qenvs import toolkit
+
+
+def _emit(obj: object, as_json: bool, text: str) -> None:
+    if as_json:
+        json.dump(obj, sys.stdout, indent=2, default=str)
+        sys.stdout.write("\n")
+    else:
+        sys.stdout.write(text + "\n")
+
+
+def _parse_kv(items: list[str] | None) -> dict[str, object]:
+    out: dict[str, object] = {}
+    for item in items or []:
+        if "=" not in item:
+            msg = f"expected key=value, got {item!r}"
+            raise argparse.ArgumentTypeError(msg)
+        k, v = item.split("=", 1)
+        try:
+            out[k] = json.loads(v)
+        except json.JSONDecodeError:
+            out[k] = v
+    return out
+
+
+def cmd_list(args: argparse.Namespace) -> int:
+    """List registered tasks."""
+    tasks = toolkit.list_tasks()
+    text = "\n".join(f"{t['id']:32s} {t['model']:24s} {t['description']}" for t in tasks)
+    _emit(tasks, args.json, text)
+    return 0
+
+
+def cmd_info(args: argparse.Namespace) -> int:
+    """Describe one task."""
+    info = toolkit.task_info(args.task)
+    lines = [f"{info['id']}  ({info['model']})", "", info["description"], "", "kwargs:"]
+    lines += [f"  {k}: {v}" for k, v in info["kwargs"].items()]
+    lines += ["", "defaults: " + json.dumps(info["defaults"]), "", "references:"]
+    lines += [f"  - {r}" for r in info["references"]]
+    _emit(info, args.json, "\n".join(lines))
+    return 0
+
+
+def cmd_run(args: argparse.Namespace) -> int:
+    """Open-loop rollout with a reference policy."""
+    out = toolkit.run_rollout(args.task, policy=args.policy, seed=args.seed, **_parse_kv(args.set))
+    text = (
+        f"{out['task']} policy={out['policy']} horizon={out['horizon']} dt={out['dt']}\n"
+        f"return={out['return']:.6f}  final={json.dumps(out['final'])}"
+    )
+    _emit(out, args.json, text)
+    return 0
+
+
+def cmd_validate(args: argparse.Namespace) -> int:
+    """Run physics checks and print errors vs tolerances."""
+    out = toolkit.validate()
+    rows = [f"{'check':30s} {'category':16s} {'error':>10s} {'tol':>8s}  status"]
+    for c in out["checks"]:
+        status = "PASS" if c["passed"] else "FAIL"
+        rows.append(
+            f"{c['name']:30s} {c['category']:16s} "
+            f"{c['error']:10.2e} {c['tolerance']:8.0e}  {status}"
+        )
+    rows.append("")
+    rows += [f"{k}: {v}" for k, v in out["environment"].items()]
+    _emit(out, args.json, "\n".join(rows))
+    return 0 if out["all_passed"] else 1
+
+
+def cmd_doctor(args: argparse.Namespace) -> int:
+    """Report installation health."""
+    out = toolkit.doctor()
+    lines = [f"{k}: {v}" for k, v in out["environment"].items()]
+    lines += ["", "optional:"] + [f"  {k}: {v}" for k, v in out["optional"].items()]
+    lines += ["", f"smoke: {out['smoke']}"]
+    _emit(out, args.json, "\n".join(lines))
+    return 0 if out["ok"] else 1
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Construct the argument parser."""
+    from qenvs import __version__
+
+    p = argparse.ArgumentParser(prog="qenvs", description=__doc__.split("\n\n")[0])
+    p.add_argument("--version", action="version", version=f"qenvs {__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    s = sub.add_parser("list", help="list registered tasks")
+    s.add_argument("--json", action="store_true")
+    s.set_defaults(func=cmd_list)
+
+    s = sub.add_parser("info", help="describe a task (physics, kwargs, references)")
+    s.add_argument("task")
+    s.add_argument("--json", action="store_true")
+    s.set_defaults(func=cmd_info)
+
+    s = sub.add_parser("run", help="open-loop rollout with a reference policy")
+    s.add_argument("task")
+    s.add_argument("--policy", choices=toolkit.POLICIES, default="pi-pulse")
+    s.add_argument("--seed", type=int, default=0)
+    s.add_argument(
+        "--set", action="append", metavar="KEY=VALUE", help="task kwargs, e.g. --set gamma1=0.1"
+    )
+    s.add_argument("--json", action="store_true")
+    s.set_defaults(func=cmd_run)
+
+    s = sub.add_parser("validate", help="run physics validation checks and report errors")
+    s.add_argument("--json", action="store_true")
+    s.set_defaults(func=cmd_validate)
+
+    s = sub.add_parser("doctor", help="check installation, backend and optional extras")
+    s.add_argument("--json", action="store_true")
+    s.set_defaults(func=cmd_doctor)
+    return p
+
+
+ERROR_CODES: dict[type[Exception], str] = {
+    KeyError: "unknown_task",
+    ValueError: "invalid_argument",
+    argparse.ArgumentTypeError: "invalid_argument",
+}
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Entry point (enables JAX x64 so reported errors reflect double precision)."""
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+    args = build_parser().parse_args(argv)
+    try:
+        return int(args.func(args))
+    except (KeyError, ValueError, argparse.ArgumentTypeError) as e:
+        code = next(c for t, c in ERROR_CODES.items() if isinstance(e, t))
+        msg = e.args[0] if e.args else str(e)
+        if getattr(args, "json", False):
+            json.dump({"error": {"code": code, "message": msg}}, sys.stdout)
+            sys.stdout.write("\n")
+        else:
+            sys.stderr.write(f"error [{code}]: {msg}\n")
+        return 1
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

qenvs/envs/__init__.py

@@ -0,0 +1,7 @@
+"""Functional (JAX-native) control environments."""
+
+from qenvs.envs.base import Environment, EnvState, TimeStep
+from qenvs.envs.battery import BatteryCharging
+from qenvs.envs.state_transfer import StateTransfer
+
+__all__ = ["BatteryCharging", "EnvState", "Environment", "StateTransfer", "TimeStep"]

qenvs/envs/base.py

@@ -0,0 +1,71 @@
+"""Functional environment API (``reset``/``step`` return ``(state, timestep)``).
+
+The design follows Jumanji (Apache-2.0): the environment object is stateless,
+all mutable data lives in :class:`EnvState`, so ``step`` can be ``jit``-ed and
+``vmap``-ed over batches of states and actions.
+"""
+
+from __future__ import annotations
+
+import abc
+
+import equinox as eqx
+from dynamiqs import QArray
+from jax import Array
+
+
+class EnvState(eqx.Module):
+    """Everything needed to advance the environment."""
+
+    qstate: QArray
+    t: Array
+    step: Array
+    key: Array
+
+
+class TimeStep(eqx.Module):
+    """Output of ``reset``/``step``.
+
+    Attributes:
+        observation: real-valued observation vector.
+        reward: scalar reward.
+        done: whether the episode terminated at this step.
+        info: dictionary of named diagnostics (observable expectation values, ...).
+    """
+
+    observation: Array
+    reward: Array
+    done: Array
+    info: dict[str, Array]
+
+
+class Environment(eqx.Module):
+    """Abstract functional environment."""
+
+    @property
+    @abc.abstractmethod
+    def action_dim(self) -> int:
+        """Dimension of the continuous action vector."""
+
+    @property
+    @abc.abstractmethod
+    def observation_dim(self) -> int:
+        """Dimension of the observation vector."""
+
+    @property
+    @abc.abstractmethod
+    def action_bound(self) -> float:
+        """Actions are clipped to ``[-action_bound, action_bound]``."""
+
+    @property
+    @abc.abstractmethod
+    def horizon(self) -> int:
+        """Number of control steps per episode."""
+
+    @abc.abstractmethod
+    def reset(self, key: Array) -> tuple[EnvState, TimeStep]:
+        """Start a new episode."""
+
+    @abc.abstractmethod
+    def step(self, state: EnvState, action: Array) -> tuple[EnvState, TimeStep]:
+        """Advance one control interval."""

qenvs/envs/battery.py

@@ -0,0 +1,127 @@
+"""Battery-charging task: maximise stored energy or ergotropy of ``model.battery_hamiltonian()``."""
+
+from __future__ import annotations
+
+import dynamiqs as dq
+import jax
+import jax.numpy as jnp
+from dynamiqs import QArray
+from dynamiqs.method import Method, Tsit5
+from jax import Array
+
+from qenvs.envs.base import Environment, EnvState, TimeStep
+from qenvs.models.base import Model
+from qenvs.observables import energy, ergotropy
+
+REWARDS = ("energy", "ergotropy")
+
+
+class BatteryCharging(Environment):
+    """Piecewise-constant charging control of a quantum battery.
+
+    Works for any :class:`~qenvs.models.Model` that implements
+    ``battery_hamiltonian()``. The reward is the per-step *gain* of the chosen
+    quantity, so the undiscounted return equals ``Q(final) - Q(initial)`` with
+    ``Q`` the stored energy or the ergotropy. ``info`` reports energy,
+    ergotropy, average charging power ``E(t) / t`` and purity every step.
+
+    Attributes:
+        model: physical model with a battery Hamiltonian.
+        horizon: number of control steps.
+        dt: duration of each step (control period; the integrator adapts inside it).
+        action_bound: maximum control amplitude.
+        reward: ``"energy"`` or ``"ergotropy"``.
+        method: dynamiqs ODE method.
+    """
+
+    model: Model
+    horizon: int = 20
+    dt: float = 0.1
+    action_bound: float = 2.0
+    reward: str = "energy"
+    method: Method = Tsit5()
+
+    def __check_init__(self) -> None:
+        if self.reward not in REWARDS:
+            msg = f"reward must be one of {REWARDS}, got {self.reward!r}"
+            raise ValueError(msg)
+
+    @property
+    def action_dim(self) -> int:
+        return self.model.n_controls
+
+    @property
+    def observation_dim(self) -> int:
+        n = self.model.dim
+        return 2 * (n * n if self.model.is_open else n) + 1
+
+    @property
+    def duration(self) -> float:
+        """Total episode duration ``horizon * dt``."""
+        return self.horizon * self.dt
+
+    def _quantity(self, qstate: QArray) -> Array:
+        hb = self.model.battery_hamiltonian()
+        return energy(qstate, hb) if self.reward == "energy" else ergotropy(qstate, hb)
+
+    def _observe(self, qstate: QArray, t: Array) -> Array:
+        x = qstate.to_jax().reshape(-1)
+        return jnp.concatenate([x.real, x.imag, jnp.atleast_1d(t / self.duration)])
+
+    def _info(self, qstate: QArray, t: Array) -> dict[str, Array]:
+        hb = self.model.battery_hamiltonian()
+        e = energy(qstate, hb)
+        rho = qstate.todm().to_jax()
+        info = {name: dq.expect(op, qstate).real for name, op in self.model.observables().items()}
+        info.update(
+            {
+                "energy": e,
+                "ergotropy": ergotropy(qstate, hb),
+                "power": jnp.where(t > 0, e / jnp.maximum(t, 1e-12), 0.0),
+                "purity": jnp.real(jnp.trace(rho @ rho)),
+            }
+        )
+        return info
+
+    def reset(self, key: Array) -> tuple[EnvState, TimeStep]:
+        qstate = self.model.initial_state()
+        t = jnp.asarray(0.0)
+        state = EnvState(qstate=qstate, t=t, step=jnp.asarray(0), key=key)
+        ts = TimeStep(
+            observation=self._observe(qstate, t),
+            reward=jnp.asarray(0.0),
+            done=jnp.asarray(False),
+            info=self._info(qstate, t),
+        )
+        return state, ts
+
+    def step(self, state: EnvState, action: Array) -> tuple[EnvState, TimeStep]:
+        action = jnp.clip(
+            jnp.asarray(action, dtype=jnp.float32), -self.action_bound, self.action_bound
+        )
+        h = self.model.hamiltonian(action)
+        qstate = self.model.evolve(state.qstate, h, self.dt, method=self.method)
+        q_prev = self._quantity(state.qstate)
+        q_new = self._quantity(qstate)
+        t = state.t + self.dt
+        step = state.step + 1
+        new_state = EnvState(qstate=qstate, t=t, step=step, key=state.key)
+        ts = TimeStep(
+            observation=self._observe(qstate, t),
+            reward=q_new - q_prev,
+            done=step >= self.horizon,
+            info=self._info(qstate, t),
+        )
+        return new_state, ts
+
+    def rollout(self, actions: Array, key: Array | None = None) -> tuple[EnvState, Array]:
+        """Open-loop rollout of ``actions`` of shape ``(horizon, action_dim)`` via ``lax.scan``."""
+        if key is None:
+            key = jax.random.PRNGKey(0)
+        state0, _ = self.reset(key)
+
+        def body(state: EnvState, a: Array) -> tuple[EnvState, Array]:
+            state, ts = self.step(state, a)
+            return state, ts.reward
+
+        return jax.lax.scan(body, state0, actions)

qenvs/envs/state_transfer.py

@@ -0,0 +1,110 @@
+"""State-transfer task: drive ``model`` from its initial state to ``target``."""
+
+from __future__ import annotations
+
+import dynamiqs as dq
+import jax
+import jax.numpy as jnp
+from dynamiqs import QArray
+from dynamiqs.method import Method, Tsit5
+from jax import Array
+
+from qenvs.envs.base import Environment, EnvState, TimeStep
+from qenvs.models.base import Model
+from qenvs.observables import fidelity
+
+
+class StateTransfer(Environment):
+    """Piecewise-constant control towards a target state.
+
+    Each step applies the action as constant control amplitudes for ``dt``. The
+    observation is the flattened real/imaginary parts of the quantum state plus
+    the normalised time. The reward is the fidelity *gain* of the step, so the
+    undiscounted return equals ``F(final) - F(initial)``.
+
+    Attributes:
+        model: physical model.
+        target: target state (ket or density matrix).
+        horizon: number of control steps.
+        dt: duration of each step.
+        action_bound: maximum control amplitude.
+        method: dynamiqs ODE method used inside each control interval. The
+            integrator step is independent of ``dt`` (adaptive by default).
+    """
+
+    model: Model
+    target: QArray
+    horizon: int = 20
+    dt: float = 0.1
+    action_bound: float = 2.0
+    method: Method = Tsit5()
+
+    @property
+    def action_dim(self) -> int:
+        return self.model.n_controls
+
+    @property
+    def observation_dim(self) -> int:
+        n = self.model.dim
+        return 2 * (n * n if self.model.is_open else n) + 1
+
+    @property
+    def duration(self) -> float:
+        """Total episode duration ``horizon * dt``."""
+        return self.horizon * self.dt
+
+    def _observe(self, qstate: QArray, t: Array) -> Array:
+        x = qstate.to_jax().reshape(-1)
+        return jnp.concatenate([x.real, x.imag, jnp.atleast_1d(t / self.duration)])
+
+    def _info(self, qstate: QArray) -> dict[str, Array]:
+        info = {name: dq.expect(op, qstate).real for name, op in self.model.observables().items()}
+        info["fidelity"] = fidelity(qstate, self.target)
+        return info
+
+    def reset(self, key: Array) -> tuple[EnvState, TimeStep]:
+        qstate = self.model.initial_state()
+        t = jnp.asarray(0.0)
+        state = EnvState(qstate=qstate, t=t, step=jnp.asarray(0), key=key)
+        ts = TimeStep(
+            observation=self._observe(qstate, t),
+            reward=jnp.asarray(0.0),
+            done=jnp.asarray(False),
+            info=self._info(qstate),
+        )
+        return state, ts
+
+    def step(self, state: EnvState, action: Array) -> tuple[EnvState, TimeStep]:
+        action = jnp.clip(
+            jnp.asarray(action, dtype=jnp.float32), -self.action_bound, self.action_bound
+        )
+        h = self.model.hamiltonian(action)
+        qstate = self.model.evolve(state.qstate, h, self.dt, method=self.method)
+        f_prev = fidelity(state.qstate, self.target)
+        f_new = fidelity(qstate, self.target)
+        t = state.t + self.dt
+        step = state.step + 1
+        new_state = EnvState(qstate=qstate, t=t, step=step, key=state.key)
+        ts = TimeStep(
+            observation=self._observe(qstate, t),
+            reward=f_new - f_prev,
+            done=step >= self.horizon,
+            info=self._info(qstate),
+        )
+        return new_state, ts
+
+    def rollout(self, actions: Array, key: Array | None = None) -> tuple[EnvState, Array]:
+        """Open-loop rollout of ``actions`` of shape ``(horizon, action_dim)``.
+
+        Returns the final state and the per-step rewards. Implemented with
+        ``lax.scan`` so it is differentiable with respect to ``actions``.
+        """
+        if key is None:
+            key = jax.random.PRNGKey(0)
+        state0, _ = self.reset(key)
+
+        def body(state: EnvState, a: Array) -> tuple[EnvState, Array]:
+            state, ts = self.step(state, a)
+            return state, ts.reward
+
+        return jax.lax.scan(body, state0, actions)

qenvs/models/__init__.py

@@ -0,0 +1,6 @@
+"""Physical models: drift + control Hamiltonians, dissipators, initial states, observables."""
+
+from qenvs.models.base import Model
+from qenvs.models.two_level import TwoLevel
+
+__all__ = ["Model", "TwoLevel"]

qenvs/models/base.py

@@ -0,0 +1,123 @@
+"""Base class for physical models.
+
+A :class:`Model` is a pure description of the physics: the drift Hamiltonian,
+the control operators the agent can drive, optional Lindblad jump operators,
+the initial state and a dictionary of named observables. It knows nothing about
+rewards, horizons or action spaces; those live in :mod:`qenvs.envs`.
+
+Every model ships with analytic reference results that the validation test-suite
+checks (see ``tests/``). That is the contract behind "physically validated".
+"""
+
+from __future__ import annotations
+
+import abc
+
+import dynamiqs as dq
+import equinox as eqx
+import jax.numpy as jnp
+from dynamiqs import QArray, TimeQArray
+from dynamiqs.method import Method, Tsit5
+from jax import Array
+
+DEFAULT_METHOD: Method = Tsit5()
+
+
+class Model(eqx.Module):
+    """Abstract physical model.
+
+    Subclasses are immutable equinox modules (JAX pytrees), so they can be
+    jitted, vmapped and differentiated through.
+    """
+
+    @property
+    @abc.abstractmethod
+    def dim(self) -> int:
+        """Hilbert-space dimension."""
+
+    @abc.abstractmethod
+    def drift(self) -> QArray:
+        """Drift (uncontrolled) Hamiltonian ``H_0``."""
+
+    @abc.abstractmethod
+    def controls(self) -> list[QArray]:
+        """Control operators ``H_k``; the total Hamiltonian is ``H_0 + sum_k u_k(t) H_k``."""
+
+    def jump_ops(self) -> list[QArray]:
+        """Lindblad jump operators. Empty list means closed (unitary) dynamics."""
+        return []
+
+    @abc.abstractmethod
+    def initial_state(self) -> QArray:
+        """Initial state (ket for closed systems, density matrix otherwise)."""
+
+    def observables(self) -> dict[str, QArray]:
+        """Named Hermitian observables reported by environments (may be empty)."""
+        return {}
+
+    def battery_hamiltonian(self) -> QArray:
+        """Energy-storage Hamiltonian ``H_B`` (ground energy 0); battery models override."""
+        msg = f"{type(self).__name__} does not define a battery Hamiltonian"
+        raise NotImplementedError(msg)
+
+    # ------------------------------------------------------------------ helpers
+    @property
+    def n_controls(self) -> int:
+        """Number of control channels."""
+        return len(self.controls())
+
+    @property
+    def is_open(self) -> bool:
+        """Whether the model has dissipation."""
+        return len(self.jump_ops()) > 0
+
+    def hamiltonian(self, amplitudes: Array) -> QArray:
+        """Constant total Hamiltonian for fixed control amplitudes of shape ``(n_controls,)``."""
+        h = self.drift()
+        for k, hk in enumerate(self.controls()):
+            h = h + amplitudes[k] * hk
+        return h
+
+    def pwc_hamiltonian(self, times: Array, amplitudes: Array) -> TimeQArray:
+        """Piecewise-constant Hamiltonian.
+
+        Args:
+            times: interval boundaries, shape ``(N + 1,)``.
+            amplitudes: control amplitudes, shape ``(n_controls, N)``.
+        """
+        h: TimeQArray | QArray = self.drift()
+        for k, hk in enumerate(self.controls()):
+            h = h + dq.pwc(times, amplitudes[k], hk)
+        return h
+
+    def evolve(
+        self,
+        state: QArray,
+        hamiltonian: QArray | TimeQArray,
+        duration: float | Array,
+        *,
+        method: Method | None = None,
+    ) -> QArray:
+        """Evolve ``state`` under ``hamiltonian`` for ``duration`` and return the final state.
+
+        Uses the Schrodinger equation for closed models and the Lindblad master
+        equation otherwise. Density matrices are always accepted.
+
+        The ODE ``method`` (adaptive Tsit5 by default) takes as many internal
+        integration steps as needed inside ``duration``: the agent's control
+        period and the integrator's time step are independent. For piecewise
+        constant Hamiltonians ``dynamiqs.method.Expm()`` is exact.
+        """
+        tsave = jnp.asarray([0.0, duration])
+        method = DEFAULT_METHOD if method is None else method
+        if self.is_open:
+            # dynamiqs 0.3.6: mesolve(..., method=Expm(), save_states=False).final_state
+            # returns the *vectorised* density matrix (n^2, 1); states[-1] is always (n, n).
+            res = dq.mesolve(
+                hamiltonian, self.jump_ops(), state, tsave, method=method, progress_meter=False
+            )
+            return res.states[-1]
+        res = dq.sesolve(
+            hamiltonian, state, tsave, method=method, save_states=False, progress_meter=False
+        )
+        return res.final_state