code: def brunt_vaisala_frequency(heights, potential_temperature, axis=0):
      r
      bv_freq_squared = brunt_vaisala_frequency_squared(heights, potential_temperature,
                                                        axis=axis)
      bv_freq_squared[bv_freq_squared.magnitude < 0] = np.nan
  
      return np.sqrt(bv_freq_squared)
docs: r"""Calculate the Brunt-Vaisala frequency.
  
      This function will calculate the Brunt-Vaisala frequency as follows:
  
      .. math:: N = \left( \frac{g}{\theta} \frac{d\theta}{dz} \right)^\frac{1}{2}
  
      This formula based off of Equations 3.75 and 3.77 in [Hobbs2006]_.
  
      This function is a wrapper for `brunt_vaisala_frequency_squared` that filters out negative
      (unstable) quanties and takes the square root.
  
      Parameters
      ----------
      heights : array-like
          One-dimensional profile of atmospheric height
      potential_temperature : array-like
          Atmospheric potential temperature
      axis : int, optional
          The axis corresponding to vertical in the potential temperature array, defaults to 0.
  
      Returns
      -------
      array-like
          Brunt-Vaisala frequency.
  
      See Also
      --------
      brunt_vaisala_frequency_squared, brunt_vaisala_period, potential_temperature