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

oafuncs/_script/data_interp.py

@@ -1,24 +1,52 @@
-#!/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
+from oafuncs.oa_tool import PEx
 from scipy.interpolate import griddata
 
-from oafuncs.oa_tool import PEx
+_has_pykdtree = importlib.util.find_spec("pykdtree.kdtree") is not None
+
+
+def _fill_nan_nearest(arr):
+    """用最近邻插值填充 NaN，优先用pykdtree加速"""
+    mask = np.isnan(arr)
+    if not mask.any():
+        return arr
+    if _has_pykdtree:
+        from pykdtree.kdtree import KDTree
+
+        valid_idx = np.array(np.where(~mask)).T
+        nan_idx = np.array(np.where(mask)).T
+        if len(valid_idx) == 0:
+            # 全是nan，直接返回
+            return arr
+        tree = KDTree(valid_idx)
+        dist, idx = tree.query(nan_idx, k=1)
+        filled = arr.copy()
+        # idx shape: (n_nan, 1)，valid_idx shape: (n_valid, ndim)
+        # valid_idx[idx].T shape: (ndim, n_nan)
+        filled[tuple(nan_idx.T)] = arr[tuple(valid_idx[idx.flatten()].T)]
+        return filled
+    else:
+        from scipy.ndimage import distance_transform_edt
+
+        idx = distance_transform_edt(mask, return_distances=False, return_indices=True)
+        return arr[tuple(idx)]
+
+
+def _data_clip(data, data_min, data_max):
+    """
+    对data进行范围裁剪，超出范围的点设为nan，并用fill_nan_nearest填充，最后再次填充极端nan。
+    """
+    arr = np.asarray(data)
+    mask = np.isnan(arr) | (arr < data_min) | (arr > data_max)
+    if np.any(mask):
+        arr = np.where(mask, np.nan, arr)
+        arr = _fill_nan_nearest(arr)
+    if np.any(np.isnan(arr)):
+        arr = _fill_nan_nearest(arr)
+    return arr
 
 
 def _interp_single_worker(*args):
@@ -39,12 +67,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 +136,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,147 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Author: Liu Kun && 16031215@qq.com
+Date: 2025-04-26 11:54:21
+LastEditors: Liu Kun && 16031215@qq.com
+LastEditTime: 2025-04-26 11:54:22
+FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\_script\\data_interp_geo.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
+from scipy.interpolate import RectBivariateSpline
+
+from oafuncs.oa_tool import PEx
+
+_has_pykdtree = importlib.util.find_spec("pykdtree.kdtree") is not None
+
+
+def fill_nan_nearest(arr):
+    """用最近邻插值填充 NaN，优先用pykdtree加速"""
+    mask = np.isnan(arr)
+    if not mask.any():
+        return arr
+    if _has_pykdtree:
+        from pykdtree.kdtree import KDTree
+
+        valid_idx = np.array(np.where(~mask)).T
+        nan_idx = np.array(np.where(mask)).T
+        if len(valid_idx) == 0:
+            # 全是nan，直接返回
+            return arr
+        tree = KDTree(valid_idx)
+        dist, idx = tree.query(nan_idx, k=1)
+        filled = arr.copy()
+        # idx shape: (n_nan, 1)，valid_idx shape: (n_valid, ndim)
+        # valid_idx[idx].T shape: (ndim, n_nan)
+        filled[tuple(nan_idx.T)] = arr[tuple(valid_idx[idx.flatten()].T)]
+        return filled
+    else:
+        from scipy.ndimage import distance_transform_edt
+
+        idx = distance_transform_edt(mask, return_distances=False, return_indices=True)
+        return arr[tuple(idx)]
+
+
+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)
+    mask = np.isnan(arr) | (arr < data_min) | (arr > data_max)
+    if np.any(mask):
+        arr = np.where(mask, np.nan, arr)
+        arr = fill_nan_nearest(arr)
+    # 最后再填充nan（极端情况）
+    if np.any(np.isnan(arr)):
+        arr = fill_nan_nearest(arr)
+    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.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: oafuncs
-Version: 0.0.98.15
+Version: 0.0.98.16
 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.16.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=W4kYZmkZxWJqKfZVdQOd8jUv0y11t6k8tsukRQt670I,6373
+oafuncs/_script/data_interp_geo.py,sha256=Fv-l8MeKhx6e7UVtzk2b6Pb2Cfcg-RNnMFmkcd0ng64,5296
 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.16.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
+oafuncs-0.0.98.16.dist-info/METADATA,sha256=yPzl7v8ASdNRSjiXasAGuxPKA9rM_zCIQBQHfmz4I0k,4273
+oafuncs-0.0.98.16.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+oafuncs-0.0.98.16.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
+oafuncs-0.0.98.16.dist-info/RECORD,,