Polygon json now has string cell ids

Also fixed a bug in boundary_polygons_auto
kharchenkolab · Nov 28, 2023 · 8899f64 · 8899f64
1 parent 6685095
commit 8899f64
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,10 @@
 - Improved algorithm for polygon estimation: fixed bugs and reduced overlaps
 - Fixed a bug in the `split` step, improved cell continuity
 
+### Changed
+
+- Cell IDs in polygon GeoJSON are now strings to match segmentation.csv
+
 ## [0.6.2] — 2023-08-15
 
 ### Fixed

diff --git a/src/data_loading/cli_wrappers.jl b/src/data_loading/cli_wrappers.jl
@@ -4,7 +4,9 @@ import LinearAlgebra: Adjoint
 using ProgressMeter
 using StatsBase: denserank
 
-Polygons = Union{Dict{Int, Matrix{T}}, Dict{String, Dict{Int, Matrix{T}}}} where T <: Real
+PolygonCollection = Dict{TI, Matrix{TV}} where {TV <: Real, TI <: Union{String, Integer}}
+PolygonStack = Dict{String, Dict{TI, Matrix{TV}}} where {TV <: Real, TI <: Union{String, Integer}}
+Polygons = Union{PolygonCollection, PolygonStack}
 
 function parse_prior_assignment(pos_data::Matrix{Float64}, prior_segmentation::Vector; col_name::Symbol, min_molecules_per_segment::Int, min_mols_per_cell::Int)
     try
@@ -64,12 +66,12 @@ function save_molecule_counts(counts::Union{AbstractDataFrame, AbstractMatrix},
 end
 
 
-function polygons_to_geojson(polygons::Dict{Int, Matrix{T}} where T <: Real)
+function polygons_to_geojson(polygons::PolygonCollection)
     geoms = [Dict("type" => "Polygon", "coordinates" => [collect.(eachrow(p))], "cell" => c) for (c,p) in polygons]
     return Dict("type" => "GeometryCollection", "geometries" => geoms);
 end
 
-polygons_to_geojson(polygons::Dict{String, Dict{Int, Matrix{T}}}) where T <:Real =
+polygons_to_geojson(polygons::PolygonStack) =
     [merge!(polygons_to_geojson(poly), Dict("z" => k)) for (k,poly) in polygons]
 
 function save_polygons_to_geojson(polygons::Polygons, file::String)

diff --git a/src/processing/data_processing/boundary_estimation.jl b/src/processing/data_processing/boundary_estimation.jl
@@ -265,7 +265,7 @@ 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})
+function boundary_polygons(pos_data::Matrix{Float64}, cell_labels::Vector{<:Integer}; cell_names::Union{Vector{String}, Nothing}=nothing)
     if size(pos_data, 1) == 3
         pos_data = pos_data[1:2,:]
     elseif size(pos_data, 1) != 2
@@ -294,28 +294,34 @@ function boundary_polygons(pos_data::Matrix{Float64}, cell_labels::Vector{<:Inte
         cell_borders[cid] = Matrix(cpd[:,cbord]')
     end
 
-    return Dict(cid => cb for (cid,cb) in enumerate(cell_borders) if !isempty(cb))
+    if cell_names == nothing
+        cell_names = 1:length(cell_borders)
+    else
+        cell_names = get.(Ref(cell_names), cell_labels, "")
+    end
+
+    return Dict(cid => cb for (cid,cb) in zip(cell_names, cell_borders) if !isempty(cb))
 end
 
-function boundary_polygons_auto(pos_data::Matrix{Float64}, assignment::Vector{<:Integer}; estimate_per_z::Bool)
+function boundary_polygons_auto(pos_data::Matrix{Float64}, assignment::Vector{<:Integer}; estimate_per_z::Bool, cell_names::Union{Vector{String}, Nothing}=nothing)
     @info "Estimating boundary polygons"
 
-    poly_joint = boundary_polygons(pos_data, assignment);
+    poly_joint = boundary_polygons(pos_data, assignment; cell_names);
 
     if !estimate_per_z || size(pos_data, 1) == 2
         return poly_joint, poly_joint
     end
 
     z_coords = @view pos_data[3,:]
     z_vals = sort(unique(z_coords))
-    if length(z_vals) > (size(pos_data, 1) / 100)
-        @warn "To many values of z. Using 2D polygons"
+    if length(z_vals) > (size(pos_data, 2) / 100)
+        @warn "To many values of z ($(length(z_vals))). 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]),
+        mask -> boundary_polygons(pos_data[:, mask], assignment[mask]; cell_names),
         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
@@ -8,7 +8,7 @@ import Base64
 
 plot_molecules!(args...; kwargs...) = plot_molecules(args...; append=true, kwargs...)
 
-plot_molecules(df_spatial::DataFrame, polygons::Dict{Int, Matrix{Float64}}; kwargs...) =
+plot_molecules(df_spatial::DataFrame, polygons::Dict{<:Union{Int, String}, Matrix{Float64}}; kwargs...) =
     plot_molecules(df_spatial, collect(values(polygons)); kwargs...)
 
 function plot_molecules(