Fix TPMS surface discovery with Autograd 1.9

microgen/shape/surface_functions.py

@@ -3,6 +3,8 @@
 import autograd.numpy as np
 from autograd.numpy import cos, sin
 
+from ._types import Field
+
 
 def gyroid(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
     """Gyroid.
@@ -309,7 +311,7 @@ def split_p(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
     )
 
 
-def honeycomb_gyroid(x: float, y: float, _: float) -> float:
+def honeycomb_gyroid(x: np.ndarray, y: np.ndarray, _: np.ndarray) -> np.ndarray:
     """Honeycomb Gyroid.
 
     .. math::
@@ -332,7 +334,7 @@ def honeycomb_gyroid(x: float, y: float, _: float) -> float:
     return sin(x) * cos(y) + sin(y) + cos(x)
 
 
-def honeycomb_schwarz_p(x: float, y: float, _: float) -> float:
+def honeycomb_schwarz_p(x: np.ndarray, y: np.ndarray, _: np.ndarray) -> np.ndarray:
     """Honeycomb Schwarz P.
 
     .. math::
@@ -355,7 +357,7 @@ def honeycomb_schwarz_p(x: float, y: float, _: float) -> float:
     return cos(x) + cos(y)
 
 
-def honeycomb_schwarz_d(x: float, y: float, _: float) -> float:
+def honeycomb_schwarz_d(x: np.ndarray, y: np.ndarray, _: np.ndarray) -> np.ndarray:
     """Honneycomb Schwarz D.
 
     .. math::
@@ -378,7 +380,7 @@ def honeycomb_schwarz_d(x: float, y: float, _: float) -> float:
     return cos(x) * cos(y) + sin(x) * sin(y) + sin(x) * cos(y) + cos(x) * sin(y)
 
 
-def honeycomb_schoen_iwp(x: float, y: float, _: float) -> float:
+def honeycomb_schoen_iwp(x: np.ndarray, y: np.ndarray, _: np.ndarray) -> np.ndarray:
     """Honneycomb Schoen IWP.
 
     .. math::
@@ -401,7 +403,7 @@ def honeycomb_schoen_iwp(x: float, y: float, _: float) -> float:
     return cos(x) * cos(y) + cos(y) + cos(x)
 
 
-def honeycomb_lidinoid(x: float, y: float, _: float) -> float:
+def honeycomb_lidinoid(x: np.ndarray, y: np.ndarray, _: np.ndarray) -> np.ndarray:
     """Honeycomb Lidinoid.
 
     .. math::
@@ -425,3 +427,23 @@ def honeycomb_lidinoid(x: float, y: float, _: float) -> float:
     return 1.1 * (sin(2 * x) * cos(y) + sin(2 * y) * sin(x) + cos(x) * sin(y)) - (
         cos(2 * x) * cos(2 * y) + cos(2 * y) + cos(2 * x)
     )
+
+
+BUILTIN_SURFACES: tuple[Field, ...] = (
+    gyroid,
+    schwarz_p,
+    schwarz_d,
+    neovius,
+    schoen_iwp,
+    schoen_frd,
+    fischer_koch_s,
+    pmy,
+    honeycomb,
+    lidinoid,
+    split_p,
+    honeycomb_gyroid,
+    honeycomb_schwarz_p,
+    honeycomb_schwarz_d,
+    honeycomb_schoen_iwp,
+    honeycomb_lidinoid,
+)

tests/shapes/test_tpms.py

@@ -3,15 +3,17 @@
 from __future__ import annotations
 
 import re
-from inspect import getmembers, isfunction, signature
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 import numpy.typing as npt
 import pytest
 
 import microgen
 
+if TYPE_CHECKING:
+    from microgen.shape._types import Field
+
 # ruff: noqa: S101 assert https://docs.astral.sh/ruff/rules/assert/
 # ruff: noqa: E501 line-too-long https://docs.astral.sh/ruff/rules/line-too-long/
 
@@ -20,19 +22,6 @@
 TEST_DEFAULT_SURFACE_FUNCTION = microgen.surface_functions.gyroid
 
 
-def _get_microgen_surface_functions() -> list[str]:
-    """List the actual TPMS surface functions in microgen.surface_functions.
-
-    Filters out non-TPMS callables exposed via re-import (autograd ``cos``/``sin``)
-    that take fewer than 3 args and would crash when invoked as ``f(x,y,z)``.
-    """
-    return [
-        name
-        for name, fn in getmembers(microgen.surface_functions, isfunction)
-        if len(signature(fn).parameters) == 3
-    ]
-
-
 @pytest.mark.parametrize("type_part", ["lower skeletal", "upper skeletal", "sheet"])
 def test_tpms_given_cadquery_vtk_shapes_volume_must_be_equivalent(
     type_part: Literal["sheet", "lower skeletal", "upper skeletal"],
@@ -89,19 +78,20 @@ def test_tpms_given_non_default_cell_size_and_repeat_cell_must_have_same_volume_
 
 @pytest.mark.parametrize(
     "surface",
-    _get_microgen_surface_functions(),
+    microgen.surface_functions.BUILTIN_SURFACES,
+    ids=lambda function: function.__name__,
 )
 @pytest.mark.parametrize("repeat_cell", [2, (2, 1, 3)])
 @pytest.mark.parametrize("cell_size", [3.0, (0.5, 1.5, 1.0)])
 def test_tpms_given_sum_volume_must_be_cube_volume(
-    surface: str,
+    surface: Field,
     repeat_cell: int | tuple[int, int, int],
     cell_size: float | tuple[float, float, float],
 ) -> None:
     """Test for the volume of the TPMS shapes generated with CadQuery and VTK."""
     # Arrange
     tpms = microgen.Tpms(
-        surface_function=getattr(microgen.surface_functions, surface),
+        surface_function=surface,
         offset=TEST_DEFAULT_OFFSET,
         repeat_cell=repeat_cell,
         cell_size=cell_size,
@@ -117,16 +107,20 @@ def test_tpms_given_sum_volume_must_be_cube_volume(
     assert np.isclose(volume, cube_volume, rtol=1e-2)
 
 
-@pytest.mark.parametrize("surface", _get_microgen_surface_functions())
+@pytest.mark.parametrize(
+    "surface",
+    microgen.surface_functions.BUILTIN_SURFACES,
+    ids=lambda function: function.__name__,
+)
 @pytest.mark.parametrize("density", [0.05, 0.5, 0.99, 1.0])
 def test_tpms_given_density_must_match_computed_density(
-    surface: str,
+    surface: Field,
     density: float,
 ) -> None:
     """Test for the density of the TPMS shapes generated with CadQuery and VTK."""
     # Arrange
     tpms = microgen.Tpms(
-        surface_function=getattr(microgen.surface_functions, surface),
+        surface_function=surface,
         density=density,
     )