optimizer: report failed structures as None instead of dropping them

src/gpuma/api.py

@@ -160,13 +160,17 @@ def optimize_ensemble_smiles(
         s.multiplicity = multiplicity
     results = optimize_structure_batch(conformers, config)
 
+    # Number by position in the input, not by position among the survivors: a
+    # conformer that failed to optimize must not renumber the ones after it.
+    optimized = [(i, s) for i, s in enumerate(results) if s is not None]
+
     if output_file:
         comments = [
-            f"Optimized conformer {i + 1} from SMILES: {smiles}" for i in range(len(results))
+            f"Optimized conformer {i + 1} from SMILES: {smiles}" for i, _ in optimized
         ]
-        save_multi_xyz(results, output_file, comments)
+        save_multi_xyz([s for _, s in optimized], output_file, comments)
 
-    return results
+    return [s for _, s in optimized]
 
 
 def optimize_batch_multi_xyz_file(
@@ -210,13 +214,17 @@ def optimize_batch_multi_xyz_file(
     )
     results = optimize_structure_batch(structures, config)
 
+    # Numbered by position in the file, so a structure that failed to optimize
+    # leaves a gap rather than shifting every later label onto the wrong input.
+    optimized = [(i, s) for i, s in enumerate(results) if s is not None]
+
     if output_file:
         comments = [
-            f"Optimized structure {i + 1} from: {input_file}" for i in range(len(results))
+            f"Optimized structure {i + 1} from: {input_file}" for i, _ in optimized
         ]
-        save_multi_xyz(results, output_file, comments)
+        save_multi_xyz([s for _, s in optimized], output_file, comments)
 
-    return results
+    return [s for _, s in optimized]
 
 
 def optimize_batch_xyz_directory(
@@ -255,14 +263,16 @@ def optimize_batch_xyz_directory(
     )
     results = optimize_structure_batch(structures, config)
 
+    # Numbered by position in the directory listing; see the note above.
+    optimized = [(i, s) for i, s in enumerate(results) if s is not None]
+
     if output_file:
         comments = [
-            f"Optimized structure {i + 1} from batch input"
-            for i in range(len(results))
+            f"Optimized structure {i + 1} from batch input" for i, _ in optimized
         ]
-        save_multi_xyz(results, output_file, comments)
+        save_multi_xyz([s for _, s in optimized], output_file, comments)
 
-    return results
+    return [s for _, s in optimized]
 
 
 __all__ = [

src/gpuma/optimizer.py

@@ -301,7 +301,7 @@ def optimize_single_structure(
 def optimize_structure_batch(
     structures: list[Structure],
     config: Config | None = None,
-) -> list[Structure]:
+) -> list[Structure | None]:
     """Optimize a list of structures and return them with updated coordinates.
 
     The optimization mode is controlled by
@@ -321,8 +321,18 @@ def optimize_structure_batch(
 
     Returns
     -------
-    list[Structure]
-        Optimized structures with coordinates and energies set.
+    list[Structure | None]
+        One entry per input, in input order. ``None`` marks a structure whose
+        optimization failed, so a single bad geometry cannot abort the run.
+
+        .. versionchanged::
+            Previously the failures were simply left out, which silently
+            shortened the list and shifted every later result onto the wrong
+            input. Callers that index results against inputs, or label output
+            by position, were quietly misattributing geometries. Iterating
+            callers now need to skip ``None``; this matches
+            :func:`gpuma.conformer_generation.embed.generate_structures`,
+            which already reports per-molecule failure this way.
 
     Raises
     ------
@@ -373,23 +383,29 @@ def optimize_structure_batch(
 def _optimize_sequential(
     structures: list[Structure],
     config: Config,
-) -> list[Structure]:
-    """Optimize structures one-by-one using ASE with a shared calculator."""
+) -> list[Structure | None]:
+    """Optimize structures one-by-one using ASE with a shared calculator.
+
+    Returns one entry per input, in input order, with ``None`` where the
+    optimization raised. Appending only the successes -- as this used to --
+    made the returned list shorter than the input and shifted every result
+    after a failure onto the wrong structure, with nothing to signal it.
+    """
     calculator = _get_cached_calculator(config)
 
     logger.info("Starting sequential optimization of %d structures", len(structures))
-    results: list[Structure] = []
+    results: list[Structure | None] = [None] * len(structures)
+    n_ok = 0
     for i, struct in enumerate(structures):
         try:
-            optimized = optimize_single_structure(struct, config, calculator)
-            results.append(optimized)
+            results[i] = optimize_single_structure(struct, config, calculator)
+            n_ok += 1
         except Exception as exc:  # pragma: no cover - defensive logging
             logger.warning("Structure %d optimization failed: %s", i + 1, exc)
-            continue
 
     logger.info(
         "Sequential optimization completed. %d/%d successful",
-        len(results),
+        n_ok,
         len(structures),
     )
     return results
@@ -398,7 +414,7 @@ def _optimize_sequential(
 def _optimize_batch(
     structures: list[Structure],
     config: Config,
-) -> list[Structure]:
+) -> list[Structure | None]:
     """Optimize structures in parallel using torch-sim batch inference."""
     import torch
     import torch_sim

src/gpuma/utils/logging_utils.py

@@ -52,7 +52,9 @@ def log_optimization_summary(
     input_structures:
         Structures that were submitted for optimization.
     results:
-        Successfully optimized structures.
+        One entry per input, in input order, with ``None`` where the
+        optimization failed. Those entries are excluded from every statistic
+        below, so the success rate reflects what actually came back.
     total_time:
         Wall-clock time in seconds.
     mode:
@@ -64,9 +66,10 @@ def log_optimization_summary(
     from ..config import resolve_model_type
 
     n_input = len(input_structures)
-    n_output = len(results)
-    energies = [s.energy for s in results if s.energy is not None]
-    atom_counts = [s.n_atoms for s in results]
+    optimized = [s for s in results if s is not None]
+    n_output = len(optimized)
+    energies = [s.energy for s in optimized if s.energy is not None]
+    atom_counts = [s.n_atoms for s in optimized]
 
     model_name = getattr(config.model, "model_name", "unknown")
     model_type = resolve_model_type(config)

tests/optimization/test_batch.py

@@ -159,3 +159,149 @@ def _copy(structs):
                 f"Sequential ({s.energy:.4f}) and batch ({b.energy:.4f}) "
                 f"energies differ by more than 1 eV"
             )
+
+
+class TestFailureAlignment:
+    """A structure that fails to optimize must not shift the others.
+
+    A stub calculator stands in for the MLIP: the property under test is how
+    _optimize_sequential handles an exception, and a real model would have to
+    be coaxed into failing on cue to show it.
+    """
+
+    @staticmethod
+    def _flaky_calculator(fail_on_n_atoms):
+        import numpy as np
+
+        class FlakyCalc:
+            def _check(self, atoms):
+                if len(atoms) == fail_on_n_atoms:
+                    raise RuntimeError("simulated failure")
+
+            def get_potential_energy(self, atoms=None, force_consistent=False):
+                self._check(atoms)
+                return 0.0
+
+            def get_forces(self, atoms=None):
+                self._check(atoms)
+                return np.zeros((len(atoms), 3))
+
+            def get_property(self, name, atoms=None, allow_calculation=True):
+                self._check(atoms)
+                return {"energy": 0.0, "forces": np.zeros((len(atoms), 3))}[name]
+
+            def calculation_required(self, atoms, properties):
+                return True
+
+        return FlakyCalc()
+
+    @staticmethod
+    def _chain(n, tag):
+        from gpuma.structure import Structure
+
+        return Structure(
+            symbols=["H"] * n,
+            coordinates=[(i * 1.0, 0.0, 0.0) for i in range(n)],
+            charge=0,
+            multiplicity=1,
+            comment=tag,
+        )
+
+    def test_failure_becomes_none_in_place(self, monkeypatch):
+        """Results stay aligned to inputs, with None marking the failure.
+
+        The old code appended only successes, so [A, B, C_fails, D] came back
+        as [A, B, D] -- and D was then written out labelled as structure 3.
+        """
+        from gpuma import optimizer as opt
+        from gpuma.config import Config
+
+        monkeypatch.setattr(
+            opt, "_get_cached_calculator", lambda config: self._flaky_calculator(10)
+        )
+        config = Config({
+            "optimization": {"batch_optimization_mode": "sequential"},
+            "technical": {"device": "cpu"},
+        })
+        inputs = [
+            self._chain(3, "A"),
+            self._chain(4, "B"),
+            self._chain(10, "C_fails"),
+            self._chain(5, "D"),
+        ]
+
+        results = opt.optimize_structure_batch(inputs, config)
+
+        assert len(results) == len(inputs)
+        assert results[2] is None
+        assert [r.comment for r in results if r is not None] == ["A", "B", "D"]
+        assert results[3].comment == "D"
+
+    def test_all_succeed_has_no_none(self, monkeypatch):
+        """The common path is unchanged."""
+        from gpuma import optimizer as opt
+        from gpuma.config import Config
+
+        monkeypatch.setattr(
+            opt, "_get_cached_calculator", lambda config: self._flaky_calculator(-1)
+        )
+        config = Config({
+            "optimization": {"batch_optimization_mode": "sequential"},
+            "technical": {"device": "cpu"},
+        })
+        inputs = [self._chain(3, "A"), self._chain(4, "B")]
+
+        results = opt.optimize_structure_batch(inputs, config)
+
+        assert [r.comment for r in results] == ["A", "B"]
+
+    def test_summary_counts_only_successes(self, monkeypatch, caplog):
+        """None entries must not be counted as optimized structures."""
+        import logging as _logging
+
+        from gpuma import optimizer as opt
+        from gpuma.config import Config
+
+        monkeypatch.setattr(
+            opt, "_get_cached_calculator", lambda config: self._flaky_calculator(10)
+        )
+        config = Config({
+            "optimization": {"batch_optimization_mode": "sequential"},
+            "technical": {"device": "cpu"},
+        })
+        inputs = [self._chain(3, "A"), self._chain(10, "B_fails")]
+
+        with caplog.at_level(_logging.INFO):
+            opt.optimize_structure_batch(inputs, config)
+
+        assert "Structures output:   1" in caplog.text
+        assert "Success rate:        1/2" in caplog.text
+
+    def test_api_labels_survivors_by_input_position(self, monkeypatch, tmp_path):
+        """The output file must not renumber structures around a failure."""
+        from gpuma import api, optimizer as opt
+        from gpuma.config import Config
+
+        monkeypatch.setattr(
+            opt, "_get_cached_calculator", lambda config: self._flaky_calculator(10)
+        )
+        inputs = [
+            self._chain(3, "A"),
+            self._chain(10, "B_fails"),
+            self._chain(5, "C"),
+        ]
+        monkeypatch.setattr(api, "read_multi_xyz", lambda *a, **kw: inputs)
+        monkeypatch.setattr(api.os.path, "isfile", lambda p: True)
+
+        out = tmp_path / "out.xyz"
+        config = Config({
+            "optimization": {"batch_optimization_mode": "sequential"},
+            "technical": {"device": "cpu"},
+        })
+        api.optimize_batch_multi_xyz_file("in.xyz", str(out), config)
+
+        text = out.read_text()
+        # The third input keeps label 3; it must not inherit the failure's 2.
+        assert "Optimized structure 1 from" in text
+        assert "Optimized structure 3 from" in text
+        assert "Optimized structure 2 from" not in text