added pqdm version for multiple processes

Author: EliHei2

src/segger/validation/xenium_explorer.py

@@ -469,3 +469,172 @@ def generate_experiment_file(
 
     with open(output_path, "w") as f:
         json.dump(experiment, f, indent=2)
+
+
+
+
+from pqdm.processes import pqdm  # or from pqdm.processes import pqdm for process backend
+import os
+
+def _process_one_cell(args):
+    seg_cell_id, seg_cell, area_low, area_high = args
+
+    if len(seg_cell) < 5:
+        return None
+
+    cell_convex_hull = generate_boundary(seg_cell)
+    if cell_convex_hull is None or not isinstance(cell_convex_hull, Polygon):
+        return None
+
+    if not (area_low <= cell_convex_hull.area <= area_high):
+        return None
+
+    # Get original vertices and remove duplicate closing vertex if present
+    cell_vertices = list(cell_convex_hull.exterior.coords)
+    if cell_vertices[0] == cell_vertices[-1]:
+        cell_vertices = cell_vertices[:-1]
+    
+    n_vertices = len(cell_vertices)
+    
+    # Sample up to 16 vertices
+    if n_vertices > 16:
+        # Evenly sample 16 vertices from original set
+        indices = np.linspace(0, n_vertices-1, 16, dtype=int)
+        sampled_vertices = [cell_vertices[i] for i in indices]
+    else:
+        sampled_vertices = cell_vertices
+    
+    # Pad with first vertex if needed
+    if len(sampled_vertices) < 16:
+        sampled_vertices += [sampled_vertices[0]] * (16 - len(sampled_vertices))
+    
+    return {
+        "seg_cell_id": seg_cell_id,
+        "cell_area": float(cell_convex_hull.area),
+        "cell_vertices": sampled_vertices,
+        "cell_num_vertices": len(sampled_vertices),
+    }
+
+
+def seg2explorer_pqdm(
+    seg_df: pd.DataFrame,
+    source_path: str,
+    output_dir: str,
+    cells_filename: str = "seg_cells",
+    analysis_filename: str = "seg_analysis",
+    xenium_filename: str = "seg_experiment.xenium",
+    analysis_df: Optional[pd.DataFrame] = None,
+    draw: bool = False,
+    cell_id_columns: str = "seg_cell_id",
+    area_low: float = 10,
+    area_high: float = 100,
+    n_jobs: int = 1
+) -> None:
+    source_path = Path(source_path)
+    storage = Path(output_dir)
+    storage.mkdir(parents=True, exist_ok=True)
+
+    grouped_by = seg_df.groupby(cell_id_columns)
+
+    # Build a lightweight iterable of work items (id, slice, thresholds)
+    # NOTE: this will still materialize each group slice, but we avoid copying the whole DF per worker.
+    work_iter = ((seg_cell_id, seg_cell, area_low, area_high) for seg_cell_id, seg_cell in grouped_by)
+
+    # Parallel map with threads (good default). Tune n_jobs.
+    # n_jobs = min(32, os.cpu_count() or 8)
+    results = pqdm(work_iter, _process_one_cell, n_jobs=n_jobs, desc="Cells", exception_behaviour="immediate")
+
+    # Collate results
+    cell_id2old_id: Dict[int, Any] = {}
+    cell_id: List[int] = []
+    polygon_num_vertices: List[List[int]] = []
+    polygon_vertices: List[List[Any]] = []
+
+    # We need a stable incremental id — use enumerate over kept results
+    kept = [r for r in results if r is not None]
+    for cell_incremental_id, r in enumerate(kept):
+        uint_cell_id = cell_incremental_id + 1
+        cell_id2old_id[uint_cell_id] = r["seg_cell_id"]
+        cell_id.append(uint_cell_id)
+        polygon_num_vertices.append(r["cell_num_vertices"])
+        polygon_vertices.append(r["cell_vertices"])
+
+    # Flatten vertices exactly as before
+    cell_polygon_vertices = get_flatten_version(polygon_vertices)
+
+    source_zarr_store = ZipStore(source_path / "cells.zarr.zip", mode="r") # added this line
+    existing_store = zarr.open(source_zarr_store, mode="r")
+    new_store = zarr.open(storage / f"{cells_filename}.zarr.zip", mode="w")
+
+    # Create polygon_sets group with the new structure
+    polygon_group = new_store.create_group("polygon_sets")
+
+    # Process cell polygons (set 1)
+    # cell_polygons = cells["polygon_vertices"][1]  # Cell polygons are at index 1
+    cell_num_vertices = polygon_num_vertices # Cell vertex counts
+
+    # Reshape cell polygons to (n_cells, 50) format
+    n_cells = cell_polygon_vertices.shape[0]
+    cell_vertices_flat = cell_polygon_vertices.reshape(n_cells, -1)[:, :33]  # Take first 50 values
+
+    set1 = polygon_group.create_group("1")
+    set1["cell_index"] = np.arange(1, n_cells + 1, dtype=np.uint32)  # 1-based indexing
+    set1["method"] = np.ones(n_cells, dtype=np.uint32)  # All method=1
+    set1["num_vertices"] = np.array(cell_num_vertices, dtype=np.int32)
+    set1["vertices"] = cell_vertices_flat.astype(np.float32)
+
+    new_store.attrs.update(existing_store.attrs)
+    new_store.attrs["number_cells"] = n_cells
+    new_store.store.close()
+
+    if analysis_df is None:
+        analysis_df = pd.DataFrame(
+            [cell_id2old_id[i] for i in cell_id], columns=[cell_id_columns]
+        )
+        analysis_df["default"] = "segger"
+
+    zarr_df = pd.DataFrame(
+        [cell_id2old_id[i] for i in cell_id], columns=[cell_id_columns]
+    )
+    clustering_df = pd.merge(zarr_df, analysis_df, how="left", on=cell_id_columns)
+    clusters_names = [col for col in analysis_df.columns if col != cell_id_columns]
+
+    clusters_dict = {
+        cluster: {
+            label: idx + 1
+            for idx, label in enumerate(
+                sorted(np.unique(clustering_df[cluster].dropna()))
+            )
+        }
+        for cluster in clusters_names
+    }
+
+    new_zarr = zarr.open(storage / f"{analysis_filename}.zarr.zip", mode="w")
+    new_zarr.create_group("/cell_groups")
+    for i, cluster in enumerate(clusters_names):
+        new_zarr["cell_groups"].create_group(str(i))
+        group_values = [clusters_dict[cluster].get(x, 0) for x in clustering_df[cluster]]
+        indices, indptr = get_indices_indptr(np.array(group_values))
+        new_zarr["cell_groups"][str(i)]["indices"] = indices
+        new_zarr["cell_groups"][str(i)]["indptr"] = indptr
+
+    new_zarr["cell_groups"].attrs.update(
+        {
+            "major_version": 1,
+            "minor_version": 0,
+            "number_groupings": len(clusters_names),
+            "grouping_names": clusters_names,
+            "group_names": [
+                sorted(clusters_dict[cluster], key=clusters_dict[cluster].get)
+                for cluster in clusters_names
+            ],
+        }
+    )
+    new_zarr.store.close()
+
+    generate_experiment_file(
+        template_path=source_path / "experiment.xenium",
+        output_path=storage / xenium_filename,
+        cells_name=cells_filename,
+        analysis_name=analysis_filename,
+    )