Merge pull request #14 from hdrake/bounded-regions

Bounded regions

Author: hdrake

examples/1_thickness_budget.ipynb

@@ -18,7 +18,7 @@ dependencies = [
     "geopandas",
     "regionmask",
     "contourpy",
-    "sectionate @ git+https://github.com/MOM6-community/sectionate.git@master"
+    "sectionate @ git+https://github.com/hdrake/sectionate.git@master"
 ]
 
 [project.urls]

examples/2_advective_heat_convergence.ipynb

@@ -18,7 +18,7 @@ def grid_boundaries_from_mask(grid, mask):
     
     RETURNS
     -------
-    i_list, j_list, lons_list, lats_list
+    i_c_list, j_c_list, lons_c_list, lats_c_list
     """
 
     symmetric = check_symmetric(grid)
@@ -31,33 +31,33 @@ def grid_boundaries_from_mask(grid, mask):
         )
         .create_contour(0.5)
     )
-    i_list = []
-    j_list = []
-    lons_list = []
-    lats_list = []
+    i_c_list = []
+    j_c_list = []
+    lons_c_list = []
+    lats_c_list = []
     for c in contours:
-        i, j = c[:-1,0], c[:-1,1]
+        i_c, j_c = c[:-1,0], c[:-1,1]
 
-        i_new, j_new = i.copy(), j.copy()
+        i_c_new, j_c_new = i_c.copy(), j_c.copy()
 
-        i_inc = np.roll(i, -1)-i
-        j_inc = np.roll(j, -1)-j
+        i_inc = np.roll(i_c, -1)-i_c
+        j_inc = np.roll(j_c, -1)-j_c
 
-        i_new[(i%1)==0.0] = (i - (i_inc<0))[(i%1)==0.0] + symmetric
-        j_new[(j%1)==0.0] = (j - (j_inc<0))[(j%1)==0.0] + symmetric
-        i_new[(i%1)==0.5] = np.floor(i[(i%1)==0.5]) + symmetric
-        j_new[(j%1)==0.5] = np.floor(j[(j%1)==0.5]) + symmetric
+        i_c_new[(i_c%1)==0.0] = (i_c - (i_inc<0))[(i_c%1)==0.0] + symmetric
+        j_c_new[(j_c%1)==0.0] = (j_c - (j_inc<0))[(j_c%1)==0.0] + symmetric
+        i_c_new[(i_c%1)==0.5] = np.floor(i_c[(i_c%1)==0.5]) + symmetric
+        j_c_new[(j_c%1)==0.5] = np.floor(j_c[(j_c%1)==0.5]) + symmetric
 
-        i_new, j_new = loop(i_new).astype(np.int64), loop(j_new).astype(np.int64)
+        i_c_new, j_c_new = loop(i_c_new).astype(np.int64), loop(j_c_new).astype(np.int64)
 
-        i_list.append(i_new)
-        j_list.append(j_new)
+        i_c_list.append(i_c_new)
+        j_c_list.append(j_c_new)
 
         idx = {
-            "xq":xr.DataArray(i_new, dims=("pt",)),
-            "yq":xr.DataArray(j_new, dims=("pt",))
+            "xq":xr.DataArray(i_c_new, dims=("pt",)),
+            "yq":xr.DataArray(j_c_new, dims=("pt",))
         }
-        lons_list.append(grid._ds["geolon_c"].isel(idx).values[:-1])
-        lats_list.append(grid._ds["geolat_c"].isel(idx).values[:-1])
+        lons_c_list.append(grid._ds["geolon_c"].isel(idx).values[:-1])
+        lats_c_list.append(grid._ds["geolat_c"].isel(idx).values[:-1])
 
-    return i_list, j_list, lons_list, lats_list
+    return i_c_list, j_c_list, lons_c_list, lats_c_list

examples/3_Arctic_heat_CM4p25.ipynb

@@ -18,10 +18,10 @@ def get_region_boundary_grid_indices(lons, lats, grid):
     
     RETURNS
     -------
-    (i, j, lons_c, lats_c, lons_uv, lats_uv)
+    (i_c, j_c, lons_c, lats_c, lons_uv, lats_uv)
 
-    i : "X"-axis grid indices of corner points
-    j : "Y"-axis grid indices of corner points
+    i_c : "X"-axis grid indices of corner points
+    j_c : "Y"-axis grid indices of corner points
     lons_c : longitudes of corner points
     lats_c : latitudes of corner points
     lons_uv : longitudes of (u,v) velocity faces
@@ -30,21 +30,15 @@ def get_region_boundary_grid_indices(lons, lats, grid):
     if (lons[0], lats[0]) != (lons[-1], lats[-1]):
         lons, lats = loop(lons), loop(lats)
 
-    i, j, lons_c, lats_c = sec.grid_section(
+    i_c, j_c, lons_c, lats_c = sec.grid_section(
         grid,
         lons,
         lats,
         topology="MOM-tripolar"
     )
-    
-    uvindices = sec.uvindices_from_qindices(grid, i, j)
-    
-    lons_uv, lats_uv = sec.uvcoords_from_uvindices(
-        grid,
-        uvindices
-    )
+    lons_uv, lats_uv = sec.uvcoords_from_qindices(grid, i_c, j_c)
 
-    return (i, j, lons_c[:-1], lats_c[:-1], lons_uv, lats_uv)
+    return (i_c, j_c, lons_c, lats_c, lons_uv, lats_uv)
 
 def mask_from_grid_boundaries(
     lons_c,

examples/4_bounded_by_named_sections.ipynb

@@ -3,12 +3,12 @@
 
 def find_indices_of_overlaps(r1, r2, closeness_threshold=5.e3, face_indices=False):
     overlaps = haversine(
-            r1.lons[:, np.newaxis], r1.lats[:, np.newaxis],
-            r2.lons[np.newaxis, :], r2.lats[np.newaxis, :]
+            r1.lons_c[:, np.newaxis], r1.lats_c[:, np.newaxis],
+            r2.lons_c[np.newaxis, :], r2.lats_c[np.newaxis, :]
         ) < closeness_threshold
 
-    r1_oidx = consecutive_lists( np.where(np.any(overlaps, axis=1))[0], mod=r1.lons.size)
-    r2_oidx = consecutive_lists( np.where(np.any(overlaps, axis=0))[0], mod=r2.lons.size)
+    r1_oidx = consecutive_lists( np.where(np.any(overlaps, axis=1))[0], mod=r1.lons_c.size)
+    r2_oidx = consecutive_lists( np.where(np.any(overlaps, axis=0))[0], mod=r2.lons_c.size)
 
     if face_indices:
         r1_oidx = [l[:-1] for l in r1_oidx]
@@ -23,8 +23,8 @@ def group_overlaps(overlaps, r1, r2, closeness_threshold=5.e3):
         for o1 in overlaps[r1.name]:
             if np.any(o1) and np.any(o2):
                 if np.any(haversine(
-                        r1.lons[o1, np.newaxis], r1.lats[o1, np.newaxis],
-                        r2.lons[np.newaxis, o2], r2.lats[np.newaxis, o2]
+                        r1.lons_c[o1, np.newaxis], r1.lats_c[o1, np.newaxis],
+                        r2.lons_c[np.newaxis, o2], r2.lats_c[np.newaxis, o2]
                     ) < closeness_threshold):
                     grouped_overlaps[r1.name][group_num] = o1
                     grouped_overlaps[r2.name][group_num] = o2
@@ -42,21 +42,21 @@ def align_boundaries_with_overlap_sections(regions, remove_gaps=True):
             )
 
 def roll_boundary_to_align_with_overlap(r, oidx, remove_gaps=True):
-    roll_idx = -consecutive_lists(oidx[0], mod=r.lons.size)[0][0]
-    new_oidx = {onum: np.mod(np.array(o) + roll_idx, r.lons.size) for (onum, o) in oidx.items()}
+    roll_idx = -consecutive_lists(oidx[0], mod=r.lons_c.size)[0][0]
+    new_oidx = {onum: np.mod(np.array(o) + roll_idx, r.lons_c.size) for (onum, o) in oidx.items()}
 
-    r.lons = np.roll(r.lons, roll_idx)
-    r.lats = np.roll(r.lats, roll_idx)
+    r.lons_c = np.roll(r.lons_c, roll_idx)
+    r.lats_c = np.roll(r.lats_c, roll_idx)
 
     if remove_gaps:
         for onum, o in new_oidx.items():
             gaps = np.array([i for i in range(np.min(o), np.max(o)+1) if i not in o])
             if any(gaps):
-                r.lons[gaps] = np.nan
-                r.lats[gaps] = np.nan
+                r.lons_c[gaps] = np.nan
+                r.lats_c[gaps] = np.nan
 
-    nan_idx = np.isnan(r.lons) | np.isnan(r.lats)
-    r.lons = r.lons[~nan_idx] 
-    r.lats = r.lats[~nan_idx]
+    nan_idx = np.isnan(r.lons_c) | np.isnan(r.lats_c)
+    r.lons_c = r.lons_c[~nan_idx] 
+    r.lats_c = r.lats_c[~nan_idx]
 
     return r