Inverting a MetPy function

[sgdecker]

MetPy provides saturation_equivalent_potential_temperature, but I'd like to compute temperature given theta-es and pressure (assuming the parcel is saturated). I tried the following code:

from scipy.optimize import minimize_scalar

from metpy.units import units
import metpy.calc as mpcalc

def theta_e_to_temp(thetae, p):
    """Given thetae and p, find temperature"""
    def objective(T):
        return mpcalc.saturation_equivalent_potential_temperature(p, T) - thetae
    
    solution = minimize_scalar(objective)
    return solution

p = 850 * units('hPa')
T = 280 * units('degK')

thetae = mpcalc.saturation_equivalent_potential_temperature(p, T)

T = theta_e_to_temp(thetae, p)

where I am hoping at the end, T (T.x really) will be within reasonable round-off error of 280 K. However, as I sort of expected, the units are an issue:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 T = theta_e_to_temp(thetae, p)

Input In [2], in theta_e_to_temp(thetae, p)
      3 def objective(T):
      4     return mpcalc.saturation_equivalent_potential_temperature(p, T) - thetae
----> 6 solution = minimize_scalar(objective)
      7 return solution

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/scipy/optimize/_minimize.py:886, in minimize_scalar(fun, bracket, bounds, args, method, tol, options)
    884     return method(fun, args=args, bracket=bracket, bounds=bounds, **options)
    885 elif meth == 'brent':
--> 886     return _minimize_scalar_brent(fun, bracket, args, **options)
    887 elif meth == 'bounded':
    888     if bounds is None:

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/scipy/optimize/_optimize.py:2498, in _minimize_scalar_brent(func, brack, args, xtol, maxiter, disp, **unknown_options)
   2495 brent = Brent(func=func, args=args, tol=tol,
   2496               full_output=True, maxiter=maxiter, disp=disp)
   2497 brent.set_bracket(brack)
-> 2498 brent.optimize()
   2499 x, fval, nit, nfev = brent.get_result(full_output=True)
   2501 success = nit < maxiter and not (np.isnan(x) or np.isnan(fval))

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/scipy/optimize/_optimize.py:2268, in Brent.optimize(self)
   2265 def optimize(self):
   2266     # set up for optimization
   2267     func = self.func
-> 2268     xa, xb, xc, fa, fb, fc, funcalls = self.get_bracket_info()
   2269     _mintol = self._mintol
   2270     _cg = self._cg

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/scipy/optimize/_optimize.py:2242, in Brent.get_bracket_info(self)
   2239 ### BEGIN core bracket_info code ###
   2240 ### carefully DOCUMENT any CHANGES in core ##
   2241 if brack is None:
-> 2242     xa, xb, xc, fa, fb, fc, funcalls = bracket(func, args=args)
   2243 elif len(brack) == 2:
   2244     xa, xb, xc, fa, fb, fc, funcalls = bracket(func, xa=brack[0],
   2245                                                xb=brack[1], args=args)

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/scipy/optimize/_optimize.py:2752, in bracket(func, xa, xb, args, grow_limit, maxiter)
   2750 _gold = 1.618034  # golden ratio: (1.0+sqrt(5.0))/2.0
   2751 _verysmall_num = 1e-21
-> 2752 fa = func(*(xa,) + args)
   2753 fb = func(*(xb,) + args)
   2754 if (fa < fb):                      # Switch so fa > fb

Input In [2], in theta_e_to_temp.<locals>.objective(T)
      3 def objective(T):
----> 4     return mpcalc.saturation_equivalent_potential_temperature(p, T) - thetae

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/metpy/xarray.py:1235, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
   1232     _mutate_arguments(bound_args, units.Quantity, lambda arg, _: arg.m)
   1234 # Evaluate inner calculation
-> 1235 result = func(*bound_args.args, **bound_args.kwargs)
   1237 # Wrap output based on match and match_unit
   1238 if match is None:

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/metpy/units.py:295, in check_units.<locals>.dec.<locals>.wrapper(*args, **kwargs)
    293 @functools.wraps(func)
    294 def wrapper(*args, **kwargs):
--> 295     _check_units_inner_helper(func, sig, defaults, dims, *args, **kwargs)
    296     return func(*args, **kwargs)

File /usr/local/python/3.8/envs/met433/lib/python3.9/site-packages/metpy/units.py:282, in _check_units_inner_helper(func, sig, defaults, dims, *args, **kwargs)
    277         if changed_version:
    278             msg = (
    279                 f'This function changed in {changed_version}--double check '
    280                 'that the function is being called properly.\n'
    281             ) + msg
--> 282     raise ValueError(msg)
    284 # Return the bound arguments for reuse
    285 return bound_args

ValueError: `saturation_equivalent_potential_temperature` given arguments with incorrect units: `temperature` requires "[temperature]" but given "none"
A xarray DataArray or numpy array `x` can be assigned a unit as follows:
    from metpy.units import units
    x = x * units("m/s")
For more information see the Units Tutorial: https://unidata.github.io/MetPy/latest/tutorials/unit_tutorial.html

I figured I'd ask here before devoting too much time to possibly reinventing the wheel. Is there an easy way around this?

[sgdecker]

Oops, I should be using root_scalar instead, and I now have the idea to wrap saturation_equivalent_potential_temperature in my own function that adds and then removes units. Don't have time to realize the idea at the moment, but I think I'm on the right track.

[dopplershift]

To make that work, inside objective you'd want to attach units before calling saturation_equivalent_potential_temperature, then convert back with .m_as('kelvin') for the return.

[sgdecker]

For completeness, here's the working code.

import numpy as np

from scipy.optimize import root_scalar

from metpy.units import units
import metpy.calc as mpcalc
import metpy.interpolate as interp


def theta_e_to_temp(thetae, p):
    """Given thetae and p, find temperature"""
    
    p1 = p * units('Pa')
    
    def objective(T):
        T1 = T * units('degK')
        result = mpcalc.saturation_equivalent_potential_temperature(p1, T1).m_as('kelvin') - thetae
        return result

    solution = root_scalar(objective, bracket=(200, 365))
    return solution


p = 85000 * units('Pa')
T = 280 * units('degK')

thetae = mpcalc.saturation_equivalent_potential_temperature(p, T)
print(thetae)

thetae = thetae.m
p = p.m
T = theta_e_to_temp(thetae, p)
print(T.root)