roms-tools 1.7.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

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 (124) hide show
  1. roms_tools/__init__.py +2 -1
  2. roms_tools/setup/boundary_forcing.py +246 -146
  3. roms_tools/setup/datasets.py +229 -69
  4. roms_tools/setup/download.py +13 -17
  5. roms_tools/setup/grid.py +777 -614
  6. roms_tools/setup/initial_conditions.py +168 -32
  7. roms_tools/setup/mask.py +115 -0
  8. roms_tools/setup/nesting.py +575 -0
  9. roms_tools/setup/plot.py +218 -63
  10. roms_tools/setup/regrid.py +4 -2
  11. roms_tools/setup/river_forcing.py +125 -29
  12. roms_tools/setup/surface_forcing.py +31 -25
  13. roms_tools/setup/tides.py +29 -14
  14. roms_tools/setup/topography.py +250 -153
  15. roms_tools/setup/utils.py +174 -44
  16. roms_tools/setup/vertical_coordinate.py +5 -16
  17. roms_tools/tests/test_setup/test_boundary_forcing.py +10 -5
  18. roms_tools/tests/test_setup/test_data/grid.zarr/.zattrs +0 -1
  19. roms_tools/tests/test_setup/test_data/grid.zarr/.zmetadata +56 -201
  20. roms_tools/tests/test_setup/test_data/grid.zarr/Cs_r/.zattrs +1 -1
  21. roms_tools/tests/test_setup/test_data/grid.zarr/Cs_w/.zattrs +1 -1
  22. roms_tools/tests/test_setup/test_data/grid.zarr/{layer_depth_rho → sigma_r}/.zarray +2 -6
  23. roms_tools/tests/test_setup/test_data/grid.zarr/sigma_r/.zattrs +7 -0
  24. roms_tools/tests/test_setup/test_data/grid.zarr/sigma_r/0 +0 -0
  25. roms_tools/tests/test_setup/test_data/grid.zarr/sigma_w/.zarray +20 -0
  26. roms_tools/tests/test_setup/test_data/grid.zarr/sigma_w/.zattrs +7 -0
  27. roms_tools/tests/test_setup/test_data/grid.zarr/sigma_w/0 +0 -0
  28. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/.zattrs +1 -2
  29. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/.zmetadata +58 -203
  30. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/Cs_r/.zattrs +1 -1
  31. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/Cs_w/.zattrs +1 -1
  32. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/h/.zattrs +1 -1
  33. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/h/0.0 +0 -0
  34. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/mask_coarse/0.0 +0 -0
  35. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/mask_rho/0.0 +0 -0
  36. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/mask_u/0.0 +0 -0
  37. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/mask_v/0.0 +0 -0
  38. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/sigma_r/.zarray +20 -0
  39. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/sigma_r/.zattrs +7 -0
  40. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/sigma_r/0 +0 -0
  41. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/sigma_w/.zarray +20 -0
  42. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/sigma_w/.zattrs +7 -0
  43. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/sigma_w/0 +0 -0
  44. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata +2 -3
  45. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/.zattrs +1 -2
  46. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/.zarray +1 -1
  47. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/0 +0 -0
  48. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zmetadata +5 -6
  49. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zarray +2 -2
  50. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zattrs +1 -2
  51. roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_tracer/0.0.0 +0 -0
  52. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/tracer_name/.zarray +2 -2
  53. roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/tracer_name/0 +0 -0
  54. roms_tools/tests/test_setup/test_datasets.py +2 -2
  55. roms_tools/tests/test_setup/test_grid.py +110 -12
  56. roms_tools/tests/test_setup/test_initial_conditions.py +2 -1
  57. roms_tools/tests/test_setup/test_nesting.py +489 -0
  58. roms_tools/tests/test_setup/test_river_forcing.py +53 -15
  59. roms_tools/tests/test_setup/test_surface_forcing.py +3 -22
  60. roms_tools/tests/test_setup/test_tides.py +2 -1
  61. roms_tools/tests/test_setup/test_topography.py +106 -1
  62. roms_tools/tests/test_setup/test_validation.py +2 -2
  63. {roms_tools-1.7.0.dist-info → roms_tools-2.1.0.dist-info}/LICENSE +1 -1
  64. {roms_tools-1.7.0.dist-info → roms_tools-2.1.0.dist-info}/METADATA +9 -4
  65. {roms_tools-1.7.0.dist-info → roms_tools-2.1.0.dist-info}/RECORD +85 -108
  66. {roms_tools-1.7.0.dist-info → roms_tools-2.1.0.dist-info}/WHEEL +1 -1
  67. roms_tools/_version.py +0 -2
  68. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_rho/.zarray +0 -24
  69. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_rho/.zattrs +0 -9
  70. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_rho/0.0.0 +0 -0
  71. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_u/.zarray +0 -24
  72. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_u/.zattrs +0 -9
  73. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_u/0.0.0 +0 -0
  74. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_v/.zarray +0 -24
  75. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_v/.zattrs +0 -9
  76. roms_tools/tests/test_setup/test_data/grid.zarr/interface_depth_v/0.0.0 +0 -0
  77. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_rho/.zattrs +0 -9
  78. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_rho/0.0.0 +0 -0
  79. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_u/.zarray +0 -24
  80. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_u/.zattrs +0 -9
  81. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_u/0.0.0 +0 -0
  82. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_v/.zarray +0 -24
  83. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_v/.zattrs +0 -9
  84. roms_tools/tests/test_setup/test_data/grid.zarr/layer_depth_v/0.0.0 +0 -0
  85. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_rho/.zarray +0 -24
  86. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_rho/.zattrs +0 -9
  87. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_rho/0.0.0 +0 -0
  88. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_u/.zarray +0 -24
  89. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_u/.zattrs +0 -9
  90. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_u/0.0.0 +0 -0
  91. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_v/.zarray +0 -24
  92. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_v/.zattrs +0 -9
  93. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/interface_depth_v/0.0.0 +0 -0
  94. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_rho/.zarray +0 -24
  95. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_rho/.zattrs +0 -9
  96. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_rho/0.0.0 +0 -0
  97. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_u/.zarray +0 -24
  98. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_u/.zattrs +0 -9
  99. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_u/0.0.0 +0 -0
  100. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_v/.zarray +0 -24
  101. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_v/.zattrs +0 -9
  102. roms_tools/tests/test_setup/test_data/grid_that_straddles_dateline.zarr/layer_depth_v/0.0.0 +0 -0
  103. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_tracer/0.0.0 +0 -0
  104. roms_tools/tests/test_setup/test_data/river_forcing.zarr/tracer_name/0 +0 -0
  105. roms_tools/tests/test_setup/test_vertical_coordinate.py +0 -91
  106. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zattrs +0 -0
  107. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zgroup +0 -0
  108. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/abs_time/.zarray +0 -0
  109. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/abs_time/.zattrs +0 -0
  110. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/abs_time/0 +0 -0
  111. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/month/.zarray +0 -0
  112. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/month/.zattrs +0 -0
  113. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/month/0 +0 -0
  114. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_name/.zarray +0 -0
  115. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_name/.zattrs +0 -0
  116. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_name/0 +0 -0
  117. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_time/.zarray +0 -0
  118. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_time/.zattrs +0 -0
  119. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_time/0 +0 -0
  120. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_volume/.zarray +0 -0
  121. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_volume/.zattrs +0 -0
  122. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_volume/0.0 +0 -0
  123. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/tracer_name/.zattrs +0 -0
  124. {roms_tools-1.7.0.dist-info → roms_tools-2.1.0.dist-info}/top_level.txt +0 -0
roms_tools/__init__.py CHANGED
@@ -5,7 +5,7 @@ try:
5
5
  __version__ = _version("roms_tools")
6
6
  except ImportError: # pragma: no cover
7
7
  # Local copy or not installed with setuptools
8
- __version__ = "999"
8
+ __version__ = "9999"
9
9
 
10
10
 
11
11
  from roms_tools.setup.grid import Grid # noqa: F401
@@ -14,6 +14,7 @@ from roms_tools.setup.surface_forcing import SurfaceForcing # noqa: F401
14
14
  from roms_tools.setup.initial_conditions import InitialConditions # noqa: F401
15
15
  from roms_tools.setup.boundary_forcing import BoundaryForcing # noqa: F401
16
16
  from roms_tools.setup.river_forcing import RiverForcing # noqa: F401
17
+ from roms_tools.setup.nesting import Nesting # noqa: F401
17
18
 
18
19
  # Configure logging when the package is imported
19
20
  logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
@@ -9,6 +9,7 @@ from roms_tools.setup.grid import Grid
9
9
  from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
10
10
  from datetime import datetime
11
11
  from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
12
+ from roms_tools.setup.vertical_coordinate import compute_depth
12
13
  from roms_tools.setup.utils import (
13
14
  get_variable_metadata,
14
15
  group_dataset,
@@ -20,7 +21,10 @@ from roms_tools.setup.utils import (
20
21
  one_dim_fill,
21
22
  nan_check,
22
23
  substitute_nans_by_fillvalue,
24
+ interpolate_from_rho_to_u,
25
+ interpolate_from_rho_to_v,
23
26
  convert_to_roms_time,
27
+ get_boundary_coords,
24
28
  _to_yaml,
25
29
  _from_yaml,
26
30
  )
@@ -68,6 +72,10 @@ class BoundaryForcing:
68
72
  Reference date for the model. Default is January 1, 2000.
69
73
  use_dask: bool, optional
70
74
  Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
75
+ bypass_validation: bool, optional
76
+ Indicates whether to skip validation checks in the processed data. When set to True,
77
+ the validation process that ensures no NaN values exist at wet points
78
+ in the processed dataset is bypassed. Defaults to False.
71
79
 
72
80
  Examples
73
81
  --------
@@ -97,6 +105,7 @@ class BoundaryForcing:
97
105
  apply_2d_horizontal_fill: bool = False
98
106
  model_reference_date: datetime = datetime(2000, 1, 1)
99
107
  use_dask: bool = False
108
+ bypass_validation: bool = False
100
109
 
101
110
  ds: xr.Dataset = field(init=False, repr=False)
102
111
 
@@ -115,8 +124,8 @@ class BoundaryForcing:
115
124
  data.extrapolate_deepest_to_bottom()
116
125
  data.apply_lateral_fill()
117
126
 
118
- variable_info = self._set_variable_info(data)
119
- bdry_coords = get_boundary_info()
127
+ self._set_variable_info(data)
128
+ self._set_boundary_info()
120
129
  ds = xr.Dataset()
121
130
 
122
131
  for direction in ["south", "east", "north", "west"]:
@@ -124,10 +133,10 @@ class BoundaryForcing:
124
133
 
125
134
  bdry_target_coords = {
126
135
  "lat": target_coords["lat"].isel(
127
- **bdry_coords["vector"][direction]
136
+ **self.bdry_coords["vector"][direction]
128
137
  ),
129
138
  "lon": target_coords["lon"].isel(
130
- **bdry_coords["vector"][direction]
139
+ **self.bdry_coords["vector"][direction]
131
140
  ),
132
141
  "straddle": target_coords["straddle"],
133
142
  }
@@ -145,11 +154,17 @@ class BoundaryForcing:
145
154
 
146
155
  # lateral regridding of vector fields
147
156
  vector_var_names = [
148
- name for name, info in variable_info.items() if info["is_vector"]
157
+ name
158
+ for name, info in self.variable_info.items()
159
+ if info["is_vector"]
149
160
  ]
150
161
  if len(vector_var_names) > 0:
151
- lon = target_coords["lon"].isel(**bdry_coords["vector"][direction])
152
- lat = target_coords["lat"].isel(**bdry_coords["vector"][direction])
162
+ lon = target_coords["lon"].isel(
163
+ **self.bdry_coords["vector"][direction]
164
+ )
165
+ lat = target_coords["lat"].isel(
166
+ **self.bdry_coords["vector"][direction]
167
+ )
153
168
  lateral_regrid = LateralRegrid(
154
169
  {"lat": lat, "lon": lon}, bdry_data.dim_names
155
170
  )
@@ -162,12 +177,16 @@ class BoundaryForcing:
162
177
  # lateral regridding of tracer fields
163
178
  tracer_var_names = [
164
179
  name
165
- for name, info in variable_info.items()
180
+ for name, info in self.variable_info.items()
166
181
  if not info["is_vector"]
167
182
  ]
168
183
  if len(tracer_var_names) > 0:
169
- lon = target_coords["lon"].isel(**bdry_coords["rho"][direction])
170
- lat = target_coords["lat"].isel(**bdry_coords["rho"][direction])
184
+ lon = target_coords["lon"].isel(
185
+ **self.bdry_coords["rho"][direction]
186
+ )
187
+ lat = target_coords["lat"].isel(
188
+ **self.bdry_coords["rho"][direction]
189
+ )
171
190
  lateral_regrid = LateralRegrid(
172
191
  {"lat": lat, "lon": lon}, bdry_data.dim_names
173
192
  )
@@ -178,9 +197,9 @@ class BoundaryForcing:
178
197
  )
179
198
 
180
199
  # rotation of velocities and interpolation to u/v points
181
- if "u" in variable_info and "v" in variable_info:
200
+ if "u" in self.variable_info and "v" in self.variable_info:
182
201
  angle = target_coords["angle"].isel(
183
- **bdry_coords["vector"][direction]
202
+ **self.bdry_coords["vector"][direction]
184
203
  )
185
204
  (processed_fields["u"], processed_fields["v"],) = rotate_velocities(
186
205
  processed_fields["u"],
@@ -190,54 +209,68 @@ class BoundaryForcing:
190
209
  )
191
210
 
192
211
  # selection of outermost margin for u/v variables
193
- for var_name in variable_info.keys():
212
+ for var_name in self.variable_info.keys():
194
213
  if var_name in processed_fields:
195
- location = variable_info[var_name]["location"]
214
+ location = self.variable_info[var_name]["location"]
196
215
  if location in ["u", "v"]:
197
216
  processed_fields[var_name] = processed_fields[
198
217
  var_name
199
- ].isel(**bdry_coords[location][direction])
218
+ ].isel(**self.bdry_coords[location][direction])
200
219
 
201
220
  if not self.apply_2d_horizontal_fill:
202
221
  self._validate_1d_fill(
203
222
  processed_fields,
204
- variable_info,
205
- bdry_coords,
206
223
  direction,
207
224
  bdry_data.dim_names["depth"],
208
225
  )
209
226
  processed_fields = apply_1d_horizontal_fill(processed_fields)
210
227
 
211
- # vertical regridding
228
+ var_names_dict = {}
212
229
  for location in ["rho", "u", "v"]:
213
- var_names = [
230
+ var_names_dict[location] = [
214
231
  name
215
- for name, info in variable_info.items()
232
+ for name, info in self.variable_info.items()
216
233
  if info["location"] == location and info["is_3d"]
217
234
  ]
218
- if len(var_names) > 0:
235
+ # compute layer depth coordinates
236
+ if len(var_names_dict["u"]) > 0 or len(var_names_dict["v"]) > 0:
237
+ self._get_vertical_coordinates(
238
+ type="layer",
239
+ direction=direction,
240
+ additional_locations=["u", "v"],
241
+ )
242
+ else:
243
+ if len(var_names_dict["rho"]) > 0:
244
+ self._get_vertical_coordinates(
245
+ type="layer", direction=direction, additional_locations=[]
246
+ )
247
+
248
+ # vertical regridding
249
+ for location in ["rho", "u", "v"]:
250
+ if len(var_names_dict[location]) > 0:
219
251
  vertical_regrid = VerticalRegrid(
220
- self.grid.ds[f"layer_depth_{location}"].isel(
221
- **bdry_coords[location][direction]
222
- ),
252
+ self.grid.ds[f"layer_depth_{location}_{direction}"],
223
253
  bdry_data.ds[bdry_data.dim_names["depth"]],
224
254
  )
225
- for var_name in var_names:
255
+ for var_name in var_names_dict[location]:
226
256
  if var_name in processed_fields:
227
257
  processed_fields[var_name] = vertical_regrid.apply(
228
258
  processed_fields[var_name]
229
259
  )
230
260
 
231
261
  # compute barotropic velocities
232
- if "u" in variable_info and "v" in variable_info:
233
- for var_name in ["u", "v"]:
262
+ if "u" in self.variable_info and "v" in self.variable_info:
263
+ self._get_vertical_coordinates(
264
+ type="interface",
265
+ direction=direction,
266
+ additional_locations=["u", "v"],
267
+ )
268
+ for location in ["u", "v"]:
234
269
  processed_fields[
235
- f"{var_name}bar"
270
+ f"{location}bar"
236
271
  ] = compute_barotropic_velocity(
237
- processed_fields[var_name],
238
- self.grid.ds[f"interface_depth_{var_name}"].isel(
239
- **bdry_coords[var_name][direction]
240
- ),
272
+ processed_fields[location],
273
+ self.grid.ds[f"interface_depth_{location}_{direction}"],
241
274
  )
242
275
 
243
276
  # Reorder dimensions
@@ -252,7 +285,8 @@ class BoundaryForcing:
252
285
  # Add global information
253
286
  ds = self._add_global_metadata(data, ds)
254
287
 
255
- self._validate(ds, variable_info, bdry_coords)
288
+ if not self.bypass_validation:
289
+ self._validate(ds)
256
290
 
257
291
  # substitute NaNs over land by a fill value to avoid blow-up of ROMS
258
292
  for var_name in ds.data_vars:
@@ -317,11 +351,15 @@ class BoundaryForcing:
317
351
  - `vector_pair`: For vector variables, this indicates the associated variable that forms the vector (e.g., 'u' and 'v').
318
352
  - `is_3d`: Indicates whether the variable is 3D (True for variables like 'temp' and 'salt') or 2D (False for 'zeta').
319
353
 
354
+ Parameters
355
+ ----------
356
+ data : object
357
+ An object that contains variable names for the data being processed. This is used to set variable information for biogeochemical data.
358
+
320
359
  Returns
321
360
  -------
322
- dict
323
- A dictionary where the keys are variable names and the values are dictionaries of metadata
324
- about each variable, including 'location', 'is_vector', 'vector_pair', and 'is_3d'.
361
+ None
362
+ This method updates the instance attribute `variable_info` with the metadata dictionary for the variables.
325
363
  """
326
364
  default_info = {
327
365
  "location": "rho",
@@ -379,7 +417,7 @@ class BoundaryForcing:
379
417
  else:
380
418
  variable_info[var_name] = {**default_info, "validate": False}
381
419
 
382
- return variable_info
420
+ object.__setattr__(self, "variable_info", variable_info)
383
421
 
384
422
  def _write_into_dataset(self, direction, processed_fields, ds=None):
385
423
  if ds is None:
@@ -414,45 +452,153 @@ class BoundaryForcing:
414
452
  "lat_v",
415
453
  "lon_v",
416
454
  ]
417
- existing_vars = [var_name for var_name in variables_to_drop if var_name in ds]
455
+ suffixes = ["", "_south", "_east", "_north", "_west"]
456
+ # Existing variables with suffixes
457
+ existing_vars = []
458
+ for var_name in variables_to_drop:
459
+ for suffix in suffixes:
460
+ full_var_name = f"{var_name}{suffix}"
461
+ if full_var_name in ds:
462
+ existing_vars.append(full_var_name)
463
+
418
464
  ds = ds.drop_vars(existing_vars)
419
465
 
420
466
  return ds
421
467
 
422
- def _get_coordinates(self, direction, point):
468
+ def _set_boundary_info(self):
469
+ """Sets boundary coordinates for rho, u, and v variables on the grid.
470
+
471
+ This method determines the boundary points for the grid variables by specifying the
472
+ indices for the south, east, north, and west boundaries. The resulting boundary
473
+ information is stored in the instance attribute `bdry_coords`.
474
+
475
+ Returns
476
+ -------
477
+ None
478
+ The method does not return a value. Instead, it updates the instance attribute
479
+ `bdry_coords`, which is a dictionary structured as follows:
480
+ - Keys: Variable types ("rho", "u", "v", "vector").
481
+ - Values: Nested dictionaries mapping each direction ("south", "east", "north", "west")
482
+ to their corresponding boundary coordinates. The coordinates are specified in terms of
483
+ grid indices for the respective variable types.
484
+ """
485
+
486
+ bdry_coords = get_boundary_coords()
487
+
488
+ object.__setattr__(self, "bdry_coords", bdry_coords)
489
+
490
+ def _get_vertical_coordinates(
491
+ self, type, direction, additional_locations=["u", "v"]
492
+ ):
423
493
  """Retrieve layer and interface depth coordinates for a specified grid boundary.
424
494
 
425
- This method extracts the layer depth and interface depth coordinates along
426
- a specified boundary (north, south, east, or west) and for a specified point
427
- type (rho, u, or v) from the grid dataset.
495
+ This method computes and updates the layer and interface depth coordinates along a specified
496
+ boundary (north, south, east, or west). It handles depth calculations for rho points and
497
+ additional specified locations (u and v).
428
498
 
429
499
  Parameters
430
500
  ----------
431
- direction : str
432
- The direction of the boundary to retrieve coordinates for. Valid options
433
- are "north", "south", "east", and "west".
434
- point : str
435
- The type of grid point to retrieve coordinates for. Valid options are
436
- "rho" for the grid's central points, "u" for the u-flux points, and "v"
437
- for the v-flux points.
501
+ type : str
502
+ The type of depth coordinate to retrieve. Valid options are:
503
+ - "layer": Retrieves layer depth coordinates.
504
+ - "interface": Retrieves interface depth coordinates.
438
505
 
439
- Returns
506
+ direction : str
507
+ The direction of the boundary to retrieve coordinates for. Valid options are:
508
+ - "north"
509
+ - "south"
510
+ - "east"
511
+ - "west"
512
+
513
+ additional_locations : list of str, optional
514
+ Specifies additional locations to compute depth coordinates for. Default is ["u", "v"].
515
+ Valid options include:
516
+ - "u": Computes depth coordinates for u points.
517
+ - "v": Computes depth coordinates for v points.
518
+
519
+ Updates
440
520
  -------
441
- xarray.DataArray, xarray.DataArray
442
- The layer depth and interface depth coordinates for the specified grid
443
- boundary and point type.
521
+ self.grid.ds : xarray.Dataset
522
+ The dataset is updated with the following vertical depth coordinates:
523
+ - f"{type}_depth_rho_{direction}": Depth coordinates at rho points.
524
+ - f"{type}_depth_u_{direction}": Depth coordinates at u points (if applicable).
525
+ - f"{type}_depth_v_{direction}": Depth coordinates at v points (if applicable).
444
526
  """
445
527
 
446
- bdry_coords = get_boundary_info()
528
+ layer_vars = []
529
+ for location in ["rho"] + additional_locations:
530
+ layer_vars.append(f"{type}_depth_{location}_{direction}")
447
531
 
448
- layer_depth = self.grid.ds[f"layer_depth_{point}"].isel(
449
- **bdry_coords[point][direction]
450
- )
451
- interface_depth = self.grid.ds[f"interface_depth_{point}"].isel(
452
- **bdry_coords[point][direction]
453
- )
532
+ if all(layer_var in self.grid.ds for layer_var in layer_vars):
533
+ # Vertical coordinate data already exists
534
+ pass
535
+
536
+ elif f"{type}_depth_rho" in self.grid.ds:
537
+ depth = self.grid.ds[f"{type}_depth_rho"]
538
+ depth.attrs["long_name"] = f"{type} depth at rho-points"
539
+ depth.attrs["units"] = "m"
540
+ self.grid.ds[f"{type}_depth_rho_{direction}"] = depth.isel(
541
+ **self.bdry_coords["rho"][direction]
542
+ )
454
543
 
455
- return layer_depth, interface_depth
544
+ if "u" in additional_locations or "v" in additional_locations:
545
+ # selection of margin consisting of 2 grid cells
546
+ depth = depth.isel(**self.bdry_coords["vector"][direction])
547
+ # interpolation
548
+ if "u" in additional_locations:
549
+ depth_u = interpolate_from_rho_to_u(depth)
550
+ depth_u.attrs["long_name"] = f"{type} depth at u-points"
551
+ depth_u.attrs["units"] = "m"
552
+ self.grid.ds[f"{type}_depth_u_{direction}"] = depth_u.isel(
553
+ **self.bdry_coords["u"][direction]
554
+ )
555
+ if "v" in additional_locations:
556
+ depth_v = interpolate_from_rho_to_v(depth)
557
+ depth_v.attrs["long_name"] = f"{type} depth at v-points"
558
+ depth_v.attrs["units"] = "m"
559
+ self.grid.ds[f"{type}_depth_v_{direction}"] = depth_v.isel(
560
+ **self.bdry_coords["v"][direction]
561
+ )
562
+ else:
563
+ if "u" in additional_locations or "v" in additional_locations:
564
+ h = self.grid.ds["h"].isel(**self.bdry_coords["vector"][direction])
565
+ else:
566
+ h = self.grid.ds["h"].isel(**self.bdry_coords["rho"][direction])
567
+ if type == "layer":
568
+ depth = compute_depth(
569
+ 0, h, self.grid.hc, self.grid.ds.Cs_r, self.grid.ds.sigma_r
570
+ )
571
+ else:
572
+ depth = compute_depth(
573
+ 0, h, self.grid.hc, self.grid.ds.Cs_w, self.grid.ds.sigma_w
574
+ )
575
+
576
+ if "u" in additional_locations or "v" in additional_locations:
577
+ depth.attrs["long_name"] = f"{type} depth at rho-points"
578
+ depth.attrs["units"] = "m"
579
+ self.grid.ds[f"{type}_depth_rho_{direction}"] = depth.isel(
580
+ **self.bdry_coords["rho"][direction]
581
+ )
582
+ # selection of margin consisting of 2 grid cells
583
+ depth = depth.isel(**self.bdry_coords["vector"][direction])
584
+ # interpolation
585
+ depth_u = interpolate_from_rho_to_u(depth)
586
+ depth_v = interpolate_from_rho_to_v(depth)
587
+ # selection of outermost margin
588
+ depth_u.attrs["long_name"] = f"{type} depth at u-points"
589
+ depth_u.attrs["units"] = "m"
590
+ self.grid.ds[f"{type}_depth_u_{direction}"] = depth_u.isel(
591
+ **self.bdry_coords["u"][direction]
592
+ )
593
+ depth_v.attrs["long_name"] = f"{type} depth at v-points"
594
+ depth_v.attrs["units"] = "m"
595
+ self.grid.ds[f"{type}_depth_v_{direction}"] = depth_v.isel(
596
+ **self.bdry_coords["v"][direction]
597
+ )
598
+ else:
599
+ depth.attrs["long_name"] = f"{type} depth at rho-points"
600
+ depth.attrs["units"] = "m"
601
+ self.grid.ds[f"{type}_depth_rho_{direction}"] = depth
456
602
 
457
603
  def _add_global_metadata(self, data, ds=None):
458
604
 
@@ -485,9 +631,7 @@ class BoundaryForcing:
485
631
 
486
632
  return ds
487
633
 
488
- def _validate_1d_fill(
489
- self, processed_fields, variable_info, bdry_coords, direction, depth_dim
490
- ):
634
+ def _validate_1d_fill(self, processed_fields, direction, depth_dim):
491
635
  """Check if any boundary is divided by land and issue a warning if so,
492
636
  suggesting the use of 2D horizontal fill for safer regridding.
493
637
 
@@ -497,15 +641,6 @@ class BoundaryForcing:
497
641
  A dictionary where keys are variable names and values are `xarray.DataArray`
498
642
  objects representing the processed data for each variable.
499
643
 
500
- variable_info : dict
501
- A dictionary containing metadata about each variable (e.g., location,
502
- whether it's a 3D variable, etc.). Used to retrieve information for
503
- validating each variable.
504
-
505
- bdry_coords : dict
506
- A dictionary containing boundary coordinates for different directions (north, south,
507
- east, west), used to slice the boundary-specific data for each variable.
508
-
509
644
  direction : str
510
645
  The boundary direction being processed (e.g., "north", "south", "east", or "west").
511
646
 
@@ -521,8 +656,8 @@ class BoundaryForcing:
521
656
 
522
657
  for var_name in processed_fields.keys():
523
658
  # Only validate variables based on "validate" flag if use_dask is False
524
- if not self.use_dask or variable_info[var_name]["validate"]:
525
- location = variable_info[var_name]["location"]
659
+ if not self.use_dask or self.variable_info[var_name]["validate"]:
660
+ location = self.variable_info[var_name]["location"]
526
661
 
527
662
  # Select the appropriate mask based on variable location
528
663
  if location == "rho":
@@ -532,9 +667,9 @@ class BoundaryForcing:
532
667
  elif location == "v":
533
668
  mask = self.grid.ds.mask_v
534
669
 
535
- mask = mask.isel(**bdry_coords[location][direction])
670
+ mask = mask.isel(**self.bdry_coords[location][direction])
536
671
 
537
- if variable_info[var_name]["is_3d"]:
672
+ if self.variable_info[var_name]["is_3d"]:
538
673
  da = processed_fields[var_name].isel({depth_dim: 0, "time": 0})
539
674
  else:
540
675
  da = processed_fields[var_name].isel({"time": 0})
@@ -553,7 +688,7 @@ class BoundaryForcing:
553
688
  f"For {var_name}, the {direction}ern boundary is divided by land. It would be safer (but slower) to use `apply_2d_horizontal_fill = True`."
554
689
  )
555
690
 
556
- def _validate(self, ds, variable_info, bdry_coords):
691
+ def _validate(self, ds):
557
692
  """Validate the dataset for NaN values at the first time step (bry_time=0) for
558
693
  specified variables. If NaN values are found at wet points, this function raises
559
694
  an error.
@@ -563,12 +698,6 @@ class BoundaryForcing:
563
698
  ds : xarray.Dataset
564
699
  The dataset to validate.
565
700
 
566
- variable_info : dict
567
- A dictionary containing metadata about the variables, including their locations (e.g., 'rho', 'u', 'v').
568
-
569
- bdry_coords : dict
570
- A dictionary containing the boundary coordinates for each variable location.
571
-
572
701
  Raises
573
702
  ------
574
703
  ValueError
@@ -580,10 +709,10 @@ class BoundaryForcing:
580
709
  Validation is performed on the initial boundary time step (`bry_time=0`) for each
581
710
  variable in the dataset.
582
711
  """
583
- for var_name in variable_info:
712
+ for var_name in self.variable_info:
584
713
  # only validate variables based on "validate" flag if use_dask is false
585
- if not self.use_dask or variable_info[var_name]["validate"]:
586
- location = variable_info[var_name]["location"]
714
+ if not self.use_dask or self.variable_info[var_name]["validate"]:
715
+ location = self.variable_info[var_name]["location"]
587
716
 
588
717
  # Select the appropriate mask based on variable location
589
718
  if location == "rho":
@@ -610,7 +739,7 @@ class BoundaryForcing:
610
739
 
611
740
  nan_check(
612
741
  ds[bdry_var_name].isel(bry_time=0),
613
- mask.isel(**bdry_coords[location][direction]),
742
+ mask.isel(**self.bdry_coords[location][direction]),
614
743
  error_message=error_message,
615
744
  )
616
745
 
@@ -696,20 +825,13 @@ class BoundaryForcing:
696
825
  field = field.load()
697
826
 
698
827
  title = field.long_name
828
+ var_name_wo_direction, direction = var_name.split("_")
829
+ location = self.variable_info[var_name_wo_direction]["location"]
699
830
 
700
831
  if "s_rho" in field.dims:
701
- if var_name.startswith(("u_", "ubar_")):
702
- point = "u"
703
- elif var_name.startswith(("v_", "vbar_")):
704
- point = "v"
705
- else:
706
- point = "rho"
707
- direction = var_name.split("_")[-1]
708
-
709
- layer_depth, interface_depth = self._get_coordinates(direction, point)
710
-
711
- field = field.assign_coords({"layer_depth": layer_depth})
712
-
832
+ field = field.assign_coords(
833
+ {"layer_depth": self.grid.ds[f"layer_depth_{location}_{direction}"]}
834
+ )
713
835
  # chose colorbar
714
836
  if var_name.startswith(("u", "v", "ubar", "vbar", "zeta")):
715
837
  vmax = max(field.max().values, -field.min().values)
@@ -727,6 +849,19 @@ class BoundaryForcing:
727
849
 
728
850
  if len(field.dims) == 2:
729
851
  if layer_contours:
852
+ if location in ["u", "v"]:
853
+ additional_locations = ["u", "v"]
854
+ else:
855
+ additional_locations = []
856
+ self._get_vertical_coordinates(
857
+ type="interface",
858
+ direction=direction,
859
+ additional_locations=additional_locations,
860
+ )
861
+
862
+ interface_depth = self.grid.ds[
863
+ f"interface_depth_{location}_{direction}"
864
+ ]
730
865
  # restrict number of layer_contours to 10 for the sake of plot clearity
731
866
  nr_layers = len(interface_depth["s_w"])
732
867
  selected_layers = np.linspace(
@@ -827,7 +962,10 @@ class BoundaryForcing:
827
962
 
828
963
  @classmethod
829
964
  def from_yaml(
830
- cls, filepath: Union[str, Path], use_dask: bool = False
965
+ cls,
966
+ filepath: Union[str, Path],
967
+ use_dask: bool = False,
968
+ bypass_validation: bool = False,
831
969
  ) -> "BoundaryForcing":
832
970
  """Create an instance of the BoundaryForcing class from a YAML file.
833
971
 
@@ -837,6 +975,10 @@ class BoundaryForcing:
837
975
  The path to the YAML file from which the parameters will be read.
838
976
  use_dask: bool, optional
839
977
  Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
978
+ bypass_validation: bool, optional
979
+ Indicates whether to skip validation checks in the processed data. When set to True,
980
+ the validation process that ensures no NaN values exist at wet points
981
+ in the processed dataset is bypassed. Defaults to False.
840
982
 
841
983
  Returns
842
984
  -------
@@ -849,51 +991,9 @@ class BoundaryForcing:
849
991
  params = _from_yaml(cls, filepath)
850
992
 
851
993
  # Create and return an instance of InitialConditions
852
- return cls(grid=grid, **params, use_dask=use_dask)
853
-
854
-
855
- def get_boundary_info():
856
- """This function provides information about the boundary points for the rho, u, and
857
- v variables on the grid, specifying the indices for the south, east, north, and west
858
- boundaries.
859
-
860
- Returns
861
- -------
862
- dict
863
- A dictionary where keys are variable types ("rho", "u", "v"), and values
864
- are nested dictionaries mapping directions ("south", "east", "north", "west")
865
- to the corresponding boundary coordinates.
866
- """
867
-
868
- # Boundary coordinates
869
- bdry_coords = {
870
- "rho": {
871
- "south": {"eta_rho": 0},
872
- "east": {"xi_rho": -1},
873
- "north": {"eta_rho": -1},
874
- "west": {"xi_rho": 0},
875
- },
876
- "u": {
877
- "south": {"eta_rho": 0},
878
- "east": {"xi_u": -1},
879
- "north": {"eta_rho": -1},
880
- "west": {"xi_u": 0},
881
- },
882
- "v": {
883
- "south": {"eta_v": 0},
884
- "east": {"xi_rho": -1},
885
- "north": {"eta_v": -1},
886
- "west": {"xi_rho": 0},
887
- },
888
- "vector": {
889
- "south": {"eta_rho": [0, 1]},
890
- "east": {"xi_rho": [-2, -1]},
891
- "north": {"eta_rho": [-2, -1]},
892
- "west": {"xi_rho": [0, 1]},
893
- },
894
- }
895
-
896
- return bdry_coords
994
+ return cls(
995
+ grid=grid, **params, use_dask=use_dask, bypass_validation=bypass_validation
996
+ )
897
997
 
898
998
 
899
999
  def apply_1d_horizontal_fill(processed_fields: dict) -> dict: