oafuncs 0.0.98.19__py3-none-any.whl → 0.0.98.21__py3-none-any.whl

oafuncs/_script/data_interp.py

@@ -1,3 +1,18 @@
+#!/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-26 19:21:31
+FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\_script\\data_interp.py
+Description:
+EditPlatform: vscode
+ComputerInfo: XPS 15 9510
+SystemInfo: Windows 11
+Python Version: 3.12
+"""
+
 from typing import List, Union
 
 import numpy as np
@@ -8,10 +23,10 @@ from oafuncs.oa_tool import PEx
 
 def _interp_single_worker(*args):
     """
-    用于PEx并行的单slice插值worker，参数为(t, z, source_data, origin_points, target_points, interpolation_method, target_shape)
+    用于PEx并行的单slice插值worker。
+    参数: data_slice, origin_points, target_points, interpolation_method, target_shape
     """
     data_slice, origin_points, target_points, interpolation_method, target_shape = args
-
     # 过滤掉包含 NaN 的点
     valid_mask = ~np.isnan(data_slice.ravel())
     valid_data = data_slice.ravel()[valid_mask]
@@ -21,20 +36,16 @@ def _interp_single_worker(*args):
         return np.full(target_shape, np.nanmean(data_slice))
 
     # 使用有效数据进行插值
-    result = griddata(valid_points, valid_data, target_points, method=interpolation_method)
-    result = result.reshape(target_shape)
+    result = griddata(valid_points, valid_data, target_points, method=interpolation_method).reshape(target_shape)
+    # 对仍为 NaN 的点用最近邻填充
+    if np.isnan(result).any():
+        nn = griddata(valid_points, valid_data, target_points, method="nearest").reshape(target_shape)
+        result[np.isnan(result)] = nn[np.isnan(result)]
 
     return result
 
 
-def interp_2d_func(
-    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:
+def interp_2d_func(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:
     """
     Perform 2D interpolation on the last two dimensions of a multi-dimensional array.
 
@@ -46,7 +57,6 @@ def interp_2d_func(
         source_data (np.ndarray): Multi-dimensional array with the last two dimensions as spatial.
         interpolation_method (str, optional): Interpolation method. Defaults to "cubic".
             >>> optional: 'linear', 'nearest', 'cubic', 'quintic', etc.
-        use_parallel (bool, optional): Enable parallel processing. Defaults to True.
 
     Returns:
         np.ndarray: Interpolated data array.
@@ -60,7 +70,7 @@ def interp_2d_func(
         >>> source_x_coordinates = np.array([7, 8, 9])
         >>> source_y_coordinates = np.array([10, 11, 12])
         >>> source_data = np.random.rand(3, 3)
-        >>> result = interp_2d(target_x_coordinates, target_y_coordinates, source_x_coordinates, source_y_coordinates, source_data)
+        >>> result = interp_2d_func(target_x_coordinates, target_y_coordinates, source_x_coordinates, source_y_coordinates, source_data)
         >>> print(result.shape)  # Expected output: (3, 3)
     """
     if len(target_y_coordinates.shape) == 1:
@@ -80,7 +90,7 @@ def interp_2d_func(
         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]")
+        raise ValueError(f"Source data has {data_dims} dimensions, but this function currently supports only up to 4.")
 
     # Reshape to 4D by adding leading dimensions of size 1 if needed
     num_dims_to_add = 4 - data_dims

oafuncs/_script/data_interp_geo.py

@@ -1,98 +1,167 @@
 from typing import List, Union
 
 import numpy as np
-from scipy.interpolate import RectBivariateSpline
+from scipy.interpolate import NearestNDInterpolator, griddata
 
 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):
+def _normalize_lon(lon, ref_lon):
     """
-    单slice插值worker，参数为(data_slice, sx, sy, tx, ty, interpolation_method, data_min, data_max)
+    将经度数组 lon 归一化到与 ref_lon 相同的区间（[-180,180] 或 [0,360]）
+    并在经度分界（如180/-180, 0/360）附近自动拓宽，避免插值断裂。
     """
-    # 兼容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
+    lon = np.asarray(lon)
+    ref_lon = np.asarray(ref_lon)
+    if np.nanmax(ref_lon) > 180:
+        lon = np.where(lon < 0, lon + 360, lon)
     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:
+        lon = np.where(lon > 180, lon - 360, lon)
+    return lon
+
+
+def _expand_lonlat_for_dateline(points, values):
+    """
+    对经度分界（如180/-180, 0/360）附近的数据进行拓宽，避免插值断裂。
+    points: (N,2) [lon,lat]
+    values: (N,)
+    返回拓宽后的 points, values
+    """
+    lon = points[:, 0]
+    lat = points[:, 1]
+    expanded_points = [points]
+    expanded_values = [values]
+    if (np.nanmax(lon) > 170) and (np.nanmin(lon) < -170):
+        expanded_points.append(np.column_stack((lon + 360, lat)))
+        expanded_points.append(np.column_stack((lon - 360, lat)))
+        expanded_values.append(values)
+        expanded_values.append(values)
+    if (np.nanmax(lon) > 350) and (np.nanmin(lon) < 10):
+        expanded_points.append(np.column_stack((lon - 360, lat)))
+        expanded_points.append(np.column_stack((lon + 360, lat)))
+        expanded_values.append(values)
+        expanded_values.append(values)
+    points_new = np.vstack(expanded_points)
+    values_new = np.concatenate(expanded_values)
+    return points_new, values_new
+
+
+def _interp_single_worker(*args):
     """
-    更平滑的二维插值，采用RectBivariateSpline实现bicubic效果，接口与interp_2d兼容。
-    支持输入2D/3D/4D数据，最后两维为空间。
-    interpolation_method: "cubic"（默认，bicubic），"linear"（双线性）
-    插值后自动裁剪并用最近邻填充超限和NaN，范围取原始数据的nanmin/nanmax
+    用于PEx并行的单slice插值worker。
+    参数: data_slice, origin_points, target_points, interpolation_method, target_shape
+    球面插值：经纬度转球面坐标后插值
     """
-    # 保证输入为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
+    data_slice, origin_points, target_points, interpolation_method, target_shape = args
+
+    # 经纬度归一化
+    origin_points = origin_points.copy()
+    target_points = target_points.copy()
+    origin_points[:, 0] = _normalize_lon(origin_points[:, 0], target_points[:, 0])
+    target_points[:, 0] = _normalize_lon(target_points[:, 0], origin_points[:, 0])
+
+    def lonlat2xyz(lon, lat):
+        lon_rad = np.deg2rad(lon)
+        lat_rad = np.deg2rad(lat)
+        x = np.cos(lat_rad) * np.cos(lon_rad)
+        y = np.cos(lat_rad) * np.sin(lon_rad)
+        z = np.sin(lat_rad)
+        return np.stack([x, y, z], axis=-1)
+
+    # 过滤掉包含 NaN 的点
+    valid_mask = ~np.isnan(data_slice.ravel())
+    valid_data = data_slice.ravel()[valid_mask]
+    valid_points = origin_points[valid_mask]
+
+    if len(valid_data) < 10:
+        return np.full(target_shape, np.nanmean(data_slice))
+
+    # 拓宽经度分界，避免如179/-181插值断裂
+    valid_points_exp, valid_data_exp = _expand_lonlat_for_dateline(valid_points, valid_data)
+
+    valid_xyz = lonlat2xyz(valid_points_exp[:, 0], valid_points_exp[:, 1])
+    target_xyz = lonlat2xyz(target_points[:, 0], target_points[:, 1])
+
+    # 使用 griddata 的 cubic 插值以获得更好平滑效果
+    result = griddata(valid_xyz, valid_data_exp, target_xyz, method=interpolation_method).reshape(target_shape)
+
+    # 用最近邻处理残余 NaN
+    if np.isnan(result).any():
+        nn_interp = NearestNDInterpolator(valid_xyz, valid_data_exp)
+        nn = nn_interp(target_xyz).reshape(target_shape)
+        result[np.isnan(result)] = nn[np.isnan(result)]
+
+    return result
+
+
+def interp_2d_func_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:
+    """
+    Perform 2D interpolation on the last two dimensions of a multi-dimensional array (spherical coordinates).
+    使用球面坐标系进行插值，适用于全球尺度的地理数据，能正确处理经度跨越日期线的情况。
+
+    Args:
+        target_x_coordinates (Union[np.ndarray, List[float]]): Target grid's longitude (-180 to 180 or 0 to 360).
+        target_y_coordinates (Union[np.ndarray, List[float]]): Target grid's latitude (-90 to 90).
+        source_x_coordinates (Union[np.ndarray, List[float]]): Original grid's longitude (-180 to 180 or 0 to 360).
+        source_y_coordinates (Union[np.ndarray, List[float]]): Original grid's latitude (-90 to 90).
+        source_data (np.ndarray): Multi-dimensional array with the last two dimensions as spatial.
+        interpolation_method (str, optional): Interpolation method. Defaults to "cubic".
+            >>> optional: 'linear', 'nearest', 'cubic', 'quintic', etc.
+
+    Returns:
+        np.ndarray: Interpolated data array.
+
+    Raises:
+        ValueError: If input shapes are invalid.
+
+    Examples:
+        >>> # 创建一个全球网格示例
+        >>> target_lon = np.arange(-180, 181, 1)  # 1度分辨率目标网格
+        >>> target_lat = np.arange(-90, 91, 1)
+        >>> source_lon = np.arange(-180, 181, 5)  # 5度分辨率源网格
+        >>> source_lat = np.arange(-90, 91, 5)
+        >>> # 创建一个简单的数据场 (例如温度场)
+        >>> source_data = np.cos(np.deg2rad(source_lat.reshape(-1, 1))) * np.cos(np.deg2rad(source_lon))
+        >>> # 插值到高分辨率网格
+        >>> result = interp_2d_geo(target_lon, target_lat, source_lon, source_lat, source_data)
+        >>> print(result.shape)  # Expected output: (181, 361)
+    """
+    # 验证输入数据范围
+    if np.nanmin(target_y_coordinates) < -90 or np.nanmax(target_y_coordinates) > 90:
+        raise ValueError("[red]Target latitude must be in range [-90, 90].[/red]")
+    if np.nanmin(source_y_coordinates) < -90 or np.nanmax(source_y_coordinates) > 90:
+        raise ValueError("[red]Source latitude must be in range [-90, 90].[/red]")
+
+    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:
+        source_x_coordinates, source_y_coordinates = np.meshgrid(source_x_coordinates, source_y_coordinates)
+
+    if source_x_coordinates.shape != source_data.shape[-2:] or source_y_coordinates.shape != source_data.shape[-2:]:
+        raise ValueError("[red]Shape of source_data does not match shape of source_x_coordinates or source_y_coordinates.[/red]")
+
+    target_points = np.column_stack((np.array(target_x_coordinates).ravel(), np.array(target_y_coordinates).ravel()))
+    origin_points = np.column_stack((np.array(source_x_coordinates).ravel(), np.array(source_y_coordinates).ravel()))
+
+    data_dims = len(source_data.shape)
     if data_dims < 2:
-        raise ValueError("Source data must have at least 2 dimensions.")
+        raise ValueError(f"[red]Source data must have at least 2 dimensions, but got {data_dims}.[/red]")
     elif data_dims > 4:
-        raise ValueError("Source data has more than 4 dimensions, not supported.")
+        raise ValueError(f"Source data has {data_dims} dimensions, but this function currently supports only up to 4.")
 
     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
+    new_shape = (1,) * num_dims_to_add + source_data.shape
+    new_src_data = source_data.reshape(new_shape)
 
-    data_min, data_max = np.nanmin(data), np.nanmax(data)
-    target_shape = ty.shape
+    t, z, y, x = new_src_data.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))
+    target_shape = target_y_coordinates.shape
+    for t_index in range(t):
+        for z_index in range(z):
+            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, params)
 
-    result = np.array(result).reshape(t, z, *target_shape)
-    result = np.squeeze(result)
-    return result
+    return np.squeeze(np.array(result).reshape(t, z, *target_shape))

oafuncs/_script/netcdf_merge.py

@@ -1,26 +1,10 @@
-#!/usr/bin/env python
-# coding=utf-8
-"""
-Author: Liu Kun && 16031215@qq.com
-Date: 2025-03-30 11:16:29
-LastEditors: Liu Kun && 16031215@qq.com
-LastEditTime: 2025-04-25 14:23:10
-FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\_script\\netcdf_merge.py
-Description
-EditPlatform: vscode
-ComputerInfo: XPS 15 9510
-SystemInfo: Windows 11
-Python Version: 3.12
-"""
-
+import logging
 import os
 from typing import List, Optional, Union
 
-import numpy as np
 import xarray as xr
 
 from oafuncs import pbar
-import logging
 
 
 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:
@@ -41,7 +25,6 @@ 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')
     """
-    from oafuncs._script.netcdf_write import save_to_nc
 
     if target_filename is None:
         target_filename = "merged.nc"
@@ -68,33 +51,68 @@ def merge_nc(file_list: Union[str, List[str]], var_name: Optional[Union[str, Lis
     # 初始化合并数据字典
     merged_data = {}
 
-    for i, file in pbar(enumerate(file_list), description="Reading files", total=len(file_list)):
+    for i, file in pbar(enumerate(file_list), "Reading files", total=len(file_list)):
         with xr.open_dataset(file) as ds:
             for var in var_names:
                 data_var = ds[var]
                 if dim_name in data_var.dims:
                     merged_data.setdefault(var, []).append(data_var)
                 elif var not in merged_data:
-                    # 判断类型，时间类型用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"):
+                    # 只负责合并，不做NaN填充，统一交由 netcdf_write.py 处理
+                    merged_data[var] = data_var
+
+    # 记录变量的填充值和缺失值信息，确保不会丢失
+    fill_values = {}
+    missing_values = {}
+    for var_name, var_data in merged_data.items():
+        if isinstance(var_data, list) and var_data:
+            # 如果是要合并的变量，检查第一个元素的属性
+            attrs = var_data[0].attrs
+            if "_FillValue" in attrs:
+                fill_values[var_name] = attrs["_FillValue"]
+            if "missing_value" in attrs:
+                missing_values[var_name] = attrs["missing_value"]
+        else:
+            # 如果是单个变量，直接检查属性
+            attrs = var_data.attrs if hasattr(var_data, "attrs") else {}
+            if "_FillValue" in attrs:
+                fill_values[var_name] = attrs["_FillValue"]
+            if "missing_value" in attrs:
+                missing_values[var_name] = attrs["missing_value"]
+
+    for var in pbar(merged_data, "Merging variables"):
         if isinstance(merged_data[var], list):
-            # 判断类型，时间类型用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)
+            # 使用compat='override'确保合并时属性不会冲突
+            merged_data[var] = xr.concat(merged_data[var], dim=dim_name, compat="override")
+            # 恢复原始填充值和缺失值属性
+            if var in fill_values:
+                merged_data[var].attrs["_FillValue"] = fill_values[var]
+            if var in missing_values:
+                merged_data[var].attrs["missing_value"] = missing_values[var]
+
+    # 合并后构建 Dataset，此时 merged_data 只包含数据变量，不包含坐标变量
+    merged_ds = xr.Dataset(merged_data)
+
+    # 自动补充坐标变量（如 time、lat、lon 等），以第一个文件为准
+    with xr.open_dataset(file_list[0]) as ds0:
+        for coord in ds0.coords:
+            # 保证坐标变量不会被覆盖，且数据类型和属性保持一致
+            if coord not in merged_ds.coords:
+                merged_ds = merged_ds.assign_coords({coord: ds0[coord]})
+
+    # 如果合并维度是坐标，检查所有文件的该坐标是否一致
+    if dim_name in merged_ds.coords and len(file_list) > 1:
+        logging.info(f"验证合并维度 {dim_name} 的一致性...")
+        for file in file_list[1:]:
+            with xr.open_dataset(file) as ds:
+                if dim_name in ds.coords and not ds[dim_name].equals(merged_ds[dim_name]):
+                    logging.warning(f"文件 {file} 的 {dim_name} 坐标与合并后的数据不一致，可能导致数据失真")
 
     if os.path.exists(target_filename):
-        # print("Warning: The target file already exists. Removing it ...")
         logging.warning("The target file already exists. Removing it ...")
         os.remove(target_filename)
 
-    save_to_nc(target_filename, xr.Dataset(merged_data))
+    merged_ds.to_netcdf(target_filename,mode='w')
 
 
 # Example usage