added color to Model

EasyScience · Oct 3, 2024 · 73cb6e5 · 73cb6e5
1 parent 847d6fc
commit 73cb6e5
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/src/easyreflectometry/model/model.py b/src/easyreflectometry/model/model.py
@@ -60,6 +60,7 @@ def __init__(
         background: Union[Parameter, Number, None] = None,
         resolution_function: Union[ResolutionFunction, None] = None,
         name: str = 'EasyModel',
+        color: str = 'black',
         unique_name: Optional[str] = None,
         interface=None,
     ):
@@ -81,6 +82,7 @@ def __init__(
 
         scale = get_as_parameter('scale', scale, DEFAULTS)
         background = get_as_parameter('background', background, DEFAULTS)
+        self.color = color
 
         super().__init__(
             name=name,
@@ -173,6 +175,7 @@ def _dict_repr(self) -> dict[str, dict[str, str]]:
                 'scale': float(self.scale.value),
                 'background': float(self.background.value),
                 'resolution': resolution,
+                'color': self.color,
                 'sample': self.sample._dict_repr,
             }
         }

diff --git a/tests/model/test_model.py b/tests/model/test_model.py
@@ -389,7 +389,7 @@ def test_repr(self):
 
         assert (
             model.__repr__()
-            == 'EasyModel:\n  scale: 1.0\n  background: 1.0e-08\n  resolution: 5.0 %\n  sample:\n    EasySample:\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n'  # noqa: E501
+            == 'EasyModel:\n  scale: 1.0\n  background: 1.0e-08\n  resolution: 5.0 %\n  color: black\n  sample:\n    EasySample:\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n'  # noqa: E501
         )
 
     def test_repr_resolution_function(self):
@@ -398,7 +398,7 @@ def test_repr_resolution_function(self):
         model.resolution_function = resolution_function
         assert (
             model.__repr__()
-            == 'EasyModel:\n  scale: 1.0\n  background: 1.0e-08\n  resolution: function of Q\n  sample:\n    EasySample:\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n'  # noqa: E501
+            == 'EasyModel:\n  scale: 1.0\n  background: 1.0e-08\n  resolution: function of Q\n  color: black\n  sample:\n    EasySample:\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n    - EasyMultilayer:\n        EasyLayerCollection:\n        - EasyLayer:\n            material:\n              EasyMaterial:\n                sld: 4.186e-6 1/Å^2\n                isld: 0.000e-6 1/Å^2\n            thickness: 10.000 Å\n            roughness: 3.300 Å\n'  # noqa: E501
         )
 
 

diff --git a/tests/model/test_model_collection.py b/tests/model/test_model_collection.py
@@ -59,7 +59,7 @@ def test_delete_model(self):
 
         # Then
         collection = ModelCollection(model_1, model_2)
-        collection.remove_model(0)
+        collection.remove(0)
 
         # Expect
         assert len(collection) == 1