diff --git a/doc/source/credits.rst b/doc/source/credits.rst index 01ec5b0..52b0ac4 100644 --- a/doc/source/credits.rst +++ b/doc/source/credits.rst @@ -2,6 +2,7 @@ Credits ======= +* Dimitris Arquin * Michael P. Howard * Mayukh Kundu * Philipp Leclercq diff --git a/src/lammpsio/topology.py b/src/lammpsio/topology.py index 5f05241..43b1583 100644 --- a/src/lammpsio/topology.py +++ b/src/lammpsio/topology.py @@ -381,12 +381,23 @@ class LabelMap(collections.abc.MutableMapping[int, str]): type_label = lammpsio.topology.LabelMap({1: "A", 2: "B"}) This creates a dictionary mapping numeric type ID labels 1 and 2 used by LAMMPS - to alphanumeric type labels "A" and "B", such as those used by HOOMD-blue. + to alphanumeric type labels "A" and "B", such as those used by HOOMD-blue. The + map can be accessed directly: + .. code-block:: python + + assert type_label[1] == "A" + + The `LabelMap` additionally supports inverse mapping from type label to type ID: + + .. code-block:: python + + assert type_label.inverse["B"] == 2 """ def __init__(self, map=None): self._map = {} + self._inverse_map = {} if map is not None: self.update(map) @@ -395,9 +406,11 @@ def __getitem__(self, key): def __setitem__(self, key, value): self._map[key] = value + self._inverse_map[value] = key def __delitem__(self, key): - del self._map[key] + value = self._map.pop(key) + del self._inverse_map[value] def __iter__(self): return iter(self._map) @@ -414,3 +427,8 @@ def types(self): def typeid(self): """tuple of int: Type IDs in map.""" return tuple(self._map.keys()) + + @property + def inverse(self) -> dict[str, int]: + """dict: Inverse map from type label to type id.""" + return self._inverse_map diff --git a/tests/test_topology.py b/tests/test_topology.py index c82d7c4..1917fea 100644 --- a/tests/test_topology.py +++ b/tests/test_topology.py @@ -204,16 +204,19 @@ def test_impropers_wrong_shape(snap_8): def test_LabelMap(): # create a simple label map label = lammpsio.topology.LabelMap({1: "typeA", 2: "typeB"}) - # check the types assert label.types == ("typeA", "typeB") # check the typeids assert label.typeid == (1, 2) # get assert label[2] == "typeB" + # get inverse + assert label.inverse["typeB"] == 2 # set label[3] = "typeC" assert label[3] == "typeC" + assert label.inverse["typeC"] == 3 # delete del label[3] assert 3 not in label + assert "typeC" not in label.inverse