diff --git a/CHANGELOG.md b/CHANGELOG.md index d089db95f..efcf9add5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,8 @@ ## Bug Fixes +- [#182](https://github.com/pybop-team/PyBOP/pull/182) - Allow square-brackets indexing of Dataset + # [v23.11](https://github.com/pybop-team/PyBOP/releases/tag/v23.11) - Initial release - Adds Pints, NLOpt, and SciPy optimisers diff --git a/pybop/_dataset.py b/pybop/_dataset.py index 1263ace3e..ae3ad125a 100644 --- a/pybop/_dataset.py +++ b/pybop/_dataset.py @@ -35,6 +35,9 @@ def __init__(self, data_dictionary): self.data = data_dictionary self.names = self.data.keys() + def __getitem__(self, key): + return self.data[key] + def __repr__(self): """ Return a string representation of the Dataset instance. diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py index b5cfa96dc..934e0f4ab 100644 --- a/tests/unit/test_dataset.py +++ b/tests/unit/test_dataset.py @@ -28,6 +28,7 @@ def test_dataset(self): # Test data structure assert dataset.data == data_dictionary + assert np.all(dataset["Time [s]"] == solution["Time [s]"].data) # Test exception for non-dictionary inputs with pytest.raises(ValueError):