oafuncs 0.0.98.0__py3-none-any.whl → 0.0.98.2__py3-none-any.whl

oafuncs/_script/netcdf_write.py

@@ -1,6 +1,6 @@
-import numpy as np
 import os
 import netCDF4 as nc
+import numpy as np
 import xarray as xr
 
 
@@ -18,108 +18,99 @@ def _numpy_to_nc_type(numpy_type):
         "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'
+    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.")
 
-    # 使用 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.")
+        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):
-    """
-    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 {}
-
-    # 检查文件存在性并根据模式决定操作
+    # 处理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"
 
-    # 打开 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
+    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
 
-        # 添加坐标
-        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())
+            dims = list(coords.keys()) if coords else []
             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)
+                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, 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)
+                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/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/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.dist-info → oafuncs-0.0.98.2.dist-info}/METADATA

@@ -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.dist-info → oafuncs-0.0.98.2.dist-info}/RECORD

@@ -5,7 +5,7 @@ oafuncs/oa_date.py,sha256=KqU-bHtC74hYsf6VgiA3i2vI__q_toOVR-whFy4cYP8,5523
 oafuncs/oa_draw.py,sha256=9hchfs7kyBo_On1FA5jliSoAwpsjEXNUIC9C20D9ugw,15050
 oafuncs/oa_file.py,sha256=goF5iRXJFFCIKhIjlkCnYYt0EYlJb_4r8AeYNZ0-SOk,16209
 oafuncs/oa_help.py,sha256=_4AZgRDq5Or0vauNvq5IDDHIBoBfdOQtzak-mG1wwAw,4537
-oafuncs/oa_nc.py,sha256=StkcZZYo6FKANuPcq7OaIKoBKj5HuFwHsrZRVKa4ULw,10489
+oafuncs/oa_nc.py,sha256=L1gqXxg93kIDsMOa87M0o-53KVmdqCipnXeF9XfzfY8,10513
 oafuncs/oa_python.py,sha256=NkopwkYFGSEuVljnTBvXCl6o2CeyRNBqRXSsUl3euEE,5192
 oafuncs/oa_tool.py,sha256=8y0X3KXNLbYchr5hBHCwhdsA_1VDEnzsZVqcM5sNeog,4484
 oafuncs/_data/hycom.png,sha256=MadKs6Gyj5n9-TOu7L4atQfTXtF9dvN9w-tdU9IfygI,10945710
@@ -14,7 +14,7 @@ oafuncs/_script/cprogressbar.py,sha256=wRU3SFPFtMI7ER26tTzg223kVKNo5RDWE9CzdIgUs
 oafuncs/_script/email.py,sha256=lL4HGKrr524-g0xLlgs-4u7x4-u7DtgNoD9AL8XJKj4,3058
 oafuncs/_script/netcdf_merge.py,sha256=_EPF9Xj4HOVC9sZpi1lt62-Aq6pMlgsgwaajEBLhW6g,5092
 oafuncs/_script/netcdf_modify.py,sha256=sGRUYNhfGgf9JV70rnBzw3bzuTRSXzBTL_RMDnDPeLQ,4552
-oafuncs/_script/netcdf_write.py,sha256=KZ0ui2nvTujAciLKx4bwQHLEkH1R3P8hCzd7BysJ0D4,5288
+oafuncs/_script/netcdf_write.py,sha256=8bVsJKRPLWJawUTAsMUiG5D9lOg_-sRMgshgtRjo_us,4598
 oafuncs/_script/parallel.py,sha256=FS9FgaByq2yb9j6nL-Y0xP1VLvp4USMLBFMRsJDoqeQ,21848
 oafuncs/_script/parallel_example_usage.py,sha256=uLvE7iwkMn9Cyq6-wk5_RpbQk7PXM9d16-26lTknW9s,2646
 oafuncs/_script/plot_dataset.py,sha256=zkSEnO_-biyagorwWXPoihts_cwuvripzEt-l9bHJ2E,13989
@@ -23,7 +23,7 @@ oafuncs/oa_down/User_Agent-list.txt,sha256=pazxSip8_lphEBOPHG902zmIBUg8sBKXgmqp_
 oafuncs/oa_down/__init__.py,sha256=kRX5eTUCbAiz3zTaQM1501paOYS_3fizDN4Pa0mtNUA,585
 oafuncs/oa_down/hycom_3hourly.py,sha256=Xi6LuD3DIiYq2gYYBVxY3OpVvfy3RkSMzNkOX2isSK8,52915
 oafuncs/oa_down/hycom_3hourly_20250407.py,sha256=DQd_NmQgmSqu7jsrfpDB7k23mkUEy9kyWs-dLUg7GDw,64472
-oafuncs/oa_down/idm.py,sha256=gZHyUJMyd9zgzZGksI-BuNuouK_JGjxQZniUgUFENPs,1801
+oafuncs/oa_down/idm.py,sha256=4z5IvgfTyIKEI1kOtqXZwN7Jnfjwp6qDBOIoVyOLp0I,1823
 oafuncs/oa_down/literature.py,sha256=2bF9gSKQbzcci9LcKE81j8JEjIJwON7jbwQB3gDDA3E,11331
 oafuncs/oa_down/test_ua.py,sha256=0IQq3NjqfNr7KkyjS_U-a4mYu-r-E7gzawwo4IfEa6Y,10851
 oafuncs/oa_down/user_agent.py,sha256=TsPcAxFmMTYAEHRFjurI1bQBJfDhcA70MdHoUPwQmks,785
@@ -36,8 +36,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.0.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
-oafuncs-0.0.98.0.dist-info/METADATA,sha256=EVqN2an2LrePsPwvCK_OTDtX6Vc5mJMECkLjbDMlmC8,4220
-oafuncs-0.0.98.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-oafuncs-0.0.98.0.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
-oafuncs-0.0.98.0.dist-info/RECORD,,
+oafuncs-0.0.98.2.dist-info/licenses/LICENSE.txt,sha256=rMtLpVg8sKiSlwClfR9w_Dd_5WubTQgoOzE2PDFxzs4,1074
+oafuncs-0.0.98.2.dist-info/METADATA,sha256=2-IktPjhTVHt-GLLgJk7J2tYMBYaNA0eY7IZyB0kxlk,4220
+oafuncs-0.0.98.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+oafuncs-0.0.98.2.dist-info/top_level.txt,sha256=bgC35QkXbN4EmPHEveg_xGIZ5i9NNPYWqtJqaKqTPsQ,8
+oafuncs-0.0.98.2.dist-info/RECORD,,