ADD: addition of subset_radar function in radar_utils

pyart/util/radar_utils.py

@@ -104,6 +104,203 @@ def to_vpt(radar, single_scan=True):
     return
 
 
+def cut_radar(radar, field_names, rng_min=None, rng_max=None, ele_min=None,
+              ele_max=None, azi_min=None, azi_max=None):
+    """
+    Cuts the radar object into new dimensions
+    Parameters
+    ----------
+    radar : radar object
+        The radar object containing the data
+    field_names : str or None
+        The fields to keep in the new radar
+    rng_min, rng_max : float
+        The range limits [m]. If None the entire coverage of the radar is
+        going to be used
+    ele_min, ele_max, azi_min, azi_max : float or None
+        The limits of the grid [deg]. If None the limits will be the limits
+        of the radar volume
+    Returns
+    -------
+    radar : radar object
+        The radar object containing only the desired data
+    """
+    radar_aux = copy.deepcopy(radar)
+
+    if (rng_min is None and rng_max is None and ele_min is None and
+            ele_max is None and azi_min is None and azi_max is None):
+        return radar_aux
+
+    if rng_min is None:
+        rng_min = 0.
+    if rng_max is None:
+        rng_max = np.max(radar_aux.range['data'])
+
+    ind_rng = np.where(np.logical_and(
+        radar_aux.range['data'] >= rng_min,
+        radar_aux.range['data'] <= rng_max))[0]
+
+    if ind_rng.size == 0:
+        print('No range bins between '+str(rng_min)+' and '+str(rng_max)+' m')
+        return None
+
+    # Determine angle limits
+    if radar_aux.scan_type == 'ppi':
+        if ele_min is None:
+            ele_min = np.min(radar_aux.fixed_angle['data'])
+        if ele_max is None:
+            ele_max = np.max(radar_aux.fixed_angle['data'])
+        if azi_min is None:
+            azi_min = np.min(radar_aux.azimuth['data'])
+        if azi_max is None:
+            azi_max = np.max(radar_aux.azimuth['data'])
+    else:
+        if ele_min is None:
+            ele_min = np.min(radar_aux.elevation['data'])
+        if ele_max is None:
+            ele_max = np.max(radar_aux.elevation['data'])
+        if azi_min is None:
+            azi_min = np.min(radar_aux.fixed_angle['data'])
+        if azi_max is None:
+            azi_max = np.max(radar_aux.fixed_angle['data'])
+
+    if radar_aux.scan_type == 'ppi':
+        # Get radar elevation angles within limits
+        ele_vec = np.sort(radar_aux.fixed_angle['data'])
+        ele_vec = ele_vec[
+            np.logical_and(ele_vec >= ele_min, ele_vec <= ele_max)]
+        if ele_vec.size == 0:
+            print('No elevation angles between '+str(ele_min)+' and ' +
+                 str(ele_max))
+            return None
+
+        # get sweeps corresponding to the desired elevation angles
+        ind_sweeps = []
+        for ele in ele_vec:
+            ind_sweeps.append(
+                np.where(radar_aux.fixed_angle['data'] == ele)[0][0])
+        radar_aux = radar_aux.extract_sweeps(ind_sweeps)
+
+        # Get indices of rays within limits
+        if azi_min < azi_max:
+            ind_rays = np.where(np.logical_and(
+                radar_aux.azimuth['data'] >= azi_min,
+                radar_aux.azimuth['data'] <= azi_max))[0]
+        else:
+            ind_rays = np.where(np.logical_or(
+                radar_aux.azimuth['data'] >= azi_min,
+                radar_aux.azimuth['data'] <= azi_max))[0]
+
+    else:
+        # Get radar azimuth angles within limits
+        azi_vec = radar_aux.fixed_angle['data']
+        if azi_min < azi_max:
+            azi_vec = np.sort(azi_vec[
+                np.logical_and(azi_vec >= azi_min, azi_vec <= azi_max)])
+        else:
+            azi_vec = azi_vec[
+                np.logical_or(azi_vec >= azi_min, azi_vec <= azi_max)]
+        if azi_vec.size == 0:
+            print('No azimuth angles between '+str(azi_min)+' and ' +
+                 str(azi_max))
+            return None
+
+        # get sweeps corresponding to the desired azimuth angles
+        ind_sweeps = []
+        for azi in azi_vec:
+            ind_sweeps.append(
+                np.where(radar_aux.fixed_angle['data'] == azi)[0][0])
+        radar_aux = radar_aux.extract_sweeps(ind_sweeps)
+
+        # Get indices of rays within limits
+        ind_rays = np.where(np.logical_and(
+            radar_aux.elevation['data'] >= ele_min,
+            radar_aux.elevation['data'] <= ele_max))[0]
+
+    # get new sweep start index and stop index
+    sweep_start_inds = copy.deepcopy(radar_aux.sweep_start_ray_index['data'])
+    sweep_end_inds = copy.deepcopy(radar_aux.sweep_end_ray_index['data'])
+
+    nrays = 0
+    ind_rays_aux = []
+    for j in range(radar_aux.nsweeps):
+        # get ray indices for this sweep
+        ind_rays_sweep = ind_rays[np.logical_and(
+            ind_rays >= sweep_start_inds[j], ind_rays <= sweep_end_inds[j])]
+        # order rays
+        if radar_aux.scan_type == 'ppi':
+            ind = np.argsort(radar_aux.azimuth['data'][ind_rays_sweep])
+            ind_rays_sweep = ind_rays_sweep[ind]
+            # avoid large gaps in data
+            azimuths = radar_aux.azimuth['data'][ind_rays_sweep]
+            azi_steps = azimuths[1:]-azimuths[:-1]
+            ind_gap = np.where(azi_steps > 2*np.median(azi_steps))[0]
+            if ind_gap.size > 0:
+                ind_rays_sweep = np.append(
+                    ind_rays_sweep[ind_gap[0]+1:],
+                    ind_rays_sweep[:ind_gap[0]+1])
+            ind_rays_aux.extend(ind_rays_sweep)
+        else:
+            ind = np.argsort(radar_aux.elevation['data'][ind_rays_sweep])
+            ind_rays_aux.extend(ind_rays_sweep[ind])
+
+        rays_in_sweep = ind_rays_sweep.size
+        radar_aux.rays_per_sweep['data'][j] = rays_in_sweep
+        if j == 0:
+            radar_aux.sweep_start_ray_index['data'][j] = 0
+        else:
+            radar_aux.sweep_start_ray_index['data'][j] = int(
+                radar_aux.sweep_end_ray_index['data'][j-1]+1)
+        radar_aux.sweep_end_ray_index['data'][j] = (
+            radar_aux.sweep_start_ray_index['data'][j]+rays_in_sweep-1)
+        nrays += rays_in_sweep
+
+    ind_rays = np.array(ind_rays_aux)
+
+    # Update metadata
+    radar_aux.range['data'] = radar_aux.range['data'][ind_rng]
+    radar_aux.time['data'] = radar_aux.time['data'][ind_rays]
+    radar_aux.azimuth['data'] = radar_aux.azimuth['data'][ind_rays]
+    radar_aux.elevation['data'] = radar_aux.elevation['data'][ind_rays]
+    radar_aux.init_gate_x_y_z()
+    radar_aux.init_gate_longitude_latitude()
+    radar_aux.init_gate_altitude()
+    radar_aux.nrays = nrays
+    radar_aux.ngates = ind_rng.size
+
+    if radar_aux.instrument_parameters is not None:
+        if 'nyquist_velocity' in radar_aux.instrument_parameters:
+            radar_aux.instrument_parameters['nyquist_velocity']['data'] = (
+                radar_aux.instrument_parameters['nyquist_velocity']['data'][
+                    ind_rays])
+        if 'pulse_width' in radar_aux.instrument_parameters:
+            radar_aux.instrument_parameters['pulse_width']['data'] = (
+                radar_aux.instrument_parameters['pulse_width']['data'][
+                    ind_rays])
+        if 'number_of_pulses' in radar_aux.instrument_parameters:
+            radar_aux.instrument_parameters['number_of_pulses']['data'] = (
+                radar_aux.instrument_parameters['number_of_pulses']['data'][
+                    ind_rays])
+
+    # Get new fields
+    if field_names is None:
+        radar_aux.fields = dict()
+    else:
+        fields_aux = copy.deepcopy(radar_aux.fields)
+        radar_aux.fields = dict()
+        for field_name in field_names:
+            if field_name not in fields_aux:
+                print('Field '+field_name+' not available')
+                continue
+
+            fields_aux[field_name]['data'] = (
+                fields_aux[field_name]['data'][:, ind_rng])
+            fields_aux[field_name]['data'] = (
+                fields_aux[field_name]['data'][ind_rays, :])
+            radar_aux.add_field(field_name, fields_aux[field_name])
+
+    return radar_aux
+
 
 def join_radar(radar1, radar2):
     """

tests/util/test_radar_utils.py

@@ -54,7 +54,30 @@ def test_to_vpt():
     assert radar.elevation['data'][0] == 90.0
     assert len(radar.instrument_parameters['prt_mode']['data']) == 108
 
-
+def test_cut_radar():
+    radar = pyart.testing.make_empty_ppi_radar(10, 36, 3)
+    field = {'data': np.ones((36 * 3, 10))}
+    radar.add_field('f1', field)
+    radar.add_field('f2', field)
+    azi_min = 10
+    azi_max = 100
+    rng_min = 200
+    rng_max = 800
+    ele_min = 0.75
+    ele_max = 0.75
+    radarcut = pyart.util.radar_utils.cut_radar(radar, ['f1'], rng_min = 200,
+                                                rng_max = 800, ele_min = 0.75,
+                                                ele_max = 0.75, azi_min = 10,
+                                                azi_max = 100)
+    # assert correct domain
+    assert(radarcut.azimuth['data'].min() >= azi_min)
+    assert(radarcut.azimuth['data'].max() <= azi_max)
+    assert(radarcut.range['data'].min() >= rng_min)
+    assert(radarcut.range['data'].max() <= rng_max)
+    assert(radarcut.elevation['data'].min() >= ele_min)
+    assert(radarcut.elevation['data'].max() <= ele_max)
+
+
 # read in example file
 radar = pyart.io.read_nexrad_archive(pyart.testing.NEXRAD_ARCHIVE_MSG31_FILE)
 def test_image_mute_radar():