From 21c31b1f6601df4da7f7f617289ce8075b1aac9d Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Fri, 9 Feb 2024 11:48:24 +0100 Subject: [PATCH 1/5] Add some rudimentary tests to test merging of 1D grids --- tests/test_partitioning.py | 64 +++++++++++++++++++++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/tests/test_partitioning.py b/tests/test_partitioning.py index 6caa98c08..26d25b9d6 100644 --- a/tests/test_partitioning.py +++ b/tests/test_partitioning.py @@ -24,6 +24,17 @@ def generate_mesh_2d(nx, ny, name="mesh2d"): return xu.Ugrid2d(*np.array(points).T, -1, np.array(connectivity), name=name) +def generate_mesh_1d(n, name="mesh1d"): + points = [ + (p,p) for p in np.linspace(0, n, n+1) + ] + connectivity = [ + [it, it+1] for it in range(n) + ] + + return xu.Ugrid1d(*np.array(points).T, -1, np.array(connectivity), name=name) + + def test_labels_to_indices(): labels = np.array([0, 1, 0, 2, 2]) indices = pt.labels_to_indices(labels) @@ -159,7 +170,7 @@ def test_merge_partitions_no_duplicates(self): assert np.bincount(merged["face_z"] == 1).all() -class TestMultiTopologyMergePartitions: +class TestMultiTopology2DMergePartitions: @pytest.fixture(autouse=True) def setup(self): grid_a = generate_mesh_2d(2, 3, "first") @@ -220,3 +231,54 @@ def test_merge_partitions__errors(self): TypeError, match="All partition topologies with name second" ): pt.merge_partitions([self.datasets[0], dataset3]) + +class TestMergeDataset1D: + @pytest.fixture(autouse=True) + def setup(self): + grid = generate_mesh_1d(6, "mesh1d") + # TODO: If partitioning implented for 1D grids, replace with that. + i_edges = [[0,1,2], [3,4,5]] + parts = [grid.isel(mesh1d_nEdges=np.array(ls)) for ls in i_edges] + + datasets = [] + for i, part in enumerate(parts): + ds = xu.UgridDataset(grids=[part]) + ds["a"] = ((part.edge_dimension), np.arange(part.n_edge)) + ds["c"] = i + datasets.append(ds) + + self.datasets = datasets + + def test_merge_partitions(self): + merged = pt.merge_partitions(self.datasets) + assert isinstance(merged, xu.UgridDataset) + assert len(merged.ugrid.grids) == 1 + # In case of non-UGRID data, it should default to the first partition: + assert merged["c"] == 0 + +class TestMultiTopology1D2DMergePartitions: + @pytest.fixture(autouse=True) + def setup(self): + grid_a = generate_mesh_2d(2, 3, "mesh2d") + grid_b = generate_mesh_1d(6, "mesh1d") + parts_a = grid_a.partition(n_part=2) + # TODO: If partitioning implented for 1D grids, replace with that. + i_edges = [[0,1,2], [3,4,5]] + parts_b = [grid_b.isel(mesh1d_nEdges=np.array(ls)) for ls in i_edges] + + datasets = [] + for i, (part_a, part_b) in enumerate(zip(parts_a, parts_b)): + ds = xu.UgridDataset(grids=[part_a, part_b]) + ds["a"] = ((part_a.face_dimension), np.arange(part_a.n_face)) + ds["b"] = ((part_b.edge_dimension), np.arange(part_b.n_edge)) + ds["c"] = i + datasets.append(ds) + + self.datasets = datasets + + def test_merge_partitions(self): + merged = pt.merge_partitions(self.datasets) + assert isinstance(merged, xu.UgridDataset) + assert len(merged.ugrid.grids) == 2 + # In case of non-UGRID data, it should default to the first partition: + assert merged["c"] == 0 \ No newline at end of file From 6551b0aeeb337a2eb7d14ec2fc900ba6b334a961 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Fri, 9 Feb 2024 11:48:52 +0100 Subject: [PATCH 2/5] Provide indexes as argument --- xugrid/ugrid/ugrid1d.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xugrid/ugrid/ugrid1d.py b/xugrid/ugrid/ugrid1d.py index f553595ab..a5f705d16 100644 --- a/xugrid/ugrid/ugrid1d.py +++ b/xugrid/ugrid/ugrid1d.py @@ -662,6 +662,7 @@ def merge_partitions(grids: Sequence["Ugrid1d"]) -> "Ugrid1d": fill_value, new_edges, name=grid.name, + indexes=grid._indexes, projected=grid.projected, crs=grid.crs, attrs=grid._attrs, From d1adfe01a67a7c9bb6c6de32904b629f9b2cf0da Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Fri, 9 Feb 2024 14:04:14 +0100 Subject: [PATCH 3/5] Add testcases for merging 1D and 1D2D datasets --- tests/test_partitioning.py | 60 ++++++++++++++++++++++++++++---------- 1 file changed, 44 insertions(+), 16 deletions(-) diff --git a/tests/test_partitioning.py b/tests/test_partitioning.py index 26d25b9d6..3be173e91 100644 --- a/tests/test_partitioning.py +++ b/tests/test_partitioning.py @@ -236,49 +236,77 @@ class TestMergeDataset1D: @pytest.fixture(autouse=True) def setup(self): grid = generate_mesh_1d(6, "mesh1d") - # TODO: If partitioning implented for 1D grids, replace with that. + # TODO: If partitioning implemented for 1D grids, replace with that. i_edges = [[0,1,2], [3,4,5]] parts = [grid.isel(mesh1d_nEdges=np.array(ls)) for ls in i_edges] - datasets = [] - for i, part in enumerate(parts): + values_parts = [np.arange(part.n_edge) for part in parts] + + datasets_partitioned = [] + for i, (part, values) in enumerate(zip(parts, values_parts)): ds = xu.UgridDataset(grids=[part]) - ds["a"] = ((part.edge_dimension), np.arange(part.n_edge)) + ds["a"] = ((part.edge_dimension), values) ds["c"] = i - datasets.append(ds) + datasets_partitioned.append(ds) - self.datasets = datasets + ds_expected = xu.UgridDataset(grids=[grid]) + ds_expected["a"] = ((grid.edge_dimension), np.concatenate(values_parts)) + ds_expected["c"] = 0 + # Assign coordinates also added during merge_partitions + coords = {grid.edge_dimension: np.arange(grid.n_edge)} + ds_expected = ds_expected.assign_coords(**coords) + + self.datasets_partitioned = datasets_partitioned + self.dataset_expected = ds_expected def test_merge_partitions(self): - merged = pt.merge_partitions(self.datasets) + merged = pt.merge_partitions(self.datasets_partitioned) assert isinstance(merged, xu.UgridDataset) assert len(merged.ugrid.grids) == 1 # In case of non-UGRID data, it should default to the first partition: assert merged["c"] == 0 + assert self.dataset_expected.ugrid.grid.equals(merged.ugrid.grid) + assert self.dataset_expected["a"].equals(merged["a"]) + assert self.dataset_expected.equals(merged) + class TestMultiTopology1D2DMergePartitions: @pytest.fixture(autouse=True) def setup(self): grid_a = generate_mesh_2d(2, 3, "mesh2d") grid_b = generate_mesh_1d(6, "mesh1d") parts_a = grid_a.partition(n_part=2) - # TODO: If partitioning implented for 1D grids, replace with that. + # TODO: If partitioning implemented for 1D grids, replace with that. i_edges = [[0,1,2], [3,4,5]] parts_b = [grid_b.isel(mesh1d_nEdges=np.array(ls)) for ls in i_edges] - datasets = [] - for i, (part_a, part_b) in enumerate(zip(parts_a, parts_b)): + values_parts_a = [np.arange(part.n_face) for part in parts_a] + values_parts_b = [np.arange(part.n_edge) for part in parts_b] + + datasets_parts = [] + for i, (part_a, part_b, values_a, values_b) in enumerate(zip(parts_a, parts_b, values_parts_a, values_parts_b)): ds = xu.UgridDataset(grids=[part_a, part_b]) - ds["a"] = ((part_a.face_dimension), np.arange(part_a.n_face)) - ds["b"] = ((part_b.edge_dimension), np.arange(part_b.n_edge)) + ds["a"] = ((part_a.face_dimension), values_a) + ds["b"] = ((part_b.edge_dimension), values_b) ds["c"] = i - datasets.append(ds) + datasets_parts.append(ds) - self.datasets = datasets + ds_expected = xu.UgridDataset(grids=[grid_a, grid_b]) + ds_expected["a"] = ((grid_a.face_dimension), np.concatenate(values_parts_a)) + ds_expected["b"] = ((grid_b.edge_dimension), np.concatenate(values_parts_b)) + ds_expected["c"] = 0 + # Assign coordinates also added during merge_partitions + coords = {grid_a.face_dimension: np.arange(grid_a.n_face), grid_b.edge_dimension: np.arange(grid_b.n_edge)} + ds_expected = ds_expected.assign_coords(**coords) + + self.datasets_parts = datasets_parts + self.dataset_expected = ds_expected def test_merge_partitions(self): - merged = pt.merge_partitions(self.datasets) + merged = pt.merge_partitions(self.datasets_parts) assert isinstance(merged, xu.UgridDataset) assert len(merged.ugrid.grids) == 2 # In case of non-UGRID data, it should default to the first partition: - assert merged["c"] == 0 \ No newline at end of file + assert merged["c"] == 0 + + assert self.dataset_expected.equals(merged) \ No newline at end of file From d62f7369caa5844499872c29678d23f30bb8b87a Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Fri, 9 Feb 2024 14:08:18 +0100 Subject: [PATCH 4/5] update changelog --- docs/changelog.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/changelog.rst b/docs/changelog.rst index e2bcafffa..34e327ea0 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -17,6 +17,8 @@ Fixed - Fixed bug in :func:`xugrid.concat` and :func:`xugrid.merge` where multiple grids were returned if grids did not point to the same object id (i.e. copies). +- Fixed bug in :meth:`xugrid.Ugrid1d.merge_partitions`, which caused + ``ValueError: indexes must be provided for attrs``. Added ~~~~~ From dbe49f115e8e38d67e781bc0b411643fd51dbc82 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Fri, 9 Feb 2024 14:11:01 +0100 Subject: [PATCH 5/5] format --- tests/test_partitioning.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/tests/test_partitioning.py b/tests/test_partitioning.py index 3be173e91..04596cf4a 100644 --- a/tests/test_partitioning.py +++ b/tests/test_partitioning.py @@ -25,13 +25,9 @@ def generate_mesh_2d(nx, ny, name="mesh2d"): def generate_mesh_1d(n, name="mesh1d"): - points = [ - (p,p) for p in np.linspace(0, n, n+1) - ] - connectivity = [ - [it, it+1] for it in range(n) - ] - + points = [(p, p) for p in np.linspace(0, n, n + 1)] + connectivity = [[it, it + 1] for it in range(n)] + return xu.Ugrid1d(*np.array(points).T, -1, np.array(connectivity), name=name) @@ -232,12 +228,13 @@ def test_merge_partitions__errors(self): ): pt.merge_partitions([self.datasets[0], dataset3]) + class TestMergeDataset1D: @pytest.fixture(autouse=True) def setup(self): grid = generate_mesh_1d(6, "mesh1d") # TODO: If partitioning implemented for 1D grids, replace with that. - i_edges = [[0,1,2], [3,4,5]] + i_edges = [[0, 1, 2], [3, 4, 5]] parts = [grid.isel(mesh1d_nEdges=np.array(ls)) for ls in i_edges] values_parts = [np.arange(part.n_edge) for part in parts] @@ -270,6 +267,7 @@ def test_merge_partitions(self): assert self.dataset_expected["a"].equals(merged["a"]) assert self.dataset_expected.equals(merged) + class TestMultiTopology1D2DMergePartitions: @pytest.fixture(autouse=True) def setup(self): @@ -277,14 +275,16 @@ def setup(self): grid_b = generate_mesh_1d(6, "mesh1d") parts_a = grid_a.partition(n_part=2) # TODO: If partitioning implemented for 1D grids, replace with that. - i_edges = [[0,1,2], [3,4,5]] + i_edges = [[0, 1, 2], [3, 4, 5]] parts_b = [grid_b.isel(mesh1d_nEdges=np.array(ls)) for ls in i_edges] - values_parts_a = [np.arange(part.n_face) for part in parts_a] + values_parts_a = [np.arange(part.n_face) for part in parts_a] values_parts_b = [np.arange(part.n_edge) for part in parts_b] datasets_parts = [] - for i, (part_a, part_b, values_a, values_b) in enumerate(zip(parts_a, parts_b, values_parts_a, values_parts_b)): + for i, (part_a, part_b, values_a, values_b) in enumerate( + zip(parts_a, parts_b, values_parts_a, values_parts_b) + ): ds = xu.UgridDataset(grids=[part_a, part_b]) ds["a"] = ((part_a.face_dimension), values_a) ds["b"] = ((part_b.edge_dimension), values_b) @@ -296,7 +296,10 @@ def setup(self): ds_expected["b"] = ((grid_b.edge_dimension), np.concatenate(values_parts_b)) ds_expected["c"] = 0 # Assign coordinates also added during merge_partitions - coords = {grid_a.face_dimension: np.arange(grid_a.n_face), grid_b.edge_dimension: np.arange(grid_b.n_edge)} + coords = { + grid_a.face_dimension: np.arange(grid_a.n_face), + grid_b.edge_dimension: np.arange(grid_b.n_edge), + } ds_expected = ds_expected.assign_coords(**coords) self.datasets_parts = datasets_parts @@ -309,4 +312,4 @@ def test_merge_partitions(self): # In case of non-UGRID data, it should default to the first partition: assert merged["c"] == 0 - assert self.dataset_expected.equals(merged) \ No newline at end of file + assert self.dataset_expected.equals(merged)