<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def brunt_vaisala_frequency_squared(heights, potential_temperature, axis=0):
 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 """

    # 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)