Rectified Multivariate Expansion Updated

MMRROOO · Mar 20, 2024 · 6e852ff · 6e852ff
1 parent c6628cd
commit 6e852ff
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/MParT/MultivariateExpansion.h b/MParT/MultivariateExpansion.h
@@ -221,4 +221,4 @@ namespace mpart{
 }
 
 
-#endif 
+#endif 
diff --git a/MParT/MultivariateExpansionWorker.h b/MParT/MultivariateExpansionWorker.h
@@ -432,7 +432,7 @@ class MultivariateExpansionWorker
             if(multiSet_.nzDims(end_idx)==dim_-1){
                 termVal = Rectifier::Evaluate(termVal)*lastVal;
             } else {
-                termVal = Rectifier::Evaluate(termVal*lastVal);
+                termVal = termVal*lastVal;
             }
         } else {
             termVal *= lastVal;
@@ -484,13 +484,18 @@ class MultivariateExpansionWorker
             lastVal = polyCache[startPos_(d) + multiSet_.nzOrders(end_idx)];
         }
         if constexpr (!std::is_same_v<Rectifier, Identity>) {
+            bool isRectified = d == dim_ - 1;
             if(wrt == dim_ - 1) { // Diagonal deriv
                 termVal = Rectifier::Evaluate(termVal)*wrtDeriv; // if wrt != d, wrtDeriv = 0
             } else if (wrt == -1) { // No deriv
-                termVal = (d == dim_ - 1) ? Rectifier::Evaluate(termVal)*lastVal : Rectifier::Evaluate(termVal*lastVal);
+                termVal = (isRectified) ? Rectifier::Evaluate(termVal)*lastVal : termVal*lastVal;
             } else { // Offdiag deriv
-                if(d != dim_ - 1) wrtVal *= lastVal; // lastVal belongs on inside and outside
-                termVal = Rectifier::Derivative(termVal*wrtVal)*termVal*wrtDeriv*lastVal;
+                if(!isRectified) {
+                    termVal *= lastVal*wrtDeriv;
+                }
+                else{
+                    termVal = Rectifier::Derivative(termVal*wrtVal)*termVal*wrtDeriv*lastVal;
+                }
             }
         } else { // Reduce to loop body
             termVal *= lastVal*wrtDeriv;
@@ -568,4 +573,4 @@ class MultivariateExpansionWorker
 
 
 
-#endif  // #ifndef MPART_MULTIVARIATEEXPANSION_H
+#endif  // #ifndef MPART_MULTIVARIATEEXPANSION_H
diff --git a/tests/Test_MultivariateExpansionWorker.cpp b/tests/Test_MultivariateExpansionWorker.cpp
@@ -179,8 +179,7 @@ TEMPLATE_TEST_CASE( "Testing multivariate expansion worker", "[MultivariateExpan
             expansion.FillCache2(&cache[0], pt, pt(dim-1), DerivativeFlags::None);
 
             eval2 = expansion.Evaluate(&cache[0], coeffs);
-
-            REQUIRE_THAT(inGrad(wrt), Matchers::WithinAbs((eval2-eval)/fdStep, fdStep*10));
+            CHECK_THAT(inGrad(wrt), Matchers::WithinAbs((eval2-eval)/fdStep, fdStep*10));
             pt(wrt) -= fdStep;
         }
     }