add more models

src/sax/models/__init__.py

@@ -22,6 +22,10 @@
     passthru,
     unitary,
 )
+from .isolators import (
+    circulator,
+    isolator,
+)
 from .mmis import (
     mmi1x2,
     mmi1x2_ideal,
@@ -31,6 +35,10 @@
 from .probes import (
     ideal_probe,
 )
+from .reflectors import (
+    mirror,
+    reflector,
+)
 from .splitters import (
     splitter_ideal,
 )
@@ -39,16 +47,22 @@
     phase_shifter,
     straight,
 )
+from .terminators import (
+    terminator,
+)
 
 __all__ = [
     "attenuator",
     "bend",
+    "circulator",
     "copier",
     "coupler",
     "coupler_ideal",
     "crossing_ideal",
     "grating_coupler",
     "ideal_probe",
+    "isolator",
+    "mirror",
     "mmi1x2",
     "mmi1x2_ideal",
     "mmi2x2",
@@ -58,8 +72,10 @@
     "model_4port",
     "passthru",
     "phase_shifter",
+    "reflector",
     "rf",
     "splitter_ideal",
     "straight",
+    "terminator",
     "unitary",
 ]

src/sax/models/isolators.py

@@ -0,0 +1,149 @@
+"""SAX Isolator and Circulator Models."""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+from pydantic import validate_call
+
+import sax
+
+
+@jax.jit
+@validate_call
+def isolator(
+    *,
+    wl: sax.FloatArrayLike = sax.WL_C,
+    insertion_loss_dB: sax.FloatArrayLike = 0.0,
+    isolation_dB: sax.FloatArrayLike = 40.0,
+) -> sax.SDict:
+    """Optical isolator model (non-reciprocal).
+
+    Transmits light in the forward direction (in0 -> out0) with low loss
+    while blocking the reverse direction (out0 -> in0).
+
+    ```{svgbob}
+    in0             out0
+     o1 ====>====== o2
+    ```
+
+    Args:
+        wl: Wavelength in micrometers.
+        insertion_loss_dB: Forward insertion loss in dB. Defaults to 0.0 dB.
+        isolation_dB: Reverse isolation in dB. Higher values mean better
+            blocking of backward-propagating light. Defaults to 40.0 dB.
+
+    Returns:
+        S-matrix dictionary for the isolator.
+
+    Examples:
+        Isolator with 1 dB insertion loss and 30 dB isolation:
+
+        ```python
+        # mkdocs: render
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import sax
+
+        sax.set_port_naming_strategy("optical")
+
+        wl = sax.wl_c()
+        s = sax.models.isolator(
+            wl=wl,
+            insertion_loss_dB=1.0,
+            isolation_dB=30.0,
+        )
+        fwd = np.abs(s[("o1", "o2")]) ** 2
+        bwd = np.abs(s[("o2", "o1")]) ** 2
+        plt.figure()
+        plt.plot(wl, 10 * np.log10(fwd), label="forward")
+        plt.plot(wl, 10 * np.log10(bwd + 1e-10), label="backward")
+        plt.xlabel("Wavelength [μm]")
+        plt.ylabel("Transmission [dB]")
+        plt.legend()
+        ```
+    """
+    one = jnp.ones_like(jnp.asarray(wl))
+    forward = jnp.asarray(10 ** (-insertion_loss_dB / 20), dtype=complex) * one
+    backward = jnp.asarray(10 ** (-isolation_dB / 20), dtype=complex) * one
+    return {
+        ("o1", "o2"): forward,
+        ("o2", "o1"): backward,
+    }
+
+
+@jax.jit
+@validate_call
+def circulator(
+    *,
+    wl: sax.FloatArrayLike = sax.WL_C,
+    insertion_loss_dB: sax.FloatArrayLike = 0.0,
+    isolation_dB: sax.FloatArrayLike = 40.0,
+) -> sax.SDict:
+    r"""Optical circulator model (non-reciprocal 3-port device).
+
+    Routes light in a circular fashion: port 1 -> port 2 -> port 3 -> port 1.
+    Light traveling in the reverse direction is strongly attenuated.
+
+    ```{svgbob}
+          o2
+          *
+         / \\
+        /   \\
+       / --> \\
+      *-------*
+     o1       o3
+    ```
+
+    Args:
+        wl: Wavelength in micrometers.
+        insertion_loss_dB: Insertion loss in dB for the forward circulation
+            path. Defaults to 0.0 dB.
+        isolation_dB: Isolation in dB between non-adjacent ports (reverse
+            direction). Defaults to 40.0 dB.
+
+    Returns:
+        S-matrix dictionary for the circulator.
+
+    Examples:
+        3-port circulator:
+
+        ```python
+        # mkdocs: render
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import sax
+
+        sax.set_port_naming_strategy("optical")
+
+        wl = sax.wl_c()
+        s = sax.models.circulator(
+            wl=wl,
+            insertion_loss_dB=0.5,
+            isolation_dB=30.0,
+        )
+        fwd_12 = np.abs(s[("o1", "o2")]) ** 2
+        fwd_23 = np.abs(s[("o2", "o3")]) ** 2
+        iso_21 = np.abs(s[("o2", "o1")]) ** 2
+        plt.figure()
+        plt.plot(wl, 10 * np.log10(fwd_12), label="o1→o2")
+        plt.plot(wl, 10 * np.log10(fwd_23), label="o2→o3")
+        plt.plot(wl, 10 * np.log10(iso_21 + 1e-10), label="o2→o1 (isolated)")
+        plt.xlabel("Wavelength [μm]")
+        plt.ylabel("Transmission [dB]")
+        plt.legend()
+        ```
+    """
+    one = jnp.ones_like(jnp.asarray(wl))
+    forward = jnp.asarray(10 ** (-insertion_loss_dB / 20), dtype=complex) * one
+    isolated = jnp.asarray(10 ** (-isolation_dB / 20), dtype=complex) * one
+    return {
+        # Forward circulation: o1 -> o2 -> o3 -> o1
+        ("o1", "o2"): forward,
+        ("o2", "o3"): forward,
+        ("o3", "o1"): forward,
+        # Reverse (isolated): o2 -> o1, o3 -> o2, o1 -> o3
+        ("o2", "o1"): isolated,
+        ("o3", "o2"): isolated,
+        ("o1", "o3"): isolated,
+    }

src/sax/models/reflectors.py

@@ -0,0 +1,126 @@
+"""SAX Reflector Models."""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+from pydantic import validate_call
+
+import sax
+
+
+@jax.jit
+@validate_call
+def reflector(
+    *,
+    wl: sax.FloatArrayLike = sax.WL_C,
+    reflection: sax.FloatArrayLike = 0.5,
+    loss_dB: sax.FloatArrayLike = 0.0,
+) -> sax.SDict:
+    r"""Partial reflector / mirror model.
+
+    A 2-port component that reflects a fraction of the input power and
+    transmits the rest (minus any loss).
+
+    ```{svgbob}
+    in0          |   out0
+     o1 ========= | ========= o2
+                 |
+              mirror
+    ```
+
+    Args:
+        wl: Wavelength in micrometers.
+        reflection: Power reflection coefficient between 0 and 1.
+            0 means full transmission, 1 means full reflection. Defaults to 0.5.
+        loss_dB: Insertion loss in dB applied to both reflected and
+            transmitted amplitudes. Defaults to 0.0 dB.
+
+    Returns:
+        S-matrix dictionary for the reflector.
+
+    Examples:
+        50% reflector:
+
+        ```python
+        # mkdocs: render
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import sax
+
+        sax.set_port_naming_strategy("optical")
+
+        wl = sax.wl_c()
+        s = sax.models.reflector(wl=wl, reflection=0.5)
+        thru = np.abs(s[("o1", "o2")]) ** 2
+        refl = np.abs(s[("o1", "o1")]) ** 2
+        plt.figure()
+        plt.plot(wl, thru, label="transmission")
+        plt.plot(wl, refl, label="reflection")
+        plt.xlabel("Wavelength [μm]")
+        plt.ylabel("Power")
+        plt.legend()
+        ```
+    """
+    one = jnp.ones_like(jnp.asarray(wl))
+    loss_amp = jnp.asarray(10 ** (-loss_dB / 20), dtype=complex) * one
+    r = jnp.asarray(reflection**0.5, dtype=complex) * loss_amp
+    t = jnp.asarray((1 - reflection) ** 0.5, dtype=complex) * loss_amp
+    p = sax.PortNamer(1, 1)
+    return sax.reciprocal(
+        {
+            (p.in0, p.in0): r,
+            (p.in0, p.out0): t,
+            (p.out0, p.out0): r,
+        }
+    )
+
+
+@jax.jit
+@validate_call
+def mirror(
+    *,
+    wl: sax.FloatArrayLike = sax.WL_C,
+    reflection: sax.FloatArrayLike = 1.0,
+    loss_dB: sax.FloatArrayLike = 0.0,
+) -> sax.SDict:
+    r"""Ideal mirror model (reflector with default 100% reflection).
+
+    ```{svgbob}
+    in0          ||
+     o1 ========= ||
+                 ||
+              mirror
+    ```
+
+    Args:
+        wl: Wavelength in micrometers.
+        reflection: Power reflection coefficient between 0 and 1.
+            Defaults to 1.0 (perfect mirror).
+        loss_dB: Insertion loss in dB. Defaults to 0.0 dB.
+
+    Returns:
+        S-matrix dictionary for the mirror.
+
+    Examples:
+        Perfect mirror:
+
+        ```python
+        # mkdocs: render
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import sax
+
+        sax.set_port_naming_strategy("optical")
+
+        wl = sax.wl_c()
+        s = sax.models.mirror(wl=wl)
+        refl = np.abs(s[("o1", "o1")]) ** 2
+        plt.figure()
+        plt.plot(wl, refl, label="reflection")
+        plt.xlabel("Wavelength [μm]")
+        plt.ylabel("Power")
+        plt.legend()
+        ```
+    """
+    return reflector(wl=wl, reflection=reflection, loss_dB=loss_dB)

src/sax/models/rf.py

@@ -233,6 +233,49 @@ def admittance(
     return sax.reciprocal(sdict)
 
 
+@partial(jax.jit, inline=True)
+def resistor(
+    *,
+    f: sax.FloatArrayLike = DEFAULT_FREQUENCY,
+    resistance: sax.FloatLike = 50,
+    z0: sax.ComplexLike = 50,
+) -> sax.SDict:
+    r"""Ideal two-port resistor model.
+
+    Args:
+        f: Frequency in Hz
+        resistance: Resistance in Ohms
+        z0: Reference impedance in Ω.
+
+    Returns:
+        S-dictionary representing the resistor element
+
+    References:
+        [@pozar2012]
+
+    Examples:
+        ```python
+        # mkdocs: render
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import sax
+
+        sax.set_port_naming_strategy("optical")
+
+        f = np.linspace(1e9, 10e9, 500)
+        s = sax.models.rf.resistor(f=f, resistance=100, z0=50)
+        plt.figure()
+        plt.plot(f / 1e9, np.abs(s[("o1", "o1")]), label="|S11|")
+        plt.plot(f / 1e9, np.abs(s[("o1", "o2")]), label="|S12|")
+        plt.plot(f / 1e9, np.abs(s[("o2", "o2")]), label="|S22|")
+        plt.xlabel("Frequency [GHz]")
+        plt.ylabel("Magnitude")
+        plt.legend()
+        ```
+    """
+    return impedance(f=f, z=resistance, z0=z0)
+
+
 @partial(jax.jit, inline=True)
 def capacitor(
     *,