import numpy as np
from . import hamiltonian_list as hl
from scipy.integrate import solve_ivp
from scipy.interpolate import CubicSpline
from scipy.optimize import root
[docs]class Hamiltonian():
r"""
A class for working with Hamiltonian systems.
This class provides a unified interface for various 1D Hamiltonian models,
including pendulum, Taylor series expansions, and resonant axisymmetric
Hamiltonians. It supports time-dependent parameters and offers methods
for evaluation, integration, and analysis of fixed points.
Parameters
----------
Htype : str
Type of the Hamiltonian. Must be one of:
* ``pendulum`` : Simple pendulum Hamiltonian
* ``taylor`` : Generalized Hamiltonian using Taylor series in J
* ``sqrt_taylor`` : Generalized Hamiltonian using Taylor series in √J
* ``axisym_res`` : Resonant axisymmetric Hamiltonian (2D)
* ``my_fun`` : User-defined Hamiltonian functions
ham : callable, optional
User-defined Hamiltonian function when Htype='my_fun'.
Signature: ham(J, theta, **kwargs)
dJdt : callable, optional
User-defined dJ/dt function when Htype='my_fun'.
Signature: dJdt(J, theta, **kwargs)
dthetadt : callable, optional
User-defined dθ/dt function when Htype='my_fun'.
Signature: dthetadt(J, theta, **kwargs)
t : (N,) numpy array, optional
Time array for time-dependent parameters. If provided, parameters
in **kwargs will be interpolated using cubic splines.
**kwargs : dict
Additional parameters passed to the Hamiltonian functions.
For time-dependent case (if ``t`` is defined), each parameter should be an array of
same length as ``t``.
Hamiltonian Types
-----------------
* **pendulum** — Simple pendulum Hamiltonian
.. math:: H = a(J-J_0)^2 + b\cos\theta
**Required parameters:** ``coef=[a, b]``, ``J0``
* **taylor** — Generalized Hamiltonian using Taylor series in J
.. math:: H = h_0 + \sum_i \left[h_i^c\cos(i\theta) + h_i^s\sin(i\theta)\right]
where :math:`h_i = a_0 + a_1 p + a_2 p^2 + \ldots` and :math:`p = (J-J_0)`
**Parameters:**
* ``n`` : (M,) int array — Fourier modes to include.
Positive integers → cosine terms :math:`h_i^c\cos(i\theta)`
Negative integers → sine terms :math:`h_i^s\sin(i\theta)`
Zero → constant term :math:`h_0`
* ``coef`` : (M, degree+1) array — Polynomial coefficients for each mode in ``n``.
Each row contains :math:`[a_0, a_1, a_2, \ldots]` for the corresponding mode.
* ``J0`` : float — Reference action for the Taylor expansion.
* **sqrt_taylor** — Generalized Hamiltonian using Taylor series in :math:`\sqrt{J}`
.. math:: H = h_0 + \sum_i \left[h_i^c\cos(i\theta) + h_i^s\sin(i\theta)\right]
where :math:`h_i = a_1 p + a_2 p^2 + a_3 p^3 + \\ldots` and :math:`p = \sqrt{J}`
**Parameters:**
* ``n`` : (M,) int array — Fourier modes to include (same format as ``taylor``).
* ``coef`` : (M, degree+1) array — Polynomial coefficients for each mode.
* **axisym_res** — Resonant axisymmetric Hamiltonian (2D)
.. math:: H = a_0 + a_R J_R + a_{R2} J_R^2 + a_z J_z + a_{z2} J_z^2 + a_{zR} J_z J_R + a_{zR}^{res} J_z J_R \cos(\theta_z - \theta_R)
**Parameters:**
``coef=[a0, aR, aR2, az, az2, azR, azR_res]``
* **my_fun** — User-defined Hamiltonian functions
**Parameters:**
``ham`` (callable), ``dJdt`` (callable), ``dthetadt`` (callable)
"""
_HTYPES = ['pendulum', 'sqrt_taylor', 'taylor', 'axisym_res', 'my_fun']
def __init__(
self,
Htype,
ham=None,
dJdt=None,
dthetadt=None,
t=None,
**kwargs):
"""
Initialise the class Hamiltonian
Parameters
----------
Htype : str
Type of the Hamiltonian. Must be one of:
* ``pendulum`` : Simple pendulum Hamiltonian
* ``taylor`` : Generalized Hamiltonian using Taylor series in J
* ``sqrt_taylor`` : Generalized Hamiltonian using Taylor series in √J
* ``axisym_res`` : Resonant axisymmetric Hamiltonian (2D)
* ``my_fun`` : User-defined Hamiltonian functions
ham : callable, optional
User-defined Hamiltonian function when Htype='my_fun'.
Signature: ham(J, theta, **kwargs)
dJdt : callable, optional
User-defined dJ/dt function when Htype='my_fun'.
Signature: dJdt(J, theta, **kwargs)
dthetadt : callable, optional
User-defined dθ/dt function when Htype='my_fun'.
Signature: dthetadt(J, theta, **kwargs)
t : (N,) numpy array, optional
Time array for time-dependent parameters. If provided, parameters
in **kwargs will be interpolated using cubic splines.
**kwargs : dict
Additional parameters passed to the Hamiltonian functions.
For time-dependent case, each parameter should be an array of
same length as t.
Raises
------
ValueError
if Htype is not in the list
"""
if Htype not in self._HTYPES:
raise ValueError(f"The Hamiltonian type (Htype) is list: {self._HTYPES}")
self._Htype = Htype
if Htype == 'pendulum':
self._ham = hl.H_pendulum
self._dJdt = hl.dJdt_pendulum
self._dthetadt = hl.dthetadt_pendulum
if Htype == 'sqrt_taylor':
self._ham = hl.H_sqrt_taylor
self._dJdt = hl.dJdt_sqrt_taylor
self._dthetadt = hl.dthetadt_sqrt_taylor
if Htype == 'taylor':
self._ham = hl.H_taylor
self._dJdt = hl.dJdt_taylor
self._dthetadt = hl.dthetadt_taylor
if Htype == 'axisym_res':
self._ham = hl.H_axisym_res
self._dJdt = hl.dJdt_axisym_res
self._dthetadt = hl.dthetadt_axisym_res
if Htype == 'my_fun':
if ham is not None:
self._ham = ham
if (dJdt is not None) and (dthetadt is not None):
self._dJdt = dJdt
self._dthetadt = dthetadt
self._times = t
if t is not None:
self.kwargs_spline = dict.fromkeys(kwargs)
_ = self.kwargs_spline.pop('n', None)
for key in self.kwargs_spline:
self.kwargs_spline[key] = CubicSpline(t, kwargs[key])
self._time_depend = True
else:
self._time_depend = False
self._kwargs = kwargs
def __call__(self, J, theta, t=None):
"""
Evaluate the Hamiltonian (wrapper for hamiltonian() method).
Parameters
----------
J : (N,) numpy array
The array of actions
theta : (N,) numpy array
The array of angles
t : float, optional
Time for time-dependent Hamiltonians., by default None
Returns
-------
float or numpy array
Hamiltonian value(s).
See Also
--------
hamiltonian : Equivalent method with same behavior.
"""
return self.hamiltonian(J, theta, t=t)
[docs] def hamiltonian(self, J, theta, t=None):
"""
Find the value of the Hamiltonian
Parameters
----------
J : (N,) numpy array
The array of actions
theta : (N,) numpy array
The array of angles
t : float, optional
time, by default None
Return
------
H : (N,) numpy array
The array of the Hamiltonian's values
"""
kwarg_t = self._kwargs
if self._time_depend:
t = self._times[0] if t is None else t
for key in self.kwargs_spline:
kwarg_t[key] = self.kwargs_spline[key](t)
return self._ham(J, theta, **kwarg_t)
[docs] def dJdt(self, J, theta, t=None):
"""
Find dJ/dt
Parameters
----------
J : (N,) numpy array
The array of actions
theta : (N,) numpy array
The array of angles
t : float, optional
time, by default None
Return
------
dJ/dt : (N,) numpy array
The array of actions' time derivative
"""
kwarg_t = self._kwargs
if self._time_depend:
t = self._times[0] if t is None else t
for key in self.kwargs_spline:
kwarg_t[key] = self.kwargs_spline[key](t)
return self._dJdt(J, theta, **kwarg_t)
[docs] def dthetadt(self, J, theta, t=None):
"""
Find the value of the Hamiltonian
Parameters
----------
J : (N,) numpy array
The array of actions
theta : (N,) numpy array
The array of angles
t : float, optional
time, by default None
Return
------
dJdt : (N,) numpy array
The array of actions' time derivative dJ/dt
"""
kwarg_t = self._kwargs
if self._time_depend:
t = self._times[0] if t is None else t
for key in self.kwargs_spline:
kwarg_t[key] = self.kwargs_spline[key](t)
return self._dthetadt(J, theta, **kwarg_t)
[docs] def derivative(self, J, theta, t=None):
"""
Find time derivatives of action and angle
dJ/dt = -dH/dθ, dθ/dt = dH/dJ
Parameters
----------
J : (N,) numpy array
The array of actions
theta : (N,) numpy array
The array of angles
t : float, optional
time, by default None
Returns
-------
dJdt : (N,) numpy array
The array of actions' time derivative dJ/dt
dthetadt : (N,) numpy array
The array of angles' time derivative dθ/dt
"""
dJdt = self.dJdt(J, theta, t=t)
dthetadt = self.dthetadt(J, theta, t=t)
return dJdt, dthetadt
[docs] def integrate(self, J0, theta0, t0=0.0, Tint=100, Nint=10000, rtol=10**-10, atol=10**-12):
"""
Integrate Hamiltonian trajectories
Parameters
----------
J0 : (N,) or (N,2) numpy array
Initial actions
For 2D Hamiltonians (``axisym_res``) array of shape (N,2) containing ``[JR, Jz]``
theta0 : (N,) numpy array
Initial angles
t0 : float, optional
Initial time. Default: 0.0
Tint : int, optional
Total integration time. Default: 100
Nint : int, optional
Number of time steps for output. Default: 10000
rtol, atol : float, optional
Relative/absolute tolerance for ``solve_ivp``. Default: 1e-10, 1e-12
Returns
-------
t_eval : (M,) array
Time points from t0 to t0+Tint.
(J, theta) : tuple of arrays
Evolution of actions and angles.
* 1D: J.shape = (N, M), theta.shape = (N, M)
* 2D: J.shape = (2, N, M) ([JR, Jz]), theta.shape = (N, M)
Notes
-----
Uses ``scipy.integrate.solve_ivp`` with DOP853 method.
"""
# nJ = J0.ndim
# Transform for solve_ivp
def ham_system(t, x):
J, theta = x.reshape((len(x)//2, 2)).T
dJdt, dthetadt = self.derivative(J, theta, t=t)
return np.column_stack((dJdt, dthetadt)).reshape(-1)
J0 = np.atleast_1d(J0)
theta0 = np.atleast_1d(theta0)
x0 = np.column_stack((J0, theta0)).reshape(-1)
t_eval = np.linspace(t0, t0+Tint, Nint)
sol = solve_ivp(
ham_system,
[t0, t0+Tint],
x0, # Initial condition
t_eval=t_eval,
method='DOP853',
vectorized=True,
rtol=rtol,
atol=atol
)
J = sol.y[::2]
theta = sol.y[1::2]
return t_eval, (J, theta)
[docs] def jacobian(self, J, theta, t=None, eps=10**-6):
"""
Compute second derivatives (Jacobian elements) of the Hamiltonian using finite differences.
The returned values correspond to:
.. math::
\\frac{\\partial^2 H}{\\partial J^2} = \\frac{d}{dJ}\\left(\\frac{d\\theta}{dt}\\right)
.. math::
\\frac{\\partial^2 H}{\\partial \\theta^2} = -\\frac{d}{d\\theta}\\left(\\frac{dJ}{dt}\\right)
.. math::
\\frac{\\partial^2 H}{\\partial J \\partial \\theta} = -\\frac{d}{dJ}\\left(\\frac{dJ}{dt}\\right)
These derivatives are useful for stability analysis and determining
the nature of fixed points (elliptic/hyperbolic).
Parameters
----------
J : (N,) numpy array
Action values.
theta : (N,) numpy array
Angle values.
t : float, optional
Time for time-dependent Hamiltonians.
eps : float, optional
Step size for finite difference approximation. Default: 1e-6
Returns
-------
d2HdJ2 : (N,) numpy array
Second derivative with respect to action.
d2Hdtheta2 : (N,) numpy array
Second derivative with respect to angle.
d2HdJdtheta : (N,) numpy array
Mixed second derivative.
Notes
-----
Uses central finite differences: :math:`\dfrac{f(x+\\epsilon) - f(x-\\epsilon)}{2\\epsilon}`
"""
d2HdJ2 = (self.dthetadt(J+eps, theta, t=t) - self.dthetadt(J-eps, theta, t=t)) / (2*eps)
d2Hdtheta2 = -(self.dJdt(J, theta+eps, t=t) - self.dJdt(J, theta-eps, t=t)) / (2*eps)
d2HdJdtheta = -(self.dJdt(J+eps, theta, t=t) - self.dJdt(J-eps, theta, t=t)) / (2*eps)
return d2HdJ2, d2Hdtheta2, d2HdJdtheta
[docs] def fix_points(self, J_range=[0.001, 0.5], Nstart=100, t=None, tol=10**-2):
"""fix_points
Parameters
----------
J_range : list, optional
range of actions, by default [0, 0.5]
Nstart : int, optional
number of initial random points, by default 100
t : float, optional
time, by default None
tol : float, optional
tolerance between roots, by default 10**-2
Return
------
fix_point: dict
J, theta, type(stable, unstable), H, librating time
"""
J0 = np.random.random(Nstart)*(J_range[1] - J_range[0]) + \
J_range[0]
theta0 = np.random.random(Nstart)*2*np.pi
x0_tab = np.zeros((Nstart, 2))
x0_tab[:, 0] = J0
x0_tab[:, 1] = theta0
J_fix = np.zeros_like(J0)*np.nan
theta_fix = np.zeros_like(J0)*np.nan
# Find derivative
def derivative_for_root(x):
J = x[0]
theta = x[1]
dJdt, dthetadt = self.derivative(J, theta, t=t)
return np.hstack((dJdt, dthetadt))
# Solve Nstart equations
for i, x0 in enumerate(x0_tab):
res = root(derivative_for_root, x0=x0)
if res.success:
J_fix[i] = res.x[0]
theta_fix[i] = res.x[1] % (2*np.pi)
# Delete not success solution
J_fix = np.delete(J_fix, np.isnan(J_fix))
theta_fix = np.delete(theta_fix, np.isnan(theta_fix))
# Delete repeated roots
i = 0
while True:
rtol = ((J_fix[i] - J_fix)**2 +
((theta_fix[i] - theta_fix + np.pi) % (2*np.pi)-np.pi)**2)**0.5
mask = (rtol < tol)
mask[i] = False
J_fix = np.delete(J_fix, mask)
theta_fix = np.delete(theta_fix, mask)
i += 1
if i >= len(J_fix):
break
d2HdJ2, d2Hdtheta2, d2HdJdtheta = self.jacobian(J_fix, theta_fix, t=t)
det_jac = (d2HdJdtheta**2 - d2HdJ2*d2Hdtheta2)
stable = det_jac < 0
T_libration = 2*np.pi / np.sqrt(-det_jac)
H_fix = self.hamiltonian(J_fix, theta_fix, t=t)
dict_fix = {'J': J_fix,
'theta': theta_fix,
'stable': stable,
'T_lib': T_libration,
'H': H_fix
}
return dict_fix
[docs] def find_J(self, H, theta, J0=0.1, t=None):
"""find_J
Find an action over a known angle and Hamiltonian
Parameters
----------
H : float
the value of the Hamiltonian
theta : (N,) numpy array
The array of actions
J0 : float, optional
the initial action, by default 0.1
t : float or None, optional
time, by default None
"""
def H_for_root(x, theta_fix=0.):
return self.hamiltonian(x, theta_fix, t=t)[0] - H
J_find = lambda x: root(H_for_root, J0, args=(x)).x
J_find = np.vectorize(J_find)
return J_find(theta)
[docs] def find_theta(self, H, J, theta0=0., t=None):
"""find_theta
Find an action over a known action and Hamiltonian
Parameters
----------
H : float
the value of the Hamiltonian
J : (N,) numpy array
The array of actions
theta0 : float, optional
the initial angle, by default 0.
t : float or None, optional
time, by default None
"""
def H_for_root(x, J0=0.):
return self.hamiltonian(J0, x, t=t)[0] - H
theta_find = lambda x: root(H_for_root, theta0, args=(x)).x
theta_find = np.vectorize(theta_find)
return theta_find(J)