1import numpy as np
2from numba import jit
3
4
5#@jit
6def gravity_spherical_harmonic(x, r_planet, mu, c, s, n_max, m_max):
7    """
8    Calculate the gravitational acceleration due to the spherical harmonics gravity model supplied.
9
10    Args:
11        - x (``array-like``): (N x 3) array of Cartesian coordinates of satellite in frame defined by gravity model.
12        - r_planet (``float``): Equatorial radius of central body.
13        - mu (``float``): Gravitational parameter of central body.
14        - c (``array-like``): Two-dimensional normalised C coefficient array: C[degree, order] = C_(n, m).
15        - s (``array-like``): Two-dimensional normalised S coefficient array: S[degree, order] = S_(n, m).
16        - n_max (``int``): Degree up to which to calculate the gravitational acceleration.
17        - m_max (``int``): Order up to which to calculate the gravitational acceleration. Cannot be higher than n_max.
18
19    Returns:
20        - acc (``array-like``): (N x 3) array of the gravitational acceleration.
21
22    Example::
23
24        r, mu, c, s, n, m = pykep.util.load_gravity_model('gravity_models/Moon/glgm3150.txt')
25        x = numpy.array([[459.5, 795.8773461, 1591.754692], [-459.5, -795.8773461, -1591.754692]])
26
27        acc_zonal = pykep.util.gravity_spherical_harmonic(x, r, mu, c, s, 20, 0)
28        acc_sqr = pykep.util.gravity_spherical_harmonic(x, r, mu, c, s, 20, 20)
29        acc_high = pykep.util.gravity_spherical_harmonic(x, r, mu, c, s, 900, 900)
30
31    .. note::
32
33        This model was taken from a report by NASA:
34        https://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/20160011252.pdf
35        This is the normalised gottlieb algorithm, as coded in MATLAB in the report and transferred to Python.
36    """
37    if not (len(x[0]) == 3):
38        raise ValueError(f"Position must be an (N x 3) array. Shape of position is ({len(x)} x {len(x[0])}).")
39
40    r = np.min(np.linalg.norm(x, axis=1))
41    if r < r_planet:
42        raise ValueError(f"Radial position is less than defined radius of central body ({r} < {r_planet}).")
43
44    if m_max > n_max:
45        raise ValueError(f"Order of model is larger than degree ({m_max} > {n_max}).")
46
47    acc = _gottlieb(x, r_planet, mu, c, s, n_max, m_max)
48
49    return acc
50
51
52def _gottlieb(x, r_planet, mu, c, s, n_max, m_max):
53    norm1, norm2, norm11, normn10, norm1m, norm2m, normn1 = _calculate_normalisation_parameters(n_max)
54
55    n_coor = x.shape[0]
56
57    r = np.sqrt(x[:, 0]**2 + x[:, 1]**2 + x[:, 2]**2)
58
59    r_inverted = 1 / r
60
61    x_r = x[:, 0] * r_inverted
62    y_r = x[:, 1] * r_inverted
63    z_r = x[:, 2] * r_inverted
64
65    rp_r = r_planet * r_inverted
66    rp_rn = np.copy(rp_r)
67
68    mu_r2 = mu * r_inverted * r_inverted
69
70    shape = (max(2, n_max + 1), max(2, n_max + 2), n_coor)
71    p = np.zeros(shape)
72
73    p[0, 0] = 1
74    p[1, 0, :] = np.sqrt(3) * z_r
75    p[1, 1] = np.sqrt(3)
76
77    # sectorial ALFs
78    for n in range(2, n_max + 1):
79        p[n, n] = norm11[n] * p[n - 1, n - 1] * (2 * n - 1)
80
81    shape = (max(3, n_max + 1), n_coor)
82    ctil = np.ones(shape)
83    stil = np.zeros(shape)
84
85    ctil[1] = x_r
86    stil[1] = y_r
87
88    sumh = np.zeros(n_coor)
89    sumgm = np.ones(n_coor)
90    sumj = np.zeros(n_coor)
91    sumk = np.zeros(n_coor)
92
93    for n in range(2, n_max + 1):
94        rp_rn *= rp_r
95
96        n2m1 = 2 * n - 1
97        nm1 = n - 1
98        nm2 = n - 2
99        np1 = n + 1
100
101        # tesseral ALFs
102        p[n, nm1] = normn1[n, nm1] * z_r * p[n, n]
103
104        # zonal ALFs
105        p[n, 0] = (n2m1 * z_r * norm1[n] * p[nm1, 0] - nm1 * norm2[n] * p[nm2, 0]) / n
106        p[n, 1] = (n2m1 * z_r * norm1m[n, 1] * p[nm1, 1] - n * norm2m[n, 1] * p[nm2, 1]) / nm1
107
108        sumhn = normn10[n] * p[n, 1] * c[n, 0]
109        sumgmn = p[n, 0] * c[n, 0] * np1
110
111        if m_max > 0:
112            for m in range(2, n - 1):
113                p[n, m] = (n2m1 * z_r * norm1m[n, m] * p[nm1, m] -
114                           (nm1 + m) * norm2m[n, m] * p[nm2, m]) / (n - m)
115
116            sumjn = np.zeros(n_coor)
117            sumkn = np.zeros(n_coor)
118
119            ctil[n] = ctil[1] * ctil[nm1] - stil[1] * stil[nm1]
120            stil[n] = stil[1] * ctil[nm1] + ctil[1] * stil[nm1]
121
122            for m in range(1, min(n, m_max) + 1):
123                mm1 = m - 1
124                mp1 = m + 1
125                mxpnm = m * p[n, m]
126
127                bnmtil = c[n, m] * ctil[m] + s[n, m] * stil[m]
128
129                sumhn += normn1[n, m] * p[n, mp1] * bnmtil
130                sumgmn += (n + m + 1) * p[n, m] * bnmtil
131
132                bnmtm1 = c[n, m] * ctil[mm1] + s[n, m] * stil[mm1]
133                anmtm1 = c[n, m] * stil[mm1] - s[n, m] * ctil[mm1]
134
135                sumjn += mxpnm * bnmtm1
136                sumkn -= mxpnm * anmtm1
137
138            sumj += rp_rn * sumjn
139            sumk += rp_rn * sumkn
140
141        sumh += rp_rn * sumhn
142        sumgm += rp_rn * sumgmn
143
144    _lambda = sumgm + z_r * sumh
145
146    acc = np.zeros(x.shape)
147
148    acc[:, 0] = - mu_r2 * (_lambda * x_r - sumj)
149    acc[:, 1] = - mu_r2 * (_lambda * y_r - sumk)
150    acc[:, 2] = - mu_r2 * (_lambda * z_r - sumh)
151
152    return acc
153
154
155@jit
156def _calculate_normalisation_parameters(n_max):
157    """
158    Calculate the normalisation parameters for the normalised Gottlieb algorithm. The mapping from the parameters defined in the report
159    (https://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/20160011252.pdf) to the names used in the code looks like
160
161        - lambda_n-1(n)       = norm1
162        - lambda_n-2(n)       = norm2
163        - lambda_n-1_n-1(n)   = norm11
164        - lambda_n-1_m(n, m)  = norm1m
165        - lambda_n-2_m(n, m)  = norm2m
166        - lambda_n_m+1(n, m)  = normn1
167        - lamdba_n_m+1(n, 0)  = normn10
168
169    Args:
170        - n_max (``int``): Maximum degree.
171
172    Returns:
173        Arrays of normalisation parameters.
174    """
175    m_max = n_max
176    n_max += 1
177
178    norm1 = np.zeros(n_max)
179    norm2 = np.zeros(n_max)
180    norm11 = np.zeros(n_max)
181    normn10 = np.zeros(n_max)
182
183    norm1m = np.zeros((n_max, m_max))
184    norm2m = np.zeros((n_max, m_max))
185    normn1 = np.zeros((n_max, m_max + 1))
186
187    for n in range(2, n_max):
188        norm1[n] = np.sqrt((2 * n + 1) / (2 * n - 1))
189        norm2[n] = np.sqrt((2 * n + 1) / (2 * n - 3))
190        norm11[n] = np.sqrt((2 * n + 1) / (2 * n)) / (2 * n - 1)
191        normn10[n] = np.sqrt((n + 1) * n / 2)
192
193        for m in range(1, n):
194            norm1m[n, m] = np.sqrt((n - m) * (2 * n + 1) / ((n + m) * (2 * n - 1)))
195            norm2m[n, m] = np.sqrt((n - m) * (n - m - 1) * (2 * n + 1) / ((n + m) * (n + m - 1) * (2 * n - 3)))
196            normn1[n, m] = np.sqrt((n + m + 1) * (n - m))
197
198    return norm1, norm2, norm11, normn10, norm1m, norm2m, normn1
199