electrostatics._quantum_region — quantum-region plane extraction

tempura/electrostatics/_quantum_region.py turns solved solver vectors into a rectangular view of the active region — the plane where dots, barriers, or Majorana modes live. Its two entry points, extract_quantum_region_plane(...) and extract_quantum_region_basis(...), are re-exported from tempura.electrostatics for use after a solve.

Why a separate step

A solved basis potential is a flat vector aligned with the finalized problem's mesh sites, not an image. To plot or analyze the field in the quantum region you need the subset of sites inside that region, arranged on a regular grid. This module extracts that grid and reshapes the vectors onto it.

Extracting the plane

extract_quantum_region_plane(...) resolves the target \(z\) plane from the region's realized \(z\) levels (by "lower", "midpoint", "upper", or an explicit value), then samples the mesh on a slab spanning the region's own bounding box rather than only the cells the region owns. It selects the sites at the target \(z\) and sorts them lexicographically by \((y, x)\). Sampling the region's bounding box — instead of just the owned footprint — is what lets a patterned quantum region still produce a complete plane: the patterned cells sit inside an otherwise rectangular grid of sites at the same \(z\). Deriving the slab from the region's own box (rather than a symmetric extent over every region) keeps it off neighboring, coarser layers and works for devices that are not centered on the origin.

It then verifies the selected sites form a complete rectangle:

\[ N_\text{sites} = n_x \cdot n_y, \]

and that the sorted \(x\) and \(y\) coordinates broadcast back to a regular grid. The result is a frozen QuantumRegionPlaneData carrying the sorted solver indices, the coordinates, the unique \(x\)/\(y\) values, the shared \(z\), \((n_x, n_y)\), and a region_mask. The rectangular-grid check is what guarantees a clean reshape downstream.

Because the plane spans the region's full bounding box, a patterned region leaves non-owned filler cells inside that rectangle. region_mask is an \((n_y, n_x)\) boolean array — the same membership test as points_inside — that flags exactly the cells the region owns, so callers can blank the filler instead of treating it as active material. For a full-footprint region every entry is True.

Reshaping the basis

Given that plane, extract_quantum_region_basis(...) indexes each gate's solver vector by the plane's sorted indices and reshapes it to \((n_y, n_x)\), returning one image per gate in the requested order. project_gate_footprints_to_plane(...) similarly samples each gate's footprint onto the same plane for overlays.

API

_quantum_region

Helpers for extracting rectangular quantum region planes from solved problems.

QuantumRegionPlaneData dataclass

Rectangular quantum region plane data extracted from a solved problem.

Attributes:

Name	Type	Description
sorted_indices	NDArray[int_]	Solver-vector indices sorted lexicographically by (y, x) on the selected plane.
coords	NDArray[float64]	Sorted (x, y, z) coordinates on the selected quantum region plane.
x_values	NDArray[float64]	Unique x coordinates spanning the rectangular plane.
y_values	NDArray[float64]	Unique y coordinates spanning the rectangular plane.
z_value	float	Shared z coordinate of the selected plane.
nx	int	Number of x samples in the plane.
ny	int	Number of y samples in the plane.
region_mask	NDArray[bool_]	(ny, nx) boolean mask marking which plane cells are actually owned by the quantum region. For a full-footprint region every cell is True; for a patterned region the plane is a complete rectangle (so basis vectors reshape cleanly) but only the stenciled cells are True, so callers never treat the bounding-box filler cells as active material.

Source code in tempura/electrostatics/_quantum_region.py

@dataclass(frozen=True)
class QuantumRegionPlaneData:
    """Rectangular quantum region plane data extracted from a solved problem.

    Attributes:
        sorted_indices: Solver-vector indices sorted lexicographically by
            ``(y, x)`` on the selected plane.
        coords: Sorted ``(x, y, z)`` coordinates on the selected quantum region plane.
        x_values: Unique x coordinates spanning the rectangular plane.
        y_values: Unique y coordinates spanning the rectangular plane.
        z_value: Shared z coordinate of the selected plane.
        nx: Number of x samples in the plane.
        ny: Number of y samples in the plane.
        region_mask: ``(ny, nx)`` boolean mask marking which plane cells are
            actually owned by the quantum region. For a full-footprint region
            every cell is ``True``; for a **patterned** region the plane is a
            complete rectangle (so basis vectors reshape cleanly) but only the
            stenciled cells are ``True``, so callers never treat the
            bounding-box filler cells as active material.
    """

    sorted_indices: NDArray[np.int_]
    coords: NDArray[np.float64]
    x_values: NDArray[np.float64]
    y_values: NDArray[np.float64]
    z_value: float
    nx: int
    ny: int
    region_mask: NDArray[np.bool_]

extract_quantum_region_basis(basis_potentials, gate_names, plane)

Reshape solved gate basis vectors onto a rectangular quantum region plane.

Parameters:

Name	Type	Description	Default
basis_potentials	Mapping[str, ndarray]	Mapping of gate name to solved potential vector.	required
gate_names	list[str]	Gate order to preserve in the returned mapping.	required
plane	QuantumRegionPlaneData	Extracted quantum region plane metadata and sorted solver indices.	required

Returns:

Type	Description
dict[str, ndarray]	Mapping of gate name to (ny, nx) quantum region basis potential arrays.

Source code in tempura/electrostatics/_quantum_region.py

def extract_quantum_region_basis(
    basis_potentials: Mapping[str, np.ndarray],
    gate_names: list[str],
    plane: QuantumRegionPlaneData,
) -> dict[str, np.ndarray]:
    """Reshape solved gate basis vectors onto a rectangular quantum region plane.

    Args:
        basis_potentials: Mapping of gate name to solved potential vector.
        gate_names: Gate order to preserve in the returned mapping.
        plane: Extracted quantum region plane metadata and sorted solver indices.

    Returns:
        Mapping of gate name to ``(ny, nx)`` quantum region basis potential arrays.
    """
    return {
        gate_name: np.asarray(basis_potentials[gate_name])[plane.sorted_indices].reshape(
            plane.ny,
            plane.nx,
        )
        for gate_name in gate_names
    }

extract_quantum_region_plane(problem, region_shapes, *, region_name='quantum_region', plane_selection='midpoint')

Return one rectangular quantum region plane extracted from problem.

Parameters:

Name	Type	Description	Default
problem	Any	Finalized Pescado problem exposing coordinates and points_inside.	required
region_shapes	Mapping[str, Any]	Mapping of region names to shapes.	required
region_name	str	Region to extract from region_shapes.	'quantum_region'
plane_selection	PlaneSelection	Which z plane to export when the realized quantum region spans multiple z coordinates. Use "lower", "midpoint", "upper", or an explicit z value.	'midpoint'

Returns:

Type	Description
QuantumRegionPlaneData	Sorted rectangular plane data for the selected z slice across the full
QuantumRegionPlaneData	device XY span.

Raises:

Type	Description
ValueError	If the quantum region is missing, empty, or the selected sampling slab does not form a rectangular plane.

Source code in tempura/electrostatics/_quantum_region.py

def extract_quantum_region_plane(
    problem: Any,
    region_shapes: Mapping[str, Any],
    *,
    region_name: str = "quantum_region",
    plane_selection: PlaneSelection = "midpoint",
) -> QuantumRegionPlaneData:
    """Return one rectangular quantum region plane extracted from ``problem``.

    Args:
        problem: Finalized Pescado problem exposing ``coordinates`` and
            ``points_inside``.
        region_shapes: Mapping of region names to shapes.
        region_name: Region to extract from ``region_shapes``.
        plane_selection: Which z plane to export when the realized quantum region spans
            multiple z coordinates. Use ``"lower"``, ``"midpoint"``,
            ``"upper"``, or an explicit z value.

    Returns:
        Sorted rectangular plane data for the selected z slice across the full
        device XY span.

    Raises:
        ValueError: If the quantum region is missing, empty, or the selected
            sampling slab does not form a rectangular plane.
    """
    try:
        region_shape = region_shapes[region_name]
    except KeyError as exc:
        raise ValueError(
            "quantum region export requires a matching region_shapes entry; "
            f"missing {region_name!r}."
        ) from exc

    indices_quantum_region = np.asarray(problem.points_inside(region_shape), dtype=int)
    coords_quantum_region = np.asarray(
        problem.coordinates[indices_quantum_region],
        dtype=float,
    )
    if coords_quantum_region.size == 0:
        raise ValueError("No solver coordinates were found inside the quantum region.")

    z_vals = np.unique(np.round(coords_quantum_region[:, 2], 12))
    z_target = _resolve_plane_z_value(z_vals, plane_selection=plane_selection)
    # Sample the quantum region's own XY/z extent to recover a rectangular
    # plane before slicing to the requested z level. Using the region's own
    # bounding box (rather than a symmetric max over every region) keeps the
    # slab off neighboring coarser layers and works for devices that are not
    # centered on the origin.
    slab = _sampling_slab_for_quantum_region_plane(region_shape)
    indices_quantum_region = np.asarray(problem.points_inside(slab), dtype=int)
    coords_quantum_region = np.asarray(
        problem.coordinates[indices_quantum_region],
        dtype=float,
    )
    if coords_quantum_region.size == 0:
        raise ValueError(
            "No solver coordinates were found inside the quantum region sampling slab."
        )

    plane_mask = np.isclose(coords_quantum_region[:, 2], z_target, atol=1e-12, rtol=0.0)
    if not np.any(plane_mask):
        raise ValueError(
            "The requested quantum region plane does not match any realized "
            f"z coordinates; requested={z_target}, available={z_vals.tolist()}."
        )
    coords_quantum_region = coords_quantum_region[plane_mask]
    indices_quantum_region = indices_quantum_region[plane_mask]

    # Sort by y first, then x, so reshape(ny, nx) produces row-major arrays
    # matching plotting/image conventions.
    order = np.lexsort((coords_quantum_region[:, 0], coords_quantum_region[:, 1]))
    coords_quantum_region = coords_quantum_region[order]
    sorted_indices_quantum_region = indices_quantum_region[order]

    # Rectangular-grid validation catches malformed or partially meshed active
    # regions before callers reshape solver vectors into plane arrays.
    x_values = np.unique(np.round(coords_quantum_region[:, 0], 12))
    y_values = np.unique(np.round(coords_quantum_region[:, 1], 12))
    nx = len(x_values)
    ny = len(y_values)
    if coords_quantum_region.shape[0] != nx * ny:
        raise ValueError(
            "The extracted quantum region coordinates do not form a rectangular grid."
        )

    x_grid = coords_quantum_region[:, 0].reshape(ny, nx)
    y_grid = coords_quantum_region[:, 1].reshape(ny, nx)
    if not np.allclose(x_grid, np.broadcast_to(x_values, (ny, nx))):
        raise ValueError(
            "The extracted quantum region x coordinates do not form a rectangular grid."
        )
    if not np.allclose(y_grid, np.broadcast_to(y_values[:, None], (ny, nx))):
        raise ValueError(
            "The extracted quantum region y coordinates do not form a rectangular grid."
        )

    # The plane is the region's full bounding box, so a patterned region leaves
    # non-owned filler cells inside it. Mark which cells the region actually
    # owns (the same membership test as ``points_inside``) so callers can mask
    # out the filler instead of treating it as active material.
    region_mask = np.asarray(region_shape(coords_quantum_region), dtype=bool).reshape(
        ny, nx
    )

    return QuantumRegionPlaneData(
        sorted_indices=sorted_indices_quantum_region,
        coords=coords_quantum_region,
        x_values=np.asarray(x_values, dtype=float),
        y_values=np.asarray(y_values, dtype=float),
        z_value=float(coords_quantum_region[0, 2]),
        nx=nx,
        ny=ny,
        region_mask=region_mask,
    )

project_gate_footprints_to_plane(gate_shapes, plane)

Project gate XY footprints onto the exported rectangular quantum region plane.

Each gate is sampled at the midpoint of its z extent so the resulting mask matches the gate's in-plane footprint while remaining aligned to the quantum region export grid.

Parameters:

Name	Type	Description	Default
gate_shapes	Mapping[str, Any]	Mapping of gate name to callable solver-owned gate shape.	required
plane	QuantumRegionPlaneData	Rectangular quantum region plane onto which the footprints should be sampled.	required

Returns:

Type	Description
dict[str, ndarray]	Mapping of gate name to (ny, nx) boolean footprint mask aligned to
dict[str, ndarray]	plane.

Source code in tempura/electrostatics/_quantum_region.py

def project_gate_footprints_to_plane(
    gate_shapes: Mapping[str, Any],
    plane: QuantumRegionPlaneData,
) -> dict[str, np.ndarray]:
    """Project gate XY footprints onto the exported rectangular quantum region plane.

    Each gate is sampled at the midpoint of its z extent so the resulting mask
    matches the gate's in-plane footprint while remaining aligned to the quantum region
    export grid.

    Args:
        gate_shapes: Mapping of gate name to callable solver-owned gate shape.
        plane: Rectangular quantum region plane onto which the footprints should be
            sampled.

    Returns:
        Mapping of gate name to ``(ny, nx)`` boolean footprint mask aligned to
        ``plane``.
    """
    if not gate_shapes:
        return {}

    coords = np.asarray(plane.coords, dtype=float)
    footprints: dict[str, np.ndarray] = {}
    for gate_name, gate_shape in gate_shapes.items():
        # The footprint projection only needs the gate's XY shape. Sampling at
        # the gate midpoint in z avoids coupling the mask to the quantum-region
        # plane height.
        bbox = getattr(gate_shape, "bbox", None)
        if bbox is None or not callable(gate_shape):
            continue

        bbox_arr = np.asarray(bbox, dtype=float)
        if bbox_arr.shape != (2, 3):
            continue
        z_sample = float(0.5 * (bbox_arr[0, 2] + bbox_arr[1, 2]))
        sample_points = coords.copy()
        sample_points[:, 2] = z_sample

        mask = np.asarray(gate_shape(sample_points), dtype=bool).reshape(-1)
        if mask.shape[0] != coords.shape[0]:
            continue
        footprints[gate_name] = mask.reshape(plane.ny, plane.nx)

    return footprints