electronic ¤

@component(ports=("c", "b", "e"))
def BJT_NPN(
    signals: Signals,
    Is: float = 1e-12,
    BetaF: float = 100.0,
    BetaR: float = 1.0,
    Vt: float = 25.85e-3,
) -> PhysicsReturn:
    """NPN BJT using the transport form of the Ebers-Moll DC model.

    Junction voltages are clipped to ``[-5, 2]`` V before exponentiation.
    For transient simulations with junction charge dynamics use
    :func:`BJT_NPN_Dynamic` instead.

    Args:
        signals: Port voltages at collector (``c``), base (``b``), and emitter (``e``).
        Is: Saturation current in amperes. Defaults to ``1e-12``.
        BetaF: Forward common-emitter current gain. Defaults to ``100``.
        BetaR: Reverse common-emitter current gain. Defaults to ``1``.
        Vt: Thermal voltage in volts. Defaults to ``25.85e-3``.

    """
    vbe = signals.b - signals.e
    vbc = signals.b - signals.c

    alpha_f = BetaF / (1.0 + BetaF)
    alpha_r = BetaR / (1.0 + BetaR)

    vbe_safe = jnp.clip(vbe, -5.0, 2.0)
    vbc_safe = jnp.clip(vbc, -5.0, 2.0)

    i_f = Is * (jnp.exp(vbe_safe / Vt) - 1.0)
    i_r = Is * (jnp.exp(vbc_safe / Vt) - 1.0)

    i_c = alpha_f * i_f - i_r
    i_e = -i_f + alpha_r * i_r
    i_b = -(i_c + i_e)

    return {"c": i_c, "b": i_b, "e": i_e}, {}

@component(ports=("c", "b", "e"))
def BJT_NPN_Dynamic(
    signals: Signals,
    Is: float = 1e-12,
    BetaF: float = 100.0,
    BetaR: float = 1.0,
    Vt: float = 25.85e-3,
    Cje: float = 1e-12,
    Cjc: float = 1e-12,
    Vje: float = 0.75,
    Vjc: float = 0.75,
    Mje: float = 0.33,
    Mjc: float = 0.33,
    Tf: float = 0.0,
    Tr: float = 0.0,
) -> PhysicsReturn:
    """NPN BJT with the Ebers-Moll DC model and first-order junction charge dynamics.

    DC currents follow the transport form of the Ebers-Moll equations. Junction
    voltages are clipped to ``[-5, 2]`` V before exponentiation to prevent
    overflow during transient solver iterations.

    Charge dynamics combine two contributions at each junction:

    - **Depletion charge** — modelled as a nonlinear junction capacitance via
      :func:`_junction_charge`, using the standard abrupt/graded junction
      formula with ideality parameters ``Vj`` and ``Mj``.
    - **Diffusion charge** — proportional to the junction current scaled by the
      transit time (``Tf`` for BE, ``Tr`` for BC``), representing minority
      carrier storage in the base.

    Args:
        signals: Port voltages at collector (``c``), base (``b``), and
            emitter (``e``).
        Is: Saturation current in amperes. Defaults to ``1e-12``.
        BetaF: Forward common-emitter current gain. Defaults to ``100``.
        BetaR: Reverse common-emitter current gain. Defaults to ``1``.
        Vt: Thermal voltage in volts. Defaults to ``25.85e-3`` (room temperature).
        Cje: Zero-bias BE junction capacitance in farads. Defaults to ``1e-12``.
        Cjc: Zero-bias BC junction capacitance in farads. Defaults to ``1e-12``.
        Vje: BE built-in junction potential in volts. Defaults to ``0.75``.
        Vjc: BC built-in junction potential in volts. Defaults to ``0.75``.
        Mje: BE junction grading coefficient. Defaults to ``0.33``.
        Mjc: BC junction grading coefficient. Defaults to ``0.33``.
        Tf: Forward transit time in seconds. Defaults to ``0`` (no BE diffusion charge).
        Tr: Reverse transit time in seconds. Defaults to ``0`` (no BC diffusion charge).

    Returns:
        A two-tuple ``(f, q)`` where:

        - **f** — DC current dict ``{"c": i_c, "b": i_b, "e": i_e}``.
        - **q** — Junction charge dict ``{"c": Q_collector, "b": Q_base, "e": Q_emitter}``,
            where ``Q_base = Q_be_total + Q_bc_total``, ``Q_collector = -Q_bc_total``,
            and ``Q_emitter = -Q_be_total``.

    """
    vbe = signals.b - signals.e
    vbc = signals.b - signals.c

    vbe_safe = jnp.clip(vbe, -5.0, 2.0)
    vbc_safe = jnp.clip(vbc, -5.0, 2.0)

    alpha_f = BetaF / (1.0 + BetaF)
    alpha_r = BetaR / (1.0 + BetaR)

    i_f = Is * (jnp.exp(vbe_safe / Vt) - 1.0)
    i_r = Is * (jnp.exp(vbc_safe / Vt) - 1.0)

    i_c = alpha_f * i_f - i_r
    i_e = -i_f + alpha_r * i_r
    i_b = -(i_c + i_e)

    f_dict = {"c": i_c, "b": i_b, "e": i_e}

    Qje_depl = _junction_charge(vbe, Cje, Vje, Mje)
    Qjc_depl = _junction_charge(vbc, Cjc, Vjc, Mjc)

    Qbe_diff = Tf * i_f
    Qbc_diff = Tr * i_r

    Q_be_total = Qje_depl + Qbe_diff
    Q_bc_total = Qjc_depl + Qbc_diff

    Q_base = Q_be_total + Q_bc_total
    Q_collector = -Q_bc_total
    Q_emitter = -Q_be_total

    return f_dict, {"c": Q_collector, "b": Q_base, "e": Q_emitter}

@component(ports=("out_p", "out_m", "in_p", "in_m"), states=("i_ctrl",), holomorphic=True)
def CCCS(signals: Signals, alpha: float = 1.0) -> PhysicsReturn:
    """Current Controlled Current Source: I_out = alpha * i_ctrl, V_in = 0 (short)."""
    eq_in = signals.in_p - signals.in_m
    i_out = alpha * signals.i_ctrl
    return {
        "out_p": i_out,
        "out_m": -i_out,
        "in_p": signals.i_ctrl,
        "in_m": -signals.i_ctrl,
        "i_ctrl": eq_in,
    }, {}

@component(ports=("out_p", "out_m", "in_p", "in_m"), states=("i_src", "i_ctrl"), holomorphic=True)
def CCVS(signals: Signals, R: float = 1.0) -> PhysicsReturn:
    """Current Controlled Voltage Source: V_out = R * i_ctrl, V_in = 0 (short)."""
    eq_in = signals.in_p - signals.in_m
    eq_out = (signals.out_p - signals.out_m) - (R * signals.i_ctrl)
    return {
        "out_p": signals.i_src,
        "out_m": -signals.i_src,
        "in_p": signals.i_ctrl,
        "in_m": -signals.i_ctrl,
        "i_src": eq_out,
        "i_ctrl": eq_in,
    }, {}

@component(ports=("p1", "p2"), port_aliases=_PN_ALIASES, holomorphic=True)
def Capacitor(signals: Signals, C: float = 1e-12) -> PhysicsReturn:
    """Q = C * V.
    Returns Charge (q) so the solver computes I = dq/dt.
    """  # noqa: D205
    v_drop = signals.p1 - signals.p2
    q_val = C * v_drop
    return {}, {"p1": q_val, "p2": -q_val}

@component(
    ports=("p1", "p2", "s1", "s2"),
    states=("i1", "i2"),
    port_aliases={"p1": "P1", "p2": "P2", "s1": "S1", "s2": "S2"},
    holomorphic=True,
)
def CoupledInductors(signals: Signals, L1: float = 1e-9, L2: float = 1e-9, k: float = 0.5) -> PhysicsReturn:
    """Two magnetically coupled inductors with ``M = k * sqrt(L1 * L2)``.

    Dots are on ``p1`` and ``s1``; currents are positive into the dotted
    terminals. The parameters should satisfy ``L1, L2 > 0`` and ``|k| <= 1``.
    At ``k = 0`` the windings are independent. For ``|k| = 1`` the inductance
    matrix is singular.
    """
    m = k * jnp.sqrt(L1 * L2)
    v1 = signals.p1 - signals.p2
    v2 = signals.s1 - signals.s2
    return (
        {"p1": signals.i1, "p2": -signals.i1, "s1": signals.i2, "s2": -signals.i2, "i1": v1, "i2": v2},
        {"i1": -(L1 * signals.i1 + m * signals.i2), "i2": -(m * signals.i1 + L2 * signals.i2)},
    )

@component(ports=("p1", "p2"), amplitude_param="I", holomorphic=True, ac_source_stamp={"p1": -1.0, "p2": 1.0})
def CurrentSource(signals: Signals, I: float = 0.0) -> PhysicsReturn:
    """Constant current source."""
    return {"p1": I, "p2": -I}, {}

@component(ports=("p1", "p2"))
def Diode(signals: Signals, Is: float = 1e-12, n: float = 1.0, Vt: float = 25.85e-3) -> PhysicsReturn:
    """Ideal diode using the Shockley equation ``I = Is * (exp(Vd / n*Vt) - 1)``.

    Junction voltage is clipped to ``[-5, 5]`` V for numerical stability.

    Args:
        signals: Port voltages at anode (``p1``) and cathode (``p2``).
        Is: Saturation current in amperes. Defaults to ``1e-12``.
        n: Ideality factor. Defaults to ``1.0``.
        Vt: Thermal voltage in volts. Defaults to ``25.85e-3``.

    """
    vd = signals.p1 - signals.p2
    # Clip for numerical stability
    vd_safe = jnp.clip(vd, -5.0, 5.0)
    i = Is * (jnp.exp(vd_safe / (n * Vt)) - 1.0)
    return {"p1": i, "p2": -i}, {}

@component(ports=("out_p", "out_m", "in_p", "in_m"), states=("i_src",), holomorphic=True)
def IdealOpAmp(signals: Signals, A: float = 1e6) -> PhysicsReturn:
    """Ideal Op Amp."""
    constraint = (signals.out_p - signals.out_m) - A * (signals.in_p - signals.in_m)
    return {
        "out_p": signals.i_src,
        "out_m": -signals.i_src,
        "in_p": 0.0,
        "in_m": 0.0,
        "i_src": constraint,
    }, {}

@component(
    ports=("p1", "p2", "s1", "s2"),
    states=("i_p",),
    port_aliases={"p1": "P1", "p2": "P2", "s1": "S1", "s2": "S2"},
    holomorphic=True,
)
def IdealTransformer(signals: Signals, n: float = 1.0) -> PhysicsReturn:
    """Ideal transformer with turns ratio ``n = N1 / N2`` (``n > 0``).

    Dots are on ``p1`` and ``s1``. It passes DC; add a magnetizing inductor
    across the primary to model low-frequency roll-off.
    """
    v1 = signals.p1 - signals.p2
    v2 = signals.s1 - signals.s2
    return {
        "p1": signals.i_p,
        "p2": -signals.i_p,
        "s1": -n * signals.i_p,
        "s2": n * signals.i_p,
        "i_p": v1 - n * v2,
    }, {}

@component(ports=("p1", "p2"), states=("i_L",), port_aliases=_PN_ALIASES, holomorphic=True)
def Inductor(signals: Signals, L: float = 1e-9) -> PhysicsReturn:
    """V = L * di/dt formulated via flux: f['i_L'] = V, q['i_L'] = -L*i_L."""
    v_drop = signals.p1 - signals.p2
    return ({"p1": signals.i_L, "p2": -signals.i_L, "i_L": v_drop}, {"i_L": -L * signals.i_L})

@component(ports=("d", "g", "s"))
def NMOS(
    signals: Signals,
    Kp: float = 2e-5,
    W: float = 10e-6,
    L: float = 1e-6,
    Vth: float = 1.0,
    lam: float = 0.0,
) -> PhysicsReturn:
    """N-channel MOSFET with square-law DC model and channel-length modulation.

    Gate current is zero (infinite input impedance).

    Args:
        signals: Port voltages at drain (``d``), gate (``g``), and source (``s``).
        Kp: Process transconductance parameter in A/V². Defaults to ``2e-5``.
        W: Gate width in metres. Defaults to ``10e-6``.
        L: Gate length in metres. Defaults to ``1e-6``.
        Vth: Threshold voltage in volts. Defaults to ``1.0``.
        lam: Channel-length modulation coefficient in V⁻¹. Defaults to ``0``.

    """
    i_ds = _nmos_current(signals.d, signals.g, signals.s, Kp, W, L, Vth, lam)
    return {"d": i_ds, "g": 0.0, "s": -i_ds}, {}

@component(ports=("d", "g", "s"))
def NMOSDynamic(
    signals: Signals,
    Kp: float = 2e-5,
    W: float = 10e-6,
    L: float = 1e-6,
    Vth: float = 1.0,
    lam: float = 0.0,
    Cox: float = 1e-3,
    Cgd_ov: float = 1e-15,
    Cgs_ov: float = 1e-15,
) -> PhysicsReturn:
    """NMOS with square-law DC model and Meyer gate capacitance model.

    Gate charge is split into bias-dependent intrinsic charge (Meyer) and
    linear overlap contributions at the drain and source.

    Args:
        signals: Port voltages at drain (``d``), gate (``g``), and source (``s``).
        Kp: Process transconductance parameter in A/V². Defaults to ``2e-5``.
        W: Gate width in metres. Defaults to ``10e-6``.
        L: Gate length in metres. Defaults to ``1e-6``.
        Vth: Threshold voltage in volts. Defaults to ``1.0``.
        lam: Channel-length modulation coefficient in V⁻¹. Defaults to ``0``.
        Cox: Gate oxide capacitance per unit area in F/m². Defaults to ``1e-3``.
        Cgd_ov: Gate-drain overlap capacitance in farads. Defaults to ``1e-15``.
        Cgs_ov: Gate-source overlap capacitance in farads. Defaults to ``1e-15``.

    """
    i_ds = _nmos_current(signals.d, signals.g, signals.s, Kp, W, L, Vth, lam)
    f_dict = {"d": i_ds, "g": 0.0, "s": -i_ds}

    vgs = signals.g - signals.s
    vds = signals.d - signals.s
    vgd = signals.g - signals.d

    WL = W * L
    Cox_total = Cox * WL
    v_over = vgs - Vth

    cutoff = vgs <= Vth
    saturation = vds >= v_over

    # Meyer Capacitance Logic
    Qg_cut = 0.0
    Qg_sat = (2.0 / 3.0) * Cox_total * v_over
    Qg_lin = 0.5 * Cox_total * v_over

    Qg = jnp.where(cutoff, Qg_cut, jnp.where(saturation, Qg_sat, Qg_lin))
    Qd = jnp.where(cutoff, 0.0, jnp.where(saturation, 0.0, -0.5 * Qg))
    Qs = -Qg - Qd

    Q_gate = Qg + Cgd_ov * vgd + Cgs_ov * vgs
    Q_drain = Qd - Cgd_ov * vgd
    Q_source = Qs - Cgs_ov * vgs

    return f_dict, {"d": Q_drain, "g": Q_gate, "s": Q_source}

@component(ports=("d", "g", "s"))
def PMOS(
    signals: Signals,
    Kp: float = 1e-5,
    W: float = 20e-6,
    L: float = 1e-6,
    Vth: float = -1.0,
    lam: float = 0.0,
) -> PhysicsReturn:
    """P-channel MOSFET with square-law DC model, formulated in terms of ``Vsg`` and ``Vsd``.

    Gate current is zero (infinite input impedance).

    Args:
        signals: Port voltages at drain (``d``), gate (``g``), and source (``s``).
        Kp: Process transconductance parameter in A/V². Defaults to ``1e-5``.
        W: Gate width in metres. Defaults to ``20e-6``.
        L: Gate length in metres. Defaults to ``1e-6``.
        Vth: Threshold voltage in volts (negative for PMOS). Defaults to ``-1.0``.
        lam: Channel-length modulation coefficient in V⁻¹. Defaults to ``0``.

    """
    vsg = signals.s - signals.g
    vsd = signals.s - signals.d

    beta = Kp * (W / L)
    vth_abs = jnp.abs(Vth)
    v_over = vsg - vth_abs

    linear_current = beta * (v_over * vsd - 0.5 * vsd**2) * (1 + lam * vsd)
    sat_current = (beta / 2.0) * (v_over**2) * (1 + lam * vsd)

    i_sd = jnp.where(vsg <= vth_abs, 0.0, jnp.where(vsd < v_over, linear_current, sat_current))

    return {"d": -i_sd, "g": 0.0, "s": i_sd}, {}

@source(
    ports=("p1", "p2"),
    states=("i_src",),
    amplitude_param="v2",
    holomorphic=True,
    ac_source_stamp={"i_src": 1.0},
    ideal_voltage_constraint_ports=("p1", "p2"),
)
def PulseVoltageSource(
    signals: Signals,
    t: float,
    v1: float = 0.0,
    v2: float = 1.0,
    td: float = 0.0,
    tr: float = 1e-9,
    tf: float = 1e-9,
    pw: float = 1e-3,
    per: float = 2e-3,
) -> PhysicsReturn:
    """SPICE-compatible PULSE voltage source: PULSE(v1 v2 td tr tf pw per).

    Parameters
    ----------
    v1  : initial (low) voltage
    v2  : pulsed (high) voltage
    td  : delay time before first edge
    tr  : rise time (v1 → v2)
    tf  : fall time (v2 → v1)
    pw  : pulse width (time at v2)
    per : period (must satisfy per >= td + tr + pw + tf)

    """
    t_shifted = jnp.where(t >= td, t - td, 0.0)
    t_in_period = jnp.mod(t_shifted, per)
    v_rising = v1 + (v2 - v1) * t_in_period / (tr + 1e-30)
    v_falling = v2 - (v2 - v1) * (t_in_period - tr - pw) / (tf + 1e-30)
    v_periodic = jnp.where(
        t_in_period < tr,
        v_rising,
        jnp.where(t_in_period < tr + pw, v2, jnp.where(t_in_period < tr + pw + tf, v_falling, v1)),
    )
    v_val = jnp.where(t < td, v1, v_periodic)
    constraint = (signals.p1 - signals.p2) - v_val
    return {"p1": signals.i_src, "p2": -signals.i_src, "i_src": constraint}, {}

@component(ports=("p1", "p2"), port_aliases=_PN_ALIASES, holomorphic=True)
def Resistor(signals: Signals, R: float = 1e3) -> PhysicsReturn:
    """Ohm's Law: I = V/R."""
    i = (signals.p1 - signals.p2) / (R + 1e-12)
    return {"p1": i, "p2": -i}, {}

@source(
    ports=("p1", "p2"),
    states=("i_src",),
    amplitude_param="V",
    holomorphic=True,
    ac_source_stamp={"i_src": 1.0},
    ideal_voltage_constraint_ports=("p1", "p2"),
)
def SmoothPulse(
    signals: Signals,
    t: float,
    V: float = 1.0,
    delay: float = 1e-9,
    tr: float = 1e-10,
) -> PhysicsReturn:
    """Sigmoid-smoothed pulse."""
    k = 10.0 / tr
    v_val = V * jnn.sigmoid(k * (t - delay))
    constraint = (signals.p1 - signals.p2) - v_val
    return {"p1": signals.i_src, "p2": -signals.i_src, "i_src": constraint}, {}

@component(
    ports=("p1", "p2"),
    states=("a1", "a2"),
    port_aliases=_PN_ALIASES,
    holomorphic=True,
)
def TransmissionLine(
    signals: Signals,
    tau: float = 1e-9,
    z0: float = 50.0,
    attenuation: float = 1.0,
) -> PhysicsReturn:
    """Matched bidirectional transmission line with exact fixed delay.

    ``a1`` and ``a2`` are the incident power-wave amplitudes at the two
    reference planes. The outgoing waves are the delayed incident wave from
    the opposite end. This stamp stays finite for an exactly lossless line
    and therefore avoids the singular ``S -> Y`` conversion of an ideal
    through connection.

    The same equations are interpreted by DC, transient, AC, and HB.
    """
    past = signals.at_delay(tau)
    b1 = attenuation * past.a2
    b2 = attenuation * past.a1
    i1 = (signals.a1 - b1) / z0
    i2 = (signals.a2 - b2) / z0
    return {
        "p1": i1,
        "p2": i2,
        "a1": signals.p1 - signals.a1 - b1,
        "a2": signals.p2 - signals.a2 - b2,
    }, {}

@component(ports=("out_p", "out_m", "ctrl_p", "ctrl_m"), holomorphic=True)
def VCCS(signals: Signals, G: float = 0.0) -> PhysicsReturn:
    """Voltage Controlled Current Source."""
    i = G * (signals.ctrl_p - signals.ctrl_m)
    return {"out_p": i, "out_m": -i, "ctrl_p": 0.0, "ctrl_m": 0.0}, {}

@component(ports=("out_p", "out_m", "ctrl_p", "ctrl_m"), states=("i_src",), holomorphic=True)
def VCVS(signals: Signals, A: float = 1.0) -> PhysicsReturn:
    """Voltage Controlled Voltage Source."""
    constraint = (signals.out_p - signals.out_m) - A * (signals.ctrl_p - signals.ctrl_m)
    return {
        "out_p": signals.i_src,
        "out_m": -signals.i_src,
        "ctrl_p": 0.0,
        "ctrl_m": 0.0,
        "i_src": constraint,
    }, {}

@component(ports=("p1", "p2", "cp", "cm"), holomorphic=True)
def VoltageControlledSwitch(signals: Signals, Ron: float = 1.0, Roff: float = 1e6, Vt: float = 0.0) -> PhysicsReturn:
    """Voltage Controlled Switch."""
    v_ctrl = signals.cp - signals.cm
    k = 10.0
    sig = jnn.sigmoid(k * (v_ctrl - Vt))

    g_on = 1.0 / Ron
    g_off = 1.0 / Roff
    g_eff = g_off + (g_on - g_off) * sig

    i = (signals.p1 - signals.p2) * g_eff
    return {"p1": i, "p2": -i, "cp": 0.0, "cm": 0.0}, {}

@source(
    ports=("p1", "p2"),
    states=("i_src",),
    amplitude_param="V",
    holomorphic=True,
    ac_source_stamp={"i_src": 1.0},
    ideal_voltage_constraint_ports=("p1", "p2"),
)
def VoltageSource(signals: Signals, t: float, V: float = 0.0, delay: float = 0.0) -> PhysicsReturn:
    """Step voltage source."""
    v_val = jnp.where(t >= delay, V, 0.0)
    constraint = (signals.p1 - signals.p2) - v_val
    return {"p1": signals.i_src, "p2": -signals.i_src, "i_src": constraint}, {}

@source(
    ports=("p1", "p2"),
    states=("i_src",),
    amplitude_param="V",
    holomorphic=True,
    ac_source_stamp={"i_src": 1.0},
    ideal_voltage_constraint_ports=("p1", "p2"),
)
def VoltageSourceAC(
    signals: Signals,
    t: float,
    V: float = 0.0,
    freq: float = 1e6,
    phase: float = 0.0,
    delay: float = 0.0,
) -> PhysicsReturn:
    """Sinusoidal voltage source."""
    omega = 2.0 * jnp.pi * freq
    v_ac = V * jnp.sin(omega * t + phase)
    v_val = jnp.where(t >= delay, v_ac, 0.0)
    constraint = (signals.p1 - signals.p2) - v_val
    return {"p1": signals.i_src, "p2": -signals.i_src, "i_src": constraint}, {}