"""Defines design space specification for tidy3d."""
from __future__ import annotations
import inspect
from typing import Any, Callable, Dict, List, Tuple, Union
import pydantic.v1 as pd
from ...components.base import TYPE_TAG_STR, Tidy3dBaseModel, cached_property
from ...components.data.sim_data import SimulationData
from ...components.simulation import Simulation
from ...log import Console, get_logging_console, log
from ...web.api.container import Batch, BatchData, Job
from .method import (
MethodBayOpt,
MethodGenAlg,
MethodOptimize,
MethodParticleSwarm,
MethodType,
)
from .parameter import ParameterAny, ParameterInt, ParameterType
from .result import Result
[docs]
class DesignSpace(Tidy3dBaseModel):
"""Manages all exploration of a parameter space within specified parameters using a supplied search method.
The ``DesignSpace`` forms the basis of the ``Design`` plugin, and receives a ``Method`` and ``Parameter`` list that
define the scope of the design space and how it should be searched. ``DesignSpace.run()`` can then be called with
a function(s) to generate different solutions from parameters suggested by the ``Method``. The ``Method`` can either
sample the design space systematically or randomly, or can optimize for a given problem through an iterative search
and evaluate approach.
Notes
-----
Schematic outline of how to use the ``Design`` plugin to explore a design space.
.. image:: ../../../../_static/img/design.png
:width: 80%
:align: center
The `Design <https://www.flexcompute.com/tidy3d/examples/notebooks/Design/>`_ notebook contains an overview of the
``Design`` plugin and is the best place to learn how to get started.
Detailed examples using the ``Design`` plugin can be found in the following notebooks:
* `All-Dielectric Structural Colors <https://www.flexcompute.com/tidy3d/examples/notebooks/AllDielectricStructuralColor/>`_
* `Bayesian Optimization of Y-Junction <https://www.flexcompute.com/tidy3d/examples/notebooks/BayesianOptimizationYJunction/>`_
* `Genetic Algorithm Reflector <https://www.flexcompute.com/tidy3d/examples/notebooks/GeneticAlgorithmReflector/>`_
* `Particle Swarm Optimizer PBS <https://www.flexcompute.com/tidy3d/examples/notebooks/ParticleSwarmOptimizedPBS/>`_
* `Particle Swarm Optimizer Bullseye Cavity <https://www.flexcompute.com/tidy3d/examples/notebooks/BullseyeCavityPSO/>`_
Example
-------
>>> import tidy3d.plugins.design as tdd
>>> param = tdd.ParameterFloat(name="x", span=(0,1))
>>> method = tdd.MethodMonteCarlo(num_points=10)
>>> design_space = tdd.DesignSpace(parameters=[param], method=method)
>>> fn = lambda x: x**2
>>> result = design_space.run(fn)
>>> df = result.to_dataframe()
>>> im = df.plot()
"""
parameters: Tuple[ParameterType, ...] = pd.Field(
(),
title="Parameters",
description="Set of parameters defining the dimensions and allowed values for the design space.",
)
method: MethodType = pd.Field(
...,
title="Search Type",
description="Specifications for the procedure used to explore the parameter space.",
discriminator=TYPE_TAG_STR, # Stops pydantic trying to validate every method whilst checking MethodType
)
task_name: str = pd.Field(
"",
title="Task Name",
description="Task name assigned to tasks along with a simulation counter in the form of {task_name}_{sim_index}_{counter} where ``sim_index`` is "
"the index of the ``Simulation`` from the pre function output. "
"If the pre function outputs a dictionary the key will be included in the task name as {task_name}_{dict_key}_{counter}. "
"Only used when pre-post functions are supplied.",
)
name: str = pd.Field(None, title="Name", description="Optional name for the design space.")
path_dir: str = pd.Field(
".",
title="Path Directory",
description="Directory where simulation data files will be locally saved to. Only used when pre and post functions are supplied.",
)
folder_name: str = pd.Field(
"default",
title="Folder Name",
description="Folder path where the simulation will be uploaded in the Tidy3D Workspace. Will use 'default' if no path is set.",
)
@cached_property
def dims(self) -> Tuple[str]:
"""dimensions defined by the design parameter names."""
return tuple(param.name for param in self.parameters)
def _package_run_results(
self,
fn_args: list[dict[str, Any]],
fn_values: List[Any],
fn_source: str,
task_names: Tuple[str] = None,
task_paths: list = None,
aux_values: List[Any] = None,
opt_output: Any = None,
) -> Result:
"""How to package results from ``method.run`` and ``method.run_batch``"""
fn_args_coords = tuple(
[[arg_dict[key] for arg_dict in fn_args] for key in fn_args[0].keys()]
)
fn_args_coords_T = list(map(list, zip(*fn_args_coords)))
return Result(
dims=tuple(fn_args[0].keys()),
values=fn_values,
coords=fn_args_coords_T,
fn_source=fn_source,
task_names=task_names,
task_paths=task_paths,
aux_values=aux_values,
optimizer=opt_output,
)
[docs]
@staticmethod
def get_fn_source(function: Callable) -> str:
"""Get the function source as a string, return ``None`` if not available."""
try:
return inspect.getsource(function)
except (TypeError, OSError):
return None
[docs]
def run(self, fn: Callable, fn_post: Callable = None, verbose: bool = True) -> Result:
"""Explore a parameter space with a supplied method using the user supplied function.
Supplied functions are used to evaluate the design space and are called within the method.
For optimization methods these functions act as the fitness function. A single function can be
supplied which will contain the preprocessing, computation, and analysis of the desired problem.
If running a Tidy3D simulation, it is recommended to split this function into a pre function, that creates a Simulation object(s),
and a post function which analyses the SimulationData produced by the pre function Simulations. This allows the DesignSpace to
manage the batching of Simulations, which varies between Method used, and saving time writing their own batching code.
It also efficiently submits simulations to the cloud servers allowing for the fastest exploration of a design space.
The ``fn`` function must take a dictionary input - this can be stored a dictionary ``def example_fn(**params)``
or left as keyword arguments ``def example_fn(arg1, arg2)`` where the keywords correspond to the ``name`` of the parameters in the design space.
If used as a pre function, the output of ``fn`` must be a float, ``Simulation``, ``Batch``, list, or dict. Supplied ``Batch`` objects are
run without modification and are run in series. A list or dict of ``Simulation`` objects is flattened into a single ``Batch`` to enable
parallel computation on the cloud. The original structure is then restored for output; all ``Simulation`` objects are replaced by ``SimulationData`` objects.
Example pre return formats and associated post inputs can be seen in the table below.
.. list-table:: Pre return formats and post input formats
:widths: 50 50
:header-rows: 1
* - fn_pre return
- fn_post call
* - 1.0
- fn_post(1.0)
* - [1,2,3]
- fn_post(1,2,3)
* - {'a': 2, 'b': 'hi'}
- fn_post(a=2, b='hi')
* - Simulation
- fn_post(SimulationData)
* - Batch
- fn_post(BatchData)
* - [Simulation, Simulation]
- fn_post(SimulationData, SimulationData)
* - [Simulation, 1.0]
- fn_post(SimulationData, 1.0)
* - [Simulation, Batch]
- fn_post(SimulationData, BatchData)
* - {'a': Simulation, 'b': Batch, 'c': 2.0}
- fn_post(a=SimulationData, b=BatchData, c=2.0)
The output of ``fn_post`` (or ``fn`` if only one function is supplied) must be a float
or a container where the first element is a ``float`` and second element is a ``list`` / ``dict`` e,g. [float {"aux_1": str}].
The float is used by the optimizers as the return of the fitness function.
The second element is for auxiliary data from the analysis that the user may want to keep.
Sampling methods (``MethodGrid`` or ``MethodMonteCarlo``) can have any return type.
Parameters
----------
fn : Callable
Function accepting arguments that correspond to the ``name`` fields
of the ``DesignSpace.parameters``.
Must return in the expected format for the ``method`` used in ``DesignSpace``,
or return an object that fn_post can accept as an input.
fn_post : Callable = None
Optional function performing postprocessing on the output of ``fn``.
It is recommended to supply fn_post when working with Simulation objects.
Must return in the expected format for the ``method`` used in ``DesignSpace``.
verbose : bool = True
Toggle the output of statements stored in the logging console.
Returns
-------
:class:`Result`
Object containing the results of the design space exploration.
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, console)
sim_names = None
sim_paths = None
else:
fn_args, fn_values, aux_values, opt_output, sim_names, sim_paths = self.run_pre_post(
fn_pre=fn, fn_post=fn_post, console=console
)
if len(sim_names) == 0:
sim_names = None
fn_source = self.get_fn_source(fn)
# Package the result
return self._package_run_results(
fn_args=fn_args,
fn_values=fn_values,
fn_source=fn_source,
aux_values=aux_values,
task_names=sim_names,
task_paths=sim_paths,
opt_output=opt_output,
)
[docs]
def run_single(self, fn: Callable, console: 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, console=console)
[docs]
def run_pre_post(self, fn_pre: Callable, fn_post: Callable, console: 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, console=console
)
fn_args, fn_values, aux_values, opt_output = self.method._run(
run_fn=handler.fn_combined, parameters=self.parameters, console=console
)
return fn_args, fn_values, aux_values, opt_output, handler.sim_names, handler.sim_paths
""" Helpers """
def _get_evaluate_fn_single(self, fn: Callable) -> list[Any]:
"""Get function that sequentially evaluates single `fn` for a list of arguments."""
def evaluate(args_list: list) -> list[Any]:
"""Evaluate a list of arguments passed to ``fn``."""
return [fn(**args) for args in args_list]
return evaluate
def _get_evaluate_fn_pre_post(
self, fn_pre: Callable, fn_post: Callable, fn_mid: Callable, console: Console
):
"""Get function that tries to use batch processing on a set of arguments."""
class Pre_Post_Handler:
def __init__(self, console):
self.sim_counter = 0
self.sim_names = []
self.sim_paths = []
self.console = console
def fn_combined(self, args_list: list[dict[str, Any]]) -> list[Any]:
"""Compute fn_pre and fn_post functions and capture other outputs."""
sim_dict = {str(idx): fn_pre(**arg_list) for idx, arg_list in enumerate(args_list)}
if not all(
isinstance(val, (int, float, Simulation, Batch, list, dict))
for val in sim_dict.values()
):
raise ValueError(
"Unrecognized output of fn_pre. Please change the return of fn_pre."
)
data, task_names, task_paths, sim_counter = fn_mid(
sim_dict, self.sim_counter, self.console
)
self.sim_names.extend(task_names)
self.sim_paths.extend(task_paths)
self.sim_counter = sim_counter
post_out = [fn_post(val) for val in data.values()]
return post_out
handler = Pre_Post_Handler(console)
return handler
def _fn_mid(
self, pre_out: dict[int, Any], sim_counter: int, console: Console
) -> Union[dict[int, Any], BatchData]:
"""A function of the output of ``fn_pre`` that gives the input to ``fn_post``."""
# Keep copy of original to use if no tidy3d simulation required
original_pre_out = pre_out.copy()
# Convert list input to dict of dict before passing through checks
was_list = False
if all(isinstance(sim, list) for sim in pre_out.values()):
pre_out = {
str(idx): {str(sub_idx): sub_list for sub_idx, sub_list in enumerate(sim_list)}
for idx, sim_list in enumerate(pre_out.values())
}
was_list = True
def _find_and_map(
search_dict: dict,
search_type: Any,
output_dict: dict,
naming_dict: dict,
previous_key: str = "",
):
"""Recursively search for search_type objects within a dictionary."""
current_key = previous_key
for key, value in search_dict.items():
if not len(previous_key):
latest_key = str(key)
else:
latest_key = f"{current_key}_{key}"
if isinstance(value, dict):
_find_and_map(value, search_type, output_dict, naming_dict, latest_key)
elif isinstance(value, search_type):
output_dict[latest_key] = value
naming_dict[latest_key] = key
simulations = {}
batches = {}
naming_keys = {}
_find_and_map(pre_out, Simulation, simulations, naming_keys)
_find_and_map(pre_out, Batch, batches, naming_keys)
# Exit fn_mid here if no td computation is required
if not len(simulations) and not len(batches):
return original_pre_out, list(), list(), sim_counter
# Create task names for simulations
named_sims = {}
translate_sims = {}
for sim_key, sim in simulations.items():
# Checks stop standard indexing keys being included in the name
suffix = f"{naming_keys[sim_key]}_{sim_counter}"
# Handle if the user does not want a task name
if len(self.task_name) > 0:
sim_name = f"{self.task_name}_{suffix}"
else:
sim_name = suffix
named_sims[sim_name] = sim
translate_sims[sim_name] = sim_key
sim_counter += 1
# Log the simulations and batches for the user
if console is not None:
# Writen like this to include batches on the same line if present
run_statement = f"{len(named_sims)} Simulations"
if len(batches) > 0:
run_statement = run_statement + f" and {len(batches)} user Batches"
console.log(f"Running {run_statement}")
# Running simulations and batches
sims_out = Batch(
simulations=named_sims,
folder_name=self.folder_name,
simulation_type="tidy3d_design",
verbose=False, # Using a custom output instead of Batch.monitor updates
).run(path_dir=self.path_dir)
batch_results = {}
for batch_key, batch in batches.items():
batch_out = batch.run(path_dir=self.path_dir)
batch_results[batch_key] = batch_out
def _return_to_dict(return_dict: dict, key: str, return_obj: Any) -> None:
"""Recursively insert items into a dict by keys split with underscore. Only works for dict or dict of dict inputs."""
split_key = key.split("_", 1)
if len(split_key) > 1:
_return_to_dict(return_dict[split_key[0]], split_key[1], return_obj)
else:
return_dict[split_key[0]] = return_obj
for (sim_name, sim), task_name, task_path in zip(
sims_out.items(), sims_out.task_ids.keys(), sims_out.task_paths.values()
):
translated_name = translate_sims[sim_name]
sim.attrs["task_name"] = task_name
sim.attrs["task_path"] = task_path
_return_to_dict(pre_out, translated_name, sim)
for batch_name, batch in batch_results.items():
_return_to_dict(pre_out, batch_name, batch)
def _remove_or_replace(search_dict: dict, attr_name: str) -> dict:
"""Recursively search through a dict replacing Sims and Batches or ignoring other items thus removing them"""
new_dict = {}
for key, value in search_dict.items():
if isinstance(value, dict):
new_sub_dict = _remove_or_replace(value, attr_name)
new_dict[key] = new_sub_dict
else:
if isinstance(value, SimulationData):
new_dict[key] = value.attrs[attr_name]
elif isinstance(value, BatchData):
if attr_name == "task_name":
new_dict[key] = list(value.task_ids.keys())
else:
new_dict[key] = list(value.task_paths.values())
return new_dict
# Build out a dict of task_name or task_path in the same shape as the original data
task_names = _remove_or_replace(pre_out.copy(), "task_name")
task_paths = _remove_or_replace(pre_out.copy(), "task_path")
# Reduce down to a list to be extended later
task_names = list(task_names.values())
task_paths = list(task_paths.values())
# Restore output to a list if a list was supplied
if was_list:
pre_out = {dict_idx: list(sub_dict.values()) for dict_idx, sub_dict in pre_out.items()}
task_names = [list(sub_dict.values()) for sub_dict in task_names]
task_paths = [list(sub_dict.values()) for sub_dict in task_paths]
return pre_out, task_names, task_paths, sim_counter
[docs]
def run_batch(
self,
fn_pre: Callable[Any, Union[Simulation, List[Simulation], Dict[str, Simulation]]],
fn_post: Callable[
Union[SimulationData, List[SimulationData], Dict[str, SimulationData]], Any
],
path_dir: str = ".",
**batch_kwargs,
) -> Result:
"""
This function has been superceded by `run`, please use `run` for batched simulations.
"""
log.warning(
"In version 2.8.0, the 'run_batch' method is replaced by 'run'."
"'fn_pre' has become 'fn', whilst 'fn_post' remains the same."
"While the original syntax will still be supported, future functionality will be added to the new method moving forward."
)
if len(batch_kwargs) > 0:
log.warning(
"'batch_kwargs' supplied here will no longer be used within the simulation."
)
new_self = self.updated_copy(path_dir=path_dir)
result = new_self.run(fn=fn_pre, fn_post=fn_post)
return result
[docs]
def estimate_cost(self, fn_pre: Callable) -> float:
"""Compute the maximum FlexCredit charge for the ``DesignSpace.run`` computation.
Require a pre function that should return a ``Simulation`` object, a ``Batch`` object, or collection of either.
The pre function is called to estimate the cost - complicated pre functions may cause long runtimes. The cost per
iteration is multiplied by the theoretical maximum number of iterations to give the maximum cost.
Parameters
----------
fn_pre : Callable
Function accepting arguments that correspond to the ``name`` fields
of the ``DesignSpace.parameters``. Should return a ``Simulation`` or ``Batch`` object, or a
``list`` / ``dict`` of these objects.
Returns
-------
float
Estimated maximum cost for the ``DesignSpace.run``.
"""
# Get output fn_pre for paramters at the lowest span / default
arg_dict = {}
for param in self.parameters:
arg_dict[param.name] = param.sample_first()
# Compute fn_pre
pre_out = fn_pre(**arg_dict)
def _estimate_sim_cost(sim):
job = Job(simulation=sim, task_name="estimate_cost")
estimate = job.estimate_cost()
job.delete() # Deleted as only a test with initial parameters
return estimate
if isinstance(pre_out, Simulation):
per_run_estimate = _estimate_sim_cost(pre_out)
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
# Accept list or dict inputs
if isinstance(pre_out, dict):
pre_out = list(pre_out.values())
sims = []
batches = []
for value in pre_out:
if isinstance(value, Simulation):
sims.append(value)
elif isinstance(value, Batch):
batches.append(value)
calculated_estimates = []
for sim in sims:
calculated_estimates.append(_estimate_sim_cost(sim))
for batch in batches:
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("Unrecognized output from pre-function, unable to estimate cost.")
# Calculate maximum number of runs for different methods
run_count = self.method._get_run_count(self.parameters)
# 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)
[docs]
def summarize(self, fn_pre: Callable = None, verbose: bool = True) -> dict[str, Any]:
"""Summarize the setup of the DesignSpace
Prints a summary of the DesignSpace including the method and associated args, the parameters,
and the maximum number of runs expected. If ``fn_pre`` is provided an estimated cost will
also be included. Additional notes are printed where relevant. All data is returned as a dict.
Parameters
----------
fn_pre : Callable = None
Function accepting arguments that correspond to the ``name`` fields
of the ``DesignSpace.parameters``. Allows for estimated cost to be included
in the summary.
verbose: bool = True
Toggle if the summary should be output to log. If False, the dict is returned silently.
Returns
-------
summary_dict: dict
Dictionary containing the summary information.
"""
# Get output console
console = get_logging_console()
# Assemble message
# If check stops it printing standard attributes
arg_values = [
f"{field}: {getattr(self.method, field)}\n"
for field in self.method.__fields__
if field not in MethodOptimize.__fields__
]
param_values = []
for param in self.parameters:
if isinstance(param, ParameterAny):
param_values.append(f"{param.name}: {param.type} {param.allowed_values}\n")
else:
param_values.append(f"{param.name}: {param.type} {param.span}\n")
run_count = self.method._get_run_count(self.parameters)
# Compile data into a dict for return
summary_dict = {
"method": self.method.type,
"method_args": "".join(arg_values),
"param_count": len(self.parameters),
"param_names": ", ".join([param.name for param in self.parameters]),
"param_vals": "".join(param_values),
"max_run_count": run_count,
}
if verbose:
console.log(
"\nSummary of DesignSpace\n\n"
f"Method: {summary_dict['method']}\n"
f"Method Args\n{summary_dict['method_args']}\n"
f"No. of Parameters: {summary_dict['param_count']}\n"
f"Parameters: {summary_dict['param_names']}\n"
f"{summary_dict['param_vals']}\n"
f"Maximum Run Count: {summary_dict['max_run_count']}\n"
)
if fn_pre is not None:
cost_estimate = self.estimate_cost(fn_pre)
summary_dict["cost_estimate"] = cost_estimate
console.log(f"Estimated Maximum Cost: {cost_estimate} FlexCredits")
# NOTE: Could then add more details regarding the output of fn_pre - confirm batching?
# Include additional notes/warnings
notes = []
if isinstance(self.method, MethodGenAlg):
notes.append(
"The maximum run count for MethodGenAlg is difficult to predict. "
"Repeated solutions are not executed, reducing the total number of simulations. "
"High crossover and mutation probabilities may result in an increased number of simulations, potentially exceeding the predicted maximum run count."
)
if isinstance(self.method, (MethodBayOpt, MethodGenAlg, MethodParticleSwarm)):
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"
)
notes.append(
"Discrete 'int' values are automatically rounded if optimizers generate 'float' predictions.\n"
)
if len(notes) > 0 and verbose:
console.log(
"Notes:",
)
for note in notes:
console.log(note)
return summary_dict
