From 6a7393da2f89e7caec9fcdc2c5794cdf98c7117b Mon Sep 17 00:00:00 2001
From: Bradley J Chambers <brad.chambers@gmail.com>
Date: Wed, 10 Feb 2021 13:13:12 -0600
Subject: [PATCH] Add tests and update to work on local vector rather than
 Dimension

---
 filters/private/Segmentation.cpp | 17 ++++---
 test/unit/SegmentationTest.cpp   | 79 ++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 7 deletions(-)

diff --git a/filters/private/Segmentation.cpp b/filters/private/Segmentation.cpp
index 17b4091bce..aacd27f3de 100644
--- a/filters/private/Segmentation.cpp
+++ b/filters/private/Segmentation.cpp
@@ -263,25 +263,28 @@ PointIdList inverseDensityImportanceSampling(PointView& view,
     // Increment knn by one to account for query point, which does not
     // contribute to density calculation.
     knn++;
-
+    std::vector<double> densities(view.size());
     for (PointRef p : view)
     {
         PointIdList indices(knn);
         std::vector<double> sqr_dists(knn);
         kdi.knnSearch(p, knn, &indices, &sqr_dists);
+        std::transform(sqr_dists.begin(), sqr_dists.end(), sqr_dists.begin(),
+                       [](double val) { return std::sqrt(val); });
         double density =
             std::accumulate(sqr_dists.begin(), sqr_dists.end(), 0.0);
-        p.setField(Id::Density, 1.0 / density);
+        densities[p.pointId()] = 1.0 / density;
     }
 
-    PointIdList ids(view.size());
+    PointIdList ids(densities.size());
     std::iota(ids.begin(), ids.end(), 0);
-    auto cmp = [&view](PointId const& i1, PointId const& i2) {
-        return view.getFieldAs<double>(Id::Density, i1) <
-               view.getFieldAs<double>(Id::Density, i2);
+    auto cmp = [&densities](PointId const& i1, PointId const& i2) {
+        return densities[i1] > densities[i2];
     };
     std::stable_sort(ids.begin(), ids.end(), cmp);
-    ids.resize(count);
+    if (count < view.size())
+        ids.resize(count);
+
     return ids;
 }
 
diff --git a/test/unit/SegmentationTest.cpp b/test/unit/SegmentationTest.cpp
index 14ff8ece0d..a7faf6ba45 100644
--- a/test/unit/SegmentationTest.cpp
+++ b/test/unit/SegmentationTest.cpp
@@ -244,3 +244,82 @@ TEST(SegmentationTest, SegmentReturns)
     EXPECT_EQ(3u, first->size());
     EXPECT_EQ(1u, second->size());
 }
+
+TEST(SegmentationTest, InverseDensityImportanceSamplingSingle)
+{
+    using namespace Segmentation;
+
+    PointIdList ids;
+
+    PointTable table;
+    PointLayoutPtr layout(table.layout());
+
+    layout->registerDim(Dimension::Id::X);
+    layout->registerDim(Dimension::Id::Y);
+    layout->registerDim(Dimension::Id::Z);
+
+    PointView src(table);
+
+    src.setField(Dimension::Id::X, 0, 0.0);
+    src.setField(Dimension::Id::Y, 0, 0.0);
+    src.setField(Dimension::Id::Z, 0, 0.0);
+    ids = inverseDensityImportanceSampling(src, 2, 1);
+    EXPECT_EQ(1u, ids.size());
+}
+
+TEST(SegmentationTest, InverseDensityImportanceSamplingDouble)
+{
+    using namespace Segmentation;
+
+    PointIdList ids;
+
+    PointTable table;
+    PointLayoutPtr layout(table.layout());
+
+    layout->registerDim(Dimension::Id::X);
+    layout->registerDim(Dimension::Id::Y);
+    layout->registerDim(Dimension::Id::Z);
+
+    PointView src(table);
+
+    src.setField(Dimension::Id::X, 0, 0.0);
+    src.setField(Dimension::Id::Y, 0, 0.0);
+    src.setField(Dimension::Id::Z, 0, 0.0);
+
+    src.setField(Dimension::Id::X, 1, 10.0);
+    src.setField(Dimension::Id::Y, 1, 10.0);
+    src.setField(Dimension::Id::Z, 1, 10.0);
+    ids = inverseDensityImportanceSampling(src, 2, 1);
+    EXPECT_EQ(2u, ids.size());
+}
+
+TEST(SegmentationTest, InverseDensityImportanceSamplingTriple)
+{
+    using namespace Segmentation;
+
+    PointIdList ids;
+
+    PointTable table;
+    PointLayoutPtr layout(table.layout());
+
+    layout->registerDim(Dimension::Id::X);
+    layout->registerDim(Dimension::Id::Y);
+    layout->registerDim(Dimension::Id::Z);
+
+    PointView src(table);
+
+    src.setField(Dimension::Id::X, 0, 0.0);
+    src.setField(Dimension::Id::Y, 0, 0.0);
+    src.setField(Dimension::Id::Z, 0, 0.0);
+
+    src.setField(Dimension::Id::X, 1, 10.0);
+    src.setField(Dimension::Id::Y, 1, 10.0);
+    src.setField(Dimension::Id::Z, 1, 10.0);
+
+    src.setField(Dimension::Id::X, 2, 0.5);
+    src.setField(Dimension::Id::Y, 2, 0.5);
+    src.setField(Dimension::Id::Z, 2, 0.5);
+    ids = inverseDensityImportanceSampling(src, 3, 1);
+    EXPECT_EQ(3u, ids.size());
+    EXPECT_EQ(point_count_t(1), ids[2]);
+}