From eddf6df384864d8732278cd837b3cf3c542d2580 Mon Sep 17 00:00:00 2001
From: Kenneth Moreland <morelandkd@ornl.gov>
Date: Fri, 3 Feb 2023 12:55:57 -0500
Subject: [PATCH] Update probe filter to interpolate any field type

Previously, the probe filter only worked on certain `Vec` sizes and
converted many types to floating point.

This change uses the extract component feature to pull data from any
array at its natural component type.

The bad part of this change is that it has to call the worklet
separately for each component in the field. That adds overhead and
probably lowers the cache efficiency. It was implemented this way
because the cell interpolation function does not work with the
recombined vecs returned from extract array.
---
 docs/changelog/interp-any-field.md     |  3 ++-
 vtkm/filter/resampling/Probe.cxx       | 35 ++++++++++++++++++++------
 vtkm/filter/resampling/worklet/Probe.h | 18 ++++++-------
 3 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/docs/changelog/interp-any-field.md b/docs/changelog/interp-any-field.md
index c4aea7be8..0a029e7e6 100644
--- a/docs/changelog/interp-any-field.md
+++ b/docs/changelog/interp-any-field.md
@@ -19,7 +19,8 @@ fallback is used.
   * `ClipWithImplicitFunction`
   * `Contour`
   * `MIRFilter`
-  * `PointAverage`
   * `NDHistogram`
   * `ParticleDensityCloudInCell`
   * `ParticleDensityNearestGridPoint`
+  * `PointAverage`
+  * `Probe`
diff --git a/vtkm/filter/resampling/Probe.cxx b/vtkm/filter/resampling/Probe.cxx
index 0c134ef3d..586f3cf13 100644
--- a/vtkm/filter/resampling/Probe.cxx
+++ b/vtkm/filter/resampling/Probe.cxx
@@ -30,15 +30,34 @@ bool DoMapField(vtkm::cont::DataSet& result,
 {
   if (field.IsPointField())
   {
-    auto resolve = [&](const auto& concrete) {
-      using T = typename std::decay_t<decltype(concrete)>::ValueType;
-      vtkm::cont::ArrayHandle<T> outputArray = worklet.ProcessPointField(
-        concrete, vtkm::cont::internal::CastInvalidValue<T>(invalidValue));
-      result.AddPointField(field.GetName(), outputArray);
+    vtkm::cont::UnknownArrayHandle inArray = field.GetData();
+    vtkm::cont::UnknownArrayHandle outArray = inArray.NewInstanceBasic();
+
+    bool called = false;
+    auto tryType = [&](auto t) {
+      using T = std::decay_t<decltype(t)>;
+      if (!called && inArray.IsBaseComponentType<T>())
+      {
+        called = true;
+        vtkm::IdComponent numComponents = inArray.GetNumberOfComponentsFlat();
+        VTKM_ASSERT(numComponents == outArray.GetNumberOfComponentsFlat());
+
+        for (vtkm::IdComponent cIndex = 0; cIndex < numComponents; ++cIndex)
+        {
+          worklet.ProcessPointField(inArray.ExtractComponent<T>(cIndex),
+                                    outArray.ExtractComponent<T>(cIndex, vtkm::CopyFlag::Off),
+                                    vtkm::cont::internal::CastInvalidValue<T>(invalidValue));
+        }
+      }
     };
-    field.GetData()
-      .CastAndCallForTypesWithFloatFallback<vtkm::TypeListField, VTKM_DEFAULT_STORAGE_LIST>(
-        resolve);
+    vtkm::ListForEach(tryType, vtkm::TypeListScalarAll{});
+    if (!called)
+    {
+      VTKM_LOG_CAST_FAIL(worklet, vtkm::TypeListScalarAll);
+      return false;
+    }
+
+    result.AddPointField(field.GetName(), outArray);
     return true;
   }
   else if (field.IsCellField())
diff --git a/vtkm/filter/resampling/worklet/Probe.h b/vtkm/filter/resampling/worklet/Probe.h
index 10ce47ac3..70872e641 100644
--- a/vtkm/filter/resampling/worklet/Probe.h
+++ b/vtkm/filter/resampling/worklet/Probe.h
@@ -217,24 +217,22 @@ public:
   };
 
   /// Intepolate the input point field data at the points of the geometry
-  template <typename T,
-            typename Storage,
+  template <typename InArrayType,
+            typename OutArrayType,
+            typename ComponentType,
             typename InputCellSetTypeList = VTKM_DEFAULT_CELL_SET_LIST>
-  vtkm::cont::ArrayHandle<T> ProcessPointField(
-    const vtkm::cont::ArrayHandle<T, Storage>& field,
-    const T& invalidValue,
-    InputCellSetTypeList icsTypes = InputCellSetTypeList()) const
+  void ProcessPointField(const InArrayType& field,
+                         const OutArrayType& result,
+                         ComponentType invalidValue,
+                         InputCellSetTypeList icsTypes = InputCellSetTypeList()) const
   {
-    vtkm::cont::ArrayHandle<T> result;
     vtkm::cont::Invoker invoke;
-    invoke(InterpolatePointField<T>(invalidValue),
+    invoke(InterpolatePointField<ComponentType>(invalidValue),
            this->CellIds,
            this->ParametricCoordinates,
            this->InputCellSet.ResetCellSetList(icsTypes),
            field,
            result);
-
-    return result;
   }
 
   vtkm::cont::ArrayHandle<vtkm::Id> GetCellIds() const { return this->CellIds; }