-
Notifications
You must be signed in to change notification settings - Fork 95
/
actionAngleHarmonic.py
132 lines (113 loc) · 4.18 KB
/
actionAngleHarmonic.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
###############################################################################
# actionAngle: a Python module to calculate actions, angles, and frequencies
#
# class: actionAngleHarmonic
#
# Calculate actions-angle coordinates for the harmonic-oscillator
#
# methods:
# __call__: returns (j)
# actionsFreqs: returns (j,omega)
# actionsFreqsAngles: returns (j,omega,a)
#
###############################################################################
import numpy
from ..util import conversion
from .actionAngle import actionAngle
class actionAngleHarmonic(actionAngle):
"""Action-angle formalism for the one-dimensional harmonic oscillator"""
def __init__(self, *args, **kwargs):
"""
Initialize an actionAngleHarmonic object.
Parameters
----------
omega : float or numpy.ndarray
Frequencies (can be Quantity).
ro : float or Quantity, optional
Distance scale for translation into internal units (default from configuration file).
vo : float or Quantity, optional
Velocity scale for translation into internal units (default from configuration file).
Notes
-----
- 2018-04-08 - Written - Bovy (Uoft)
"""
actionAngle.__init__(self, ro=kwargs.get("ro", None), vo=kwargs.get("vo", None))
if not "omega" in kwargs: # pragma: no cover
raise OSError("Must specify omega= for actionAngleHarmonic")
self._omega = conversion.parse_frequency(
kwargs.get("omega"), ro=self._ro, vo=self._vo
)
return None
def _evaluate(self, *args, **kwargs):
"""
Evaluate the action for the harmonic oscillator
Parameters
----------
Either:
a) x,vx:
1) floats: phase-space value for single object (each can be a Quantity)
2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity)
Returns
-------
float or numpy.ndarray
action
Notes
-----
- 2018-04-08 - Written - Bovy (UofT)
"""
if len(args) == 2: # x,vx
x, vx = args
return (vx**2.0 / self._omega + self._omega * x**2.0) / 2.0
else: # pragma: no cover
raise ValueError("actionAngleHarmonic __call__ input not understood")
def _actionsFreqs(self, *args, **kwargs):
"""
Evaluate the action and frequency for the harmonic oscillator
Parameters
----------
Either:
a) x,vx:
1) floats: phase-space value for single object (each can be a Quantity)
2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity)
Returns
-------
tuple
(action,frequency)
Notes
-----
- 2018-04-08 - Written - Bovy (UofT)
"""
if len(args) == 2: # x,vx
x, vx = args
return (
(vx**2.0 / self._omega + self._omega * x**2.0) / 2.0,
self._omega * numpy.ones_like(x),
)
else: # pragma: no cover
raise ValueError("actionAngleHarmonic __call__ input not understood")
def _actionsFreqsAngles(self, *args, **kwargs):
"""
Evaluate the action, frequency, and angle for the harmonic oscillator
Parameters
----------
Either:
a) x,vx:
1) floats: phase-space value for single object (each can be a Quantity)
2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity)
Returns
-------
tuple
(action,frequency,angle)
Notes
-----
- 2018-04-08 - Written - Bovy (UofT)
"""
if len(args) == 2: # x,vx
x, vx = args
return (
(vx**2.0 / self._omega + self._omega * x**2.0) / 2.0,
self._omega * numpy.ones_like(x),
numpy.arctan2(self._omega * x, vx),
)
else: # pragma: no cover
raise ValueError("actionAngleHarmonic __call__ input not understood")