oafuncs 0.0.98.19__tar.gz → 0.0.98.21__tar.gz

{oafuncs-0.0.98.19/oafuncs.egg-info → oafuncs-0.0.98.21}/PKG-INFO

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

{oafuncs-0.0.98.19 → oafuncs-0.0.98.21}/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-0.0.98.21/oafuncs/_script/data_interp_geo.py

@@ -0,0 +1,167 @@
+from typing import List, Union
+
+import numpy as np
+from scipy.interpolate import NearestNDInterpolator, griddata
+
+from oafuncs.oa_tool import PEx
+
+
+def _normalize_lon(lon, ref_lon):
+    """
+    将经度数组 lon 归一化到与 ref_lon 相同的区间（[-180,180] 或 [0,360]）
+    并在经度分界（如180/-180, 0/360）附近自动拓宽，避免插值断裂。
+    """
+    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:
+        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):
+    """
+    用于PEx并行的单slice插值worker。
+    参数: data_slice, origin_points, target_points, interpolation_method, target_shape
+    球面插值：经纬度转球面坐标后插值
+    """
+    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(f"[red]Source data must have at least 2 dimensions, but got {data_dims}.[/red]")
+    elif data_dims > 4:
+        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 + source_data.shape
+    new_src_data = source_data.reshape(new_shape)
+
+    t, z, y, x = new_src_data.shape
+
+    params = []
+    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)
+
+    return np.squeeze(np.array(result).reshape(t, z, *target_shape))

oafuncs-0.0.98.21/oafuncs/_script/netcdf_merge.py

@@ -0,0 +1,122 @@
+import logging
+import os
+from typing import List, Optional, Union
+
+import xarray as xr
+
+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:
+        Merge variables from multiple NetCDF files along a specified dimension and write to a new file.
+        If var_name is a string, it is considered a single variable; if it is a list and has only one element, it is also a single variable;
+        If the list has more than one element, it is a multi-variable; if var_name is None, all variables are merged.
+
+    Parameters:
+        file_list: List of NetCDF file paths or a single file path as a string
+        var_name: Name of the variable to be extracted or a list of variable names, default is None, which means all variables are extracted
+        dim_name: Dimension name used for merging
+        target_filename: Target file name after merging
+
+    Example:
+        merge(file_list, var_name='u', dim_name='time', target_filename='merged.nc')
+        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"
+
+    # 确保目标路径存在
+    target_dir = os.path.dirname(target_filename)
+    if target_dir and not os.path.exists(target_dir):
+        os.makedirs(target_dir)
+
+    if isinstance(file_list, str):
+        file_list = [file_list]
+
+    # 初始化变量名列表
+    if var_name is None:
+        with xr.open_dataset(file_list[0]) as ds:
+            var_names = list(ds.variables.keys())
+    elif isinstance(var_name, str):
+        var_names = [var_name]
+    elif isinstance(var_name, list):
+        var_names = var_name
+    else:
+        raise ValueError("var_name must be a string, a list of strings, or None")
+
+    # 初始化合并数据字典
+    merged_data = {}
+
+    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:
+                    # 只负责合并，不做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):
+            # 使用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):
+        logging.warning("The target file already exists. Removing it ...")
+        os.remove(target_filename)
+
+    merged_ds.to_netcdf(target_filename,mode='w')
+
+
+# Example usage
+if __name__ == "__main__":
+    files_to_merge = ["file1.nc", "file2.nc", "file3.nc"]
+    output_path = "merged_output.nc"
+    merge_nc(files_to_merge, var_name=None, dim_name="time", target_filename=output_path)