changes to comp_z

pyart/retrieve/comp_z.py

@@ -13,7 +13,13 @@
 from pyart.core import Radar
 
 
-def composite_reflectivity(radar, field="reflectivity", gatefilter=None):
+def composite_reflectivity(
+    radar,
+    field="reflectivity",
+    gatefilter=None,
+    same_nyquist=False,
+    nyquist_vector_idx=0,
+):
     """
     Composite Reflectivity
 
@@ -41,6 +47,15 @@ def composite_reflectivity(radar, field="reflectivity", gatefilter=None):
     gatefilter : GateFilter
         GateFilter instance. None will result in no gatefilter mask being
         applied to data.
+    same_nyquist : bool
+        During a volume scan (i.e., file) the PRF (Nyquist velocity) can change.
+        This can create odd artifacts when data quality is low on certain scans.
+        To avoid this, only the max of scans sharing the reference sweep's
+        Nyquist (+/- 1 m/s) is taken. Default is off (False); set to True to
+        enable this filtering (requires the radar to have Nyquist velocity).
+    nyquist_vector_idx : int
+        Index of the reference sweep whose Nyquist the other sweeps are matched
+        to when same_nyquist is True. Default is 0 (the first sweep).
 
     Returns
     -------
@@ -67,6 +82,10 @@ def composite_reflectivity(radar, field="reflectivity", gatefilter=None):
         z = radar.get_field(sweep, field)
         z_dtype = z.dtype
 
+        # get the nyquist (only needed when filtering sweeps by matching nyquist)
+        if same_nyquist:
+            nyquist = np.asarray([np.round(radar.get_nyquist_vel(sweep=sweep))])
+
         # Use gatefilter
         if gatefilter is not None:
             mask_sweep = gatefilter.gate_excluded[sweep_slice, :]
@@ -111,8 +130,20 @@ def composite_reflectivity(radar, field="reflectivity", gatefilter=None):
         # if first sweep, create new dim, otherwise concat them up
         if sweep == minimum_sweep:
             z_stack = copy.deepcopy(z[np.newaxis, :, :])
+            if same_nyquist:
+                nyquist_stack = copy.deepcopy(nyquist[np.newaxis, :])
         else:
             z_stack = np.concatenate([z_stack, z[np.newaxis, :, :]])
+            if same_nyquist:
+                nyquist_stack = np.concatenate([nyquist_stack, nyquist[np.newaxis, :]])
+
+    # only stack up sweeps with the same nyquist
+    if same_nyquist:
+        left = np.where(nyquist_stack >= nyquist_stack[nyquist_vector_idx] - 1)[0]
+        right = np.where(nyquist_stack <= nyquist_stack[nyquist_vector_idx] + 1)[0]
+        same_ny = np.intersect1d(left, right)
+
+        z_stack = z_stack[same_ny]
 
     # now that the stack is made, take max across vertical
     compz = z_stack.max(axis=0).astype(z_dtype)

tests/retrieve/test_comp_z.py

@@ -53,7 +53,7 @@ def test_composite_z():
     z_new["data"] = z.astype("float32")
     radar.add_field("reflectivity", z_new, replace_existing=True)
     compz = pyart.retrieve.composite_reflectivity(
-        radar, field=ref_field, gatefilter=gatefilter
+        radar, field=ref_field, gatefilter=gatefilter, same_nyquist=False
     )
     assert_equal(compz.fields["composite_reflectivity"]["data"].max(), 40)
 
@@ -80,7 +80,7 @@ def test_composite_z():
     z_new["data"] = z.astype("float32")
     radar.add_field("reflectivity", z_new, replace_existing=True)
     compz = pyart.retrieve.composite_reflectivity(
-        radar, field=ref_field, gatefilter=gatefilter
+        radar, field=ref_field, gatefilter=gatefilter, same_nyquist=False
     )
 
     # choose a random az
@@ -89,3 +89,38 @@ def test_composite_z():
         compz.fields["composite_reflectivity"]["data"][random_az, :],
         np.arange(0, z.shape[1]),
     )
+
+
+def test_composite_z_same_nyquist():
+    # NEXRAD VCPs mix Nyquist velocities across sweeps; same_nyquist keeps only
+    # the sweeps whose Nyquist matches the reference sweep (nyquist_vector_idx).
+    radar = pyart.io.read(pyart.testing.NEXRAD_ARCHIVE_MSG31_FILE)
+    ref_field = "reflectivity"
+
+    ny0 = np.round(radar.get_nyquist_vel(sweep=0))
+    # pick a sweep whose Nyquist does NOT match sweep 0 (the default reference)
+    target = next(
+        sweep
+        for sweep in radar.sweep_number["data"]
+        if abs(np.round(radar.get_nyquist_vel(sweep=sweep)) - ny0) > 1
+    )
+
+    z = np.zeros(radar.fields[ref_field]["data"].shape)
+    s_idx = radar.sweep_start_ray_index["data"][target]
+    e_idx = radar.sweep_end_ray_index["data"][target] + 1
+    z[s_idx:e_idx, :] = 40
+    z_new = copy.deepcopy(radar.fields[ref_field])
+    z_new["data"] = z.astype("float32")
+    radar.add_field(ref_field, z_new, replace_existing=True)
+
+    # same_nyquist=False uses every sweep, so the 40 dBZ layer is included
+    compz_all = pyart.retrieve.composite_reflectivity(
+        radar, field=ref_field, same_nyquist=False
+    )
+    assert_equal(compz_all.fields["composite_reflectivity"]["data"].max(), 40)
+
+    # same_nyquist=True drops the mismatched sweep, so the layer is excluded
+    compz_matched = pyart.retrieve.composite_reflectivity(
+        radar, field=ref_field, same_nyquist=True
+    )
+    assert_equal(compz_matched.fields["composite_reflectivity"]["data"].max(), 0)