Repository: openradar/PyDDA (language: python)

Issue:
Mass continuity cost function not minimizing all the way
I am running pydda on data from two radars and comparing its output to that of another program which I know works well. The output for u and v wind is great, but the magnitude of w is maybe 3 or 4 times too low, and the shape is inaccurate in some areas.

When I initialize w as a constant field of something other than 0, the minimization function does not get very far from that point (so if I initialize at 5, w stays fairly close to there for all points). Initializing u or v as something other than 0 does not produce this result. They still are able to go back to the correct values.

I am only using the radial velocity and mass continuity cost functions, and I have narrowed this down to be an issue with the cost function for mass continuity. That would explain why changing initialization of u or v does not affect the end state (those vectors are dominated by the radial velocity cost), but changing initialization of w does (those vectors are dominated by mass continuity).

I have further confirmed this by optimizing purely u and v with only the radial velocity cost - the u and v output is very accurate. I then held u and v constant and optimized only w with only mass continuity, and the w output was about the same as if I had run the program normally.

The l-lfgs-b function occasionally returns an “ABNORMAL_TERMINATION_IN_LNSRCH” warning, depending on various factors, such as the parameters passed to l-lfgs-b, the weights for the cost functions, and the wind field initialization. It is worth noting that the erroneous outputs I’m getting for w (with a field initialized at 0,0,0) are still erroneous whether the minimization converges or returns the aforementioned warning. As far as I recall, this warning was not output when the mass continuity weight was zero (not 100% sure on that though).

I will continue looking into this for a little while longer, but does anyone have thoughts as to what I could do to fix this?
little mistake?
https://github.com/openradar/PyDDA/blob/master/pydda/cost_functions/cost_functions.py

Line 804: dvdy = np.gradient(v, dy, axis=2)
-> dvdy = np.gradient(v, dy, axis=1)


Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index c9904759..c0c18584 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,10 +7,6 @@ env:
    
 matrix:
     include:
-    - python: 3.6
-      env:
-        - PYTHON_VERSION="3.6"
-        - DOC_BUILD="true"
     - python: 3.7
       sudo: yes
       dist: xenial
@@ -21,8 +17,13 @@ matrix:
       sudo: yes
       dist: xenial
       env:
-        - PYTHON_VERSION="3.8"      
-install: 
+        - PYTHON_VERSION="3.8"
+    - python: 3.9
+      sudo: yes
+      dist: xenial
+      env:
+        - PYTHON_VERSION="3.9"      
+install:
     - source continuous_integration/install.sh
 script: 
     - eval xvfb-run pytest
diff --git a/REQUIREMENTS.txt b/REQUIREMENTS.txt
index 77edcbce..a1c93a9a 100644
--- a/REQUIREMENTS.txt
+++ b/REQUIREMENTS.txt
@@ -6,3 +6,4 @@ wheels
 pandas
 dask
 six
+
diff --git a/continuous_integration/environment-3.5.yml b/continuous_integration/environment-3.5.yml
deleted file mode 100644
index aff2e5aa..00000000
--- a/continuous_integration/environment-3.5.yml
+++ /dev/null
@@ -1,18 +0,0 @@
-name: testenv
-channels:
-  - conda-forge
-  - defaults
-dependencies:
-  - python=3.5
-  - numpy
-  - scipy
-  - matplotlib
-  - netcdf4
-  - pytest
-  - pytest-mpl
-  - pytest-cov
-  - cartopy
-  - arm_pyart
-  - eccodes
-  - dask
-  - distributed
diff --git a/continuous_integration/environment-3.8.yml b/continuous_integration/environment-3.8.yml
index 38561564..965327a1 100644
--- a/continuous_integration/environment-3.8.yml
+++ b/continuous_integration/environment-3.8.yml
@@ -16,5 +16,8 @@ dependencies:
   - eccodes
   - dask
   - distributed
+  - jax
+  - tensorflow==2.6
+  - tensorflow-probability
   - pip:
      - arm_pyart
diff --git a/continuous_integration/environment-3.6.yml b/continuous_integration/environment-3.9.yml
similarity index 67%
rename from continuous_integration/environment-3.6.yml
rename to continuous_integration/environment-3.9.yml
index d828b56a..e1aa2331 100644
--- a/continuous_integration/environment-3.6.yml
+++ b/continuous_integration/environment-3.9.yml
@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python=3.6
+  - python=3.9
+  - pip
   - numpy
   - scipy
   - matplotlib
@@ -12,8 +13,12 @@ dependencies:
   - pytest-mpl
   - pytest-cov
   - cartopy
-  - arm_pyart
   - cfgrib
   - eccodes
   - dask
   - distributed
+  - jax
+  - tensorflow==2.6
+  - tensorflow-probability
+  - pip:
+     - arm_pyart
diff --git a/doc/source/index.rst b/doc/source/index.rst
index 97c69a35..a0b066af 100644
--- a/doc/source/index.rst
+++ b/doc/source/index.rst
@@ -33,7 +33,9 @@ of parallelism we recommend:
 ::
 
 While PyDDA will work on less than this, you may run into performance issues.
-In addition, we do not support Python versions less than 3.6. If you have an older version installed, PyDDA may work just fine but we will not provide support for any issues unless you are using at least Python 3.6.
+In addition, we do not support Python versions less than 3.6.
+If you have an older version installed, PyDDA may work just fine but we will
+not provide support for any issues unless you are using at least Python 3.6.
 
 =========================
 Installation instructions
@@ -93,7 +95,16 @@ just type in the following commands assuming you have the above dependencies ins
  cd PyDDA
  python setup.py install
 ::
- 
+
+Finally, PyDDA now supports using `Jax <jax.readthedocs.io>` and `TensorFlow <tensorflow.org>`
+for solving the three dimensional wind field. PyDDA requries TensorFlow 2.6 and the
+tensorflow-probability package for TensorFlow to be enabled. Both the Jax and
+TensorFlow-based engines now use automatic differentiation to solve for the gradients
+of each cost function. This therefore will create gradients that are less susceptible
+to boundary artifacts and rounding errors. In addition, both of these packages can
+utilize CUDA-enabled GPUs for much faster processing. These two
+dependencies are optional as the user can still use PyDDA with the SciPy ecosystem.
+
 Contents:
 
 .. toctree::
diff --git a/examples/hurricane_florence.py b/examples/hurricane_florence.py
index c6c89add..6b5473bd 100644
--- a/examples/hurricane_florence.py
+++ b/examples/hurricane_florence.py
@@ -25,8 +25,8 @@
             'hrrr.t06z.wrfprsf00.grib2')
 urllib.request.urlretrieve(hrrr_url, 'test.grib2')
 
-grid_mhx = pyart.io.read_grid(u'grid_mhx.nc')
-grid_ltx = pyart.io.read_grid(u'grid_ltx.nc')
+grid_mhx = pyart.io.read_grid(u'grid_mhx_5km.nc')
+grid_ltx = pyart.io.read_grid(u'grid_ltx_5km.nc')
 
 grid_mhx = pydda.constraints.add_hrrr_constraint_to_grid(grid_mhx,
                                                          'test.grib2')
diff --git a/examples/plot_examples.py b/examples/plot_examples.py
index b69328c9..0d2fa388 100644
--- a/examples/plot_examples.py
+++ b/examples/plot_examples.py
@@ -11,41 +11,41 @@
 
 import pyart
 import pydda
+import numpy as np
 from matplotlib import pyplot as plt
 
 
 berr_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
 cpol_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR1)
 
-sounding = pyart.io.read_arm_sonde(
-    pydda.tests.SOUNDING_PATH)
-
 
 # Load sounding data and insert as an intialization
-u_init, v_init, w_init = pydda.initialization.make_wind_field_from_profile(
-        cpol_grid, sounding[1], vel_field='corrected_velocity')
+u_init, v_init, w_init = 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=1500.0,
-                                          Cz=0, 
-                                          frz=5000.0, filt_iterations=2,
-                                          mask_outside_opt=True, upper_bc=1)
+                                          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")
 # Plot a horizontal cross section
 plt.figure(figsize=(9, 9))
 pydda.vis.plot_horiz_xsection_barbs(Grids, background_field='reflectivity',
                                     level=6,
-                                    w_vel_contours=[3, 6, 9, 12, 15],
+                                    w_vel_contours=[5, 10, 15],
                                     barb_spacing_x_km=5.0,
                                     barb_spacing_y_km=15.0)
 plt.show()
+plt.savefig("Darwin_horiz.png")
 
 # Plot a vertical X-Z cross section
 plt.figure(figsize=(9, 9))
 pydda.vis.plot_xz_xsection_barbs(Grids, background_field='reflectivity',
                                  level=40,
-                                 w_vel_contours=[3, 6, 9, 12, 15],
+                                 w_vel_contours=[5, 10, 15],
                                  barb_spacing_x_km=10.0,
                                  barb_spacing_z_km=2.0)
 plt.show()
@@ -54,7 +54,7 @@
 plt.figure(figsize=(9, 9))
 pydda.vis.plot_yz_xsection_barbs(Grids, background_field='reflectivity', 
                                  level=40,
-                                 w_vel_contours=[3, 6, 9, 12, 15],
+                                 w_vel_contours=[5, 10, 15],
                                  barb_spacing_y_km=10.0,
                                  barb_spacing_z_km=2.0)
-plt.show()
+plt.savefig('Darwin.png')
diff --git a/examples/plot_examples_nested.py b/examples/plot_examples_nested.py
deleted file mode 100644
index 5cfc8aa5..00000000
--- a/examples/plot_examples_nested.py
+++ /dev/null
@@ -1,70 +0,0 @@
-"""
-Example on retrieving and plotting winds on a distributed cluster
------------------------------------------------------------------
-
-This is a simple example for how to retrieve winds using the
-nested grid features of PyDDA.
-
-Author: Robert C. Jackson
-
-"""
-
-import pyart
-import pydda
-from matplotlib import pyplot as plt
-from distributed import LocalCluster, Client
-
-# Needed so that distributed doesn't run all of your code when the worker 
-# starts!
-if __name__ == '__main__':
-
-    berr_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
-    cpol_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR1)
-
-    sounding = pyart.io.read_arm_sonde(pydda.tests.SOUNDING_PATH)
-
-    # Load sounding data and insert as an intialization
-    u_init, v_init, w_init = pydda.initialization.make_wind_field_from_profile(
-        cpol_grid, sounding[1], vel_field='corrected_velocity')
-
-    # Start our dask distributed cluster. You can use any distributed cluster
-    # for this...a LocalCluster is used here for the sake of being able to run
-    # this example locally.
-    cluster = LocalCluster(n_workers=2)
-    print(cluster)
-    client = Client(cluster)
-    print(client)
-
-    # Start the wind retrieval. This example only uses the mass continuity
-    # and data weighting constraints.
-    Grids = pydda.retrieval.get_dd_wind_field_nested(
-        [berr_grid, cpol_grid], u_init,  v_init, w_init, client, Co=1.0,
-        Cm=1500.0, Cz=0, frz=5000.0,
-        filt_iterations=2, mask_outside_opt=True, upper_bc=1)
-
-    # Plot a horizontal cross section
-    plt.figure(figsize=(9, 9))
-    pydda.vis.plot_horiz_xsection_barbs(Grids, background_field='reflectivity', 
-                                        level=6,
-                                        w_vel_contours=[3, 6, 9, 12, 15],
-                                        barb_spacing_x_km=5.0,
-                                        barb_spacing_y_km=15.0)
-    plt.show()
-
-    # Plot a vertical X-Z cross section
-    plt.figure(figsize=(9, 9))
-    pydda.vis.plot_xz_xsection_barbs(Grids, background_field='reflectivity', 
-                                     level=40,
-                                     w_vel_contours=[3, 6, 9, 12, 15],
-                                     barb_spacing_x_km=10.0,
-                                     barb_spacing_z_km=2.0)
-    plt.show()
-
-    # Plot a vertical Y-Z cross section
-    plt.figure(figsize=(9, 9))
-    pydda.vis.plot_yz_xsection_barbs(Grids, background_field='reflectivity',
-                                     level=40,
-                                     w_vel_contours=[3, 6, 9, 12, 15],
-                                     barb_spacing_y_km=10.0,
-                                     barb_spacing_z_km=2.0)
-    plt.show()
diff --git a/pydda/__init__.py b/pydda/__init__.py
index ec2fac74..7527f672 100644
--- a/pydda/__init__.py
+++ b/pydda/__init__.py
@@ -11,4 +11,23 @@
 from . import tests
 from . import constraints
 
-__version__ = '0.5.2'
+__version__ = '1.0.0'
+
+print("Welcome to PyDDA 1.0")
+print("Detecting Jax...")
+try:
+    import jax
+    print("Jax engine enabled!")
+except ImportError:
+    print("Jax is not installed on your system, unable to use Jax engine.")
+
+print("Detecting TensorFlow...")
+try:
+    import tensorflow
+    print("TensorFlow detected.")
+    import tensorflow_probability
+    print("TensorFlow-probability detected. TensorFlow engine enabled!")
+except (ImportError, AttributeError):
+    print("Unable to load TensorFlow or tensorflow-probability. " +
+            "TensorFlow engine disabled.")
+
diff --git a/pydda/cost_functions/__init__.py b/pydda/cost_functions/__init__.py
index 50ed67f8..7284c55a 100644
--- a/pydda/cost_functions/__init__.py
+++ b/pydda/cost_functions/__init__.py
@@ -77,18 +77,18 @@
 """
 
 
-from .cost_functions import calculate_radial_vel_cost_function
-from .cost_functions import calculate_fall_speed
-from .cost_functions import calculate_grad_radial_vel
-from .cost_functions import calculate_mass_continuity
-from .cost_functions import calculate_mass_continuity_gradient
-from .cost_functions import calculate_smoothness_cost
-from .cost_functions import calculate_smoothness_gradient
-from .cost_functions import calculate_background_gradient
-from .cost_functions import calculate_background_cost
-from .cost_functions import calculate_vertical_vorticity_cost
-from .cost_functions import calculate_vertical_vorticity_gradient
-from .cost_functions import calculate_model_cost
-from .cost_functions import calculate_model_gradient
-from .cost_functions import calculate_point_cost, calculate_point_gradient
+from ._cost_functions_numpy import calculate_radial_vel_cost_function
+from ._cost_functions_numpy import calculate_fall_speed
+from ._cost_functions_numpy import calculate_grad_radial_vel
+from ._cost_functions_numpy import calculate_mass_continuity
+from ._cost_functions_numpy import calculate_mass_continuity_gradient
+from ._cost_functions_numpy import calculate_smoothness_cost
+from ._cost_functions_numpy import calculate_smoothness_gradient
+from ._cost_functions_numpy import calculate_background_gradient
+from ._cost_functions_numpy import calculate_background_cost
+from ._cost_functions_numpy import calculate_vertical_vorticity_cost
+from ._cost_functions_numpy import calculate_vertical_vorticity_gradient
+from ._cost_functions_numpy import calculate_model_cost
+from ._cost_functions_numpy import calculate_model_gradient
+from ._cost_functions_numpy import calculate_point_cost, calculate_point_gradient
 from .cost_functions import J_function, grad_J
diff --git a/pydda/cost_functions/_cost_functions_jax.py b/pydda/cost_functions/_cost_functions_jax.py
new file mode 100644
index 00000000..504277c6
--- /dev/null
+++ b/pydda/cost_functions/_cost_functions_jax.py
@@ -0,0 +1,695 @@
+import numpy as np
+try:
+    import jax
+    import jax.numpy as jnp
+
+    from jax import jit
+    from jax import float0
+    JAX_AVAILABLE = True
+except ImportError:
+    JAX_AVAILABLE = False
+
+def calculate_radial_vel_cost_function(vrs, azs, els, u, v,
+                                       w, wts, rmsVr, weights, coeff=1.0):
+    """
+    Calculates the cost function due to difference of the wind field from
+    radar radial velocities. For more information on this cost function, see
+    Potvin et al. (2012) and Shapiro et al. (2009).
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    azs: List of float arrays
+        List of azimuths from each radar
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    wts: List of float arrays
+        Float array containing fall speed from radar.
+    rmsVr: float
+        The sum of squares of velocity/num_points. Use for normalization
+        of data weighting coefficient
+    weights: n_radars x_bins x y_bins float array
+        Data weights for each pair of radars
+    coeff: float
+        Constant for cost function
+
+    Returns
+    -------
+    J_o: float
+         Observational cost function
+
+    References
+    -----------
+
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+
+    """
+    # Jax version of the cost function
+    J_o = 0
+    lambda_o = coeff / (rmsVr * rmsVr)
+    for i in range(len(vrs)):
+        v_ar = (jnp.cos(els[i]) * jnp.sin(azs[i]) * u +
+                jnp.cos(els[i]) * jnp.cos(azs[i]) * v +
+                jnp.sin(els[i]) * (w - jnp.abs(wts[i])))
+        the_weight = jnp.asarray(weights[i])
+        J_o += lambda_o * jnp.sum(jnp.square(vrs[i] - v_ar) * the_weight)
+    return J_o
+
+
+def calculate_grad_radial_vel(vrs, els, azs, u, v, w,
+                              wts, weights, rmsVr, coeff=1.0, upper_bc=True):
+    """
+    Calculates the gradient of the cost function due to difference of wind
+    field from radar radial velocities.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    azs: List of azimuths
+        List of azimuths from each radar
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    coeff: float
+        Constant for cost function
+    vel_name: str
+        Background velocity field name
+    weights: n_radars x_bins x y_bins float array
+        Data weights for each pair of radars
+
+    Returns
+    -------
+    y: 1-D float array
+         Gradient vector of observational cost function.
+
+    More information
+    ----------------
+    The gradient is calculated using Jax's vector Jacobian product.
+
+    # Use zero for all masked values since we don't want to add them into
+    # the cost function
+    """
+    primals, fun_vjp = jax.vjp(calculate_radial_vel_cost_function,
+            vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff)
+    _, _, _, p_x1, p_y1, p_z1, _, _, _, _ = fun_vjp(1.0)
+
+    # Impermeability condition
+    p_z1 = p_z1.at[0, :, :].set(0)
+    if (upper_bc is True):
+        p_z1 = p_z1.at[-1, :, :].set(0)
+    y = jnp.stack((p_x1, p_y1, p_z1), axis=0)
+    return np.copy(y.flatten())
+
+
+def calculate_smoothness_cost(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5):
+    """
+    Calculates the smoothness cost function by taking the Laplacian of the
+    wind field.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    Cx: float
+        Constant controlling smoothness in x-direction
+    Cy: float
+        Constant controlling smoothness in y-direction
+    Cz: float
+        Constant controlling smoothness in z-direction
+
+    Returns
+    -------
+    Js: float
+        value of smoothness cost function
+    """
+    dudx = jnp.gradient(u, dx, axis=2)
+    dudy = jnp.gradient(u, dy, axis=1)
+    dudz = jnp.gradient(u, dz, axis=0)
+    dvdx = jnp.gradient(v, dx, axis=2)
+    dvdy = jnp.gradient(v, dy, axis=1)
+    dvdz = jnp.gradient(v, dz, axis=0)
+    dwdx = jnp.gradient(w, dx, axis=2)
+    dwdy = jnp.gradient(w, dy, axis=1)
+    dwdz = jnp.gradient(w, dz, axis=0)
+
+    x_term = Cx * (
+            jnp.gradient(dudx, dx, axis=2) ** 2 + 
+            jnp.gradient(dvdx, dx, axis=1) ** 2 +
+            jnp.gradient(dwdx, dx, axis=2) ** 2)
+    y_term = Cy * (
+            jnp.gradient(dudy, dy, axis=2) ** 2 +
+            jnp.gradient(dvdy, dy, axis=1) ** 2 +
+            jnp.gradient(dwdy, dy, axis=2) ** 2)
+    z_term = Cz * (
+            jnp.gradient(dudz, dz, axis=2) ** 2 +
+            jnp.gradient(dvdz, dz, axis=1) ** 2 +
+            jnp.gradient(dwdz, dz, axis=2) ** 2)
+    return np.asanyarray(jnp.sum(x_term + y_term + z_term))
+
+
+def calculate_smoothness_gradient(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5,
+                                  upper_bc=True):
+    """
+    Calculates the gradient of the smoothness cost function
+    by taking the Laplacian of the Laplacian of the wind field.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    Cx: float
+        Constant controlling smoothness in x-direction
+    Cy: float
+        Constant controlling smoothness in y-direction
+    Cz: float
+        Constant controlling smoothness in z-direction
+
+    Returns
+    -------
+    y: float array
+        value of gradient of smoothness cost function
+    """
+    du = np.zeros(w.shape)
+    dv = np.zeros(w.shape)
+    dw = np.zeros(w.shape)
+    grad_u = np.zeros(w.shape)
+    grad_v = np.zeros(w.shape)
+    grad_w = np.zeros(w.shape)
+    scipy.ndimage.filters.laplace(u, du, mode='wrap')
+    scipy.ndimage.filters.laplace(v, dv, mode='wrap')
+    scipy.ndimage.filters.laplace(w, dw, mode='wrap')
+    scipy.ndimage.filters.laplace(du, grad_u, mode='wrap')
+    scipy.ndimage.filters.laplace(dv, grad_v, mode='wrap')
+    scipy.ndimage.filters.laplace(dw, grad_w, mode='wrap')
+
+    # Impermeability condition
+    grad_w[0, :, :] = 0
+    if (upper_bc is True):
+        grad_w[-1, :, :] = 0
+    y = np.stack([grad_u * Cx * 2, grad_v * Cy * 2, grad_w * Cz * 2], axis=0)
+    return y.flatten()
+
+
+def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
+    """
+    Calculates the cost function related to point observations. A mean square error cost
+    function term is applied to points that are within the sphere of influence
+    whose radius is determined by *roi*.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    x:  Float array
+        X coordinates of grid centers
+    y:  Float array
+        Y coordinates of grid centers
+    z:  Float array
+        Z coordinated of grid centers
+    point_list: list of dicts
+        List of point constraints.
+        Each member is a dict with keys of "u", "v", to correspond
+        to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation.
+
+        In addition, "site_id" gives the METAR code (or name) to the station.
+    Cp: float
+        The weighting coefficient of the point cost function.
+    roi: float
+        Radius of influence of observations
+
+    Returns
+    -------
+    J: float
+        The cost function related to the difference between 
+        wind field and points.
+    """
+    J = 0.0
+    for the_point in point_list:
+        the_box = jnp.where(np.logical_and.reduce(
+            (jnp.abs(x - the_point["x"]) < roi,
+                jnp.abs(y - the_point["y"]) < roi,
+             jnp.abs(z - the_point["z"]) < roi)))
+        J += jnp.sum(
+                ((u[the_box] - the_point["u"]) ** 2 + 
+                    (v[the_box] - the_point["v"]) ** 2))
+
+    return J * Cp
+
+
+def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
+    """
+    Calculates the gradient of the cost function related to point observations.
+    A mean square error cost function term is applied to points that are within the sphere of influence
+    whose radius is determined by *roi*.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    x: Float array
+        X coordinates of grid centers
+    y: Float array
+        Y coordinates of grid centers
+    z: Float array
+        Z coordinated of grid centers
+    point_list: list of dicts
+        List of point constraints. Each member is a dict with keys of "u", "v",
+        to correspond to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation.
+
+        In addition, "site_id" gives the METAR code (or name) to the station.
+    Cp: float
+        The weighting coefficient of the point cost function.
+    roi: float
+        Radius of influence of observations
+
+    Returns
+    -------
+    gradJ: float array
+        The gradient of the cost function related to the difference between wind field and points.
+
+    """
+
+    primals, fun_vjp = jax.vjp(
+            calculate_point_gradient, u, v, x, y, z, point_list, Cp, roi)
+    grad_u, grad_v, _, _, _, _, _, _, _ = fun_vjp(1.0)
+
+    gradJ_w = jnp.zeros_like(grad_u)
+    gradJ = jnp.stack([gradJ_u, gradJ_v, gradJ_w], axis=0)
+    return np.copy(gradJ.flatten()) * Cp
+
+
+def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1):
+    """
+    Calculates the mass continuity cost function by taking the divergence
+    of the wind field.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    dx: float
+        Grid spacing in x direction.
+    dy: float
+        Grid spacing in y direction.
+    dz: float
+        Grid spacing in z direction.
+    z: Float array (1D)
+        1D Float array with heights of grid
+    coeff: float
+        Constant controlling contribution of mass continuity to cost function
+    anel: int
+        = 1 use anelastic approximation, 0=don't
+
+    Returns
+    -------
+    J: float
+        value of mass continuity cost function
+    """
+    # Jax version of the cost function
+    dudx = jnp.gradient(u, dx, axis=2)
+    dvdy = jnp.gradient(v, dy, axis=1)
+    dwdz = jnp.gradient(w, dz, axis=0)
+
+    if (anel == 1):
+        rho = jnp.exp(-z / 10000.0)
+        drho_dz = jnp.gradient(rho, dz, axis=0)
+        anel_term = w / rho * drho_dz
+    else:
+        anel_term = jnp.zeros(w.shape)
+    return coeff * jnp.sum(
+            jnp.square(dudx + dvdy + dwdz + anel_term)) / 2.0
+
+
+def calculate_mass_continuity_gradient(u, v, w, z, dx,
+                                       dy, dz, coeff=1500.0, anel=1,
+                                       upper_bc=True):
+    """
+    Calculates the gradient of mass continuity cost function. This is done by
+    taking the negative gradient of the divergence of the wind field.
+
+    All grids must have the same grid specification.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    z: Float array (1D)
+        1D Float array with heights of grid
+    dx: float
+        Grid spacing in x direction.
+    dy: float
+        Grid spacing in y direction.
+    dz: float
+        Grid spacing in z direction.
+    coeff: float
+        Constant controlling contribution of mass continuity to cost function
+    anel: int
+        = 1 use anelastic approximation, 0=don't
+
+    Returns
+    -------
+    y: float array
+        value of gradient of mass continuity cost function
+    """
+    primals, fun_vjp = jax.vjp(
+            calculate_mass_continuity, u, v, w, z, dx, dy, dz, coeff, anel)
+    grad_u, grad_v, grad_w, _, _, _, _, _, _ = fun_vjp(1.0)
+
+    # Impermeability condition
+    grad_w = grad_w.at[0, :, :].set(0)
+    if (upper_bc is True):
+        grad_w = grad_w.at[-1, :, :].set(0)
+    y = jnp.stack([grad_u, grad_v, grad_w], axis=0)
+    return y.flatten().copy()
+
+
+def calculate_background_cost(u, v, w, weights, u_back, v_back, Cb=0.01):
+    """
+    Calculates the background cost function. The background cost function is
+    simply the sum of the squared differences between the wind field and the
+    background wind field multiplied by the weighting coefficient.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: Float array
+        Weights for each point to consider into cost function
+    u_back: 1D float array
+        Zonal winds vs height from sounding
+    w_back: 1D float array
+        Meridional winds vs height from sounding
+    Cb: float
+        Weight of background constraint to total cost function
+
+    Returns
+    -------
+    cost: float
+        value of background cost function
+    """
+    the_shape = u.shape
+    cost = 0
+    for i in range(the_shape[0]):
+        cost += (Cb * jnp.sum(jnp.square(u[i] - u_back[i]) * (weights[i]) +
+                              jnp.square(v[i] - v_back[i]) * (weights[i])))
+    return cost
+
+
+def calculate_background_gradient(u, v, w, weights, u_back, v_back, Cb=0.01):
+    """
+    Calculates the gradient of the background cost function. For each u, v
+    this is given as 2*coefficent*(analysis wind - background wind).
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: Float array
+        Weights for each point to consider into cost function
+    u_back: 1D float array
+        Zonal winds vs height from sounding
+    w_back: 1D float array
+        Meridional winds vs height from sounding
+    Cb: float
+        Weight of background constraint to total cost function
+
+    Returns
+    -------
+    y: float array
+        value of gradient of background cost function
+    """
+    primals, fun_vjp = jax.vjp(
+            calculate_background_cost, u, v, w, weights, u_back, v_back, Cb)
+    u_grad, v_grad, w_grad, _, _, _, _ = fun_vjp(1.0)
+    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    return y.flatten().copy()
+
+
+def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt,
+                                      coeff=1e-5):
+    """
+    Calculates the cost function due to deviance from vertical vorticity
+    equation. For more information of the vertical vorticity cost function,
+    see Potvin et al. (2012) and Shapiro et al. (2009).
+
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    dx: float array
+        Spacing in x grid
+    dy: float array
+        Spacing in y grid
+    dz: float array
+        Spacing in z grid
+    coeff: float
+        Weighting coefficient
+    Ut: float
+        U component of storm motion
+    Vt: float
+        V component of storm motion
+
+    Returns
+    -------
+    Jv: float
+        Value of vertical vorticity cost function.
+
+    References
+    ----------
+
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+
+    """
+    # Jax version of the cost function
+    dvdz = jnp.gradient(v, dz, axis=0)
+    dudz = jnp.gradient(u, dz, axis=0)
+    dwdz = jnp.gradient(w, dz, axis=0)
+    dvdx = jnp.gradient(v, dx, axis=2)
+    dwdy = jnp.gradient(w, dy, axis=1)
+    dwdx = jnp.gradient(w, dx, axis=2)
+    dudx = jnp.gradient(u, dx, axis=2)
+    dvdy = jnp.gradient(v, dy, axis=1)
+    dudy = jnp.gradient(u, dy, axis=1)
+    zeta = dvdx - dudy
+    dzeta_dx = jnp.gradient(zeta, dx, axis=2)
+    dzeta_dy = jnp.gradient(zeta, dy, axis=1)
+    dzeta_dz = jnp.gradient(zeta, dz, axis=0)
+    jv_array = ((u - Ut) * dzeta_dx + (v - Vt) * dzeta_dy +
+                w * dzeta_dz + (dvdz * dwdx - dudz * dwdy) +
+                zeta * (dudx + dvdy))
+    return jnp.sum(coeff * jv_array ** 2)
+
+
+def calculate_vertical_vorticity_gradient(u, v, w, dx, dy, dz, Ut, Vt,
+                                          coeff=1e-5):
+    """
+    Calculates the gradient of the cost function due to deviance from vertical
+    vorticity equation. This is done by taking the functional derivative of
+    the vertical vorticity cost function.
+
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    dx: float array
+        Spacing in x grid
+    dy: float array
+        Spacing in y grid
+    dz: float array
+        Spacing in z grid
+    Ut: float
+        U component of storm motion
+    Vt: float
+        V component of storm motion
+    coeff: float
+        Weighting coefficient
+
+    Returns
+    -------
+    Jv: 1D float array
+        Value of the gradient of the vertical vorticity cost function.
+
+    References
+    ----------
+
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+
+
+    """
+    # Jax version of the gradient cost function
+    primals, fun_vjp = jax.vjp(
+            calculate_vertical_vorticity_cost, u, v, w, dx, dy,
+            dz, Ut, Vt, coeff)
+    u_grad, v_grad, w_grad, _, _, _, _, _, _ = fun_vjp(1.0)
+    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    return y.flatten().copy()
+
+
+def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model,
+                         coeff=1.0):
+    """
+    Calculates the cost function for the model constraint.
+    This is calculated simply as the sum of squares of the differences
+    between the model wind field and the analysis wind field. Vertical
+    velocities are not factored into this cost function as there is typically
+    a high amount of uncertainty in model derived vertical velocities.
+
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    weights: list of 3D arrays
+        Float array showing how much each point from model weighs into
+        constraint.
+    u_model: list of 3D arrays
+        Float array with u component of wind field from model
+    v_model: list of 3D arrays
+        Float array with v component of wind field from model
+    w_model: list of 3D arrays
+        Float array with w component of wind field from model
+    coeff: float
+        Weighting coefficient
+
+    Returns
+    -------
+    Jv: float
+        Value of model cost function
+    """
+
+    cost = 0
+    for i in range(len(u_model)):
+        cost += (coeff * jnp.sum(jnp.square(u - u_model[i]) * weights[i] +
+                                 jnp.square(v - v_model[i]) * weights[i]))
+    return cost
+
+
+def calculate_model_gradient(u, v, w, weights, u_model,
+                             v_model, w_model, coeff=1.0):
+    """
+    Calculates the cost function for the model constraint.
+    This is calculated simply as twice the differences
+    between the model wind field and the analysis wind field for each u, v.
+    Vertical velocities are not factored into this cost function as there is
+    typically a high amount of uncertainty in model derived vertical
+    velocities. Therefore, the gradient for all of the w's will be 0.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: list of 3D float arrays
+        Weights for each point to consider into cost function
+    u_model: list of 3D float arrays
+        Zonal wind field from model
+    v_model: list of 3D float arrays
+        Meridional wind field from model
+    w_model: list of 3D float arrays
+        Vertical wind field from model
+    coeff: float
+        Weight of background constraint to total cost function
+
+    Returns
+    -------
+    y: float array
+        value of gradient of background cost function
+    """
+    primals, fun_vjp = jax.vjp(
+            calculate_model_cost, u, v, w, dx, dy, dz, Ut, Vt, coeff)
+    u_grad, v_grad, w_grad, _, _, _, _, _, _ = fun_vjp(1.0)
+    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    return y.flatten().copy()
diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py
new file mode 100644
index 00000000..c136de4d
--- /dev/null
+++ b/pydda/cost_functions/_cost_functions_numpy.py
@@ -0,0 +1,766 @@
+import numpy as np
+import scipy
+
+def calculate_radial_vel_cost_function(vrs, azs, els, u, v,
+                                       w, wts, rmsVr, weights, coeff=1.0):
+    """
+    Calculates the cost function due to difference of the wind field from
+    radar radial velocities. For more information on this cost function, see
+    Potvin et al. (2012) and Shapiro et al. (2009).
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+    Parameters
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    azs: List of float arrays
+        List of azimuths from each radar
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    wts: List of float arrays
+        Float array containing fall speed from radar.
+    rmsVr: float
+        The sum of squares of velocity/num_points. Use for normalization
+        of data weighting coefficient
+    weights: n_radars x_bins x y_bins float array
+        Data weights for each pair of radars
+    coeff: float
+        Constant for cost function
+    Returns
+    -------
+    J_o: float
+         Observational cost function
+    References
+    -----------
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+    """
+
+    J_o = 0
+    lambda_o = coeff / (rmsVr * rmsVr)
+    for i in range(len(vrs)):
+        v_ar = (np.cos(els[i]) * np.sin(azs[i]) * u +
+                np.cos(els[i]) * np.cos(azs[i]) * v +
+                np.sin(els[i]) * (w - np.abs(wts[i])))
+        the_weight = weights[i]
+        J_o += lambda_o * np.sum(np.square(vrs[i] - v_ar) * the_weight)
+
+    return J_o
+
+
+def calculate_grad_radial_vel(vrs, els, azs, u, v, w,
+                              wts, weights, rmsVr, coeff=1.0, upper_bc=True):
+    """
+    Calculates the gradient of the cost function due to difference of wind
+    field from radar radial velocities.
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+    Parameters
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    azs: List of azimuths
+        List of azimuths from each radar
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    coeff: float
+        Constant for cost function
+    vel_name: str
+        Background velocity field name
+    weights: n_radars x_bins x y_bins float array
+        Data weights for each pair of radars
+    Returns
+    -------
+    y: 1-D float array
+         Gradient vector of observational cost function.
+
+    More information
+    ----------------
+    The gradient is calculated by taking the functional derivative of the
+    cost function. For more information on functional derivatives, see the
+    Euler-Lagrange Equation:
+    https://en.wikipedia.org/wiki/Euler%E2%80%93Lagrange_equation
+    """
+
+    # Use zero for all masked values since we don't want to add them into
+    # the cost function
+
+    p_x1 = np.zeros(vrs[0].shape)
+    p_y1 = np.zeros(vrs[0].shape)
+    p_z1 = np.zeros(vrs[0].shape)
+    lambda_o = coeff / (rmsVr * rmsVr)
+
+    for i in range(len(vrs)):
+        v_ar = (np.cos(els[i]) * np.sin(azs[i]) * u +
+                np.cos(els[i]) * np.cos(azs[i]) * v +
+                np.sin(els[i]) * (w - np.abs(wts[i])))
+
+        x_grad = (2 * (v_ar - vrs[i]) * np.cos(els[i]) *
+                  np.sin(azs[i]) * weights[i]) * lambda_o
+        y_grad = (2 * (v_ar - vrs[i]) * np.cos(els[i]) *
+                  np.cos(azs[i]) * weights[i]) * lambda_o
+        z_grad = (2 * (v_ar - vrs[i]) * np.sin(els[i]) * weights[i]) * lambda_o
+
+        p_x1 += x_grad
+        p_y1 += y_grad
+        p_z1 += z_grad
+
+    # Impermeability condition
+    p_z1[0, :, :] = 0
+    if (upper_bc is True):
+        p_z1[-1, :, :] = 0
+    y = np.stack((p_x1, p_y1, p_z1), axis=0)
+    return y.flatten()
+
+
+def calculate_smoothness_cost(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5):
+    """
+    Calculates the smoothness cost function by taking the Laplacian of the
+    wind field.
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    Cx: float
+        Constant controlling smoothness in x-direction
+    Cy: float
+        Constant controlling smoothness in y-direction
+    Cz: float
+        Constant controlling smoothness in z-direction
+    Returns
+    -------
+    Js: float
+        value of smoothness cost function
+    """
+    du = np.zeros(w.shape)
+    dv = np.zeros(w.shape)
+    dw = np.zeros(w.shape)
+    scipy.ndimage.filters.laplace(u, du, mode='wrap')
+    scipy.ndimage.filters.laplace(v, dv, mode='wrap')
+    scipy.ndimage.filters.laplace(w, dw, mode='wrap')
+    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,
+                                  upper_bc=True):
+    """
+    Calculates the gradient of the smoothness cost function
+    by taking the Laplacian of the Laplacian of the wind field.
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    Cx: float
+        Constant controlling smoothness in x-direction
+    Cy: float
+        Constant controlling smoothness in y-direction
+    Cz: float
+        Constant controlling smoothness in z-direction
+    Returns
+    -------
+    y: float array
+        value of gradient of smoothness cost function
+    """
+    du = np.zeros(w.shape)
+    dv = np.zeros(w.shape)
+    dw = np.zeros(w.shape)
+    grad_u = np.zeros(w.shape)
+    grad_v = np.zeros(w.shape)
+    grad_w = np.zeros(w.shape)
+    scipy.ndimage.filters.laplace(u, du, mode='wrap')
+    scipy.ndimage.filters.laplace(v, dv, mode='wrap')
+    scipy.ndimage.filters.laplace(w, dw, mode='wrap')
+    scipy.ndimage.filters.laplace(du, grad_u, mode='wrap')
+    scipy.ndimage.filters.laplace(dv, grad_v, mode='wrap')
+    scipy.ndimage.filters.laplace(dw, grad_w, mode='wrap')
+
+    # Impermeability condition
+    grad_w[0, :, :] = 0
+    if (upper_bc is True):
+        grad_w[-1, :, :] = 0
+    y = np.stack([grad_u * Cx * 2, grad_v * Cy * 2, grad_w * Cz * 2], axis=0)
+    return y.flatten()
+
+
+def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
+    """
+    Calculates the cost function related to point observations. A mean square error cost
+    function term is applied to points that are within the sphere of influence
+    whose radius is determined by *roi*.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    x:  Float array
+        X coordinates of grid centers
+    y:  Float array
+        Y coordinates of grid centers
+    z:  Float array
+        Z coordinated of grid centers
+    point_list: list of dicts
+        List of point constraints.
+        Each member is a dict with keys of "u", "v", to correspond
+        to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation.
+        In addition, "site_id" gives the METAR code (or name) to the station.
+    Cp: float
+        The weighting coefficient of the point cost function.
+    roi: float
+        Radius of influence of observations
+    Returns
+    -------
+    J: float
+        The cost function related to the difference between wind field and points.
+    """
+    J = 0.0
+    for the_point in point_list:
+        # Instead of worrying about whole domain, just find points in radius of influence
+        # Since we know that the weight will be zero outside the sphere of influence anyways
+        the_box = np.where(np.logical_and.reduce(
+            (np.abs(x - the_point["x"]) < roi, np.abs(y - the_point["y"]) < roi,
+             np.abs(z - the_point["z"]) < roi)))
+        J += np.sum(((u[the_box] - the_point["u"]) ** 2 + (v[the_box] - the_point["v"]) ** 2))
+
+    return J * Cp
+
+
+def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
+    """
+    Calculates the gradient of the cost function related to point observations.
+    A mean square error cost function term is applied to points that are within the sphere of influence
+    whose radius is determined by *roi*.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    x: Float array
+        X coordinates of grid centers
+    y: Float array
+        Y coordinates of grid centers
+    z: Float array
+        Z coordinated of grid centers
+    point_list: list of dicts
+        List of point constraints. Each member is a dict with keys of "u", "v",
+        to correspond to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation.
+        In addition, "site_id" gives the METAR code (or name) to the station.
+    Cp: float
+        The weighting coefficient of the point cost function.
+    roi: float
+        Radius of influence of observations
+    Returns
+    -------
+    gradJ: float array
+        The gradient of the cost function related to the difference between wind field and points.
+    """
+
+    gradJ_u = np.zeros_like(u)
+    gradJ_v = np.zeros_like(v)
+    gradJ_w = np.zeros_like(u)
+
+    for the_point in point_list:
+        the_box = np.where(np.logical_and.reduce(
+            (np.abs(x - the_point["x"]) < roi, np.abs(y - the_point["y"]) < roi,
+             np.abs(z - the_point["z"]) < roi)))
+        gradJ_u[the_box] += 2 * (u[the_box] - the_point["u"])
+        gradJ_v[the_box] += 2 * (v[the_box] - the_point["v"])
+
+    gradJ = np.stack([gradJ_u, gradJ_v, gradJ_w], axis=0).flatten()
+    return gradJ * Cp
+
+
+def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1):
+    """
+    Calculates the mass continuity cost function by taking the divergence
+    of the wind field.
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    dx: float
+        Grid spacing in x direction.
+    dy: float
+        Grid spacing in y direction.
+    dz: float
+        Grid spacing in z direction.
+    z: Float array (1D)
+        1D Float array with heights of grid
+    coeff: float
+        Constant controlling contribution of mass continuity to cost function
+    anel: int
+        = 1 use anelastic approximation, 0=don't
+    Returns
+    -------
+    J: float
+        value of mass continuity cost function
+    """
+    dudx = np.gradient(u, dx, axis=2)
+    dvdy = np.gradient(v, dy, axis=1)
+    dwdz = np.gradient(w, dz, axis=0)
+
+    if (anel == 1):
+        rho = np.exp(-z / 10000.0)
+        drho_dz = np.gradient(rho, dz, axis=0)
+        anel_term = w / rho * drho_dz
+    else:
+        anel_term = np.zeros(w.shape)
+    div = dudx + dvdy + dwdz + anel_term
+
+    return coeff * np.sum(np.square(div)) / 2.0
+
+
+def calculate_mass_continuity_gradient(u, v, w, z, dx,
+                                       dy, dz, coeff=1500.0, anel=1,
+                                       upper_bc=True):
+    """
+    Calculates the gradient of mass continuity cost function. This is done by
+    taking the negative gradient of the divergence of the wind field.
+    All grids must have the same grid specification.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    z: Float array (1D)
+        1D Float array with heights of grid
+    dx: float
+        Grid spacing in x direction.
+    dy: float
+        Grid spacing in y direction.
+    dz: float
+        Grid spacing in z direction.
+    coeff: float
+        Constant controlling contribution of mass continuity to cost function
+    anel: int
+        = 1 use anelastic approximation, 0=don't
+    Returns
+    -------
+    y: float array
+        value of gradient of mass continuity cost function
+    """
+    dudx = np.gradient(u, dx, axis=2)
+    dvdy = np.gradient(v, dy, axis=1)
+    dwdz = np.gradient(w, dz, axis=0)
+    if (anel == 1):
+        rho = np.exp(-z / 10000.0)
+        drho_dz = np.gradient(rho, dz, axis=0)
+        anel_term = w / rho * drho_dz
+    else:
+        anel_term = 0
+
+    div = dudx + dvdy + dwdz + anel_term
+
+    grad_u = -np.gradient(div, dx, axis=2) * coeff
+    grad_v = -np.gradient(div, dy, axis=1) * coeff
+    grad_w = -np.gradient(div, dz, axis=0) * coeff
+
+    # Impermeability condition
+    grad_w[0, :, :] = 0
+    if(upper_bc is True):
+       grad_w[-1, :, :] = 0
+    y = np.stack([grad_u, grad_v, grad_w], axis=0)
+    return y.flatten()
+
+
+def calculate_fall_speed(grid, refl_field=None, frz=4500.0):
+    """
+    Estimates fall speed based on reflectivity.
+    Uses methodology of Mike Biggerstaff and Dan Betten
+    Parameters
+    ----------
+    Grid: Py-ART Grid
+        Py-ART Grid containing reflectivity to calculate fall speed from
+    refl_field: str
+        String containing name of reflectivity field. None will automatically
+        determine the name.
+    frz: float
+        Height of freezing level in m
+    Returns
+    -------
+    3D float array:
+        Float array of terminal velocities
+    """
+    # Parse names of velocity field
+    if refl_field is None:
+        refl_field = pyart.config.get_field_name('reflectivity')
+
+    refl = grid.fields[refl_field]['data']
+    grid_z = grid.point_z['data']
+    term_vel = np.zeros(refl.shape)
+    A = np.zeros(refl.shape)
+    B = np.zeros(refl.shape)
+    rho = np.exp(-grid_z / 10000.0)
+    A[np.logical_and(grid_z < frz, refl < 55)] = -2.6
+    B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107
+    A[np.logical_and(grid_z < frz,
+                     np.logical_and(refl >= 55, refl < 60))] = -2.5
+    B[np.logical_and(grid_z < frz,
+                     np.logical_and(refl >= 55, refl < 60))] = 0.013
+    A[np.logical_and(grid_z < frz, refl > 60)] = -3.95
+    B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148
+    A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817
+    B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063
+    A[np.logical_and(grid_z >= frz,
+                     np.logical_and(refl >= 33, refl < 49))] = -2.5
+    B[np.logical_and(grid_z >= frz,
+                     np.logical_and(refl >= 33, refl < 49))] = 0.013
+    A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95
+    B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148
+
+    fallspeed = A * np.power(10, refl * B) * np.power(1.2 / rho, 0.4)
+    del A, B, rho
+    return fallspeed
+
+
+def calculate_background_cost(u, v, w, weights, u_back, v_back, Cb=0.01):
+    """
+    Calculates the background cost function. The background cost function is
+    simply the sum of the squared differences between the wind field and the
+    background wind field multiplied by the weighting coefficient.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: Float array
+        Weights for each point to consider into cost function
+    u_back: 1D float array
+        Zonal winds vs height from sounding
+    w_back: 1D float array
+        Meridional winds vs height from sounding
+    Cb: float
+        Weight of background constraint to total cost function
+    Returns
+    -------
+    cost: float
+        value of background cost function
+    """
+    the_shape = u.shape
+    cost = 0
+    for i in range(the_shape[0]):
+        cost += (Cb * np.sum(np.square(u[i] - u_back[i]) * (weights[i]) +
+                             np.square(v[i] - v_back[i]) * (weights[i])))
+    return cost
+
+
+def calculate_background_gradient(u, v, w, weights, u_back, v_back, Cb=0.01):
+    """
+    Calculates the gradient of the background cost function. For each u, v
+    this is given as 2*coefficent*(analysis wind - background wind).
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: Float array
+        Weights for each point to consider into cost function
+    u_back: 1D float array
+        Zonal winds vs height from sounding
+    w_back: 1D float array
+        Meridional winds vs height from sounding
+    Cb: float
+        Weight of background constraint to total cost function
+    Returns
+    -------
+    y: float array
+        value of gradient of background cost function
+    """
+    the_shape = u.shape
+    u_grad = np.zeros(the_shape)
+    v_grad = np.zeros(the_shape)
+    w_grad = np.zeros(the_shape)
+
+    for i in range(the_shape[0]):
+        u_grad[i] = Cb * 2 * (u[i] - u_back[i]) * (weights[i])
+        v_grad[i] = Cb * 2 * (v[i] - v_back[i]) * (weights[i])
+
+    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    return y.flatten()
+
+
+def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt,
+                                      coeff=1e-5):
+    """
+    Calculates the cost function due to deviance from vertical vorticity
+    equation. For more information of the vertical vorticity cost function,
+    see Potvin et al. (2012) and Shapiro et al. (2009).
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    dx: float array
+        Spacing in x grid
+    dy: float array
+        Spacing in y grid
+    dz: float array
+        Spacing in z grid
+    coeff: float
+        Weighting coefficient
+    Ut: float
+        U component of storm motion
+    Vt: float
+        V component of storm motion
+    Returns
+    -------
+    Jv: float
+        Value of vertical vorticity cost function.
+    References
+    ----------
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+    """
+    dvdz = np.gradient(v, dz, axis=0)
+    dudz = np.gradient(u, dz, axis=0)
+    dvdx = np.gradient(v, dx, axis=2)
+    dwdy = np.gradient(w, dy, axis=1)
+    dwdx = np.gradient(w, dx, axis=2)
+    dudx = np.gradient(u, dx, axis=2)
+    dvdy = np.gradient(v, dy, axis=2)
+    dudy = np.gradient(u, dy, axis=1)
+    zeta = dvdx - dudy
+    dzeta_dx = np.gradient(zeta, dx, axis=2)
+    dzeta_dy = np.gradient(zeta, dy, axis=1)
+    dzeta_dz = np.gradient(zeta, dz, axis=0)
+    jv_array = ((u - Ut) * dzeta_dx + (v - Vt) * dzeta_dy +
+                w * dzeta_dz + (dvdz * dwdx - dudz * dwdy) +
+                zeta * (dudx + dvdy))
+    return np.sum(coeff * jv_array ** 2)
+
+
+def calculate_vertical_vorticity_gradient(u, v, w, dx, dy, dz, Ut, Vt,
+                                          coeff=1e-5):
+    """
+    Calculates the gradient of the cost function due to deviance from vertical
+    vorticity equation. This is done by taking the functional derivative of
+    the vertical vorticity cost function.
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    dx: float array
+        Spacing in x grid
+    dy: float array
+        Spacing in y grid
+    dz: float array
+        Spacing in z grid
+    Ut: float
+        U component of storm motion
+    Vt: float
+        V component of storm motion
+    coeff: float
+        Weighting coefficient
+    Returns
+    -------
+    Jv: 1D float array
+        Value of the gradient of the vertical vorticity cost function.
+    References
+    ----------
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+    """
+
+    # First derivatives
+    dvdz = np.gradient(v, dz, axis=0)
+    dwdy = np.gradient(w, dy, axis=1)
+    dudx = np.gradient(u, dx, axis=2)
+    dvdy = np.gradient(v, dy, axis=1)
+    dvdx = np.gradient(v, dx, axis=2)
+    dwdx = np.gradient(w, dx, axis=2)
+    dudz = np.gradient(u, dz, axis=0)
+    dudy = np.gradient(u, dy, axis=1)
+
+    zeta = dvdx - dudy
+    dzeta_dx = np.gradient(zeta, dx, axis=2)
+    dzeta_dy = np.gradient(zeta, dy, axis=1)
+    dzeta_dz = np.gradient(zeta, dz, axis=0)
+
+    # Second deriviatives
+    dwdydz = np.gradient(dwdy, dz, axis=0)
+    dwdxdz = np.gradient(dwdx, dz, axis=0)
+    dudzdy = np.gradient(dudz, dy, axis=1)
+    dvdxdy = np.gradient(dvdx, dy, axis=1)
+    dudx2 = np.gradient(dudx, dx, axis=2)
+    dudxdy = np.gradient(dudx, dy, axis=1)
+    dudxdz = np.gradient(dudx, dz, axis=0)
+    dudy2 = np.gradient(dudx, dy, axis=1)
+
+    dzeta_dt = ((u - Ut) * dzeta_dx + (v - Vt) * dzeta_dy + w * dzeta_dz +
+                (dvdz * dwdx - dudz * dwdy) + zeta * (dudx + dvdy))
+
+    # Now we intialize our gradient value
+    u_grad = np.zeros(u.shape)
+    v_grad = np.zeros(v.shape)
+    w_grad = np.zeros(w.shape)
+
+    # Vorticity Advection
+    u_grad += dzeta_dx + (Ut - u) * dudxdy + (Vt - v) * dudxdy
+    v_grad += dzeta_dy + (Vt - v) * dvdxdy + (Ut - u) * dvdxdy
+    w_grad += dzeta_dz
+
+    # Tilting term
+    u_grad += dwdydz
+    v_grad += dwdxdz
+    w_grad += dudzdy - dudxdz
+
+    # Stretching term
+    u_grad += -dudxdy + dudy2 - dzeta_dx
+    u_grad += -dudx2 + dudxdy - dzeta_dy
+
+    # Multiply by 2*dzeta_dt according to chain rule
+    u_grad = u_grad * 2 * dzeta_dt * coeff
+    v_grad = v_grad * 2 * dzeta_dt * coeff
+    w_grad = w_grad * 2 * dzeta_dt * coeff
+
+    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    return y.flatten()
+
+
+def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model,
+                         coeff=1.0):
+    """
+    Calculates the cost function for the model constraint.
+    This is calculated simply as the sum of squares of the differences
+    between the model wind field and the analysis wind field. Vertical
+    velocities are not factored into this cost function as there is typically
+    a high amount of uncertainty in model derived vertical velocities.
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    weights: list of 3D arrays
+        Float array showing how much each point from model weighs into
+        constraint.
+    u_model: list of 3D arrays
+        Float array with u component of wind field from model
+    v_model: list of 3D arrays
+        Float array with v component of wind field from model
+    w_model: list of 3D arrays
+        Float array with w component of wind field from model
+    coeff: float
+        Weighting coefficient
+    Returns
+    -------
+    Jv: float
+        Value of model cost function
+    """
+
+    cost = 0
+    for i in range(len(u_model)):
+        cost += (coeff * np.sum(np.square(u - u_model[i]) * weights[i] +
+                                np.square(v - v_model[i]) * weights[i]))
+    return cost
+
+
+def calculate_model_gradient(u, v, w, weights, u_model,
+                             v_model, w_model, coeff=1.0):
+    """
+    Calculates the cost function for the model constraint.
+    This is calculated simply as twice the differences
+    between the model wind field and the analysis wind field for each u, v.
+    Vertical velocities are not factored into this cost function as there is
+    typically a high amount of uncertainty in model derived vertical
+    velocities. Therefore, the gradient for all of the w's will be 0.
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: list of 3D float arrays
+        Weights for each point to consider into cost function
+    u_model: list of 3D float arrays
+        Zonal wind field from model
+    v_model: list of 3D float arrays
+        Meridional wind field from model
+    w_model: list of 3D float arrays
+        Vertical wind field from model
+    coeff: float
+        Weight of background constraint to total cost function
+    Returns
+    -------
+    y: float array
+        value of gradient of background cost function
+    """
+    the_shape = u.shape
+    u_grad = np.zeros(the_shape)
+    v_grad = np.zeros(the_shape)
+    w_grad = np.zeros(the_shape)
+    for i in range(len(u_model)):
+        u_grad += coeff * 2 * (u - u_model[i]) * weights[i]
+        v_grad += coeff * 2 * (v - v_model[i]) * weights[i]
+
+    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    return y.flatten()
+
diff --git a/pydda/cost_functions/_cost_functions_tensorflow.py b/pydda/cost_functions/_cost_functions_tensorflow.py
new file mode 100644
index 00000000..ff9bc777
--- /dev/null
+++ b/pydda/cost_functions/_cost_functions_tensorflow.py
@@ -0,0 +1,808 @@
+import numpy as np
+try:
+    import tensorflow as tf
+    TENSORFLOW_AVAILABLE = True
+except ImportError:
+    TENSORFLOW_AVAILABLE = False
+
+def calculate_radial_vel_cost_function(vrs, azs, els, u, v,
+                                       w, wts, rmsVr, weights, coeff=1.0):
+    """
+    Calculates the cost function due to difference of the wind field from
+    radar radial velocities. For more information on this cost function, see
+    Potvin et al. (2012) and Shapiro et al. (2009).
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    azs: List of float arrays
+        List of azimuths from each radar
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    wts: List of float arrays
+        Float array containing fall speed from radar.
+    rmsVr: float
+        The sum of squares of velocity/num_points. Use for normalization
+        of data weighting coefficient
+    weights: n_radars x_bins x y_bins float array
+        Data weights for each pair of radars
+    coeff: float
+        Constant for cost function
+
+    Returns
+    -------
+    J_o: float
+         Observational cost function
+
+    References
+    -----------
+
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+    """
+
+    J_o = 0.
+    lambda_o = coeff / (rmsVr * rmsVr)
+
+    for i in range(len(vrs)):
+        v_ar = (tf.math.cos(els[i]) * tf.math.sin(azs[i]) * u +
+                tf.math.cos(els[i]) * tf.math.cos(azs[i]) * v +
+                tf.math.sin(els[i]) * (w - tf.math.abs(wts[i])))
+        J_o += lambda_o * tf.reduce_sum(
+            tf.math.square(vrs[i] - v_ar) * weights[i])
+    return J_o
+
+
+def calculate_grad_radial_vel(vrs, els, azs, u, v, w,
+                              wts, weights, rmsVr, coeff=1.0, upper_bc=True):
+    """
+    Calculates the gradient of the cost function due to difference of wind
+    field from radar radial velocities.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    azs: List of azimuths
+        List of azimuths from each radar
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    coeff: float
+        Constant for cost function
+    vel_name: str
+        Background velocity field name
+    weights: n_radars x_bins x y_bins float array
+        Data weights for each pair of radars
+    upper_bc: bool
+        Set to true to impose w=0 at top of domain.
+
+    Returns
+    -------
+    y: 1-D float array
+         Gradient vector of observational cost function.
+
+    More information
+    ----------------
+    The gradient is calculated by taking the functional derivative of the
+    cost function. For more information on functional derivatives, see the
+    Euler-Lagrange Equation:
+
+    https://en.wikipedia.org/wiki/Euler%E2%80%93Lagrange_equation
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        tape.watch(w)
+        loss = calculate_radial_vel_cost_function(vrs, azs, els,
+                                                  u, v, w, wts, rmsVr, weights, coeff)
+    vars = {'u': u, 'v': v, 'w': w}
+    grad = tape.gradient(loss, vars)
+    p_x1 = grad['u']
+    p_y1 = grad['v']
+    p_z1 = grad['w']
+
+    # Impermeability condition
+    p_z1 = tf.concat(
+        [tf.zeros((1, u.shape[1], u.shape[2])), p_z1[1:, :, :]], axis=0)
+    if (upper_bc is True):
+        p_z1 = tf.concat(
+            [p_z1[:-1, :, :],
+             tf.zeros((1, u.shape[1], u.shape[2]))], axis=0)
+    y = tf.stack((p_x1, p_y1, p_z1), axis=0)
+    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):
+    """
+    Calculates the smoothness cost function by taking the Laplacian of the
+    wind field.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    Cx: float
+        Constant controlling smoothness in x-direction
+    Cy: float
+        Constant controlling smoothness in y-direction
+    Cz: float
+        Constant controlling smoothness in z-direction
+
+    Returns
+    -------
+    Js: float
+        value of smoothness cost function
+    """
+    dudx = _tf_gradient(u, dx, axis=2)
+    dudy = _tf_gradient(u, dy, axis=1)
+    dudz = _tf_gradient(u, dz, axis=0)
+    dvdx = _tf_gradient(v, dx, axis=2)
+    dvdy = _tf_gradient(v, dy, axis=1)
+    dvdz = _tf_gradient(v, dz, axis=0)
+    dwdx = _tf_gradient(w, dx, axis=2)
+    dwdy = _tf_gradient(w, dy, axis=1)
+    dwdz = _tf_gradient(w, dz, axis=0)
+
+    x_term = Cx * (_tf_gradient(dudx, dx, axis=2) ** 2 + _tf_gradient(dvdx, dx, axis=1) ** 2 +
+                   _tf_gradient(dwdx, dx, axis=2) ** 2)
+    y_term = Cy * (_tf_gradient(dudy, dy, axis=2) ** 2 + _tf_gradient(dvdy, dy, axis=1) ** 2 +
+                   _tf_gradient(dwdy, dy, axis=2) ** 2)
+    z_term = Cz * (_tf_gradient(dudz, dz, axis=2) ** 2 + _tf_gradient(dvdz, dz, axis=1) ** 2 +
+                   _tf_gradient(dwdz, dz, axis=2) ** 2)
+    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,
+                                  upper_bc=True):
+    """
+    Calculates the gradient of the smoothness cost function
+    by taking the Laplacian of the Laplacian of the wind field.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    Cx: float
+        Constant controlling smoothness in x-direction
+    Cy: float
+        Constant controlling smoothness in y-direction
+    Cz: float
+        Constant controlling smoothness in z-direction
+    upper_bc: bool
+        Set to true to impose w=0 at top of domain.
+
+    Returns
+    -------
+    y: float array
+        value of gradient of smoothness cost function
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        tape.watch(w)
+        loss = calculate_smoothness_cost(u, v, w,
+                                         dx, dy, dz, Cx=Cx, Cy=Cy, Cz=Cz)
+
+    vars = {'u': u, 'v': v, 'w': w}
+    grad = tape.gradient(loss, vars)
+    p_x1 = grad['u']
+    p_y1 = grad['v']
+    p_z1 = grad['w']
+
+    # Impermeability condition
+    p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])),
+                      p_z1[1:, :, :]], axis=0)
+    if (upper_bc is True):
+        p_z1 = tf.concat(
+            [p_z1[:-1, :, :], tf.zeros((1, u.shape[1], u.shape[2]), )], axis=0)
+    y = tf.stack((p_x1, p_y1, p_z1), axis=0)
+    return tf.reshape(y, (3 * np.prod(u.shape),))
+
+
+def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
+    """
+    Calculates the cost function related to point observations. A mean square error cost
+    function term is applied to points that are within the sphere of influence
+    whose radius is determined by *roi*.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    x:  Float array
+        X coordinates of grid centers
+    y:  Float array
+        Y coordinates of grid centers
+    z:  Float array
+        Z coordinated of grid centers
+    point_list: list of dicts
+        List of point constraints.
+        Each member is a dict with keys of "u", "v", to correspond
+        to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation.
+
+        In addition, "site_id" gives the METAR code (or name) to the station.
+    Cp: float
+        The weighting coefficient of the point cost function.
+    roi: float
+        Radius of influence of observations
+
+    Returns
+    -------
+    J: float
+        The cost function related to the difference between wind field and points.
+
+    """
+    J = 0.0
+
+    for the_point in point_list:
+        # Instead of worrying about whole domain, just find points in radius of influence
+        # Since we know that the weight will be zero outside the sphere of influence anyways
+        xp = tf.ones_like(x) * the_point["x"]
+        yp = tf.ones_like(y) * the_point["y"]
+        zp = tf.ones_like(z) * the_point["z"]
+        up = tf.ones_like(u) * the_point["u"]
+        vp = tf.ones_like(v) * the_point["v"]
+
+        the_box = tf.where(tf.math.logical_and(
+            tf.math.logical_and(tf.math.abs(x - xp) < roi, tf.math.abs(y - yp) < roi),
+            tf.math.abs(z - zp) < roi), 1., 0.)
+        J += tf.math.reduce_sum(((u - up) ** 2 + (v - vp) ** 2) * the_box)
+
+    return J * Cp
+
+
+def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0, upper_bc=True):
+    """
+    Calculates the gradient of the cost function related to point observations.
+    A mean square error cost function term is applied to points that are within the sphere of influence
+    whose radius is determined by *roi*.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    x: Float array
+        X coordinates of grid centers
+    y: Float array
+        Y coordinates of grid centers
+    z: Float array
+        Z coordinated of grid centers
+    point_list: list of dicts
+        List of point constraints. Each member is a dict with keys of "u", "v",
+        to correspond to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation.
+
+        In addition, "site_id" gives the METAR code (or name) to the station.
+    Cp: float
+        The weighting coefficient of the point cost function.
+    roi: float
+        Radius of influence of observations
+    upper_bc: bool
+        Set to true to impose w=0 at top of domain.
+    Returns
+    -------
+    gradJ: float array
+        The gradient of the cost function related to the difference between wind field and points.
+
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        loss = calculate_point_cost(u, v, x, y, z, point_list)
+
+    vars = {'u': u, 'v': v}
+    grad = tape.gradient(loss, vars)
+    gradJ_u = grad['u']
+    gradJ_v = grad['v']
+    gradJ_w = tf.zeros_like(gradJ_u)
+    gradJ = tf.stack([gradJ_u, gradJ_v, gradJ_w], axis=0)
+    gradJ = tf.reshape(gradJ, (3 * np.prod(u.shape),))
+    return gradJ * Cp
+
+
+def _tf_gradient(x, dx, axis):
+    if axis == 0:
+        fd = tf.experimental.numpy.diff(
+            tf.concat([x, tf.expand_dims(x[-1, :, :], 0)], axis=0), axis=0) / dx
+        bd = tf.experimental.numpy.diff(
+            tf.concat([tf.expand_dims(x[0, :, :], 0), x], axis=0), axis=0) / dx
+    elif axis == 1:
+        fd = tf.experimental.numpy.diff(
+            tf.concat([x, tf.expand_dims(x[:, -1, :], 1)], axis=1), axis=1) / dx
+        bd = tf.experimental.numpy.diff(
+            tf.concat([tf.expand_dims(x[:, 0, :], 1), x], axis=1), axis=1) / dx
+    elif axis == 2:
+        fd = tf.experimental.numpy.diff(
+            tf.concat([x, tf.expand_dims(x[:, :, -1], 2)], axis=2), axis=2) / dx
+        bd = tf.experimental.numpy.diff(
+            tf.concat([tf.expand_dims(x[:, :, 0], 2), x], axis=2), axis=2) / dx
+
+    cd = (fd + bd) / 2
+
+    if axis == 0:
+        return tf.concat([tf.expand_dims(fd[0, :, :], 0), cd[1:-1, :, :], tf.expand_dims(bd[-1, :, :], 0)], axis=0)
+    elif axis == 1:
+        return tf.concat([tf.expand_dims(fd[:, 0, :], 1), cd[:, 1:-1, :], tf.expand_dims(bd[:, -1, :], 1)], axis=1)
+    elif axis == 2:
+        return tf.concat([tf.expand_dims(fd[:, :, 0], 2), cd[:, :, 1:-1], tf.expand_dims(bd[:, :, -1], 2)], axis=2)
+
+
+def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1):
+    """
+    Calculates the mass continuity cost function by taking the divergence
+    of the wind field.
+
+    All arrays in the given lists must have the same dimensions and represent
+    the same spatial coordinates.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    dx: float
+        Grid spacing in x direction.
+    dy: float
+        Grid spacing in y direction.
+    dz: float
+        Grid spacing in z direction.
+    z: Float array (1D)
+        1D Float array with heights of grid
+    coeff: float
+        Constant controlling contribution of mass continuity to cost function
+    anel: int
+        = 1 use anelastic approximation, 0=don't
+
+    Returns
+    -------
+    J: float
+        value of mass continuity cost function
+
+    """
+    # Jax version of the cost function
+    rho = tf.math.exp(-z / 10000.0)
+    dudx = _tf_gradient(u, dx, axis=2)
+    dvdy = _tf_gradient(v, dy, axis=1)
+    dwdz = _tf_gradient(w, dz, axis=0)
+
+    if (anel == 1):
+        drho_dz = _tf_gradient(rho, dz, axis=0)
+        anel_term = w / rho * drho_dz
+    else:
+        anel_term = tf.zeros(w.shape)
+    return coeff * tf.math.reduce_sum(
+        tf.math.square(dudx + dvdy + dwdz + anel_term)) / 2.0
+
+
+def calculate_mass_continuity_gradient(u, v, w, z, dx,
+                                       dy, dz, coeff=1500.0, anel=1,
+                                       upper_bc=True):
+    """
+    Calculates the gradient of mass continuity cost function. This is done by
+    taking the negative gradient of the divergence of the wind field.
+
+    All grids must have the same grid specification.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    z: Float array (1D)
+        1D Float array with heights of grid
+    dx: float
+        Grid spacing in x direction.
+    dy: float
+        Grid spacing in y direction.
+    dz: float
+        Grid spacing in z direction.
+    coeff: float
+        Constant controlling contribution of mass continuity to cost function
+    anel: int
+        = 1 use anelastic approximation, 0=don't
+    upper_bc: bool
+        Set to true to impose w=0 at top of domain.
+
+    Returns
+    -------
+    y: float array
+        value of gradient of mass continuity cost function
+
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        tape.watch(w)
+        loss = calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff)
+
+    vars = {'u': u, 'v': v, 'w': w}
+    grad = tape.gradient(loss, vars)
+    p_x1 = grad['u']
+    p_y1 = grad['v']
+    p_z1 = grad['w']
+
+    # Impermeability condition
+    p_z1 = tf.concat(
+        [tf.zeros((1, u.shape[1], u.shape[2])),
+         p_z1[1:, :, :]], axis=0)
+    if (upper_bc is True):
+        p_z1 = tf.concat([p_z1[:-1, :, :],
+                          tf.zeros((1, u.shape[1], u.shape[2]))], axis=0)
+    y = tf.stack((p_x1, p_y1, p_z1), axis=0)
+    return tf.reshape(y, (3 * np.prod(u.shape),))
+
+
+def calculate_background_cost(u, v, weights, u_back, v_back, Cb=0.01):
+    """
+    Calculates the background cost function. The background cost function is
+    simply the sum of the squared differences between the wind field and the
+    background wind field multiplied by the weighting coefficient.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: Float array
+        Weights for each point to consider into cost function
+    u_back: 1D float array
+        Zonal winds vs height from sounding
+    w_back: 1D float array
+        Meridional winds vs height from sounding
+    Cb: float
+        Weight of background constraint to total cost function
+
+    Returns
+    -------
+    cost: float
+        value of background cost function
+    """
+    the_shape = u.shape
+    cost = tf.constant(0., dtype=tf.float32)
+
+    for i in range(the_shape[0]):
+        cost += tf.math.reduce_sum(
+            Cb * tf.math.square(u[i] - u_back[i]) * weights[i] + \
+            tf.math.square(v[i] - v_back[i]) * weights[i])
+    return cost
+
+
+def calculate_background_gradient(u, v, weights, u_back, v_back, Cb=0.01):
+    """
+    Calculates the gradient of the background cost function. For each u, v
+    this is given as 2*coefficent*(analysis wind - background wind).
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: Float array
+        Weights for each point to consider into cost function
+    u_back: 1D float array
+        Zonal winds vs height from sounding
+    w_back: 1D float array
+        Meridional winds vs height from sounding
+    Cb: float
+        Weight of background constraint to total cost function
+    upper_bc: bool
+        Set to true to impose w=0 at top of domain.
+
+    Returns
+    -------
+    y: float array
+        value of gradient of background cost function
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        loss = calculate_background_cost(u, v, weights, u_back, v_back, Cb=Cb)
+
+    vars = {'u': u, 'v': v}
+    grad = tape.gradient(loss, vars)
+    p_x1 = grad['u']
+    p_y1 = grad['v']
+    p_z1 = tf.zeros_like(p_x1)
+
+    y = tf.stack((p_x1, p_y1, p_z1), axis=0)
+    return tf.reshape(y, (3 * np.prod(u.shape),))
+
+
+def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt,
+                                      coeff=1):
+    """
+    Calculates the cost function due to deviance from vertical vorticity
+    equation. For more information of the vertical vorticity cost function,
+    see Potvin et al. (2012) and Shapiro et al. (2009).
+
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    dx: float array
+        Spacing in x grid
+    dy: float array
+        Spacing in y grid
+    dz: float array
+        Spacing in z grid
+    coeff: float
+        Weighting coefficient
+    Ut: float
+        U component of storm motion
+    Vt: float
+        V component of storm motion
+
+    Returns
+    -------
+    Jv: float
+        Value of vertical vorticity cost function.
+
+    References
+    ----------
+
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+
+    """
+
+    dvdz = _tf_gradient(v, dz, axis=0)
+    dudz = _tf_gradient(u, dz, axis=0)
+    dwdz = _tf_gradient(w, dz, axis=0)
+    dvdx = _tf_gradient(v, dx, axis=2)
+    dwdy = _tf_gradient(w, dy, axis=1)
+    dwdx = _tf_gradient(w, dx, axis=2)
+    dudx = _tf_gradient(u, dx, axis=2)
+    dvdy = _tf_gradient(v, dy, axis=1)
+    dudy = _tf_gradient(u, dy, axis=1)
+    print(dvdx[0, 0, :])
+    print(v[0, 0, :])
+    zeta = dvdx - dudy
+    dzeta_dx = _tf_gradient(zeta, dx, axis=2)
+    dzeta_dy = _tf_gradient(zeta, dy, axis=1)
+    dzeta_dz = _tf_gradient(zeta, dz, axis=0)
+    jv_array = ((u - Ut) * dzeta_dx + (v - Vt) * dzeta_dy +
+                w * dzeta_dz + (dvdz * dwdx - dudz * dwdy) +
+                zeta * (dudx + dvdy))
+    return tf.math.reduce_sum(coeff * tf.math.square(jv_array))
+
+
+# Using Jax version of function
+def calculate_vertical_vorticity_gradient(u, v, w, dx, dy, dz, Ut, Vt,
+                                          coeff=1e-5, upper_bc=True):
+    """
+    Calculates the gradient of the cost function due to deviance from vertical
+    vorticity equation. This is done by taking the functional derivative of
+    the vertical vorticity cost function.
+
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    dx: float array
+        Spacing in x grid
+    dy: float array
+        Spacing in y grid
+    dz: float array
+        Spacing in z grid
+    Ut: float
+        U component of storm motion
+    Vt: float
+        V component of storm motion
+    coeff: float
+        Weighting coefficient
+    upper_bc: bool
+        Set to true to impose w=0 at top of domain.
+
+    Returns
+    -------
+    Jv: 1D float array
+        Value of the gradient of the vertical vorticity cost function.
+
+    References
+    ----------
+
+    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
+    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
+    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
+    https://doi.org/10.1175/JTECH-D-11-00019.1
+
+    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
+    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
+    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
+
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        tape.watch(w)
+        loss = calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt, coeff)
+    vars = {'u': u, 'v': v, 'w': w}
+    grad = tape.gradient(loss, vars)
+    p_x1 = grad['u']
+    p_y1 = grad['v']
+    p_z1 = grad['w']
+
+    # Impermeability condition
+    p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])),
+                      p_z1[1:, :, :]], axis=0)
+    if (upper_bc is True):
+        p_z1 = tf.concat([p_z1[:-1, :, :],
+                          tf.zeros((1, u.shape[1], u.shape[2]))], axis=0)
+    y = tf.stack((p_x1, p_y1, p_z1), axis=0)
+    return tf.reshape(y, (3 * np.prod(u.shape),))
+
+
+def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model,
+                         coeff=1.0):
+    """
+    Calculates the cost function for the model constraint.
+    This is calculated simply as the sum of squares of the differences
+    between the model wind field and the analysis wind field. Vertical
+    velocities are not factored into this cost function as there is typically
+    a high amount of uncertainty in model derived vertical velocities.
+
+    Parameters
+    ----------
+    u: 3D array
+        Float array with u component of wind field
+    v: 3D array
+        Float array with v component of wind field
+    w: 3D array
+        Float array with w component of wind field
+    weights: list of 3D arrays
+        Float array showing how much each point from model weighs into
+        constraint.
+    u_model: list of 3D arrays
+        Float array with u component of wind field from model
+    v_model: list of 3D arrays
+        Float array with v component of wind field from model
+    w_model: list of 3D arrays
+        Float array with w component of wind field from model
+    coeff: float
+        Weighting coefficient
+
+    Returns
+    -------
+    Jv: float
+        Value of model cost function
+    """
+
+    cost = 0
+    for i in range(len(u_model)):
+        cost += (coeff * tf.math.reduce_sum(tf.math.square(u - u_model[i]) * weights[i] +
+                                            tf.math.square(v - v_model[i]) * weights[i]))
+    return cost
+
+
+def calculate_model_gradient(u, v, w, weights, u_model,
+                             v_model, w_model, coeff=1.0, upper_bc=True):
+    """
+    Calculates the cost function for the model constraint.
+    This is calculated simply as twice the differences
+    between the model wind field and the analysis wind field for each u, v.
+    Vertical velocities are not factored into this cost function as there is
+    typically a high amount of uncertainty in model derived vertical
+    velocities. Therefore, the gradient for all of the w's will be 0.
+
+    Parameters
+    ----------
+    u: Float array
+        Float array with u component of wind field
+    v: Float array
+        Float array with v component of wind field
+    w: Float array
+        Float array with w component of wind field
+    weights: list of 3D float arrays
+        Weights for each point to consider into cost function
+    u_model: list of 3D float arrays
+        Zonal wind field from model
+    v_model: list of 3D float arrays
+        Meridional wind field from model
+    w_model: list of 3D float arrays
+        Vertical wind field from model
+    coeff: float
+        Weight of background constraint to total cost function
+
+    Returns
+    -------
+    y: float array
+        value of gradient of background cost function
+    """
+
+    with tf.GradientTape() as tape:
+        tape.watch(u)
+        tape.watch(v)
+        tape.watch(w)
+        loss = calculate_model_cost(
+            u, v, w, weights, u_model, v_model, w_model, coeff=coeff)
+
+    vars = {'u': u, 'v': v, 'w': w}
+    grad = tape.gradient(loss, vars)
+    p_x1 = grad['u']
+    p_y1 = grad['v']
+    p_z1 = tf.zeros(p_x1.shape)
+
+    # Impermeability condition
+    p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])),
+                      p_z1[1:, :, :]], axis=0)
+    if (upper_bc is True):
+        p_z1 = tf.concat([p_z1[:-1, :, :],
+                          tf.zeros((1, u.shape[1], u.shape[2]))], axis=0)
+    y = tf.stack((p_x1, p_y1, p_z1), axis=0)
+    return tf.reshape(y, (3 * np.prod(u.shape),))
diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py
index e1356aa5..9d8748e1 100644
--- a/pydda/cost_functions/cost_functions.py
+++ b/pydda/cost_functions/cost_functions.py
@@ -1,1003 +1,447 @@
-import numpy as np
-import pyart
-import scipy.ndimage.filters
-
-
-def J_function(winds, parameters):
-    """
-    Calculates the total cost function. This typically does not need to be
-    called directly as get_dd_wind_field is a wrapper around this function and
-    :py:func:`pydda.cost_functions.grad_J`.
-    In order to add more terms to the cost function, modify this
-    function and :py:func:`pydda.cost_functions.grad_J`.
-
-    Parameters
-    ----------
-    winds: 1-D float array
-        The wind field, flattened to 1-D for f_min. The total size of the
-        array will be a 1D array of 3*nx*ny*nz elements.
-    parameters: DDParameters
-        The parameters for the cost function evaluation as specified by the
-        :py:func:`pydda.retrieval.DDParameters` class.
-
-    Returns
-    -------
-    J: float
-        The value of the cost function
-    """
-    winds = np.reshape(winds,
-                       (3, parameters.grid_shape[0], parameters.grid_shape[1],
-                        parameters.grid_shape[2]))
-
-    Jvel = calculate_radial_vel_cost_function(
-         parameters.vrs, parameters.azs, parameters.els,
-         winds[0], winds[1], winds[2], parameters.wts, rmsVr=parameters.rmsVr,
-         weights=parameters.weights, coeff=parameters.Co)
-
-    if(parameters.Cm > 0):
-        Jmass = calculate_mass_continuity(
-            winds[0], winds[1], winds[2], parameters.z,
-            parameters.dx, parameters.dy, parameters.dz,
-            coeff=parameters.Cm)
-    else:
-        Jmass = 0
-
-    if(parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
-        Jsmooth = calculate_smoothness_cost(
-            winds[0], winds[1], winds[2], Cx=parameters.Cx,
-            Cy=parameters.Cy, Cz=parameters.Cz)
-    else:
-        Jsmooth = 0
-
-    if(parameters.Cb > 0):
-        Jbackground = calculate_background_cost(
-            winds[0], winds[1], winds[2], parameters.bg_weights,
-            parameters.u_back, parameters.v_back, parameters.Cb)
-    else:
-        Jbackground = 0
-
-    if(parameters.Cv > 0):
-        Jvorticity = calculate_vertical_vorticity_cost(
-            winds[0], winds[1], winds[2], parameters.dx,
-            parameters.dy, parameters.dz, parameters.Ut,
-            parameters.Vt, coeff=parameters.Cv)
-    else:
-        Jvorticity = 0
-
-    if(parameters.Cmod > 0):
-        Jmod = calculate_model_cost(
-            winds[0], winds[1], winds[2],
-            parameters.model_weights, parameters.u_model,
-            parameters.v_model,
-            parameters.w_model, coeff=parameters.Cmod)
-    else:
-        Jmod = 0
-
-    if parameters.Cpoint > 0:
-        Jpoint = calculate_point_cost(
-            winds[0], winds[1], parameters.x, parameters.y, parameters.z,
-            parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi)
-    else:
-        Jpoint = 0
-
-    if(parameters.print_out is True):
-        print(('| Jvel    | Jmass   | Jsmooth |   Jbg   | Jvort   | Jmodel  | Jpoint  |' +
-               ' Max w  '))
-        print(('|' + "{:9.4f}".format(Jvel) + '|' +
-               "{:9.4f}".format(Jmass) + '|' +
-               "{:9.4f}".format(Jsmooth) + '|' +
-               "{:9.4f}".format(Jbackground) + '|' +
-               "{:9.4f}".format(Jvorticity) + '|' +
-               "{:9.4f}".format(Jmod) + '|' +
-               "{:9.4f}".format(Jpoint)) + '|' +
-               "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2]))))
-
-    return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + Jpoint
-
-
-def grad_J(winds, parameters):
-    """
-    Calculates the gradient of the cost function. This typically does not need
-    to be called directly as get_dd_wind_field is a wrapper around this
-    function and :py:func:`pydda.cost_functions.J_function`.
-    In order to add more terms to the cost function,
-    modify this function and :py:func:`pydda.cost_functions.grad_J`.
-
-    Parameters
-    ----------
-    winds: 1-D float array
-        The wind field, flattened to 1-D for f_min
-    parameters: DDParameters
-        The parameters for the cost function evaluation as specified by the
-        :py:func:`pydda.retrieve.DDParameters` class.
-
-    Returns
-    -------
-    grad: 1D float array
-        Gradient vector of cost function
-    """
-    winds = np.reshape(winds,
-                       (3, parameters.grid_shape[0],
-                        parameters.grid_shape[1], parameters.grid_shape[2]))
-    grad = calculate_grad_radial_vel(
-        parameters.vrs, parameters.els, parameters.azs,
-        winds[0], winds[1], winds[2], parameters.wts, parameters.weights,
-        parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc)
-
-    if(parameters.Cm > 0):
-        grad += calculate_mass_continuity_gradient(
-            winds[0], winds[1], winds[2], parameters.z,
-            parameters.dx, parameters.dy, parameters.dz,
-            coeff=parameters.Cm, upper_bc=parameters.upper_bc)
-
-    if(parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
-        grad += calculate_smoothness_gradient(
-            winds[0], winds[1], winds[2], Cx=parameters.Cx,
-            Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc)
-
-    if(parameters.Cb > 0):
-        grad += calculate_background_gradient(
-            winds[0], winds[1], winds[2], parameters.bg_weights,
-            parameters.u_back, parameters.v_back, parameters.Cb,
-            upper_bc=parameters.upper_bc)
-
-    if(parameters.Cv > 0):
-        grad += calculate_vertical_vorticity_gradient(
-            winds[0], winds[1], winds[2], parameters.dx,
-            parameters.dy, parameters.dz, parameters.Ut,
-            parameters.Vt, coeff=parameters.Cv)
-
-    if(parameters.Cmod > 0):
-        grad += calculate_model_gradient(
-            winds[0], winds[1], winds[2],
-            parameters.model_weights, parameters.u_model, parameters.v_model,
-            parameters.w_model, coeff=parameters.Cmod)
-
-    if parameters.Cpoint > 0:
-        grad += calculate_point_gradient(
-            winds[0], winds[1], parameters.x, parameters.y, parameters.z,
-            parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi)
-
-    if(parameters.print_out is True):
-        print('Norm of gradient: ' + str(np.linalg.norm(grad, np.inf)))
-
-    return grad
-
-
-def calculate_radial_vel_cost_function(vrs, azs, els, u, v,
-                                       w, wts, rmsVr, weights, coeff=1.0):
-    """
-    Calculates the cost function due to difference of the wind field from
-    radar radial velocities. For more information on this cost function, see
-    Potvin et al. (2012) and Shapiro et al. (2009).
-
-    All arrays in the given lists must have the same dimensions and represent
-    the same spatial coordinates.
-
-    Parameters
-    ----------
-    vrs: List of float arrays
-        List of radial velocities from each radar
-    els: List of float arrays
-        List of elevations from each radar
-    azs: List of float arrays
-        List of azimuths from each radar
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    wts: List of float arrays
-        Float array containing fall speed from radar.
-    rmsVr: float
-        The sum of squares of velocity/num_points. Use for normalization
-        of data weighting coefficient
-    weights: n_radars x_bins x y_bins float array
-        Data weights for each pair of radars
-    coeff: float
-        Constant for cost function
-
-    Returns
-    -------
-    J_o: float
-         Observational cost function
-
-    References
-    -----------
-
-    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
-    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
-    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
-    https://doi.org/10.1175/JTECH-D-11-00019.1
-
-    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
-    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
-    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
-    """
-
-    J_o = 0
-    lambda_o = coeff / (rmsVr * rmsVr)
-    for i in range(len(vrs)):
-        v_ar = (np.cos(els[i])*np.sin(azs[i])*u +
-                np.cos(els[i])*np.cos(azs[i])*v +
-                np.sin(els[i])*(w - np.abs(wts[i])))
-        the_weight = weights[i]
-        the_weight[els[i].mask] = 0
-        the_weight[azs[i].mask] = 0
-        the_weight[vrs[i].mask] = 0
-        the_weight[wts[i].mask] = 0
-        J_o += lambda_o*np.sum(np.square(vrs[i] - v_ar)*the_weight)
-
-    return J_o
-
-
-def calculate_grad_radial_vel(vrs, els, azs, u, v, w,
-                              wts, weights, rmsVr, coeff=1.0, upper_bc=True):
-    """
-    Calculates the gradient of the cost function due to difference of wind
-    field from radar radial velocities.
-
-    All arrays in the given lists must have the same dimensions and represent
-    the same spatial coordinates.
-
-    Parameters
-    ----------
-    vrs: List of float arrays
-        List of radial velocities from each radar
-    els: List of float arrays
-        List of elevations from each radar
-    azs: List of azimuths
-        List of azimuths from each radar
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    coeff: float
-        Constant for cost function
-    vel_name: str
-        Background velocity field name
-    weights: n_radars x_bins x y_bins float array
-        Data weights for each pair of radars
-
-    Returns
-    -------
-    y: 1-D float array
-         Gradient vector of observational cost function. 
-         
-    More information
-    ----------------
-    The gradient is calculated by taking the functional derivative of the
-    cost function. For more information on functional derivatives, see the
-    Euler-Lagrange Equation:
-
-    https://en.wikipedia.org/wiki/Euler%E2%80%93Lagrange_equation
-
-    """
-
-    # Use zero for all masked values since we don't want to add them into
-    # the cost function
-
-    p_x1 = np.zeros(vrs[0].shape)
-    p_y1 = np.zeros(vrs[0].shape)
-    p_z1 = np.zeros(vrs[0].shape)
-    lambda_o = coeff / (rmsVr * rmsVr)
-
-    for i in range(len(vrs)):
-        v_ar = (np.cos(els[i])*np.sin(azs[i])*u +
-                np.cos(els[i])*np.cos(azs[i])*v +
-                np.sin(els[i])*(w - np.abs(wts[i])))
-
-        x_grad = (2*(v_ar - vrs[i]) * np.cos(els[i]) *
-                  np.sin(azs[i]) * weights[i]) * lambda_o
-        y_grad = (2*(v_ar - vrs[i]) * np.cos(els[i]) *
-                  np.cos(azs[i]) * weights[i]) * lambda_o
-        z_grad = (2*(v_ar - vrs[i]) * np.sin(els[i]) * weights[i]) * lambda_o
-
-        x_grad[els[i].mask] = 0
-        y_grad[els[i].mask] = 0
-        z_grad[els[i].mask] = 0
-        x_grad[azs[i].mask] = 0
-        y_grad[azs[i].mask] = 0
-        z_grad[azs[i].mask] = 0
-
-        x_grad[els[i].mask] = 0
-        x_grad[azs[i].mask] = 0
-        x_grad[vrs[i].mask] = 0
-        x_grad[wts[i].mask] = 0
-        y_grad[els[i].mask] = 0
-        y_grad[azs[i].mask] = 0
-        y_grad[vrs[i].mask] = 0
-        y_grad[wts[i].mask] = 0
-        z_grad[els[i].mask] = 0
-        z_grad[azs[i].mask] = 0
-        z_grad[vrs[i].mask] = 0
-        z_grad[wts[i].mask] = 0
-
-        p_x1 += x_grad
-        p_y1 += y_grad
-        p_z1 += z_grad
-
-    # Impermeability condition
-    p_z1[0, :, :] = 0
-    if(upper_bc is True):
-        p_z1[-1, :, :] = 0
-    y = np.stack((p_x1, p_y1, p_z1), axis=0)
-    return y.flatten()
-
-
-def calculate_smoothness_cost(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5):
-    """
-    Calculates the smoothness cost function by taking the Laplacian of the
-    wind field.
-
-    All arrays in the given lists must have the same dimensions and represent
-    the same spatial coordinates.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    Cx: float
-        Constant controlling smoothness in x-direction
-    Cy: float
-        Constant controlling smoothness in y-direction
-    Cz: float
-        Constant controlling smoothness in z-direction
-
-    Returns
-    -------
-    Js: float
-        value of smoothness cost function
-    """
-    du = np.zeros(w.shape)
-    dv = np.zeros(w.shape)
-    dw = np.zeros(w.shape)
-    scipy.ndimage.filters.laplace(u, du, mode='wrap')
-    scipy.ndimage.filters.laplace(v, dv, mode='wrap')
-    scipy.ndimage.filters.laplace(w, dw, mode='wrap')
-    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,
-                                  upper_bc=True):
-    """
-    Calculates the gradient of the smoothness cost function
-    by taking the Laplacian of the Laplacian of the wind field.
-
-    All arrays in the given lists must have the same dimensions and represent
-    the same spatial coordinates.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    Cx: float
-        Constant controlling smoothness in x-direction
-    Cy: float
-        Constant controlling smoothness in y-direction
-    Cz: float
-        Constant controlling smoothness in z-direction
-
-    Returns
-    -------
-    y: float array
-        value of gradient of smoothness cost function
-    """
-    du = np.zeros(w.shape)
-    dv = np.zeros(w.shape)
-    dw = np.zeros(w.shape)
-    grad_u = np.zeros(w.shape)
-    grad_v = np.zeros(w.shape)
-    grad_w = np.zeros(w.shape)
-    scipy.ndimage.filters.laplace(u, du, mode='wrap')
-    scipy.ndimage.filters.laplace(v, dv, mode='wrap')
-    scipy.ndimage.filters.laplace(w, dw, mode='wrap')
-    scipy.ndimage.filters.laplace(du, grad_u, mode='wrap')
-    scipy.ndimage.filters.laplace(dv, grad_v, mode='wrap')
-    scipy.ndimage.filters.laplace(dw, grad_w, mode='wrap')
-
-    # Impermeability condition
-    grad_w[0, :, :] = 0
-    if(upper_bc is True):
-        grad_w[-1, :, :] = 0
-    y = np.stack([grad_u*Cx*2, grad_v*Cy*2, grad_w*Cz*2], axis=0)
-    return y.flatten()
-
-
-def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
-    """
-    Calculates the cost function related to point observations. A mean square error cost
-    function term is applied to points that are within the sphere of influence
-    whose radius is determined by *roi*.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    x:  Float array
-        X coordinates of grid centers
-    y:  Float array
-        Y coordinates of grid centers
-    z:  Float array
-        Z coordinated of grid centers
-    point_list: list of dicts
-        List of point constraints.
-        Each member is a dict with keys of "u", "v", to correspond
-        to each component of the wind field and "x", "y", "z"
-        to correspond to the location of the point observation.
-
-        In addition, "site_id" gives the METAR code (or name) to the station.
-    Cp: float
-        The weighting coefficient of the point cost function.
-    roi: float
-        Radius of influence of observations
-
-    Returns
-    -------
-    J: float
-        The cost function related to the difference between wind field and points.
-
-    """
-    J = 0.0
-    for the_point in point_list:
-        # Instead of worrying about whole domain, just find points in radius of influence
-        # Since we know that the weight will be zero outside the sphere of influence anyways
-        the_box = np.where(np.logical_and.reduce(
-            (np.abs(x - the_point["x"]) < roi, np.abs(y - the_point["y"]) < roi,
-                                            np.abs(z - the_point["z"]) < roi)))
-        J += np.sum(((u[the_box] - the_point["u"])**2 + (v[the_box] - the_point["v"])**2))
-
-    return J * Cp
-
-
-def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0):
-    """
-    Calculates the gradient of the cost function related to point observations.
-    A mean square error cost function term is applied to points that are within the sphere of influence
-    whose radius is determined by *roi*.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    x: Float array
-        X coordinates of grid centers
-    y: Float array
-        Y coordinates of grid centers
-    z: Float array
-        Z coordinated of grid centers
-    point_list: list of dicts
-        List of point constraints. Each member is a dict with keys of "u", "v",
-        to correspond to each component of the wind field and "x", "y", "z"
-        to correspond to the location of the point observation.
-
-        In addition, "site_id" gives the METAR code (or name) to the station.
-    Cp: float
-        The weighting coefficient of the point cost function.
-    roi: float
-        Radius of influence of observations
-
-    Returns
-    -------
-    gradJ: float array
-        The gradient of the cost function related to the difference between wind field and points.
-
-    """
-
-    gradJ_u = np.zeros_like(u)
-    gradJ_v = np.zeros_like(v)
-    gradJ_w = np.zeros_like(u)
-
-    for the_point in point_list:
-        the_box = np.where(np.logical_and.reduce(
-            (np.abs(x - the_point["x"]) < roi, np.abs(y - the_point["y"]) < roi,
-             np.abs(z - the_point["z"]) < roi)))
-        gradJ_u[the_box] += 2 * (u[the_box] - the_point["u"])
-        gradJ_v[the_box] += 2 * (v[the_box] - the_point["v"])
-
-    gradJ = np.stack([gradJ_u, gradJ_v, gradJ_w], axis=0).flatten()
-    return gradJ * Cp
-
-
-def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1):
-    """
-    Calculates the mass continuity cost function by taking the divergence
-    of the wind field.
-
-    All arrays in the given lists must have the same dimensions and represent
-    the same spatial coordinates.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    dx: float
-        Grid spacing in x direction.
-    dy: float
-        Grid spacing in y direction.
-    dz: float
-        Grid spacing in z direction.
-    z: Float array (1D)
-        1D Float array with heights of grid
-    coeff: float
-        Constant controlling contribution of mass continuity to cost function
-    anel: int
-        = 1 use anelastic approximation, 0=don't
-
-    Returns
-    -------
-    J: float
-        value of mass continuity cost function
-    """
-    dudx = np.gradient(u, dx, axis=2)
-    dvdy = np.gradient(v, dy, axis=1)
-    dwdz = np.gradient(w, dz, axis=0)
-
-    if(anel == 1):
-        rho = np.exp(-z/10000.0)
-        drho_dz = np.gradient(rho, dz, axis=0)
-        anel_term = w/rho*drho_dz
-    else:
-        anel_term = np.zeros(w.shape)
-    return coeff*np.sum(np.square(dudx + dvdy + dwdz + anel_term))/2.0
-
-
-def calculate_mass_continuity_gradient(u, v, w, z, dx,
-                                       dy, dz, coeff=1500.0, anel=1,
-                                       upper_bc=True):
-    """
-    Calculates the gradient of mass continuity cost function. This is done by
-    taking the negative gradient of the divergence of the wind field.
-
-    All grids must have the same grid specification.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    z: Float array (1D)
-        1D Float array with heights of grid
-    dx: float
-        Grid spacing in x direction.
-    dy: float
-        Grid spacing in y direction.
-    dz: float
-        Grid spacing in z direction.
-    coeff: float
-        Constant controlling contribution of mass continuity to cost function
-    anel: int
-        = 1 use anelastic approximation, 0=don't
-
-    Returns
-    -------
-    y: float array
-        value of gradient of mass continuity cost function
-    """
-    dudx = np.gradient(u, dx, axis=2)
-    dvdy = np.gradient(v, dy, axis=1)
-    dwdz = np.gradient(w, dz, axis=0)
-    if(anel == 1):
-        rho = np.exp(-z/10000.0)
-        drho_dz = np.gradient(rho, dz, axis=0)
-        anel_term = w/rho*drho_dz
-    else:
-        anel_term = 0
-
-    div2 = dudx + dvdy + dwdz + anel_term
-
-    grad_u = -np.gradient(div2, dx, axis=2)*coeff
-    grad_v = -np.gradient(div2, dy, axis=1)*coeff
-    grad_w = -np.gradient(div2, dz, axis=0)*coeff
-
-    # Impermeability condition
-    grad_w[0, :, :] = 0
-    if(upper_bc is True):
-        grad_w[-1, :, :] = 0
-    y = np.stack([grad_u, grad_v, grad_w], axis=0)
-    return y.flatten()
-
-
-def calculate_fall_speed(grid, refl_field=None, frz=4500.0):
-    """
-    Estimates fall speed based on reflectivity.
-
-    Uses methodology of Mike Biggerstaff and Dan Betten
-
-    Parameters
-    ----------
-    Grid: Py-ART Grid
-        Py-ART Grid containing reflectivity to calculate fall speed from
-    refl_field: str
-        String containing name of reflectivity field. None will automatically
-        determine the name.
-    frz: float
-        Height of freezing level in m
-
-    Returns
-    -------
-    3D float array:
-        Float array of terminal velocities
-
-    """
-    # Parse names of velocity field
-    if refl_field is None:
-        refl_field = pyart.config.get_field_name('reflectivity')
-
-    refl = grid.fields[refl_field]['data']
-    grid_z = grid.point_z['data']
-    term_vel = np.zeros(refl.shape)
-    A = np.zeros(refl.shape)
-    B = np.zeros(refl.shape)
-    rho = np.exp(-grid_z/10000.0)
-    A[np.logical_and(grid_z < frz, refl < 55)] = -2.6
-    B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107
-    A[np.logical_and(grid_z < frz,
-                     np.logical_and(refl >= 55, refl < 60))] = -2.5
-    B[np.logical_and(grid_z < frz,
-                     np.logical_and(refl >= 55, refl < 60))] = 0.013
-    A[np.logical_and(grid_z < frz, refl > 60)] = -3.95
-    B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148
-    A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817
-    B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063
-    A[np.logical_and(grid_z >= frz,
-                     np.logical_and(refl >= 33, refl < 49))] = -2.5
-    B[np.logical_and(grid_z >= frz,
-                     np.logical_and(refl >= 33, refl < 49))] = 0.013
-    A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95
-    B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148
-
-    fallspeed = A*np.power(10, refl*B)*np.power(1.2/rho, 0.4)
-    del A, B, rho
-    return fallspeed
-
-
-def calculate_background_cost(u, v, w, weights, u_back, v_back, Cb=0.01):
-    """
-    Calculates the background cost function. The background cost function is
-    simply the sum of the squared differences between the wind field and the
-    background wind field multiplied by the weighting coefficient.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    weights: Float array
-        Weights for each point to consider into cost function
-    u_back: 1D float array
-        Zonal winds vs height from sounding
-    w_back: 1D float array
-        Meridional winds vs height from sounding
-    Cb: float
-        Weight of background constraint to total cost function
-
-    Returns
-    -------
-    cost: float
-        value of background cost function
-    """
-    the_shape = u.shape
-    cost = 0
-    for i in range(the_shape[0]):
-        cost += (Cb*np.sum(np.square(u[i]-u_back[i])*(weights[i]) +
-                           np.square(v[i]-v_back[i])*(weights[i])))
-    return cost
-
-
-def calculate_background_gradient(u, v, w, weights, u_back, v_back, Cb=0.01):
-    """
-    Calculates the gradient of the background cost function. For each u, v
-    this is given as 2*coefficent*(analysis wind - background wind).
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    weights: Float array
-        Weights for each point to consider into cost function
-    u_back: 1D float array
-        Zonal winds vs height from sounding
-    w_back: 1D float array
-        Meridional winds vs height from sounding
-    Cb: float
-        Weight of background constraint to total cost function
-
-    Returns
-    -------
-    y: float array
-        value of gradient of background cost function
-    """
-    the_shape = u.shape
-    u_grad = np.zeros(the_shape)
-    v_grad = np.zeros(the_shape)
-    w_grad = np.zeros(the_shape)
-
-    for i in range(the_shape[0]):
-        u_grad[i] = Cb*2*(u[i]-u_back[i])*(weights[i])
-        v_grad[i] = Cb*2*(v[i]-v_back[i])*(weights[i])
-
-    y = np.stack([u_grad, v_grad, w_grad], axis=0)
-    return y.flatten()
-
-
-def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt,
-                                      coeff=1e-5):
-    """
-    Calculates the cost function due to deviance from vertical vorticity
-    equation. For more information of the vertical vorticity cost function,
-    see Potvin et al. (2012) and Shapiro et al. (2009).
-
-    Parameters
-    ----------
-    u: 3D array
-        Float array with u component of wind field
-    v: 3D array
-        Float array with v component of wind field
-    w: 3D array
-        Float array with w component of wind field
-    dx: float array
-        Spacing in x grid
-    dy: float array
-        Spacing in y grid
-    dz: float array
-        Spacing in z grid
-    coeff: float
-        Weighting coefficient
-    Ut: float
-        U component of storm motion
-    Vt: float
-        V component of storm motion
-
-    Returns
-    -------
-    Jv: float
-        Value of vertical vorticity cost function.
-
-    References
-    ----------
-
-    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
-    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
-    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
-    https://doi.org/10.1175/JTECH-D-11-00019.1
-
-    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
-    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
-    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
-    """
-    dvdz = np.gradient(v, dz, axis=0)
-    dudz = np.gradient(u, dz, axis=0)
-    dvdx = np.gradient(v, dx, axis=2)
-    dwdy = np.gradient(w, dy, axis=1)
-    dwdx = np.gradient(w, dx, axis=2)
-    dudx = np.gradient(u, dx, axis=2)
-    dvdy = np.gradient(v, dy, axis=2)
-    dudy = np.gradient(u, dy, axis=1)
-    zeta = dvdx - dudy
-    dzeta_dx = np.gradient(zeta, dx, axis=2)
-    dzeta_dy = np.gradient(zeta, dy, axis=1)
-    dzeta_dz = np.gradient(zeta, dz, axis=0)
-    jv_array = ((u - Ut) * dzeta_dx + (v - Vt) * dzeta_dy +
-                w * dzeta_dz + (dvdz * dwdx - dudz * dwdy) +
-                zeta * (dudx + dvdy))
-    return np.sum(coeff*jv_array**2)
-
-
-def calculate_vertical_vorticity_gradient(u, v, w, dx, dy, dz, Ut, Vt,
-                                          coeff=1e-5):
-    """
-    Calculates the gradient of the cost function due to deviance from vertical
-    vorticity equation. This is done by taking the functional derivative of
-    the vertical vorticity cost function.
-
-    Parameters
-    ----------
-    u: 3D array
-        Float array with u component of wind field
-    v: 3D array
-        Float array with v component of wind field
-    w: 3D array
-        Float array with w component of wind field
-    dx: float array
-        Spacing in x grid
-    dy: float array
-        Spacing in y grid
-    dz: float array
-        Spacing in z grid
-    Ut: float
-        U component of storm motion
-    Vt: float
-        V component of storm motion
-    coeff: float
-        Weighting coefficient
-
-    Returns
-    -------
-    Jv: 1D float array
-        Value of the gradient of the vertical vorticity cost function.
-
-    References
-    ----------
-
-    Potvin, C.K., A. Shapiro, and M. Xue, 2012: Impact of a Vertical Vorticity
-    Constraint in Variational Dual-Doppler Wind Analysis: Tests with Real and
-    Simulated Supercell Data. J. Atmos. Oceanic Technol., 29, 32–49,
-    https://doi.org/10.1175/JTECH-D-11-00019.1
-
-    Shapiro, A., C.K. Potvin, and J. Gao, 2009: Use of a Vertical Vorticity
-    Equation in Variational Dual-Doppler Wind Analysis. J. Atmos. Oceanic
-    Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
-    """
-
-    # First derivatives
-    dvdz = np.gradient(v, dz, axis=0)
-    dwdy = np.gradient(w, dy, axis=1)
-    dudx = np.gradient(u, dx, axis=2)
-    dvdy = np.gradient(v, dy, axis=1)
-    dvdx = np.gradient(v, dx, axis=2)
-    dwdx = np.gradient(w, dx, axis=2)
-    dudz = np.gradient(u, dz, axis=0)
-    dudy = np.gradient(u, dy, axis=1)
-
-    zeta = dvdx - dudy
-    dzeta_dx = np.gradient(zeta, dx, axis=2)
-    dzeta_dy = np.gradient(zeta, dy, axis=1)
-    dzeta_dz = np.gradient(zeta, dz, axis=0)
-
-    # Second deriviatives
-    dwdydz = np.gradient(dwdy, dz, axis=0)
-    dwdxdz = np.gradient(dwdx, dz, axis=0)
-    dudzdy = np.gradient(dudz, dy, axis=1)
-    dvdxdy = np.gradient(dvdx, dy, axis=1)
-    dudx2 = np.gradient(dudx, dx, axis=2)
-    dudxdy = np.gradient(dudx, dy, axis=1)
-    dudxdz = np.gradient(dudx, dz, axis=0)
-    dudy2 = np.gradient(dudx, dy, axis=1)
-
-    dzeta_dt = ((u - Ut)*dzeta_dx + (v - Vt)*dzeta_dy + w*dzeta_dz +
-                (dvdz*dwdx - dudz*dwdy) + zeta*(dudx + dvdy))
-
-    # Now we intialize our gradient value
-    u_grad = np.zeros(u.shape)
-    v_grad = np.zeros(v.shape)
-    w_grad = np.zeros(w.shape)
-
-    # Vorticity Advection
-    u_grad += dzeta_dx + (Ut - u)*dudxdy + (Vt - v)*dudxdy
-    v_grad += dzeta_dy + (Vt - v)*dvdxdy + (Ut - u)*dvdxdy
-    w_grad += dzeta_dz
-
-    # Tilting term
-    u_grad += dwdydz
-    v_grad += dwdxdz
-    w_grad += dudzdy - dudxdz
-
-    # Stretching term
-    u_grad += -dudxdy + dudy2 - dzeta_dx
-    u_grad += -dudx2 + dudxdy - dzeta_dy
-
-    # Multiply by 2*dzeta_dt according to chain rule
-    u_grad = u_grad*2*dzeta_dt*coeff
-    v_grad = v_grad*2*dzeta_dt*coeff
-    w_grad = w_grad*2*dzeta_dt*coeff
-
-    y = np.stack([u_grad, v_grad, w_grad], axis=0)
-    return y.flatten()
-
-
-def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model,
-                         coeff=1.0):
-    """
-    Calculates the cost function for the model constraint.
-    This is calculated simply as the sum of squares of the differences
-    between the model wind field and the analysis wind field. Vertical
-    velocities are not factored into this cost function as there is typically
-    a high amount of uncertainty in model derived vertical velocities.
-
-    Parameters
-    ----------
-    u: 3D array
-        Float array with u component of wind field
-    v: 3D array
-        Float array with v component of wind field
-    w: 3D array
-        Float array with w component of wind field
-    weights: list of 3D arrays
-        Float array showing how much each point from model weighs into
-        constraint.
-    u_model: list of 3D arrays
-        Float array with u component of wind field from model
-    v_model: list of 3D arrays
-        Float array with v component of wind field from model
-    w_model: list of 3D arrays
-        Float array with w component of wind field from model
-    coeff: float
-        Weighting coefficient
-
-    Returns
-    -------
-    Jv: float
-        Value of model cost function
-    """
-
-    cost = 0
-    for i in range(len(u_model)):
-        cost += (coeff*np.sum(np.square(u-u_model[i])*weights[i] +
-                              np.square(v-v_model[i])*weights[i]))
-    return cost
-
-
-def calculate_model_gradient(u, v, w, weights, u_model,
-                             v_model, w_model, coeff=1.0):
-    """
-    Calculates the cost function for the model constraint.
-    This is calculated simply as twice the differences
-    between the model wind field and the analysis wind field for each u, v.
-    Vertical velocities are not factored into this cost function as there is
-    typically a high amount of uncertainty in model derived vertical
-    velocities. Therefore, the gradient for all of the w's will be 0.
-
-    Parameters
-    ----------
-    u: Float array
-        Float array with u component of wind field
-    v: Float array
-        Float array with v component of wind field
-    w: Float array
-        Float array with w component of wind field
-    weights: list of 3D float arrays
-        Weights for each point to consider into cost function
-    u_model: list of 3D float arrays
-        Zonal wind field from model
-    v_model: list of 3D float arrays
-        Meridional wind field from model
-    w_model: list of 3D float arrays
-        Vertical wind field from model
-    coeff: float
-        Weight of background constraint to total cost function
-
-    Returns
-    -------
-    y: float array
-        value of gradient of background cost function
-    """
-    the_shape = u.shape
-    u_grad = np.zeros(the_shape)
-    v_grad = np.zeros(the_shape)
-    w_grad = np.zeros(the_shape)
-    for i in range(len(u_model)):
-        u_grad += coeff*2*(u-u_model[i])*weights[i]
-        v_grad += coeff*2*(v-v_model[i])*weights[i]
-
-    y = np.stack([u_grad, v_grad, w_grad], axis=0)
-    return y.flatten()
+import numpy as np
+#Adding jax inport statements
+try:
+    import tensorflow as tf
+    TENSORFLOW_AVAILABLE = True
+except ImportError:
+    TENSORFLOW_AVAILABLE = False
+
+try:
+    from jax.config import config
+    config.update("jax_enable_x64", True)
+    from jax import float0
+    import jax.numpy as jnp
+    JAX_AVAILABLE = True
+except ImportError:
+    JAX_AVAILABLE = False
+
+import pyart
+import scipy.ndimage.filters
+
+#Added to incorpeate JAX within the cost functions
+from . import _cost_functions_jax
+from . import _cost_functions_numpy
+from . import _cost_functions_tensorflow
+
+def J_function(winds, parameters):
+    """
+    Calculates the total cost function. This typically does not need to be
+    called directly as get_dd_wind_field is a wrapper around this function and
+    :py:func:`pydda.cost_functions.grad_J`.
+    In order to add more terms to the cost function, modify this
+    function and :py:func:`pydda.cost_functions.grad_J`.
+
+    Parameters
+    ----------
+    winds: 1-D float array
+        The wind field, flattened to 1-D for f_min. The total size of the
+        array will be a 1D array of 3*nx*ny*nz elements.
+    parameters: DDParameters
+        The parameters for the cost function evaluation as specified by the
+        :py:func:`pydda.retrieval.DDParameters` class.
+
+    Returns
+    -------
+    J: float
+        The value of the cost function
+    """
+    if parameters.engine == "tensorflow":
+        if not TENSORFLOW_AVAILABLE:
+            raise ImportError("Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!")
+        winds = tf.reshape(winds,
+                           (3, parameters.grid_shape[0], parameters.grid_shape[1],
+                            parameters.grid_shape[2]))
+        #Had to change to float because Jax returns device array (use np.float_())
+        Jvel = _cost_functions_tensorflow.calculate_radial_vel_cost_function(
+             parameters.vrs, parameters.azs, parameters.els,
+            winds[0], winds[1], winds[2],  parameters.wts, rmsVr=parameters.rmsVr,
+             weights=parameters.weights, coeff=parameters.Co)
+        #print("apples Jvel", Jvel)
+
+        if(parameters.Cm > 0):
+            #Had to change to float because Jax returns device array (use np.float_())
+            Jmass = _cost_functions_tensorflow.calculate_mass_continuity(winds[0], winds[1], winds[2],
+                parameters.z, parameters.dx, parameters.dy, parameters.dz,
+                coeff=parameters.Cm)
+        else:
+            Jmass = 0
+
+        if(parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
+            Jsmooth = _cost_functions_tensorflow.calculate_smoothness_cost(
+                winds[0], winds[1], winds[2], parameters.dx, parameters.dy, parameters.dz,
+                Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz)
+        else:
+            Jsmooth = 0
+
+        if(parameters.Cb > 0):
+            Jbackground = _cost_functions_tensorflow.calculate_background_cost(
+                winds[0], winds[1], parameters.bg_weights,
+                parameters.u_back, parameters.v_back, parameters.Cb)
+        else:
+            Jbackground = 0
+
+        if(parameters.Cv > 0):
+            #Had to change to float because Jax returns device array (use np.float_())
+            Jvorticity = _cost_functions_tensorflow.calculate_vertical_vorticity_cost(
+                winds[0], winds[1], winds[2], parameters.dx,
+                parameters.dy, parameters.dz, parameters.Ut,
+                parameters.Vt, coeff=parameters.Cv)
+        else:
+            Jvorticity = 0
+
+        if(parameters.Cmod > 0):
+            Jmod = _cost_functions_tensorflow.calculate_model_cost(
+                winds[0], winds[1], winds[2],
+                parameters.model_weights, parameters.u_model,
+                parameters.v_model,
+                parameters.w_model, coeff=parameters.Cmod)
+        else:
+            Jmod = 0
+
+        if parameters.Cpoint > 0:
+            Jpoint = _cost_functions_tensorflow.calculate_point_cost(
+                winds[0], winds[1], parameters.x, parameters.y, parameters.z,
+                parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi)
+        else:
+            Jpoint = 0
+    elif parameters.engine == "scipy":
+        winds = np.reshape(winds,
+                           (3, parameters.grid_shape[0], parameters.grid_shape[1],
+                            parameters.grid_shape[2]))
+        # Had to change to float because Jax returns device array (use np.float_())
+        Jvel = _cost_functions_numpy.calculate_radial_vel_cost_function(
+            parameters.vrs, parameters.azs, parameters.els,
+            winds[0], winds[1], winds[2], parameters.wts, rmsVr=parameters.rmsVr,
+            weights=parameters.weights, coeff=parameters.Co)
+        # print("apples Jvel", Jvel)
+
+        if (parameters.Cm > 0):
+            # Had to change to float because Jax returns device array (use np.float_())
+            Jmass = _cost_functions_numpy.calculate_mass_continuity(winds[0], winds[1], winds[2],
+                                                                    parameters.z, parameters.dx, parameters.dy,
+                                                                    parameters.dz, coeff=parameters.Cm)
+        else:
+            Jmass = 0
+
+        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], 
+                Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz)
+        else:
+            Jsmooth = 0
+
+        if (parameters.Cb > 0):
+            Jbackground = _cost_functions_numpy.calculate_background_cost(
+                winds[0], winds[1], parameters.bg_weights,
+                parameters.u_back, parameters.v_back, parameters.Cb)
+        else:
+            Jbackground = 0
+
+        if (parameters.Cv > 0):
+            # Had to change to float because Jax returns device array (use np.float_())
+            Jvorticity = _cost_functions_numpy.calculate_vertical_vorticity_cost(
+                winds[0], winds[1], winds[2], parameters.dx,
+                parameters.dy, parameters.dz, parameters.Ut,
+                parameters.Vt, coeff=parameters.Cv)
+        else:
+            Jvorticity = 0
+
+        if (parameters.Cmod > 0):
+            Jmod = _cost_functions_numpy.calculate_model_cost(
+                winds[0], winds[1], winds[2],
+                parameters.model_weights, parameters.u_model,
+                parameters.v_model, parameters.w_model, coeff=parameters.Cmod)
+        else:
+            Jmod = 0
+
+        if parameters.Cpoint > 0:
+            Jpoint = _cost_functions_numpy.calculate_point_cost(
+                winds[0], winds[1], parameters.x, parameters.y, parameters.z,
+                parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi)
+        else:
+            Jpoint = 0
+    elif parameters.engine == "jax":
+        if not JAX_AVAILABLE:
+            raise ImportError("Jax is needed in order to use the Jax-based PyDDA!")
+
+        winds = np.reshape(winds,
+                           (3, parameters.grid_shape[0], parameters.grid_shape[1],
+                            parameters.grid_shape[2]))
+        # Had to change to float because Jax returns device array (use np.float_())
+        Jvel = _cost_functions_jax.calculate_radial_vel_cost_function(
+            parameters.vrs, parameters.azs, parameters.els,
+            winds[0], winds[1], winds[2], parameters.wts, rmsVr=parameters.rmsVr,
+            weights=parameters.weights, coeff=parameters.Co)
+        # print("apples Jvel", Jvel)
+
+        if (parameters.Cm > 0):
+            # Had to change to float because Jax returns device array (use np.float_())
+            Jmass = _cost_functions_jax.calculate_mass_continuity(winds[0], winds[1], winds[2],
+                                                                 parameters.z, parameters.dx, parameters.dy,
+                                                                 parameters.dz, coeff=parameters.Cm)
+        else:
+            Jmass = 0
+
+        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)
+        else:
+            Jsmooth = 0
+
+        if (parameters.Cb > 0):
+            Jbackground = _cost_functions_jax.calculate_background_cost(
+                winds[0], winds[1], parameters.bg_weights,
+                parameters.u_back, parameters.v_back, parameters.Cb)
+        else:
+            Jbackground = 0
+
+        if (parameters.Cv > 0):
+            # Had to change to float because Jax returns device array (use np.float_())
+            Jvorticity = _cost_functions_jax.calculate_vertical_vorticity_cost(
+                winds[0], winds[1], winds[2], parameters.dx,
+                parameters.dy, parameters.dz, parameters.Ut,
+                parameters.Vt, coeff=parameters.Cv)
+        else:
+            Jvorticity = 0
+
+        if (parameters.Cmod > 0):
+            Jmod = _cost_functions_jax.calculate_model_cost(
+                winds[0], winds[1], winds[2],
+                parameters.model_weights, parameters.u_model,
+                parameters.v_model, parameters.w_model, coeff=parameters.Cmod)
+        else:
+            Jmod = 0
+
+        if parameters.Cpoint > 0:
+            Jpoint = _cost_functions_jax.calculate_point_cost(
+                winds[0], winds[1], parameters.x, parameters.y, parameters.z,
+                parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi)
+        else:
+            Jpoint = 0
+    if(parameters.Nfeval % 10 == 0):
+        print(('Nfeval | Jvel    | Jmass   | Jsmooth |   Jbg   | Jvort   | Jmodel  | Jpoint  |' +
+                ' Max w  '))
+        print(( "{:7d}".format(int(parameters.Nfeval)) + '|' + "{:9.4f}".format(float(Jvel)) + '|' +
+                "{:9.4f}".format(float(Jmass)) + '|' +
+                "{:9.4f}".format(float(Jsmooth)) + '|' +
+                "{:9.4f}".format(float(Jbackground)) + '|' +
+                "{:9.4f}".format(float(Jvorticity)) + '|' +
+                "{:9.4f}".format(float(Jmod)) + '|' +
+                "{:9.4f}".format(float(Jpoint)) + '|' +
+                "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2])))))
+
+    parameters.Nfeval += 1
+    #print("The cost functions print", Jvel + Jmass)
+
+    return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + Jpoint
+
+
+def grad_J(winds, parameters):
+    """
+    Calculates the gradient of the cost function. This typically does not need
+    to be called directly as get_dd_wind_field is a wrapper around this
+    function and :py:func:`pydda.cost_functions.J_function`.
+    In order to add more terms to the cost function,
+    modify this function and :py:func:`pydda.cost_functions.grad_J`.
+
+    Parameters
+    ----------
+    winds: 1-D float array
+        The wind field, flattened to 1-D for f_min
+    parameters: DDParameters
+        The parameters for the cost function evaluation as specified by the
+        :py:func:`pydda.retrieve.DDParameters` class.
+
+    Returns
+    -------
+    grad: 1D float array
+        Gradient vector of cost function
+    """
+    if parameters.engine == "tensorflow":
+        if not TENSORFLOW_AVAILABLE:
+            raise ImportError("Tensorflow 2.5 or greater is needed in order to use TensorFlow-based PyDDA!")
+        winds = tf.reshape(winds,
+                           (3, parameters.grid_shape[0],
+                            parameters.grid_shape[1], parameters.grid_shape[2]))
+        grad = _cost_functions_tensorflow.calculate_grad_radial_vel(
+            parameters.vrs, parameters.els, parameters.azs,
+            winds[0], winds[1], winds[2], parameters.wts, parameters.weights,
+            parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc)
+
+        if(parameters.Cm > 0):
+            grad += _cost_functions_tensorflow.calculate_mass_continuity_gradient(
+                winds[0], winds[1], winds[2],
+                parameters.z,
+                parameters.dx, parameters.dy, parameters.dz,
+                coeff=parameters.Cm, upper_bc=parameters.upper_bc)
+
+        if(parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0):
+            grad += _cost_functions_tensorflow.calculate_smoothness_gradient(
+                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):
+            grad += _cost_functions_tensorflow.calculate_background_gradient(
+                winds[0], winds[1], parameters.bg_weights,
+                parameters.u_back, parameters.v_back, parameters.Cb)
+
+        if(parameters.Cv > 0):
+            grad += _cost_functions_tensorflow.calculate_vertical_vorticity_gradient(
+                winds[0], winds[1], winds[2], parameters.dx,
+                parameters.dy, parameters.dz, parameters.Ut,
+                parameters.Vt, coeff=parameters.Cv, upper_bc=parameters.upper_bc).numpy()
+
+        if(parameters.Cmod > 0):
+            grad += _cost_functions_tensorflow.calculate_model_gradient(
+                winds[0], winds[1], winds[2],
+                parameters.model_weights, parameters.u_model, parameters.v_model,
+                parameters.w_model, coeff=parameters.Cmod)
+
+        if parameters.Cpoint > 0:
+            grad += _cost_functions_tensorflow.calculate_point_gradient(
+                winds[0], winds[1], parameters.x, parameters.y, parameters.z,
+                parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, upper_bc=parameters.upper_bc)
+    elif parameters.engine == "scipy":
+        winds = np.reshape(winds,
+                           (3, parameters.grid_shape[0],
+                            parameters.grid_shape[1], parameters.grid_shape[2]))
+        grad = _cost_functions_numpy.calculate_grad_radial_vel(
+            parameters.vrs, parameters.els, parameters.azs,
+            winds[0], winds[1], winds[2], parameters.wts, parameters.weights,
+            parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc)
+
+        if (parameters.Cm > 0):
+            grad += _cost_functions_numpy.calculate_mass_continuity_gradient(
+                winds[0], winds[1], winds[2],
+                parameters.z,
+                parameters.dx, parameters.dy, parameters.dz,
+                coeff=parameters.Cm, upper_bc=parameters.upper_bc)
+
+        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], 
+                Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc)
+
+        if (parameters.Cb > 0):
+            grad += _cost_functions_numpy.calculate_background_gradient(
+                winds[0], winds[1], parameters.bg_weights,
+                parameters.u_back, parameters.v_back, parameters.Cb)
+
+        if (parameters.Cv > 0):
+            grad += _cost_functions_numpy.calculate_vertical_vorticity_gradient(
+                winds[0], winds[1], winds[2], parameters.dx,
+                parameters.dy, parameters.dz, parameters.Ut,
+                parameters.Vt, coeff=parameters.Cv, upper_bc=parameters.upper_bc).numpy()
+
+        if (parameters.Cmod > 0):
+            grad += _cost_functions_numpy.calculate_model_gradient(
+                winds[0], winds[1], winds[2],
+                parameters.model_weights, parameters.u_model, parameters.v_model,
+                parameters.w_model, coeff=parameters.Cmod)
+
+        if parameters.Cpoint > 0:
+            grad += _cost_functions_numpy.calculate_point_gradient(
+                winds[0], winds[1], parameters.x, parameters.y, parameters.z,
+                parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, upper_bc=parameters.upper_bc)
+    elif parameters.engine == "jax":
+        winds = jnp.reshape(winds,
+                           (3, parameters.grid_shape[0],
+                            parameters.grid_shape[1], parameters.grid_shape[2]))
+        grad = _cost_functions_jax.calculate_grad_radial_vel(
+            parameters.vrs, parameters.els, parameters.azs,
+            winds[0], winds[1], winds[2], parameters.wts, parameters.weights,
+            parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc)
+
+        if (parameters.Cm > 0):
+            grad += _cost_functions_jax.calculate_mass_continuity_gradient(
+                winds[0], winds[1], winds[2],
+                parameters.z,
+                parameters.dx, parameters.dy, parameters.dz,
+                coeff=parameters.Cm, upper_bc=parameters.upper_bc)
+
+        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], 
+                Cx=parameters.Cx, Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc)
+
+        if (parameters.Cb > 0):
+            grad += _cost_functions_jax.calculate_background_gradient(
+                winds[0], winds[1], parameters.bg_weights,
+                parameters.u_back, parameters.v_back, parameters.Cb)
+
+        if (parameters.Cv > 0):
+            grad += _cost_functions_jax.calculate_vertical_vorticity_gradient(
+                winds[0], winds[1], winds[2], parameters.dx,
+                parameters.dy, parameters.dz, parameters.Ut,
+                parameters.Vt, coeff=parameters.Cv, upper_bc=parameters.upper_bc).numpy()
+
+        if (parameters.Cmod > 0):
+            grad += _cost_functions_jax.calculate_model_gradient(
+                winds[0], winds[1], winds[2],
+                parameters.model_weights, parameters.u_model, parameters.v_model,
+                parameters.w_model, coeff=parameters.Cmod)
+
+        if parameters.Cpoint > 0:
+            grad += _cost_functions_jax.calculate_point_gradient(
+                winds[0], winds[1], parameters.x, parameters.y, parameters.z,
+                parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, upper_bc=parameters.upper_bc)
+
+    if(parameters.Nfeval % 10 == 0):
+        print("The gradient of the cost functions is", str(np.linalg.norm(grad, 2)))
+
+    return grad
+
+
+def calculate_fall_speed(grid, refl_field=None, frz=4500.0):
+    """
+    Estimates fall speed based on reflectivity.
+
+    Uses methodology of Mike Biggerstaff and Dan Betten
+
+    Parameters
+    ----------
+    Grid: Py-ART Grid
+        Py-ART Grid containing reflectivity to calculate fall speed from
+    refl_field: str
+        String containing name of reflectivity field. None will automatically
+        determine the name.
+    frz: float
+        Height of freezing level in m
+
+    Returns
+    -------
+    3D float array:
+        Float array of terminal velocities
+
+    """
+    # Parse names of velocity field
+    if refl_field is None:
+        refl_field = pyart.config.get_field_name('reflectivity')
+
+    refl = grid.fields[refl_field]['data']
+    grid_z = grid.point_z['data']
+    term_vel = np.zeros(refl.shape)
+    A = np.zeros(refl.shape)
+    B = np.zeros(refl.shape)
+    rho = np.exp(-grid_z/10000.0)
+    A[np.logical_and(grid_z < frz, refl < 55)] = -2.6
+    B[np.logical_and(grid_z < frz, refl < 55)] = 0.0107
+    A[np.logical_and(grid_z < frz,
+                     np.logical_and(refl >= 55, refl < 60))] = -2.5
+    B[np.logical_and(grid_z < frz,
+                     np.logical_and(refl >= 55, refl < 60))] = 0.013
+    A[np.logical_and(grid_z < frz, refl > 60)] = -3.95
+    B[np.logical_and(grid_z < frz, refl > 60)] = 0.0148
+    A[np.logical_and(grid_z >= frz, refl < 33)] = -0.817
+    B[np.logical_and(grid_z >= frz, refl < 33)] = 0.0063
+    A[np.logical_and(grid_z >= frz,
+                     np.logical_and(refl >= 33, refl < 49))] = -2.5
+    B[np.logical_and(grid_z >= frz,
+                     np.logical_and(refl >= 33, refl < 49))] = 0.013
+    A[np.logical_and(grid_z >= frz, refl > 49)] = -3.95
+    B[np.logical_and(grid_z >= frz, refl > 49)] = 0.0148
+
+    fallspeed = A * np.power(10, refl * B) * np.power(1.2 / rho, 0.4)
+    del A, B, rho
+    return np.ma.masked_invalid(fallspeed)
diff --git a/pydda/retrieval/__init__.py b/pydda/retrieval/__init__.py
index 72193c73..1b7089fc 100644
--- a/pydda/retrieval/__init__.py
+++ b/pydda/retrieval/__init__.py
@@ -16,7 +16,6 @@
     :toctree: generated/
 
     get_dd_wind_field
-    get_dd_wind_field_nested
     get_bca
     DDParameters
 
@@ -25,4 +24,3 @@
 from .wind_retrieve import get_dd_wind_field
 from .wind_retrieve import get_bca
 from .wind_retrieve import DDParameters
-from .nesting import get_dd_wind_field_nested
diff --git a/pydda/retrieval/nesting.py b/pydda/retrieval/nesting.py
deleted file mode 100644
index 8bbe7ac9..00000000
--- a/pydda/retrieval/nesting.py
+++ /dev/null
@@ -1,295 +0,0 @@
-import numpy as np
-import pyart
-import os
-import gc
-import glob
-
-from distributed import Client, wait
-from scipy.interpolate import griddata
-from copy import deepcopy
-from datetime import datetime
-from .wind_retrieve import get_dd_wind_field
-
-
-# Reduces the resolution of a PyART grid
-def _reduce_pyart_grid_res(Grid, skip_factor):
-    Grid2 = deepcopy(Grid)
-    field_dict = {}
-    for field_name in Grid2.fields.keys():
-        field_dict[field_name] = Grid2.fields[field_name].copy()
-        field_dict[field_name]["data"] = Grid2.fields[field_name]["data"][
-            :, ::skip_factor, ::skip_factor]
-
-    x = Grid2.x
-    x["data"] = x["data"][::skip_factor]
-    y = Grid2.y
-    y["data"] = y["data"][::skip_factor]
-    z = Grid2.z
-    z["data"] = z["data"]
-    metadata = Grid2.metadata
-    origin_latitude = Grid2.origin_latitude
-    origin_longitude = Grid2.origin_longitude
-    origin_altitude = Grid2.origin_altitude
-    projection = Grid2.projection
-    radar_latitude = Grid2.radar_latitude
-    radar_longitude = Grid2.radar_longitude
-    radar_altitude = Grid2.radar_altitude
-    radar_time = Grid2.radar_time
-    radar_name = Grid2.radar_name
-    gtime = Grid2.time
-    new_grid = pyart.core.Grid(
-        gtime, field_dict, metadata, origin_latitude, origin_longitude,
-        origin_altitude, x, y, z, projection, radar_latitude, radar_longitude,
-        radar_altitude, radar_time, radar_name)
-    del Grid2
-    return new_grid
-
-
-# Splits a Py-ART Grid
-def _split_pyart_grid(Grid, split_factor, axis=1):
-    grid_splits = []
-    split_field = {}
-    Grid2 = deepcopy(Grid)
-    for field_name in Grid2.fields.keys():
-        if isinstance(Grid2.fields[field_name]["data"], np.ma.MaskedArray):
-            no_mask = Grid2.fields[field_name]["data"].filled(np.nan).copy()
-        else:
-            no_mask = Grid2.fields[field_name]["data"].copy()
-        split_field[field_name] = np.array_split(
-            no_mask, split_factor, axis=axis)
-        if isinstance(Grid2.fields[field_name]["data"], np.ma.MaskedArray):
-            split_field[field_name] = [np.ma.masked_where(
-                np.isnan(arr), arr) for arr in split_field[field_name]]
-    x = Grid2.x
-    y = Grid2.y
-    z = Grid2.z
-    x_split = np.array_split(x["data"], split_factor)
-    y_split = np.array_split(y["data"], split_factor)
-    z_split = np.array_split(z["data"], split_factor)
-    gtime = Grid2.time
-    metadata = Grid2.metadata
-    origin_latitude = Grid2.origin_latitude
-    origin_longitude = Grid2.origin_longitude
-    origin_altitude = Grid2.origin_altitude
-    projection = Grid2.projection
-    radar_latitude = Grid2.radar_latitude
-    radar_longitude = Grid2.radar_longitude
-    radar_altitude = Grid2.radar_altitude
-    radar_time = Grid2.radar_time
-    radar_name = Grid2.radar_name
-    for i in range(split_factor):
-        grid_dic = {}
-
-        for field_name in Grid2.fields.keys():
-            grid_dic[field_name] = Grid2.fields[field_name].copy()
-            grid_dic[field_name]["data"] = split_field[field_name][i]
-        x_dic = x.copy()
-        y_dic = y.copy()
-        z_dic = z.copy()
-        if(axis == 1):
-            y_dic["data"] = y_split[i]
-        elif(axis == 2):
-            x_dic["data"] = x_split[i]
-        elif(axis == 0):
-            z_dic["data"] = z_split[i]
-
-        new_grid = pyart.core.Grid(
-            gtime, grid_dic, metadata, origin_latitude, origin_longitude,
-            origin_altitude, x_dic, y_dic, z_dic, projection, radar_latitude,
-            radar_longitude, radar_altitude, radar_time, radar_name)
-        grid_splits.append(new_grid)
-
-    return grid_splits
-
-
-# Concatenates Py-ART Grids
-def _concatenate_pyart_grids(grid_list, axis=1):
-    new_grid = deepcopy(grid_list[0])
-    for field_name in new_grid.fields.keys():
-        new_grid.fields[field_name]["data"] = np.ma.concatenate(
-            [x.fields[field_name]["data"] for x in grid_list], axis=axis)
-    if(axis == 2):
-        new_grid.x["data"] = np.ma.concatenate(
-            [x.x["data"] for x in grid_list])
-        new_grid.nx = np.sum([x.nx for x in grid_list])
-    elif(axis == 1):
-        new_grid.y["data"] = np.ma.concatenate(
-            [x.y["data"] for x in grid_list])
-        new_grid.ny = np.sum([x.ny for x in grid_list])
-    elif(axis == 0):
-        new_grid.z["data"] = np.ma.concatenate(
-            [x.z["data"] for x in grid_list])
-        new_grid.nz = np.sum([x.nz for x in grid_list])
-    return new_grid
-
-
-# Procedure: 1. Do first pass of retrieval on reduced resolution grid
-# 2. Then, we use the reduced resolution retrieval as an input to the
-# high resolution retrieval in each region
-# Finally, we check for continuity at the boundaries
-def get_dd_wind_field_nested(grid_list, u_init, v_init, w_init, client,
-                             reduction_factor=2, num_splits=2, **kwargs):
-    """
-    This function performs a wind retrieval using a nested domain.
-    This is useful for grids that are larger than about 500 by 500
-    by 40 points, since the use of larger grids on a single machine
-    will exceed memory limitations.
-
-    This procedure relies on a dask distributed cluster to be set up.
-    The retrieval is first performed at a resolution that is coarser
-    than the analysis grid by reduction_factor. This provides the
-    initial state for the nested loop.
-
-    The domain is split into num_splits**2 sub-domains for the nested
-    retrieval step, and each nested retrieval is mapped onto a distributed
-    worker for parallel processing. If NumPy and SciPy are already set up to
-    use parallel numerical analysis libraries, it is recommended that a single
-    machine be dedicated to each nest rather than a single core
-    for best performance.
-
-    Parameters
-    ==========
-    grid_list: list
-       A list of Py-ART grids for each radar to use in the retrieval.
-    u_init: 3D NumPy array
-       The initial guess of the zonal wind field. This has to be in the same
-       shape as the analysis grid.
-    v_init: 3D NumPy array
-       The initial guess of the meridional wind field. This has to be in the
-       same shape as the analysis grid.
-    w_init: 3D NumPy array
-       The initial guess of the vertical wind field. This has to be in the same
-       shape as the analysis grid.
-    client: dask distributed Client
-       The distributed Client that is linked to a distributed cluster. The
-       :cluster must be running before get_dd_wind_field_nested is called.
-       The retrieval on each nest will be mapped onto each worker. Since
-       the optimization loop already takes advantage of parallelism, it's
-       best to allow at least 16 cores per one worker.
-    reduction_factor: int
-       How much to reduce the factor of the analysis grid by when doing the
-       initial retrieval on the entire grid.
-    num_splits: int
-       The number of splits to make through each axis when doing the nesting.
-
-    **kwargs: dict
-        This function will take the same keyword arguments as
-        get_dd_wind_field, as these arguments are passed into each call of
-        get_dd_wind_field. See get_dd_wind_field for more information on the
-    """
-    # First, we do retrieval on whole grid with fraction of resolution
-    grid_lo_res_list = [_reduce_pyart_grid_res(G, reduction_factor)
-                        for G in grid_list]
-
-    first_pass = get_dd_wind_field(
-        grid_lo_res_list, u_init[::, ::reduction_factor, ::reduction_factor],
-        v_init[::, ::reduction_factor, ::reduction_factor],
-        w_init[::, ::reduction_factor, ::reduction_factor], **kwargs)
-
-    # Take the first pass field and regrid to analysis field
-    reduced_x = first_pass[0].point_x["data"].flatten()
-    reduced_y = first_pass[0].point_y["data"].flatten()
-    reduced_z = first_pass[0].point_z["data"].flatten()
-    x = grid_list[0].point_x["data"].flatten()
-    y = grid_list[0].point_y["data"].flatten()
-    z = grid_list[0].point_z["data"].flatten()
-    u_init_new = griddata((reduced_z, reduced_y, reduced_x),
-                          first_pass[0].fields["u"]["data"].flatten(),
-                          (z, y, x), method='nearest')
-    v_init_new = griddata((reduced_z, reduced_y, reduced_x),
-                          first_pass[0].fields["v"]["data"].flatten(),
-                          (z, y, x), method='nearest')
-    w_init_new = griddata((reduced_z, reduced_y, reduced_x),
-                          first_pass[0].fields["w"]["data"].flatten(),
-                          (z, y, x), method='nearest')
-    u_init_new = np.reshape(u_init_new, u_init.shape)
-    v_init_new = np.reshape(v_init_new, v_init.shape)
-    w_init_new = np.reshape(w_init_new, w_init.shape)
-
-    # Finally, split the analysis into num_splits**2 pieces and save
-    # as temporary files
-    tempfile_name_base = datetime.now().strftime('%y%m%d.%H%M%S')
-    tiny_grids = []
-    k = 0
-    for G in grid_list:
-        cur_list = []
-        split_grids_x = _split_pyart_grid(G, num_splits, axis=2)
-        i = 0
-        for sgrid in split_grids_x:
-            g_list = _split_pyart_grid(sgrid, num_splits)
-            grid_fns = []
-            j = 0
-            for g in g_list:
-                fn = (tempfile_name_base + str(k) + '.' +
-                      str(i) + '.' + str(j) + '.nc')
-                pyart.io.write_grid((tempfile_name_base + str(k) +
-                                     '.' + str(i) + '.' + str(j) + '.nc'), g)
-                j = j + 1
-                grid_fns.append(fn)
-            cur_list.append(grid_fns)
-            i = i + 1
-        del split_grids_x, g_list
-
-        k = k + 1
-        tiny_grids.append(cur_list)
-
-    # Temporarily save the tiny grids and free up memory...we want to
-    # load these when we are running it on the cluster
-
-    u_init_split_x = np.array_split(u_init_new, num_splits, axis=2)
-    u_init_split = [np.array_split(ux, num_splits, axis=1)
-                    for ux in u_init_split_x]
-    w_init_split_x = np.array_split(w_init_new, num_splits, axis=2)
-    w_init_split = [np.array_split(wx, num_splits, axis=1)
-                    for wx in w_init_split_x]
-    v_init_split_x = np.array_split(v_init_new, num_splits, axis=2)
-    v_init_split = [np.array_split(vx, num_splits, axis=1)
-                    for vx in v_init_split_x]
-
-    # Clear out unneeded variables (do not need lo-res grids in memory anymore)
-    del u_init_split_x, w_init_split_x, v_init_split_x
-    del first_pass, reduced_x, reduced_y, reduced_z, x, y, z, grid_lo_res_list
-    gc.collect()
-
-    # Serial just for testing, need to use dask in future
-    tiny_retrieval = []
-
-    def do_tiny_retrieval(i, j):
-        tgrids = [pyart.io.read_grid(tiny_grids[k][i][j])
-                  for k in range(len(grid_list))]
-        new_grids = get_dd_wind_field(
-            tgrids, u_init_split[i][j], v_init_split[i][j],
-            w_init_split[i][j], **kwargs)
-        del tgrids
-        gc.collect()
-        return new_grids
-
-    futures_array = []
-    for i in range(num_splits):
-        for j in range(num_splits):
-            futures_array.append(client.submit(do_tiny_retrieval, i, j))
-
-    print("Waiting for nested grid to be retrieved...")
-    wait(futures_array)
-    tiny_retrieval2 = client.gather(futures_array)
-    tiny_retrieval = []
-
-    for i in range(num_splits):
-        new_grid_list = []
-
-        for j in range(len(grid_list)):
-            new_grid_list.append(_concatenate_pyart_grids(
-                [tiny_retrieval2[k+i*num_splits][j]
-                 for k in range(0, num_splits)], axis=1))
-        tiny_retrieval.append(new_grid_list)
-
-    new_grid_list = []
-    for i in range(len(grid_list)):
-        new_grid_list.append(_concatenate_pyart_grids(
-            [tiny_retrieval[k][i] for k in range(num_splits)], axis=2))
-
-    tempfile_list = glob.glob(tempfile_name_base + "*")
-    for fn in tempfile_list:
-        os.remove(fn)
-
-    return new_grid_list
diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py
index c1d4d194..c35b28b5 100644
--- a/pydda/retrieval/wind_retrieve.py
+++ b/pydda/retrieval/wind_retrieve.py
@@ -1,654 +1,1131 @@
-"""
-Created on Mon Aug  7 09:17:40 2017
-
-@author: rjackson
-"""
-
-import pyart
-import numpy as np
-import time
-import cartopy.crs as ccrs
-import math
-
-from .. import cost_functions
-from ..cost_functions import J_function, grad_J
-from scipy.optimize import fmin_l_bfgs_b
-from scipy.interpolate import interp1d
-from scipy.signal import savgol_filter
-from matplotlib import pyplot as plt
-from copy import deepcopy
-from .angles import add_azimuth_as_field, add_elevation_as_field
-
-
-class DDParameters(object):
-    """
-    This is a helper class for inserting more arguments into the :func:`pydda.cost_functions.J_function` and
-    :func:`pydda.cost_functions.grad_J` function. Since these cost functions take numerous parameters, this class
-    will store the needed parameters as one positional argument for easier readability of the code.
-
-    In addition, class members can be added here so that those contributing more constraints to the variational
-    framework can add any parameters they may need.
-
-    Attributes
-    ----------
-    vrs: List of float arrays
-        List of radial velocities from each radar
-    azs: List of float arrays
-        List of azimuths from each radar
-    els: List of float arrays
-        List of elevations from each radar
-    wts: List of float arrays
-        Float array containing fall speed from radar.
-    u_back: 1D float array (number of vertical levels)
-        Background u wind
-    v_back: 1D float array (number of vertical levels)
-        Background u wind
-    u_model: list of 3D float arrays
-        U from each model integrated into the retrieval
-    v_model: list of 3D float arrays
-        V from each model integrated into the retrieval
-    w_model:
-        W from each model integrated into the retrieval
-    Co: float
-        Weighting coefficient for data constraint.
-    Cm: float
-        Weighting coefficient for mass continuity constraint.
-    Cx: float
-        Smoothing coefficient for x-direction
-    Cy: float
-        Smoothing coefficient for y-direction
-    Cz: float
-        Smoothing coefficient for z-direction
-    Cb: float
-        Coefficient for sounding constraint
-    Cv: float
-        Weight for cost function related to vertical vorticity equation.
-    Cmod: float
-        Coefficient for model constraint
-    Cpoint: float
-        Coefficient for point constraint
-    Ut: float
-        Prescribed storm motion. This is only needed if Cv is not zero.
-    Vt: float
-        Prescribed storm motion. This is only needed if Cv is not zero.
-    grid_shape:
-        Shape of wind grid
-    dx:
-        Spacing of grid in x direction
-    dy:
-        Spacing of grid in y direction
-    dz:
-        Spacing of grid in z direction
-    x:
-        E-W grid levels in m
-    y:
-        N-S grid levels in m
-    z:
-        Grid vertical levels in m
-    rmsVr: float
-        The sum of squares of velocity/num_points. Use for normalization
-        of data weighting coefficient
-    weights: n_radars by z_bins by y_bins x x_bins float array
-        Data weights for each pair of radars
-    bg_weights: z_bins by y_bins x x_bins float array
-        Data weights for sounding constraint
-    model_weights: n_models by z_bins by y_bins by x_bins float array
-        Data weights for each model.
-    point_list: list or None
-        point_list: list of dicts
-        List of point constraints. Each member is a dict with keys of "u", "v",
-        to correspond to each component of the wind field and "x", "y", "z"
-        to correspond to the location of the point observation in the Grid's
-        Cartesian coordinates.
-    roi: float
-        The radius of influence of each point observation in m.
-    upper_bc: bool
-        True to enforce w=0 at top of domain (impermeability condition),
-        False to not enforce impermeability at top of domain
-    """
-    def __init__(self):
-        self.Ut = np.nan
-        self.Vt = np.nan
-        self.rmsVr = np.nan
-        self.grid_shape = None
-        self.Cmod = np.nan
-        self.Cpoint = np.nan
-        self.u_back = None
-        self.v_back = None
-        self.wts = []
-        self.vrs = []
-        self.azs = []
-        self.els = []
-        self.weights = []
-        self.bg_weights = []
-        self.model_weights = []
-        self.u_model = []
-        self.v_model = []
-        self.w_model = []
-        self.x = None
-        self.y = None
-        self.z = None
-        self.dx = np.nan
-        self.dy = np.nan
-        self.dz = np.nan
-        self.Co = 1.0
-        self.Cm = 1500.0
-        self.Cx = 0.0
-        self.Cy = 0.0
-        self.Cz = 0.0
-        self.Cb = 0.0
-        self.Cv = 0.0
-        self.Cmod = 0.0
-        self.Cpoint = 0.0
-        self.Ut = 0.0
-        self.Vt = 0.0
-        self.upper_bc = True
-        self.roi = 1000.0
-        self.frz = 4500.0
-        self.point_list = []
-
-
-def get_dd_wind_field(Grids, u_init, v_init, w_init, points=None, vel_name=None,
-                      refl_field=None, u_back=None, v_back=None, z_back=None,
-                      frz=4500.0, Co=1.0, Cm=1500.0, Cx=0.0,
-                      Cy=0.0, Cz=0.0, Cb=0.0, Cv=0.0, Cmod=0.0, Cpoint=0.0,
-                      Ut=None, Vt=None, filt_iterations=2,
-                      mask_outside_opt=False, weights_obs=None,
-                      weights_model=None, weights_bg=None,
-                      max_iterations=200, mask_w_outside_opt=True,
-                      filter_window=9, filter_order=3, min_bca=30.0,
-                      max_bca=150.0, upper_bc=True, model_fields=None,
-                      output_cost_functions=True, roi=1000.0):
-    """
-    This function takes in a list of Py-ART Grid objects and derives a
-    wind field. Every Py-ART Grid in Grids must have the same grid
-    specification. 
-
-    In order for the model data constraint to be used,
-    the model data must be added as a field to at least one of the
-    grids in Grids. This involves interpolating the model data to the
-    Grids' coordinates. There are helper functions for this for WRF
-    and HRRR data in :py:func:`pydda.constraints`:
-
-    :py:func:`make_constraint_from_wrf`
-
-    :py:func:`add_hrrr_constraint_to_grid`
-
-    Parameters
-    ==========
-
-    Grids: list of Py-ART Grids
-        The list of Py-ART grids to take in corresponding to each radar.
-        All grids must have the same shape, x coordinates, y coordinates
-        and z coordinates.
-    u_init: 3D ndarray
-        The intial guess for the zonal wind field, input as a 3D array
-        with the same shape as the fields in Grids.
-    v_init: 3D ndarray
-        The intial guess for the meridional wind field, input as a 3D array
-        with the same shape as the fields in Grids.
-    w_init: 3D ndarray
-        The intial guess for the vertical wind field, input as a 3D array
-        with the same shape as the fields in Grids.
-    points: None or list of dicts
-        Point observations as returned by :func:`pydda.constraints.get_iem_obs`. Set
-        to None to disable.
-    vel_name: string
-        Name of radial velocity field. Setting to None will have PyDDA attempt
-        to automatically detect the velocity field name.
-    refl_field: string
-        Name of reflectivity field. Setting to None will have PyDDA attempt 
-        to automatically detect the reflectivity field name.
-    u_back: 1D array
-        Background zonal wind field from a sounding as a function of height.
-        This should be given in the sounding's vertical coordinates.
-    v_back: 1D array
-        Background meridional wind field from a sounding as a function of
-        height. This should be given in the sounding's vertical coordinates.
-    z_back: 1D array
-        Heights corresponding to background wind field levels in meters. This
-        is given in the sounding's original coordinates.
-    frz: float
-        Freezing level used for fall speed calculation in meters.
-    Co: float
-        Weight for cost function related to observed radial velocities.
-    Cm: float
-        Weight for cost function related to the mass continuity equation.
-    Cx: float
-        Weight for cost function related to smoothness in x direction
-    Cy: float
-        Weight for cost function related to smoothness in y direction
-    Cz: float
-        Weight for cost function related to smoothness in z direction
-    Cv: float
-        Weight for cost function related to vertical vorticity equation.
-    Cmod: float
-        Weight for cost function related to custom constraints.
-    Cpoint: float
-        Weight for cost function related to point observations.
-    weights_obs: list of floating point arrays or None
-        List of weights for each point in grid from each radar in Grids.
-        Set to None to let PyDDA determine this automatically.
-    weights_model: list of floating point arrays or None
-        List of weights for each point in grid from each custom field in
-        model_fields. Set to None to let PyDDA determine this automatically.
-    weights_bg: list of floating point arrays or None
-        List of weights for each point in grid from the sounding. Set to None
-        to let PyDDA determine this automatically.
-    Ut: float
-        Prescribed storm motion in zonal direction.
-        This is only needed if Cv is not zero.
-    Vt: float
-        Prescribed storm motion in meridional direction.
-        This is only needed if Cv is not zero.
-    filt_iterations: int
-        If this is greater than 0, PyDDA will run a low pass filter on
-        the retrieved wind field and then do the optimization step for
-        filt_iterations iterations. Set to 0 to disable the low pass filter.
-    mask_outside_opt: bool
-        If set to true, wind values outside the multiple doppler lobes will
-        be masked, i.e. if less than 2 radars provide coverage for a given
-        point.
-    max_iterations: int
-        The maximum number of iterations to run the optimization loop for.
-    mask_w_outside_opt: bool
-        If set to true, vertical winds outside the multiple doppler lobes will
-        be masked, i.e. if less than 2 radars provide coverage for a given
-        point.
-    filter_window: int
-        Window size to use for the low pass filter. A larger window will
-        increase the number of points factored into the polynomial fit for
-        the filter, and hence will increase the smoothness.
-    filter_order: int
-        The order of the polynomial to use for the low pass filter. Higher
-        order polynomials allow for the retention of smaller scale features
-        but may also not remove enough noise.
-    min_bca: float
-        Minimum beam crossing angle in degrees between two radars. 30.0 is the
-        typical value used in many publications.
-    max_bca: float
-        Minimum beam crossing angle in degrees between two radars. 150.0 is the
-        typical value used in many publications.
-    upper_bc: bool
-        Set this to true to enforce w = 0 at the top of the atmosphere. This is
-        commonly called the impermeability condition.
-    model_fields: list of strings
-        The list of fields in the first grid in Grids that contain the custom
-        data interpolated to the Grid's grid specification. Helper functions
-        to create such gridded fields for HRRR and NetCDF WRF data exist
-        in ::pydda.constraints::. PyDDA will look for fields named U_(model
-        field name), V_(model field name), and W_(model field name). For
-        example, if you have U_hrrr, V_hrrr, and W_hrrr, then specify ["hrrr"]
-        into model_fields.
-    output_cost_functions: bool
-        Set to True to output the value of each cost function every
-        10 iterations.
-    roi: float
-        Radius of influence for the point observations. The point observation will
-        not hold any weight outside this radius.
-
-    Returns
-    =======
-    new_grid_list: list
-        A list of Py-ART grids containing the derived wind fields. These fields
-        are displayable by the visualization module.
-    """
-
-    # We have to have a prescribed storm motion for vorticity constraint
-    if(Ut is None or Vt is None):
-        if(Cv != 0.0):
-            raise ValueError(('Ut and Vt cannot be None if vertical ' +
-                              'vorticity constraint is enabled!'))
-
-    if not isinstance(Grids, list):
-        raise ValueError('Grids has to be a list!')
-
-    parameters = DDParameters()
-    parameters.Ut = Ut
-    parameters.Vt = Vt
-    
-    # Ensure that all Grids are on the same coordinate system
-    prev_grid = Grids[0]
-    for g in Grids:
-        if not np.allclose(
-            g.x['data'], prev_grid.x['data'], atol=10):
-            raise ValueError('Grids do not have equal x coordinates!')
-
-        if not np.allclose(
-            g.y['data'], prev_grid.y['data'], atol=10):
-            raise ValueError('Grids do not have equal y coordinates!')
-
-        if not np.allclose(
-            g.z['data'], prev_grid.z['data'], atol=10):
-            raise ValueError('Grids do not have equal z coordinates!')
-
-        if not g.origin_latitude['data'] == prev_grid.origin_latitude['data']:
-            raise ValueError(("Grids have unequal origin lat/lons!"))
-
-        prev_grid = g
-
-    # Disable background constraint if none provided
-    if(u_back is None or v_back is None):
-        parameters.u_back = np.zeros(u_init.shape[0])
-        parameters.v_back = np.zeros(v_init.shape[0])
-    else:
-        # Interpolate sounding to radar grid
-        print('Interpolating sounding to radar grid')
-        u_interp = interp1d(z_back, u_back, bounds_error=False)
-        v_interp = interp1d(z_back, v_back, bounds_error=False)
-        parameters.u_back = u_interp(Grids[0].z['data'])
-        parameters.v_back = v_interp(Grids[0].z['data'])
-        print('Interpolated U field:')
-        print(parameters["u_back"])
-        print('Interpolated V field:')
-        print(parameters["v_back"])
-        print('Grid levels:')
-        print(Grids[0].z['data'])
-
-    # Parse names of velocity field
-    if refl_field is None:
-        refl_field = pyart.config.get_field_name('reflectivity')
-
-    # Parse names of velocity field
-    if vel_name is None:
-        vel_name = pyart.config.get_field_name('corrected_velocity')
-    winds = np.stack([u_init, v_init, w_init])
-
-
-    # Set up wind fields and weights from each radar
-    parameters.weights = np.zeros(
-        (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]))
-
-    parameters.bg_weights = np.zeros(v_init.shape)
-    if(model_fields is not None):
-        parameters.model_weights = np.ones(
-            (len(model_fields), u_init.shape[0], u_init.shape[1],
-             u_init.shape[2]))
-    else:
-        parameters.model_weights = np.zeros(
-            (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]))
-
-    if(model_fields is None):
-        if(Cmod != 0.0):
-            raise ValueError(
-                 'Cmod must be zero if model fields are not specified!')
-
-    bca = np.zeros(
-        (len(Grids), len(Grids), u_init.shape[1], u_init.shape[2]))
-    sum_Vr = np.zeros(len(Grids))
-
-    for i in range(len(Grids)):
-        parameters.wts.append(cost_functions.calculate_fall_speed(Grids[i],
-                                                       refl_field=refl_field, frz=frz))
-        add_azimuth_as_field(Grids[i], dz_name=refl_field)
-        add_elevation_as_field(Grids[i], dz_name=refl_field)
-        parameters.vrs.append(Grids[i].fields[vel_name]['data'])
-        parameters.azs.append(Grids[i].fields['AZ']['data']*np.pi/180)
-        parameters.els.append(Grids[i].fields['EL']['data']*np.pi/180)
-
-    if(len(Grids) > 1):
-        for i in range(len(Grids)):
-            for j in range(i+1, len(Grids)):
-                print(("Calculating weights for radars " + str(i) +
-                       " and " + str(j)))
-                bca[i, j] = get_bca(Grids[i].radar_longitude['data'],
-                                    Grids[i].radar_latitude['data'],
-                                    Grids[j].radar_longitude['data'],
-                                    Grids[j].radar_latitude['data'],
-                                    Grids[i].point_x['data'][0],
-                                    Grids[i].point_y['data'][0],
-                                    Grids[i].get_projparams())
-
-                for k in range(parameters.vrs[i].shape[0]):
-                    if(weights_obs is None):
-                        cur_array = parameters.weights[i, k]
-                        cur_array[np.logical_and(
-                            ~parameters.vrs[i][k].mask,
-                            np.logical_and(
-                                bca[i, j] >= math.radians(min_bca),
-                                bca[i, j] <= math.radians(max_bca)))] += 1
-                        parameters.weights[i, k] = cur_array
-                    else:
-                        parameters.weights[i, k] = weights_obs[i][k, :, :]
-
-                    if(weights_obs is None):
-                        cur_array = parameters.weights[j, k]
-                        cur_array[np.logical_and(
-                            ~parameters.vrs[j][k].mask,
-                            np.logical_and(
-                                bca[i, j] >= math.radians(min_bca),
-                                bca[i, j] <= math.radians(max_bca)))] += 1
-                        parameters.weights[j, k] = cur_array
-                    else:
-                        parameters.weights[j, k] = weights_obs[j][k, :, :]
-
-                    if(weights_bg is None):
-                        cur_array = parameters.bg_weights[k]
-                        cur_array[np.logical_or(
-                            bca[i, j] >= math.radians(min_bca),
-                            bca[i, j] <= math.radians(max_bca))] = 1
-                        cur_array[parameters.vrs[i][k].mask] = 0
-                        parameters.bg_weights[i] = cur_array
-                    else:
-                        parameters.bg_weights[i] = weights_bg[i]
-
-        print("Calculating weights for models...")
-        coverage_grade = parameters.weights.sum(axis=0)
-        coverage_grade = coverage_grade/coverage_grade.max()
-
-        # Weigh in model input more when we have no coverage
-        # Model only weighs 1/(# of grids + 1) when there is full
-        # Coverage
-        if(model_fields is not None):
-            if(weights_model is None):
-                for i in range(len(model_fields)):
-                    parameters.model_weights[i] = 1 - (coverage_grade/(len(Grids)+1))
-            else:
-                for i in range(len(model_fields)):
-                    parameters.model_weights[i] = weights_model[i]
-    else:
-        if weights_obs is None:
-            parameters.weights[0] = np.where(~parameters.vrs[0].mask, 1, 0)
-        else:
-            parameters.weights[0] = weights_obs[0]
-
-        if weights_bg is None:
-            parameters.bg_weights = np.where(~parameters.vrs[0].mask, 0, 1)
-        else:
-            parameters.bg_weights = weights_bg
-
-
-    parameters.weights[parameters.weights > 0] = 1
-    parameters.bg_weights[parameters.bg_weights > 0] = 1
-    sum_Vr = np.nansum(np.square(parameters.vrs * parameters.weights))
-    parameters.rmsVr = np.sqrt(np.nansum(sum_Vr) / np.nansum(parameters.weights))
-
-    del bca
-    parameters.grid_shape = u_init.shape
-    # Parse names of velocity field
-
-    winds = winds.flatten()
-
-    print("Starting solver ")
-    parameters.dx = np.diff(Grids[0].x['data'], axis=0)[0]
-    parameters.dy = np.diff(Grids[0].y['data'], axis=0)[0]
-    parameters.dz = np.diff(Grids[0].z['data'], axis=0)[0]
-    print('rmsVR = ' + str(parameters.rmsVr))
-    print('Total points: %d' % parameters.weights.sum())
-    parameters.z = Grids[0].point_z['data']
-    parameters.x = Grids[0].point_x['data']
-    parameters.y = Grids[0].point_y['data']
-    bt = time.time()
-
-    # First pass - no filter
-    wprevmax = 99
-    wcurrmax = w_init.max()
-    iterations = 0
-    bounds = [(-x, x) for x in 100*np.ones(winds.shape)]
-
-    if(model_fields is not None):
-        for the_field in model_fields:
-            u_field = ("U_" + the_field)
-            v_field = ("V_" + the_field)
-            w_field = ("W_" + the_field)
-            parameters.u_model.append(Grids[0].fields[u_field]["data"])
-            parameters.v_model.append(Grids[0].fields[v_field]["data"])
-            parameters.w_model.append(Grids[0].fields[w_field]["data"])
-
-    parameters.Co = Co
-    parameters.Cm = Cm
-    parameters.Cx = Cx
-    parameters.Cy = Cy
-    parameters.Cz = Cz
-    parameters.Cb = Cb
-    parameters.Cv = Cv
-    parameters.Cmod = Cmod
-    parameters.Cpoint = Cpoint
-    parameters.roi = roi
-    parameters.upper_bc = upper_bc
-    parameters.points = points
-    parameters.point_list = points
-
-    while(iterations < max_iterations and
-          (abs(wprevmax-wcurrmax) > 0.02)):
-        wprevmax = wcurrmax
-        parameters.print_out = False
-        winds = fmin_l_bfgs_b(J_function, winds, args=(parameters,),
-                              maxiter=10, pgtol=1e-3, bounds=bounds,
-                              fprime=grad_J, disp=0, iprint=-1)
-        parameters.print_out = True
-        if output_cost_functions is True:
-            J_function(winds[0], parameters)
-            grad_J(winds[0], parameters)
-        winds = np.reshape(
-            winds[0], (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2]))
-        iterations = iterations+10
-        print('Iterations before filter: ' + str(iterations))
-        wcurrmax = winds[2].max()
-        winds = np.stack([winds[0], winds[1], winds[2]])
-        winds = winds.flatten()
-
-    if filt_iterations > 0:
-        print('Applying low pass filter to wind field...')
-        winds = np.reshape(winds, (3, parameters.grid_shape[0], parameters.grid_shape[1],
-                                   parameters.grid_shape[2]))
-        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=0)
-        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=1)
-        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=2)
-        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=0)
-        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=1)
-        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=2)
-        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=0)
-        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=1)
-        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=2)
-        winds = np.stack([winds[0], winds[1], winds[2]])
-        winds = winds.flatten()
-        iterations = 0
-        while(iterations < filt_iterations):
-            winds = fmin_l_bfgs_b(
-                J_function, winds, args=(parameters,),
-                maxiter=10, pgtol=1e-3, bounds=bounds,
-                fprime=grad_J, disp=0, iprint=-1)
-            parameters.print_out = False
-            winds = np.reshape(
-                winds[0], (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2]))
-            iterations = iterations+1
-            print('Iterations after filter: ' + str(iterations))
-            winds = np.stack([winds[0], winds[1], winds[2]])
-            winds = winds.flatten()
-
-    print("Done! Time = " + "{:2.1f}".format(time.time() - bt))
-
-    # First pass - no filter
-    the_winds = np.reshape(
-        winds, (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2]))
-    u = the_winds[0]
-    v = the_winds[1]
-    w = the_winds[2]
-    where_mask = np.sum(parameters.weights, axis=0) + \
-                 np.sum(parameters.model_weights, axis=0)
-    
-    u = np.ma.array(u)
-    w = np.ma.array(w)
-    v = np.ma.array(v)
-
-    if mask_outside_opt is True:
-        u = np.ma.masked_where(where_mask < 1, u)
-        v = np.ma.masked_where(where_mask < 1, v)
-        w = np.ma.masked_where(where_mask < 1, w)
-
-    if mask_w_outside_opt is True:
-        w = np.ma.masked_where(where_mask < 1, w)
-
-    u_field = deepcopy(Grids[0].fields[vel_name])
-    u_field['data'] = u
-    u_field['standard_name'] = 'u_wind'
-    u_field['long_name'] = 'meridional component of wind velocity'
-    u_field['min_bca'] = min_bca
-    u_field['max_bca'] = max_bca
-    v_field = deepcopy(Grids[0].fields[vel_name])
-    v_field['data'] = v
-    v_field['standard_name'] = 'v_wind'
-    v_field['long_name'] = 'zonal component of wind velocity'
-    v_field['min_bca'] = min_bca
-    v_field['max_bca'] = max_bca
-    w_field = deepcopy(Grids[0].fields[vel_name])
-    w_field['data'] = w
-    w_field['standard_name'] = 'w_wind'
-    w_field['long_name'] = 'vertical component of wind velocity'
-    w_field['min_bca'] = min_bca
-    w_field['max_bca'] = max_bca
-
-    new_grid_list = []
-
-    for grid in Grids:
-        temp_grid = deepcopy(grid)
-        temp_grid.add_field('u', u_field, replace_existing=True)
-        temp_grid.add_field('v', v_field, replace_existing=True)
-        temp_grid.add_field('w', w_field, replace_existing=True)
-        new_grid_list.append(temp_grid)
-
-    return new_grid_list
-
-
-def get_bca(rad1_lon, rad1_lat, rad2_lon, rad2_lat, x, y, projparams):
-    """
-    This function gets the beam crossing angle between two lat/lon pairs.
-
-    Parameters
-    ==========
-    rad1_lon: float
-        The longitude of the first radar.
-    rad1_lat: float
-        The latitude of the first radar.
-    rad2_lon: float
-        The longitude of the second radar.
-    rad2_lat: float
-        The latitude of the second radar.
-    x: nD float array
-        The Cartesian x coordinates of the grid
-    y: nD float array
-        The Cartesian y corrdinates of the grid
-    projparams: Py-ART projparams
-        The projection parameters of the Grid
-
-    Returns
-    =======
-    bca: nD float array
-        The beam crossing angle between the two radars in radians.
-
-    """
-
-    rad1 = pyart.core.geographic_to_cartesian(rad1_lon, rad1_lat, projparams)
-    rad2 = pyart.core.geographic_to_cartesian(rad2_lon, rad2_lat, projparams)
-    # Create grid with Radar 1 in center
-
-    x = x-rad1[0]
-    y = y-rad1[1]
-    rad2 = np.array(rad2) - np.array(rad1)
-    a = np.sqrt(np.multiply(x, x) + np.multiply(y, y))
-    b = np.sqrt(pow(x-rad2[0], 2) + pow(y-rad2[1], 2))
-    c = np.sqrt(rad2[0]*rad2[0] + rad2[1]*rad2[1])
-    theta_1 = np.arccos(x/a)
-    theta_2 = np.arccos((x-rad2[1])/b)
-    return np.arccos((a*a+b*b-c*c)/(2*a*b))
+"""
+Created on Mon Aug  7 09:17:40 2017
+
+@author: rjackson
+"""
+
+import pyart
+import numpy as np
+import time
+import cartopy.crs as ccrs
+import math
+import sys
+
+from scipy.interpolate import interp1d
+from scipy.optimize import fmin_l_bfgs_b
+from scipy.signal import savgol_filter
+try:
+    import tensorflow_probability as tfp
+    import tensorflow as tf
+    TENSORFLOW_AVAILABLE = True
+except (ImportError, AttributeError):
+    TENSORFLOW_AVAILABLE = False
+
+try:
+    import jax.numpy as jnp
+    JAX_AVAILABLE = True
+except ImportError:
+    JAX_AVAILABLE = False
+
+#imports changed to local import path to run on computer
+from ..cost_functions import J_function, grad_J, calculate_fall_speed
+from copy import deepcopy
+from .angles import add_azimuth_as_field, add_elevation_as_field
+_wprevmax = np.empty(0)
+_wcurrmax = np.empty(0)
+iterations = 0
+
+
+class DDParameters(object):
+    """
+    This is a helper class for inserting more arguments into the :func:`pydda.cost_functions.J_function` and
+    :func:`pydda.cost_functions.grad_J` function. Since these cost functions take numerous parameters, this class
+    will store the needed parameters as one positional argument for easier readability of the code.
+
+    In addition, class members can be added here so that those contributing more constraints to the variational
+    framework can add any parameters they may need.
+
+    Attributes
+    ----------
+    vrs: List of float arrays
+        List of radial velocities from each radar
+    azs: List of float arrays
+        List of azimuths from each radar
+    els: List of float arrays
+        List of elevations from each radar
+    wts: List of float arrays
+        Float array containing fall speed from radar.
+    u_back: 1D float array (number of vertical levels)
+        Background u wind
+    v_back: 1D float array (number of vertical levels)
+        Background u wind
+    u_model: list of 3D float arrays
+        U from each model integrated into the retrieval
+    v_model: list of 3D float arrays
+        V from each model integrated into the retrieval
+    w_model:
+        W from each model integrated into the retrieval
+    Co: float
+        Weighting coefficient for data constraint.
+    Cm: float
+        Weighting coefficient for mass continuity constraint.
+    Cx: float
+        Smoothing coefficient for x-direction
+    Cy: float
+        Smoothing coefficient for y-direction
+    Cz: float
+        Smoothing coefficient for z-direction
+    Cb: float
+        Coefficient for sounding constraint
+    Cv: float
+        Weight for cost function related to vertical vorticity equation.
+    Cmod: float
+        Coefficient for model constraint
+    Cpoint: float
+        Coefficient for point constraint
+    Ut: float
+        Prescribed storm motion. This is only needed if Cv is not zero.
+    Vt: float
+        Prescribed storm motion. This is only needed if Cv is not zero.
+    grid_shape:
+        Shape of wind grid
+    dx:
+        Spacing of grid in x direction
+    dy:
+        Spacing of grid in y direction
+    dz:
+        Spacing of grid in z direction
+    x:
+        E-W grid levels in m
+    y:
+        N-S grid levels in m
+    z:
+        Grid vertical levels in m
+    rmsVr: float
+        The sum of squares of velocity/num_points. Use for normalization
+        of data weighting coefficient
+    weights: n_radars by z_bins by y_bins x x_bins float array
+        Data weights for each pair of radars
+    bg_weights: z_bins by y_bins x x_bins float array
+        Data weights for sounding constraint
+    model_weights: n_models by z_bins by y_bins by x_bins float array
+        Data weights for each model.
+    point_list: list or None
+        point_list: list of dicts
+        List of point constraints. Each member is a dict with keys of "u", "v",
+        to correspond to each component of the wind field and "x", "y", "z"
+        to correspond to the location of the point observation in the Grid's
+        Cartesian coordinates.
+    roi: float
+        The radius of influence of each point observation in m.
+    upper_bc: bool
+        True to enforce w=0 at top of domain (impermeability condition),
+        False to not enforce impermeability at top of domain
+    """
+    def __init__(self):
+        self.Ut = np.nan
+        self.Vt = np.nan
+        self.rmsVr = np.nan
+        self.grid_shape = None
+        self.Cmod = np.nan
+        self.Cpoint = np.nan
+        self.u_back = None
+        self.v_back = None
+        self.wts = []
+        self.vrs = []
+        self.azs = []
+        self.els = []
+        self.weights = []
+        self.bg_weights = []
+        self.model_weights = []
+        self.u_model = []
+        self.v_model = []
+        self.w_model = []
+        self.x = None
+        self.y = None
+        self.z = None
+        self.dx = np.nan
+        self.dy = np.nan
+        self.dz = np.nan
+        self.Co = 1.0
+        self.Cm = 1500.0
+        self.Cx = 0.0
+        self.Cy = 0.0
+        self.Cz = 0.0
+        self.Cb = 0.0
+        self.Cv = 0.0
+        self.Cmod = 0.0
+        self.Cpoint = 0.0
+        self.Ut = 0.0
+        self.Vt = 0.0
+        self.upper_bc = True
+        self.roi = 1000.0
+        self.frz = 4500.0
+        self.Nfeval = 0.0
+        self.engine = "scipy"
+        self.point_list = []
+
+
+def _get_dd_wind_field_scipy(Grids, u_init, v_init, w_init, engine,
+                             points=None, vel_name=None,
+                             refl_field=None, u_back=None, v_back=None, z_back=None,
+                             frz=4500.0, Co=1.0, Cm=1500.0, Cx=0.0,
+                             Cy=0.0, Cz=0.0, Cb=0.0, Cv=0.0, Cmod=0.0, Cpoint=0.0,
+                             Ut=None, Vt=None, low_pass_filter=True,
+                             mask_outside_opt=False, weights_obs=None,
+                             weights_model=None, weights_bg=None,
+                             max_iterations=1000, mask_w_outside_opt=True,
+                             filter_window=5, filter_order=3, min_bca=30.0,
+                             max_bca=150.0, upper_bc=True, model_fields=None,
+                             output_cost_functions=True, roi=1000.0,
+                             wind_tol=0.1):
+    global _wcurrmax
+    global _wprevmax
+    global iterations
+
+    # We have to have a prescribed storm motion for vorticity constraint
+    if (Ut is None or Vt is None):
+        if (Cv != 0.0):
+            raise ValueError(('Ut and Vt cannot be None if vertical ' +
+                              'vorticity constraint is enabled!'))
+
+    if not isinstance(Grids, list):
+        raise ValueError('Grids has to be a list!')
+
+    parameters = DDParameters()
+    parameters.Ut = Ut
+    parameters.Vt = Vt
+    parameters.engine = engine
+
+    # Ensure that all Grids are on the same coordinate system
+    prev_grid = Grids[0]
+    for g in Grids:
+        if not np.allclose(
+                g.x['data'], prev_grid.x['data'], atol=10):
+            raise ValueError('Grids do not have equal x coordinates!')
+
+        if not np.allclose(
+                g.y['data'], prev_grid.y['data'], atol=10):
+            raise ValueError('Grids do not have equal y coordinates!')
+
+        if not np.allclose(
+                g.z['data'], prev_grid.z['data'], atol=10):
+            raise ValueError('Grids do not have equal z coordinates!')
+
+        if not g.origin_latitude['data'] == prev_grid.origin_latitude['data']:
+            raise ValueError(("Grids have unequal origin lat/lons!"))
+
+        prev_grid = g
+
+    # Disable background constraint if none provided
+    if (u_back is None or v_back is None):
+        parameters.u_back = np.zeros(u_init.shape[0])
+        parameters.v_back = np.zeros(v_init.shape[0])
+    else:
+        # Interpolate sounding to radar grid
+        print('Interpolating sounding to radar grid')
+        u_interp = interp1d(z_back, u_back, bounds_error=False)
+        v_interp = interp1d(z_back, v_back, bounds_error=False)
+        parameters.u_back = u_interp(Grids[0].z['data'])
+        parameters.v_back = v_interp(Grids[0].z['data'])
+        print('Interpolated U field:')
+        print(parameters["u_back"])
+        print('Interpolated V field:')
+        print(parameters["v_back"])
+        print('Grid levels:')
+        print(Grids[0].z['data'])
+
+    # Parse names of velocity field
+    if refl_field is None:
+        refl_field = pyart.config.get_field_name('reflectivity')
+
+    # Parse names of velocity field
+    if vel_name is None:
+        vel_name = pyart.config.get_field_name('corrected_velocity')
+    winds = np.stack([u_init, v_init, w_init])
+
+    # Set up wind fields and weights from each radar
+    parameters.weights = np.zeros(
+        (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]))
+
+    parameters.bg_weights = np.zeros(v_init.shape)
+    if (model_fields is not None):
+        parameters.model_weights = np.ones(
+            (len(model_fields), u_init.shape[0], u_init.shape[1],
+             u_init.shape[2]))
+    else:
+        parameters.model_weights = np.zeros(
+            (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]))
+
+    if (model_fields is None):
+        if (Cmod != 0.0):
+            raise ValueError(
+                'Cmod must be zero if model fields are not specified!')
+
+    bca = np.zeros(
+        (len(Grids), len(Grids), u_init.shape[1], u_init.shape[2]))
+    sum_Vr = np.zeros(len(Grids))
+
+    for i in range(len(Grids)):
+        parameters.wts.append(np.ma.masked_invalid(
+            calculate_fall_speed(Grids[i], refl_field=refl_field, frz=frz)))
+        add_azimuth_as_field(Grids[i], dz_name=refl_field)
+        add_elevation_as_field(Grids[i], dz_name=refl_field)
+        parameters.vrs.append(np.ma.masked_invalid(
+            Grids[i].fields[vel_name]['data']))
+        parameters.azs.append(np.ma.masked_invalid(
+            Grids[i].fields['AZ']['data'] * np.pi / 180))
+        parameters.els.append(np.ma.masked_invalid(
+            Grids[i].fields['EL']['data'] * np.pi / 180))
+
+    if (len(Grids) > 1):
+        for i in range(len(Grids)):
+            for j in range(i + 1, len(Grids)):
+                print(("Calculating weights for radars " + str(i) +
+                       " and " + str(j)))
+                bca[i, j] = get_bca(Grids[i].radar_longitude['data'],
+                                    Grids[i].radar_latitude['data'],
+                                    Grids[j].radar_longitude['data'],
+                                    Grids[j].radar_latitude['data'],
+                                    Grids[i].point_x['data'][0],
+                                    Grids[i].point_y['data'][0],
+                                    Grids[i].get_projparams())
+
+                for k in range(parameters.vrs[i].shape[0]):
+                    if (weights_obs is None):
+                        valid = np.logical_and.reduce((
+                            ~parameters.vrs[i][k].mask,
+                            ~parameters.wts[i][k].mask,
+                            ~parameters.azs[i][k].mask,
+                            ~parameters.els[i][k].mask))
+                        valid = np.logical_and.reduce((valid,
+                            np.isfinite(parameters.vrs[i][k]),
+                            np.isfinite(parameters.wts[i][k]),
+                            np.isfinite(parameters.azs[i][k]),
+                            np.isfinite(parameters.els[i][k])))
+                        cur_array = parameters.weights[i, k]
+                        cur_array[np.logical_and(
+                            valid, np.logical_and(
+                                bca[i, j] >= math.radians(min_bca),
+                                bca[i, j] <= math.radians(max_bca)))] += 1
+                        cur_array[~valid] = 0
+                        parameters.weights[i, k] = cur_array
+                    else:
+                        parameters.weights[i, k] = weights_obs[i][k, :, :]
+
+                    if (weights_obs is None):
+                        valid = np.logical_and.reduce((
+                            ~parameters.vrs[j][k].mask,
+                            ~parameters.wts[j][k].mask,
+                            ~parameters.azs[j][k].mask,
+                            ~parameters.els[j][k].mask))
+                        valid = np.logical_and.reduce((valid,
+                            np.isfinite(parameters.vrs[j][k]),
+                            np.isfinite(parameters.wts[j][k]),
+                            np.isfinite(parameters.azs[j][k]),
+                            np.isfinite(parameters.els[j][k])))
+                        cur_array = parameters.weights[j, k]
+                        cur_array[np.logical_and(
+                            valid, np.logical_and(
+                                bca[i, j] >= math.radians(min_bca),
+                                bca[i, j] <= math.radians(max_bca)))] += 1
+                        cur_array[~valid] = 0
+                        parameters.weights[j, k] = cur_array
+                    else:
+                        parameters.weights[j, k] = weights_obs[j][k, :, :]
+
+                    if (weights_bg is None):
+                        valid = np.logical_and.reduce((
+                            ~parameters.vrs[j][k].mask,
+                            ~parameters.wts[j][k].mask,
+                            ~parameters.azs[j][k].mask,
+                            ~parameters.els[j][k].mask))
+                        valid = np.logical_and.reduce((
+                            valid, np.isfinite(parameters.vrs[j][k]),
+                            np.isfinite(parameters.wts[j][k]),
+                            np.isfinite(parameters.azs[j][k]),
+                            np.isfinite(parameters.els[j][k])))
+                        cur_array = parameters.bg_weights[k]
+                        cur_array[np.logical_or(
+                            bca[i, j] >= math.radians(min_bca),
+                            bca[i, j] <= math.radians(max_bca))] = 1
+                        cur_array[~valid] = 0
+                        parameters.bg_weights[i] = cur_array
+                    else:
+                        parameters.bg_weights[i] = weights_bg[i]
+
+        print("Calculating weights for models...")
+        coverage_grade = parameters.weights.sum(axis=0)
+        coverage_grade = coverage_grade / coverage_grade.max()
+
+        # Weigh in model input more when we have no coverage
+        # Model only weighs 1/(# of grids + 1) when there is full
+        # Coverage
+        if (model_fields is not None):
+            if (weights_model is None):
+                for i in range(len(model_fields)):
+                    parameters.model_weights[i] = 1 - (coverage_grade / (len(Grids) + 1))
+            else:
+                for i in range(len(model_fields)):
+                    parameters.model_weights[i] = weights_model[i]
+    else:
+        if weights_obs is None:
+            parameters.weights[0] = np.where(~parameters.vrs[0].mask, 1, 0)
+        else:
+            parameters.weights[0] = weights_obs[0]
+
+        if weights_bg is None:
+            parameters.bg_weights = np.where(~parameters.vrs[0].mask, 0, 1)
+        else:
+            parameters.bg_weights = weights_bg
+
+    parameters.vrs = [x.filled(-9999.) for x in parameters.vrs]
+    parameters.azs = [x.filled(-9999.) for x in parameters.azs]
+    parameters.els = [x.filled(-9999.) for x in parameters.els]
+    parameters.wts = [x.filled(-9999.) for x in parameters.wts]
+    parameters.weights[~np.isfinite(parameters.weights)] = 0
+    parameters.bg_weights[~np.isfinite(parameters.bg_weights)] = 0
+    parameters.weights[parameters.weights > 0] = 1
+    parameters.bg_weights[parameters.bg_weights > 0] = 1
+    sum_Vr = np.nansum(np.square(parameters.vrs * parameters.weights))
+    parameters.rmsVr = np.sqrt(np.nansum(sum_Vr) / np.nansum(parameters.weights))
+
+    del bca
+    parameters.grid_shape = u_init.shape
+    # Parse names of velocity field
+
+    winds = winds.flatten()
+
+    print("Starting solver ")
+    parameters.dx = np.diff(Grids[0].x['data'], axis=0)[0]
+    parameters.dy = np.diff(Grids[0].y['data'], axis=0)[0]
+    parameters.dz = np.diff(Grids[0].z['data'], axis=0)[0]
+    print('rmsVR = ' + str(parameters.rmsVr))
+    print('Total points: %d' % parameters.weights.sum())
+    parameters.z = Grids[0].point_z['data']
+    parameters.x = Grids[0].point_x['data']
+    parameters.y = Grids[0].point_y['data']
+    bt = time.time()
+
+    # First pass - no filter
+    wprevmax = 99
+    wcurrmax = w_init.max()
+    print("The max of w_init is", wcurrmax)
+    iterations = 0
+    bounds = [(-x, x) for x in 100 * np.ones(winds.shape)]
+
+    if (model_fields is not None):
+        for the_field in model_fields:
+            u_field = ("U_" + the_field)
+            v_field = ("V_" + the_field)
+            w_field = ("W_" + the_field)
+            parameters.u_model.append(Grids[0].fields[u_field]["data"])
+            parameters.v_model.append(Grids[0].fields[v_field]["data"])
+            parameters.w_model.append(Grids[0].fields[w_field]["data"])
+
+    parameters.Co = Co
+    parameters.Cm = Cm
+    parameters.Cx = Cx
+    parameters.Cy = Cy
+    parameters.Cz = Cz
+    parameters.Cb = Cb
+    parameters.Cv = Cv
+    parameters.Cmod = Cmod
+    parameters.Cpoint = Cpoint
+    parameters.roi = roi
+    parameters.upper_bc = upper_bc
+    parameters.points = points
+    parameters.point_list = points
+    _wprevmax = np.zeros(parameters.grid_shape)
+    _wcurrmax = np.zeros(parameters.grid_shape)
+    iterations = 0
+    def _vert_velocity_callback(x):
+        global _wprevmax
+        global _wcurrmax
+        global iterations
+        
+        if iterations % 10 > 0:
+            iterations = iterations + 1
+            return False
+         
+        wind = np.reshape(
+            x, (3, parameters.grid_shape[0],
+                parameters.grid_shape[1],
+                parameters.grid_shape[2]))
+        _wcurrmax = wind[2]
+        if iterations == 0:
+            _wprevmax = _wcurrmax
+            iterations = iterations + 1
+            return False
+        diff = np.abs(_wprevmax - _wcurrmax)
+        diff = np.where(parameters.bg_weights == 0, diff, np.nan)
+        delta = np.nanmax(diff)
+        if delta < wind_tol:
+            return True
+        _wprevmax = _wcurrmax
+        iterations = iterations + 1
+        print("Max change in w: %4.3f" % delta)
+        return False
+
+    parameters.print_out = False
+    winds = fmin_l_bfgs_b(J_function, winds, args=(parameters,),
+                          maxiter=max_iterations, pgtol=1e-8, bounds=bounds,
+                          fprime=grad_J, disp=0, iprint=-1, 
+                          callback=_vert_velocity_callback)
+    winds = np.reshape(
+        winds[0], 
+        (3, parameters.grid_shape[0], 
+            parameters.grid_shape[1], parameters.grid_shape[2]))
+    parameters.print_out = True
+    winds = np.stack([winds[0], winds[1], winds[2]])
+    winds = winds.flatten()
+        # """
+
+    if low_pass_filter == True:
+        print('Applying low pass filter to wind field...')
+        winds = np.reshape(winds, (3, parameters.grid_shape[0], parameters.grid_shape[1],
+                                   parameters.grid_shape[2]))
+        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=0)
+        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=1)
+        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=2)
+        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=0)
+        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=1)
+        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=2)
+        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=0)
+        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=1)
+        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=2)
+        winds = np.stack([winds[0], winds[1], winds[2]])
+        winds = winds.flatten()
+
+    print("Done! Time = " + "{:2.1f}".format(time.time() - bt))
+
+    # First pass - no filter
+    the_winds = np.reshape(
+        winds, (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2]))
+    u = the_winds[0]
+    v = the_winds[1]
+    w = the_winds[2]
+    where_mask = np.sum(parameters.weights, axis=0) + \
+                 np.sum(parameters.model_weights, axis=0)
+
+    u = np.ma.array(u)
+    w = np.ma.array(w)
+    v = np.ma.array(v)
+
+    if mask_outside_opt is True:
+        u = np.ma.masked_where(where_mask < 1, u)
+        v = np.ma.masked_where(where_mask < 1, v)
+        w = np.ma.masked_where(where_mask < 1, w)
+
+    if mask_w_outside_opt is True:
+        w = np.ma.masked_where(where_mask < 1, w)
+
+    u_field = deepcopy(Grids[0].fields[vel_name])
+    u_field['data'] = u
+    u_field['standard_name'] = 'u_wind'
+    u_field['long_name'] = 'meridional component of wind velocity'
+    u_field['min_bca'] = min_bca
+    u_field['max_bca'] = max_bca
+    v_field = deepcopy(Grids[0].fields[vel_name])
+    v_field['data'] = v
+    v_field['standard_name'] = 'v_wind'
+    v_field['long_name'] = 'zonal component of wind velocity'
+    v_field['min_bca'] = min_bca
+    v_field['max_bca'] = max_bca
+    w_field = deepcopy(Grids[0].fields[vel_name])
+    w_field['data'] = w
+    w_field['standard_name'] = 'w_wind'
+    w_field['long_name'] = 'vertical component of wind velocity'
+    w_field['min_bca'] = min_bca
+    w_field['max_bca'] = max_bca
+
+    new_grid_list = []
+
+    for grid in Grids:
+        temp_grid = deepcopy(grid)
+        temp_grid.add_field('u', u_field, replace_existing=True)
+        temp_grid.add_field('v', v_field, replace_existing=True)
+        temp_grid.add_field('w', w_field, replace_existing=True)
+        new_grid_list.append(temp_grid)
+
+    return new_grid_list
+
+
+def _get_dd_wind_field_tensorflow(Grids, u_init, v_init, w_init, points=None, vel_name=None,
+                                  refl_field=None, u_back=None, v_back=None, z_back=None,
+                                  frz=4500.0, Co=1.0, Cm=1500.0, Cx=0.0,
+                                  Cy=0.0, Cz=0.0, Cb=0.0, Cv=0.0, Cmod=0.0, Cpoint=0.0,
+                                  Ut=None, Vt=None, low_pass_filter=True,
+                                  mask_outside_opt=False, weights_obs=None,
+                                  weights_model=None, weights_bg=None,
+                                  max_iterations=1000, mask_w_outside_opt=True,
+                                  filter_window=5, filter_order=3, min_bca=30.0,
+                                  max_bca=150.0, upper_bc=True, model_fields=None,
+                                  output_cost_functions=True, roi=1000.0,
+                                  parallel_iterations=1, wind_tol=0.1):
+    if not TENSORFLOW_AVAILABLE:
+        raise ImportError("Tensorflow >=2.5 and tensorflow-probability " + 
+                "need to be installed in order to use the tensorflow engine.")
+    # We have to have a prescribed storm motion for vorticity constraint
+    if (Ut is None or Vt is None):
+        if (Cv != 0.0):
+            raise ValueError(('Ut and Vt cannot be None if vertical ' +
+                              'vorticity constraint is enabled!'))
+
+    if not isinstance(Grids, list):
+        raise ValueError('Grids has to be a list!')
+
+    parameters = DDParameters()
+    parameters.Ut = Ut
+    parameters.Vt = Vt
+    parameters.engine = "tensorflow"
+    # Ensure that all Grids are on the same coordinate system
+    prev_grid = Grids[0]
+    for g in Grids:
+        if not np.allclose(
+                g.x['data'], prev_grid.x['data'], atol=10):
+            raise ValueError('Grids do not have equal x coordinates!')
+
+        if not np.allclose(
+                g.y['data'], prev_grid.y['data'], atol=10):
+            raise ValueError('Grids do not have equal y coordinates!')
+
+        if not np.allclose(
+                g.z['data'], prev_grid.z['data'], atol=10):
+            raise ValueError('Grids do not have equal z coordinates!')
+
+        if not g.origin_latitude['data'] == prev_grid.origin_latitude['data']:
+            raise ValueError(("Grids have unequal origin lat/lons!"))
+
+        prev_grid = g
+
+    # Disable background constraint if none provided
+    if (u_back is None or v_back is None):
+        parameters.u_back = tf.zeros(u_init.shape[0])
+        parameters.v_back = tf.zeros(v_init.shape[0])
+    else:
+        # Interpolate sounding to radar grid
+        print('Interpolating sounding to radar grid')
+        u_interp = interp1d(z_back, u_back, bounds_error=False)
+        v_interp = interp1d(z_back, v_back, bounds_error=False)
+        parameters.u_back = tf.constant(
+            u_interp(Grids[0].z['data'].filled(np.nan)), dtype=tf.float32)
+        parameters.v_back = tf.constant(
+            v_interp(Grids[0].z['data'].filled(np.nan)), dtype=tf.float32)
+        print('Interpolated U field:')
+        print(parameters.u_back)
+        print('Interpolated V field:')
+        print(parameters.v_back)
+        print('Grid levels:')
+        print(Grids[0].z['data'])
+
+    # Parse names of velocity field
+    if refl_field is None:
+        refl_field = pyart.config.get_field_name('reflectivity')
+
+    # Parse names of velocity field
+    if vel_name is None:
+        vel_name = pyart.config.get_field_name('corrected_velocity')
+    winds = np.stack([u_init, v_init, w_init])
+    winds = winds.astype(np.float32)
+
+    # Set up wind fields and weights from each radar
+    parameters.weights = np.zeros(
+        (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2]),
+        dtype=np.float32)
+
+    parameters.bg_weights = np.zeros(v_init.shape)
+    if (model_fields is not None):
+        parameters.model_weights = np.ones(
+            (len(model_fields), u_init.shape[0], u_init.shape[1],
+             u_init.shape[2]), dtype=np.float32)
+    else:
+        parameters.model_weights = np.zeros(
+            (1, u_init.shape[0], u_init.shape[1], u_init.shape[2]),
+            dtype=np.float32)
+
+    if (model_fields is None):
+        if (Cmod != 0.0):
+            raise ValueError(
+                'Cmod must be zero if model fields are not specified!')
+
+    bca = np.zeros(
+        (len(Grids), len(Grids), u_init.shape[1], u_init.shape[2]),
+        dtype=np.float32)
+
+    for i in range(len(Grids)):
+        parameters.wts.append(np.ma.masked_invalid(
+            calculate_fall_speed(Grids[i], refl_field=refl_field, frz=frz)))
+        add_azimuth_as_field(Grids[i], dz_name=refl_field)
+        add_elevation_as_field(Grids[i], dz_name=refl_field)
+        parameters.vrs.append(np.ma.masked_invalid(
+            Grids[i].fields[vel_name]['data']))
+        parameters.azs.append(np.ma.masked_invalid(
+            Grids[i].fields['AZ']['data'] * np.pi / 180))
+        parameters.els.append(np.ma.masked_invalid(
+            Grids[i].fields['EL']['data'] * np.pi / 180))
+
+    if (len(Grids) > 1):
+        for i in range(len(Grids)):
+            for j in range(i + 1, len(Grids)):
+                print(("Calculating weights for radars " + str(i) +
+                       " and " + str(j)))
+                bca[i, j] = get_bca(Grids[i].radar_longitude['data'],
+                                    Grids[i].radar_latitude['data'],
+                                    Grids[j].radar_longitude['data'],
+                                    Grids[j].radar_latitude['data'],
+                                    Grids[i].point_x['data'][0],
+                                    Grids[i].point_y['data'][0],
+                                    Grids[i].get_projparams())
+
+                for k in range(parameters.vrs[i].shape[0]):
+                    if (weights_obs is None):
+                        cur_array = parameters.weights[i, k]
+                        valid = np.logical_and.reduce((
+                            ~parameters.vrs[i][k].mask,
+                            ~parameters.wts[i][k].mask,
+                            ~parameters.azs[i][k].mask,
+                            ~parameters.els[i][k].mask))
+                        valid = np.logical_and.reduce((valid,
+                                                       np.isfinite(parameters.vrs[i][k]),
+                                                       np.isfinite(parameters.wts[i][k]),
+                                                       np.isfinite(parameters.azs[i][k]),
+                                                       np.isfinite(parameters.els[i][k])))
+                        cur_array[np.logical_and(valid,
+                                                 np.logical_and(
+                                                     bca[i, j] >= math.radians(min_bca),
+                                                     bca[i, j] <= math.radians(max_bca)))] = 1
+                        cur_array[~valid] = 0
+                        parameters.weights[i, k] = cur_array
+                    else:
+                        parameters.weights[i, k] = weights_obs[i][k, :, :]
+
+                    if (weights_obs is None):
+                        cur_array = parameters.weights[j, k]
+                        valid = np.logical_and.reduce((
+                            ~parameters.vrs[j][k].mask,
+                            ~parameters.wts[j][k].mask,
+                            ~parameters.azs[j][k].mask,
+                            ~parameters.els[j][k].mask))
+                        valid = np.logical_and.reduce((valid,
+                                                       np.isfinite(parameters.vrs[j][k]),
+                                                       np.isfinite(parameters.wts[j][k]),
+                                                       np.isfinite(parameters.azs[j][k]),
+                                                       np.isfinite(parameters.els[j][k])))
+                        cur_array[np.logical_and(valid,
+                                                 np.logical_and(
+                                                     bca[i, j] >= math.radians(min_bca),
+                                                     bca[i, j] <= math.radians(max_bca)))] = 1
+                        cur_array[~valid] = 0
+                        parameters.weights[j, k] = cur_array
+                    else:
+                        parameters.weights[j, k] = weights_obs[j][k, :, :]
+
+                    if (weights_bg is None):
+                        cur_array = parameters.bg_weights[k]
+                        valid = np.logical_and.reduce((
+                            ~parameters.vrs[i][k].mask,
+                            ~parameters.wts[i][k].mask,
+                            ~parameters.azs[i][k].mask,
+                            ~parameters.els[i][k].mask))
+                        valid = np.logical_and.reduce((valid,
+                                                       np.isfinite(parameters.vrs[i][k]),
+                                                       np.isfinite(parameters.wts[i][k]),
+                                                       np.isfinite(parameters.azs[i][k]),
+                                                       np.isfinite(parameters.els[i][k])))
+                        cur_array[np.logical_or(
+                            bca[i, j] >= math.radians(min_bca),
+                            bca[i, j] <= math.radians(max_bca))] = 1
+                        cur_array[~valid] = 0
+                        parameters.bg_weights[i] = cur_array
+                    else:
+                        parameters.bg_weights[i] = weights_bg[i]
+
+        print("Calculating weights for models...")
+        coverage_grade = parameters.weights.sum(axis=0)
+        coverage_grade = coverage_grade / coverage_grade.max()
+
+        # Weigh in model input more when we have no coverage
+        # Model only weighs 1/(# of grids + 1) when there is full
+        # Coverage
+        if (model_fields is not None):
+            if (weights_model is None):
+                for i in range(len(model_fields)):
+                    parameters.model_weights[i] = 1 - (coverage_grade / (len(Grids) + 1))
+            else:
+                for i in range(len(model_fields)):
+                    parameters.model_weights[i] = weights_model[i]
+    else:
+        if weights_obs is None:
+            parameters.weights[0] = np.where(np.isfinite(parameters.vrs[0]), 1, 0)
+        else:
+            parameters.weights[0] = weights_obs[0]
+
+        if weights_bg is None:
+            parameters.bg_weights = np.where(np.isfinite(parameters.vrs[0]), 0, 1)
+        else:
+            parameters.bg_weights = weights_bg
+
+    parameters.vrs = [tf.constant(x.filled(-9999),
+                                  dtype=tf.float32) for x in parameters.vrs]
+    parameters.azs = [tf.constant(x.filled(-9999)
+                                  , dtype=tf.float32) for x in parameters.azs]
+    parameters.els = [tf.constant(x.filled(-9999)
+                                  , dtype=tf.float32) for x in parameters.els]
+    parameters.wts = [tf.constant(x.filled(-9999)
+                                  , dtype=tf.float32) for x in parameters.wts]
+    parameters.model_weights = tf.constant(parameters.model_weights,
+                                           dtype=tf.float32)
+    parameters.weights[~np.isfinite(parameters.weights)] = 0
+    parameters.weights[parameters.weights > 0] = 1
+    parameters.weights = tf.constant(parameters.weights, dtype=tf.float32)
+    parameters.bg_weights[parameters.bg_weights > 0] = 1
+    parameters.bg_weights = tf.constant(parameters.bg_weights, dtype=tf.float32)
+    sum_Vr = tf.experimental.numpy.nansum(
+        tf.square(parameters.vrs * parameters.weights))
+    parameters.rmsVr = np.sqrt(
+        np.nansum(sum_Vr) / tf.experimental.numpy.nansum(parameters.weights))
+
+    del bca
+    parameters.grid_shape = u_init.shape
+    # Parse names of velocity field
+
+    winds = winds.flatten()
+    winds = tf.Variable(winds, name="winds")
+
+    print("Starting solver ")
+    parameters.dx = np.diff(Grids[0].x['data'], axis=0)[0]
+    parameters.dy = np.diff(Grids[0].y['data'], axis=0)[0]
+    parameters.dz = np.diff(Grids[0].z['data'], axis=0)[0]
+    print('rmsVR = ' + str(parameters.rmsVr))
+    print('Total points: %d' % tf.reduce_sum(parameters.weights))
+    parameters.z = tf.constant(Grids[0].point_z['data'], dtype=tf.float32)
+    parameters.x = tf.constant(Grids[0].point_x['data'], dtype=tf.float32)
+    parameters.y = tf.constant(Grids[0].point_y['data'], dtype=tf.float32)
+    bt = time.time()
+
+    # First pass - no filter
+    wcurrmax = w_init.max()
+    print("The max of w_init is", wcurrmax)
+    iterations = 0
+    bounds = [(-x, x) for x in 100. * np.ones(winds.shape)]
+
+    if (model_fields is not None):
+        for the_field in model_fields:
+            u_field = ("U_" + the_field)
+            v_field = ("V_" + the_field)
+            w_field = ("W_" + the_field)
+            parameters.u_model.append(
+                tf.constant(Grids[0].fields[u_field]["data"]))
+            parameters.v_model.append(
+                tf.constant(Grids[0].fields[v_field]["data"]))
+            parameters.w_model.append(
+                tf.constant(Grids[0].fields[w_field]["data"]))
+
+    parameters.Co = Co
+    parameters.Cm = Cm
+    parameters.Cx = Cx
+    parameters.Cy = Cy
+    parameters.Cz = Cz
+    parameters.Cb = Cb
+    parameters.Cv = Cv
+    parameters.Cmod = Cmod
+    parameters.Cpoint = Cpoint
+    parameters.roi = roi
+    parameters.upper_bc = upper_bc
+    parameters.points = points
+    parameters.point_list = points
+    loss_and_gradient = lambda x: (J_function(x, parameters), grad_J(x, parameters))
+
+    winds = tfp.optimizer.lbfgs_minimize(
+        loss_and_gradient, initial_position=winds,
+        tolerance=1e-8, x_tolerance=wind_tol,
+        max_iterations=max_iterations, parallel_iterations=parallel_iterations)
+    print(winds)
+    winds = np.reshape(
+        winds.position.numpy(), (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2]))
+    iterations = iterations + 10
+    print('Iterations before filter: ' + str(iterations))
+    wcurrmax = winds[2].max()
+    winds = np.stack([winds[0], winds[1], winds[2]])
+    winds = winds.flatten()
+    # """
+
+    if low_pass_filter:
+        print('Applying low pass filter to wind field...')
+        winds = np.asarray(winds)
+        winds = np.reshape(winds,
+                           (3, parameters.grid_shape[0], parameters.grid_shape[1],
+                            parameters.grid_shape[2]))
+        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=0)
+        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=1)
+        winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=2)
+        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=0)
+        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=1)
+        winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=2)
+        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=0)
+        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=1)
+        winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=2)
+        winds = np.stack([winds[0], winds[1], winds[2]])
+        winds = winds.flatten()
+
+    print("Done! Time = " + "{:2.1f}".format(time.time() - bt))
+
+    # First pass - no filter
+    the_winds = np.reshape(
+        winds, (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2]))
+    u = the_winds[0]
+    v = the_winds[1]
+    w = the_winds[2]
+    where_mask = np.sum(parameters.weights, axis=0) + \
+                 np.sum(parameters.model_weights, axis=0)
+
+    u = np.ma.array(u)
+    w = np.ma.array(w)
+    v = np.ma.array(v)
+
+    if mask_outside_opt is True:
+        u = np.ma.masked_where(where_mask < 1, u)
+        v = np.ma.masked_where(where_mask < 1, v)
+        w = np.ma.masked_where(where_mask < 1, w)
+
+    if mask_w_outside_opt is True:
+        w = np.ma.masked_where(where_mask < 1, w)
+
+    u_field = deepcopy(Grids[0].fields[vel_name])
+    u_field['data'] = u
+    u_field['standard_name'] = 'u_wind'
+    u_field['long_name'] = 'meridional component of wind velocity'
+    u_field['min_bca'] = min_bca
+    u_field['max_bca'] = max_bca
+    v_field = deepcopy(Grids[0].fields[vel_name])
+    v_field['data'] = v
+    v_field['standard_name'] = 'v_wind'
+    v_field['long_name'] = 'zonal component of wind velocity'
+    v_field['min_bca'] = min_bca
+    v_field['max_bca'] = max_bca
+    w_field = deepcopy(Grids[0].fields[vel_name])
+    w_field['data'] = w
+    w_field['standard_name'] = 'w_wind'
+    w_field['long_name'] = 'vertical component of wind velocity'
+    w_field['min_bca'] = min_bca
+    w_field['max_bca'] = max_bca
+
+    new_grid_list = []
+
+    for grid in Grids:
+        temp_grid = deepcopy(grid)
+        temp_grid.add_field('u', u_field, replace_existing=True)
+        temp_grid.add_field('v', v_field, replace_existing=True)
+        temp_grid.add_field('w', w_field, replace_existing=True)
+        new_grid_list.append(temp_grid)
+
+    return new_grid_list
+
+
+def get_dd_wind_field(Grids, u_init, v_init, w_init, engine="scipy", **kwargs):
+    """
+    This function takes in a list of Py-ART Grid objects and derives a
+    wind field. Every Py-ART Grid in Grids must have the same grid
+    specification. 
+
+    In order for the model data constraint to be used,
+    the model data must be added as a field to at least one of the
+    grids in Grids. This involves interpolating the model data to the
+    Grids' coordinates. There are helper functions for this for WRF
+    and HRRR data in :py:func:`pydda.constraints`:
+
+    :py:func:`make_constraint_from_wrf`
+
+    :py:func:`add_hrrr_constraint_to_grid`
+
+    Parameters
+    ==========
+
+    Grids: list of Py-ART Grids
+        The list of Py-ART grids to take in corresponding to each radar.
+        All grids must have the same shape, x coordinates, y coordinates
+        and z coordinates.
+    u_init: 3D ndarray
+        The intial guess for the zonal wind field, input as a 3D array
+        with the same shape as the fields in Grids.
+    v_init: 3D ndarray
+        The intial guess for the meridional wind field, input as a 3D array
+        with the same shape as the fields in Grids.
+    w_init: 3D ndarray
+        The intial guess for the vertical wind field, input as a 3D array
+        with the same shape as the fields in Grids.
+    engine: str (one of "scipy", "tensorflow", "jax")
+        Setting this flag will use the solver based off of SciPy, TensorFlow, or Jax.
+        Using Tensorflow or Jax expands PyDDA's capabiability to take advantage of GPU-based systems.
+        In addition, these two implementations use automatic differentation to calcualte the gradient
+        of the cost function in order to optimize the gradient calculation.
+        TensorFlow 2.6 and tensorflow-probability are required for the TensorFlow-basedengine.
+        The latest version of Jax is required for the Jax-based engine.
+    points: None or list of dicts
+        Point observations as returned by :func:`pydda.constraints.get_iem_obs`. Set
+        to None to disable.
+    vel_name: string
+        Name of radial velocity field. Setting to None will have PyDDA attempt
+        to automatically detect the velocity field name.
+    refl_field: string
+        Name of reflectivity field. Setting to None will have PyDDA attempt 
+        to automatically detect the reflectivity field name.
+    u_back: 1D array
+        Background zonal wind field from a sounding as a function of height.
+        This should be given in the sounding's vertical coordinates.
+    v_back: 1D array
+        Background meridional wind field from a sounding as a function of
+        height. This should be given in the sounding's vertical coordinates.
+    z_back: 1D array
+        Heights corresponding to background wind field levels in meters. This
+        is given in the sounding's original coordinates.
+    frz: float
+        Freezing level used for fall speed calculation in meters.
+    Co: float
+        Weight for cost function related to observed radial velocities.
+    Cm: float
+        Weight for cost function related to the mass continuity equation.
+    Cx: float
+        Weight for cost function related to smoothness in x direction
+    Cy: float
+        Weight for cost function related to smoothness in y direction
+    Cz: float
+        Weight for cost function related to smoothness in z direction
+    Cv: float
+        Weight for cost function related to vertical vorticity equation.
+    Cmod: float
+        Weight for cost function related to custom constraints.
+    Cpoint: float
+        Weight for cost function related to point observations.
+    weights_obs: list of floating point arrays or None
+        List of weights for each point in grid from each radar in Grids.
+        Set to None to let PyDDA determine this automatically.
+    weights_model: list of floating point arrays or None
+        List of weights for each point in grid from each custom field in
+        model_fields. Set to None to let PyDDA determine this automatically.
+    weights_bg: list of floating point arrays or None
+        List of weights for each point in grid from the sounding. Set to None
+        to let PyDDA determine this automatically.
+    Ut: float
+        Prescribed storm motion in zonal direction.
+        This is only needed if Cv is not zero.
+    Vt: float
+        Prescribed storm motion in meridional direction.
+        This is only needed if Cv is not zero.
+    filter_winds: bool
+        If this is True, PyDDA will run a low pass filter on
+        the retrieved wind field. Set to False to disable the low pass filter.
+    mask_outside_opt: bool
+        If set to true, wind values outside the multiple doppler lobes will
+        be masked, i.e. if less than 2 radars provide coverage for a given
+        point.
+    max_iterations: int
+        The maximum number of iterations to run the optimization loop for.
+    mask_w_outside_opt: bool
+        If set to true, vertical winds outside the multiple doppler lobes will
+        be masked, i.e. if less than 2 radars provide coverage for a given
+        point.
+    filter_window: int
+        Window size to use for the low pass filter. A larger window will
+        increase the number of points factored into the polynomial fit for
+        the filter, and hence will increase the smoothness.
+    filter_order: int
+        The order of the polynomial to use for the low pass filter. Higher
+        order polynomials allow for the retention of smaller scale features
+        but may also not remove enough noise.
+    min_bca: float
+        Minimum beam crossing angle in degrees between two radars. 30.0 is the
+        typical value used in many publications.
+    max_bca: float
+        Minimum beam crossing angle in degrees between two radars. 150.0 is the
+        typical value used in many publications.
+    upper_bc: bool
+        Set this to true to enforce w = 0 at the top of the atmosphere. This is
+        commonly called the impermeability condition.
+    model_fields: list of strings
+        The list of fields in the first grid in Grids that contain the custom
+        data interpolated to the Grid's grid specification. Helper functions
+        to create such gridded fields for HRRR and NetCDF WRF data exist
+        in ::pydda.constraints::. PyDDA will look for fields named U_(model
+        field name), V_(model field name), and W_(model field name). For
+        example, if you have U_hrrr, V_hrrr, and W_hrrr, then specify ["hrrr"]
+        into model_fields.
+    output_cost_functions: bool
+        Set to True to output the value of each cost function every
+        10 iterations.
+    roi: float
+        Radius of influence for the point observations. The point observation will
+        not hold any weight outside this radius.
+    parallel_iterations: int
+        The number of iterations to run in parallel in the optimization loop.
+        This is only for the TensorFlow-based engine.
+    wind_tol: float
+        Stop iterations after maximum change in winds is less than this value.
+
+    Returns
+    =======
+    new_grid_list: list
+        A list of Py-ART grids containing the derived wind fields. These fields
+        are displayable by the visualization module.
+    """
+
+    if engine.lower() == "scipy" or engine.lower() == "jax":
+        return _get_dd_wind_field_scipy(Grids, u_init, v_init, w_init, engine, **kwargs)
+    elif engine.lower() == "tensorflow":
+        return _get_dd_wind_field_tensorflow(Grids, u_init, v_init, w_init, **kwargs)
+    else:
+        raise NotImplementedError("Engine %s is not supported." % engine)
+
+
+def get_bca(rad1_lon, rad1_lat, rad2_lon, rad2_lat, x, y, projparams):
+    """
+    This function gets the beam crossing angle between two lat/lon pairs.
+
+    Parameters
+    ==========
+    rad1_lon: float
+        The longitude of the first radar.
+    rad1_lat: float
+        The latitude of the first radar.
+    rad2_lon: float
+        The longitude of the second radar.
+    rad2_lat: float
+        The latitude of the second radar.
+    x: nD float array
+        The Cartesian x coordinates of the grid
+    y: nD float array
+        The Cartesian y corrdinates of the grid
+    projparams: Py-ART projparams
+        The projection parameters of the Grid
+
+    Returns
+    =======
+    bca: nD float array
+        The beam crossing angle between the two radars in radians.
+
+    """
+
+    rad1 = pyart.core.geographic_to_cartesian(rad1_lon, rad1_lat, projparams)
+    rad2 = pyart.core.geographic_to_cartesian(rad2_lon, rad2_lat, projparams)
+    # Create grid with Radar 1 in center
+
+    x = x - rad1[0]
+    y = y - rad1[1]
+    rad2 = np.array(rad2) - np.array(rad1)
+    a = np.sqrt(np.multiply(x, x) + np.multiply(y, y))
+    b = np.sqrt(pow(x-rad2[0], 2) + pow(y-rad2[1], 2))
+    c = np.sqrt(rad2[0]*rad2[0] + rad2[1]*rad2[1])
+    inp_array1 = x / a
+    inp_array1 = np.where(inp_array1 < -1, -1, inp_array1)
+    inp_array1 = np.where(inp_array1 > 1, 1, inp_array1)
+    theta_1 = np.arccos(inp_array1)
+    inp_array2 = (x - rad2[1]) / b
+    inp_array2 = np.where(inp_array2 < -1, -1, inp_array2)
+    inp_array2 = np.where(inp_array2 > 1, 1, inp_array2)
+    theta_2 = np.arccos(inp_array2)
+    inp_array3 = (a*a + b*b - c*c) / (2 * a *b)
+    inp_array3 = np.where(inp_array3 < -1, -1, inp_array3)
+    inp_array3 = np.where(inp_array3 > 1, 1, inp_array3)
+    
+    return np.ma.masked_invalid(np.arccos(inp_array3))
diff --git a/pydda/vis/barb_plot.py b/pydda/vis/barb_plot.py
index 97f8080e..2580e11d 100644
--- a/pydda/vis/barb_plot.py
+++ b/pydda/vis/barb_plot.py
@@ -8,7 +8,12 @@
 
 from .. import retrieval
 from matplotlib.axes import Axes
-from cartopy.mpl.geoaxes import GeoAxes
+try:
+    from cartopy.mpl.geoaxes import GeoAxes
+    CARTOPY_AVAILABLE = True
+except ImportError:
+    CARTOPY_AVAILABLE = False
+
 GeoAxes._pcolormesh_patched = Axes.pcolormesh
 
 def plot_horiz_xsection_barbs(Grids, ax=None,
@@ -146,9 +151,8 @@ def plot_horiz_xsection_barbs(Grids, ax=None,
 
     if(u_vel_contours is not None):
         u_filled = np.ma.filled(u[level, :, :], fill_value=np.nan)
-        cs = ax.contourf(grid_x[level, :, :], grid_y[level, :, :],
-                         u_filled, levels=u_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[level, :, :], grid_y[level, :, :],
+                        u_filled, levels=u_vel_contours, linewidths=2)
         cs.set_clim([np.min(u_vel_contours), np.max(u_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -158,9 +162,8 @@ def plot_horiz_xsection_barbs(Grids, ax=None,
 
     if(v_vel_contours is not None):
         v_filled = np.ma.filled(v[level, :, :], fill_value=np.nan)
-        cs = ax.contourf(grid_x[level, :, :], grid_y[level, :, :],
-                         v_filled, levels=u_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[level, :, :], grid_y[level, :, :],
+                         v_filled, levels=u_vel_contours, linewidths=2)
         cs.set_clim([np.min(v_vel_contours), np.max(v_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -170,9 +173,8 @@ def plot_horiz_xsection_barbs(Grids, ax=None,
 
     if(w_vel_contours is not None):
         w_filled = np.ma.filled(w[level, :, :], fill_value=np.nan)
-        cs = ax.contourf(grid_x[level, :, :], grid_y[level, :, :],
-                         w_filled, levels=w_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[level, :, :], grid_y[level, :, :],
+                        w_filled, levels=w_vel_contours, linewidths=2)
         cs.set_clim([np.min(w_vel_contours), np.max(w_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -183,9 +185,8 @@ def plot_horiz_xsection_barbs(Grids, ax=None,
     if(wind_vel_contours is not None):
         vel = np.ma.sqrt(u[level, :, :]**2 + v[level, :, :]**2)
         #vel = vel.filled(fill_value=np.nan)
-        cs = ax.contourf(grid_x[level, :, :], grid_y[level, :, :],
-                         vel, levels=wind_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[level, :, :], grid_y[level, :, :],
+                        vel, levels=wind_vel_contours, linewidths=2)
         cs.set_clim([np.min(wind_vel_contours), np.max(wind_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -241,7 +242,6 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
                                   colorbar_contour_flag=False,
                                   bg_grid_no=0, barb_spacing_x_km=10.0,
                                   barb_spacing_y_km=10.0,
-                                  contour_alpha=0.7,
                                   coastlines=True, gridlines=True):
     """
     This procedure plots a horizontal cross section of winds from wind fields
@@ -301,8 +301,6 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
         The spacing in km between each wind barb in the x direction.
     barb_spacing_y_km: float
         The spacing in km between each wind barb in the y direction.
-    contour_alpha: float
-        Alpha (transparency) of velocity contours. 0 = transparent, 1 = opaque
     coastlines: bool
         Set to true to display coastlines
     gridlines: bool
@@ -314,6 +312,9 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
         Axis handle to output axis
     """
 
+    if not CARTOPY_AVAILABLE:
+        raise ModuleNotFoundError("Cartopy needs to be installed in order to use plotting module!")
+
     if(bg_grid_no > -1):
         grid_bg = Grids[bg_grid_no].fields[background_field]['data']
     else:
@@ -367,9 +368,9 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
         u_filled = np.ma.masked_where(u[level, :, :] < np.min(u_vel_contours), 
                                       u[level, :, :])
         try:
-            cs = ax.contourf(grid_lon[:, :], grid_lat[:, :],
-                             u_filled, levels=u_vel_contours, linewidths=2,
-                             alpha=contour_alpha, zorder=2, extend='both')
+            cs = ax.contour(grid_lon[:, :], grid_lat[:, :],
+                            u_filled, levels=u_vel_contours, linewidths=2,
+                            zorder=2, extend='both')
             cs.set_clim([np.min(u_vel_contours), np.max(u_vel_contours)])
             cs.cmap.set_under(color='white', alpha=0)
             cs.cmap.set_bad(color='white', alpha=0)
@@ -386,9 +387,9 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
         v_filled = np.ma.masked_where(v[level, :, :] < np.min(v_vel_contours), 
                                       v[level, :, :])
         try:
-            cs = ax.contourf(grid_lon[:, :], grid_lat[:, :],
-                             v_filled, levels=u_vel_contours, linewidths=2,
-                             alpha=contour_alpha, zorder=2, extend='both')
+            cs = ax.contour(grid_lon[:, :], grid_lat[:, :],
+                            v_filled, levels=u_vel_contours, linewidths=2,
+                            zorder=2, extend='both')
             cs.set_clim([np.min(v_vel_contours), np.max(v_vel_contours)])
             cs.cmap.set_under(color='white', alpha=0)
             cs.cmap.set_bad(color='white', alpha=0)
@@ -404,9 +405,9 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
         w_filled = np.ma.masked_where(w[level, :, :] < np.min(w_vel_contours), 
                                       w[level, :, :])
         try:
-            cs = ax.contourf(grid_lon[::, ::], grid_lat[::, ::],
-                             w_filled, levels=w_vel_contours, linewidths=2,
-                             alpha=contour_alpha, zorder=2, extend='both')
+            cs = ax.contour(grid_lon[::, ::], grid_lat[::, ::],
+                            w_filled, levels=w_vel_contours, linewidths=2,
+                            zorder=2, extend='both')
             cs.set_clim([np.min(w_vel_contours), np.max(w_vel_contours)])
             cs.cmap.set_under(color='white', alpha=0)
             cs.cmap.set_bad(color='white', alpha=0)
@@ -424,9 +425,8 @@ def plot_horiz_xsection_barbs_map(Grids, ax=None,
         vel = np.ma.sqrt(u[level, :, :]**2 + v[level, :, :]**2)
         vel = vel.filled(fill_value=np.nan)
         try:
-            cs = ax.contourf(grid_x[level, :, :], grid_y[level, :, :],
-                             vel, levels=wind_vel_contours, linewidths=2,
-                             alpha=contour_alpha)
+            cs = ax.contour(grid_x[level, :, :], grid_y[level, :, :],
+                            vel, levels=wind_vel_contours, linewidths=2)
             cs.cmap.set_under(color='white', alpha=0)
             cs.cmap.set_bad(color='white', alpha=0)
 
@@ -496,8 +496,7 @@ def plot_xz_xsection_barbs(Grids, ax=None,
                            colorbar_flag=True,
                            colorbar_contour_flag=False, 
                            bg_grid_no=0, barb_spacing_x_km=10.0,
-                           barb_spacing_z_km=1.0,
-                           contour_alpha=0.7):
+                           barb_spacing_z_km=1.0):
     """
     This procedure plots a cross section of winds from wind fields
     generated by PyDDA in the X-Z plane using barbs.
@@ -564,6 +563,8 @@ def plot_xz_xsection_barbs(Grids, ax=None,
         Axis handle to output axis
     """
 
+    if not CARTOPY_AVAILABLE:
+        raise ModuleNotFoundError("Cartopy needs to be installed in order to use plotting module!")
     if(bg_grid_no > -1):
         grid_bg = Grids[bg_grid_no].fields[background_field]['data']
     else:
@@ -608,9 +609,8 @@ def plot_xz_xsection_barbs(Grids, ax=None,
 
     if(u_vel_contours is not None):
         u_filled = np.ma.filled(u[::, level, ::], fill_value=0)
-        cs = ax.contourf(grid_x[::, level, ::], grid_h[::, level, ::],
-                         u_filled, levels=u_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[::, level, ::], grid_h[::, level, ::],
+                        u_filled, levels=u_vel_contours, linewidths=2)
         cs.set_clim([np.min(u_vel_contours), np.max(u_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -620,9 +620,8 @@ def plot_xz_xsection_barbs(Grids, ax=None,
 
     if(v_vel_contours is not None):
         v_filled = np.ma.filled(w[::, level, ::], fill_value=0)
-        cs = ax.contourf(grid_x[::, level, ::], grid_h[::, level, ::],
-                         v_filled, levels=v_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[::, level, ::], grid_h[::, level, ::],
+                        v_filled, levels=v_vel_contours, linewidths=2)
         cs.set_clim([np.min(v_vel_contours), np.max(v_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -632,9 +631,8 @@ def plot_xz_xsection_barbs(Grids, ax=None,
 
     if(w_vel_contours is not None):
         w_filled = np.ma.filled(w[::, level, ::], fill_value=0)
-        cs = ax.contourf(grid_x[::, level, ::], grid_h[::, level, ::],
-                         w_filled, levels=w_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[::, level, ::], grid_h[::, level, ::],
+                        w_filled, levels=w_vel_contours, linewidths=2)
         cs.set_clim([np.min(w_vel_contours), np.max(w_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -645,9 +643,8 @@ def plot_xz_xsection_barbs(Grids, ax=None,
     if(wind_vel_contours is not None):
         vel = np.ma.sqrt(u[:, level, :]**2 + v[:, level, :]**2)
         vel = vel.filled(fill_value=np.nan)
-        cs = ax.contourf(grid_x[:, level, :], grid_h[:, level, :],
-                         vel, levels=wind_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_x[:, level, :], grid_h[:, level, :],
+                        vel, levels=wind_vel_contours, linewidths=2)
         cs.set_clim([np.min(wind_vel_contours), np.max(wind_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -685,8 +682,7 @@ def plot_yz_xsection_barbs(Grids, ax=None,
                            colorbar_flag=True,
                            colorbar_contour_flag=False,
                            bg_grid_no=0, barb_spacing_y_km=10.0,
-                           barb_spacing_z_km=1.0,
-                           contour_alpha=0.7):
+                           barb_spacing_z_km=1.0):
     """
     This procedure plots a cross section of winds from wind fields
     generated by PyDDA in the Y-Z plane using barbs.
@@ -796,9 +792,8 @@ def plot_yz_xsection_barbs(Grids, ax=None,
 
     if(u_vel_contours is not None):
         u_filled = np.ma.filled(u[:, :, level], fill_value=0)
-        cs = plt.contourf(grid_y[:, :, level], grid_h[:, :, level],
-                          u_filled, levels=u_vel_contours, linewidths=2,
-                          alpha=contour_alpha)
+        cs = plt.contour(grid_y[:, :, level], grid_h[:, :, level],
+                         u_filled, levels=u_vel_contours, linewidths=2)
         cs.set_clim([np.min(u_vel_contours), np.max(u_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -808,9 +803,8 @@ def plot_yz_xsection_barbs(Grids, ax=None,
 
     if(v_vel_contours is not None):
         v_filled = np.ma.filled(v[:, :, level], fill_value=0)
-        cs = plt.contourf(grid_y[:, :, level], grid_h[:, :, level],
-                          v_filled, levels=w_vel_contours, linewidths=2,
-                          alpha=contour_alpha)
+        cs = plt.contour(grid_y[:, :, level], grid_h[:, :, level],
+                          v_filled, levels=w_vel_contours, linewidths=2)
         cs.set_clim([np.min(v_vel_contours), np.max(v_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -820,9 +814,8 @@ def plot_yz_xsection_barbs(Grids, ax=None,
 
     if(w_vel_contours is not None):
         w_filled = np.ma.filled(w[::, ::, level], fill_value=0)
-        cs = plt.contourf(grid_y[::, ::, level], grid_h[::, ::, level],
-                          w_filled, levels=w_vel_contours, linewidths=2,
-                          alpha=contour_alpha)
+        cs = plt.contour(grid_y[::, ::, level], grid_h[::, ::, level],
+                          w_filled, levels=w_vel_contours, linewidths=2)
         cs.set_clim([np.min(w_vel_contours), np.max(w_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
@@ -833,9 +826,8 @@ def plot_yz_xsection_barbs(Grids, ax=None,
     if(wind_vel_contours is not None):
         vel = np.ma.sqrt(u[:, :, level]**2 + v[:, :, level]**2)
         vel = vel.filled(fill_value=np.nan)
-        cs = ax.contourf(grid_y[:, :, level], grid_h[:, :, level],
-                         vel, levels=wind_vel_contours, linewidths=2,
-                         alpha=contour_alpha)
+        cs = ax.contour(grid_y[:, :, level], grid_h[:, :, level],
+                        vel, levels=wind_vel_contours, linewidths=2)
         cs.set_clim([np.min(wind_vel_contours), np.max(wind_vel_contours)])
         cs.cmap.set_under(color='white', alpha=0)
         cs.cmap.set_bad(color='white', alpha=0)
diff --git a/pydda/vis/quiver_plot.py b/pydda/vis/quiver_plot.py
index 2d42b4bd..89c238c1 100644
--- a/pydda/vis/quiver_plot.py
+++ b/pydda/vis/quiver_plot.py
@@ -7,7 +7,12 @@
 
 from .. import retrieval
 from matplotlib.axes import Axes
-from cartopy.mpl.geoaxes import GeoAxes
+try:
+    from cartopy.mpl.geoaxes import GeoAxes
+    CARTOPY_AVAILABLE = True
+except ImportError:
+    CARTOPY_AVAILABLE = False
+
 GeoAxes._pcolormesh_patched = Axes.pcolormesh
 
 def plot_horiz_xsection_quiver(Grids, ax=None,
@@ -138,7 +143,8 @@ def plot_horiz_xsection_quiver(Grids, ax=None,
     """
 
     grid_bg = Grids[bg_grid_no].fields[background_field]['data']
-
+    if not CARTOPY_AVAILABLE:
+        raise ModuleNotFoundError("Cartopy needs to be installed in order to use plotting module!")
     if(vmin is None):
         vmin = grid_bg.min()
 
@@ -388,7 +394,8 @@ def plot_horiz_xsection_quiver_map(Grids, ax=None,
     ax: matplotlib axis
         Axis handle to output axis
     """
-
+    if not CARTOPY_AVAILABLE:
+        raise ModuleNotFoundError("Cartopy needs to be installed in order to use plotting module!")
     if(bg_grid_no > -1):
         grid_bg = Grids[bg_grid_no].fields[background_field]['data']
     else:
diff --git a/pydda/vis/streamline_plot.py b/pydda/vis/streamline_plot.py
index 98da9f95..ac5c31fb 100644
--- a/pydda/vis/streamline_plot.py
+++ b/pydda/vis/streamline_plot.py
@@ -9,7 +9,12 @@
 
 from .. import retrieval
 from matplotlib.axes import Axes
-from cartopy.mpl.geoaxes import GeoAxes
+try:
+    from cartopy.mpl.geoaxes import GeoAxes
+    CARTOPY_AVAILABLE = True
+except ImportError:
+    CARTOPY_AVAILABLE = False
+
 GeoAxes._pcolormesh_patched = Axes.pcolormesh
 
 def plot_horiz_xsection_streamlines(Grids, ax=None,
@@ -297,6 +302,8 @@ def plot_horiz_xsection_streamlines_map(Grids, ax=None,
     ax: matplotlib axis
         Axis handle to output axis
     """
+    if not CARTOPY_AVAILABLE:
+        raise ModuleNotFoundError("Cartopy needs to be installed in order to use plotting module!")
 
     if(bg_grid_no > -1):
         grid_bg = Grids[bg_grid_no].fields[background_field]['data']
diff --git a/setup.py b/setup.py
index b352cacf..0326ca8a 100644
--- a/setup.py
+++ b/setup.py
@@ -40,9 +40,9 @@
 LONG_DESCRIPTION = "\n".join(DOCLINES[2:])
 LICENSE = 'BSD'
 PLATFORMS = "Linux, Windows, OSX"
-MAJOR = 0
-MINOR = 5
-MICRO = 2
+MAJOR = 1
+MINOR = 0
+MICRO = 0
 
 #SCRIPTS = glob.glob('scripts/*')
 #TEST_SUITE = 'nose.collector'


Test patch:
diff --git a/pydda/tests/test_plots.py b/pydda/tests/test_plots.py
index ffac39a9..c9294af0 100644
--- a/pydda/tests/test_plots.py
+++ b/pydda/tests/test_plots.py
@@ -5,7 +5,11 @@
 import pyart
 import pytest
 import matplotlib.pyplot as plt
-import cartopy.crs as ccrs
+try:
+    import cartopy.crs as ccrs
+    CARTOPY_AVAILABLE = True
+except:
+    CARTOPY_AVAILABLE = False
 
 
 @pytest.mark.mpl_image_compare(tolerance=50)
@@ -44,6 +48,7 @@ def test_plot_horiz_yz_xsection_barbs():
     return fig
 
 
+@pytest.mark.skipif(CARTOPY_AVAILABLE=False)
 @pytest.mark.mpl_image_compare(tolerance=50)
 def test_plot_horiz_xsection_barbs_map():
     berr_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
@@ -68,7 +73,7 @@ def test_plot_horiz_xsection_streamlines():
     return fig
 
 
-
+@pytest.mark.skipif(CARTOPY_AVAILABLE=False)
 @pytest.mark.mpl_image_compare(tolerance=60)
 def test_plot_horiz_xsection_streamlines_map():
     berr_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
@@ -120,6 +125,7 @@ def test_plot_horiz_xsection_quiver():
     return fig
 
 
+@pytest.mark.skipif(CARTOPY_AVAILABLE=False)
 @pytest.mark.mpl_image_compare(tolerance=50)
 def test_plot_horiz_xsection_quiver_map():
     berr_grid = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
diff --git a/pydda/tests/test_retrieval.py b/pydda/tests/test_retrieval.py
index f8e0b408..3e5049f6 100644
--- a/pydda/tests/test_retrieval.py
+++ b/pydda/tests/test_retrieval.py
@@ -8,8 +8,21 @@
 
 import pydda
 import pyart
+import pytest
 import numpy as np
 
+try:
+    import tensorflow as tf
+    TF_AVAILABLE = True
+except ImportError:
+    TF_AVAILABLE = False
+
+try:
+    import jax
+    JAX_AVAILABLE = True
+except ImportError:
+    JAX_AVAILABLE = False
+
 from distributed import Client, LocalCluster
 from copy import deepcopy
 
@@ -49,20 +62,22 @@ def test_make_updraft_from_convergence_field():
     assert np.ma.max(new_w > 3)
 
 
-def test_twpice_case():
+@pytest.mark.skipif(~JAX_AVAILABLE, reason="Jax not installed")
+def test_twpice_case_jax():
     """ Use a test case from TWP-ICE """
     Grid0 = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
     Grid1 = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR1)
     sounding = pyart.io.read_arm_sonde(pydda.tests.SOUNDING_PATH)
 
     u_init, v_init, w_init = pydda.initialization.make_wind_field_from_profile(
-        Grid0, sounding[1], vel_field='corrected_velocity')
+        Grid1, sounding[1], vel_field='corrected_velocity')
 
     Grids = pydda.retrieval.get_dd_wind_field(
         [Grid0, Grid1], u_init, v_init, w_init, Co=100, Cm=1500.0,
+        wind_tol=0.1,
         Cz=0, Cmod=0.0, vel_name='corrected_velocity',
-        refl_field='reflectivity', frz=5000.0,
-        filt_iterations=0, mask_outside_opt=True, upper_bc=1)
+        refl_field='reflectivity', frz=5000.0, engine="jax",
+        mask_outside_opt=True, upper_bc=1)
 
     # In this test grid, we expect the mean flow to be to the southeast
     # Maximum updrafts should be at least 10 m/s
@@ -74,28 +89,31 @@ def test_twpice_case():
     assert v_mean < 0
     assert w_max > 10
 
-    # Now we will test the nesting. Do the same retrieval, and make sure
-    # that we get the same result within a prescribed tolerance
-    cluster = LocalCluster(n_workers=2, processes=True)
-    client = Client(cluster)
-    Grids2 = pydda.retrieval.get_dd_wind_field_nested(
-        [Grid0, Grid1], u_init, v_init, w_init, client, Co=100, Cm=1500.0,
-        Cz=0, Cmod=0.0, vel_name='corrected_velocity',
-        refl_field='reflectivity', frz=5000.0,
-        filt_iterations=0, mask_outside_opt=True, upper_bc=1)
-
-    # Make sure features are correlated between both versions. No reason
-    # to expect the same answer, but features should be correlated
-    # Nesting tends to make the updrafts a bit better resolved, so expect
-    # less of an outright correlation (but still strong)
-    assert np.corrcoef(Grids2[0].fields["u"]["data"].flatten(),
-                       Grids[0].fields["u"]["data"].flatten())[0, 1] > 0.9
-    assert np.corrcoef(Grids2[0].fields["v"]["data"].flatten(),
-                       Grids[0].fields["v"]["data"].flatten())[0, 1] > 0.9
-    assert np.corrcoef(Grids2[0].fields["w"]["data"].flatten(),
-                       Grids[0].fields["w"]["data"].flatten())[0, 1] > 0.5
-    cluster.close()
-    client.close()
+
+@pytest.mark.skipif(~TF_AVAILABLE, reason="TensorFlow not installed")
+def test_twpice_case_tensorflow():
+    """ Use a test case from TWP-ICE """
+    Grid0 = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR0)
+    Grid1 = pyart.io.read_grid(pydda.tests.EXAMPLE_RADAR1)
+    sounding = pyart.io.read_arm_sonde(pydda.tests.SOUNDING_PATH)
+
+    u_init, v_init, w_init = pydda.initialization.make_wind_field_from_profile(
+        Grid1, sounding[1], vel_field='corrected_velocity')
+    Grids = pydda.retrieval.get_dd_wind_field(
+        [Grid0, Grid1], u_init, v_init, w_init, Co=100, Cm=1500.0,
+        Cz=0, Cmod=0.0, vel_name='corrected_velocity', wind_tol=0.1,
+        refl_field='reflectivity', frz=5000.0, engine="tensorflow",
+        mask_outside_opt=True, upper_bc=1)
+
+    # In this test grid, we expect the mean flow to be to the southeast
+    # Maximum updrafts should be at least 10 m/s
+    u_mean = np.nanmean(Grids[0].fields['u']['data'])
+    v_mean = np.nanmean(Grids[0].fields['v']['data'])
+    w_max = np.max(Grids[0].fields['v']['data'])
+
+    assert u_mean > 0
+    assert v_mean < 0
+    assert w_max > 10
 
 
 def test_smoothing():
@@ -115,7 +133,7 @@ def test_smoothing():
     w = np.zeros((20, 40, 40))
     new_grids = pydda.retrieval.get_dd_wind_field(
         [Grid], u, v, w, Co=0.0, Cx=1e-4, Cy=1e-4, Cm=0.0, Cmod=0.0,
-        mask_outside_opt=False, filt_iterations=0, vel_name='one_field',
+        mask_outside_opt=False, vel_name='one_field',
         refl_field='one_field')
     new_u = new_grids[0].fields['u']['data']
     new_v = new_grids[0].fields['v']['data']
@@ -126,7 +144,6 @@ def test_smoothing():
                                                   Cx=1e-2, Cy=1e-2, Cm=0.0,
                                                   Cmod=0.0,
                                                   mask_outside_opt=False,
-                                                  filt_iterations=0,
                                                   vel_name='one_field',
                                                   refl_field='one_field')
     new_u2 = new_grids[0].fields['u']['data']
@@ -158,7 +175,7 @@ def test_model_constraint():
 
     new_grids = pydda.retrieval.get_dd_wind_field(
         [Grid0], u_init, v_init, w_init, Co=0.0, Cx=0.0, Cy=0.0, Cm=0.0,
-        Cmod=1.0, mask_outside_opt=False, filt_iterations=0,
+        Cmod=1.0, mask_outside_opt=False,
         vel_name='corrected_velocity', refl_field='reflectivity',
         model_fields=['fakemodel'])