fix mzi lattice example

[joamatab]

trying to fix #251

Still not working, marking it as draft

@lucas-flexcompute

[lucas-flexcompute]

The plugin seems to be working correctly, but I get strange results when using the jnp.interp function inside a model. Here's a simple example: when the circuit model is calculated with the same wavelengths used in the original simulation, the result is correct:

def bend_model(cross_section: gf.typings.CrossSectionSpec = "xs_sc"):
    component = gf.components.bend_euler(cross_section=cross_section)
    s = gt.write_sparameters(
        component=component,
        filepath=PATH.sparameters_repo / "bend_filter.npz",
        layer_stack=layer_stack,
    )
    wavelengths = s.pop("wavelengths")

    @jax.jit
    def _model(wl=1.55):
        s11 = jnp.interp(wl, wavelengths, s["o1@0,o1@0"])
        s21 = jnp.interp(wl, wavelengths, s["o2@0,o1@0"])
        return {
            ("o1", "o1"): s11,
            ("o1", "o2"): s21,
            ("o2", "o1"): s21,
            ("o2", "o2"): s11,
        }

    return _model

c = gf.Component(name="bend")
ref = c.add_ref(gf.components.bend_euler(cross_section=cross_section))
c.add_ports(ref.ports)
x, _ = sax.circuit(
    c.get_netlist(), {"bend_euler": bend_model(cross_section=cross_section)}
)

wl = np.linspace(1.5, 1.6, 11)
s = x(wl=wl)
plt.plot(wl, jnp.abs(s[("o1", "o2")]) ** 2)
plt.ylabel("S21")
plt.xlabel("λ (µm)")

But when I try to use a more refined wavelength set wl = np.linspace(1.5, 1.6, 101), the interpolated points tend to zero as they get farther from the originally sampled points:

Any ideas to what might be causing this?

[joamatab]

this worked well for me

i did not see any issue when increasing the number of points

import sax
import gdsfactory as gf 
from gplugins.common.config import PATH
import gplugins.tidy3d as gt
import numpy as np 
import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt 

from gdsfactory.pdk import get_active_pdk


pdk = get_active_pdk()
layer_stack = pdk.get_layer_stack()
cross_section = gf.cross_section.xs_sc

def bend_model(cross_section: gf.typings.CrossSectionSpec = "xs_sc"):
    component = gf.components.bend_euler(cross_section=cross_section)
    s = gt.write_sparameters(
        component=component,
        filepath=PATH.sparameters_repo / "bend_filter.npz",
        layer_stack=layer_stack,
    )
    wavelengths = s.pop("wavelengths")

    @jax.jit
    def _model(wl=1.55):
        s11 = jnp.interp(wl, wavelengths, s["o1@0,o1@0"])
        s21 = jnp.interp(wl, wavelengths, s["o2@0,o1@0"])
        return {
            ("o1", "o1"): s11,
            ("o1", "o2"): s21,
            ("o2", "o1"): s21,
            ("o2", "o2"): s11,
        }

    return _model

c = gf.Component(name="bend")
ref = c.add_ref(gf.components.bend_euler(cross_section=cross_section))
c.add_ports(ref.ports)
x, _ = sax.circuit(
    c.get_netlist(), {"bend_euler": bend_model(cross_section=cross_section)}
)

# wl = np.linspace(1.5, 1.6, 11) # works well
wl = np.linspace(1.5, 1.6, 101) # works well as well
s = x(wl=wl)
plt.plot(wl, jnp.abs(s[("o1", "o2")]) ** 2)
plt.ylabel("S21")
plt.xlabel("λ (µm)")
plt.show()

@flaport

[flaport]

Hi, @lucas-flexcompute ,

The issue is that you're interpolating an S-parameter, which is complex valued. jnp.interp only works on real valued arrays (and it will take the real value of the array if complex).

Interpolating S-parameters is kinda difficult, as the phase wraps around, luckily in this case it seems the wavelength step is small enough to be able to unwrap the phase. I adapted the code you provided to correctly interpolate the phase:

def bend_model(cross_section: gf.typings.CrossSectionSpec = "xs_sc"):
    component = gf.components.bend_euler(cross_section=cross_section)
    s = gt.write_sparameters(
        component=component,
        filepath=PATH.sparameters_repo / "bend_filter.npz",
        layer_stack=layer_stack,
    )
    wavelengths = s.pop("wavelengths")
    amp = {k: jnp.abs(v) for k, v in s.items()}
    phi = {k: jnp.unwrap(jnp.angle(v)) for k, v in s.items()}

    @jax.jit
    def _model(wl=1.55):
        amp11 = jnp.interp(wl, wavelengths, amp["o1@0,o1@0"])
        amp21 = jnp.interp(wl, wavelengths, amp["o2@0,o1@0"])
        phi11 = jnp.interp(wl, wavelengths, phi["o1@0,o1@0"])
        phi21 = jnp.interp(wl, wavelengths, phi["o2@0,o1@0"])
        s11 = amp11 * jnp.exp(1j*phi11)
        s21 = amp21 * jnp.exp(1j*phi21)
        return {
            ("o1", "o1"): s11,
            ("o1", "o2"): s21,
            ("o2", "o1"): s21,
            ("o2", "o2"): s11,
        }

    return _model

c = gf.Component(name="bend")
ref = c.add_ref(gf.components.bend_euler(cross_section=cross_section))
c.add_ports(ref.ports)
x, _ = sax.circuit(
    c.get_netlist(), {"bend_euler": bend_model(cross_section=cross_section)}
)

wl = np.linspace(1.5, 1.6, 101)
s = x(wl=wl)
plt.plot(wl, 10*jnp.log10(jnp.abs(s[("o1", "o2")]) ** 2))
plt.ylabel("S21 (amp, dB)")
plt.xlabel("λ (µm)")
plt.twinx()
plt.plot(wl, jnp.unwrap(jnp.angle(s[("o1", "o2")])), color="C1")
plt.ylabel("S21 (phi, rad)")

PS: In cases where using small wavelength steps to be able to unwrap the phase is infeasible / expensive you can look into sax.grouped_interp to unwrap a phase for wavelengths containing a 'small' step and a 'big' step.

[lucas-flexcompute]

Ok, so that's the problem with the notebook, then, thanks @flaport! I guess the safest solution here would be to avoid interpolation altogether, but that means fixing the wavelength from the start. What do you think @joamatab?

[flaport]

@lucas-flexcompute , I think interpolating is safe as long as you use enough frequency points, so I would recommend that. Adding more wavelengths is relatively cheap for FDTD anyway.

[lucas-flexcompute]

Using the complex interpolation by magnitude/phase works:

[joamatab]

Awesome, now it looks even better than before!

[joamatab]

Thank you Floris and Lucas, we made another PR for this

#255