pymsm.py

"""PyMSM - Python implementation for calculating geomagnetic rigidity cutoff."""
import datetime
import os
import sys
from math import sqrt, acos, asin, cos, log10

import numpy as np
from IRBEM import MagFields, Coords


class MapDB:
    """Database manager for pre-calculated rigidity maps.
    
    All instances share the same map dictionary for efficiency.
    """
    maps = {}
    
    def __init__(self, mapdir=None):
        """Initialize MapDB with optional custom map directory.
        
        Args:
            mapdir: Optional custom directory path for map files.
        """
        self.mapdir = mapdir
    
    def getMap(self, year, kp, ut):
        """Return the pre-calculated rigidity map for given year, kp and ut.
        
        Args:
            year: Year as string.
            kp: Kp index as string.
            ut: Universal time as string.
            
        Returns:
            Tuple of (Lm, Rc, Rc*Lm^2) arrays.
        """
        if self.mapdir is None:
            mdir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "MAPS")
        else:
            mdir = self.mapdir
        tag = year + kp + ut
        if tag not in MapDB.maps:
            file = os.path.join(mdir, year, f"AVKP{kp}T{ut}.AVG")
            print(file)
            MapDB.maps[tag] = self.readMap(file)
        return MapDB.maps[tag]
    
    def readMap(self, file):
        """Read in the map file.
        
        Map dimensions: nlat = 37, nlon = 73
        Column 3 contains Lm, column 5 contains Rc.
        
        Args:
            file: Path to the map file.
            
        Returns:
            Tuple of (Lm, Rc, Rc*Lm^2) arrays.
        """
        try:
            lm, rc = np.loadtxt(file, skiprows=0, usecols=(3, 5), unpack=True)
            lm = lm.reshape((37, 73)).transpose()
            rc = rc.reshape((37, 73)).transpose()
            return lm, rc, rc * lm**2
        except Exception as e:
            print(e, file=sys.stderr)
            raise
            
class PyMSM:
    """Main class for calculating geomagnetic rigidity cutoff and transmission functions."""
    
    def __init__(self, times, positions, kps=None, rc=None, mapdir=None):
        """Initialize PyMSM with observation times, positions, and parameters.
        
        Args:
            times: 1D array of date and time in ISO string format.
            positions: 2D array of locations in GDZ coordinates, e.g., [[alt0, lat0, lon0], [alt1, lat1, lon1], ...].
            kps: 1D array of Kp indices corresponding to the times.
            rc: Optional 1D array of rigidity cutoffs in GV for transmission factor calculation.
                Default: [0.1, 0.2, 0.5, 1., 1.5, 2., 2.5, 3., 3.5, 4., 4.5, 5., 5.5,
                          6., 6.5, 7., 7.5, 8., 9., 10., 11., 12., 13, 14., 15., 16., 17.,
                          20., 25., 20., 30., 40., 50., 60.]
            mapdir: Optional path to alternative/user precalculated maps.
        """
        self.cyears = ['1955', '1960', '1965', '1970', '1975', '1980', '1985', '1990',
                       '1995', '2000', '2005', '2010', '2015', '2020', '2025']
        self.cuts = ['00', '03', '06', '09', '12', '15', '18', '21']
        self.ckps = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'X']
        self.kps = kps
        if rc is None:
            # The rigidity values to be used for calculating the transmission function
            self.rc = [0.1, 0.2, 0.5, 1., 1.5, 2., 2.5, 3., 3.5, 4., 4.5, 5., 5.5,
                       6., 6.5, 7., 7.5, 8., 9., 10., 11., 12., 13, 14., 15., 16., 17.,
                       20., 25., 20., 30., 40., 50., 60.]
        else:
            self.rc = rc
        t = []
        for tm in times:
            t.append(datetime.datetime.fromisoformat(tm))
        self.times = t
        self.model = MagFields(options=[0, 30, 0, 0, 0], kext=4, verbose=False)
        # Positions are in GDZ
        self.coords = Coords().coords_transform(times, positions, 'GDZ', 'RLL')
        self.radius = self.coords[:, 0]
        self.coords = Coords().coords_transform(times, self.coords, 'RLL', 'GDZ')
        self.lla = {}
        self.lla['x1'] = self.coords[:, 0]
        self.lla['x2'] = self.coords[:, 1]
        self.lla['x3'] = self.coords[:, 2]
        self.lla['dateTime'] = self.times
        self.maginput = {'Kp': kps * 10.}
        self.model.make_lstar(self.lla, self.maginput)
        # The (B, L) for the input times and positions
        self.lm = np.abs(self.model.make_lstar_output['Lm'])
        self.bm = np.abs(self.model.make_lstar_output['blocal'])
        self.dbMgr = MapDB(mapdir)
    
    def getTransmissionFunctions(self):
        """Return all relevant results for the specified (times, locations) series.
        
        Returns:
            Tuple of (Lm, Bm, Mlat, Rcv, ES, TF) where:
                Lm: McIlwain's L-parameter (np.array).
                Bm: Magnetic field intensity at the location (np.array).
                Mlat: Magnetic latitude (np.array).
                Rcv: Vertical rigidity cutoff (np.array).
                ES: Earth's shadowing factor (np.array).
                TF: Transmission function 2D array [len(times) x len(rc)].
        """
        # First obtain the interpolated vertical cutoffs
        Rcv = self.getRc()
        # Second, get the magnetic latitude
        Mlats = self.getEMLat()
        # Third, calculate transmission factors
        TF = self.getTransfact(Mlats, Rcv)
        # Fourth, get the Earth shadowing factors
        ES = self.facshadow(self.radius)
        return self.lm, self.bm, Mlats, Rcv, ES, TF
        
    def getRc(self):
        """Calculate the vertical cutoff rigidities for the specified series of (times, locations).
        
        Returns:
            Rcv: Array of vertical cutoff rigidity values in units of GV.
        """ 

        t_utc = self.lla['dateTime']
        rclist = []
        local = {}

        for i in range(len(t_utc)):
            year = t_utc[i].year
            if year < 1955:
                year = 1955
            if year > 2025:
                year = 2025
            iy = int((year - 1955) / 5)
            ir = (year - 1955) % 5
            cyear = self.cyears[iy]
            iu = int((t_utc[i].hour + t_utc[i].minute / 60. + 1.5) / 3.)
            # UT=1 corresponds to ut: 1.5 - 4.5 hrs
            if iu > 7:
                iu = 0
            cut = self.cuts[iu]
            ckp = self.ckps[self.kps[i]]
            self.dbMgr.getMap(cyear, ckp, cut)
            mkey = cyear + ckp + cut
            # In first map
            local['x1'] = 450.
            local['x2'] = self.lla['x2'][i]
            local['x3'] = self.lla['x3'][i]
            local['dateTime'] = datetime.datetime(int(cyear), 1, 1, int(cut)).isoformat()
            maginput = {'Kp': self.maginput['Kp'][i]}
            self.model.make_lstar(local, maginput)
            lm = np.abs(self.model.make_lstar_output['Lm'])
            rc = self.getRC450km(mkey, self.lla['x2'][i], self.lla['x3'][i], lm)
            
            w = (ir + (t_utc[i].timetuple().tm_yday + (t_utc[i].hour + t_utc[i].minute / 60.) / 24.) / 365.) / 5.
            # The next map at +5 years if required
            if 1955 < year < 2025 and w < 1.:
                cyear = self.cyears[iy + 1]
                self.dbMgr.getMap(cyear, ckp, cut)
                mkey = cyear + ckp + cut
                # In 2nd map
                local['dateTime'] = datetime.datetime(int(cyear), 1, 1, int(cut)).isoformat()
                self.model.make_lstar(local, maginput)
                lm1 = np.abs(self.model.make_lstar_output['Lm'])
                rc1 = self.getRC450km(mkey, self.lla['x2'][i], self.lla['x3'][i], lm1)
                lm = lm * w + (1. - w) * lm1
                rc = rc * w + (1. - w) * rc1  # Now apply altitude interpolation
            lmr = self.lm[i]  # For real time and altitude
            rcr = rc * (lm / lmr)**2  # Scaled by LM^2
            
            # Further Radial Distance Adjustment according to email from Don on 09/12/2013:
            # "If you examine the cutoff interpolation FORTRAN code in detail,
            # you may notice that there is an adjustment in the 'L' value altitude
            # interpolation process at the end of Subroutine LINT5X5 in a section
            # labeled 'Radial Distance Adjustment'. While the 'L' interpolation
            # equation has the basic form of L**-2, when this exact form is used to
            # extend to geosynchronous altitude, the cutoff values extrapolated from
            # the near Earth low altitudes are too high; approximately 0.3 GV at the
            # magnetic equator. (See figure 7 of Shea & Smart, JGR, 72, 3447, 1967)
            # We incorporated a 'patch' (actually an 'ad-hoc' exponential function
            # that makes adjustments so at 6.6 earth radii the vertical cutoff rigidity
            # for local noon at the magnetic equator under extremely quiet magnetic
            # conditions is about zero (or extremely small).
            # This 'ad-hoc' exponential function is not going to be reliable
            # beyond geosynchronous distances."
            radist = self.radius[i]
            rcorr = log10(radist * radist) / 14.  # Note: Don used 11, but 14 works better
            rcr -= rcorr
            if rcr < 0.:
                rcr = 0.
            try:
                rct = rcr[0]
            except (IndexError, TypeError):
                rct = rcr
            rclist.append(rct)

        return np.array(rclist)          
            
    def getRC450km(self, mkey, lat, lon, lm):
        """Interpolation to obtain Rc for the given location at 450km.
        
        Should be called after the mkey map has been prepared (e.g., after dbMgr.getMap()).
        
        Args:
            mkey: Map key (cyear + ckp + cut).
            lat: Latitude in degrees.
            lon: Longitude in degrees.
            lm: L shell number of the position at 450km altitude.
            
        Returns:
            Vertical rigidity cutoff in GV at 450km altitude.
        """
        # Get the left-top box corner indices
        i, j = self.getGridIdx(lat, lon)

        # Get the Lm and Rc from the maps
        # Left-top corner
        rclm_LT = self.dbMgr.maps[mkey][2][i, j]
        # Right-top corner
        rclm_RT = self.dbMgr.maps[mkey][2][i + 1, j]
        # Left-bottom corner
        rclm_LB = self.dbMgr.maps[mkey][2][i, j + 1]
        # Right-bottom corner
        rclm_RB = self.dbMgr.maps[mkey][2][i + 1, j + 1]

        if any(map(lambda x: x == 99.99, (rclm_LT, rclm_RT, rclm_LB, rclm_RB, lm))):
            lm = rclm_LT = rclm_RT = rclm_LB = rclm_RB = 99.99

        # Get the weights
        wl, wr, wt, wb = self.getWeights(lat, lon)

        rclm_l = wt * rclm_LT + wb * rclm_LB
        rclm_r = wt * rclm_RT + wb * rclm_RB
        rclm = wl * rclm_l + wr * rclm_r

        return rclm / lm**2
     
    def getWeights(self, lat, lon):
        """Calculate interpolation weights for a given lat/lon position.
        
        Args:
            lat: Latitude in degrees.
            lon: Longitude in degrees.
            
        Returns:
            Tuple of (wl, wr, wt, wb) weights.
        """
        # Weights in longitude
        wr = (lon % 5) / 5.  # Right side of the box
        wl = 1.0 - wr  # Left side of the box
        # Weights in latitude
        wt = ((lat + 90) % 5) / 5.  # Top side of the box
        wb = 1. - wt  # Bottom side of the box
        if lat == 90.:
            wt = 1.
            wb = 0.
        return wl, wr, wt, wb
        
    def getGridIdx(self, lat, lon):
        """Get grid indices for the given lat/lon position.
        
        Args:
            lat: Latitude in degrees.
            lon: Longitude in degrees.
            
        Returns:
            Tuple of (ix, iy) grid indices.
        """
        ix = int(lon / 5)
        if ix > 71:
            ix = 0
        iy = int(18 - lat / 5)
        if iy < 0:
            iy = 0
        if iy > 35:
            iy = 35
        return ix, iy 

    def getEMLat(self):
        """Calculate the equivalent magnetic latitude of the given locations.
        
        Gets corrected geomagnetic latitude (GMLATC) at sub-satellite point.
        Then, gets invariant latitude at satellite position:
            INVARIANT LAT = ACOS(1.0/SQRT(L))
        Selects the smaller value for magnetic latitude.
        Always uses absolute value of equivalent magnetic latitude.
        
        Returns:
            Array of equivalent magnetic latitudes in radians.
        """ 
        # Calculate the corrected magnetic latitude
        # GDZ -> MAG (Cartesian coordinates)
        mpos = Coords().coords_transform(self.times, self.coords, 'GDZ', 'MAG')
        
        # mpos are in Cartesian coordinates
        gmlatcr = []
        for mp in mpos:
            r = sqrt(mp[0]**2 + mp[1]**2 + mp[2]**2)
            gmlatcr.append(asin(mp[2] / r))
        emlats = []
        for i in range(len(self.lm)):
            glmdar = 0.0
            if self.lm[i] > 1.:
                glmdar = acos(1.0 / sqrt(self.lm[i]))
            if abs(gmlatcr[i]) < abs(glmdar):
                glmdar = abs(gmlatcr[i])
            emlats.append(glmdar)
        return np.array(emlats)    
    

    def calculateRInv(self, B, L):
        """Calculate the invariant radial distance (R) and magnetic latitude (lambda).
        
        Based on the method from:
        Roberts, C. S. (1964), Coordinates for the study of particles trapped in
        the Earth's magnetic field: A method of converting from B, L to R, l
        coordinates, J. Geophys. Res., 69, 5089-5090.
        
        Note: Not currently used. Need to compare lambda vs emlats.
        
        Args:
            B: Magnetic field values (numpy 1D array).
            L: L-shell values (numpy 1D array).
            
        Returns:
            Tuple of (R, lambda) as numpy 1D arrays.
        """
        a = [1.25992106, -0.19842592, -0.04686632, -0.01314096,
             -0.00308824, 0.00082777, -0.00105877, 0.00183142]
        Md = 31165.3  # nT*Re^3
        if L < 0.:
            return -1., -1.
        p = np.pow(np.pow(L, 3.) * B / Md, -1. / 3.)
        
        s = 0.
        for i in range(8):
            s += a[i] * np.pow(p, i)
        ps = p * s
        for i in np.nonzero(ps > 1.):
            ps[i] = 1.
        
        lamb = np.degrees(np.acos(np.sqrt(ps)))
        R = L * ps
        for i in np.nonzero((p < 0.) | (p > 1.)):
            R[i] = -1.
            lamb[i] = -1.
        return R, lamb
    
    def getTransfact(self, mlat, rcv):
        """Compute the angle-averaged transmission function at the given RCs.
        
        Averages over arrival directions.
        
        Args:
            mlat: Magnetic latitude in radians (numpy 1D array or list).
            rcv: Vertical cut-off (numpy 1D array or list).
        
        Returns:
            Transmission function at the specified rigidities (numpy 2D array).
        """

        N = len(self.rc)
        fac = np.empty(N)
        facs = np.empty(shape=(len(mlat), N))
        rcv = np.array(rcv)
        fac.fill(1.0)
        for i in range(len(rcv)):
            if rcv[i] <= 0.1:
                facs[i] = fac
            else:
                facs[i] = self.calcTF(mlat[i], rcv[i])
        return facs
            
    def calcTF(self, mlat, rcv):
        """Calculate transmission function for a given magnetic latitude and cutoff.
        
        Args:
            mlat: Magnetic latitude in radians.
            rcv: Vertical cutoff rigidity in GV.
            
        Returns:
            Array of transmission factors for each rigidity in self.rc.
        """
        # Table of One-Minus-Cos-Angles-Over-2
        omcao2 = np.array([0., .067, .146, .25, .5, .75, .854, .933, 1.])
        ang = np.array([0.01, .5236, .785, 1.047, 1.571, 2.094, 2.356, 2.618, 3.1416])
        nangle = 9
        
        cosa = np.cos(ang)
        cosl = cos(mlat)
        cut = 4. * rcv / (1.0 + np.sqrt(1.0 + cosa * cosl**3))**2
        
        N = len(self.rc)
        fac = np.empty(N)
        for ir in range(N):
            fac[ir] = 0.
            if self.rc[ir] >= cut[0]:
                fac[ir] = 1.0
                for ia in range(1, nangle):
                    # Find the angular location where the cutoff goes over the rc
                    if self.rc[ir] <= cut[ia]:
                        fac[ir] = omcao2[ia - 1] + (self.rc[ir] - cut[ia - 1]) * \
                                  (omcao2[ia] - omcao2[ia - 1]) / (cut[ia] - cut[ia - 1])
                        break
        return fac
                        
    def facshadow(self, R):
        """Calculate correction factor for Earth's shadow on the spacecraft.
        
        Uses simple geometrical optics.
        
        Args:
            R: Radius in Earth radii (list or 1D numpy array).
            
        Returns:
            Shadow correction factor.
        """
        fac = 1. - .5 * (1. - np.sqrt(R**2 - 1.) / R)
        return fac