# Source code for gala.dynamics.actionangle

"""
Utilities for estimating actions and angles for an arbitrary orbit in an
arbitrary potential.
"""

# Standard library
import time
import warnings

# Third-party
import numpy as np
from astropy import log as logger
from scipy.linalg import solve
from scipy.optimize import minimize, leastsq

# Project
from ..potential import HarmonicOscillatorPotential, IsochronePotential

__all__ = ['generate_n_vectors', 'fit_isochrone',
'fit_harmonic_oscillator', 'fit_toy_potential', 'check_angle_sampling',
'find_actions']

[docs]def generate_n_vectors(N_max, dx=1, dy=1, dz=1, half_lattice=True):
r"""
Generate integer vectors, :math:\boldsymbol{n}, with
:math:|\boldsymbol{n}| < N_{\rm max}.

If half_lattice=True, only return half of the three-dimensional
lattice. If the set N = {(i,j,k)} defines the lattice, we restrict to
the cases such that (k > 0), (k = 0, j > 0), and
(k = 0, j = 0, i > 0).

.. todo::

Return shape should be (3,N) to be consistent.

Parameters
----------
N_max : int
Maximum norm of the integer vector.
dx : int
Step size in x direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
dy : int
Step size in y direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
dz : int
Step size in z direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
half_lattice : bool (optional)
Only return half of the 3D lattice.

Returns
-------
vecs : :class:numpy.ndarray
A 2D array of integers with :math:|\boldsymbol{n}| < N_{\rm max}
with shape (N,3).

"""
vecs = np.meshgrid(np.arange(-N_max, N_max+1, dx),
np.arange(-N_max, N_max+1, dy),
np.arange(-N_max, N_max+1, dz))
vecs = np.vstack(map(np.ravel, vecs)).T
vecs = vecs[np.linalg.norm(vecs, axis=1) <= N_max]

if half_lattice:
ix = ((vecs[:, 2] > 0) |
((vecs[:, 2] == 0) &
(vecs[:, 1] > 0)) |
((vecs[:, 2] == 0) &
(vecs[:, 1] == 0) &
(vecs[:, 0] > 0)))
vecs = vecs[ix]

vecs = np.array(sorted(vecs, key=lambda x: (x, x, x)))
return vecs

[docs]def fit_isochrone(orbit, m0=2E11, b0=1., minimize_kwargs=None):
r"""
Fit the toy Isochrone potential to the sum of the energy residuals relative
to the mean energy by minimizing the function

.. math::

f(m,b) = \sum_i (\frac{1}{2}v_i^2 + \Phi_{\rm iso}(x_i\,|\,m,b) - <E>)^2

TODO: This should fail if the Hamiltonian associated with the orbit has
a frame other than StaticFrame

Parameters
----------
orbit : ~gala.dynamics.Orbit
m0 : numeric (optional)
Initial mass guess.
b0 : numeric (optional)
Initial b guess.
minimize_kwargs : dict (optional)
Keyword arguments to pass through to scipy.optimize.minimize.

Returns
-------
m : float
Best-fit scale mass for the Isochrone potential.
b : float
Best-fit core radius for the Isochrone potential.

"""
pot = orbit.hamiltonian.potential
if pot is None:
raise ValueError("The orbit object must have an associated potential")

w = np.squeeze(orbit.w(pot.units))
if w.ndim > 2:
raise ValueError("Input orbit object must be a single orbit.")

def f(p, w):
logm, logb = p
potential = IsochronePotential(m=np.exp(logm), b=np.exp(logb),
units=pot.units)
H = (potential.value(w[:3]).decompose(pot.units).value +
0.5*np.sum(w[3:]**2, axis=0))
return np.sum(np.squeeze(H - np.mean(H))**2)

logm0 = np.log(m0)
logb0 = np.log(b0)

if minimize_kwargs is None:
minimize_kwargs = dict()
minimize_kwargs['x0'] = np.array([logm0, logb0])
res = minimize(f, args=(w,), **minimize_kwargs)

if not res.success:
raise ValueError("Failed to fit toy potential to orbit.")

logm, logb = np.abs(res.x)
m = np.exp(logm)
b = np.exp(logb)

return IsochronePotential(m=m, b=b, units=pot.units)

[docs]def fit_harmonic_oscillator(orbit, omega0=[1., 1, 1], minimize_kwargs=None):
r"""
Fit the toy harmonic oscillator potential to the sum of the energy
residuals relative to the mean energy by minimizing the function

.. math::

f(\boldsymbol{\omega}) = \sum_i (\frac{1}{2}v_i^2 + \Phi_{\rm sho}(x_i\,|\,\boldsymbol{\omega}) - <E>)^2

TODO: This should fail if the Hamiltonian associated with the orbit has
a frame other than StaticFrame

Parameters
----------
orbit : ~gala.dynamics.Orbit
omega0 : array_like (optional)
Initial frequency guess.
minimize_kwargs : dict (optional)
Keyword arguments to pass through to scipy.optimize.minimize.

Returns
-------
omegas : float
Best-fit harmonic oscillator frequencies.

"""
omega0 = np.atleast_1d(omega0)

pot = orbit.hamiltonian.potential
if pot is None:
raise ValueError("The orbit object must have an associated potential")

w = np.squeeze(orbit.w(pot.units))
if w.ndim > 2:
raise ValueError("Input orbit object must be a single orbit.")

def f(omega, w):
potential = HarmonicOscillatorPotential(omega=omega, units=pot.units)
H = (potential.value(w[:3]).decompose(pot.units).value +
0.5*np.sum(w[3:]**2, axis=0))
return np.sum(np.squeeze(H - np.mean(H))**2)

if minimize_kwargs is None:
minimize_kwargs = dict()
minimize_kwargs['x0'] = omega0
res = minimize(f, args=(w,), **minimize_kwargs)

if not res.success:
raise ValueError("Failed to fit toy potential to orbit.")

best_omega = np.abs(res.x)
return HarmonicOscillatorPotential(omega=best_omega, units=pot.units)

[docs]def fit_toy_potential(orbit, force_harmonic_oscillator=False):
"""
Fit a best fitting toy potential to the orbit provided. If the orbit is a
tube (loop) orbit, use the Isochrone potential. If the orbit is a box
potential, use the harmonic oscillator potential. An option is available to
force using the harmonic oscillator (force_harmonic_oscillator).

See the docstrings for ~gala.dynamics.fit_isochrone() and
~gala.dynamics.fit_harmonic_oscillator() for more information.

Parameters
----------
orbit : ~gala.dynamics.Orbit
force_harmonic_oscillator : bool (optional)
Force using the harmonic oscillator potential as the toy potential.

Returns
-------
potential : :class:~gala.potential.IsochronePotential or :class:~gala.potential.HarmonicOscillatorPotential
The best-fit potential object.

"""
circulation = orbit.circulation()
if np.any(circulation == 1) and not force_harmonic_oscillator:  # tube orbit
logger.debug("===== Tube orbit =====")
logger.debug("Using Isochrone toy potential")

toy_potential = fit_isochrone(orbit)
logger.debug("Best m={}, b={}".format(toy_potential.parameters['m'],
toy_potential.parameters['b']))

else:  # box orbit
logger.debug("===== Box orbit =====")
logger.debug("Using triaxial harmonic oscillator toy potential")

toy_potential = fit_harmonic_oscillator(orbit)
logger.debug("Best omegas ({})"
.format(toy_potential.parameters['omega']))

[docs]def check_angle_sampling(nvecs, angles):
"""
Returns a list of the index of elements of n which do not have adequate
toy angle coverage. The criterion is that we must have at least one sample
in each Nyquist box when we project the toy angles along the vector n.

Parameters
----------
nvecs : array_like
Array of integer vectors.
angles : array_like
Array of angles.

Returns
-------
failed_nvecs : :class:numpy.ndarray
Array of all integer vectors that failed checks. Has shape (N,3).
failures : :class:numpy.ndarray
Array of flags that designate whether this failed needing a longer
integration window (0) or finer sampling (1).

"""

failed_nvecs = []
failures = []

for i, vec in enumerate(nvecs):
# N = np.linalg.norm(vec)
# X = np.dot(angles,vec)
X = (angles*vec[:, None]).sum(axis=0)
diff = float(np.abs(X.max() - X.min()))

if diff < (2.*np.pi):
warnings.warn("Need a longer integration window for mode {0}"
.format(vec))
failed_nvecs.append(vec.tolist())
# P.append(2.*np.pi - diff)
failures.append(0)

elif (diff/len(X)) > np.pi:
warnings.warn("Need a finer sampling for mode {0}"
.format(str(vec)))
failed_nvecs.append(vec.tolist())
# P.append(np.pi - diff/len(X))
failures.append(1)

return np.array(failed_nvecs), np.array(failures)

def _action_prepare(aa, N_max, dx, dy, dz, sign=1., throw_out_modes=False):
"""
Given toy actions and angles, aa, compute the matrix A and
vector b to solve for the vector of "true" actions and generating
function values, x (see Equations 12-14 in Sanders & Binney (2014)).

.. todo::

Wrong shape for aa -- should be (6,n) as usual...

Parameters
----------
aa : array_like
Shape (6,ntimes) array of toy actions and angles.
N_max : int
Maximum norm of the integer vector.
dx : int
Step size in x direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
dy : int
Step size in y direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
dz : int
Step size in z direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
sign : numeric (optional)
Vector that defines direction of circulation about the axes.
"""

# unroll the angles so they increase continuously instead of wrap
angles = np.unwrap(aa[3:])

# generate integer vectors for fourier modes
nvecs = generate_n_vectors(N_max, dx, dy, dz)

# make sure we have enough angle coverage
modes, P = check_angle_sampling(nvecs, angles)

# throw out modes?
# if throw_out_modes:
#     nvecs = np.delete(nvecs, (modes,P), axis=0)

n = len(nvecs) + 3
b = np.zeros(shape=(n, ))
A = np.zeros(shape=(n, n))

# top left block matrix: identity matrix summed over timesteps
A[:3, :3] = aa.shape*np.identity(3)

actions = aa[:3]
angles = aa[3:]

# top right block matrix: transpose of C_nk matrix (Eq. 12)
C_T = 2.*nvecs.T * np.sum(np.cos(np.dot(nvecs, angles)), axis=-1)
A[:3,3:] = C_T
A[3:, :3] = C_T.T

# lower right block matrix: C_nk dotted with C_nk^T
cosv = np.cos(np.dot(nvecs, angles))
A[3:,3:] = 4.*np.dot(nvecs, nvecs.T)*np.einsum('it,jt->ij', cosv, cosv)

# b vector first three is just sum of toy actions
b[:3] = np.sum(actions, axis=1)

# rest of the vector is C dotted with actions
b[3:] = 2*np.sum(np.dot(nvecs, actions)*np.cos(np.dot(nvecs, angles)),
axis=1)

return A, b, nvecs

def _angle_prepare(aa, t, N_max, dx, dy, dz, sign=1.):
"""
Given toy actions and angles, aa, compute the matrix A and
vector b to solve for the vector of "true" angles, frequencies, and
generating function derivatives, x (see Appendix of
Sanders & Binney (2014)).

.. todo::

Wrong shape for aa -- should be (6,n) as usual...

Parameters
----------
aa : array_like
Shape (6,ntimes) array of toy actions and angles.
t : array_like
Array of times.
N_max : int
Maximum norm of the integer vector.
dx : int
Step size in x direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
dy : int
Step size in y direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
dz : int
Step size in z direction. Set to 1 for odd and even terms, set
to 2 for just even terms.
sign : numeric (optional)
Vector that defines direction of circulation about the axes.
"""

# unroll the angles so they increase continuously instead of wrap
angles = np.unwrap(aa[3:])

# generate integer vectors for fourier modes
nvecs = generate_n_vectors(N_max, dx, dy, dz)

# make sure we have enough angle coverage
modes, P = check_angle_sampling(nvecs, angles)

# TODO: throw out modes?
# if(throw_out_modes):
#     n_vectors = np.delete(n_vectors,check_each_direction(n_vectors,angs),axis=0)

nv = len(nvecs)
n = 3 + 3 + 3*nv # angle(0)'s, freqs, 3 derivatives of Sn

b = np.zeros(shape=(n,))
A = np.zeros(shape=(n, n))

# top left block matrix: identity matrix summed over timesteps
A[:3, :3] = aa.shape*np.identity(3)

# identity matrices summed over times
A[:3, 3:6] = A[3:6, :3] = np.sum(t)*np.identity(3)
A[3:6, 3:6] = np.sum(t*t)*np.identity(3)

# S1,2,3
A[6:6+nv, 0] = -2.*np.sum(np.sin(np.dot(nvecs, angles)), axis=1)
A[6+nv:6+2*nv, 1] = A[6:6+nv, 0]
A[6+2*nv:6+3*nv, 2] = A[6:6+nv, 0]

# t*S1,2,3
A[6:6+nv, 3] = -2.*np.sum(t[None, :]*np.sin(np.dot(nvecs, angles)),
axis=1)
A[6+nv:6+2*nv, 4] = A[6:6+nv, 3]
A[6+2*nv:6+3*nv, 5] = A[6:6+nv, 3]

# lower right block structure: S dot S^T
sinv = np.sin(np.dot(nvecs, angles))
SdotST = np.einsum('it,jt->ij', sinv, sinv)
A[6:6+nv, 6:6+nv] = A[6+nv:6+2*nv, 6+nv:6+2*nv] = \
A[6+2*nv:6+3*nv, 6+2*nv:6+3*nv] = 4*SdotST

# top rectangle
A[:6, :] = A[:, :6].T

b[:3] = np.sum(angles.T, axis=0)
b[3:6] = np.sum(t[:, None]*angles.T, axis=0)
b[6:6+nv] = -2.*np.sum(angles*np.sin(np.dot(nvecs, angles)), axis=1)
b[6+nv:6+2*nv] = -2.*np.sum(angles*np.sin(np.dot(nvecs, angles)),
axis=1)
b[6+2*nv:6+3*nv] = -2.*np.sum(angles*np.sin(np.dot(nvecs, angles)),
axis=1)

return A, b, nvecs

def _single_orbit_find_actions(orbit, N_max, toy_potential=None,
force_harmonic_oscillator=False):
"""
Find approximate actions and angles for samples of a phase-space orbit,
w, at times t. Uses toy potentials with known, analytic action-angle
transformations to approximate the true coordinates as a Fourier sum.

This code is adapted from Jason Sanders'
genfunc <https://github.com/jlsanders/genfunc>_

.. todo::

Wrong shape for w -- should be (6,n) as usual...

Parameters
----------
orbit : ~gala.dynamics.Orbit
N_max : int
Maximum integer Fourier mode vector length, |n|.
toy_potential : Potential (optional)
Fix the toy potential class.
force_harmonic_oscillator : bool (optional)
Force using the harmonic oscillator potential as the toy potential.
"""

if orbit.norbits > 1:
raise ValueError("must be a single orbit")

if toy_potential is None:
toy_potential = fit_toy_potential(
orbit, force_harmonic_oscillator=force_harmonic_oscillator)

else:
logger.debug("Using *fixed* toy potential: {}"
.format(toy_potential.parameters))

if isinstance(toy_potential, IsochronePotential):
orbit_align = orbit.align_circulation_with_z()
w = orbit_align.w()

dxyz = (1, 2, 2)
circ = np.sign(w[0, 0]*w[4, 0]-w[1, 0]*w[3, 0])
sign = np.array([1., circ, 1.])
orbit = orbit_align
elif isinstance(toy_potential, HarmonicOscillatorPotential):
dxyz = (2, 2, 2)
sign = 1.
w = orbit.w()
else:
raise ValueError("Invalid toy potential.")

t = orbit.t.value

# Now find toy actions and angles
aaf = toy_potential.action_angle(orbit)

if aaf.ndim > 2:
aa = np.vstack((aaf.value[..., 0], aaf.value[..., 0]))
else:
aa = np.vstack((aaf.value, aaf.value))

if np.any(np.isnan(aa)):
ix = ~np.any(np.isnan(aa), axis=0)
aa = aa[:, ix]
t = t[ix]
warnings.warn("NaN value in toy actions or angles!")
if sum(ix) > 1:
raise ValueError("Too many NaN value in toy actions or angles!")

t1 = time.time()
A, b, nvecs = _action_prepare(aa, N_max, dx=dxyz, dy=dxyz, dz=dxyz)
actions = np.array(solve(A,b))
logger.debug("Action solution found for N_max={}, size {} symmetric"
" matrix in {} seconds"
.format(N_max, len(actions), time.time()-t1))

t1 = time.time()
A, b, nvecs = _angle_prepare(aa, t, N_max, dx=dxyz,
dy=dxyz, dz=dxyz, sign=sign)
angles = np.array(solve(A, b))
logger.debug("Angle solution found for N_max={}, size {} symmetric"
" matrix in {} seconds"
.format(N_max, len(angles), time.time()-t1))

# Just some checks
if len(angles) > len(aa):
warnings.warn("More unknowns than equations!")

J = actions[:3]  # * sign
theta = angles[:3]
freqs = angles[3:6]  # * sign

return dict(actions=J*aaf.unit, angles=theta*aaf.unit,
freqs=freqs*aaf.unit,
Sn=actions[3:], dSn_dJ=angles[6:], nvecs=nvecs)

[docs]def find_actions(orbit, N_max, force_harmonic_oscillator=False, toy_potential=None):
r"""
Find approximate actions and angles for samples of a phase-space orbit.
Uses toy potentials with known, analytic action-angle transformations to
approximate the true coordinates as a Fourier sum.

This code is adapted from Jason Sanders'
genfunc <https://github.com/jlsanders/genfunc>_

Parameters
----------
orbit : ~gala.dynamics.Orbit
N_max : int
Maximum integer Fourier mode vector length, :math:|\boldsymbol{n}|.
force_harmonic_oscillator : bool (optional)
Force using the harmonic oscillator potential as the toy potential.
toy_potential : Potential (optional)
Fix the toy potential class.

Returns
-------
aaf : dict
A Python dictionary containing the actions, angles, frequencies, and
value of the generating function and derivatives for each integer
vector. Each value of the dictionary is a :class:numpy.ndarray or
:class:astropy.units.Quantity.

"""

if orbit.norbits == 1:
return _single_orbit_find_actions(
orbit, N_max,
force_harmonic_oscillator=force_harmonic_oscillator,
toy_potential=toy_potential)

else:
norbits = orbit.norbits
actions = np.zeros((3, norbits))
angles = np.zeros((3, norbits))
freqs = np.zeros((3, norbits))
for n in range(norbits):
aaf = _single_orbit_find_actions(
orbit[:, n], N_max,
force_harmonic_oscillator=force_harmonic_oscillator,
toy_potential=toy_potential)
actions[n] = aaf['actions'].value
angles[n] = aaf['angles'].value
freqs[n] = aaf['freqs'].value

return dict(actions=actions*aaf['actions'].unit,
angles=angles*aaf['angles'].unit,
freqs=freqs*aaf['freqs'].unit,
Sn=actions[3:], dSn=angles[6:], nvecs=aaf['nvecs'])

# def solve_hessian(relative_actions, relative_freqs):
#     """ Use ordinary least squares to solve for the Hessian, given a
#         set of actions and frequencies relative to the parent orbit.
#     """

# def compute_hessian(t, w, actions_kwargs={}):
#     """ Compute the Hessian (in action-space) of the given orbit

#     """

#     N = dJ.shape

#     Y = np.ravel(dF)
#     A = np.zeros((3*N,9))
#     A[::3,:3] = dJ
#     A[1::3,3:6] = dJ
#     A[2::3,6:9] = dJ

#     # Solve for 'parameters' - the Hessian elements
#     X,res,rank,s = np.linalg.lstsq(A, Y)

#     # Symmetrize
#     D0 = X.reshape(3,3)
#     D0[0,1] = D0[1,0] = (D0[0,1] + D0[1,0])/2.
#     D0[0,2] = D0[2,0] = (D0[0,2] + D0[2,0])/2.
#     D0[1,2] = D0[2,1] = (D0[1,2] + D0[2,1])/2.

#     print("Residual: " + str(res))

#     return D0,np.linalg.eigh(D0) # symmetric matrix