fix: Handle initial zero gradients better

Author: aseyboldt

src/adapt_strategy.rs

@@ -352,7 +352,7 @@ impl<F: CpuLogpFunc> AdaptStrategy for ExpWindowDiagAdapt<F> {
             state
                 .grad
                 .iter()
-                .map(|&grad| grad.abs().recip().clamp(LOWER_LIMIT, UPPER_LIMIT))
+                .map(|&grad| grad.abs().clamp(LOWER_LIMIT, UPPER_LIMIT).recip())
                 .map(|var| if var.is_finite() { Some(var) } else { Some(1.) }),
         );
     }

src/cpu_potential.rs

@@ -165,7 +165,16 @@ impl<F: CpuLogpFunc, M: MassMatrix> Hamiltonian for EuclideanPotential<F, M> {
         }
         self.update_potential_gradient(&mut state)
             .map_err(|e| NutsError::LogpFailure(Box::new(e)))?;
-        Ok(state)
+        if state
+            .grad
+            .iter()
+            .cloned()
+            .any(|val| (val == 0f64) | !val.is_finite())
+        {
+            Err(NutsError::BadInitGrad())
+        } else {
+            Ok(state)
+        }
     }
 
     fn randomize_momentum<R: rand::Rng + ?Sized>(&self, state: &mut Self::State, rng: &mut R) {

src/nuts.rs

@@ -13,13 +13,17 @@ use crate::math::logaddexp;
 #[cfg(feature = "arrow")]
 use crate::SamplerArgs;
 
+#[non_exhaustive]
 #[derive(Error, Debug)]
 pub enum NutsError {
     #[error("Logp function returned error: {0}")]
     LogpFailure(Box<dyn std::error::Error + Send + Sync>),
 
     #[error("Could not serialize sample stats")]
     SerializeFailure(),
+
+    #[error("Could not initialize state because of bad initial gradient.")]
+    BadInitGrad(),
 }
 
 pub type Result<T> = std::result::Result<T, NutsError>;