Resolve #528

scico/optimize/_admmaux.py

@@ -404,19 +404,24 @@ def internal_init(self, admm: soa.ADMM):
                     f"or scico.linop.Identity; got {type(admm.f.A)}."
                 )
 
+        self.ndims = admm.f.A.ndims if isinstance(admm.f.A, CircularConvolve) else None
         super().internal_init(admm)
 
         self.real_result = is_real_dtype(admm.C_list[0].input_dtype)
 
         # All of the C operators are assumed to be linear and shift invariant
         # but this is not checked.
         lhs_op_list = [
-            rho * CircularConvolve.from_operator(C.gram_op)
+            rho * CircularConvolve.from_operator(C.gram_op, ndims=self.ndims)
             for rho, C in zip(admm.rho_list, admm.C_list)
         ]
         A_lhs = reduce(lambda a, b: a + b, lhs_op_list)
         if self.admm.f is not None:
-            A_lhs += 2.0 * admm.f.scale * CircularConvolve.from_operator(admm.f.A.gram_op)
+            A_lhs += (
+                2.0
+                * admm.f.scale
+                * CircularConvolve.from_operator(admm.f.A.gram_op, ndims=self.ndims)
+            )
 
         self.A_lhs = A_lhs
 

scico/test/optimize/test_admm.py

@@ -363,23 +363,29 @@ def test_admm_quadratic_matrix(self):
         assert (snp.linalg.norm(self.grdA(x) - self.grdb) / snp.linalg.norm(self.grdb)) < 1e-5
 
 
+@pytest.mark.parametrize("extra_axis", (False, True))
+@pytest.mark.parametrize("center", (None, [-1.0, 2.5]))
 class TestCircularConvolveSolve:
-    def setup_method(self, method):
+
+    @pytest.fixture(scope="function", autouse=True)
+    def setup_and_teardown(self, extra_axis, center):
         np.random.seed(12345)
         Nx = 8
-        x = np.pad(np.ones((Nx, Nx), dtype=np.float32), Nx)
+        x = snp.pad(snp.ones((Nx, Nx), dtype=np.float32), Nx)
         Npsf = 3
         psf = snp.ones((Npsf, Npsf), dtype=np.float32) / (Npsf**2)
+        if extra_axis:
+            x = x[np.newaxis]
+            psf = psf[np.newaxis]
         self.A = linop.CircularConvolve(
-            h=psf,
-            input_shape=x.shape,
-            input_dtype=np.float32,
+            h=psf, input_shape=x.shape, ndims=2, input_dtype=np.float32, h_center=center
         )
         self.y = self.A(x)
         λ = 1e-2
         self.f = loss.SquaredL2Loss(y=self.y, A=self.A)
         self.g_list = [λ * functional.L1Norm()]
         self.C_list = [linop.FiniteDifference(input_shape=x.shape, circular=True)]
+        yield
 
     def test_admm(self):
         maxiter = 50
@@ -406,6 +412,7 @@ def test_admm(self):
             x0=self.A.adj(self.y),
             subproblem_solver=CircularConvolveSolver(),
         )
+        assert admm_dft.subproblem_solver.A_lhs.ndims == 2
         x_dft = admm_dft.solve()
         np.testing.assert_allclose(x_dft, x_lin, atol=1e-4, rtol=0)
         assert metric.mse(x_lin, x_dft) < 1e-9