Merge pull request #31 from invrs-io/sym

Fix batch dims for symmetry

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v0.6.3"
+current_version = "v0.6.4"
 commit = true
 commit_args = "--no-verify"
 tag = true

README.md

@@ -1,5 +1,5 @@
 # totypes - Custom types for topology optimization
-`v0.6.3`
+`v0.6.4`
 
 ## Overview
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "totypes"
-version = "v0.6.3"
+version = "v0.6.4"
 description = "Custom datatypes useful in a topology optimization context"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"

src/totypes/__init__.py

@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.6.3"
+__version__ = "v0.6.4"
 __author__ = "Martin F. Schubert <[email protected]>"
 
 __all__ = ["json_utils", "symmetry", "types"]

src/totypes/symmetry.py

@@ -36,13 +36,13 @@ def symmetrize(
 def _reflection_ne_sw(array: jnp.ndarray) -> jnp.ndarray:
     """Transform `array` to have reflection symmetry about the ne-sw axis."""
     assert array.shape[-2] == array.shape[-1]
-    return (array + array[..., ::-1, ::-1].T) / 2
+    return (array + jnp.swapaxes(array[..., ::-1, ::-1], -2, -1)) / 2
 
 
 def _reflection_nw_se(array: jnp.ndarray) -> jnp.ndarray:
     """Transform `array` to have reflection symmetry about the nw-se axis."""
     assert array.shape[-2] == array.shape[-1]
-    return (array + array.T) / 2
+    return (array + jnp.swapaxes(array, -2, -1)) / 2
 
 
 def _reflection_n_s(array: jnp.ndarray) -> jnp.ndarray:

tests/test_symmetry.py

@@ -147,3 +147,17 @@ def test_multiple_symmetry(self, obj):
             ]
         )
         _assert_array_equal(result, expected)
+
+    @parameterized.expand(
+        [
+            [symmetry.REFLECTION_E_W],
+            [symmetry.REFLECTION_N_S],
+            [symmetry.REFLECTION_NE_SW],
+            [symmetry.REFLECTION_NW_SE],
+            [symmetry.ROTATION_180],
+            [symmetry.ROTATION_90],
+        ]
+    )
+    def test_with_batch(self, sym):
+        arr = jnp.ones((1, 10, 13, 13))
+        symmetry.symmetrize(arr, (sym,))