Example-07: FMA (differentiable)

# In this example, FMA indicator derivative is used for chaos identification

# Import

import numpy

from tqdm import tqdm

import jax
from jax import jit
from jax import vmap
from jax import jacrev

# Test symplectic mapping and corresponding inverse

from tohubohu.util import forward4D

# Frequency factory

from tohubohu import exponential
from tohubohu import frequency

# FMA factory

from tohubohu import fma

# Plotting

from matplotlib import pyplot as plt
from matplotlib import colormaps

cmap = colormaps.get_cmap('viridis')
cmap.set_bad(color='lightgray')

# Set data type

jax.config.update("jax_enable_x64", True)

# Set device

device, *_ = jax.devices('gpu')
jax.config.update('jax_default_device', device)

# Set mapping parameters

nux, nuy = 0.168, 0.201
mux, muy = 2*jax.numpy.pi*nux, 2*jax.numpy.pi*nuy
cx, sx, cy, sy = jax.numpy.cos(mux), jax.numpy.sin(mux), jax.numpy.cos(muy), jax.numpy.sin(muy)
mu = 0.0

k = jax.numpy.asarray([cx, sx, cy, sy, mu])

# Set initial grid in (qx, qy) plane

n = 1001

qx = jax.numpy.linspace(0.0, 0.6, n)
qy = jax.numpy.linspace(0.0, 0.6, n)
qs = jax.numpy.stack(jax.numpy.meshgrid(qx, qy, indexing='ij')).swapaxes(-1, 0).reshape(n*n, -1)
ps = jax.numpy.full_like(qs, 1.0E-12)
xs = jax.numpy.hstack([qs, ps])

# Set indicator

ws = exponential(2**10)

length = 2**1
fn = jit(fma(length, ws, forward4D))

@jit
def gn(x):
    return jax.numpy.log10(1.0E-16 + jax.numpy.sqrt(jax.numpy.sum(jax.numpy.diff(fn(x, k).T)**2)))

# Evaluate indicator

out = vmap(gn)(xs)

# Winsorize data

data = numpy.array(out)
data[data < -12.0] = -12.0
data[data > -3.0] = -3.0
data = data.reshape(n, n)

# Plot

plt.figure(figsize=(8, 8))
plt.imshow(data, aspect='equal', vmin=-12.0, vmax=-3.0, origin='lower', cmap=cmap, interpolation='nearest', extent=(0., 0.6, 0., 0.6))
plt.xlabel('qx')
plt.ylabel('qy')
plt.tight_layout()
plt.show()

# Set FMA indicator derivative

@jit
def hn(x):
    return jax.numpy.log10(1.0E-16 + jax.numpy.linalg.norm(jacrev(lambda x: jax.numpy.sqrt(jax.numpy.sum(jax.numpy.diff(fn(x, k).T)**2)))(x)))

# Set initial grid in (qx, qy) plane

n = 1001

qx = jax.numpy.linspace(0.0, 0.6, n)
qy = jax.numpy.linspace(0.0, 0.6, n)
qs = jax.numpy.stack(jax.numpy.meshgrid(qx, qy, indexing='ij')).swapaxes(-1, 0).reshape(n*n, -1)
ps = jax.numpy.full_like(qs, 1.0E-12)
xs = jax.numpy.hstack([qs, ps])
xs = jax.numpy.array_split(xs, n)

# Evaluate indicator

xb, *xr = xs
fj = jit(vmap(hn))
out = [fj(xb)]

for xb in tqdm(xr):
    out.append(fj(xb))
out = jax.numpy.concatenate(out)

# Winsorize data

data = numpy.array(out)
data[data < -16.0] = -16.0
data[data > 16.0] = 16
data = data.reshape(n, n)

# Plot

plt.figure(figsize=(8, 8))
plt.imshow(data, aspect='equal', vmin=-16.0, vmax=16.0, origin='lower', cmap=cmap, interpolation='nearest', extent=(0., 0.6, 0., 0.6))
plt.xlabel('qx')
plt.ylabel('qy')
plt.tight_layout()
plt.show()

100%|█████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [11:07<00:00,  1.50it/s]