From 7efa69b40df3986a5f956877fc629a85ecf49ec0 Mon Sep 17 00:00:00 2001
From: flaport <floris.laporte@gmail.com>
Date: Thu, 7 Mar 2024 13:27:36 -0800
Subject: [PATCH] custom mmi

---
 cspdk/models.py | 72 ++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 54 insertions(+), 18 deletions(-)

diff --git a/cspdk/models.py b/cspdk/models.py
index 987cffa..7ae6b01 100644
--- a/cspdk/models.py
+++ b/cspdk/models.py
@@ -5,7 +5,7 @@
 import jax
 import jax.numpy as jnp
 import sax
-from gplugins.sax.models import bend, coupler, grating_coupler, straight
+from gplugins.sax.models import bend, grating_coupler, straight
 
 from cspdk.tech import LAYER
 
@@ -106,7 +106,6 @@ def _check_cross_section(cross_section):
         raise err
 
 
-@partial(jax.jit, static_argnames=["cross_section"])
 def _straight(
     *,
     wl: float = 1.55,
@@ -150,18 +149,59 @@ def _straight(
 # MMIs
 ################
 
-mmi1x2_rc = coupler
-mmi2x2_rc = coupler
-mmi1x2_sc = coupler
-mmi2x2_sc = coupler
-mmi1x2_ro = coupler
-mmi2x2_ro = coupler
-mmi1x2_so = coupler
-mmi2x2_so = coupler
-mmi1x2_no = coupler
-mmi2x2_no = coupler
-mmi1x2_nc = coupler
-mmi2x2_nc = coupler
+
+def _mmi_amp(wl=1.55, wl0=1.55, fwhm=0.2, loss_dB=0.3):
+    max_power = 10 ** (-abs(loss_dB) / 10)
+    print(max_power)
+    f = 1 / wl
+    f0 = 1 / wl0
+    f1 = 1 / (wl0 + fwhm / 2)
+    f2 = 1 / (wl0 - fwhm / 2)
+    _fwhm = f2 - f1
+
+    sigma = _fwhm / (2 * jnp.sqrt(2 * jnp.log(2)))
+    power = jnp.exp(-((f - f0) ** 2) / (2 * sigma**2))
+    power = max_power * power / power.max() / 2
+    amp = jnp.sqrt(power)
+    return amp
+
+
+def mmi1x2(wl=1.55, wl0=1.55, fwhm=0.2, loss_dB=0.3):
+    thru = _mmi_amp(wl=wl, wl0=wl0, fwhm=fwhm, loss_dB=loss_dB)
+    return sax.reciprocal(
+        {
+            ("o1", "o2"): thru,
+            ("o1", "o3"): thru,
+        }
+    )
+
+
+def mmi2x2(wl=1.55, wl0=1.55, fwhm=0.2, loss_dB=0.3, shift=0.005):
+    thru = _mmi_amp(wl=wl, wl0=wl0, fwhm=fwhm, loss_dB=loss_dB)
+    cross = 1j * _mmi_amp(wl=wl, wl0=wl0 + shift, fwhm=fwhm, loss_dB=loss_dB)
+    return sax.reciprocal(
+        {
+            ("o1", "o3"): thru,
+            ("o1", "o4"): cross,
+            ("o2", "o3"): cross,
+            ("o2", "o4"): thru,
+        }
+    )
+
+
+mmi1x2_rc = mmi1x2
+mmi1x2_sc = mmi1x2
+mmi1x2_ro = mmi1x2
+mmi1x2_so = mmi1x2
+mmi1x2_no = mmi1x2
+mmi1x2_nc = mmi1x2
+
+mmi2x2_nc = mmi2x2
+mmi2x2_no = mmi2x2
+mmi2x2_so = mmi2x2
+mmi2x2_ro = mmi2x2
+mmi2x2_rc = mmi2x2
+mmi2x2_sc = mmi2x2
 
 ##############################
 # grating couplers Rectangular
@@ -248,7 +288,3 @@ def _crossing(wl=1.5):
     pad=pad,
     heater=heater,
 )
-
-
-if __name__ == "__main__":
-    print(coupler())