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 "scipp/dataset/dataset.h"
       8             : #include "scipp/dataset/string.h"
       9             : #include "scipp/variable/bins.h"
      10             : #include "scipp/variable/except.h"
      11             : 
      12             : namespace scipp::dataset {
      13             : 
      14             : namespace bins_view_detail {
      15             : template <class T> class BinsCommon {
      16             : public:
      17       25701 :   explicit BinsCommon(Variable var) : m_var(std::move(var)) {}
      18       41545 :   auto indices() const { return std::get<0>(get()); }
      19       63284 :   auto dim() const { return std::get<1>(get()); }
      20       89584 :   auto &buffer() const { return m_var.bin_buffer<T>(); }
      21        4092 :   auto &buffer() { return m_var.bin_buffer<T>(); }
      22             : 
      23             : protected:
      24       37648 :   auto make(const variable::Variable &view) const {
      25       37648 :     return make_bins_no_validate(this->indices(), this->dim(), view);
      26             :   }
      27        3898 :   auto check_and_get_buf(const Variable &var) const {
      28        3898 :     const auto &[i, d, buf] = var.constituents<Variable>();
      29        3903 :     core::expect::equals(i, this->indices());
      30        3895 :     core::expect::equals(d, this->dim());
      31        7790 :     return buf;
      32        3897 :   }
      33           0 :   [[nodiscard]] const Variable &get_var() const noexcept { return m_var; }
      34             : 
      35             : private:
      36      104829 :   auto get() const { return m_var.constituents<T>(); }
      37             :   Variable m_var;
      38             : };
      39             : 
      40             : template <class T, class MapGetter> class BinsMapView : public BinsCommon<T> {
      41             :   struct make_value {
      42             :     const BinsMapView *view;
      43       21741 :     template <class Value> auto operator()(const Value &value) const {
      44       21741 :       if (value.dims().contains(view->dim()))
      45       21723 :         return view->make(value);
      46             :       else
      47          18 :         return copy(value);
      48             :     }
      49             :   };
      50             :   struct make_item {
      51             :     const BinsMapView *view;
      52       21739 :     template <class Item> auto operator()(const Item &item) const {
      53       21739 :       return std::pair(item.first, make_value{view}(item.second));
      54             :     }
      55             :   };
      56             :   using MapView =
      57             :       std::decay_t<decltype(std::declval<MapGetter>()(std::declval<T>()))>;
      58             : 
      59             : public:
      60             :   using key_type = typename MapView::key_type;
      61             :   using mapped_type = typename MapView::mapped_type;
      62       42262 :   BinsMapView(const BinsCommon<T> base, MapGetter map)
      63       42262 :       : BinsCommon<T>(base), m_map(map) {}
      64           6 :   scipp::index size() const noexcept { return mapView().size(); }
      65       11105 :   auto operator[](const key_type &key) const {
      66       11105 :     return this->make(mapView()[key]);
      67             :   }
      68           7 :   void erase(const key_type &key) { return mapView().erase(key); }
      69          58 :   auto extract(const key_type &key) {
      70          58 :     return this->make(mapView().extract(key));
      71             :   }
      72        3896 :   void set(const key_type &key, const Variable &var) {
      73        3896 :     mapView().set(key, this->check_and_get_buf(var));
      74        3892 :   }
      75       25076 :   auto begin() const noexcept {
      76       25076 :     return mapView().begin().transform(make_item{this});
      77             :   }
      78       25076 :   auto end() const noexcept {
      79       25076 :     return mapView().end().transform(make_item{this});
      80             :   }
      81           8 :   auto keys_begin() const noexcept { return mapView().keys_begin(); }
      82           8 :   auto keys_end() const noexcept { return mapView().keys_end(); }
      83           2 :   auto values_begin() const noexcept {
      84           2 :     return mapView().values_begin().transform(make_value{this});
      85             :   }
      86           2 :   auto values_end() const noexcept {
      87           2 :     return mapView().values_end().transform(make_value{this});
      88             :   }
      89           2 :   auto items_begin() const noexcept { return begin(); }
      90           2 :   auto items_end() const noexcept { return end(); }
      91        5598 :   bool contains(const key_type &key) const noexcept {
      92        5598 :     return mapView().contains(key);
      93             :   }
      94        5834 :   scipp::index count(const key_type &key) const noexcept {
      95        5834 :     return mapView().count(key);
      96             :   }
      97         129 :   void set_aligned(const key_type &key, const bool aligned) {
      98         129 :     mapView().set_aligned(key, aligned);
      99         129 :   }
     100           0 :   bool is_edges(const key_type &, std::optional<Dim>) const noexcept {
     101           0 :     return false; // event-coordinates are never edges
     102             :   }
     103           0 :   bool operator==(const BinsMapView &other) const noexcept {
     104           0 :     return mapView() == other.mapView();
     105             :   }
     106           0 :   bool operator!=(const BinsMapView &other) const noexcept {
     107           0 :     return mapView() != other.mapView();
     108             :   }
     109             : 
     110           0 :   friend std::string dict_keys_to_string(const BinsMapView &view) {
     111           0 :     return dict_keys_to_string(view.mapView());
     112             :   }
     113           0 :   friend std::string to_string(const BinsMapView &view) {
     114           0 :     return to_string(view.mapView());
     115             :   }
     116           0 :   friend BinsMapView copy(const BinsMapView &view) {
     117           0 :     return BinsMapView{BinsCommon<T>{copy(view.get_var())}, view.m_map};
     118             :   }
     119             : 
     120             : private:
     121       72715 :   decltype(auto) mapView() const { return m_map(this->buffer()); }
     122        4090 :   decltype(auto) mapView() { return m_map(this->buffer()); }
     123             :   MapGetter m_map;
     124             : };
     125             : 
     126             : template <class T> class Bins : public BinsCommon<T> {
     127             : public:
     128             :   using BinsCommon<T>::BinsCommon;
     129        8511 :   auto data() const { return this->make(this->buffer().data()); }
     130             :   // TODO how to handle const-ness?
     131           2 :   void setData(const Variable &var) {
     132           2 :     this->buffer().setData(this->check_and_get_buf(var));
     133           2 :   }
     134        4969 :   auto meta() const { return BinsMapView(*this, get_meta); }
     135       20533 :   auto coords() const { return BinsMapView(*this, get_coords); }
     136        8383 :   auto attrs() const { return BinsMapView(*this, get_attrs); }
     137        8377 :   auto masks() const { return BinsMapView(*this, get_masks); }
     138        8358 :   auto &name() const { return this->buffer().name(); }
     139           0 :   auto drop_coords(const scipp::span<const Dim> coord_names) const {
     140           0 :     auto result = *this;
     141           0 :     for (const auto &name : coord_names)
     142           0 :       result.coords().erase(name);
     143           0 :   }
     144           0 :   auto drop_masks(const scipp::span<const std::string> mask_names) const {
     145           0 :     auto result = *this;
     146           0 :     for (const auto &name : mask_names)
     147           0 :       result.masks().erase(name);
     148           0 :   }
     149           0 :   auto drop_attrs(const scipp::span<const Dim> attr_names) const {
     150           0 :     auto result = *this;
     151           0 :     for (const auto &name : attr_names)
     152           0 :       result.attrs().erase(name);
     153           0 :   }
     154             : };
     155             : } // namespace bins_view_detail
     156             : 
     157             : /// Return helper for accessing bin data and coords as non-owning views
     158             : ///
     159             : /// Usage:
     160             : /// auto data = bins_view<DataArray>(var).data();
     161             : /// auto coord = bins_view<DataArray>(var).coords()[dim];
     162             : ///
     163             : /// The returned objects are variables referencing data in `var`. They do not
     164             : /// own or share ownership of any data.
     165       25701 : template <class T> auto bins_view(Variable var) {
     166       25701 :   return bins_view_detail::Bins<T>(std::move(var));
     167             : }
     168             : 
     169             : } // namespace scipp::dataset