stdlib: Add Vector2d to PyStats

Change-Id: Icb2f691abf88ef4bac8d277e421329edb000209b

src/python/m5/ext/pystats/statistic.py

@@ -142,6 +142,64 @@ class Vector(Statistic):
         return sum(float(self.value[key]) for key in self.values)
 
 
+class Vector2d(Statistic):
+    """
+    A 2D vector of scalar values.
+    """
+
+    value: Dict[Union[str, int, float], Vector]
+
+    def __init__(
+        self,
+        value: Dict[Union[str, int, float], Vector],
+        type: Optional[str] = None,
+        description: Optional[str] = None,
+    ):
+        assert (
+            len({vector.size() for vector in value.values()}) == 1
+        ), "All the Vectors in the 2d Vector are not of equal length."
+
+        super().__init__(
+            value=value,
+            type=type,
+            description=description,
+        )
+
+    def x_size(self) -> int:
+        """Returns the number of elements in the x dimension."""
+        assert self.value is not None
+        return len(self.value)
+
+    def y_size(self) -> int:
+        """Returns the number of elements in the y dimension."""
+        assert self.value is not None
+        return len(self.value[0])
+
+    def size(self) -> int:
+        """Returns the total number of elements."""
+        return self.x_size() * self.y_size()
+
+    def total(self) -> int:
+        """The total (sum) of all the entries in the 2d vector/"""
+        assert self.value is not None
+        total = 0
+        for vector in self.value.values():
+            for scalar in vector.values():
+                total += scalar.value
+        return total
+
+    def __getitem__(self, index: Union[str, int, float]) -> Vector:
+        assert self.value is not None
+        # In the case of string, we cast strings to integers of floats if they
+        # are numeric. This avoids users having to cast strings to integers.
+        if isinstance(index, str):
+            if index.isindex():
+                index = int(index)
+            elif index.isnumeric():
+                index = float(index)
+        return self.value[index]
+
+
 class Distribution(Vector):
     """
     A statistic type that stores information relating to distributions. Each

src/python/m5/stats/gem5stats.py

@@ -138,6 +138,8 @@ def __get_statistic(statistic: _m5.stats.Info) -> Optional[Statistic]:
         pass
     elif isinstance(statistic, _m5.stats.VectorInfo):
         return __get_vector(statistic)
+    elif isinstance(statistic, _m5.stats.Vector2dInfo):
+        return __get_vector2d(statistic)
 
     return None
 
@@ -193,7 +195,6 @@ def __get_distribution(statistic: _m5.stats.DistInfo) -> Distribution:
 def __get_vector(statistic: _m5.stats.VectorInfo) -> Vector:
     vec: Dict[Union[str, int, float], Scalar] = {}
 
-
     for index in range(statistic.size):
         # All the values in a Vector are Scalar values
         value = statistic.value[index]
@@ -231,6 +232,42 @@ def __get_vector(statistic: _m5.stats.VectorInfo) -> Vector:
     )
 
 
+def __get_vector2d(statistic: _m5.stats.Vector2dInfo) -> Vector2d:
+    # All the values in a 2D Vector are Scalar values
+    description = statistic.desc
+    x_size = statistic.x_size
+    y_size = statistic.y_size
+
+    vector_rep: Dict[Union[str, int, float], Vector] = {}
+    for x_index in range(x_size):
+        x_index_string = x_index
+        if x_index in statistic.subnames:
+            x_index_string = str(statistic.subnames[x_index])
+
+        x_desc = description
+        if x_index in statistic.subdescs:
+            x_desc = str(statistic.subdescs[x_index])
+        x_vec: Dict[str, Scalar] = {}
+        for y_index in range(y_size):
+            y_index_val = y_index
+            if y_index in statistic.ysubnames:
+                y_index_val = str(statistic.subnames[y_index])
+
+            x_vec[y_index_val] = Scalar(
+                value=statistic.value[x_index * y_size + y_index],
+                unit=statistic.unit,
+                datatype=StorageType["f64"],
+            )
+
+        vector_rep[x_index_string] = Vector(
+            x_vec,
+            type="Vector",
+            description=x_desc,
+        )
+
+    return Vector2d(value=vector_rep, type="Vector2d", description=description)
+
+
 def _prepare_stats(group: _m5.stats.Group):
     """
     Prepares the statistics for dumping.

src/python/pybind11/stats.cc

@@ -76,6 +76,7 @@ cast_stat_info(const statistics::Info *info)
      */
     TRY_CAST(statistics::FormulaInfo);
     TRY_CAST(statistics::VectorInfo);
+    TRY_CAST(statistics::Vector2dInfo);
     TRY_CAST(statistics::DistInfo);
 
     return py::cast(info);
@@ -183,6 +184,17 @@ pybind_init_stats(py::module_ &m_native)
             [](const statistics::VectorInfo &info) { return info.total(); })
         ;
 
+    py::class_<statistics::Vector2dInfo, statistics::Info,
+               std::unique_ptr<statistics::Vector2dInfo, py::nodelete>>(
+                    m, "Vector2dInfo")
+        .def_readonly("x_size", &statistics::Vector2dInfo::x)
+        .def_readonly("y_size", &statistics::Vector2dInfo::y)
+        .def_readonly("subnames", &statistics::Vector2dInfo::subnames)
+        .def_readonly("subdescs", &statistics::Vector2dInfo::subdescs)
+        .def_readonly("ysubnames", &statistics::Vector2dInfo::y_subnames)
+        .def_readonly("value", &statistics::Vector2dInfo::cvec)
+        ;
+
     py::class_<statistics::FormulaInfo, statistics::VectorInfo,
                std::unique_ptr<statistics::FormulaInfo, py::nodelete>>(
                       m, "FormulaInfo")