invrs-opt 0.5.2__tar.gz → 0.7.0__tar.gz

{invrs_opt-0.5.2/src/invrs_opt.egg-info → invrs_opt-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.5.2
+Version: 0.7.0
 Summary: Algorithms for inverse design
 Author-email: "Martin F. Schubert" <mfschubert@gmail.com>
 Maintainer-email: "Martin F. Schubert" <mfschubert@gmail.com>
@@ -34,6 +34,7 @@ Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
 Requires-Dist: requests
+Requires-Dist: optax
 Requires-Dist: scipy
 Requires-Dist: totypes
 Requires-Dist: types-requests
@@ -49,7 +50,7 @@ Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.2`
+`v0.7.0`
 
 ## Overview
 

{invrs_opt-0.5.2 → invrs_opt-0.7.0}/README.md

@@ -1,5 +1,5 @@
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.2`
+`v0.7.0`
 
 ## Overview
 

{invrs_opt-0.5.2 → invrs_opt-0.7.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_opt"
-version = "v0.5.2"
+version = "v0.7.0"
 description = "Algorithms for inverse design"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"
@@ -20,6 +20,7 @@ dependencies = [
     "jaxlib",
     "numpy",
     "requests",
+    "optax",
     "scipy",
     "totypes",
     "types-requests",

invrs_opt-0.7.0/src/invrs_opt/__init__.py

@@ -0,0 +1,21 @@
+"""invrs_opt - Optimization algorithms for inverse design.
+
+Copyright (c) 2023 The INVRS-IO authors.
+"""
+
+__version__ = "v0.7.0"
+__author__ = "Martin F. Schubert <mfschubert@gmail.com>"
+
+from invrs_opt import parameterization as parameterization
+
+from invrs_opt.optimizers.lbfgsb import (
+    density_lbfgsb as density_lbfgsb,
+    lbfgsb as lbfgsb,
+    levelset_lbfgsb as levelset_lbfgsb,
+)
+
+from invrs_opt.optimizers.wrapped_optax import (
+    density_wrapped_optax as density_wrapped_optax,
+    levelset_wrapped_optax as levelset_wrapped_optax,
+    wrapped_optax as wrapped_optax,
+)

{invrs_opt-0.5.2 → invrs_opt-0.7.0}/src/invrs_opt/experimental/client.py

@@ -4,15 +4,14 @@ Copyright (c) 2023 The INVRS-IO authors.
 """
 
 import json
-import requests
 import time
 from typing import Any, Dict, Optional
 
+import requests
 from totypes import json_utils
 
-from invrs_opt import base
 from invrs_opt.experimental import labels
-
+from invrs_opt.optimizers import base
 
 PyTree = Any
 StateToken = str

{invrs_opt-0.5.2/src/invrs_opt → invrs_opt-0.7.0/src/invrs_opt/optimizers}/base.py

@@ -4,8 +4,12 @@ Copyright (c) 2023 The INVRS-IO authors.
 """
 
 import dataclasses
+import inspect
 from typing import Any, Protocol
 
+import optax  # type: ignore[import-untyped]
+from totypes import json_utils
+
 PyTree = Any
 
 
@@ -44,3 +48,9 @@ class Optimizer:
     init: InitFn
     params: ParamsFn
     update: UpdateFn
+
+
+# Register all optax state types for serialization.
+for name, obj in inspect.getmembers(optax):
+    if name.endswith("State") and isinstance(obj, type):
+        json_utils.register_custom_type(obj)