def get_var(backend_model, var, dims=None, sparse=False): """ Return output for variable `var` as a pandas.Series (1d), pandas.Dataframe (2d), or xarray.DataArray (3d and higher). Parameters ---------- var : variable name as string, e.g. 'resource' dims : list, optional indices as strings, e.g. ('loc_techs', 'timesteps'); if not given, they are auto-detected sparse : bool, optional; default = False If extracting Pyomo Param data, the output sparse array includes inputs the user left as NaN replaced with the default value for that Param. """ try: var_container = getattr(backend_model, var) except AttributeError: raise exceptions.BackendError("Variable {} inexistent.".format(var)) if not dims: if var + "_index" == var_container.index_set().name: dims = [i.name for i in var_container.index_set().set_tuple] else: dims = [var_container.index_set().name] if sparse: result = pd.DataFrame.from_dict(var_container.extract_values_sparse(), orient="index") else: result = pd.DataFrame.from_dict(var_container.extract_values(), orient="index") if result.empty: raise exceptions.BackendError("Variable {} has no data.".format(var)) result = result[0] # Get the only column in the dataframe if len(dims) > 1: result.index = pd.MultiIndex.from_tuples(result.index, names=dims) if len(result.index.names) == 1: result = result.sort_index() result.index.name = dims[0] elif len(result.index.names) == 2: # if len(dims) is 2, we already have a well-formed DataFrame result = result.unstack(level=0) result = result.sort_index() else: # len(dims) >= 3 result = xr.DataArray.from_series(result) return result
def get_var(backend_model, var, dims=None, sparse=False, expr=False): """ Return output for variable `var` as a pandas.Series (1d), pandas.Dataframe (2d), or xarray.DataArray (3d and higher). Parameters ---------- var : variable name as string, e.g. 'resource' dims : list, optional indices as strings, e.g. ('loc_techs', 'timesteps'); if not given, they are auto-detected sparse : bool, optional; default = False If extracting Pyomo Param data, the output sparse array includes inputs the user left as NaN replaced with the default value for that Param. expr : bool, optional If True, treat var as a pyomo expression, which requires calculating the result of the expression before translating into nd data structure """ try: var_container = getattr(backend_model, var) except AttributeError: raise exceptions.BackendError("Variable {} inexistent.".format(var)) if not dims: if var + "_index" == var_container.index_set().name: dims = [i.name for i in var_container.index_set().subsets()] else: dims = [var_container.index_set().name] if sparse and not expr: if invalid(var_container.default()): result = pd.Series(var_container._data).apply( lambda x: po.value(x) if not invalid(x) else np.nan) else: result = pd.Series(var_container.extract_values_sparse()) else: if expr: result = pd.Series(var_container._data).apply(po.value) else: result = pd.Series(var_container.extract_values()) if result.empty: raise exceptions.BackendError("Variable {} has no data.".format(var)) result = result.rename_axis(index=dims) return xr.DataArray.from_series(result)