xradar icon indicating copy to clipboard operation
xradar copied to clipboard

ADD: Apply accessor

Open syedhamidali opened this issue 6 months ago • 12 comments

  • [x] Closes #202
  • [x] Tests added
  • [x] Changes are documented in history.md

This PR introduces an apply method to the xradar accessor, allowing users to apply functions to DataTree, Dataset, and DataArray objects seamlessly. This enhancement follows the discussion on generalizing the approach using accessors and aims to provide a more consistent and intuitive API.

Applying to DataTree

import xradar as xd
from open_radar_data import DATASETS

# Fetch the sample radar file
filename = DATASETS.fetch("sample_sgp_data.nc")

# Open the radar file into a DataTree object
dtree = xd.io.open_cfradial1_datatree(filename)

# Define a function to calculate rain rate from reflectivity
def calculate_rain_rate(ds, ref_field='DBZH'):
    def _rain_rate(dbz, a=200.0, b=1.6):
        Z = 10.0 ** (dbz / 10.0)
        return (Z / a) ** (1.0 / b)

    ds['RAIN_RATE'] = _rain_rate(ds[ref_field])
    ds['RAIN_RATE'].attrs = {'units': 'mm/h', 'long_name': 'Rain Rate'}
    return ds

# Apply the function across all sweeps
dtree = dtree.xradar.apply(calculate_rain_rate, ref_field='corrected_reflectivity_horizontal')

# Inspect the first sweep to see the added RAIN_RATE field
print(dtree['sweep_0'].data_vars)
Data variables:
    corrected_reflectivity_horizontal  (azimuth, range) float32 1MB -5.617 .....
    reflectivity_horizontal            (azimuth, range) float32 1MB ...
    recalculated_diff_phase            (azimuth, range) float32 1MB ...
    specific_attenuation               (azimuth, range) float32 1MB ...
    unf_dp_phase_shift                 (azimuth, range) float32 1MB ...
    mean_doppler_velocity              (azimuth, range) float32 1MB ...
    diff_phase                         (azimuth, range) float32 1MB ...
    rain_rate_A                        (azimuth, range) float32 1MB ...
    norm_coherent_power                (azimuth, range) float32 1MB ...
    dp_phase_shift                     (azimuth, range) float32 1MB ...
    diff_reflectivity                  (azimuth, range) float32 1MB ...
    proc_dp_phase_shift                (azimuth, range) float32 1MB ...
    copol_coeff                        (azimuth, range) float32 1MB ...
    sweep_number                       int32 4B ...
    sweep_fixed_angle                  float64 8B ...
    sweep_mode                         <U20 80B 'azimuth_surveillance'
    RAIN_RATE                          (azimuth, range) float32 1MB 0.01625 ....

Applying to Dataset

# Extract a single sweep as a Dataset
ds = dtree['sweep_0'].to_dataset()

# Apply the function to the Dataset using the xradar accessor
ds = ds.xradar.apply(calculate_rain_rate, ref_field='corrected_reflectivity_horizontal')

# Inspect the Dataset to see the new RAIN_RATE field
print(ds['RAIN_RATE'])
<xarray.DataArray 'RAIN_RATE' (azimuth: 400, range: 667)> Size: 1MB
array([[0.01624733, 0.04791909, 0.00855976, ..., 0.02479567, 0.02921897,
               nan],
       [0.01765691, 0.05320087, 0.00698366, ..., 0.01000767,        nan,
               nan],
       [0.01667295, 0.05207646, 0.00778836, ...,        nan,        nan,
               nan],
       ...,
       [0.01503456, 0.04075643, 0.0061091 , ..., 0.01069403,        nan,
        0.01479976],
       [0.0104445 , 0.06000188, 0.00770995, ..., 0.00208536,        nan,
               nan],
       [0.01611996, 0.05063309, 0.01133788, ...,        nan, 0.00524289,
               nan]], dtype=float32)
Coordinates:
    time       (azimuth) datetime64[ns] 3kB 2011-05-20T06:42:11.039436300 ......
  * range      (range) float64 5kB 0.0 60.0 120.0 ... 3.99e+04 3.996e+04
  * azimuth    (azimuth) float64 3kB 0.8281 1.719 2.594 ... 358.1 359.0 360.0
    elevation  (azimuth) float64 3kB ...
    latitude   float64 8B ...
    longitude  float64 8B ...
    altitude   float64 8B ...
Attributes:
    units:      mm/h
    long_name:  Rain Rate

Applying to DataArray

# Extract the reflectivity field as a DataArray
da = ds['corrected_reflectivity_horizontal']

# Define a function to add a constant value to the DataArray
def add_constant(da, constant=10):
    return da + constant

# Apply the function using the xradar accessor
da_modified = da.xradar.apply(add_constant, constant=5)

# Inspect the modified DataArray
print(da_modified)
<xarray.DataArray 'corrected_reflectivity_horizontal' (azimuth: 400, range: 667)> Size: 1MB
array([[ -0.6171875,   6.8984375,  -5.0703125, ...,   2.3203125,
          3.4609375,         nan],
       [ -0.0390625,   7.625    ,  -6.484375 , ...,  -3.984375 ,
                nan,         nan],
       [ -0.4375   ,   7.4765625,  -5.7265625, ...,         nan,
                nan,         nan],
       ...,
       [ -1.15625  ,   5.7734375,  -7.4140625, ...,  -3.5234375,
                nan,  -1.265625 ],
       [ -3.6875   ,   8.4609375,  -5.796875 , ..., -14.8828125,
                nan,         nan],
       [ -0.671875 ,   7.28125  ,  -3.1171875, ...,         nan,
         -8.4765625,         nan]], dtype=float32)
Coordinates:
    time       (azimuth) datetime64[ns] 3kB 2011-05-20T06:42:11.039436300 ......
  * range      (range) float64 5kB 0.0 60.0 120.0 ... 3.99e+04 3.996e+04
  * azimuth    (azimuth) float64 3kB 0.8281 1.719 2.594 ... 358.1 359.0 360.0
    elevation  (azimuth) float64 3kB ...
    latitude   float64 8B ...
    longitude  float64 8B ...
    altitude   float64 8B ...

syedhamidali avatar Aug 31 '24 08:08 syedhamidali