Alaea 2.4.0__tar.gz → 2.4.2__tar.gz

{alaea-2.4.0 → alaea-2.4.2}/Alaea/S_DateTime.py

@@ -123,9 +123,9 @@ class Ttime:
 
         if TIME is not None and fmt is not None:
             if isinstance(TIME, list):
-                self.Times = [datetime.datetime.strptime(TIME[i], fmt) for i in range(len(TIME))]
+                self.Times = [datetime.strptime(TIME[i], fmt) for i in range(len(TIME))]
             else:
-                self.Times = [datetime.datetime.strptime(TIME, fmt)]
+                self.Times = [datetime.strptime(TIME, fmt)]
         elif TIME is not None:
             if hasattr(TIME[0], 'dtype'):
                 if isinstance(TIME, list):

{alaea-2.4.0 → alaea-2.4.2}/Alaea/baysalt/erosion_coastline.py

@@ -14,10 +14,11 @@
 
 这里的“erosion（侵蚀）”更接近一种海岸线附近的掩膜修正 / 邻域填补：
 """
-
-import itertools
+import contextlib
+import os
 import time
-from collections import deque
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from typing import List
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -30,9 +31,13 @@ from Alaea.Blogging import Blog
 from Alaea.tools.Read import read_json, write_json
 
 __all__ = [
+    'LOGGER',
     'erosion_cal_id',
     'erosion_via_id',
-    'LOGGER',
+    'erosion_cal_id_multi_times',
+    'erosion_via_id_multi_times',
+    'erosion_via_id_multi_vars',
+    'erosion_via_id_multi_times_multi_vars',
 ]
 
 BLOG = Blog('erosion_coastline')
@@ -47,7 +52,7 @@ def erosion_cal_id(
         k,
         judge_num,
         method='np.nan',
-):
+) -> dict:
     """
     计算“哪些点需要被回填”，以及“每个点应该参考哪些邻居点”。
 
@@ -90,79 +95,73 @@ def erosion_cal_id(
     lon_mesh, lat_mesh = np.meshgrid(g_lon, g_lat)
     # 后面 KDTree 就是在这个二维坐标集合上做最近邻搜索
     point = np.column_stack((lon_mesh.ravel(), lat_mesh.ravel()))
-    # 下面两个 deque 用来保存所有“候选邻居点”的二维索引
-    fp_lon = deque([])
-    fp_lat = deque([])
-
-    # 建立 KDTree：适合大量最近邻查询
-    # 这里 point 里装的是所有格点坐标，所以相当于“让每个格点去找离自己最近的 k 个格点”
-    ckt = spt.cKDTree(point)  # 用C写的查找类，执行速度更快
-    # 这里不再逐点循环，而是一次性把所有格点的近邻都查出来
-    LOGGER.debug('Kdtree建立完毕，开始查找每个格点的近邻点')
-
-    # 对所有格点一次性查询最近邻：
-    # distances.shape 约为 (n_points, k)
-    # sequences.shape 约为 (n_points, k)
-    # sequences 里保存的是最近邻在 point 数组中的一维编号
-    distances, sequences = ckt.query(point, k)
-
-    # point 中每个元素对应一个展开后的一维索引 [0, 1, 2, ...]
-    idx_all = np.arange(point.shape[0])
-    # 把展开索引还原回二维索引：
-    # 行号 = index // 每行列数
-    lat_i_all = idx_all // lon_number
-    # 列号 = index % 每行列数
-    lon_i_all = idx_all % lon_number
 
     # 先找出哪些格点属于“待处理点”
     # 不同数据源中，水点/无效点/边界点的表达方式不同，所以这里分三种判断方法
     if method == 'np.nan':
-        mask = np.isnan(value[lat_i_all, lon_i_all])
+        mask = np.isnan(value)
     elif method == 'nc.mask':
-        mask = value.mask[lat_i_all, lon_i_all]
+        mask = value.mask
     elif method == 'txt.mask':
-        mask = (value[lat_i_all, lon_i_all] == 0)
+        mask = (value == 0)
     else:
-        mask = np.zeros_like(lat_i_all, dtype=bool)
+        mask = np.zeros((lat_number, lon_number), dtype=bool)
 
-    # 只保留那些“需要处理”的点在 point 里的编号
-    valid_idx = idx_all[mask]
+    target_lat, target_lon = np.where(mask)
+    if target_lat.size == 0:
+        return {
+            'change_lat': [],
+            'change_lon': [],
+            'value_lat': [],
+            'value_lon': [],
+        }
 
-    # 取出这些目标点各自对应的 k 个最近邻编号
-    neigh = sequences[valid_idx]
+    valid_idx = target_lat * lon_number + target_lon
+    target_point = point[valid_idx]
 
-    # 展平后，后面统一换算成二维索引
-    neigh = neigh.reshape(-1)
+    # 建立 KDTree：适合大量最近邻查询
+    ckt = spt.cKDTree(point)  # 用C写的查找类，执行速度更快
+    LOGGER.debug('Kdtree建立完毕，开始查找待处理格点的近邻点')
 
-    # 邻居点的一维编号 -> 二维行列编号
-    lat_j = neigh // lon_number
-    lon_j = neigh % lon_number
+    # 只对待处理点查找最近邻，而不是对所有格点都查
+    query_kwargs = {'k': k}
+    with contextlib.suppress(Exception):
+        query_kwargs['workers'] = -1
+    _, neigh = ckt.query(target_point, **query_kwargs)
 
-    # 这里会把所有目标点的 k 个邻居顺序拼在一起
-    # 因此后面需要再按 k reshape 回去，恢复“每 k 个属于同一个目标点”的结构
-    fp_lat.extend(lat_j.tolist())
-    fp_lon.extend(lon_j.tolist())
-    LOGGER.debug('所有格点的近邻点查找完毕，开始获取行列索引')
+    # 当 k == 1 时，scipy 会返回一维数组，这里统一成二维
+    neigh = np.atleast_2d(neigh)
+    if neigh.shape[0] != target_lat.size:
+        neigh = neigh.reshape(target_lat.size, -1)
 
-    # 这里 fp_lat / fp_lon 已经是索引，不再需要 find_index
-    lat_id = list(fp_lat)
-    lon_id = list(fp_lon)
+    # 邻居点的一维编号 -> 二维行列编号
+    lat_id = neigh // lon_number
+    lon_id = neigh % lon_number
 
     LOGGER.debug('行列索引获取完毕，开始判断每个目标点的邻居里有效点数量')
-    
+
     # 根据不同类型的数据，进一步判断：
     # 某个目标点的 k 个邻居里，是否有足够多的有效点可以拿来回填
-    lonli_arr = np.array([0])
-    latli_arr = np.array([0])
     if method == 'np.nan':
         latli_arr, lonli_arr = judge_nan_npnan(lat_id, lon_id, value, k, judge_num)
-    elif method =='nc.mask':
+    elif method == 'nc.mask':
         latli_arr, lonli_arr = judge_nan_ncmask(lat_id, lon_id, value, k, judge_num)
     elif method == 'txt.mask':
         latli_arr, lonli_arr = judge_nan_txtmask(lat_id, lon_id, value, k, judge_num)
+    else:
+        latli_arr = np.empty((0, k), dtype=int)
+        lonli_arr = np.empty((0, k), dtype=int)
 
     LOGGER.debug('有效邻居点数量判断完毕，开始生成索引')
 
+    if latli_arr.size == 0:
+        return {
+            'change_lat': [],
+            'change_lon': [],
+            'value_lat': [],
+            'value_lon': [],
+        }
+
     # reshape 之后，每一行对应“一个目标点 + 它的 k-1 或 k 个邻居”
     # 约定每行第 1 列是目标点本身，后面各列是参考邻居
     _change_lat = latli_arr[:, 0].tolist()
@@ -191,27 +190,23 @@ def judge_nan_ncmask(latli, lonli, value, k, judge_num):
     - 如果这 k 个点里“未被 mask 的有效点”数量不少于 judge_num，
       就保留这一组索引
     """
-    # 先转成数组，便于 reshape 成 [目标点数, k] 的二维结构
-    latli_1 = np.array(latli)
-    lonli_1 = np.array(lonli)
-    # 每一行表示一个目标点对应的 k 个最近邻索引
-    latli_2 = np.reshape(latli_1, [-1, k])
-    lonli_2 = np.reshape(lonli_1, [-1, k])
+    latli_2 = np.asarray(latli)
+    lonli_2 = np.asarray(lonli)
+    if latli_2.size == 0:
+        return latli_2.reshape(0, k), lonli_2.reshape(0, k)
+
     # mask=True 说明该位置是无效/被遮罩的点
     mask_arr = value.mask
     # 取出每组 k 个邻居的 mask 情况
-    test_arr = mask_arr[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
+    test_arr = mask_arr[latli_2, lonli_2]
 
     # 统计每组中“无效点”的数量
-    sum_arr = np.sum(test_arr, axis=1)  #
+    sum_arr = np.sum(test_arr, axis=1)
 
     # 无效点数 <= k - judge_num，等价于有效点数 >= judge_num
-    select_id = np.where(sum_arr <= k-judge_num)[0].tolist()
+    select_mask = sum_arr <= (k - judge_num)
 
-    latli_r = latli_2[select_id[:]]
-    lonli_r = lonli_2[select_id[:]]
-
-    return latli_r, lonli_r
+    return latli_2[select_mask], lonli_2[select_mask]
 
 
 def judge_nan_npnan(latli, lonli, value, k, judge_num):
@@ -221,26 +216,22 @@ def judge_nan_npnan(latli, lonli, value, k, judge_num):
     逻辑与 judge_nan_ncmask 一样，只不过这里通过 np.isnan(value)
     判断某个点是否无效。
     """
-    # 每 k 个索引是一组，后面要恢复成二维结构
-    latli_1 = np.array(latli)
-    lonli_1 = np.array(lonli)
-    # reshape 后每一行对应 1 个目标点的 k 个邻居
-    latli_2 = np.reshape(latli_1, [-1, k])
-    lonli_2 = np.reshape(lonli_1, [-1, k])
+    latli_2 = np.asarray(latli)
+    lonli_2 = np.asarray(lonli)
+    if latli_2.size == 0:
+        return latli_2.reshape(0, k), lonli_2.reshape(0, k)
+
     # NaN 为 True，表示该点不能直接作为有效值参与平均
     mask_arr = np.isnan(value)
 
     # 检查每组 k 个邻居中哪些位置是 NaN
-    test_arr = mask_arr[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
-    sum_arr = np.sum(test_arr, axis=1)  #
+    test_arr = mask_arr[latli_2, lonli_2]
+    sum_arr = np.sum(test_arr, axis=1)
 
     # 只有有效值数量 >= judge_num 的目标点，才保留下来
-    select_id = np.where(sum_arr <= k-judge_num)[0].tolist()
+    select_mask = sum_arr <= (k - judge_num)
 
-    latli_r = latli_2[select_id[:]]
-    lonli_r = lonli_2[select_id[:]]
-
-    return latli_r, lonli_r
+    return latli_2[select_mask], lonli_2[select_mask]
 
 
 def judge_nan_txtmask(latli, lonli, value, k, judge_num):
@@ -253,25 +244,21 @@ def judge_nan_txtmask(latli, lonli, value, k, judge_num):
 
     当一组邻居中陆点数量 >= judge_num 时，保留该组。
     """
-    # 同样先恢复成 [目标点数, k] 的结构
-    latli_1 = np.array(latli)
-    lonli_1 = np.array(lonli)
-    latli_2 = np.reshape(latli_1, [-1, k])
-    lonli_2 = np.reshape(lonli_1, [-1, k])
+    latli_2 = np.asarray(latli)
+    lonli_2 = np.asarray(lonli)
+    if latli_2.size == 0:
+        return latli_2.reshape(0, k), lonli_2.reshape(0, k)
 
     # 取出每组邻居的实际值；在 mask 中通常 0=水，1=陆
-    test_arr = value[latli_2.tolist(), lonli_2.tolist()]  # 判断找出来的点里有哪些需要处理
+    test_arr = value[latli_2, lonli_2]
 
     # 每组求和，相当于统计这一组里有多少个陆点
-    sum_arr = np.sum(test_arr, axis=1)  #
+    sum_arr = np.sum(test_arr, axis=1)
 
     # 只有陆点数达到阈值，才把该目标点纳入后续腐蚀/填补
-    select_id = np.where(sum_arr >= judge_num)[0].tolist()
-
-    latli_r = latli_2[select_id[:]]
-    lonli_r = lonli_2[select_id[:]]
+    select_mask = sum_arr >= judge_num
 
-    return latli_r, lonli_r
+    return latli_2[select_mask], lonli_2[select_mask]
 
 
 def erosion_via_id(_weight, value):
@@ -286,6 +273,7 @@ def erosion_via_id(_weight, value):
     支持两种输入：
     - 2D: [lat, lon]
     - 3D: [time, lat, lon]
+        对 3D 输入会一次性批量处理整个 time 维，而不是逐时次循环。
 
     :param _weight: erosion_cal_id 返回的字典
     :param value: 二维或三维格点值
@@ -298,30 +286,151 @@ def erosion_via_id(_weight, value):
     _value_lat = _weight['value_lat']
     _value_lon = _weight['value_lon']
 
+    if len(_change_lat) == 0:
+        return value.copy()
+
     # 拷贝一份，避免直接改原数组
     new = value.copy()
+    
+    # 统一转成整数索引数组：
+    # value_lat_arr / value_lon_arr 形状约为 [n_target, n_neighbor]
+    # change_lat_arr / change_lon_arr 形状约为 [n_target]
+    # 用 np.intp 可避免高级索引时的额外类型转换
+    value_lat_arr = np.asarray(_value_lat, dtype=np.intp)
+    value_lon_arr = np.asarray(_value_lon, dtype=np.intp)
+    change_lat_arr = np.asarray(_change_lat, dtype=np.intp)
+    change_lon_arr = np.asarray(_change_lon, dtype=np.intp)
+
     # 3D 情况：常见于 [time, lat, lon]
     if value.ndim == 3:
-        for i in range(value.shape[0]):
-            # 取出当前时刻下，每个目标点对应的一组邻居值
-            replace = value[i, _value_lat, _value_lon]
-            # 对每组邻居求平均，得到每个目标点的新值
-            replace = np.nanmean(replace, axis=1)
-            # 写回当前时刻的目标点位置
-            new[i, _change_lat, _change_lon] = replace[:]
+        # 一次性取出所有时刻、所有目标点对应的邻居值，形状约为 [time, n_target, n_neighbor]
+        replace = value[:, value_lat_arr, value_lon_arr]
+        # 沿邻居维求平均，得到每个时刻、每个目标点的新值
+        replace = np.nanmean(replace, axis=2)
+        # 一次性写回，避免 Python 层逐时次循环
+        new[:, change_lat_arr, change_lon_arr] = replace
 
     # 2D 情况：直接对一个平面进行处理
     elif value.ndim == 2:
         # 取出每个目标点对应的一组邻居值
-        replace = value[_value_lat, _value_lon]
+        replace = value[value_lat_arr, value_lon_arr]
         # 每组邻居求平均
         replace = np.nanmean(replace, axis=1)
         # 把平均值写回目标点
-        new[_change_lat, _change_lon] = replace[:]
+        new[change_lat_arr, change_lon_arr] = replace
 
     return new
 
 
+def erosion_cal_id_multi_times(_times: int,
+                               g_lon,
+                               g_lat,
+                               value,
+                               k,
+                               judge_num,
+                               method='np.nan',
+                               ) -> List:
+    """
+    多次迭代执行 erosion_cal_id，适用于需要连续多次回填的情况
+    :param _times: 迭代次数
+    :param g_lon: 经度数组 传递给 erosion_cal_id 的参数
+    :param g_lat: 纬度数组 传递给 erosion_cal_id 的参数
+    :param value: 初始格点值 传递给 erosion_cal_id 的参数
+    :param k: 每次 erosion_cal_id 传递的参数
+    :param judge_num: 每次 erosion_cal_id 传递的参数
+    :param method: 每次 erosion_cal_id 传递的参数
+    :return: 每次迭代生成的索引列表
+    """
+    weights = []
+    for _ in range(_times):
+        weight = erosion_cal_id(g_lon, g_lat, value, k, judge_num, method)
+        weights.append(weight)
+        # 每次迭代后，更新 value，供下一次迭代使用
+        value = erosion_via_id(weight, value)
+    return weights
+
+
+def erosion_via_id_multi_times(weights: List, 
+                               value):
+    """
+    根据 erosion_cal_id_multi_times 生成的索引列表，连续多次执行回填
+    :param weights: erosion_cal_id_multi_times 返回的索引列表
+    :param value: 初始的格点值，后续迭代会在上一次回填的基础上继续修改
+    """
+    new_value = value.copy()
+    for weight in weights:
+        new_value = erosion_via_id(weight, new_value)
+    return new_value
+
+def erosion_via_id_multi_vars(weights: dict,
+                              viss: dict,
+                              _workers=-1,
+                              _method_multi: str = 'thread',
+                              ) -> dict:
+    """
+    根据 erosion_cal_id_multi_times 生成的索引列表，对多个变量执行回填
+    :param weights: erosion_cal_id 返回的索引列表
+    :param viss: 字典，key 是变量名，value 是对应的格点值数组
+    :param _workers: 并行线程数，默认为 -1（使用所有可用线程）
+    :param _method_multi: 并行方法，'thread' 或 'process'
+    :return: 回填后的变量字典
+    """
+    keys = list(viss.keys())
+    key_count = len(keys)
+
+    if key_count == 0:
+        return {}
+    elif key_count == 1:
+        viss[keys[0]] = erosion_via_id(weights, viss[keys[0]])
+    elif key_count > 1:
+        cpu_count = os.cpu_count() or 1
+        # _workers cpu_count key_count 取最小值，避免过度并行导致性能下降
+        if _workers < 0:
+            _workers = min(cpu_count, key_count)
+        else:            
+            _workers = min(_workers, cpu_count, key_count)
+
+        # 并行处理多个变量
+        if _method_multi == 'thread':
+            with ThreadPoolExecutor(max_workers=_workers) as executor:
+                futures = []
+                for key in keys:
+                    future = executor.submit(erosion_via_id, weights, viss[key])
+                    futures.append((key, future))
+                for key, future in futures:
+                    viss[key] = future.result()
+
+        elif _method_multi == 'process':
+            with ProcessPoolExecutor(max_workers=_workers) as ex:
+                future_map = {
+                    key: ex.submit(erosion_via_id, weights, viss[key])
+                    for key in keys
+                }
+                for key in keys:
+                    viss[key] = future_map[key].result()
+
+    return viss
+
+
+def erosion_via_id_multi_times_multi_vars(weights: List,
+                                          viss: dict,
+                                          _workers=-1,
+                                          _method_multi: str = 'thread',
+                                          ) -> dict:
+    """
+    根据 erosion_cal_id_multi_times 生成的索引列表，对多个变量连续多次执行回填
+    :param weights: erosion_cal_id_multi_times 返回的索引列表
+    :param viss: 字典，key 是变量名，value 是对应的格点值数组
+    :param _workers: 并行线程数，默认为 -1（使用所有可用线程）
+    :param _method_multi: 并行方法，'thread' 或 'process'
+    :return: 回填后的变量字典
+    """
+    new_viss = viss.copy()
+    for weight in weights:
+        new_viss = erosion_via_id_multi_vars(weight, new_viss, _workers=_workers, _method_multi=_method_multi)
+    return new_viss
+
+
 def example_cal_weight(methoood):  # sourcery skip: extract-duplicate-method
     """
     示例函数：演示如何针对不同类型的数据源计算 weight，并执行一次回填。

{alaea-2.4.0 → alaea-2.4.2/Alaea.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Alaea
-Version: 2.4.0
+Version: 2.4.2
 Summary: Some simple tools by Christmas
 Author: Christmas
 Author-email: 273519355@qq.com

{alaea-2.4.0/Alaea.egg-info → alaea-2.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Alaea
-Version: 2.4.0
+Version: 2.4.2
 Summary: Some simple tools by Christmas
 Author: Christmas
 Author-email: 273519355@qq.com

{alaea-2.4.0 → alaea-2.4.2}/setup.py

@@ -8,7 +8,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="Alaea",
-    version="2.4.0",
+    version="2.4.2",
     author="Christmas",
     author_email="273519355@qq.com",
     description="Some simple tools by Christmas",