Source code for SDF.data_model.array_data_2d

import re
from ast import literal_eval
from itertools import islice
from xml.etree.ElementTree import Element

import numpy as np

from SDF.data_model._helper_functions import pop_element_text, pop_element_attribute, element_is_empty
from SDF.data_model.abstract import Data


class ArrayData2D(Data):
    """Wraps a 2-dimensional Numpy ndarray, represents SDF <data> elements of dataset type 'mc'"""

    def __init__(self, data: np.ndarray):
        if not isinstance(data, np.ndarray):
            raise TypeError(f"Expected a numpy ndarray, got {type(data)}")
        if data.ndim != 2:
            raise ValueError(f"Expected a 2D array, got {data.ndim} dimensions (shape: {data.shape})")

        self.__array = data

    @property
    def data(self):
        """The wrapped array"""
        return self.__array

    @property
    def type_for_xml(self) -> str:
        return "mc"

    def to_xml_element(self) -> Element:
        element = Element("data")

        # shape
        rows, cols = self.data.shape
        element.attrib["rows"] = str(rows)
        element.attrib["cols"] = str(cols)

        # data to 1D string iterable
        if np.issubdtype(self.data.dtype, np.integer):
            element.attrib["type"] = "int"
            text_iter = map(str, self.data.flatten())
        elif np.issubdtype(self.data.dtype, np.floating):
            element.attrib["type"] = "float"
            text_iter = map(str, self.data.flatten())
        else:
            raise ValueError(f"Cannot handle data of dtype {self.data.dtype}")

        # format data to XML text
        element.text = "\n".join(" ".join(islice(text_iter, 0, cols)) for _ in range(rows))

        return element

    @classmethod
    def from_xml_element(cls, element: Element) -> "ArrayData2D":
        if element.tag != "data":
            raise ValueError(f"Expected a <data> element, got {element.tag}")

        text = pop_element_text(element)
        xml_dtype = pop_element_attribute(element, "type")

        # read shape
        if "rows" in element.attrib and "cols" in element.attrib:
            # remove "shape" attribute of old multi-column SDF class
            if "shape" in element.attrib:
                pop_element_attribute(element, "shape")

            # <data rows="..." cols="..." ... >
            rows = int(pop_element_attribute(element, "rows"))
            cols = int(pop_element_attribute(element, "cols"))
            if rows == 1:
                shape = (cols,)
            elif cols == 1:
                shape = (rows,)
            else:
                shape = (rows, cols)
        elif "shape" in element.attrib:
            # <data shape="(rows, cols)" ... >, as generated by old multi-column SDF class
            shape_str = pop_element_attribute(element, "shape")
            if re.match(r"\([0-9]+, ?[0-9]+\)", shape_str):  # like '(50, 10)'
                rows, cols = literal_eval(shape_str)
                shape = (rows, cols)
            else:
                raise ValueError(f"Expected a 2D shape string like '(50, 10)', got '{shape_str}'")
        else:
            raise ValueError("No shape information found")

        # read data
        if xml_dtype == "int":
            data = np.fromiter(map(int, text.split()), dtype=int)
        elif xml_dtype == "float":
            data = np.fromiter(map(float, text.split()), dtype=float)
        else:
            raise ValueError(f"Expected attribute 'type' to be 'int', 'float' or 'hex', got '{xml_dtype}'")

        if not element_is_empty(element):
            raise ValueError("Element is not empty")

        return cls(data.reshape(shape))

    def __repr__(self):
        return f"{self.__class__.__name__}({self.data!r})"

    def __eq__(self, other):
        if isinstance(other, ArrayData2D):
            return (self.data.shape == other.data.shape
                    and np.allclose(self.data, other.data, equal_nan=True))
        return False