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/core/dict.h"
       8             : #include "scipp/core/sizes.h"
       9             : #include "scipp/core/slice.h"
      10             : #include "scipp/dataset/sized_dict_forward.h"
      11             : #include "scipp/units/dim.h"
      12             : #include "scipp/units/unit.h"
      13             : #include "scipp/variable/logical.h"
      14             : #include "scipp/variable/variable.h"
      15             : 
      16             : namespace scipp::dataset {
      17             : template <class Mapping>
      18       29537 : Mapping slice_map(const Sizes &sizes, const Mapping &map, const Slice &params) {
      19       29537 :   Mapping out;
      20       60479 :   for (const auto &[key, value] : map) {
      21       30946 :     if (value.dims().contains(params.dim())) {
      22       19964 :       if (value.dims()[params.dim()] == sizes[params.dim()]) {
      23       18826 :         out.insert_or_assign(key, value.slice(params));
      24             :       } else { // bin edge
      25        1138 :         if (params.stride() != 1)
      26           1 :           throw except::SliceError(
      27             :               "Object has bin-edges along dimension " +
      28             :               to_string(params.dim()) + " so slicing with stride " +
      29             :               std::to_string(params.stride()) + " != 1 is not valid.");
      30        1796 :         const auto end = params.end() == -1               ? params.begin() + 2
      31        1075 :                          : params.begin() == params.end() ? params.end()
      32         416 :                                                           : params.end() + 1;
      33        1137 :         out.insert_or_assign(
      34        1137 :             key, value.slice(Slice{params.dim(), params.begin(), end}));
      35             :       }
      36       10982 :     } else if (params == Slice{}) {
      37           3 :       out.insert_or_assign(key, value);
      38             :     } else {
      39       10979 :       out.insert_or_assign(key, value.as_const());
      40             :     }
      41             :   }
      42       29533 :   return out;
      43           4 : }
      44             : 
      45             : // Used to disambiguate overloaded constructors that accept either a SizedDict
      46             : // or a SizedDict::holder_type.
      47             : // E.g. DataArray
      48             : struct AutoSizeTag {};
      49             : 
      50             : /// Dict with fixed dimensions.
      51             : ///
      52             : /// Values must have dimensions and those dimensions must be a subset
      53             : /// of the sizes stored in SizedDict. This is used, e.g., to ensure
      54             : /// that coords are valid for a data array.
      55             : template <class Key, class Value> class SizedDict {
      56             : public:
      57             :   using key_type = Key;
      58             :   using mapped_type = Value;
      59             :   using holder_type = core::Dict<key_type, mapped_type>;
      60             : 
      61       47048 :   SizedDict() = default;
      62             :   SizedDict(const Sizes &sizes,
      63             :             std::initializer_list<std::pair<const Key, Value>> items,
      64             :             bool readonly = false);
      65             :   SizedDict(AutoSizeTag,
      66             :             std::initializer_list<std::pair<const Key, Value>> items,
      67             :             bool readonly = false);
      68             :   SizedDict(Sizes sizes, holder_type items, bool readonly = false);
      69             :   SizedDict(AutoSizeTag, holder_type items, bool readonly = false);
      70             :   SizedDict(const SizedDict &other);
      71             :   SizedDict(SizedDict &&other) noexcept;
      72             :   SizedDict &operator=(const SizedDict &other);
      73             :   SizedDict &operator=(SizedDict &&other) noexcept;
      74             : 
      75             :   /// Return the number of coordinates in the view.
      76       42927 :   [[nodiscard]] index size() const noexcept { return scipp::size(m_items); }
      77             :   /// Return true if there are 0 coordinates in the view.
      78         782 :   [[nodiscard]] bool empty() const noexcept { return m_items.empty(); }
      79             :   /// Return the number of elements that space is currently allocated for.
      80         918 :   [[nodiscard]] index capacity() const noexcept { return m_items.capacity(); }
      81           0 :   void reserve(const index new_capacity) { m_items.reserve(new_capacity); }
      82             : 
      83             :   bool contains(const Key &k) const;
      84             :   scipp::index count(const Key &k) const;
      85             : 
      86             :   const mapped_type &operator[](const Key &key) const;
      87             :   const mapped_type &at(const Key &key) const;
      88             :   // Note that the non-const versions return by value, to avoid breakage of
      89             :   // invariants.
      90             :   mapped_type operator[](const Key &key);
      91             :   mapped_type at(const Key &key);
      92             :   Dim dim_of(const Key &key) const;
      93             : 
      94        3115 :   auto find(const Key &k) const noexcept { return m_items.find(k); }
      95        3052 :   auto find(const Key &k) noexcept { return m_items.find(k); }
      96             : 
      97             :   /// Return const iterator to the beginning of all items.
      98      571771 :   auto begin() const noexcept { return m_items.begin(); }
      99        3356 :   auto begin() noexcept { return m_items.begin(); }
     100             :   /// Return const iterator to the end of all items.
     101      574886 :   auto end() const noexcept { return m_items.end(); }
     102        6889 :   auto end() noexcept { return m_items.end(); }
     103             : 
     104             :   auto items_begin() const && = delete;
     105             :   /// Return const iterator to the beginning of all items.
     106        9232 :   auto items_begin() const &noexcept { return begin(); }
     107             :   auto items_end() const && = delete;
     108             :   /// Return const iterator to the end of all items.
     109        9232 :   auto items_end() const &noexcept { return end(); }
     110             : 
     111             :   auto keys_begin() const && = delete;
     112             :   /// Return const iterator to the beginning of all keys.
     113       35622 :   auto keys_begin() const &noexcept { return m_items.keys_begin(); }
     114             :   auto keys_end() const && = delete;
     115             :   /// Return const iterator to the end of all keys.
     116       35624 :   auto keys_end() const &noexcept { return m_items.keys_end(); }
     117             : 
     118             :   auto values_begin() const && = delete;
     119             :   /// Return const iterator to the beginning of all values.
     120        3254 :   auto values_begin() const &noexcept { return m_items.values_begin(); }
     121             :   auto values_end() const && = delete;
     122             :   /// Return const iterator to the end of all values.
     123        3278 :   auto values_end() const &noexcept { return m_items.values_end(); }
     124             : 
     125             :   bool operator==(const SizedDict &other) const;
     126             :   bool operator!=(const SizedDict &other) const;
     127             : 
     128      972952 :   [[nodiscard]] const Sizes &sizes() const noexcept { return m_sizes; }
     129         146 :   [[nodiscard]] const auto &items() const noexcept { return m_items; }
     130             : 
     131             :   void setSizes(const Sizes &sizes);
     132             :   void set(const key_type &key, mapped_type coord);
     133             :   void erase(const key_type &key);
     134             :   mapped_type extract(const key_type &key);
     135             :   mapped_type extract(const key_type &key, const mapped_type &default_value);
     136             : 
     137             :   SizedDict slice(const Slice &params) const;
     138             :   SizedDict slice_coords(const Slice &params) const;
     139             :   void validateSlice(const Slice &s, const SizedDict &dict) const;
     140             :   [[maybe_unused]] SizedDict &setSlice(const Slice &s, const SizedDict &dict);
     141             : 
     142             :   [[nodiscard]] SizedDict
     143             :   rename_dims(const std::vector<std::pair<Dim, Dim>> &names,
     144             :               const bool fail_on_unknown = true) const;
     145             : 
     146             :   void set_readonly() noexcept;
     147             :   [[nodiscard]] bool is_readonly() const noexcept;
     148             :   [[nodiscard]] SizedDict as_const() const;
     149             :   [[nodiscard]] SizedDict merge_from(const SizedDict &other) const;
     150             : 
     151             :   bool item_applies_to(const Key &key, const Dimensions &dims) const;
     152             :   bool is_edges(const Key &key, std::optional<Dim> dim = std::nullopt) const;
     153             : 
     154             :   void set_aligned(const Key &key, bool aligned);
     155             : 
     156             : protected:
     157             :   Sizes m_sizes;
     158             :   holder_type m_items;
     159             :   bool m_readonly{false};
     160             : };
     161             : 
     162             : /// Returns the union of all masks with irreducible dimension `dim`.
     163             : ///
     164             : /// Irreducible means that a reduction operation must apply these masks since
     165             : /// they depend on the reduction dimension. Returns an invalid (empty) variable
     166             : /// if there is no irreducible mask.
     167             : template <class Masks>
     168       64764 : [[nodiscard]] Variable irreducible_mask(const Masks &masks, const Dim dim) {
     169       64764 :   Variable union_;
     170       71511 :   for (const auto &mask : masks)
     171        6747 :     if (mask.second.dims().contains(dim))
     172        5090 :       union_ = union_.is_valid() ? union_ | mask.second : copy(mask.second);
     173       64764 :   return union_;
     174           0 : }
     175             : 
     176             : template <class Key, class Value>
     177             : bool equals_nan(const SizedDict<Key, Value> &a, const SizedDict<Key, Value> &b);
     178             : 
     179             : template <class Key, class Value>
     180             : core::Dict<Key, Value> union_(const SizedDict<Key, Value> &a,
     181             :                               const SizedDict<Key, Value> &b,
     182             :                               std::string_view opname);
     183             : 
     184             : /// Return intersection of dicts, i.e., all items with matching names that
     185             : /// have matching content.
     186             : template <class Key, class Value>
     187             : core::Dict<Key, Value> intersection(const SizedDict<Key, Value> &a,
     188             :                                     const SizedDict<Key, Value> &b);
     189             : 
     190        2353 : constexpr auto get_data = [](auto &&a) -> decltype(auto) { return a.data(); };
     191          14 : constexpr auto get_sizes = [](auto &&a) -> decltype(auto) { return a.sizes(); };
     192       14327 : constexpr auto get_meta = [](auto &&a) -> decltype(auto) { return a.meta(); };
     193       31356 : constexpr auto get_coords = [](auto &&a) -> decltype(auto) {
     194       31356 :   return a.coords();
     195             : };
     196       19071 : constexpr auto get_attrs = [](auto &&a) -> decltype(auto) { return a.attrs(); };
     197       19096 : constexpr auto get_masks = [](auto &&a) -> decltype(auto) { return a.masks(); };
     198             : 
     199             : } // namespace scipp::dataset