QoL, tidying and bug fixing

Added optimizer packages to pyproject
Made estimate_cost robust to errors like Batch.estimate_cost
Added console to design.py which gets passed to methods
Optimizer methods now output progress to console.log and can be disabled with verbose
Optimizers now try and except relevant module loading
Tests for all new additions and tidying of code

pyproject.toml

@@ -73,6 +73,11 @@ jax = { version = "0.4.25", extras = [
   "cpu",
 ], source = "jaxsource", optional = true }
 
+# design
+bayesian-optimization = { version = "*", optional = true }
+pygad = { version = "*", optional = true }
+pyswarms = { version = "*", optional = true }
+
 # scikit-rf
 scikit-rf = { version = "*", optional = true }
 
@@ -104,7 +109,6 @@ signac = { version = "*", optional = true }
 flax = { version = ">=0.8.2", optional = true }
 sax = { version = "^0.11", optional = true }
 vtk = { version = ">=9.2.6", optional = true }
-pyswarms = { version = "*", optional = true }
 sphinxemoji = { version = "*", optional = true }
 devsim = { version = "*", optional = true }
 cma = { version = "*", optional = true }

tests/test_plugins/test_design.py

@@ -1,5 +1,7 @@
 """Test the parameter sweep plugin."""
 
+import sys
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
@@ -14,12 +16,12 @@
     grid=tdd.MethodGrid(),
     monte_carlo=tdd.MethodMonteCarlo(num_points=3, rng_seed=1),
     custom=tdd.MethodRandomCustom(
-        num_points=10, sampler=qmc.Halton, sampler_kwargs={"d": 3, "seed": 1}
+        num_points=5, sampler=qmc.Halton, sampler_kwargs={"d": 3, "seed": 1}
     ),
-    random=tdd.MethodRandom(num_points=5, rng_seed=1),  # TODO: remove this if not used
+    random=tdd.MethodRandom(num_points=3, rng_seed=1),  # TODO: remove this if not used
     bay_opt=tdd.MethodBayOpt(initial_iter=2, n_iter=2, rng_seed=1),
     gen_alg=tdd.MethodGenAlg(solutions_per_pop=4, n_generations=2, n_parents_mating=2, rng_seed=1),
-    part_swarm=tdd.MethodParticleSwarm(n_particles=5, n_iter=2, rng_seed=1),
+    part_swarm=tdd.MethodParticleSwarm(n_particles=3, n_iter=2, rng_seed=1),
 )
 
 
@@ -44,6 +46,16 @@ def __getitem__(self, task_name):
     return BatchDataEmulated(data_dict=data_dict, task_ids=task_ids, task_paths=task_paths)
 
 
+def emulated_estimate_cost(self, verbose=True):
+    # Test both value and failed None returns
+    value = np.random.random()
+
+    if value < 0.5:
+        return value
+    else:
+        return None
+
+
 @pytest.mark.parametrize("sweep_method", SWEEP_METHODS.values())
 def test_sweep(sweep_method, monkeypatch):
     # Problem, simulate scattering cross section of sphere ensemble
@@ -55,6 +67,8 @@ def test_sweep(sweep_method, monkeypatch):
 
     monkeypatch.setattr(web.Batch, "run", emulated_batch_run)
 
+    monkeypatch.setattr(web.Job, "estimate_cost", emulated_estimate_cost)
+
     # STEP1: define your design function (inputs and outputs)
 
     # Non td function testing
@@ -253,7 +267,7 @@ def scs_post_dict_const(sim_dict):
 
     # Check some summaries
     design_space.summary()
-    # design_space.summary(fn_pre=scs_pre)
+    design_space.summary(fn_pre=scs_pre)
 
     # STEP3: Run your design problem
 
@@ -276,9 +290,7 @@ def scs_post_dict_const(sim_dict):
 
     # Try functions that include td objects
     # Test that estimate_cost outputs and fails for combined function output
-    # estimate = design_space.estimate_cost(scs_pre)
-
-    # assert estimate > 0
+    estimate = design_space.estimate_cost(scs_pre)
 
     with pytest.raises(ValueError):
         estimate = design_space.estimate_cost(scs_combined)
@@ -290,10 +302,9 @@ def scs_post_dict_const(sim_dict):
     assert td_sweep1.values == td_sweep2.values
 
     # Try with batch output from pre
-    # estimate = design_space.estimate_cost(scs_pre_batch)
+    estimate = design_space.estimate_cost(scs_pre_batch)
     td_batch = design_space.run(scs_pre_batch, scs_post_batch)
-
-    # assert estimate > 0
+    td_batch_run_batch = design_space.run_batch(scs_pre_batch, scs_post_batch, path_dir="")
 
     # # Test user specified batching works with combined function
     td_batch_combined = design_space.run(scs_combined_batch)
@@ -306,6 +317,7 @@ def scs_post_dict_const(sim_dict):
     assert "0_0" not in td_sim_list.task_ids[0] and "0_3" not in td_sim_list.task_ids[4]
 
     # Test with dict of sims
+    estimate = design_space.estimate_cost(scs_pre_dict)
     td_sim_dict = design_space.run(scs_pre_dict, scs_post_dict)
 
     # Check naming is including dict keys
@@ -457,6 +469,13 @@ def float_non_td_complex_return_post(res):
     assert ts_sim_complex_df["test4"][0] == 3.14
 
 
+method_module_convert = {
+    "MethodBayOpt": "bayes_opt",
+    "MethodGenAlg": "pygad",
+    "MethodParticleSwarm": "pyswarms.single.global_best",
+}
+
+
 @pytest.mark.parametrize(
     "sweep_method",
     [SWEEP_METHODS["bay_opt"], SWEEP_METHODS["gen_alg"], SWEEP_METHODS["part_swarm"]],
@@ -563,6 +582,13 @@ def bad_format_non_td_aux_post(res):
     # Test that plots have been produced and stored
     assert ts_sim_aux.opt_output is not None
 
+    # Create an import error to test optimizer can error if relevant package is not installed
+    # Placed at the end of the test so that module is loaded for other checks
+
+    with pytest.raises(ImportError):
+        sys.modules[method_module_convert[sweep_method.type]] = None
+        import_fail = design_space.run(float_non_td_pre, float_non_td_aux_post)
+
 
 def test_method_custom_validators():
     """Test the MethodRandomCustom validation performs as expected."""

tidy3d/plugins/design/design.py

@@ -49,6 +49,7 @@ class DesignSpace(Tidy3dBaseModel):
         ...,
         title="Search Type",
         description="Specifications for the procedure used to explore the parameter space.",
+        discriminator="type",  # Stops pydantic trying to validate every method whilst checking MethodType
     )
 
     name: str = pd.Field(None, title="Name", description="Optional name for the design space.")
@@ -111,7 +112,7 @@ def get_fn_source(function: Callable) -> str:
         except (TypeError, OSError):
             return None
 
-    def run(self, fn: Callable, fn_post: Callable = None) -> Result:
+    def run(self, fn: Callable, fn_post: Callable = None, verbose: bool = True) -> Result:
         """Run the design problem on a user defined function of the design parameters.
 
         Parameters
@@ -127,14 +128,18 @@ def run(self, fn: Callable, fn_post: Callable = None) -> Result:
             Can be converted to ``pandas.DataFrame`` with ``.to_dataframe()``.
         """
 
+        # Get the console
+        # Method.run checks for console is None instead of being passed console and verbose
+        console = get_logging_console() if verbose else None
+
         # Run based on how many functions the user provides
         if fn_post is None:
-            fn_args, fn_values, aux_values, opt_output = self.run_single(fn)
+            fn_args, fn_values, aux_values, opt_output = self.run_single(fn, console)
             sim_names = None
 
         else:
             fn_args, fn_values, aux_values, opt_output, sim_names = self.run_pre_post(
-                fn_pre=fn, fn_post=fn_post
+                fn_pre=fn, fn_post=fn_post, console=console
             )
 
             sim_names = tuple(sim_names)
@@ -154,20 +159,20 @@ def run(self, fn: Callable, fn_post: Callable = None) -> Result:
             opt_output=opt_output,
         )
 
-    def run_single(self, fn: Callable) -> Tuple(list[dict], list, list[Any]):
+    def run_single(self, fn: Callable, console) -> Tuple(list[dict], list, list[Any]):
         """Run a single function of parameter inputs."""
         evaluate_fn = self._get_evaluate_fn_single(fn=fn)
-        return self.method.run(run_fn=evaluate_fn, parameters=self.parameters)
+        return self.method.run(run_fn=evaluate_fn, parameters=self.parameters, console=console)
 
-    def run_pre_post(self, fn_pre: Callable, fn_post: Callable) -> Tuple(
+    def run_pre_post(self, fn_pre: Callable, fn_post: Callable, console) -> Tuple(
         list[dict], list[dict], list[Any]
     ):
         """Run a function with tidy3d implicitly called in between."""
         handler = self._get_evaluate_fn_pre_post(
             fn_pre=fn_pre, fn_post=fn_post, fn_mid=self._fn_mid
         )
         fn_args, fn_values, aux_values, opt_output = self.method.run(
-            run_fn=handler.evaluate, parameters=self.parameters
+            run_fn=handler.evaluate, parameters=self.parameters, console=console
         )
         return fn_args, fn_values, aux_values, opt_output, handler.sim_names
 
@@ -362,6 +367,7 @@ def estimate_cost(self, fn_pre: Callable) -> float:
         """
         Compute the maximum FlexCredit charge for the entire design space computation.
         Require a pre function that should return a Simulation object, Batch object, or collection of either.
+        The pre function is called to estimate the cost - complicated pre functions may cause long runtimes.
         """
         # Get output fn_pre for paramters at the lowest span / default
         arg_dict = {}
@@ -378,8 +384,7 @@ def _estimate_sim_cost(sim):
             job = Job(simulation=sim, task_name="estimate_cost")
 
             estimate = job.estimate_cost()
-
-            job.delete()
+            job.delete()  # Deleted as only a test with initial parameters
 
             return estimate
 
@@ -388,6 +393,7 @@ def _estimate_sim_cost(sim):
 
         elif isinstance(pre_out, Batch):
             per_run_estimate = pre_out.estimate_cost()
+            pre_out.delete()  # Deleted as only a test with initial parameters
 
         elif isinstance(pre_out, (list, dict)):
             # Iterate through container to get simulations and batches and sum cost
@@ -404,20 +410,30 @@ def _estimate_sim_cost(sim):
                 elif isinstance(value, Batch):
                     batches.append(value)
 
-            per_run_estimate = 0
+            calculated_estimates = []
             for sim in sims:
-                per_run_estimate += _estimate_sim_cost(sim)
+                calculated_estimates.append(_estimate_sim_cost(sim))
 
             for batch in batches:
-                per_run_estimate += batch.estimate_cost()
+                calculated_estimates.append(batch.estimate_cost())
+                batch.delete()  # Deleted as only a test with initial parameters
+
+            if None in calculated_estimates:
+                per_run_estimate = None
+            else:
+                per_run_estimate = sum(calculated_estimates)
 
         else:
             raise ValueError("Unrecognised output from pre-function, unable to estimate cost.")
 
         # Calculate maximum number of runs for different methods
         run_count = self._get_run_count()
 
-        return round(per_run_estimate * run_count, 3)
+        # For if tidy3d server cannot determine the estimate
+        if per_run_estimate is None:
+            return None
+        else:
+            return round(per_run_estimate * run_count, 3)
 
     def summary(self, fn_pre: Callable = None):
         """Summarise the setup of the DesignSpace"""
@@ -465,11 +481,11 @@ def summary(self, fn_pre: Callable = None):
             if any(isinstance(param, ParameterInt) for param in self.parameters):
                 if any(isinstance(param, ParameterAny) for param in self.parameters):
                     notes.append(
-                        "Discrete ParameterAny values are automatically converted to Int values to be optimized.\n"
+                        "Discrete ParameterAny values are automatically converted to int values to be optimized.\n"
                     )
 
                 notes.append(
-                    "Discrete Int values are automatically rounded as optimizers generate float predictions.\n"
+                    "Discrete int values are automatically rounded as optimizers generate float predictions.\n"
                 )
 
         if len(notes) > 0: