How to use the metpy.calc.coriolis_parameter function in MetPy

def test_coriolis_warning():
    """Test that warning is raise when latitude larger than pi radians."""
    with pytest.warns(UserWarning):
        coriolis_parameter(50)
    with pytest.warns(UserWarning):
        coriolis_parameter(-50)

def test_coriolis_force():
    """Test basic coriolis force calculation."""
    lat = np.array([-90., -30., 0., 30., 90.]) * units.degrees
    cor = coriolis_parameter(lat)
    values = np.array([-1.4584232E-4, -.72921159E-4, 0, .72921159E-4,
                       1.4584232E-4]) * units('s^-1')
    assert_almost_equal(cor, values, 7)

def test_coriolis_warning():
    """Test that warning is raise when latitude larger than pi radians."""
    with pytest.warns(UserWarning):
        coriolis_parameter(50)
    with pytest.warns(UserWarning):
        coriolis_parameter(-50)

def test_intertial_advective_wind_4d(data_4d):
    """Test intertial_advective_wind on a 4D (time, pressure, y, x) grid."""
    f = coriolis_parameter(data_4d.latitude)
    u_g, v_g = geostrophic_wind(data_4d.height, f, data_4d.dx, data_4d.dy)
    u_i, v_i = inertial_advective_wind(u_g, v_g, u_g, v_g, data_4d.dx, data_4d.dy,
                                       data_4d.latitude)
    u_i_truth = np.array([[[[-4.74579332, -6.36486064, -7.20354171, -11.08307751],
                            [-1.88515129, -4.33855679, -6.82871465, -9.38096911],
                            [2.308649, -6.93391208, -14.06293133, -20.60786775],
                            [-0.92388354, -13.76737076, -17.9039117, -23.71419254]],
                           [[-2.60558413, -3.48755492, -3.62050089, -4.18871134],
                            [-3.36965812, -2.57689219, -2.66529828, -3.34582207],
                            [-0.56309499, -2.3322732, -4.37379768, -6.6663065],
                            [1.70092943, -3.59623514, -5.94640587, -7.50380432]],
                           [[-1.60508844, -2.30572073, -2.39044749, -2.59511279],
                            [-2.18854472, -1.47967397, -1.57319604, -2.24386278],
                            [-1.10582176, -1.24627092, -2.02075175, -3.314856],
                            [-0.25911941, -1.62294229, -1.75103256, -1.21885814]]],
                          [[[-6.69345313, -6.73506869, -7.9082287, -12.43972804],

# Digging Trough
# Z = digging_300hPa_trough(parameter='hght')
# T = digging_300hPa_trough(parameter='temp')


######################################################################
# Set geographic parameters for analytic grid to then
#

lats = np.linspace(35, 50, 101)
lons = np.linspace(260, 290, 101)
lon, lat = np.meshgrid(lons, lats)

# Calculate Geostrophic Wind from Analytic Heights
f = mpcalc.coriolis_parameter(lat * units('degrees'))
dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)
ugeo, vgeo = mpcalc.geostrophic_wind(Z*units.meter, f, dx, dy, dim_order='yx')

# Get the wind direction for each point
wdir = mpcalc.wind_direction(ugeo, vgeo)

# Compute the Gradient Wind via an approximation
dydx = mpcalc.first_derivative(Z, delta=dx, axis=1)
d2ydx2 = mpcalc.first_derivative(dydx, delta=dx, axis=1)
R = ((1 + dydx.m**2)**(3. / 2.)) / d2ydx2.m

geo_mag = mpcalc.wind_speed(ugeo, vgeo)
grad_mag = geo_mag.m - (geo_mag.m**2) / (f.magnitude * R)

ugrad, vgrad = mpcalc.wind_components(grad_mag * units('m/s'), wdir)

#########################################################################
# Calculations
# ------------
#
# Most of the calculations in `metpy.calc` will accept DataArrays by converting them
# into their corresponding unit arrays. While this may often work without any issues, we must
# keep in mind that because the calculations are working with unit arrays and not DataArrays:
#
# - The calculations will return unit arrays rather than DataArrays
# - Broadcasting must be taken care of outside of the calculation, as it would only recognize
#   dimensions by order, not name
#
# As an example, we calculate geostropic wind at 500 hPa below:

lat, lon = xr.broadcast(y, x)
f = mpcalc.coriolis_parameter(lat)
dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat, initstring=data_crs.proj4_init)
heights = data['height'].metpy.loc[{'time': time[0], 'vertical': 500. * units.hPa}]
u_geo, v_geo = mpcalc.geostrophic_wind(heights, f, dx, dy)
print(u_geo)
print(v_geo)

#########################################################################
# Also, a limited number of calculations directly support xarray DataArrays or Datasets (they
# can accept *and* return xarray objects). Right now, this includes
#
# - Derivative functions
#     - ``first_derivative``
#     - ``second_derivative``
#     - ``gradient``
#     - ``laplacian``
# - Cross-section functions

def f(heights, **kwargs):
        c = metpy.calc.coriolis_parameter(lat * _ureg.degrees)
        if dim_order == "yx":
            dy, dx = kwargs['deltas']
        else:
            dx, dy = kwargs['deltas']

        return metpy.calc.geostrophic_wind(xr.DataArray(heights), c, dx, dy,
                                           dim_order=kwargs['dim_order'])

height = height_var[0, 0, :, :].squeeze()
u_wind = u_wind_var[0, 0, :, :].squeeze() * units('m/s')
v_wind = v_wind_var[0, 0, :, :].squeeze() * units('m/s')

# Convert number of hours since the reference time into an actual date
time = num2date(time_var[:].squeeze(), time_var.units)

# Combine 1D latitude and longitudes into a 2D grid of locations
lon_2d, lat_2d = np.meshgrid(lon, lat)

# Smooth height data
height = ndimage.gaussian_filter(height, sigma=1.5, order=0)

# Set up some constants based on our projection, including the Coriolis parameter and
# grid spacing, converting lon/lat spacing to Cartesian
f = mpcalc.coriolis_parameter(np.deg2rad(lat_2d)).to('1/s')
dx, dy = mpcalc.lat_lon_grid_deltas(lon_2d, lat_2d)

# In MetPy 0.5, geostrophic_wind() assumes the order of the dimensions is (X, Y),
# so we need to transpose from the input data, which are ordered lat (y), lon (x).
# Once we get the components,transpose again so they match our original data.
geo_wind_u, geo_wind_v = mpcalc.geostrophic_wind(height * units.m, f, dx, dy)
geo_wind_u = geo_wind_u
geo_wind_v = geo_wind_v

# Calculate ageostrophic wind components
ageo_wind_u = u_wind - geo_wind_u
ageo_wind_v = v_wind - geo_wind_v

# Create new figure
fig = plt.figure(figsize=(15, 10), facecolor='black')

lev_500 = np.where(ds.variables['isobaric'][:] == 500)[0][0]

hght_500 = ds.variables['Geopotential_height_isobaric'][0, lev_500, :, :]
hght_500 = ndimage.gaussian_filter(hght_500, sigma=3, order=0) * units.meter

uwnd_500 = units('m/s') * ds.variables['u-component_of_wind_isobaric'][0, lev_500, :, :]
vwnd_500 = units('m/s') * ds.variables['v-component_of_wind_isobaric'][0, lev_500, :, :]

#######################################
# Begin Data Calculations
# -----------------------

dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat)

f = mpcalc.coriolis_parameter(np.deg2rad(lat)).to(units('1/sec'))

avor = mpcalc.vorticity(uwnd_500, vwnd_500, dx, dy, dim_order='yx') + f

avor = ndimage.gaussian_filter(avor, sigma=3, order=0) * units('1/s')

vort_adv = mpcalc.advection(avor, [uwnd_500, vwnd_500], (dx, dy), dim_order='yx') * 1e9

#######################################
# Map Creation
# ------------

# Set up Coordinate System for Plot and Transforms
dproj = ds.variables['LambertConformal_Projection']
globe = ccrs.Globe(ellipse='sphere', semimajor_axis=dproj.earth_radius,
                   semiminor_axis=dproj.earth_radius)
datacrs = ccrs.LambertConformal(central_latitude=dproj.latitude_of_projection_origin,