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 <functional>
       8             : 
       9             : #include "scipp/core/flags.h"
      10             : 
      11             : #include "scipp/variable/variable.h"
      12             : 
      13             : namespace scipp::variable {
      14             : 
      15             : /// Abstract base class for "variable makers", used by VariableFactory to
      16             : /// dynamically create variables with given type.
      17             : class SCIPP_VARIABLE_EXPORT AbstractVariableMaker {
      18             : protected:
      19           0 :   auto unreachable() const {
      20           0 :     return std::logic_error("This code path should never be reached");
      21             :   }
      22             : 
      23             : public:
      24             :   using parent_list = std::vector<std::reference_wrapper<const Variable>>;
      25             : 
      26         245 :   virtual ~AbstractVariableMaker() = default;
      27             :   virtual bool is_bins() const = 0;
      28             :   virtual Variable create(const DType elem_dtype, const Dimensions &dims,
      29             :                           const units::Unit &unit, const bool variances,
      30             :                           const parent_list &parents) const = 0;
      31             :   virtual Dim elem_dim(const Variable &var) const = 0;
      32             :   virtual DType elem_dtype(const Variable &var) const = 0;
      33             :   virtual units::Unit elem_unit(const Variable &var) const = 0;
      34             :   virtual void expect_can_set_elem_unit(const Variable &var,
      35             :                                         const units::Unit &u) const = 0;
      36             :   virtual void set_elem_unit(Variable &var, const units::Unit &u) const = 0;
      37           0 :   virtual bool has_masks(const Variable &) const { return false; }
      38             :   virtual bool has_variances(const Variable &var) const = 0;
      39           0 :   virtual const Variable &data(const Variable &) const { throw unreachable(); }
      40           0 :   virtual Variable data(Variable &) const { throw unreachable(); }
      41           0 :   virtual core::ElementArrayViewParams array_params(const Variable &) const {
      42           0 :     throw unreachable();
      43             :   }
      44             :   virtual Variable empty_like(const Variable &prototype,
      45             :                               const std::optional<Dimensions> &shape,
      46             :                               const Variable &sizes) const = 0;
      47       33450 :   [[nodiscard]] virtual Variable apply_event_masks(const Variable &var,
      48             :                                                    const FillValue) const {
      49       33450 :     return var;
      50             :   }
      51             : 
      52             :   [[nodiscard]] virtual Variable
      53           6 :   irreducible_event_mask([[maybe_unused]] const Variable &var) const {
      54           6 :     return Variable{};
      55             :   }
      56             : };
      57             : 
      58             : SCIPP_VARIABLE_EXPORT bool is_bins(const Variable &var);
      59             : 
      60             : /// Dynamic factory for variables.
      61             : ///
      62             : /// The factory can be used for creating variables with a dtype that is not
      63             : /// known in the current module, e.g., dtype<bucket<Dataset>> can be used in
      64             : /// scipp::variable. The main purpose of this is the implementation of
      65             : /// `transform`.
      66             : class SCIPP_VARIABLE_EXPORT VariableFactory {
      67             : private:
      68             :   using parent_list = typename AbstractVariableMaker::parent_list;
      69             :   DType bin_dtype(const parent_list &vars) const noexcept;
      70             : 
      71             : public:
      72           5 :   VariableFactory() = default;
      73             :   VariableFactory(const VariableFactory &) = delete;
      74             :   VariableFactory &operator=(const VariableFactory &) = delete;
      75             :   void emplace(const DType key, std::unique_ptr<AbstractVariableMaker> makes);
      76             :   bool contains(const DType key) const noexcept;
      77             :   bool is_bins(const Variable &var) const;
      78             :   template <class... Parents>
      79     1334963 :   Variable create(const DType elem_dtype, const Dimensions &dims,
      80             :                   const units::Unit &unit, const bool with_variances,
      81             :                   const Parents &...parents) const {
      82     1334963 :     const auto parents_ = parent_list{parents...};
      83     1334963 :     const auto key = bin_dtype(parents_);
      84     1334963 :     return m_makers.at(key == dtype<void> ? elem_dtype : key)
      85     2669925 :         ->create(elem_dtype, dims, unit, with_variances, parents_);
      86     1334963 :   }
      87             :   Dim elem_dim(const Variable &var) const;
      88             :   DType elem_dtype(const Variable &var) const;
      89             :   units::Unit elem_unit(const Variable &var) const;
      90             :   void expect_can_set_elem_unit(const Variable &var,
      91             :                                 const units::Unit &u) const;
      92             :   void set_elem_unit(Variable &var, const units::Unit &u) const;
      93             :   bool has_masks(const Variable &var) const;
      94             :   bool has_variances(const Variable &var) const;
      95     5598518 :   template <class T, class Var> auto values(Var &&var) const {
      96     5598518 :     if (!is_bins(var))
      97     5488438 :       return var.template values<T>();
      98      110080 :     const auto &maker = *m_makers.at(var.dtype());
      99      110080 :     auto &&data = maker.data(var);
     100      110080 :     return ElementArrayView(maker.array_params(var),
     101      110080 :                             data.template values<T>().data());
     102       41110 :   }
     103       26880 :   template <class T, class Var> auto variances(Var &&var) const {
     104       26880 :     if (!is_bins(var))
     105       26422 :       return var.template variances<T>();
     106         458 :     const auto &maker = *m_makers.at(var.dtype());
     107         458 :     auto &&data = maker.data(var);
     108         458 :     return ElementArrayView(maker.array_params(var),
     109         458 :                             data.template variances<T>().data());
     110          88 :   }
     111             :   Variable empty_like(const Variable &prototype,
     112             :                       const std::optional<Dimensions> &shape,
     113             :                       const Variable &sizes = {});
     114             :   /// Return a binned variable where masked elements are replaced by fill.
     115             :   /// Coords and attrs of the input are not propagated to the output.
     116             :   [[nodiscard]] Variable apply_event_masks(const Variable &var,
     117             :                                            const FillValue fill) const;
     118             :   [[nodiscard]] Variable irreducible_event_mask(const Variable &var) const;
     119             : 
     120             : private:
     121             :   std::map<DType, std::unique_ptr<AbstractVariableMaker>> m_makers;
     122             : };
     123             : 
     124             : /// Return the global variable factory instance
     125             : SCIPP_VARIABLE_EXPORT VariableFactory &variableFactory();
     126             : 
     127             : } // namespace scipp::variable