Improve readability of ResidueRepresentations repr

src/sceptr/model.py

@@ -104,6 +104,9 @@ def __init__(self, representation_array: ndarray, compartment_mask: ndarray) ->
         self.representation_array = representation_array
         self.compartment_mask = compartment_mask
 
+    def __repr__(self) -> str:
+        return f"ResidueRepresentations[num_tcrs: {self.representation_array.shape[0]}, rep_dim: {self.representation_array.shape[2]}]"
+
 
 class Sceptr:
     """

tests/test_residue_representations.py

@@ -0,0 +1,14 @@
+import numpy as np
+import pytest
+from sceptr.model import ResidueRepresentations
+
+
+def test_repr(res_reps):
+    assert res_reps.__repr__() == "ResidueRepresentations[num_tcrs: 3, rep_dim: 64]"
+
+
+@pytest.fixture
+def res_reps() -> ResidueRepresentations:
+    rep_array = np.zeros((3, 10, 64))
+    comp_mask = np.zeros_like(rep_array, dtype=int)
+    return ResidueRepresentations(rep_array, comp_mask)