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             : #include "scipp/variable/subspan_view.h"
       6             : #include "scipp/core/eigen.h"
       7             : #include "scipp/core/except.h"
       8             : #include "scipp/variable/cumulative.h"
       9             : #include "scipp/variable/shape.h"
      10             : #include "scipp/variable/transform.h"
      11             : #include "scipp/variable/util.h"
      12             : 
      13             : namespace scipp::variable {
      14             : 
      15             : namespace {
      16             : 
      17             : template <class T>
      18      320241 : auto make_subspans(T *base, const Variable &indices,
      19             :                    const scipp::index stride) {
      20      320241 :   if (stride != 1)
      21           0 :     throw std::logic_error(
      22             :         "span only supports stride=1, this should be "
      23             :         "unreachable due to an earlier check, may want to generalize this "
      24             :         "later to support in particular stride=0 for broadcasted buffers");
      25             :   return variable::transform<scipp::index_pair>(
      26             :       indices,
      27             :       overloaded{core::transform_flags::expect_no_variance_arg<0>,
      28      320241 :                  [](const units::Unit &) { return units::one; },
      29     1229166 :                  [base, stride](const auto &offset) {
      30     2458332 :                    return scipp::span(base + stride * offset.first,
      31     1229166 :                                       base + stride * offset.second);
      32             :                  }},
      33      320241 :       "make_subspans");
      34             : }
      35             : 
      36             : /// Return Variable containing spans with extents given by indices over given
      37             : /// dimension as elements.
      38             : template <class T, class Var>
      39      316763 : Variable subspan_view(Var &var, const Dim dim, const Variable &indices) {
      40      316763 :   auto subspans =
      41             :       make_subspans(var.template values<T>().data(), indices, var.stride(dim));
      42      316762 :   if (var.has_variances())
      43        3479 :     subspans.setVariances(make_subspans(var.template variances<T>().data(),
      44             :                                         indices, var.stride(dim)));
      45      316762 :   subspans.setUnit(var.unit());
      46      316762 :   return subspans;
      47           0 : }
      48             : 
      49             : template <class... Ts, class... Args>
      50      316763 : auto invoke_subspan_view(const DType dtype, Args &&...args) {
      51      316763 :   Variable ret;
      52      633525 :   if (!((scipp::dtype<Ts> == dtype
      53      633533 :              ? (ret = subspan_view<Ts>(std::forward<Args>(args)...), true)
      54             :              : false) ||
      55             :         ...))
      56           0 :     throw except::TypeError("Unsupported dtype.");
      57      316762 :   return ret;
      58           1 : }
      59             : 
      60             : template <class Var, class... Args>
      61      316764 : Variable subspan_view_impl(Var &var, const Dim dim, Args &&...args) {
      62      316764 :   if (var.stride(dim) != 1)
      63           1 :     throw except::DimensionError(
      64             :         "View over subspan can only be created for contiguous "
      65             :         "range of data.");
      66             :   return invoke_subspan_view<double, float, int64_t, int32_t, bool,
      67             :                              core::time_point, std::string, Eigen::Vector3d>(
      68      316763 :       var.dtype(), var, dim, args...);
      69             : }
      70             : 
      71         399 : auto make_range(const scipp::index num, const scipp::index stride,
      72             :                 const Dim dim) {
      73         798 :   return cumsum(broadcast(stride * units::one, {dim, num}), dim,
      74         798 :                 CumSumMode::Exclusive);
      75             : }
      76             : 
      77      215732 : Variable make_indices(const Variable &var, const Dim dim) {
      78      215732 :   auto dims = var.dims();
      79      215732 :   dims.erase(dim);
      80      215732 :   auto start = scipp::index(0) * units::one;
      81      216131 :   for (const auto &label : dims) {
      82         399 :     const auto stride = var.stride(label);
      83         399 :     start = start + make_range(dims[label], stride, label);
      84             :   }
      85      647196 :   return zip(start, start + var.dims()[dim] * units::one);
      86      215732 : }
      87             : 
      88             : } // namespace
      89             : 
      90             : /// Return Variable containing mutable spans over given dimension as elements.
      91       47665 : Variable subspan_view(Variable &var, const Dim dim) {
      92       95328 :   return subspan_view(var, dim, make_indices(var, dim));
      93             : }
      94             : /// Return Variable containing const spans over given dimension as elements.
      95      168067 : Variable subspan_view(const Variable &var, const Dim dim) {
      96      336134 :   return subspan_view(var, dim, make_indices(var, dim));
      97             : }
      98             : 
      99       78191 : Variable subspan_view(Variable &var, const Dim dim, const Variable &indices) {
     100       78191 :   return subspan_view_impl(var, dim, indices);
     101             : }
     102      238573 : Variable subspan_view(const Variable &var, const Dim dim,
     103             :                       const Variable &indices) {
     104      238573 :   return subspan_view_impl(var, dim, indices);
     105             : }
     106             : 
     107             : } // namespace scipp::variable