tn_noise_learning

Tensor-network learning of quantum noise channels.

Pipeline: Topology → LPDO → NoiseDataset → Trainer → History.

PAULI_EIGENSTATES module-attribute

PAULI_EIGENSTATES = array(
    [
        [[1, 0], [0, 0]],
        [[0, 0], [0, 1]],
        [[0.5, 0.5], [0.5, 0.5]],
        [[0.5, -0.5], [-0.5, 0.5]],
        [[0.5, -0.5j], [0.5j, 0.5]],
        [[0.5, 0.5j], [-0.5j, 0.5]],
    ],
    dtype=complex,
)

States |0>, |1>, |+>, |−>, |+i>, |−i>; effect 2 * basis + bit is outcome bit of basis Z, X or Y.

Batch

Bases: NamedTuple

One training batch: per-site state and effect matrices, each of shape (num_qubits, batch, 2, 2).

History dataclass

History(
    metrics: dict[str, list[float]] = dict(),
    params: list[list[ndarray]] = list(),
    best_epoch: int | None = None,
    num_sites: int | None = None,
)

Metric curves and optional parameter snapshots recorded by Trainer.

PARAMETER	DESCRIPTION
metrics	one list per metric name, one value per epoch. TYPE: dict[str, list[float]] DEFAULT: dict()
params	Kraus tensors after each epoch, filled only when TrainingConfig.snapshot_params is set. TYPE: list[list[ndarray]] DEFAULT: list()
best_epoch	1-based epoch whose parameters were written back to the model. TYPE: int | None DEFAULT: None
num_sites	number of qubits, used to draw the trace-preserving reference 2**num_sites. TYPE: int | None DEFAULT: None

num_epochs property

num_epochs: int

Number of recorded epochs.

append

append(
    values: dict[str, Any],
    tensors: list[ndarray] | None = None,
) -> None

Record one epoch.

Source code in src/tn_noise_learning/history.py

def append(self, values: dict[str, Any], tensors: list[np.ndarray] | None = None) -> None:
    """Record one epoch."""
    for name, value in values.items():
        self.metrics.setdefault(name, []).append(float(value))
    if tensors is not None:
        self.params.append([np.asarray(t) for t in tensors])

load classmethod

load(path: str | Path) -> History

Read a history written by save.

Source code in src/tn_noise_learning/history.py

@classmethod
def load(cls, path: str | Path) -> History:
    """Read a history written by [`save`][tn_noise_learning.history.History.save]."""
    group = zarr.open_group(str(path), mode="r")
    metrics = {
        str(name): read_array(group, f"metrics/{name}").tolist() for name in read_attr(group, "metric_names")
    }
    params = []
    if "params" in group:
        params_group = zarr.open_group(str(path), mode="r", path="params")
        for epoch in sorted(params_group.keys(), key=lambda k: int(k.split("_")[1])):
            epoch_group = zarr.open_group(str(path), mode="r", path=f"params/{epoch}")
            sites = sorted(epoch_group.keys(), key=lambda k: int(k.split("_")[1]))
            params.append([read_array(epoch_group, site) for site in sites])
    best = read_attr(group, "best_epoch")
    num_sites = read_attr(group, "num_sites")
    return cls(metrics, params, None if best is None else int(best), None if num_sites is None else int(num_sites))

plot

plot(path: str | Path | None = None) -> Figure

Plot the metric curves with plot_history.

Source code in src/tn_noise_learning/history.py

def plot(self, path: str | Path | None = None) -> Figure:
    """Plot the metric curves with [`plot_history`][tn_noise_learning.plotting.plot_history]."""
    from .plotting import plot_history  # noqa: PLC0415  (avoids a circular import)

    return plot_history(self, path)

save

save(path: str | Path) -> None

Write to a zarr store at path.

Source code in src/tn_noise_learning/history.py

def save(self, path: str | Path) -> None:
    """Write to a zarr store at ``path``."""
    group = zarr.open_group(str(path), mode="w")
    group.attrs.update(best_epoch=self.best_epoch, num_sites=self.num_sites, metric_names=list(self.metrics))
    for name, values in self.metrics.items():
        group.create_array(f"metrics/{name}", shape=(len(values),), dtype="float64")[:] = np.asarray(values)
    for epoch, tensors in enumerate(self.params):
        for site, t in enumerate(tensors):
            group.create_array(f"params/epoch_{epoch}/site_{site}", shape=t.shape, dtype=t.dtype)[:] = t

LPDO dataclass

LPDO(topology: Topology, tensors: list[ndarray])

Quantum channel as an LPDO on a topology.

PARAMETER	DESCRIPTION
topology	site connectivity. TYPE: Topology
tensors	one Kraus tensor per site in canonical axis order (*bonds to neighbours ascending, kraus, out, in). TYPE: list[ndarray]

bond_dim property

bond_dim: int

Largest virtual bond dimension (1 when there are no edges).

kraus_dim property

kraus_dim: int

Largest Kraus dimension over sites.

num_sites property

num_sites: int

Number of qubits.

__post_init__

__post_init__() -> None

Validate tensor count, axis counts, physical dimensions and bond sizes.

Source code in src/tn_noise_learning/lpdo.py

def __post_init__(self) -> None:
    """Validate tensor count, axis counts, physical dimensions and bond sizes."""
    if len(self.tensors) != self.topology.num_sites:
        raise ValueError(f"Expected {self.topology.num_sites} tensors, got {len(self.tensors)}.")
    self.tensors = [np.asarray(t) for t in self.tensors]
    for j, t in enumerate(self.tensors):
        expected = self.topology.degree(j) + 3
        if t.ndim != expected or t.shape[-2:] != (2, 2):
            raise ValueError(
                f"Site {j}: expected shape (*{self.topology.degree(j)} bonds, kraus, 2, 2), got {t.shape}."
            )
    for i, j in self.topology.edges:
        if self.bond_size(i, j) != self.bond_size(j, i):
            raise ValueError(f"Bond sizes on edge ({i}, {j}) disagree.")

bond_size

bond_size(site: int, neighbor: int) -> int

Size of the bond leg of site pointing at neighbor.

Source code in src/tn_noise_learning/lpdo.py

def bond_size(self, site: int, neighbor: int) -> int:
    """Size of the bond leg of ``site`` pointing at ``neighbor``."""
    return self.tensors[site].shape[self.topology.neighbors(site).index(neighbor)]

distance

distance(other: LPDO) -> float

Squared normalised Frobenius distance ‖Λ − Λ'‖_F² / 4**num_sites.

Source code in src/tn_noise_learning/lpdo.py

def distance(self, other: LPDO) -> float:
    """Squared normalised Frobenius distance ``‖Λ − Λ'‖_F² / 4**num_sites``."""
    return float(qtn.tensor_network_distance(self.network(), other.network()) / 2**self.num_sites) ** 2

from_group classmethod

from_group(group: Group) -> LPDO

Read an LPDO written by to_group.

Source code in src/tn_noise_learning/lpdo.py

@classmethod
def from_group(cls, group: zarr.Group) -> LPDO:
    """Read an LPDO written by [`to_group`][tn_noise_learning.lpdo.LPDO.to_group]."""
    edges = tuple((int(i), int(j)) for i, j in read_attr(group, "edges"))
    topology = Topology(int(read_attr(group, "num_sites")), edges)
    return cls(topology, [read_array(group, f"site_{j}") for j in range(topology.num_sites)])

from_network classmethod

from_network(
    topology: Topology, kraus_tn: TensorNetwork
) -> LPDO

Rebuild from a Kraus network with the index names of kraus_inds.

Source code in src/tn_noise_learning/lpdo.py

@classmethod
def from_network(cls, topology: Topology, kraus_tn: qtn.TensorNetwork) -> LPDO:
    """Rebuild from a Kraus network with the index names of [`kraus_inds`][tn_noise_learning.lpdo.kraus_inds]."""
    return cls(
        topology,
        [
            np.asarray(kraus_tn[site_tag(j)].transpose(*kraus_inds(topology, j)).data)
            for j in range(topology.num_sites)
        ],
    )

identity classmethod

identity(topology: Topology) -> LPDO

Identity channel: one Kraus operator 𝟙 per site, all bonds of size one.

Source code in src/tn_noise_learning/lpdo.py

@classmethod
def identity(cls, topology: Topology) -> LPDO:
    """Identity channel: one Kraus operator ``𝟙`` per site, all bonds of size one."""
    return cls(
        topology,
        [
            np.eye(2, dtype=complex).reshape((1,) * (topology.degree(j) + 1) + (2, 2))
            for j in range(topology.num_sites)
        ],
    )

kraus_network

kraus_network() -> qtn.TensorNetwork

Tensor network of the Kraus half, tagged I{j} and V.

Source code in src/tn_noise_learning/lpdo.py

def kraus_network(self) -> qtn.TensorNetwork:
    """Tensor network of the Kraus half, tagged ``I{j}`` and ``V``."""
    return qtn.TensorNetwork(
        [
            qtn.Tensor(t, inds=kraus_inds(self.topology, j), tags=(site_tag(j), KRAUS_TAG))
            for j, t in enumerate(self.tensors)
        ]
    )

load classmethod

load(path: str | Path) -> LPDO

Load from a zarr store written by save.

Source code in src/tn_noise_learning/lpdo.py

@classmethod
def load(cls, path: str | Path) -> LPDO:
    """Load from a zarr store written by [`save`][tn_noise_learning.lpdo.LPDO.save]."""
    return cls.from_group(zarr.open_group(str(path), mode="r"))

network

network() -> qtn.TensorNetwork

Full LPDO network with open legs kn, k (ket) and bn, b (bra) per site.

Source code in src/tn_noise_learning/lpdo.py

def network(self) -> qtn.TensorNetwork:
    """Full LPDO network with open legs ``kn, k`` (ket) and ``bn, b`` (bra) per site."""
    return full_network(self.kraus_network(), self.topology)

random classmethod

random(
    topology: Topology,
    bond_dim: int,
    kraus_dim: int,
    seed: int | None = None,
) -> LPDO

Complex Gaussian Kraus tensors scaled so that the expected trace is 2**num_sites.

Source code in src/tn_noise_learning/lpdo.py

@classmethod
def random(cls, topology: Topology, bond_dim: int, kraus_dim: int, seed: int | None = None) -> LPDO:
    """Complex Gaussian Kraus tensors scaled so that the expected trace is ``2**num_sites``."""
    rng = np.random.default_rng(seed)
    scale = np.sqrt(2.0 / (4.0 * kraus_dim * bond_dim ** (len(topology.edges) / topology.num_sites)))
    tensors = []
    for j in range(topology.num_sites):
        shape = (*(bond_dim,) * topology.degree(j), kraus_dim, 2, 2)
        tensors.append(scale * (rng.standard_normal(shape) + 1j * rng.standard_normal(shape)) / np.sqrt(2))
    return cls(topology, tensors)

save

save(path: str | Path) -> None

Save to a zarr store at path.

Source code in src/tn_noise_learning/lpdo.py

def save(self, path: str | Path) -> None:
    """Save to a zarr store at ``path``."""
    self.to_group(zarr.open_group(str(path), mode="w"))

to_dense

to_dense() -> np.ndarray

Superoperator matrix Σ_m K_m ⊗ conj(K_m) with site 0 most significant (small systems only).

Source code in src/tn_noise_learning/lpdo.py

def to_dense(self) -> np.ndarray:
    """Superoperator matrix ``Σ_m K_m ⊗ conj(K_m)`` with site 0 most significant (small systems only)."""
    if self.num_sites > MAX_DENSE_SITES:
        raise ValueError(f"to_dense is limited to {MAX_DENSE_SITES} sites, got {self.num_sites}.")
    sites = range(self.num_sites)
    rows = [output_ind(j) for j in sites] + [output_ind(j, dagger=True) for j in sites]
    cols = [input_ind(j) for j in sites] + [input_ind(j, dagger=True) for j in sites]
    return np.asarray(self.network().contract(optimize="auto-hq").to_dense(rows, cols))

to_group

to_group(group: Group) -> None

Write topology and tensors into an open zarr group.

Source code in src/tn_noise_learning/lpdo.py

def to_group(self, group: zarr.Group) -> None:
    """Write topology and tensors into an open zarr group."""
    group.attrs["num_sites"] = self.num_sites
    group.attrs["edges"] = [list(e) for e in self.topology.edges]
    for j, t in enumerate(self.tensors):
        group.create_array(f"site_{j}", shape=t.shape, dtype=t.dtype, overwrite=True)[:] = t

trace

trace() -> float

Trace of the channel's Choi-like network, 2**num_sites for a trace-preserving channel.

Source code in src/tn_noise_learning/lpdo.py

def trace(self) -> float:
    """Trace of the channel's Choi-like network, ``2**num_sites`` for a trace-preserving channel."""
    connect = {output_ind(j): output_ind(j, dagger=True) for j in range(self.num_sites)}
    connect |= {input_ind(j): input_ind(j, dagger=True) for j in range(self.num_sites)}
    return float(np.real(self.network().reindex(connect).contract(optimize="auto-hq")))

NoiseDataset dataclass

NoiseDataset(
    inputs: ndarray,
    outcomes: ndarray,
    states: ndarray,
    povms: ndarray,
    circuit: QuantumCircuit | None = None,
    ideal_channel: LPDO | None = None,
)

Shots of a process-tomography experiment on product states with product measurements.

PARAMETER	DESCRIPTION
inputs	prepared-state labels, shape (num_shots, num_qubits). TYPE: ndarray
outcomes	measured-effect labels, shape (num_shots, num_qubits). TYPE: ndarray
states	single-qubit input states per qubit, shape (num_qubits, num_states, 2, 2). TYPE: ndarray
povms	single-qubit effects per qubit, shape (num_qubits, num_effects, 2, 2). TYPE: ndarray
circuit	ideal circuit applied before the noise channel, or None for the identity. TYPE: QuantumCircuit | None DEFAULT: None
ideal_channel	ground-truth channel, if known, to report the distance to during training. TYPE: LPDO | None DEFAULT: None

num_qubits property

num_qubits: int

Number of qubits.

num_shots property

num_shots: int

Number of recorded shots.

__len__

__len__() -> int

Number of recorded shots.

Source code in src/tn_noise_learning/dataset.py

def __len__(self) -> int:
    """Number of recorded shots."""
    return self.num_shots

__post_init__

__post_init__() -> None

Normalise dtypes and validate shapes and label ranges.

Source code in src/tn_noise_learning/dataset.py

def __post_init__(self) -> None:
    """Normalise dtypes and validate shapes and label ranges."""
    self.inputs = np.asarray(self.inputs, dtype=np.int8)
    self.outcomes = np.asarray(self.outcomes, dtype=np.int8)
    self.states = np.asarray(self.states, dtype=complex)
    self.povms = np.asarray(self.povms, dtype=complex)
    n = self.num_qubits
    if self.inputs.shape != self.outcomes.shape or self.inputs.ndim != 2:
        raise ValueError("inputs and outcomes must both have shape (num_shots, num_qubits).")
    if (
        self.states.shape[0] != n
        or self.states.shape[-2:] != (2, 2)
        or self.povms.shape[0] != n
        or self.povms.shape[-2:] != (2, 2)
    ):
        raise ValueError("states and povms must have shape (num_qubits, num_labels, 2, 2).")
    if self.inputs.min(initial=0) < 0 or self.inputs.max(initial=0) >= self.states.shape[1]:
        raise ValueError("inputs contain labels outside the range of states.")
    if self.outcomes.min(initial=0) < 0 or self.outcomes.max(initial=0) >= self.povms.shape[1]:
        raise ValueError("outcomes contain labels outside the range of povms.")
    if self.circuit is not None and self.circuit.num_qubits != n:
        raise ValueError(f"Circuit acts on {self.circuit.num_qubits} qubits, dataset has {n}.")
    if self.ideal_channel is not None and self.ideal_channel.num_sites != n:
        raise ValueError(f"Ideal channel has {self.ideal_channel.num_sites} sites, dataset has {n} qubits.")

batches

batches(
    batch_size: int,
    *,
    shuffle: bool = True,
    seed: int = 0,
    drop_last: bool = True,
) -> Iterator[Batch]

Iterate over batches; with drop_last every batch has exactly batch_size shots (one JIT trace).

Source code in src/tn_noise_learning/dataset.py

def batches(
    self, batch_size: int, *, shuffle: bool = True, seed: int = 0, drop_last: bool = True
) -> Iterator[Batch]:
    """Iterate over batches; with ``drop_last`` every batch has exactly ``batch_size`` shots (one JIT trace)."""
    if batch_size > self.num_shots:
        raise ValueError(f"batch_size {batch_size} exceeds the {self.num_shots} shots in the dataset.")
    order = np.random.default_rng(seed).permutation(self.num_shots) if shuffle else np.arange(self.num_shots)
    stop = self.num_shots - (self.num_shots % batch_size if drop_last else 0)
    for start in range(0, stop, batch_size):
        yield self.gather(order[start : start + batch_size])

gather

gather(indices: ndarray) -> Batch

State and effect matrices of the given shots, shape (num_qubits, len(indices), 2, 2) each.

Source code in src/tn_noise_learning/dataset.py

def gather(self, indices: np.ndarray) -> Batch:
    """State and effect matrices of the given shots, shape ``(num_qubits, len(indices), 2, 2)`` each."""
    sites = np.arange(self.num_qubits)[:, None]
    return Batch(self.states[sites, self.inputs[indices].T], self.povms[sites, self.outcomes[indices].T])

load classmethod

load(path: str | Path) -> NoiseDataset

Read a dataset written by save.

Source code in src/tn_noise_learning/dataset.py

@classmethod
def load(cls, path: str | Path) -> NoiseDataset:
    """Read a dataset written by [`save`][tn_noise_learning.dataset.NoiseDataset.save]."""
    group = zarr.open_group(str(path), mode="r")
    arrays = {name: read_array(group, name) for name in ARRAY_NAMES}
    circuit = None
    if "circuit_qasm2" in group.attrs:
        circuit = qasm2.loads(
            str(group.attrs["circuit_qasm2"]), custom_instructions=qasm2.LEGACY_CUSTOM_INSTRUCTIONS
        )
    ideal = None
    if "ideal_channel" in group:
        ideal = LPDO.from_group(zarr.open_group(str(path), mode="r", path="ideal_channel"))
    return cls(**arrays, circuit=circuit, ideal_channel=ideal)

save

save(path: str | Path) -> None

Write the dataset to a zarr store at path (overwriting).

Source code in src/tn_noise_learning/dataset.py

def save(self, path: str | Path) -> None:
    """Write the dataset to a zarr store at ``path`` (overwriting)."""
    group = zarr.open_group(str(path), mode="w")
    group.attrs.update(
        format_version=FORMAT_VERSION, num_qubits=self.num_qubits, created=datetime.now(UTC).isoformat()
    )
    if self.circuit is not None:
        group.attrs["circuit_qasm2"] = qasm2.dumps(self.circuit)
    for name in ARRAY_NAMES:
        array = getattr(self, name)
        chunks = (min(self.num_shots, 65536), self.num_qubits) if name in ("inputs", "outcomes") else array.shape
        group.create_array(name, shape=array.shape, dtype=array.dtype, chunks=chunks)[:] = array
    if self.ideal_channel is not None:
        self.ideal_channel.to_group(group.create_group("ideal_channel"))

split

split(
    test_fraction: float = 0.1, seed: int = 0
) -> tuple[NoiseDataset, NoiseDataset]

Random disjoint (train, test) split.

Source code in src/tn_noise_learning/dataset.py

def split(self, test_fraction: float = 0.1, seed: int = 0) -> tuple[NoiseDataset, NoiseDataset]:
    """Random disjoint ``(train, test)`` split."""
    if not 0 < test_fraction < 1:
        raise ValueError(f"test_fraction must be in (0, 1), got {test_fraction}.")
    permutation = np.random.default_rng(seed).permutation(self.num_shots)
    num_test = max(1, round(test_fraction * self.num_shots))
    return self.subset(permutation[num_test:]), self.subset(permutation[:num_test])

subset

subset(indices: ndarray) -> NoiseDataset

Dataset restricted to the given shots; states, povms, circuit and ideal channel are shared.

Source code in src/tn_noise_learning/dataset.py

def subset(self, indices: np.ndarray) -> NoiseDataset:
    """Dataset restricted to the given shots; states, povms, circuit and ideal channel are shared."""
    return replace(self, inputs=self.inputs[indices], outcomes=self.outcomes[indices])

Topology dataclass

Topology(
    num_sites: int, edges: tuple[tuple[int, int], ...]
)

Graph of qubit sites 0..num_sites-1 and the edges that carry virtual bonds.

Site integers are the qubit order of the dataset and of the Qiskit circuit. Edges are stored normalised as sorted (i, j) pairs with i < j, without duplicates.

PARAMETER	DESCRIPTION
num_sites	number of qubits. TYPE: int
edges	pairs of sites connected by a bond. TYPE: tuple[tuple[int, int], ...]

__post_init__

__post_init__() -> None

Validate sites and normalise the edge list.

Source code in src/tn_noise_learning/topology.py

def __post_init__(self) -> None:
    """Validate sites and normalise the edge list."""
    if self.num_sites < 1:
        raise ValueError(f"num_sites must be positive, got {self.num_sites}.")
    normalised = set()
    for i, j in self.edges:
        if i == j:
            raise ValueError(f"Self-loop on site {i} is not allowed.")
        if not (0 <= i < self.num_sites and 0 <= j < self.num_sites):
            raise ValueError(f"Edge ({i}, {j}) is outside the site range 0..{self.num_sites - 1}.")
        normalised.add((min(i, j), max(i, j)))
    object.__setattr__(self, "edges", tuple(sorted(normalised)))

chain classmethod

chain(num_sites: int) -> Topology

Open 1D chain 0 - 1 - ... - (num_sites - 1).

Source code in src/tn_noise_learning/topology.py

@classmethod
def chain(cls, num_sites: int) -> Topology:
    """Open 1D chain ``0 - 1 - ... - (num_sites - 1)``."""
    return cls(num_sites, tuple(qtn.edges_1d_chain(num_sites)))

degree

degree(site: int) -> int

Number of bonds attached to site.

Source code in src/tn_noise_learning/topology.py

def degree(self, site: int) -> int:
    """Number of bonds attached to ``site``."""
    return len(self.neighbors(site))

from_edges classmethod

from_edges(
    edges: Iterable[tuple[int, int]],
    num_sites: int | None = None,
) -> Topology

Custom topology from an edge list; num_sites defaults to the largest site index plus one.

Source code in src/tn_noise_learning/topology.py

@classmethod
def from_edges(cls, edges: Iterable[tuple[int, int]], num_sites: int | None = None) -> Topology:
    """Custom topology from an edge list; ``num_sites`` defaults to the largest site index plus one."""
    pairs = tuple((int(i), int(j)) for i, j in edges)
    if num_sites is None:
        num_sites = 1 + max((max(e) for e in pairs), default=0)
    return cls(num_sites, pairs)

neighbors

neighbors(site: int) -> tuple[int, ...]

Sites bonded to site, in ascending order.

Source code in src/tn_noise_learning/topology.py

def neighbors(self, site: int) -> tuple[int, ...]:
    """Sites bonded to ``site``, in ascending order."""
    return tuple(sorted(j if i == site else i for i, j in self.edges if site in (i, j)))

square classmethod

square(rows: int, cols: int) -> Topology

Open 2D square lattice, with site (r, c) mapped row-major to r * cols + c.

Source code in src/tn_noise_learning/topology.py

@classmethod
def square(cls, rows: int, cols: int) -> Topology:
    """Open 2D square lattice, with site ``(r, c)`` mapped row-major to ``r * cols + c``."""
    edges = [(r1 * cols + c1, r2 * cols + c2) for (r1, c1), (r2, c2) in qtn.edges_2d_square(rows, cols)]
    return cls(rows * cols, tuple(edges))

Trainer

Trainer(config: TrainingConfig | None = None)

Fit an LPDO to a NoiseDataset.

Source code in src/tn_noise_learning/trainer.py

def __init__(self, config: TrainingConfig | None = None) -> None:
    self.config = config or TrainingConfig()

fit

fit(
    model: LPDO,
    train: NoiseDataset,
    test: NoiseDataset | None = None,
    ideal_channel: LPDO | None = None,
) -> History

Train model in place and return the training history.

PARAMETER	DESCRIPTION
model	initial LPDO; on return it holds the parameters of the best epoch. TYPE: LPDO
train	training shots; its circuit is applied inside the loss. TYPE: NoiseDataset
test	held-out shots; split off train when omitted. TYPE: NoiseDataset | None DEFAULT: None
ideal_channel	ground truth to report distance_to_ideal; defaults to train.ideal_channel. TYPE: LPDO | None DEFAULT: None

Source code in src/tn_noise_learning/trainer.py

def fit(
    self,
    model: LPDO,
    train: NoiseDataset,
    test: NoiseDataset | None = None,
    ideal_channel: LPDO | None = None,
) -> History:
    """Train ``model`` in place and return the training history.

    Args:
        model: initial LPDO; on return it holds the parameters of the best epoch.
        train: training shots; its circuit is applied inside the loss.
        test: held-out shots; split off ``train`` when omitted.
        ideal_channel: ground truth to report ``distance_to_ideal``; defaults to ``train.ideal_channel``.
    """
    cfg = self.config
    if test is None:
        train, test = train.split(cfg.test_fraction, cfg.seed)
    ideal = train.ideal_channel if ideal_channel is None else ideal_channel
    if cfg.tp_warmstart_steps > 0:
        model.tensors = project_trace_preserving(model, cfg.tp_warmstart_steps).tensors

    params, skeleton = qtn.pack(model.kraus_network())
    rebuild = lambda p: LPDO.from_network(model.topology, qtn.unpack(p, skeleton))  # noqa: E731
    term_names, weights_at, step, evaluate_loss, opt_state = self._compile(model, params, skeleton, train)
    history = History(num_sites=model.num_sites)
    selection = _Selection(cfg.patience, cfg.tolerance, params)
    global_step = 0
    epoch = 0
    try:
        for epoch in range(1, cfg.epochs + 1):
            sums = np.zeros(len(term_names) + 1)
            for batch in train.batches(cfg.batch_size, shuffle=True, seed=cfg.seed + epoch):
                params, opt_state, total, values = step(
                    params, opt_state, batch.states, batch.povms, weights_at(global_step)
                )
                sums += np.concatenate([[float(total)], np.asarray(values)])
                global_step += 1
                if cfg.tp_project_every and global_step % cfg.tp_project_every == 0:
                    projected = project_trace_preserving(rebuild(params), cfg.tp_project_steps)
                    params = dict(zip(params, qtn.pack(projected.kraus_network())[0].values(), strict=True))
            num_batches = len(train) // cfg.batch_size
            test_loss = np.mean(
                [
                    float(evaluate_loss(params, b.states, b.povms, weights_at(global_step))[0])
                    for b in test.batches(min(cfg.batch_size, len(test)), shuffle=False, drop_last=False)
                ]
            )
            metrics = {"train_loss": sums[0] / num_batches, "test_loss": float(test_loss)}
            metrics |= {name: sums[i + 1] / num_batches for i, name in enumerate(term_names)}
            current = rebuild(params)
            metrics |= self._model_metrics(current, ideal)
            history.append(metrics, current.tensors if cfg.snapshot_params else None)
            if epoch % cfg.log_every == 0:
                logger.info(
                    "epoch %3d/%d  %s", epoch, cfg.epochs, "  ".join(f"{k}={v:.4g}" for k, v in metrics.items())
                )
            if cfg.monitor not in metrics:
                raise KeyError(f"monitor {cfg.monitor!r} is not a recorded metric: {sorted(metrics)}")
            if selection.update(metrics[cfg.monitor], params, epoch):
                break
    except KeyboardInterrupt:
        logger.warning("Training interrupted at epoch %d; keeping the best model so far.", epoch)

    model.tensors = rebuild(selection.best_params).tensors
    history.best_epoch = selection.best_epoch
    return history

TrainingConfig dataclass

TrainingConfig(
    epochs: int = 50,
    batch_size: int = 256,
    learning_rate: ScalarOrSchedule = (
        lambda: optax.exponential_decay(
            0.01, 200, 0.8, end_value=0.0001
        )
    )(),
    optimizer: Callable[
        [ScalarOrSchedule], GradientTransformation
    ] = optax.adam,
    tp_weight: Schedule = 100.0,
    tp_warmstart_steps: int = 500,
    tp_project_every: int | None = None,
    tp_project_steps: int = 50,
    extra_loss_terms: dict[
        str, tuple[LossTerm, Schedule]
    ] = dict(),
    metrics: dict[str, Metric] = dict(),
    monitor: str = "test_loss",
    patience: int = 10,
    tolerance: float = 0.0001,
    test_fraction: float = 0.1,
    seed: int = 0,
    snapshot_params: bool = False,
    log_every: int = 1,
    contract_optimize: str = "auto-hq",
)

Hyper-parameters of Trainer.

PARAMETER	DESCRIPTION
epochs	maximum number of passes over the training set. TYPE: int DEFAULT: 50
batch_size	shots per gradient step; every batch has exactly this size. TYPE: int DEFAULT: 256
learning_rate	constant, optax schedule, or callable of the step index; defaults to an exponential decay from 1e-2 by a factor 0.8 every 200 steps down to 1e-4. TYPE: ScalarOrSchedule DEFAULT: (lambda: exponential_decay(0.01, 200, 0.8, end_value=0.0001))()
optimizer	factory turning the learning rate into an optax optimizer. TYPE: Callable[[ScalarOrSchedule], GradientTransformation] DEFAULT: adam
tp_weight	weight of the trace-preservation penalty, constant or callable of the step index. The trained trace exceeds 2**n by about 1/(2 tp_weight) relative; weights far above 1000 make Adam stiff. TYPE: Schedule DEFAULT: 100.0
tp_warmstart_steps	L-BFGS-B steps moving the initial model onto the trace-preserving manifold (0 disables). TYPE: int DEFAULT: 500
tp_project_every	re-run the L-BFGS-B projection every this many steps (None disables). TYPE: int | None DEFAULT: None
tp_project_steps	L-BFGS-B steps per projection. TYPE: int DEFAULT: 50
extra_loss_terms	name -> (term, weight) added to the loss; a term maps (lpdo_tn, probs) to a scalar. TYPE: dict[str, tuple[LossTerm, Schedule]] DEFAULT: dict()
metrics	name -> metric(model) evaluated on the current model after every epoch. TYPE: dict[str, Metric] DEFAULT: dict()
monitor	metric used for early stopping and model selection. TYPE: str DEFAULT: 'test_loss'
patience	epochs without improvement of monitor before training stops. TYPE: int DEFAULT: 10
tolerance	minimum decrease of monitor that counts as an improvement. TYPE: float DEFAULT: 0.0001
test_fraction	held-out fraction used when fit receives no test set. TYPE: float DEFAULT: 0.1
seed	seed for the train/test split and batch shuffling. TYPE: int DEFAULT: 0
snapshot_params	record the Kraus tensors after every epoch in the history. TYPE: bool DEFAULT: False
log_every	log one line every this many epochs. TYPE: int DEFAULT: 1
contract_optimize	cotengra path optimizer for the loss contraction. TYPE: str DEFAULT: 'auto-hq'

apply_circuit

apply_circuit(
    states_tn: TensorNetwork, circuit: QuantumCircuit
) -> qtn.TensorNetwork

Evolve a raw state network (legs p, q) to U ρ U† with legs k, b, one tensor per site.

Source code in src/tn_noise_learning/circuits.py

def apply_circuit(states_tn: qtn.TensorNetwork, circuit: QuantumCircuit) -> qtn.TensorNetwork:
    """Evolve a raw state network (legs ``p, q``) to ``U ρ U†`` with legs ``k, b``, one tensor per site."""
    tn = qtn.TensorNetwork([states_tn, circuit_network(circuit)], check_collisions=False)
    for j in range(circuit.num_qubits):
        tn.contract_tags_(site_tag(j))
    return tn

circuit_network

circuit_network(
    circuit: QuantumCircuit,
) -> qtn.TensorNetwork

One tensor per site mapping (p_j, q_j) to (k_j, b_j) such that the whole network is U · U†.

Only one- and two-qubit unitary gates are supported; barriers are skipped.

Source code in src/tn_noise_learning/circuits.py

def circuit_network(circuit: QuantumCircuit) -> qtn.TensorNetwork:
    """One tensor per site mapping ``(p_j, q_j)`` to ``(k_j, b_j)`` such that the whole network is ``U · U†``.

    Only one- and two-qubit unitary gates are supported; barriers are skipped.
    """
    num_sites = circuit.num_qubits
    identity = np.eye(4, dtype=complex).reshape(2, 2, 2, 2)
    tn = qtn.TensorNetwork(
        [
            qtn.Tensor(
                identity,
                inds=(raw_ind(j), raw_ind(j, dagger=True), input_ind(j), input_ind(j, dagger=True)),
                tags=site_tag(j),
            )
            for j in range(num_sites)
        ]
    )
    for instruction in circuit.data:
        operation = instruction.operation
        if operation.name == "barrier":
            continue
        sites = [circuit.find_bit(q).index for q in instruction.qubits]
        if len(sites) > 2:
            raise ValueError(
                f"Gate {operation.name} acts on {len(sites)} qubits; only one- and two-qubit gates are supported."
            )
        try:
            gate = np.asarray(Operator(operation).data, dtype=complex)
        except Exception as error:
            raise ValueError(f"Operation {operation.name} is not a unitary gate.") from error
        pieces = [(sites[0], gate[None])] if len(sites) == 1 else list(zip(sites, _schmidt_pair(gate), strict=True))
        bond_ket, bond_bra = qtn.rand_uuid("s"), qtn.rand_uuid("s")
        for site, piece in pieces:
            old_ket, old_bra = qtn.rand_uuid("s"), qtn.rand_uuid("s")
            tn.reindex_({input_ind(site): old_ket, input_ind(site, dagger=True): old_bra})
            tn.add_tensor(qtn.Tensor(piece, inds=(bond_ket, input_ind(site), old_ket), tags=site_tag(site)))
            tn.add_tensor(
                qtn.Tensor(piece.conj(), inds=(bond_bra, input_ind(site, dagger=True), old_bra), tags=site_tag(site))
            )
    for j in range(num_sites):
        tn.contract_tags_(site_tag(j))
    tn.fuse_multibonds_()
    tn.squeeze_()
    largest = max(t.size for t in tn)
    if largest > LARGE_TENSOR_SIZE:
        logger.warning(
            "Circuit network has a site tensor with %d elements; the circuit is deep for an uncompressed decomposition",
            largest,
        )
    return tn

depolarizing

depolarizing(topology: Topology, p: float) -> LPDO

Single-qubit depolarizing channel (1 − p) ρ + (p / 3) Σ_{P≠I} P ρ P on every site.

Source code in src/tn_noise_learning/channels.py

def depolarizing(topology: Topology, p: float) -> LPDO:
    """Single-qubit depolarizing channel ``(1 − p) ρ + (p / 3) Σ_{P≠I} P ρ P`` on every site."""
    return pauli_channel(topology, [1 - p, p / 3, p / 3, p / 3])

full_network

full_network(
    kraus_tn: TensorNetwork, topology: Topology
) -> qtn.TensorNetwork

Join a Kraus network with its conjugate; works on JAX tracers inside a jitted loss.

Source code in src/tn_noise_learning/lpdo.py

def full_network(kraus_tn: qtn.TensorNetwork, topology: Topology) -> qtn.TensorNetwork:
    """Join a Kraus network with its conjugate; works on JAX tracers inside a jitted loss."""
    dagger = kraus_tn.H.reindex(dagger_rename(topology)).retag({KRAUS_TAG: DAGGER_TAG})
    return qtn.TensorNetwork([kraus_tn, dagger])

negative_log_likelihood

negative_log_likelihood(
    lpdo_tn: TensorNetwork, probs: ndarray
) -> jnp.ndarray

-mean(log p) over the batch, the KL divergence up to a data-dependent constant.

Source code in src/tn_noise_learning/losses.py

def negative_log_likelihood(lpdo_tn: qtn.TensorNetwork, probs: jnp.ndarray) -> jnp.ndarray:  # noqa: ARG001
    """``-mean(log p)`` over the batch, the KL divergence up to a data-dependent constant."""
    return -jnp.mean(jnp.log(probs))

outcome_distribution

outcome_distribution(
    channel: LPDO,
    circuit: QuantumCircuit | None,
    state_labels: Labels,
    basis_labels: Labels,
    *,
    states: ndarray,
    povms: ndarray,
    optimize: str = "auto-hq",
) -> np.ndarray

Joint probability of every bit string for one setting, shape (2,) * num_qubits with qubit 0 first.

Source code in src/tn_noise_learning/simulate.py

def outcome_distribution(
    channel: LPDO,
    circuit: QuantumCircuit | None,
    state_labels: Labels,
    basis_labels: Labels,
    *,
    states: np.ndarray,
    povms: np.ndarray,
    optimize: str = "auto-hq",
) -> np.ndarray:
    """Joint probability of every bit string for one setting, shape ``(2,) * num_qubits`` with qubit 0 first."""
    n = channel.num_sites
    if n > MAX_SAMPLED_SITES:
        raise ValueError(f"Exact sampling enumerates 2**n outcomes; n must be at most {MAX_SAMPLED_SITES}, got {n}.")
    prep = qtn.TensorNetwork(
        [
            qtn.Tensor(states[j, state_labels[j]], inds=(raw_ind(j), raw_ind(j, dagger=True)), tags=site_tag(j))
            for j in range(n)
        ]
    )
    if circuit is None:
        prep.reindex_(
            {raw_ind(j): input_ind(j) for j in range(n)}
            | {raw_ind(j, dagger=True): input_ind(j, dagger=True) for j in range(n)}
        )
    else:
        prep = apply_circuit(prep, circuit)
    measurement = qtn.TensorNetwork(
        [
            qtn.Tensor(
                povms[j, 2 * b : 2 * b + 2], inds=(f"o{j}", output_ind(j, dagger=True), output_ind(j)), tags=site_tag(j)
            )
            for j, b in enumerate(basis_labels)
        ]
    )
    tn = qtn.TensorNetwork([channel.network(), prep, measurement], check_collisions=False)
    probs = np.real(tn.contract(output_inds=[f"o{j}" for j in range(n)], optimize=optimize).data)
    probs = np.clip(probs, 0, None)
    return probs / probs.sum()

pauli_channel

pauli_channel(
    topology: Topology, probs: Iterable[float] | ndarray
) -> LPDO

Independent single-qubit Pauli channels ρ → Σ_P p_P P ρ P.

PARAMETER	DESCRIPTION
topology	site connectivity (all bonds have size one). TYPE: Topology
probs	(p_I, p_X, p_Y, p_Z) shared by all sites, or one row per site with shape (num_sites, 4). TYPE: Iterable[float] | ndarray

Source code in src/tn_noise_learning/channels.py

def pauli_channel(topology: Topology, probs: Iterable[float] | np.ndarray) -> LPDO:
    """Independent single-qubit Pauli channels ``ρ → Σ_P p_P P ρ P``.

    Args:
        topology: site connectivity (all bonds have size one).
        probs: ``(p_I, p_X, p_Y, p_Z)`` shared by all sites, or one row per site with shape ``(num_sites, 4)``.
    """
    probs = np.broadcast_to(np.asarray(probs, dtype=float), (topology.num_sites, 4))
    if np.any(probs < 0) or not np.allclose(probs.sum(axis=1), 1.0):
        raise ValueError("Pauli probabilities must be non-negative and sum to one on every site.")
    tensors = [_place(np.sqrt(p)[:, None, None] * PAULI, topology, j) for j, p in enumerate(probs)]
    return LPDO(topology, tensors)

pauli_eigenstates

pauli_eigenstates(num_sites: int) -> np.ndarray

The six Pauli eigenstates on every site, shape (num_sites, 6, 2, 2).

Source code in src/tn_noise_learning/channels.py

def pauli_eigenstates(num_sites: int) -> np.ndarray:
    """The six Pauli eigenstates on every site, shape ``(num_sites, 6, 2, 2)``."""
    return np.broadcast_to(PAULI_EIGENSTATES, (num_sites, 6, 2, 2)).copy()

plot_history

plot_history(
    history: History, path: str | Path | None = None
) -> Figure

One panel per metric: loss (train and test), NLL, TP penalty, trace, distance to ideal, then custom metrics.

Source code in src/tn_noise_learning/plotting.py

def plot_history(history: History, path: str | Path | None = None) -> Figure:
    """One panel per metric: loss (train and test), NLL, TP penalty, trace, distance to ideal, then custom metrics."""
    try:
        import matplotlib.pyplot as plt  # noqa: PLC0415  (optional dependency)
    except ImportError as error:
        raise ImportError(
            "Plotting requires matplotlib; install the extra with `pip install tn-noise-learning[plot]`."
        ) from error

    metrics = history.metrics
    custom = [name for name in metrics if name not in {*STANDARD_PANELS, "train_loss", "test_loss"}]
    panels = [p for p in STANDARD_PANELS if p == "loss" or p in metrics] + custom
    epochs = range(1, history.num_epochs + 1)

    fig, axes = plt.subplots(1, len(panels), figsize=(3.6 * len(panels), 3.2), squeeze=False)
    for ax, name in zip(axes[0], panels, strict=True):
        if name == "loss":
            ax.plot(epochs, metrics["train_loss"], label="train")
            if "test_loss" in metrics:
                ax.plot(epochs, metrics["test_loss"], label="test")
            ax.legend()
        else:
            ax.plot(epochs, metrics[name])
        if name == "trace" and history.num_sites is not None:
            ax.axhline(2**history.num_sites, color="gray", linestyle="--", linewidth=1)
        if name in LOG_SCALE:
            ax.set_yscale("log")
        if history.best_epoch is not None:
            ax.axvline(history.best_epoch, color="gray", linestyle=":", linewidth=1)
        ax.set_title(name.replace("_", " "))
        ax.set_xlabel("epoch")
    fig.tight_layout()
    if path is not None:
        fig.savefig(path, dpi=150)
    return fig

probabilities

probabilities(
    lpdo_tn: TensorNetwork,
    states: ndarray | Array,
    povms: ndarray | Array,
    circuit_tn: TensorNetwork | None = None,
    optimize: str = "auto-hq",
) -> jnp.ndarray

Transition probabilities Tr[E Λ(U ρ U†)] for a batch of (state, effect) pairs, clipped from below.

Source code in src/tn_noise_learning/losses.py

def probabilities(
    lpdo_tn: qtn.TensorNetwork,
    states: np.ndarray | Array,
    povms: np.ndarray | Array,
    circuit_tn: qtn.TensorNetwork | None = None,
    optimize: str = "auto-hq",
) -> jnp.ndarray:
    """Transition probabilities ``Tr[E Λ(U ρ U†)]`` for a batch of ``(state, effect)`` pairs, clipped from below."""
    parts = [lpdo_tn, product_state_network(states, raw=circuit_tn is not None), measurement_network(povms)]
    if circuit_tn is not None:
        parts.append(circuit_tn)
    tn = qtn.TensorNetwork(parts, check_collisions=False)
    probs = tn.contract(output_inds=[BATCH_IND], optimize=optimize, backend="jax").data
    return jnp.clip(jnp.real(probs), min=PROBABILITY_FLOOR)

product_state_network

product_state_network(
    states: ndarray | Array, *, raw: bool = True
) -> qtn.TensorNetwork

Batched product states states[j, exp] of shape (num_sites, batch, 2, 2) as site tensors.

Each tensor has legs (exp, ket_j, bra_j): p{j}, q{j} with raw=True (to be evolved by a circuit), otherwise k{j}, b{j}.

Source code in src/tn_noise_learning/circuits.py

def product_state_network(states: np.ndarray | Array, *, raw: bool = True) -> qtn.TensorNetwork:
    """Batched product states ``states[j, exp]`` of shape ``(num_sites, batch, 2, 2)`` as site tensors.

    Each tensor has legs ``(exp, ket_j, bra_j)``: ``p{j}, q{j}`` with ``raw=True`` (to be evolved by a circuit),
    otherwise ``k{j}, b{j}``.
    """
    ket, bra = (raw_ind, raw_ind) if raw else (input_ind, input_ind)
    return qtn.TensorNetwork(
        [qtn.Tensor(s, inds=(BATCH_IND, ket(j), bra(j, dagger=True)), tags=site_tag(j)) for j, s in enumerate(states)]
    )

project_trace_preserving

project_trace_preserving(model: LPDO, steps: int) -> LPDO

Minimise the trace-preservation penalty alone with L-BFGS-B, starting from model.

Source code in src/tn_noise_learning/trainer.py

def project_trace_preserving(model: LPDO, steps: int) -> LPDO:
    """Minimise the trace-preservation penalty alone with L-BFGS-B, starting from ``model``."""
    topology = model.topology
    optimizer = qtn.TNOptimizer(
        model.kraus_network(),
        loss_fn=lambda tn: tp_penalty(full_network(tn, topology)),
        optimizer="L-BFGS-B",
        autodiff_backend="jax",
        jit_fn=True,
        progbar=False,
    )
    return LPDO.from_network(topology, optimizer.optimize(steps))

sample_shots

sample_shots(
    channel: LPDO,
    circuit: QuantumCircuit | None = None,
    *,
    num_settings: int = 100,
    shots_per_setting: int = 100,
    seed: int | None = None,
    states: ndarray | None = None,
    povms: ndarray | None = None,
    settings: Sequence[Setting] | None = None,
) -> NoiseDataset

Sample a dataset from channel applied after circuit.

PARAMETER	DESCRIPTION
channel	ground-truth channel; stored in the dataset as ideal_channel. TYPE: LPDO
circuit	ideal circuit applied before the channel (None for the identity). TYPE: QuantumCircuit | None DEFAULT: None
num_settings	number of random (preparation, basis) settings when settings is not given. TYPE: int DEFAULT: 100
shots_per_setting	shots drawn per setting. TYPE: int DEFAULT: 100
seed	seed for settings, shots and the final shuffle. TYPE: int | None DEFAULT: None
states	input states per qubit, shape (num_qubits, num_states, 2, 2); default Pauli eigenstates. TYPE: ndarray | None DEFAULT: None
povms	effects per qubit as consecutive basis pairs; default Pauli eigenstates. TYPE: ndarray | None DEFAULT: None
settings	explicit settings, overriding num_settings. TYPE: Sequence[Setting] | None DEFAULT: None

Source code in src/tn_noise_learning/simulate.py

def sample_shots(
    channel: LPDO,
    circuit: QuantumCircuit | None = None,
    *,
    num_settings: int = 100,
    shots_per_setting: int = 100,
    seed: int | None = None,
    states: np.ndarray | None = None,
    povms: np.ndarray | None = None,
    settings: Sequence[Setting] | None = None,
) -> NoiseDataset:
    """Sample a dataset from ``channel`` applied after ``circuit``.

    Args:
        channel: ground-truth channel; stored in the dataset as ``ideal_channel``.
        circuit: ideal circuit applied before the channel (``None`` for the identity).
        num_settings: number of random ``(preparation, basis)`` settings when ``settings`` is not given.
        shots_per_setting: shots drawn per setting.
        seed: seed for settings, shots and the final shuffle.
        states: input states per qubit, shape ``(num_qubits, num_states, 2, 2)``; default Pauli eigenstates.
        povms: effects per qubit as consecutive basis pairs; default Pauli eigenstates.
        settings: explicit settings, overriding ``num_settings``.
    """
    rng = np.random.default_rng(seed)
    n = channel.num_sites
    states = pauli_eigenstates(n) if states is None else np.asarray(states, dtype=complex)
    povms = pauli_eigenstates(n) if povms is None else np.asarray(povms, dtype=complex)
    _check_bases(povms)
    if settings is None:
        settings = [
            (rng.integers(0, states.shape[1], n), rng.integers(0, povms.shape[1] // 2, n)) for _ in range(num_settings)
        ]
    bits = np.array(np.unravel_index(np.arange(2**n), (2,) * n)).T
    inputs, outcomes = [], []
    for raw_prep, raw_basis in settings:
        prep, basis = np.asarray(raw_prep), np.asarray(raw_basis)
        counts = rng.multinomial(
            shots_per_setting,
            outcome_distribution(channel, circuit, prep, basis, states=states, povms=povms).ravel(),
        )
        for flat in np.flatnonzero(counts):
            inputs.append(np.repeat(prep[None], counts[flat], axis=0))
            outcomes.append(np.repeat((2 * basis + bits[flat])[None], counts[flat], axis=0))
    order = rng.permutation(len(settings) * shots_per_setting)
    return NoiseDataset(
        np.concatenate(inputs)[order], np.concatenate(outcomes)[order], states, povms, circuit, ideal_channel=channel
    )

tp_penalty

tp_penalty(
    lpdo_tn: TensorNetwork, probs: ndarray | None = None
) -> jnp.ndarray

Squared relative distance of the output-traced channel from the identity, ‖Tr_out Λ − 𝟙‖_F² / 2**n.

Squaring keeps the gradient proportional to the violation, so large weights act as a smooth quadratic constraint. The likelihood alone rewards inflating the channel's trace, which this term counteracts; at weight w the equilibrium trace exceeds 2**n by a factor 1 + 1/(2w).

Source code in src/tn_noise_learning/losses.py

def tp_penalty(lpdo_tn: qtn.TensorNetwork, probs: jnp.ndarray | None = None) -> jnp.ndarray:  # noqa: ARG001
    """Squared relative distance of the output-traced channel from the identity, ``‖Tr_out Λ − 𝟙‖_F² / 2**n``.

    Squaring keeps the gradient proportional to the violation, so large weights act as a smooth quadratic
    constraint. The likelihood alone rewards inflating the channel's trace, which this term counteracts; at
    weight ``w`` the equilibrium trace exceeds ``2**n`` by a factor ``1 + 1/(2w)``.
    """
    num_sites = len(lpdo_tn.tensors) // 2
    traced = lpdo_tn.reindex({output_ind(j): output_ind(j, dagger=True) for j in range(num_sites)})
    identity = qtn.TensorNetwork(
        [qtn.Tensor(jnp.eye(2), inds=(input_ind(j), input_ind(j, dagger=True))) for j in range(num_sites)]
    )
    distance = jnp.real(qtn.tensor_network_distance(traced, identity, backend="jax"))
    return distance**2 / 2**num_sites

two_qubit_depolarizing

two_qubit_depolarizing(
    topology: Topology,
    p: float,
    edges: Iterable[tuple[int, int]] | None = None,
) -> LPDO

Two-qubit depolarizing channel (1 − p) ρ + p · Tr[ρ] 𝟙/4 on pairwise-disjoint edges, identity elsewhere.

PARAMETER	DESCRIPTION
topology	site connectivity. TYPE: Topology
p	depolarizing probability; the sixteen two-qubit Paulis get weights 1 − 15p/16 (identity) and p/16. TYPE: float
edges	edges that carry the channel (default: all topology edges). They must be topology edges and must not share sites. TYPE: Iterable[tuple[int, int]] | None DEFAULT: None

Source code in src/tn_noise_learning/channels.py

def two_qubit_depolarizing(topology: Topology, p: float, edges: Iterable[tuple[int, int]] | None = None) -> LPDO:
    """Two-qubit depolarizing channel ``(1 − p) ρ + p · Tr[ρ] 𝟙/4`` on pairwise-disjoint edges, identity elsewhere.

    Args:
        topology: site connectivity.
        p: depolarizing probability; the sixteen two-qubit Paulis get weights ``1 − 15p/16`` (identity) and ``p/16``.
        edges: edges that carry the channel (default: all topology edges). They must be topology edges and
            must not share sites.
    """
    edges = tuple(topology.edges) if edges is None else tuple((min(e), max(e)) for e in edges)
    sites = [s for e in edges for s in e]
    if len(set(sites)) != len(sites):
        raise ValueError("Edges carrying the two-qubit channel must be pairwise disjoint.")
    if any(e not in topology.edges for e in edges):
        raise ValueError(f"Edges {edges} must all be edges of the topology.")
    weights = np.full((4, 4), p / 16.0)
    weights[0, 0] = 1 - 15 * p / 16.0
    first = np.einsum("am,mij->amij", np.eye(4), PAULI)
    second = np.einsum("am,mij->amij", np.sqrt(weights), PAULI)
    tensors = [_place(np.eye(2, dtype=complex)[None], topology, j) for j in range(topology.num_sites)]
    for i, j in edges:
        tensors[i] = _place(first, topology, i, partner=j)
        tensors[j] = _place(second, topology, j, partner=i)
    return LPDO(topology, tensors)