shnitsel.core.xrhelpers ======================= .. py:module:: shnitsel.core.xrhelpers Functions --------- .. autoapisummary:: shnitsel.core.xrhelpers.replace_total shnitsel.core.xrhelpers.midx_combs shnitsel.core.xrhelpers.flatten_midx shnitsel.core.xrhelpers.flatten_levels shnitsel.core.xrhelpers.expand_midx shnitsel.core.xrhelpers.assign_levels shnitsel.core.xrhelpers.open_frames shnitsel.core.xrhelpers.save_frames shnitsel.core.xrhelpers.split_for_saving shnitsel.core.xrhelpers.save_split shnitsel.core.xrhelpers.mgroupby shnitsel.core.xrhelpers.msel shnitsel.core.xrhelpers.sel_trajs shnitsel.core.xrhelpers.sel_trajids shnitsel.core.xrhelpers.unstack_trajs shnitsel.core.xrhelpers.stack_trajs Module Contents --------------- .. py:function:: replace_total(da, to_replace, value) Replaces each occurence of `to_replace` in `da` with the corresponding element of `value`. Replacement must be total, i.e. every element of `da` must be in `to_replace`. This permits a change of dtype between `to_replace` and `value`. This function is based on the snippets at https://github.com/pydata/xarray/issues/6377 :param da: An xr.DataArray :param to_replace: Values to replace :param value: Values with which to replace them :rtype: An xr.DataArray with dtype matching `value`. .. py:function:: midx_combs(values, name = None) .. py:function:: flatten_midx(obj, idx_name, renamer = None) .. 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 with the new level values replacing the old level values. :raises ValueError: If levels are provided in both keyword and dictionary form. .. py:function:: open_frames(path) Opens a NetCDF4 file saved by shnitsel-tools, specially interpreting certain attributes. :param path: The path of the file to open. :rtype: An :py:class:`xarray.Dataset` with any MultiIndex restored. :raises FileNotFoundError: If there is is nothing at ``path``, or ``path`` is not a file. :raises ValueError (or other exception): Raised by the underlying `h5netcdf `_ engine if the file is corrupted. .. py:function:: save_frames(frames, path, complevel=9) Save a ``Dataset``, presumably (but not necessarily) consisting of frames of trajectories, to a file at ``path``. :param frames (omit if using accessor): The ``Dataset`` to save :param path: The path at which to save it :param complevel: The level of ``gzip`` compression which will be applied to all variables in the ``Dataset``, by default 9 :param optional: The level of ``gzip`` compression which will be applied to all variables in the ``Dataset``, by default 9 .. rubric:: Notes This function/accessor method wraps :py:meth:`xarray.Dataset.to_netcdf` but not :py:func:`numpy.any`. .. py:function:: split_for_saving(frames, bytes_per_chunk=50000000.0) .. py:function:: save_split(frames, path_template, bytes_per_chunk=50000000.0, complevel=9, ignore_errors=False) .. 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. :returns: * An :py:class:`xarray.Dataset` with independent ``trajid`` and ``time`` * *dimensions.* .. 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.