From 5af193f5940dc860894846f30cf16c9c1c66ed8c Mon Sep 17 00:00:00 2001
From: Joaquin Matres <4514346+joamatab@users.noreply.github.com>
Date: Sun, 24 Mar 2024 10:52:58 -0700
Subject: [PATCH] fix simulations

---
 pyproject.toml                             |  2 +-
 ubcpdk/__init__.py                         | 10 ++--
 ubcpdk/models.py                           | 57 ++++++++--------------
 ubcpdk/samples/test_circuit_simulations.py | 34 +++++++++++++
 4 files changed, 61 insertions(+), 42 deletions(-)
 create mode 100644 ubcpdk/samples/test_circuit_simulations.py

diff --git a/pyproject.toml b/pyproject.toml
index 5def08a8..6b2f4edf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@ classifiers = [
 ]
 dependencies = [
   "gdsfactory==7.22.2",
-  "gplugins[tidy3d,sax,schematic]>=0.10.2,<0.11"
+  "gplugins[tidy3d,sax,schematic]>=0.11,<0.12"
 ]
 description = "ubcpdk pdk"
 keywords = ["python"]
diff --git a/ubcpdk/__init__.py b/ubcpdk/__init__.py
index ddd184f7..defe5039 100644
--- a/ubcpdk/__init__.py
+++ b/ubcpdk/__init__.py
@@ -3,8 +3,9 @@
 from gdsfactory.config import PATH as GPATH
 from gdsfactory.get_factories import get_cells
 from gdsfactory.pdk import Pdk
+from gplugins.sax.models import get_models
 
-from ubcpdk import components, data, tech
+from ubcpdk import components, data, models, tech
 from ubcpdk.config import CONFIG, PATH
 from ubcpdk.tech import LAYER, LAYER_STACK, LAYER_VIEWS, cross_sections
 
@@ -29,6 +30,7 @@
     name="ubcpdk",
     cells=cells,
     cross_sections=cross_sections,
+    models=get_models(models),
     layers=dict(LAYER),
     layer_stack=LAYER_STACK,
     layer_views=LAYER_VIEWS,
@@ -40,6 +42,6 @@
 
 
 if __name__ == "__main__":
-    f = PDK.cells
-    for k, _v in f.items():
-        print(k)
+    m = get_models(models)
+    for model in m.keys():
+        print(model)
diff --git a/ubcpdk/models.py b/ubcpdk/models.py
index ffc8f8ff..b406d0df 100644
--- a/ubcpdk/models.py
+++ b/ubcpdk/models.py
@@ -2,50 +2,33 @@
 
 from functools import partial
 
-from gplugins.sax.models import (
-    attenuator,
-    bend,
-    coupler,
-    grating_coupler,
-    mmi1x2,
-    mmi2x2,
-    phase_shifter,
-)
-from gplugins.sax.models import straight as _straight
+import gplugins.sax.models as sm
 
 nm = 1e-3
 
 
-straight = partial(_straight, wl0=1.55, neff=2.4, ng=4.2)
-bend_euler = partial(bend, loss=0.03)
+straight = partial(sm.straight, wl0=1.55, neff=2.4, ng=4.2)
+bend_euler_sc = bend_euler = partial(sm.bend, loss=0.03)
 
-gc_te1550 = partial(grating_coupler, loss=6, bandwidth=35 * nm, wl0=1.55)
-gc_te1550_broadband = partial(grating_coupler, loss=6, bandwidth=50 * nm, wl0=1.55)
-gc_tm1550 = partial(grating_coupler, loss=6, bandwidth=35 * nm, wl0=1.55)
-
-gc_te1310_broadband = partial(grating_coupler, loss=6, bandwidth=50 * nm, wl0=1.31)
-gc_te1310 = partial(grating_coupler, loss=6, bandwidth=35 * nm, wl0=1.31)
-
-
-models = dict(
-    attenuator=attenuator,
-    bend_euler=bend,
-    coupler=coupler,
-    mmi1x2=mmi1x2,
-    mmi2x2=mmi2x2,
-    phase_shifter=phase_shifter,
-    straight=straight,
-    taper=straight,
-    gc_te1550=gc_te1550,
-    gc_te1550_broadband=gc_te1550_broadband,
-    gc_tm1550=gc_tm1550,
-    gc_te1310_broadband=gc_te1310_broadband,
-    gc_te1310=gc_te1310,
-)
+################
+# grating couplers
+################
+gc_te1550 = partial(sm.grating_coupler, loss=6, bandwidth=35 * nm, wl0=1.55)
+gc_te1550_broadband = partial(sm.grating_coupler, loss=6, bandwidth=50 * nm, wl0=1.55)
+gc_tm1550 = partial(sm.grating_coupler, loss=6, bandwidth=35 * nm, wl0=1.55)
+gc_te1310_broadband = partial(sm.grating_coupler, loss=6, bandwidth=50 * nm, wl0=1.31)
+gc_te1310 = partial(sm.grating_coupler, loss=6, bandwidth=35 * nm, wl0=1.31)
 
+################
+# MMIs
+################
+mmi1x2 = partial(sm.mmi1x2, wl0=1.55, fwhm=0.2, loss_dB=0.3)
+mmi2x2 = partial(sm.mmi2x2, wl0=1.55, fwhm=0.2, loss_dB=0.3)
+ebeam_y_1550 = mmi1x2
+coupler = sm.coupler
 
 if __name__ == "__main__":
     import gplugins.sax as gs
 
-    gs.plot_model(grating_coupler)
-    # gs.plot_model(coupler)
+    gs.plot_model(gc_te1550)
+    gs.plot_model(coupler)
diff --git a/ubcpdk/samples/test_circuit_simulations.py b/ubcpdk/samples/test_circuit_simulations.py
new file mode 100644
index 00000000..f81d2e59
--- /dev/null
+++ b/ubcpdk/samples/test_circuit_simulations.py
@@ -0,0 +1,34 @@
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import sax
+
+import ubcpdk
+from ubcpdk import PDK
+
+
+def test_mzi():
+    c = ubcpdk.components.mzi(delta_length=20)
+    netlist = c.get_netlist()
+    models = PDK.models
+    circuit, _ = sax.circuit(netlist, models=models)  # type: ignore
+    wl = jnp.linspace(1.5, 1.6)
+
+    S = circuit(wl=wl)
+    assert S
+
+
+if __name__ == "__main__":
+    c = ubcpdk.components.mzi(delta_length=20)
+    netlist = c.get_netlist()
+    models = PDK.models
+    circuit, _ = sax.circuit(netlist, models=models)  # type: ignore
+    wl = jnp.linspace(1.5, 1.6)
+
+    S = circuit(wl=wl)
+    plt.figure(figsize=(14, 4))
+    plt.title("MZI")
+    plt.plot(1e3 * wl, jnp.abs(S["o1", "o2"]) ** 2)  # type: ignore
+    plt.xlabel("λ [nm]")
+    plt.ylabel("T")
+    plt.grid(True)
+    plt.show()