initial commit

backend/data_types.py

@@ -0,0 +1,134 @@
+"""
+Core data types for argonode.
+
+DataField mirrors Gwyddion's GwyDataField structure:
+  xres, yres  – pixel dimensions
+  xreal, yreal – physical dimensions in metres
+  xoff, yoff   – position offset in metres
+  si_unit_xy   – lateral unit string (e.g. "m", "nm")
+  si_unit_z    – height/value unit string (e.g. "m", "V", "A")
+  domain       – "spatial" or "frequency" (set by FFT nodes)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import numpy as np
+
+
+@dataclass
+class DataField:
+    data: np.ndarray        # shape (yres, xres), dtype float64
+    xres: int = 0
+    yres: int = 0
+    xreal: float = 1e-6    # physical width in metres
+    yreal: float = 1e-6    # physical height in metres
+    xoff: float = 0.0
+    yoff: float = 0.0
+    si_unit_xy: str = "m"
+    si_unit_z: str = "m"
+    domain: str = "spatial"  # "spatial" or "frequency"
+
+    def __post_init__(self) -> None:
+        self.data = np.asarray(self.data, dtype=np.float64)
+        if self.data.ndim != 2:
+            raise ValueError(f"DataField.data must be 2-D, got shape {self.data.shape}")
+        self.yres, self.xres = self.data.shape
+
+    def copy(self) -> "DataField":
+        """Return a deep copy with independent data array."""
+        return DataField(
+            data=self.data.copy(),
+            xres=self.xres,
+            yres=self.yres,
+            xreal=self.xreal,
+            yreal=self.yreal,
+            xoff=self.xoff,
+            yoff=self.yoff,
+            si_unit_xy=self.si_unit_xy,
+            si_unit_z=self.si_unit_z,
+            domain=self.domain,
+        )
+
+    def replace(self, **kwargs) -> "DataField":
+        """Return a copy with selected fields replaced. data is deep-copied unless provided."""
+        base = {
+            "data": self.data.copy(),
+            "xres": self.xres,
+            "yres": self.yres,
+            "xreal": self.xreal,
+            "yreal": self.yreal,
+            "xoff": self.xoff,
+            "yoff": self.yoff,
+            "si_unit_xy": self.si_unit_xy,
+            "si_unit_z": self.si_unit_z,
+            "domain": self.domain,
+        }
+        base.update(kwargs)
+        return DataField(**base)
+
+    @property
+    def dx(self) -> float:
+        """Physical pixel size in x (metres)."""
+        return self.xreal / self.xres if self.xres else 1.0
+
+    @property
+    def dy(self) -> float:
+        """Physical pixel size in y (metres)."""
+        return self.yreal / self.yres if self.yres else 1.0
+
+
+# ---------------------------------------------------------------------------
+# Utility helpers shared across nodes
+# ---------------------------------------------------------------------------
+
+def datafield_to_uint8(df: DataField, colormap: str = "gray") -> np.ndarray:
+    """
+    Normalize a DataField to a uint8 (H, W, 3) RGB array using matplotlib colormap.
+    Returns shape (H, W, 3) uint8.
+    """
+    import matplotlib.cm as cm
+    import matplotlib.colors as mcolors
+
+    data = df.data
+    dmin, dmax = data.min(), data.max()
+    if dmax > dmin:
+        normalized = (data - dmin) / (dmax - dmin)
+    else:
+        normalized = np.zeros_like(data)
+
+    cmap = cm.get_cmap(colormap)
+    rgba = cmap(normalized)                     # (H, W, 4) float [0,1]
+    rgb = (rgba[:, :, :3] * 255).astype(np.uint8)
+    return rgb
+
+
+def image_to_uint8(image: np.ndarray) -> np.ndarray:
+    """
+    Convert an IMAGE (float or uint8, 2-D or 3-D) to uint8 (H,W,3) or (H,W) for PIL.
+    """
+    if image.dtype == np.uint8:
+        return image
+    # float — normalize to [0, 255]
+    imin, imax = image.min(), image.max()
+    if imax > imin:
+        out = (image - imin) / (imax - imin) * 255.0
+    else:
+        out = np.zeros_like(image)
+    return out.astype(np.uint8)
+
+
+def encode_preview(arr: np.ndarray) -> str:
+    """
+    Encode a uint8 numpy array as a base64 data URI (PNG).
+    arr: (H, W) grayscale or (H, W, 3) RGB, uint8.
+    """
+    import base64
+    import io
+    from PIL import Image
+
+    img = Image.fromarray(arr)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"