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.
Files changed (57) hide show
  1. {oafuncs-0.0.97.16/oafuncs.egg-info → oafuncs-0.0.97.17}/PKG-INFO +1 -1
  2. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_modify.py +10 -2
  3. oafuncs-0.0.97.17/oafuncs/oa_cmap.py +278 -0
  4. oafuncs-0.0.97.17/oafuncs/oa_data.py +232 -0
  5. oafuncs-0.0.97.17/oafuncs/oa_date.py +160 -0
  6. oafuncs-0.0.97.17/oafuncs/oa_down/hycom_3hourly.py +1182 -0
  7. 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
  8. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/idm.py +4 -4
  9. oafuncs-0.0.97.17/oafuncs/oa_draw.py +336 -0
  10. oafuncs-0.0.97.17/oafuncs/oa_file.py +407 -0
  11. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_help.py +10 -0
  12. oafuncs-0.0.97.17/oafuncs/oa_nc.py +283 -0
  13. oafuncs-0.0.97.17/oafuncs/oa_python.py +124 -0
  14. oafuncs-0.0.97.17/oafuncs/oa_tool.py +119 -0
  15. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17/oafuncs.egg-info}/PKG-INFO +1 -1
  16. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/SOURCES.txt +2 -1
  17. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/setup.py +1 -1
  18. oafuncs-0.0.97.16/oafuncs/oa_cmap.py +0 -162
  19. oafuncs-0.0.97.16/oafuncs/oa_data.py +0 -293
  20. oafuncs-0.0.97.16/oafuncs/oa_date.py +0 -126
  21. oafuncs-0.0.97.16/oafuncs/oa_draw.py +0 -236
  22. oafuncs-0.0.97.16/oafuncs/oa_file.py +0 -461
  23. oafuncs-0.0.97.16/oafuncs/oa_nc.py +0 -250
  24. oafuncs-0.0.97.16/oafuncs/oa_python.py +0 -98
  25. oafuncs-0.0.97.16/oafuncs/oa_tool.py +0 -80
  26. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/LICENSE.txt +0 -0
  27. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/MANIFEST.in +0 -0
  28. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/README.md +0 -0
  29. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/__init__.py +0 -0
  30. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_data/OAFuncs.png +0 -0
  31. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_data/hycom_3hourly.png +0 -0
  32. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/cprogressbar.py +0 -0
  33. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/email.py +0 -0
  34. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_merge.py +0 -0
  35. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/netcdf_write.py +0 -0
  36. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/parallel.py +0 -0
  37. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/parallel_example_usage.py +0 -0
  38. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/_script/plot_dataset.py +0 -0
  39. /oafuncs-0.0.97.16/oafuncs/_script/replace_file_concent.py → /oafuncs-0.0.97.17/oafuncs/_script/replace_file_content.py +0 -0
  40. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/User_Agent-list.txt +0 -0
  41. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/__init__.py +0 -0
  42. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/literature.py +0 -0
  43. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/test_ua.py +0 -0
  44. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_down/user_agent.py +0 -0
  45. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/__init__.py +0 -0
  46. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/roms/__init__.py +0 -0
  47. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/roms/test.py +0 -0
  48. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/wrf/__init__.py +0 -0
  49. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_model/wrf/little_r.py +0 -0
  50. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/__init__.py +0 -0
  51. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/meteorological.py +0 -0
  52. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/ocean.py +0 -0
  53. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs/oa_sign/scientific.py +0 -0
  54. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/dependency_links.txt +0 -0
  55. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/requires.txt +0 -0
  56. {oafuncs-0.0.97.16 → oafuncs-0.0.97.17}/oafuncs.egg-info/top_level.txt +0 -0
  57. {oafuncs-0.0.97.16 → 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.16
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
@@ -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