Merge pull request #90 from rcjackson/yz_quiver_fix

FIX: Yz quiver plot bug and weight calculation bug from last radar

doc/environment_docs.yml

@@ -6,6 +6,7 @@ dependencies:
   - python
   - pip
   - arm_pyart
+  - cython
   - pytest
   - pytest-mpl
   - pytest-cov

examples/hurricane_florence.py

@@ -42,10 +42,11 @@
 
 grid_mhx = pydda.constraints.add_hrrr_constraint_to_grid(grid_mhx,
                                                          H.grib)
-u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(
+grid_mhx = pydda.initialization.make_constant_wind_field(
     grid_mhx, (0.0, 0.0, 0.0))
-out_grids = pydda.retrieval.get_dd_wind_field(
-    [grid_mhx, grid_ltx], u_init, v_init, w_init, Co=1.0, Cm=1.0, Cmod=1.0,
+out_grids, _ = pydda.retrieval.get_dd_wind_field(
+    [grid_mhx, grid_ltx], Co=1e-3, Cm=1.0, Cmod=1e-3, Cx=1, Cy=1, Cz=1,
+    max_iterations=100, 
     mask_outside_opt=True, vel_name='corrected_velocity', engine="tensorflow",
     model_fields=["hrrr"])
 

examples/plot_examples.py

@@ -20,17 +20,17 @@
 
 
 # Load sounding data and insert as an intialization
-u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(
+cpol_grid = pydda.initialization.make_constant_wind_field(
         cpol_grid, (0.0, 0.0, 0.0), vel_field='corrected_velocity')
 
 # Start the wind retrieval. This example only uses the mass continuity
 # and data weighting constraints.
-Grids = pydda.retrieval.get_dd_wind_field([berr_grid, cpol_grid], u_init,
-                                          v_init, w_init, Co=1.0, Cm=256.0,
-                                          Cx=0.0, Cy=0., Cz=0.0, Cb=0., 
-                                          frz=5000.0, filter_window=5,
-                                          mask_outside_opt=True, upper_bc=1,
-                                          wind_tol=0.5, engine="tensorflow")
+Grids, _ = pydda.retrieval.get_dd_wind_field([berr_grid, cpol_grid], 
+                                             Co=1.0, Cm=256.0,
+                                             Cx=0.0, Cy=0., Cz=0.0, Cb=0., 
+                                             frz=5000.0, filter_window=5,
+                                             mask_outside_opt=True, upper_bc=1,
+                                             wind_tol=0.5, engine="tensorflow")
 # Plot a horizontal cross section
 plt.figure(figsize=(9, 9))
 pydda.vis.plot_horiz_xsection_barbs(Grids, background_field='reflectivity',

examples/plot_fun_with_constraints.py

@@ -31,7 +31,7 @@
 # Load our radar data
 sounding = pyart.io.read_arm_sonde(
     pydda.tests.SOUNDING_PATH)
-u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(
+berr_grid = pydda.initialization.make_constant_wind_field(
     berr_grid, (0.0, 0.0, 0.0))
 
 # Let's make a plot on a map
@@ -45,10 +45,10 @@
 u_back = sounding[1].u_wind
 v_back = sounding[1].v_wind
 z_back = sounding[1].height
-u_init, v_init, w_init = pydda.initialization.make_wind_field_from_profile(cpol_grid, sounding[1])
+cpol_grid = pydda.initialization.make_wind_field_from_profile(cpol_grid, sounding[1])
 
-new_grids = pydda.retrieval.get_dd_wind_field([cpol_grid, berr_grid],
-                                    u_init, v_init, w_init,
+new_grids, _ = pydda.retrieval.get_dd_wind_field([cpol_grid, berr_grid],
+
                                     u_back=u_back, v_back=v_back, z_back=z_back,
                                     Co=10.0, Cm=4096.0, frz=5000.0, Cb=1e-6,
                                     mask_outside_opt=False, wind_tol=0.2,
@@ -59,10 +59,9 @@
     new_grids,  bg_grid_no=-1, level=50, w_vel_contours=[1, 3, 5, 8])
 plt.show()
 # Let's see what happens when we use a zero initialization
-u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(
+cpol_grid = pydda.initialization.make_constant_wind_field(
     berr_grid, (0.0, 0.0, 0.0))    
-new_grids = pydda.retrieval.get_dd_wind_field([cpol_grid, berr_grid],
-                                    u_init, v_init, w_init,
+new_grids, _ = pydda.retrieval.get_dd_wind_field([cpol_grid, berr_grid],
                                     u_back=u_back, v_back=v_back, z_back=z_back,
                                     Co=1.0, Cm=128.0, frz=5000.0, Cb=1e-6,
                                     mask_outside_opt=False, wind_tol=0.2,
@@ -75,8 +74,7 @@
 plt.show()
 
 # Or, let's make the radar data more important!
-new_grids = pydda.retrieval.get_dd_wind_field([cpol_grid, berr_grid],
-                                    u_init, v_init, w_init,
+new_grids, _ = pydda.retrieval.get_dd_wind_field([cpol_grid, berr_grid],
                                     Co=100.0, Cm=128.0, frz=5000.0,
                                     mask_outside_opt=False, wind_tol=0.2,
                                     engine="tensorflow")

examples/plot_sydney_tornado.py

@@ -43,12 +43,14 @@
 grid4 = pyart.io.read_grid(grid4_path)
 
 # Set initialization and do retrieval
-u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(grid1, vel_field='VRADH_corr')
-new_grids = pydda.retrieval.get_dd_wind_field([grid1, grid2, grid3, grid4],
-                                              u_init, v_init, w_init, Co=1e-2, Cm=256.0, Cx=1e3, Cy=1e3, Cz=1e3,
-                                              vel_name='VRADH_corr', refl_field='DBZH', 
-                                              mask_outside_opt=True, wind_tol=0.1,
-                                              engine='tensorflow')
+grid1 = pydda.initialization.make_constant_wind_field(grid1, vel_field='VRADH_corr')
+new_grids, _ = pydda.retrieval.get_dd_wind_field([grid1, grid2, grid3, grid4],
+                                                 Co=1e-2, Cm=256.0, Cx=1e-4, Cy=1e-4,
+                                                 Cz=1e-4,
+                                                 vel_name='VRADH_corr', refl_field='DBZH', 
+                                                 mask_outside_opt=True, wind_tol=0.1,
+                                                 max_iterations=200,
+                                                 engine='jax')
 # Make a neat plot
 fig = plt.figure(figsize=(10,7))
 ax = pydda.vis.plot_horiz_xsection_quiver_map(new_grids, background_field='DBZH', level=3,

pydda/__init__.py

@@ -11,7 +11,7 @@
 from . import tests
 from . import constraints
 
-__version__ = '1.2.0'
+__version__ = '1.3.0'
 
 print("Welcome to PyDDA %s" % __version__)
 print("Detecting Jax...")
@@ -24,7 +24,7 @@
 print("Detecting TensorFlow...")
 try:
     import tensorflow
-    print("TensorFlow detected.")
+    print("TensorFlow detected. Checking for tensorflow-probability...")
     import tensorflow_probability
     print("TensorFlow-probability detected. TensorFlow engine enabled!")
 except (ImportError, AttributeError) as e:

pydda/cost_functions/_cost_functions_numpy.py

@@ -132,7 +132,7 @@ def calculate_grad_radial_vel(vrs, els, azs, u, v, w,
     return y.flatten()
 
 
-def calculate_smoothness_cost(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5):
+def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5):
     """
     Calculates the smoothness cost function by taking the Laplacian of the
     wind field.
@@ -163,10 +163,13 @@ def calculate_smoothness_cost(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5):
     scipy.ndimage.laplace(u, du, mode='wrap')
     scipy.ndimage.laplace(v, dv, mode='wrap')
     scipy.ndimage.laplace(w, dw, mode='wrap')
+    du = du / dx
+    dv = dv / dy
+    dw = dw / dz 
     return np.sum(Cx * du ** 2 + Cy * dv ** 2 + Cz * dw ** 2)
 
 
-def calculate_smoothness_gradient(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5,
+def calculate_smoothness_gradient(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5,
                                   upper_bc=True):
     """
     Calculates the gradient of the smoothness cost function
@@ -201,9 +204,15 @@ def calculate_smoothness_gradient(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5,
     scipy.ndimage.laplace(u, du, mode='wrap')
     scipy.ndimage.laplace(v, dv, mode='wrap')
     scipy.ndimage.laplace(w, dw, mode='wrap')
+    du = du / dx
+    dv = dv / dy
+    dz = dv / dz
     scipy.ndimage.laplace(du, grad_u, mode='wrap')
     scipy.ndimage.laplace(dv, grad_v, mode='wrap')
     scipy.ndimage.laplace(dw, grad_w, mode='wrap')
+    grad_u = grad_u / dx
+    grad_v = grad_v / dy
+    grad_w = grad_w / dz
 
     # Impermeability condition
     grad_w[0, :, :] = 0

pydda/cost_functions/_cost_functions_tensorflow.py

@@ -141,7 +141,7 @@ def calculate_grad_radial_vel(vrs, els, azs, u, v, w,
     return tf.reshape(y, (3 * np.prod(u.shape),))
 
 
-def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5):
+def calculate_smoothness_cost(u, v, w, dx=1, dy=1, dz=1, Cx=1e-5, Cy=1e-5, Cz=1e-5):
     """
     Calculates the smoothness cost function by taking the Laplacian of the
     wind field.
@@ -188,7 +188,7 @@ def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5):
     return tf.math.reduce_sum(x_term + y_term + z_term)
 
 
-def calculate_smoothness_gradient(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5,
+def calculate_smoothness_gradient(u, v, w, dx=1, dy=1, dz=1, Cx=1e-5, Cy=1e-5, Cz=1e-5,
                                   upper_bc=True, lower_bc=True):
     """
     Calculates the gradient of the smoothness cost function

pydda/cost_functions/cost_functions.py

@@ -125,7 +125,7 @@ def J_function(winds, parameters):
 
         if (parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
             Jsmooth = _cost_functions_numpy.calculate_smoothness_cost(
-                winds[0], winds[1], winds[2], 
+                winds[0], winds[1], winds[2], parameters.dx, parameters.dy, parameters.dz,
                 Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz)
         else:
             Jsmooth = 0
@@ -184,8 +184,8 @@ def J_function(winds, parameters):
 
         if (parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
             Jsmooth = _cost_functions_jax.calculate_smoothness_cost(
-                winds[0], winds[1], winds[2], 
-                Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz)
+                winds[0], winds[1], winds[2], parameters.dx, parameters.dy,
+                parameters.dz, Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz)
         else:
             Jsmooth = 0
 
@@ -324,7 +324,7 @@ def grad_J(winds, parameters):
 
         if (parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
             grad += _cost_functions_numpy.calculate_smoothness_gradient(
-                winds[0], winds[1], winds[2], 
+                winds[0], winds[1], winds[2],  parameters.dx, parameters.dy, parameters.dz,
                 Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc)
 
         if (parameters.Cb > 0):
@@ -366,7 +366,7 @@ def grad_J(winds, parameters):
 
         if (parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
             grad += _cost_functions_jax.calculate_smoothness_gradient(
-                winds[0], winds[1], winds[2], 
+                winds[0], winds[1], winds[2], parameters.dx, parameters.dy, parameters.dz,
                 Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc)
 
         if (parameters.Cb > 0):