code: def brunt_vaisala_frequency_squared(heights, potential_temperature, axis=0):
      r
      # Ensure validity of temperature units
      potential_temperature = potential_temperature.to('K')
  
      # Calculate and return the square of Brunt-Vaisala frequency
      return mpconsts.g / potential_temperature * first_derivative(potential_temperature,
                                                                   x=heights, axis=axis)
docs: r"""Calculate the square of the Brunt-Vaisala frequency.
  
      Brunt-Vaisala frequency squared (a measure of atmospheric stability) is given by the
      formula:
  
      .. math:: N^2 = \frac{g}{\theta} \frac{d\theta}{dz}
  
      This formula is based off of Equations 3.75 and 3.77 in [Hobbs2006]_.
  
      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
          The square of the Brunt-Vaisala frequency.
  
      See Also
      --------
      brunt_vaisala_frequency, brunt_vaisala_period, potential_temperature