Allow square-brackets access to dataset objects

CHANGELOG.md

@@ -28,6 +28,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

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.

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):