oafuncs 0.0.98.41__tar.gz → 0.0.98.43__tar.gz

{oafuncs-0.0.98.41/oafuncs.egg-info → oafuncs-0.0.98.43}/PKG-INFO

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

{oafuncs-0.0.98.41 → oafuncs-0.0.98.43}/oafuncs/_script/netcdf_write.py

@@ -374,75 +374,6 @@ def save_to_nc(file, data, varname=None, coords=None, mode="w", convert_dtype='i
         raise RuntimeError(f"netCDF4 保存失败: {str(e)}") from e
 
 
-def _compress_netcdf(src_path, dst_path=None, tolerance=1e-10, preserve_mask_values=True):
-    """
-    压缩 NetCDF 文件，使用 scale_factor/add_offset 压缩数据。
-    若 dst_path 省略，则自动生成新文件名，写出后删除原文件并将新文件改回原名。
-    压缩后验证数据是否失真。
-
-    参数：
-      - src_path: 原始 NetCDF 文件路径
-      - dst_path: 压缩后的文件路径（可选）
-      - tolerance: 数据验证的允许误差范围（默认 1e-10）
-      - preserve_mask_values: 是否保留掩码区域的原始值（True）或将其替换为缺省值（False）
-    """
-    # 判断是否要替换原文件
-    delete_orig = dst_path is None
-    if delete_orig:
-        dst_path = src_path.replace(".nc", "_compress.nc")
-    # 打开原始文件并保存压缩文件
-    ds = xr.open_dataset(src_path)
-    save_to_nc(dst_path, ds, convert_dtype='int32',scale_offset_switch=True, compile_switch=True, preserve_mask_values=preserve_mask_values)
-    ds.close()
-
-    # 验证压缩后的数据是否失真
-    original_ds = xr.open_dataset(src_path)
-    compressed_ds = xr.open_dataset(dst_path)
-    # 更详细地验证数据
-    for var in original_ds.data_vars:
-        original_data = original_ds[var].values
-        compressed_data = compressed_ds[var].values
-        # 跳过非数值类型变量
-        if not np.issubdtype(original_data.dtype, np.number):
-            continue
-        # 获取掩码（如果存在）
-        original_mask = None
-        if hasattr(original_data, "mask") and np.ma.is_masked(original_data):  # 修正：确保是有效的掩码数组
-            original_mask = original_data.mask.copy()
-        # 检查有效数据是否在允许误差范围内
-        valid_mask = np.isfinite(original_data)
-        if original_mask is not None:
-            valid_mask &= ~original_mask
-        if np.any(valid_mask):
-            if np.issubdtype(original_data.dtype, np.floating):
-                diff = np.abs(original_data[valid_mask] - compressed_data[valid_mask])
-                max_diff = np.max(diff)
-                if max_diff > tolerance:
-                    print(f"警告: 变量 {var} 的压缩误差 {max_diff} 超出容许范围 {tolerance}")
-                    if max_diff > tolerance * 10:  # 严重偏差时抛出错误
-                        raise ValueError(f"变量 {var} 的数据在压缩后严重失真 (max_diff={max_diff})")
-            elif np.issubdtype(original_data.dtype, np.integer):
-                # 整数类型应该完全相等
-                if not np.array_equal(original_data[valid_mask], compressed_data[valid_mask]):
-                    raise ValueError(f"变量 {var} 的整数数据在压缩后不一致")
-        # 如果需要保留掩码区域值，检查掩码区域的值
-        if preserve_mask_values and original_mask is not None and np.any(original_mask):
-            # 确保掩码区域的原始值被正确保留
-            # 修正：掩码数组可能存在数据类型不匹配问题，添加安全检查
-            try:
-                mask_diff = np.abs(original_data[original_mask] - compressed_data[original_mask])
-                if np.any(mask_diff > tolerance):
-                    print(f"警告: 变量 {var} 的掩码区域数据在压缩后发生变化")
-            except Exception as e:
-                print(f"警告: 变量 {var} 的掩码区域数据比较失败: {str(e)}")
-    original_ds.close()
-    compressed_ds.close()
-
-    # 替换原文件
-    if delete_orig:
-        os.remove(src_path)
-        os.rename(dst_path, src_path)
-
 
 # 测试用例
 if __name__ == "__main__":

{oafuncs-0.0.98.41 → oafuncs-0.0.98.43}/oafuncs/_script/plot_dataset.py

@@ -197,12 +197,12 @@ def select_colormap_and_levels(data_range: Tuple[float, float], plot_type: str)
         num_levels = 128
 
     if data_range[0] * data_range[1] < 0:
-        cmap = oafuncs.oa_cmap.get("diverging_1")
+        cmap = oafuncs.oa_cmap.get(diverging_cmap)
         bdy = max(abs(data_range[0]), abs(data_range[1]))
         norm = mpl.colors.TwoSlopeNorm(vmin=-bdy, vcenter=0, vmax=bdy)
         levels = np.linspace(-bdy, bdy, num_levels)
     else:
-        cmap = oafuncs.oa_cmap.get("cool_1") if data_range[0] < 0 else oafuncs.oa_cmap.get("warm_1")
+        cmap = oafuncs.oa_cmap.get(negative_cmap) if data_range[0] < 0 else oafuncs.oa_cmap.get(positive_cmap)
         norm = mpl.colors.Normalize(vmin=data_range[0], vmax=data_range[1])
         levels = np.linspace(data_range[0], data_range[1], num_levels)
 
@@ -320,9 +320,14 @@ def get_xyzt_names(ds_in, xyzt_dims):
     return x_dim, y_dim, z_dim, t_dim
 
 
-def func_plot_dataset(ds_in: Union[xr.Dataset, xr.DataArray], output_dir: str, xyzt_dims: Tuple[str, str, str, str] = None, plot_type: str = "contourf", fixed_colorscale: bool = False) -> None:
+def func_plot_dataset(ds_in: Union[xr.Dataset, xr.DataArray], output_dir: str, cmap='diverging_3', pcmap='warm_3', ncmap='cool_3', xyzt_dims: Tuple[str, str, str, str] = None, plot_type: str = "contourf", fixed_colorscale: bool = False) -> None:
     """Plot variables from a NetCDF file and save the plots to the specified directory."""
     os.makedirs(output_dir, exist_ok=True)
+    
+    global diverging_cmap, positive_cmap, negative_cmap
+    diverging_cmap = cmap
+    positive_cmap = pcmap
+    negative_cmap = ncmap
 
     # Main processing function
     try:

{oafuncs-0.0.98.41 → oafuncs-0.0.98.43}/oafuncs/oa_data.py

@@ -1,11 +1,11 @@
-from typing import Any, List, Union
+from typing import Any, List, Union, Literal
 
 import numpy as np
 import xarray as xr
 from rich import print
 
 
-__all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile"]
+__all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile", "mask_land_ocean"]
 
 
 def ensure_list(input_value: Any) -> List[str]:
@@ -188,5 +188,83 @@ def mask_shapefile(
         return None
 
 
+
+def _normalize_lon(lon: np.ndarray) -> np.ndarray:
+    """将经度转换到 [-180, 180)。"""
+    lon = np.asarray(lon, dtype=float)
+    return np.where(lon >= 180, lon - 360, lon)
+
+
+def _land_sea_mask(
+    lon: np.ndarray,
+    lat: np.ndarray,
+    keep: Literal["land", "ocean"],
+) -> np.ndarray:
+    """
+    根据 1-D 或 2-D 经纬度返回布尔掩膜。
+    True 表示该位置 *保留*，False 表示该位置将被掩掉。
+    """
+    from global_land_mask import globe
+    
+    lon = _normalize_lon(lon)
+    lat = np.asarray(lat, dtype=float)
+
+    # 如果输入是 1-D，则网格化；2-D 则直接使用
+    if lon.ndim == 1 and lat.ndim == 1:
+        lon_2d, lat_2d = np.meshgrid(lon, lat)
+    elif lon.ndim == 2 and lat.ndim == 2:
+        lon_2d, lat_2d = lon, lat
+    else:
+        raise ValueError("经纬度必须是同维度的 1-D 或 2-D 数组")
+
+    is_ocean = globe.is_ocean(lat_2d, lon_2d)
+
+    if keep == "land":
+        mask = ~is_ocean
+    elif keep == "ocean":
+        mask = is_ocean
+    else:
+        raise ValueError("keep 只能是 'land' 或 'ocean'")
+
+    return mask
+
+
+def mask_land_ocean(
+    data: xr.DataArray | xr.Dataset,
+    lon: np.ndarray,
+    lat: np.ndarray,
+    *,  # 强制关键字参数
+    keep: Literal["land", "ocean"] = "land",
+) -> xr.DataArray | xr.Dataset:
+    """
+    根据海陆分布掩膜 xarray 对象。
+
+    Parameters
+    ----------
+    data : xr.DataArray 或 xr.Dataset
+        至少包含 'lat' 和 'lon' 维度/坐标的数组。
+    lon : array_like
+        经度，可以是 1-D 或 2-D。
+    lat : array_like
+        纬度，可以是 1-D 或 2-D。
+    keep : {'land', 'ocean'}, optional
+        指定要保留的部分，默认为 'land'。
+
+    Returns
+    -------
+    掩膜后的 xr.DataArray / xr.Dataset
+    """
+    mask = _land_sea_mask(lon, lat, keep)
+
+    # 用 apply_ufunc 自动对齐并广播掩膜
+    return xr.apply_ufunc(
+        lambda x, m: x.where(m),
+        data,
+        xr.DataArray(mask, dims=("lat", "lon")),
+        dask="parallelized",
+        keep_attrs=True,
+    )
+
+
 if __name__ == "__main__":
     pass

{oafuncs-0.0.98.41 → oafuncs-0.0.98.43}/oafuncs/oa_nc.py

@@ -270,6 +270,9 @@ def draw(
     output_directory: Optional[str] = None,
     dataset: Optional[xr.Dataset] = None,
     file_path: Optional[str] = None,
+    cmap='diverging_3',
+    pcmap='warm_3',
+    ncmap='cool_3',
     dims_xyzt: Union[List[str], Tuple[str, str, str, str]] = None,
     plot_style: str = "contourf",
     use_fixed_colorscale: bool = False,
@@ -296,18 +299,18 @@ def draw(
         raise ValueError("dimensions must be a list or tuple")
 
     if dataset is not None:
-        func_plot_dataset(dataset, output_directory, dims_xyzt, plot_style, use_fixed_colorscale)
+        func_plot_dataset(dataset, output_directory, cmap, pcmap, ncmap, dims_xyzt, plot_style, use_fixed_colorscale)
     elif file_path is not None:
         if check(file_path):
             ds = xr.open_dataset(file_path)
-            func_plot_dataset(ds, output_directory, dims_xyzt, plot_style, use_fixed_colorscale)
+            func_plot_dataset(ds, output_directory, cmap, pcmap, ncmap, dims_xyzt, plot_style, use_fixed_colorscale)
         else:
             print(f"[red]Invalid file: {file_path}[/red]")
     else:
         print("[red]No dataset or file provided.[/red]")
 
 
-def compress(src_path, dst_path=None,convert_dtype='int16'):
+def compress(src_path, dst_path=None,convert_dtype='int32'):
     """
     压缩 NetCDF 文件，使用 scale_factor/add_offset 压缩数据。
     若 dst_path 省略，则自动生成新文件名，写出后删除原文件并将新文件改回原名。

{oafuncs-0.0.98.41 → oafuncs-0.0.98.43/oafuncs.egg-info}/PKG-INFO

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

{oafuncs-0.0.98.41 → oafuncs-0.0.98.43}/setup.py

@@ -18,7 +18,7 @@ URL = "https://github.com/Industry-Pays/OAFuncs"
 EMAIL = "liukun0312@stu.ouc.edu.cn"
 AUTHOR = "Kun Liu"
 REQUIRES_PYTHON = ">=3.10.0"  # 2025/03/13
-VERSION = "0.0.98.41"
+VERSION = "0.0.98.43"
 
 # What packages are required for this module to be executed?
 REQUIRED = [