Fix number of steps for single step structure

vasp-dev · Aug 31, 2023 · 6edc039 · 6edc039
1 parent 0b9b2ff
commit 6edc039
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/py4vasp/_data/structure.py b/src/py4vasp/_data/structure.py
@@ -277,7 +277,11 @@ def number_atoms(self):
     @base.data_access
     def number_steps(self):
         """Return the number of structures in the trajectory."""
-        return len(self._raw_data.positions[self._slice])
+        if self._is_trajectory:
+            range_ = range(len(self._raw_data.positions))
+            return len(range_[self._slice])
+        else:
+            return 1
 
     def _topology(self):
         return data.Topology.from_data(self._raw_data.topology)

diff --git a/tests/data/test_structure.py b/tests/data/test_structure.py
@@ -289,6 +289,10 @@ def test_plot_Fe3O4(Fe3O4, not_core):
         check_plot_structure(Fe3O4[steps])
 
 
+def test_plot_Ca3AsBr3(Ca3AsBr3, not_core):
+    check_plot_structure(Ca3AsBr3)
+
+
 def check_plot_structure(structure):
     obj = viewer3d.Viewer3d
     cm_init = patch.object(obj, "__init__", autospec=True, return_value=None)