roms-tools 2.0.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 (54) hide show
  1. roms_tools/__init__.py +2 -1
  2. roms_tools/setup/boundary_forcing.py +21 -30
  3. roms_tools/setup/datasets.py +13 -21
  4. roms_tools/setup/grid.py +253 -139
  5. roms_tools/setup/initial_conditions.py +21 -3
  6. roms_tools/setup/mask.py +50 -4
  7. roms_tools/setup/nesting.py +575 -0
  8. roms_tools/setup/plot.py +214 -55
  9. roms_tools/setup/river_forcing.py +125 -29
  10. roms_tools/setup/surface_forcing.py +21 -8
  11. roms_tools/setup/tides.py +21 -3
  12. roms_tools/setup/topography.py +168 -35
  13. roms_tools/setup/utils.py +127 -21
  14. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata +2 -3
  15. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/.zattrs +1 -2
  16. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/.zarray +1 -1
  17. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/0 +0 -0
  18. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zmetadata +5 -6
  19. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zarray +2 -2
  20. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_tracer/.zattrs +1 -2
  21. roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_tracer/0.0.0 +0 -0
  22. roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/tracer_name/.zarray +2 -2
  23. roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/tracer_name/0 +0 -0
  24. roms_tools/tests/test_setup/test_datasets.py +2 -2
  25. roms_tools/tests/test_setup/test_nesting.py +489 -0
  26. roms_tools/tests/test_setup/test_river_forcing.py +50 -13
  27. roms_tools/tests/test_setup/test_surface_forcing.py +1 -0
  28. roms_tools/tests/test_setup/test_validation.py +2 -2
  29. {roms_tools-2.0.0.dist-info → roms_tools-2.1.0.dist-info}/METADATA +8 -4
  30. {roms_tools-2.0.0.dist-info → roms_tools-2.1.0.dist-info}/RECORD +51 -50
  31. {roms_tools-2.0.0.dist-info → roms_tools-2.1.0.dist-info}/WHEEL +1 -1
  32. roms_tools/_version.py +0 -2
  33. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_tracer/0.0.0 +0 -0
  34. roms_tools/tests/test_setup/test_data/river_forcing.zarr/tracer_name/0 +0 -0
  35. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zattrs +0 -0
  36. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/.zgroup +0 -0
  37. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/abs_time/.zarray +0 -0
  38. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/abs_time/.zattrs +0 -0
  39. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/abs_time/0 +0 -0
  40. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/month/.zarray +0 -0
  41. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/month/.zattrs +0 -0
  42. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/month/0 +0 -0
  43. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_name/.zarray +0 -0
  44. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_name/.zattrs +0 -0
  45. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_name/0 +0 -0
  46. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_time/.zarray +0 -0
  47. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_time/.zattrs +0 -0
  48. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_time/0 +0 -0
  49. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_volume/.zarray +0 -0
  50. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_volume/.zattrs +0 -0
  51. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/river_volume/0.0 +0 -0
  52. /roms_tools/tests/test_setup/test_data/{river_forcing.zarr → river_forcing_with_bgc.zarr}/tracer_name/.zattrs +0 -0
  53. {roms_tools-2.0.0.dist-info → roms_tools-2.1.0.dist-info}/LICENSE +0 -0
  54. {roms_tools-2.0.0.dist-info → roms_tools-2.1.0.dist-info}/top_level.txt +0 -0
roms_tools/setup/grid.py CHANGED
@@ -1,5 +1,4 @@
1
1
  import time
2
- import copy
3
2
  import logging
4
3
  from dataclasses import dataclass, field, asdict
5
4
 
@@ -16,13 +15,12 @@ from roms_tools.setup.utils import (
16
15
  interpolate_from_rho_to_u,
17
16
  interpolate_from_rho_to_v,
18
17
  get_target_coords,
18
+ gc_dist,
19
19
  )
20
20
  from roms_tools.setup.vertical_coordinate import sigma_stretch, compute_depth
21
21
  from roms_tools.setup.utils import extract_single_value, save_datasets
22
22
  from pathlib import Path
23
23
 
24
- RADIUS_OF_EARTH = 6371315.0 # in m
25
-
26
24
 
27
25
  @dataclass(frozen=True, kw_only=True)
28
26
  class Grid:
@@ -380,25 +378,33 @@ class Grid:
380
378
  else:
381
379
  ds[coarse_var] = coarse_field
382
380
 
381
+ del fine_field, coarse_field
382
+
383
383
  ds["mask_coarse"] = xr.where(ds["mask_coarse"] > 0.5, 1, 0).astype(np.int32)
384
384
 
385
385
  for fine_var, coarse_var in d.items():
386
- ds[coarse_var].attrs[
387
- "long_name"
388
- ] = f"{ds[fine_var].attrs['long_name']} on coarsened grid"
386
+ long_name = ds[fine_var].attrs.get(
387
+ "long_name", ds[fine_var].attrs.get("Long_name", "")
388
+ )
389
+ ds[coarse_var].attrs["long_name"] = f"{long_name} on coarsened grid"
389
390
  ds[coarse_var].attrs["units"] = ds[fine_var].attrs["units"]
390
391
 
391
392
  object.__setattr__(self, "ds", ds)
392
393
 
393
- def plot(self, bathymetry: bool = False, title: str = None) -> None:
394
+ def plot(
395
+ self, bathymetry: bool = False, title: str = None, with_dim_names: bool = False
396
+ ) -> None:
394
397
  """Plot the grid.
395
398
 
396
399
  Parameters
397
400
  ----------
398
- bathymetry : bool
401
+ bathymetry : bool, optional
399
402
  Whether or not to plot the bathymetry. Default is False.
400
403
  title : str, optional
401
404
  The title of the plot. If not provided, it will be set to a default.
405
+ with_dim_names : bool, optional
406
+ Whether or not to plot the dimension names. Default is False.
407
+
402
408
 
403
409
  Returns
404
410
  -------
@@ -425,12 +431,18 @@ class Grid:
425
431
  field=field,
426
432
  straddle=self.straddle,
427
433
  title=title,
434
+ with_dim_names=with_dim_names,
428
435
  kwargs=kwargs,
429
436
  )
430
437
  else:
431
438
  if title is None:
432
439
  title = "ROMS grid"
433
- _plot(self.ds, straddle=self.straddle, title=title)
440
+ _plot(
441
+ self.ds,
442
+ straddle=self.straddle,
443
+ title=title,
444
+ with_dim_names=with_dim_names,
445
+ )
434
446
 
435
447
  def plot_vertical_coordinate(
436
448
  self,
@@ -631,7 +643,7 @@ class Grid:
631
643
  hc = 300.0
632
644
 
633
645
  grid.update_vertical_coordinate(
634
- N=N, theta_s=theta_s, theta_b=theta_b, hc=hc, verbose=verbose
646
+ N=N, theta_s=theta_s, theta_b=theta_b, hc=hc, verbose=True
635
647
  )
636
648
  else:
637
649
  object.__setattr__(grid, "theta_s", ds.attrs["theta_s"].item())
@@ -726,20 +738,38 @@ class Grid:
726
738
  yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
727
739
 
728
740
  @classmethod
729
- def from_yaml(cls, filepath: Union[str, Path], verbose: bool = False) -> "Grid":
741
+ def from_yaml(
742
+ cls,
743
+ filepath: Union[str, Path],
744
+ section_name: str = "Grid",
745
+ verbose: bool = False,
746
+ ) -> "Grid":
730
747
  """Create an instance of the class from a YAML file.
731
748
 
732
749
  Parameters
733
750
  ----------
734
751
  filepath : Union[str, Path]
735
752
  The path to the YAML file from which the parameters will be read.
736
- verbose: bool, optional
753
+ section_name : str, optional
754
+ The name of the YAML section containing the grid configuration. Defaults to "Grid".
755
+ verbose : bool, optional
737
756
  Indicates whether to print grid generation steps with timing. Defaults to False.
738
757
 
739
758
  Returns
740
759
  -------
741
760
  Grid
742
- An instance of the Grid class.
761
+ An instance of the Grid class initialized with the parameters from the YAML file.
762
+
763
+ Raises
764
+ ------
765
+ ValueError
766
+ If the ROMS-Tools version is not found in the YAML file or if the specified section
767
+ does not exist in the file.
768
+
769
+ Warnings
770
+ --------
771
+ Issues a warning if the ROMS-Tools version in the YAML header does not match the
772
+ currently installed version.
743
773
  """
744
774
 
745
775
  filepath = Path(filepath)
@@ -759,8 +789,8 @@ class Grid:
759
789
  continue
760
790
  if "roms_tools_version" in doc:
761
791
  header_data = doc
762
- elif "Grid" in doc:
763
- grid_data = doc["Grid"]
792
+ elif section_name in doc:
793
+ grid_data = doc[section_name]
764
794
 
765
795
  if header_data is None:
766
796
  raise ValueError("Version of ROMS-Tools not found in the YAML file.")
@@ -782,11 +812,12 @@ class Grid:
782
812
  raise ValueError("No Grid configuration found in the YAML file.")
783
813
  return cls(**grid_data, verbose=verbose)
784
814
 
785
- # override __repr__ method to only print attributes that are actually set
786
815
  def __repr__(self) -> str:
816
+ """Return a string representation of the object with non-None attributes,
817
+ excluding 'ds'."""
787
818
  cls = self.__class__
788
819
  cls_name = cls.__name__
789
- # Create a dictionary of attribute names and values, filtering out those that are not set and 'ds'
820
+ # Filter attributes to exclude 'ds' and those with None values
790
821
  attr_dict = {
791
822
  k: v for k, v in self.__dict__.items() if k != "ds" and v is not None
792
823
  }
@@ -794,6 +825,23 @@ class Grid:
794
825
  return f"{cls_name}({attr_str})"
795
826
 
796
827
  def _create_horizontal_grid(self) -> xr.Dataset():
828
+ """Create the horizontal grid based on a Mercator projection and store it in the
829
+ 'ds' attribute.
830
+
831
+ Parameters
832
+ ----------
833
+ None
834
+
835
+ Returns
836
+ -------
837
+ xr.Dataset
838
+ The created horizontal grid dataset, including coordinates, grid metrics, angles, and metadata.
839
+
840
+ Notes
841
+ -----
842
+ - Longitude values are adjusted to fall within the range [0, 360].
843
+ - Grid rotation and translation are applied based on the specified parameters.
844
+ """
797
845
  if self.verbose:
798
846
  start_time = time.time()
799
847
  logging.info("=== Creating the horizontal grid ===")
@@ -831,7 +879,24 @@ class Grid:
831
879
  object.__setattr__(self, "ds", ds)
832
880
 
833
881
  def _add_global_metadata(self, ds):
882
+ """Add global metadata and attributes to the dataset.
834
883
 
884
+ Parameters
885
+ ----------
886
+ ds : xr.Dataset
887
+ Dataset to which global metadata and attributes will be added.
888
+
889
+ Returns
890
+ -------
891
+ xr.Dataset
892
+ The dataset with added global metadata, including grid type, tool version,
893
+ grid dimensions, center coordinates, and rotation.
894
+
895
+ Notes
896
+ -----
897
+ - The "spherical" attribute indicates the grid type and is set to "T" (spherical).
898
+ - The ROMS-Tools version is included as "roms_tools_version". If unavailable, it defaults to "unknown".
899
+ """
835
900
  ds["spherical"] = xr.DataArray(np.array("T", dtype="S1"))
836
901
  ds["spherical"].attrs["Long_name"] = "Grid type logical switch"
837
902
  ds["spherical"].attrs["option_T"] = "spherical"
@@ -854,6 +919,11 @@ class Grid:
854
919
  return ds
855
920
 
856
921
  def _raise_if_domain_size_too_large(self):
922
+ """Raise a ValueError if the domain size exceeds the allowable threshold.
923
+
924
+ Checks if either the x or y domain size exceeds 20,000 km and raises an error
925
+ with appropriate details if the threshold is surpassed.
926
+ """
857
927
  threshold = 20000
858
928
  if self.size_x > threshold or self.size_y > threshold:
859
929
  raise ValueError(
@@ -863,7 +933,20 @@ class Grid:
863
933
  )
864
934
 
865
935
  def _make_initial_lon_lat_ds(self):
866
- # Mercator projection around the equator
936
+ """Generate initial longitude and latitude arrays with Mercator projection
937
+ around the equator.
938
+
939
+ Returns
940
+ -------
941
+ dict
942
+ A dictionary containing the following arrays:
943
+ - lon, lat: 2D arrays of longitudes and latitudes at cell centers.
944
+ - lonu, latu: 2D arrays of longitudes and latitudes at u-points.
945
+ - lonv, latv: 2D arrays of longitudes and latitudes at v-points.
946
+ - lonq, latq: 2D arrays of longitudes and latitudes at cell corners.
947
+ """
948
+
949
+ r_earth = 6371315.0
867
950
 
868
951
  # initially define the domain to be longer in x-direction (dimension "length")
869
952
  # than in y-direction (dimension "width") to keep grid distortion minimal
@@ -874,35 +957,31 @@ class Grid:
874
957
  domain_length, domain_width = self.size_x * 1e3, self.size_y * 1e3 # in m
875
958
  nl, nw = self.nx, self.ny
876
959
 
877
- domain_length_in_degrees = domain_length / RADIUS_OF_EARTH
878
- domain_width_in_degrees = domain_width / RADIUS_OF_EARTH
960
+ domain_length_in_degrees = domain_length / r_earth
961
+ domain_width_in_degrees = domain_width / r_earth
879
962
 
880
- # 1d array describing the longitudes at cell centers
881
- x = np.arange(-0.5, nl + 1.5, 1)
882
- lon_array_1d_in_degrees = (
883
- domain_length_in_degrees * x / nl - domain_length_in_degrees / 2
963
+ # Generate 1D longitude arrays at cell centers and corners
964
+ lon_array_1d_in_degrees = domain_length_in_degrees * (
965
+ np.arange(-0.5, nl + 1.5) / nl - 0.5
884
966
  )
885
- # 1d array describing the longitudes at cell corners (or vorticity points "q")
886
- xq = np.arange(-1, nl + 2, 1)
887
- lonq_array_1d_in_degrees_q = (
888
- domain_length_in_degrees * xq / nl - domain_length_in_degrees / 2
967
+ lonq_array_1d_in_degrees_q = domain_length_in_degrees * (
968
+ np.arange(-1, nl + 2) / nl - 0.5
889
969
  )
890
970
 
891
- # convert degrees latitude to y-coordinate using Mercator projection
971
+ # Mercator projection for latitude
892
972
  y1 = np.log(np.tan(np.pi / 4 - domain_width_in_degrees / 4))
893
973
  y2 = np.log(np.tan(np.pi / 4 + domain_width_in_degrees / 4))
894
974
 
895
- # linearly space points in y-space
896
- y = (y2 - y1) * np.arange(-0.5, nw + 1.5, 1) / nw + y1
897
- yq = (y2 - y1) * np.arange(-1, nw + 2) / nw + y1
898
-
899
- # inverse Mercator projections
900
- lat_array_1d_in_degrees = np.arctan(np.sinh(y))
901
- latq_array_1d_in_degrees = np.arctan(np.sinh(yq))
975
+ # Generate 1D latitude arrays with inverse Mercator projection
976
+ lat_array_1d_in_degrees = np.arctan(
977
+ np.sinh((y2 - y1) * (np.arange(-0.5, nw + 1.5) / nw) + y1)
978
+ )
979
+ latq_array_1d_in_degrees = np.arctan(
980
+ np.sinh((y2 - y1) * (np.arange(-1, nw + 2) / nw) + y1)
981
+ )
902
982
 
903
- # 2d grid at cell centers
983
+ # 2D grids for cell centers and corners
904
984
  lon, lat = np.meshgrid(lon_array_1d_in_degrees, lat_array_1d_in_degrees)
905
- # 2d grid at cell corners
906
985
  lonq, latq = np.meshgrid(lonq_array_1d_in_degrees_q, latq_array_1d_in_degrees)
907
986
 
908
987
  if self.size_y > self.size_x:
@@ -917,7 +996,7 @@ class Grid:
917
996
  lonq = np.transpose(np.flip(lonq, 0))
918
997
  latq = np.transpose(np.flip(latq, 0))
919
998
 
920
- # infer longitudes and latitudes at u- and v-points
999
+ # Inference for u- and v-point coordinates
921
1000
  lonu = 0.5 * (lon[:, :-1] + lon[:, 1:])
922
1001
  latu = 0.5 * (lat[:, :-1] + lat[:, 1:])
923
1002
  lonv = 0.5 * (lon[:-1, :] + lon[1:, :])
@@ -937,6 +1016,22 @@ class Grid:
937
1016
  return coords
938
1017
 
939
1018
  def _create_grid_ds(self, coords):
1019
+ """Create an xarray Dataset with grid coordinates and metrics.
1020
+
1021
+ Parameters
1022
+ ----------
1023
+ coords : dict
1024
+ Dictionary containing:
1025
+ - lon, lat, lonu, latu, lonv, latv : 1d arrays of coordinates (degrees)
1026
+ - angle : 2d array (radians)
1027
+ - pm, pn : 2d arrays (meter^-1)
1028
+
1029
+ Returns
1030
+ -------
1031
+ xarray.Dataset
1032
+ Dataset with variables: lon_rho, lat_rho, lon_u, lat_u, lon_v, lat_v,
1033
+ angle, f (Coriolis parameter), pm, pn.
1034
+ """
940
1035
 
941
1036
  ds = xr.Dataset()
942
1037
 
@@ -1076,116 +1171,129 @@ def _translate(coords, tra_lat, tra_lon):
1076
1171
 
1077
1172
 
1078
1173
  def _rot_sphere(lon, lat, rot):
1079
- (n, m) = np.shape(lon)
1080
- # convert rotation angle from degrees to radians
1174
+ """Rotate longitude and latitude coordinates on a sphere.
1175
+
1176
+ Parameters
1177
+ ----------
1178
+ lon : ndarray
1179
+ 2D array of longitudes in radians.
1180
+ lat : ndarray
1181
+ 2D array of latitudes in radians.
1182
+ rot : float
1183
+ Rotation angle in degrees.
1184
+
1185
+ Returns
1186
+ -------
1187
+ tuple
1188
+ Rotated longitude and latitude arrays (lon, lat) in radians.
1189
+ """
1190
+ # Convert rotation angle from degrees to radians
1081
1191
  rot = rot * np.pi / 180
1082
1192
 
1083
- # translate into Cartesian coordinates x,y,z
1084
- # conventions: (lon,lat) = (0,0) corresponds to (x,y,z) = ( 0,-r, 0)
1085
- # (lon,lat) = (0,90) corresponds to (x,y,z) = ( 0, 0, r)
1193
+ # Convert spherical coordinates to Cartesian coordinates (x, y, z)
1086
1194
  x1 = np.sin(lon) * np.cos(lat)
1087
1195
  y1 = np.cos(lon) * np.cos(lat)
1088
1196
  z1 = np.sin(lat)
1089
1197
 
1090
- # We will rotate these points around the small circle defined by
1091
- # the intersection of the sphere and the plane that
1092
- # is orthogonal to the line through (lon,lat) (0,0) and (180,0)
1093
-
1094
- # The rotation is in that plane around its intersection with
1095
- # aforementioned line.
1096
-
1097
- # Since the plane is orthogonal to the y-axis (in my definition at least),
1098
- # Rotations in the plane of the small circle maintain constant y and are around
1099
- # (x,y,z) = (0,y1,0)
1100
-
1198
+ # Calculate the radial distance in the x-z plane
1101
1199
  rp1 = np.sqrt(x1**2 + z1**2)
1102
1200
 
1103
- ap1 = np.pi / 2 * np.ones((n, m))
1104
- ap1[np.abs(x1) > 1e-7] = np.arctan(
1105
- np.abs(z1[np.abs(x1) > 1e-7] / x1[np.abs(x1) > 1e-7])
1106
- )
1201
+ # Compute azimuthal angle
1202
+ ap1 = np.arctan2(np.abs(z1), np.abs(x1))
1107
1203
  ap1[x1 < 0] = np.pi - ap1[x1 < 0]
1108
1204
  ap1[z1 < 0] = -ap1[z1 < 0]
1109
1205
 
1206
+ # Apply rotation to the azimuthal angle
1110
1207
  ap2 = ap1 + rot
1111
1208
  x2 = rp1 * np.cos(ap2)
1112
1209
  y2 = y1
1113
1210
  z2 = rp1 * np.sin(ap2)
1114
1211
 
1115
- lon = np.pi / 2 * np.ones((n, m))
1116
- lon[abs(y2) > 1e-7] = np.arctan(
1117
- np.abs(x2[np.abs(y2) > 1e-7] / y2[np.abs(y2) > 1e-7])
1118
- )
1119
- lon[y2 < 0] = np.pi - lon[y2 < 0]
1120
- lon[x2 < 0] = -lon[x2 < 0]
1212
+ # Recompute longitude and latitude
1213
+ lon_rot = np.arctan2(np.abs(x2), np.abs(y2))
1214
+ lon_rot[y2 < 0] = np.pi - lon_rot[y2 < 0]
1215
+ lon_rot[x2 < 0] = -lon_rot[x2 < 0]
1121
1216
 
1122
1217
  pr2 = np.sqrt(x2**2 + y2**2)
1123
- lat = np.pi / 2 * np.ones((n, m))
1124
- lat[np.abs(pr2) > 1e-7] = np.arctan(
1125
- np.abs(z2[np.abs(pr2) > 1e-7] / pr2[np.abs(pr2) > 1e-7])
1126
- )
1127
- lat[z2 < 0] = -lat[z2 < 0]
1218
+ lat_rot = np.arctan2(np.abs(z2), pr2)
1219
+ lat_rot[z2 < 0] = -lat_rot[z2 < 0]
1128
1220
 
1129
- return (lon, lat)
1221
+ return lon_rot, lat_rot
1130
1222
 
1131
1223
 
1132
1224
  def _tra_sphere(lon, lat, tra):
1133
- (n, m) = np.shape(lon)
1134
- tra = tra * np.pi / 180 # translation in latitude direction
1225
+ """Translate longitude and latitude coordinates on a sphere in the latitude
1226
+ direction.
1135
1227
 
1136
- # translate into x,y,z
1137
- # conventions: (lon,lat) = (0,0) corresponds to (x,y,z) = ( 0,-r, 0)
1138
- # (lon,lat) = (0,90) corresponds to (x,y,z) = ( 0, 0, r)
1139
- x1 = np.sin(lon) * np.cos(lat)
1140
- y1 = np.cos(lon) * np.cos(lat)
1141
- z1 = np.sin(lat)
1228
+ Parameters
1229
+ ----------
1230
+ lon : ndarray
1231
+ 2D array of longitudes in radians.
1232
+ lat : ndarray
1233
+ 2D array of latitudes in radians.
1234
+ tra : float
1235
+ Translation angle in degrees.
1142
1236
 
1143
- # We will rotate these points around the small circle defined by
1144
- # the intersection of the sphere and the plane that
1145
- # is orthogonal to the line through (lon,lat) (90,0) and (-90,0)
1237
+ Returns
1238
+ -------
1239
+ tuple
1240
+ Translated longitude and latitude arrays (lon, lat) in radians.
1241
+ """
1146
1242
 
1147
- # The rotation is in that plane around its intersection with
1148
- # aforementioned line.
1243
+ # Convert translation angle from degrees to radians
1244
+ tra = tra * np.pi / 180
1149
1245
 
1150
- # Since the plane is orthogonal to the x-axis (in my definition at least),
1151
- # Rotations in the plane of the small circle maintain constant x and are around
1152
- # (x,y,z) = (x1,0,0)
1246
+ # Convert spherical coordinates to Cartesian coordinates (x, y, z)
1247
+ x1 = np.sin(lon) * np.cos(lat)
1248
+ y1 = np.cos(lon) * np.cos(lat)
1249
+ z1 = np.sin(lat)
1153
1250
 
1251
+ # Radial distance in the y-z plane
1154
1252
  rp1 = np.sqrt(y1**2 + z1**2)
1155
1253
 
1156
- ap1 = np.pi / 2 * np.ones((n, m))
1157
- ap1[np.abs(y1) > 1e-7] = np.arctan(
1158
- np.abs(z1[np.abs(y1) > 1e-7] / y1[np.abs(y1) > 1e-7])
1159
- )
1254
+ # Compute azimuthal angle in the y-z plane
1255
+ ap1 = np.arctan2(np.abs(z1), np.abs(y1))
1160
1256
  ap1[y1 < 0] = np.pi - ap1[y1 < 0]
1161
1257
  ap1[z1 < 0] = -ap1[z1 < 0]
1162
1258
 
1259
+ # Apply translation in the azimuthal angle
1163
1260
  ap2 = ap1 + tra
1164
- x2 = x1
1165
1261
  y2 = rp1 * np.cos(ap2)
1166
1262
  z2 = rp1 * np.sin(ap2)
1167
1263
 
1168
- ## transformation from (x,y,z) to (lat,lon)
1169
- lon = np.pi / 2 * np.ones((n, m))
1170
- lon[np.abs(y2) > 1e-7] = np.arctan(
1171
- np.abs(x2[np.abs(y2) > 1e-7] / y2[np.abs(y2) > 1e-7])
1172
- )
1173
- lon[y2 < 0] = np.pi - lon[y2 < 0]
1174
- lon[x2 < 0] = -lon[x2 < 0]
1264
+ # Convert back to spherical coordinates
1265
+ lon_rot = np.arctan2(np.abs(x1), np.abs(y2))
1266
+ lon_rot[y2 < 0] = np.pi - lon_rot[y2 < 0]
1267
+ lon_rot[x1 < 0] = -lon_rot[x1 < 0]
1175
1268
 
1176
- pr2 = np.sqrt(x2**2 + y2**2)
1177
- lat = np.pi / (2 * np.ones((n, m)))
1178
- lat[np.abs(pr2) > 1e-7] = np.arctan(
1179
- np.abs(z2[np.abs(pr2) > 1e-7] / pr2[np.abs(pr2) > 1e-7])
1180
- )
1181
- lat[z2 < 0] = -lat[z2 < 0]
1269
+ pr2 = np.sqrt(x1**2 + y2**2)
1270
+ lat_rot = np.arctan2(np.abs(z2), pr2)
1271
+ lat_rot[z2 < 0] = -lat_rot[z2 < 0]
1182
1272
 
1183
- return (lon, lat)
1273
+ return lon_rot, lat_rot
1184
1274
 
1185
1275
 
1186
1276
  def _compute_coordinate_metrics(coords):
1187
- """Compute the curvilinear coordinate metrics pn and pm, defined as 1/grid
1188
- spacing."""
1277
+ """Compute the reciprocal of grid spacing (`pn` and `pm`) in the latitude and
1278
+ longitude directions.
1279
+
1280
+ Parameters
1281
+ ----------
1282
+ coords : dict
1283
+ A dictionary containing coordinate arrays 'lonu', 'latu', 'lonv', and 'latv' for the u- and v-velocity points.
1284
+
1285
+ Returns
1286
+ -------
1287
+ pn : ndarray
1288
+ The metric for the latitude direction (1/dy).
1289
+
1290
+ pm : ndarray
1291
+ The metric for the longitude direction (1/dx).
1292
+
1293
+ Notes
1294
+ -----
1295
+ Boundary values of `pn` and `pm` are copied from adjacent interior values.
1296
+ """
1189
1297
 
1190
1298
  # pm = 1/dx
1191
1299
  pmu = gc_dist(
@@ -1193,9 +1301,11 @@ def _compute_coordinate_metrics(coords):
1193
1301
  coords["latu"][:, :-1],
1194
1302
  coords["lonu"][:, 1:],
1195
1303
  coords["latu"][:, 1:],
1304
+ input_in_degrees=False,
1196
1305
  )
1197
- pm = 0 * coords["lon"]
1306
+ pm = np.zeros_like(coords["lon"])
1198
1307
  pm[:, 1:-1] = pmu
1308
+ # Handle boundary conditions
1199
1309
  pm[:, 0] = pm[:, 1]
1200
1310
  pm[:, -1] = pm[:, -2]
1201
1311
  pm = 1 / pm
@@ -1206,9 +1316,11 @@ def _compute_coordinate_metrics(coords):
1206
1316
  coords["latv"][:-1, :],
1207
1317
  coords["lonv"][1:, :],
1208
1318
  coords["latv"][1:, :],
1319
+ input_in_degrees=False,
1209
1320
  )
1210
- pn = 0 * coords["lon"]
1321
+ pn = np.zeros_like(coords["lon"])
1211
1322
  pn[1:-1, :] = pnv
1323
+ # Handle boundary conditions
1212
1324
  pn[0, :] = pn[1, :]
1213
1325
  pn[-1, :] = pn[-2, :]
1214
1326
  pn = 1 / pn
@@ -1216,44 +1328,46 @@ def _compute_coordinate_metrics(coords):
1216
1328
  return pn, pm
1217
1329
 
1218
1330
 
1219
- def gc_dist(lon1, lat1, lon2, lat2):
1220
- # Distance between 2 points along a great circle
1221
- # lat and lon in radians!!
1222
- # 2008, Jeroen Molemaker, UCLA
1223
-
1224
- dlat = lat2 - lat1
1225
- dlon = lon2 - lon1
1226
-
1227
- dang = 2 * np.arcsin(
1228
- np.sqrt(
1229
- np.sin(dlat / 2) ** 2 + np.cos(lat2) * np.cos(lat1) * np.sin(dlon / 2) ** 2
1230
- )
1231
- ) # haversine function
1232
-
1233
- dis = RADIUS_OF_EARTH * dang
1331
+ def _compute_angle(coords):
1332
+ """Compute angles of the local grid's positive x-axis relative to east.
1234
1333
 
1235
- return dis
1334
+ The angle is computed for each grid cell using the latitude and longitude
1335
+ differences between neighboring grid points. The result is wrapped to
1336
+ the range [-π, π] and adjusted based on longitude and latitude conditions.
1236
1337
 
1338
+ Parameters
1339
+ ----------
1340
+ coords : dict
1341
+ A dictionary containing 'latu' (latitudes) and 'lonu' (longitudes) arrays.
1237
1342
 
1238
- def _compute_angle(coords):
1239
- """Compute angles of local grid positive x-axis relative to east."""
1343
+ Returns
1344
+ -------
1345
+ ang : ndarray
1346
+ An array of angles (in radians) of the local grid's positive x-axis
1347
+ relative to east for each grid point.
1348
+ """
1240
1349
 
1350
+ # Compute differences in latitudes and longitudes
1241
1351
  dellat = coords["latu"][:, 1:] - coords["latu"][:, :-1]
1242
1352
  dellon = coords["lonu"][:, 1:] - coords["lonu"][:, :-1]
1243
- dellon[dellon > np.pi] = dellon[dellon > np.pi] - 2 * np.pi
1244
- dellon[dellon < -np.pi] = dellon[dellon < -np.pi] + 2 * np.pi
1245
- dellon = dellon * np.cos(0.5 * (coords["latu"][:, 1:] + coords["latu"][:, :-1]))
1246
1353
 
1247
- ang = copy.copy(coords["lon"])
1248
- ang_s = np.arctan(dellat / (dellon + 1e-16))
1249
- ang_s[(dellon < 0) & (dellat < 0)] = ang_s[(dellon < 0) & (dellat < 0)] - np.pi
1250
- ang_s[(dellon < 0) & (dellat >= 0)] = ang_s[(dellon < 0) & (dellat >= 0)] + np.pi
1251
- ang_s[ang_s > np.pi] = ang_s[ang_s > np.pi] - np.pi
1252
- ang_s[ang_s < -np.pi] = ang_s[ang_s < -np.pi] + np.pi
1354
+ # Normalize longitude differences to the range [-π, π]
1355
+ dellon = (dellon + np.pi) % (2 * np.pi) - np.pi
1356
+ dellon *= np.cos(0.5 * (coords["latu"][:, 1:] + coords["latu"][:, :-1]))
1357
+
1358
+ # Compute the angle in radians
1359
+ ang_s = np.arctan2(dellat, dellon)
1360
+
1361
+ # Adjust angles based on longitude and latitude conditions
1362
+ ang_s[(dellon < 0) & (dellat < 0)] -= np.pi
1363
+ ang_s[(dellon < 0) & (dellat >= 0)] += np.pi
1364
+ ang_s = np.mod(ang_s + np.pi, 2 * np.pi) - np.pi # Ensure angles are in [-π, π]
1253
1365
 
1366
+ # Create output array and set angles
1367
+ ang = np.zeros_like(coords["lon"])
1254
1368
  ang[:, 1:-1] = ang_s
1255
- ang[:, 0] = ang[:, 1]
1256
- ang[:, -1] = ang[:, -2]
1369
+ ang[:, 0] = ang[:, 1] # Set first column to the second column
1370
+ ang[:, -1] = ang[:, -2] # Set last column to the second-to-last column
1257
1371
 
1258
1372
  return ang
1259
1373