Add bucket and orbit utilities

Author: ghiggi

Diff for: gpm/bucket/analysis.py

@@ -98,7 +98,6 @@ def get_swath_indices(df):
     df_along.columns = ["granule_id", "along_track_id"]
 
     # We will assign new x indices so that each granule's along-track block is contiguous.
-    # We also build a full list of x indices.
     x_index_list = []
     # Allocate an array to hold the new x value for each row in df.
     x_values = np.empty(len(df), dtype=int)

Diff for: gpm/etc/colormaps/categorical/surfaceTypeIndexPalette.yaml

@@ -1,6 +1,26 @@
 ---
 colormap_type: ListedColormap
 color_space: hex
-color_palette: ["#1f78b4", "#E0FFFF", "#006400", "#228B22", "#32CD32", "#98FB98", "#F4A460", "#7D7F80", "#9E9E9E", "#BDBDBD", "#E0E0E0", "#87CEFA", "#5F9EA0", "#ff4c00", "#F5DEB3", "#B0E0E6", "#708090", "#B0BEC5"]
+color_palette:
+  [
+    "#1f78b4",
+    "#E0FFFF",
+    "#006400",
+    "#228B22",
+    "#32CD32",
+    "#98FB98",
+    "#F4A460",
+    "#7D7F80",
+    "#9E9E9E",
+    "#BDBDBD",
+    "#E0E0E0",
+    "#87CEFA",
+    "#5F9EA0",
+    "#ff4c00",
+    "#F5DEB3",
+    "#B0E0E6",
+    "#708090",
+    "#B0BEC5",
+  ]
 auxiliary:
   category: [surfaceTypeIndex]

Diff for: gpm/tests/test_utils/test_orbit.py

@@ -27,6 +27,7 @@
 """This module tests the dataframe utilities functions."""
 import numpy as np
 import pytest
+
 from gpm.utils.orbit import adjust_short_sequences
 
 
@@ -67,36 +68,39 @@ def test_replace_short_sequence_end(self):
         """Replace a short ending sequence with the previous sequence value."""
         # The ending sequence [2] is short (< min_size=2) and should be replaced by 1.
         arr = [1, 1, 1, 2]
-        result = adjust_short_sequences(arr, min_size=2)  
+        result = adjust_short_sequences(arr, min_size=2)
         expected = np.array([1, 1, 1, 1])
         np.testing.assert_array_equal(result, expected)
 
     def test_min_size_one(self):
         """Return unchanged array when min_size is 1 (all sequences allowed)."""
         arr = [1, 2, 3, 4]
-        result = adjust_short_sequences(arr, min_size=1)  
+        result = adjust_short_sequences(arr, min_size=1)
         np.testing.assert_array_equal(result, arr)
 
     def test_non_1d_input(self):
         """Raise ValueError for non 1D array input."""
         arr = [[1, 1, 1], [2, 2, 2]]
         with pytest.raises(ValueError):
             adjust_short_sequences(arr, min_size=2)
-        
-    @pytest.mark.parametrize("arr,expected", [
-        ( 
-            [1, -1, -1, -1, 1, 1, 1, 1, -1], # Short sequences at the edges
-            [-1, -1, -1, -1, 1, 1, 1, 1, 1]
-        ), 
-        ( 
-            [1, -1, -1, -1, 1, 1, 1, 1, -1, -1], # Short sequence at start  
-            [-1, -1, -1, -1, 1, 1, 1, 1, -1, -1]
-        ),
-        (
-            [1, 1, -1, 1, 1, 1, 1, 1, -1], #Short sequence in the middle and at the end 
-            [1, 1, 1, 1, 1, 1, 1, 1, 1]
-        ),
-    ])
+
+    @pytest.mark.parametrize(
+        ("arr", "expected"),
+        [
+            (
+                [1, -1, -1, -1, 1, 1, 1, 1, -1],  # Short sequences at the edges
+                [-1, -1, -1, -1, 1, 1, 1, 1, 1],
+            ),
+            (
+                [1, -1, -1, -1, 1, 1, 1, 1, -1, -1],  # Short sequence at start
+                [-1, -1, -1, -1, 1, 1, 1, 1, -1, -1],
+            ),
+            (
+                [1, 1, -1, 1, 1, 1, 1, 1, -1],  # Short sequence in the middle and at the end
+                [1, 1, 1, 1, 1, 1, 1, 1, 1],
+            ),
+        ],
+    )
     def test_edge_cases(self, arr, expected):
         """Test various edge cases with short sequences at the edges."""
         result = adjust_short_sequences(arr, min_size=2)

Diff for: gpm/utils/orbit.py

@@ -26,27 +26,27 @@
 # -----------------------------------------------------------------------------.
 """This module contains utilities for orbit processing."""
 
-import numpy as np 
-import pandas as pd 
+import numpy as np
+import pandas as pd
 import xarray as xr
 
 
 def adjust_short_sequences(arr, min_size):
     """
     Replace value of short sequences of consecutive identical values.
-    
+
     The function examines contiguous sequences of identical elements in the input
     array.
     If a sequence is shorter than `min_size`, its values are replaced with
-    the value of the adjacent longer sequence, working outward from the first 
+    the value of the adjacent longer sequence, working outward from the first
     valid sequence.
-    
+
     Parameters
     ----------
     arr : array_like
         The input array of values.
     min_size : int
-        The minimum number of consecutive identical elements to not be modified. 
+        The minimum number of consecutive identical elements to not be modified.
         Shorter sequences will be replaced with the previous sequence value.
 
     Returns
@@ -59,33 +59,33 @@ def adjust_short_sequences(arr, min_size):
     arr = np.asarray(arr)
     if arr.ndim != 1:
         raise ValueError("Input must be a 1D array.")
-    
+
     # If array is empty or has only one element, return as is
     if len(arr) <= 1:
         return arr
- 
+
     # Create a copy to modify
     result = arr.copy()
-   
+
     # Find boundaries of sequences (where values change)
     change_indices = np.where(np.diff(result) != 0)[0]
     sequence_starts = np.concatenate(([0], change_indices + 1))
     sequence_ends = np.concatenate((change_indices + 1, [len(result)]))
-    
+
     # Find the first valid sequence (length >= min_size)
     first_valid_idx = None
     valid_value = None
-    
-    for i, (start, end) in enumerate(zip(sequence_starts, sequence_ends)):
+
+    for i, (start, end) in enumerate(zip(sequence_starts, sequence_ends, strict=False)):
         if end - start >= min_size:
             first_valid_idx = i
             valid_value = result[start]
             break
-    
+
     if first_valid_idx is None:
         # If no valid sequence found, return original array
         return result
-        
+
     # Process sequences after the first valid sequence
     for i in range(first_valid_idx + 1, len(sequence_starts)):
         start = sequence_starts[i]
@@ -94,15 +94,15 @@ def adjust_short_sequences(arr, min_size):
             result[start:end] = valid_value
         else:
             valid_value = result[start]
-            
+
     # Process sequences before the first valid sequence (in reverse)
     valid_value = result[sequence_starts[first_valid_idx]]  # Reset to first valid sequence value
     for i in range(first_valid_idx - 1, -1, -1):
         start = sequence_starts[i]
         end = sequence_ends[i]
         if end - start < min_size:
             result[start:end] = valid_value
-            
+
     # Return array with replaced values
     return result
 
@@ -111,31 +111,31 @@ def get_orbit_direction(lats, n_tol=1):
     """
     Infer the satellite orbit direction from latitude values.
 
-    This function determines the orbit direction by computing the sign 
-    of the differences between consecutive latitude values. 
-    
-    A positive sign (+1) indicates an ascending orbit (increasing latitude), 
-    while a negative sign (-1) indicates a descending orbit (decreasing latitude). 
-    
+    This function determines the orbit direction by computing the sign
+    of the differences between consecutive latitude values.
+
+    A positive sign (+1) indicates an ascending orbit (increasing latitude),
+    while a negative sign (-1) indicates a descending orbit (decreasing latitude).
+
     Any zero differences are replaced by the nearest nonzero direction.
-    Additionally, short sequences of direction changes - those lasting fewer 
-    than `n_tol` consecutive data points - are adjusted to reduce 
+    Additionally, short sequences of direction changes - those lasting fewer
+    than `n_tol` consecutive data points - are adjusted to reduce
     the influence of geolocation errors.
 
     Parameters
     ----------
     lats : array_like
         1-dimensional array of latitude values corresponding to the satellite's orbit.
     n_tol : int, optional
-        The minimum number of consecutive data points required to confirm a change 
+        The minimum number of consecutive data points required to confirm a change
         in direction.
         Sequences shorter than this threshold will be smoothed. Default is 1.
 
     Returns
     -------
     numpy.ndarray
-        A 1-dimensional array of the same length as `lats` containing 
-        the inferred orbit direction. 
+        A 1-dimensional array of the same length as `lats` containing
+        the inferred orbit direction.
         A value of +1 denotes an ascending orbit and -1 denotes a descending orbit.
 
     Examples
@@ -147,11 +147,11 @@ def get_orbit_direction(lats, n_tol=1):
     # Get direction (1 for ascending, -1 for descending)
     directions = np.sign(np.diff(lats))
     directions = np.append(directions[0], directions)  # Include starting point
-    # Set 0 to NaN and infill values 
+    # Set 0 to NaN and infill values
     directions = directions.astype(float)
     directions[directions == 0] = np.nan
     if np.all(np.isnan(directions)):
-        raise ValueError("Invalid orbit.") 
+        raise ValueError("Invalid orbit.")
     directions = pd.Series(directions).ffill().bfill().to_numpy()
     # Remove short consecutive changes in direction to account for geolocation errors
     directions = adjust_short_sequences(directions, min_size=n_tol)
@@ -162,19 +162,19 @@ def get_orbit_direction(lats, n_tol=1):
 
 def get_orbit_mode(ds):
     # Retrieve latitude coordinates
-    if "SClatitude" in ds: 
+    if "SClatitude" in ds:
         lats = ds["SClatitude"]
-    elif "scLat" in ds: 
+    elif "scLat" in ds:
         lats = ds["scLat"]
-    else: 
+    else:
         # Define cross_track idx defining the orbit coordinates
-        # - TODO: conically scanning vs cross-track scanning 
+        # - TODO: conically scanning vs cross-track scanning
         # - TODO: Use SClatitude  ?
-        # Remove invalid outer cross-track_id if selecting 0 or -1 
+        # Remove invalid outer cross-track_id if selecting 0 or -1
         # from gpm.visualization.orbit import remove_invalid_outer_cross_track
         # ds, _ = remove_invalid_outer_cross_track(ds)
         idx = int(ds["cross_track"].shape[0] / 2)
         lats = ds["lat"].isel({"cross_track": idx}).to_numpy()
     directions = get_orbit_direction(lats, n_tol=10)
     orbit_mode = xr.DataArray(directions, dims=["along_track"])
-    return orbit_mode
+    return orbit_mode