-
Notifications
You must be signed in to change notification settings - Fork 95
/
NumericalPotentialDerivativesMixin.py
170 lines (154 loc) · 6.75 KB
/
NumericalPotentialDerivativesMixin.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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
###############################################################################
# NumericalPotentialDerivativesMixin: helper class to add numerical derivs
###############################################################################
class NumericalPotentialDerivativesMixin:
"""Mixin to add numerical derivatives to a Potential class, use as, e.g.,
.. highlight:: python
.. code-block:: python
class PotWithNumericalDerivs(Potential,NumericalPotentialDerivativesMixin):
def __init__(self,*args,**kwargs):
NumericalPotentialDerivativesMixin.__init__(self,kwargs) # *not* **kwargs!
# Remainder of initialization
...
def _evaluate(self,R,z,phi=0.,t=0.):
# Evaluate the potential
# All forces and second derivatives then computed by NumericalPotentialDerivativesMixin
to add numerical derivatives to a new potential class ``PotWithNumericalDerivs`` that only implements the potential itself, but not the forces. The class may implement any of the forces or second derivatives, all non-implemented forces/second-derivatives will be computed numerically by adding this Mixin
The step used to compute the first (force) and second derivatives can be controlled at object instantiation by the keyword arguments ``dR``, ``dz``, ``dphi`` (for the forces; 1e-8 default) and ``dR2``, ``dz2``, and ``dphi2`` (for the second derivaives; 1e-4 default)
"""
def __init__(self, kwargs): # no **kwargs to get a reference, not a copy!
# For first derivatives
self._dR = kwargs.pop("dR", 1e-8)
self._dphi = kwargs.pop("dphi", 1e-8)
self._dz = kwargs.pop("dz", 1e-8)
# For second derivatives
self._dR2 = kwargs.pop("dR2", 1e-4)
self._dphi2 = kwargs.pop("dphi2", 1e-4)
self._dz2 = kwargs.pop("dz2", 1e-4)
def _Rforce(self, R, z, phi=0.0, t=0.0):
# Do forward difference because R cannot be negative
RplusdR = R + self._dR
Rplus2dR = R + 2.0 * self._dR
dR = (Rplus2dR - R) / 2.0
return (
1.5 * self._evaluate(R, z, phi=phi, t=t)
- 2.0 * self._evaluate(RplusdR, z, phi=phi, t=t)
+ 0.5 * self._evaluate(Rplus2dR, z, phi=phi, t=t)
) / dR
def _zforce(self, R, z, phi=0.0, t=0.0):
# Central difference to get derivative at z=0 right
zplusdz = z + self._dz
zminusdz = z - self._dz
dz = zplusdz - zminusdz
return (
self._evaluate(R, zminusdz, phi=phi, t=t)
- self._evaluate(R, zplusdz, phi=phi, t=t)
) / dz
def _phitorque(self, R, z, phi=0.0, t=0.0):
if not self.isNonAxi:
return 0.0
# Central difference
phiplusdphi = phi + self._dphi
phiminusdphi = phi - self._dphi
dphi = phiplusdphi - phiminusdphi
return (
self._evaluate(R, z, phi=phiminusdphi, t=t)
- self._evaluate(R, z, phi=phiplusdphi, t=t)
) / dphi
def _R2deriv(self, R, z, phi=0.0, t=0.0):
# Do forward difference because R cannot be negative
RplusdR = R + self._dR2
Rplus2dR = R + 2.0 * self._dR2
Rplus3dR = R + 3.0 * self._dR2
dR = (Rplus3dR - R) / 3.0
return (
2.0 * self._evaluate(R, z, phi=phi, t=t)
- 5.0 * self._evaluate(RplusdR, z, phi=phi, t=t)
+ 4.0 * self._evaluate(Rplus2dR, z, phi=phi, t=t)
- 1.0 * self._evaluate(Rplus3dR, z, phi=phi, t=t)
) / dR**2.0
def _z2deriv(self, R, z, phi=0.0, t=0.0):
# Central derivative
zplusdz = z + self._dz2
zminusdz = z - self._dz2
dz = (zplusdz - zminusdz) / 2.0
return (
self._evaluate(R, zplusdz, phi=phi, t=t)
+ self._evaluate(R, zminusdz, phi=phi, t=t)
- 2.0 * self._evaluate(R, z, phi=phi, t=t)
) / dz**2.0
def _phi2deriv(self, R, z, phi=0.0, t=0.0):
if not self.isNonAxi:
return 0.0
# Central derivative
phiplusdphi = phi + self._dphi2
phiminusdphi = phi - self._dphi2
dphi = (phiplusdphi - phiminusdphi) / 2.0
return (
self._evaluate(R, z, phi=phiplusdphi, t=t)
+ self._evaluate(R, z, phi=phiminusdphi, t=t)
- 2.0 * self._evaluate(R, z, phi=phi, t=t)
) / dphi**2.0
def _Rzderiv(self, R, z, phi=0.0, t=0.0):
# Do forward difference in R because R cannot be negative
RplusdR = R + self._dR2
Rplus2dR = R + 2.0 * self._dR2
dR = (Rplus2dR - R) / 2.0
zplusdz = z + self._dz2
zminusdz = z - self._dz2
dz = zplusdz - zminusdz
return (
(
-1.5 * self._evaluate(R, zplusdz, phi=phi, t=t)
+ 2.0 * self._evaluate(RplusdR, zplusdz, phi=phi, t=t)
- 0.5 * self._evaluate(Rplus2dR, zplusdz, phi=phi, t=t)
+ 1.5 * self._evaluate(R, zminusdz, phi=phi, t=t)
- 2.0 * self._evaluate(RplusdR, zminusdz, phi=phi, t=t)
+ 0.5 * self._evaluate(Rplus2dR, zminusdz, phi=phi, t=t)
)
/ dR
/ dz
)
def _Rphideriv(self, R, z, phi=0.0, t=0.0):
if not self.isNonAxi:
return 0.0
# Do forward difference in R because R cannot be negative
RplusdR = R + self._dR2
Rplus2dR = R + 2.0 * self._dR2
dR = (Rplus2dR - R) / 2.0
phiplusdphi = phi + self._dphi2
phiminusdphi = phi - self._dphi2
dphi = phiplusdphi - phiminusdphi
return (
(
-1.5 * self._evaluate(R, z, phi=phiplusdphi, t=t)
+ 2.0 * self._evaluate(RplusdR, z, phi=phiplusdphi, t=t)
- 0.5 * self._evaluate(Rplus2dR, z, phi=phiplusdphi, t=t)
+ 1.5 * self._evaluate(R, z, phi=phiminusdphi, t=t)
- 2.0 * self._evaluate(RplusdR, z, phi=phiminusdphi, t=t)
+ 0.5 * self._evaluate(Rplus2dR, z, phi=phiminusdphi, t=t)
)
/ dR
/ dphi
)
def _phizderiv(self, R, z, phi=0.0, t=0.0):
if not self.isNonAxi:
return 0.0
# Central derivative
phiplusdphi = phi + self._dphi2
phiminusdphi = phi - self._dphi2
dphi = (phiplusdphi - phiminusdphi) / 2.0
zplusdz = z + self._dz2
zminusdz = z - self._dz2
dz = zplusdz - zminusdz
return (
(
self._evaluate(R, zplusdz, phi=phiplusdphi, t=t)
- self._evaluate(R, zplusdz, phi=phiminusdphi, t=t)
- self._evaluate(R, zminusdz, phi=phiplusdphi, t=t)
+ self._evaluate(R, zminusdz, phi=phiminusdphi, t=t)
)
/ dz
/ dphi
/ 2.0
)