Prototype new flow

Author: daquinteroflex

tidy3d/plugins/smatrix/run.py

@@ -15,52 +15,96 @@
 DEFAULT_DATA_DIR = "."
 
 
-def create_batch(
-    modeler: ComponentModelerType,
-    path_dir: str = DEFAULT_DATA_DIR,
-    parent_batch_id: str = None,
-    group_id: str = None,
-    file_name: str = "batch.hdf5",
-    **kwargs,
-) -> Batch:
-    """Creates a simulation Batch from a component modeler and saves it to a file.
+def compose_simulation_data_index(
+    port_task_map: dict[port, str]
+) -> IndexSimulation:
+    port_data_dict = {}
+    for port, task in port_task_map.items():
+        # get simulationdata for each port
+        port_data_dict[port] = sim_data_i
+
+    return IndexSimulationData(
+        index=port_data_dict.keys(),
+        data=port_data_dict.values()
+    )
+
+def compose_terminal_modeler_data(
+    modeler: TerminalComponentModeler,
+    port_task_map: dict[port, str]
+) -> TerminalComponentModelerData:
+    """Assembles `TerminalComponentModelerData` from simulation results.
+
+    This function maps the simulation data from a completed batch run back to the
+    ports of the terminal component modeler.
 
     Args:
-        modeler: The component modeler that defines the set of simulations.
-        path_dir: Directory where the batch file will be saved.
-        file_name: Name for the HDF5 file where the batch is stored.
-        **kwargs: Additional keyword arguments passed to the `Batch` constructor.
+        modeler: The `TerminalComponentModeler` used to generate the simulations.
 
     Returns:
-        The configured `Batch` object ready for execution.
+        A `TerminalComponentModelerData` object containing the results mapped to
+        their respective ports.
     """
-    filepath = os.path.join(path_dir, file_name)
+    port_simulation_data = compose_simulation_data_index(port_task_map)
+    return TerminalComponentModelerData(modeler=modeler, data=port_simulation_data)
 
-    if parent_batch_id is not None:
-        parent_task_dict = dict()
-        for key in modeler.sim_dict.keys():
-            parent_task_dict[key] = (parent_batch_id,)
-    else:
-        parent_task_dict = None
+def compose_component_modeler_data(
+    modeler: ComponentModeler,
+    port_task_map: dict[port, str]
+) -> ComponentModelerData:
+    """Assembles `ComponentModelerData` from simulation results.
 
-    if group_id is not None:
-        group_id_dict = dict()
-        for key in modeler.sim_dict.keys():
-            group_id_dict[key] = (group_id,)
-    else:
-        group_id_dict = None
+    This function maps the simulation data from a completed batch run back to the
+    ports of the component modeler.
 
-    batch = Batch(
-        simulations=modeler.sim_dict,
-        parent_tasks=parent_task_dict,
-        group_ids=group_id_dict,
-        **kwargs,
-    )
-    batch.to_file(filepath)
-    return batch
+    Args:
+        modeler: The `ComponentModeler` used to generate the simulations.
+        batch_data: The results obtained from running the simulation `Batch`.
 
+    Returns:
+        A `ComponentModelerData` object containing the results mapped to
+        their respective ports.
+    """
+    port_simulation_data = compose_simulation_data_index(port_task_map)
+    return ComponentModelerData(modeler=modeler, data=port_simulation_data)
 
-def compose_terminal_modeler_data(
+
+def compose_modeler_data(
+    modeler_file: str,
+    port_task_map: dict[port, str]
+) -> ComponentModelerDataType:
+    """Selects the correct composer based on the modeler type and creates the data object.
+
+    This method acts as a dispatcher, inspecting the type of `modeler` to determine
+    which composer function (`compose_component_modeler_data` or
+    `compose_terminal_modeler_data`) to invoke.
+
+    Args:
+        modeler: The component modeler, which can be either a `ComponentModeler` or
+            a `TerminalComponentModeler`.
+        batch_data: The results obtained from running the simulation `Batch`.
+
+    Returns:
+        The appropriate `ComponentModelerDataType` object containing the simulation results.
+
+    Raises:
+        TypeError: If the provided `modeler` is not a recognized type.
+    """
+    json_str = Tidy3dBaseModel._json_string_from_hdf5(modeler_file)
+    model_dict = json.loads(json_str)
+    modeler = model_dict["type"]
+
+    if modeler_type == "ComponentModeler":
+        modeler = ComponentModeler.from_file(task_file)
+        modeler_data = compose_component_modeler_data(modeler=modeler, port_task_map=port_task_map)
+    elif modeler_type == "TerminalComponentModeler":
+        modeler = TerminalComponentModeler.from_file(task_file)
+        modeler_data = compose_terminal_modeler_data(modeler=modeler, port_task_map=port_task_map)
+    else:
+        raise TypeError(f"Unsupported modeler type: {type(modeler).__name__}")
+    return modeler_data
+
+
+def compose_terminal_modeler_data_from_batch_data(
     modeler: TerminalComponentModeler,
     batch_data: BatchData,
 ) -> TerminalComponentModelerData:
@@ -83,9 +127,8 @@ def compose_terminal_modeler_data(
     return TerminalComponentModelerData(modeler=modeler, data=port_simulation_data)
 
 
-def compose_component_modeler_data(
+def compose_component_modeler_data_from_batch_data(
     modeler: ComponentModeler,
-    simulation_data_list: Optional[list[SimulationData]] = None,
     batch_data: Optional[BatchData] = None,
 ) -> ComponentModelerData:
     """Assembles `ComponentModelerData` from simulation results.
@@ -101,18 +144,13 @@ def compose_component_modeler_data(
         A `ComponentModelerData` object containing the results mapped to
         their respective ports.
     """
-    if simulation_data_list:
-        pass
-    elif batch_data:
-        ports = [modeler.get_task_name(port=port_i) for port_i in modeler.ports]
-        data = [batch_data[modeler.get_task_name(port=port_i)] for port_i in modeler.ports]
-        port_simulation_data = IndexSimulationData(ports=ports, data=data)
-    else:
-        pass
+    ports = [modeler.get_task_name(port=port_i) for port_i in modeler.ports]
+    data = [batch_data[modeler.get_task_name(port=port_i)] for port_i in modeler.ports]
+    port_simulation_data = IndexSimulationData(ports=ports, data=data)
     return ComponentModelerData(modeler=modeler, data=port_simulation_data)
 
 
-def compose_modeler_data(
+def compose_modeler_data_from_batch_data(
     modeler: ComponentModelerType,
     batch_data: BatchData,
 ) -> ComponentModelerDataType:
@@ -134,15 +172,59 @@ def compose_modeler_data(
         TypeError: If the provided `modeler` is not a recognized type.
     """
     if isinstance(modeler, ComponentModeler):
-        modeler_data = compose_component_modeler_data(modeler=modeler, batch_data=batch_data)
+        modeler_data = compose_component_modeler_data_from_batch_data(modeler=modeler, batch_data=batch_data)
     elif isinstance(modeler, TerminalComponentModeler):
-        modeler_data = compose_terminal_modeler_data(modeler=modeler, batch_data=batch_data)
+        modeler_data = compose_terminal_modeler_data_from_batch_data(modeler=modeler, batch_data=batch_data)
     else:
         raise TypeError(f"Unsupported modeler type: {type(modeler).__name__}")
 
     return modeler_data
 
 
+def create_batch(
+    modeler: ComponentModelerType,
+    path_dir: str = DEFAULT_DATA_DIR,
+    parent_batch_id: str = None,
+    group_id: str = None,
+    file_name: str = "batch.hdf5",
+    **kwargs,
+) -> Batch:
+    """Creates a simulation Batch from a component modeler and saves it to a file.
+
+    Args:
+        modeler: The component modeler that defines the set of simulations.
+        path_dir: Directory where the batch file will be saved.
+        file_name: Name for the HDF5 file where the batch is stored.
+        **kwargs: Additional keyword arguments passed to the `Batch` constructor.
+
+    Returns:
+        The configured `Batch` object ready for execution.
+    """
+    filepath = os.path.join(path_dir, file_name)
+
+    if parent_batch_id is not None:
+        parent_task_dict = dict()
+        for key in modeler.sim_dict.keys():
+            parent_task_dict[key] = (parent_batch_id,)
+    else:
+        parent_task_dict = None
+
+    if group_id is not None:
+        group_id_dict = dict()
+        for key in modeler.sim_dict.keys():
+            group_id_dict[key] = (group_id,)
+    else:
+        group_id_dict = None
+
+    batch = Batch(
+        simulations=modeler.sim_dict,
+        parent_tasks=parent_task_dict,
+        group_ids=group_id_dict,
+        **kwargs,
+    )
+    batch.to_file(filepath)
+    return batch
+
 def run(
     modeler: ComponentModelerType,
     path_dir: str = DEFAULT_DATA_DIR,