Van der Pol equation

import jax.numpy as jnp
from prax import Oscillator
from jax.config import config; config.update("jax_enable_x64", True)

import matplotlib.pyplot as plt

Matplotlib is building the font cache; this may take a moment.

Exception ignored in: <matplotlib.ft2font.FT2Font object at 0x7ff14bdd3ea0>
Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/prax/envs/latest/lib/python3.8/site-packages/matplotlib/font_manager.py", line 1092, in addfont
    font = ft2font.FT2Font(path)
KeyboardInterrupt: 

class VanderPol(Oscillator):
    def __init__(self, mu, dt=0.01, eps=10**-5):
        super().__init__(n_dim=2, dt=dt, eps=eps)
        self.mu = mu

    def forward(self, state):
        x, y = state
        vx = y
        vy = self.mu * (1.0 - x*x) * y - x
        return jnp.array([vx, vy])

model = VanderPol(mu=0.2)

init_val = jnp.array([0.1, 0.2])
model.find_periodic_orbit(init_val)
print(model.period) # 6.3088767
plt.plot(model.ts, model.periodic_orbit)

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

6.308876695632844

[<matplotlib.lines.Line2D at 0x16a79b250>,
 <matplotlib.lines.Line2D at 0x16ad3b550>]

model.calc_phase_response()
plt.plot(model.ts, model.phase_response_curve)

[<matplotlib.lines.Line2D at 0x16b542580>,
 <matplotlib.lines.Line2D at 0x16b5d71c0>]