From f32fe4b21dece217dc34fae1372d7427a2a27f5b Mon Sep 17 00:00:00 2001
From: Viktor Petukhov <vpetukhov@altoslabs.com>
Date: Thu, 11 Apr 2024 08:50:04 -0700
Subject: [PATCH] Fixed problems with empty polygons; validated dimensions of
 input data

---
 .../data_processing/boundary_estimation.jl    | 26 ++++++++++++++-----
 src/reporting/plots.jl                        |  4 +--
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/processing/data_processing/boundary_estimation.jl b/src/processing/data_processing/boundary_estimation.jl
index ccb753c..ab33211 100644
--- a/src/processing/data_processing/boundary_estimation.jl
+++ b/src/processing/data_processing/boundary_estimation.jl
@@ -40,7 +40,7 @@ end
 function grid_borders_per_label(grid_labels::Matrix{<:Unsigned})
     d_cols = [-1 0 1 0]
     d_rows = [0 1 0 -1]
-    borders_per_label = [Array{UInt32, 1}[] for _ in 1:maximum(grid_labels)]
+    borders_per_label = [Vector{UInt32}[] for _ in 1:maximum(grid_labels)]
 
     for row in 1:size(grid_labels, 1)
         for col in 1:size(grid_labels, 2)
@@ -254,6 +254,10 @@ function boundary_polygons_from_grid(grid_labels::Matrix{<:Unsigned}; grid_step:
     borders_per_label = grid_borders_per_label(grid_labels);
     polys = Matrix{Float64}[]
     for bords in borders_per_label
+        if size(bords, 1) == 0
+            push!(polys, Float64[;;])
+            continue
+        end
         c_coords =Float64.(hcat(bords...));
         tess = adjacency_list(c_coords, adjacency_type=:triangulation, filter=false, return_tesselation=true)[1];
         poly_ids = extract_triangle_verts(tess) |> extract_border_edges |> border_edges_to_poly;
@@ -266,6 +270,14 @@ end
 boundary_polygons(bm_data::BmmData) = boundary_polygons(position_data(bm_data), bm_data.assignment)
 
 function boundary_polygons(pos_data::Matrix{Float64}, cell_labels::Vector{<:Integer})
+    if size(pos_data, 1) > 2
+        pos_data = pos_data[1:2,:]
+    end
+
+    if size(pos_data, 2) < 3
+        return Dict{Int, Matrix{Float64}}()
+    end
+
     points = normalize_points(pos_data)
     points = [IndexedPoint2D(p[1], p[2], i) for (i,p) in enumerate(eachcol(points))];
 
@@ -291,8 +303,8 @@ function boundary_polygons(pos_data::Matrix{Float64}, cell_labels::Vector{<:Inte
     return Dict(cid => cb for (cid,cb) in enumerate(cell_borders) if !isempty(cb))
 end
 
-function boundary_polygons_auto(pos_data::Matrix{Float64}, assignment::Vector{<:Integer}; estimate_per_z::Bool)
-    @info "Estimating boundary polygons"
+function boundary_polygons_auto(pos_data::Matrix{Float64}, assignment::Vector{<:Integer}; estimate_per_z::Bool, verbose=true)
+    verbose && @info "Estimating boundary polygons"
 
     poly_joint = boundary_polygons(pos_data, assignment);
 
@@ -302,14 +314,16 @@ function boundary_polygons_auto(pos_data::Matrix{Float64}, assignment::Vector{<:
 
     z_coords = @view pos_data[3,:]
     z_vals = sort(unique(z_coords))
-    if length(z_vals) > (size(pos_data, 1) / 100)
+    if length(z_vals) > (size(pos_data, 2) / 100)
         @warn "To many values of z. Using 2D polygons"
         return poly_joint, poly_joint
     end
 
     mask_per_z = [(z_coords .≈ z) for z in z_vals]
-    poly_per_z = progress_map(
-        mask -> boundary_polygons(pos_data[mask,:], assignment[mask]),
+
+    pmap = verbose ? progress_map : map
+    poly_per_z = pmap(
+        mask -> boundary_polygons(pos_data[:,mask], assignment[mask]),
         mask_per_z
     );
     poly_per_z = Dict("$k" => p for (k,p) in zip(z_vals, poly_per_z))
diff --git a/src/reporting/plots.jl b/src/reporting/plots.jl
index 85708e6..dacee17 100644
--- a/src/reporting/plots.jl
+++ b/src/reporting/plots.jl
@@ -22,7 +22,7 @@ function plot_molecules(
 
     noise_args_default = (marker=:xcross, markersize=(0.75 * markersize), strokewidth=0, color="black");
     axis_args_default = (xticklabelsize=12, yticklabelsize=12, xticksvisible=ticksvisible, xticklabelsvisible=ticksvisible, yticksvisible=ticksvisible, yticklabelsvisible=ticksvisible);
-    legend_args_default = (bgcolor=Colors.RGBA(1, 1, 1, 0.85),);
+    legend_args_default = (backgroundcolor=Colors.RGBA(1, 1, 1, 0.85),);
     polygon_args_default = (strokecolor="black", color="transparent", strokewidth=poly_strokewidth, label="")
 
     noise_kwargs = update_args(update_args(noise_args_default, Dict(kwargs...)), noise_kwargs)
@@ -98,7 +98,7 @@ function plot_molecules(
     end
 
     if length(polygons) > 0
-        MK.poly!([MK.Point2.(eachrow(p .+ [offset[1] offset[2]])) for p in polygons]; polygon_kwargs...)
+        MK.poly!([MK.Point2.(eachrow(p .+ [offset[1] offset[2]])) for p in polygons if Base.size(p, 1) > 0]; polygon_kwargs...)
         # We can also do `for p in polygons MK.lines!(p[:,1], p[:,2], color="black") end`, but this is 10+ times slower
     end