Minor improvements

src/ensemble/ensemble.py

@@ -366,19 +366,21 @@ def run_on_HPC(self, enX, batch_size=None, **kwargs):
         """Run the members through the simulator's HPC queue, ``batch_size`` at a time; needs the queue hooks on the wrapper."""
         import pipt.misc_tools.analysis_tools as at
 
-        list_member_index = list(range(self.ne))
+        # Use the size of the input ensemble, not self.ne: for multilevel, enX holds a single level only
+        ne = len(enX)
+        list_member_index = list(range(ne))
 
         # Split the ensemble into batches of 500
         if batch_size >= 1000:
             self.logger.info(f'Cannot run batch size of {batch_size}. Set to 1000')
             batch_size = 1000
         en_pred = []
         batch_en = [np.arange(start, start + batch_size) for start in
-                    np.arange(0, self.ne - batch_size, batch_size)]
-        if len(batch_en): # if self.ne is less than batch_size
-            batch_en.append(np.arange(batch_en[-1][-1]+1, self.ne))
+                    np.arange(0, ne - batch_size, batch_size)]
+        if len(batch_en): # if ne is less than batch_size
+            batch_en.append(np.arange(batch_en[-1][-1]+1, ne))
         else:
-            batch_en.append(np.arange(0, self.ne))
+            batch_en.append(np.arange(0, ne))
         for n_e in batch_en:
             _ = [self.sim.run_fwd_sim(state, member_index, nosim=True) for state, member_index in
                     zip([enX[curr_n] for curr_n in n_e], [list_member_index[curr_n] for curr_n in n_e])]

src/popt/ensembles/ensemble_gaussian.py

@@ -119,8 +119,11 @@ def gradient(self, x, *args, **kwargs):
             dF = self.enF[levelID] - np.repeat(self.stateF, nr)
             ne = dF.shape[0]
 
+            # Members with non-finite objective (failed simulation/eCalc) are skipped
+            ok = self._finite_members(dF)
+
             dx = self.enX[:, start_idx:start_idx + ne] - self.stateX[:, None]
-            grad_ml[levelID] = np.squeeze(dx @ dF) / ne
+            grad_ml[levelID] = np.squeeze(dx[:, ok] @ dF[ok]) / ok.sum()
             start_idx += ne
 
         if 'multilevel' in self.keys_en:
@@ -141,6 +144,16 @@ def gradient(self, x, *args, **kwargs):
 
         return grad
 
+    @staticmethod
+    def _finite_members(dF):
+        '''Mask of ensemble members with a finite objective difference. Raises if none are finite.'''
+        ok = np.isfinite(dF)
+        if not ok.any():
+            raise RuntimeError('Objective function is non-finite for all ensemble members')
+        if not ok.all():
+            print(f'Warning: skipping {np.count_nonzero(~ok)} of {ok.size} ensemble members with non-finite objective')
+        return ok
+
     def hessian(self, x=None, *args, **kwargs):
         '''
         Ensemble-based Hessian.
@@ -184,9 +197,12 @@ def hessian(self, x=None, *args, **kwargs):
             dF = self.enF[levelID] - np.repeat(self.stateF, nr)
             ne = self.enF[levelID].shape[0]
 
+            # Members with non-finite objective (failed simulation/eCalc) are skipped
+            ok = self._finite_members(dF)
+
             # Vectorized Hessian estimate for this level.
-            dx = self.enX[:, index:index + ne] - self.stateX[:, None]
-            hess_ml[levelID] = (dx * dF) @ dx.T / ne - self.covX * np.mean(dF)
+            dx = self.enX[:, index:index + ne][:, ok] - self.stateX[:, None]
+            hess_ml[levelID] = (dx * dF[ok]) @ dx.T / ok.sum() - self.covX * np.mean(dF[ok])
             index += ne
 
         if 'multilevel' in self.keys_en:

src/popt/optimization_methods/optimizer_base.py

@@ -756,7 +756,8 @@ def _set_base_restart_state(self, state):
         self.fk_old = state['fk_old']
         self.iteration = state['iteration']
         self.epf_iteration = state['epf_iteration']
-        self.ftol = state['ftol']
+        # A tolerance given in the options overrides the saved one, so it can be changed on restart
+        self.ftol = self.options.get('ftol', state['ftol']) if self.options else state['ftol']
         self.xtol = state['xtol']
         self.epf = state['epf']
         self.conv_msg = state.get('conv_msg', '')