shnitsel.data.multi_indices =========================== .. py:module:: shnitsel.data.multi_indices Attributes ---------- .. autoapisummary:: shnitsel.data.multi_indices.DatasetOrArray Functions --------- .. autoapisummary:: shnitsel.data.multi_indices.midx_combs shnitsel.data.multi_indices.flatten_midx shnitsel.data.multi_indices.flatten_levels shnitsel.data.multi_indices.expand_midx shnitsel.data.multi_indices.assign_levels shnitsel.data.multi_indices.mgroupby shnitsel.data.multi_indices.msel shnitsel.data.multi_indices.sel_trajs shnitsel.data.multi_indices.sel_trajids shnitsel.data.multi_indices.unstack_trajs shnitsel.data.multi_indices.stack_trajs shnitsel.data.multi_indices.mdiff Module Contents --------------- .. py:data:: DatasetOrArray .. py:function:: midx_combs(values, name = None) Helper function to create a Multi-index based dimension coordinate for an xarray from all (unordered) pairwise combinations of entries in `values` :param values: The source values to generate pairwise combinations for :type values: pd.core.indexes.base.Index | list :param name: Optionally a name for the resulting combination dimension. Defaults to None. :type name: str | None, optional :raises ValueError: If no name was provided and the name could not be extracted from the `values` parameter :returns: The resulting coordinates object. :rtype: xr.Coordinates .. py:function:: flatten_midx(obj, idx_name, renamer = None) Function to flatten a multi-index into a flat index. Has the option to provide a custom renaming function :param obj: The object with the index intended to be flattened :type obj: xr.Dataset | xr.DataArray :param idx_name: The name of the index to flatten. :type idx_name: str :param renamer: An optional function to carry out the renaming of the combined entry from individual entries. Defaults to None. :type renamer: callable | None, optional :returns: The refactored object without the original index coordinates but with a combined index instead :rtype: xr.Dataset | xr.DataArray .. py:function:: flatten_levels(obj, idx_name, levels, new_name = None, position = 0, renamer = None) .. py:function:: expand_midx(obj, midx_name, level_name, value) .. py:function:: assign_levels(obj, levels = None, **levels_kwargs) Assign new values to levels of MultiIndexes in ``obj`` :param obj: An ``xarray`` object with at least one MultiIndex :param levels: A mapping whose keys are the names of the levels and whose values are the levels to assign. The mapping will be passed to :py:meth:`xarray.DataArray.assign_coords` (or the :py:class:`xarray.Dataset` equivalent). :param optional: A mapping whose keys are the names of the levels and whose values are the levels to assign. The mapping will be passed to :py:meth:`xarray.DataArray.assign_coords` (or the :py:class:`xarray.Dataset` equivalent). :rtype: A new object (of the same type as `obj`) with the new level values replacing the old level values. :raises ValueError: If levels are provided in both keyword and dictionary form. .. py:function:: mgroupby(obj, levels) Group a Dataset or DataArray by several levels of a MultiIndex it contains. :param obj: The :py:mod:`xr` object to group :param levels: Names of MultiIndex levels all belonging to the *same* MultiIndex :returns: * The grouped object, which behaves as documented at :py:meth:`xr.Dataset.groupby` * and `xr.DataArray.groupby` with the caveat that the specified levels have been * *"flattened" into a single Multiindex level of tuples.* :raises ValueError: If no MultiIndex is found, or if the named levels belong to different MultiIndexes. .. warning:: The function does not currently check whether the levels specified are really levels of a MultiIndex, as opposed to names of non-MultiIndex indexes. .. py:function:: msel(obj, **kwargs) .. py:function:: sel_trajs(frames, trajids_or_mask, invert=False) Select trajectories using a list of trajectories IDs or a boolean mask :param frames: The :py:class:`xr.Dataset` from which a selection is to be drawn :param trajids_or_mask: Either - A sequences of integers representing trajectory IDs to be included, in which case the trajectories **may not be returned in the order specified**. - Or a sequence of booleans, each indicating whether the trajectory with an ID in the corresponding entry in the ``Dataset``'s ``trajid_`` coordinate should be included :param invert: Whether to invert the selection, i.e. return those trajectories not specified, by default False :param optional: Whether to invert the selection, i.e. return those trajectories not specified, by default False :rtype: A new :py:class:`xr.Dataset` containing only the specified trajectories :raises NotImplementedError: when an attempt is made to index an :py:class:`xr.Datset` without a ``trajid_`` dimension/coordinate using a boolean mask :raises TypeError: If ``trajids_or_mask`` has a dtype other than integer or boolean .. py:function:: sel_trajids(frames, trajids, invert=False) Will not generally return trajectories in order given .. py:function:: unstack_trajs(frames) Unstack the ``frame`` MultiIndex so that ``trajid`` and ``time`` become separate dims. Wraps the :py:meth:`xarray.Dataset.unstack` method. :param frames: An :py:class:`xarray.Dataset` with a ``frame`` dimension associated with a MultiIndex coordinate with levels named ``trajid`` and ``time``. The Dataset may also have a ``trajid_`` dimension used for variables and coordinates that store information pertaining to each trajectory in aggregate; this will be aligned along the ``trajid`` dimension of the unstacked Dataset. :param DatasetOrArray: An :py:class:`xarray.Dataset` with a ``frame`` dimension associated with a MultiIndex coordinate with levels named ``trajid`` and ``time``. The Dataset may also have a ``trajid_`` dimension used for variables and coordinates that store information pertaining to each trajectory in aggregate; this will be aligned along the ``trajid`` dimension of the unstacked Dataset. :returns: * An :py:class:`xarray.Dataset` with independent ``trajid`` and ``time`` * dimensions. Same type as `frames` .. py:function:: stack_trajs(unstacked) Stack the ``trajid`` and ``time`` dims of an unstacked Dataset into a MultiIndex along a new dimension called ``frame``. Wraps the :py:meth:`xarray.Dataset.stack` method. :param frames: An :py:class:`xarray.Dataset` with independent ``trajid`` and ``time`` dimensions. :returns: * An :py:class:`xarray.Dataset` with a ``frame`` dimension associated with * a MultiIndex coordinate with levels named ``trajid`` and ``time``. Those variables * and coordinates which only depended on one of ``trajid`` * or ``time`` but not the other in the unstacked Dataset, will be aligned along new * dimensions named ``trajid_`` and ``time_``. The new dimensions ``trajid_`` and * ``time_`` will be independent of the ``frame`` dimension and its ``trajid`` and * ``time`` levels. .. py:function:: mdiff(da) Take successive differences along the 'frame' dimension :param da: An ``xarray.DataArray`` with a 'frame' dimension corresponding to a ``pandas.MultiIndex`` of which the innermost level is 'time'. :returns: * An ``xarray.DataArray`` with the same shape, dimension names etc., * *but with the data of the (i)th frame replaced by the difference between* * *the original (i+1)th and (i)th frames, with zeros filling in for both the* * *initial frame and any frame for which time = 0, to avoid taking differences* * *between the last and first frames of successive trajectories.*