fix: rf-migraton propagate batch-kwargs

Author: daquinteroflex

tests/test_plugins/smatrix/test_run_functions.py

@@ -99,8 +99,6 @@ def test_create_batch(monkeypatch, tmp_path):
     result_batch = create_batch(modeler=dummy_modeler, path_dir=str(tmp_path))
     mock_batch_class.assert_called_once_with(
         simulations=dummy_modeler.sim_dict,
-        parent_tasks=None,
-        group_ids=None,
     )
     mock_batch_instance.to_file.assert_called_once_with(os.path.join(str(tmp_path), "batch.hdf5"))
     assert result_batch == mock_batch_instance
@@ -110,25 +108,16 @@ def test_create_batch(monkeypatch, tmp_path):
     mock_batch_instance.to_file.reset_mock()
 
     # Test with parent_batch_id and group_id
-    parent_id = "parent123"
-    group_id = "groupABC"
     file_name = "custom_batch.hdf5"
     result_batch = create_batch(
         modeler=dummy_modeler,
         path_dir=str(tmp_path),
-        parent_batch_id=parent_id,
-        group_id=group_id,
         file_name=file_name,
         some_kwarg="value",
     )
 
-    expected_parent_tasks = dict.fromkeys(dummy_modeler.sim_dict.keys(), (parent_id,))
-    expected_group_ids = dict.fromkeys(dummy_modeler.sim_dict.keys(), (group_id,))
-
     mock_batch_class.assert_called_once_with(
         simulations=dummy_modeler.sim_dict,
-        parent_tasks=expected_parent_tasks,
-        group_ids=expected_group_ids,
         some_kwarg="value",
     )
     mock_batch_instance.to_file.assert_called_once_with(os.path.join(str(tmp_path), file_name))

tidy3d/plugins/smatrix/run.py

@@ -2,7 +2,6 @@
 
 import json
 import os
-from typing import Optional
 
 from tidy3d.components.base import Tidy3dBaseModel
 from tidy3d.plugins.smatrix.component_modelers.modal import ModalComponentModeler
@@ -250,8 +249,6 @@ def compose_modeler_data_from_batch_data(
 def create_batch(
     modeler: ComponentModelerType,
     path_dir: str = DEFAULT_DATA_DIR,
-    parent_batch_id: Optional[str] = None,
-    group_id: Optional[str] = None,
     file_name: str = "batch.hdf5",
     **kwargs,
 ) -> Batch:
@@ -279,33 +276,16 @@ def create_batch(
     """
     filepath = os.path.join(path_dir, file_name)
 
-    if parent_batch_id is not None:
-        parent_task_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 = {}
-        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,
+    modeler: ComponentModelerType, path_dir: str = DEFAULT_DATA_DIR, **kwargs
 ) -> ComponentModelerDataType:
     """Execute the full simulation workflow for a given component modeler.
 
@@ -320,14 +300,16 @@ def run(
         The component modeler defining the simulations to be run.
     path_dir : str, optional
         The directory where the batch file will be saved. Defaults to ".".
+    **kwargs
+        Extra keyword arguments propagated to the Batch creation.
 
     Returns
     -------
     ComponentModelerDataType
         An object containing the processed simulation data, ready for
         S-parameter extraction and analysis.
     """
-    batch = create_batch(modeler=modeler, path_dir=path_dir)
+    batch = create_batch(modeler=modeler, path_dir=path_dir, **kwargs)
     batch_data = batch.run()
     modeler_data = compose_modeler_data_from_batch_data(modeler=modeler, batch_data=batch_data)
     return modeler_data