oafuncs 0.0.98.15__py3-none-any.whl → 0.0.98.17__py3-none-any.whl

oafuncs/_script/data_interp.py

@@ -1,18 +1,4 @@
-#!/usr/bin/env python
-# coding=utf-8
-"""
-Author: Liu Kun && 16031215@qq.com
-Date: 2025-04-25 16:22:52
-LastEditors: Liu Kun && 16031215@qq.com
-LastEditTime: 2025-04-25 16:22:52
-FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\_script\\data_interp.py
-Description:
-EditPlatform: vscode
-ComputerInfo: XPS 15 9510
-SystemInfo: Windows 11
-Python Version: 3.12
-"""
-
+import importlib.util
 from typing import List, Union
 
 import numpy as np
@@ -20,6 +6,85 @@ from scipy.interpolate import griddata
 
 from oafuncs.oa_tool import PEx
 
+# 检查 pykdtree 是否可用
+_has_pykdtree = importlib.util.find_spec("pykdtree.kdtree") is not None
+
+
+def _fill_nan_nearest(arr: np.ndarray) -> np.ndarray:
+    """
+    用最近邻填充 NaN（只支持2D数组）
+    """
+    # 基础检查：如果输入为None，直接返回None
+    if arr is None:
+        return None
+
+    # 确保是2D ndarray
+    arr = np.asarray(arr)
+    if arr.ndim != 2:
+        raise ValueError(f"_fill_nan_nearest 只支持2D数组，但输入的维度是 {arr.ndim}")
+
+    # 保存原始dtype并转为float
+    orig_dtype = arr.dtype
+    arr = arr.astype(float, copy=True)  # 使用copy=True确保不修改原数据
+
+    # 检查是否有NaN需要填充
+    mask = np.isnan(arr)
+    if not mask.any():
+        return arr.copy()
+
+    try:
+        valid = np.array(np.where(~mask)).T
+        invalid = np.array(np.where(mask)).T
+
+        # 如果有效点为空，直接返回原数据
+        if valid.shape[0] == 0:
+            return arr.copy()
+
+        # 使用KDTree进行最近邻填充
+        if _has_pykdtree:
+            from pykdtree.kdtree import KDTree
+
+            tree = KDTree(valid)
+            _, idx = tree.query(invalid, k=1)
+            filled = arr.copy()
+            filled[tuple(invalid.T)] = arr[tuple(valid[idx.flatten()].T)]
+        else:
+            # 备用方法：使用scipy的distance_transform_edt
+            from scipy.ndimage import distance_transform_edt
+
+            idx = distance_transform_edt(mask, return_distances=False, return_indices=True)
+            filled = arr[tuple(idx)]
+
+        return filled.astype(orig_dtype)
+    except Exception as e:
+        import warnings
+
+        warnings.warn(f"Error in _fill_nan_nearest: {e}, shape={arr.shape}")
+        return arr.copy()  # 发生异常返回原始数据
+
+
+def _data_clip(data: np.ndarray, data_min, data_max) -> np.ndarray:
+    """
+    将数据裁剪至 [data_min, data_max]，超出或 NaN 用最近邻填补。
+    支持 1~4D。
+    """
+    arr = np.array(data, copy=True)  # 使用副本避免修改原数据
+    ndims = arr.ndim
+    if ndims != 2:
+        raise ValueError(f"_data_clip 只支持1~4维数组，但输入的维度是 {ndims}")
+    dtype = arr.dtype
+
+    # 检查是否需要裁剪
+    mask = np.isnan(arr) | (arr < data_min) | (arr > data_max)
+    if not np.any(mask):
+        return arr.astype(dtype)
+
+    # 将超出范围的值设为NaN
+    arr[mask] = np.nan
+
+    return _fill_nan_nearest(arr).astype(dtype)
+
+
 
 def _interp_single_worker(*args):
     """
@@ -39,12 +104,9 @@ def _interp_single_worker(*args):
     result = griddata(valid_points, valid_data, target_points, method=interpolation_method)
     result = result.reshape(target_shape)
 
-    # 检查插值结果中是否仍有 NaN，如果有，用最近邻插值填充
-    if np.any(np.isnan(result)):
-        # 使用最近邻方法填充剩余的 NaN 值
-        nan_mask = np.isnan(result)
-        result_nn = griddata(valid_points, valid_data, target_points[nan_mask.ravel()], method="nearest")
-        result.ravel()[nan_mask.ravel()] = result_nn
+    # 第二步：用data_clip裁剪并填充
+    data_min, data_max = np.nanmin(data_slice), np.nanmax(data_slice)
+    result = _data_clip(result, data_min, data_max)
 
     return result
 
@@ -111,13 +173,13 @@ def interp_2d_func(
 
     t, z, y, x = new_src_data.shape
 
-    paras = []
+    params = []
     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))
+            params.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)
+        result = excutor.run(_interp_single_worker, params)
 
     return np.squeeze(np.array(result).reshape(t, z, *target_shape))

oafuncs/_script/data_interp_geo.py

@@ -0,0 +1,98 @@
+from typing import List, Union
+
+import numpy as np
+from scipy.interpolate import RectBivariateSpline
+
+from oafuncs.oa_tool import PEx
+from oafuncs.oa_data import data_clip
+from oafuncs._script.data_interp import _fill_nan_nearest
+
+
+def _interp_single_worker(*args):
+    """
+    单slice插值worker，参数为(data_slice, sx, sy, tx, ty, interpolation_method, data_min, data_max)
+    """
+    # 兼容PEx调用方式：args为tuple或list
+    if len(args) == 1 and isinstance(args[0], (tuple, list)):
+        args = args[0]
+    data_slice, sx, sy, tx, ty, interpolation_method, data_min, data_max = args
+    # 处理nan
+    if np.isnan(data_slice).any():
+        mask = np.isnan(data_slice)
+        if mask.any():
+            data_slice = _fill_nan_nearest(data_slice)
+    x1d = np.unique(sx[0, :])
+    y1d = np.unique(sy[:, 0])
+    if sx.shape != (len(y1d), len(x1d)) or sy.shape != (len(y1d), len(x1d)):
+        from scipy.interpolate import griddata
+
+        grid_points = np.column_stack((sx.ravel(), sy.ravel()))
+        grid_values = data_slice.ravel()
+        data_slice = griddata(grid_points, grid_values, (x1d[None, :], y1d[:, None]), method="linear")
+    if interpolation_method == "linear":
+        kx = ky = 1
+    else:
+        kx = ky = 3
+    interp_func = RectBivariateSpline(y1d, x1d, data_slice, kx=kx, ky=ky)
+    out = interp_func(ty[:, 0], tx[0, :])
+    # 优化裁剪逻辑：超出范围的点设为nan，再用fill_nan_nearest填充
+    arr = np.asarray(out)
+    arr = data_clip(arr,data_min,data_max)
+    return arr
+
+
+def interp_2d_geo(
+    target_x_coordinates: Union[np.ndarray, List[float]],
+    target_y_coordinates: Union[np.ndarray, List[float]],
+    source_x_coordinates: Union[np.ndarray, List[float]],
+    source_y_coordinates: Union[np.ndarray, List[float]],
+    source_data: np.ndarray,
+    interpolation_method: str = "cubic",
+) -> np.ndarray:
+    """
+    更平滑的二维插值，采用RectBivariateSpline实现bicubic效果，接口与interp_2d兼容。
+    支持输入2D/3D/4D数据，最后两维为空间。
+    interpolation_method: "cubic"（默认，bicubic），"linear"（双线性）
+    插值后自动裁剪并用最近邻填充超限和NaN，范围取原始数据的nanmin/nanmax
+    """
+    # 保证输入为ndarray
+    tx = np.asarray(target_x_coordinates)
+    ty = np.asarray(target_y_coordinates)
+    sx = np.asarray(source_x_coordinates)
+    sy = np.asarray(source_y_coordinates)
+    data = np.asarray(source_data)
+
+    if ty.ndim == 1:
+        tx, ty = np.meshgrid(tx, ty)
+    if sy.ndim == 1:
+        sx, sy = np.meshgrid(sx, sy)
+
+    if sx.shape != data.shape[-2:] or sy.shape != data.shape[-2:]:
+        raise ValueError("Shape of source_data does not match shape of source_x_coordinates or source_y_coordinates.")
+
+    data_dims = data.ndim
+    if data_dims < 2:
+        raise ValueError("Source data must have at least 2 dimensions.")
+    elif data_dims > 4:
+        raise ValueError("Source data has more than 4 dimensions, not supported.")
+
+    num_dims_to_add = 4 - data_dims
+    new_shape = (1,) * num_dims_to_add + data.shape
+    data4d = data.reshape(new_shape)
+    t, z, ny, nx = data4d.shape
+
+    data_min, data_max = np.nanmin(data), np.nanmax(data)
+    target_shape = ty.shape
+
+    # 并行参数准备
+    params = []
+    for ti in range(t):
+        for zi in range(z):
+            params.append((data4d[ti, zi], sx, sy, tx, ty, interpolation_method, data_min, data_max))
+
+    with PEx() as excutor:
+        result = excutor.run(_interp_single_worker, params)
+
+    result = np.array(result).reshape(t, z, *target_shape)
+    result = np.squeeze(result)
+    return result

oafuncs/oa_data.py

@@ -23,7 +23,7 @@ from rich import print
 from scipy.interpolate import interp1d
 
 
-__all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile"]
+__all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile", "data_clip"]
 
 
 def ensure_list(input_value: Any) -> List[str]:
@@ -114,7 +114,9 @@ def interp_along_dim(
     return np.apply_along_axis(apply_interp_extrap, interpolation_axis, source_data)
 
 
-
+def data_clip(data: np.ndarray, data_min: float, data_max: float) -> np.ndarray:
+    from ._script.data_interp import _data_clip
+    _data_clip(data, data_min, data_max)
 
 
 def interp_2d(

{oafuncs-0.0.98.15.dist-info → oafuncs-0.0.98.17.dist-info}/METADATA

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

{oafuncs-0.0.98.15.dist-info → oafuncs-0.0.98.17.dist-info}/RECORD

@@ -1,6 +1,6 @@
 oafuncs/__init__.py,sha256=T_-VtnWWllV3Q91twT5Yt2sUapeA051QbPNnBxmg9nw,1456
 oafuncs/oa_cmap.py,sha256=DimWT4Bg7uE5Lx8hSw1REp7whpsR2pFRStAwk1cowEM,11494
-oafuncs/oa_data.py,sha256=hngaxAi_r6PsHXzSeT3DMY_QdChWjuBMPOZNFvWU388,8442
+oafuncs/oa_data.py,sha256=NOp7Rmjjypb17CNjk3Ksz08R7LygCEP6K4mna9mSXa8,8626
 oafuncs/oa_date.py,sha256=WhM6cyD4G3IeghjLTHhAMtlvJbA7kwQG2sHnxdTgyso,6303
 oafuncs/oa_draw.py,sha256=Wj2QBgyIPpV_dxaDrH10jqj_puK9ZM9rd-si-3VrsrE,17631
 oafuncs/oa_file.py,sha256=j9gXJgPOJsliu4IOUc4bc-luW4yBvQyNCEmMyDVjUwQ,16404
@@ -11,7 +11,8 @@ oafuncs/oa_tool.py,sha256=rpPkLqWhqMmqlCc5wjL8qMTg3gThCkSrYJckbX_0iJc,8631
 oafuncs/_data/hycom.png,sha256=MadKs6Gyj5n9-TOu7L4atQfTXtF9dvN9w-tdU9IfygI,10945710
 oafuncs/_data/oafuncs.png,sha256=o3VD7wm-kwDea5E98JqxXl04_78cBX7VcdUt7uQXGiU,3679898
 oafuncs/_script/cprogressbar.py,sha256=UIgGcLFs-6IgWlITuBLaQqrpt4OAK3Mst5RlCiNfZdQ,15772
-oafuncs/_script/data_interp.py,sha256=70U-Jsoxd5g-7dEQt4IaDRuRV-M_1lEKMGOci86vSVE,5431
+oafuncs/_script/data_interp.py,sha256=_k8EMSiFxutrqEVTLsL4mEPE6ssYq8bzmBlksCZ9nAE,7428
+oafuncs/_script/data_interp_geo.py,sha256=X89KxLYhpltWi0Sf96gIhBL3r1M5aExd_JCmgBmmvUc,3742
 oafuncs/_script/email.py,sha256=lL4HGKrr524-g0xLlgs-4u7x4-u7DtgNoD9AL8XJKj4,3058
 oafuncs/_script/netcdf_merge.py,sha256=9hCyxfeUHnBzs50_0v0jzVfxpMxTX4dNTo0pmsp_T6g,4226
 oafuncs/_script/netcdf_modify.py,sha256=sGRUYNhfGgf9JV70rnBzw3bzuTRSXzBTL_RMDnDPeLQ,4552
@@ -38,8 +39,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.15.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
-oafuncs-0.0.98.15.dist-info/METADATA,sha256=jTAHHAY0xOxy2z13wYD4x4JOOSfsUxrz1WVyzXNfi9o,4273
-oafuncs-0.0.98.15.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-oafuncs-0.0.98.15.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
-oafuncs-0.0.98.15.dist-info/RECORD,,
+oafuncs-0.0.98.17.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
+oafuncs-0.0.98.17.dist-info/METADATA,sha256=8-0Gp7bgVD7qR4sfSQtKmoCl5yV8f7YBnEROZV28oJY,4273
+oafuncs-0.0.98.17.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+oafuncs-0.0.98.17.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
+oafuncs-0.0.98.17.dist-info/RECORD,,