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             : 
       6             : #include "scipp/units/unit.h"
       7             : 
       8             : #include "scipp/core/spatial_transforms.h"
       9             : #include "scipp/core/time_point.h"
      10             : 
      11             : #include "scipp/variable/operations.h"
      12             : #include "scipp/variable/structures.h"
      13             : #include "scipp/variable/util.h"
      14             : #include "scipp/variable/variable.h"
      15             : #include "scipp/variable/variable_factory.h"
      16             : 
      17             : #include "scipp/dataset/dataset.h"
      18             : 
      19             : #include "bind_data_access.h"
      20             : #include "bind_operators.h"
      21             : #include "bind_slice_methods.h"
      22             : #include "dim.h"
      23             : #include "pybind11.h"
      24             : #include "rename.h"
      25             : 
      26             : using namespace scipp;
      27             : using namespace scipp::variable;
      28             : 
      29             : namespace py = pybind11;
      30             : 
      31             : template <class T> struct GetElements {
      32          42 :   static Variable apply(Variable &var, const std::string &key) {
      33          42 :     return var.elements<T>(key);
      34             :   }
      35             : };
      36             : 
      37             : template <class T> struct SetElements {
      38          17 :   static void apply(Variable &var, const std::string &key,
      39             :                     const Variable &elems) {
      40          17 :     copy(elems, var.elements<T>(key));
      41          16 :   }
      42             : };
      43             : 
      44           3 : template <class T> void bind_alignment_functions(py::class_<T> &variable) {
      45             :   // We use a separate setter instead of making the 'aligned' property writable
      46             :   // in order to reduce the chance of accidentally setting the flag on
      47             :   // temporary variables.
      48           3 :   variable.def_property_readonly(
      49        1382 :       "aligned", [](const Variable &self) { return self.is_aligned(); },
      50             :       R"(Alignment flag for coordinates.
      51             : 
      52             : Indicates whether a coordinate is aligned.
      53             : Aligned coordinates must match between the operands of binary operations while
      54             : unaligned coordinates are dropped on mismatch.
      55             : 
      56             : This flag is only meaningful when the variable is contained in a ``coords``
      57             : ``dict``.
      58             : 
      59             : It cannot be set on a variable directly;
      60             : instead, use :meth:`sc.Coords.set_aligned`.
      61             : 
      62             : For *binned* coordinates of a binned data array ``da``,
      63             : ``da.bins.coords[name].aligned`` should always be ``True``.
      64             : The alignment w.r.t. the events can be queried via
      65             : ``da.bins.coords[name].bins.aligned`` and set via
      66             : ``da.bins.coords.set_aligned(name, aligned)``.
      67             : )");
      68           3 : }
      69             : 
      70             : void bind_init(py::class_<Variable> &cls);
      71             : 
      72           3 : void init_variable(py::module &m) {
      73             :   // Needed to let numpy arrays keep alive the scipp buffers.
      74             :   // VariableConcept must ALWAYS be passed to Python by its handle.
      75             :   py::class_<VariableConcept, VariableConceptHandle> variable_concept(
      76           3 :       m, "_VariableConcept");
      77             : 
      78           0 :   py::class_<Variable> variable(m, "Variable", py::dynamic_attr(),
      79             :                                 R"(
      80             : Array of values with dimension labels and a unit, optionally including an array
      81           3 : of variances.)");
      82             : 
      83           3 :   bind_init(variable);
      84           3 :   variable.def("_rename_dims", &rename_dims<Variable>)
      85           3 :       .def_property_readonly("dtype", &Variable::dtype);
      86             : 
      87           3 :   bind_common_operators(variable);
      88             : 
      89           3 :   bind_astype(variable);
      90             : 
      91           3 :   bind_slice_methods(variable);
      92             : 
      93           3 :   bind_comparison<Variable>(variable);
      94           3 :   bind_comparison_scalars(variable);
      95             : 
      96           3 :   bind_in_place_binary<Variable>(variable);
      97           3 :   bind_in_place_binary_scalars(variable);
      98             : 
      99           3 :   bind_binary<Variable>(variable);
     100           3 :   bind_binary<DataArray>(variable);
     101           3 :   bind_binary_scalars(variable);
     102           3 :   bind_reverse_binary_scalars(variable);
     103             : 
     104           3 :   bind_unary(variable);
     105             : 
     106           3 :   bind_boolean_unary(variable);
     107           3 :   bind_logical<Variable>(variable);
     108             : 
     109           3 :   bind_data_properties(variable);
     110           3 :   bind_alignment_functions(variable);
     111             : 
     112           3 :   m.def(
     113             :       "islinspace",
     114          91 :       [](const Variable &x,
     115             :          const std::optional<std::string> &dim = std::nullopt) {
     116         151 :         return scipp::variable::islinspace(x, dim.has_value() ? Dim{*dim}
     117         151 :                                                               : x.dim());
     118             :       },
     119           6 :       py::arg("x"), py::arg("dim") = py::none(),
     120           0 :       py::call_guard<py::gil_scoped_release>());
     121             : 
     122             :   using structured_t =
     123             :       std::tuple<Eigen::Vector3d, Eigen::Matrix3d, Eigen::Affine3d,
     124             :                  scipp::core::Quaternion, scipp::core::Translation>;
     125           3 :   m.def("_element_keys", element_keys);
     126           3 :   m.def("_get_elements", [](Variable &self, const std::string &key) {
     127             :     return core::callDType<GetElements>(
     128          84 :         structured_t{}, variableFactory().elem_dtype(self), self, key);
     129             :   });
     130           3 :   m.def("_set_elements", [](Variable &self, const std::string &key,
     131             :                             const Variable &elems) {
     132          17 :     core::callDType<SetElements>(
     133          17 :         structured_t{}, variableFactory().elem_dtype(self), self, key, elems);
     134          16 :   });
     135           3 : }