oafuncs 0.0.97.13__py3-none-any.whl → 0.0.97.15__py3-none-any.whl

oafuncs/{oa_tool → _script}/parallel.py

@@ -2,10 +2,10 @@
 # coding=utf-8
 """
 Author: Liu Kun && 16031215@qq.com
-Date: 2025-01-11 19:47:08
+Date: 2025-04-04 20:19:23
 LastEditors: Liu Kun && 16031215@qq.com
-LastEditTime: 2025-03-18 19:21:36
-FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\oa_tool\\parallel.py
+LastEditTime: 2025-04-04 20:19:23
+FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\_script\\parallel.py
 Description:
 EditPlatform: vscode
 ComputerInfo: XPS 15 9510
@@ -13,6 +13,8 @@ SystemInfo: Windows 11
 Python Version: 3.12
 """
 
+
+
 import contextlib
 import logging
 import multiprocessing as mp

oafuncs/oa_cmap.py

@@ -1,27 +1,15 @@
-#!/usr/bin/env python
-# coding=utf-8
-"""
-Author: Liu Kun && 16031215@qq.com
-Date: 2024-09-17 16:55:11
-LastEditors: Liu Kun && 16031215@qq.com
-LastEditTime: 2024-11-21 13:14:24
-FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\oa_cmap.py
-Description:
-EditPlatform: vscode
-ComputerInfo: XPS 15 9510
-SystemInfo: Windows 11
-Python Version: 3.11
-"""
+from typing import List, Optional, Union
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 from rich import print
 
-__all__ = ["show", "to_color", "create", "create_rgbtxt", "get"]
+__all__ = ["show", "to_color", "create", "get"]
+
 
 # ** 将cmap用填色图可视化（官网摘抄函数）
-def show(colormaps):
+def show(colormaps: Union[str, mpl.colors.Colormap, List[Union[str, mpl.colors.Colormap]]]) -> None:
     """
     Description:
         Helper function to plot data with associated colormap.
@@ -31,107 +19,84 @@ def show(colormaps):
         cmap = ListedColormap(["darkorange", "gold", "lawngreen", "lightseagreen"])
         show([cmap]); show("viridis"); show(["viridis", "cividis"])
     """
-    if isinstance(colormaps, str) or isinstance(colormaps, mpl.colors.Colormap):
+    if not isinstance(colormaps, list):
         colormaps = [colormaps]
     np.random.seed(19680801)
     data = np.random.randn(30, 30)
     n = len(colormaps)
     fig, axs = plt.subplots(1, n, figsize=(n * 2 + 2, 3), constrained_layout=True, squeeze=False)
-    for [ax, cmap] in zip(axs.flat, colormaps):
+    for ax, cmap in zip(axs.flat, colormaps):
         psm = ax.pcolormesh(data, cmap=cmap, rasterized=True, vmin=-4, vmax=4)
         fig.colorbar(psm, ax=ax)
     plt.show()
 
 
 # ** 将cmap转为list，即多个颜色的列表
-def to_color(cmap, n=256):
+def to_color(cmap_name: str, n: int = 256) -> List[tuple]:
     """
     Description:
         Convert a colormap to a list of colors
     Parameters:
-        cmap : str; the name of the colormap
-        n    : int, optional; the number of colors
+        cmap_name : str; the name of the colormap
+        n         : int, optional; the number of colors
     Return:
         out_colors : list of colors
     Example:
         out_colors = to_color('viridis', 256)
     """
-    c_map = mpl.colormaps.get_cmap(cmap)
-    out_colors = [c_map(i) for i in np.linspace(0, 1, n)]
-    return out_colors
+    cmap = mpl.colormaps.get_cmap(cmap_name)
+    return [cmap(i) for i in np.linspace(0, 1, n)]
 
 
 # ** 自制cmap，多色，可带位置
-def create(colors: list, nodes=None, under=None, over=None):  # 利用颜色快速配色
+def create(color_list: Optional[List[Union[str, tuple]]] = None, rgb_file_path: Optional[str] = None, positions: Optional[List[float]] = None, under_color: Optional[Union[str, tuple]] = None, over_color: Optional[Union[str, tuple]] = None, delimiter: str = ",") -> mpl.colors.Colormap:
     """
     Description:
-        Create a custom colormap
+        Create a custom colormap from a list of colors or an RGB txt document.
     Parameters:
-        colors : list of colors
-        nodes  : list of positions
-        under  : color
-        over   : color
+        color_list  : list of colors (optional, required if rgb_file_path is None)
+        rgb_file_path : str, the path of txt file (optional, required if color_list is None)
+        positions   : list of positions (optional, for color_list)
+        under_color : color (optional)
+        over_color  : color (optional)
+        delimiter   : str, optional, default is ','; the delimiter of RGB values in txt file
     Return:
         cmap : colormap
     Example:
-        cmap = create(['#C2B7F3','#B3BBF2','#B0CBF1','#ACDCF0','#A8EEED'])
-        cmap = create(['aliceblue','skyblue','deepskyblue'],[0.0,0.5,1.0])
+        cmap = create(color_list=['#C2B7F3','#B3BBF2','#B0CBF1','#ACDCF0','#A8EEED'])
+        cmap = create(color_list=['aliceblue','skyblue','deepskyblue'], positions=[0.0,0.5,1.0])
+        cmap = create(rgb_file_path='path/to/file.txt', delimiter=',')
     """
+    if rgb_file_path:
+        with open(rgb_file_path) as fid:
+            data = fid.readlines()
+        n = len(data)
+        rgb = np.zeros((n, 3))
+        for i in np.arange(n):
+            rgb[i][0] = data[i].split(delimiter)[0]
+            rgb[i][1] = data[i].split(delimiter)[1]
+            rgb[i][2] = data[i].split(delimiter)[2]
+        max_rgb = np.max(rgb)
+        if max_rgb > 2:  # if the value is greater than 2, it is normalized to 0-1
+            rgb = rgb / 255.0
+        cmap_color = mpl.colors.ListedColormap(rgb, name="my_color")
+    elif color_list:
+        if positions is None:  # 自动分配比例
+            cmap_color = mpl.colors.LinearSegmentedColormap.from_list("mycmap", color_list)
+        else:  # 按提供比例分配
+            cmap_color = mpl.colors.LinearSegmentedColormap.from_list("mycmap", list(zip(positions, color_list)))
+    else:
+        raise ValueError("Either 'color_list' or 'rgb_file_path' must be provided.")
 
-    if nodes is None:  # 采取自动分配比例
-        cmap_color = mpl.colors.LinearSegmentedColormap.from_list("mycmap", colors)
-    else:  # 按照提供比例分配
-        cmap_color = mpl.colors.LinearSegmentedColormap.from_list("mycmap", list(zip(nodes, colors)))
-    if under is not None:
-        cmap_color.set_under(under)
-    if over is not None:
-        cmap_color.set_over(over)
+    if under_color is not None:
+        cmap_color.set_under(under_color)
+    if over_color is not None:
+        cmap_color.set_over(over_color)
     return cmap_color
 
 
-# ** 根据RGB的txt文档制作色卡（利用Grads调色盘）
-def create_rgbtxt(rgbtxt_file,split_mark=','):  # 根据RGB的txt文档制作色卡/根据rgb值制作
-    """
-    Description
-    -----------
-    Make a color card according to the RGB txt document, each line in the txt file is an RGB value, separated by commas, such as: 251,251,253
-    
-    Parameters
-    ----------
-    rgbtxt_file : str, the path of txt file
-    split_mark  : str, optional, default is ','; the split mark of rgb value
-
-    Returns
-    -------
-    cmap : colormap
-
-    Example
-    -------
-    cmap=create_rgbtxt(path,split_mark=',')
-    
-    txt example
-    -----------
-    251,251,253
-    225,125,25
-    250,205,255
-    """
-    with open(rgbtxt_file) as fid:
-        data = fid.readlines()
-    n = len(data)
-    rgb = np.zeros((n, 3))
-    for i in np.arange(n):
-        rgb[i][0] = data[i].split(split_mark)[0]
-        rgb[i][1] = data[i].split(split_mark)[1]
-        rgb[i][2] = data[i].split(split_mark)[2]
-    max_rgb = np.max(rgb)
-    if max_rgb > 2:  # if the value is greater than 2, it is normalized to 0-1
-        rgb = rgb / 255.0
-    my_cmap = mpl.colors.ListedColormap(rgb, name="my_color")
-    return my_cmap
-
-
 # ** 选择cmap
-def get(cmap_name=None, query=False):
+def get(cmap_name: Optional[str] = None, query: bool = False) -> Optional[mpl.colors.Colormap]:
     """
     Description:
         Choosing a colormap from the list of available colormaps or a custom colormap
@@ -147,43 +112,26 @@ def get(cmap_name=None, query=False):
         cmap = get('warm_1')
         cmap = get('colorful_1')
     """
-
     my_cmap_dict = {
         "diverging_1": create(["#4e00b3", "#0000FF", "#00c0ff", "#a1d3ff", "#DCDCDC", "#FFD39B", "#FF8247", "#FF0000", "#FF5F9E"]),
         "cool_1": create(["#4e00b3", "#0000FF", "#00c0ff", "#a1d3ff", "#DCDCDC"]),
         "warm_1": create(["#DCDCDC", "#FFD39B", "#FF8247", "#FF0000", "#FF5F9E"]),
-        # "land_1": create_custom(["#3E6436", "#678A59", "#91A176", "#B8A87D", "#D9CBB2"], under="#A6CEE3", over="#FFFFFF"),
-        # "ocean_1": create_custom(["#126697", "#2D88B3", "#4EA1C9", "#78B9D8", "#A6CEE3"], under="#8470FF", over="#3E6436"), 
-        # "ocean_land_1": create_custom(
-        #     [
-        #         "#126697",  # 深蓝（深海）
-        #         "#2D88B3",  # 蓝
-        #         "#4EA1C9",  # 蓝绿
-        #         "#78B9D8",  # 浅蓝（浅海）
-        #         "#A6CEE3",  # 浅蓝（近岸）
-        #         "#AAAAAA",  # 灰色（0值，海平面）
-        #         "#D9CBB2",  # 沙质土壤色（陆地开始）
-        #         "#B8A87D",  # 浅棕
-        #         "#91A176",  # 浅绿
-        #         "#678A59",  # 中绿
-        #         "#3E6436",  # 深绿（高山）
-        #     ]
-        # ),
         "colorful_1": create(["#6d00db", "#9800cb", "#F2003C", "#ff4500", "#ff7f00", "#FE28A2", "#FFC0CB", "#DDA0DD", "#40E0D0", "#1a66f2", "#00f7fb", "#8fff88", "#E3FF00"]),
     }
+
     if query:
         print("Available cmap names:")
-        print('-' * 20)
-        print('Defined by myself:')
-        for key, _ in my_cmap_dict.items():
-            print(key)
-        print('-' * 20)
-        print('Matplotlib built-in:')
-        print(mpl.colormaps())
         print("-" * 20)
-    
+        print("Defined by myself:")
+        print("\n".join(my_cmap_dict.keys()))
+        print("-" * 20)
+        print("Matplotlib built-in:")
+        print("\n".join(mpl.colormaps.keys()))
+        print("-" * 20)
+        return None
+
     if cmap_name is None:
-        return
+        return None
 
     if cmap_name in my_cmap_dict:
         return my_cmap_dict[cmap_name]
@@ -191,9 +139,8 @@ def get(cmap_name=None, query=False):
         try:
             return mpl.colormaps.get_cmap(cmap_name)
         except ValueError:
-            # raise ValueError(f"Unknown cmap name: {cmap_name}")
             print(f"Unknown cmap name: {cmap_name}\nNow return 'rainbow' as default.")
-            return mpl.colormaps.get_cmap("rainbow")
+            return mpl.colormaps.get_cmap("rainbow")  # 默认返回 'rainbow'
 
 
 if __name__ == "__main__":
@@ -209,7 +156,7 @@ if __name__ == "__main__":
 
     # ** 测试根据RGB的txt文档制作色卡
     file_path = "E:/python/colorbar/test.txt"
-    cmap_rgb = create_rgbtxt(file_path)
+    cmap_rgb = create(rgb_file_path=file_path)
 
     # ** 测试将cmap转为list
     out_colors = to_color("viridis", 256)

oafuncs/oa_data.py

@@ -17,14 +17,13 @@ import itertools
 import multiprocessing as mp
 from concurrent.futures import ThreadPoolExecutor
 
+
 import numpy as np
 import salem
 import xarray as xr
-from scipy.interpolate import griddata
-from scipy.interpolate import interp1d
-from typing import Iterable
+from scipy.interpolate import griddata, interp1d
 
-__all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile", "pbar"]
+__all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile"]
 
 
 def ensure_list(input_data):
@@ -255,26 +254,6 @@ def mask_shapefile(data: np.ndarray, lons: np.ndarray, lats: np.ndarray, shapefi
         return None
 
 
-def pbar(iterable: Iterable, prefix: str = "", color: str = "cyan", cmap: str = None, **kwargs) -> Iterable:
-    """
-    快速创建进度条的封装函数
-    :param iterable: 可迭代对象
-    :param prefix: 进度条前缀
-    :param color: 基础颜色
-    :param cmap: 渐变色名称
-    :param kwargs: 其他ColorProgressBar支持的参数
-    
-    example:
-    from oafuncs.oa_data import pbar
-    from time import sleep
-    for i in pbar(range(100), prefix="Processing", color="green", cmap="viridis"):
-        sleep(0.1)
-    """
-    from ._script.cprogressbar import ColorProgressBar  # 从progressbar.py导入类
-
-    return ColorProgressBar(iterable=iterable, prefix=prefix, color=color, cmap=cmap, **kwargs)
-
-
 if __name__ == "__main__":
     pass
     """ import time

oafuncs/oa_date.py

@@ -4,7 +4,7 @@
 Author: Liu Kun && 16031215@qq.com
 Date: 2025-03-27 16:56:57
 LastEditors: Liu Kun && 16031215@qq.com
-LastEditTime: 2025-03-27 16:56:57
+LastEditTime: 2025-04-04 12:58:15
 FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\oa_date.py
 Description:
 EditPlatform: vscode
@@ -15,36 +15,37 @@ Python Version: 3.12
 
 import calendar
 import datetime
+from typing import List, Optional
 
-__all__ = ["get_days_in_month", "generate_hour_list", "adjust_time"]
+__all__ = ["month_days", "hour_range", "adjust_time", "timeit"]
 
 
-def get_days_in_month(year, month):
+def month_days(year: int, month: int) -> int:
     return calendar.monthrange(year, month)[1]
 
 
-def generate_hour_list(start_date, end_date, interval_hours=6):
+def hour_range(start: str, end: str, interval: int = 6) -> List[str]:
     """
     Generate a list of datetime strings with a specified interval in hours.
 
     Args:
-        start_date (str): Start date in the format "%Y%m%d%H".
-        end_date (str): End date in the format "%Y%m%d%H".
-        interval_hours (int): Interval in hours between each datetime.
+        start (str): Start date in the format "%Y%m%d%H".
+        end (str): End date in the format "%Y%m%d%H".
+        interval (int): Interval in hours between each datetime.
 
     Returns:
         list: List of datetime strings in the format "%Y%m%d%H".
     """
-    date_s = datetime.datetime.strptime(start_date, "%Y%m%d%H")
-    date_e = datetime.datetime.strptime(end_date, "%Y%m%d%H")
+    date_s = datetime.datetime.strptime(start, "%Y%m%d%H")
+    date_e = datetime.datetime.strptime(end, "%Y%m%d%H")
     date_list = []
     while date_s <= date_e:
         date_list.append(date_s.strftime("%Y%m%d%H"))
-        date_s += datetime.timedelta(hours=interval_hours)
+        date_s += datetime.timedelta(hours=interval)
     return date_list
 
 
-def adjust_time(initial_time, amount, time_unit="hours", output_format=None):
+def adjust_time(initial_time: str, amount: int, time_unit: str = "hours", output_format: Optional[str] = None) -> str:
     """
     Adjust a given initial time by adding a specified amount of time.
 
@@ -88,3 +89,38 @@ def adjust_time(initial_time, amount, time_unit="hours", output_format=None):
         elif time_unit == "days":
             default_format = "%Y%m%d"
         return time_obj.strftime(default_format)
+
+
+class timeit:
+    """
+    A decorator to measure the execution time of a function.
+
+    Usage:
+        @timeit(log=True, print_time=True)
+        def my_function():
+            # Function code here
+
+    Args:
+        log (bool): Whether to log the execution time to a file. Defaults to False.
+        print_time (bool): Whether to print the execution time to the console. Defaults to True.
+    """
+
+    def __init__(self, func, log: bool = False, print_time: bool = True):
+        self.func = func
+        self.log = log
+        self.print_time = print_time
+
+    def __call__(self, *args, **kwargs):
+        start_time = datetime.datetime.now()
+        result = self.func(*args, **kwargs)
+        end_time = datetime.datetime.now()
+        elapsed_time = (end_time - start_time).total_seconds()
+
+        if self.print_time:
+            print(f"Function '{self.func.__name__}' executed in {elapsed_time:.2f} seconds.")
+
+        if self.log:
+            with open("execution_time.log", "a") as log_file:
+                log_file.write(f"{datetime.datetime.now()} - Function '{self.func.__name__}' executed in {elapsed_time:.2f} seconds.\n")
+
+        return result

oafuncs/oa_down/hycom_3hourly.py

@@ -40,7 +40,7 @@ from oafuncs.oa_nc import modify as modify_nc
 
 warnings.filterwarnings("ignore", category=RuntimeWarning, message="Engine '.*' loading failed:.*")
 
-__all__ = ["draw_time_range", "download", "how_to_use", "get_time_list"]
+__all__ = ["draw_time_range", "download"]
 
 
 def _get_initial_data():
@@ -278,7 +278,7 @@ def draw_time_range(pic_save_folder=None):
     plt.close()
 
 
-def get_time_list(time_s, time_e, delta, interval_type="hour"):
+def _get_time_list(time_s, time_e, delta, interval_type="hour"):
     """
     Description: get a list of time strings from time_s to time_e with a specified interval
     Args:
@@ -1029,7 +1029,7 @@ def _download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_mi
         _prepare_url_to_download(var, lon_min, lon_max, lat_min, lat_max, ymdh_time_s, None, depth, level, store_path, dataset_name, version_name, check)
     elif int(ymdh_time_s) < int(ymdh_time_e):
         print("Downloading a series of files...")
-        time_list = get_time_list(ymdh_time_s, ymdh_time_e, interval_hour, "hour")
+        time_list = _get_time_list(ymdh_time_s, ymdh_time_e, interval_hour, "hour")
         with Progress() as progress:
             task = progress.add_task(f"[cyan]{bar_desc}", total=len(time_list))
             if ftimes == 1:
@@ -1048,7 +1048,7 @@ def _download_hourly_func(var, time_s, time_e, lon_min=0, lon_max=359.92, lat_mi
                             _done_callback(feature, progress, task, len(time_list), counter_lock)
             else:
                 # new_time_list = get_time_list(ymdh_time_s, ymdh_time_e, 3 * ftimes, "hour")
-                new_time_list = get_time_list(ymdh_time_s, ymdh_time_e, interval_hour * ftimes, "hour")
+                new_time_list = _get_time_list(ymdh_time_s, ymdh_time_e, interval_hour * ftimes, "hour")
                 total_num = len(new_time_list)
                 if num_workers is None or num_workers <= 1:
                     # 串行方式
@@ -1248,55 +1248,6 @@ def download(var, time_s, time_e=None, lon_min=0, lon_max=359.92, lat_min=-80, l
     print("[bold #ecdbfe]=" * mark_len)
 
 
-def how_to_use():
-    print("""
-    # 1. Choose the dataset and version according to the time:
-    # 1.1 Use function to query
-    You can use the function check_time_in_dataset_and_version(time_input=20241101) to find the dataset and version  according to the time.
-    Then, you can see the dataset and version in the output.
-    # 1.2 Draw a picture to see
-    You can draw a picture to see the time range of each dataset and version.
-    Using the function draw_time_range(pic_save_folder=None) to draw the picture.
-
-    # 2. Get the base url according to the dataset, version, var and year:
-    # 2.1 Dataset and version were found in step 1
-    # 2.2 Var: u, v, temp, salt, ssh, u_b, v_b, temp_b, salt_b
-    # 2.3 Year: 1994-2024(current year)
-
-    # 3. Get the query_dict according to the var, lon_min, lon_max, lat_min, lat_max, depth, level_num, time_str_ymdh:
-    # 3.1 Var: u, v, temp, salt, ssh, u_b, v_b, temp_b, salt_b
-    # 3.2 Lon_min, lon_max, lat_min, lat_max: float
-    # 3.3 Depth: 0-5000m, if you wanna get single depth data, you can set the depth
-    # 3.4 Level_num: 1-40, if you wanna get single level data, you can set the level_num
-    # 3.5 Time_str_ymdh: '2024110112', the hour normally is 00, 03, 06, 09, 12, 15, 18, 21, besides 1 hourly data
-    # 3.6 Use the function to get the query_dict
-    # 3.7 Note: If you wanna get the full depth or full level data, you can needn't set the depth or level_num
-
-    # 4. Get the submit url according to the dataset, version, var, year, query_dict:
-    # 4.1 Use the function to get the submit url
-    # 4.2 You can use the submit url to download the data
-
-    # 5. Download the data according to the submit url:
-    # 5.1 Use the function to download the data
-    # 5.2 You can download the data of single time or a series of time
-    # 5.3 Note: If you wanna download a series of data, you can set the ymdh_time_s and ymdh_time_e different
-    # 5.4 Note: The time resolution is 3 hours
-
-    # 6. Direct download the data:
-    # 6.1 Use the function to direct download the data
-    # 6.2 You can set the dataset_name and version_name by yourself
-    # 6.3 Note: If you do not set the dataset_name and version_name, the dataset and version will be chosen according to the download_time
-    # 6.4 Note: If you set the dataset_name and version_name, please ensure the dataset_name and version_name are correct
-    # 6.5 Note: If you just set one of the dataset_name and version_name, both the dataset and version will be chosen according to the download_time
-
-    # 7. Simple use:
-    # 7.1 You can use the function: download(var, ymdh_time_s, ymdh_time_e, lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90,  depth=None, level_num=None, store_path=None, dataset_name=None, version_name=None)
-    # 7.2 You can download the data of single time or a series of time
-    # 7.3 The parameters you must set are var, ymdh_time_s, ymdh_time_e
-    # 7.4 Example: download('u', '2024110112', '2024110212', lon_min=0, lon_max=359.92, lat_min=-80, lat_max=90,  depth=None, level_num=None, store_path=None, dataset_name=None, version_name=None)
-    """)
-
-
 if __name__ == "__main__":
     download_dict = {
         "water_u": {"simple_name": "u", "download": 1},