refactoring hw

mosgim/data/__init__.py

@@ -0,0 +1,25 @@
+# loader.py
+from loader import Loader
+from loader import LoaderHDF
+from loader import LoaderTxt
+from loader import LoaderRinex
+
+#tec_prepare.py
+#classses
+from tec_prepare import DataSourceType
+from tec_prepare import MagneticCoordType
+from tec_prepare import ProcessingType
+
+#functions
+from tec_prepare import process_data
+from tec_prepare import get_continuos_intervals
+from tec_prepare import process_intervals
+from tec_prepare import combine_data
+from tec_prepare import get_chunk_indexes
+from tec_prepare import calc_mag
+from tec_prepare import calc_mag_ref
+from tec_prepare import save_data
+from tec_prepare import get_data
+from tec_prepare import calculate_seed_mag_coordinates_parallel
+
+from tec_prepare import sites

mosgim/data/loader.py

@@ -1,14 +1,16 @@
 import os
 import time
 import numpy as np
+import re
 import h5py
 import concurrent.futures
 from datetime import datetime
 from warnings import warn
 from collections import defaultdict
 from pathlib import Path
 from concurrent.futures import ProcessPoolExecutor
-from numpy import cos, sin, sqrt, arctan, arctan2, rad2deg
+from numpy import( cos, sin, sqrt, 
+arctan, arctan2, rad2deg)
 
 from mosgim.geo.geo import HM
 from mosgim.geo.geo import sub_ionospheric
@@ -23,12 +25,12 @@ def __init__(self):
 
 class LoaderTxt(Loader):
 
-    def __init__(self, root_dir):
+    def __init__(self, root_dir:Path):
         super().__init__()
         self.dformat = "%Y-%m-%dT%H:%M:%S"
         self.root_dir = root_dir
 
-    def get_files(self, rootdir):
+    def get_files(self, rootdir:Path)->defaultdict[str,list[str]]:
         """
         Root directroy must contain folders with site name 
         Inside subfolders are *.dat files for every satellite
@@ -48,7 +50,7 @@ def get_files(self, rootdir):
             result[site].sort()
         return result
 
-    def load_data(self, filepath):
+    def load_data(self, filepath:Path)->tuple[np.array,Path]:
         convert = lambda x: datetime.strptime(x.decode("utf-8"), self.dformat)
         data = np.genfromtxt(filepath, 
                              comments='#', 
@@ -62,10 +64,8 @@ def load_data(self, filepath):
         #data = append_fields(data, 'sec_of_day', tt, np.float)
         return data, filepath
 
-    def __load_data_pool(self, filepath):
-        return self.load_data(filepath), filepath
 
-    def generate_data(self, sites=[]):
+    def generate_data(self, sites:list[str]=[]):
         files = self.get_files(self.root_dir)
         print(f'Collected {len(files)} sites')
         self.not_found_sites = sites[:]
@@ -84,7 +84,7 @@ def generate_data(self, sites=[]):
                     print(f'{sat_file} not processed. Reason: {e}')
             print(f'{site} contribute {count} files, takes {time.time() - st}')
 
-    def generate_data_pool(self, sites=[], nworkers=1):
+    def generate_data_pool(self, sites:list[str]=[], nworkers:int=1):
         files = self.get_files(self.root_dir)
         print(f'Collected {len(files)} sites')
         self.not_found_sites = sites[:]
@@ -103,17 +103,17 @@ def generate_data_pool(self, sites=[], nworkers=1):
                     except Exception as e:
                         print(f'{sat_file} not processed. Reason: {e}')
                 print(site)
-            for v in concurrent.futures.as_completed(queue):
-                yield v.result()
+            for cur_future in concurrent.futures.as_completed(queue):
+                yield cur_future.result()
 
 
 class LoaderHDF(Loader):
 
-    def __init__(self, hdf_path):
+    def __init__(self, hdf_path:Path):
         super().__init__()
         self.hdf_path = hdf_path
 
-    def get_files(self):
+    def get_files(self)->list[Path]:
         result = []
         for subdir, _, files in os.walk(self.hdf_path):
             for filename in files:
@@ -125,10 +125,10 @@ def get_files(self):
             raise ValueError(msg)
         return result
 
-    def __get_hdf_file(self):
+    def __get_hdf_file(self)->Path:
         return self.get_files()[0]
 
-    def generate_data(self, sites=[]):
+    def generate_data(self, sites:list[str]=[]):
         hdf_file = h5py.File(self.__get_hdf_file(), 'r')
         self.not_found_sites = sites[:]
         for site in hdf_file:
@@ -148,7 +148,7 @@ def generate_data(self, sites=[]):
                 ts = sat_data['timestamp'][:]
                 ipp_lat, ipp_lon = sub_ionospheric(slat, slon, HM, az, el)
 
-                arr['datetime'] = np.array([datetime.utcfromtimestamp(float(t)) for t in ts])
+                arr['datetime'] = np.array([datetime.fromtimestamp(float(t)) for t in ts])
                 arr['el'] = np.rad2deg(el)
                 arr['ipp_lat'] = np.rad2deg(ipp_lat)
                 arr['ipp_lon'] = np.rad2deg(ipp_lon)
@@ -160,22 +160,22 @@ def generate_data(self, sites=[]):
 
 class LoaderRinex(Loader):
 
-    def __init__(self, rinex_path, nav_path):
+    def __init__(self, rinex_path:Path, nav_path:Path):
         super().__init__()
         self.rinex_path = rinex_path
         self.nav_path = nav_path
         self.o_regx = re.compile((r'.[1-9][1-9]o'))
 
-    def __is_ofile(self, f):
+    def __is_ofile(self, file:Path)->bool:
         is_ofile = False
-        is_ofile = is_ofile or bool(self.o_regx.match(f[-4:])) 
-        is_ofile = is_ofile or f.endswith('.rnx')
-        is_ofile = is_ofile or f.endswith('.RNX')
+        is_ofile = is_ofile or bool(self.o_regx.match(file[-4:])) 
+        is_ofile = is_ofile or file.endswith('.rnx')
+        is_ofile = is_ofile or file.endswith('.RNX')
         return is_ofile
 
-    def get_files(self):
+    def get_files(self)->dict[str,Path]:
         result = dict()
-        for subdir, _, files in os.walk(rootdir):
+        for subdir, _, files in os.walk(self.rootdir):
             for filename in files:
                 filepath = Path(subdir) / filename
                 if not self.__is_ofile(filename):

mosgim/data/tec_prepare.py

@@ -7,9 +7,10 @@
 from collections import defaultdict
 from enum import Enum
 from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
 
-from mosgim.geo.geomag import geo2mag
-from mosgim.geo.geomag import geo2modip
+from mosgim.geo import geo2mag
+from mosgim.geo import geo2modip
 from mosgim.utils.time_util import sec_of_day, sec_of_interval
 
 sites = ['019b', '7odm', 'ab02', 'ab06', 'ab09', 'ab11', 'ab12', 'ab13',
@@ -73,7 +74,7 @@ class ProcessingType(Enum):
     def __str__(self):
         return self.value
 
-def process_data(data_generator):
+def process_data(data_generator)->defaultdict[str,list[any]]:
     all_data = defaultdict(list)
     count = 0
     for data, data_id in data_generator:
@@ -99,15 +100,15 @@ def process_data(data_generator):
     return all_data
 
 
-def get_continuos_intervals(data, maxgap=30, maxjump=1):
+def get_continuos_intervals(data:dict, maxgap:int=30, maxjump:int=1)->tuple[bool,list[tuple[int,int]]]:
     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):
+def getContInt(times:list[int], tec:float, lon:float, lat:float, el:float,  maxgap:int=30, maxjump:int=1)->tuple[bool,list[tuple[int,int]]]:
     r = np.array(range(len(times)))
     idx = np.isfinite(tec) & np.isfinite(lon) & np.isfinite(lat) & np.isfinite(el) & (el > 10.)
     r = r[idx]
@@ -127,8 +128,8 @@ def getContInt(times, tec, lon, lat, el,  maxgap=30, maxjump=1):
             intervals.append((beginning, last))
     return idx, intervals
 
-def process_intervals(data, maxgap, maxjump, derivative, 
-                      short = 3600, sparse = 600):
+def process_intervals(data:dict, maxgap:int, maxjump:int, derivative:bool, 
+                      short:int = 3600, sparse:int = 600)->defaultdict[str,list[any]]:
     result = defaultdict(list)
     tt = sec_of_day(data['datetime'])        
     idx, intervals = getContInt(tt, 
@@ -163,7 +164,7 @@ def process_intervals(data, maxgap, maxjump, derivative,
         result['ref'].append(data_ref)
     return result
 
-def combine_data(all_data, nchunks=1):
+def combine_data(all_data:dict, nchunks:int=1)->list[dict[str,any]]:
     count = 0
     for arr in all_data['dtec']:
         count += arr.shape[0]
@@ -198,7 +199,7 @@ def combine_data(all_data, nchunks=1):
     return combs
 
 
-def get_chunk_indexes(size, nchunks):
+def get_chunk_indexes(size:int, nchunks:int)->list[tuple[int,int]]:
     if nchunks > 1:
         step = int(size / nchunks)
         ichunks = [(i-step, i) for i in range(step, size, step)]
@@ -211,24 +212,24 @@ def get_chunk_indexes(size, nchunks):
         return [(0, size)]
 
 
-def calc_mag(comb, g2m):
+def calc_mag(comb:dict, g2m:function):
     colat, mlt = \
         g2m(np.pi/2 - rad(comb['lat']), rad(comb['lon']), comb['time'])  
     return colat, mlt
 
-def calc_mag_ref(comb, g2m):
+def calc_mag_ref(comb:dict, g2m:function):
     rcolat, rmlt = \
         g2m(np.pi/2 - rad(comb['rlat']), rad(comb['rlon']), comb['rtime'])  
     return  rcolat, rmlt
 
-def calc_mag_coordinates(comb):
+def calc_mag_coordinates(comb:dict)->dict[str,tuple[int,int]]:
     comb['colat_mdip'], comb['mlt_mdip'] = calc_mag(comb, geo2modip)  
     comb['rcolat_mdip'], comb['rmlt_mdip'] = calc_mag_ref(comb, geo2modip)
     comb['colat_mag'], comb['mlt_mag'] = calc_mag(comb, geo2mag)
     comb['rcolat_mag'], comb['rmlt_mag'] = calc_mag_ref(comb, geo2mag)
     return comb
 
-def calculate_seed_mag_coordinates_parallel(chunks, nworkers=3):
+def calculate_seed_mag_coordinates_parallel(chunks:list, nworkers=3):
     if len(chunks) < 1:
         return None
     if len(chunks) == 1:
@@ -262,7 +263,7 @@ def calculate_seed_mag_coordinates_parallel(chunks, nworkers=3):
     return comb
 
 
-def save_data(comb, modip_file, mag_file, day_date):
+def save_data(comb:dict, modip_file:Path, mag_file:Path, day_date:datetime):
     mags = [MagneticCoordType.mdip, MagneticCoordType.mag]
     for mtype, filename in zip(mags, [modip_file, mag_file]):
         data = get_data(comb, mtype, day_date)
@@ -271,7 +272,7 @@ def save_data(comb, modip_file, mag_file, day_date):
                 day = day_date,
                 **data)    
 
-def get_data(comb, mtype, day_date):
+def get_data(comb:dict, mtype, day_date:datetime)->dict[str,any]:
     postf = str(mtype)
     data = dict(time = sec_of_interval(comb['time'], day_date), 
                 mlt = comb['mlt_' + postf ], 

mosgim/geo/__init__.py

@@ -0,0 +1,7 @@
+# geomag.py
+from geomag import geo2mag
+from geomag import geo2modip
+from geomag import make_inclination
+
+#geo.py
+from geo import sub_ionospheric

mosgim/geo/geo.py

@@ -3,7 +3,7 @@
 RE = 6371.2
 HM=300
 
-def subsol(year, doy, ut):
+def sub_sol(year:int, doy:int, ut:int)->float:
     '''Finds subsolar geocentric longitude and latitude.
 
 
@@ -111,7 +111,7 @@ def subsol(year, doy, ut):
     return deg2rad(sbsllon), pi / 2 - deg2rad(sbsllat)
 
 
-def sub_ionospheric(s_lat, s_lon, hm, az, el, R=RE):
+def sub_ionospheric(s_lat:float, s_lon:float, hm:float, az:float, el:float, R:float=RE)->tuple[float,float]:
     """
     Calculates subionospheric point and delatas from site
     Parameters:

mosgim/geo/geomag.py

@@ -1,7 +1,7 @@
 import pyIGRF.calculate as calculate
 import numpy as np
-
-from .geo import subsol
+from datetime import datetime
+from .geo import sub_sol
 from mosgim.utils.time_util import sec_of_day
 # GEOMAGNETIC AND MODIP COORDINATES SECTION
 
@@ -19,12 +19,12 @@
 
 
 @np.vectorize
-def geo2mag(theta, phi, date):
+def geo2mag(theta:float, phi:float, date:datetime)->tuple[float,float]:
     ut = sec_of_day(date)
     doy = date.timetuple().tm_yday
     year = date.year
 
-    phi_sbs, theta_sbs = subsol(year, doy, ut)
+    phi_sbs, theta_sbs = sub_sol(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_mag = GEOGRAPHIC_TRANSFORM.dot(r_sbs)
@@ -51,7 +51,7 @@ def geo2mag(theta, phi, date):
     return theta_m, mlt
 
 
-def inclination(lat, lon, alt=300., year=2005.):
+def make_inclination(lat:float, lon:float, alt:int=300., year:int=2005.)->float:
     """
     :return
          I is inclination (+ve down)
@@ -66,9 +66,9 @@ def inclination(lat, lon, alt=300., year=2005.):
     return i
 
 @np.vectorize
-def geo2modip(theta, phi, date):
+def geo2modip(theta:float, phi:float, date:datetime)->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
+    I = make_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)))
     ut = sec_of_day(date)
     phi_sbs = np.deg2rad(180. - ut*15./3600)