new

mosgim/data/loader.py

@@ -14,52 +14,75 @@
 from mosgim.geo.geo import sub_ionospheric
 
 
-class Loader():
+class Loader:
+    """Base class for loading TEC data from different sources."""
 
-    def __init__(self):
-        self.FIELDS = ['datetime', 'el', 'ipp_lat', 'ipp_lon', 'tec']
-        self.DTYPE = (object, float, float, float, float)
-        self.not_found_sites = []
+    def __init__(self) -> None:
+        self.fields = ['datetime', 'el', 'ipp_lat', 'ipp_lon', 'tec']
+        self.dtype = (object, float, float, float, float)
+        self.not_found_sites: list[str] = []
 
 class LoaderTxt(Loader):
+    """Loader for text-based TEC data files."""
 
-    def __init__(self, root_dir):
+    def __init__(self, root_dir: Path) -> None:
+        """
+        Initialize text loader.
+
+        Args:
+            root_dir: Root directory containing TEC data files
+        """
         super().__init__()
-        self.dformat = "%Y-%m-%dT%H:%M:%S"
+        self.date_format = "%Y-%m-%dT%H:%M:%S"
         self.root_dir = root_dir
 
-    def get_files(self, rootdir):
+    def get_files(self, root_dir: Path) -> dict[str, list[Path]]:
         """
-        Root directroy must contain folders with site name 
-        Inside subfolders are *.dat files for every satellite
+        Get all data files from root directory.
+
+        Root directory must contain folders with site name.
+        Inside subfolders are *.dat files for every satellite.
+
+        Args:
+            root_dir: Directory to search for files
+
+        Returns:
+            Dictionary mapping site names to lists of data file paths
+
+        Raises:
+            ValueError: If site name in filename doesn't match directory
         """
         result = defaultdict(list)
-        for subdir, _, files in os.walk(rootdir):
+        for subdir, _, files in os.walk(root_dir):
             for filename in files:
                 filepath = Path(subdir) / filename
                 if str(filepath).endswith(".dat"):
                     site = filename[:4]
                     if site != subdir[-4:]:
-                        raise ValueError(f'{site} in {subdir}. wrong site name')
+                        raise ValueError(f'Site {site} in {subdir} has wrong site name')
                     result[site].append(filepath)
                 else:
-                    warn(f'{filepath} in {subdir} is not data file')
+                    warn(f'{filepath} in {subdir} is not a data file')
         for site in result:
             result[site].sort()
         return result
 
-    def load_data(self, filepath):
-        convert = lambda x: datetime.strptime(x.decode("utf-8"), self.dformat)
+    def load_data(self, filepath: Path) -> tuple[np.ndarray, Path]:
+        """
+        Load data from a single file.
+
+        Args:
+            filepath: Path to data file
+
+        Returns:
+            Tuple of (data array, filepath)
+        """
+        convert = lambda x: datetime.strptime(x.decode("utf-8"), self.date_format)
         data = np.genfromtxt(filepath, 
-                             comments='#', 
-                             names=self.FIELDS, 
-                             dtype=self.DTYPE,
-                             converters={"datetime": convert},  
-                             #unpack=True
-                             )
-
-        #tt = sec_of_day(data['datetime'])
-        #data = append_fields(data, 'sec_of_day', tt, np.float)
+                            comments='#', 
+                            names=self.fields, 
+                            dtype=self.dtype,
+                            converters={"datetime": convert})
         return data, filepath
 
     def __load_data_pool(self, filepath):
@@ -142,7 +165,7 @@ def generate_data(self, sites=[]):
             for sat in hdf_file[site]:
                 sat_data = hdf_file[site][sat]
                 arr = np.empty((len(sat_data['tec']),), 
-                               list(zip(self.FIELDS,self.DTYPE)))
+                               list(zip(self.fields,self.dtype)))
                 el = sat_data['elevation'][:]
                 az = sat_data['azimuth'][:]
                 ts = sat_data['timestamp'][:]

mosgim/data/tec_prepare.py

@@ -99,42 +99,70 @@ def process_data(data_generator):
     return all_data
 
 
-def get_continuos_intervals(data, maxgap=30, maxjump=1):
-    return getContInt(data['sec_of_day'][:], 
-                      data['tec'][:], 
-                      data['ipp_lon'][:], data['ipp_lat'][:], 
-                      data['el'][:],  
-                      maxgap=maxgap, maxjump=maxjump)
-
-
-def getContInt(times, tec, lon, lat, el,  maxgap=30, maxjump=1):
-    r = np.array(range(len(times)))
-    idx = np.isfinite(tec) & np.isfinite(lon) & np.isfinite(lat) & np.isfinite(el) & (el > 10.)
-    r = r[idx]
+def get_continuous_intervals(
+    times: np.ndarray, 
+    tec: np.ndarray, 
+    lon: np.ndarray, 
+    lat: np.ndarray, 
+    elevation: np.ndarray,
+    max_time_gap: float = 30,
+    max_tec_jump: float = 1,
+    min_elevation: float = 10.0
+) -> tuple[np.ndarray, list[tuple[int, int]]]:
+    """
+    Find continuous intervals in TEC data.
+
+    Args:
+        times: Array of timestamps
+        tec: Array of TEC values
+        lon: Array of longitudes
+        lat: Array of latitudes
+        elevation: Array of elevation angles
+        max_time_gap: Maximum allowed time gap between measurements in seconds
+        max_tec_jump: Maximum allowed TEC value jump between measurements
+        min_elevation: Minimum elevation angle in degrees
+
+    Returns:
+        Tuple of (valid indices array, list of continuous intervals)
+    """
+    indices = np.array(range(len(times)))
+    valid_data = (
+        np.isfinite(tec) & 
+        np.isfinite(lon) & 
+        np.isfinite(lat) & 
+        np.isfinite(elevation) & 
+        (elevation > min_elevation)
+    )
+    valid_indices = indices[valid_data]
     intervals = []
-    if len(r) == 0:
-        return idx, intervals
-    beginning = r[0]
-    last = r[0]
-    last_time = times[last]
-    for i in r[1:]:
-        if abs(times[i] - last_time) > maxgap or abs(tec[i] - tec[last]) > maxjump:
-            intervals.append((beginning, last))
-            beginning = i
-        last = i
-        last_time = times[last]
-        if i == r[-1]:
-            intervals.append((beginning, last))
-    return idx, intervals
+
+    if len(valid_indices) == 0:
+        return valid_data, intervals
+
+    interval_start = valid_indices[0]
+    last_index = valid_indices[0]
+    last_time = times[last_index]
+
+    for current_index in valid_indices[1:]:
+        if (abs(times[current_index] - last_time) > max_time_gap or 
+            abs(tec[current_index] - tec[last_index]) > max_tec_jump):
+            intervals.append((interval_start, last_index))
+            interval_start = current_index
+        last_index = current_index
+        last_time = times[last_index]
+        if current_index == valid_indices[-1]:
+            intervals.append((interval_start, last_index))
+
+    return valid_data, intervals
 
 def process_intervals(data, maxgap, maxjump, derivative, 
                       short = 3600, sparse = 600):
     result = defaultdict(list)
     tt = sec_of_day(data['datetime'])        
-    idx, intervals = getContInt(tt, 
+    idx, intervals = get_continuous_intervals(tt, 
                                 data['tec'], data['ipp_lon'],
                                 data['ipp_lat'], data['el'],  
-                                maxgap=maxgap, maxjump=maxjump)
+                                max_time_gap=maxgap, max_tec_jump=maxjump)
     #_, intervals = get_continuos_intervals(data, maxgap=maxgap, maxjump=maxjump)
     for start, fin in intervals:
 

mosgim/geo/geo.py

@@ -1,3 +1,5 @@
+"""Geometric calculations for solar and ionospheric positions."""
+
 from numpy import sin, cos, pi, arctan2, arcsin, floor, deg2rad
 
 RE = 6371.2

mosgim/geo/geomag.py

@@ -19,30 +19,37 @@
 
 
 @np.vectorize
-def geo2mag(theta, phi, date):
+def geo2mag(theta: float, phi: float, date) -> tuple[float, float]:
+
     ut = sec_of_day(date)
     doy = date.timetuple().tm_yday
     year = date.year
 
+    # Calculate subsolar point position
     phi_sbs, theta_sbs = subsol(year, doy, ut)
-    r_sbs = np.array([np.sin(theta_sbs) * np.cos(phi_sbs), np.sin(theta_sbs) * np.sin(phi_sbs), np.cos(theta_sbs)])
+    r_sbs = np.array([np.sin(theta_sbs) * np.cos(phi_sbs), 
+                      np.sin(theta_sbs) * np.sin(phi_sbs), 
+                      np.cos(theta_sbs)])
 
+    # Transform subsolar point to magnetic coordinates
     r_sbs_mag = GEOGRAPHIC_TRANSFORM.dot(r_sbs)
     theta_sbs_m = np.arccos(r_sbs_mag[2])
     phi_sbs_m = np.arctan2(r_sbs_mag[1], r_sbs_mag[0])
     if phi_sbs_m < 0.:
         phi_sbs_m = phi_sbs_m + 2. * np.pi
 
-
-    r = np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)])
+    # Transform input point to magnetic coordinates
+    r = np.array([np.sin(theta) * np.cos(phi), 
+                  np.sin(theta) * np.sin(phi), 
+                  np.cos(theta)])
     r_mag = GEOGRAPHIC_TRANSFORM.dot(r)
     theta_m = np.arccos(r_mag[2])
     phi_m = np.arctan2(r_mag[1], r_mag[0])
     if phi_m < 0.:
         phi_m = phi_m + 2. * np.pi
 
-
-    mlt = phi_m - phi_sbs_m + np.pi # np.radians(15.) * ut /3600. + phi_m + POLE_PHI  
+    # Calculate magnetic local time
+    mlt = phi_m - phi_sbs_m + np.pi
     if mlt < 0.:
         mlt = mlt + 2. * np.pi
     if mlt > 2. * np.pi:
@@ -59,26 +66,43 @@ def inclination(lat, lon, alt=300., year=2005.):
     if lon < 0:
         lon = lon + 360.      
     FACT = 180./np.pi
-    REm = 6371.2
-    x, y, z, f = calculate.igrf12syn(year, 2, REm + alt, lat, lon) # 2 stands for geocentric coordinates
+    REm = 6371.2  # Earth's mean radius in km
+
+    # Calculate magnetic field components using IGRF
+    x, y, z, f = calculate.igrf12syn(year, 2, REm + alt, lat, lon)  # 2 for geocentric coordinates
+
+    # Calculate horizontal component and inclination
     h = np.sqrt(x * x + y * y)
     i = FACT * np.arctan2(z, h)
     return i
 
 @np.vectorize
-def geo2modip(theta, phi, date):
+def geo2modip(theta: float, phi: float, date) -> tuple[float, float]:
+
     year = date.year
-    I = inclination(lat=np.rad2deg(np.pi/2 - theta), lon=np.rad2deg(phi), alt=300., year=year) # alt=300 for modip300
+
+    # Calculate magnetic inclination and modified dip latitude
+    I = inclination(lat=np.rad2deg(np.pi/2 - theta), 
+                   lon=np.rad2deg(phi), 
+                   alt=300., 
+                   year=year)  # alt=300 for modip300
     theta_m = np.pi/2 - np.arctan2(np.deg2rad(I), np.sqrt(np.cos(np.pi/2 - theta)))
+
+    # Calculate magnetic local time
     ut = sec_of_day(date)
-    phi_sbs = np.deg2rad(180. - ut*15./3600)
+    phi_sbs = np.deg2rad(180. - ut*15./3600)  # Subsolar longitude
+
+    # Normalize angles to [0, 2π]
     if phi_sbs < 0.:
         phi_sbs = phi_sbs + 2. * np.pi
     if phi < 0.:
         phi = phi + 2. * np.pi
+
+    # Calculate MLT
     mlt = phi - phi_sbs + np.pi
     if mlt < 0.:
         mlt = mlt + 2. * np.pi
     if mlt > 2. * np.pi:
         mlt = mlt - 2. * np.pi
+
     return theta_m, mlt