diff --git a/tests/test_exporter.py b/tests/test_exporter.py index b76f4cb..644ef06 100644 --- a/tests/test_exporter.py +++ b/tests/test_exporter.py @@ -1,5 +1,6 @@ import json import os +from pathlib import Path import pytest from osgeo import ogr @@ -50,9 +51,54 @@ def test_fragments_gpgk_export(ga_fragments, tmp_path): ) # minus dummy -def test_fragments_gpkg_conversion_smoke(ga_fragment_path): - exporter = GeopackageExporter(ga_fragment_path, "fragment_export.gpkg") +def test_fragments_to_gpgk_export(ga_fragments, tmp_path): + path = str(tmp_path / "only_fragment_export.gpkg") + ga_fragments.fragments.to_gpkg(path, "fragments") + s = ogr.Open(path) + layer = s.GetLayer("fragments") + assert layer.GetFeatureCount() == ( + ga_fragments.fragments.id.size - 1 + ) # minus dummy + + +def test_fragments_to_geojson_export(ga_fragments, tmp_path): + path = str(tmp_path / "only_fragment_export.json") + ga_fragments.fragments.to_geojson(path, use_ogr=True) + s = ogr.Open(path) + layer = s.GetLayer(Path(path).name) + assert layer.GetFeatureCount() == ( + ga_fragments.fragments.id.size - 1 + ) # minus dummy + with open(path) as file: + data = json.load(file) + assert len(data["features"][0]["properties"]) == 2 + assert "node_id" in data["features"][0]["properties"] + + path = str(tmp_path / "only_fragment_export_2.json") + ga_fragments.fragments.to_geojson(path, use_ogr=False) + s = ogr.Open(path) + layer = s.GetLayer(Path(path).stem) + # We export more properties and use different layer name + assert layer.GetFeatureCount() == ( + ga_fragments.fragments.id.size - 1 + ) # minus dummy + with open(path) as file: + data = json.load(file) + assert len(data["features"][0]["properties"]) == 3 + assert "node_id" in data["features"][0]["properties"] + assert "model_type" in data["features"][0]["properties"] + + +def test_fragments_gpkg_conversion(ga_fragments, ga_fragment_path, tmp_path): + exporter = GeopackageExporter( + ga_fragment_path, str(tmp_path / "fragment_export.gpkg") + ) exporter.export() + s = ogr.Open(str(tmp_path / "fragment_export.gpkg")) + layer = s.GetLayer("fragment") + assert layer.GetFeatureCount() == ( + ga_fragments.fragments.id.size - 1 + ) # minus dummy def test_meta_data_gpgk_export(ga, tmp_path): @@ -157,6 +203,7 @@ def test_export_null(ga, tmp_path): ("grid", "grid_2D_open_water"), ("grid", "grid_2D_groundwater"), ("flowlines", "flowlines"), + ("fragments", "fragments"), ], ) def test_export_method(ga_export, export_method, expected_filename): diff --git a/tests/test_gridadmin.py b/tests/test_gridadmin.py index d77eaa7..f8661b8 100644 --- a/tests/test_gridadmin.py +++ b/tests/test_gridadmin.py @@ -572,3 +572,29 @@ def test_lines_filter_id_in(self): ) expected = self.parser.lines.data["line_coords"][:, trues] np.testing.assert_equal(filtered, expected) + + +class FragmentFilterTests(unittest.TestCase): + def setUp(self): + grid_admin_h5_file = os.path.join( + test_file_dir, "fragments", "gridadmin_fragments.h5" + ) + self.parser = GridH5Admin(grid_admin_h5_file) + + def test_fragments_filter_id_eq(self): + + filtered = self.parser.fragments.filter(id=3).data["node_id"] + (trues,) = np.where(self.parser.fragments.data["id"] == 3) + expected = self.parser.fragments.data["node_id"][trues] + np.testing.assert_equal(filtered, expected) + + def test_fragments_filter_id_in(self): + filtered = self.parser.fragments.filter(id__in=list(range(3, 7))).data[ + "node_id" + ] + (trues,) = np.where( + (self.parser.fragments.data["id"] >= 3) + & (self.parser.fragments.data["id"] < 7) + ) + expected = self.parser.fragments.data["node_id"][trues] + np.testing.assert_equal(filtered, expected) diff --git a/threedigrid/admin/exporter_constants.py b/threedigrid/admin/exporter_constants.py index 7e96825..6335ee0 100644 --- a/threedigrid/admin/exporter_constants.py +++ b/threedigrid/admin/exporter_constants.py @@ -179,6 +179,8 @@ "z_coordinate", ] +FRAGMENTS_EXPORT_FIELDS = ["id", "node_id"] + DEFAULT_EXPORT_FIELDS = { "Lines": "ALL", "Pipes": PIPES_EXPORT_FIELDS, @@ -194,4 +196,5 @@ "Breaches": BREACHES_EXPORT_FIELDS, "Levees": LEVEES_EXPORT_FIELDS, "Pumps": PUMPS_EXPORT_FIELDS, + "Fragments": FRAGMENTS_EXPORT_FIELDS, } diff --git a/threedigrid/admin/prepare.py b/threedigrid/admin/prepare.py index bf0b0a0..e915b55 100644 --- a/threedigrid/admin/prepare.py +++ b/threedigrid/admin/prepare.py @@ -508,6 +508,20 @@ def export_pumps(self): dest, indent=self._indent ) + def export_fragments(self): + if not self.ga.fragments.id.size == 0: + logger.info( + "[*] Model {} does not have fragments, " + "skipping export fragments...".format(self.ga.model_name) + ) + return + + dest_ow = os.path.join(self._dest, "fragments" + self._extension) + getattr( + self.ga.fragments.reproject_to(self._epsg), + self._export_method, + )(dest_ow, indent=self._indent) + def export_levees(self): """ writes shapefile of levees diff --git a/threedigrid/admin/serializers.py b/threedigrid/admin/serializers.py index afbe9b3..9ee7ac7 100644 --- a/threedigrid/admin/serializers.py +++ b/threedigrid/admin/serializers.py @@ -5,7 +5,7 @@ import numpy as np from threedigrid.admin import constants -from threedigrid.geo_utils import raise_import_exception, transform_bbox +from threedigrid.geo_utils import raise_import_exception, transform_bbox, transform_xys from threedigrid.orm.base.encoder import NumpyEncoder from threedigrid.orm.base.models import Model @@ -127,6 +127,26 @@ def geos_iter(self): ) properties = fill_properties(self.fields, data, i, model_type) yield geojson.Feature(geometry=polygon, properties=properties) + elif content_type == "fragments": + for i in range(data["id"].shape[-1]): + coords = np.round( + data["coords"][i].reshape(2, -1).astype("float64"), + constants.LONLAT_DIGITS, + ) + if ( + self._model.reproject_to_epsg is not None + and self._model.reproject_to_epsg != self._model.epsg_code + ): + # Pick reproject_to_epsg or original model epsg_code + coords = transform_xys( + np.array(coords[0]), + np.array(coords[1]), + self._model.epsg_code, + self._model.reproject_to_epsg, + ) + polygon = geojson.Polygon(coords.T.tolist()) + properties = fill_properties(self.fields, data, i, model_type) + yield geojson.Feature(geometry=polygon, properties=properties) elif content_type == "levees": for i in range(data["id"].shape[-1]): coords = np.round(