Merge pull request #60 from invrs-io/fix

Serialize fmmax types

.bumpversion.toml

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

README.md

@@ -1,5 +1,5 @@
 # invrs-gym - A collection of inverse design challenges
-`v0.3.0`
+`v0.3.1`
 
 ## Overview
 The `invrs_gym` package is an open-source gym containing a diverse set of photonic design challenges, which are relevant for a wide range of applications such as AR/VR, optical networking, LIDAR, and others.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_gym"
-version = "v0.3.0"
+version = "v0.3.1"
 description = "A collection of inverse design challenges"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"

src/invrs_gym/__init__.py

@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.3.0"
+__version__ = "v0.3.1"
 __author__ = "Martin F. Schubert <[email protected]>"
 
 from invrs_gym import challenges as challenges

src/invrs_gym/challenges/base.py

@@ -6,9 +6,10 @@
 import abc
 from typing import Any, Callable, Dict, Tuple
 
+import fmmax
 import jax
 import jax.numpy as jnp
-from totypes import types
+from totypes import json_utils, types
 
 AuxDict = Dict[str, Any]
 PyTree = Any
@@ -74,3 +75,8 @@ def distance_to_target(self, response: Any) -> jnp.ndarray:
     @abc.abstractmethod
     def metrics(self, response: Any, params: PyTree, aux: AuxDict) -> AuxDict:
         """Compute metrics for a component response and associated quantities."""
+
+
+# Several challenges use the `fmmax` simulator, and contain `fmmax` custom objects in
+# their responses. Ensure these are serializable by registering with totypes.
+json_utils.register_custom_type(fmmax.basis.Expansion)