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/core/element/bin_detail.h"
       6             : #include "scipp/core/element/util.h" // fill_zeros
       7             : #include "scipp/core/subbin_sizes.h"
       8             : 
       9             : #include "scipp/variable/accumulate.h"
      10             : #include "scipp/variable/bin_detail.h"
      11             : #include "scipp/variable/bin_util.h"
      12             : #include "scipp/variable/subspan_view.h"
      13             : #include "scipp/variable/transform.h"
      14             : #include "scipp/variable/util.h"
      15             : 
      16             : namespace scipp::variable::bin_detail {
      17             : 
      18             : /// Index of the bin (given by `edges`) containing a coord value.
      19             : ///
      20             : /// 0 if the coord is less than the first edge, nbin-1 if greater or equal last
      21             : /// edge. Assumes both `edges` and `coord` are sorted.
      22         726 : Variable begin_edge(const Variable &coord, const Variable &edges) {
      23         726 :   auto indices = makeVariable<scipp::index>(coord.dims(), units::none);
      24         726 :   const auto dim = edges.dims().inner();
      25         726 :   if (indices.dims()[dim] == 0)
      26          10 :     return indices;
      27         716 :   auto bin = copy(indices.slice({dim, 0}));
      28         716 :   accumulate_in_place(bin, indices, coord, subspan_view(edges, dim),
      29             :                       core::element::begin_edge, "scipp.bin.begin_edge");
      30         716 :   return indices;
      31         716 : }
      32             : 
      33             : /// End bin
      34             : ///
      35             : /// 1 if the coord is
      36             : /// nbin if the coord is greater than the last edge. Assumes both `edges` and
      37             : /// `coord` are sorted.
      38         454 : Variable end_edge(const Variable &coord, const Variable &edges) {
      39         454 :   auto indices = makeVariable<scipp::index>(coord.dims(), units::none);
      40         454 :   const auto dim = edges.dims().inner();
      41         454 :   if (indices.dims()[dim] == 0)
      42           0 :     return indices;
      43         454 :   auto bin = copy(indices.slice({dim, 0}));
      44         454 :   accumulate_in_place(bin, indices, coord, subspan_view(edges, dim),
      45             :                       core::element::end_edge, "scipp.bin.end_edge");
      46         454 :   return indices;
      47         454 : }
      48             : 
      49        8377 : Variable cumsum_exclusive_subbin_sizes(const Variable &var) {
      50             :   return transform<core::SubbinSizes>(
      51             :       var,
      52        8377 :       overloaded{[](const units::Unit &u) { return u; },
      53       29044 :                  [](const auto &sizes) { return sizes.cumsum_exclusive(); }},
      54        8377 :       "scipp.bin.cumsum_exclusive");
      55             : }
      56             : 
      57        8377 : Variable sum_subbin_sizes(const Variable &var) {
      58             :   return transform<core::SubbinSizes>(
      59             :       var,
      60        8377 :       overloaded{[](const units::Unit &u) { return u; },
      61       29044 :                  [](const auto &sizes) { return sizes.sum(); }},
      62        8377 :       "scipp.bin.sum_subbin_sizes");
      63             : }
      64             : 
      65        8377 : std::vector<scipp::index> flatten_subbin_sizes(const Variable &var,
      66             :                                                const scipp::index length) {
      67        8377 :   std::vector<scipp::index> flat(length * var.dims().volume());
      68        8377 :   auto it = flat.begin();
      69       37421 :   for (const auto &sub : var.values<core::SubbinSizes>()) {
      70       29044 :     std::copy_n(sub.sizes().begin(), sub.sizes().size(), it + sub.offset());
      71       29044 :     it += length;
      72        8377 :   }
      73       16754 :   return flat;
      74           0 : }
      75             : 
      76        6955 : Variable subbin_sizes_cumsum_exclusive(const Variable &var, const Dim dim) {
      77        6955 :   if (var.dims()[dim] == 0)
      78          20 :     return copy(var);
      79        6935 :   auto cumulative = copy(var.slice({dim, 0}));
      80        6935 :   fill_zeros(cumulative);
      81        6935 :   auto out = copy(var);
      82        6935 :   accumulate_in_place(cumulative, out,
      83             :                       core::element::subbin_sizes_exclusive_scan,
      84             :                       "scipp.bin.subbin_sizes_cumsum_exclusive");
      85        6935 :   return out;
      86        6935 : }
      87             : 
      88       15332 : void subbin_sizes_add_intersection(Variable &a, const Variable &b) {
      89       15332 :   transform_in_place(a, b, core::element::subbin_sizes_add_intersection,
      90             :                      "scipp.bin.subbin_sizes_add_intersection");
      91       15332 : }
      92             : 
      93             : } // namespace scipp::variable::bin_detail