Switch to using linestrings and multilinestrings

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)

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"