oafuncs 0.0.98.0__tar.gz → 0.0.98.2__tar.gz

{oafuncs-0.0.98.0/oafuncs.egg-info → oafuncs-0.0.98.2}/PKG-INFO

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

oafuncs-0.0.98.2/oafuncs/_script/netcdf_write.py

@@ -0,0 +1,116 @@
+import os
+import netCDF4 as nc
+import numpy as np
+import xarray as xr
+
+
+def _numpy_to_nc_type(numpy_type):
+    """将NumPy数据类型映射到NetCDF数据类型"""
+    numpy_to_nc = {
+        "float32": "f4",
+        "float64": "f8",
+        "int8": "i1",
+        "int16": "i2",
+        "int32": "i4",
+        "int64": "i8",
+        "uint8": "u1",
+        "uint16": "u2",
+        "uint32": "u4",
+        "uint64": "u8",
+    }
+    numpy_type_str = str(numpy_type) if not isinstance(numpy_type, str) else numpy_type
+    return numpy_to_nc.get(numpy_type_str, "f4")
+
+
+def _calculate_scale_and_offset(data, n=16):
+    if not isinstance(data, np.ndarray):
+        raise ValueError("Input data must be a NumPy array.")
+
+    data_min = np.nanmin(data)
+    data_max = np.nanmax(data)
+
+    if np.isnan(data_min) or np.isnan(data_max):
+        raise ValueError("Input data contains NaN values.")
+
+    scale_factor = (data_max - data_min) / (2**n - 1)
+    add_offset = data_min + 2 ** (n - 1) * scale_factor
+    return scale_factor, add_offset
+
+
+def save_to_nc(file, data, varname=None, coords=None, mode="w", scale_offset_switch=True, compile_switch=True):
+    # 处理xarray对象的情况（当varname和coords都为None时）
+    if varname is None and coords is None:
+        if not isinstance(data, (xr.DataArray, xr.Dataset)):
+            raise ValueError("When varname and coords are not provided, data must be an xarray object")
+
+        encoding = {}
+        if isinstance(data, xr.DataArray):
+            if data.name is None:
+                data = data.rename("data")
+            varname = data.name
+            encoding[varname] = {"zlib": compile_switch, "complevel": 4}
+            if scale_offset_switch:
+                scale, offset = _calculate_scale_and_offset(data.values)
+                encoding[varname].update({"dtype": "int16", "scale_factor": scale, "add_offset": offset, "_FillValue": -32767})
+            else:
+                encoding[varname].update({"dtype": "float32", "_FillValue": np.nan})
+        else:  # Dataset情况
+            for var in data.data_vars:
+                encoding[var] = {"zlib": compile_switch, "complevel": 4}
+                if scale_offset_switch:
+                    scale, offset = _calculate_scale_and_offset(data[var].values)
+                    encoding[var].update({"dtype": "int16", "scale_factor": scale, "add_offset": offset, "_FillValue": -32767})
+                else:
+                    encoding[var].update({"dtype": "float32", "_FillValue": np.nan})
+
+        try:
+            data.to_netcdf(file, mode=mode, encoding=encoding)
+            return
+        except Exception as e:
+            raise RuntimeError(f"Failed to save xarray object: {str(e)}") from e
+
+    # 处理普通numpy数组的情况
+    if mode == "w" and os.path.exists(file):
+        os.remove(file)
+    elif mode == "a" and not os.path.exists(file):
+        mode = "w"
+
+    try:
+        with nc.Dataset(file, mode, format="NETCDF4") as ncfile:
+            # 创建维度并写入坐标
+            if coords is not None:
+                for dim, values in coords.items():
+                    if dim not in ncfile.dimensions:
+                        ncfile.createDimension(dim, len(values))
+                        var = ncfile.createVariable(dim, _numpy_to_nc_type(values.dtype), (dim,))
+                        var[:] = values
+
+            # 创建变量
+            dims = list(coords.keys()) if coords else []
+            if scale_offset_switch:
+                scale, offset = _calculate_scale_and_offset(data)
+                var = ncfile.createVariable(varname, "i2", dims, fill_value=-32767, zlib=compile_switch)
+                var.scale_factor = scale
+                var.add_offset = offset
+            else:
+                dtype = _numpy_to_nc_type(data.dtype)
+                var = ncfile.createVariable(varname, dtype, dims, zlib=compile_switch)
+
+            var[:] = data
+    except Exception as e:
+        raise RuntimeError(f"Failed to save netCDF4 file: {str(e)}") from e
+
+
+if __name__ == "__main__":
+    # Example usage
+    data = xr.open_dataset(r"F:\roms_rst.nc")["u"]
+    save_to_nc(r"F:\test.nc", data)
+
+    # xarray测试
+    data = xr.DataArray(np.random.rand(10, 20), dims=("x", "y"), name="temperature")
+    save_to_nc(r"F:\test_xarray.nc", data)
+
+    # numpy测试
+    arr = np.random.rand(5, 3)
+    coords = {"x": np.arange(5), "y": np.arange(3)}
+    save_to_nc(r"F:\test_numpy.nc", arr, varname="data", coords=coords)

{oafuncs-0.0.98.0 → oafuncs-0.0.98.2}/oafuncs/oa_down/idm.py

@@ -52,6 +52,7 @@ def downloader(task_url, folder_path, file_name, idm_engine=r"D:\Programs\Intern
     time_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     time_str = time_str.center(100, " ")
     print(f"[bold purple]{time_str}")
-    print(f"[green]IDM Downloader: {file_name}[/green]".center(100, " "))
+    info = f'IDM Downloader: {file_name}'.center(100, " ")
+    print(f"[green]{info}[/green]")
     print("[purple]*" * 100)
     # print("\n")

{oafuncs-0.0.98.0 → oafuncs-0.0.98.2}/oafuncs/oa_nc.py

@@ -11,7 +11,7 @@ __all__ = ["save", "merge", "modify", "rename", "check", "convert_longitude", "i
 
 def save(
     file_path: str,
-    data: Union[np.ndarray, xr.DataArray],
+    data: Union[np.ndarray, xr.DataArray, xr.Dataset],
     variable_name: Optional[str] = None,
     coordinates: Optional[dict] = None,
     write_mode: str = "w",
@@ -23,7 +23,7 @@ def save(
 
     Args:
         file_path (str): File path to save the NetCDF file.
-        data (Union[np.ndarray, xr.DataArray]): Data to be written.
+        data (Union[np.ndarray, xr.DataArray, xr.Dataset]): Data to be written.
         variable_name (Optional[str]): Variable name for the data.
         coordinates (Optional[dict]): Coordinates, where keys are dimension names and values are coordinate data.
         write_mode (str): Write mode, 'w' for write, 'a' for append. Default is 'w'.

{oafuncs-0.0.98.0 → oafuncs-0.0.98.2/oafuncs.egg-info}/PKG-INFO

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

{oafuncs-0.0.98.0 → oafuncs-0.0.98.2}/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.9.0"  # 2025/03/13
-VERSION = "0.0.98.0"  # 下次用98.0，98.1已经被用过了
+VERSION = "0.0.98.2"
 
 # What packages are required for this module to be executed?
 REQUIRED = [

oafuncs-0.0.98.0/oafuncs/_script/netcdf_write.py

@@ -1,125 +0,0 @@
-import numpy as np
-import os
-import netCDF4 as nc
-import xarray as xr
-
-
-def _numpy_to_nc_type(numpy_type):
-    """将NumPy数据类型映射到NetCDF数据类型"""
-    numpy_to_nc = {
-        "float32": "f4",
-        "float64": "f8",
-        "int8": "i1",
-        "int16": "i2",
-        "int32": "i4",
-        "int64": "i8",
-        "uint8": "u1",
-        "uint16": "u2",
-        "uint32": "u4",
-        "uint64": "u8",
-    }
-    # 确保传入的是字符串类型，如果不是，则转换为字符串
-    numpy_type_str = str(numpy_type) if not isinstance(numpy_type, str) else numpy_type
-    return numpy_to_nc.get(numpy_type_str, "f4")  # 默认使用 'float32'
-
-
-def _calculate_scale_and_offset(data, n=16):
-    if not isinstance(data, np.ndarray):
-        raise ValueError("Input data must be a NumPy array.")
-
-    # 使用 nan_to_num 来避免 NaN 值对 min 和 max 的影响
-    data_min = np.nanmin(data)
-    data_max = np.nanmax(data)
-
-    if np.isnan(data_min) or np.isnan(data_max):
-        raise ValueError("Input data contains NaN values, which are not allowed.")
-
-    scale_factor = (data_max - data_min) / (2**n - 1)
-    add_offset = data_min + 2 ** (n - 1) * scale_factor
-
-    return scale_factor, add_offset
-
-
-def save_to_nc(file, data, varname=None, coords=None, mode="w", scale_offset_switch=True, compile_switch=True):
-    """
-    Description:
-        Write data to NetCDF file
-    Parameters:
-        file: str, file path
-        data: data
-        varname: str, variable name
-        coords: dict, coordinates, key is the dimension name, value is the coordinate data
-        mode: str, write mode, 'w' for write, 'a' for append
-        scale_offset_switch: bool, whether to use scale_factor and add_offset, default is True
-        compile_switch: bool, whether to use compression parameters, default is True
-    Example:
-        save(r'test.nc', data, 'u', {'time': np.linspace(0, 120, 100), 'lev': np.linspace(0, 120, 50)}, 'a')
-    """
-    # 设置压缩参数
-    kwargs = {"zlib": True, "complevel": 4} if compile_switch else {}
-
-    # 检查文件存在性并根据模式决定操作
-    if mode == "w" and os.path.exists(file):
-        os.remove(file)
-    elif mode == "a" and not os.path.exists(file):
-        mode = "w"
-
-    # 打开 NetCDF 文件
-    with nc.Dataset(file, mode, format="NETCDF4") as ncfile:
-        # 如果 data 是 DataArray 并且没有提供 varname 和 coords
-        if varname is None and coords is None and isinstance(data, xr.DataArray):
-            encoding = {}
-            for var in data.data_vars:
-                scale_factor, add_offset = _calculate_scale_and_offset(data[var].values)
-                encoding[var] = {
-                    "zlib": True,
-                    "complevel": 4,
-                    "dtype": "int16",
-                    "scale_factor": scale_factor,
-                    "add_offset": add_offset,
-                    "_FillValue": -32767,
-                }
-            data.to_netcdf(file, mode=mode, encoding=encoding)
-            return
-
-        # 添加坐标
-        for dim, coord_data in coords.items():
-            if dim in ncfile.dimensions:
-                if len(coord_data) != len(ncfile.dimensions[dim]):
-                    raise ValueError(f"Length of coordinate '{dim}' does not match the dimension length.")
-                else:
-                    ncfile.variables[dim][:] = np.array(coord_data)
-            else:
-                ncfile.createDimension(dim, len(coord_data))
-                var = ncfile.createVariable(dim, _numpy_to_nc_type(coord_data.dtype), (dim,), **kwargs)
-                var[:] = np.array(coord_data)
-
-                # 如果坐标数据有属性，则添加到 NetCDF 变量
-                if isinstance(coord_data, xr.DataArray) and coord_data.attrs:
-                    for attr_name, attr_value in coord_data.attrs.items():
-                        var.setncattr(attr_name, attr_value)
-
-        # 添加或更新变量
-        if varname in ncfile.variables:
-            if data.shape != ncfile.variables[varname].shape:
-                raise ValueError(f"Shape of data does not match the variable shape for '{varname}'.")
-            ncfile.variables[varname][:] = np.array(data)
-        else:
-            # 创建变量
-            dim_names = tuple(coords.keys())
-            if scale_offset_switch:
-                scale_factor, add_offset = _calculate_scale_and_offset(np.array(data))
-                dtype = "i2"
-                var = ncfile.createVariable(varname, dtype, dim_names, fill_value=-32767, **kwargs)
-                var.setncattr("scale_factor", scale_factor)
-                var.setncattr("add_offset", add_offset)
-            else:
-                dtype = _numpy_to_nc_type(data.dtype)
-                var = ncfile.createVariable(varname, dtype, dim_names, **kwargs)
-            var[:] = np.array(data)
-
-        # 添加属性
-        if isinstance(data, xr.DataArray) and data.attrs:
-            for key, value in data.attrs.items():
-                if key not in ["scale_factor", "add_offset", "_FillValue", "missing_value"] or not scale_offset_switch:
-                    var.setncattr(key, value)