Tidy up results.

Author: patrick-kidger

lineax/_solver/bicgstab.py

@@ -187,14 +187,14 @@ def body_fun(carry):
 
         if self.max_steps is None:
             result = RESULTS.where(
-                (num_steps == max_steps), RESULTS.singular, RESULTS.successful
+                num_steps == max_steps, RESULTS.singular, RESULTS.successful
             )
-        else:
+        elif has_scale:
             result = RESULTS.where(
-                (num_steps == self.max_steps),
-                RESULTS.max_steps_reached if has_scale else RESULTS.successful,
-                RESULTS.successful,
+                num_steps == max_steps, RESULTS.max_steps_reached, RESULTS.successful
             )
+        else:
+            result = RESULTS.successful
         # breakdown is only an issue if we did not converge
         breakdown = breakdown_occurred(omega, alpha, rho) & not_converged(
             residual, diff, solution

lineax/_solver/cg.py

@@ -226,18 +226,16 @@ def cheap_r():
             cond_fun, body_fun, initial_value
         )
 
-        if (self.max_steps is None) or (max_steps < self.max_steps):
+        if self.max_steps is None:
             result = RESULTS.where(
-                num_steps == max_steps,
-                RESULTS.singular,
-                RESULTS.successful,
+                num_steps == max_steps, RESULTS.singular, RESULTS.successful
             )
-        else:
+        elif has_scale:
             result = RESULTS.where(
-                num_steps == max_steps,
-                RESULTS.max_steps_reached if has_scale else RESULTS.successful,
-                RESULTS.successful,
+                num_steps == max_steps, RESULTS.max_steps_reached, RESULTS.successful
             )
+        else:
+            result = RESULTS.successful
 
         if is_nsd and not self._normal:
             solution = -(solution**ω).ω

lineax/_solver/gmres.py

@@ -220,14 +220,15 @@ def body_fun(carry):
 
         if self.max_steps is None:
             result = RESULTS.where(
-                (num_steps == max_steps), RESULTS.singular, RESULTS.successful
+                num_steps == max_steps, RESULTS.singular, RESULTS.successful
             )
-        else:
+        elif has_scale:
             result = RESULTS.where(
-                (num_steps == self.max_steps),
-                RESULTS.max_steps_reached if has_scale else RESULTS.successful,
-                RESULTS.successful,
+                num_steps == max_steps, RESULTS.max_steps_reached, RESULTS.successful
             )
+        else:
+            result = RESULTS.successful
+
         result = RESULTS.where(
             stagnation_counter >= self.stagnation_iters, RESULTS.stagnation, result
         )

lineax/_solver/lsmr.py

@@ -327,19 +327,19 @@ def beta_zero(alpha, beta, u, v):
             "cond_A": loop_state["condA"],
             "norm_x": self.norm(loop_state["x"]),
         }
-        if (self.max_steps is None) or (max_steps < self.max_steps):
+
+        if self.max_steps is None:
             result = RESULTS.where(
-                loop_state["itn"] == max_steps,
-                RESULTS.singular,
-                RESULTS.successful,
+                loop_state["itn"] == max_steps, RESULTS.singular, RESULTS.successful
             )
-        else:
+        elif has_scale:
             result = RESULTS.where(
                 loop_state["itn"] == max_steps,
-                RESULTS.max_steps_reached if has_scale else RESULTS.successful,
+                RESULTS.max_steps_reached,
                 RESULTS.successful,
             )
-
+        else:
+            result = RESULTS.successful
         result = RESULTS.where(loop_state["istop"] < 3, RESULTS.successful, result)
         result = RESULTS.where(loop_state["istop"] == 3, RESULTS.conlim, result)