Deltares / xugrid

Xarray and unstructured grids
https://deltares.github.io/xugrid/
MIT License
61 stars 8 forks source link

support get mesh/grid names #141

Closed xldeltares closed 1 year ago

xldeltares commented 1 year ago

This issue is from @hboisgon: in hydromt mesh model, there is needs to get an overview of what grids are available in a UgridDataset. For now it is done as follows. It would be nice if mesh_names could be supported in xugrid.

@property
def mesh_grids(self) -> Dict:
    """Dictionnary of grid names and Ugrid topologies in mesh."""
    grids = dict()
    if self.mesh is not None:
        for grid in self.mesh.ugrid.grids:
            grids[grid.name] = grid

    return grids

@property
def mesh_names(self) -> List[str]:
    """List of grid names in mesh."""
    if self.mesh is not None:
        return list(self.mesh_grids.keys())
    else:
        return []
Huite commented 1 year ago

This seems like a very reasonable feature. I'm not sold yet on the names: "grid" and "mesh" are somewhat synonym, so that mesh_grid is a pleonasm. I'm happy with uds.ugrid.names or uda.ugrid.name.

Perhaps I should've reserved the grids property for a dict instead of a list as this is what xarray does for basically all the attributes: https://docs.xarray.dev/en/stable/api.html#attributes

But that is an awkward breaking change now... I'll try to find some inspiration for an appropriate term here.

Huite commented 1 year ago

I've added the name, names, and topology properties. You can now (with current main) do:

uds.ugrid.name  # -> str, will only work with a single grid contained in uds
uds.ugrid.names  # -> list[str]
uds.ugrid.topology # -> dict[str, Ugrid1d | Ugrid2d]

I've mostly gone with "topology" this more or resembles what I've used in the "ugrid_roles" accessor which works on ordinary xarray objects. (Xarray is also inconsistent: dims on the dataset will return a dict, dims on the DataArray will return a tuple...)

xldeltares commented 1 year ago

hi @Huite , thanks for the quick implementation. In hydromt, @hboisgon also implemented the method mesh_datasets as follows, which I am reviewing at the moment:

    def mesh_datasets(self) -> Dict[str, xu.UgridDataset]:
        """Dictionnary of grid names and corresponding UgridDataset topology and data variables in mesh."""  # noqa: E501
        datasets = dict()
        if self._mesh is not None:
            for grid in self.mesh.ugrid.grids:
                datasets[grid.name] = self.get_mesh(
                    grid_name=grid.name, include_data=True
                )

        return datasets

    def get_mesh(self, grid_name: str, include_data: bool = False) -> xu.UgridDataset:
        """
        Return a specific ugrid dataset from mesh based on grid_name.

        If include_data is True, the data variables (including optional attribtues)
        for that specific grid are also included.

        Parameters
        ----------
        grid_name : str
            Name of the grid to return.
        include_data : bool, optional
            If True, also include data variables, by default False.

        Returns
        -------
        uds: xu.UgridDataset
            UGrid dataset with or without data variables.
        """
        if self.mesh is None:
            raise ValueError("Mesh is not set, please use set_mesh first.")
        if grid_name not in self.mesh_names:
            raise ValueError(f"Grid {grid_name} not found in mesh.")
        if include_data:
            grid = self.mesh_grids[grid_name]
            # FIXME xugrid #138 #140
            uds = xu.UgridDataset(grid.to_dataset(optional_attributes=True))
            uds.ugrid.grid.set_crs(grid.crs)
            # Look for data_vars that are defined on grid_name
            for var in self.mesh.data_vars:
                if hasattr(self.mesh[var], "ugrid"):
                    if self.mesh[var].ugrid.grid.name == grid_name:
                        uds[var] = self.mesh[var]
                # additionnal topology properties
                elif var.startswith(grid_name):
                    uds[var] = self.mesh[var]
                # else is global property (not grid specific)
        else:
            uds = xu.UgridDataset(self.mesh_grids[grid_name].to_dataset())

        return uds

Based on that, would you see fit for a uds.ugrid.dataset # -> dict[str, UgridDataset] in xugrid? or this is better to keep in hydromt? Thanks!