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/dataset/data_array.h"
       6             : #include "scipp/dataset/dataset_util.h"
       7             : #include "scipp/variable/misc_operations.h"
       8             : 
       9             : #include "dataset_operations_common.h"
      10             : 
      11             : namespace scipp::dataset {
      12             : 
      13             : namespace {
      14         176 : template <class T> void expect_writable(const T &dict) {
      15         176 :   if (dict.is_readonly())
      16           4 :     throw except::DataArrayError("Read-only flag is set, cannot set new data.");
      17         172 : }
      18             : 
      19      988520 : template <class T> auto copy_shared(const std::shared_ptr<T> &obj) {
      20      988520 :   return obj ? std::make_shared<T>(*obj) : obj;
      21             : }
      22             : } // namespace
      23             : 
      24      247130 : DataArray::DataArray(const DataArray &other, const AttrPolicy attrPolicy)
      25      247130 :     : m_name(other.m_name), m_data(copy_shared(other.m_data)),
      26      247130 :       m_coords(copy_shared(other.m_coords)),
      27      247130 :       m_masks(copy_shared(other.m_masks)),
      28      247130 :       m_attrs(attrPolicy == AttrPolicy::Keep ? copy_shared(other.m_attrs)
      29             :                                              : std::make_shared<Attrs>()),
      30      247130 :       m_readonly(false) {}
      31             : 
      32      247130 : DataArray::DataArray(const DataArray &other)
      33      247130 :     : DataArray(other, AttrPolicy::Keep) {}
      34             : 
      35       18822 : DataArray::DataArray(Variable data, Coords coords, Masks masks, Attrs attrs,
      36       18822 :                      const std::string_view name)
      37       18822 :     : m_name(name), m_data(std::make_shared<Variable>(std::move(data))),
      38       18822 :       m_coords(std::make_shared<Coords>(std::move(coords))),
      39       18822 :       m_masks(std::make_shared<Masks>(std::move(masks))),
      40       18822 :       m_attrs(std::make_shared<Attrs>(std::move(attrs))) {
      41       18822 :   const Sizes sizes(dims());
      42       18822 :   m_coords->setSizes(sizes);
      43       18822 :   m_masks->setSizes(sizes);
      44       18822 :   m_attrs->setSizes(sizes);
      45       18822 : }
      46             : 
      47      138813 : DataArray::DataArray(Variable data, typename Coords::holder_type coords,
      48             :                      typename Masks::holder_type masks,
      49             :                      typename Attrs::holder_type attrs,
      50      138813 :                      const std::string_view name)
      51      138813 :     : m_name(name), m_data(std::make_shared<Variable>(std::move(data))),
      52      138813 :       m_coords(std::make_shared<Coords>(dims(), std::move(coords))),
      53      138810 :       m_masks(std::make_shared<Masks>(dims(), std::move(masks))),
      54      138818 :       m_attrs(std::make_shared<Attrs>(dims(), std::move(attrs))) {}
      55             : 
      56          67 : DataArray &DataArray::operator=(const DataArray &other) {
      57          67 :   if (this == &other) {
      58           8 :     return *this;
      59             :   }
      60          59 :   check_nested_in_assign(*this, other);
      61          55 :   return *this = DataArray(other);
      62             : }
      63             : 
      64       48273 : DataArray &DataArray::operator=(DataArray &&other) {
      65       48273 :   if (this == &other) {
      66           0 :     return *this;
      67             :   }
      68       48273 :   check_nested_in_assign(*this, other);
      69       48273 :   m_name = std::move(other.m_name);
      70       48273 :   m_data = std::move(other.m_data);
      71       48273 :   m_coords = std::move(other.m_coords);
      72       48273 :   m_masks = std::move(other.m_masks);
      73       48273 :   m_attrs = std::move(other.m_attrs);
      74       48273 :   m_readonly = other.m_readonly;
      75       48273 :   return *this;
      76             : }
      77             : 
      78         364 : void DataArray::setData(const Variable &data) {
      79             :   // Return early on self assign to avoid exceptions from Python inplace ops
      80         364 :   if (m_data->is_same(data))
      81         188 :     return;
      82         176 :   expect_writable(*this);
      83         518 :   core::expect::equals(static_cast<Sizes>(dims()),
      84         345 :                        static_cast<Sizes>(data.dims()));
      85         171 :   *m_data = data;
      86             : }
      87             : 
      88             : /// Return true if the dataset proxies have identical content.
      89        5903 : bool operator==(const DataArray &a, const DataArray &b) {
      90        5903 :   if (a.has_variances() != b.has_variances())
      91          20 :     return false;
      92        5883 :   if (a.coords() != b.coords())
      93         102 :     return false;
      94        5781 :   if (a.masks() != b.masks())
      95          44 :     return false;
      96        5737 :   if (a.attrs() != b.attrs())
      97          54 :     return false;
      98        5683 :   return a.data() == b.data();
      99             : }
     100             : 
     101        1059 : bool operator!=(const DataArray &a, const DataArray &b) {
     102        1059 :   return !operator==(a, b);
     103             : }
     104             : 
     105         325 : bool equals_nan(const DataArray &a, const DataArray &b) {
     106         325 :   if (!equals_nan(a.coords(), b.coords()))
     107          34 :     return false;
     108         291 :   if (!equals_nan(a.masks(), b.masks()))
     109          14 :     return false;
     110         277 :   if (!equals_nan(a.attrs(), b.attrs()))
     111          16 :     return false;
     112         261 :   return equals_nan(a.data(), b.data());
     113             : }
     114             : 
     115             : /// Return the name of the data array.
     116             : ///
     117             : /// If part of a dataset, the name of the array is equal to the key of this item
     118             : /// in the dataset. Note that comparison operations ignore the name.
     119       74911 : const std::string &DataArray::name() const { return m_name; }
     120             : 
     121          21 : void DataArray::setName(const std::string_view name) { m_name = name; }
     122             : 
     123       47274 : Coords DataArray::meta() const {
     124       47274 :   auto out = attrs().merge_from(coords());
     125       47273 :   out.set_readonly();
     126       47273 :   return out;
     127             : }
     128             : 
     129        9085 : DataArray DataArray::slice(const Slice &s) const {
     130       18069 :   auto out = DataArray{m_data->slice(s), m_coords->slice_coords(s),
     131       17968 :                        m_masks->slice(s), m_attrs->slice(s), m_name};
     132        8984 :   out.m_readonly = true;
     133        8984 :   return out;
     134             : }
     135             : 
     136          18 : void DataArray::validateSlice(const Slice &s, const DataArray &array) const {
     137          18 :   expect::coords_are_superset(slice(s), array, "");
     138          18 :   data().validateSlice(s, array.data());
     139          18 :   masks().validateSlice(s, array.masks());
     140          15 : }
     141             : 
     142          45 : DataArray &DataArray::setSlice(const Slice &s, const DataArray &array) {
     143             :   // validateSlice, but not masks as otherwise repeated
     144          47 :   expect::coords_are_superset(slice(s), array, "");
     145          43 :   data().validateSlice(s, array.data());
     146             :   // Apply changes
     147          40 :   masks().setSlice(s, array.masks());
     148          36 :   return setSlice(s, array.data());
     149             : }
     150             : 
     151          44 : DataArray &DataArray::setSlice(const Slice &s, const Variable &var) {
     152          44 :   data().setSlice(s, var);
     153          43 :   return *this;
     154             : }
     155             : 
     156           1 : DataArray DataArray::view() const {
     157           1 :   DataArray out;
     158           1 :   out.m_data = m_data;     // share data
     159           1 :   out.m_coords = m_coords; // share coords
     160           1 :   out.m_masks = m_masks;   // share masks
     161           1 :   out.m_attrs = m_attrs;   // share attrs
     162           1 :   out.m_name = m_name;
     163           1 :   return out;
     164           0 : }
     165             : 
     166        7528 : DataArray DataArray::view_with_coords(const Coords &coords,
     167             :                                       const std::string &name,
     168             :                                       const bool readonly) const {
     169        7528 :   DataArray out;
     170        7528 :   out.m_data = m_data; // share data
     171        7528 :   const Sizes sizes(dims());
     172        7528 :   typename Coords::holder_type selected;
     173       28424 :   for (const auto &[dim, coord] : coords)
     174       20896 :     if (coords.item_applies_to(dim, dims()))
     175       20780 :       selected.insert_or_assign(dim, coord.as_const());
     176        7528 :   const bool readonly_coords = true;
     177             :   out.m_coords =
     178        7528 :       std::make_shared<Coords>(sizes, std::move(selected), readonly_coords);
     179        7528 :   out.m_masks = m_masks; // share masks
     180        7528 :   out.m_attrs = m_attrs; // share attrs
     181        7528 :   out.m_name = name;
     182        7528 :   out.m_readonly = readonly;
     183       15056 :   return out;
     184        7528 : }
     185             : 
     186             : DataArray
     187          10 : DataArray::drop_coords(const scipp::span<const Dim> coord_names) const {
     188          10 :   DataArray result = *this;
     189          24 :   for (const auto &name : coord_names)
     190          14 :     result.coords().erase(name);
     191          10 :   return result;
     192           0 : }
     193             : 
     194             : DataArray
     195           8 : DataArray::drop_masks(const scipp::span<const std::string> mask_names) const {
     196           8 :   DataArray result = *this;
     197          20 :   for (const auto &name : mask_names)
     198          12 :     result.masks().erase(name);
     199           8 :   return result;
     200           0 : }
     201             : 
     202           8 : DataArray DataArray::drop_attrs(const scipp::span<const Dim> attr_names) const {
     203           8 :   DataArray result = *this;
     204          20 :   for (const auto &name : attr_names)
     205          12 :     result.attrs().erase(name);
     206           8 :   return result;
     207           0 : }
     208             : 
     209         281 : DataArray DataArray::rename_dims(const std::vector<std::pair<Dim, Dim>> &names,
     210             :                                  const bool fail_on_unknown) const {
     211         553 :   return DataArray(m_data->rename_dims(names, fail_on_unknown),
     212         557 :                    m_coords->rename_dims(names, false),
     213         561 :                    m_masks->rename_dims(names, false),
     214         839 :                    m_attrs->rename_dims(names, false));
     215             : }
     216             : 
     217           2 : DataArray DataArray::as_const() const {
     218           4 :   auto out = DataArray(data().as_const(), coords().as_const(),
     219           6 :                        masks().as_const(), attrs().as_const(), name());
     220           2 :   out.m_readonly = true;
     221           2 :   return out;
     222             : }
     223             : 
     224         177 : bool DataArray::is_readonly() const noexcept { return m_readonly; }
     225             : 
     226             : } // namespace scipp::dataset