"""Gough-James algebra on resolved SLH triples: series product, concatenation, feedback.
Every function returns a new :class:`~quchip.engine.ir.ResolvedSLH` on the same
Hilbert space as its operands and does not change them. Hamiltonian corrections
that a composition generates are appended as network-origin static terms.
Scattering stays scalar. Traced values pass through the array module that they
arrive in.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING, Any, Sequence
import numpy as np
from scipy.linalg import block_diag
from quchip.engine.ir import CanonicalOperator, CollapseTerm, HamiltonianProgram, ResolvedSLH, SLHChannel, StaticTerm
from quchip.engine.reference import ReferencePlane
from quchip.utils.jax_utils import contains_tracer, is_jax_array, maybe_concrete_scalar, select_array_module
if TYPE_CHECKING:
from quchip.chip.port_network import NetworkPort
__all__ = ["series_product", "concatenate", "feedback_reduce"]
_SINGULAR_TOLERANCE = 1e-12
@dataclass(frozen=True)
class ModeResponse:
"""A removed harmonic pole attached to its original quantum-port locations."""
label: str
ports: tuple[str, ...]
freq: Any
weight: Any
internal_rate: Any
@dataclass(frozen=True)
class CompiledChannel:
"""One exposure of a compiled network with its couplings by port and its reference plane."""
exposure: NetworkPort
coupling: Mapping[str, Any]
reference: ReferencePlane
@dataclass(frozen=True)
class CompiledNetwork:
"""The Markovian core of a port network, compiled for one chip."""
channels: tuple[CompiledChannel, ...]
scattering: Any
generated_pairs: tuple[tuple[str, str, Any], ...]
support: np.ndarray
output_network: Any = None
def cascade_hamiltonian(products: Sequence[tuple[Any, Any, Any]]) -> Any | None:
"""Return ``sum (L_d^dagger c L_u - h.c.) / 2i`` over ``(L_d, L_u, c)`` matrices on one space.
A concrete result is ``None`` when each element is below ``64 n ε`` times
the magnitude of the products that form it. This conservative threshold
removes cancellation residue, e.g. the Lamb shift ``γ sin(π)`` of a giant
atom, and also a genuine term that small. Each element has its own
threshold, so the cancellation of one pair does not hide the small rate of
another pair. A traced result keeps every term.
"""
traced = contains_tracer(products)
xp = select_array_module(traced)
values: Any = 0.0
scale: Any = 0.0
for left, right, coefficient in products:
left, right = xp.asarray(left), xp.asarray(right)
product = xp.conj(xp.swapaxes(left, -1, -2)) @ (xp.asarray(coefficient) * right)
values = values + (product - xp.conj(xp.swapaxes(product, -1, -2))) / (2j)
if not traced:
size = np.abs(coefficient) * (np.swapaxes(np.abs(left), -1, -2) @ np.abs(right))
scale = scale + (size + np.swapaxes(size, -1, -2)) / 2
if not traced and np.all(np.abs(values) <= 64 * np.finfo(float).eps * values.shape[-1] * scale):
return None
return values
[docs]
def series_product(first: ResolvedSLH, second: ResolvedSLH) -> ResolvedSLH:
"""Feed every output of ``first`` into the matching input of ``second``.
Implements ``G1 ▷ G2 = (S2 S1, L2 + S2 L1, H1 + H2 + (L2† S2 L1 − h.c.) / 2i)``.
Paired channels must agree in accessibility and carrier. Reference runs are
joined leg by leg, so the outbound run of ``first`` and the inbound run of
``second`` must be empty.
Parameters
----------
first, second : ResolvedSLH
Systems with equal channel counts on one Hilbert space.
Returns
-------
ResolvedSLH
The cascaded system.
Raises
------
ValueError
Channel counts differ, Hilbert spaces differ, or a joined leg carries a
reference run.
"""
_require_common_space(first, second)
if len(first.channels) != len(second.channels):
raise ValueError(
"series_product needs equal channel counts; "
f"got {len(first.channels)} outputs feeding {len(second.channels)} inputs."
)
for upstream, downstream in zip(first.channels, second.channels, strict=True):
_require_joinable("series_product", upstream, downstream)
template = _template(first)
xp, ((s1, l1), (s2, l2)) = _arrays(first, second)
routed = xp.tensordot(s2, l1, axes=1)
couplings = l2 + routed
cross = xp.sum(_dag(l2, xp) @ routed, axis=0)
correction = _correction_term(cross, template, xp, tag="slh-series")
support = (np.asarray(second.support, dtype=int) @ np.asarray(first.support, dtype=int)) > 0
channels = tuple(
_channel(
replace(downstream, input_occupation=upstream.input_occupation),
couplings[index],
template,
reference=ReferencePlane(inbound=upstream.reference.inbound, outbound=downstream.reference.outbound),
)
for index, (upstream, downstream) in enumerate(zip(first.channels, second.channels, strict=True))
)
return ResolvedSLH(
scattering=s2 @ s1,
hamiltonian=_hamiltonian((first, second), correction),
channels=channels,
support=support,
)
[docs]
def concatenate(*systems: ResolvedSLH, prefixes: Sequence[str] | None = None) -> ResolvedSLH:
"""Place systems side by side with block-diagonal scattering.
Implements ``G1 ⊞ G2 = (diag(S1, S2), (L1, L2), H1 + H2)``. Channel keys must
stay unique. To rename each system's keys as ``prefix/key``, pass
``prefixes``. Exposed channels are ordered before hidden channels.
Parameters
----------
*systems : ResolvedSLH
One or more systems on one Hilbert space.
prefixes : sequence of str, optional
One prefix per system.
Returns
-------
ResolvedSLH
The concatenated system.
Raises
------
ValueError
No systems, mismatched prefixes, duplicate keys, or differing Hilbert spaces.
"""
if not systems:
raise ValueError("concatenate needs at least one system.")
_require_common_space(*systems)
if prefixes is not None and len(prefixes) != len(systems):
raise ValueError(f"concatenate got {len(prefixes)} prefixes for {len(systems)} systems.")
entries: list[SLHChannel] = []
for system_index, system in enumerate(systems):
for channel in system.channels:
key = channel.key if prefixes is None else f"{prefixes[system_index]}/{channel.key}"
entries.append(replace(channel, key=key))
keys = [channel.key for channel in entries]
duplicates = sorted({key for key in keys if keys.count(key) > 1})
if duplicates:
raise ValueError(f"concatenate needs prefixes= to disambiguate duplicate channel keys {duplicates}.")
order = np.asarray(
[index for index, channel in enumerate(entries) if channel.accessibility == "exposed"]
+ [index for index, channel in enumerate(entries) if channel.accessibility == "hidden"],
dtype=int,
)
if any(system.output_network is not None for system in systems):
raise ValueError("Compose the physical PortNetwork before resolving branched output reference sections.")
xp = _array_module(systems)
assemble_blocks = block_diag
if xp is not np:
from jax.scipy.linalg import block_diag as assemble_blocks
grid = np.ix_(order, order)
scattering = assemble_blocks(*(xp.asarray(system.scattering, dtype=complex) for system in systems))[grid]
support = block_diag(*(np.asarray(system.support, dtype=bool) for system in systems))[grid]
channels = tuple(replace(entries[index], coupling_operator=entries[index].coupling) for index in order)
return ResolvedSLH(
scattering=scattering,
hamiltonian=_hamiltonian(systems, None),
channels=channels,
support=support,
)
[docs]
def feedback_reduce(slh: ResolvedSLH, *, output: str | SLHChannel, input: str | SLHChannel) -> ResolvedSLH:
"""Connect one output channel back into one input channel.
For output ``x`` and input ``y`` with ``g = (1 − S_xy)⁻¹``, this implements
``S̃ = S_x̄ȳ + S_x̄y g S_xȳ``, ``L̃ = L_x̄ + S_x̄y g L_x`` and
``H̃ = H + ((Σ_j L_j† S_jy) g L_x − h.c.) / 2i``. When ``x`` and ``y`` differ,
the remaining input of ``x`` and output of ``y`` form one channel, keyed
``"<x>-><y>"`` and is placed at the earlier of the two positions.
Parameters
----------
slh : ResolvedSLH
The system to close.
output, input : str or SLHChannel
The output channel that is fed back and the input channel that receives it.
Returns
-------
ResolvedSLH
The reduced system with one channel fewer.
Raises
------
ValueError
The loop is concretely singular, the two channels differ in accessibility,
or a closed leg carries a reference run.
KeyError
An unknown channel key.
"""
x = _index(slh, output)
y = _index(slh, input)
closed_output, closed_input = slh.channels[x], slh.channels[y]
_require_joinable("feedback_reduce", closed_output, closed_input)
template = _template(slh)
xp, ((scattering, couplings),) = _arrays(slh)
loop = scattering[x, y]
if not contains_tracer(loop) and abs(1.0 - complex(np.asarray(loop))) < _SINGULAR_TOLERANCE:
raise ValueError(
f"Feedback from {closed_output.key!r} into {closed_input.key!r} is singular: 1 - S[x, y] vanishes."
)
gain = 1.0 / (1.0 - loop)
merged = min(x, y) if x != y else None
positions = [index for index in range(len(slh.channels)) if index not in (x, y) or index == merged]
row_index = np.asarray([y if index == merged else index for index in positions], dtype=int)
column_index = np.asarray([x if index == merged else index for index in positions], dtype=int)
channels = [
replace(
closed_input,
key=f"{closed_output.key}->{closed_input.key}",
input_occupation=closed_output.input_occupation,
reference=ReferencePlane(inbound=closed_output.reference.inbound, outbound=closed_input.reference.outbound),
)
if index == merged
else slh.channels[index]
for index in positions
]
reduced_scattering = (
scattering[np.ix_(row_index, column_index)]
+ xp.outer(scattering[row_index, y], scattering[x, column_index]) * gain
)
reduced_couplings = couplings[row_index] + (scattering[row_index, y] * gain)[:, None, None] * couplings[x]
inbound = xp.sum(_dag(couplings, xp) * scattering[:, y][:, None, None], axis=0)
correction = _correction_term(inbound @ (gain * couplings[x]), template, xp, tag="slh-feedback")
support = np.asarray(slh.support, dtype=bool)
reduced_support = support[np.ix_(row_index, column_index)] | np.outer(
support[row_index, y], support[x, column_index]
)
return ResolvedSLH(
scattering=reduced_scattering,
hamiltonian=_hamiltonian((slh,), correction),
channels=tuple(
_channel(channel, reduced_couplings[index], template, reference=channel.reference)
for index, channel in enumerate(channels)
),
support=reduced_support,
)
def _index(slh: ResolvedSLH, channel: str | SLHChannel) -> int:
key = channel.key if isinstance(channel, SLHChannel) else channel
for index, candidate in enumerate(slh.channels):
if candidate.key == key:
return index
raise KeyError(f"No channel {key!r}; available: {[item.key for item in slh.channels]}")
def _template(slh: ResolvedSLH) -> CanonicalOperator:
if not slh.channels:
raise ValueError("SLH algebra needs at least one channel to define the Hilbert space.")
return slh.channels[0].coupling
def _require_common_space(*systems: ResolvedSLH) -> None:
layouts = {
(template.shape, template.dims, template.basis, template.subsystem_labels)
for template in (_template(system) for system in systems)
}
if len(layouts) != 1:
raise ValueError(f"SLH algebra needs operands on one Hilbert space; got layouts {sorted(map(str, layouts))}.")
def _arrays(*systems: ResolvedSLH) -> tuple[Any, list[tuple[Any, Any]]]:
"""Return the shared array module and each system's dense ``(S, stacked L)``."""
if any(system.output_network is not None for system in systems):
raise ValueError("Compose the physical PortNetwork before resolving branched output reference sections.")
xp = _array_module(systems)
return xp, [
(
xp.asarray(system.scattering, dtype=complex),
xp.stack([xp.asarray(operator.to_dense(), dtype=complex) for operator in system.L]),
)
for system in systems
]
def _array_module(systems: Sequence[ResolvedSLH]) -> Any:
"""Return JAX when any operand payload is traced or already a JAX array, else NumPy."""
leaves = [leaf for system in systems for leaf in (system.scattering, *(operator.values for operator in system.L))]
return select_array_module(contains_tracer(leaves) or any(is_jax_array(leaf) for leaf in leaves))
def _require_joinable(operation: str, upstream: SLHChannel, downstream: SLHChannel) -> None:
"""Reject joining an output leg to an input leg whose metadata cannot be merged."""
prefix = f"{operation} cannot join {upstream.key!r} to {downstream.key!r}: "
if upstream.accessibility != downstream.accessibility:
raise ValueError(prefix + "both channels must share accessibility.")
if not _same_carrier(upstream.carrier, downstream.carrier):
raise ValueError(prefix + "carriers are not statically known to be equal.")
if downstream.input_occupation is not None:
raise ValueError(prefix + "the connected input cannot retain an independent thermal bath declaration.")
if upstream.reference.outbound or downstream.reference.inbound:
raise ValueError(prefix + "reference runs must be empty on the joined legs.")
def _same_carrier(first: Any, second: Any) -> bool:
"""Allow concrete roundoff in static SLH composition; traced carriers require identity."""
first_concrete = maybe_concrete_scalar(first)
second_concrete = maybe_concrete_scalar(second)
if first_concrete is not None and second_concrete is not None:
return bool(np.isclose(first_concrete, second_concrete, rtol=1e-12, atol=1e-12))
return first is second
def _dag(matrices: Any, xp: Any) -> Any:
return xp.conj(xp.swapaxes(matrices, -1, -2))
def _correction_term(cross: Any, template: CanonicalOperator, xp: Any, *, tag: str) -> StaticTerm | None:
values = (cross - _dag(cross, xp)) / 2j
if not contains_tracer(values) and np.allclose(np.asarray(values), 0.0):
return None
operator = CanonicalOperator.from_dense(
values, dims=template.dims, basis=template.basis, subsystem_labels=template.subsystem_labels, tag=tag
)
return StaticTerm(operator=operator, origin="network", metadata={"operation": tag})
def _hamiltonian(systems: Sequence[ResolvedSLH], correction: StaticTerm | None) -> HamiltonianProgram:
static = tuple(term for system in systems for term in system.H.static_terms)
dynamic = tuple(term for system in systems for term in system.H.dynamic_terms)
if correction is not None:
static = (*static, correction)
return HamiltonianProgram(static_terms=static, dynamic_terms=dynamic)
def _channel(source: SLHChannel, values: Any, template: CanonicalOperator, *, reference: ReferencePlane) -> SLHChannel:
coupling = CanonicalOperator.from_dense(
values,
dims=template.dims,
basis=template.basis,
subsystem_labels=template.subsystem_labels,
tag=f"slh:{source.key}",
)
collapse = CollapseTerm(
operator=coupling,
rate=1.0,
source=source.key,
channel="resolved_slh",
frame_frequency=source.collapse.frame_frequency,
)
return SLHChannel(
key=source.key,
accessibility=source.accessibility,
collapse=collapse,
coupling_operator=coupling,
reference=reference,
input_occupation=source.input_occupation,
)
