Line data    Source code

       1             : // SPDX-License-Identifier: BSD-3-Clause
       2             : // Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
       3             : /// @file
       4             : /// @author Matthew Andrew
       5             : 
       6             : #include "scipp/dataset/util.h"
       7             : #include "scipp/common/overloaded.h"
       8             : #include "scipp/core/element/arg_list.h"
       9             : #include "scipp/dataset/sized_dict_forward.h"
      10             : #include "scipp/variable/accumulate.h"
      11             : #include "scipp/variable/reduction.h"
      12             : #include "scipp/variable/structure_array_model.h"
      13             : #include "scipp/variable/util.h"
      14             : #include "scipp/variable/variable.h"
      15             : #include "scipp/variable/variable_concept.h"
      16             : 
      17             : #include "dataset_operations_common.h"
      18             : 
      19             : using namespace scipp::variable;
      20             : namespace scipp {
      21             : 
      22             : namespace {
      23             : scipp::index
      24             : size_of_elements(const Variable &view, SizeofTag tag,
      25             :                  const std::optional<std::pair<Dim, double>> &scale_in_dim);
      26             : 
      27             : scipp::index
      28        1615 : size_of_impl(const Variable &var, const SizeofTag tag,
      29             :              const std::optional<std::pair<Dim, double>> &scale_in_dim) {
      30             :   const auto object_size =
      31        1615 :       static_cast<scipp::index>(sizeof(Variable)) + var.data().object_size();
      32        1615 :   return size_of_elements(var, tag, scale_in_dim) + object_size;
      33             : }
      34             : 
      35             : /// Return the size in memory of a DataArray object. The aligned coord is
      36             : /// optional because for a DataArray owned by a dataset aligned coords are
      37             : /// assumed to be owned by the dataset as they can apply to multiple arrays.
      38             : scipp::index
      39         306 : size_of_impl(const DataArray &da, const SizeofTag tag,
      40             :              const std::optional<std::pair<Dim, double>> &scale_in_dim,
      41             :              bool include_aligned_coords = true) {
      42             :   auto size = static_cast<scipp::index>(
      43             :       sizeof(DataArray) + sizeof(dataset::Coords) + sizeof(dataset::Attrs) +
      44         306 :       sizeof(dataset::Masks) + da.coords().capacity() + da.attrs().capacity() +
      45         306 :       da.masks().capacity());
      46         306 :   size += size_of_impl(da.data(), tag, scale_in_dim);
      47         318 :   for (const auto &coord : da.attrs()) {
      48          12 :     size += size_of_impl(coord.second, tag, scale_in_dim);
      49             :   }
      50         620 :   for (const auto &mask : da.masks()) {
      51         314 :     size += size_of_impl(mask.second, tag, scale_in_dim);
      52             :   }
      53         306 :   if (include_aligned_coords) {
      54         644 :     for (const auto &coord : da.coords()) {
      55         426 :       size += size_of_impl(coord.second, tag, scale_in_dim);
      56             :     }
      57             :   }
      58         306 :   return size;
      59             : }
      60             : 
      61             : scipp::index
      62          48 : size_of_impl(const Dataset &ds, const SizeofTag tag,
      63             :              const std::optional<std::pair<Dim, double>> &scale_in_dim) {
      64          48 :   auto size = static_cast<scipp::index>(sizeof(Dataset) + ds.capacity());
      65         136 :   for (const auto &data : ds) {
      66          88 :     size += size_of_impl(data, tag, scale_in_dim, false);
      67          88 :   }
      68          56 :   for (const auto &coord : ds.coords()) {
      69           8 :     size += size_of_impl(coord.second, tag, scale_in_dim);
      70             :   }
      71          48 :   return size;
      72             : }
      73             : 
      74             : struct SizeOfKernel {
      75             :   using types = std::tuple<std::tuple<scipp::index, std::string>,
      76             :                            std::tuple<scipp::index, scipp::index>>;
      77             : 
      78     1000076 :   void operator()(scipp::index &out, const std::string &str) {
      79     1000076 :     if (const auto str_address =
      80     1000076 :             reinterpret_cast<const char *>(std::addressof(str));
      81     2000152 :         str.data() > str_address &&
      82     1000076 :         str.data() + str.size() < str_address + sizeof(std::string)) {
      83             :       // Small string optimization: The characters are located
      84             :       // in the string's internal buffer.
      85          67 :       out += sizeof(std::string);
      86             :     } else {
      87             :       // A long string: The characters are in a separate
      88             :       // array on the heap.
      89     1000009 :       out += sizeof(std::string) + str.size();
      90             :     }
      91     1000076 :   }
      92          25 :   void operator()(scipp::index &out, const scipp::index &s) { out += s; }
      93             : };
      94             : 
      95             : struct SizeOfContainerKernel {
      96             :   using types = std::tuple<std::tuple<scipp::index, Variable>,
      97             :                            std::tuple<scipp::index, DataArray>,
      98             :                            std::tuple<scipp::index, Dataset>>;
      99             : 
     100             :   SizeofTag tag;
     101             : 
     102          15 :   template <class T> void operator()(scipp::index &out, const T &x) const {
     103          15 :     out += size_of(x, tag);
     104          15 :   }
     105             : };
     106             : 
     107             : template <class T>
     108          98 : scipp::index size_of_bins(const Variable &view, const SizeofTag tag) {
     109          98 :   const auto &[indices, dim, buffer] = view.constituents<T>();
     110          98 :   std::pair<Dim, double> scale_in_dim{Dim::Invalid, 1.0};
     111          98 :   if (tag == SizeofTag::ViewOnly) {
     112          49 :     const auto &[begin, end] = unzip(indices);
     113          49 :     const auto sizes = sum(end - begin).template value<scipp::index>();
     114             :     // avoid division by zero
     115          49 :     scale_in_dim.second =
     116          49 :         sizes == 0 ? 0.0 : sizes / static_cast<double>(buffer.dims()[dim]);
     117          49 :     scale_in_dim.first = dim;
     118          49 :   }
     119          98 :   const auto indices_volume = tag == SizeofTag::Underlying
     120          98 :                                   ? indices.data().size()
     121          49 :                                   : indices.dims().volume();
     122          98 :   return indices_volume * sizeof(scipp::index_pair) +
     123             :          sizeof(
     124          98 :              variable::StructureArrayModel<scipp::index_pair, scipp::index>) +
     125         196 :          size_of_impl(buffer, tag, scale_in_dim);
     126          98 : }
     127             : 
     128             : template <class Op>
     129          27 : auto accumulate_size_of(const Variable &view, const SizeofTag tag,
     130             :                         const Op &op) {
     131          27 :   auto size = makeVariable<scipp::index>(Shape{}, Values{0});
     132          27 :   if (tag == SizeofTag::Underlying) {
     133          28 :     const Variable full(core::Dimensions{Dim::X, view.data().size()},
     134          14 :                         view.data_handle());
     135          14 :     accumulate_in_place(size, full, op, "size_of");
     136          14 :   } else {
     137          13 :     accumulate_in_place(size, view, op, "size_of");
     138             :   }
     139          54 :   return size.value<scipp::index>();
     140          27 : }
     141             : 
     142             : scipp::index
     143        1615 : size_of_elements(const Variable &view, const SizeofTag tag,
     144             :                  const std::optional<std::pair<Dim, double>> &scale_in_dim) {
     145        1615 :   if (view.dtype() == dtype<bucket<Variable>>) {
     146          26 :     return size_of_bins<Variable>(view, tag);
     147             :   }
     148        1589 :   if (view.dtype() == dtype<bucket<DataArray>>) {
     149          68 :     return size_of_bins<DataArray>(view, tag);
     150             :   }
     151        1521 :   if (view.dtype() == dtype<bucket<Dataset>>) {
     152           4 :     return size_of_bins<Dataset>(view, tag);
     153             :   }
     154        1517 :   if (view.dtype() == dtype<std::string>) {
     155          16 :     return accumulate_size_of(view, tag, SizeOfKernel{});
     156             :   }
     157        2991 :   if (view.dtype() == dtype<Variable> || view.dtype() == dtype<DataArray> ||
     158        2991 :       view.dtype() == dtype<Dataset>) {
     159          11 :     return accumulate_size_of(view, tag, SizeOfContainerKernel{tag});
     160             :   }
     161             : 
     162        1490 :   const auto value_size = view.data().dtype_size();
     163        1490 :   const auto variance_scale = view.has_variances() ? 2 : 1;
     164             :   const auto data_size =
     165        1490 :       tag == SizeofTag::Underlying ? view.data().size() : view.dims().volume();
     166             :   const auto extra_scale =
     167        1914 :       scale_in_dim.has_value() && view.dims().contains(scale_in_dim->first)
     168        1914 :           ? scale_in_dim->second
     169        1490 :           : 1.0;
     170             :   return static_cast<scipp::index>(
     171        1490 :       static_cast<double>(data_size * value_size * variance_scale) *
     172        1490 :       extra_scale);
     173             : }
     174             : } // namespace
     175             : 
     176         523 : scipp::index size_of(const Variable &var, const SizeofTag tag) {
     177         523 :   return size_of_impl(var, tag, std::nullopt);
     178             : }
     179         150 : scipp::index size_of(const DataArray &da, const SizeofTag tag,
     180             :                      const bool include_aligned_coords) {
     181         150 :   return size_of_impl(da, tag, std::nullopt, include_aligned_coords);
     182             : }
     183          44 : scipp::index size_of(const Dataset &ds, const SizeofTag tag) {
     184          44 :   return size_of_impl(ds, tag, std::nullopt);
     185             : }
     186             : } // namespace scipp
     187             : 
     188             : namespace scipp::dataset {
     189         130 : DataArray strip_edges_along(const DataArray &da, const Dim dim) {
     190         130 :   auto out = da;
     191         330 :   for (const auto &[name, var] : da.coords())
     192         200 :     if (core::is_edges(da.dims(), var.dims(), dim))
     193           3 :       out.coords().erase(name);
     194         166 :   for (const auto &[name, var] : da.masks())
     195          36 :     if (core::is_edges(da.dims(), var.dims(), dim))
     196           2 :       out.masks().erase(name);
     197         132 :   for (const auto &[name, var] : da.attrs())
     198           2 :     if (core::is_edges(da.dims(), var.dims(), dim))
     199           2 :       out.attrs().erase(name);
     200         130 :   return out;
     201           0 : }
     202             : 
     203          28 : Dataset strip_edges_along(const Dataset &ds, const Dim dim) {
     204             :   auto out = apply_to_items(
     205          64 :       ds, [](auto &&...args) { return strip_edges_along(args...); }, dim);
     206          84 :   for (const auto &[name, var] : ds.coords())
     207          56 :     if (!core::is_edges(ds.sizes(), var.dims(), dim))
     208          55 :       out.setCoord(name, var);
     209          28 :   return out;
     210           0 : }
     211             : 
     212             : } // namespace scipp::dataset