diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index b67ce33c7e..7fe0594e26 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -976,8 +976,10 @@ def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True) [data_short_names["Cycle"], i * np.ones_like(cycle.t)] ) for j, step in enumerate(cycle.steps): + repeated = j > 0 and cycle.steps[j - 1].t[-1] == step.t[0] + first = slice(1, None) if repeated else slice(None) data_short_names["Step"] = np.concatenate( - [data_short_names["Step"], j * np.ones_like(step.t)] + [data_short_names["Step"], j * np.ones_like(step.t[first])] ) return data_short_names diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index b81cb19afe..6b33403db6 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -881,7 +881,7 @@ def test_get_data_cycles_steps(self): data = sol.get_data_dict("c") np.testing.assert_array_equal(data["Cycle"], 0) np.testing.assert_array_equal( - data["Step"], np.concatenate([np.zeros(50), np.ones(50)]) + data["Step"], np.concatenate([np.zeros(50), np.ones(49)]) ) def test_pickle_first_states_across_processes(self, tmp_path):