From 2f00baf0c62fd4a174bc0e4e126c0305f4c4976f Mon Sep 17 00:00:00 2001
From: Jan Petykiewicz <jan@mpxd.net>
Date: Thu, 18 Jul 2024 01:03:51 -0700
Subject: [PATCH] fixup cylindrical wg example

---
 examples/tcyl.py | 40 +++++++++++++++++++++-------------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/examples/tcyl.py b/examples/tcyl.py
index 590a260..e6dc15f 100644
--- a/examples/tcyl.py
+++ b/examples/tcyl.py
@@ -3,7 +3,7 @@ import numpy
 from numpy.linalg import norm
 
 from meanas.fdmath import vec, unvec
-from meanas.fdfd import waveguide_mode, functional, scpml
+from meanas.fdfd import waveguide_cyl, functional, scpml
 from meanas.fdfd.solvers import generic as generic_solver
 
 import gridlock
@@ -37,29 +37,34 @@ def test1(solver=generic_solver):
     xyz_max = numpy.array([800, y_max, z_max]) + (pml_thickness + 2) * dx
 
     # Coordinates of the edges of the cells.
-    half_edge_coords = [numpy.arange(dx/2, m + dx/2, step=dx) for m in xyz_max]
+    half_edge_coords = [numpy.arange(dx / 2, m + dx / 2, step=dx) for m in xyz_max]
     edge_coords = [numpy.hstack((-h[::-1], h)) for h in half_edge_coords]
     edge_coords[0] = numpy.array([-dx, dx])
 
     # #### Create the grid and draw the device ####
     grid = gridlock.Grid(edge_coords)
     epsilon = grid.allocate(n_air**2, dtype=numpy.float32)
-    grid.draw_cuboid(epsilon, center=center, dimensions=[8e3, w, th], eps=n_wg**2)
+    grid.draw_cuboid(epsilon, center=center, dimensions=[8e3, w, th], foreground=n_wg**2)
 
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (1, 2):
         for p in (-1, 1):
-            dxes = scmpl.stretch_with_scpml(dxes, omega=omega, axis=a, polarity=p,
-                                            thickness=pml_thickness)
+            dxes = scpml.stretch_with_scpml(
+                dxes,
+                omega=omega,
+                axis=a,
+                polarity=p,
+                thickness=pml_thickness,
+                )
 
     wg_args = {
         'omega': omega,
         'dxes': [(d[1], d[2]) for d in dxes],
-        'epsilon': vec(g.transpose([1, 2, 0]) for g in epsilon),
+        'epsilon': vec(epsilon.transpose([0, 2, 3, 1])),
         'r0': r0,
     }
 
-    wg_results = waveguide_mode.solve_waveguide_mode_cylindrical(mode_number=0, **wg_args)
+    wg_results = waveguide_cyl.solve_mode(mode_number=0, **wg_args)
 
     E = wg_results['E']
 
@@ -70,20 +75,17 @@ def test1(solver=generic_solver):
     '''
     Plot results
     '''
-    def pcolor(v):
+    def pcolor(fig, ax, v, title):
         vmax = numpy.max(numpy.abs(v))
-        pyplot.pcolor(v.T, cmap='seismic', vmin=-vmax, vmax=vmax)
-        pyplot.axis('equal')
-        pyplot.colorbar()
+        mappable = ax.pcolormesh(v.T, cmap='seismic', vmin=-vmax, vmax=vmax)
+        ax.set_aspect('equal', adjustable='box')
+        ax.set_title(title)
+        ax.figure.colorbar(mappable)
 
-    pyplot.figure()
-    pyplot.subplot(2, 2, 1)
-    pcolor(numpy.real(E[0][:, :]))
-    pyplot.subplot(2, 2, 2)
-    pcolor(numpy.real(E[1][:, :]))
-    pyplot.subplot(2, 2, 3)
-    pcolor(numpy.real(E[2][:, :]))
-    pyplot.subplot(2, 2, 4)
+    fig, axes = pyplot.subplots(2, 2)
+    pcolor(fig, axes[0][0], numpy.real(E[0]), 'Ex')
+    pcolor(fig, axes[0][1], numpy.real(E[1]), 'Ey')
+    pcolor(fig, axes[1][0], numpy.real(E[2]), 'Ez')
     pyplot.show()