VolumetricGridLookupField: generalize index type, add constructor (#475)

This PR generifies the index term. It also adds a constructor which is used by #456.

include/gz/math/VolumetricGridLookupField.hh

@@ -32,7 +32,7 @@ namespace gz
   {
     // Inline bracket to help doxygen filtering.
     inline namespace GZ_MATH_VERSION_NAMESPACE {
-      template<typename T>
+      template<typename T, typename I = std::size_t>
       /// \brief Lookup table for a volumetric dataset. This class is used to
       /// lookup indices for a large dataset that's organized in a grid. This
       /// class is not meant to be used with non-grid like data sets. The grid
@@ -47,7 +47,7 @@ namespace gz
         private: AxisIndex<T> y_indices_by_lon;
 
         private: std::vector<
-          std::vector<std::vector<std::optional<std::size_t>>>> index_table;
+          std::vector<std::vector<std::optional<I>>>> index_table;
 
         /// \brief Constructor
         /// \param[in] _cloud The cloud of points to use to construct the grid.
@@ -143,6 +143,53 @@ namespace gz
           return interpolators;
         }
 
+        /// \brief Constructor
+        /// \param[in] _cloud The cloud of points to use to construct the grid.
+        /// \param[in] _indices A series of indices these points correspond to.
+        public: VolumetricGridLookupField(
+          const std::vector<Vector3<T>> &_cloud,
+          const std::vector<I> &_indices)
+        {
+          assert(_indices.size() == _cloud.size());
+          // NOTE: This part of the code assumes an exact grid of points.
+          // The grid may be distorted or the stride between different points
+          // may not be the same, but fundamentally the data is structured in
+          // a grid-like structure. It keeps track of the axis indices for
+          // each point in the grid.
+          for(auto pt : _cloud)
+          {
+            x_indices_by_lat.AddIndexIfNotFound(pt.X());
+            y_indices_by_lon.AddIndexIfNotFound(pt.Y());
+            z_indices_by_depth.AddIndexIfNotFound(pt.Z());
+          }
+
+          int num_x = x_indices_by_lat.GetNumUniqueIndices();
+          int num_y = y_indices_by_lon.GetNumUniqueIndices();
+          int num_z = z_indices_by_depth.GetNumUniqueIndices();
+
+          index_table.resize(num_z);
+          for(int i = 0; i < num_z; ++i)
+          {
+            index_table[i].resize(num_y);
+            for(int j = 0; j < num_y; ++j)
+            {
+              index_table[i][j].resize(num_x);
+            }
+          }
+
+          for(std::size_t i = 0; i < _cloud.size(); ++i)
+          {
+            const auto &pt = _cloud[i];
+            const std::size_t x_index =
+              x_indices_by_lat.GetIndex(pt.X()).value();
+            const std::size_t y_index =
+              y_indices_by_lon.GetIndex(pt.Y()).value();
+            const std::size_t z_index =
+              z_indices_by_depth.GetIndex(pt.Z()).value();
+            index_table[z_index][y_index][x_index] = _indices[i];
+          }
+        }
+
         /// \brief Estimates the values for a grid given a list of values to
         /// interpolate. This method uses Trilinear interpolation.
         /// \param[in] _pt The point to estimate for.

src/VolumetricGridLookupField_TEST.cc

@@ -48,6 +48,29 @@ TEST(VolumetricGridLookupField, CheckInterpolationExact)
   }
 }
 
+TEST(VolumetricGridLookupField, CustomIndices)
+{
+  // Test custom index.
+  std::vector<Vector3d> cloud;
+  std::vector<std::size_t> indices;
+  const double stride_x = 1, stride_y = 5, stride_z = 10;
+  for(double x = 0; x < 300; x += stride_x)
+  {
+    for(double y = 0; y < 300; y += stride_y)
+    {
+      for(double z = 0; z < 300; z += stride_z)
+      {
+        cloud.emplace_back(x, y, z);
+      }
+    }
+  }
+  indices.resize(cloud.size(), 0);
+  VolumetricGridLookupField<double> scalarIndex(cloud, indices);
+  auto val = scalarIndex.GetInterpolators(cloud[0]);
+  ASSERT_EQ(val.size(), 1UL);
+  ASSERT_EQ(val[0].index, 0);
+}
+
 TEST(VolumetricGridLookupField, CheckInterpolationBoxEightPoints)
 {
   std::vector<Vector3d> cloud;