import numpy as np
import matplotlib.pyplot as plt

e_0 = 8.85E-14
k_ox = 3.9
e_ox = e_0 * k_ox

mu = 350                # 350 cm^2 / (V*s)
t_ox = 100E-8           # 100A

C_ox = e_ox / t_ox
W = 1.2
L = 0.6

beta = mu * C_ox * W/L

V_g_range = np.arange(6)
V_t = 0.7
V_dd = np.linspace(0, 5, 50)

def I_ds(V_gs, V_ds, V_t, beta):
    if V_gs < V_t:
        return np.zeros_like(V_ds)

    V_dsat = V_gs - V_t

    return np.where(V_ds < V_dsat,
                    beta * (V_gs - V_t - V_ds/2) * V_ds, 
                    beta * (V_gs - V_t)**2 / 2)

for V_s in (0, 2):
    for V_g in V_g_range:
        V_gs = V_g - V_s
        V_ds = V_dd - V_s
        plt.plot(V_ds,
                 1000*I_ds(V_gs, V_ds, V_t, beta),
                 label="$V_{gs}$ = "+str(V_gs))

    plt.xlabel("$V_{ds}$[V]")
    plt.ylabel("$I_{ds}$[mA]")
    plt.legend()
    plt.show()