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/common/numeric.h"
       6             : #include "scipp/common/overloaded.h"
       7             : 
       8             : #include "scipp/variable/creation.h"
       9             : #include "scipp/variable/misc_operations.h"
      10             : #include "scipp/variable/reduction.h"
      11             : #include "scipp/variable/util.h"
      12             : 
      13             : #include "scipp/dataset/dataset.h"
      14             : #include "scipp/dataset/except.h"
      15             : 
      16             : #include "../variable/operations_common.h"
      17             : #include "dataset_operations_common.h"
      18             : 
      19             : namespace scipp::dataset {
      20             : 
      21           3 : auto union_(const Dataset &a, const Dataset &b) {
      22           3 :   std::map<std::string, DataArray> out;
      23             : 
      24           6 :   for (const auto &item : a)
      25           3 :     out.emplace(item.name(), item);
      26           5 :   for (const auto &item : b) {
      27           3 :     if (const auto it = a.find(item.name()); it != a.end())
      28           4 :       core::expect::equals(item, *it);
      29             :     else
      30           2 :       out.emplace(item.name(), item);
      31           3 :   }
      32           2 :   return out;
      33           1 : }
      34             : 
      35           5 : Dataset merge(const Dataset &a, const Dataset &b) {
      36           8 :   return Dataset(union_(a, b), union_(a.coords(), b.coords(), "merge"));
      37             : }
      38             : 
      39             : /// Return a copy of dict-like objects as a core::Dict.
      40       12438 : template <class Mapping> auto copy_map(const Mapping &map) {
      41       12438 :   core::Dict<typename Mapping::key_type, typename Mapping::mapped_type> out;
      42       27906 :   for (const auto &[key, item] : map)
      43       15468 :     out.insert_or_assign(key, copy(item));
      44       12438 :   return out;
      45           0 : }
      46             : 
      47        7482 : Coords copy(const Coords &coords) { return {coords.sizes(), copy_map(coords)}; }
      48       17394 : Masks copy(const Masks &masks) { return {masks.sizes(), copy_map(masks)}; }
      49             : 
      50             : /// Return a deep copy of a DataArray.
      51        1685 : DataArray copy(const DataArray &array, const AttrPolicy attrPolicy) {
      52             :   // When data is copied we generally need to copy masks, since masks are
      53             :   // typically modified when data is modified.
      54             :   return DataArray(
      55        3370 :       copy(array.data()), copy(array.coords()), copy(array.masks()),
      56        3370 :       attrPolicy == AttrPolicy::Keep ? copy(array.attrs()) : Attrs{},
      57        5055 :       array.name());
      58             : }
      59             : 
      60             : /// Return a deep copy of a Dataset.
      61         290 : Dataset copy(const Dataset &dataset, const AttrPolicy attrPolicy) {
      62         580 :   Dataset out{{}, copy(dataset.coords())};
      63         587 :   for (const auto &item : dataset) {
      64         297 :     out.setData(item.name(), copy(item, attrPolicy));
      65         297 :   }
      66         290 :   return out;
      67           0 : }
      68             : 
      69             : namespace {
      70             : template <class T>
      71           8 : void copy_item(const DataArray &from, T &&to, const AttrPolicy attrPolicy) {
      72          14 :   for (const auto &[name, mask] : from.masks())
      73           6 :     copy(mask, to.masks()[name]);
      74           8 :   if (attrPolicy == AttrPolicy::Keep)
      75           6 :     for (const auto &[dim, attr] : from.attrs())
      76           2 :       copy(attr, to.attrs()[dim]);
      77           8 :   copy(from.data(), to.data());
      78           8 : }
      79             : } // namespace
      80             : 
      81             : /// Copy data array to output data array
      82           5 : DataArray &copy(const DataArray &array, DataArray &out,
      83             :                 const AttrPolicy attrPolicy) {
      84          29 :   for (const auto &[dim, coord] : array.coords())
      85          24 :     copy(coord, out.coords()[dim]);
      86           5 :   copy_item(array, out, attrPolicy);
      87           5 :   return out;
      88             : }
      89             : 
      90             : /// Copy data array to output data array
      91           2 : DataArray copy(const DataArray &array, DataArray &&out,
      92             :                const AttrPolicy attrPolicy) {
      93           2 :   copy(array, out, attrPolicy);
      94           2 :   return std::move(out);
      95             : }
      96             : 
      97             : /// Copy dataset to output dataset
      98           3 : Dataset &copy(const Dataset &dataset, Dataset &out,
      99             :               const AttrPolicy attrPolicy) {
     100          27 :   for (const auto &[dim, coord] : dataset.coords())
     101          24 :     copy(coord, out.coords()[dim]);
     102           6 :   for (const auto &array : dataset)
     103           3 :     copy_item(array, out[array.name()], attrPolicy);
     104           3 :   return out;
     105             : }
     106             : 
     107             : /// Copy dataset to output dataset
     108           0 : Dataset copy(const Dataset &dataset, Dataset &&out,
     109             :              const AttrPolicy attrPolicy) {
     110           0 :   copy(dataset, out, attrPolicy);
     111           0 :   return std::move(out);
     112             : }
     113             : 
     114             : /// Return data of data array, applying masks along dim if applicable.
     115             : ///
     116             : /// Only in the latter case a copy is returned. Masked values are replaced by
     117             : /// fill_value. If not provided, values are replaced by zero.
     118       47846 : Variable masked_data(const DataArray &array, const Dim dim,
     119             :                      const std::optional<Variable> &fill_value) {
     120       47846 :   const auto mask = irreducible_mask(array.masks(), dim);
     121       47846 :   if (mask.is_valid()) {
     122          28 :     const auto &data = array.data();
     123          28 :     const auto fill = fill_value.value_or(zero_like(array.data()));
     124          28 :     return where(mask, fill, data);
     125          28 :   } else
     126       47818 :     return array.data();
     127       47846 : }
     128             : 
     129             : namespace {
     130           0 : template <class Dict> auto strip_(const Dict &dict, const Dim dim) {
     131           0 :   Dict stripped(dict.sizes(), {});
     132           0 :   for (const auto &[key, value] : dict)
     133           0 :     if (value.dims().contains(dim))
     134           0 :       stripped.set(key, value);
     135           0 :   return stripped;
     136           0 : }
     137             : } // namespace
     138             : 
     139           0 : DataArray strip_if_broadcast_along(const DataArray &a, const Dim dim) {
     140           0 :   return {a.data(), strip_(a.coords(), dim), strip_(a.masks(), dim),
     141           0 :           strip_(a.attrs(), dim), a.name()};
     142             : }
     143             : 
     144           0 : Dataset strip_if_broadcast_along(const Dataset &d, const Dim dim) {
     145           0 :   Dataset stripped;
     146           0 :   stripped.setCoords(strip_(d.coords(), dim));
     147           0 :   for (auto &&item : d)
     148           0 :     if (item.dims().contains(dim))
     149           0 :       stripped.setData(item.name(), strip_if_broadcast_along(item, dim));
     150           0 :   return stripped;
     151           0 : }
     152             : 
     153             : } // namespace scipp::dataset