oafuncs 0.0.97.15__tar.gz → 0.0.97.17__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (57) hide show
  1. {oafuncs-0.0.97.15/oafuncs.egg-info → oafuncs-0.0.97.17}/PKG-INFO +1 -1
  2. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/cprogressbar.py +42 -20
  3. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_modify.py +10 -2
  4. oafuncs-0.0.97.17/oafuncs/oa_cmap.py +278 -0
  5. oafuncs-0.0.97.17/oafuncs/oa_data.py +232 -0
  6. oafuncs-0.0.97.17/oafuncs/oa_date.py +160 -0
  7. oafuncs-0.0.97.17/oafuncs/oa_down/hycom_3hourly.py +1182 -0
  8. oafuncs-0.0.97.15/oafuncs/oa_down/hycom_3hourly.py → oafuncs-0.0.97.17/oafuncs/oa_down/hycom_3hourly_20250407.py +9 -7
  9. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_down/idm.py +4 -4
  10. oafuncs-0.0.97.17/oafuncs/oa_draw.py +336 -0
  11. oafuncs-0.0.97.17/oafuncs/oa_file.py +407 -0
  12. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_help.py +10 -0
  13. oafuncs-0.0.97.17/oafuncs/oa_nc.py +283 -0
  14. oafuncs-0.0.97.17/oafuncs/oa_python.py +124 -0
  15. oafuncs-0.0.97.17/oafuncs/oa_tool.py +119 -0
  16. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17/oafuncs.egg-info}/PKG-INFO +1 -1
  17. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs.egg-info/SOURCES.txt +2 -1
  18. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/setup.py +1 -1
  19. oafuncs-0.0.97.15/oafuncs/oa_cmap.py +0 -162
  20. oafuncs-0.0.97.15/oafuncs/oa_data.py +0 -293
  21. oafuncs-0.0.97.15/oafuncs/oa_date.py +0 -126
  22. oafuncs-0.0.97.15/oafuncs/oa_draw.py +0 -236
  23. oafuncs-0.0.97.15/oafuncs/oa_file.py +0 -461
  24. oafuncs-0.0.97.15/oafuncs/oa_nc.py +0 -250
  25. oafuncs-0.0.97.15/oafuncs/oa_python.py +0 -98
  26. oafuncs-0.0.97.15/oafuncs/oa_tool.py +0 -83
  27. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/LICENSE.txt +0 -0
  28. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/MANIFEST.in +0 -0
  29. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/README.md +0 -0
  30. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/__init__.py +0 -0
  31. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_data/OAFuncs.png +0 -0
  32. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_data/hycom_3hourly.png +0 -0
  33. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/email.py +0 -0
  34. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_merge.py +0 -0
  35. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_write.py +0 -0
  36. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/parallel.py +0 -0
  37. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/parallel_example_usage.py +0 -0
  38. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/_script/plot_dataset.py +0 -0
  39. /oafuncs-0.0.97.15/oafuncs/_script/replace_file_concent.py → /oafuncs-0.0.97.17/oafuncs/_script/replace_file_content.py +0 -0
  40. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_down/User_Agent-list.txt +0 -0
  41. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_down/__init__.py +0 -0
  42. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_down/literature.py +0 -0
  43. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_down/test_ua.py +0 -0
  44. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_down/user_agent.py +0 -0
  45. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_model/__init__.py +0 -0
  46. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_model/roms/__init__.py +0 -0
  47. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_model/roms/test.py +0 -0
  48. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_model/wrf/__init__.py +0 -0
  49. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_model/wrf/little_r.py +0 -0
  50. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_sign/__init__.py +0 -0
  51. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_sign/meteorological.py +0 -0
  52. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_sign/ocean.py +0 -0
  53. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs/oa_sign/scientific.py +0 -0
  54. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs.egg-info/dependency_links.txt +0 -0
  55. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs.egg-info/requires.txt +0 -0
  56. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/oafuncs.egg-info/top_level.txt +0 -0
  57. {oafuncs-0.0.97.15 → oafuncs-0.0.97.17}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: oafuncs
3
- Version: 0.0.97.15
3
+ Version: 0.0.97.17
4
4
  Summary: Oceanic and Atmospheric Functions
5
5
  Home-page: https://github.com/Industry-Pays/OAFuncs
6
6
  Author: Kun Liu
@@ -153,6 +153,7 @@ class ColorProgressBar:
153
153
  update_interval: float = 0.1,
154
154
  bar_length: int = None,
155
155
  speed_estimate_period: float = 30.0,
156
+ next_line: bool = False,
156
157
  ):
157
158
  self.iterable = iterable
158
159
  self.description = description
@@ -161,7 +162,7 @@ class ColorProgressBar:
161
162
  self.update_interval = update_interval
162
163
  self.bar_length = bar_length
163
164
  self.speed_estimate_period = speed_estimate_period
164
- self.mark_num = random.randint(0, 13) # 随机选择填充字符的数量
165
+ self.next_line = next_line
165
166
 
166
167
  # 线程安全锁
167
168
  self._lock = threading.RLock()
@@ -187,6 +188,10 @@ class ColorProgressBar:
187
188
  self._is_terminal = hasattr(self._file, "isatty") and self._file.isatty()
188
189
  self._is_jupyter = "ipykernel" in sys.modules
189
190
 
191
+ # 输出样式
192
+ filled_list = ["▊", "█", "▓", "▒", "░", "#", "=", ">", "▌", "▍", "▎", "▏", "*"]
193
+ self.filled = random.choice(filled_list)
194
+
190
195
  def _generate_gradient(self) -> Optional[List[str]]:
191
196
  """生成渐变色列表(修复内置colormap支持)"""
192
197
  try:
@@ -253,11 +258,7 @@ class ColorProgressBar:
253
258
 
254
259
  def _format_bar(self, progress: float, width: int) -> str:
255
260
  """格式化进度条显示"""
256
- filled_list = ["▊", "█", "▓", "▒", "░", "#", "=", ">", "▌", "▍", "▎", "▏", "*"]
257
- # filled = "░"
258
- # filled = filled_list[2] # 选择中间的填充字符
259
- # filled = random.choice(filled_list) # 随机选择填充字符
260
- filled = filled_list[min(self.mark_num, len(filled_list) - 1)] # 随机选择填充字符
261
+ filled = self.filled
261
262
  empty = " "
262
263
  # 为其他信息保留更多空间
263
264
  max_width = max(10, width - 60) # 至少保留10个字符的进度条
@@ -363,7 +364,10 @@ class ColorProgressBar:
363
364
  self._file.write(f"{line}\n")
364
365
  elif self._is_terminal:
365
366
  # 标准终端环境,覆盖同一行
366
- self._file.write(f"\r{line}")
367
+ if self.next_line:
368
+ self._file.write(f"\r{line}\n")
369
+ else:
370
+ self._file.write(f"\r{line}")
367
371
  else:
368
372
  # 非交互式环境,仅在完成时输出
369
373
  if self.task.finished:
@@ -395,7 +399,7 @@ class ColorProgressBar:
395
399
 
396
400
  finally:
397
401
  # 完成后进行清理
398
- if self._is_terminal and not self._is_jupyter:
402
+ if self._is_terminal and not self._is_jupyter and not self.next_line:
399
403
  self._file.write("\n")
400
404
  self._file.flush()
401
405
 
@@ -406,28 +410,46 @@ class ColorProgressBar:
406
410
  return [to_hex(cmap(i)) for i in np.linspace(0, 1, n)]
407
411
 
408
412
 
413
+ def cpbar(
414
+ iterable,
415
+ description="Working...",
416
+ total=None,
417
+ completed=0,
418
+ color="cyan",
419
+ cmap=None,
420
+ update_interval=0.1,
421
+ bar_length=None,
422
+ speed_estimate_period=30.0,
423
+ next_line=False,
424
+ ):
425
+ """便捷函数,返回 ColorProgressBar 对象"""
426
+ return ColorProgressBar(
427
+ iterable=iterable,
428
+ description=description,
429
+ total=total,
430
+ completed=completed,
431
+ color=color,
432
+ cmap=cmap,
433
+ update_interval=update_interval,
434
+ bar_length=bar_length,
435
+ speed_estimate_period=speed_estimate_period,
436
+ next_line=next_line,
437
+ )
438
+
439
+
409
440
  # 验证示例
410
441
  if __name__ == "__main__":
411
- for _ in ColorProgressBar(range(100), description="Diverging:", cmap="diverging_1"):
412
- # print("\rTesting...", end="")
442
+ for _ in cpbar(range(20), description="Diverging:", cmap="diverging_1", next_line=True):
443
+ print("Processing...")
413
444
  time.sleep(0.1)
414
445
 
415
- for _ in ColorProgressBar(range(100), description="Colorful:", cmap="colorful_1"):
416
- time.sleep(0.1)
417
-
418
- for _ in ColorProgressBar(range(100), description="Viridis:", cmap="viridis"):
446
+ for _ in ColorProgressBar(range(20), description="Viridis:", cmap="viridis"):
419
447
  time.sleep(0.1)
420
448
 
421
449
  # 使用自定义渐变色
422
450
  for _ in ColorProgressBar(range(50), description="Custom_cmap:", cmap=["#FF0000", "#0000FF"]):
423
451
  time.sleep(0.1)
424
452
 
425
- for _ in ColorProgressBar(range(50), description="Default:"):
426
- time.sleep(0.1)
427
-
428
- for _ in ColorProgressBar(range(50), description="Custom_color:", color="#FF00FF"):
429
- time.sleep(0.1)
430
-
431
453
  # 测试无法获取长度的迭代器
432
454
  def infinite_generator():
433
455
  i = 0
@@ -42,8 +42,13 @@ def _modify_var(nc_file_path, variable_name, new_value):
42
42
  raise ValueError(f"File '{nc_file_path}' is not a valid NetCDF file.")
43
43
  if not variable_name:
44
44
  raise ValueError("Variable name cannot be empty or None.")
45
+
46
+ # 自动尝试将 new_value 转换为 numpy.ndarray
45
47
  if not isinstance(new_value, np.ndarray):
46
- raise TypeError("New value must be a numpy.ndarray.")
48
+ try:
49
+ new_value = np.array(new_value)
50
+ except Exception:
51
+ raise TypeError("New value must be a numpy.ndarray or convertible to numpy.ndarray.")
47
52
 
48
53
  try:
49
54
  with nc.Dataset(nc_file_path, "r+") as dataset:
@@ -51,7 +56,10 @@ def _modify_var(nc_file_path, variable_name, new_value):
51
56
  raise ValueError(f"Variable '{variable_name}' not found in the NetCDF file.")
52
57
  variable = dataset.variables[variable_name]
53
58
  if variable.shape != new_value.shape:
54
- raise ValueError(f"Shape mismatch: Variable '{variable_name}' has shape {variable.shape}, but new value has shape {new_value.shape}.")
59
+ try:
60
+ new_value = new_value.reshape(variable.shape)
61
+ except ValueError:
62
+ raise ValueError(f"Shape mismatch: Variable '{variable_name}' has shape {variable.shape}, but new value has shape {new_value.shape}. Reshaping failed.")
55
63
  variable[:] = new_value
56
64
  print(f"[green]Successfully modified variable '{variable_name}' in '{nc_file_path}'.[/green]")
57
65
  return True
@@ -0,0 +1,278 @@
1
+ from typing import List, Optional, Union
2
+
3
+ import matplotlib as mpl
4
+ import matplotlib.pyplot as plt
5
+ import numpy as np
6
+
7
+ __all__ = ["show", "to_color", "create", "get"]
8
+
9
+
10
+ # ** 将cmap用填色图可视化(官网摘抄函数)
11
+ def show(colormaps: Union[str, mpl.colors.Colormap, List[Union[str, mpl.colors.Colormap]]]) -> None:
12
+ """Helper function to plot data with associated colormap.
13
+
14
+ This function creates a visualization of one or more colormaps by applying them
15
+ to randomly generated data in a pcolormesh plot.
16
+
17
+ Parameters
18
+ ----------
19
+ colormaps : Union[str, mpl.colors.Colormap, List[Union[str, mpl.colors.Colormap]]]
20
+ List of colormaps, or a single colormap; can be a string name or a colormap object.
21
+
22
+ Returns
23
+ -------
24
+ None
25
+ This function displays the plot but does not return any value.
26
+
27
+ Examples
28
+ --------
29
+ >>> cmap = matplotlib.colors.ListedColormap(["darkorange", "gold", "lawngreen", "lightseagreen"])
30
+ >>> show([cmap])
31
+ >>> show("viridis")
32
+ >>> show(["viridis", "cividis"])
33
+ """
34
+ # Convert single colormap to list for uniform processing
35
+ if not isinstance(colormaps, list):
36
+ colormaps = [colormaps]
37
+
38
+ # Create a formatted list of colormap names for display
39
+ cmap_names = []
40
+ for cmap in colormaps:
41
+ if isinstance(cmap, str):
42
+ cmap_names.append(cmap)
43
+ elif hasattr(cmap, "name"):
44
+ cmap_names.append(cmap.name)
45
+ else:
46
+ cmap_names.append("unnamed_colormap")
47
+
48
+ print(f"Visualizing {len(colormaps)} colormap(s): {', '.join(cmap_names)}")
49
+
50
+ # Generate random data with fixed seed for reproducibility
51
+ np.random.seed(19680801)
52
+ data = np.random.randn(30, 30)
53
+
54
+ # Create subplots based on number of colormaps
55
+ n = len(colormaps)
56
+ fig, axs = plt.subplots(1, n, figsize=(n * 2 + 2, 3), constrained_layout=True, squeeze=False)
57
+
58
+ # Plot each colormap
59
+ for ax, cmap in zip(axs.flat, colormaps):
60
+ psm = ax.pcolormesh(data, cmap=cmap, rasterized=True, vmin=-4, vmax=4)
61
+ fig.colorbar(psm, ax=ax)
62
+
63
+ # Set title if colormap has a name
64
+ if isinstance(cmap, str):
65
+ ax.set_title(cmap)
66
+ elif hasattr(cmap, "name") and cmap.name:
67
+ ax.set_title(cmap.name)
68
+
69
+ print("Displaying colormap visualization...")
70
+ plt.show()
71
+
72
+
73
+ # ** 将cmap转为list,即多个颜色的列表
74
+ def to_color(colormap_name: str, num_colors: int = 256) -> List[tuple]:
75
+ """Convert a colormap to a list of colors.
76
+
77
+ Args:
78
+ colormap_name (str): The name of the colormap.
79
+ num_colors (int, optional): The number of colors. Defaults to 256.
80
+
81
+ Returns:
82
+ List[tuple]: List of RGBA colors.
83
+
84
+ Raises:
85
+ ValueError: If the colormap name is not recognized.
86
+
87
+ Examples:
88
+ >>> out_colors = to_color('viridis', 256)
89
+ """
90
+ try:
91
+ cmap = mpl.colormaps.get_cmap(colormap_name)
92
+ return [cmap(i) for i in np.linspace(0, 1, num_colors)]
93
+ except (ValueError, TypeError):
94
+ error_msg = f"Invalid colormap name: {colormap_name}"
95
+ print(error_msg)
96
+ raise ValueError(error_msg)
97
+
98
+
99
+ # ** 自制cmap,多色,可带位置
100
+ def create(color_list: Optional[List[Union[str, tuple]]] = None, rgb_file: Optional[str] = None, color_positions: Optional[List[float]] = None, below_range_color: Optional[Union[str, tuple]] = None, above_range_color: Optional[Union[str, tuple]] = None, value_delimiter: str = ",") -> mpl.colors.Colormap:
101
+ """Create a custom colormap from a list of colors or an RGB txt document.
102
+
103
+ Args:
104
+ color_list (Optional[List[Union[str, tuple]]]): List of colors. Required if rgb_file is None.
105
+ rgb_file (Optional[str]): The path of txt file. Required if color_list is None.
106
+ color_positions (Optional[List[float]]): List of positions for color_list. Must have same length as color_list.
107
+ below_range_color (Optional[Union[str, tuple]]): Color for values below the colormap range.
108
+ above_range_color (Optional[Union[str, tuple]]): Color for values above the colormap range.
109
+ value_delimiter (str, optional): The delimiter of RGB values in txt file. Defaults to ",".
110
+
111
+ Returns:
112
+ mpl.colors.Colormap: Created colormap.
113
+
114
+ Raises:
115
+ ValueError: If neither color_list nor rgb_file is provided.
116
+ ValueError: If color_positions is provided but has different length than color_list.
117
+ FileNotFoundError: If rgb_file does not exist.
118
+ ValueError: If the RGB file format is invalid.
119
+
120
+ Examples:
121
+ >>> cmap = create(color_list=['#C2B7F3','#B3BBF2','#B0CBF1','#ACDCF0','#A8EEED'])
122
+ >>> cmap = create(color_list=['aliceblue','skyblue','deepskyblue'], color_positions=[0.0,0.5,1.0])
123
+ >>> cmap = create(rgb_file='path/to/file.txt', value_delimiter=',')
124
+ """
125
+ # Input validation
126
+ if rgb_file is None and color_list is None:
127
+ error_msg = "Either 'color_list' or 'rgb_file' must be provided."
128
+ print(error_msg)
129
+ raise ValueError(error_msg)
130
+
131
+ if color_positions is not None and color_list is not None:
132
+ if len(color_positions) != len(color_list):
133
+ error_msg = f"'color_positions' must have the same length as 'color_list' (positions: {len(color_positions)}, colors: {len(color_list)})"
134
+ print(error_msg)
135
+ raise ValueError(error_msg)
136
+ if not all(0 <= pos <= 1 for pos in color_positions):
137
+ error_msg = "All position values must be between 0 and 1"
138
+ print(error_msg)
139
+ raise ValueError(error_msg)
140
+ if color_positions != sorted(color_positions):
141
+ error_msg = f"Position values must be in ascending order: {color_positions}"
142
+ print(error_msg)
143
+ raise ValueError(error_msg)
144
+
145
+ if rgb_file:
146
+ try:
147
+ print(f"Reading RGB data from {rgb_file}...")
148
+
149
+ with open(rgb_file) as fid:
150
+ data = [line.strip() for line in fid if line.strip() and not line.strip().startswith("#")]
151
+
152
+ if not data:
153
+ error_msg = f"RGB file is empty or contains only comments: {rgb_file}"
154
+ print(error_msg)
155
+ raise ValueError(error_msg)
156
+
157
+ n = len(data)
158
+ rgb = np.zeros((n, 3))
159
+
160
+ for i in np.arange(n):
161
+ try:
162
+ parts = data[i].split(value_delimiter)
163
+ if len(parts) < 3:
164
+ error_msg = f"Line {i + 1}: Expected at least 3 values, got {len(parts)}"
165
+ print(error_msg)
166
+ raise ValueError(error_msg)
167
+
168
+ rgb[i][0] = float(parts[0])
169
+ rgb[i][1] = float(parts[1])
170
+ rgb[i][2] = float(parts[2])
171
+ except (ValueError, IndexError) as e:
172
+ error_msg = f"Error parsing RGB values at line {i + 1}: {e}"
173
+ print(error_msg)
174
+ raise ValueError(error_msg)
175
+
176
+ max_rgb = np.max(rgb)
177
+ # Normalize RGB values if they are in 0-255 range
178
+ if max_rgb > 2:
179
+ rgb = rgb / 255.0
180
+ cmap_color = mpl.colors.ListedColormap(rgb, name="my_color")
181
+ print(f"Successfully created colormap from {rgb_file}")
182
+ except FileNotFoundError:
183
+ error_msg = f"RGB file not found: {rgb_file}"
184
+ print(error_msg)
185
+ raise FileNotFoundError(error_msg)
186
+ else:
187
+ # Create colormap from color list
188
+ if color_positions is None:
189
+ cmap_color = mpl.colors.LinearSegmentedColormap.from_list("mycmap", color_list)
190
+ else:
191
+ cmap_color = mpl.colors.LinearSegmentedColormap.from_list("mycmap", list(zip(color_positions, color_list)))
192
+ print(f"Successfully created colormap from {len(color_list)} colors")
193
+
194
+ # Set below/above range colors if provided
195
+ if below_range_color is not None:
196
+ cmap_color.set_under(below_range_color)
197
+ print(f"Set below-range color to {below_range_color}")
198
+ if above_range_color is not None:
199
+ cmap_color.set_over(above_range_color)
200
+ print(f"Set above-range color to {above_range_color}")
201
+
202
+ return cmap_color
203
+
204
+
205
+ # ** 选择cmap
206
+ def get(colormap_name: Optional[str] = None, show_available: bool = False) -> Optional[mpl.colors.Colormap]:
207
+ """Choose a colormap from the list of available colormaps or a custom colormap.
208
+
209
+ Args:
210
+ colormap_name (Optional[str], optional): The name of the colormap. Defaults to None.
211
+ show_available (bool, optional): Whether to query the available colormap names. Defaults to False.
212
+
213
+ Returns:
214
+ Optional[mpl.colors.Colormap]: Selected colormap or None if show_available is True or colormap_name is None.
215
+
216
+ Examples:
217
+ >>> cmap = get('viridis')
218
+ >>> cmap = get('diverging_1')
219
+ >>> cmap = get('cool_1')
220
+ >>> cmap = get('warm_1')
221
+ >>> cmap = get('colorful_1')
222
+ """
223
+ my_cmap_dict = {
224
+ "diverging_1": ["#4e00b3", "#0000FF", "#00c0ff", "#a1d3ff", "#DCDCDC", "#FFD39B", "#FF8247", "#FF0000", "#FF5F9E"],
225
+ "cool_1": ["#4e00b3", "#0000FF", "#00c0ff", "#a1d3ff", "#DCDCDC"],
226
+ "warm_1": ["#DCDCDC", "#FFD39B", "#FF8247", "#FF0000", "#FF5F9E"],
227
+ "colorful_1": ["#6d00db", "#9800cb", "#F2003C", "#ff4500", "#ff7f00", "#FE28A2", "#FFC0CB", "#DDA0DD", "#40E0D0", "#1a66f2", "#00f7fb", "#8fff88", "#E3FF00"],
228
+ }
229
+
230
+ if show_available:
231
+ print("Available cmap names:")
232
+ print("-" * 20)
233
+ print("Defined by myself:")
234
+ for name in my_cmap_dict.keys():
235
+ print(f" • {name}")
236
+ print("-" * 20)
237
+ print("Matplotlib built-in:")
238
+ # 将Matplotlib内置cmap分批次打印,每行5个
239
+ built_in_cmaps = list(mpl.colormaps.keys())
240
+ for i in range(0, len(built_in_cmaps), 5):
241
+ print(" • " + ", ".join(built_in_cmaps[i : i + 5]))
242
+ print("-" * 20)
243
+ return None
244
+
245
+ if colormap_name is None:
246
+ return None
247
+
248
+ if colormap_name in my_cmap_dict:
249
+ print(f"Using custom colormap: {colormap_name}")
250
+ return create(my_cmap_dict[colormap_name])
251
+ else:
252
+ try:
253
+ cmap = mpl.colormaps.get_cmap(colormap_name)
254
+ print(f"Using matplotlib colormap: {colormap_name}")
255
+ return cmap
256
+ except ValueError:
257
+ print(f"Warning: Unknown cmap name: {colormap_name}")
258
+ print("Using rainbow as default.")
259
+ return mpl.colormaps.get_cmap("rainbow") # 默认返回 'rainbow'
260
+
261
+
262
+ if __name__ == "__main__":
263
+ # ** 测试自制cmap
264
+ colors = ["#C2B7F3", "#B3BBF2", "#B0CBF1", "#ACDCF0", "#A8EEED"]
265
+ color_nodes = [0.0, 0.2, 0.4, 0.6, 1.0]
266
+ custom_cmap = create(colors, color_nodes)
267
+ show([custom_cmap])
268
+
269
+ # ** 测试自制diverging型cmap
270
+ diverging_cmap = create(["#4e00b3", "#0000FF", "#00c0ff", "#a1d3ff", "#DCDCDC", "#FFD39B", "#FF8247", "#FF0000", "#FF5F9E"])
271
+ show([diverging_cmap])
272
+
273
+ # ** 测试根据RGB的txt文档制作色卡
274
+ rgb_file_path = "E:/python/colorbar/test.txt"
275
+ cmap_from_rgb = create(rgb_file=rgb_file_path)
276
+
277
+ # ** 测试将cmap转为list
278
+ viridis_colors = to_color("viridis", 256)
@@ -0,0 +1,232 @@
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+ """
4
+ Author: Liu Kun && 16031215@qq.com
5
+ Date: 2024-09-17 17:12:47
6
+ LastEditors: Liu Kun && 16031215@qq.com
7
+ LastEditTime: 2024-12-13 19:11:08
8
+ FilePath: \\Python\\My_Funcs\\OAFuncs\\oafuncs\\oa_data.py
9
+ Description:
10
+ EditPlatform: vscode
11
+ ComputerInfo: XPS 15 9510
12
+ SystemInfo: Windows 11
13
+ Python Version: 3.11
14
+ """
15
+
16
+ import itertools
17
+ import multiprocessing as mp
18
+ from concurrent.futures import ThreadPoolExecutor
19
+ from typing import Any, List, Union
20
+
21
+ import numpy as np
22
+ import salem
23
+ import xarray as xr
24
+ from rich import print
25
+ from scipy.interpolate import griddata, interp1d
26
+
27
+ __all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile"]
28
+
29
+
30
+ def ensure_list(input_value: Any) -> List[str]:
31
+ """
32
+ Ensure the input is converted into a list.
33
+
34
+ Args:
35
+ input_value (Any): The input which can be a list, a string, or any other type.
36
+
37
+ Returns:
38
+ List[str]: A list containing the input or the string representation of the input.
39
+ """
40
+ if isinstance(input_value, list):
41
+ return input_value
42
+ elif isinstance(input_value, str):
43
+ return [input_value]
44
+ else:
45
+ return [str(input_value)]
46
+
47
+
48
+ def interp_along_dim(
49
+ target_coordinates: np.ndarray,
50
+ source_coordinates: Union[np.ndarray, List[float]],
51
+ source_data: np.ndarray,
52
+ interpolation_axis: int = -1,
53
+ interpolation_method: str = "linear",
54
+ extrapolation_method: str = "linear",
55
+ ) -> np.ndarray:
56
+ """
57
+ Perform interpolation and extrapolation along a specified dimension.
58
+
59
+ Args:
60
+ target_coordinates (np.ndarray): 1D array of target coordinate points.
61
+ source_coordinates (Union[np.ndarray, List[float]]): Source coordinate points (1D or ND array).
62
+ source_data (np.ndarray): Source data array to interpolate.
63
+ interpolation_axis (int, optional): Axis to perform interpolation on. Defaults to -1.
64
+ interpolation_method (str, optional): Interpolation method. Defaults to "linear".
65
+ extrapolation_method (str, optional): Extrapolation method. Defaults to "linear".
66
+
67
+ Returns:
68
+ np.ndarray: Interpolated data array.
69
+
70
+ Raises:
71
+ ValueError: If input dimensions or shapes are invalid.
72
+
73
+ Examples:
74
+ >>> target_coordinates = np.array([1, 2, 3])
75
+ >>> source_coordinates = np.array([0, 1, 2, 3])
76
+ >>> source_data = np.array([10, 20, 30, 40])
77
+ >>> result = interp_along_dim(target_coordinates, source_coordinates, source_data)
78
+ >>> print(result) # Expected output: [20.0, 30.0]
79
+ """
80
+ target_coordinates = np.asarray(target_coordinates)
81
+ if target_coordinates.ndim != 1:
82
+ raise ValueError("[red]target_coordinates must be a 1D array.[/red]")
83
+
84
+ source_coordinates = np.asarray(source_coordinates)
85
+ source_data = np.asarray(source_data)
86
+
87
+ if source_data.ndim == 1 and source_coordinates.ndim == 1:
88
+ if len(source_coordinates) != len(source_data):
89
+ raise ValueError("[red]For 1D data, source_coordinates and source_data must have the same length.[/red]")
90
+
91
+ interpolator = interp1d(source_coordinates, source_data, kind=interpolation_method, fill_value="extrapolate", bounds_error=False)
92
+ return interpolator(target_coordinates)
93
+
94
+ if source_coordinates.ndim == 1:
95
+ shape = [1] * source_data.ndim
96
+ shape[interpolation_axis] = source_coordinates.shape[0]
97
+ source_coordinates = np.reshape(source_coordinates, shape)
98
+ source_coordinates = np.broadcast_to(source_coordinates, source_data.shape)
99
+ elif source_coordinates.shape != source_data.shape:
100
+ raise ValueError("[red]source_coordinates and source_data must have the same shape.[/red]")
101
+
102
+ def apply_interp_extrap(arr: np.ndarray) -> np.ndarray:
103
+ xp = np.moveaxis(source_coordinates, interpolation_axis, 0)
104
+ xp = xp[:, 0] if xp.ndim > 1 else xp
105
+ arr = np.moveaxis(arr, interpolation_axis, 0)
106
+ interpolator = interp1d(xp, arr, kind=interpolation_method, fill_value="extrapolate", bounds_error=False)
107
+ interpolated = interpolator(target_coordinates)
108
+ if extrapolation_method != interpolation_method:
109
+ mask_extrap = (target_coordinates < xp.min()) | (target_coordinates > xp.max())
110
+ if np.any(mask_extrap):
111
+ extrap_interpolator = interp1d(xp, arr, kind=extrapolation_method, fill_value="extrapolate", bounds_error=False)
112
+ interpolated[mask_extrap] = extrap_interpolator(target_coordinates[mask_extrap])
113
+ return np.moveaxis(interpolated, 0, interpolation_axis)
114
+
115
+ return np.apply_along_axis(apply_interp_extrap, interpolation_axis, source_data)
116
+
117
+
118
+ def interp_2d(
119
+ target_x_coordinates: Union[np.ndarray, List[float]],
120
+ target_y_coordinates: Union[np.ndarray, List[float]],
121
+ source_x_coordinates: Union[np.ndarray, List[float]],
122
+ source_y_coordinates: Union[np.ndarray, List[float]],
123
+ source_data: np.ndarray,
124
+ interpolation_method: str = "linear",
125
+ use_parallel: bool = True,
126
+ ) -> np.ndarray:
127
+ """
128
+ Perform 2D interpolation on the last two dimensions of a multi-dimensional array.
129
+
130
+ Args:
131
+ target_x_coordinates (Union[np.ndarray, List[float]]): Target grid's x-coordinates.
132
+ target_y_coordinates (Union[np.ndarray, List[float]]): Target grid's y-coordinates.
133
+ source_x_coordinates (Union[np.ndarray, List[float]]): Original grid's x-coordinates.
134
+ source_y_coordinates (Union[np.ndarray, List[float]]): Original grid's y-coordinates.
135
+ source_data (np.ndarray): Multi-dimensional array with the last two dimensions as spatial.
136
+ interpolation_method (str, optional): Interpolation method. Defaults to "linear".
137
+ use_parallel (bool, optional): Enable parallel processing. Defaults to True.
138
+
139
+ Returns:
140
+ np.ndarray: Interpolated data array.
141
+
142
+ Raises:
143
+ ValueError: If input shapes are invalid.
144
+
145
+ Examples:
146
+ >>> target_x_coordinates = np.array([1, 2, 3])
147
+ >>> target_y_coordinates = np.array([4, 5, 6])
148
+ >>> source_x_coordinates = np.array([7, 8, 9])
149
+ >>> source_y_coordinates = np.array([10, 11, 12])
150
+ >>> source_data = np.random.rand(3, 3)
151
+ >>> result = interp_2d(target_x_coordinates, target_y_coordinates, source_x_coordinates, source_y_coordinates, source_data)
152
+ >>> print(result.shape) # Expected output: (3, 3)
153
+ """
154
+
155
+ def interp_single(data_slice: np.ndarray, target_points: np.ndarray, origin_points: np.ndarray, method: str) -> np.ndarray:
156
+ return griddata(origin_points, data_slice.ravel(), target_points, method=method).reshape(target_y_coordinates.shape)
157
+
158
+ if len(target_y_coordinates.shape) == 1:
159
+ target_x_coordinates, target_y_coordinates = np.meshgrid(target_x_coordinates, target_y_coordinates)
160
+ if len(source_y_coordinates.shape) == 1:
161
+ source_x_coordinates, source_y_coordinates = np.meshgrid(source_x_coordinates, source_y_coordinates)
162
+
163
+ if source_x_coordinates.shape != source_data.shape[-2:] or source_y_coordinates.shape != source_data.shape[-2:]:
164
+ raise ValueError("[red]Shape of source_data does not match shape of source_x_coordinates or source_y_coordinates.[/red]")
165
+
166
+ target_points = np.column_stack((np.array(target_y_coordinates).ravel(), np.array(target_x_coordinates).ravel()))
167
+ origin_points = np.column_stack((np.array(source_y_coordinates).ravel(), np.array(source_x_coordinates).ravel()))
168
+
169
+ if use_parallel:
170
+ with ThreadPoolExecutor(max_workers=mp.cpu_count() - 2) as executor:
171
+ if len(source_data.shape) == 2:
172
+ interpolated_data = list(executor.map(interp_single, [source_data], [target_points], [origin_points], [interpolation_method]))
173
+ elif len(source_data.shape) == 3:
174
+ interpolated_data = list(executor.map(interp_single, [source_data[i] for i in range(source_data.shape[0])], [target_points] * source_data.shape[0], [origin_points] * source_data.shape[0], [interpolation_method] * source_data.shape[0]))
175
+ elif len(source_data.shape) == 4:
176
+ index_combinations = list(itertools.product(range(source_data.shape[0]), range(source_data.shape[1])))
177
+ interpolated_data = list(executor.map(interp_single, [source_data[i, j] for i, j in index_combinations], [target_points] * len(index_combinations), [origin_points] * len(index_combinations), [interpolation_method] * len(index_combinations)))
178
+ interpolated_data = np.array(interpolated_data).reshape(source_data.shape[0], source_data.shape[1], *target_y_coordinates.shape)
179
+ else:
180
+ if len(source_data.shape) == 2:
181
+ interpolated_data = interp_single(source_data, target_points, origin_points, interpolation_method)
182
+ elif len(source_data.shape) == 3:
183
+ interpolated_data = np.stack([interp_single(source_data[i], target_points, origin_points, interpolation_method) for i in range(source_data.shape[0])])
184
+ elif len(source_data.shape) == 4:
185
+ interpolated_data = np.stack([np.stack([interp_single(source_data[i, j], target_points, origin_points, interpolation_method) for j in range(source_data.shape[1])]) for i in range(source_data.shape[0])])
186
+
187
+ return np.squeeze(np.array(interpolated_data))
188
+
189
+
190
+ def mask_shapefile(
191
+ data_array: np.ndarray,
192
+ longitudes: np.ndarray,
193
+ latitudes: np.ndarray,
194
+ shapefile_path: str,
195
+ ) -> Union[xr.DataArray, None]:
196
+ """
197
+ Mask a 2D data array using a shapefile.
198
+
199
+ Args:
200
+ data_array (np.ndarray): 2D array of data to be masked.
201
+ longitudes (np.ndarray): 1D array of longitudes.
202
+ latitudes (np.ndarray): 1D array of latitudes.
203
+ shapefile_path (str): Path to the shapefile used for masking.
204
+
205
+ Returns:
206
+ Union[xr.DataArray, None]: Masked xarray DataArray or None if an error occurs.
207
+
208
+ Raises:
209
+ FileNotFoundError: If the shapefile does not exist.
210
+ ValueError: If the data dimensions do not match the coordinates.
211
+
212
+ Examples:
213
+ >>> data_array = np.random.rand(10, 10)
214
+ >>> longitudes = np.linspace(-180, 180, 10)
215
+ >>> latitudes = np.linspace(-90, 90, 10)
216
+ >>> shapefile_path = "path/to/shapefile.shp"
217
+ >>> masked_data = mask_shapefile(data_array, longitudes, latitudes, shapefile_path)
218
+ >>> print(masked_data) # Expected output: Masked DataArray
219
+
220
+ """
221
+ try:
222
+ shp_f = salem.read_shapefile(shapefile_path)
223
+ data_da = xr.DataArray(data_array, coords=[("latitude", latitudes), ("longitude", longitudes)])
224
+ masked_data = data_da.salem.roi(shape=shp_f)
225
+ return masked_data
226
+ except Exception as e:
227
+ print(f"[red]An error occurred: {e}[/red]")
228
+ return None
229
+
230
+
231
+ if __name__ == "__main__":
232
+ pass