Guard against division by zero in JAX where clauses.

torax/_src/core_profiles/getters.py

@@ -684,17 +684,32 @@ def get_updated_ions(
       * ion_properties.dilution_factor_edge,
   )
 
+  # Guard the dead branch of jnp.where against division by zero.
+  # JAX evaluates both branches, so we must use a safe denominator
+  # to prevent NaN propagation when Z_impurity is zero or degenerate.
+  # dilution_factor == 1.0 indicates no impurities (Z_eff=1.0), so we set
+  # safe_Z_impurity to an arbitrary non-zero value, i.e. 1.0.
+  safe_Z_impurity = jnp.where(
+      ion_properties.dilution_factor == 1.0,
+      1.0,
+      ion_properties.Z_impurity,
+  )
   n_impurity_value = jnp.where(
       ion_properties.dilution_factor == 1.0,
       0.0,
-      (n_e.value - n_i.value * Z_i) / ion_properties.Z_impurity,
+      (n_e.value - n_i.value * Z_i) / safe_Z_impurity,
   )
 
+  safe_Z_impurity_edge = jnp.where(
+      ion_properties.dilution_factor_edge == 1.0,
+      1.0,
+      ion_properties.Z_impurity_face[-1],
+  )
   n_impurity_right_face_constraint = jnp.where(
       ion_properties.dilution_factor_edge == 1.0,
       0.0,
       (n_e.right_face_constraint - n_i.right_face_constraint * Z_i_face[-1])
-      / ion_properties.Z_impurity_face[-1],
+      / safe_Z_impurity_edge,
   )
 
   n_impurity = cell_variable.CellVariable(

torax/_src/pedestal_model/set_pped_tpedratio_nped.py

@@ -88,12 +88,25 @@ def _call_implementation(
     Z_eff_ped = jnp.take(Z_eff, ped_idx)
     Z_i_ped = jnp.take(Z_i, ped_idx)
     Z_impurity_ped = jnp.take(Z_impurity, ped_idx)
-    dilution_factor_ped = formulas.calculate_main_ion_dilution_factor(
-        Z_i_ped, Z_impurity_ped, Z_eff_ped
+    dilution_factor_ped = jnp.where(
+        Z_eff_ped == 1.0,
+        1.0,
+        formulas.calculate_main_ion_dilution_factor(
+            Z_i_ped, Z_impurity_ped, Z_eff_ped
+        ),
     )
     # Calculate n_i and n_impurity.
     n_i_ped = dilution_factor_ped * n_e_ped
-    n_impurity_ped = (n_e_ped - Z_i_ped * n_i_ped) / Z_impurity_ped
+    # Guard against division by zero when Z_impurity_ped is zero or
+    # degenerate (pure plasma case).
+    safe_Z_impurity_ped = jnp.where(
+        Z_eff_ped == 1.0, 1.0, Z_impurity_ped
+    )
+    n_impurity_ped = jnp.where(
+        Z_eff_ped == 1.0,
+        0.0,
+        (n_e_ped - Z_i_ped * n_i_ped) / safe_Z_impurity_ped,
+    )
     # Assumption that impurity is at the same temperature as the ion AND
     # the pressure P = P_e + P_i + P_imp.
     # P = T_e*n_e + T_i*n_i + T_i*n_imp.

torax/_src/physics/charge_states.py

@@ -112,7 +112,11 @@ class ChargeStateInfo:
 
   @property
   def Z_mixture(self) -> array_typing.FloatVector:
-    return self.Z2_avg / self.Z_avg
+    # Guard against division by zero when Z_avg is zero (e.g. no impurities).
+    # This also prevents NaN propagation through jnp.where dead branches that
+    # evaluate this property.
+    safe_Z_avg = jnp.where(self.Z_avg == 0.0, 1.0, self.Z_avg)
+    return jnp.where(self.Z_avg == 0.0, 0.0, self.Z2_avg / safe_Z_avg)
 
 
 # pylint: disable=invalid-name