oafuncs 0.0.97.16__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.
- {oafuncs-0.0.97.16/oafuncs.egg-info → oafuncs-0.0.97.17}/PKG-INFO +1 -1
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_modify.py +10 -2
- oafuncs-0.0.97.17/oafuncs/oa_cmap.py +278 -0
- oafuncs-0.0.97.17/oafuncs/oa_data.py +232 -0
- oafuncs-0.0.97.17/oafuncs/oa_date.py +160 -0
- oafuncs-0.0.97.17/oafuncs/oa_down/hycom_3hourly.py +1182 -0
- oafuncs-0.0.97.16/oafuncs/oa_down/hycom_3hourly.py → oafuncs-0.0.97.17/oafuncs/oa_down/hycom_3hourly_20250407.py +9 -7
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/idm.py +4 -4
- oafuncs-0.0.97.17/oafuncs/oa_draw.py +336 -0
- oafuncs-0.0.97.17/oafuncs/oa_file.py +407 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_help.py +10 -0
- oafuncs-0.0.97.17/oafuncs/oa_nc.py +283 -0
- oafuncs-0.0.97.17/oafuncs/oa_python.py +124 -0
- oafuncs-0.0.97.17/oafuncs/oa_tool.py +119 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17/oafuncs.egg-info}/PKG-INFO +1 -1
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/SOURCES.txt +2 -1
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/setup.py +1 -1
- oafuncs-0.0.97.16/oafuncs/oa_cmap.py +0 -162
- oafuncs-0.0.97.16/oafuncs/oa_data.py +0 -293
- oafuncs-0.0.97.16/oafuncs/oa_date.py +0 -126
- oafuncs-0.0.97.16/oafuncs/oa_draw.py +0 -236
- oafuncs-0.0.97.16/oafuncs/oa_file.py +0 -461
- oafuncs-0.0.97.16/oafuncs/oa_nc.py +0 -250
- oafuncs-0.0.97.16/oafuncs/oa_python.py +0 -98
- oafuncs-0.0.97.16/oafuncs/oa_tool.py +0 -80
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/LICENSE.txt +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/MANIFEST.in +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/README.md +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/__init__.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_data/OAFuncs.png +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_data/hycom_3hourly.png +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/cprogressbar.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/email.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_merge.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_write.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/parallel.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/parallel_example_usage.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/plot_dataset.py +0 -0
- /oafuncs-0.0.97.16/oafuncs/_script/replace_file_concent.py → /oafuncs-0.0.97.17/oafuncs/_script/replace_file_content.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/User_Agent-list.txt +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/__init__.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/literature.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/test_ua.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/user_agent.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/__init__.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/roms/__init__.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/roms/test.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/wrf/__init__.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/wrf/little_r.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/__init__.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/meteorological.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/ocean.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/scientific.py +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/dependency_links.txt +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/requires.txt +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/top_level.txt +0 -0
- {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/setup.cfg +0 -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
|
-
|
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
|
-
|
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
|