Deltares / xugrid

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

ENH: method to rename grid topology #143

Closed hboisgon closed 1 year ago

hboisgon commented 1 year ago

It might be useful to have a method to be able to rename the grid.name in xugrid (eg if models have different expected names for their grids).

If it is useful, I implemented the following in hydromt:

def rename_mesh(mesh: Union[xu.UgridDataArray, xu.UgridDataset], name: str):
    """
    Rename all grid variables in mesh according to UGRID conventions.
    Note: adapted from xugrid.ugrid.grid.rename to also work on
    UgridDataset and UgridDataArray
    """
    # Get the old and the new names. Their keys are the same.
    old_attrs = mesh.ugrid.grid._attrs
    new_attrs = conventions.default_topology_attrs(
        name, mesh.ugrid.grid.topology_dimension
    )

    # The attrs will have some roles joined together, e.g. node_coordinates
    # will contain x and y as "mesh2d_node_x mesh2d_node_y".
    name_dict = {mesh.ugrid.grid.name: name}
    skip = ("cf_role", "long_name", "topology_dimension")
    for key, value in old_attrs.items():
        if key in new_attrs and key not in skip:
            split_new = new_attrs[key].split()
            split_old = value.split()
            if len(split_new) != len(split_old):
                raise ValueError(
                    f"Number of entries does not match on {key}: "
                    f"{split_new} versus {split_old}"
                )
            for name_key, name_value in zip(split_old, split_new):
                name_dict[name_key] = name_value

    new = mesh.copy()
    new.ugrid.grid.name = name
    new.ugrid.grid._attrs = new_attrs
    new.ugrid.grid._indexes = {
        k: name_dict[v] for k, v in new.ugrid.grid._indexes.items()
    }

    to_rename = tuple(new.data_vars) + tuple(new.coords) + tuple(new.dims)
    new = new.rename({k: v for k, v in name_dict.items() if k in to_rename})

    return new
Huite commented 1 year ago

Hi @hboisgon,

I've pushed an implementation (via .ugrid.rename()) to main. Note that you can only rename a UgridDataset with a string if the dataset contains a single topology (just like .grid will raise an error if there are multiple grids). In case of multiple topologies, you will need to provide a dictionary giving the old and new names instead.