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 <algorithm>
       6             : 
       7             : #include "scipp/core/dimensions.h"
       8             : 
       9             : #include "scipp/variable/arithmetic.h"
      10             : #include "scipp/variable/bins.h"
      11             : #include "scipp/variable/creation.h"
      12             : #include "scipp/variable/except.h"
      13             : #include "scipp/variable/shape.h"
      14             : #include "scipp/variable/util.h"
      15             : #include "scipp/variable/variable_concept.h"
      16             : #include "scipp/variable/variable_factory.h"
      17             : 
      18             : using namespace scipp::core;
      19             : 
      20             : namespace scipp::variable {
      21             : 
      22       37531 : Variable broadcast(const Variable &var, const Dimensions &dims) {
      23       37531 :   return var.broadcast(dims);
      24             : }
      25             : 
      26             : namespace {
      27          32 : auto get_bin_sizes(const scipp::span<const Variable> vars) {
      28          32 :   std::vector<Variable> sizes;
      29          32 :   sizes.reserve(vars.size());
      30         104 :   for (const auto &var : vars)
      31          72 :     sizes.emplace_back(bin_sizes(var));
      32          32 :   return sizes;
      33           0 : }
      34             : } // namespace
      35             : 
      36        2600 : Variable concat(const scipp::span<const Variable> vars, const Dim dim) {
      37        2600 :   if (vars.empty())
      38           1 :     throw std::invalid_argument("Cannot concat empty list.");
      39             :   const auto it =
      40        2599 :       std::find_if(vars.begin(), vars.end(),
      41        2818 :                    [dim](const auto &var) { return var.dims().contains(dim); });
      42        2599 :   Dimensions dims;
      43             :   // Expand dims for inputs that do not contain dim already. Favor order given
      44             :   // by first input, if not found add as outer dim.
      45        2599 :   if (it == vars.end()) {
      46         124 :     dims = vars.front().dims();
      47         124 :     dims.add(dim, 1);
      48             :   } else {
      49        2475 :     dims = it->dims();
      50        2475 :     dims.resize(dim, 1);
      51             :   }
      52        2599 :   std::vector<Variable> tmp;
      53        2599 :   scipp::index size = 0;
      54        7910 :   for (const auto &var : vars) {
      55        5312 :     if (var.dims().contains(dim))
      56        4949 :       tmp.emplace_back(var);
      57             :     else
      58         363 :       tmp.emplace_back(broadcast(var, dims));
      59        5311 :     size += tmp.back().dims()[dim];
      60             :   }
      61        2598 :   dims.resize(dim, size);
      62        2598 :   Variable out;
      63        2598 :   if (is_bins(vars.front())) {
      64          32 :     out = empty_like(vars.front(), {}, concat(get_bin_sizes(vars), dim));
      65             :   } else {
      66        2566 :     out = empty_like(vars.front(), dims);
      67             :   }
      68        2598 :   scipp::index offset = 0;
      69        7884 :   for (const auto &var : tmp) {
      70        5310 :     const auto extent = var.dims()[dim];
      71        5334 :     out.data().copy(var, out.slice({dim, offset, offset + extent}));
      72        5286 :     offset += extent;
      73             :   }
      74        5148 :   return out;
      75        2648 : }
      76             : 
      77          78 : Variable resize(const Variable &var, const Dim dim, const scipp::index size,
      78             :                 const FillValue fill) {
      79          78 :   auto dims = var.dims();
      80          78 :   dims.resize(dim, size);
      81         234 :   return special_like(broadcast(Variable(var, Dimensions{}), dims), fill);
      82          78 : }
      83             : 
      84             : /// Return new variable resized to given shape.
      85             : ///
      86             : /// For bucket variables the values of `shape` are interpreted as bucket sizes
      87             : /// to RESERVE and the buffer is also resized accordingly. The emphasis is on
      88             : /// "reserve", i.e., buffer size and begin indices are set up accordingly, but
      89             : /// end=begin is set, i.e., the buckets are empty, but may be grown up to the
      90             : /// requested size. For normal (non-bucket) variable the values of `shape` are
      91             : /// ignored, i.e., only `shape.dims()` is used to determine the shape of the
      92             : /// output.
      93           0 : Variable resize(const Variable &var, const Variable &shape) {
      94           0 :   return {shape.dims(), var.data().makeDefaultFromParent(shape)};
      95             : }
      96             : 
      97         243 : Variable fold(const Variable &view, const Dim from_dim,
      98             :               const Dimensions &to_dims) {
      99         243 :   return view.fold(from_dim, to_dims);
     100             : }
     101             : 
     102       13735 : Variable flatten(const Variable &view,
     103             :                  const scipp::span<const Dim> &from_labels, const Dim to_dim) {
     104       13735 :   if (from_labels.empty()) {
     105           9 :     auto out(view);
     106           9 :     out.unchecked_dims().addInner(to_dim, 1);
     107           9 :     out.unchecked_strides().push_back(1);
     108           9 :     return out;
     109           9 :   }
     110       13726 :   const auto &labels = view.dims().labels();
     111       13726 :   auto it = std::search(labels.begin(), labels.end(), from_labels.begin(),
     112             :                         from_labels.end());
     113       13726 :   if (it == labels.end())
     114           3 :     throw except::DimensionError("Can only flatten a contiguous set of "
     115           6 :                                  "dimensions in the correct order");
     116       13723 :   scipp::index size = 1;
     117       13723 :   auto to = std::distance(labels.begin(), it);
     118       13723 :   auto out(view);
     119       39498 :   for (const auto &from : from_labels) {
     120       29962 :     size *= out.dims().size(to);
     121       29962 :     if (from == from_labels.back()) {
     122        9536 :       out.unchecked_dims().resize(from, size);
     123        9536 :       out.unchecked_dims().replace_key(from, to_dim);
     124             :     } else {
     125       20426 :       if (out.strides()[to] != out.dims().size(to + 1) * out.strides()[to + 1])
     126        8374 :         return flatten(copy(view), from_labels, to_dim);
     127       16239 :       out.unchecked_dims().erase(from);
     128       16239 :       out.unchecked_strides().erase(to);
     129             :     }
     130             :   }
     131        9536 :   return out;
     132       13723 : }
     133             : 
     134        5885 : Variable transpose(const Variable &var, const scipp::span<const Dim> dims) {
     135        5885 :   return var.transpose(dims);
     136             : }
     137             : 
     138             : std::vector<scipp::Dim>
     139        9879 : dims_for_squeezing(const core::Sizes &data_dims,
     140             :                    const std::optional<scipp::span<const Dim>> selected_dims) {
     141        9879 :   if (selected_dims.has_value()) {
     142        8521 :     for (const auto &dim : *selected_dims) {
     143         119 :       if (const auto size = data_dims[dim]; size != 1)
     144          15 :         throw except::DimensionError("Cannot squeeze '" + to_string(dim) +
     145          20 :                                      "' of length " + std::to_string(size) +
     146          10 :                                      ", must be of length 1.");
     147             :     }
     148        8402 :     return std::vector<Dim>{selected_dims->begin(), selected_dims->end()};
     149             :   } else {
     150        1472 :     std::vector<Dim> length_1_dims;
     151        1472 :     length_1_dims.reserve(data_dims.size());
     152        4460 :     for (const auto &dim : data_dims) {
     153        1516 :       if (data_dims[dim] == 1) {
     154         170 :         length_1_dims.push_back(dim);
     155             :       }
     156             :     }
     157        1472 :     return length_1_dims;
     158        1472 :   }
     159             : }
     160             : 
     161        9842 : Variable squeeze(const Variable &var,
     162             :                  const std::optional<scipp::span<const Dim>> dims) {
     163        9842 :   auto squeezed = var;
     164       10084 :   for (const auto &dim : dims_for_squeezing(var.dims(), dims)) {
     165         242 :     squeezed = squeezed.slice({dim, 0});
     166        9839 :   }
     167        9839 :   return squeezed;
     168           3 : }
     169             : 
     170             : } // namespace scipp::variable