From db4690ed6d1cc39d7fab710925f0c7ad004e7a7f Mon Sep 17 00:00:00 2001 From: Daniel Garcia-Sanchez Date: Sat, 20 Oct 2018 14:26:59 +0200 Subject: [PATCH] Squeeze the dimensions of FullMatrix in initialize_container() --- source/base/hdf5.cc | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/source/base/hdf5.cc b/source/base/hdf5.cc index 84343ec9c3..cc507af1c4 100644 --- a/source/base/hdf5.cc +++ b/source/base/hdf5.cc @@ -328,8 +328,27 @@ namespace HDF5 Container>::type initialize_container(std::vector dimensions) { - AssertDimension(dimensions.size(), 2); - return Container(dimensions[0], dimensions[1]); + // If the rank is higher than 2, then remove single-dimensional entries + // from the shape defined by dimensions. This is equivalent to the squeeze + // function of python/numpy. For example the following code would convert + // the vector {1,3,1,2} to {3,2} + std::vector squeezed_dimensions; + + if (dimensions.size() > 2) + { + for (auto &&dimension : dimensions) + { + if (dimension > 1) + squeezed_dimensions.push_back(dimension); + } + } + else + { + squeezed_dimensions = dimensions; + } + + AssertDimension(squeezed_dimensions.size(), 2); + return Container(squeezed_dimensions[0], squeezed_dimensions[1]); } -- 2.39.5