oafuncs 0.0.90__py2.py3-none-any.whl → 0.0.92__py2.py3-none-any.whl

oafuncs/oa_down/hycom_3hourly.py

@@ -26,19 +26,24 @@ from threading import Lock
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import xarray as xr
 import requests
 from rich import print
 from rich.progress import Progress
+import netCDF4 as nc
 
 from oafuncs.oa_down.user_agent import get_ua
 from oafuncs.oa_file import file_size, mean_size
+from oafuncs.oa_nc import check as check_nc
+from oafuncs.oa_nc import modify as modify_nc
+from oafuncs.oa_down.idm import downloader as idm_downloader
 
 warnings.filterwarnings("ignore", category=RuntimeWarning, message="Engine '.*' loading failed:.*")
 
 __all__ = ["draw_time_range", "download", "how_to_use", "get_time_list"]
 
 
-def get_initial_data():
+def _get_initial_data():
     global variable_info, data_info, var_group, single_var_group
     # ----------------------------------------------
     # variable
@@ -305,14 +310,14 @@ def get_time_list(time_s, time_e, delta, interval_type="hour"):
     return dt_list
 
 
-def transform_time(time_str):
+def _transform_time(time_str):
     # old_time = '2023080203'
     # time_new = '2023-08-02T03%3A00%3A00Z'
     time_new = f"{time_str[:4]}-{time_str[4:6]}-{time_str[6:8]}T{time_str[8:10]}%3A00%3A00Z"
     return time_new
 
 
-def get_query_dict(var, lon_min, lon_max, lat_min, lat_max, time_str_ymdh, time_str_end=None, mode="single_depth", depth=None, level_num=None):
+def _get_query_dict(var, lon_min, lon_max, lat_min, lat_max, time_str_ymdh, time_str_end=None, mode="single_depth", depth=None, level_num=None):
     query_dict = {
         "var": variable_info[var]["var_name"],
         "north": lat_max,
@@ -331,11 +336,11 @@ def get_query_dict(var, lon_min, lon_max, lat_min, lat_max, time_str_ymdh, time_
     }
 
     if time_str_end is not None:
-        query_dict["time_start"] = transform_time(time_str_ymdh)
-        query_dict["time_end"] = transform_time(time_str_end)
+        query_dict["time_start"] = _transform_time(time_str_ymdh)
+        query_dict["time_end"] = _transform_time(time_str_end)
         query_dict["timeStride"] = 1
     else:
-        query_dict["time"] = transform_time(time_str_ymdh)
+        query_dict["time"] = _transform_time(time_str_ymdh)
 
     def get_nearest_level_index(depth):
         level_depth = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 125.0, 150.0, 200.0, 250.0, 300.0, 350.0, 400.0, 500.0, 600.0, 700.0, 800.0, 900.0, 1000.0, 1250.0, 1500.0, 2000.0, 2500.0, 3000.0, 4000.0, 5000]
@@ -360,7 +365,7 @@ def get_query_dict(var, lon_min, lon_max, lat_min, lat_max, time_str_ymdh, time_
     return query_dict
 
 
-def check_time_in_dataset_and_version(time_input, time_end=None):
+def _check_time_in_dataset_and_version(time_input, time_end=None):
     # 判断是处理单个时间点还是时间范围
     is_single_time = time_end is None
 
@@ -417,8 +422,8 @@ def check_time_in_dataset_and_version(time_input, time_end=None):
         if is_single_time:
             return True
         else:
-            base_url_s = get_base_url(d_list[0], v_list[0], "u", str(time_start))
-            base_url_e = get_base_url(d_list[0], v_list[0], "u", str(time_end))
+            base_url_s = _get_base_url(d_list[0], v_list[0], "u", str(time_start))
+            base_url_e = _get_base_url(d_list[0], v_list[0], "u", str(time_end))
             if base_url_s == base_url_e:
                 return True
             else:
@@ -429,7 +434,7 @@ def check_time_in_dataset_and_version(time_input, time_end=None):
         return False
 
 
-def ensure_time_in_specific_dataset_and_version(dataset_name, version_name, time_input, time_end=None):
+def _ensure_time_in_specific_dataset_and_version(dataset_name, version_name, time_input, time_end=None):
     # 根据时间长度补全时间格式
     if len(str(time_input)) == 8:
         time_input = str(time_input) + "00"
@@ -468,7 +473,7 @@ def ensure_time_in_specific_dataset_and_version(dataset_name, version_name, time
         return False
 
 
-def direct_choose_dataset_and_version(time_input, time_end=None):
+def _direct_choose_dataset_and_version(time_input, time_end=None):
     # 假设 data_info 是一个字典，包含了数据集和版本的信息
     # 示例结构：data_info['hourly']['dataset'][dataset_name]['version'][version_name]['time_range']
 
@@ -507,7 +512,7 @@ def direct_choose_dataset_and_version(time_input, time_end=None):
     return dataset_name_out, version_name_out
 
 
-def get_base_url(dataset_name, version_name, var, ymdh_str):
+def _get_base_url(dataset_name, version_name, var, ymdh_str):
     year_str = int(ymdh_str[:4])
     url_dict = data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["url"]
     classification_method = data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["classification"]
@@ -548,160 +553,272 @@ def get_base_url(dataset_name, version_name, var, ymdh_str):
     return base_url
 
 
-def get_submit_url(dataset_name, version_name, var, ymdh_str, query_dict):
-    base_url = get_base_url(dataset_name, version_name, var, ymdh_str)
+def _get_submit_url(dataset_name, version_name, var, ymdh_str, query_dict):
+    base_url = _get_base_url(dataset_name, version_name, var, ymdh_str)
     if isinstance(query_dict["var"], str):
         query_dict["var"] = [query_dict["var"]]
     target_url = base_url + "&".join(f"var={var}" for var in query_dict["var"]) + "&" + "&".join(f"{key}={value}" for key, value in query_dict.items() if key != "var")
     return target_url
 
 
-def clear_existing_file(file_full_path):
+def _clear_existing_file(file_full_path):
     if os.path.exists(file_full_path):
         os.remove(file_full_path)
         print(f"{file_full_path} has been removed")
 
 
-def check_existing_file(file_full_path, min_size):
+def _check_existing_file(file_full_path, avg_size):
     if os.path.exists(file_full_path):
         print(f"[bold #FFA54F]{file_full_path} exists")
         fsize = file_size(file_full_path)
-        if min_size:
-            if fsize < min_size:
-                print(f"[bold #FFA54F]{file_full_path} ({fsize:.2f} KB) may be incomplete")
-                # clear_existing_file(file_full_path)
+        delta_size_ratio = (fsize - avg_size) / avg_size
+        if abs(delta_size_ratio) > 0.025:
+            if check_nc(file_full_path):
+                # print(f"File size is abnormal but can be opened normally, file size: {fsize:.2f} KB")
+                if not _check_ftime(file_full_path,if_print=True):
+                    return False
+                else:
+                    return True
+            else:
+                print(f"File size is abnormal and cannot be opened, {file_full_path}: {fsize:.2f} KB")
+                return False
+        else:
+            if not _check_ftime(file_full_path,if_print=True):
                 return False
             else:
                 return True
-        if fsize < 5:
-            print(f"[bold #FFA54F]{file_full_path} ({fsize:.2f} KB) may be incomplete")
-            # clear_existing_file(file_full_path)
-            return False
-        else:
-            return True
     else:
-        # print(f'{file_full_path} does not exist')
         return False
 
 
-def download_file(target_url, store_path, file_name, check=False):
-    # Check if the file exists
-    fname = Path(store_path) / file_name
-    file_name_split = file_name.split("_")
-    file_name_split = file_name_split[:-1]
-    # same_file = f"{file_name_split[0]}_{file_name_split[1]}*nc"
-    same_file = "_".join(file_name_split) + "*nc"
-
+def _get_mean_size30(store_path, same_file):
     if same_file not in fsize_dict.keys():
-        # print(f'Same file name: {same_file}')
-        fsize_dict[same_file] = {"size": 0, "count": 0}
+            # print(f'Same file name: {same_file}')
+            fsize_dict[same_file] = {"size": 0, "count": 0}
 
     if fsize_dict[same_file]["count"] < 30 or fsize_dict[same_file]["size"] == 0:
         # 更新30次文件最小值，后续认为可以代表所有文件，不再更新占用时间
         fsize_mean = mean_size(store_path, same_file, max_num=30)
-        set_min_size = fsize_mean * 0.8
+        set_min_size = fsize_mean * 0.95
         fsize_dict[same_file]["size"] = set_min_size
         fsize_dict[same_file]["count"] += 1
     else:
         set_min_size = fsize_dict[same_file]["size"]
-    if check:
-        if check_existing_file(fname, set_min_size):
-            count_dict["skip"] += 1
-            return
-    clear_existing_file(fname)
+    return set_min_size
 
-    # -----------------------------------------------
-    print(f"[bold #f0f6d0]Requesting {file_name}...")
-    # 创建会话
-    s = requests.Session()
-    download_success = False
-    request_times = 0
 
-    def calculate_wait_time(time_str, target_url):
-        # 定义正则表达式，匹配YYYYMMDDHH格式的时间
-        time_pattern = r"\d{10}"
+def _get_mean_size_move(same_file, current_file):
+    # 获取锁
+    with fsize_dict_lock: # 全局锁，确保同一时间只能有一个线程访问
+        # 初始化字典中的值，如果文件不在字典中
+        if same_file not in fsize_dict.keys():
+            fsize_dict[same_file] = {"size_list": [], "mean_size": 1.0}
 
-        # 定义两个字符串
-        # str1 = 'HYCOM_water_u_2018010100-2018010112.nc'
-        # str2 = 'HYCOM_water_u_2018010100.nc'
+        tolerance_ratio = 0.025  # 容忍的阈值比例
+        current_file_size = file_size(current_file)
+
+        # 如果列表不为空，则计算平均值，否则保持为1
+        if fsize_dict[same_file]["size_list"]:
+            fsize_dict[same_file]["mean_size"] = sum(fsize_dict[same_file]["size_list"]) / len(fsize_dict[same_file]["size_list"])
+            fsize_dict[same_file]["mean_size"] = max(fsize_dict[same_file]["mean_size"], 1.0)
+        else:
+            fsize_dict[same_file]["mean_size"] = 1.0
+
+        size_difference_ratio = (current_file_size - fsize_dict[same_file]["mean_size"]) / fsize_dict[same_file]["mean_size"]
+
+        if abs(size_difference_ratio) > tolerance_ratio:
+            if check_nc(current_file):
+                # print(f"File size is abnormal but can be opened normally, file size: {current_file_size:.2f} KB")
+                # 文件可以正常打开，但大小异常，保留当前文件大小
+                fsize_dict[same_file]["size_list"] = [current_file_size]
+                fsize_dict[same_file]["mean_size"] = current_file_size
+            else:
+                _clear_existing_file(current_file)
+                print(f"File size is abnormal, may need to be downloaded again, file size: {current_file_size:.2f} KB")
+        else:
+            # 添加当前文件大小到列表中，并更新计数
+            fsize_dict[same_file]["size_list"].append(current_file_size)
 
-        # 使用正则表达式查找时间
-        times_in_str = re.findall(time_pattern, time_str)
+    # 返回调整后的平均值，这里根据您的需求，返回的是添加新值之前的平均值
+    return fsize_dict[same_file]["mean_size"]
 
-        # 计算每个字符串中的时间数量
-        num_times_str = len(times_in_str)
 
-        if num_times_str > 1:
-            delta_t = datetime.datetime.strptime(times_in_str[1], "%Y%m%d%H") - datetime.datetime.strptime(times_in_str[0], "%Y%m%d%H")
-            delta_t = delta_t.total_seconds() / 3600
-            delta_t = delta_t / 3 + 1
+def _check_ftime(nc_file, tname="time", if_print=False):
+    if not os.path.exists(nc_file):
+        return False
+    nc_file = str(nc_file)
+    try:
+        ds = xr.open_dataset(nc_file)
+        real_time = ds[tname].values[0]
+        ds.close()
+        real_time = str(real_time)[:13]
+        real_time = real_time.replace("-", "").replace("T", "")
+        # -----------------------------------------------------
+        f_time = re.findall(r"\d{10}", nc_file)[0]
+        if real_time == f_time:
+            return True
         else:
-            delta_t = 1
-        # 单个要素最多等待5分钟，不宜太短，太短可能请求失败；也不宜太长，太长可能会浪费时间
-        num_var = int(target_url.count("var="))
-        if num_var <= 0:
-            num_var = 1
-        return int(delta_t * 5 * 60 * num_var)
-
-    max_timeout = calculate_wait_time(file_name, target_url)
-    print(f"[bold #912dbc]Max timeout: {max_timeout} seconds")
-
-    # print(f'Download_start_time: {datetime.datetime.now()}')
-    download_time_s = datetime.datetime.now()
-    order_list = ["1st", "2nd", "3rd", "4th", "5th", "6th", "7th", "8th", "9th", "10th"]
-    while not download_success:
-        if request_times >= 10:
-            # print(f'下载失败，已重试 {request_times} 次\n可先跳过，后续再试')
-            print(f"[bold #ffe5c0]Download failed after {request_times} times\nYou can skip it and try again later")
-            count_dict["fail"] += 1
-            break
-        if request_times > 0:
-            # print(f'\r正在重试第 {request_times} 次', end="")
-            print(f"[bold #ffe5c0]Retrying the {order_list[request_times-1]} time...")
-        # 尝试下载文件
-        try:
-            headers = {"User-Agent": get_ua()}
-            """ response = s.get(target_url, headers=headers, timeout=random.randint(5, max_timeout))
-            response.raise_for_status()  # 如果请求返回的不是200，将抛出HTTPError异常
-
-            # 保存文件
-            with open(filename, 'wb') as f:
-                f.write(response.content) """
-
-            response = s.get(target_url, headers=headers, stream=True, timeout=random.randint(5, max_timeout))  # 启用流式传输
-            response.raise_for_status()  # 如果请求返回的不是200，将抛出HTTPError异常
-            # 保存文件
-            with open(fname, "wb") as f:
-                print(f"[bold #96cbd7]Downloading {file_name}...")
-                for chunk in response.iter_content(chunk_size=1024):
-                    if chunk:
-                        f.write(chunk)
-
-            f.close()
-
-            # print(f'\r文件 {fname} 下载成功', end="")
-            if os.path.exists(fname):
-                download_success = True
-                download_time_e = datetime.datetime.now()
-                download_delta = download_time_e - download_time_s
-                print(f"[#3dfc40]File [bold #dfff73]{fname} [#3dfc40]has been downloaded successfully, Time: [#39cbdd]{download_delta}")
-                count_dict["success"] += 1
-                # print(f'Download_end_time: {datetime.datetime.now()}')
-
-        except requests.exceptions.HTTPError as errh:
-            print(f"Http Error: {errh}")
-        except requests.exceptions.ConnectionError as errc:
-            print(f"Error Connecting: {errc}")
-        except requests.exceptions.Timeout as errt:
-            print(f"Timeout Error: {errt}")
-        except requests.exceptions.RequestException as err:
-            print(f"OOps: Something Else: {err}")
-
-        time.sleep(3)
-        request_times += 1
-
-
-def check_hour_is_valid(ymdh_str):
+            if if_print:
+                print(f"[bold #daff5c]File time error, file/real time: [bold blue]{f_time}/{real_time}")
+            return False
+    except Exception as e:
+        if if_print:
+            print(f"[bold #daff5c]File time check failed, {nc_file}: {e}")
+        return False
+
+
+def _correct_time(nc_file):
+    # 打开NC文件
+    dataset = nc.Dataset(nc_file)
+
+    # 读取时间单位
+    time_units = dataset.variables["time"].units
+
+    # 关闭文件
+    dataset.close()
+
+    # 解析时间单位字符串以获取时间原点
+    origin_str = time_units.split("since")[1].strip()
+    origin_datetime = datetime.datetime.strptime(origin_str, "%Y-%m-%d %H:%M:%S")
+
+    # 从文件名中提取日期字符串
+    given_date_str = re.findall(r"\d{10}", str(nc_file))[0]
+
+    # 将提取的日期字符串转换为datetime对象
+    given_datetime = datetime.datetime.strptime(given_date_str, "%Y%m%d%H")
+
+    # 计算给定日期与时间原点之间的差值（以小时为单位）
+    time_difference = (given_datetime - origin_datetime).total_seconds()
+    if "hours" in time_units:
+        time_difference /= 3600
+    elif "days" in time_units:
+        time_difference /= 3600 * 24
+
+    # 修改NC文件中的时间变量
+    modify_nc(nc_file, "time", None, time_difference)
+
+
+
+def _download_file(target_url, store_path, file_name, check=False):
+    # Check if the file exists
+    fname = Path(store_path) / file_name
+    file_name_split = file_name.split("_")
+    file_name_split = file_name_split[:-1]
+    # same_file = f"{file_name_split[0]}_{file_name_split[1]}*nc"
+    same_file = "_".join(file_name_split) + "*nc"
+    
+    if check:        
+        if same_file not in fsize_dict.keys(): # 对第一个文件单独进行检查，因为没有大小可以对比
+            check_nc(fname,if_delete=True)
+
+        # set_min_size = _get_mean_size30(store_path, same_file) # 原方案，只30次取平均值；若遇变化，无法判断
+        get_mean_size = _get_mean_size_move(same_file, fname)
+    
+        if _check_existing_file(fname, get_mean_size):
+            count_dict["skip"] += 1
+            return
+    _clear_existing_file(fname)
+
+    if not use_idm:
+        # -----------------------------------------------
+        print(f"[bold #f0f6d0]Requesting {file_name} ...")
+        # 创建会话
+        s = requests.Session()
+        download_success = False
+        request_times = 0
+
+        def calculate_wait_time(time_str, target_url):
+            # 定义正则表达式，匹配YYYYMMDDHH格式的时间
+            time_pattern = r"\d{10}"
+
+            # 定义两个字符串
+            # str1 = 'HYCOM_water_u_2018010100-2018010112.nc'
+            # str2 = 'HYCOM_water_u_2018010100.nc'
+
+            # 使用正则表达式查找时间
+            times_in_str = re.findall(time_pattern, time_str)
+
+            # 计算每个字符串中的时间数量
+            num_times_str = len(times_in_str)
+
+            if num_times_str > 1:
+                delta_t = datetime.datetime.strptime(times_in_str[1], "%Y%m%d%H") - datetime.datetime.strptime(times_in_str[0], "%Y%m%d%H")
+                delta_t = delta_t.total_seconds() / 3600
+                delta_t = delta_t / 3 + 1
+            else:
+                delta_t = 1
+            # 单个要素最多等待5分钟，不宜太短，太短可能请求失败；也不宜太长，太长可能会浪费时间
+            num_var = int(target_url.count("var="))
+            if num_var <= 0:
+                num_var = 1
+            return int(delta_t * 5 * 60 * num_var)
+
+        max_timeout = calculate_wait_time(file_name, target_url)
+        print(f"[bold #912dbc]Max timeout: {max_timeout} seconds")
+
+        # print(f'Download_start_time: {datetime.datetime.now()}')
+        download_time_s = datetime.datetime.now()
+        order_list = ["1st", "2nd", "3rd", "4th", "5th", "6th", "7th", "8th", "9th", "10th"]
+        while not download_success:
+            if request_times >= 10:
+                # print(f'下载失败，已重试 {request_times} 次\n可先跳过，后续再试')
+                print(f"[bold #ffe5c0]Download failed after {request_times} times\nYou can skip it and try again later")
+                count_dict["fail"] += 1
+                break
+            if request_times > 0:
+                # print(f'\r正在重试第 {request_times} 次', end="")
+                print(f"[bold #ffe5c0]Retrying the {order_list[request_times-1]} time...")
+            # 尝试下载文件
+            try:
+                headers = {"User-Agent": get_ua()}
+                """ response = s.get(target_url, headers=headers, timeout=random.randint(5, max_timeout))
+                response.raise_for_status()  # 如果请求返回的不是200，将抛出HTTPError异常
+
+                # 保存文件
+                with open(filename, 'wb') as f:
+                    f.write(response.content) """
+
+                response = s.get(target_url, headers=headers, stream=True, timeout=random.randint(5, max_timeout))  # 启用流式传输
+                response.raise_for_status()  # 如果请求返回的不是200，将抛出HTTPError异常
+                # 保存文件
+                with open(fname, "wb") as f:
+                    print(f"[bold #96cbd7]Downloading {file_name} ...")
+                    for chunk in response.iter_content(chunk_size=1024):
+                        if chunk:
+                            f.write(chunk)
+
+                f.close()
+                
+                if not _check_ftime(fname):
+                    _correct_time(fname)
+
+                # print(f'\r文件 {fname} 下载成功', end="")
+                if os.path.exists(fname):
+                    download_success = True
+                    download_time_e = datetime.datetime.now()
+                    download_delta = download_time_e - download_time_s
+                    print(f"[#3dfc40]File [bold #dfff73]{fname} [#3dfc40]has been downloaded successfully, Time: [#39cbdd]{download_delta}")
+                    count_dict["success"] += 1
+                    # print(f'Download_end_time: {datetime.datetime.now()}')
+
+            except requests.exceptions.HTTPError as errh:
+                print(f"Http Error: {errh}")
+            except requests.exceptions.ConnectionError as errc:
+                print(f"Error Connecting: {errc}")
+            except requests.exceptions.Timeout as errt:
+                print(f"Timeout Error: {errt}")
+            except requests.exceptions.RequestException as err:
+                print(f"OOps: Something Else: {err}")
+
+            time.sleep(3)
+            request_times += 1
+    else:
+        idm_downloader(target_url, store_path, file_name, given_idm_engine)
+        idm_download_list.append(fname)
+        print(f"[bold #3dfc40]File [bold #dfff73]{fname} [#3dfc40]has been submit to IDM for downloading")
+
+
+def _check_hour_is_valid(ymdh_str):
     # hour should be 00, 03, 06, 09, 12, 15, 18, 21
     hh = int(str(ymdh_str[-2:]))
     if hh in [0, 3, 6, 9, 12, 15, 18, 21]:
@@ -710,9 +827,9 @@ def check_hour_is_valid(ymdh_str):
         return False
 
 
-def check_dataset_version(dataset_name, version_name, download_time, download_time_end=None):
+def _check_dataset_version(dataset_name, version_name, download_time, download_time_end=None):
     if dataset_name is not None and version_name is not None:
-        just_ensure = ensure_time_in_specific_dataset_and_version(dataset_name, version_name, download_time, download_time_end)
+        just_ensure = _ensure_time_in_specific_dataset_and_version(dataset_name, version_name, download_time, download_time_end)
         if just_ensure:
             return dataset_name, version_name
         else:
@@ -725,7 +842,7 @@ def check_dataset_version(dataset_name, version_name, download_time, download_ti
         download_time_str = download_time_str + "00"
 
     # 检查小时是否有效（如果需要的话）
-    if download_time_end is None and not check_hour_is_valid(download_time_str):
+    if download_time_end is None and not _check_hour_is_valid(download_time_str):
         print("Please ensure the hour is 00, 03, 06, 09, 12, 15, 18, 21")
         raise ValueError("The hour is invalid")
 
@@ -733,18 +850,18 @@ def check_dataset_version(dataset_name, version_name, download_time, download_ti
     if download_time_end is not None:
         if len(str(download_time_end)) == 8:
             download_time_end = str(download_time_end) + "21"
-        have_data = check_time_in_dataset_and_version(download_time_str, download_time_end)
+        have_data = _check_time_in_dataset_and_version(download_time_str, download_time_end)
         if have_data:
-            return direct_choose_dataset_and_version(download_time_str, download_time_end)
+            return _direct_choose_dataset_and_version(download_time_str, download_time_end)
     else:
-        have_data = check_time_in_dataset_and_version(download_time_str)
+        have_data = _check_time_in_dataset_and_version(download_time_str)
         if have_data:
-            return direct_choose_dataset_and_version(download_time_str)
+            return _direct_choose_dataset_and_version(download_time_str)
 
     return None, None
 
 
-def get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end=None):
+def _get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end=None):
     # year_str = str(download_time)[:4]
     ymdh_str = str(download_time)
     if depth is not None and level_num is not None:
@@ -760,19 +877,19 @@ def get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max
     else:
         # print("Full depth or full level data will be downloaded...")
         which_mode = "full"
-    query_dict = get_query_dict(var, lon_min, lon_max, lat_min, lat_max, download_time, download_time_end, which_mode, depth, level_num)
-    submit_url = get_submit_url(dataset_name, version_name, var, ymdh_str, query_dict)
+    query_dict = _get_query_dict(var, lon_min, lon_max, lat_min, lat_max, download_time, download_time_end, which_mode, depth, level_num)
+    submit_url = _get_submit_url(dataset_name, version_name, var, ymdh_str, query_dict)
     return submit_url
 
 
-def prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90, download_time="2024083100", download_time_end=None, depth=None, level_num=None, store_path=None, dataset_name=None, version_name=None, check=False):
+def _prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90, download_time="2024083100", download_time_end=None, depth=None, level_num=None, store_path=None, dataset_name=None, version_name=None, check=False):
     print("[bold #ecdbfe]-" * 160)
     download_time = str(download_time)
     if download_time_end is not None:
         download_time_end = str(download_time_end)
-        dataset_name, version_name = check_dataset_version(dataset_name, version_name, download_time, download_time_end)
+        dataset_name, version_name = _check_dataset_version(dataset_name, version_name, download_time, download_time_end)
     else:
-        dataset_name, version_name = check_dataset_version(dataset_name, version_name, download_time)
+        dataset_name, version_name = _check_dataset_version(dataset_name, version_name, download_time)
     if dataset_name is None and version_name is None:
         count_dict["no_data"] += 1
         if download_time_end is not None:
@@ -787,11 +904,11 @@ def prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_max
     if isinstance(var, list):
         if len(var) == 1:
             var = var[0]
-            submit_url = get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end)
+            submit_url = _get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end)
             file_name = f"HYCOM_{variable_info[var]['var_name']}_{download_time}.nc"
             if download_time_end is not None:
                 file_name = f"HYCOM_{variable_info[var]['var_name']}_{download_time}-{download_time_end}.nc"  # 这里时间不能用下划线，不然后续处理查找同一变量文件会出问题
-            download_file(submit_url, store_path, file_name, check)
+            _download_file(submit_url, store_path, file_name, check)
         else:
             if download_time < "2024081012":
                 varlist = [_ for _ in var]
@@ -804,7 +921,7 @@ def prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_max
                         continue
 
                     var = current_group[0]
-                    submit_url = get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end)
+                    submit_url = _get_submit_url_var(var, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end)
                     file_name = f"HYCOM_{variable_info[var]['var_name']}_{download_time}.nc"
                     old_str = f'var={variable_info[var]["var_name"]}'
                     new_str = f'var={variable_info[var]["var_name"]}'
@@ -816,17 +933,17 @@ def prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_max
                         file_name = f"HYCOM_{key}_{download_time}.nc"
                         if download_time_end is not None:
                             file_name = f"HYCOM_{key}_{download_time}-{download_time_end}.nc"  # 这里时间不能用下划线，不然后续处理查找同一变量文件会出问题
-                    download_file(submit_url, store_path, file_name, check)
+                    _download_file(submit_url, store_path, file_name, check)
             else:
                 for v in var:
-                    submit_url = get_submit_url_var(v, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end)
+                    submit_url = _get_submit_url_var(v, depth, level_num, lon_min, lon_max, lat_min, lat_max, dataset_name, version_name, download_time, download_time_end)
                     file_name = f"HYCOM_{variable_info[v]['var_name']}_{download_time}.nc"
                     if download_time_end is not None:
                         file_name = f"HYCOM_{variable_info[v]['var_name']}_{download_time}-{download_time_end}.nc"
-                    download_file(submit_url, store_path, file_name, check)
+                    _download_file(submit_url, store_path, file_name, check)
 
 
-def convert_full_name_to_short_name(full_name):
+def _convert_full_name_to_short_name(full_name):
     for var, info in variable_info.items():
         if full_name == info["var_name"] or full_name == info["standard_name"] or full_name == var:
             return var
@@ -836,7 +953,7 @@ def convert_full_name_to_short_name(full_name):
     return False
 
 
-def download_task(var, time_str, time_str_end, lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, check):
+def _download_task(var, time_str, time_str_end, lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, check):
     """
     # 并行下载任务
     # 这个函数是为了并行下载而设置的，是必须的，直接调用direct_download并行下载会出问题
@@ -847,10 +964,10 @@ def download_task(var, time_str, time_str_end, lon_min, lon_max, lat_min, lat_ma
     因此，即使多个任务同时执行，也不会出现数据交互错乱的问题。
     """
 
-    prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, time_str_end, depth, level, store_path, dataset_name, version_name, check)
+    _prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, time_str_end, depth, level, store_path, dataset_name, version_name, check)
 
 
-def done_callback(future, progress, task, total, counter_lock):
+def _done_callback(future, progress, task, total, counter_lock):
     """
     # 并行下载任务的回调函数
     # 这个函数是为了并行下载而设置的，是必须的，直接调用direct_download并行下载会出问题
@@ -866,7 +983,7 @@ def done_callback(future, progress, task, total, counter_lock):
         progress.update(task, advance=1, description=f"[cyan]Downloading... {parallel_counter}/{total}")
 
 
-def download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90, depth=None, level=None, store_path=None, dataset_name=None, version_name=None, num_workers=None, check=False, ftimes=1):
+def _download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90, depth=None, level=None, store_path=None, dataset_name=None, version_name=None, num_workers=None, check=False, ftimes=1):
     """
     Description:
     Download the data of single time or a series of time
@@ -895,7 +1012,7 @@ def download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_min
         parallel_counter = 0
         counter_lock = Lock()  # 创建一个锁，线程安全的计数器
     if ymdh_time_s == ymdh_time_e:
-        prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, ymdh_time_s, None, depth, level, store_path, dataset_name, version_name, check)
+        _prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, ymdh_time_s, None, depth, level, store_path, dataset_name, version_name, check)
     elif int(ymdh_time_s) < int(ymdh_time_e):
         print("Downloading a series of files...")
         time_list = get_time_list(ymdh_time_s, ymdh_time_e, 3, "hour")
@@ -905,16 +1022,16 @@ def download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_min
                 if num_workers is None or num_workers <= 1:
                     # 串行方式
                     for i, time_str in enumerate(time_list):
-                        prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, None, depth, level, store_path, dataset_name, version_name, check)
+                        _prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, None, depth, level, store_path, dataset_name, version_name, check)
                         progress.update(task, advance=1, description=f"[cyan]Downloading... {i+1}/{len(time_list)}")
                 else:
                     # 并行方式
                     with ThreadPoolExecutor(max_workers=num_workers) as executor:
-                        futures = [executor.submit(download_task, var, time_str, None, lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, check) for time_str in time_list]
+                        futures = [executor.submit(_download_task, var, time_str, None, lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, check) for time_str in time_list]
                         """ for i, future in enumerate(futures):
                             future.add_done_callback(lambda _: progress.update(task, advance=1, description=f"[cyan]Downloading... {i+1}/{len(time_list)}")) """
                         for feature in as_completed(futures):
-                            done_callback(feature, progress, task, len(time_list), counter_lock)
+                            _done_callback(feature, progress, task, len(time_list), counter_lock)
             else:
                 new_time_list = get_time_list(ymdh_time_s, ymdh_time_e, 3 * ftimes, "hour")
                 total_num = len(new_time_list)
@@ -923,21 +1040,21 @@ def download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_min
                     for i, time_str in enumerate(new_time_list):
                         time_str_end_index = int(min(len(time_list) - 1, int(i * ftimes + ftimes - 1)))
                         time_str_end = time_list[time_str_end_index]
-                        prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, time_str_end, depth, level, store_path, dataset_name, version_name, check)
+                        _prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, time_str_end, depth, level, store_path, dataset_name, version_name, check)
                         progress.update(task, advance=1, description=f"[cyan]Downloading... {i+1}/{total_num}")
                 else:
                     # 并行方式
                     with ThreadPoolExecutor(max_workers=num_workers) as executor:
-                        futures = [executor.submit(download_task, var, new_time_list[i], time_list[int(min(len(time_list) - 1, int(i * ftimes + ftimes - 1)))], lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, check) for i in range(total_num)]
+                        futures = [executor.submit(_download_task, var, new_time_list[i], time_list[int(min(len(time_list) - 1, int(i * ftimes + ftimes - 1)))], lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, check) for i in range(total_num)]
                         """ for i, future in enumerate(futures):
                             future.add_done_callback(lambda _: progress.update(task, advance=1, description=f"[cyan]Downloading... {i+1}/{total_num}")) """
                         for feature in as_completed(futures):
-                            done_callback(feature, progress, task, len(time_list), counter_lock)
+                            _done_callback(feature, progress, task, len(time_list), counter_lock)
     else:
         print("Please ensure the time_s is no more than time_e")
 
 
-def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90, depth=None, level=None, store_path=None, dataset_name=None, version_name=None, num_workers=None, check=False, ftimes=1):
+def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90, depth=None, level=None, store_path=None, dataset_name=None, version_name=None, num_workers=None, check=False, ftimes=1, idm_engine=None):
     """
     Description:
         Download the data of single time or a series of time
@@ -958,11 +1075,12 @@ def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, l
         num_workers: int, the number of workers, default is None, if not set, the number of workers will be 1; suggest not to set the number of workers too large
         check: bool, whether to check the existing file, default is False, if set to True, the existing file will be checked and not downloaded again; else, the existing file will be covered
         ftimes: int, the number of time in one file, default is 1, if set to 1, the data of single time will be downloaded; the maximum is 8, if set to 8, the data of 8 times will be downloaded in one file
+        idm_engine: str, the IDM engine, default is None, if set, the IDM will be used to download the data; example: "D:\\Programs\\Internet Download Manager\\IDMan.exe"
 
     Returns:
         None
     """
-    get_initial_data()
+    _get_initial_data()
 
     # 打印信息并处理数据集和版本名称
     if dataset_name is None and version_name is None:
@@ -980,11 +1098,11 @@ def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, l
 
     if isinstance(var, list):
         if len(var) == 1:
-            var = convert_full_name_to_short_name(var[0])
+            var = _convert_full_name_to_short_name(var[0])
         else:
-            var = [convert_full_name_to_short_name(v) for v in var]
+            var = [_convert_full_name_to_short_name(v) for v in var]
     elif isinstance(var, str):
-        var = convert_full_name_to_short_name(var)
+        var = _convert_full_name_to_short_name(var)
     else:
         raise ValueError("The var is invalid")
     if var is False:
@@ -1005,8 +1123,8 @@ def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, l
         os.makedirs(str(store_path), exist_ok=True)
 
     if num_workers is not None:
-        num_workers = max(min(num_workers, 10), 1)
-
+        num_workers = max(min(num_workers, 10), 1) # 暂时不限制最大值，再检查的时候可以多开一些线程
+        # num_workers = int(max(num_workers, 1))
     time_s = str(time_s)
     if len(time_s) == 8:
         time_s += "00"
@@ -1025,8 +1143,37 @@ def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, l
 
     global fsize_dict
     fsize_dict = {}
-
-    download_hourly_func(var, time_s, time_e, lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, num_workers, check, ftimes)
+    
+    global fsize_dict_lock
+    fsize_dict_lock = Lock()
+    
+    global use_idm, given_idm_engine, idm_download_list
+    if idm_engine is not None:
+        use_idm = True
+        given_idm_engine = idm_engine
+        idm_download_list = []
+    else:
+        use_idm = False
+    
+
+    _download_hourly_func(var, time_s, time_e, lon_min, lon_max, lat_min, lat_max, depth, level, store_path, dataset_name, version_name, num_workers, check, ftimes)
+    
+    if idm_download_list:
+        for f in idm_download_list:
+            wait_success = 0
+            success = False
+            while not success:
+                if check_nc(f):
+                    _correct_time(f)
+                    success = True
+                    count_dict["success"] += 1
+                else:
+                    wait_success += 1
+                    time.sleep(3)
+                    if wait_success >= 20:
+                        success = True
+                        # print(f'{f} download failed')
+                        count_dict["fail"] += 1
 
     count_dict["total"] = count_dict["success"] + count_dict["fail"] + count_dict["skip"] + count_dict["no_data"]
 
@@ -1094,9 +1241,9 @@ def how_to_use():
 
 
 if __name__ == "__main__":
-    time_s, time_e = "2024101012", "2024101018"
+    time_s, time_e = "2018010800", "2024083121"
     merge_name = f"{time_s}_{time_e}"  # 合并后的文件名
-    root_path = r"G:\Data\HYCOM\3hourly_test"
+    root_path = r"G:\Data\HYCOM\3hourly"
     location_dict = {"west": 105, "east": 130, "south": 15, "north": 45}
     download_dict = {
         "water_u": {"simple_name": "u", "download": 1},
@@ -1116,10 +1263,11 @@ if __name__ == "__main__":
     # if you wanna download all depth or level, set both False
     depth = None  # or 0-5000 meters
     level = None  # or 1-40 levels
-    num_workers = 3
+    num_workers = 1
 
     check = True
     ftimes = 1
+    idm_engine = r"D:\Programs\Internet Download Manager\IDMan.exe"
 
     download_switch, single_var = True, False
     combine_switch = False
@@ -1130,9 +1278,9 @@ if __name__ == "__main__":
     if download_switch:
         if single_var:
             for var_name in var_list:
-                download(var=var_name, time_s=time_s, time_e=time_e, store_path=Path(root_path), lon_min=location_dict["west"], lon_max=location_dict["east"], lat_min=location_dict["south"], lat_max=location_dict["north"], num_workers=num_workers, check=check, depth=depth, level=level, ftimes=ftimes)
+                download(var=var_name, time_s=time_s, time_e=time_e, store_path=Path(root_path), lon_min=location_dict["west"], lon_max=location_dict["east"], lat_min=location_dict["south"], lat_max=location_dict["north"], num_workers=num_workers, check=check, depth=depth, level=level, ftimes=ftimes, idm_engine=idm_engine)
         else:
-            download(var=var_list, time_s=time_s, time_e=time_e, store_path=Path(root_path), lon_min=location_dict["west"], lon_max=location_dict["east"], lat_min=location_dict["south"], lat_max=location_dict["north"], num_workers=num_workers, check=check, depth=depth, level=level, ftimes=ftimes)
+            download(var=var_list, time_s=time_s, time_e=time_e, store_path=Path(root_path), lon_min=location_dict["west"], lon_max=location_dict["east"], lat_min=location_dict["south"], lat_max=location_dict["north"], num_workers=num_workers, check=check, depth=depth, level=level, ftimes=ftimes, idm_engine=idm_engine)
 
     """ if combine_switch or copy_switch:
         time_list = get_time_list(time_s, time_e, 3, 'hour')