oafuncs 0.0.98.6__py3-none-any.whl → 0.0.98.8__py3-none-any.whl

oafuncs/oa_data.py

@@ -13,9 +13,7 @@ SystemInfo: Windows 11
 Python Version: 3.11
 """
 
-import itertools
-import multiprocessing as mp
-from concurrent.futures import ThreadPoolExecutor
+
 from typing import Any, List, Union
 
 import numpy as np
@@ -24,6 +22,8 @@ import xarray as xr
 from rich import print
 from scipy.interpolate import griddata, interp1d
 
+from oafuncs.oa_tool import PEx
+
 __all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile"]
 
 
@@ -115,6 +115,15 @@ def interp_along_dim(
     return np.apply_along_axis(apply_interp_extrap, interpolation_axis, source_data)
 
 
+def _interp_single_worker(*args):
+    """
+    用于PEx并行的单slice插值worker，参数为(t, z, source_data, origin_points, target_points, interpolation_method, target_shape)
+    """
+    data_slice, origin_points, target_points, interpolation_method, target_shape = args
+
+    return griddata(origin_points, data_slice.ravel(), target_points, method=interpolation_method).reshape(target_shape)
+
+
 def interp_2d(
     target_x_coordinates: Union[np.ndarray, List[float]],
     target_y_coordinates: Union[np.ndarray, List[float]],
@@ -122,7 +131,6 @@ def interp_2d(
     source_y_coordinates: Union[np.ndarray, List[float]],
     source_data: np.ndarray,
     interpolation_method: str = "linear",
-    use_parallel: bool = True,
 ) -> np.ndarray:
     """
     Perform 2D interpolation on the last two dimensions of a multi-dimensional array.
@@ -151,10 +159,6 @@ def interp_2d(
         >>> result = interp_2d(target_x_coordinates, target_y_coordinates, source_x_coordinates, source_y_coordinates, source_data)
         >>> print(result.shape)  # Expected output: (3, 3)
     """
-
-    def interp_single(data_slice: np.ndarray, target_points: np.ndarray, origin_points: np.ndarray, method: str) -> np.ndarray:
-        return griddata(origin_points, data_slice.ravel(), target_points, method=method).reshape(target_y_coordinates.shape)
-
     if len(target_y_coordinates.shape) == 1:
         target_x_coordinates, target_y_coordinates = np.meshgrid(target_x_coordinates, target_y_coordinates)
     if len(source_y_coordinates.shape) == 1:
@@ -166,25 +170,30 @@ def interp_2d(
     target_points = np.column_stack((np.array(target_y_coordinates).ravel(), np.array(target_x_coordinates).ravel()))
     origin_points = np.column_stack((np.array(source_y_coordinates).ravel(), np.array(source_x_coordinates).ravel()))
 
-    if use_parallel:
-        with ThreadPoolExecutor(max_workers=mp.cpu_count() - 2) as executor:
-            if len(source_data.shape) == 2:
-                interpolated_data = list(executor.map(interp_single, [source_data], [target_points], [origin_points], [interpolation_method]))
-            elif len(source_data.shape) == 3:
-                interpolated_data = list(executor.map(interp_single, [source_data[i] for i in range(source_data.shape[0])], [target_points] * source_data.shape[0], [origin_points] * source_data.shape[0], [interpolation_method] * source_data.shape[0]))
-            elif len(source_data.shape) == 4:
-                index_combinations = list(itertools.product(range(source_data.shape[0]), range(source_data.shape[1])))
-                interpolated_data = list(executor.map(interp_single, [source_data[i, j] for i, j in index_combinations], [target_points] * len(index_combinations), [origin_points] * len(index_combinations), [interpolation_method] * len(index_combinations)))
-                interpolated_data = np.array(interpolated_data).reshape(source_data.shape[0], source_data.shape[1], *target_y_coordinates.shape)
-    else:
-        if len(source_data.shape) == 2:
-            interpolated_data = interp_single(source_data, target_points, origin_points, interpolation_method)
-        elif len(source_data.shape) == 3:
-            interpolated_data = np.stack([interp_single(source_data[i], target_points, origin_points, interpolation_method) for i in range(source_data.shape[0])])
-        elif len(source_data.shape) == 4:
-            interpolated_data = np.stack([np.stack([interp_single(source_data[i, j], target_points, origin_points, interpolation_method) for j in range(source_data.shape[1])]) for i in range(source_data.shape[0])])
-
-    return np.squeeze(np.array(interpolated_data))
+    data_dims = len(source_data.shape)
+    # Ensure source_data is 4D for consistent processing (t, z, y, x)
+    if data_dims < 2:
+        raise ValueError(f"[red]Source data must have at least 2 dimensions, but got {data_dims}.[/red]")
+    elif data_dims > 4:
+        # Or handle cases with more than 4 dimensions if necessary
+        raise ValueError(f"[red]Source data has {data_dims} dimensions, but this function currently supports only up to 4.[/red]")
+
+    # Reshape to 4D by adding leading dimensions of size 1 if needed
+    num_dims_to_add = 4 - data_dims
+    new_shape = (1,) * num_dims_to_add + source_data.shape
+    new_src_data = source_data.reshape(new_shape)
+
+    t, z, _, _ = new_src_data.shape
+
+    paras = []
+    target_shape = target_y_coordinates.shape
+    for t_index in range(t):
+        for z_index in range(z):
+            paras.append((new_src_data[t_index, z_index], origin_points, target_points, interpolation_method, target_shape))
+    excutor = PEx()
+    result = excutor.run(_interp_single_worker, paras)
+
+    return np.squeeze(np.array(result))
 
 
 def mask_shapefile(

oafuncs/oa_down/hycom_3hourly.py

@@ -19,7 +19,6 @@ import logging
 import os
 import random
 import re
-import time
 import warnings
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
@@ -728,9 +727,15 @@ class _HycomDownloader:
 
                     elapsed = datetime.datetime.now() - start
                     # logging.info(f"File {file_name} downloaded, Time: {elapsed}")
-                    logging.info(f"Saving {file_name}, Time: {elapsed}")
+                    logging.info(f"Saving {file_name} ...")
+                    logging.info(f"Timing {elapsed} ...")
                     self.count["success"] += 1
                     count_dict["success"] += 1
+                    # 输出一条绿色的成功消息
+                    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
+                    # print(f"{timestamp} - INFO - ", end="") # Output log prefix without newline
+                    # print("[bold #3dfc40]Success")
+                    print(f"{timestamp} - RESULT - [bold #3dfc40]Success")
                     return
 
             except Exception as e:
@@ -744,6 +749,10 @@ class _HycomDownloader:
                     logging.error(f"Giving up on {file_name}")
                     self.count["fail"] += 1
                     count_dict["fail"] += 1
+                    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
+                    # print(f"{timestamp} - ERROR - ", end="")
+                    # print("[bold red]Failed")
+                    print(f"{timestamp} - RESULT - [bold red]Failure")
                     return
 
     async def run(self):
@@ -757,8 +766,7 @@ class _HycomDownloader:
 
 def _download_file(target_url, store_path, file_name, cover=False):
     save_path = Path(store_path) / file_name
-    file_name_split = file_name.split("_")
-    file_name_split = file_name_split[:-1]
+    file_name_split = file_name.split("_")[:-1]
     same_file = "_".join(file_name_split) + "*nc"
 
     if match_time is not None:
@@ -775,10 +783,12 @@ def _download_file(target_url, store_path, file_name, cover=False):
                 print(f"[bold green]{file_name} is correct")
         return
 
-    if not cover and os.path.exists(save_path):
-        print(f"[bold #FFA54F]{save_path} exists, skipping ...")
-        count_dict["skip"] += 1
-        return
+    # if not cover and os.path.exists(save_path):
+    #     print(f"[bold #FFA54F]{save_path} exists, skipping ...")
+    #     count_dict["skip"] += 1
+    #     return
+    if cover and os.path.exists(save_path):
+        _clear_existing_file(save_path)
 
     if same_file not in fsize_dict.keys():
         check_nc(save_path, delete_if_invalid=True, print_messages=False)
@@ -786,6 +796,7 @@ def _download_file(target_url, store_path, file_name, cover=False):
     get_mean_size = _get_mean_size_move(same_file, save_path)
 
     if _check_existing_file(save_path, get_mean_size):
+        print(f"[bold #FFA54F]{save_path} exists, skipping ...")
         count_dict["skip"] += 1
         return
 
@@ -794,12 +805,12 @@ def _download_file(target_url, store_path, file_name, cover=False):
     if not use_idm:
         python_downloader = _HycomDownloader([(target_url, save_path)])
         asyncio.run(python_downloader.run())
-        time.sleep(3 + random.uniform(0, 10))
+        # time.sleep(3 + random.uniform(0, 10))
     else:
         idm_downloader(target_url, store_path, file_name, given_idm_engine)
         idm_download_list.append(save_path)
         # print(f"[bold #3dfc40]File [bold #dfff73]{save_path} [#3dfc40]has been submit to IDM for downloading")
-        time.sleep(3 + random.uniform(0, 10))
+        # time.sleep(3 + random.uniform(0, 10))
 
 
 def _check_hour_is_valid(ymdh_str):
@@ -890,7 +901,7 @@ def _prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_ma
         else:
             if download_time < "2024081012":
                 varlist = [_ for _ in var]
-                for key, value in pbar(var_group.items(), description=f"Var_group {download_time} ->", total=len(var_group), cmap="bwr", next_line=True):
+                for key, value in pbar(var_group.items(), description=f"Var Group {download_time} ->", total=len(var_group), color="#d7feb9", next_line=True):
                     current_group = []
                     for v in varlist:
                         if v in value:
@@ -912,7 +923,7 @@ def _prepare_url_to_download(var, lon_min=0, lon_max=359.92, lat_min=-80, lat_ma
                             file_name = f"HYCOM_{key}_{download_time}-{download_time_end}.nc"
                     _download_file(submit_url, store_path, file_name, cover)
             else:
-                for v in pbar(var,description=f'Var {download_time} ->', total=len(var), cmap='bwr', next_line=True):
+                for v in pbar(var, description=f"Var {download_time} ->", total=len(var), color="#d7feb9", next_line=True):
                     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:

oafuncs/oa_draw.py

@@ -15,6 +15,7 @@ Python Version: 3.11
 
 import warnings
 
+import cv2
 import cartopy.crs as ccrs
 import cartopy.feature as cfeature
 import matplotlib as mpl
@@ -132,6 +133,70 @@ def gif(image_paths: list[str], output_gif_name: str, frame_duration: float = 20
     return
 
 
+def movie(image_files, output_video_path, fps):
+    """
+    从图像文件列表创建视频。
+
+    Args:
+        image_files (list): 按顺序排列的图像文件路径列表。
+        output_video_path (str): 输出视频文件的路径 (例如 'output.mp4')。
+        fps (int): 视频的帧率。
+    """
+    if not image_files:
+        print("错误：图像文件列表为空。")
+        return
+
+    # 读取第一张图片以获取帧尺寸
+    try:
+        frame = cv2.imread(image_files[0])
+        if frame is None:
+            print(f"错误：无法读取第一张图片：{image_files[0]}")
+            return
+        height, width, layers = frame.shape
+        size = (width, height)
+        print(f"视频尺寸设置为：{size}")
+    except Exception as e:
+        print(f"读取第一张图片时出错：{e}")
+        return
+
+    # 选择编解码器并创建VideoWriter对象
+    # 对于 .mp4 文件，常用 'mp4v' 或 'avc1'
+    # 对于 .avi 文件，常用 'XVID' 或 'MJPG'
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # 或者尝试 'avc1', 'XVID' 等
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, size)
+
+    if not out.isOpened():
+        print(f"错误：无法打开视频文件进行写入：{output_video_path}")
+        print("请检查编解码器 ('fourcc') 是否受支持以及路径是否有效。")
+        return
+
+    print(f"开始将图像写入视频：{output_video_path}...")
+    for i, filename in enumerate(image_files):
+        try:
+            frame = cv2.imread(filename)
+            if frame is None:
+                print(f"警告：跳过无法读取的图像：{filename}")
+                continue
+            # 确保帧尺寸与初始化时相同，如果需要可以调整大小
+            current_height, current_width, _ = frame.shape
+            if (current_width, current_height) != size:
+                print(f"警告：图像 {filename} 的尺寸 ({current_width}, {current_height}) 与初始尺寸 {size} 不同。将调整大小。")
+                frame = cv2.resize(frame, size)
+
+            out.write(frame)
+            # 打印进度（可选）
+            if (i + 1) % 50 == 0 or (i + 1) == len(image_files):
+                print(f"已处理 {i + 1}/{len(image_files)} 帧")
+
+        except Exception as e:
+            print(f"处理图像 {filename} 时出错：{e}")
+            continue  # 跳过有问题的帧
+
+    # 释放资源
+    out.release()
+    print(f"视频创建成功：{output_video_path}")
+
+
 def add_lonlat_unit(longitudes: list[float] = None, latitudes: list[float] = None, decimal_places: int = 2) -> tuple[list[str], list[str]] | list[str]:
     """Convert longitude and latitude values to formatted string labels.
 

{oafuncs-0.0.98.6.dist-info → oafuncs-0.0.98.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: oafuncs
-Version: 0.0.98.6
+Version: 0.0.98.8
 Summary: Oceanic and Atmospheric Functions
 Home-page: https://github.com/Industry-Pays/OAFuncs
 Author: Kun Liu
@@ -27,6 +27,7 @@ Requires-Dist: requests
 Requires-Dist: bs4
 Requires-Dist: httpx
 Requires-Dist: matplotlib
+Requires-Dist: opencv-python
 Requires-Dist: netCDF4
 Requires-Dist: xlrd
 Requires-Dist: geopandas

{oafuncs-0.0.98.6.dist-info → oafuncs-0.0.98.8.dist-info}/RECORD

@@ -1,8 +1,8 @@
 oafuncs/__init__.py,sha256=T_-VtnWWllV3Q91twT5Yt2sUapeA051QbPNnBxmg9nw,1456
 oafuncs/oa_cmap.py,sha256=DimWT4Bg7uE5Lx8hSw1REp7whpsR2pFRStAwk1cowEM,11494
-oafuncs/oa_data.py,sha256=0AbQ8_7vf9ecZaui6hmUjubkWRJxs4TGcdhJaPdbmP8,10958
+oafuncs/oa_data.py,sha256=7heyoFOBt_xqe0YSiUdO6tOpmySm0FuG1fHSoAO1NJI,10271
 oafuncs/oa_date.py,sha256=KqU-bHtC74hYsf6VgiA3i2vI__q_toOVR-whFy4cYP8,5523
-oafuncs/oa_draw.py,sha256=9hchfs7kyBo_On1FA5jliSoAwpsjEXNUIC9C20D9ugw,15050
+oafuncs/oa_draw.py,sha256=Wj2QBgyIPpV_dxaDrH10jqj_puK9ZM9rd-si-3VrsrE,17631
 oafuncs/oa_file.py,sha256=goF5iRXJFFCIKhIjlkCnYYt0EYlJb_4r8AeYNZ0-SOk,16209
 oafuncs/oa_help.py,sha256=_4AZgRDq5Or0vauNvq5IDDHIBoBfdOQtzak-mG1wwAw,4537
 oafuncs/oa_nc.py,sha256=L1gqXxg93kIDsMOa87M0o-53KVmdqCipnXeF9XfzfY8,10513
@@ -21,7 +21,7 @@ oafuncs/_script/plot_dataset.py,sha256=zkSEnO_-biyagorwWXPoihts_cwuvripzEt-l9bHJ
 oafuncs/_script/replace_file_content.py,sha256=eCFZjnZcwyRvy6b4mmIfBna-kylSZTyJRfgXd6DdCjk,5982
 oafuncs/oa_down/User_Agent-list.txt,sha256=pHaMlElMvZ8TG4vf4BqkZYKqe0JIGkr4kCN0lM1Y9FQ,514295
 oafuncs/oa_down/__init__.py,sha256=kRX5eTUCbAiz3zTaQM1501paOYS_3fizDN4Pa0mtNUA,585
-oafuncs/oa_down/hycom_3hourly.py,sha256=zU7eNFyg5oa-5MVrCeCVyZfoYHukY7MA0-DmvCj5yIk,54723
+oafuncs/oa_down/hycom_3hourly.py,sha256=wWV14-OB9_LMmjUiZr3YXWBdKKwAyGXNa3Up7fSiWwk,55553
 oafuncs/oa_down/hycom_3hourly_proxy.py,sha256=1eaoJGI_m-7w4ZZ3n7NGxkZaeFdsm0d3U-hyw8RFNbc,54563
 oafuncs/oa_down/idm.py,sha256=4z5IvgfTyIKEI1kOtqXZwN7Jnfjwp6qDBOIoVyOLp0I,1823
 oafuncs/oa_down/literature.py,sha256=2bF9gSKQbzcci9LcKE81j8JEjIJwON7jbwQB3gDDA3E,11331
@@ -37,8 +37,8 @@ oafuncs/oa_sign/__init__.py,sha256=QKqTFrJDFK40C5uvk48GlRRbGFzO40rgkYwu6dYxatM,5
 oafuncs/oa_sign/meteorological.py,sha256=8091SHo2L8kl4dCFmmSH5NGVHDku5i5lSiLEG5DLnOQ,6489
 oafuncs/oa_sign/ocean.py,sha256=xrW-rWD7xBWsB5PuCyEwQ1Q_RDKq2KCLz-LOONHgldU,5932
 oafuncs/oa_sign/scientific.py,sha256=a4JxOBgm9vzNZKpJ_GQIQf7cokkraV5nh23HGbmTYKw,5064
-oafuncs-0.0.98.6.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
-oafuncs-0.0.98.6.dist-info/METADATA,sha256=g6aPVBYo1PEpyW-4-0o1HgXZiNJZKGg1p3YEP2skSAY,4242
-oafuncs-0.0.98.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-oafuncs-0.0.98.6.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
-oafuncs-0.0.98.6.dist-info/RECORD,,
+oafuncs-0.0.98.8.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
+oafuncs-0.0.98.8.dist-info/METADATA,sha256=Le6ieydYuvZciK8CaFHPBUy47T3geyA48Z6nzSbLMwQ,4272
+oafuncs-0.0.98.8.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+oafuncs-0.0.98.8.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
+oafuncs-0.0.98.8.dist-info/RECORD,,

{oafuncs-0.0.98.6.dist-info → oafuncs-0.0.98.8.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (78.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any