Repository: QCHackers/tqec (python)

Issue:
Implement a grid template
**Is your feature request related to a problem? Please describe.**
For the moment, we can limit ourselves to regular grids of templates (that will likely represent logical qubits). There is nothing in the code that is explicitly here to build such a grid.

**Describe the solution you'd like**
Constructing a `GridTemplate` should be quite easy with the `ComposedTemplate` class.

**Describe alternatives you've considered**
Having a specific implementation for a grid, just like `StackedTemplate` and `ShiftedTemplate` are not implemented using `ComposedTemplate` (because it does not provide the necessary features) but rather directly as a subclass of `Template` with a specific behaviour. Turns out that a grid should be describable quite easily with a `ComposedTemplate`, so this alternative has not been explored.


Pull request description:
Implement `TemplateGrid`
This PR implements a generic `TemplateGrid` that can be used to generate a grid of templates representing logical qubits (for example).

Closes #244 

Patch:
diff --git a/src/tqec/templates/constructions/grid.py b/src/tqec/templates/constructions/grid.py
new file mode 100644
index 00000000..3d45f502
--- /dev/null
+++ b/src/tqec/templates/constructions/grid.py
@@ -0,0 +1,37 @@
+from tqec.enums import BELOW_OF, RIGHT_OF
+from tqec.templates.base import Template, TemplateWithIndices
+from tqec.templates.composed import ComposedTemplate
+
+
+class TemplateGrid(ComposedTemplate):
+    def __init__(self, rows: int, cols: int, template: Template) -> None:
+        plaquettes_indices_per_qubit = template.expected_plaquettes_number
+
+        super().__init__([])
+
+        template_indices: list[list[int]] = []
+        for i in range(rows):
+            template_indices.append([])
+            for j in range(cols):
+                twi = TemplateWithIndices(
+                    template,
+                    list(
+                        range(
+                            1 + (i * cols + j) * plaquettes_indices_per_qubit,
+                            1 + (i * cols + j + 1) * plaquettes_indices_per_qubit,
+                        )
+                    ),
+                )
+                template_indices[i].append(self.add_template(twi))
+
+        # First column
+        for i in range(1, rows):
+            self.add_relation(
+                template_indices[i][0], BELOW_OF, template_indices[i - 1][0]
+            )
+        # Other templates
+        for i in range(rows):
+            for j in range(1, cols):
+                self.add_relation(
+                    template_indices[i][j], RIGHT_OF, template_indices[i][j - 1]
+                )
diff --git a/src/tqec/templates/constructions/qubit.py b/src/tqec/templates/constructions/qubit.py
index b50a0858..f8b16dc9 100644
--- a/src/tqec/templates/constructions/qubit.py
+++ b/src/tqec/templates/constructions/qubit.py
@@ -2,13 +2,86 @@ from __future__ import annotations
 
 from tqec.enums import ABOVE_OF, BELOW_OF, LEFT_OF, RIGHT_OF, TemplateOrientation
 from tqec.exceptions import TQECException
-from tqec.templates.atomic.rectangle import AlternatingRectangleTemplate
+from tqec.templates.atomic.rectangle import (
+    AlternatingRectangleTemplate,
+    RawRectangleTemplate,
+)
 from tqec.templates.atomic.square import AlternatingSquareTemplate
 from tqec.templates.base import TemplateWithIndices
 from tqec.templates.composed import ComposedTemplate
 from tqec.templates.scale import Dimension, FixedDimension
 
 
+class DenseQubitSquareTemplate(ComposedTemplate):
+    def __init__(self, dim: Dimension) -> None:
+        """An error-corrected qubit.
+
+        The below text represents this template for an input ``dimension = FixedDimension(4)`` ::
+
+            1  5  6  5  6  2
+            7  9 10  9 10 11
+            8 10  9 10  9 12
+            7  9 10  9 10 11
+            8 10  9 10  9 12
+            3 13 14 13 14  4
+
+        Args:
+            dim: dimension of the error-corrected qubit.
+        """
+        # nsone: non-scalable one
+        nsone = FixedDimension(1)
+
+        # 0  4  4  4  4  1
+        # 5  6  6  6  6  7
+        # 5  6  6  6  6  7
+        # 5  6  6  6  6  7
+        # 5  6  6  6  6  7
+        # 2  8  8  8  8  3
+        _templates = [
+            TemplateWithIndices(RawRectangleTemplate([[0]]), [1]),
+            TemplateWithIndices(RawRectangleTemplate([[0]]), [2]),
+            TemplateWithIndices(RawRectangleTemplate([[0]]), [3]),
+            TemplateWithIndices(RawRectangleTemplate([[0]]), [4]),
+            # Top rectangle, containing plaquettes of type 5 and 6
+            TemplateWithIndices(AlternatingRectangleTemplate(dim, nsone), [5, 6]),
+            # Left rectangle, containing plaquettes of type 7 and 8
+            TemplateWithIndices(AlternatingRectangleTemplate(nsone, dim), [7, 8]),
+            # Central square, containing plaquettes of types 9 and 10
+            TemplateWithIndices(AlternatingSquareTemplate(dim), [9, 10]),
+            # Right rectangle, containing plaquettes of type 11 and 12
+            TemplateWithIndices(AlternatingRectangleTemplate(nsone, dim), [11, 12]),
+            # Bottom rectangle, containing plaquettes of type 13 and 14
+            TemplateWithIndices(AlternatingRectangleTemplate(dim, nsone), [13, 14]),
+        ]
+        _relations = [
+            (5, BELOW_OF, 0),
+            (2, BELOW_OF, 5),
+            (4, RIGHT_OF, 0),
+            (1, RIGHT_OF, 4),
+            (6, RIGHT_OF, 5),
+            (7, RIGHT_OF, 6),
+            (8, BELOW_OF, 6),
+            (3, RIGHT_OF, 8),
+        ]
+        super().__init__(_templates)
+        for source, relpos, target in _relations:
+            self.add_relation(source, relpos, target)
+
+    def get_midline_plaquettes(
+        self, orientation: TemplateOrientation = TemplateOrientation.HORIZONTAL
+    ) -> list[tuple[int, int]]:
+        midline_shape, iteration_shape = self.shape.x, self.shape.y
+        if midline_shape % 2 == 1:
+            raise TQECException(
+                "Midline is not defined for odd "
+                + f"{'height' if orientation == TemplateOrientation.HORIZONTAL else 'width'}."
+            )
+        midline = midline_shape // 2 - 1
+        if orientation == TemplateOrientation.VERTICAL:
+            return [(row, midline) for row in range(iteration_shape)]
+        return [(midline, column) for column in range(iteration_shape)]
+
+
 class QubitSquareTemplate(ComposedTemplate):
     def __init__(self, dim: Dimension) -> None:
         """An error-corrected qubit.


Test patch:
diff --git a/src/tqec/templates/constructions/grid_test.py b/src/tqec/templates/constructions/grid_test.py
new file mode 100644
index 00000000..e72b14d6
--- /dev/null
+++ b/src/tqec/templates/constructions/grid_test.py
@@ -0,0 +1,51 @@
+import numpy
+
+from tqec.templates.atomic.rectangle import RawRectangleTemplate
+from tqec.templates.constructions.grid import TemplateGrid
+from tqec.templates.constructions.qubit import DenseQubitSquareTemplate
+from tqec.templates.scale import FixedDimension
+
+
+def test_grid_simple():
+    template = RawRectangleTemplate([[0]])
+    grid = TemplateGrid(2, 2, template)
+
+    assert grid.expected_plaquettes_number == 4
+    arr = grid.instantiate([1, 2, 3, 4])
+    numpy.testing.assert_allclose(arr, [[1, 2], [3, 4]])
+
+    grid.scale_to(18)
+    assert grid.expected_plaquettes_number == 4
+    arr = grid.instantiate([1, 2, 3, 4])
+    numpy.testing.assert_allclose(arr, [[1, 2], [3, 4]])
+
+
+def test_grid_qubits():
+    template = DenseQubitSquareTemplate(FixedDimension(2))
+    assert template.expected_plaquettes_number == 14
+    template_arr = template.instantiate(list(range(1, 15)))
+
+    grid = TemplateGrid(3, 2, template)
+    assert grid.expected_plaquettes_number == 14 * 3 * 2
+    grid_arr = grid.instantiate(list(range(1, 1 + 14 * 2 * 3)))
+
+    plaquettes_per_template = template.expected_plaquettes_number
+    expected_arr = numpy.vstack(
+        (
+            numpy.hstack((template_arr, template_arr + plaquettes_per_template)),
+            numpy.hstack(
+                (
+                    template_arr + 2 * plaquettes_per_template,
+                    template_arr + 3 * plaquettes_per_template,
+                )
+            ),
+            numpy.hstack(
+                (
+                    template_arr + 4 * plaquettes_per_template,
+                    template_arr + 5 * plaquettes_per_template,
+                )
+            ),
+        )
+    )
+    assert grid_arr.shape == expected_arr.shape
+    numpy.testing.assert_allclose(grid_arr, expected_arr)