diff --git a/cf_xarray/geometry.py b/cf_xarray/geometry.py index 3e7a8105..a6e2fc5a 100644 --- a/cf_xarray/geometry.py +++ b/cf_xarray/geometry.py @@ -314,7 +314,7 @@ def cf_to_lines(ds: xr.Dataset): It has the same dimension as the ``node_count`` or the coordinates variables, or ``'features'`` if those were not present in ``ds``. """ - from shapely.geometry import LineString, MultiLineString + from shapely import linestrings, multilinestrings # Shorthand for convenience geo = ds.geometry_container.attrs @@ -346,18 +346,30 @@ def cf_to_lines(ds: xr.Dataset): else: part_node_count = ds[part_node_count_name] - j = 0 # The index of the first node. - geoms = np.empty(part_node_count.shape, dtype=object) - node_count_cumsum = np.insert(np.cumsum(node_count), 0, 0) + # create all the line segments + line_indices = np.repeat(np.arange(len(node_count)), node_count) + lines = linestrings(xy, indices=line_indices) + + # contruct an array where every row corresponds to a geometry object + # and each value indicates whether the line at that index should + # be included in that object. + maxes = np.cumsum(part_node_count) + mins = np.insert(maxes[:-1], 0, 0) + dummy = np.broadcast_to(np.cumsum(node_count), (len(part_node_count), len(node_count))) + is_multi = np.where( + (dummy <= np.expand_dims(maxes, axis=1)) & (dummy > np.expand_dims(mins, axis=1)), + np.ones_like(dummy), + np.zeros_like(dummy) + ) - # i is the feature index, n its number of nodes - for i, n in enumerate(part_node_count.values): - subset = node_count_cumsum[(j <= node_count_cumsum) & (node_count_cumsum <= n)] - indexes = np.stack([subset[:-1], np.roll(subset, 2)[:-1]]).transpose() - if indexes.shape == (0, 2): - geoms[i] = LineString(xy[j : j + n, :]) - else: - geoms[i] = MultiLineString([xy[ii[0] : ii[1], :] for ii in indexes]) - j += n + # create multilinestrings for each set of lines. Note that we create multiline strings + # even for cases where there is only one line. We will pick which ones we need later on + multiline_indices = np.where(is_multi == 1)[0] + multilines = multilinestrings(lines, indices=multiline_indices) + + # construct the final geometry by pulling items from the lines or multilines based on + # the number of lines that are mapped to each object. + _, index, counts = np.unique(multiline_indices, return_counts=True, return_index=True) + geoms = np.where(counts == 1, np.array(lines[index]), multilines) return xr.DataArray(geoms, dims=part_node_count.dims, coords=part_node_count.coords) diff --git a/cf_xarray/tests/test_geometry.py b/cf_xarray/tests/test_geometry.py index dac90e52..4c29fab8 100644 --- a/cf_xarray/tests/test_geometry.py +++ b/cf_xarray/tests/test_geometry.py @@ -180,9 +180,9 @@ def test_cf_to_shapely_for_line(geometry_line_ds): @requires_shapely def test_cf_to_shapely_errors(geometry_ds): - in_ds, expected = geometry_ds - in_ds.geometry_container.attrs["geometry_type"] = "line" - with pytest.raises(NotImplementedError, match="Only point geometries conversion"): + in_ds, _ = geometry_ds + in_ds.geometry_container.attrs["geometry_type"] = "polygon" + with pytest.raises(NotImplementedError, match="Polygon geometry"): cfxr.cf_to_shapely(in_ds) in_ds.geometry_container.attrs["geometry_type"] = "punkt"