diff --git a/tests/test_components/test_autograd.py b/tests/test_components/test_autograd.py index d2b39ccb5..f7011d0bb 100644 --- a/tests/test_components/test_autograd.py +++ b/tests/test_components/test_autograd.py @@ -12,6 +12,7 @@ import tidy3d as td from tidy3d.web import run, run_async +from tidy3d.web.api.autograd.autograd import run from ..utils import run_emulated @@ -341,7 +342,7 @@ def test_autograd_objective(use_emulated_run, structure_key, monitor_key): """Test an objective function through tidy3d autograd.""" # for logging output - td.config.logging_level = "ERROR" + td.config.logging_level = "INFO" fn_dict = get_functions(structure_key, monitor_key) make_sim = fn_dict["sim"] @@ -414,14 +415,8 @@ def objective(*args): def test_autograd_async(use_emulated_run_async, structure_key, monitor_key): """Test an objective function through tidy3d autograd.""" - # # import here so it uses emulated run - # from importlib import reload - # from tidy3d.web.api.autograd import autograd - - # reload(autograd) - # for logging output - td.config.logging_level = "ERROR" + td.config.logging_level = "INFO" fn_dict = get_functions(structure_key, monitor_key) make_sim = fn_dict["sim"] @@ -442,3 +437,42 @@ def objective(*args): val, grad = ag.value_and_grad(objective)(params0) print(val, grad) assert npa.all(grad != 0.0), "some gradients are 0" + + +def test_autograd_speed_num_structures(use_emulated_run): + """Test an objective function through tidy3d autograd.""" + + num_structures_test = 10 + + import time + + # for logging output + td.config.logging_level = "ERROR" + + fn_dict = get_functions(ALL_KEY, ALL_KEY) + make_sim = fn_dict["sim"] + + monitor_key = "mode" + structure_key = "size_element" + monitor, postprocess = make_monitors()[monitor_key] + + def make_sim(*args): + structure = make_structures(*args)[structure_key] + structures = num_structures_test * [structure] + return SIM_BASE.updated_copy(structures=structures, monitors=[monitor]) + + def objective(*args): + """Objective function.""" + sim = make_sim(*args) + data = run(sim, task_name="autograd_test", verbose=False) + value = postprocess(data[monitor_key]) + return value + + # if speed test, get the profile + with cProfile.Profile() as pr: + t = time.time() + val, grad = ag.value_and_grad(objective)(params0) + t2 = time.time() - t + pr.print_stats(sort="cumtime") + pr.dump_stats("results.prof") + print(f"{num_structures_test} structures took {t2:.2e} seconds") diff --git a/tests/test_plugins/test_invdes.py b/tests/test_plugins/test_invdes.py index b6067839d..27c848285 100644 --- a/tests/test_plugins/test_invdes.py +++ b/tests/test_plugins/test_invdes.py @@ -437,10 +437,10 @@ def test_result_data(use_emulated_run_autograd): sim_data_last = result.sim_data_last(task_name="last") -def test_result_data_multi(use_emulated_run_autograd_async, tmp_path): +def test_result_data_multi(use_emulated_run_autograd_async, use_emulated_run_autograd, tmp_path): result_multi = make_result_multi(use_emulated_run_autograd_async) sim_last = result_multi.sim_last - sim_data_last = result_multi.sim_data_last(task_name="last") + sim_data_last = result_multi.sim_data_last() def test_result_empty(): diff --git a/tidy3d/plugins/invdes/design.py b/tidy3d/plugins/invdes/design.py index aafa4fc8c..82b5a41a8 100644 --- a/tidy3d/plugins/invdes/design.py +++ b/tidy3d/plugins/invdes/design.py @@ -7,7 +7,7 @@ import autograd.numpy as npa import tidy3d as td -from tidy3d.web import run_autograd +import tidy3d.web as web from tidy3d.components.autograd import get_static from .base import InvdesBaseModel @@ -130,7 +130,7 @@ def to_simulation_data(self, params: npa.ndarray, **kwargs) -> td.SimulationData """Convert the ``InverseDesign`` to a ``td.Simulation`` and run it.""" simulation = self.to_simulation(params=params) kwargs.setdefault("task_name", self.task_name) - sim_data = run_autograd(simulation, verbose=self.verbose, **kwargs) + sim_data = web.run(simulation, verbose=self.verbose, **kwargs) return sim_data @@ -201,14 +201,11 @@ def to_simulation(self, params: npa.ndarray) -> dict[str, td.Simulation]: simulation_list = [design.to_simulation(params) for design in self.designs] return dict(zip(self.task_names, simulation_list)) - def to_simulation_data(self, params: npa.ndarray, **kwargs) -> dict[str, td.SimulationData]: + def to_simulation_data(self, params: npa.ndarray, **kwargs) -> web.BatchData: """Convert the ``InverseDesignMulti`` to a set of ``td.Simulation``s and run async.""" simulations = self.to_simulation(params) kwargs.setdefault("verbose", self.verbose) - batch_data = td.web.run_async_autograd(simulations, **kwargs) - - # TODO: should we yield each one here instead? for memory saving? - return {task_name: batch_data[task_name] for task_name in self.task_names} + return web.run_async(simulations, **kwargs) InverseDesignType = typing.Union[InverseDesign, InverseDesignMulti] diff --git a/tidy3d/plugins/invdes/result.py b/tidy3d/plugins/invdes/result.py index 1cbd63cf0..45cd703e0 100644 --- a/tidy3d/plugins/invdes/result.py +++ b/tidy3d/plugins/invdes/result.py @@ -104,24 +104,22 @@ def get_last(self, key: str) -> typing.Any: def get_sim(self, index: int = -1) -> typing.Union[td.Simulation, typing.List[td.Simulation]]: """Get the simulation at a specific index in the history (list of sims if multi).""" - params = self.get(key="params", index=index) - return self.design.to_simulation(np.array(params)) + params = np.array(self.get(key="params", index=index)) + return self.design.to_simulation(params=params) - def get_sim_data( - self, task_name: str, index: int = -1, **kwargs - ) -> typing.Union[td.SimulationData, typing.List[td.SimulationData]]: + def get_sim_data(self, index: int = -1, **kwargs) -> typing.Union[td.SimulationData, typing.List[td.SimulationData]]: """Get the simulation data at a specific index in the history (list of simdata if multi).""" - params = self.get(key="params", index=index) - return self.design.to_simulation_data(params=params, task_name=task_name, **kwargs) + params = np.array(self.get(key="params", index=index)) + return self.design.to_simulation_data(params=params, **kwargs) @property def sim_last(self) -> typing.Union[td.Simulation, typing.List[td.Simulation]]: """The last simulation.""" return self.get_sim(index=-1) - def sim_data_last(self, task_name: str, **kwargs) -> td.SimulationData: + def sim_data_last(self, **kwargs) -> td.SimulationData: """Run the last simulation and return its data.""" - return self.get_sim_data(index=-1, task_name=task_name, **kwargs) + return self.get_sim_data(index=-1, **kwargs) def plot_optimization(self): """Plot the optimization progress from the history.""" diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index 0ca865390..ae191692c 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -20,7 +20,10 @@ # keys for data into auxiliary dictionary AUX_KEY_SIM_DATA_ORIGINAL = "sim_data" AUX_KEY_SIM_DATA_FWD = "sim_data_fwd_adjoint" -ISSUE_URL = "https://github.com/flexcompute/tidy3d/issues/new?assignees=&labels=feature&projects=&template=autograd_bug.md" +ISSUE_URL = ( + "https://github.com/flexcompute/tidy3d/issues/new?" + "assignees=tylerflex&labels=adjoint&projects=&template=autograd_bug.md" +) URL_LINK = f"[blue underline][link={ISSUE_URL}]'{ISSUE_URL}'[/link][/blue underline]" @@ -56,7 +59,6 @@ def run( """User-facing ``web.run`` function, compatible with ``autograd`` differentiation.""" if isinstance(simulation, td.Simulation): try: - dict()["a"] return _run( simulation=simulation, task_name=task_name, @@ -102,7 +104,13 @@ def run_async( ) -> BatchData: """User-facing ``run_async`` function.""" - if all(isinstance(sim, td.Simulation) for sim in simulations.values()): + def is_valid_for_autograd(simulations: dict[str, td.Simulation]) -> bool: + """Check whether the supplied simulations dict can use autograd run.""" + if not isinstance(simulations, dict): + return False + return all(isinstance(sim, td.Simulation) for sim in simulations.values()) + + if is_valid_for_autograd(simulations): try: return _run_async( simulations=simulations,