From 96b29b37cd6087617fe84ffd6a2bb2d00bbe7a2a Mon Sep 17 00:00:00 2001
From: Christopher Rabotin <christopher.rabotin@gmail.com>
Date: Mon, 18 Nov 2024 00:19:04 -0700
Subject: [PATCH] Reduce NaN checks to return branch; fix erroneous bugfix in
 step size (was preventing negative steps!)

---
 src/propagators/instance.rs | 40 ++++++++++++++++++++++++++-----------
 src/propagators/options.rs  | 11 ----------
 2 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/src/propagators/instance.rs b/src/propagators/instance.rs
index 3ceb11fd..c32fd2a0 100644
--- a/src/propagators/instance.rs
+++ b/src/propagators/instance.rs
@@ -443,14 +443,6 @@ where
                 self.details.step = self.step_size;
                 return Ok(((self.details.step), next_state));
             } else {
-                if next_state.iter().any(|x| x.is_nan()) {
-                    return Err(PropagationError::PropMathError {
-                        source: MathError::DomainError {
-                            value: f64::NAN,
-                            msg: "try another integration method, or decrease step size; part of state vector is",
-                        },
-                    });
-                }
                 // Compute the error estimate.
                 self.details.error =
                     self.prop
@@ -462,6 +454,14 @@ where
                     || step_size_s <= self.prop.opts.min_step.to_seconds()
                     || self.details.attempts >= self.prop.opts.attempts
                 {
+                    if next_state.iter().any(|x| x.is_nan()) {
+                        return Err(PropagationError::PropMathError {
+                            source: MathError::DomainError {
+                                value: f64::NAN,
+                                msg: "try another integration method, or decrease step size; part of state vector is",
+                            },
+                        });
+                    }
                     if self.details.attempts >= self.prop.opts.attempts {
                         warn!(
                             "Could not further decrease step size: maximum number of attempts reached ({})",
@@ -473,16 +473,28 @@ where
                     if self.details.error < self.prop.opts.tolerance {
                         // Let's increase the step size for the next iteration.
                         // Error is less than tolerance, let's attempt to increase the step for the next iteration.
-                        let proposed_step_s = 0.9
+                        let proposed_step = 0.9
                             * step_size_s
                             * (self.prop.opts.tolerance / self.details.error)
                                 .powf(1.0 / f64::from(self.prop.method.order()));
 
-                        step_size_s = self.prop.opts.bound_proposed_step(proposed_step_s);
+                        step_size_s = if proposed_step > self.prop.opts.max_step.to_seconds() {
+                            self.prop.opts.max_step.to_seconds()
+                        } else {
+                            proposed_step
+                        };
                     }
-                    assert!(step_size_s > 0.0);
                     // In all cases, let's update the step size to whatever was the adapted step size
                     self.step_size = step_size_s * Unit::Second;
+                    if self.step_size.abs() < self.prop.opts.min_step {
+                        // Custom signum in case the step size becomes zero.
+                        let signum = if self.step_size.is_negative() {
+                            -1.0
+                        } else {
+                            1.0
+                        };
+                        self.step_size = self.prop.opts.min_step * signum;
+                    }
                     return Ok((self.details.step, next_state));
                 } else {
                     // Error is too high and we aren't using the smallest step, and we haven't hit the max number of attempts.
@@ -493,7 +505,11 @@ where
                         * (self.prop.opts.tolerance / self.details.error)
                             .powf(1.0 / f64::from(self.prop.method.order() - 1));
 
-                    step_size_s = self.prop.opts.bound_proposed_step(proposed_step_s);
+                    step_size_s = if proposed_step_s < self.prop.opts.min_step.to_seconds() {
+                        self.prop.opts.min_step.to_seconds()
+                    } else {
+                        proposed_step_s
+                    };
                     // Note that we don't set self.step_size, that will be updated right before we return
                 }
             }
diff --git a/src/propagators/options.rs b/src/propagators/options.rs
index aaab90ce..8ec73d49 100644
--- a/src/propagators/options.rs
+++ b/src/propagators/options.rs
@@ -144,17 +144,6 @@ impl IntegratorOptions {
         }
         self.min_step = min_step;
     }
-
-    /// Returns a proposed step in seconds that is within the bounds of these integrator options.
-    pub(crate) fn bound_proposed_step(&self, proposed_step_s: f64) -> f64 {
-        if proposed_step_s > self.max_step.to_seconds() {
-            self.max_step.to_seconds()
-        } else if proposed_step_s < self.min_step.to_seconds() {
-            self.min_step.to_seconds()
-        } else {
-            proposed_step_s
-        }
-    }
 }
 
 impl fmt::Display for IntegratorOptions {