Merge pull request gdsfactory#1122 from simbilod/meep_visualization

expose plot options in meep solver

Former-commit-id: b636bf2 [formerly 6b8a101]
Former-commit-id: 8ed9cf34dfc6966344d5f0b1431881a99fdee3b6

gdsfactory/simulation/gmeep/write_sparameters_meep.py

@@ -139,6 +139,7 @@ def write_sparameters_meep(
     decay_by: float = 1e-3,
     is_3d: bool = False,
     z: float = 0,
+    plot_args: Dict = None,
     **settings,
 ) -> Dict:
     r"""Returns Sparameters and writes them to npz filepath.
@@ -231,6 +232,8 @@ def write_sparameters_meep(
         ymargin_bot: south distance from component to PML.
         is_3d: if True runs in 3D (much slower).
         z: for 2D plot.
+        plot_args: if animate or not run, customization keyword arguments passed to
+          `plot2D()` (i.e. `labels`, `eps_parameters`, `boundary_parameters`, `field_parameters`, etc.)
 
     keyword Args:
         extend_ports_length: to extend ports beyond the PML (um).
@@ -258,6 +261,8 @@ def write_sparameters_meep(
     component = gf.get_component(component)
     layer_stack = layer_stack or get_layer_stack()
 
+    plot_args = plot_args or {}
+
     for setting in settings:
         if setting not in settings_get_simulation:
             raise ValueError(f"{setting!r} not in {settings_get_simulation}")
@@ -333,10 +338,11 @@ def write_sparameters_meep(
                 output_plane=mp.Volume(
                     size=mp.Vector3(sim.cell_size.x, sim.cell_size.y, 0),
                     center=mp.Vector3(0, 0, z),
-                )
+                ),
+                **plot_args,
             )
         else:
-            sim.plot2D(plot_eps_flag=True)
+            sim.plot2D(plot_eps_flag=True, **plot_args)
         return sim
 
     if filepath.exists():
@@ -365,6 +371,7 @@ def sparameter_calculation(
         wavelength_stop: float = wavelength_stop,
         wavelength_points: int = wavelength_points,
         animate: bool = animate,
+        plot_args: Dict = plot_args,
         dispersive: bool = dispersive,
         decay_by: float = decay_by,
         **settings,
@@ -395,18 +402,27 @@ def sparameter_calculation(
         termination = [mp.stop_when_energy_decayed(dt=50, decay_by=decay_by)]
 
         if animate:
-            sim.use_output_directory()
-            animate = mp.Animate2D(
-                sim,
-                fields=mp.Ez,
-                realtime=True,
-                field_parameters={
+
+            # Defaults for animation
+            if "field_parameters" not in plot_args:
+                plot_args["field_parameters"] = {
                     "alpha": 0.8,
                     "cmap": "RdBu",
                     "interpolation": "none",
-                },
-                eps_parameters={"contour": True},
-                normalize=True,
+                }
+            if "eps_parameters" not in plot_args:
+                plot_args["eps_parameters"] = {"contour": True}
+            if "fields" not in plot_args:
+                plot_args["fields"] = mp.Ez
+            if "realtime" not in plot_args:
+                plot_args["realtime"] = True
+            if "normalize" not in plot_args:
+                plot_args["normalize"] = True
+
+            sim.use_output_directory()
+            animate = mp.Animate2D(
+                sim,
+                **plot_args,
             )
             sim.run(mp.at_every(1, animate), until_after_sources=termination)
             animate.to_mp4(30, f"{component.name}_{port_source_name}.mp4")
@@ -531,8 +547,33 @@ def sparameter_calculation(
     sim_settings = dict(
         wavelength_start=wavelength_start, wavelength_stop=wavelength_stop
     )
-    c = gf.components.mmi1x2(cross_section=gf.cross_section.strip)
-    sp = write_sparameters_meep(c, run=True, is_3d=False, **sim_settings)
+    # c = gf.components.mmi1x2(cross_section=gf.cross_section.strip)
+    c = gf.components.straight(length=2)
+    import matplotlib.pyplot as plt
+
+    def func(x):
+        result = np.where(np.abs(x) > 1e-10, np.abs(x) ** 2, -10)
+        return np.log10(result, out=result, where=result > 0)
+
+    sp = write_sparameters_meep(
+        c,
+        run=True,
+        animate=True,
+        is_3d=False,
+        plot_args={
+            "eps_parameters": {"contour": True},
+            "field_parameters": {
+                "alpha": 0.8,
+                "cmap": "RdBu",
+                "interpolation": "none",
+                "post_process": func,
+            },
+            "realtime": False,
+        },
+        overwrite=True,
+        **sim_settings,
+    )
+    plt.show()
 
     # from gdsfactory.simulation.add_simulation_markers import add_simulation_markers
     # import gdsfactory.simulation as sim