updated tests.utils.make_Universe

ALlow setting of additional TopologyAttrs in extras.
Either just add the TopolAttr (as before, as a list) or use a dict with names/values.

Author: orbeckst

basicrta/tests/utils.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Tuple, Union, Dict, Any
 from contextlib import contextmanager
 import os
 
@@ -34,7 +34,7 @@ def work_in(dirname):
 
 
 def make_Universe(
-    extras: Tuple[str] = tuple(),
+    extras: Union[Tuple[str], Dict[str, Any]] = tuple(),
     size: Tuple[int, int, int] = (125, 25, 5),
     n_frames: int = 0,
     velocities: bool = False,
@@ -51,10 +51,12 @@ def make_Universe(
 
     Parameters
     ----------
-    extras : tuple of strings, optional
-      extra attributes to add to Universe:
-      u = make_Universe(('masses', 'charges'))
-      Creates a lightweight Universe with only masses and charges.
+    extras : tuple of strings or dict, optional
+      extra attributes to add to Universe. Can be:
+      - tuple of strings: u = make_Universe(('masses', 'charges'))
+        Creates a lightweight Universe with default values for these attributes.
+      - dict with values: u = make_Universe({'resnames': ['TRP', 'VAL'], 'resids': [313, 314]})
+        Creates a Universe with specific values for these attributes.
     size : tuple of int, optional
       number of elements of the Universe (n_atoms, n_residues, n_segments)
     n_frames : int
@@ -86,10 +88,18 @@ def make_Universe(
         velocities=velocities,
         forces=forces,
     )
+    
+    # Handle topology attributes
     if extras is None:
         extras = []
-    for ex in extras:
-        u.add_TopologyAttr(ex)
+    elif isinstance(extras, dict):
+        # Dict format: {'resnames': ['TRP', 'VAL'], 'resids': [313, 314]}
+        for attr_name, attr_values in extras.items():
+            u.add_TopologyAttr(attr_name, attr_values)
+    else:
+        # Tuple/list format: ('masses', 'charges') - uses default values
+        for ex in extras:
+            u.add_TopologyAttr(ex)
 
     if trajectory:
         pos = np.arange(3 * n_atoms * n_frames).reshape(n_frames, n_atoms, 3)