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

oafuncs/_script/netcdf_merge.py

@@ -1,9 +1,12 @@
 import os
 from typing import List, Optional, Union
-from dask.diagnostics import ProgressBar
+
+import numpy as np
 import xarray as xr
+from dask.diagnostics import ProgressBar
 from oafuncs import pbar
 
+
 def merge_nc(file_list: Union[str, List[str]], var_name: Optional[Union[str, List[str]]] = None, dim_name: Optional[str] = None, target_filename: Optional[str] = None) -> None:
     """
     Description:
@@ -22,7 +25,7 @@ def merge_nc(file_list: Union[str, List[str]], var_name: Optional[Union[str, Lis
         merge(file_list, var_name=['u', 'v'], dim_name='time', target_filename='merged.nc')
         merge(file_list, var_name=None, dim_name='time', target_filename='merged.nc')
     """
-    
+
     if target_filename is None:
         target_filename = "merged.nc"
 
@@ -55,11 +58,19 @@ def merge_nc(file_list: Union[str, List[str]], var_name: Optional[Union[str, Lis
                 if dim_name in data_var.dims:
                     merged_data.setdefault(var, []).append(data_var)
                 elif var not in merged_data:
-                    merged_data[var] = data_var.fillna(0)  # 用0填充NaN值
+                    # 判断类型，时间类型用NaT填充
+                    if np.issubdtype(data_var.dtype, np.datetime64):
+                        merged_data[var] = data_var.fillna(np.datetime64("NaT"))
+                    else:
+                        merged_data[var] = data_var.fillna(0)
 
     for var in pbar(merged_data, description="Merging variables", color="#9b45d1"):
         if isinstance(merged_data[var], list):
-            merged_data[var] = xr.concat(merged_data[var], dim=dim_name).fillna(0)
+            # 判断类型，时间类型用NaT填充
+            if np.issubdtype(merged_data[var][0].dtype, np.datetime64):
+                merged_data[var] = xr.concat(merged_data[var], dim=dim_name).fillna(np.datetime64("NaT"))
+            else:
+                merged_data[var] = xr.concat(merged_data[var], dim=dim_name).fillna(0)
         # print(f"Variable '{var}' merged: min={merged_data[var].min().values:.3f}, max={merged_data[var].max().values:.3f}, mean={merged_data[var].mean().values:.3f}")
 
     # 修改写入数据部分，支持压缩并设置基数和比例因子

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,31 @@ 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))
+    
+    with PEx() as excutor:
+        result = excutor.run(_interp_single_worker, paras)
+
+    return np.squeeze(np.array(result))
 
 
 def mask_shapefile(

oafuncs/oa_date.py

@@ -64,7 +64,6 @@ def hour_range(start_time: str, end_time: str, hour_interval: int = 6) -> List[s
         date_s += datetime.timedelta(hours=hour_interval)
     return date_list
 
-
 def adjust_time(base_time: str, time_delta: int, delta_unit: str = "hours", output_format: Optional[str] = None) -> str:
     """
     Adjust a given base time by adding a specified time delta.
@@ -89,7 +88,12 @@ def adjust_time(base_time: str, time_delta: int, delta_unit: str = "hours", outp
     """
     # Normalize the input time to "yyyymmddHHMMSS" format
     time_format = "%Y%m%d%H%M%S"
+    if len(base_time) == 4:
+        base_time += "0101"
+    elif len(base_time) == 6:
+        base_time += "01"
     base_time = base_time.ljust(14, "0")
+
     time_obj = datetime.datetime.strptime(base_time, time_format)
 
     # Add the specified amount of time
@@ -101,8 +105,19 @@ def adjust_time(base_time: str, time_delta: int, delta_unit: str = "hours", outp
         time_obj += datetime.timedelta(hours=time_delta)
     elif delta_unit == "days":
         time_obj += datetime.timedelta(days=time_delta)
+    elif delta_unit == "months":
+        # Handle month addition separately
+        month = time_obj.month - 1 + time_delta
+        year = time_obj.year + month // 12
+        month = month % 12 + 1
+        day = min(time_obj.day, month_days(year, month))
+        time_obj = time_obj.replace(year=year, month=month, day=day)
+    elif delta_unit == "years":
+        # Handle year addition separately
+        year = time_obj.year + time_delta
+        time_obj = time_obj.replace(year=year)
     else:
-        raise ValueError("Invalid time unit. Use 'seconds', 'minutes', 'hours', or 'days'.")
+        raise ValueError("Invalid time unit. Use 'seconds', 'minutes', 'hours', 'days', 'months', or 'years'.")
 
     # Determine the output format
     if output_format:
@@ -116,6 +131,10 @@ def adjust_time(base_time: str, time_delta: int, delta_unit: str = "hours", outp
             default_format = "%Y%m%d%H"
         elif delta_unit == "days":
             default_format = "%Y%m%d"
+        elif delta_unit == "months":
+            default_format = "%Y%m"
+        elif delta_unit == "years":
+            default_format = "%Y"
         return time_obj.strftime(default_format)
 
 

{oafuncs-0.0.98.7.dist-info → oafuncs-0.0.98.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: oafuncs
-Version: 0.0.98.7
+Version: 0.0.98.9
 Summary: Oceanic and Atmospheric Functions
 Home-page: https://github.com/Industry-Pays/OAFuncs
 Author: Kun Liu

{oafuncs-0.0.98.7.dist-info → oafuncs-0.0.98.9.dist-info}/RECORD

@@ -1,7 +1,7 @@
 oafuncs/__init__.py,sha256=T_-VtnWWllV3Q91twT5Yt2sUapeA051QbPNnBxmg9nw,1456
 oafuncs/oa_cmap.py,sha256=DimWT4Bg7uE5Lx8hSw1REp7whpsR2pFRStAwk1cowEM,11494
-oafuncs/oa_data.py,sha256=0AbQ8_7vf9ecZaui6hmUjubkWRJxs4TGcdhJaPdbmP8,10958
-oafuncs/oa_date.py,sha256=KqU-bHtC74hYsf6VgiA3i2vI__q_toOVR-whFy4cYP8,5523
+oafuncs/oa_data.py,sha256=SqvG2mRVSTmr5j7H3Cn-ILgDo8iwSAl6W5NLpVpBx3c,10288
+oafuncs/oa_date.py,sha256=WhM6cyD4G3IeghjLTHhAMtlvJbA7kwQG2sHnxdTgyso,6303
 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
@@ -12,7 +12,7 @@ oafuncs/_data/hycom.png,sha256=MadKs6Gyj5n9-TOu7L4atQfTXtF9dvN9w-tdU9IfygI,10945
 oafuncs/_data/oafuncs.png,sha256=o3VD7wm-kwDea5E98JqxXl04_78cBX7VcdUt7uQXGiU,3679898
 oafuncs/_script/cprogressbar.py,sha256=wRU3SFPFtMI7ER26tTzg223kVKNo5RDWE9CzdIgUsuE,15771
 oafuncs/_script/email.py,sha256=lL4HGKrr524-g0xLlgs-4u7x4-u7DtgNoD9AL8XJKj4,3058
-oafuncs/_script/netcdf_merge.py,sha256=_EPF9Xj4HOVC9sZpi1lt62-Aq6pMlgsgwaajEBLhW6g,5092
+oafuncs/_script/netcdf_merge.py,sha256=ktmTOgGfLHBNdS4HBc6xFDfO8B7E4DT7d1e6Dtare9Y,5596
 oafuncs/_script/netcdf_modify.py,sha256=sGRUYNhfGgf9JV70rnBzw3bzuTRSXzBTL_RMDnDPeLQ,4552
 oafuncs/_script/netcdf_write.py,sha256=iO1Qv9bp6RLiw1D8Nrv7tX_8X-diUZaX3Nxhk6pJ5Nw,8556
 oafuncs/_script/parallel.py,sha256=dRT7w_rBnR3mZkUlO6v6j05SwBTQpTccOna5CXI5Msg,8196
@@ -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.7.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
-oafuncs-0.0.98.7.dist-info/METADATA,sha256=HotYRm4-mVHvJqg_pe71FgQpZ7mYKjd6WGr2PBbIeJY,4272
-oafuncs-0.0.98.7.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-oafuncs-0.0.98.7.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
-oafuncs-0.0.98.7.dist-info/RECORD,,
+oafuncs-0.0.98.9.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
+oafuncs-0.0.98.9.dist-info/METADATA,sha256=Ns5ahXYSJiK2OmggCSbzoFC5XHMnLsFdC_KMnVrAnGA,4272
+oafuncs-0.0.98.9.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+oafuncs-0.0.98.9.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
+oafuncs-0.0.98.9.dist-info/RECORD,,

{oafuncs-0.0.98.7.dist-info → oafuncs-0.0.98.9.dist-info}/WHEEL

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