pyamps.sh_utils — pyAMPS 1.0.0 documentation

# Source code for pyamps.sh_utils

    """ tools that are useful for spherical harmonic analysis

        SHkeys       -- class to contain n and m - the indices of the spherical harmonic terms
        nterms       -- function which calculates the number of terms in a
                        real expansion of a poloidal (internal + external) and toroidal expansion
        legendre -- calculate associated legendre functions - with option for Schmidt semi-normalization

    MIT License

    Copyright (c) 2017 Karl M. Laundal

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
    """

    from __future__ import absolute_import, division
    import numpy as np
    import apexpy
    from .mlt_utils import mlon_to_mlt
    from builtins import range

    d2r = np.pi/180

    DEFAULT = object()
    refre = 6371.2 # reference radius

    class SHkeys(object):
        """ container for n and m in spherical harmonics

            keys = SHkeys(Nmax, Mmax)

            keys will behave as a tuple of tuples, more or less
            keys['n'] will return a list of the n's
            keys['m'] will return a list of the m's
            keys[3] will return the fourth n,m tuple

            keys is also iterable

        """

        def __init__(self, Nmax, Mmax):
            keys = []
            for n in range(Nmax + 1):
                for m in range(Mmax + 1):
                    keys.append((n, m))

            self.keys = tuple(keys)
            self.make_arrays()

        def __getitem__(self, index):
            if index == 'n':
                return [key[0] for key in self.keys]
            if index == 'm':
                return [key[1] for key in self.keys]

            return self.keys[index]

        def __iter__(self):
            for key in self.keys:
                yield key

        def __len__(self):
            return len(self.keys)

        def __repr__(self):
            return ''.join(['n, m\n'] + [str(key)[1:-1] + '\n' for key in self.keys])[:-1]

        def __str__(self):
            return ''.join(['n, m\n'] + [str(key)[1:-1] + '\n' for key in self.keys])[:-1]

        def setNmin(self, nmin):
            """ set minimum n """
            self.keys = tuple([key for key in self.keys if key[0] >= nmin])
            self.make_arrays()
            return self

        def MleN(self):
            """ set m <= n """
            self.keys = tuple([key for key in self.keys if abs(key[1]) <= key[0]])
            self.make_arrays()
            return self

        def Mge(self, limit):
            """ set m >= limit  """
            self.keys = tuple([key for key in self.keys if abs(key[1]) >= limit])
            self.make_arrays()
            return self

        def NminusModd(self):
            """ remove keys if n - m is even """
            self.keys = tuple([key for key in self.keys if (key[0] - abs(key[1])) % 2 == 1])
            self.make_arrays()
            return self

        def NminusMeven(self):
            """ remove keys if n - m is odd """
            self.keys = tuple([key for key in self.keys if (key[0] - abs(key[1])) % 2 == 0])
            self.make_arrays()
            return self

        def negative_m(self):
            """ add negative m to the keys """
            keys = []
            for key in self.keys:
                keys.append(key)
                if key[1] != 0:
                    keys.append((key[0], -key[1]))

            self.keys = tuple(keys)
            self.make_arrays()

            return self

        def make_arrays(self):
            """ prepare arrays with shape ( 1, len(keys) )
                these are used when making G matrices
            """

            if len(self) > 0:
                self.m = np.array(self)[:, 1][np.newaxis, :]
                self.n = np.array(self)[:, 0][np.newaxis, :]
            else:
                self.m = np.array([])[np.newaxis, :]
                self.n = np.array([])[np.newaxis, :]

    def nterms(NT = 0, MT = 0, NVi = 0, MVi = 0, NVe = 0, MVe = 0):
        """ return number of coefficients in an expansion in real spherical harmonics of
            toroidal magnetic potential truncated at NT, MT
            poloidal magnetic potential truncated at NVi, MVi for internal sources
            poloidal magnetic potential truncated at NVe, MVe for external sources
        """

        return len(SHkeys(NT , MT ).setNmin(1).MleN().Mge(0)) + \
               len(SHkeys(NT , MT ).setNmin(1).MleN().Mge(1)) + \
               len(SHkeys(NVe, MVe).setNmin(1).MleN().Mge(0)) + \
               len(SHkeys(NVe, MVe).setNmin(1).MleN().Mge(1)) + \
               len(SHkeys(NVi, MVi).setNmin(1).MleN().Mge(0)) + \
               len(SHkeys(NVi, MVi).setNmin(1).MleN().Mge(1))

    def legendre(nmax, mmax, theta, schmidtnormalize = True, keys = None):
        """ Calculate associated Legendre function P and its derivative

            Algorithm from "Spacecraft Attitude Determination and Control" by James Richard Wertz

            Parameters
            ----------
            nmax : int
                highest spherical harmonic degree
            mmax : int
                hightest spherical harmonic order
            theta : array, float
                colatitude in degrees (shape is not preserved)
            schmidtnormalize : bool, optional
                True if Schmidth seminormalization is wanted, False otherwise. Default True
            keys : SHkeys, optional
                If this parameter is set, an array will be returned instead of a dict.
                The array will be (N, 2M), where N is the number of elements in `theta`, and
                M is the number of keys. The first M columns represents a matrix of P values,
                and the last M columns represent values of dP/dtheta

            Returns
            -------
            P : dict
                dictionary of Legendre function evalulated at theta. Dictionary keys are spherical harmonic
                wave number tuples (n, m), and values will have shape (N, 1), where N is number of
                elements in `theta`.
            dP : dict
                dictionary of Legendre function derivatives evaluated at theta. Dictionary keys are spherical
                harmonic wave number tuples (n, m), and values will have shape (N, 1), where N is number of
                elements in theta.
            PdP : array (only if keys != None)
                if keys != None, PdP is returned instaed of P and dP. PdP is an (N, 2M) array, where N is
                the number of elements in `theta`, and M is the number of keys. The first M columns represents
                a matrix of P values, and the last M columns represent values of dP/dtheta

        """

        theta = theta.flatten()[:, np.newaxis]

        P = {}
        dP = {}
        sinth = np.sin(d2r*theta)
        costh = np.cos(d2r*theta)

        if schmidtnormalize:
            S = {}
            S[0, 0] = 1.

        # initialize the functions:
        for n in range(nmax +1):
            for m in range(nmax + 1):
                P[n, m] = np.zeros_like(theta, dtype = np.float64)
                dP[n, m] = np.zeros_like(theta, dtype = np.float64)

        P[0, 0] = np.ones_like(theta, dtype = np.float64)
        for n in range(1, nmax +1):
            for m in range(0, min([n + 1, mmax + 1])):
                # do the legendre functions and derivatives
                if n == m:
                    P[n, n]  = sinth * P[n - 1, m - 1]
                    dP[n, n] = sinth * dP[n - 1, m - 1] + costh * P[n - 1, n - 1]
                else:

                    if n == 1:
                        Knm = 0.
                        P[n, m]  = costh * P[n -1, m]
                        dP[n, m] = costh * dP[n - 1, m] - sinth * P[n - 1, m]

                    elif n > 1:
                        Knm = ((n - 1)**2 - m**2) / ((2*n - 1)*(2*n - 3))
                        P[n, m]  = costh * P[n -1, m] - Knm*P[n - 2, m]
                        dP[n, m] = costh * dP[n - 1, m] - sinth * P[n - 1, m] - Knm * dP[n - 2, m]

                if schmidtnormalize:
                    # compute Schmidt normalization
                    if m == 0:
                        S[n, 0] = S[n - 1, 0] * (2.*n - 1)/n
                    else:
                        S[n, m] = S[n, m - 1] * np.sqrt((n - m + 1)*(int(m == 1) + 1.)/(n + m))

        if schmidtnormalize:
            # now apply Schmidt normalization
            for n in range(1, nmax + 1):
                for m in range(0, min([n + 1, mmax + 1])):
                    P[n, m]  *= S[n, m]
                    dP[n, m] *= S[n, m]

        if keys is None:
            return P, dP
        else:
            Pmat  = np.hstack(tuple(P[key] for key in keys))
            dPmat = np.hstack(tuple(dP[key] for key in keys))

            return np.hstack((Pmat, dPmat))

    def getG0(glat, glon, height, time, epoch = 2015., h_R = 110.):
        """ calculate the G matrix for the constant term in the AMPS model. The constant term is the
            term that depends only on the spherical harmonic coefficients that are not scaled by
            external parameters. This G matrix can be used to produce the full matrix.

            The structure of the matrix is such that G0.dot(m), the product of the matrix with the
            first 1/19 of the model vector, will be model values of the eastward, northward, and
            upward components of the magnetic field, stacked in a column vector.

            glat, glon, time, and height must all have the same number of elements (let's call this N)

            Parameters
            ----------
            glat : array
                Geodetic latitude (degrees)
            glon : array
                Geographic/geodetic longitude (degrees)
            height : array
                Geodetic heights, in km
            time : array
                Array of datetimes correspondign to each point. This is needed to calculate
                magnetic local time.
            epoch : float, optional
                The epoch used for conversion to apex coordinates. Default 2015.
            h_R : float, optional
                Reference height used in conversion to modified apex coordinates. Default 110 km.

            Returns
            -------
            G0 : array
                an 3N by M matrix, where N is the number of elements in the input coordinates. There will be
                3 times as many rows G0, since there are 3 components. Partionining G0 in thirds, from top to
                bottom, gives the parts that correspond to east, north, and up, respectively. M
                is the number of terms in the spherical harmonic expansion of B

        """
        glat   = np.asarray(glat).flatten()
        glon   = np.asarray(glon).flatten()
        height = np.asarray(height).flatten()

        # convert to magnetic coords and get base vectors
        a = apexpy.Apex(epoch, refh = h_R)
        qlat, qlon = a.geo2qd(  glat.flatten(), glon.flatten(), height.flatten())
        alat, alon = a.geo2apex(glat.flatten(), glon.flatten(), height.flatten())
        f1, f2, f3, g1, g2, g3, d1, d2, d3, e1, e2, e3 = a.basevectors_apex(qlat, qlon, height, coords  = 'qd')
        f1e = f1[0].reshape(-1, 1) # base vector components as column vectors
        f1n = f1[1].reshape(-1, 1)
        f2e = f2[0].reshape(-1, 1)
        f2n = f2[1].reshape(-1, 1)
        d1e = d1[0].reshape(-1, 1)
        d1n = d1[1].reshape(-1, 1)
        d2e = d2[0].reshape(-1, 1)
        d2n = d2[1].reshape(-1, 1)

        # calculate magnetic local time
        phi = mlon_to_mlt(qlon, time, a.year)[:, np.newaxis]*15 # multiply by 15 to get degrees

        # turn the coordinate arrays into column vectors:
        alat, qlat, h = map(lambda x: x.flatten()[:, np.newaxis], [alat, qlat, height])

        # truncation levels:
        NT, MT, NV, MV = 65, 3, 45, 3

        # generate spherical harmonic keys
        keys = {} # dictionary of spherical harmonic keys
        keys['cos_T'] = SHkeys(NT, MT).setNmin(1).MleN().Mge(0)
        keys['sin_T'] = SHkeys(NT, MT).setNmin(1).MleN().Mge(1)
        keys['cos_V'] = SHkeys(NV, MV).setNmin(1).MleN().Mge(0)
        keys['sin_V'] = SHkeys(NV, MV).setNmin(1).MleN().Mge(1)
        m_cos_V = keys['cos_V'].m
        m_sin_V = keys['sin_V'].m
        m_cos_T = keys['cos_T'].m
        m_sin_T = keys['sin_T'].m

        nV = np.hstack((keys['cos_V'].n, keys['sin_V'].n))

        # generate Legendre matrices - first get dicts of arrays, and then stack them in the appropriate fashion
        legendre_T = legendre(NT, MT, 90 - alat, keys = keys['cos_T'])
        legendre_V = legendre(NV, MV, 90 - qlat, keys = keys['cos_V'])
        P_cos_T  =  legendre_T[:, :len(keys['cos_T']) ] # split
        dP_cos_T = -legendre_T[:,  len(keys['cos_T']):]
        P_cos_V  =  legendre_V[:, :len(keys['cos_V']) ] # split
        dP_cos_V = -legendre_V[:,  len(keys['cos_V']):]
        P_sin_T  =  P_cos_T[ :, keys['cos_T'].m.flatten() != 0]
        dP_sin_T =  dP_cos_T[:, keys['cos_T'].m.flatten() != 0]
        P_sin_V  =  P_cos_V[ :, keys['cos_V'].m.flatten() != 0]
        dP_sin_V =  dP_cos_V[:, keys['cos_V'].m.flatten() != 0]

        # trig matrices:
        cos_T  =  np.cos(phi * d2r * m_cos_T)
        sin_T  =  np.sin(phi * d2r * m_sin_T)
        cos_V  =  np.cos(phi * d2r * m_cos_V)
        sin_V  =  np.sin(phi * d2r * m_sin_V)
        dcos_T = -np.sin(phi * d2r * m_cos_T)
        dsin_T =  np.cos(phi * d2r * m_sin_T)
        dcos_V = -np.sin(phi * d2r * m_cos_V)
        dsin_V =  np.cos(phi * d2r * m_sin_V)

        cos_qlat   = np.cos(qlat * d2r)
        cos_alat   = np.cos(alat * d2r)

        sinI  = 2 * np.sin( alat * d2r )/np.sqrt(4 - 3*cos_alat**2)

        r  = refre + h
        Rtor  = refre/r

        F = f1e*f2n - f1n*f2e

        # matrix with horizontal spherical harmonic functions in QD coordinates
        V        = np.hstack((P_cos_V * cos_V, P_sin_V * sin_V ))

        # matrices with partial derivatives in QD coordinates:
        dV_dqlon  = np.hstack(( P_cos_V * dcos_V * m_cos_V,  P_sin_V * dsin_V * m_sin_V ))
        dV_dqlat  = np.hstack((dP_cos_V *  cos_V          , dP_sin_V *  sin_V           ))

        # matrices with partial derivatives in MA coordinates:
        dT_dalon  = np.hstack(( P_cos_T * dcos_T * m_cos_T,  P_sin_T * dsin_T * m_sin_T))
        dT_dalat  = np.hstack((dP_cos_T *  cos_T          , dP_sin_T *  sin_T * m_sin_T))

        # Toroidal field components
        B_T_e  =   -d1n * dT_dalon / cos_alat + d2n * dT_dalat / sinI
        B_T_n  =    d1e * dT_dalon / cos_alat - d2e * dT_dalat / sinI
        B_T_u  =    np.zeros(B_T_n.shape)

        # Poloidal field components:
        B_V_e = (-f2n / (cos_qlat * r) * dV_dqlon + f1n * dV_dqlat / r) * refre * Rtor ** (nV + 1)
        B_V_n = ( f2e / (cos_qlat * r) * dV_dqlon - f1e * dV_dqlat / r) * refre * Rtor ** (nV + 1)
        B_V_u = np.sqrt(F) * V  * (nV + 1) * Rtor ** (nV + 2)

        # combine:
        G     = np.hstack((np.vstack((B_T_e  ,
                                      B_T_n  ,
                                      B_T_u  )),   np.vstack((B_V_e,
                                                              B_V_n,
                                                              B_V_u))
                          ))

        return G

    def get_ground_field_G0(qdlat, mlt, height, current_height):
        """ calculate the G matrix for the constant term in the AMPS model needed to calculate
            corresponding ground magnetic field perturbations. The constant term is the
            term that depends only on the spherical harmonic coefficients that are not scaled by
            external parameters. This G matrix can be used to produce the full matrix.

            Parameters
            ----------
            qdlat : array
                Quasi-dipole latitude (degrees)
            MLT : array
                Magnetic local time (MLT)
            height : float
                Geodetic height, in km (0 <= height <= current_height)
            current_height : float
                height of the current

            Returns
            -------
            G0 : array
                an 3N by M matrix, where N is the number of elements in the input coordinates. There will be
                3 times as many rows G0, since there are 3 components. Partionining G0 in thirds, from top to
                bottom, gives the parts that correspond to east, north, and up, respectively. M
                is the number of terms in the spherical harmonic expansion of B

        """
        qdlat  = qdlat.flatten() [:, np.newaxis]
        mlt    = mlt.flatten()   [:, np.newaxis]

        # convert mlt to degreees
        phi = mlt*15 # multiply by 15 to get degrees

        # truncation levels:
        N, M = 45, 3

        # generate spherical harmonic keys
        keys = {} # dictionary of spherical harmonic keys
        keys['cos'] = SHkeys(N, M).setNmin(1).MleN().Mge(0)
        keys['sin'] = SHkeys(N, M).setNmin(1).MleN().Mge(1)
        m_cos = keys['cos'].m
        m_sin = keys['sin'].m

        n = np.hstack((keys['cos'].n, keys['sin'].n))

        # generate Legendre matrices - first get dicts of arrays, and then stack them in the appropriate fashion
        legendre_arr = legendre(N, M, 90 - qdlat, keys = keys['cos'])
        P_cos  =  legendre_arr[: , :len(keys['cos']) ] # split
        dP_cos = -legendre_arr[: ,  len(keys['cos']):]
        P_sin  =  P_cos       [: , keys['cos'].m.flatten() != 0]
        dP_sin =  dP_cos      [: , keys['cos'].m.flatten() != 0]

        # trig matrices:
        cos  =  np.cos(phi * d2r * m_cos)
        sin  =  np.sin(phi * d2r * m_sin)
        dcos = -np.sin(phi * d2r * m_cos)
        dsin =  np.cos(phi * d2r * m_sin)

        cos_qlat   = np.cos(qdlat * d2r)

        r = refre + height
        r_ratio_horizontal = (r / (refre + current_height)) ** (n    ) * (refre / (refre + current_height)) ** (n + 1) * (n + 1) / n
        r_ratio_vertical   = (r / (refre + current_height)) ** (n - 1) * (refre / (refre + current_height)) ** (n + 2) * (n + 1)

        # matrix with trig functions
        trigs   = np.hstack((cos, sin ))

        # matrix with derivative of trig functions
        trigs_d = np.hstack(( dcos * m_cos, dsin * m_sin ))

        # matrices with P and dP stacked
        P  = np.hstack((P_cos , P_sin ))
        dP = np.hstack((dP_cos, dP_sin))

        G0east  = P  * trigs_d * r_ratio_horizontal / cos_qlat
        G0north = dP * trigs   * r_ratio_horizontal
        G0up    = P  * trigs   * r_ratio_vertical

        # stack vertically and return
        return np.vstack((G0east, G0north, G0up))