Factor out dual basis calculation

src/qldpc/abstract/rings.py

@@ -806,12 +806,10 @@ def _remove_zero_rows(matrices: list[galois.FieldArray]) -> list[galois.FieldArr
             1. The column of the first nonzero value in the pivot_row of each component.
             2. The column that will contain the pivot when we recombine the components.
             """
-            pivot_cols = [
-                num_cols
-                if not np.any(row := matrix[pivot_row])
-                else int(np.argmax(row.view(np.ndarray).astype(bool)))
-                for matrix in matrices
+            pivot_rows_as_bools = [
+                matrix[pivot_row].view(np.ndarray).astype(bool) for matrix in matrices
             ]
+            pivot_cols = qldpc.math.first_nonzero_cols(pivot_rows_as_bools)
             pivot_col = min(pivot_cols)
 
             """
@@ -1163,7 +1161,7 @@ def __init__(self, pci: RingMember, *, seed: np.random.Generator | None = None)
         self.size = math.isqrt(np.linalg.matrix_rank(self.pci_reg) // self.degree)
 
         self.power_basis = self._get_power_basis(seed)
-        self.power_basis_dual = self._get_dual_basis(self.power_basis)
+        self.power_basis_dual = qldpc.math.get_dual_basis(self.power_basis, validate=False)
 
         self.extended_field = galois.GF(self.field.order**self.degree)
         self.embedded_scalars, self.embedded_power_basis = self._get_center_embeddings()
@@ -1262,17 +1260,6 @@ def _random_nonzero_vec(self, length: int, seed: np.random.Generator) -> galois.
             pass  # pragma: no cover
         return vector
 
-    def _get_dual_basis(self, basis: galois.FieldArray) -> galois.FieldArray:
-        """Construct a dual basis, for which dual_basis @ basis.T = identity_matrix.
-
-        Assumes that the provided basis is a wide matrix with full rank.
-        """
-        pivot_cols = qldpc.math.first_nonzero_cols(basis.row_reduce())[: len(basis)]
-        linearly_independent_cols = basis[:, pivot_cols].view(type(basis))
-        dual_basis = type(basis).Zeros(basis.shape)
-        dual_basis[:, pivot_cols] = np.linalg.inv(linearly_independent_cols).T
-        return dual_basis.view(type(basis))
-
     def _get_center_embeddings(self) -> tuple[galois.FieldArray, galois.FieldArray]:
         r"""Construct embeddings of elements in the center Z(S) into GF(p^{kd}) ≅ GF(q^d).
 

src/qldpc/math.py

@@ -80,12 +80,38 @@ def symplectic_weight(vectors: npt.NDArray[np.int_]) -> int:
     return np.count_nonzero(vectors_x | vectors_z, axis=-1).reshape(vectors.shape[:-1])
 
 
-def first_nonzero_cols(matrix: npt.NDArray[np.generic]) -> npt.NDArray[np.int_]:
-    """Get the first nonzero column for every row in a matrix."""
-    _matrix = np.asanyarray(matrix)
+def first_nonzero_cols(
+    matrix: npt.NDArray[np.generic] | Sequence[npt.NDArray[np.generic]],
+) -> npt.NDArray[np.int_]:
+    """Get the first nonzero column for every row in a matrix.
+
+    If all columns are zero in a particular row, the column valus is set to the number of columns.
+    If the array has more than two dimensions, return for each "row" r, the first "column" c for
+    which array[r, c] is not all zero.
+    """
+    _matrix = np.asarray(matrix)
+    if _matrix.ndim < 2:
+        raise ValueError(f"Cannot identify nonzero columns in array with dimension {_matrix.ndim}")
     if _matrix.size == 0:
         return np.array([], dtype=int)
-    return np.argmax(_matrix.view(np.ndarray).astype(bool), axis=-1)
+    nonzero_mask = np.any(_matrix.view(np.ndarray).astype(bool), axis=tuple(range(2, _matrix.ndim)))
+    has_any_nonzero_in_row = np.any(nonzero_mask, axis=1)
+    first_nonzero_col_index = np.argmax(nonzero_mask, axis=1)
+    first_nonzero_col_index[~has_any_nonzero_in_row] = nonzero_mask.shape[1]
+    return first_nonzero_col_index.astype(np.int_, copy=False)
+
+
+def get_dual_basis(basis: galois.FieldArray, *, validate: bool = True) -> galois.FieldArray:
+    """Construct a dual basis, for which dual_basis @ basis.T = identity_matrix."""
+    if validate and (
+        basis.shape[0] > basis.shape[1] or np.linalg.matrix_rank(basis) != basis.shape[0]
+    ):
+        raise ValueError("A dual basis can only be found for wide matrices of full rank")
+    pivot_cols = first_nonzero_cols(basis.row_reduce())[: len(basis)]
+    linearly_independent_cols = basis[:, pivot_cols].view(type(basis))
+    dual_basis = np.zeros(basis.shape, dtype=int).view(type(basis))
+    dual_basis[:, pivot_cols] = np.linalg.inv(linearly_independent_cols).T
+    return dual_basis.view(type(basis))
 
 
 def block_matrix(

src/qldpc/math_test.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import galois
 import numpy as np
 import pytest
 import stim
@@ -43,12 +44,40 @@ def test_vectors() -> None:
     vectors_conj = np.array([[1, 0], [2, -1]], dtype=int)
     assert np.array_equal(qldpc.math.symplectic_weight(vectors), [1, 1])
     assert np.array_equal(qldpc.math.symplectic_conjugate(vectors), vectors_conj)
-
-    assert np.array_equal(qldpc.math.first_nonzero_cols(np.empty(0, dtype=int)), [])
     assert np.array_equal(qldpc.math.first_nonzero_cols(vectors), [1, 0])
     assert np.array_equal(qldpc.math.first_nonzero_cols(vectors_conj), [0, 0])
 
 
+def test_nonzero_cols() -> None:
+    """Edge cases in finding the pivot columns."""
+    with pytest.raises(ValueError, match="Cannot identify nonzero columns"):
+        empty_array = np.array([], dtype=int)
+        qldpc.math.first_nonzero_cols(empty_array)
+
+    empty_matrix = np.array([], ndmin=2, dtype=int)
+    assert qldpc.math.first_nonzero_cols(empty_matrix).size == 0
+
+    zero_matrix = np.zeros((1, 1), dtype=int)
+    assert np.array_equal(qldpc.math.first_nonzero_cols(zero_matrix), [1])
+
+    tensor = np.ones((1, 1, 1), dtype=int)
+    assert np.array_equal(qldpc.math.first_nonzero_cols(tensor), [0])
+
+
+def test_dual_basis(pytestconfig: pytest.Config) -> None:
+    """Construst dual bases."""
+    np.random.seed(pytestconfig.getoption("randomly_seed"))
+
+    field = galois.GF(2)
+    basis = field.Random((4, 5)).row_reduce()
+    basis = basis[qldpc.math.first_nonzero_cols(basis) < basis.shape[1]]
+    dual_basis = qldpc.math.get_dual_basis(basis)
+    assert np.array_equal(dual_basis @ basis.T, field.Identity(len(basis)))
+
+    with pytest.raises(ValueError, match="wide matrices of full rank"):
+        qldpc.math.get_dual_basis(field.Random((2, 1)))
+
+
 def test_block_matrix() -> None:
     eye = np.eye(2, dtype=float)
     zero = np.zeros_like(eye)