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 Simon Heybrock
       5             : #pragma once
       6             : 
       7             : #include <boost/container/small_vector.hpp>
       8             : #include <vector>
       9             : 
      10             : #include "scipp/variable/creation.h"
      11             : #include <scipp/dataset/dataset.h>
      12             : 
      13             : namespace scipp::dataset {
      14             : 
      15             : /// Implementation detail of GroupBy.
      16             : ///
      17             : /// Stores the actual grouping details, independent of the container type.
      18             : class SCIPP_DATASET_EXPORT GroupByGrouping {
      19             : public:
      20             :   using group = boost::container::small_vector<Slice, 4>;
      21          69 :   GroupByGrouping(const Dim sliceDim, Variable key, std::vector<group> groups)
      22          69 :       : m_sliceDim(sliceDim), m_key(std::move(key)),
      23          69 :         m_groups(std::move(groups)) {}
      24             : 
      25         128 :   scipp::index size() const noexcept { return scipp::size(m_groups); }
      26             :   Dim sliceDim() const noexcept { return m_sliceDim; }
      27         208 :   Dim dim() const noexcept { return m_key.dims().inner(); }
      28          76 :   const Variable &key() const noexcept { return m_key; }
      29         115 :   const std::vector<group> &groups() const noexcept { return m_groups; }
      30             : 
      31             : private:
      32             :   Dim m_sliceDim;
      33             :   Variable m_key;
      34             :   std::vector<group> m_groups;
      35             : };
      36             : 
      37             : /// Helper class for implementing "split-apply-combine" functionality.
      38             : template <class T> class SCIPP_DATASET_EXPORT GroupBy {
      39             : public:
      40          69 :   GroupBy(const T &data, GroupByGrouping &&grouping)
      41          69 :       : m_data(data), m_grouping(std::move(grouping)) {}
      42             : 
      43         128 :   scipp::index size() const noexcept { return m_grouping.size(); }
      44         208 :   Dim dim() const noexcept { return m_grouping.dim(); }
      45          76 :   const Variable &key() const noexcept { return m_grouping.key(); }
      46         115 :   const std::vector<GroupByGrouping::group> &groups() const noexcept {
      47         115 :     return m_grouping.groups();
      48             :   }
      49             : 
      50             :   T concat(const Dim reductionDim) const;
      51             :   T mean(const Dim reductionDim) const;
      52             :   T sum(const Dim reductionDim) const;
      53             :   T nansum(const Dim reductionDim) const;
      54             :   T all(const Dim reductionDim) const;
      55             :   T any(const Dim reductionDim) const;
      56             :   T max(const Dim reductionDim) const;
      57             :   T nanmax(const Dim reductionDim) const;
      58             :   T min(const Dim reductionDim) const;
      59             :   T nanmin(const Dim reductionDim) const;
      60             : 
      61             : private:
      62             :   T makeReductionOutput(const Dim reductionDim, const FillValue fill) const;
      63             :   template <class Op>
      64             :   T reduce(Op op, const Dim reductionDim, const FillValue fill) const;
      65             : 
      66             :   T m_data;
      67             :   GroupByGrouping m_grouping;
      68             : };
      69             : 
      70             : SCIPP_DATASET_EXPORT GroupBy<DataArray> groupby(const DataArray &dataset,
      71             :                                                 const Dim dim);
      72             : SCIPP_DATASET_EXPORT GroupBy<DataArray>
      73             : groupby(const DataArray &dataset, const Dim dim, const Variable &bins);
      74             : 
      75             : SCIPP_DATASET_EXPORT GroupBy<Dataset> groupby(const Dataset &dataset,
      76             :                                               const Dim dim);
      77             : SCIPP_DATASET_EXPORT GroupBy<Dataset>
      78             : groupby(const Dataset &dataset, const Dim dim, const Variable &bins);
      79             : 
      80             : SCIPP_DATASET_EXPORT GroupBy<DataArray> groupby(const DataArray &dataset,
      81             :                                                 const Variable &variable,
      82             :                                                 const Variable &bins);
      83             : 
      84             : SCIPP_DATASET_EXPORT GroupBy<Dataset>
      85             : groupby(const Dataset &dataset, const Variable &variable, const Variable &bins);
      86             : 
      87             : } // namespace scipp::dataset