From 6edc039c20a6b1d7c9960e5aec9ad39d12028715 Mon Sep 17 00:00:00 2001 From: Martin Schlipf Date: Thu, 31 Aug 2023 20:46:50 +0200 Subject: [PATCH] Fix number of steps for single step structure --- src/py4vasp/_data/structure.py | 6 +++++- tests/data/test_structure.py | 4 ++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/py4vasp/_data/structure.py b/src/py4vasp/_data/structure.py index 30a3a286..bb49cdaf 100644 --- 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 index 6c5a2282..f75e95bb 100644 --- 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)