spencerahill / aospy

Python package for automated analysis and management of gridded climate data
Apache License 2.0
83 stars 12 forks source link

Calculations on only one variable #303

Closed jdossgollin closed 5 years ago

jdossgollin commented 5 years ago

When running calculations which use only one variable I run into

WARNING:root:Skipping aospy calculation `<aospy.Calc instance: meridional_gradient, cmip5_dynamical, reanalysis_model, reanalysis>` due to error with the following traceback:
Traceback (most recent call last):
  File "/Users/james/Documents/Github/2018-gcm-lfv/src/aospy/automate.py", line 272, in _compute_or_skip_on_error
    return calc.compute(**compute_kwargs)
  File "/Users/james/Documents/Github/2018-gcm-lfv/src/aospy/calc.py", line 437, in compute
    data = self._get_all_data(self.start_date, self.end_date)
  File "/Users/james/Documents/Github/2018-gcm-lfv/src/aospy/calc.py", line 335, in _get_all_data
    self.dtype_in_vert)]
  File "/Users/james/Documents/Github/2018-gcm-lfv/src/aospy/calc.py", line 32, in _replace_pressure
    for arg in arguments:
TypeError: 'Var' object is not iterable

calculations.py has the following content:

import xarray as xr
import numpy as np
from aospy.internal_names import LAT_STR, LON_STR

def _weighted_average(obj: xr.DataArray, dim: xr.DataArray, weights: np.ndarray):
    """
    xarray wrapper to np.average
    """
    return xr.apply_ufunc(
        np.average,
        obj,
        input_core_dims=[[dim]],
        kwargs={"axis": -1, "weights": weights,}
    )

def meridional_gradient(temp: xr.DataArray) -> xr.DataArray:
    """
    Calculate the equator to pole temperature contrast
    """
    weights = np.cos(np.deg2rad(temp[LAT_STR].values))
    zonal_mean = temp.mean(dim=LON_STR)
    return zonal_mean.differentiate(coord=LAT_STR)

and variables.py has

from aospy import Var
from .calculations import meridional_gradient

temp = Var(
    name="temp",
    alt_names=["air", "temperature"],
    def_time=True,
    def_lat=True,
    def_lon=True,
    def_vert=True,
    description="Air temperature",
    valid_range=(150, 350),
)
meridional_gradient = Var(
    name="meridional_gradient",
    def_time=True,
    def_lat=True,
    def_lon=False,
    def_vert=True,
    description="meridional temperature gradient",
    func=meridional_gradient,
    variables=(temp),
)

Interestingly, if I change things so that meridional_gradient takes in a second argument def meridional_gradient(temp: xr.DataArray, discard: xr.DataArray) -> xr.DataArray: and change the variables accordingly variables=(temp, temp), it works. Is there a way to run calculations that just take in a single variable?

spencerkclark commented 5 years ago

I think you may just need to change (temp) to (temp, ). The comma indicates that you want the result to be a length-1 tuple (otherwise Python just evaluates the expression in parentheses, which in the case of (temp) is just temp).

spencerahill commented 5 years ago

You can also specify them as a list rather than a tuple. That way, you don't have to worry about remembering the comma for length-1 cases: [temp] rather than (temp,)

jdossgollin commented 5 years ago

That works, thanks -- can I close this?