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/variable/arithmetic.h"
       8             : #include "scipp/variable/subspan_view.h"
       9             : #include "scipp/variable/transform.h"
      10             : #include "scipp/variable/variable_factory.h"
      11             : 
      12             : namespace scipp::variable {
      13             : 
      14             : namespace transform_subspan_detail {
      15       47723 : static constexpr auto erase = [](Dimensions dims, const Dim dim) {
      16       47723 :   dims.erase(dim);
      17       47723 :   return dims;
      18             : };
      19             : 
      20      142827 : static constexpr auto maybe_subspan = [](const Variable &var, const Dim dim) {
      21      142827 :   if (var.dims().contains(dim))
      22       47723 :     return subspan_view(var, dim);
      23       95104 :   return var;
      24             : };
      25             : } // namespace transform_subspan_detail
      26             : 
      27             : template <class... Types, class Op, class... Var>
      28             : [[nodiscard]] Variable
      29       47609 : transform_subspan_impl(const DType type, const Dim dim, const scipp::index size,
      30             :                        Op op, const std::string_view &name, Var... var) {
      31             :   using namespace transform_subspan_detail;
      32             : 
      33       47609 :   auto dims =
      34       95332 :       merge(var.dims().contains(dim) ? erase(var.dims(), dim) : var.dims()...);
      35       47609 :   dims.addInner(dim, size);
      36       47609 :   const bool variance =
      37             :       std::is_base_of_v<core::transform_flags::expect_variance_arg_t<0>, Op> ||
      38             :       (std::is_base_of_v<
      39       47609 :            core::transform_flags::expect_in_variance_if_out_variance_t, Op> &&
      40       47609 :        (var.has_variances() || ...));
      41       47609 :   Variable out =
      42       47609 :       variableFactory().create(type, dims, op(var.unit()...), variance);
      43             : 
      44       47613 :   in_place<false>::transform_data(type_tuples<Types...>(op), op, name,
      45             :                                   subspan_view(out, dim),
      46             :                                   maybe_subspan(var, dim)...);
      47       95216 :   return out;
      48       47610 : }
      49             : 
      50             : /// Non-element-wise transform.
      51             : ///
      52             : /// This is a specialized version of transform, handling the case of inputs (and
      53             : /// output) that differ along one of their dimensions. Applications are mixing
      54             : /// of events and dense data, as well as operations that change the length of a
      55             : /// dimension (such as rebin). The syntax for the user-provided operator is
      56             : /// special and differs from that of `transform` and `transform_in_place`, and
      57             : /// also the tuple of supported type combinations is special:
      58             : /// 1. The overload for the transform of the unit is as for `transform`, i.e.,
      59             : ///    returns the new unit.
      60             : /// 2. The overload handling the data has one extra argument. This additional
      61             : ///    (first) argument is the "out" argument, as used in `transform_in_place`.
      62             : /// 3. The tuple of supported type combinations must include the type of the out
      63             : ///    argument as the first type in the inner tuples. The output type must be
      64             : ///    passed at runtime as the first argument. `transform_subspan` DOES NOT
      65             : ///    INITIALIZE the output array, i.e., `Op` must take care of initializing
      66             : ///    the respective subspans. This is done for improved performance, avoiding
      67             : ///    streaming/writing to memory twice. Optionally, initialization can be
      68             : ///    requested by passing an operator inheriting
      69             : ///    core::transform_flags::zero_output_t.
      70             : /// 4. The output type and the type of non-events inputs that depend on `dim`
      71             : ///    must be specified as `span<T>`. The user-provided lambda is called with a
      72             : ///    span of values for these arguments.
      73             : /// 5. Use the flag transform_flags::expect_variance_arg<0> to control whether
      74             : ///    the output should have variances or not.
      75             : template <class... Types, class Op>
      76             : [[nodiscard]] Variable
      77             : transform_subspan(const DType type, const Dim dim, const scipp::index size,
      78             :                   const Variable &var1, const Variable &var2, Op op,
      79             :                   const std::string_view &name = "operation") {
      80             :   return transform_subspan_impl<Types...>(type, dim, size, op, name, var1,
      81             :                                           var2);
      82             : }
      83             : 
      84             : template <class... Types, class Op>
      85             : [[nodiscard]] Variable
      86       47609 : transform_subspan(const DType type, const Dim dim, const scipp::index size,
      87             :                   const Variable &var1, const Variable &var2,
      88             :                   const Variable &var3, Op op,
      89             :                   const std::string_view &name = "operation") {
      90             :   return transform_subspan_impl<Types...>(type, dim, size, op, name, var1, var2,
      91       47611 :                                           var3);
      92             : }
      93             : 
      94             : } // namespace scipp::variable