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 Owen Arnold, Simon Heybrock
       5             : #include <algorithm>
       6             : 
       7             : #include "scipp/units/dim.h"
       8             : #include "scipp/variable/comparison.h"
       9             : #include "scipp/variable/reduction.h"
      10             : #include "scipp/variable/slice.h"
      11             : #include "scipp/variable/string.h"
      12             : #include "scipp/variable/util.h"
      13             : #include "scipp/variable/variable.h"
      14             : 
      15             : namespace scipp::variable {
      16             : 
      17             : namespace {
      18             : 
      19         294 : scipp::index get_count(const Variable &coord, const Dim dim,
      20             :                        const Variable &value, const bool ascending) {
      21         792 :   return (ascending ? sum(less_equal(coord, value), dim)
      22         498 :                     : sum(greater_equal(coord, value), dim))
      23         588 :       .value<scipp::index>();
      24             : }
      25             : 
      26         278 : scipp::index get_index(const Variable &coord, const Dim dim,
      27             :                        const Variable &value, const bool ascending,
      28             :                        const bool edges) {
      29         278 :   auto i = get_count(coord, dim, value, edges == ascending);
      30         278 :   i = edges ? i - 1 : coord.dims()[dim] - i;
      31         278 :   return std::clamp<scipp::index>(0, i, coord.dims()[dim]);
      32             : }
      33             : 
      34         332 : const Variable &get_1d_coord(const Variable &coord) {
      35         332 :   if (coord.dims().ndim() != 1)
      36           2 :     throw except::DimensionError("Multi-dimensional coordinates cannot be used "
      37           4 :                                  "for label-based indexing.");
      38         330 :   return coord;
      39             : }
      40             : 
      41         178 : auto get_coord(const Variable &coord, const Dim dim) {
      42         178 :   get_1d_coord(coord);
      43         176 :   if (coord.dims()[dim] == 1)
      44             :     // Need this because issorted returns false for length-1 variables.
      45           2 :     return std::tuple(coord, true);
      46         174 :   const bool ascending = allsorted(coord, dim, SortOrder::Ascending);
      47         174 :   const bool descending = allsorted(coord, dim, SortOrder::Descending);
      48         174 :   if (!(ascending ^ descending))
      49           2 :     throw std::runtime_error("Coordinate must be monotonically increasing or "
      50           4 :                              "decreasing for label-based indexing.");
      51         172 :   return std::tuple(coord, ascending);
      52             : }
      53             : 
      54         459 : void expect_same_unit(const Variable &coord, const Variable &value,
      55             :                       const std::string &name) {
      56         459 :   if (coord.unit() != value.unit()) {
      57          21 :     throw except::UnitError("The unit of the slice " + name + " (" +
      58          28 :                             to_string(value.unit()) +
      59          14 :                             ") does not match the unit of the coordinate (" +
      60          28 :                             to_string(coord.unit()) + ").");
      61             :   }
      62         452 : }
      63             : 
      64         614 : void expect_valid_dtype(const Variable &var, const bool is_range,
      65             :                         const std::string &name) {
      66         614 :   if (is_range && !is_total_orderable(var.dtype())) {
      67           4 :     throw except::TypeError(
      68           8 :         "The dtype of the slice " + name + " (" + to_string(var.dtype()) +
      69             :         ") cannot be used for label-based slicing because it does not"
      70           8 :         " define an order.");
      71             :   }
      72         610 : }
      73             : 
      74         515 : void expect_valid_slice_value(const Variable &coord, const Variable &value,
      75             :                               const bool is_range,
      76             :                               const std::string_view name) {
      77         515 :   if (value.is_valid()) {
      78         475 :     core::expect::equals(Dimensions{}, value.dims());
      79         473 :     expect_same_unit(coord, value, std::string(name));
      80         456 :     expect_valid_dtype(value, is_range, std::string(name));
      81             :   }
      82         500 : }
      83             : } // namespace
      84             : 
      85         173 : std::tuple<Dim, scipp::index> get_slice_params(const Sizes &dims,
      86             :                                                const Variable &coord_,
      87             :                                                const Variable &value) {
      88         173 :   expect_valid_slice_value(coord_, value, false, "key");
      89         170 :   const auto dim = coord_.dims().inner();
      90         170 :   if (dims[dim] + 1 == coord_.dims()[dim]) {
      91          16 :     const auto &[coord, ascending] = get_coord(coord_, dim);
      92          16 :     return std::tuple{dim, get_count(coord, dim, value, ascending) - 1};
      93          16 :   } else {
      94         154 :     auto eq = equal(get_1d_coord(coord_), value);
      95         154 :     if (sum(eq, dim).template value<scipp::index>() != 1)
      96          15 :       throw except::SliceError("Coord " + to_string(dim) +
      97          10 :                                " does not contain unique point with value " +
      98          20 :                                to_string(value) + '\n');
      99         149 :     auto values = eq.template values<bool>();
     100         149 :     auto it = std::find(values.begin(), values.end(), true);
     101         149 :     return {dim, std::distance(values.begin(), it)};
     102         154 :   }
     103             : }
     104             : 
     105             : std::tuple<Dim, scipp::index, scipp::index>
     106         174 : get_slice_params(const Sizes &dims, const Variable &coord_,
     107             :                  const Variable &begin, const Variable &end) {
     108         174 :   expect_valid_slice_value(coord_, begin, true, "begin");
     109         168 :   expect_valid_slice_value(coord_, end, true, "end");
     110         162 :   expect_valid_dtype(coord_, true, "coord");
     111         162 :   const auto dim = coord_.dims().inner();
     112         162 :   const auto &[coord, ascending] = get_coord(coord_, dim);
     113         158 :   scipp::index first = 0;
     114         158 :   scipp::index last = dims[dim];
     115         158 :   const auto bin_edges = last + 1 == coord.dims()[dim];
     116         158 :   if (begin.is_valid())
     117         135 :     first = get_index(coord, dim, begin, ascending, bin_edges);
     118         158 :   if (end.is_valid())
     119         143 :     last = get_index(coord, dim, end, ascending, false);
     120         316 :   return {dim, first, std::min(dims[dim], last)};
     121         158 : }
     122             : 
     123             : } // namespace scipp::variable