#
# Copyright 2021 Quantum Benchmark Inc.
#
"""
Writing Custom Noise Sources
============================
"""
#%%
# This example will demonstrate how to write a custom noise source for the simulator.
# See also :doc:`this guide<../../guides/run/simulator>` for more general information.
import numpy as np
import trueq as tq
import trueq.simulation as tqs
#%%
# A custom noise source can be implemented by subclassing
# :py:class:`~trueq.simulation.noise_source.NoiseSource` and overriding its
# :py:meth:`~trueq.simulation.noise_source.NoiseSource.apply` method.
#
# The :py:meth:`~trueq.simulation.noise_source.NoiseSource.apply` method will be passed
# three arguments: ``cycle_wrapper``, ``state``, and ``circuit_cache``. The job of apply
# is to mutate ``state`` based on the contents of ``cycle_wrapper``. The ``state``
# refers to "the state of simulation" and can either be a
# :py:class:`~trueq.math.tensor.StateTensor` (which can store either a pure state or a
# mixed state), or a :py:class:`~trueq.math.tensor.OperatorTensor` (which can store
# either a unitary or a superoperator). Each of these four possibilities share the
# method :py:meth:`~trueq.math.tensor.Tensor.apply_matrix` which can always accept a
# unitary or superoperator as input, and for example upgrades pure states to mixed
# states when necessary, and so the distinction between these various types of states
# becomes irrelevant to the implementation of a noise source.
#
# Note that ``cycle_wrapper`` is not a :py:class:`~trueq.Cycle` instance as one
# might expect, but instead a thin wrapper called
# :py:class:`~trueq.simulation.match.CycleWrapper` which in turn wraps each cycle
# operation as :py:class:`~trueq.simulation.match.OpWrapper`.
#
# These wrapper types are required to store metadata on each operation (also discussed
# further below). However, noise sources will rarely need to deal with these wrapper
# objects directly. This is because noise sources (by strongly encouraged convention,
# though not by contract) own a :py:class:`~trueq.simulation.match.Match` instance which
# takes ownership of the mechanics of these wrappers. This is seen in the example below,
# where the :py:class:`~trueq.simulation.match.Match.iter_gates` method is used to yield
# gate objects with their corresponding gate labels.
# predefine noise matrices to be applied to qubits following ideal gates
# note that rowstack_subsys is the superoperator convention required by the simulator
p = 0.01
s = tq.math.Superop.from_kraus([np.sqrt(1 - p) * np.eye(2), np.sqrt(p) * tq.Gate.x.mat])
s1 = s.rowstack_subsys()
p = 0.02
s = tq.math.Superop.from_kraus(
[np.sqrt(1 - p) * np.eye(4), np.sqrt(p) * tq.math.random_unitary(4)]
)
s2 = s.rowstack_subsys()
class ExampleNoise(tqs.NoiseSource):
def __init__(self, match=None):
# this simulator hardcodes the subsystem dimension to 2
super().__init__(dim=2, match=match)
def make_circuit_cache(self, circuit):
# this return will be made available to apply() as circuit_cache for every
# cycle in the circuit
return set(circuit.labels)
def apply(self, cycle_wrapper, state, circuit_cache):
# in this method we need to mutate state by inspecting cycle_wrapper
used_labels = set()
# loop through the gates in the cycle and mutate the state
# note that we must set noise_only=False since we are simulating each gate
for labels, gate in self.match.iter_gates(cycle_wrapper, noise_only=False):
used_labels.update(labels)
# first do ideal simulation of this gate
state.apply_matrix(labels, gate.mat)
# now add some noise to each qubit the gate acts on
if len(labels) == 1:
state.apply_matrix(labels, s1)
elif len(labels) == 2:
state.apply_matrix(labels, s2)
# apply a 20 degree Z rotation to every qubit without a gate in this cycle
for label in circuit_cache.difference(used_labels):
state.apply_matrix((label,), tq.Gate.rp("Z", 20))
# %%
# Now we can instantiate a simulator that uses this noise source and do any simulator
# calculation. Here, we use the simulator to predict the output of a hypothetical
# KNR experiment.
sim = tq.Simulator().append_noise_source(ExampleNoise())
cycle = {(0, 1): tq.Gate.cnot, 2: tq.Gate.x}
fit = sim.predict_knr(cycle, twirl=tq.Twirl("P", range(5)))
fit.plot.knr_heatmap()
# %%
# Caching
# -------
#
# There are two distinct caches available to a noise source. The first is a private
# noise source attribute that is called ``_cache`` (by convention, not by contract) in
# built-in noise models which persists throughout the lifetime of the noise source. For
# example, a gate-dependent noise model may consider caching gate noise if it is
# expensive to recompute everytime the same gate is encountered.
#
# The second cache comes as the third argument to the apply
# method. This cache is instantiated by
# :py:meth:`~trueq.simulation.noise_source.NoiseSource.make_circuit_cache` just before a
# new circuit is simulated and is presented as the third argument to the apply method
# for every cycle within that circuit. It typically stores information about the
# circuit that is not available in every cycle, such as all of the labels the circuit
# acts on, or which measurements it performs at the end.
# %%
# Wrappers and metadata
# ---------------------
#
# The variable ``cycle_wrapper`` in the example above has type
# :py:class:`~trueq.simulation.match.CycleWrapper` which is a thin wrapper for a cycle
# and also wraps every cycle operation with
# :py:class:`~trueq.simulation.match.OpWrapper`\. During simulation, everytime a new
# cycle is encountered it is wrapped once, and the same wrapped cycle instance is passed
# to all noise sources. The cycle wrapper exists to enable efficient looping logic used
# by :py:class:`~trueq.simulation.match.Match` and to persist operation wrappers for
# multiple noise sources\. The operation wrapper exists to store
# two important pieces of information:
#
# 1. Whether some noise source has previously simulated the same operation of the
# cycle. This is stored as a boolean called ``has_been_simulated`` and is set to
# true whenever a match method is called with the argument ``noise_only=False``, as
# it is in the above example.
# 2. Whether no noise sources (except the ideal final simulation if relevant) should
# touch the operation again. This is stored as a boolean called ``no_more_noise``
# and is set to true whenever it is yielded by a match method whose ``exclusive``
# value is true.