diff --git a/src/parcels/_core/utils/sgrid.py b/src/parcels/_core/utils/sgrid.py index 2ee70bfce..2db714964 100644 --- a/src/parcels/_core/utils/sgrid.py +++ b/src/parcels/_core/utils/sgrid.py @@ -15,6 +15,7 @@ import re from collections.abc import Hashable, Iterable from dataclasses import dataclass +from textwrap import indent from typing import Any, Literal, Protocol, Self, cast, overload import xarray as xr @@ -26,6 +27,10 @@ Dim = str +def _indent_lines(lst: list[str], indent: int = 2): + return [indent * " " + line for line in lst] + + class Padding(enum.Enum): NONE = "none" LOW = "low" @@ -143,6 +148,9 @@ def __init__( def __repr__(self) -> str: return repr_from_dunder_dict(self) + def __str__(self) -> str: + return _grid2d_to_ascii(self) + def __eq__(self, other: Any) -> bool: if not isinstance(other, Grid2DMetadata): return NotImplemented @@ -256,6 +264,9 @@ def __init__( def __repr__(self) -> str: return repr_from_dunder_dict(self) + def __str__(self) -> str: + return _grid3d_to_ascii(self) + def __eq__(self, other: Any) -> bool: if not isinstance(other, Grid3DMetadata): return NotImplemented @@ -333,6 +344,9 @@ def load(cls, s: str) -> Self: padding = Padding(match.group(3).lower()) return cls(dim1, dim2, padding) + def to_diagram(self) -> str: + return "\n".join(_face_node_padding_to_text(self)) + def dump_mappings(parts: Iterable[DimDimPadding | Dim]) -> str: """Takes in a list of edge-node-padding tuples and serializes them into a string @@ -505,6 +519,181 @@ def get_unique_names(grid: Grid2DMetadata | Grid3DMetadata) -> set[str]: return dims +def _face_node_padding_to_text(obj: DimDimPadding) -> list[str]: + """Return ASCII diagram lines showing a face-node padding relationship. + + Produces a symbolic 5-node diagram like the image below, matching the + four padding modes:: + + face:node (padding:none) + ●─────●─────●─────●─────● + 1 1 2 2 3 3 4 4 5 + + face:node (padding:low) + ─────●─────●─────●─────●─────● + 1 1 2 2 3 3 4 4 5 5 + + face:node (padding:high) + ●─────●─────●─────●─────●───── + 1 1 2 2 3 3 4 4 5 5 + + face:node (padding:both) + ─────●─────●─────●─────●─────●───── + 1 1 2 2 3 3 4 4 5 5 6 + """ + FACE_WIDTH = 5 # dashes per face segment + padding = obj.padding + + bars = { + Padding.NONE: "x-x-x-x-x", + Padding.LOW: "-x-x-x-x-x", + Padding.HIGH: "x-x-x-x-x-", + Padding.BOTH: "-x-x-x-x-x-", + } + bar = bars[obj.padding] + node_count = 0 + face_count = 0 + bar_rendered = "" + label = "" + for char in bar: + if char == "x": + bar_rendered += "●" + label += str(node_count) + node_count += 1 + elif char == "-": + bar_rendered += "─" * FACE_WIDTH + label += str(face_count).center(FACE_WIDTH) + face_count += 1 + + return [ + f"{obj.dim1}:{obj.dim2} (padding:{padding.value})", + f" {bar_rendered}", + f" {label.rstrip()}", + ] + + +_TEXT_GRID2D_WITHOUT_Z = """ +Staggered grid layout (symbolic 3x3 nodes): + + ↑ Y + | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n --→ X + + n = node ({n1}, {n2}) + u = x-face ({u}) + v = y-face ({v}) + · = cell centre""" + +_TEXT_GRID2D_WITH_Z = """ +Staggered grid layout (symbolic 3x3 nodes): + + ↑ Y ↑ Z + | | + n --u-- n --u-- n w + | | | | + v · v · v · + | | | | + n --u-- n --u-- n w + | | | | + v · v · v · + | | | | + n --u-- n --u-- n --→ X w + + n = node ({n1}, {n2}) + u = x-face ({u}) + v = y-face ({v}) + w = z-node ({w}) + · = cell centre""" + +_TEXT_GRID3D = """ +Staggered grid layout (XY cross-section; Z-faces not shown): + + ↑ Y + | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n --→ X + + n = node ({n1}, {n2}, {n3}) + u = x-face ({u}) + v = y-face ({v}) + w = z-face ({w}) [not shown in cross-section] + · = cell centre""" + + +def _grid2d_to_ascii(grid: Grid2DMetadata) -> str: + fd = grid.face_dimensions + nd = grid.node_dimensions + lines = [ + "Grid2DMetadata", + f" X-axis: face={fd[0].dim1!r} node={nd[0]!r} padding={fd[0].padding.value}", + f" Y-axis: face={fd[1].dim1!r} node={nd[1]!r} padding={fd[1].padding.value}", + ] + if grid.vertical_dimensions: + vd = grid.vertical_dimensions[0] + lines.append(f" Z-axis: face={vd.dim1!r} node={vd.dim2!r} padding={vd.padding.value}") + if grid.node_coordinates: + lines.append(f" Coordinates: {grid.node_coordinates[0]}, {grid.node_coordinates[1]}") + + format_kwargs = dict(n1=nd[0], n2=nd[1], u=fd[0].dim1, v=fd[1].dim1) + + if grid.vertical_dimensions: + format_kwargs["w"] = grid.vertical_dimensions[0].dim2 + lines += indent(_TEXT_GRID2D_WITH_Z, " ").format(**format_kwargs).split("\n") + else: + lines += indent(_TEXT_GRID2D_WITHOUT_Z, " ").format(**format_kwargs).split("\n") + + lines += ["", " Axis padding:", ""] + lines += _indent_lines(_face_node_padding_to_text(fd[0])) + lines += [""] + lines += _indent_lines(_face_node_padding_to_text(fd[1])) + if grid.vertical_dimensions: + lines += [""] + lines += _indent_lines(_face_node_padding_to_text(grid.vertical_dimensions[0])) + return "\n".join(lines) + + +def _grid3d_to_ascii(grid: Grid3DMetadata) -> str: + vd = grid.volume_dimensions + nd = grid.node_dimensions + lines = [ + "Grid3DMetadata", + f" X-axis: face={vd[0].dim1!r} node={nd[0]!r} padding={vd[0].padding.value}", + f" Y-axis: face={vd[1].dim1!r} node={nd[1]!r} padding={vd[1].padding.value}", + f" Z-axis: face={vd[2].dim1!r} node={nd[2]!r} padding={vd[2].padding.value}", + ] + if grid.node_coordinates: + lines.append(f" Coordinates: {', '.join(grid.node_coordinates)}") + + lines += ( + indent(_TEXT_GRID3D, " ") + .format(n1=nd[0], n2=nd[1], n3=nd[2], u=vd[0].dim1, v=vd[1].dim1, w=vd[2].dim1) + .split("\n") + ) + + lines += ["", " Axis padding:", ""] + lines += _indent_lines(_face_node_padding_to_text(vd[0])) + lines += [""] + lines += _indent_lines(_face_node_padding_to_text(vd[1])) + lines += [""] + lines += _indent_lines(_face_node_padding_to_text(vd[2])) + return "\n".join(lines) + + def _attach_sgrid_metadata(ds, grid: Grid2DMetadata | Grid3DMetadata): """Copies the dataset and attaches the SGRID metadata in 'grid' variable. Modifies 'conventions' attribute.""" ds = ds.copy() diff --git a/tests/utils/test_sgrid.py b/tests/utils/test_sgrid.py index c3f3d666a..43588a987 100644 --- a/tests/utils/test_sgrid.py +++ b/tests/utils/test_sgrid.py @@ -348,3 +348,216 @@ def test_rename_dataset(ds): assert "XC_updated" in ds_new.dims assert "XC" not in ds_new.dims assert "XC_updated" == grid_new.face_dimensions[0].dim1 + + +@pytest.mark.parametrize( + ("metadata, expected"), + [ + ( + create_example_grid2dmetadata(with_vertical_dimensions=False, with_node_coordinates=False), + """Grid2DMetadata + X-axis: face='face_dimension1' node='node_dimension1' padding=low + Y-axis: face='face_dimension2' node='node_dimension2' padding=low + + Staggered grid layout (symbolic 3x3 nodes): + + ↑ Y + | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n --→ X + + n = node (node_dimension1, node_dimension2) + u = x-face (face_dimension1) + v = y-face (face_dimension2) + · = cell centre + + Axis padding: + + face_dimension1:node_dimension1 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + face_dimension2:node_dimension2 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4""", + ), + ( + create_example_grid2dmetadata(with_vertical_dimensions=True, with_node_coordinates=True), + """Grid2DMetadata + X-axis: face='face_dimension1' node='node_dimension1' padding=low + Y-axis: face='face_dimension2' node='node_dimension2' padding=low + Z-axis: face='vertical_dimensions_dim1' node='vertical_dimensions_dim2' padding=low + Coordinates: node_coordinates_var1, node_coordinates_var2 + + Staggered grid layout (symbolic 3x3 nodes): + + ↑ Y ↑ Z + | | + n --u-- n --u-- n w + | | | | + v · v · v · + | | | | + n --u-- n --u-- n w + | | | | + v · v · v · + | | | | + n --u-- n --u-- n --→ X w + + n = node (node_dimension1, node_dimension2) + u = x-face (face_dimension1) + v = y-face (face_dimension2) + w = z-node (vertical_dimensions_dim2) + · = cell centre + + Axis padding: + + face_dimension1:node_dimension1 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + face_dimension2:node_dimension2 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + vertical_dimensions_dim1:vertical_dimensions_dim2 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4""", + ), + ( + create_example_grid3dmetadata(with_node_coordinates=False), + """Grid3DMetadata + X-axis: face='face_dimension1' node='node_dimension1' padding=low + Y-axis: face='face_dimension2' node='node_dimension2' padding=low + Z-axis: face='face_dimension3' node='node_dimension3' padding=low + + Staggered grid layout (XY cross-section; Z-faces not shown): + + ↑ Y + | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n --→ X + + n = node (node_dimension1, node_dimension2, node_dimension3) + u = x-face (face_dimension1) + v = y-face (face_dimension2) + w = z-face (face_dimension3) [not shown in cross-section] + · = cell centre + + Axis padding: + + face_dimension1:node_dimension1 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + face_dimension2:node_dimension2 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + face_dimension3:node_dimension3 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4""", + ), + ( + create_example_grid3dmetadata(with_node_coordinates=True), + """Grid3DMetadata + X-axis: face='face_dimension1' node='node_dimension1' padding=low + Y-axis: face='face_dimension2' node='node_dimension2' padding=low + Z-axis: face='face_dimension3' node='node_dimension3' padding=low + Coordinates: node_coordinates_var1, node_coordinates_var2, node_coordinates_dim3 + + Staggered grid layout (XY cross-section; Z-faces not shown): + + ↑ Y + | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n + | | | + v · v · v + | | | + n --u-- n --u-- n --→ X + + n = node (node_dimension1, node_dimension2, node_dimension3) + u = x-face (face_dimension1) + v = y-face (face_dimension2) + w = z-face (face_dimension3) [not shown in cross-section] + · = cell centre + + Axis padding: + + face_dimension1:node_dimension1 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + face_dimension2:node_dimension2 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4 + + face_dimension3:node_dimension3 (padding:low) + ─────●─────●─────●─────●─────● + 0 0 1 1 2 2 3 3 4 4""", + ), + ], +) +def test_grid_str(metadata, expected): + actual = str(metadata) + assert actual == expected + + +@pytest.mark.parametrize( + ("face_node_padding", "expected_lines"), + [ + ( + sgrid.DimDimPadding("face", "node", sgrid.Padding.LOW), + [ + "face:node (padding:low)", + " ─────●─────●─────●─────●─────●", + " 0 0 1 1 2 2 3 3 4 4", + ], + ), + ( + sgrid.DimDimPadding("face", "node", sgrid.Padding.HIGH), + [ + "face:node (padding:high)", + " ●─────●─────●─────●─────●─────", + " 0 0 1 1 2 2 3 3 4 4", + ], + ), + ( + sgrid.DimDimPadding("face", "node", sgrid.Padding.BOTH), + [ + "face:node (padding:both)", + " ─────●─────●─────●─────●─────●─────", + " 0 0 1 1 2 2 3 3 4 4 5", + ], + ), + ( + sgrid.DimDimPadding("face", "node", sgrid.Padding.NONE), + [ + "face:node (padding:none)", + " ●─────●─────●─────●─────●", + " 0 0 1 1 2 2 3 3 4", + ], + ), + ], +) +def test_face_node_padding_to_diagram(face_node_padding: sgrid.DimDimPadding, expected_lines: list[str]): + actual = face_node_padding.to_diagram() + lines = actual.split("\n") + assert lines == expected_lines