qadence 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

qadence/operations.py

@@ -31,6 +31,7 @@ from qadence.blocks.analog import (
     WaitBlock,
 )
 from qadence.blocks.block_to_tensor import block_to_tensor
+from qadence.blocks.primitive import ProjectorBlock
 from qadence.blocks.utils import (
     add,  # noqa
     block_is_commuting_hamiltonian,
@@ -117,9 +118,6 @@ class X(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> X:
-        return self
-
 
 class Y(PrimitiveBlock):
     """The Y gate."""
@@ -141,9 +139,6 @@ class Y(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> Y:
-        return self
-
 
 class Z(PrimitiveBlock):
     """The Z gate."""
@@ -165,17 +160,36 @@ class Z(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> Z:
-        return self
 
+class Projector(ProjectorBlock):
+    """The projector operator."""
+
+    name = OpName.PROJ
+
+    def __init__(
+        self,
+        ket: str,
+        bra: str,
+        qubit_support: int | tuple[int, ...],
+    ):
+        super().__init__(ket=ket, bra=bra, qubit_support=qubit_support)
+
+    @property
+    def generator(self) -> None:
+        raise ValueError("Property `generator` not available for non-unitary operator.")
+
+    @property
+    def eigenvalues_generator(self) -> None:
+        raise ValueError("Property `eigenvalues_generator` not available for non-unitary operator.")
 
-class N(PrimitiveBlock):
+
+class N(Projector):
     """The N = (1/2)(I-Z) operator."""
 
     name = OpName.N
 
-    def __init__(self, target: int):
-        super().__init__((target,))
+    def __init__(self, target: int, state: str = "1"):
+        super().__init__(ket=state, bra=state, qubit_support=(target,))
 
     @property
     def generator(self) -> None:
@@ -189,9 +203,6 @@ class N(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([0, 1], dtype=cdouble)
 
-    def dagger(self) -> N:
-        return self
-
 
 class S(PrimitiveBlock):
     """The S / Phase gate."""
@@ -296,9 +307,6 @@ class I(PrimitiveBlock):
     def __ascii__(self, console: Console) -> Padding:
         return Padding("──────", (1, 1, 1, 1))
 
-    def dagger(self) -> I:
-        return I(*self.qubit_support)
-
 
 TPauliBlock = Union[X, Y, Z, I, N]
 
@@ -320,9 +328,6 @@ class H(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return torch.tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> H:
-        return H(*self.qubit_support)
-
 
 class Zero(PrimitiveBlock):
     name = OpName.ZERO
@@ -363,9 +368,6 @@ class Zero(PrimitiveBlock):
     def __pow__(self, power: int) -> AbstractBlock:
         return self
 
-    def dagger(self) -> Zero:
-        return Zero()
-
 
 class RX(ParametricBlock):
     """The Rx gate."""
@@ -668,7 +670,7 @@ class CNOT(ControlBlock):
     name = OpName.CNOT
 
     def __init__(self, control: int, target: int) -> None:
-        self.generator = kron((I(control) - Z(control)) * 0.5, X(target) - I(target))
+        self.generator = kron(N(control), X(target) - I(target))
         super().__init__((control,), X(target))
 
     @property
@@ -691,17 +693,12 @@ class CNOT(ControlBlock):
             tree.add(self._block_title)
         return tree
 
-    def dagger(self) -> CNOT:
-        return CNOT(*self.qubit_support)
-
 
 class MCZ(ControlBlock):
     name = OpName.MCZ
 
     def __init__(self, control: tuple[int, ...], target: int) -> None:
-        self.generator = kron(
-            *[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Z(target) - I(target)
-        )
+        self.generator = kron(*[N(qubit) for qubit in control], Z(target) - I(target))
         super().__init__(control, Z(target))
 
     @property
@@ -724,9 +721,6 @@ class MCZ(ControlBlock):
             tree.add(self._block_title)
         return tree
 
-    def dagger(self) -> MCZ:
-        return MCZ(self.qubit_support[:-1], self.qubit_support[-1])
-
 
 class CZ(MCZ):
     """The CZ gate."""
@@ -736,9 +730,6 @@ class CZ(MCZ):
     def __init__(self, control: int, target: int) -> None:
         super().__init__((control,), target)
 
-    def dagger(self) -> CZ:
-        return CZ(self.qubit_support[-2], self.qubit_support[-1])
-
 
 class MCRX(ParametricControlBlock):
     name = OpName.MCRX
@@ -749,7 +740,7 @@ class MCRX(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(*[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], X(target))
+        self.generator = kron(*[N(qubit) for qubit in control], X(target))
         super().__init__(control, RX(target, parameter))
 
     @classmethod
@@ -792,7 +783,7 @@ class MCRY(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(*[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Y(target))
+        self.generator = kron(*[N(qubit) for qubit in control], Y(target))
         super().__init__(control, RY(target, parameter))
 
     @classmethod
@@ -821,7 +812,7 @@ class CRY(MCRY):
         self,
         control: int,
         target: int,
-        parameter: Parameter | TNumber | sympy.Expr | str,
+        parameter: TParameter,
     ):
         super().__init__((control,), target, parameter)
 
@@ -835,7 +826,7 @@ class MCRZ(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(*[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Z(target))
+        self.generator = kron(*[N(qubit) for qubit in control], Z(target))
         super().__init__(control, RZ(target, parameter))
 
     @classmethod
@@ -878,10 +869,10 @@ class CSWAP(ControlBlock):
         if isinstance(control, tuple):
             control = control[0]
 
-        a00m = 0.5 * (Z(control) - I(control))
-        a00p = -0.5 * (Z(control) + I(control))
-        a11 = 0.5 * (Z(target1) - I(target1))
-        a22 = -0.5 * (Z(target2) + I(target2))
+        a00m = -N(target=control)
+        a00p = -N(target=control, state="0")
+        a11 = -N(target=target1)
+        a22 = -N(target=target2, state="0")
         a12 = 0.5 * (chain(X(target1), Z(target1)) + X(target1))
         a21 = 0.5 * (chain(Z(target2), X(target2)) + X(target2))
         no_effect = kron(a00m, I(target1), I(target2))
@@ -906,9 +897,6 @@ class CSWAP(ControlBlock):
     def nqubits(self) -> int:
         return 3
 
-    def dagger(self) -> CSWAP:
-        return CSWAP(*self.qubit_support)
-
 
 class T(PrimitiveBlock):
     """The T gate."""
@@ -994,9 +982,6 @@ class SWAP(PrimitiveBlock):
         s = f"{self.name}({c}, {t})"
         return s if self.tag is None else (s + rf" \[tag: {self.tag}]")
 
-    def dagger(self) -> SWAP:
-        return SWAP(*self.qubit_support)
-
 
 class AnalogSWAP(HamEvo):
     """
@@ -1031,9 +1016,7 @@ class MCPHASE(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(
-            *[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Z(target) - I(target)
-        )
+        self.generator = kron(*[N(qubit) for qubit in control], Z(target) - I(target))
         super().__init__(control, PHASE(target, parameter))
 
     @classmethod
@@ -1082,14 +1065,9 @@ class Toffoli(ControlBlock):
     name = OpName.TOFFOLI
 
     def __init__(self, control: tuple[int, ...], target: int) -> None:
-        self.generator = kron(
-            *[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], X(target) - I(target)
-        )
+        self.generator = kron(*[N(qubit) for qubit in control], X(target) - I(target))
         super().__init__(control, X(target))
 
-    def dagger(self) -> Toffoli:
-        return Toffoli(self.qubit_support[:-1], self.qubit_support[-1])
-
     @property
     def n_qubits(self) -> int:
         return len(self.qubit_support)
@@ -1129,6 +1107,7 @@ def _cast(T: Any, val: Any) -> Any:
 def wait(
     duration: TNumber | sympy.Basic,
     qubit_support: str | QubitSupport | tuple = "global",
+    add_pattern: bool = True,
 ) -> WaitBlock:
     """Constructs a [`WaitBlock`][qadence.blocks.analog.WaitBlock].
 
@@ -1142,7 +1121,7 @@ def wait(
     """
     q = _cast(QubitSupport, qubit_support)
     ps = ParamMap(duration=duration)
-    return WaitBlock(parameters=ps, qubit_support=q)
+    return WaitBlock(parameters=ps, qubit_support=q, add_pattern=add_pattern)
 
 
 def entangle(
@@ -1160,6 +1139,7 @@ def AnalogRot(
     delta: float | str | Parameter = 0,
     phase: float | str | Parameter = 0,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """General analog rotation operation.
 
@@ -1174,18 +1154,20 @@ def AnalogRot(
         ConstantAnalogRotation
     """
     q = _cast(QubitSupport, qubit_support)
-    if isinstance(duration, str):
-        duration = Parameter(duration)
-    alpha = duration * sympy.sqrt(omega**2 + delta**2) / 1000  # type: ignore [operator]
-
+    duration = Parameter(duration)
+    omega = Parameter(omega)
+    delta = Parameter(delta)
+    phase = Parameter(phase)
+    alpha = duration * sympy.sqrt(omega**2 + delta**2) / 1000
     ps = ParamMap(alpha=alpha, duration=duration, omega=omega, delta=delta, phase=phase)
-    return ConstantAnalogRotation(parameters=ps, qubit_support=q)
+    return ConstantAnalogRotation(parameters=ps, qubit_support=q, add_pattern=add_pattern)
 
 
 def _analog_rot(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple,
     phase: float,
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     q = _cast(QubitSupport, qubit_support)
     # assuming some arbitrary omega = π rad/μs
@@ -1200,12 +1182,13 @@ def _analog_rot(
     # and compute omega like this:
     # omega = alpha / duration * 1000
     ps = ParamMap(alpha=alpha, duration=duration, omega=omega, delta=0, phase=phase)
-    return ConstantAnalogRotation(parameters=ps, qubit_support=q)
+    return ConstantAnalogRotation(parameters=ps, qubit_support=q, add_pattern=add_pattern)
 
 
 def AnalogRX(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """Analog X rotation.
 
@@ -1223,12 +1206,13 @@ def AnalogRX(
     Returns:
         ConstantAnalogRotation
     """
-    return _analog_rot(angle, qubit_support, phase=0)
+    return _analog_rot(angle, qubit_support, phase=0, add_pattern=add_pattern)
 
 
 def AnalogRY(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """Analog Y rotation.
 
@@ -1245,12 +1229,13 @@ def AnalogRY(
     Returns:
         ConstantAnalogRotation
     """
-    return _analog_rot(angle, qubit_support, phase=-np.pi / 2)
+    return _analog_rot(angle, qubit_support, phase=-np.pi / 2, add_pattern=add_pattern)
 
 
 def AnalogRZ(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """Analog Z rotation. Shorthand for [`AnalogRot`][qadence.operations.AnalogRot]:
     ```
@@ -1263,7 +1248,7 @@ def AnalogRZ(
     delta = np.pi
     duration = alpha / delta * 1000
     ps = ParamMap(alpha=alpha, duration=duration, omega=0, delta=delta, phase=0.0)
-    return ConstantAnalogRotation(qubit_support=q, parameters=ps)
+    return ConstantAnalogRotation(qubit_support=q, parameters=ps, add_pattern=add_pattern)
 
 
 # gate sets
@@ -1288,4 +1273,4 @@ analog_gateset = [
     entangle,
     wait,
 ]
-non_unitary_gateset = [Zero, N]
+non_unitary_gateset = [Zero, N, Projector]

qadence/parameters.py

@@ -8,8 +8,9 @@ import numpy as np
 import sympy
 from sympy import *
 from sympy import Array, Basic, Expr, Symbol, sympify
+from sympy.physics.quantum.dagger import Dagger
 from sympytorch import SymPyModule
-from torch import Tensor, rand, tensor
+from torch import Tensor, heaviside, no_grad, rand, tensor
 
 from qadence.logger import get_logger
 from qadence.types import TNumber
@@ -19,6 +20,7 @@ __all__ = ["FeatureParameter", "Parameter", "VariationalParameter"]
 
 logger = get_logger(__file__)
 
+dagger_expression = Dagger
 
 ParameterJSONSchema = {
     "$schema": "http://json-schema.org/draft-07/schema#",
@@ -197,7 +199,13 @@ def torchify(expr: Expr) -> SymPyModule:
     Returns:
         A torchified, differentiable Expression.
     """
-    extra_funcs = {sympy.core.numbers.ImaginaryUnit: 1.0j}
+
+    def heaviside_func(x: Tensor, _: Any) -> Tensor:
+        with no_grad():
+            res = heaviside(x, tensor(0.5))
+        return res
+
+    extra_funcs = {sympy.core.numbers.ImaginaryUnit: 1.0j, sympy.Heaviside: heaviside_func}
     return SymPyModule(expressions=[sympy.N(expr)], extra_funcs=extra_funcs)
 
 

qadence/register.py

@@ -11,12 +11,16 @@ import numpy as np
 from deepdiff import DeepDiff
 from networkx.classes.reportviews import EdgeView, NodeView
 
+from qadence.analog import IdealDevice, RydbergDevice
 from qadence.types import LatticeTopology
 
 # Modules to be automatically added to the qadence namespace
 __all__ = ["Register"]
 
 
+DEFAULT_DEVICE = IdealDevice()
+
+
 def _scale_node_positions(graph: nx.Graph, min_distance: float, spacing: float) -> None:
     scaled_nodes = {}
     scale_factor = spacing / min_distance
@@ -27,8 +31,14 @@ def _scale_node_positions(graph: nx.Graph, min_distance: float, spacing: float)
 
 
 class Register:
-    def __init__(self, support: nx.Graph | int, spacing: float | None = 1.0):
-        """A 2D register of qubits which includes their coordinates.
+    def __init__(
+        self,
+        support: nx.Graph | int,
+        spacing: float | None = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
+    ):
+        """
+        A 2D register of qubits which includes their coordinates.
 
         It is needed for e.g. analog computing.
         The coordinates are ignored in backends that don't need them. The easiest
@@ -52,21 +62,16 @@ class Register:
         reg.draw()
         ```
         """
-        self.graph = support if isinstance(support, nx.Graph) else alltoall_graph(support)
+        if device_specs is not None and not isinstance(device_specs, RydbergDevice):
+            raise ValueError("Device specs are not valid. Please pass a `RydbergDevice` instance.")
 
-        # Auxiliary complete graph
-        support = self.graph.nodes
-        all_edges = list(filter(lambda x: x[0] < x[1], product(support, support)))
-        self.complete_graph = nx.Graph()
-        self.complete_graph.add_nodes_from(support)
-        self.complete_graph.add_edges_from(all_edges)
+        self.device_specs = device_specs
+
+        self.graph = support if isinstance(support, nx.Graph) else alltoall_graph(support)
 
         if spacing is not None and self.min_distance != 0.0:
             _scale_node_positions(self.graph, self.min_distance, spacing)
 
-        pos_values = nx.get_node_attributes(self.graph, "pos")
-        nx.set_node_attributes(self.complete_graph, pos_values, "pos")
-
     @property
     def n_qubits(self) -> int:
         return len(self.graph)
@@ -77,27 +82,43 @@ class Register:
         coords: list[tuple],
         lattice: LatticeTopology | str = LatticeTopology.ARBITRARY,
         spacing: float | None = None,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
     ) -> Register:
         graph = nx.Graph()
         for i, pos in enumerate(coords):
             graph.add_node(i, pos=pos)
-        return cls(graph, spacing)
+        return cls(graph, spacing, device_specs)
 
     @classmethod
-    def line(cls, n_qubits: int, spacing: float = 1.0) -> Register:
-        return cls(line_graph(n_qubits), spacing)
+    def line(
+        cls,
+        n_qubits: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
+    ) -> Register:
+        return cls(line_graph(n_qubits), spacing, device_specs)
 
     @classmethod
-    def circle(cls, n_qubits: int, spacing: float = 1.0) -> Register:
+    def circle(
+        cls,
+        n_qubits: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
+    ) -> Register:
         graph = nx.grid_2d_graph(n_qubits, 1, periodic=True)
         graph = nx.relabel_nodes(graph, {(i, 0): i for i in range(n_qubits)})
         coords = nx.circular_layout(graph)
         values = {i: {"pos": pos} for i, pos in coords.items()}
         nx.set_node_attributes(graph, values)
-        return cls(graph, spacing)
+        return cls(graph, spacing, device_specs)
 
     @classmethod
-    def square(cls, qubits_side: int, spacing: float = 1.0) -> Register:
+    def square(
+        cls,
+        qubits_side: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
+    ) -> Register:
         n_points = 4 * (qubits_side - 1)
 
         def gen_points() -> np.ndarray:
@@ -122,35 +143,52 @@ class Register:
         graph = nx.relabel_nodes(graph, {(i, 0): i for i in range(n_points)})
         values = {i: {"pos": point} for i, point in zip(graph.nodes, gen_points())}
         nx.set_node_attributes(graph, values)
-        return cls(graph, spacing)
+        return cls(graph, spacing, device_specs)
 
     @classmethod
-    def all_to_all(cls, n_qubits: int, spacing: float = 1.0) -> Register:
-        return cls(alltoall_graph(n_qubits), spacing)
+    def all_to_all(
+        cls,
+        n_qubits: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
+    ) -> Register:
+        return cls(alltoall_graph(n_qubits), spacing, device_specs)
 
     @classmethod
     def rectangular_lattice(
-        cls, qubits_row: int, qubits_col: int, spacing: float = 1.0
+        cls,
+        qubits_row: int,
+        qubits_col: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
     ) -> Register:
         graph = nx.grid_2d_graph(qubits_col, qubits_row)
         values = {i: {"pos": node} for (i, node) in enumerate(graph.nodes)}
         graph = nx.relabel_nodes(graph, {(i, j): k for k, (i, j) in enumerate(graph.nodes)})
         nx.set_node_attributes(graph, values)
-        return cls(graph, spacing)
+        return cls(graph, spacing, device_specs)
 
     @classmethod
     def triangular_lattice(
-        cls, n_cells_row: int, n_cells_col: int, spacing: float = 1.0
+        cls,
+        n_cells_row: int,
+        n_cells_col: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
     ) -> Register:
-        return cls(triangular_lattice_graph(n_cells_row, n_cells_col), spacing)
+        return cls(triangular_lattice_graph(n_cells_row, n_cells_col), spacing, device_specs)
 
     @classmethod
     def honeycomb_lattice(
-        cls, n_cells_row: int, n_cells_col: int, spacing: float = 1.0
+        cls,
+        n_cells_row: int,
+        n_cells_col: int,
+        spacing: float = 1.0,
+        device_specs: RydbergDevice = DEFAULT_DEVICE,
     ) -> Register:
         graph = nx.hexagonal_lattice_graph(n_cells_row, n_cells_col)
         graph = nx.relabel_nodes(graph, {(i, j): k for k, (i, j) in enumerate(graph.nodes)})
-        return cls(graph, spacing)
+        return cls(graph, spacing, device_specs)
 
     @classmethod
     def lattice(cls, topology: LatticeTopology | str, *args: Any, **kwargs: Any) -> Register:
@@ -166,25 +204,29 @@ class Register:
         return self.graph.nodes[item]
 
     @property
-    def support(self) -> set:
-        return set(self.graph.nodes)
-
-    @property
-    def coords(self) -> dict:
-        return {i: tuple(node.get("pos", ())) for i, node in self.graph.nodes.items()}
+    def nodes(self) -> NodeView:
+        return self.graph.nodes
 
     @property
     def edges(self) -> EdgeView:
         return self.graph.edges
 
     @property
-    def all_edges(self) -> EdgeView:
-        return self.complete_graph.edges
+    def support(self) -> set:
+        return set(self.nodes)
+
+    @property
+    def coords(self) -> dict:
+        return {i: tuple(node.get("pos", ())) for i, node in self.nodes.items()}
+
+    @property
+    def all_node_pairs(self) -> EdgeView:
+        return list(filter(lambda x: x[0] < x[1], product(self.support, self.support)))
 
     @property
     def distances(self) -> dict:
         coords = self.coords
-        return {edge: dist(coords[edge[0]], coords[edge[1]]) for edge in self.all_edges}
+        return {edge: dist(coords[edge[0]], coords[edge[1]]) for edge in self.all_node_pairs}
 
     @property
     def edge_distances(self) -> dict:
@@ -197,21 +239,27 @@ class Register:
         value: float = min(self.distances.values()) if len(distances) > 0 else 0.0
         return value
 
-    @property
-    def nodes(self) -> NodeView:
-        return self.graph.nodes
-
     def rescale_coords(self, scaling: float) -> Register:
         g = deepcopy(self.graph)
         _scale_node_positions(g, min_distance=1.0, spacing=scaling)
-        return Register(g, spacing=None)
+        return Register(g, spacing=None, device_specs=self.device_specs)
 
     def _to_dict(self) -> dict:
-        return {"graph": nx.node_link_data(self.graph)}
+        return {
+            "graph": nx.node_link_data(self.graph),
+            "device_specs": self.device_specs._to_dict(),
+        }
 
     @classmethod
     def _from_dict(cls, d: dict) -> Register:
-        return cls(nx.node_link_graph(d["graph"]))
+        device_dict = d.get("device_specs", None)
+        if device_dict is None:
+            device_dict = DEFAULT_DEVICE._to_dict()
+
+        return cls(
+            support=nx.node_link_graph(d["graph"]),
+            device_specs=RydbergDevice._from_dict(device_dict),
+        )
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Register):

qadence/transpile/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from .apply_fn import apply_fn_to_blocks
 from .block import (
     chain_single_qubit_ops,
     repeat,

qadence/transpile/apply_fn.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from qadence.blocks import AbstractBlock, CompositeBlock, add, chain, kron
+from qadence.blocks.analog import AnalogChain
+
+COMPOSE_FN_DICT = {
+    "ChainBlock": chain,
+    "AnalogChain": chain,
+    "KronBlock": kron,
+    "AnalogKron": kron,
+    "AddBlock": add,
+}
+
+
+def apply_fn_to_blocks(
+    input_block: AbstractBlock, block_fn: Callable, *args: Any, **kwargs: Any
+) -> AbstractBlock:
+    """
+    Recurses through the block tree and applies a given function to all the leaf blocks.
+
+    Arguments:
+        input_block: tree of blocks on which to apply the recurse.
+        block_fn: callable function to apply to each leaf block.
+        args: any positional arguments to pass to the leaf function.
+        kwargs: any keyword arguments to pass to the leaf function.
+    """
+
+    if isinstance(input_block, (CompositeBlock, AnalogChain)):
+        parsed_blocks = [
+            apply_fn_to_blocks(block, block_fn, *args, **kwargs) for block in input_block.blocks
+        ]
+        compose_fn = COMPOSE_FN_DICT[type(input_block).__name__]
+        output_block = compose_fn(*parsed_blocks)  # type: ignore [arg-type]
+    else:
+        # AnalogKrons are considered as a leaf block
+        output_block = block_fn(input_block, *args, **kwargs)
+
+    return output_block