converters.hpp

// Copyright (c) 2024--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.
 
#pragma once
 
#include <nda/nda.hpp>
#include "./make_numpy_proxy_from_array.hpp"
 
namespace c2py {
 
  // -----------------------------------
  // array_view
  // -----------------------------------
 
  template <int R, typename Layout>
  bool numpy_check_layout(PyObject *obj) {
    EXPECTS(PyArray_Check(obj));
    PyArrayObject *arr = (PyArrayObject *)(obj);
    return Layout::template mapping<R>::is_stride_order_valid(PyArray_DIMS(arr), PyArray_STRIDES(arr));
  }
 
  template <typename T, int R, typename Layout, char Algebra>
    requires(has_npy_type<std::decay_t<T>>)
  struct py_converter<nda::basic_array_view<T, R, Layout, Algebra>> {
 
    using view_t = nda::basic_array_view<T, R, Layout, Algebra>;
    using U      = std::decay_t<T>;
    static_assert(not std::is_same_v<U, pyref>, "Not implemented"); // would require to take care of the incref...
    // However, it works for PyObject *
 
    // ------------ tp_name ---------------
 
    static std::string tp_name() {
      std::ostringstream out;
      out << "ndarray[" << python_typename<T>() << ", " << R << "]";
      return out.str();
    }
 
    // --------- C -> PY --------
 
    static PyObject *c2py(view_t v) {
      auto p = nda::python::make_numpy_proxy_from_array_or_view(v);
      return p.to_python();
    }
 
    // --------- PY -> C is possible --------
 
    static bool is_convertible(PyObject *obj, bool raise_python_exception, bool allow_lower_rank = false, bool require_c_order = true) {
      // has_npy_type<T> is true (static_assert at top)
      // check the rank and type. First protects the rest
      if (not PyArray_Check(obj)) {
        if (raise_python_exception) PyErr_SetString(PyExc_TypeError, "Cannot convert to array_view : Python object is not a Numpy array");
        return false;
      }
      PyArrayObject *arr = (PyArrayObject *)(obj);
 
      auto r = PyArray_NDIM(arr);
      if (allow_lower_rank ? r < R : r != R) {
        if (raise_python_exception)
          PyErr_SetString(
             PyExc_TypeError,
             ("Cannot convert to array_view : Rank is not correct. Expected " + std::to_string(R) + "\n Got " + std::to_string(PyArray_NDIM(arr)))
                .c_str());
        return false;
      }
 
      if (not allow_lower_rank && r == R) {
        if (has_npy_type<U> && (PyArray_TYPE(arr) != npy_type<U>)) {
          if (raise_python_exception) PyErr_SetString(PyExc_TypeError, "Cannot convert to array_view : Type mismatch");
          return false;
        }
      }
 
      if (require_c_order and not numpy_check_layout<R, Layout>(obj)) {
        if (raise_python_exception) PyErr_SetString(PyExc_TypeError, "Cannot convert to array_view : Numpy array is not in C order");
        return false;
      }
 
      if constexpr (not std::is_const_v<T>) {
        if (not PyArray_ISWRITEABLE(arr)) {
          if (raise_python_exception) PyErr_SetString(PyExc_TypeError, "Cannot convert to mutable array_view : Numpy array is read-only");
          return false;
        }
      }
 
      return true;
    }
 
    // --------- PY -> C --------
 
    static view_t py2c(PyObject *obj) {
      auto p = make_numpy_proxy(obj);
      EXPECTS(p.extents.size() >= R);
      EXPECTS(p.element_type == npy_type<T> or p.extents.size() > R);
      EXPECTS((numpy_check_layout<R, Layout>(obj)));
 
      std::array<long, R> extents, strides;
      for (int u = 0; u < R; ++u) {
        extents[u] = p.extents[u];
        strides[u] = p.strides[u] / sizeof(T);
      }
      return view_t{{extents, strides}, static_cast<T *>(p.data)};
    }
  };
 
  // -----------------------------------
  // array
  // -----------------------------------
  template <typename T, int R, char Algebra>
  struct py_converter<nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>>> {
 
    // T can be a npy type c2py::has_npy_type<T> == true or NOT (then we need to convert using c2py)
    static_assert(not std::is_same_v<T, pyref>, "Not implemented");
    static_assert(not std::is_same_v<T, PyObject *>, "Not implemented");
 
    using array_t                 = nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>>;
    using view_t                  = nda::basic_array_view<T, R, nda::C_stride_layout, Algebra>;
    using converter_T             = py_converter<std::decay_t<T>>;
    using converter_view_T        = py_converter<view_t>;
    using converter_view_pyobject = py_converter<nda::basic_array_view<PyObject *, R, nda::C_stride_layout, Algebra>>;
 
    // ------------ tp_name ---------------
 
    static std::string tp_name() {
      std::ostringstream out;
      out << "ndarray[" << python_typename<T>() << ", " << R << "]";
      return out.str();
    }
 
    // --------- C -> PY --------
 
    template <typename A>
    static PyObject *c2py(A &&src) {
      static_assert(std::is_same_v<std::decay_t<A>, nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>>>,
                    "Logic Error in array c2py conversion");
      auto p = nda::python::make_numpy_proxy_from_array_or_view(std::forward<A>(src));
      return p.to_python();
    }
 
    // --------- PY -> C is possible --------
 
    static PyObject *make_numpy(PyObject *obj) {
      return PyArray_FromAny(obj, PyArray_DescrFromType(npy_type<T>), R, R, NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_ENSURECOPY, NULL); // new ref
    }
 
    static bool is_convertible(PyObject *obj, bool raise_python_exception) {
      // if obj is not an numpy, we try to make a numpy with the proper type
      if (not PyArray_Check(obj) or (PyArray_Check(obj) and has_npy_type<T> and (PyArray_TYPE((PyArrayObject *)(obj)) != npy_type<T>))) {
        c2py::pyref numpy_obj = make_numpy(obj);
        if (PyErr_Occurred()) {
          if (!raise_python_exception) PyErr_Clear();
          return false;
        }
        return is_convertible(numpy_obj, raise_python_exception);
      }
 
      if constexpr (has_npy_type<T>) {
        // in this case, we convert to a view and then copy to the array, so the condition is the same
        // use a const view converter to skip the writability check.
        using const_view_t         = nda::basic_array_view<const T, R, nda::C_stride_layout, Algebra>;
        return py_converter<const_view_t>::is_convertible(obj, raise_python_exception, false /*allow_lower_rank*/, false /*require_c_order*/);
      } else {
        // T is a type requiring conversion.
        // First I see whether I can convert it to a array of PyObject* ( i.e. it is an array of object and it has the proper rank...)
        bool res = converter_view_pyobject::is_convertible(obj, raise_python_exception, true /*allow_lower_rank*/);
        if (not res) return false;
 
        // Check if all elements are convertible
        // CAUTION: numpy array might be of higher rank!
        // We Extract the first R extents from the numpy proxy
        // and use __getitem__ to check convertibility of
        // each such element
        auto p = make_numpy_proxy(obj);
        std::array<long, R> shape;
        for (int i = 0; i < R; ++i) shape[i] = p.extents[i];
        auto l = [obj](auto... i) -> bool {
          pyref subobj = PyObject_GetItem(obj, pyref::make_tuple(PyLong_FromLong(i)...));
          return converter_T::is_convertible(subobj, false);
        };
        long n_elements = 1;
        for (int i = 0; i < R; ++i) n_elements *= shape[i];
        long n_convertible = sum(nda::array_adapter{shape, l});
 
        if (n_convertible != n_elements) {
          if (raise_python_exception) PyErr_SetString(PyExc_TypeError, "Cannot convert to array. One element can not be converted to C++.");
          return false;
        }
        return true;
      }
    }
 
    // --------- PY -> C --------
 
    static array_t py2c(PyObject *obj) {
 
      // if obj is not an numpy, we make a numpy and rerun
      if (not PyArray_Check(obj) or (PyArray_Check(obj) and has_npy_type<T> and (PyArray_TYPE((PyArrayObject *)(obj)) != npy_type<T>))) {
 
        c2py::pyref numpy_obj = make_numpy(obj);
        EXPECTS(not PyErr_Occurred());
        return py2c(numpy_obj);
      }
 
      if constexpr (has_npy_type<T>) {
        if (not numpy_check_layout<R, nda::C_layout>(obj) or not PyArray_ISWRITEABLE((PyArrayObject *)(obj))) {
          c2py::pyref obj_copy = make_numpy(obj);
          return array_t{converter_view_T::py2c(obj_copy)};
        }
        return converter_view_T::py2c(obj);
      } else {
        auto p = make_numpy_proxy(obj);
        std::array<long, R> shape;
        for (int i = 0; i < R; ++i) shape[i] = p.extents[i];
        auto l = [obj](auto... i) {
          pyref subobj = PyObject_GetItem(obj, pyref::make_tuple(PyLong_FromLong(i)...));
          return converter_T::py2c(subobj);
        };
        array_t res = nda::array_adapter{shape, l};
        if (PyErr_Occurred()) PyErr_Print();
        return res;
      }
    }
  };
 
  template <typename T, int R, char Algebra>
    requires(has_npy_type<std::decay_t<T>>)
  struct py_converter<nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>> const &> {
    using array_t = nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>>;
    static PyObject *c2py(array_t const &a, [[maybe_unused]] PyObject *guardian) {
      return cxx2py(nda::make_const_view(a));
    } // guardian is not used, as nda has its own shared handle which takes care of the ownership.
  };
 
  template <typename T, int R, char Algebra>
    requires(has_npy_type<std::decay_t<T>>)
  struct py_converter<nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>> &> {
    using array_t = nda::basic_array<T, R, nda::C_layout, Algebra, nda::heap<>>;
    using view_t  = nda::basic_array_view<T, R, nda::C_layout, Algebra>;
    static PyObject *c2py(array_t &a, [[maybe_unused]] PyObject *guardian) {
      return cxx2py(view_t(a));
      // guardian is not used, as nda has its own shared handle which takes care of the ownership.return arr;
    }
  };
 
  template <typename E>
    requires(nda::is_expression<std::decay_t<E>>)
  struct py_converter<E> {
    static PyObject *c2py(E const &ex) { return cxx2py(nda::make_regular(ex)); }
  };
 
} // namespace c2py