"""One backend-neutral expression tree for authored scalar and operator physics."""
from __future__ import annotations
import inspect
from collections.abc import Iterable, Iterator
from dataclasses import dataclass, field, fields, is_dataclass, replace
from math import prod
from numbers import Number
from typing import Any, Mapping
import jax.numpy as jnp
import numpy as np
from quchip.utils.jax_utils import concrete_array_module
from quchip.utils.values import copy_value, value_fingerprint
[docs]
class UnboundParameterError(ValueError):
"""Numerical materialization started without all required values."""
[docs]
def is_opaque_callable(value: Any) -> bool:
"""Return whether *value* is a callable authoring function, not a matrix-like object."""
return callable(value) and getattr(value, "shape", None) is None
[docs]
@dataclass(frozen=True)
class PhysicsExpr:
"""Authored scalar and operator algebra, independent of numerical values.
Parameters
----------
kind : str
Expression node kind, such as ``"op"``, ``"add"``, ``"parameter"``
or ``"matrix"``.
args : tuple, default ()
Node payload in the representation that ``kind`` requires.
labels : tuple[str, ...], default ()
Ordered endpoint labels for this expression's support.
"""
kind: str
args: tuple[Any, ...] = ()
labels: tuple[str, ...] = ()
_bindings: Mapping[str, Any] = field(default_factory=dict, compare=False, repr=False)
[docs]
@classmethod
def parameter(
cls,
*,
scope: str,
name: str,
symbol: str | None = None,
unit: str | None = None,
) -> "PhysicsExpr":
"""Create a symbolic declared-parameter leaf.
Parameters
----------
scope, name : str
Dotted owner scope and declared field name.
symbol, unit : str or None
Optional display symbol and unit metadata.
"""
return cls("parameter", (f"{scope}.{name}", symbol or name, unit))
[docs]
@classmethod
def literal(cls, value: Any) -> "PhysicsExpr":
"""Create a literal scalar leaf.
Parameters
----------
value : scalar
Numerical constant copied into the expression.
"""
return cls("literal", (copy_value(value, readonly=True),))
[docs]
@classmethod
def from_matrix(
cls,
value: Any,
*,
labels: tuple[str, ...],
dims: tuple[int, ...],
name: str | None = None,
excitation_changes: Iterable[int] | None = None,
) -> "PhysicsExpr":
"""Create a named backend-neutral matrix contribution.
Parameters
----------
value : array-like
Matrix payload with one local block per ``dims`` entry.
labels, dims : tuple
Matching ordered endpoint labels and positive dimensions.
name : str or None
Optional display name.
excitation_changes : iterable of int or None
Total energy-level changes that the matrix can carry (column minus
row in the captured energy bases). Band decomposition treats all
other total changes as structural zeros, even for traced payloads.
``None`` declares no structure.
"""
if len(labels) != len(dims):
raise ValueError("Matrix labels and dimensions must have the same length.")
changes = None if excitation_changes is None else frozenset(int(c) for c in excitation_changes)
return cls("matrix", (copy_value(value, readonly=True), tuple(dims), name, changes), tuple(labels))
[docs]
@classmethod
def from_function(
cls,
function: Any,
*arguments: Any,
labels: tuple[str, ...],
dims: tuple[int, ...],
name: str | None = None,
) -> "PhysicsExpr":
"""Create an opaque matrix-valued contribution from a pure function.
The function runs only during numerical materialization. Display keeps
its declared name and arguments, such as ``X(a, b)``, and does not show
the implementation as symbolic algebra.
Parameters
----------
function : callable
Pure callable evaluated during numerical materialization.
*arguments : Any
Expression arguments passed to ``function``.
labels, dims : tuple
Matching endpoint labels and dimensions.
name : str or None
Display name, required for anonymous callables.
"""
if len(labels) != len(dims):
raise ValueError("Function labels and dimensions must have the same length.")
if not callable(function):
raise TypeError("function must be callable.")
display_name = name or getattr(function, "__name__", None)
if not display_name or display_name == "<lambda>":
raise ValueError("Anonymous functions require a symbolic name.")
return cls(
"function",
(function, tuple(dims), display_name, *(ensure_expr(arg) for arg in arguments)),
tuple(labels),
)
[docs]
@classmethod
def from_state(
cls,
value: Any,
*,
labels: tuple[str, ...],
dims: tuple[int, ...],
name: str,
) -> "PhysicsExpr":
"""Create a named authored ket contribution.
Parameters
----------
value : array-like
Ket payload.
labels, dims : tuple
Matching endpoint labels and dimensions.
name : str
Display name.
"""
return cls("state", (copy_value(value, readonly=True), tuple(dims), name), tuple(labels))
[docs]
@classmethod
def from_state_function(
cls,
function: Any,
*arguments: Any,
labels: tuple[str, ...],
dims: tuple[int, ...],
name: str,
) -> "PhysicsExpr":
"""Create an opaque callable ket contribution.
Parameters
----------
function : callable
Callable that returns a ket during materialization.
*arguments : Any
Expression arguments passed to ``function``.
labels, dims : tuple
Matching endpoint labels and dimensions.
name : str
Display name.
"""
if not callable(function):
raise TypeError("function must be callable.")
return cls(
"state_function",
(function, tuple(dims), name, *(ensure_expr(arg) for arg in arguments)),
tuple(labels),
)
[docs]
@classmethod
def from_signal(cls, signal: Any, *, name: str = "f") -> "PhysicsExpr":
"""Create a scalar time-function leaf backed by an engine signal.
Parameters
----------
signal : object
Engine signal object evaluated at runtime.
name : str, default ``"f"``
Display name for the signal node.
"""
return cls("signal", (signal, name))
[docs]
def embed(self, labels: tuple[str, ...], dims: tuple[int, ...]) -> "PhysicsExpr":
"""Embed this local contribution into an ordered composite Hilbert space.
Parameters
----------
labels, dims : tuple
Complete ordered composite labels and matching dimensions.
"""
if len(labels) != len(dims):
raise ValueError("Composite labels and dimensions must have the same length.")
missing = set(self.labels) - set(labels)
if missing:
raise ValueError(f"Cannot embed labels absent from the composite space: {sorted(missing)}")
return PhysicsExpr("embed", (self, tuple(labels), tuple(dims)), tuple(labels))
[docs]
def with_bindings(self, bindings: Mapping[str, Any]) -> "PhysicsExpr":
"""Attach default values that only direct numerical inspection uses.
Parameters
----------
bindings : mapping[str, Any]
Dotted parameter paths to numerical values.
"""
return replace(self, _bindings=copy_value(dict(bindings), readonly=True))
[docs]
def parameter_paths(self) -> tuple[str, ...]:
"""Return referenced dotted parameter paths in authored order."""
return tuple(dict.fromkeys(
node.args[0] for node in _walk_expr(self) if node.kind == "parameter"
))
[docs]
def numeric_values(self) -> tuple[Any, ...]:
"""Return bound and matrix payloads for tracer-safe cache decisions."""
values: list[Any] = []
for node in _walk_expr(self):
values.extend(node._bindings.values())
if node.kind in ("literal", "matrix", "state", "signal"):
values.append(node.args[0])
if node.kind == "level" and not isinstance(node.args[0], PhysicsExpr):
values.append(node.args[0])
return tuple(values)
@property
def shape(self) -> tuple[int, int]:
"""Matrix shape implied by this operator expression's static support."""
if not self.labels:
raise AttributeError("Scalar expressions do not have a matrix shape.")
if self.kind == "level":
return self.args[0].shape
if self.kind == "op":
dimension = self.args[1].dimension
return (dimension, dimension)
if self.kind in ("matrix", "function"):
dimension = prod(self.args[1])
return (dimension, dimension)
if self.kind in ("state", "state_function"):
return (prod(self.args[1]), 1)
if self.kind == "embed":
dimension = prod(self.args[2])
return (dimension, dimension)
if self.kind == "tensor":
dimension = self.args[0].shape[0] * self.args[1].shape[0]
return (dimension, dimension)
for arg in self.args:
if isinstance(arg, PhysicsExpr) and arg.labels:
return arg.shape
raise AttributeError("Operator shape is not available for this expression.")
def _binary(self, other: Any, kind: str) -> "PhysicsExpr":
rhs = ensure_expr(other)
if kind in ("add", "sub"):
if bool(self.labels) != bool(rhs.labels):
raise TypeError(
"Cannot add a scalar directly to an operator; write the identity explicitly."
)
if self.labels != rhs.labels:
raise TypeError(
"Addition and subtraction require operands with the same endpoint support "
f"(got {self.labels!r} and {rhs.labels!r})."
)
return PhysicsExpr(
kind,
(self, rhs),
tuple(dict.fromkeys(self.labels + rhs.labels)),
)
def __add__(self, other: Any) -> "PhysicsExpr":
return self._binary(other, "add")
def __radd__(self, other: Any) -> "PhysicsExpr":
return ensure_expr(other)._binary(self, "add")
def __sub__(self, other: Any) -> "PhysicsExpr":
return self._binary(other, "sub")
def __rsub__(self, other: Any) -> "PhysicsExpr":
return ensure_expr(other)._binary(self, "sub")
def __matmul__(self, other: Any) -> "PhysicsExpr":
rhs = ensure_expr(other)
if not self.labels or self.labels != rhs.labels:
raise TypeError("Cannot use @ for operators on different endpoints; use * across endpoints.")
return self._binary(rhs, "matmul")
def __mul__(self, other: Any) -> "PhysicsExpr":
rhs = ensure_expr(other)
if self.labels and rhs.labels:
if set(self.labels) & set(rhs.labels):
raise TypeError(
"Cannot use * for operators on the same endpoint or with overlapping "
"endpoint support; use @ on one endpoint."
)
return PhysicsExpr(
"tensor",
(self, rhs),
tuple(dict.fromkeys(self.labels + rhs.labels)),
)
if self.labels or rhs.labels:
scalar, operator = (rhs, self) if self.labels else (self, rhs)
return PhysicsExpr("scale", (scalar, operator), operator.labels)
return PhysicsExpr("mul", (self, rhs))
def __rmul__(self, other: Any) -> "PhysicsExpr":
return ensure_expr(other).__mul__(self)
def __truediv__(self, other: Any) -> "PhysicsExpr":
rhs = ensure_expr(other)
if rhs.labels:
raise TypeError("Division by an operator is not defined.")
return self * PhysicsExpr("pow", (rhs, PhysicsExpr.literal(-1)))
def __rtruediv__(self, other: Any) -> "PhysicsExpr":
if self.labels:
raise TypeError("Division by an operator is not defined.")
return ensure_expr(other) * PhysicsExpr(
"pow", (self, PhysicsExpr.literal(-1))
)
def __pow__(self, other: Any) -> "PhysicsExpr":
rhs = ensure_expr(other)
if self.labels or rhs.labels:
raise TypeError("Use @ for operator powers; ** is scalar-only.")
return PhysicsExpr("pow", (self, rhs))
def __neg__(self) -> "PhysicsExpr":
return -1 * self
[docs]
def latex(self) -> str:
"""Render the authored expression with familiar mathematical notation."""
return _latex(self)
def _repr_latex_(self) -> str:
return f"${self.latex()}$"
def __str__(self) -> str:
return self.latex()
[docs]
def matrix(
self,
bindings: Mapping[str, Any] | None = None,
*,
t: Any | None = None,
backend: Any = None,
) -> Any:
"""Materialize this expression and return its dense numerical array.
Parameters
----------
bindings : mapping or None, optional
Values for unresolved parameter paths.
t : scalar or None, keyword-only
Time for signal evaluation.
backend : backend or None, keyword-only
Backend lowering object. ``None`` selects the default.
"""
if backend is None:
from quchip.backend import get_default_backend
backend = get_default_backend()
native = materialize_expr(
self,
backend,
bindings=bindings,
t=t,
)
return backend.to_array(native)
[docs]
class ParameterNamespace:
"""Attribute view that exposes one owner's declared fields as symbolic leaves."""
__slots__ = ("_scope", "_fields")
def __init__(self, scope: str, fields: Mapping[str, Any]) -> None:
self._scope = scope
self._fields = fields
def __getattr__(self, name: str) -> PhysicsExpr:
if name == "thermal_population" and "thermal_occupation" in self._fields:
from quchip.utils.deprecation import warn_renamed
warn_renamed(name, "thermal_occupation")
name = "thermal_occupation"
try:
spec = self._fields[name]
except KeyError as exc:
raise AttributeError(f"No declared parameter {name!r} on {self._scope!r}.") from exc
return PhysicsExpr.parameter(
scope=self._scope,
name=name,
symbol=spec.symbol,
unit=spec.unit,
)
def __dir__(self) -> list[str]:
return sorted(self._fields)
[docs]
def ensure_expr(value: Any) -> PhysicsExpr:
"""Coerce a scalar value into the shared expression tree."""
if isinstance(value, PhysicsExpr):
return value
if isinstance(value, (int, float, complex)) or getattr(value, "ndim", None) == 0:
return PhysicsExpr.literal(value)
raise TypeError(f"Expected a scalar or PhysicsExpr, got {type(value).__name__}.")
[docs]
def as_operator_expr(
value: Any,
*,
labels: tuple[str, ...],
dims: tuple[int, ...],
name: str,
arguments: tuple[Any, ...] = (),
owner: Any | None = None,
scope: str | None = None,
allowed: Mapping[str, Any] | None = None,
) -> PhysicsExpr:
"""Normalize symbolic, matrix, or opaque callable operator authorship.
Parameters
----------
value : PhysicsExpr, callable, or matrix
Authored operator representation.
labels, dims : tuple
Ordered endpoint labels and matching dimensions.
name : str
Display name for opaque or matrix contributions.
arguments : tuple, default ()
Arguments for an opaque callable.
owner : object or None, default None
Object whose fields bind the positional argument names of a callable.
Do not combine with explicit ``arguments``.
scope : str or None, default None
Dotted-path prefix for owner-bound parameters, such as a device label.
Required with ``owner`` when the callable has arguments.
allowed : mapping or None, default None
Permitted owner-field names, given as mapping keys whose values are not
used. ``None`` accepts all named owner fields.
"""
expected = (prod(dims), prod(dims))
if isinstance(value, PhysicsExpr):
if not value.labels:
raise TypeError("An operator expression must carry at least one subsystem label.")
if value.labels != labels:
raise TypeError(
f"Operator expression support {value.labels!r} does not match {labels!r}."
)
if value.shape != expected:
raise ValueError(
f"Operator expression has shape {value.shape}, expected {expected}."
)
return value
if is_opaque_callable(value):
arguments, bindings = _resolve_callable_arguments(
value, arguments=arguments, owner=owner, scope=scope, allowed=allowed
)
return PhysicsExpr.from_function(
value,
*arguments,
labels=labels,
dims=dims,
name=name,
).with_bindings(bindings)
if getattr(value, "shape", None) != expected:
raise TypeError(
f"Operator must be a PhysicsExpr, callable, or matrix with shape {expected}; "
f"got {type(value).__name__} with shape {getattr(value, 'shape', None)}."
)
return PhysicsExpr.from_matrix(value, labels=labels, dims=dims, name=name)
[docs]
def as_scalar_expr(
value: Any,
*,
name: str,
arguments: tuple[Any, ...] = (),
owner: Any | None = None,
scope: str | None = None,
allowed: Mapping[str, Any] | None = None,
) -> PhysicsExpr:
"""Normalize symbolic, numeric, or opaque callable scalar authorship.
Parameters
----------
value : PhysicsExpr, scalar, or callable
Authored scalar representation.
name : str
Display name for opaque callables.
arguments : tuple, default ()
Arguments for an opaque callable.
owner : object or None, default None
Object whose fields bind the positional argument names of a callable.
Do not combine with explicit ``arguments``.
scope : str or None, default None
Dotted-path prefix for owner-bound parameters, such as a device label.
Required with ``owner`` when the callable has arguments.
allowed : mapping or None, default None
Permitted owner-field names, given as mapping keys whose values are not
used. ``None`` accepts all named owner fields.
"""
if isinstance(value, PhysicsExpr):
if value.labels:
raise TypeError("A scalar expression cannot carry subsystem labels.")
return value
if is_opaque_callable(value):
arguments, bindings = _resolve_callable_arguments(
value, arguments=arguments, owner=owner, scope=scope, allowed=allowed
)
return PhysicsExpr.from_function(
value,
*arguments,
labels=(),
dims=(),
name=name,
).with_bindings(bindings)
return ensure_expr(value)
[docs]
def as_state_expr(
value: Any,
*,
labels: tuple[str, ...],
dims: tuple[int, ...],
name: str,
arguments: tuple[Any, ...] = (),
owner: Any | None = None,
scope: str | None = None,
allowed: Mapping[str, Any] | None = None,
) -> PhysicsExpr:
"""Normalize an authored ket array or opaque callable without evaluating it.
Parameters
----------
value : PhysicsExpr, callable, or array
Authored ket representation.
labels, dims : tuple
Ordered endpoint labels and matching dimensions.
name : str
Display name for the state contribution.
arguments : tuple, default ()
Arguments for an opaque callable.
owner : object or None, default None
Object whose fields bind the positional argument names of a callable.
Do not combine with explicit ``arguments``.
scope : str or None, default None
Dotted-path prefix for owner-bound parameters, such as a device label.
Required with ``owner`` when the callable has arguments.
allowed : mapping or None, default None
Permitted owner-field names, given as mapping keys whose values are not
used. ``None`` accepts all named owner fields.
"""
expected = (prod(dims), 1)
if isinstance(value, PhysicsExpr):
if value.labels != labels:
raise ValueError(
f"State expression support {value.labels!r} does not match {labels!r}."
)
if value.shape != expected:
raise ValueError(
f"State expression has shape {value.shape}, expected {expected}."
)
return value
if is_opaque_callable(value):
arguments, bindings = _resolve_callable_arguments(
value, arguments=arguments, owner=owner, scope=scope, allowed=allowed
)
return PhysicsExpr.from_state_function(
value,
*arguments,
labels=labels,
dims=dims,
name=name,
).with_bindings(bindings)
shape = getattr(value, "shape", None)
if shape not in (expected, (expected[0],)):
raise TypeError(
f"State must be a PhysicsExpr, callable, or ket with shape {expected} or {(expected[0],)}; "
f"got {type(value).__name__} with shape {shape}."
)
return PhysicsExpr.from_state(value, labels=labels, dims=dims, name=name)
def _resolve_callable_arguments(
function: Any,
*,
arguments: tuple[Any, ...],
owner: Any | None,
scope: str | None,
allowed: Mapping[str, Any] | None = None,
) -> tuple[tuple[PhysicsExpr, ...], dict[str, Any]]:
"""Resolve explicit callable arguments or bind their names to an owner."""
if not is_opaque_callable(function):
return (), {}
if arguments:
if owner is not None:
raise TypeError("Pass explicit callable arguments or an owner, not both.")
return tuple(ensure_expr(argument) for argument in arguments), {}
if owner is None or scope is None:
if inspect.signature(function).parameters:
raise TypeError("Opaque functions with parameters require an owner and scope.")
return (), {}
fields = getattr(type(owner), "__quchip_param_fields__", {})
resolved_arguments: list[PhysicsExpr] = []
bindings: dict[str, Any] = {}
for item in inspect.signature(function).parameters.values():
if item.kind not in (item.POSITIONAL_ONLY, item.POSITIONAL_OR_KEYWORD):
raise TypeError("Opaque functions may only declare positional parameters.")
if allowed is not None and item.name not in allowed:
raise ValueError(
f"Opaque function argument {item.name!r} is not declared; "
f"available fields are {sorted(allowed)}."
)
if not hasattr(owner, item.name):
raise ValueError(f"Opaque function argument {item.name!r} is not a field on {scope!r}.")
spec = fields.get(item.name)
path = f"{scope}.{item.name}"
resolved_arguments.append(
PhysicsExpr.parameter(
scope=scope,
name=item.name,
symbol=getattr(spec, "symbol", None) or item.name,
unit=getattr(spec, "unit", None),
)
)
bindings[path] = getattr(owner, item.name)
return tuple(resolved_arguments), bindings
def _walk_expr(expr: PhysicsExpr) -> Iterator[PhysicsExpr]:
"""Yield an expression tree in authored preorder."""
yield expr
for arg in expr.args:
if isinstance(arg, PhysicsExpr):
yield from _walk_expr(arg)
[docs]
def split_dynamic_hamiltonian(expr: PhysicsExpr) -> tuple[tuple[PhysicsExpr, PhysicsExpr], ...]:
"""Split a linear Hamiltonian sum into scalar-signal and operator factors.
Signal algebra can be nonlinear and can contain more than one signal leaf.
Only the quantum operator must be linear. Each additive term must be one
operator-valued factor times a scalar expression that depends on at least
one delivered signal.
"""
if not isinstance(expr, PhysicsExpr) or not expr.labels:
raise TypeError("A drive Hamiltonian must return an operator-valued PhysicsExpr.")
def additive_terms(node: PhysicsExpr) -> list[tuple[int, PhysicsExpr]]:
if node.kind == "add":
return additive_terms(node.args[0]) + additive_terms(node.args[1])
if node.kind == "sub":
return additive_terms(node.args[0]) + [
(-sign, term) for sign, term in additive_terms(node.args[1])
]
return [(1, node)]
terms: list[tuple[PhysicsExpr, PhysicsExpr]] = []
for sign, term in additive_terms(expr):
if term.kind != "scale":
raise TypeError(
"Each drive Hamiltonian term must multiply a delivered signal by a quantum operator."
)
scalar, operator = term.args
if scalar.labels or not operator.labels:
raise TypeError("Drive Hamiltonian terms must be scalar-signal times operator.")
if not any(node.kind == "signal" for node in _walk_expr(scalar)):
raise TypeError("Drive Hamiltonian scalar factors must depend on the delivered signal.")
terms.append((scalar if sign > 0 else -scalar, operator))
return tuple(terms)
[docs]
def scalar_signal_program(expr: PhysicsExpr) -> Any:
"""Lower scalar signal algebra into the backend-neutral signal program."""
from quchip.engine.ir import Add, Constant, Multiply, Scale, SignalPower
values = _bound_values(expr)
def lower(node: PhysicsExpr) -> Any:
if node.labels:
raise TypeError("Signal-program lowering accepts scalar expressions only.")
if node.kind == "signal":
return node.args[0]
if node.kind == "literal":
return Constant(node.args[0])
if node.kind == "parameter":
try:
return Constant(values[node.args[0]])
except KeyError as exc:
raise UnboundParameterError(
f"Missing numerical binding: {node.args[0]}"
) from exc
if node.kind == "add":
return Add((lower(node.args[0]), lower(node.args[1])))
if node.kind == "sub":
return Add((lower(node.args[0]), Scale(lower(node.args[1]), -1.0)))
if node.kind == "mul":
return Multiply((lower(node.args[0]), lower(node.args[1])))
if node.kind == "pow":
exponent = node.args[1]
if exponent.kind != "literal":
raise TypeError("Signal powers require a literal exponent.")
return SignalPower(lower(node.args[0]), exponent.args[0])
raise TypeError(f"Unsupported scalar signal expression kind {node.kind!r}.")
return lower(expr)
def _bound_values(expr: PhysicsExpr, bindings: Mapping[str, Any] | None = None) -> dict[str, Any]:
values: dict[str, Any] = {}
for node in _walk_expr(expr):
values.update(node._bindings)
if bindings is not None:
values.update(bindings)
return values
def _same_authored_value(first: Any, second: Any) -> bool:
"""Compare known authored structure and values, using identity for tracers."""
from quchip.utils.jax_utils import contains_tracer
if first is second:
return True
if type(first) is not type(second):
return False
if is_dataclass(first) and not isinstance(first, type):
return all(_same_authored_value(getattr(first, item.name), getattr(second, item.name))
for item in fields(first) if not (isinstance(first, PhysicsExpr) and item.name == "_bindings"))
if isinstance(first, (tuple, list)):
return len(first) == len(second) and all(_same_authored_value(a, b) for a, b in zip(first, second))
if contains_tracer((first, second)):
return False
try:
return value_fingerprint(first) == value_fingerprint(second)
except (ValueError, RecursionError):
return False
def _matching_local_basis(
level: PhysicsExpr, bases: Mapping[str, Any] | None, bindings: Mapping[str, Any],
) -> Any | None:
record = None if bases is None else bases.get(level.labels[0])
if record is None or not _same_authored_value(level.args[0], record.authored_hamiltonian):
return None
hamiltonian = record.authored_hamiltonian
if isinstance(hamiltonian, PhysicsExpr):
expected = _bound_values(hamiltonian)
if not all(path in bindings and path in expected and _same_authored_value(bindings[path], expected[path])
for path in hamiltonian.parameter_paths()):
return None
return record
[docs]
def declared_excitation_changes(expr: Any) -> frozenset[int] | None:
"""Return the total energy-level changes a matrix contribution declares, if any."""
if isinstance(expr, PhysicsExpr) and expr.kind == "matrix":
return expr.args[3]
return None
def _authored_entry_support(expr: Any) -> np.ndarray | None:
"""Return a bound on the nonzero entries of authored operator algebra, or ``None``.
Scalar coefficients do not change the bound, so it also holds for traced values.
"""
import jax
from quchip.utils.jax_utils import contains_tracer
if isinstance(expr, PhysicsExpr):
if expr.kind == "scale":
return _authored_entry_support(expr.args[1])
if expr.kind in ("add", "sub", "matmul", "tensor"):
left, right = (_authored_entry_support(arg) for arg in expr.args)
if left is None or right is None:
return None
if expr.kind in ("add", "sub"):
return left | right
return left @ right if expr.kind == "matmul" else np.kron(left, right)
if expr.kind not in ("op", "matrix") or contains_tracer(expr.numeric_values()):
return None
with jax.ensure_compile_time_eval():
expr = materialize_array(expr)
if contains_tracer(expr):
return None
values = np.asarray(expr)
return values != 0 if values.ndim == 2 else None
[docs]
def is_energy_diagonal(expr: Any, bases: Mapping[str, Any]) -> bool:
"""Return whether authored algebra guarantees zero energy-change weight in captured bases."""
from quchip.devices.spaces import ChargeSpace, FockSpace, PhaseGridSpace
if not isinstance(expr, PhysicsExpr):
return False
bindings = _bound_values(expr)
def diagonal(node: PhysicsExpr) -> bool:
if not node.labels:
return True
if node.kind == "level":
return _matching_local_basis(node, bases, bindings) is not None
if node.kind == "op":
return node.args[0] == "I" and type(node.args[1]) in (FockSpace, ChargeSpace, PhaseGridSpace)
if node.kind in ("add", "sub", "scale", "tensor", "matmul", "embed"):
return all(diagonal(arg) for arg in node.args if isinstance(arg, PhysicsExpr))
return False
return diagonal(expr)
[docs]
def materialize_expr(
expr: Any,
backend: Any,
*,
bindings: Mapping[str, Any] | None = None,
t: Any | None = None,
local_bases: Mapping[str, Any] | None = None,
) -> Any:
"""Lower symbolic physics, and pass an already-native contribution through unchanged."""
if not isinstance(expr, PhysicsExpr):
return expr
values = _bound_values(expr, bindings)
missing = [path for path in expr.parameter_paths() if path not in values]
if missing:
raise UnboundParameterError("Missing numerical bindings: " + ", ".join(missing))
combine = getattr(backend, "linear_combination", None)
def lower(node: PhysicsExpr) -> Any:
if node.kind == "literal":
return node.args[0]
if node.kind == "parameter":
return values[node.args[0]]
if node.kind == "signal":
if t is None:
raise ValueError("t is required to materialize a time-dependent expression.")
from quchip.engine.ir import evaluate_signal_program
return evaluate_signal_program(node.args[0], t, xp=backend.array_module)
if node.kind == "matrix":
value, dims, *_ = node.args
return backend.from_array(
value,
dims=[list(dims), list(dims)],
)
if node.kind == "state":
value, dims, _name = node.args
return backend.from_array(
value,
dims=[list(dims), [1]],
)
if node.kind == "function":
function, dims, _name, *arguments = node.args
value = function(*(lower(argument) for argument in arguments))
if not dims:
return value
return backend.from_array(value, dims=[list(dims), list(dims)])
if node.kind == "state_function":
function, dims, _name, *arguments = node.args
value = function(*(lower(argument) for argument in arguments))
return backend.from_array(value, dims=[list(dims), [1]])
if node.kind == "op":
name, space = node.args
return space.operator(name, backend)
if node.kind == "level":
import jax
from quchip.engine.basis import resolve_local_basis
record = _matching_local_basis(node, local_bases, values)
if record is None:
# Concrete authored values give a concrete basis, also inside a JAX trace.
with jax.ensure_compile_time_eval():
record = resolve_local_basis(materialize_array(node.args[0], bindings=values),
authored_hamiltonian=node.args[0])
if record.kind == "native" and record.energy_to_solver() is None:
return backend.number(record.native_dim)
return backend.from_array(record.authored_level_operator(),
dims=[[record.native_dim], [record.native_dim]])
if node.kind == "embed":
local = lower(node.args[0])
labels, dims = node.args[1:]
support = tuple(labels.index(label) for label in node.args[0].labels)
if len(support) == 1:
return backend.embed(local, support[0], dims)
if len(support) == 2:
return backend.embed_two_body(local, support[0], support[1], dims)
if support == tuple(range(len(dims))):
return local
raise ValueError(f"Cannot embed a contribution with support {support}.")
left = lower(node.args[0])
right = lower(node.args[1])
# A backend that combines operators itself can form concrete sums
# and scalings without device programs.
if node.kind in ("add", "sub"):
sign = 1 if node.kind == "add" else -1
if combine is not None and not (_is_scalar(left) or _is_scalar(right)):
return combine(((1, left), (sign, right)))
return left + right if sign == 1 else left - right
if node.kind == "matmul":
return backend.matmul(left, right)
if node.kind == "tensor":
return backend.tensor(left, right)
if node.kind in ("scale", "mul"):
if combine is not None and _is_scalar(left) != _is_scalar(right):
scalar, operator = (left, right) if _is_scalar(left) else (right, left)
return combine(((scalar, operator),))
return left * right
if node.kind == "pow":
exponent = node.args[1]
if exponent.kind == "literal" and exponent.args[0] == -1:
return 1.0 / left
return left ** right
raise TypeError(f"Unknown PhysicsExpr kind {node.kind!r}.")
return lower(expr)
def _is_scalar(value: Any) -> bool:
"""Whether a lowered value is a number rather than an operator."""
return isinstance(value, Number) or getattr(value, "shape", None) == ()
class _ArrayLowerer:
"""Minimal operator algebra for backend-independent JAX materialization.
Concrete operands combine on the host, so a constant operator compiles no
XLA programs; traced operands stay in JAX.
"""
array_module = jnp
@staticmethod
def from_array(value: Any, dims: Any = None) -> Any:
del dims
if hasattr(value, "to_jax"):
value = value.to_jax()
elif hasattr(value, "full"):
value = value.full()
return concrete_array_module(value).asarray(value, dtype=complex)
@staticmethod
def to_array(value: Any) -> Any:
if hasattr(value, "to_jax"):
value = value.to_jax()
elif hasattr(value, "full"):
value = value.full()
return jnp.asarray(value)
@staticmethod
def destroy(dimension: int) -> Any:
return np.diag(np.sqrt(np.arange(1, dimension)), 1).astype(complex)
@staticmethod
def create(dimension: int) -> Any:
return np.diag(np.sqrt(np.arange(1, dimension)), -1).astype(complex)
@staticmethod
def number(dimension: int) -> Any:
return np.diag(np.arange(dimension, dtype=complex))
@staticmethod
def identity(dimension: int) -> Any:
return np.eye(dimension, dtype=complex)
@staticmethod
def dag(value: Any) -> Any:
return concrete_array_module(value).asarray(value).conj().T
@staticmethod
def matmul(left: Any, right: Any) -> Any:
return left @ right
@staticmethod
def tensor(left: Any, right: Any) -> Any:
return concrete_array_module(left, right).kron(left, right)
@staticmethod
def embed(local: Any, target: int, dims: tuple[int, ...]) -> Any:
xp = concrete_array_module(local)
factors = [xp.eye(dim, dtype=complex) for dim in dims]
factors[target] = local
result = factors[0]
for factor in factors[1:]:
result = xp.kron(result, factor)
return result
@staticmethod
def embed_two_body(local: Any, first: int, second: int, dims: tuple[int, ...]) -> Any:
from quchip.backend._dims import _embed_array
return _embed_array(local, (first, second), dims, concrete_array_module(local))
_ARRAY_LOWERER = _ArrayLowerer()
[docs]
def materialize_array(
expr: Any,
*,
bindings: Mapping[str, Any] | None = None,
t: Any | None = None,
) -> Any:
"""Materialize authored physics as a backend-independent JAX array."""
return _ARRAY_LOWERER.to_array(
materialize_expr(expr, _ARRAY_LOWERER, bindings=bindings, t=t)
)
[docs]
def materialize_scalar(
expr: Any,
*,
bindings: Mapping[str, Any] | None = None,
) -> Any:
"""Materialize a scalar expression without coercing traced values."""
return materialize_expr(expr, _ARRAY_LOWERER, bindings=bindings)
def _latex(expr: PhysicsExpr, parent_precedence: int = 0) -> str:
if expr.kind == "level":
return rf"\hat \ell_{{{expr.labels[0]}}}"
if expr.kind == "literal":
value = expr.args[0]
if isinstance(value, complex):
return _complex_latex(value)
return f"{value:g}" if isinstance(value, (int, float)) else str(value)
if expr.kind == "parameter":
path, symbol, _unit = expr.args
scope = path.rsplit(".", 1)[0]
return _scoped_symbol(symbol, scope)
if expr.kind == "signal":
_signal, name = expr.args
return rf"{name}\!\left(t\right)"
if expr.kind == "matrix":
name = expr.args[2]
if name is not None:
return name
return rf"\hat H_{{{','.join(expr.labels)}}}"
if expr.kind == "state":
_value, _dims, name = expr.args
return name
if expr.kind == "function":
_function, _dims, name, *arguments = expr.args
rendered = ", ".join(_latex(argument) for argument in arguments)
return rf"{name}\!\left({rendered}\right)"
if expr.kind == "state_function":
_function, _dims, name, *arguments = expr.args
rendered = ", ".join(_latex(argument) for argument in arguments)
return rf"{name}\!\left({rendered}\right)"
if expr.kind == "op":
name, _space = expr.args
scope = expr.labels[0]
symbols = {
"a": r"\hat a",
"adag": r"\hat a^\dagger",
"n": r"\hat n",
"I": r"\hat I",
"sigma_x": r"\hat\sigma_x",
"sigma_y": r"\hat\sigma_y",
"sigma_z": r"\hat\sigma_z",
"sigma_plus": r"\hat\sigma_+",
"sigma_minus": r"\hat\sigma_-",
}
symbol = symbols.get(name, rf"\hat{{\mathrm{{{name}}}}}")
return _scoped_symbol(symbol, scope)
if expr.kind == "embed":
return _latex(expr.args[0], parent_precedence)
if expr.kind in ("add", "sub"):
precedence = 1
elif expr.kind in ("scale", "mul", "tensor"):
precedence = 2
else:
precedence = 3
left = _latex(expr.args[0], precedence)
right = _latex(expr.args[1], precedence + (1 if expr.kind == "sub" else 0))
if expr.kind == "add":
text = f"{left} + {right}"
elif expr.kind == "sub":
text = f"{left} - {right}"
elif expr.kind in ("matmul", "tensor", "mul", "scale"):
text = rf"{left}\,{right}"
elif expr.kind == "pow":
text = f"{left}^{{{right}}}"
else:
raise TypeError(f"Unknown PhysicsExpr kind {expr.kind!r}.")
return f"({text})" if precedence < parent_precedence else text
def _complex_latex(value: complex) -> str:
"""Render a complex literal with ``i`` as the imaginary unit."""
if value.imag == 0:
return f"{value.real:g}"
imaginary = {1.0: "i", -1.0: "-i"}.get(value.imag, f"{value.imag:g}i")
if value.real == 0:
return imaginary
sign = "" if imaginary.startswith("-") else "+"
return f"({value.real:g}{sign}{imaginary})"
def _scoped_symbol(symbol: str, scope: str) -> str:
"""Attach an owner scope without producing nested LaTeX subscripts."""
if "_" not in symbol:
return f"{symbol}_{{{scope}}}"
base, subscript = symbol.split("_", 1)
subscript = subscript.removeprefix("{").removesuffix("}")
return f"{base}_{{{subscript},{scope}}}"
