Source code for eradiate.xarray._helpers

import typing as t

import pandas as pd
import xarray as xr

XarrayObj = t.Union[xr.DataArray, xr.Dataset]

[docs]def unstack_mdistant_grid(obj: XarrayObj) -> XarrayObj: """ Reindex data produced by a :meth:`.MultiDistantMeasure.grid()` measure on (VAA, VZA) dimensions. This effectively reshapes the data, which is initially indexed on film width and height dimensions (the latter having a single element). Parameters ---------- obj : DataArray or Dataset An xarray object produced by processing a :meth:`.MultiDistantMeasure.grid()` measure, indexed on film dimensions. Returns ------- DataArray or Dataset An object of the same type, reindexed on VZA and VAA coordinates. """ # Collect metadata vza_attrs = obj["vza"].attrs vaa_attrs = obj["vaa"].attrs # Build new index idx = pd.MultiIndex.from_arrays( (obj.vza.values.flatten(), obj.vaa.values.flatten()), names=("vza", "vaa") ) # Reindex object result = ( obj.drop_vars( # First remove the target coords, as well as the indexing dim and associated coords # Also remove the x film coordinate (irrelevant after unstacking) ("vza", "vaa", "x_index", "x") ) .reindex(x_index=idx) .unstack() ) # Apply the new index and unstack # Reapply metadata result["vza"].attrs = vza_attrs result["vaa"].attrs = vaa_attrs return result