Last few APIs

keenon · Sep 23, 2024 · 2df6d00 · 2df6d00
1 parent 9f67b47
commit 2df6d00
Show file tree

Hide file tree

Showing 6 changed files with 95 additions and 1 deletion.
diff --git a/dart/biomechanics/SubjectOnDisk.cpp b/dart/biomechanics/SubjectOnDisk.cpp
@@ -110,6 +110,20 @@ proto::MissingGRFReason missingGRFReasonToProto(MissingGRFReason reason)
       return proto::MissingGRFReason::manualReview;
     case footContactDetectedButNoForce:
       return proto::MissingGRFReason::footContactDetectedButNoForce;
+    case tooHighMarkerRMS:
+      return proto::MissingGRFReason::tooHighMarkerRMS;
+    case hasInputOutliers:
+      return proto::MissingGRFReason::hasInputOutliers;
+    case hasNoForcePlateData:
+      return proto::MissingGRFReason::hasNoForcePlateData;
+    case velocitiesStillTooHighAfterFiltering:
+      return proto::MissingGRFReason::velocitiesStillTooHighAfterFiltering;
+    case copOutsideConvexFootError:
+      return proto::MissingGRFReason::copOutsideConvexFootError;
+    case zeroForceFrame:
+      return proto::MissingGRFReason::zeroForceFrame;
+    case extendedToNearestPeakForce:
+      return proto::MissingGRFReason::extendedToNearestPeakForce;
   }
   return proto::MissingGRFReason::notMissingGRF;
 }
@@ -142,6 +156,20 @@ MissingGRFReason missingGRFReasonFromProto(proto::MissingGRFReason reason)
       return manualReview;
     case proto::MissingGRFReason::footContactDetectedButNoForce:
       return footContactDetectedButNoForce;
+    case proto::MissingGRFReason::tooHighMarkerRMS:
+      return tooHighMarkerRMS;
+    case proto::MissingGRFReason::hasInputOutliers:
+      return hasInputOutliers;
+    case proto::MissingGRFReason::hasNoForcePlateData:
+      return hasNoForcePlateData;
+    case proto::MissingGRFReason::velocitiesStillTooHighAfterFiltering:
+      return velocitiesStillTooHighAfterFiltering;
+    case proto::MissingGRFReason::copOutsideConvexFootError:
+      return copOutsideConvexFootError;
+    case proto::MissingGRFReason::zeroForceFrame:
+      return zeroForceFrame;
+    case proto::MissingGRFReason::extendedToNearestPeakForce:
+      return extendedToNearestPeakForce;
       // These are just here to keep Clang from complaining
     case proto::MissingGRFReason_INT_MIN_SENTINEL_DO_NOT_USE_:
       return notMissingGRF;
@@ -3294,6 +3322,27 @@ SubjectOnDiskHeader::getTrials()
   return mTrials;
 }
 
+void SubjectOnDiskHeader::filterTrials(std::vector<bool> keepTrials)
+{
+  if (keepTrials.size() != mTrials.size())
+  {
+    std::cout << "SubjectOnDisk::writeSubject() passed bad info: keepTrials "
+                 "size is "
+              << keepTrials.size() << " but we have " << mTrials.size()
+              << " trials" << std::endl;
+    return;
+  }
+  std::vector<std::shared_ptr<SubjectOnDiskTrial>> newTrials;
+  for (int i = 0; i < keepTrials.size(); i++)
+  {
+    if (keepTrials[i])
+    {
+      newTrials.push_back(mTrials[i]);
+    }
+  }
+  mTrials = newTrials;
+}
+
 void SubjectOnDiskHeader::trimToProcessingPasses(int numPasses)
 {
   mPasses.resize(numPasses);

diff --git a/dart/biomechanics/SubjectOnDisk.hpp b/dart/biomechanics/SubjectOnDisk.hpp
@@ -437,6 +437,7 @@ class SubjectOnDiskHeader
   std::vector<std::shared_ptr<SubjectOnDiskPassHeader>> getProcessingPasses();
   std::shared_ptr<SubjectOnDiskTrial> addTrial();
   std::vector<std::shared_ptr<SubjectOnDiskTrial>> getTrials();
+  void filterTrials(std::vector<bool> keepTrials);
   void trimToProcessingPasses(int numPasses);
   void setTrials(std::vector<std::shared_ptr<SubjectOnDiskTrial>> trials);
   void recomputeColumnNames();

diff --git a/dart/biomechanics/enums.hpp b/dart/biomechanics/enums.hpp
@@ -7,6 +7,7 @@ namespace biomechanics {
 enum MissingGRFReason
 {
   notMissingGRF,
+  // These are the legacy filter's reasons
   measuredGrfZeroWhenAccelerationNonZero,
   unmeasuredExternalForceDetected,
   footContactDetectedButNoForce,
@@ -17,7 +18,15 @@ enum MissingGRFReason
   missingBlip,
   shiftGRF,
   manualReview,
-  interpolatedClippedGRF
+  interpolatedClippedGRF,
+  // These are the new filter's reasons
+  tooHighMarkerRMS,
+  hasInputOutliers,
+  hasNoForcePlateData,
+  velocitiesStillTooHighAfterFiltering,
+  copOutsideConvexFootError,
+  zeroForceFrame,
+  extendedToNearestPeakForce
 };
 
 enum MissingGRFStatus

diff --git a/dart/proto/SubjectOnDisk.proto b/dart/proto/SubjectOnDisk.proto
@@ -14,6 +14,13 @@ enum MissingGRFReason { notMissingGRF = 0;
                         interpolatedClippedGRF = 9;
                         manualReview = 10;
                         footContactDetectedButNoForce = 11;
+                        tooHighMarkerRMS = 12;
+                        hasInputOutliers = 13;
+                        hasNoForcePlateData = 14;
+                        velocitiesStillTooHighAfterFiltering = 15;
+                        copOutsideConvexFootError = 16;
+                        zeroForceFrame = 17;
+                        extendedToNearestPeakForce = 18;
                       };
 
 enum ProcessingPassType {

diff --git a/python/_nimblephysics/biomechanics/DynamicsFitter.cpp b/python/_nimblephysics/biomechanics/DynamicsFitter.cpp
@@ -21,6 +21,7 @@ void DynamicsFitter(py::module& m)
 {
 
   py::enum_<dart::biomechanics::MissingGRFReason>(m, "MissingGRFReason")
+      // These values are used by the legacy filter
       .value(
           "notMissingGRF", dart::biomechanics::MissingGRFReason::notMissingGRF)
       .value(
@@ -50,6 +51,29 @@ void DynamicsFitter(py::module& m)
       .value(
           "footContactDetectedButNoForce",
           dart::biomechanics::MissingGRFReason::footContactDetectedButNoForce)
+      // These values are used by the new filter
+      .value(
+          "tooHighMarkerRMS",
+          dart::biomechanics::MissingGRFReason::tooHighMarkerRMS)
+      .value(
+          "hasInputOutliers",
+          dart::biomechanics::MissingGRFReason::hasInputOutliers)
+      .value(
+          "hasNoForcePlateData",
+          dart::biomechanics::MissingGRFReason::hasNoForcePlateData)
+      .value(
+          "velocitiesStillTooHighAfterFiltering",
+          dart::biomechanics::MissingGRFReason::
+              velocitiesStillTooHighAfterFiltering)
+      .value(
+          "copOutsideConvexFootError",
+          dart::biomechanics::MissingGRFReason::copOutsideConvexFootError)
+      .value(
+          "zeroForceFrame",
+          dart::biomechanics::MissingGRFReason::zeroForceFrame)
+      .value(
+          "extendedToNearestPeakForce",
+          dart::biomechanics::MissingGRFReason::extendedToNearestPeakForce)
       .export_values();
 
   py::enum_<dart::biomechanics::MissingGRFStatus>(m, "MissingGRFStatus")

diff --git a/python/_nimblephysics/biomechanics/SubjectOnDisk.cpp b/python/_nimblephysics/biomechanics/SubjectOnDisk.cpp
@@ -980,6 +980,10 @@ Note that these are specified in the local body frame, acting on the body at its
             .def(
                 "getTrials",
                 &dart::biomechanics::SubjectOnDiskHeader::getTrials)
+            .def(
+                "filterTrials",
+                &dart::biomechanics::SubjectOnDiskHeader::filterTrials,
+                ::py::arg("keepTrials"))
             .def(
                 "trimToProcessingPasses",
                 &dart::biomechanics::SubjectOnDiskHeader::