Bifurcation diagram

import jax; jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from jax import vmap
from jaxkuramoto import theory
from jaxkuramoto.distribution import Cauchy, Normal, Uniform

import matplotlib.pyplot as plt

loc, gamma = 0.0, 1.0

dist_cauchy = Cauchy(loc=loc, gamma=gamma)
dist_normal = Normal(loc=loc, scale=gamma)
dist_uniform = Uniform(low=-gamma, high=gamma)

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Ks = jnp.arange(0.0, 5.0, 0.01)
orderparams_cauchy = theory.orderparam(Ks, dist_cauchy)
orderparams_normal = theory.orderparam(Ks, dist_normal)
orderparams_uniform = theory.orderparam(Ks, dist_uniform)

plt.figure(figsize=(8, 6))
plt.rcParams.update({"font.size": 16})
plt.xlabel("K")
plt.ylabel("Order parameter")
plt.xlim(0, 5)
plt.plot(Ks, orderparams_cauchy, label="Cauchy")
plt.plot(Ks, orderparams_normal, label="Normal")
plt.plot(Ks, orderparams_uniform, label="Uniform")
plt.legend()

<matplotlib.legend.Legend at 0x7f390c0accd0>