Package Reference#

Kuramoto model#

class jaxkuramoto.Kuramoto(omegas: jax.Array, K: float)[source]#

Kuramoto model.

orderparameter(t: float, thetas: jax.Array) → float[source]#

Order parameter.

Parameters

  • t (float) – time

  • thetas (jnp.ndarray) – Oscillator phases.

Returns

Order parameter.

Return type

float

vector_fn(t, thetas: jax.Array) → jax.Array[source]#

Vector field of Kuramoto model.

Parameters

  • t (float) – time

  • thetas (jnp.ndarray) – Oscillator phases.

Returns

Derivatives of oscillator phases.

Return type

jnp.ndarray

class jaxkuramoto.SakaguchiKuramoto(omegas: jax.Array, K: float, alpha: float)[source]#

Sakaguchi-Kuramoto model.

orderparameter(t: float, thetas: jax.Array) → float[source]#

Order parameter.

Parameters

  • t (float) – time

  • thetas (jnp.ndarray) – Oscillator phases.

Returns

Order parameter.

Return type

float

vector_fn(t, thetas: jax.Array) → jax.Array[source]#

Vector field of Kuramoto-Sakaguchi model.

Parameters

  • t (float) – time

  • thetas (jnp.ndarray) – Oscillator phases.

Returns

Derivatives of oscillator phases.

Return type

jnp.ndarray

Solver#

ODE solver#

jaxkuramoto.solver.euler(func: Callable[[float, jax.Array], jax.Array], t: float, dt: float, state: jax.Array)[source]#

Euler method.

Parameters

  • func – A function of the form func(t, x) -> dx/dt.

  • t – Current time.

  • dt – Time step.

  • state – Current state.

Returns

Next state.

jaxkuramoto.solver.runge_kutta(func: Callable[[float, jax.Array], jax.Array], t: float, dt: float, state: jax.Array)[source]#

Runge Kutta method.

Parameters

  • func – A function of the form func(t, x) -> dx/dt.

  • t – Current time.

  • dt – Time step.

  • state – Current state.

Returns

Next state.

Integral#

jaxkuramoto.solver.integral_fn(func, a, minval, maxval, n) → float[source]#

Integrate a function from minval to maxval using the trapezoidal rule.

Parameters

  • func – A function of the form func(x, a) -> y.

  • a – The parameter of the function.

  • minval – The lower bound of the integral.

  • maxval – The upper bound of the integral.

  • n – The number of trapezoids to use.

Returns

The value of the integral.

Fixed point solver#

jaxkuramoto.solver.fixed_point(func, a, x_guess, eps=1e-06)[source]#

Find the fixed point of a function using Newton’s method.

Parameters

  • func – A function of the form func(a, x) -> y.

  • a – The parameter of the function.

  • x_guess – The initial guess for the fixed point.

  • eps – The tolerance for the fixed point solver.

Returns

The fixed point.

ODE Integration#

jaxkuramoto.odeint(vector_fn: Callable[[float, jax.Array], jax.Array], solver: Callable[[Callable[[float, jax.Array], jax.Array], float, float, jax.Array], jax.Array], t0: float, t1: float, dt: float, init_state: jax.Array, observable_fn: Optional[Callable[[jax.Array], jax.Array]] = None)[source]#

Integrate the ODE.

Parameters

  • vector_fn (VECTOR_FN) – Vector field of the ODE.

  • solver (SOLVER) – Solver of the ODE.

  • t0 (float) – Initial time.

  • t1 (float) – Final time.

  • dt (float) – Time step.

  • init_state (jnp.ndarray) – Initial state.

  • observable_fn (OBSERVABLE_FN, optional) – Function to calculate observables. Defaults to None.

Returns

Solution class.

Return type

Solution

Raises

ValueError – If t0 >= t1 or dt <= 0.

Solution#

class jaxkuramoto.Solution(t0: float, t1: float, ts: jax.Array, init_state: jax.Array, final_state: jax.Array, observables: jax.Array)[source]#

Solution class.

t0#

Initial time.

Type

float

t1#

Final time.

Type

float

ts#

Time array.

Type

jnp.ndarray

init_state#

Initial state.

Type

jnp.ndarray

final_state#

Final state.

Type

jnp.ndarray

observables#

Observables.

Type

jnp.ndarray

Theory#

jaxkuramoto.theory.orderparam(K, dist: jaxkuramoto.distribution.base.Distribution, n=1000, r_guess=1.0, eps=1e-06)[source]#

Solve the self-consistent equation for the Kuramoto model and return the order parameter.

Parameters

  • K – The coupling strength.

  • dist – A function of the form pdf_fn(x) -> y.

  • n – Number of trapezoids to use in the integral.

  • r_guess – The initial guess for the order parameter.

  • eps – The tolerance for the fixed point solver.

Returns

Order parameter.

jaxkuramoto.theory.critical_point(pdf_fn, loc=0.0)[source]#

Find the critical coupling strength for the Kuramoto model.

Parameters

  • pdf_fn – A function of the form pdf_fn(x) -> y.

  • loc – The center of the distribution.

Returns

Critical coupling strength Kc.

Notes

pdf_fn must be normalized, symmetric, and have a single maximum at loc.

class jaxkuramoto.theory.OttAntonsen(dist: jaxkuramoto.distribution.base.Distribution, K: float)[source]#

Ott-Antonsen reduction of the Kuramoto model.