oafuncs 0.0.98.5__py3-none-any.whl → 0.0.98.7__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,1230 @@
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+ """
4
+ Author: Liu Kun && 16031215@qq.com
5
+ Date: 2025-04-17 15:16:02
6
+ LastEditors: Liu Kun && 16031215@qq.com
7
+ LastEditTime: 2025-04-17 15:16:04
8
+ FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\oa_down\\hycom_3hourly copy.py
9
+ Description:
10
+ EditPlatform: vscode
11
+ ComputerInfo: XPS 15 9510
12
+ SystemInfo: Windows 11
13
+ Python Version: 3.12
14
+ """
15
+
16
+
17
+
18
+ import asyncio
19
+ import datetime
20
+ import logging
21
+ import os
22
+ import random
23
+ import re
24
+ import time
25
+ import warnings
26
+ from concurrent.futures import ThreadPoolExecutor, as_completed
27
+ from pathlib import Path
28
+ from threading import Lock
29
+
30
+ import httpx
31
+ import matplotlib.pyplot as plt
32
+ import netCDF4 as nc
33
+ import numpy as np
34
+ import pandas as pd
35
+ import xarray as xr
36
+ from rich import print
37
+ from rich.progress import Progress
38
+
39
+ from oafuncs.oa_down.idm import downloader as idm_downloader
40
+ from oafuncs.oa_down.user_agent import get_ua
41
+ from oafuncs.oa_file import file_size
42
+ from oafuncs.oa_nc import check as check_nc
43
+ from oafuncs.oa_nc import modify as modify_nc
44
+
45
+ logging.getLogger("httpx").setLevel(logging.WARNING) # 关闭 httpx 的 INFO 日志,只显示 WARNING 及以上
46
+
47
+
48
+ warnings.filterwarnings("ignore", category=RuntimeWarning, message="Engine '.*' loading failed:.*")
49
+
50
+ __all__ = ["draw_time_range", "download"]
51
+
52
+
53
+ def _get_initial_data():
54
+ global variable_info, data_info, var_group, single_var_group
55
+ # ----------------------------------------------
56
+ # variable
57
+ variable_info = {
58
+ "u": {"var_name": "water_u", "standard_name": "eastward_sea_water_velocity"},
59
+ "v": {"var_name": "water_v", "standard_name": "northward_sea_water_velocity"},
60
+ "temp": {"var_name": "water_temp", "standard_name": "sea_water_potential_temperature"},
61
+ "salt": {"var_name": "salinity", "standard_name": "sea_water_salinity"},
62
+ "ssh": {"var_name": "surf_el", "standard_name": "sea_surface_elevation"},
63
+ "u_b": {"var_name": "water_u_bottom", "standard_name": "eastward_sea_water_velocity_at_sea_floor"},
64
+ "v_b": {"var_name": "water_v_bottom", "standard_name": "northward_sea_water_velocity_at_sea_floor"},
65
+ "temp_b": {"var_name": "water_temp_bottom", "standard_name": "sea_water_potential_temperature_at_sea_floor"},
66
+ "salt_b": {"var_name": "salinity_bottom", "standard_name": "sea_water_salinity_at_sea_floor"},
67
+ }
68
+ # ----------------------------------------------
69
+ # time resolution
70
+ data_info = {"yearly": {}, "monthly": {}, "daily": {}, "hourly": {}}
71
+
72
+ # hourly data
73
+ # dataset: GLBv0.08, GLBu0.08, GLBy0.08
74
+ data_info["hourly"]["dataset"] = {"GLBv0.08": {}, "GLBu0.08": {}, "GLBy0.08": {}, "ESPC_D": {}}
75
+
76
+ # version
77
+ # version of GLBv0.08: 53.X, 56.3, 57.2, 92.8, 57.7, 92.9, 93.0
78
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"] = {"53.X": {}, "56.3": {}, "57.2": {}, "92.8": {}, "57.7": {}, "92.9": {}, "93.0": {}}
79
+ # version of GLBu0.08: 93.0
80
+ data_info["hourly"]["dataset"]["GLBu0.08"]["version"] = {"93.0": {}}
81
+ # version of GLBy0.08: 93.0
82
+ data_info["hourly"]["dataset"]["GLBy0.08"]["version"] = {"93.0": {}}
83
+ # version of ESPC_D: V02
84
+ data_info["hourly"]["dataset"]["ESPC_D"]["version"] = {"V02": {}}
85
+
86
+ # info details
87
+ # time range
88
+ # GLBv0.08
89
+ # 在网页上提交超过范围的时间,会返回该数据集实际时间范围,从而纠正下面的时间范围
90
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["53.X"]["time_range"] = {"time_start": "1994010112", "time_end": "2015123109"}
91
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["56.3"]["time_range"] = {"time_start": "2014070112", "time_end": "2016093009"}
92
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["57.2"]["time_range"] = {"time_start": "2016050112", "time_end": "2017020109"}
93
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["92.8"]["time_range"] = {"time_start": "2017020112", "time_end": "2017060109"}
94
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["57.7"]["time_range"] = {"time_start": "2017060112", "time_end": "2017100109"}
95
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["92.9"]["time_range"] = {"time_start": "2017100112", "time_end": "2018032009"}
96
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["93.0"]["time_range"] = {"time_start": "2018010112", "time_end": "2020021909"}
97
+ # GLBu0.08
98
+ data_info["hourly"]["dataset"]["GLBu0.08"]["version"]["93.0"]["time_range"] = {"time_start": "2018091912", "time_end": "2018120909"}
99
+ # GLBy0.08
100
+ data_info["hourly"]["dataset"]["GLBy0.08"]["version"]["93.0"]["time_range"] = {"time_start": "2018120412", "time_end": "2024090509"}
101
+ # ESPC-D
102
+ data_info["hourly"]["dataset"]["ESPC_D"]["version"]["V02"]["time_range"] = {"time_start": "2024081012", "time_end": "2030010100"}
103
+
104
+ # classification method
105
+ # year_different: the data of different years is stored in different files
106
+ # same_path: the data of different years is stored in the same file
107
+ # var_different: the data of different variables is stored in different files
108
+ # var_year_different: the data of different variables and years is stored in different files
109
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["53.X"]["classification"] = "year_different"
110
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["56.3"]["classification"] = "same_path"
111
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["57.2"]["classification"] = "same_path"
112
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["92.8"]["classification"] = "var_different"
113
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["57.7"]["classification"] = "same_path"
114
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["92.9"]["classification"] = "var_different"
115
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["93.0"]["classification"] = "var_different"
116
+ data_info["hourly"]["dataset"]["GLBu0.08"]["version"]["93.0"]["classification"] = "var_different"
117
+ data_info["hourly"]["dataset"]["GLBy0.08"]["version"]["93.0"]["classification"] = "var_year_different"
118
+ data_info["hourly"]["dataset"]["ESPC_D"]["version"]["V02"]["classification"] = "single_var_year_different"
119
+
120
+ # download info
121
+ # base url
122
+ # GLBv0.08 53.X
123
+ url_53x = {}
124
+ for y_53x in range(1994, 2016):
125
+ # r'https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_53.X/data/2013?'
126
+ url_53x[str(y_53x)] = rf"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_53.X/data/{y_53x}?"
127
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["53.X"]["url"] = url_53x
128
+ # GLBv0.08 56.3
129
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["56.3"]["url"] = r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_56.3?"
130
+ # GLBv0.08 57.2
131
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["57.2"]["url"] = r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_57.2?"
132
+ # GLBv0.08 92.8
133
+ url_928 = {
134
+ "uv3z": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_92.8/uv3z?",
135
+ "ts3z": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_92.8/ts3z?",
136
+ "ssh": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_92.8/ssh?",
137
+ }
138
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["92.8"]["url"] = url_928
139
+ # GLBv0.08 57.7
140
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["57.7"]["url"] = r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_57.7?"
141
+ # GLBv0.08 92.9
142
+ url_929 = {
143
+ "uv3z": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_92.9/uv3z?",
144
+ "ts3z": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_92.9/ts3z?",
145
+ "ssh": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_92.9/ssh?",
146
+ }
147
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["92.9"]["url"] = url_929
148
+ # GLBv0.08 93.0
149
+ url_930_v = {
150
+ "uv3z": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_93.0/uv3z?",
151
+ "ts3z": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_93.0/ts3z?",
152
+ "ssh": r"https://ncss.hycom.org/thredds/ncss/GLBv0.08/expt_93.0/ssh?",
153
+ }
154
+ data_info["hourly"]["dataset"]["GLBv0.08"]["version"]["93.0"]["url"] = url_930_v
155
+ # GLBu0.08 93.0
156
+ url_930_u = {
157
+ "uv3z": r"https://ncss.hycom.org/thredds/ncss/GLBu0.08/expt_93.0/uv3z?",
158
+ "ts3z": r"https://ncss.hycom.org/thredds/ncss/GLBu0.08/expt_93.0/ts3z?",
159
+ "ssh": r"https://ncss.hycom.org/thredds/ncss/GLBu0.08/expt_93.0/ssh?",
160
+ }
161
+ data_info["hourly"]["dataset"]["GLBu0.08"]["version"]["93.0"]["url"] = url_930_u
162
+ # GLBy0.08 93.0
163
+ uv3z_930_y = {}
164
+ ts3z_930_y = {}
165
+ ssh_930_y = {}
166
+ for y_930_y in range(2018, 2025):
167
+ uv3z_930_y[str(y_930_y)] = rf"https://ncss.hycom.org/thredds/ncss/GLBy0.08/expt_93.0/uv3z/{y_930_y}?"
168
+ ts3z_930_y[str(y_930_y)] = rf"https://ncss.hycom.org/thredds/ncss/GLBy0.08/expt_93.0/ts3z/{y_930_y}?"
169
+ ssh_930_y[str(y_930_y)] = rf"https://ncss.hycom.org/thredds/ncss/GLBy0.08/expt_93.0/ssh/{y_930_y}?"
170
+ # GLBy0.08 93.0 data time range in each year: year-01-01 12:00 to year+1-01-01 09:00
171
+ url_930_y = {
172
+ "uv3z": uv3z_930_y,
173
+ "ts3z": ts3z_930_y,
174
+ "ssh": ssh_930_y,
175
+ }
176
+ data_info["hourly"]["dataset"]["GLBy0.08"]["version"]["93.0"]["url"] = url_930_y
177
+ # ESPC-D-V02
178
+ u3z_espc_d_v02_y = {}
179
+ v3z_espc_d_v02_y = {}
180
+ t3z_espc_d_v02_y = {}
181
+ s3z_espc_d_v02_y = {}
182
+ ssh_espc_d_v02_y = {}
183
+ for y_espc_d_v02 in range(2024, 2030):
184
+ u3z_espc_d_v02_y[str(y_espc_d_v02)] = rf"https://ncss.hycom.org/thredds/ncss/ESPC-D-V02/u3z/{y_espc_d_v02}?"
185
+ v3z_espc_d_v02_y[str(y_espc_d_v02)] = rf"https://ncss.hycom.org/thredds/ncss/ESPC-D-V02/v3z/{y_espc_d_v02}?"
186
+ t3z_espc_d_v02_y[str(y_espc_d_v02)] = rf"https://ncss.hycom.org/thredds/ncss/ESPC-D-V02/t3z/{y_espc_d_v02}?"
187
+ s3z_espc_d_v02_y[str(y_espc_d_v02)] = rf"https://ncss.hycom.org/thredds/ncss/ESPC-D-V02/s3z/{y_espc_d_v02}?"
188
+ ssh_espc_d_v02_y[str(y_espc_d_v02)] = rf"https://ncss.hycom.org/thredds/ncss/ESPC-D-V02/ssh/{y_espc_d_v02}?"
189
+ url_espc_d_v02_y = {
190
+ "u3z": u3z_espc_d_v02_y,
191
+ "v3z": v3z_espc_d_v02_y,
192
+ "t3z": t3z_espc_d_v02_y,
193
+ "s3z": s3z_espc_d_v02_y,
194
+ "ssh": ssh_espc_d_v02_y,
195
+ }
196
+ data_info["hourly"]["dataset"]["ESPC_D"]["version"]["V02"]["url"] = url_espc_d_v02_y
197
+ # ----------------------------------------------
198
+ var_group = {
199
+ "uv3z": ["u", "v", "u_b", "v_b"],
200
+ "ts3z": ["temp", "salt", "temp_b", "salt_b"],
201
+ "ssh": ["ssh"],
202
+ }
203
+ # ----------------------------------------------
204
+ single_var_group = {
205
+ "u3z": ["u"],
206
+ "v3z": ["v"],
207
+ "t3z": ["temp"],
208
+ "s3z": ["salt"],
209
+ "ssh": ["ssh"],
210
+ }
211
+
212
+ return variable_info, data_info, var_group, single_var_group
213
+
214
+
215
+ def draw_time_range(pic_save_folder=None):
216
+ if pic_save_folder is not None:
217
+ os.makedirs(pic_save_folder, exist_ok=True)
218
+ # Converting the data into a format suitable for plotting
219
+ data = []
220
+ for dataset, versions in data_info["hourly"]["dataset"].items():
221
+ for version, time_range in versions["version"].items():
222
+ t_s = time_range["time_range"]["time_start"]
223
+ t_e = time_range["time_range"]["time_end"]
224
+ if len(t_s) == 8:
225
+ t_s = t_s + "00"
226
+ if len(t_e) == 8:
227
+ t_e = t_e + "21"
228
+ t_s, t_e = t_s + "0000", t_e + "0000"
229
+ data.append(
230
+ {
231
+ "dataset": dataset,
232
+ "version": version,
233
+ "start_date": pd.to_datetime(t_s),
234
+ "end_date": pd.to_datetime(t_e),
235
+ }
236
+ )
237
+
238
+ # Creating a DataFrame
239
+ df = pd.DataFrame(data)
240
+
241
+ # Plotting with combined labels for datasets and versions on the y-axis
242
+ plt.figure(figsize=(12, 6))
243
+
244
+ # Combined labels for datasets and versions
245
+ combined_labels = [f"{dataset}_{version}" for dataset, version in zip(df["dataset"], df["version"])]
246
+
247
+ colors = plt.cm.viridis(np.linspace(0, 1, len(combined_labels)))
248
+
249
+ # Assigning a color to each combined label
250
+ label_colors = {label: colors[i] for i, label in enumerate(combined_labels)}
251
+
252
+ # Plotting each time range
253
+ k = 1
254
+ for _, row in df.iterrows():
255
+ plt.plot([row["start_date"], row["end_date"]], [k, k], color=label_colors[f"{row['dataset']}_{row['version']}"], linewidth=6)
256
+ # plt.text(row['end_date'], k,
257
+ # f"{row['version']}", ha='right', color='black')
258
+ ymdh_s = row["start_date"].strftime("%Y-%m-%d %H")
259
+ ymdh_e = row["end_date"].strftime("%Y-%m-%d %H")
260
+ # if k == 1 or k == len(combined_labels):
261
+ if k == 1:
262
+ plt.text(row["start_date"], k + 0.125, f"{ymdh_s}", ha="left", color="black")
263
+ plt.text(row["end_date"], k + 0.125, f"{ymdh_e}", ha="right", color="black")
264
+ else:
265
+ plt.text(row["start_date"], k + 0.125, f"{ymdh_s}", ha="right", color="black")
266
+ plt.text(row["end_date"], k + 0.125, f"{ymdh_e}", ha="left", color="black")
267
+ k += 1
268
+
269
+ # Setting the y-axis labels
270
+ plt.yticks(range(1, len(combined_labels) + 1), combined_labels)
271
+ plt.xlabel("Time")
272
+ plt.ylabel("Dataset - Version")
273
+ plt.title("Time Range of Different Versions of Datasets")
274
+ plt.xticks(rotation=45)
275
+ plt.grid(True)
276
+ plt.tight_layout()
277
+ if pic_save_folder:
278
+ plt.savefig(Path(pic_save_folder) / "HYCOM_time_range.png")
279
+ print(f"[bold green]HYCOM_time_range.png has been saved in {pic_save_folder}")
280
+ else:
281
+ plt.savefig("HYCOM_time_range.png")
282
+ print("[bold green]HYCOM_time_range.png has been saved in the current folder")
283
+ print(f"Curren folder: {os.getcwd()}")
284
+ # plt.show()
285
+ plt.close()
286
+
287
+
288
+ def _get_time_list(time_s, time_e, delta, interval_type="hour"):
289
+ """
290
+ Description: get a list of time strings from time_s to time_e with a specified interval
291
+ Args:
292
+ time_s: start time string, e.g. '2023080203' for hours or '20230802' for days
293
+ time_e: end time string, e.g. '2023080303' for hours or '20230803' for days
294
+ delta: interval of hours or days
295
+ interval_type: 'hour' for hour interval, 'day' for day interval
296
+ Returns:
297
+ dt_list: a list of time strings
298
+ """
299
+ time_s, time_e = str(time_s), str(time_e)
300
+ if interval_type == "hour":
301
+ time_format = "%Y%m%d%H"
302
+ delta_type = "hours"
303
+ elif interval_type == "day":
304
+ time_format = "%Y%m%d"
305
+ delta_type = "days"
306
+ # Ensure time strings are in the correct format for days
307
+ time_s = time_s[:8]
308
+ time_e = time_e[:8]
309
+ else:
310
+ raise ValueError("interval_type must be 'hour' or 'day'")
311
+
312
+ dt = datetime.datetime.strptime(time_s, time_format)
313
+ dt_list = []
314
+ while dt.strftime(time_format) <= time_e:
315
+ dt_list.append(dt.strftime(time_format))
316
+ dt += datetime.timedelta(**{delta_type: delta})
317
+ return dt_list
318
+
319
+
320
+ def _transform_time(time_str):
321
+ # old_time = '2023080203'
322
+ # time_new = '2023-08-02T03%3A00%3A00Z'
323
+ time_new = f"{time_str[:4]}-{time_str[4:6]}-{time_str[6:8]}T{time_str[8:10]}%3A00%3A00Z"
324
+ return time_new
325
+
326
+
327
+ 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):
328
+ query_dict = {
329
+ "var": variable_info[var]["var_name"],
330
+ "north": lat_max,
331
+ "west": lon_min,
332
+ "east": lon_max,
333
+ "south": lat_min,
334
+ "horizStride": 1,
335
+ "time": None,
336
+ "time_start": None,
337
+ "time_end": None,
338
+ "timeStride": None,
339
+ "vertCoord": None,
340
+ "vertStride": None,
341
+ "addLatLon": "true",
342
+ "accept": "netcdf4",
343
+ }
344
+
345
+ if time_str_end is not None:
346
+ query_dict["time_start"] = _transform_time(time_str_ymdh)
347
+ query_dict["time_end"] = _transform_time(time_str_end)
348
+ query_dict["timeStride"] = 1
349
+ else:
350
+ query_dict["time"] = _transform_time(time_str_ymdh)
351
+
352
+ def get_nearest_level_index(depth):
353
+ 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]
354
+ return min(range(len(level_depth)), key=lambda i: abs(level_depth[i] - depth))
355
+
356
+ if var not in ["ssh", "u_b", "v_b", "temp_b", "salt_b"] and var in ["u", "v", "temp", "salt"]:
357
+ if mode == "depth":
358
+ if depth < 0 or depth > 5000:
359
+ print("Please ensure the depth is in the range of 0-5000 m")
360
+ query_dict["vertCoord"] = get_nearest_level_index(depth) + 1
361
+ elif mode == "level":
362
+ if level_num < 1 or level_num > 40:
363
+ print("Please ensure the level_num is in the range of 1-40")
364
+ query_dict["vertCoord"] = max(1, min(level_num, 40))
365
+ elif mode == "full":
366
+ query_dict["vertStride"] = 1
367
+ else:
368
+ raise ValueError("Invalid mode. Choose from 'depth', 'level', or 'full'")
369
+
370
+ query_dict = {k: v for k, v in query_dict.items() if v is not None}
371
+
372
+ return query_dict
373
+
374
+
375
+ def _check_time_in_dataset_and_version(time_input, time_end=None):
376
+ # 判断是处理单个时间点还是时间范围
377
+ is_single_time = time_end is None
378
+
379
+ # 如果是单个时间点,初始化时间范围
380
+ if is_single_time:
381
+ time_start = int(time_input)
382
+ time_end = time_start
383
+ time_input_str = str(time_input)
384
+ else:
385
+ time_start = int(time_input)
386
+ time_end = int(time_end)
387
+ time_input_str = f"{time_input}-{time_end}"
388
+
389
+ # 根据时间长度补全时间格式
390
+ if len(str(time_start)) == 8:
391
+ time_start = str(time_start) + "00"
392
+ if len(str(time_end)) == 8:
393
+ time_end = str(time_end) + "21"
394
+ time_start, time_end = int(time_start), int(time_end)
395
+
396
+ d_list = []
397
+ v_list = []
398
+ trange_list = []
399
+ have_data = False
400
+
401
+ # 遍历数据集和版本
402
+ for dataset_name in data_info["hourly"]["dataset"].keys():
403
+ for version_name in data_info["hourly"]["dataset"][dataset_name]["version"].keys():
404
+ time_s, time_e = list(data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["time_range"].values())
405
+ time_s, time_e = str(time_s), str(time_e)
406
+ if len(time_s) == 8:
407
+ time_s = time_s + "00"
408
+ if len(time_e) == 8:
409
+ time_e = time_e + "21"
410
+ # 检查时间是否在数据集的时间范围内
411
+ if is_single_time:
412
+ if time_start >= int(time_s) and time_start <= int(time_e):
413
+ d_list.append(dataset_name)
414
+ v_list.append(version_name)
415
+ trange_list.append(f"{time_s}-{time_e}")
416
+ have_data = True
417
+ else:
418
+ if time_start >= int(time_s) and time_end <= int(time_e):
419
+ d_list.append(dataset_name)
420
+ v_list.append(version_name)
421
+ trange_list.append(f"{time_s}-{time_e}")
422
+ have_data = True
423
+
424
+ if have_data:
425
+ if match_time is None:
426
+ print(f"[bold red]Time {time_input_str} included in:")
427
+ dv_num = 1
428
+ for d, v, trange in zip(d_list, v_list, trange_list):
429
+ print(f"{dv_num} -> [bold blue]{d} - {v} : {trange}")
430
+ dv_num += 1
431
+ if is_single_time:
432
+ return True
433
+ else:
434
+ base_url_s = _get_base_url(d_list[0], v_list[0], "u", str(time_start))
435
+ base_url_e = _get_base_url(d_list[0], v_list[0], "u", str(time_end))
436
+ if base_url_s == base_url_e:
437
+ return True
438
+ else:
439
+ print(f"[bold red]{time_start} to {time_end} is in different datasets or versions, so you can't download them together")
440
+ return False
441
+ else:
442
+ print(f"[bold red]Time {time_input_str} has no data")
443
+ return False
444
+
445
+
446
+ def _ensure_time_in_specific_dataset_and_version(dataset_name, version_name, time_input, time_end=None):
447
+ # 根据时间长度补全时间格式
448
+ if len(str(time_input)) == 8:
449
+ time_input = str(time_input) + "00"
450
+ time_start = int(time_input)
451
+ if time_end is not None:
452
+ if len(str(time_end)) == 8:
453
+ time_end = str(time_end) + "21"
454
+ time_end = int(time_end)
455
+ else:
456
+ time_end = time_start
457
+
458
+ # 检查指定的数据集和版本是否存在
459
+ if dataset_name not in data_info["hourly"]["dataset"]:
460
+ print(f"[bold red]Dataset {dataset_name} not found.")
461
+ return False
462
+ if version_name not in data_info["hourly"]["dataset"][dataset_name]["version"]:
463
+ print(f"[bold red]Version {version_name} not found in dataset {dataset_name}.")
464
+ return False
465
+
466
+ # 获取指定数据集和版本的时间范围
467
+ time_range = data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["time_range"]
468
+ time_s, time_e = list(time_range.values())
469
+ time_s, time_e = str(time_s), str(time_e)
470
+ if len(time_s) == 8:
471
+ time_s = time_s + "00"
472
+ if len(time_e) == 8:
473
+ time_e = time_e + "21"
474
+ time_s, time_e = int(time_s), int(time_e)
475
+
476
+ # 检查时间是否在指定数据集和版本的时间范围内
477
+ if time_start >= time_s and time_end <= time_e:
478
+ print(f"[bold blue]Time {time_input} to {time_end} is within dataset {dataset_name} and version {version_name}.")
479
+ return True
480
+ else:
481
+ print(f"[bold red]Time {time_input} to {time_end} is not within dataset {dataset_name} and version {version_name}.")
482
+ return False
483
+
484
+
485
+ def _direct_choose_dataset_and_version(time_input, time_end=None):
486
+ # 假设 data_info 是一个字典,包含了数据集和版本的信息
487
+ # 示例结构:data_info['hourly']['dataset'][dataset_name]['version'][version_name]['time_range']
488
+
489
+ if len(str(time_input)) == 8:
490
+ time_input = str(time_input) + "00"
491
+
492
+ # 如果 time_end 是 None,则将 time_input 的值赋给它
493
+ if time_end is None:
494
+ time_end = time_input
495
+
496
+ # 处理开始和结束时间,确保它们是完整的 ymdh 格式
497
+ time_start, time_end = int(str(time_input)[:10]), int(str(time_end)[:10])
498
+
499
+ dataset_name_out, version_name_out = None, None
500
+
501
+ for dataset_name in data_info["hourly"]["dataset"].keys():
502
+ for version_name in data_info["hourly"]["dataset"][dataset_name]["version"].keys():
503
+ [time_s, time_e] = list(data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["time_range"].values())
504
+ time_s, time_e = str(time_s), str(time_e)
505
+ if len(time_s) == 8:
506
+ time_s = time_s + "00"
507
+ if len(time_e) == 8:
508
+ time_e = time_e + "21"
509
+ time_s, time_e = int(time_s), int(time_e)
510
+
511
+ # 检查时间是否在数据集版本的时间范围内
512
+ if time_start >= time_s and time_end <= time_e:
513
+ dataset_name_out, version_name_out = dataset_name, version_name
514
+
515
+ if dataset_name_out is not None and version_name_out is not None:
516
+ if match_time is None:
517
+ # print(f"[bold purple]dataset: {dataset_name_out}, version: {version_name_out} is chosen")
518
+ print(f"[bold purple]Chosen dataset: {dataset_name_out} - {version_name_out}")
519
+
520
+ # 如果没有找到匹配的数据集和版本,会返回 None
521
+ return dataset_name_out, version_name_out
522
+
523
+
524
+ def _get_base_url(dataset_name, version_name, var, ymdh_str):
525
+ year_str = int(ymdh_str[:4])
526
+ url_dict = data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["url"]
527
+ classification_method = data_info["hourly"]["dataset"][dataset_name]["version"][version_name]["classification"]
528
+ if classification_method == "year_different":
529
+ base_url = url_dict[str(year_str)]
530
+ elif classification_method == "same_path":
531
+ base_url = url_dict
532
+ elif classification_method == "var_different":
533
+ base_url = None
534
+ for key, value in var_group.items():
535
+ if var in value:
536
+ base_url = url_dict[key]
537
+ break
538
+ if base_url is None:
539
+ print("Please ensure the var is in [u,v,temp,salt,ssh,u_b,v_b,temp_b,salt_b]")
540
+ elif classification_method == "var_year_different":
541
+ if dataset_name == "GLBy0.08" and version_name == "93.0":
542
+ mdh_str = ymdh_str[4:]
543
+ # GLBy0.08 93.0
544
+ # data time range in each year: year-01-01 12:00 to year+1-01-01 09:00
545
+ if "010100" <= mdh_str <= "010109":
546
+ year_str = int(ymdh_str[:4]) - 1
547
+ else:
548
+ year_str = int(ymdh_str[:4])
549
+ base_url = None
550
+ for key, value in var_group.items():
551
+ if var in value:
552
+ base_url = url_dict[key][str(year_str)]
553
+ break
554
+ if base_url is None:
555
+ print("Please ensure the var is in [u,v,temp,salt,ssh,u_b,v_b,temp_b,salt_b]")
556
+ elif classification_method == "single_var_year_different":
557
+ base_url = None
558
+ if dataset_name == "ESPC_D" and version_name == "V02":
559
+ mdh_str = ymdh_str[4:]
560
+ # ESPC-D-V02
561
+ if "010100" <= mdh_str <= "010109":
562
+ year_str = int(ymdh_str[:4]) - 1
563
+ else:
564
+ year_str = int(ymdh_str[:4])
565
+ for key, value in single_var_group.items():
566
+ if var in value:
567
+ base_url = url_dict[key][str(year_str)]
568
+ break
569
+ if base_url is None:
570
+ print("Please ensure the var is in [u,v,temp,salt,ssh]")
571
+ return base_url
572
+
573
+
574
+ def _get_submit_url(dataset_name, version_name, var, ymdh_str, query_dict):
575
+ base_url = _get_base_url(dataset_name, version_name, var, ymdh_str)
576
+ if isinstance(query_dict["var"], str):
577
+ query_dict["var"] = [query_dict["var"]]
578
+ 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")
579
+ return target_url
580
+
581
+
582
+ def _clear_existing_file(file_full_path):
583
+ if os.path.exists(file_full_path):
584
+ os.remove(file_full_path)
585
+ print(f"{file_full_path} has been removed")
586
+
587
+
588
+ def _check_existing_file(file_full_path, avg_size):
589
+ if os.path.exists(file_full_path):
590
+ print(f"[bold #FFA54F]{file_full_path} exists")
591
+ fsize = file_size(file_full_path)
592
+ delta_size_ratio = (fsize - avg_size) / avg_size
593
+ if abs(delta_size_ratio) > 0.025:
594
+ if check_nc(file_full_path):
595
+ return True
596
+ else:
597
+ # print(f"File size is abnormal and cannot be opened, {file_full_path}: {fsize:.2f} KB")
598
+ return False
599
+ else:
600
+ return True
601
+ else:
602
+ return False
603
+
604
+
605
+ def _get_mean_size_move(same_file, current_file):
606
+ with fsize_dict_lock:
607
+ if same_file not in fsize_dict.keys():
608
+ fsize_dict[same_file] = {"size_list": [], "mean_size": 1.0}
609
+
610
+ tolerance_ratio = 0.025
611
+ current_file_size = file_size(current_file)
612
+
613
+ if fsize_dict[same_file]["size_list"]:
614
+ fsize_dict[same_file]["mean_size"] = sum(fsize_dict[same_file]["size_list"]) / len(fsize_dict[same_file]["size_list"])
615
+ fsize_dict[same_file]["mean_size"] = max(fsize_dict[same_file]["mean_size"], 1.0)
616
+ else:
617
+ fsize_dict[same_file]["mean_size"] = 1.0
618
+
619
+ size_difference_ratio = (current_file_size - fsize_dict[same_file]["mean_size"]) / fsize_dict[same_file]["mean_size"]
620
+
621
+ if abs(size_difference_ratio) > tolerance_ratio:
622
+ if check_nc(current_file, print_messages=False):
623
+ fsize_dict[same_file]["size_list"] = [current_file_size]
624
+ fsize_dict[same_file]["mean_size"] = current_file_size
625
+ else:
626
+ _clear_existing_file(current_file)
627
+ # print(f"File size is abnormal, may need to be downloaded again, file size: {current_file_size:.2f} KB")
628
+ else:
629
+ fsize_dict[same_file]["size_list"].append(current_file_size)
630
+
631
+ return fsize_dict[same_file]["mean_size"]
632
+
633
+
634
+ def _check_ftime(nc_file, tname="time", if_print=False):
635
+ if not os.path.exists(nc_file):
636
+ return False
637
+ nc_file = str(nc_file)
638
+ try:
639
+ ds = xr.open_dataset(nc_file)
640
+ real_time = ds[tname].values[0]
641
+ ds.close()
642
+ real_time = str(real_time)[:13]
643
+ real_time = real_time.replace("-", "").replace("T", "")
644
+ f_time = re.findall(r"\d{10}", nc_file)[0]
645
+ if real_time == f_time:
646
+ return True
647
+ else:
648
+ if if_print:
649
+ print(f"[bold #daff5c]File time error, file/real time: [bold blue]{f_time}/{real_time}")
650
+ return False
651
+ except Exception as e:
652
+ if if_print:
653
+ print(f"[bold #daff5c]File time check failed, {nc_file}: {e}")
654
+ return False
655
+
656
+
657
+ def _correct_time(nc_file):
658
+ dataset = nc.Dataset(nc_file)
659
+ time_units = dataset.variables["time"].units
660
+ dataset.close()
661
+ origin_str = time_units.split("since")[1].strip()
662
+ origin_datetime = datetime.datetime.strptime(origin_str, "%Y-%m-%d %H:%M:%S")
663
+ given_date_str = re.findall(r"\d{10}", str(nc_file))[0]
664
+ given_datetime = datetime.datetime.strptime(given_date_str, "%Y%m%d%H")
665
+ time_difference = (given_datetime - origin_datetime).total_seconds()
666
+ if "hours" in time_units:
667
+ time_difference /= 3600
668
+ elif "days" in time_units:
669
+ time_difference /= 3600 * 24
670
+ modify_nc(nc_file, "time", None, time_difference)
671
+
672
+
673
+ class _HycomDownloader:
674
+ def __init__(self, tasks, delay_range=(3, 6), timeout_factor=120, max_var_count=5, max_retries=3):
675
+ """
676
+ :param tasks: List of (url, save_path)
677
+ """
678
+ self.tasks = tasks
679
+ self.delay_range = delay_range
680
+ self.timeout_factor = timeout_factor
681
+ self.max_var_count = max_var_count
682
+ self.max_retries = max_retries
683
+ self.count = {"success": 0, "fail": 0}
684
+
685
+ def user_agent(self):
686
+ return get_ua()
687
+
688
+ async def _download_one(self, url, save_path):
689
+ file_name = os.path.basename(save_path)
690
+ headers = {"User-Agent": self.user_agent()}
691
+ var_count = min(max(url.count("var="), 1), self.max_var_count)
692
+ timeout_max = self.timeout_factor * var_count
693
+ timeout = random.randint(timeout_max // 2, timeout_max)
694
+
695
+ retry = 0
696
+ while retry <= self.max_retries:
697
+ if proxy_txt_path:
698
+ from .read_proxy import get_valid_proxy
699
+
700
+ proxy_one = get_valid_proxy(proxy_txt_path)
701
+ if proxy_one:
702
+ if proxy_one.startswith("http://"):
703
+ proxy_one = proxy_one[7:]
704
+ elif proxy_one.startswith("https://"):
705
+ proxy_one = proxy_one[8:]
706
+ else:
707
+ proxy_one = None
708
+ if proxy_one:
709
+ proxy = f"http://{proxy_one}"
710
+ mounts = {
711
+ "http://": httpx.AsyncHTTPTransport(proxy=proxy),
712
+ "https://": httpx.AsyncHTTPTransport(proxy=proxy),
713
+ }
714
+ else:
715
+ proxy = None
716
+ mounts = None
717
+ try:
718
+ await asyncio.sleep(random.uniform(*self.delay_range))
719
+ start = datetime.datetime.now()
720
+
721
+ async with httpx.AsyncClient(
722
+ timeout=httpx.Timeout(timeout),
723
+ limits=httpx.Limits(max_connections=2, max_keepalive_connections=2),
724
+ transport=httpx.AsyncHTTPTransport(retries=2),
725
+ # proxy=proxy,
726
+ mounts=mounts,
727
+ ) as client:
728
+ print(f"[bold #f0f6d0]Requesting {file_name} (Attempt {retry + 1}) ...")
729
+ response = await client.get(url, headers=headers, follow_redirects=True)
730
+ response.raise_for_status()
731
+ if not response.content:
732
+ raise ValueError("Empty response received")
733
+
734
+ print(f"[bold #96cbd7]Downloading {file_name} ...")
735
+ with open(save_path, "wb") as f:
736
+ async for chunk in response.aiter_bytes(32 * 1024):
737
+ f.write(chunk)
738
+
739
+ elapsed = datetime.datetime.now() - start
740
+ print(f"[#3dfc40]File [bold #dfff73]{file_name} [#3dfc40]downloaded, Time: [#39cbdd]{elapsed}")
741
+ self.count["success"] += 1
742
+ count_dict["success"] += 1
743
+ return
744
+
745
+ except Exception as e:
746
+ print(f"[bold red]Failed ({type(e).__name__}): {e}")
747
+ if retry < self.max_retries:
748
+ backoff = 2**retry
749
+ print(f"[yellow]Retrying in {backoff:.1f}s ...")
750
+ await asyncio.sleep(backoff)
751
+ retry += 1
752
+ else:
753
+ print(f"[red]Giving up on {file_name}")
754
+ self.count["fail"] += 1
755
+ count_dict["fail"] += 1
756
+ return
757
+
758
+ async def run(self):
759
+ print(f"📥 Starting download of {len(self.tasks)} files ...")
760
+ for url, save_path in self.tasks:
761
+ await self._download_one(url, save_path)
762
+
763
+ print("\n✅ All tasks completed.")
764
+ print(f"✔️ Success: {self.count['success']} | ❌ Fail: {self.count['fail']}")
765
+
766
+
767
+ def _download_file(target_url, store_path, file_name, cover=False):
768
+ save_path = Path(store_path) / file_name
769
+ file_name_split = file_name.split("_")
770
+ file_name_split = file_name_split[:-1]
771
+ same_file = "_".join(file_name_split) + "*nc"
772
+
773
+ if match_time is not None:
774
+ if check_nc(save_path, print_messages=False):
775
+ if not _check_ftime(save_path, if_print=True):
776
+ if match_time:
777
+ _correct_time(save_path)
778
+ count_dict["skip"] += 1
779
+ else:
780
+ _clear_existing_file(save_path)
781
+ count_dict["no_data"] += 1
782
+ else:
783
+ count_dict["skip"] += 1
784
+ print(f"[bold green]{file_name} is correct")
785
+ return
786
+
787
+ if not cover and os.path.exists(save_path):
788
+ print(f"[bold #FFA54F]{save_path} exists, skipping ...")
789
+ count_dict["skip"] += 1
790
+ return
791
+
792
+ if same_file not in fsize_dict.keys():
793
+ check_nc(save_path, delete_if_invalid=True, print_messages=False)
794
+
795
+ get_mean_size = _get_mean_size_move(same_file, save_path)
796
+
797
+ if _check_existing_file(save_path, get_mean_size):
798
+ count_dict["skip"] += 1
799
+ return
800
+
801
+ _clear_existing_file(save_path)
802
+
803
+ if not use_idm:
804
+ python_downloader = _HycomDownloader([(target_url, save_path)])
805
+ asyncio.run(python_downloader.run())
806
+ time.sleep(3 + random.uniform(0, 10))
807
+ else:
808
+ idm_downloader(target_url, store_path, file_name, given_idm_engine)
809
+ idm_download_list.append(save_path)
810
+ # print(f"[bold #3dfc40]File [bold #dfff73]{save_path} [#3dfc40]has been submit to IDM for downloading")
811
+ time.sleep(3 + random.uniform(0, 10))
812
+
813
+
814
+ def _check_hour_is_valid(ymdh_str):
815
+ hh = int(str(ymdh_str[-2:]))
816
+ if hh in [0, 3, 6, 9, 12, 15, 18, 21]:
817
+ return True
818
+ else:
819
+ return False
820
+
821
+
822
+ def _check_dataset_version(dataset_name, version_name, download_time, download_time_end=None):
823
+ if dataset_name is not None and version_name is not None:
824
+ just_ensure = _ensure_time_in_specific_dataset_and_version(dataset_name, version_name, download_time, download_time_end)
825
+ if just_ensure:
826
+ return dataset_name, version_name
827
+ else:
828
+ return None, None
829
+
830
+ download_time_str = str(download_time)
831
+
832
+ if len(download_time_str) == 8:
833
+ download_time_str = download_time_str + "00"
834
+
835
+ if download_time_end is None and not _check_hour_is_valid(download_time_str):
836
+ print("Please ensure the hour is 00, 03, 06, 09, 12, 15, 18, 21")
837
+ raise ValueError("The hour is invalid")
838
+
839
+ if download_time_end is not None:
840
+ if len(str(download_time_end)) == 8:
841
+ download_time_end = str(download_time_end) + "21"
842
+ have_data = _check_time_in_dataset_and_version(download_time_str, download_time_end)
843
+ if have_data:
844
+ return _direct_choose_dataset_and_version(download_time_str, download_time_end)
845
+ else:
846
+ have_data = _check_time_in_dataset_and_version(download_time_str)
847
+ if have_data:
848
+ return _direct_choose_dataset_and_version(download_time_str)
849
+
850
+ return None, None
851
+
852
+
853
+ 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):
854
+ ymdh_str = str(download_time)
855
+ if depth is not None and level_num is not None:
856
+ print("Please ensure the depth or level_num is None")
857
+ print("Progress will use the depth")
858
+ which_mode = "depth"
859
+ elif depth is not None and level_num is None:
860
+ print(f"Data of single depth (~{depth} m) will be downloaded...")
861
+ which_mode = "depth"
862
+ elif level_num is not None and depth is None:
863
+ print(f"Data of single level ({level_num}) will be downloaded...")
864
+ which_mode = "level"
865
+ else:
866
+ which_mode = "full"
867
+ query_dict = _get_query_dict(var, lon_min, lon_max, lat_min, lat_max, download_time, download_time_end, which_mode, depth, level_num)
868
+ submit_url = _get_submit_url(dataset_name, version_name, var, ymdh_str, query_dict)
869
+ return submit_url
870
+
871
+
872
+ 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, cover=False):
873
+ print("[bold #ecdbfe]-" * mark_len)
874
+ download_time = str(download_time)
875
+ if download_time_end is not None:
876
+ download_time_end = str(download_time_end)
877
+ dataset_name, version_name = _check_dataset_version(dataset_name, version_name, download_time, download_time_end)
878
+ else:
879
+ dataset_name, version_name = _check_dataset_version(dataset_name, version_name, download_time)
880
+ if dataset_name is None and version_name is None:
881
+ count_dict["no_data"] += 1
882
+ if download_time_end is not None:
883
+ count_dict["no_data_list"].append(f"{download_time}-{download_time_end}")
884
+ else:
885
+ count_dict["no_data_list"].append(download_time)
886
+ return
887
+
888
+ if isinstance(var, str):
889
+ var = [var]
890
+
891
+ if isinstance(var, list):
892
+ if len(var) == 1:
893
+ var = var[0]
894
+ 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)
895
+ file_name = f"HYCOM_{variable_info[var]['var_name']}_{download_time}.nc"
896
+ if download_time_end is not None:
897
+ file_name = f"HYCOM_{variable_info[var]['var_name']}_{download_time}-{download_time_end}.nc"
898
+ _download_file(submit_url, store_path, file_name, cover)
899
+ else:
900
+ if download_time < "2024081012":
901
+ varlist = [_ for _ in var]
902
+ for key, value in var_group.items():
903
+ current_group = []
904
+ for v in varlist:
905
+ if v in value:
906
+ current_group.append(v)
907
+ if len(current_group) == 0:
908
+ continue
909
+
910
+ var = current_group[0]
911
+ 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)
912
+ file_name = f"HYCOM_{variable_info[var]['var_name']}_{download_time}.nc"
913
+ old_str = f"var={variable_info[var]['var_name']}"
914
+ new_str = f"var={variable_info[var]['var_name']}"
915
+ if len(current_group) > 1:
916
+ for v in current_group[1:]:
917
+ new_str = f"{new_str}&var={variable_info[v]['var_name']}"
918
+ submit_url = submit_url.replace(old_str, new_str)
919
+ file_name = f"HYCOM_{key}_{download_time}.nc"
920
+ if download_time_end is not None:
921
+ file_name = f"HYCOM_{key}_{download_time}-{download_time_end}.nc"
922
+ _download_file(submit_url, store_path, file_name, cover)
923
+ else:
924
+ for v in var:
925
+ 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)
926
+ file_name = f"HYCOM_{variable_info[v]['var_name']}_{download_time}.nc"
927
+ if download_time_end is not None:
928
+ file_name = f"HYCOM_{variable_info[v]['var_name']}_{download_time}-{download_time_end}.nc"
929
+ _download_file(submit_url, store_path, file_name, cover)
930
+
931
+
932
+ def _convert_full_name_to_short_name(full_name):
933
+ for var, info in variable_info.items():
934
+ if full_name == info["var_name"] or full_name == info["standard_name"] or full_name == var:
935
+ return var
936
+ print("[bold #FFE4E1]Please ensure the var is in:\n[bold blue]u,v,temp,salt,ssh,u_b,v_b,temp_b,salt_b")
937
+ print("or")
938
+ print("[bold blue]water_u, water_v, water_temp, salinity, surf_el, water_u_bottom, water_v_bottom, water_temp_bottom, salinity_bottom")
939
+ return False
940
+
941
+
942
+ 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, cover):
943
+ _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, cover)
944
+
945
+
946
+ def _done_callback(future, progress, task, total, counter_lock):
947
+ global parallel_counter
948
+ with counter_lock:
949
+ parallel_counter += 1
950
+ progress.update(task, advance=1, description=f"[cyan]{bar_desc} {parallel_counter}/{total}")
951
+
952
+
953
+ 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, cover=False, interval_hour=3):
954
+ ymdh_time_s, ymdh_time_e = str(time_s), str(time_e)
955
+ if num_workers is not None and num_workers > 1:
956
+ global parallel_counter
957
+ parallel_counter = 0
958
+ counter_lock = Lock()
959
+ if ymdh_time_s == ymdh_time_e:
960
+ _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, cover)
961
+ elif int(ymdh_time_s) < int(ymdh_time_e):
962
+ if match_time is None:
963
+ print("*" * mark_len)
964
+ print("Downloading a series of files...")
965
+ time_list = _get_time_list(ymdh_time_s, ymdh_time_e, interval_hour, "hour")
966
+ with Progress() as progress:
967
+ task = progress.add_task(f"[cyan]{bar_desc}", total=len(time_list))
968
+ if num_workers is None or num_workers <= 1:
969
+ for i, time_str in enumerate(time_list):
970
+ _prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, time_str, None, depth, level, store_path, dataset_name, version_name, cover)
971
+ progress.update(task, advance=1, description=f"[cyan]{bar_desc} {i + 1}/{len(time_list)}")
972
+ else:
973
+ with ThreadPoolExecutor(max_workers=num_workers) as executor:
974
+ 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, cover) for time_str in time_list]
975
+ for feature in as_completed(futures):
976
+ _done_callback(feature, progress, task, len(time_list), counter_lock)
977
+ else:
978
+ print("[bold red]Please ensure the time_s is no more than time_e")
979
+
980
+
981
+ def download(
982
+ variables,
983
+ start_time,
984
+ end_time=None,
985
+ lon_min=0,
986
+ lon_max=359.92,
987
+ lat_min=-80,
988
+ lat_max=90,
989
+ depth=None,
990
+ level=None,
991
+ output_dir=None,
992
+ dataset=None,
993
+ version=None,
994
+ workers=None,
995
+ overwrite=False,
996
+ idm_path=None,
997
+ validate_time=None,
998
+ interval_hours=3,
999
+ proxy_txt=None,
1000
+ ):
1001
+ """
1002
+ Download data for a single time or a series of times.
1003
+
1004
+ Parameters:
1005
+ variables (str or list): Variable names to download. Examples include:
1006
+ 'u', 'v', 'temp', 'salt', 'ssh', 'u_b', 'v_b', 'temp_b', 'salt_b'
1007
+ or their full names like 'water_u', 'water_v', etc.
1008
+ start_time (str): Start time in the format 'YYYYMMDDHH' or 'YYYYMMDD'.
1009
+ If hour is included, it must be one of [00, 03, 06, 09, 12, 15, 18, 21].
1010
+ end_time (str, optional): End time in the format 'YYYYMMDDHH' or 'YYYYMMDD'.
1011
+ If not provided, only data for the start_time will be downloaded.
1012
+ lon_min (float, optional): Minimum longitude. Default is 0.
1013
+ lon_max (float, optional): Maximum longitude. Default is 359.92.
1014
+ lat_min (float, optional): Minimum latitude. Default is -80.
1015
+ lat_max (float, optional): Maximum latitude. Default is 90.
1016
+ depth (float, optional): Depth in meters. If specified, data for a single depth
1017
+ will be downloaded. Suggested range: [0, 5000].
1018
+ level (int, optional): Vertical level number. If specified, data for a single
1019
+ level will be downloaded. Suggested range: [1, 40].
1020
+ output_dir (str, optional): Directory to save downloaded files. If not provided,
1021
+ files will be saved in the current working directory.
1022
+ dataset (str, optional): Dataset name. Examples: 'GLBv0.08', 'GLBu0.08', etc.
1023
+ If not provided, the dataset will be chosen based on the time range.
1024
+ version (str, optional): Dataset version. Examples: '53.X', '56.3', etc.
1025
+ If not provided, the version will be chosen based on the time range.
1026
+ workers (int, optional): Number of parallel workers. Default is 1. Maximum is 10.
1027
+ overwrite (bool, optional): Whether to overwrite existing files. Default is False.
1028
+ idm_path (str, optional): Path to the Internet Download Manager (IDM) executable.
1029
+ If provided, IDM will be used for downloading.
1030
+ validate_time (bool, optional): Time validation mode. Default is None.
1031
+ - None: Only download data.
1032
+ - True: Modify the real time in the data to match the file name.
1033
+ - False: Check if the real time matches the file name. If not, delete the file.
1034
+ interval_hours (int, optional): Time interval in hours for downloading data.
1035
+ Default is 3. Examples: 3, 6, etc.
1036
+
1037
+ Returns:
1038
+ None
1039
+
1040
+ Example:
1041
+ >>> download(
1042
+ variables='u',
1043
+ start_time='2024083100',
1044
+ end_time='2024090100',
1045
+ lon_min=0,
1046
+ lon_max=359.92,
1047
+ lat_min=-80,
1048
+ lat_max=90,
1049
+ depth=None,
1050
+ level=None,
1051
+ output_dir=None,
1052
+ dataset=None,
1053
+ version=None,
1054
+ workers=4,
1055
+ overwrite=False,
1056
+ idm_path=None,
1057
+ validate_time=None,
1058
+ interval_hours=3,
1059
+ )
1060
+ """
1061
+ from oafuncs.oa_tool import pbar
1062
+
1063
+ _get_initial_data()
1064
+
1065
+ if dataset is None and version is None:
1066
+ if validate_time is None:
1067
+ print("Dataset and version will be chosen based on the time range.")
1068
+ print("If multiple datasets or versions exist, the latest one will be used.")
1069
+ elif dataset is None:
1070
+ print("Please provide a dataset name if specifying a version.")
1071
+ elif version is None:
1072
+ print("Please provide a version if specifying a dataset name.")
1073
+ else:
1074
+ print("Using the specified dataset and version.")
1075
+
1076
+ if isinstance(variables, list):
1077
+ if len(variables) == 1:
1078
+ variables = _convert_full_name_to_short_name(variables[0])
1079
+ else:
1080
+ variables = [_convert_full_name_to_short_name(v) for v in variables]
1081
+ elif isinstance(variables, str):
1082
+ variables = _convert_full_name_to_short_name(variables)
1083
+ else:
1084
+ raise ValueError("Invalid variable(s) provided.")
1085
+ if variables is False:
1086
+ raise ValueError("Invalid variable(s) provided.")
1087
+ if not (0 <= lon_min <= 359.92 and 0 <= lon_max <= 359.92 and -80 <= lat_min <= 90 and -80 <= lat_max <= 90):
1088
+ raise ValueError("Longitude or latitude values are out of range.")
1089
+
1090
+ if output_dir is None:
1091
+ output_dir = str(Path.cwd())
1092
+ else:
1093
+ os.makedirs(output_dir, exist_ok=True)
1094
+
1095
+ if workers is not None:
1096
+ workers = max(min(workers, 10), 1)
1097
+ start_time = str(start_time)
1098
+ if len(start_time) == 8:
1099
+ start_time += "00"
1100
+ if end_time is None:
1101
+ end_time = start_time[:]
1102
+ else:
1103
+ end_time = str(end_time)
1104
+ if len(end_time) == 8:
1105
+ end_time += "21"
1106
+
1107
+ global count_dict
1108
+ count_dict = {"success": 0, "fail": 0, "skip": 0, "no_data": 0, "total": 0, "no_data_list": []}
1109
+
1110
+ global fsize_dict
1111
+ fsize_dict = {}
1112
+
1113
+ global fsize_dict_lock
1114
+ fsize_dict_lock = Lock()
1115
+
1116
+ global use_idm, given_idm_engine, idm_download_list, bar_desc
1117
+ if idm_path is not None:
1118
+ use_idm = True
1119
+ workers = 1
1120
+ given_idm_engine = idm_path
1121
+ idm_download_list = []
1122
+ bar_desc = "Submitting to IDM ..."
1123
+ else:
1124
+ use_idm = False
1125
+ bar_desc = "Downloading ..."
1126
+
1127
+ global match_time
1128
+ match_time = validate_time
1129
+
1130
+ global mark_len
1131
+ mark_len = 100
1132
+
1133
+ global proxy_txt_path
1134
+ proxy_txt_path = proxy_txt
1135
+
1136
+ if validate_time is not None:
1137
+ workers = 1
1138
+ print("*" * mark_len)
1139
+ print("[bold red]Only checking the time of existing files.")
1140
+ bar_desc = "Checking time ..."
1141
+
1142
+ _download_hourly_func(
1143
+ variables,
1144
+ start_time,
1145
+ end_time,
1146
+ lon_min,
1147
+ lon_max,
1148
+ lat_min,
1149
+ lat_max,
1150
+ depth,
1151
+ level,
1152
+ output_dir,
1153
+ dataset,
1154
+ version,
1155
+ workers,
1156
+ overwrite,
1157
+ int(interval_hours),
1158
+ )
1159
+
1160
+ if idm_path is not None:
1161
+ print("[bold #ecdbfe]*" * mark_len)
1162
+ print(f"[bold #3dfc40]{'All files have been submitted to IDM for downloading'.center(mark_len, '*')}")
1163
+ print("[bold #ecdbfe]*" * mark_len)
1164
+ if idm_download_list:
1165
+ remain_list = idm_download_list.copy()
1166
+ for _ in pbar(range(len(idm_download_list)), cmap="diverging_1", description="Downloading: "):
1167
+ success = False
1168
+ while not success:
1169
+ for f in remain_list:
1170
+ if check_nc(f, print_messages=False):
1171
+ count_dict["success"] += 1
1172
+ success = True
1173
+ remain_list.remove(f)
1174
+ break
1175
+
1176
+ count_dict["total"] = count_dict["success"] + count_dict["fail"] + count_dict["skip"] + count_dict["no_data"]
1177
+ print("[bold #ecdbfe]=" * mark_len)
1178
+ print(f"[bold #ff80ab]Total: {count_dict['total']}\nSuccess: {count_dict['success']}\nFail: {count_dict['fail']}\nSkip: {count_dict['skip']}\nNo data: {count_dict['no_data']}")
1179
+ print("[bold #ecdbfe]=" * mark_len)
1180
+ if count_dict["fail"] > 0:
1181
+ print("[bold #be5528]Please try again to download the failed data later.")
1182
+ if count_dict["no_data"] > 0:
1183
+ print(f"[bold #f90000]{count_dict['no_data']} data entries do not exist in any dataset or version.")
1184
+ for no_data in count_dict["no_data_list"]:
1185
+ print(f"[bold #d81b60]{no_data}")
1186
+ print("[bold #ecdbfe]=" * mark_len)
1187
+
1188
+
1189
+ if __name__ == "__main__":
1190
+ download_dict = {
1191
+ "water_u": {"simple_name": "u", "download": 1},
1192
+ "water_v": {"simple_name": "v", "download": 1},
1193
+ "surf_el": {"simple_name": "ssh", "download": 1},
1194
+ "water_temp": {"simple_name": "temp", "download": 1},
1195
+ "salinity": {"simple_name": "salt", "download": 1},
1196
+ "water_u_bottom": {"simple_name": "u_b", "download": 0},
1197
+ "water_v_bottom": {"simple_name": "v_b", "download": 0},
1198
+ "water_temp_bottom": {"simple_name": "temp_b", "download": 0},
1199
+ "salinity_bottom": {"simple_name": "salt_b", "download": 0},
1200
+ }
1201
+
1202
+ var_list = [var_name for var_name in download_dict.keys() if download_dict[var_name]["download"]]
1203
+
1204
+ single_var = False
1205
+
1206
+ options = {
1207
+ "variables": var_list,
1208
+ "start_time": "2018010100",
1209
+ "end_time": "2019063000",
1210
+ "output_dir": r"G:\Data\HYCOM\china_sea\hourly_24",
1211
+ "lon_min": 105,
1212
+ "lon_max": 135,
1213
+ "lat_min": 10,
1214
+ "lat_max": 45,
1215
+ "workers": 1,
1216
+ "overwrite": False,
1217
+ "depth": None,
1218
+ "level": None,
1219
+ "validate_time": None,
1220
+ # "idm_path": r"D:\Programs\Internet Download Manager\IDMan.exe",
1221
+ "interval_hours": 24,
1222
+ "proxy_txt": None,
1223
+ }
1224
+
1225
+ if single_var:
1226
+ for var_name in var_list:
1227
+ options["variables"] = var_name
1228
+ download(**options)
1229
+ else:
1230
+ download(**options)