Merge pull request #215 from sparks-baird/lower-tri

fixup lower_tri test
sparks-baird · Jul 1, 2023 · b34d0a2 · b34d0a2
2 parents 3c802a7 + 645e769
commit b34d0a2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/customization_test.py b/tests/customization_test.py
@@ -119,8 +119,10 @@ def test_distance_mask():
 def test_lower_tri_mask():
     xc = XtalConverter(mask_types=["lower_tri"])
     imgs = xc.xtal2png(example_structures)
-    if not np.all(xc.data[np.tril(xc.data[0, 0])] == 0):
-        raise ValueError("Lower triangle mask not applied correctly")
+    for d in xc.data:
+        for sq in d:
+            if not np.all(sq[np.tril_indices(sq.shape[0])] == 0):
+                raise ValueError("Lower triangle mask not applied correctly")
 
     return imgs