From f315abf0d9c3229a77824f23c73a0e63fcd80adc Mon Sep 17 00:00:00 2001
From: Tim Widrick <twmacro@gmail.com>
Date: Sat, 2 Sep 2023 17:12:38 -0400
Subject: [PATCH] in ytools.compmat: raise ValueError for arrays > 2 dimensions

---
 pyyeti/tests/test_ytools.py | 5 +++++
 pyyeti/ytools.py            | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/pyyeti/tests/test_ytools.py b/pyyeti/tests/test_ytools.py
index c92349b..7c1a409 100644
--- a/pyyeti/tests/test_ytools.py
+++ b/pyyeti/tests/test_ytools.py
@@ -219,9 +219,14 @@ def test_compmat():
 
     with pytest.raises(ValueError):
         ytools.compmat(A, A[:2, :2])
+
     with pytest.raises(ValueError):
         ytools.compmat(A, B, method="badmethod")
 
+    a = np.zeros((3, 3, 3))
+    with pytest.raises(ValueError):
+        ytools.compmat(a, a)
+
 
 def test_max_complex_vector_sum():
     x = [3.0, 1.0 + 2.0j]
diff --git a/pyyeti/ytools.py b/pyyeti/ytools.py
index dccfcf0..46e7bc4 100644
--- a/pyyeti/ytools.py
+++ b/pyyeti/ytools.py
@@ -1568,6 +1568,8 @@ def compmat(a, b, filterval=0.0, method="abs", pdiff_tol=0, verbose=5):
     ------
     ValueError
         If `a` and `b` are different sizes
+    ValueError
+        If number of dimensions of `a` and `b` > 2
 
     Notes
     -----
@@ -1645,6 +1647,9 @@ def compmat(a, b, filterval=0.0, method="abs", pdiff_tol=0, verbose=5):
             f"matrix sizes do not match: a.shape={a.shape}, b.shape={b.shape}"
         )
 
+    if a.ndim > 2:
+        raise ValueError(f"arrays have {a.ndim} dimensions; must be <= 2")
+
     if np.iscomplexobj(a) ^ np.iscomplexobj(b):
         # if only one is complex:
         if np.isrealobj(a):