🛠️ [WIP] First attempts at Graphs

xtl.common.labels:LabelFmt
- New enum for determining how a `Label` will be formatted

xtl.common.data
- New type `TData` for defining any `Data*` container
- Fixed a type annotation issue in `DataGrid2D.z`

xtl.plots.base
- New `GraphConfig` model for storing certain global options
- New `GraphBase` model that should be inherited
- New `Graph1D` model with simple implementation of a 1D plot

Author: dtriand

.gitignore

@@ -5,7 +5,6 @@ src/xtl/scripts
 src/xtl/chemistry
 src/xtl/chemistry2
 src/xtl/crystallization2
-src/xtl/plots
 
 ### Temporary
 docs_/

src/xtl/common/data.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union
 from typing_extensions import Self
 
 from numpydantic import NDArray, Shape
@@ -8,6 +8,9 @@
 from xtl.common.labels import Label
 
 
+TData = Union['Data0D', 'Data1D', 'Data2D', 'Data3D', 'DataGrid2D']
+
+
 class Data0D(BaseModel):
     """
     A class to represent 0D data, i.e. one series of values.
@@ -86,7 +89,7 @@ class DataGrid2D(BaseModel):
 
     x: Data0D
     y: Data0D
-    z: NDArray[Shape['*', '*'], Number]
+    z: NDArray[Shape['*, *'], Number]
 
     @model_validator(mode='after')
     def check_array_shapes(self) -> Self:

src/xtl/common/labels.py

@@ -1,12 +1,20 @@
+from enum import Enum
 from typing import Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 
+class LabelFmt(Enum):
+    VALUE = 'value'
+    DESC = 'desc'
+    REPR = 'repr'
+    LATEX = 'latex'
+
+
 class Label(BaseModel):
     __pydantic_config__ = ConfigDict(validate_assignment=True, extra='forbid')
 
     value: str
     desc: Optional[str] = Field(default=None, repr=False)
-    repr: Optional[str] = Field(default=None, repr=False)
+    repr: Optional[str] = Field(default=None, repr=False)  # rename: ascii
     latex: Optional[str] = Field(default=None, repr=False)

src/xtl/plots/__init__.py

@@ -0,0 +1,72 @@
+from typing import Any, Optional
+
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+import matplotlib.pyplot as plt
+from pydantic import BaseModel, ConfigDict, Field, validate_call
+
+from xtl.common.data import TData, Data0D, Data1D
+from xtl.common.labels import Label, LabelFmt
+
+
+class GraphConfig(BaseModel):
+    model_config = ConfigDict(validate_assignment=True, extra='forbid')
+
+    label_fmt: LabelFmt = LabelFmt.VALUE
+
+
+class GraphBase(BaseModel):
+    model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True,
+                              extra='forbid')
+
+    data: TData
+    config: GraphConfig = Field(default_factory=GraphConfig)
+
+    # Fix: Should axis & figure be excluded from the model serialization?
+    axis: Optional[Axes] = Field(default_factory=lambda: plt.gca(), alias='ax')
+    figure: Optional[Figure] = Field(default_factory=lambda: plt.gcf(), alias='fig')
+
+    title: Optional[str] = None
+
+    def get_label(self, data: Data0D) -> Optional[str]:
+        value = getattr(data.label, self.config.label_repr.value)
+        units = getattr(data.units, self.config.label_repr.value)
+
+        v = value or data.label.value
+        u = units or data.units.value
+        if self.config.label_repr in [LabelFmt.VALUE, LabelFmt.LATEX]:
+            if u:
+                return f'{v} ({u})'
+            return v
+        elif self.config.label_repr == LabelFmt.REPR:
+            if u:
+                return f'{v} [{u}]'
+            return v
+        return None
+
+    def plot(self):
+        raise NotImplementedError
+
+
+class Graph1D(GraphBase):
+    data: Data1D
+
+    @property
+    def xlabel(self) -> str:
+        return self.get_label(self.data.x)
+
+    @property
+    def ylabel(self) -> str:
+        return self.get_label(self.data.y)
+
+    @validate_call
+    def plot(self, xlabel: Optional[str] = None, ylabel: Optional[str] = None):
+        if xlabel:
+            self.data.x.label = Label(value=xlabel)
+        if ylabel:
+            self.data.y.label = Label(value=ylabel)
+
+        self.axis.plot(self.data.x.data, self.data.y.data)
+
+        self.axis.set_xlabel(self.xlabel)
+        self.axis.set_ylabel(self.ylabel)