from math import pi
import numpy as np
import warnings
__all__ = ["integrate", "dissipation", "energy", "energy_variation", "power"]
def _averaging_wrapper(instance_with_state_to_average, average="spacetime"):
"""
Average with respect to time, space, spacetime, or discretization.
Parameters
----------
instance_with_state_to_average : Orbit
average : str
Which dimensions/values to average with respect to.
Returns
-------
ndarray or float :
The averages either along certain dimensions or all dimensions.
"""
if average == "space":
return (
1.0 / instance_with_state_to_average.x
) * instance_with_state_to_average.state.mean(axis=1)
elif average == "time":
return (
1.0 / instance_with_state_to_average.t
) * instance_with_state_to_average.state.mean(axis=0)
elif average == "spacetime":
# numpy average is over flattened array by default
return (
(1.0 / instance_with_state_to_average.t)
* (1.0 / instance_with_state_to_average.x)
* instance_with_state_to_average.state.mean()
)
elif average == "discrete_spacetime":
# numpy average is over flattened array by default
return (
(1.0 / instance_with_state_to_average.n)
* (1.0 / instance_with_state_to_average.m)
* instance_with_state_to_average.state.mean()
)
else:
return instance_with_state_to_average.state
[docs]def dissipation(orbit_instance, average="spacetime"):
"""
Average energy dissipation or corresponding field.
Parameters
----------
float or ndarray :
Spatiotemporal dissipation returned as a field, or averaged along an axis (ndarray) or spacetime (float).
Notes
-----
Dissipation $= u_xx^2$.
"""
return _averaging_wrapper(
orbit_instance.dx(order=2).transform(to="field") ** 2, average=average
)
[docs]def energy(orbit_instance, average="spacetime"):
"""
Average energy or corresponding field
Parameters
----------
float or ndarray :
Spatiotemporal power returned as a field, or averaged along an axis (ndarray) or spacetime (float).
Notes
-----
Energy = 1/2 u^2
"""
return _averaging_wrapper(
0.5 * orbit_instance.transform(to="field") ** 2, average=average
)
[docs]def energy_variation(orbit_instance, average="spacetime"):
"""
The field u_t * u whose spatial average should equal power - dissipation.
Parameters
----------
float or ndarray :
Spatiotemporal energy variation returned as a field, or averaged along an axis (ndarray) or spacetime (float).
Returns
-------
Field equivalent to u_t * u.
"""
return _averaging_wrapper(
orbit_instance.transform(to="field")
* orbit_instance.dt().transform(to="field"),
average=average,
)
[docs]def power(orbit_instance, average="spacetime"):
"""
Average power or corresponding field
Parameters
----------
float or ndarray :
Spatiotemporal power returned as a field, or averaged along an axis (ndarray) or spacetime (float).
Notes
-----
power = u_x^2
"""
return _averaging_wrapper(
orbit_instance.dx().transform(to="field") ** 2, average=average
)
[docs]def integrate(orbit_, **kwargs):
"""
Exponential time-differencing Runge-Kutta 4th order integration scheme.
Parameters
----------
orbit_ : Orbit
The Orbit which contains the initial starting point for the integration
kwargs : dict
verbose : bool
starting_point :
The row to select as the initial value.
integration_time : float
The total amount of time to integrate, if not provided then defaults to orbit's period (for reproduction
of orbit, typically).
step_size : float
The integration step
return_trajectory : bool
Whether to store and return all integration steps or only the endpoint.
Returns
-------
Orbit :
Orbit instance with either an integrated trajectory or its final value as a state.
Notes
-----
Adapter https://epubs.siam.org/doi/abs/10.1137/S1064827502410633?journalCode=sjoce3
By default, when input is an instance of relative periodic orbit then shift is calculated off of the
integrated trajectory. This will lead to plotting issues so unless desired, you should convert to the
base orbit type first.
"""
warnings.simplefilter(action="ignore", category=RuntimeWarning)
from scipy.fft import rfftfreq
verbose = kwargs.get("verbose", False)
orbit_ = orbit_.transform(to="spatial_modes")
integration_time = kwargs.get("integration_time", orbit_.t)
start_point = kwargs.get("starting_point", -1)
# Take the last row (t=0) or first row (t=T) so this works for relative periodic solutions as well.
orbit_t = orbit_.__class__(
state=orbit_.state[start_point, :].reshape(1, -1),
basis=orbit_.basis,
parameters=orbit_.parameters,
).transform(to="spatial_modes")
# stepsize
step_size = kwargs.get("step_size", 0.01)
m, x = orbit_t.m, orbit_t.x
q = (2 * pi * m / x) * rfftfreq(m)[1:-1].reshape(1, -1)
q = np.concatenate((q, q))
# Because N = 1, this is just the spatial matrices, negative sign b.c. other side of equation.
lin_diag = (q ** 2 - q ** 4).reshape(-1, 1)
E = np.exp(step_size * lin_diag)
E2 = np.exp(step_size * lin_diag / 2.0)
n_roots = 16
roots_of_unity = np.exp(
1.0j * pi * (np.arange(1, n_roots + 1, 1) - 0.5) / n_roots
).reshape(1, n_roots)
# Matrix quantities for exponential time differencing.
LR = step_size * np.tile(lin_diag, (1, n_roots)) + np.tile(
roots_of_unity, (orbit_t.shapes()[2][1], 1)
)
Q = step_size * np.real(np.mean((np.exp(LR / 2.0) - 1.0) / LR, axis=1))
f1 = step_size * np.real(
np.mean((-4.0 - LR + np.exp(LR) * (4.0 - 3.0 * LR + LR ** 2)) / LR ** 3, axis=1)
)
f2 = step_size * np.real(
np.mean((2.0 + LR + np.exp(LR) * (-2.0 + LR)) / LR ** 3, axis=1)
)
f3 = step_size * np.real(
np.mean((-4.0 - 3.0 * LR - LR ** 2 + np.exp(LR) * (4.0 - LR)) / LR ** 3, axis=1)
)
Q = Q.reshape(1, -1)
f1 = f1.reshape(1, -1)
f2 = f2.reshape(1, -1)
f3 = f3.reshape(1, -1)
E = E.reshape(1, -1)
E2 = E2.reshape(1, -1)
if orbit_t.__class__.__name__ == "AntisymmetricOrbitKS":
Q[: -orbit_t.m, :] = 0
f1[: -orbit_t.m, :] = 0
f2[: -orbit_t.m, :] = 0
f3[: -orbit_t.m, :] = 0
E[: -orbit_t.m, :] = 0
E2[: -orbit_t.m, :] = 0
u = orbit_t.transform(to="field").state
v = orbit_t.transform(to="spatial_modes")
nmax = int(integration_time / step_size)
if verbose:
print("Integration progress [", end="")
if kwargs.get("return_trajectory", True):
u = np.zeros([nmax, orbit_t.shapes()[0][1]])
for step in range(1, nmax + 1):
Nv = -0.5 * (v.transform(to="field") ** 2).dx(
computation_basis="spatial_modes", return_basis="spatial_modes"
)
a = v * E2 + Nv * Q
Na = -0.5 * (a.transform(to="field") ** 2).dx(
computation_basis="spatial_modes", return_basis="spatial_modes"
)
b = v * E2 + Na * Q
Nb = -0.5 * (b.transform(to="field") ** 2).dx(
computation_basis="spatial_modes", return_basis="spatial_modes"
)
c = a * E2 + (2 * Nb - Nv) * Q
Nc = -0.5 * (c.transform(to="field") ** 2).dx(
computation_basis="spatial_modes", return_basis="spatial_modes"
)
v = v * E + Nv * f1 + (2 * (Na + Nb)) * f2 + Nc * f3
if orbit_.__class__.__name__ in ["AntisymmetricOrbitKS", "EquilibriumOrbitKS"]:
v.real = 0
if kwargs.get("return_trajectory", True):
u[-step, :] = v.transform(to="field").state.ravel()
else:
u = v.transform(to="field").state
if not np.mod(step, nmax // 25) and verbose:
print("#", end="")
if verbose:
print("]", end="")
# By default do not assign spatial shift S.
if kwargs.get("return_trajectory", True):
int_orbit = orbit_.__class__(
state=u.reshape(nmax, -1),
basis="field",
parameters=(integration_time, orbit_.x, 0),
)
else:
int_orbit = orbit_.__class__(
state=u.reshape(1, -1),
basis="field",
parameters=(integration_time, orbit_.x, 0),
)
warnings.resetwarnings()
return int_orbit