ADD: addition of subset_radar function in radar_utils (#1227)

* added cut_radar function in util * added cut_radar function in util * renamed cut_radar to subset_radar * MNT: Update ci.yml. * MNT: Updating environments. * MNT: Duplicate 'latest' * MNT: Revert back os runner check. * MNT: Add bash to each run Co-authored-by: zssherman <[email protected]>
ARM-DOE · Aug 12, 2022 · a623736 · a623736
1 parent 7e71ca6
commit a623736
Show file tree

Hide file tree

Showing 4 changed files with 235 additions and 12 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -28,7 +28,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.7", "3.8", "3.9", "3.10"]
-        os: [macOS, ubuntu, Windows]
+        os: [macos, ubuntu, windows]
 
     steps:
       - uses: actions/checkout@v2
@@ -49,42 +49,44 @@ jobs:
         shell: bash -l {0}
         run: |
           echo "$RUNNER_OS"
-          echo "$RUNNER_OS" == 'Windows'
           if [ "$RUNNER_OS" == 'Windows' ]; then
                echo 'Trmm will not be installed for Windows'
           elif [ "$RUNNER_OS" == 'Linux' ]; then
                mamba install -c conda-forge trmm_rsl
-               export RSL_PATH=$CONDA/pkgs/
-          elif [ "$RUNNER_OS" == 'macOS' ]
-          then
+               RSL_PATH=$(find "$CONDA_PKGS_DIR" -name '*trmm_rsl-1.49-*' ! -name '*.tar.bz2')
+               echo "RSL_PATH=$(echo $RSL_PATH)" >> $GITHUB_ENV
+          elif [ "$RUNNER_OS" == 'macOS' ]; then
                mamba install -c conda-forge trmm_rsl
-               export RSL_PATH=$CONDA/pkgs/
+               RSL_PATH=$(find "$CONDA_PKGS_DIR" -name '*trmm_rsl-1.49-*' ! -name '*.tar.bz2')
+               echo "RSL_PATH=$(echo $RSL_PATH)" >> $GITHUB_ENV
           else
                mamba install -c conda-forge trmm_rsl
-               export RSL_PATH=$CONDA/pkgs/
+               RSL_PATH=$(find "$CONDA_PKGS_DIR" -name '*trmm_rsl-1.49-*' ! -name '*.tar.bz2')
+               echo "RSL_PATH=$(echo $RSL_PATH)" >> $GITHUB_ENV
           fi
           conda list
-          echo "$RSL_PATH"
           
       - name: Fetch all history for all tags and branches
         run: |
            git fetch --prune --unshallow
 
       - name: Install PyART
-        env:
-          RSL_PATH: ${{ env.CONDA_PKGS_DIR }}/
+        shell: bash -l {0}
         run: |
+          echo $RSL_PATH
           python -m pip install -e . --no-deps --force-reinstall
 
       - name: Run Linting
+        shell: bash -l {0}
         run: |
           python -m flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
           python -m flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
       
       - name: Run Tests
         id: run_tests
+        shell: bash -l {0}
         env:
-          RSL_PATH: ${{ env.CONDA_PKGS_DIR }}/
+          RSL_PATH: $RSL_PATH
         run: |
           python -m pytest -v --cov=./ --cov-report=xml
 

diff --git a/pyart/util/__init__.py b/pyart/util/__init__.py
@@ -16,6 +16,7 @@
 from .xsect import cross_section_ppi, cross_section_rhi
 from .hildebrand_sekhon import estimate_noise_hs74
 from .radar_utils import is_vpt, to_vpt, join_radar, image_mute_radar
+from .radar_utils import subset_radar
 from .simulated_vel import simulated_vel_from_profile
 from .sigmath import texture_along_ray, rolling_window
 from .sigmath import texture, angular_texture_2d

diff --git a/pyart/util/radar_utils.py b/pyart/util/radar_utils.py
@@ -104,6 +104,203 @@ def to_vpt(radar, single_scan=True):
     return
 
 
+def subset_radar(radar, field_names, rng_min=None, rng_max=None, ele_min=None,
+              ele_max=None, azi_min=None, azi_max=None):
+    """
+    Creates a subset of the radar object along 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):
     """

diff --git a/tests/util/test_radar_utils.py b/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_subset_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.subset_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 and correct fields
+    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)
+    assert(list(radarcut.fields) == ['f1'])
+
 # read in example file
 radar = pyart.io.read_nexrad_archive(pyart.testing.NEXRAD_ARCHIVE_MSG31_FILE)
 def test_image_mute_radar():