import matplotlib.pyplot as plt
import numpy as np
from differt2d.logic import activation, hard_sigmoid, sigmoid
from jax import grad, vmap

x = np.linspace(-5, +5, 200)

_, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=[6.4, 8])

for name, function in [("sigmoid", sigmoid), ("hard_sigmoid", hard_sigmoid)]:
    def f(x):
        return activation(x, alpha=1.5, function=function)

    y = f(x)
    dydx = vmap(grad(f))(x)
    _ = ax1.plot(x, y, "--", label=f"{name}")
    _ = ax2.plot(x, dydx, "-", label=f"{name}")

ax2.set_xlabel("$x$")
ax1.set_ylabel("$f(x)$")
ax2.set_ylabel(r"$\frac{\partial f(x)}{\partial x}$")
plt.legend()
plt.tight_layout()
plt.show()  # doctest: +SKIP