🚀 New data containers

xtl.common.labels:Label
- Refactored into a pydantic.BaseModel

xtl.common.data
- New pydantic.BaseModels for storing 0D, 1D, 2D and 3D numpy data
- Additional validation is performed on the shape of the stored arrays

pyproject.toml & requirements.txt
- Added pydantic and numpydantic on the list of dependencies

Author: dtriand

pyproject.toml

@@ -38,8 +38,10 @@ dependencies = [
     "hdf5plugin>=4.2.0",
     "matplotlib>=3.8.0",
     "numpy>=1.26.0",
+    "numpydantic>=1.6.7",
     "pandas>=2.2.3",
     "pint>=0.19.2",
+    "pydantic>=2.10.6",
     "pyfai>=2023.9.0",
     "pyxray>=1.7.0",
     "requests>=2.31.0",

requirements.txt

@@ -8,8 +8,10 @@ gemmi==0.5.8
 hdf5plugin==4.2.0
 matplotlib==3.8.0
 numpy==1.26.0
+numpydantic==1.6.7
 pandas==2.2.3
 pint==0.19.2
+pydantic==2.10.6
 pyfai==2023.9.0
 pyxray==1.7.0
 requests==2.31.0

src/xtl/common/data.py

@@ -0,0 +1,98 @@
+from typing import Optional
+from typing_extensions import Self
+
+from numpydantic import NDArray, Shape
+from numpydantic.dtype import Number
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from xtl.common.labels import Label
+
+
+class Data0D(BaseModel):
+    """
+    A class to represent 0D data, i.e. one series of values.
+    """
+    model_config = ConfigDict(validate_assignment=True, extra='forbid')
+
+    data: NDArray[Shape['*'], Number]
+    label: Label
+    units: Optional[Label] = None
+
+
+class Data1D(BaseModel):
+    """
+    A class to represent 1D data, i.e. y = f(x).
+    """
+    model_config = ConfigDict(validate_assignment=True, extra='forbid')
+
+    x: Data0D
+    y: Data0D
+    
+    @model_validator(mode='after')
+    def check_array_shapes(self) -> Self:
+        if self.x.data.shape != self.y.data.shape:
+            raise ValueError('x and y shapes must be equal')
+        return self
+
+
+class Data2D(BaseModel):
+    """
+    A class to represent 2D data, i.e. z = f(x, y).
+
+    """
+    model_config = ConfigDict(validate_assignment=True, extra='forbid')
+
+    x: Data0D
+    y: Data0D
+    z: Data0D
+
+    @model_validator(mode='after')
+    def check_array_shapes(self) -> Self:
+        if self.x.data.shape != self.y.data.shape:
+            raise ValueError('x and y shapes must be equal')
+        if self.x.data.shape != self.z.data.shape:
+            raise ValueError('x and z shapes must be equal')
+        return self
+
+
+class Data3D(BaseModel):
+    """
+    A class to represent 3D data, i.e. w = f(x, y, z).
+    """
+    model_config = ConfigDict(validate_assignment=True, extra='forbid')
+
+    x: Data0D
+    y: Data0D
+    z: Data0D
+    w: Data0D
+
+    @model_validator(mode='after')
+    def check_array_shapes(self) -> Self:
+        if self.x.data.shape != self.y.data.shape:
+            raise ValueError('x and y shapes must be equal')
+        if self.x.data.shape != self.z.data.shape:
+            raise ValueError('x and z shapes must be equal')
+        if self.x.data.shape != self.w.data.shape:
+            raise ValueError('x and w shapes must be equal')
+        return self
+
+
+class DataGrid2D(BaseModel):
+    """
+    A class to represent 2D data in a grid format, i.e. z = f(x, y), where x and y are
+    series of coordinates. Useful for calculating 2D functions on a grid.
+    """
+    model_config = ConfigDict(validate_assignment=True, extra='forbid')
+
+    x: Data0D
+    y: Data0D
+    z: NDArray[Shape['*', '*'], Number]
+
+    @model_validator(mode='after')
+    def check_array_shapes(self) -> Self:
+        if self.x.data.shape[0] != self.z.shape[0]:
+            raise ValueError('x and z shapes must be equal')
+        if self.y.data.shape[0] != self.z.shape[1]:
+            raise ValueError('y and z shapes must be equal')
+        return self
+

src/xtl/common/labels.py

@@ -1,10 +1,12 @@
-from dataclasses import dataclass
 from typing import Optional
 
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Label(BaseModel):
+    __pydantic_config__ = ConfigDict(validate_assignment=True, extra='forbid')
 
-@dataclass
-class Label:
     value: str
-    description: Optional[str] = None
-    repr: Optional[str] = None
-    latex: Optional[str] = None
+    desc: Optional[str] = Field(default=None, repr=False)
+    repr: Optional[str] = Field(default=None, repr=False)
+    latex: Optional[str] = Field(default=None, repr=False)