diff --git a/examples/vqe_qudit_example.py b/examples/vqe_qudit_example.py new file mode 100644 index 00000000..88179557 --- /dev/null +++ b/examples/vqe_qudit_example.py @@ -0,0 +1,171 @@ +""" +VQE on QuditCircuits. + +This example shows how to run a simple VQE on a qudit system using +`tensorcircuit.QuditCircuit`. We build a compact ansatz using single-qudit +rotations in selected two-level subspaces and RXX-type entanglers, then +optimize the energy of a Hermitian "clock–shift" Hamiltonian: + + H(d) = - J * (X_c \otimes X_c) - h * (Z_c \otimes I + I \otimes Z_c) + +where, for local dimension `d`, +- Z_c = (Z + Z^\dagger)/2 is the Hermitian "clock" observable with Z = diag(1, \omega, \omega^2, ..., \omega^{d-1}) +- X_c = (S + S^\dagger)/2 is the Hermitian "shift" observable with S the cyclic shift +- \omega = exp(2\pi i/d) + +The code defaults to a 2-qutrit (d=3) problem but can be changed via CLI flags. +""" + +# import os, sys +# +# base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +# if base_dir not in sys.path: +# sys.path.insert(0, base_dir) + +import time +import argparse +import tensorcircuit as tc + +tc.set_backend("jax") +tc.set_dtype("complex128") + + +def vqe_forward(param, *, nqudits: int, d: int, nlayers: int, J: float, h: float): + """Build a QuditCircuit ansatz and compute ⟨H⟩. + + Ansatz: + [ for L in 1...nlayers ] + - On each site q: + RX(q; θ_Lq^(01)) ∘ RY(q; θ_Lq^(12)) ∘ RZ(q; φ_Lq^(0)) + (subspace indices shown as superscripts) + - Entangle neighboring pairs with RXX on subspaces (0,1) + """ + if d < 3: + raise ValueError("This example assblumes d >= 3 (qutrit or higher).") + + S = tc.quditgates._x_matrix_func(d) + Z = tc.quditgates._z_matrix_func(d) + Sdag = tc.backend.adjoint(S) + Zdag = tc.backend.adjoint(Z) + + c = tc.QuditCircuit(nqudits, dim=d) + + pairs = [(i, i + 1) for i in range(nqudits - 1)] + + it = iter(param) + + for _ in range(nlayers): + for q in range(nqudits): + c.rx(q, theta=next(it), j=0, k=1) + c.ry(q, theta=next(it), j=1, k=2) + c.rz(q, theta=next(it), j=0) + + for i, j in pairs: + c.rxx(i, j, theta=next(it), j1=0, k1=1, j2=0, k2=1) + + # H = -J * 1/2 (S_i S_j^\dagger + S_i^\dagger S_j) - h * 1/2 (Z + Z^\dagger) + energy = 0.0 + for i, j in pairs: + e_ij = 0.5 * ( + c.expectation((S, [i]), (Sdag, [j])) + c.expectation((Sdag, [i]), (S, [j])) + ) + energy += -J * tc.backend.real(e_ij) + for q in range(nqudits): + zq = 0.5 * (c.expectation((Z, [q])) + c.expectation((Zdag, [q]))) + energy += -h * tc.backend.real(zq) + return tc.backend.real(energy) + + +def build_param_shape(nqudits: int, d: int, nlayers: int): + # Per layer per qudit: RX^(01), RY^(12) (or dummy), RZ^(0) = 3 params + # Per layer entanglers: len(pairs) parameters + pairs = nqudits - 1 + per_layer = 3 * nqudits + pairs + return (nlayers * per_layer,) + + +def main(): + parser = argparse.ArgumentParser( + description="VQE on QuditCircuit (clock–shift model)" + ) + parser.add_argument( + "--d", type=int, default=3, help="Local dimension per site (>=3)" + ) + parser.add_argument("--nqudits", type=int, default=2, help="Number of sites") + parser.add_argument("--nlayers", type=int, default=3, help="Ansatz depth (layers)") + parser.add_argument( + "--J", type=float, default=1.0, help="Coupling strength for XcXc term" + ) + parser.add_argument( + "--h", type=float, default=0.6, help="Field strength for Zc terms" + ) + parser.add_argument("--steps", type=int, default=200, help="Optimization steps") + parser.add_argument("--lr", type=float, default=0.05, help="Learning rate") + args = parser.parse_args() + + assert args.d >= 3, "d must be >= 3" + + shape = build_param_shape(args.nqudits, args.d, args.nlayers) + param = tc.backend.random_uniform(shape, boundaries=(-0.1, 0.1), seed=42) + + try: + import optax + + optimizer = optax.adam(args.lr) + vgf = tc.backend.jit( + tc.backend.value_and_grad( + lambda p: vqe_forward( + p, + nqudits=args.nqudits, + d=args.d, + nlayers=args.nlayers, + J=args.J, + h=args.h, + ) + ) + ) + opt_state = optimizer.init(param) + + @tc.backend.jit + def train_step(p, opt_state): + loss, grads = vgf(p) + updates, opt_state = optimizer.update(grads, opt_state, p) + p = optax.apply_updates(p, updates) + return p, opt_state, loss + + print("Starting VQE optimization (optax/adam)...") + loss = None + for i in range(args.steps): + t0 = time.time() + param, opt_state, loss = train_step(param, opt_state) + # ensure sync for accurate timing + _ = float(loss) + if i % 20 == 0: + dt = time.time() - t0 + print(f"Step {i:4d} loss={loss:.6f} dt/step={dt:.4f}s") + print("Final loss:", float(loss) if loss is not None else "n/a") + + except ModuleNotFoundError: + print("Optax not available; using naive gradient descent.") + value_and_grad = tc.backend.value_and_grad( + lambda p: vqe_forward( + p, + nqudits=args.nqudits, + d=args.d, + nlayers=args.nlayers, + J=args.J, + h=args.h, + ) + ) + lr = args.lr + loss = None + for i in range(args.steps): + loss, grads = value_and_grad(param) + param = param - lr * grads + if i % 20 == 0: + print(f"Step {i:4d} loss={float(loss):.6f}") + print("Final loss:", float(loss) if loss is not None else "n/a") + + +if __name__ == "__main__": + main() diff --git a/tensorcircuit/__init__.py b/tensorcircuit/__init__.py index b6c4ef09..c4aea55c 100644 --- a/tensorcircuit/__init__.py +++ b/tensorcircuit/__init__.py @@ -23,8 +23,10 @@ runtime_contractor, ) # prerun of set hooks from . import gates +from . import quditgates from . import basecircuit from .gates import Gate +from .quditcircuit import QuditCircuit from .circuit import Circuit, expectation from .mpscircuit import MPSCircuit from .densitymatrix import DMCircuit as DMCircuit_reference diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py index f996cf08..ffa50498 100644 --- a/tensorcircuit/basecircuit.py +++ b/tensorcircuit/basecircuit.py @@ -480,28 +480,12 @@ def measure_jit( def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]: r""" - Returns the tensornetwor nodes for the amplitude of the circuit given a computational-basis label ``l``. - For a state simulator, it computes :math:`\langle l \vert \psi\rangle`; - for a density-matrix simulator, it computes :math:`\mathrm{Tr}(\rho \vert l\rangle\langle l\vert)`. + Returns the tensornetwor nodes for the amplitude of the circuit given the bitstring l. + For state simulator, it computes :math:`\langle l\vert \psi\rangle`, + for density matrix simulator, it computes :math:`Tr(\rho \vert l\rangle \langle 1\vert)` Note how these two are different up to a square operation. - :Example: - - >>> c = tc.Circuit(2) - >>> c.X(0) - >>> c.amplitude("10") # d=2, per-qubit digits - array(1.+0.j, dtype=complex64) - >>> c.CNOT(0, 1) - >>> c.amplitude("11") - array(1.+0.j, dtype=complex64) - - For qudits (d>2, d<=36): - >>> c = tc.Circuit(3, dim=12) - >>> c.amplitude("0A2") # base-12 string, A stands for 10 - - :param l: Basis label. - - If a string: it must be a base-d string of length ``nqubits``, using 0–9A–Z (A=10,…,Z=35) when ``d<=36``. - - If a tensor/array/list: it should contain per-site integers in ``[0, d-1]`` with length ``nqubits``. + :param l: The bitstring of 0 and 1s. :type l: Union[str, Tensor] :return: The tensornetwork nodes for the amplitude of the circuit. :rtype: List[Gate] diff --git a/tensorcircuit/quditcircuit.py b/tensorcircuit/quditcircuit.py new file mode 100644 index 00000000..9ae0f013 --- /dev/null +++ b/tensorcircuit/quditcircuit.py @@ -0,0 +1,618 @@ +""" +Quantum circuit: the state simulator. +Supports qudit (3 <= dim <= 36) systems. +For string-encoded samples/counts, digits use 0–9A–Z where A=10, …, Z=35. +""" + +from functools import partial +from typing import Any, Dict, List, Optional, Tuple, Sequence, Union, Literal + +import numpy as np +import tensornetwork as tn + +from .gates import Gate +from .utils import arg_alias +from .basecircuit import BaseCircuit +from .circuit import Circuit +from .quantum import QuOperator, QuVector +from .quditgates import SINGLE_BUILDERS, TWO_BUILDERS + + +Tensor = Any +SAMPLE_FORMAT = Literal["sample_bin", "count_dict_bin"] + + +class QuditCircuit: + r""" + ``QuditCircuit`` class. + + Qudit quick example (d=3): + .. code-block:: python + + c = tc.QuditCircuit(2, d=3) + c.h(0) + c.x(1) + c.csum(0, 1) + c.sample(1024, format="count_dict_bin") + # For d <= 36, string samples use base-d characters 0–9A–Z (A=10, ...). + """ + + is_dm = False + + def __init__( + self, + nqubits: int, + dim: int, + inputs: Optional[Tensor] = None, + mps_inputs: Optional[QuOperator] = None, + split: Optional[Dict[str, Any]] = None, + ): + self._set_dim(dim=dim) + self._nqubits = nqubits + + self._circ = Circuit( + nqubits=nqubits, + dim=dim, + inputs=inputs, + mps_inputs=mps_inputs, + split=split, + ) + self._omega = np.exp(2j * np.pi / self._d) + self.circuit_param = self._circ.circuit_param + + def _set_dim(self, dim: int) -> None: + if not isinstance(dim, int) or dim <= 2: + raise ValueError( + f"QuditCircuit is only for qudits (dim>=3). " + f"You passed dim={dim}. For qubits, please use `Circuit` instead." + ) + # Require integer d>=2; current string-encoded IO supports d<=36 (0–9A–Z digits). + if dim > 36: + raise NotImplementedError( + "The Qudit interface is only supported for dimension < 36 now." + ) + self._d = dim + + @property + def dim(self) -> int: + """dimension of the qudit circuit""" + return self._d + + @property + def nqubits(self) -> int: + """qudit number of the circuit""" + return self._nqubits + + def _apply_gate(self, *indices: int, name: str, **kwargs: Any) -> None: + """ + Apply a quantum gate (unitary) to one or two qudits in the circuit. + + The gate matrix is looked up by name in either ``SINGLE_BUILDERS`` (for single-qudit gates) + or ``TWO_BUILDERS`` (for two-qudit gates). The matrix is built by the registered builder, + then applied to the circuit at the given indices. + + :param indices: The qudit indices the gate should act on. + - One index → single-qudit gate. + - Two indices → two-qudit gate. + :type indices: int + :param name: The name of the gate (must exist in the chosen builder set). + :type name: str + :param kwargs: Extra parameters for the gate matched against the builder signature. + :type kwargs: Any + :raises ValueError: If ``name`` is not found, + or if the number of indices does not match the gate type (single vs two). + """ + if len(indices) == 1 and name in SINGLE_BUILDERS: + sig, builder = SINGLE_BUILDERS[name] + extras = tuple(kwargs.get(k) for k in sig if k != "none") + builder_kwargs = {k: v for k, v in zip(sig, extras)} + mat = builder(self._d, self._omega, **builder_kwargs) + self._circ.unitary(*indices, unitary=mat, name=name, dim=self._d) # type: ignore + elif len(indices) == 2 and name in TWO_BUILDERS: + sig, builder = TWO_BUILDERS[name] + extras = tuple(kwargs.get(k) for k in sig if k != "none") + builder_kwargs = {k: v for k, v in zip(sig, extras)} + mat = builder(self._d, self._omega, **builder_kwargs) + self._circ.unitary( # type: ignore + *indices, unitary=mat, name=name, dim=self._d + ) + else: + raise ValueError(f"Unsupported gate/arity: {name} on {len(indices)} qudits") + + def any(self, *indices: int, unitary: Tensor, name: str = "any") -> None: + """ + Apply a quantum gate (unitary) to one or two qudits in the circuit. + + :param indices: The qudit indices the gate should act on. + :type indices: int + :param unitary: The unitary matrix to apply to the qudit(s). + :type unitary: Tensor + :param name: The name to record for this gate. + :type name: str + """ + self._circ.unitary(*indices, unitary=unitary, name=name, dim=self._d) # type: ignore + + unitary = any + + def i(self, index: int) -> None: + """ + Apply the identity (I) gate on the given qudit index. + + :param index: Qudit index to apply the gate on. + :type index: int + """ + self._apply_gate(index, name="I") + + I = i + + def x(self, index: int) -> None: + """ + Apply the X gate on the given qudit index. + + :param index: Qudit index to apply the gate on. + :type index: int + """ + self._apply_gate(index, name="X") + + X = x + + # def y(self, index: int) -> None: + # """ + # Apply the Y gate on the given qudit index. + # + # :param index: Qudit index to apply the gate on. + # :type index: int + # """ + # self._apply_gate(index, name="Y") + + def z(self, index: int) -> None: + """ + Apply the Z gate on the given qudit index. + + :param index: Qudit index to apply the gate on. + :type index: int + """ + self._apply_gate(index, name="Z") + + Z = z + + def h(self, index: int) -> None: + """ + Apply the Hadamard-like (H) gate on the given qudit index. + + :param index: Qudit index to apply the gate on. + :type index: int + """ + self._apply_gate(index, name="H") + + H = h + + def u8( + self, index: int, gamma: float = 2.0, z: float = 1.0, eps: float = 0.0 + ) -> None: + """ + Apply the U8 parameterized single-qudit gate. + + :param index: Qudit index to apply the gate on. + :type index: int + :param gamma: Gate parameter ``gamma``. + :type gamma: float + :param z: Gate parameter ``z``. + :type z: float + :param eps: Gate parameter ``eps``. + :type eps: float + """ + self._apply_gate(index, name="U8", extra=(gamma, z, eps)) + + U8 = u8 + + def rx(self, index: int, theta: float, j: int = 0, k: int = 1) -> None: + """ + Apply the single-qudit RX rotation on ``index``. + + :param index: Qudit index to apply the gate on. + :type index: int + :param theta: Rotation angle. + :type theta: float + :param j: Source level of the rotation subspace. + :type j: int + :param k: Target level of the rotation subspace. + :type k: int + """ + self._apply_gate(index, name="RX", theta=theta, j=j, k=k) + + RX = rx + + def ry(self, index: int, theta: float, j: int = 0, k: int = 1) -> None: + """ + Apply the single-qudit RY rotation on ``index``. + + :param index: Qudit index to apply the gate on. + :type index: int + :param theta: Rotation angle. + :type theta: float + :param j: Source level of the rotation subspace. + :type j: int + :param k: Target level of the rotation subspace. + :type k: int + """ + self._apply_gate(index, name="RY", theta=theta, j=j, k=k) + + RY = ry + + def rz(self, index: int, theta: float, j: int = 0) -> None: + """ + Apply the single-qudit RZ rotation on ``index``. + + :param index: Qudit index to apply the gate on. + :type index: int + :param theta: Rotation angle around Z. + :type theta: float + :param j: Level where the phase rotation is applied. + :type j: int + """ + self._apply_gate(index, name="RZ", theta=theta, j=j) + + RZ = rz + + def rxx( + self, + *indices: int, + theta: float, + j1: int = 0, + k1: int = 1, + j2: int = 0, + k2: int = 1, + ) -> None: + """ + Apply a two-qudit RXX-type interaction on the given indices. + + :param indices: Two qudit indices. + :type indices: int + :param theta: Interaction strength/angle. + :type theta: float + :param j1: Source level of the first qudit subspace. + :type j1: int + :param k1: Target level of the first qudit subspace. + :type k1: int + :param j2: Source level of the second qudit subspace. + :type j2: int + :param k2: Target level of the second qudit subspace. + :type k2: int + """ + self._apply_gate(*indices, name="RXX", theta=theta, j1=j1, k1=k1, j2=j2, k2=k2) + + RXX = rxx + + def rzz( + self, + *indices: int, + theta: float, + j1: int = 0, + k1: int = 1, + j2: int = 0, + k2: int = 1, + ) -> None: + """ + Apply a two-qudit RZZ-type interaction on the given indices. + + :param indices: Two qudit indices. + :type indices: int + :param theta: Interaction strength/angle. + :type theta: float + :param j1: Source level of the first qudit subspace. + :type j1: int + :param k1: Target level of the first qudit subspace. + :type k1: int + :param j2: Source level of the second qudit subspace. + :type j2: int + :param k2: Target level of the second qudit subspace. + :type k2: int + """ + self._apply_gate(*indices, name="RZZ", theta=theta, j1=j1, k1=k1, j2=j2, k2=k2) + + RZZ = rzz + + def cphase(self, *indices: int, cv: Optional[int] = None) -> None: + """ + Apply a controlled phase (CPHASE) gate. + + :param indices: Two qudit indices (control, target). + :type indices: int + :param cv: Optional control value. If ``None``, defaults to ``1``. + :type cv: Optional[int] + """ + self._apply_gate(*indices, name="CPHASE", cv=cv) + + CPHASE = cphase + + def csum(self, *indices: int, cv: Optional[int] = None) -> None: + """ + Apply a controlled-sum (CSUM) gate. + + :param indices: Two qudit indices (control, target). + :type indices: int + :param cv: Optional control value. If ``None``, defaults to ``1``. + :type cv: Optional[int] + """ + self._apply_gate(*indices, name="CSUM", cv=cv) + + cnot, CSUM, CNOT = csum, csum, csum + + # Functional + def wavefunction(self, form: str = "default") -> tn.Node.tensor: + return self._circ.wavefunction(form) + + state = wavefunction + + def get_quoperator(self) -> QuOperator: + """ + Get the ``QuOperator`` MPO like representation of the circuit unitary without contraction. + + :return: ``QuOperator`` object for the circuit unitary (open indices for the input state) + :rtype: QuOperator + """ + return self._circ.quoperator() + + quoperator = get_quoperator + + get_circuit_as_quoperator = get_quoperator + get_state_as_quvector = BaseCircuit.quvector + + def matrix(self) -> Tensor: + """ + Get the unitary matrix for the circuit irrespective with the circuit input state. + + :return: The circuit unitary matrix + :rtype: Tensor + """ + return self._circ.matrix() + + # def measure_reference( + # self, *index: int, with_prob: bool = False + # ) -> Tuple[str, float]: + # return self._circ.measure_reference(*index, with_prob=with_prob) + + def expectation( + self, + *ops: Tuple[tn.Node, List[int]], + reuse: bool = True, + enable_lightcone: bool = False, + nmc: int = 1000, + status: Optional[Tensor] = None, + **kws: Any, + ) -> Tensor: + return self._circ.expectation( + *ops, + reuse=reuse, + enable_lightcone=enable_lightcone, + noise_conf=None, + nmc=nmc, + status=status, + **kws, + ) + + def measure_jit( + self, *index: int, with_prob: bool = False, status: Optional[Tensor] = None + ) -> Tuple[Tensor, Tensor]: + """ + Take measurement on the given site indices (computational basis). + This method is jittable! + + :param index: Measure on which site (wire) index. + :type index: int + :param with_prob: If true, theoretical probability is also returned. + :type with_prob: bool, optional + :param status: external randomness, with shape [index], defaults to None + :type status: Optional[Tensor] + :return: The sample output and probability (optional) of the quantum line. + :rtype: Tuple[Tensor, Tensor] + """ + return self._circ.measure_jit(*index, with_prob=with_prob, status=status) + + measure = measure_jit + + def amplitude(self, l: Union[str, Tensor]) -> Tensor: + r""" + Returns the amplitude of the circuit given the bitstring l. + For state simulator, it computes :math:`\langle l\vert \psi\rangle`, + for density matrix simulator, it computes :math:`Tr(\rho \vert l\rangle \langle 1\vert)` + Note how these two are different up to a square operation. + + :Example: + + >>> c = tc.QuditCircuit(2, dim=3) + >>> c.x(0) + >>> c.x(1) + >>> c.amplitude("20") + array(1.+0.j, dtype=complex64) + >>> c.csum(0, 1, cv=2) + >>> c.amplitude("21") + array(1.+0.j, dtype=complex64) + + :param l: The bitstring of base-d characters. + :type l: Union[str, Tensor] + :return: The amplitude of the circuit. + :rtype: tn.Node.tensor + """ + return self._circ.amplitude(l) + + def probability(self) -> Tensor: + """ + Get the ``d^n`` length probability vector over the computational basis. + + :return: Probability vector of shape ``[dim**n]``. + :rtype: Tensor + """ + return self._circ.probability() + + @partial(arg_alias, alias_dict={"format": ["format_"]}) + def sample( + self, + batch: Optional[int] = None, + allow_state: bool = False, + readout_error: Optional[Sequence[Any]] = None, + format: Optional[SAMPLE_FORMAT] = None, + random_generator: Optional[Any] = None, + status: Optional[Tensor] = None, + jittable: bool = True, + ) -> Any: + r""" + batched sampling from state or circuit tensor network directly + + :param batch: number of samples, defaults to None + :type batch: Optional[int], optional + :param allow_state: if true, we sample from the final state + if memory allows, True is preferred, defaults to False + :type allow_state: bool, optional + :param readout_error: readout_error, defaults to None + :type readout_error: Optional[Sequence[Any]]. Tensor, List, Tuple + :param format: sample format, defaults to None as backward compatibility + check the doc in :py:meth:`tensorcircuit.quantum.measurement_results` + Six formats of measurement counts results: + + "sample_bin": # [np.array([1, 0]), np.array([1, 0])] + + "count_vector": # np.array([2, 0, 0, 0]) + + "count_dict_bin": # {"00": 2, "01": 0, "10": 0, "11": 0} + for cases d\in [11, 36], use 0–9A–Z digits (e.g., 'A' -> 10, …, 'Z' -> 35); + + :type format: Optional[str] + :param random_generator: random generator, defaults to None + :type random_generator: Optional[Any], optional + :param status: external randomness given by tensor uniformly from [0, 1], + if set, can overwrite random_generator, shape [batch] for `allow_state=True` + and shape [batch, nqubits] for `allow_state=False` using perfect sampling implementation + :type status: Optional[Tensor] + :param jittable: when converting to count, whether keep the full size. if false, may be conflict + external jit, if true, may fail for large scale system with actual limited count results + :type jittable: bool, defaults true + :return: List (if batch) of tuple (binary configuration tensor and corresponding probability) + if the format is None, and consistent with format when given + :rtype: Any + """ + if format in ["sample_int", "count_tuple", "count_dict_int", "count_vector"]: + raise NotImplementedError( + "`int` representation is not friendly for d-dimensional systems." + ) + return self._circ.sample( + batch=batch, + allow_state=allow_state, + readout_error=readout_error, + format=format, + random_generator=random_generator, + status=status, + jittable=jittable, + ) + + def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor: + """ + remaining wavefunction or density matrix on sites in ``left``, with other sites + fixed to given digits (0..d-1) as indicated by ``traceout`` + + :param traceout: can be jitted + :type traceout: Tensor + :param left: cannot be jitted + :type left: Tuple + :return: _description_ + :rtype: Tensor + """ + return self._circ.projected_subsystem( + traceout=traceout, + left=left, + ) + + def replace_inputs(self, inputs: Tensor) -> None: + """ + Replace the input state with the circuit structure unchanged. + + :param inputs: Input wavefunction. + :type inputs: Tensor + """ + return self._circ.replace_inputs(inputs) + + def mid_measurement(self, index: int, keep: int = 0) -> Tensor: + """ + Middle measurement in z-basis on the circuit, note the wavefunction output is not normalized + with ``mid_measurement`` involved, one should normalize the state manually if needed. + This is a post-selection method as keep is provided as a prior. + + :param index: The index of qubit that the Z direction postselection applied on. + :type index: int + :param keep: the post-selected digit in {0, ..., d-1}, defaults to be 0. + :type keep: int, optional + """ + return self._circ.mid_measurement(index, keep=keep) + + mid_measure = mid_measurement + post_select = mid_measurement + post_selection = mid_measurement + + def get_quvector(self) -> QuVector: + """ + Get the representation of the output state in the form of ``QuVector`` + while maintaining the circuit uncomputed + + :return: ``QuVector`` representation of the output state from the circuit + :rtype: QuVector + """ + return self._circ.quvector() + + quvector = get_quvector + + def replace_mps_inputs(self, mps_inputs: QuOperator) -> None: + """ + Replace the input state in MPS representation while keep the circuit structure unchanged. + + :Example: + >>> c = tc.QuditCircuit(2, dim=3) + >>> c.x(0) + >>> + >>> c2 = tc.QuditCircuit(2, dim=3, mps_inputs=c.quvector()) + >>> c2.x(0) + >>> c2.wavefunction() + array([0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j], dtype=complex64) + >>> + >>> c3 = tc.QuditCircuit(2, dim=3) + >>> c3.x(0) + >>> c3.replace_mps_inputs(c.quvector()) + >>> c3.wavefunction() + array([0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j], dtype=complex64) + + :param mps_inputs: (Nodes, dangling Edges) for a MPS like initial wavefunction. + :type mps_inputs: Tuple[Sequence[Gate], Sequence[Edge]] + """ + return self._circ.replace_mps_inputs(mps_inputs) + + def expectation_before( + self, + *ops: Tuple[tn.Node, List[int]], + reuse: bool = True, + **kws: Any, + ) -> List[tn.Node]: + """ + Get the tensor network in the form of a list of nodes + for the expectation calculation before the real contraction + + :param reuse: _description_, defaults to True + :type reuse: bool, optional + :raises ValueError: _description_ + :return: _description_ + :rtype: List[tn.Node] + """ + return self._circ.expectation_before(*ops, reuse=reuse, **kws) + + def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]: + r""" + Returns the tensornetwor nodes for the amplitude of the circuit given the bitstring l. + For state simulator, it computes :math:`\langle l\vert \psi\rangle`, + for density matrix simulator, it computes :math:`Tr(\rho \vert l\rangle \langle 1\vert)` + Note how these two are different up to a square operation. + + :param l: The bitstring of 0 and 1s. + :type l: Union[str, Tensor] + :return: The tensornetwork nodes for the amplitude of the circuit. + :rtype: List[Gate] + """ + return self._circ.amplitude_before(l) diff --git a/tensorcircuit/quditgates.py b/tensorcircuit/quditgates.py new file mode 100644 index 00000000..a4e17ed8 --- /dev/null +++ b/tensorcircuit/quditgates.py @@ -0,0 +1,558 @@ +from typing import Any, Optional, Tuple + +import numpy as np + +from .cons import backend, dtypestr +from .gates import num_to_tensor + +Tensor = Any + +SINGLE_BUILDERS = { + "I": (("none",), lambda d, omega, **kw: _i_matrix_func(d)), + "X": (("none",), lambda d, omega, **kw: _x_matrix_func(d)), + "Z": (("none",), lambda d, omega, **kw: _z_matrix_func(d, omega)), + "H": (("none",), lambda d, omega, **kw: _h_matrix_func(d, omega)), + "RX": ( + ("theta", "j", "k"), + lambda d, omega, **kw: _rx_matrix_func(d, kw["theta"], kw["j"], kw["k"]), + ), + "RY": ( + ("theta", "j", "k"), + lambda d, omega, **kw: _ry_matrix_func(d, kw["theta"], kw["j"], kw["k"]), + ), + "RZ": ( + ("theta", "j"), + lambda d, omega, **kw: _rz_matrix_func(d, kw["theta"], kw["j"]), + ), + "U8": ( + ("gamma", "z", "eps"), + lambda d, omega, **kw: _u8_matrix_func( + d, kw["gamma"], kw["z"], kw["eps"], omega + ), + ), +} + +TWO_BUILDERS = { + "RXX": ( + ("theta", "j1", "k1", "j2", "k2"), + lambda d, omega, **kw: _rxx_matrix_func( + d, kw["theta"], kw["j1"], kw["k1"], kw["j2"], kw["k2"] + ), + ), + "RZZ": (("theta",), lambda d, omega, **kw: _rzz_matrix_func(d, kw["theta"])), + "CPHASE": (("cv",), lambda d, omega, **kw: _cphase_matrix_func(d, kw["cv"], omega)), + "CSUM": (("cv",), lambda d, omega, **kw: _csum_matrix_func(d, kw["cv"])), +} + + +def _is_prime(n: int) -> bool: + """ + Check whether a number is prime. + + :param n: Integer to test. + :type n: int + :return: ``True`` if ``n`` is prime, else ``False``. + :rtype: bool + """ + if n < 2: + return False + if n in (2, 3, 5, 7): + return True + if n % 2 == 0 or n % 3 == 0: + return False + + r = int(n**0.5) + 1 + for i in range(5, r, 6): + if n % i == 0 or n % (i + 2) == 0: + return False + return True + + +def _i_matrix_func(d: int) -> Tensor: + """ + Identity matrix of size ``d``. + + :param d: Qudit dimension. + :type d: int + :return: ``(d, d)`` identity matrix. + :rtype: Tensor + """ + return backend.eye(d, dtype=dtypestr) + + +def _x_matrix_func(d: int) -> Tensor: + r""" + Generalized Pauli-X on a ``d``-level system. + + .. math:: X_d\lvert j \rangle = \lvert (j+1) \bmod d \rangle + + :param d: Qudit dimension. + :type d: int + :return: ``(d, d)`` matrix for :math:`X_d`. + :rtype: Tensor + """ + m = np.roll(np.eye(d), shift=1, axis=0) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _z_matrix_func(d: int, omega: Optional[complex] = None) -> Tensor: + r""" + Generalized Pauli-Z on a ``d``-level system. + + .. math:: Z_d\lvert j \rangle = \omega^{j}\lvert j \rangle,\quad \omega=e^{2\pi i/d} + + :param d: Qudit dimension. + :type d: int + :param omega: Optional primitive ``d``-th root of unity. Defaults to :math:`e^{2\pi i/d}`. + :type omega: Optional[complex] + :return: ``(d, d)`` matrix for :math:`Z_d`. + :rtype: Tensor + """ + omega = np.exp(2j * np.pi / d) if omega is None else omega + m = np.diag(omega ** np.arange(d)) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _h_matrix_func(d: int, omega: Optional[complex] = None) -> Tensor: + r""" + Discrete Fourier transform (Hadamard-like) on ``d`` levels. + + .. math:: H_d\lvert j \rangle = \frac{1}{\sqrt{d}} \sum_{k=0}^{d-1} \omega^{jk}\lvert k \rangle + + :param d: Qudit dimension. + :type d: int + :param omega: Optional primitive ``d``-th root of unity. Defaults to :math:`e^{2\pi i/d}`. + :type omega: Optional[complex] + :return: ``(d, d)`` matrix for :math:`H_d`. + :rtype: Tensor + """ + omega = np.exp(2j * np.pi / d) if omega is None else omega + j, k = np.arange(d), np.arange(d) + m = omega ** np.outer(j, k) / np.sqrt(d) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _s_matrix_func(d: int, omega: Optional[complex] = None) -> Tensor: + r""" + Diagonal phase gate ``S_d`` on ``d`` levels. + + .. math:: S_d\lvert j \rangle = \omega^{j(j+p_d)/2}\lvert j \rangle,\quad p_d = (d \bmod 2) + + :param d: Qudit dimension. + :type d: int + :param omega: Optional primitive ``d``-th root of unity. Defaults to :math:`e^{2\pi i/d}`. + :type omega: Optional[complex] + :return: ``(d, d)`` diagonal matrix for :math:`S_d`. + :rtype: Tensor + """ + omega = np.exp(2j * np.pi / d) if omega is None else omega + pd = 0 if d % 2 == 0 else 1 + + j = np.arange(d) + m = np.diag(omega ** ((j * (j + pd)) / 2)) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _check_rotation_indices( + d: int, *indices: int, distinct_pairs: bool = False +) -> None: + """ + Validate that indices are within [0, d-1] and optionally form distinct pairs. + + :param d: Qudit dimension. + :type d: int + :param indices: Indices to validate. + :type indices: int + :param distinct_pairs: If True, enforce that (indices[0], indices[1]) + ≠ (indices[2], indices[3]) for 4 indices. + :type distinct_pairs: bool + :raises ValueError: If indices are invalid. + """ + for idx in indices: + if not (0 <= idx < d): + raise ValueError(f"Index {idx} must satisfy 0 <= index < d (d={d}).") + + if len(indices) == 2 and indices[0] == indices[1]: + raise ValueError("Rotation requires two distinct levels: j != k.") + + if distinct_pairs and len(indices) == 4: + j1, k1, j2, k2 = indices + if j1 == k1 and j2 == k2: + raise ValueError( + "Selected basis states must be different: (j1, j2) ≠ (k1, k2)." + ) + + +def _two_level_projectors( + d: int, j: int, k: Optional[int] = None +) -> Tuple[Tensor, ...]: + r""" + Construct projectors for single- or two-level subspaces in a ``d``-level qudit. + + :param d: Qudit dimension. + :type d: int + :param j: First level index. + :type j: int + :param k: Optional second level index. If None, only projectors for ``j`` are returned. + :type k: Optional[int] + :return: + - If ``k is None``: ``(I, Pjj)`` + - Else: ``(I, Pjj, Pkk, Pjk, Pkj)`` + :rtype: Tuple[Tensor, ...] + """ + I = backend.eye(d, dtype=dtypestr) + ej = I[:, j] + Pjj = backend.outer_product(ej, ej) + + if k is None: + return I, Pjj + + ek = I[:, k] + Pkk = backend.outer_product(ek, ek) + Pjk = backend.outer_product(ej, ek) + Pkj = backend.outer_product(ek, ej) + return I, Pjj, Pkk, Pjk, Pkj + + +def _rx_matrix_func(d: int, theta: float, j: int = 0, k: int = 1) -> Tensor: + r""" + Rotation-X (``RX``) gate on a selected two-level subspace of a qudit. + + Acts like the qubit :math:`RX(\theta)` on levels ``j`` and ``k``, identity elsewhere. + + :param d: Qudit dimension. + :type d: int + :param theta: Rotation angle :math:`\theta`. + :type theta: float + :param j: First level index. + :type j: int + :param k: Second level index. + :type k: int + :return: ``(d, d)`` matrix for :math:`RX(\theta)` on the ``j,k`` subspace. + :rtype: Tensor + """ + _check_rotation_indices(d, j, k) + I, Pjj, Pkk, Pjk, Pkj = _two_level_projectors(d, j, k) + theta = num_to_tensor(theta) + c = backend.cos(theta / 2.0) + s = backend.sin(theta / 2.0) + return I + (c - 1.0) * (Pjj + Pkk) + (-1j * s) * (Pjk + Pkj) + + +def _ry_matrix_func(d: int, theta: float, j: int = 0, k: int = 1) -> Tensor: + r""" + Rotation-Y (``RY``) gate on a selected two-level subspace of a qudit. + + :param d: Qudit dimension. + :type d: int + :param theta: Rotation angle :math:`\theta`. + :type theta: float + :param j: First level index. + :type j: int + :param k: Second level index. + :type k: int + :return: ``(d, d)`` matrix for :math:`RY(\theta)` on the ``j,k`` subspace. + :rtype: Tensor + """ + _check_rotation_indices(d, j, k) + I, Pjj, Pkk, Pjk, Pkj = _two_level_projectors(d, j, k) + theta = num_to_tensor(theta) + c = backend.cos(theta / 2.0) + s = backend.sin(theta / 2.0) + return I + (c - 1.0) * (Pjj + Pkk) - s * Pjk + s * Pkj + + +def _rz_matrix_func(d: int, theta: float, j: int = 0) -> Tensor: + r""" + Rotation-Z (``RZ``) gate for qudits. + + For qubits it reduces to the usual :math:`RZ(\theta)`. For general ``d``, it + applies a phase :math:`e^{i\theta}` to level ``j`` and leaves others unchanged. + + :param d: Qudit dimension. + :type d: int + :param theta: Rotation angle :math:`\theta`. + :type theta: float + :param j: Level index receiving the phase. + :type j: int + :return: ``(d, d)`` diagonal matrix implementing :math:`RZ(\theta)` on level ``j``. + :rtype: Tensor + """ + I, Pjj = _two_level_projectors(d, j, k=None) + theta = num_to_tensor(theta) + phase = backend.exp(1j * theta) + return I + (phase - 1.0) * Pjj + + +def _swap_matrix_func(d: int) -> Tensor: + """ + SWAP gate for two qudits of dimension ``d``. + + Exchanges basis states ``|i⟩|j⟩ → |j⟩|i⟩``. + + :param d: Qudit dimension (for each register). + :type d: int + :return: ``(d*d, d*d)`` matrix representing SWAP. + :rtype: Tensor + """ + D = d * d + I = np.eye(D, dtype=dtypestr) + m = I.reshape(d, d, d, d).transpose(1, 0, 2, 3).reshape(D, D) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _rzz_matrix_func( + d: int, theta: float, j1: int = 0, k1: int = 1, j2: int = 0, k2: int = 1 +) -> Tensor: + r""" + Two-qudit ``RZZ(\theta)`` on a selected two-state subspace. + + Acts like a qubit :math:`RZZ(\theta)=\exp(-i\,\tfrac{\theta}{2}\,\sigma_z)` on the + two-dimensional subspace spanned by ``|j1, j2⟩`` and ``|k1, k2⟩``, + and as identity elsewhere. The resulting block is diagonal with phases + :math:`\mathrm{diag}(e^{-i\theta/2},\, e^{+i\theta/2})`. + + :param d: Dimension of each qudit (assumed equal). + :type d: int + :param theta: Rotation angle. + :type theta: float + :param j1: Level on qudit-1 for the first basis state. + :type j1: int + :param k1: Level on qudit-1 for the second basis state. + :type k1: int + :param j2: Level on qudit-2 for the first basis state. + :type j2: int + :param k2: Level on qudit-2 for the second basis state. + :type k2: int + :return: ``(d*d, d*d)`` matrix representing subspace :math:`RZZ(\theta)`. + :rtype: Tensor + :raises ValueError: If indices are out of range or select the same basis state. + """ + _check_rotation_indices(d, j1, k1, j2, k2, distinct_pairs=True) + idx_a = j1 * d + j2 + idx_b = k1 * d + k2 + theta = num_to_tensor(theta) + phase_minus = backend.exp(-1j * theta / 2.0) + phase_plus = backend.exp(+1j * theta / 2.0) + + I = backend.eye(d * d, dtype=dtypestr) + ea = I[:, idx_a] + eb = I[:, idx_b] + Paa = backend.outer_product(ea, ea) + Pbb = backend.outer_product(eb, eb) + return I + (phase_minus - 1.0) * Paa + (phase_plus - 1.0) * Pbb + + +def _rxx_matrix_func( + d: int, theta: float, j1: int = 0, k1: int = 1, j2: int = 0, k2: int = 1 +) -> Tensor: + r""" + Two-qudit ``RXX(\theta)`` on a selected two-state subspace. + + Acts like a qubit :math:`RXX` on the subspace spanned by ``|j1, j2⟩`` and ``|k1, k2⟩``. + + :param d: Dimension of each qudit (assumed equal). + :type d: int + :param theta: Rotation angle. + :type theta: float + :param j1: Level on qudit-1. + :type j1: int + :param k1: Level on qudit-1. + :type k1: int + :param j2: Level on qudit-2. + :type j2: int + :param k2: Level on qudit-2. + :type k2: int + :return: ``(d*d, d*d)`` matrix representing :math:`RXX(\theta)` on the selected subspace. + :rtype: Tensor + """ + _check_rotation_indices(d, j1, k1, j2, k2, distinct_pairs=True) + idx_a = j1 * d + j2 + idx_b = k1 * d + k2 + theta = num_to_tensor(theta) + c = backend.cos(theta / 2.0) + s = backend.sin(theta / 2.0) + + I = backend.eye(d * d, dtype=dtypestr) + ea = I[:, idx_a] + eb = I[:, idx_b] + Paa = backend.outer_product(ea, ea) + Pbb = backend.outer_product(eb, eb) + Pab = backend.outer_product(ea, eb) + Pba = backend.outer_product(eb, ea) + return I + (c - 1.0) * (Paa + Pbb) + (-1j * s) * (Pab + Pba) + + +def _u8_matrix_func( + d: int, + gamma: int = 2, + z: int = 1, + eps: int = 0, + omega: Optional[complex] = None, +) -> Tensor: + r""" + ``U8`` diagonal single-qudit gate for prime dimensions. + + See ref: Howard, Mark, and Jiri Vala. + "Qudit versions of the qubit π/8 gate." Physical Review A 86, no. 2 (2012): 022316. + https://doi.org/10.1103/PhysRevA.86.022316 + + This gate is the qudit analogue of the qubit :math:`\pi/8` gate, defined in + *Howard & Campbell, Phys. Rev. A 86, 022316 (2012)*. It is diagonal in the + computational basis with exponents determined by modular polynomials in the + parameters :math:`\gamma, z, \epsilon`. These gates, together with Pauli and + Clifford operations, generate the full single-qudit Clifford hierarchy. + + - For :math:`d=2`, this reduces (up to global phase) to the standard qubit + :math:`\pi/8` gate. + - For :math:`d=3`, the exponents live in :math:`\mathbb{Z}_9` and the + primitive ninth root :math:`\zeta = e^{2\pi i/9}` is used. + - For prime :math:`d>3`, the construction uses the modular inverse of 12 in + :math:`\mathbb{Z}_d`. + + :param d: Qudit dimension (must be prime). + :type d: int + :param gamma: Shear parameter :math:`\gamma' \in \mathbb{Z}_d`. + If ``gamma = 0``, the gate is a diagonal Clifford. + If ``gamma ≠ 0``, the gate is a genuine non-Clifford (analogue of :math:`\pi/8`). + :type gamma: int + :param z: Displacement parameter :math:`z' \in \mathbb{Z}_d`, + which sets the symplectic part of the associated Clifford. + :type z: int + :param eps: Phase offset parameter :math:`\epsilon' \in \mathbb{Z}_d`. + It only contributes a global phase factor :math:`\omega^{\epsilon'}`. + :type eps: int + :param omega: Optional primitive :math:`d`-th root of unity (complex). + Defaults to :math:`e^{2\pi i/d}` for d>3, and :math:`e^{2\pi i/9}` for d=3. + :type omega: Optional[complex] + :return: ``(d, d)`` diagonal matrix of dtype ``npdtype``. + :rtype: Tensor + :raises ValueError: If ``d`` is not prime; if 12 has no modular inverse + mod ``d`` (for ``d>3``); or if the computed exponents do not sum to + 0 mod ``d`` (or 0 mod 3 for ``d=3``). + """ + if not _is_prime(d): + raise ValueError( + f"Dimension d={d} is not prime, U8 gate requires a prime dimension." + ) + + omega = np.exp(2j * np.pi / d) if omega is None else omega + + gamma = int(gamma) % d + z = int(z) % d + eps = int(eps) % d + + if d == 3: + vks = [0, (6 * z + 2 * gamma + 3 * eps) % 9, (6 * z + 1 * gamma + 6 * eps) % 9] + if sum(vks) % 3 != 0: + raise ValueError( + f"Sum of v_k's is not 0 mod 3. Got {sum(vks) % 3}. Check parameters." + ) + + zeta = np.exp(2j * np.pi / 9) + m = np.diag([zeta**v for v in vks]) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + try: + inv_12 = pow(12, -1, d) + except ValueError: + raise ValueError( + f"Inverse of 12 mod {d} does not exist. Choose a prime d that does not divide 12." + ) + + vks = [0] * d + for k in range(1, d): + term_inner = ((6 * z) % d + ((2 * k - 3) % d) * gamma) % d + term = (gamma + (k * term_inner) % d) % d + vk = ((inv_12 * (k % d)) % d) * term % d + vk = (vk + (eps * (k % d)) % d) % d + vks[k] = vk + + if sum(vks) % d != 0: + raise ValueError( + f"Sum of v_k's is not 0 mod {d}. Got {sum(vks) % d}. Check parameters." + ) + + omega = np.exp(2j * np.pi / d) if omega is None else omega + m = np.diag([omega**v for v in vks]) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _cphase_matrix_func( + d: int, cv: Optional[int] = None, omega: Optional[complex] = None +) -> Tensor: + r""" + Qudit controlled-phase (``CPHASE``) gate. + Implements ``|r⟩|s⟩ → ω^{rs}|r⟩|s⟩``; optionally condition on a specific control value ``cv``. + ┌─ ─┐ + │ I_d 0 0 ... 0 │ + │ 0 Z_d 0 ... 0 │ + SUMZ_d = │ 0 0 Z_d^2 ... 0 │ + │ . . . . . │ + │ 0 0 0 ... Z_d^{d-1} │ + └ ─┘ + + :param d: Qudit dimension (for each register). + :type d: int + :param cv: Optional control value in ``[0, d-1]``. If ``None``, builds the full SUMZ block-diagonal. + :type cv: Optional[int] + :param omega: Optional primitive ``d``-th root of unity for ``Z_d``. + :type omega: Optional[complex] + :return: ``(d*d, d*d)`` matrix representing the controlled-phase. + :rtype: Tensor + :raises ValueError: If ``cv`` is provided and is outside ``[0, d-1]``. + """ + omega = np.exp(2j * np.pi / d) if omega is None else omega + r = np.arange(d).reshape(-1, 1) + s = np.arange(d).reshape(1, -1) + + if cv is None: + phase = omega ** (r * s) + else: + if not (0 <= cv < d): + raise ValueError(f"cv must be in [0, {d - 1}], got {cv}") + phase = 1 + (r == cv) * (omega**s - 1) + + diag = np.ravel(phase) + m = np.diag(diag) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + +def _csum_matrix_func(d: int, cv: Optional[int] = None) -> Tensor: + r""" + Qudit controlled-sum (``CSUM`` / ``SUMX``) gate. + Implements ``|r⟩|s⟩ → |r⟩|r+s (\bmod d)⟩``; optionally condition on a specific control value ``cv``. + ┌─ ─┐ + │ I_d 0 0 ... 0 │ + │ 0 X_d 0 ... 0 │ + SUMX_d = │ 0 0 X_d^2 ... 0 │ + │ . . . . . │ + │ 0 0 0 ... X_d^{d-1} │ + └ ─┘ + + :param d: Qudit dimension (for each register). + :type d: int + :param cv: Optional control value in ``[0, d-1]``. If ``None``, builds the full SUMX block-diagonal. + :type cv: Optional[int] + :return: ``(d*d, d*d)`` matrix representing the controlled-sum. + :rtype: Tensor + :raises ValueError: If ``cv`` is provided and is outside ``[0, d-1]``. + """ + I = np.eye(d) + + if cv is None: + blocks = [np.roll(I, shift=r, axis=0) for r in range(d)] + m = np.block( + [ + [blocks[r] if r == c else np.zeros((d, d)) for c in range(d)] + for r in range(d) + ] + ) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) + + if not (0 <= cv < d): + raise ValueError(f"cv must be in [0, {d - 1}], got {cv}") + + X = np.roll(I, shift=1, axis=0) + m = np.kron(I, I) + np.kron(np.outer(I[:, cv], I[:, cv]), (X - I)) + return backend.cast(backend.convert_to_tensor(m), dtype=dtypestr) diff --git a/tests/test_quditcircuit.py b/tests/test_quditcircuit.py new file mode 100644 index 00000000..054389d3 --- /dev/null +++ b/tests/test_quditcircuit.py @@ -0,0 +1,447 @@ +# pylint: disable=invalid-name + +import os +import sys + +import numpy as np +import pytest +from pytest_lazyfixture import lazy_fixture as lf + +# see https://stackoverflow.com/questions/56307329/how-can-i-parametrize-tests-to-run-with-different-fixtures-in-pytest + +thisfile = os.path.abspath(__file__) +modulepath = os.path.dirname(os.path.dirname(thisfile)) + +sys.path.insert(0, modulepath) +import tensorcircuit as tc + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")]) +def test_basics(backend): + c = tc.QuditCircuit(2, 3) + c.x(0) + np.testing.assert_allclose(tc.backend.numpy(c.amplitude("10")), np.array(1.0)) + c.csum(0, 1) + np.testing.assert_allclose(tc.backend.numpy(c.amplitude("11")), np.array(1.0)) + c.csum(0, 1) + np.testing.assert_allclose(tc.backend.numpy(c.amplitude("12")), np.array(1.0)) + + c = tc.QuditCircuit(2, 3) + c.x(0) + c.x(0) + np.testing.assert_allclose(tc.backend.numpy(c.amplitude("20")), np.array(1.0)) + c.csum(0, 1, cv=1) + np.testing.assert_allclose(tc.backend.numpy(c.amplitude("21")), np.array(0.0)) + c.csum(0, 1, cv=2) + np.testing.assert_allclose(tc.backend.numpy(c.amplitude("21")), np.array(1.0)) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")]) +def test_measure(backend): + c = tc.QuditCircuit(2, 3) + c.h(0) + c.x(1) + c.csum(0, 1) + assert c.measure(1)[0] in [0, 1, 2] + + +@pytest.mark.parametrize("backend", [lf("jaxb")]) +def test_large_scale_sample(backend): + L = 30 + d = 3 + c = tc.QuditCircuit(L, d) + c.h(0) + for i in range(L - 1): + c.csum(i, i + 1) + + batch = 1024 + results = c.sample( + allow_state=False, batch=batch, format="count_dict_bin", jittable=False + ) + + k0, k1, k2 = "0" * L, "1" * L, "2" * L + c0, c1, c2 = results.get(k0, 0), results.get(k1, 0), results.get(k2, 0) + assert c0 + c1 + c2 == batch + + probs = np.array([c0, c1, c2], dtype=float) / batch + np.testing.assert_allclose(probs, np.ones(3) / 3, rtol=0.2, atol=0.0) + + for a, b in [(c0, c1), (c1, c2), (c0, c2)]: + ratio = (a + 1e-12) / (b + 1e-12) + assert 0.8 <= ratio <= 1.25 + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")]) +def test_expectation(backend): + c = tc.QuditCircuit(2, 3) + c.h(0) + np.testing.assert_allclose( + tc.backend.numpy(c.expectation((tc.quditgates._z_matrix_func(3), [0]))), + 0, + atol=1e-7, + ) + + +def test_complex128(highp, tfb): + c = tc.QuditCircuit(2, 3) + c.h(1) + c.rx(0, theta=1j) + c.wavefunction() + np.testing.assert_allclose( + c.expectation((tc.quditgates._z_matrix_func(3), [1])), 0, atol=1e-15 + ) + + +def test_single_qubit(): + c = tc.QuditCircuit(1, 10) + c.h(0) + w = c.state()[0] + np.testing.assert_allclose( + w, + np.array( + [ + 1, + ] + * 10 + ) + / np.sqrt(10), + atol=1e-5, + ) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")]) +def test_expectation_between_two_states_qudit(backend): + dim = 3 + X3 = tc.quditgates._x_matrix_func(dim) + # Y3 = tc.quditgates._y_matrix_func(dim) # ZX/i + Z3 = tc.quditgates._z_matrix_func(dim) + H3 = tc.quditgates._h_matrix_func(dim) + X3_dag = np.conjugate(X3.T) + + # e0 = np.array([1.0, 0.0, 0.0], dtype=np.complex64) + # e1 = np.array([0.0, 1.0, 0.0], dtype=np.complex64) + # val = tc.expectation((tc.gates.Gate(Y3), [0]), ket=e0, bra=e1, dim=dim) + # omega = np.exp(2j * np.pi / dim) + # expected = omega / 1j + # np.testing.assert_allclose(tc.backend.numpy(val), expected, rtol=1e-6, atol=1e-6) + + c = tc.QuditCircuit(3, dim) + c.unitary(0, unitary=tc.gates.Gate(H3)) + c.ry(1, theta=0.8, j=0, k=1) + state = c.wavefunction() + x1z2 = [(tc.gates.Gate(X3), [0]), (tc.gates.Gate(Z3), [1])] + e1 = c.expectation(*x1z2) + e2 = tc.expectation(*x1z2, ket=state, bra=state, normalization=True, dim=dim) + np.testing.assert_allclose(tc.backend.numpy(e2), tc.backend.numpy(e1)) + + c = tc.QuditCircuit(3, dim) + c.unitary(0, unitary=tc.gates.Gate(H3)) + c.ry(1, theta=0.8 + 0.7j, j=0, k=1) + state = c.wavefunction() + e1 = c.expectation(*x1z2) / (tc.backend.norm(state) ** 2) + e2 = tc.expectation(*x1z2, ket=state, normalization=True, dim=dim) + np.testing.assert_allclose(tc.backend.numpy(e2), tc.backend.numpy(e1)) + + c1 = tc.QuditCircuit(2, dim) + c1.unitary(1, unitary=tc.gates.Gate(X3)) + s1 = c1.state() + + c2 = tc.QuditCircuit(2, dim) + c2.unitary(0, unitary=tc.gates.Gate(X3)) + s2 = c2.state() + + c3 = tc.QuditCircuit(2, dim) + c3.unitary(1, unitary=tc.gates.Gate(H3)) + s3 = c3.state() + + x1x2_fixed = [(tc.gates.Gate(X3), [0]), (tc.gates.Gate(X3_dag), [1])] + e = tc.expectation(*x1x2_fixed, ket=s1, bra=s2, dim=dim) + np.testing.assert_allclose(tc.backend.numpy(e), 1.0) + + e2 = tc.expectation(*x1x2_fixed, ket=s3, bra=s2, dim=dim) + np.testing.assert_allclose(tc.backend.numpy(e2), 1.0 / np.sqrt(3)) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("cpb")]) +def test_any_inputs_state_qudit_true_gates(backend): + dim = 3 + Xd = tc.quditgates._x_matrix_func(dim) + Zd = tc.quditgates._z_matrix_func(dim) + omega = np.exp(2j * np.pi / dim) + + def idx(j0, j1): + return dim * j0 + j1 + + vec = np.zeros(dim * dim, dtype=np.complex64) + vec[idx(2, 0)] = 1.0 + c = tc.QuditCircuit(2, dim, inputs=tc.array_to_tensor(vec)) + c.unitary(0, unitary=tc.gates.Gate(Xd)) + z0 = c.expectation((tc.gates.Gate(Zd), [0])) + np.testing.assert_allclose(tc.backend.numpy(z0), 1.0 + 0j, rtol=1e-6, atol=1e-6) + + vec = np.zeros(dim * dim, dtype=np.complex64) + vec[idx(0, 0)] = 1.0 + c = tc.QuditCircuit(2, dim, inputs=tc.array_to_tensor(vec)) + c.unitary(0, unitary=tc.gates.Gate(Xd)) + z0 = c.expectation((tc.gates.Gate(Zd), [0])) + np.testing.assert_allclose(tc.backend.numpy(z0), omega, rtol=1e-6, atol=1e-6) + + vec = np.zeros(dim * dim, dtype=np.complex64) + vec[idx(1, 0)] = 1.0 + c = tc.QuditCircuit(2, dim, inputs=tc.array_to_tensor(vec)) + c.unitary(0, unitary=tc.gates.Gate(Xd)) + z0 = c.expectation((tc.gates.Gate(Zd), [0])) + np.testing.assert_allclose(tc.backend.numpy(z0), omega**2, rtol=1e-6, atol=1e-6) + + vec = np.zeros(dim * dim, dtype=np.complex64) + vec[idx(0, 0)] = 1 / np.sqrt(2) + vec[idx(1, 0)] = 1 / np.sqrt(2) + c = tc.QuditCircuit(2, dim, inputs=tc.array_to_tensor(vec)) + c.unitary(0, unitary=tc.gates.Gate(Xd)) + z0 = c.expectation((tc.gates.Gate(Zd), [0])) + np.testing.assert_allclose(tc.backend.numpy(z0), -0.5 + 0j, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("cpb")]) +def test_postselection(backend): + c = tc.QuditCircuit(3, 3) + c.h(1) + c.h(2) + c.mid_measurement(1, 1) + c.mid_measurement(2, 1) + s = c.wavefunction() + np.testing.assert_allclose(tc.backend.numpy(s[4]).real, 1.0 / 3.0, rtol=1e-6) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")]) +def test_unitary(backend): + c = tc.QuditCircuit(2, dim=3, inputs=np.eye(9)) + c.x(0) + c.z(1) + answer = tc.backend.numpy( + np.kron(tc.quditgates._x_matrix_func(3), tc.quditgates._z_matrix_func(3)) + ) + np.testing.assert_allclose( + tc.backend.numpy(c.wavefunction().reshape([9, 9])), answer, atol=1e-5 + ) + + +def test_probability(): + c = tc.QuditCircuit(2, 3) + c.h(0) + c.h(1) + np.testing.assert_allclose(c.probability(), np.ones(9) / 9, atol=1e-5) + + +def test_circuit_add_demo(): + dim = 3 + c = tc.QuditCircuit(2, dim=dim) + c.x(0) # |00> -> |10> + c2 = tc.QuditCircuit(2, dim=dim, mps_inputs=c.quvector()) + c2.x(0) # |00> -> |20> + answer = np.zeros(dim * dim, dtype=np.complex64) + answer[dim * 2 + 0] = 1.0 + np.testing.assert_allclose(c2.wavefunction(), answer, atol=1e-5) + c3 = tc.QuditCircuit(2, dim=dim) + c3.x(0) + c3.replace_mps_inputs(c.quvector()) + np.testing.assert_allclose(c3.wavefunction(), answer, atol=1e-5) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_circuit_matrix(backend): + c = tc.QuditCircuit(2, 3) + c.x(1) + c.csum(0, 1) + + U = c.matrix() + U_np = tc.backend.numpy(U) + + row_10 = U_np[3] + expected_row_10 = np.zeros(9, dtype=row_10.dtype) + expected_row_10[4] = 1.0 + np.testing.assert_allclose(row_10, expected_row_10, atol=1e-5) + + state = tc.backend.numpy(c.state()) + expected_state = np.zeros(9, dtype=state.dtype) + expected_state[1] = 1.0 + np.testing.assert_allclose(state, expected_state, atol=1e-5) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_batch_sample(backend): + c = tc.QuditCircuit(3, 3) + c.h(0) + c.csum(0, 1) + print(c.sample()) + print(c.sample(batch=8, status=np.random.uniform(size=[8, 3]))) + print(c.sample(batch=8)) + print(c.sample(random_generator=tc.backend.get_random_state(42))) + print(c.sample(allow_state=True)) + print(c.sample(batch=8, allow_state=True)) + print( + c.sample( + batch=8, allow_state=True, random_generator=tc.backend.get_random_state(42) + ) + ) + print( + c.sample( + batch=8, + allow_state=True, + status=np.random.uniform(size=[8]), + format="sample_bin", + ) + ) + print( + c.sample( + batch=8, + allow_state=False, + status=np.random.uniform(size=[8, 3]), + format="sample_bin", + ) + ) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_sample_format(backend): + c = tc.QuditCircuit(2, 3) + c.h(0) + c.csum(0, 1) + key = tc.backend.get_random_state(42) + for allow_state in [False, True]: + print("allow_state: ", allow_state) + for batch in [None, 1, 3]: + print(" batch: ", batch) + for format_ in [ + None, + "sample_bin", + "count_dict_bin", + ]: + print(" format: ", format_) + print( + " ", + c.sample( + batch=batch, + allow_state=allow_state, + format_=format_, + random_generator=key, + ), + ) + + +def test_sample_representation(): + _ALPHBET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" + c = tc.QuditCircuit(1, 36) + (result,) = c.sample(1, format="count_dict_bin").keys() + assert result == _ALPHBET[0] + + for i in range(1, 35): + c.x(0) + (result,) = c.sample(1, format="count_dict_bin").keys() + assert result == _ALPHBET[i] + + +def test_quditcircuit_set_dim_validation(): + with pytest.raises(ValueError): + tc.QuditCircuit(1, 2) + with pytest.raises(ValueError): + tc.QuditCircuit(1, 2.5) # type: ignore[arg-type] + + +@pytest.mark.parametrize("backend", [lf("jaxb"), lf("tfb"), lf("torchb")]) +def test_qudit_minimal_ad_qudit(backend): + """Minimal AD test on a single-qudit (d=3) circuit. + We differentiate the expectation ⟨Z⟩ w.r.t. a single RY parameter and + compare to a finite-difference estimate. + """ + + dim = 3 + + def energy(theta): + c = tc.QuditCircuit(1, dim) + # rotate on the (0,1) subspace so that the observable is sensitive to theta + c.ry(0, theta=theta, j=0, k=1) + # measure Z on site 0 (qudit Z for d=3) + E = c.expectation((tc.quditgates._z_matrix_func(dim), [0])) + return tc.backend.real(E) + + # backend autodiff gradient + grad_energy = tc.backend.grad(energy) + + theta0 = tc.num_to_tensor(0.37) + g = grad_energy(theta0) + + # finite-difference check + eps = 1e-3 + num = (energy(theta0 + eps) - energy(theta0 - eps)) / (2 * eps) + + np.testing.assert_allclose( + tc.backend.numpy(g), tc.backend.numpy(num), rtol=1e-2, atol=1e-3 + ) + + +@pytest.mark.parametrize("backend", [lf("jaxb"), lf("tfb"), lf("torchb")]) +def test_qudit_minimal_jit_qudit(backend): + """Minimal JIT test: jit-compiled energy matches eager energy.""" + + dim = 3 + + def energy(theta): + c = tc.QuditCircuit(1, dim) + c.ry(0, theta=theta, j=0, k=1) + E = c.expectation((tc.quditgates._z_matrix_func(dim), [0])) + return tc.backend.real(E) + + jit_energy = tc.backend.jit(energy) + + theta0 = tc.num_to_tensor(-0.91) + e_eager = energy(theta0) + e_jit = jit_energy(theta0) + + np.testing.assert_allclose( + tc.backend.numpy(e_jit), tc.backend.numpy(e_eager), rtol=1e-6, atol=1e-6 + ) + + +@pytest.mark.parametrize("backend", [lf("jaxb"), lf("tfb"), lf("torchb")]) +def test_qudit_minimal_vmap_qudit(backend): + """Minimal VMAP test: vectorized energies equal per-element eager results.""" + + dim = 3 + + def energy(theta): + c = tc.QuditCircuit(1, dim) + c.ry(0, theta=theta, j=0, k=1) + E = c.expectation((tc.quditgates._z_matrix_func(dim), [0])) + return tc.backend.real(E) + + venergy = tc.backend.vmap(energy) + + thetas = tc.array_to_tensor(np.linspace(-1.0, 1.0, 7)) + vvals = venergy(thetas) + eager_vals = np.array([energy(t) for t in thetas]) + + np.testing.assert_allclose( + tc.backend.numpy(vvals), eager_vals, rtol=1e-6, atol=1e-6 + ) + + +def test_qudit_paths_and_sampling_wrappers(): + c = tc.QuditCircuit(2, 3) + c.x(0) + c.rzz(0, 1, theta=np.float64(0.2), j1=0, k1=1, j2=0, k2=1) + c.cphase(0, 1, cv=1) + qo = c.get_quoperator() + assert qo is not None + _ = c.sample(allow_state=False, batch=1, format="count_dict_bin") + for bad in ["sample_int", "count_tuple", "count_dict_int", "count_vector"]: + with pytest.raises(NotImplementedError): + c.sample(allow_state=False, batch=1, format=bad) + + +def test_quditcircuit_amplitude_before_wrapper(): + c = tc.QuditCircuit(2, 3) + c.x(0) + nodes = c.amplitude_before("00") + assert isinstance(nodes, list) + assert len(nodes) == 5 # one gate (X on qudit 0) -> single node in the traced path diff --git a/tests/test_quditgates.py b/tests/test_quditgates.py new file mode 100644 index 00000000..a0612cdb --- /dev/null +++ b/tests/test_quditgates.py @@ -0,0 +1,371 @@ +import sys +import os +import numpy as np + +import pytest +from pytest_lazyfixture import lazy_fixture as lf + +thisfile = os.path.abspath(__file__) +modulepath = os.path.dirname(os.path.dirname(thisfile)) + +sys.path.insert(0, modulepath) + +import tensorcircuit as tc + +from tensorcircuit.quditgates import ( + _i_matrix_func, + _x_matrix_func, + _z_matrix_func, + _h_matrix_func, + _s_matrix_func, + _rx_matrix_func, + _ry_matrix_func, + _rz_matrix_func, + _swap_matrix_func, + _rzz_matrix_func, + _rxx_matrix_func, + _u8_matrix_func, + _cphase_matrix_func, + _csum_matrix_func, + _is_prime, +) + + +def is_unitary(M): + M = tc.backend.numpy(M) + Mc = M.astype(np.complex128, copy=False) + I = np.eye(M.shape[0], dtype=np.complex128) + return np.allclose(Mc.conj().T @ Mc, I, atol=1e-5, rtol=1e-5) and np.allclose( + Mc @ Mc.conj().T, I, atol=1e-5, rtol=1e-5 + ) + + +@pytest.mark.parametrize("d", [2, 3, 4, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_I_X_Z_shapes_and_unitarity(d, backend, highp): + I = _i_matrix_func(d) + X = _x_matrix_func(d) + Z = _z_matrix_func(d) + assert I.shape == (d, d) and X.shape == (d, d) and Z.shape == (d, d) + assert is_unitary(X) + assert is_unitary(Z) + np.testing.assert_allclose(tc.backend.numpy(I), np.eye(d), atol=1e-5) + + +@pytest.mark.parametrize("d", [2, 3, 4]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_X_is_right_cyclic_shift(d, backend, highp): + X = _x_matrix_func(d) + X = tc.backend.numpy(X) + for j in range(d): + v = np.zeros(d) + v[j] = 1 + out = X @ v + expected = np.zeros(d) + expected[(j + 1) % d] = 1 + np.testing.assert_allclose(out, expected, atol=1e-5) + + +@pytest.mark.parametrize("d", [2, 3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_Z_diagonal_and_value(d, backend, highp): + omega = np.exp(2j * np.pi / d) + Z = _z_matrix_func(d, omega) + np.testing.assert_allclose( + tc.backend.numpy(Z), np.diag([omega**j for j in range(d)]), atol=1e-5 + ) + assert is_unitary(Z) + + +# @pytest.mark.parametrize("d", [2, 3, 5]) +# def test_Y_equals_ZX_over_i(d, highp): +# Y = _y_matrix_func(d) +# ZX_over_i = (_z_matrix_func(d) @ _x_matrix_func(d)) / 1j +# np.testing.assert_allclose(Y, ZX_over_i, atol=1e-5) +# assert is_unitary(Y) + + +@pytest.mark.parametrize("d", [2, 3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_H_is_fourier_like_and_unitary(d, backend, highp): + H = _h_matrix_func(d) + assert H.shape == (d, d) + assert is_unitary(H) + omega = np.exp(2j * np.pi / d) + F = (1 / np.sqrt(d)) * np.array( + [[omega ** (j * k) for k in range(d)] for j in range(d)] + ).T + np.testing.assert_allclose(tc.backend.numpy(H), F, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("d", [2, 3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_S_is_diagonal(d, backend, highp): + S = _s_matrix_func(d) + np.testing.assert_allclose(S, np.diag(np.diag(S)), atol=1e-5) + + +@pytest.mark.parametrize("d", [3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_RX_RY_only_affect_subspace(d, backend, highp): + theta = 0.7 + j, k = 0, 1 + RX = _rx_matrix_func(d, theta, j, k) + RY = _ry_matrix_func(d, theta, j, k) + assert is_unitary(RX) and is_unitary(RY) + RX, RY = tc.backend.numpy(RX), tc.backend.numpy(RY) + for t in range(d): + if t not in (j, k): + e = np.zeros(d) + e[t] = 1 + outx = RX @ e + outy = RY @ e + np.testing.assert_allclose(outx, e, atol=1e-5) + np.testing.assert_allclose(outy, e, atol=1e-5) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_RZ_phase_on_single_level(backend, highp): + d, theta, j = 5, 1.234, 2 + RZ = _rz_matrix_func(d, theta, j) + assert is_unitary(RZ) + diag = np.ones(d, dtype=np.complex64) + diag[j] = np.exp(1j * theta) + assert np.allclose(RZ, np.diag(diag), atol=1e-5) + + +@pytest.mark.parametrize("d", [2, 3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_SWAP_permutation(d, backend, highp): + SW = _swap_matrix_func(d) + D = d * d + assert SW.shape == (D, D) + assert is_unitary(SW) + SW = tc.backend.numpy(SW) + for i in range(min(d, 3)): + for j in range(min(d, 3)): + v = np.zeros(D) + v[i * d + j] = 1 + out = SW @ v + exp = np.zeros(D) + exp[j * d + i] = 1 + np.testing.assert_allclose(out, exp, atol=1e-5) + + +@pytest.mark.parametrize("d", [2, 3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_RZZ_diagonal(d, backend, highp): + theta = 0.37 + RZZ = _rzz_matrix_func(d, theta, j1=0, k1=1, j2=0, k2=1) + assert is_unitary(RZZ) + + D = d * d + I = np.eye(D, dtype=np.complex128) + + idx_a = 0 * d + 0 + idx_b = 1 * d + 1 + + for t in range(D): + if t not in (idx_a, idx_b): + np.testing.assert_allclose(RZZ[t], I[t], atol=1e-5) + + np.testing.assert_allclose(RZZ[idx_a, idx_a], np.exp(-1j * theta / 2), atol=1e-5) + np.testing.assert_allclose(RZZ[idx_b, idx_b], np.exp(+1j * theta / 2), atol=1e-5) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_RXX_selected_block(backend, highp): + d = 4 + theta = 0.81 + j1, k1 = 0, 2 + j2, k2 = 1, 3 + RXX = _rxx_matrix_func(d, theta, j1, k1, j2, k2) + assert is_unitary(RXX) + D = d * d + I = np.eye(D) + idx_a = j1 * d + j2 + idx_b = k1 * d + k2 + for t in range(D): + for s in range(D): + if {t, s} & {idx_a, idx_b}: + continue + np.testing.assert_allclose(RXX[t, s], I[t, s], atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("d", [3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_CPHASE_blocks(d, backend, highp): + omega = np.exp(2j * np.pi / d) + Z = _z_matrix_func(d, omega) + M = _cphase_matrix_func(d, cv=None, omega=omega) + for a in range(d): + rs = a * d + block = M[rs : rs + d, rs : rs + d] + Za = np.linalg.matrix_power(Z, a) + np.testing.assert_allclose(block, Za, atol=1e-5) + assert is_unitary(M) + + cv = 1 + M2 = _cphase_matrix_func(d, cv=cv, omega=omega) + for a in range(d): + rs = a * d + block = M2[rs : rs + d, rs : rs + d] + if a == cv: + np.testing.assert_allclose(block, Z, atol=1e-5) + else: + np.testing.assert_allclose(block, np.eye(d), atol=1e-5) + + +@pytest.mark.parametrize("d", [3, 5]) +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_CSUM_blocks(d, backend, highp): + X = _x_matrix_func(d) + M = _csum_matrix_func(d, cv=None) + for a in range(d): + rs = a * d + block = M[rs : rs + d, rs : rs + d] + Xa = np.linalg.matrix_power(X, a) + np.testing.assert_allclose(block, Xa, atol=1e-5) + assert is_unitary(M) + + cv = 2 % d + M2 = _csum_matrix_func(d, cv=cv) + for a in range(d): + rs = a * d + block = M2[rs : rs + d, rs : rs + d] + if a == cv: + np.testing.assert_allclose(block, X, atol=1e-5) + else: + np.testing.assert_allclose(block, np.eye(d), atol=1e-5) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_CSUM_mapping_small_d(backend, highp): + d = 3 + M = _csum_matrix_func(d) + M = tc.backend.numpy(M) + for r in range(d): + for s in range(d): + v = np.zeros(d * d) + v[r * d + s] = 1 + out = M @ v + exp = np.zeros(d * d) + exp[r * d + ((r + s) % d)] = 1 + np.testing.assert_allclose(out, exp, atol=1e-5) + + +def test_rotation_index_errors(highp): + d = 4 + with pytest.raises(ValueError): + _rx_matrix_func(d, 0.1, j=-1, k=1) + with pytest.raises(ValueError): + _ry_matrix_func(d, 0.1, j=0, k=4) + with pytest.raises(ValueError): + _rx_matrix_func(d, 0.1, j=2, k=2) + + +def test_CPHASE_CSUM_cv_range(highp): + d = 5 + with pytest.raises(ValueError): + _cphase_matrix_func(d, cv=-1) + with pytest.raises(ValueError): + _cphase_matrix_func(d, cv=d) + with pytest.raises(ValueError): + _csum_matrix_func(d, cv=-1) + with pytest.raises(ValueError): + _csum_matrix_func(d, cv=d) + + +def test__is_prime_edge_and_composites(): + assert _is_prime(2) is True + assert _is_prime(3) is True + assert _is_prime(4) is False + assert _is_prime(5) is True + assert _is_prime(25) is False + assert _is_prime(29) is True + + +def test_two_qudit_builders_index_validation(): + d = 3 + theta = 0.1 + with pytest.raises(ValueError): + _rzz_matrix_func(d, theta, j1=0, k1=1, j2=0, k2=3) + with pytest.raises(ValueError): + _rxx_matrix_func(d, theta, j1=0, k1=1, j2=3, k2=0) + with pytest.raises(ValueError): + _rzz_matrix_func(d, theta, j1=0, k1=0, j2=1, k2=1) + with pytest.raises(ValueError): + _rxx_matrix_func(d, theta, j1=2, k1=2, j2=0, k2=0) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_u8_prime_dimension_and_qubit_case(backend): + with pytest.raises(ValueError): + _u8_matrix_func(d=4) + with pytest.raises(ValueError): + _u8_matrix_func(d=9, gamma=1, z=0, eps=0) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_u8_qutrit_correct_phases_and_gamma_zero_allowed(backend): + d = 3 + U3 = _u8_matrix_func(d=d, gamma=2, z=1, eps=0) + zeta = np.exp(2j * np.pi / 9) + expected3 = np.diag([zeta**0, zeta**1, zeta**8]) + assert np.allclose(tc.backend.numpy(U3), expected3, atol=1e-12) + + U3_g0 = _u8_matrix_func(d=d, gamma=0, z=1, eps=2) + U3_g0 = tc.backend.numpy(U3_g0) + assert U3_g0.shape == (3, 3) + assert np.allclose(U3_g0, np.diag(np.diag(U3_g0)), atol=1e-12) + assert np.allclose(U3_g0.conj().T @ U3_g0, np.eye(3), atol=1e-12) + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_u8_p_greater_than_3_matches_closed_form(backend): + d = 5 + gamma, z, eps = 2, 1, 3 + + inv_12 = pow(12, -1, d) # 12 \equiv 2 (mod 5), inverse = 3 + vks = [0] * d + for k in range(1, d): + k_ = k % d + term_inner = ((6 * z) % d + ((2 * k_ - 3) % d) * gamma % d) % d + term = (gamma + (k_ * term_inner) % d) % d + vk = ((inv_12 * k_) % d) * term % d + vk = (vk + (eps * k_) % d) % d + vks[k] = vk + + omega = np.exp(2j * np.pi / d) + expected5 = np.diag([omega**v for v in vks]) + + U5 = _u8_matrix_func(d=d, gamma=gamma, z=z, eps=eps) + assert np.allclose(tc.backend.numpy(U5), expected5, atol=1e-12) + assert sum(vks) % d == 0 + + +@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")]) +def test_u8_parameter_normalization_and_custom_omega(backend): + d = 5 + U_modded = _u8_matrix_func(d=d, gamma=2, z=1, eps=3) + U_unnormalized = _u8_matrix_func( + d=d, gamma=7, z=-4, eps=13 + ) # 7\equiv 2, -4\equiv 1, 13\equiv 3 (mod 5) + assert np.allclose(U_modded, U_unnormalized, atol=1e-12) + + d = 7 + gamma, z, eps = 3, 2, 1 + inv_12 = pow(12, -1, d) + vks = [0] * d + for k in range(1, d): + k_ = k % d + term_inner = ((6 * z) % d + ((2 * k_ - 3) % d) * gamma % d) % d + term = (gamma + (k_ * term_inner) % d) % d + vk = ((inv_12 * k_) % d) * term % d + vk = (vk + (eps * k_) % d) % d + vks[k] = vk + + omega_custom = np.exp(2j * np.pi / d) * np.exp(0j) + U7_custom = _u8_matrix_func(d=d, gamma=gamma, z=z, eps=eps, omega=omega_custom) + expected7_custom = np.diag([omega_custom**v for v in vks]) + assert np.allclose(tc.backend.numpy(U7_custom), expected7_custom, atol=1e-12)