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 : #pragma once
6 :
7 : #include "scipp/core/dict.h"
8 : #include "scipp/dataset/dataset.h"
9 : #include "scipp/dataset/except.h"
10 : #include "scipp/variable/arithmetic.h"
11 :
12 : namespace scipp::dataset {
13 : template <bool ApplyToData, class Func, class... Args>
14 50374 : DataArray apply_or_copy_dim_impl(const DataArray &da, Func func, const Dim dim,
15 : Args &&...args) {
16 262291 : const auto copy_independent = [&](const auto &mapping, const bool share) {
17 151122 : typename std::decay_t<decltype(mapping)>::holder_type out;
18 211917 : for (auto &&[d, var] : mapping)
19 60795 : if (!var.dims().contains(dim))
20 4227 : out.insert_or_assign(d, share ? var : copy(var));
21 151122 : return out;
22 0 : };
23 50374 : auto coords = copy_independent(da.coords(), true);
24 50374 : auto attrs = copy_independent(da.attrs(), true);
25 50374 : auto masks = copy_independent(da.masks(), false);
26 :
27 : if constexpr (ApplyToData) {
28 2781 : return DataArray(func(da.data(), dim, args...), std::move(coords),
29 8343 : std::move(masks), std::move(attrs), da.name());
30 : } else {
31 47593 : return DataArray(func(da, dim, std::forward<Args>(args)...),
32 142779 : std::move(coords), std::move(masks), std::move(attrs),
33 285560 : da.name());
34 : }
35 50376 : }
36 :
37 : /// Helper for creating operations that return an object with modified data with
38 : /// a dropped dimension or different dimension extent.
39 : ///
40 : /// Examples are mostly reduction operations such as `sum` (dropping a
41 : /// dimension), or `resize` (altering a dimension extent). Creates new data
42 : /// array by applying `func` to data and dropping coords/masks/attrs depending
43 : /// on dim.
44 : template <class Func, class... Args>
45 2781 : DataArray apply_to_data_and_drop_dim(const DataArray &a, Func func,
46 : const Dim dim, Args &&...args) {
47 : return apply_or_copy_dim_impl<true>(a, func, dim,
48 2781 : std::forward<Args>(args)...);
49 : }
50 :
51 : /// Helper for creating operations that return an object with a dropped
52 : /// dimension or different dimension extent.
53 : ///
54 : /// In contrast to `apply_to_data_and_drop_dim`, `func` is applied to the input
55 : /// array, not just its data. This is useful for more complex operations such as
56 : /// `histogram`, which require access to coords when computing output data.
57 : template <class Func, class... Args>
58 47593 : DataArray apply_and_drop_dim(const DataArray &a, Func func, const Dim dim,
59 : Args &&...args) {
60 : return apply_or_copy_dim_impl<false>(a, func, dim,
61 47593 : std::forward<Args>(args)...);
62 : }
63 :
64 : template <class Func, class... Args>
65 : DataArray apply_to_items(const DataArray &d, Func func, Args &&...args) {
66 : return func(d, std::forward<Args>(args)...);
67 : }
68 :
69 : template <class... Args>
70 : bool copy_attr(const Variable &attr, const Dim dim, const Args &...) {
71 : return !attr.dims().contains(dim);
72 : }
73 : template <class... Args> bool copy_attr(const Variable &, const Args &...) {
74 : return true;
75 : }
76 :
77 : template <class Func, class... Args>
78 203 : Dataset apply_to_items(const Dataset &d, Func func, Args &&...args) {
79 203 : Dataset result;
80 428 : for (const auto &data : d)
81 225 : result.setDataInit(data.name(), func(data, std::forward<Args>(args)...));
82 406 : return std::move(result).or_empty();
83 203 : }
84 :
85 : /// Return a copy of map-like objects such as Coords with `func` applied to each
86 : /// item.
87 : ///
88 : /// If `func` returns an invalid object it will not be inserted into the output
89 : /// map. This can be used to drop/filter items.
90 : template <class OutMapping, class Mapping, class Func>
91 32304 : auto transform_map(const Mapping &map, Func func) {
92 32304 : OutMapping out;
93 61751 : for (const auto &[key, item] : map) {
94 29447 : auto transformed = func(item);
95 29447 : if (transformed.is_valid())
96 29413 : out.insert_or_assign(key, std::move(transformed));
97 : }
98 32304 : return out;
99 0 : }
100 :
101 10768 : template <class T, class Func> DataArray transform(const T &a, Func func) {
102 : return DataArray(
103 : func(a.data()),
104 2310 : transform_map<typename Coords::holder_type>(a.coords(), func),
105 2310 : transform_map<typename Masks::holder_type>(a.masks(), func),
106 15390 : transform_map<typename Attrs::holder_type>(a.attrs(), func), a.name());
107 : }
108 :
109 : [[nodiscard]] DataArray strip_if_broadcast_along(const DataArray &a,
110 : const Dim dim);
111 : [[nodiscard]] Dataset strip_if_broadcast_along(const Dataset &d, const Dim dim);
112 :
113 : // Helpers for reductions for DataArray and Dataset, which include masks.
114 : [[nodiscard]] Variable mean(const Variable &var, const Dim dim,
115 : const Masks &masks);
116 : [[nodiscard]] Variable nanmean(const Variable &var, const Dim dim,
117 : const Masks &masks);
118 : [[nodiscard]] Variable sum(const Variable &var, const Dim dim,
119 : const Masks &masks);
120 : [[nodiscard]] Variable nansum(const Variable &var, const Masks &masks);
121 : [[nodiscard]] Variable nansum(const Variable &var, const Dim dim,
122 : const Masks &masks);
123 : [[nodiscard]] Variable max(const Variable &var, const Dim dim,
124 : const Masks &masks);
125 : [[nodiscard]] Variable nanmax(const Variable &var, const Dim dim,
126 : const Masks &masks);
127 : [[nodiscard]] Variable min(const Variable &var, const Dim dim,
128 : const Masks &masks);
129 : [[nodiscard]] Variable nanmin(const Variable &var, const Dim dim,
130 : const Masks &masks);
131 : [[nodiscard]] Variable all(const Variable &var, const Dim dim,
132 : const Masks &masks);
133 : [[nodiscard]] Variable any(const Variable &var, const Dim dim,
134 : const Masks &masks);
135 :
136 : [[nodiscard]] Variable
137 : masked_data(const DataArray &array, const Dim dim,
138 : const std::optional<Variable> &fill_value = std::nullopt);
139 :
140 : } // namespace scipp::dataset
|
