Fix common issues in custom types

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Fixed
 
 - Fix Python library linkage for Debug build on Windows ([#514](https://github.com/stack-of-tasks/eigenpy/pull/514))
+- Fix np.ones when dtype is a custom user type ([#525](https://github.com/stack-of-tasks/eigenpy/pull/525))
 
 ## [3.10.1] - 2024-10-30
 

include/eigenpy/ufunc.hpp

@@ -1,5 +1,5 @@
 //
-// Copyright (c) 2020-2021 INRIA
+// Copyright (c) 2020-2025 INRIA
 // code aptapted from
 // https://github.com/numpy/numpy/blob/41977b24ae011a51f64faa75cb524c7350fdedd9/numpy/core/src/umath/_rational_tests.c.src
 //
@@ -151,6 +151,7 @@ EIGENPY_REGISTER_BINARY_OPERATOR(greater_equal, >=)
   }
 
 EIGENPY_REGISTER_UNARY_OPERATOR(negative, -)
+EIGENPY_REGISTER_UNARY_OPERATOR(square, x *)
 
 }  // namespace internal
 
@@ -258,6 +259,7 @@ void registerCommonUfunc() {
 
   // Unary operators
   EIGENPY_REGISTER_UNARY_UFUNC(negative, type_code, Scalar, Scalar);
+  EIGENPY_REGISTER_UNARY_UFUNC(square, type_code, Scalar, Scalar);
 
   Py_DECREF(numpy);
 }

include/eigenpy/user-type.hpp

@@ -132,34 +132,43 @@ struct SpecialMethods<T, NPY_USERDEF> {
       eigenpy::Exception("Cannot retrieve the type stored in the array.");
       return -1;
     }
+
     PyArrayObject* py_array = static_cast<PyArrayObject*>(array);
     PyArray_Descr* descr = PyArray_DTYPE(py_array);
     PyTypeObject* array_scalar_type = descr->typeobj;
     PyTypeObject* src_obj_type = Py_TYPE(src_obj);
 
+    T& dest = *static_cast<T*>(dest_ptr);
     if (array_scalar_type != src_obj_type) {
-      std::stringstream ss;
-      ss << "The input type is of wrong type. ";
-      ss << "The expected type is " << bp::type_info(typeid(T)).name()
-         << std::endl;
-      eigenpy::Exception(ss.str());
-      return -1;
-    }
+      long long src_value = PyLong_AsLongLong(src_obj);
+      if (src_value == -1 && PyErr_Occurred()) {
+        std::stringstream ss;
+        ss << "The input type is of wrong type. ";
+        ss << "The expected type is " << bp::type_info(typeid(T)).name()
+           << std::endl;
+        eigenpy::Exception(ss.str());
+        return -1;
+      }
+
+      dest = T(src_value);
 
-    bp::extract<T&> extract_src_obj(src_obj);
-    if (!extract_src_obj.check()) {
-      std::stringstream ss;
-      ss << "The input type is of wrong type. ";
-      ss << "The expected type is " << bp::type_info(typeid(T)).name()
-         << std::endl;
-      eigenpy::Exception(ss.str());
-      return -1;
+    } else {
+      bp::extract<T&> extract_src_obj(src_obj);
+      if (!extract_src_obj.check()) {
+        std::cout << "if (!extract_src_obj.check())" << std::endl;
+        std::stringstream ss;
+        ss << "The input type is of wrong type. ";
+        ss << "The expected type is " << bp::type_info(typeid(T)).name()
+           << std::endl;
+        eigenpy::Exception(ss.str());
+        return -1;
+      }
+
+      const T& src = extract_src_obj();
+      T& dest = *static_cast<T*>(dest_ptr);
+      dest = src;
     }
 
-    const T& src = extract_src_obj();
-    T& dest = *static_cast<T*>(dest_ptr);
-    dest = src;
-
     return 0;
   }
 

src/register.cpp

@@ -64,7 +64,6 @@ int Register::registerNewType(
   }
 
   PyArray_DescrProto* descr_ptr = new PyArray_DescrProto();
-  Py_SET_TYPE(descr_ptr, &PyArrayDescr_Type);
   PyArray_DescrProto& descr = *descr_ptr;
   descr.typeobj = py_type_ptr;
   descr.kind = 'V';
@@ -92,6 +91,7 @@ int Register::registerNewType(
   funcs.fill = fill;
   funcs.fillwithscalar = fillwithscalar;
   //      f->cast = cast;
+  Py_SET_TYPE(descr_ptr, &PyArrayDescr_Type);
 
   const int code = call_PyArray_RegisterDataType(descr_ptr);
   assert(code >= 0 && "The return code should be positive");

unittest/python/test_user_type.py

@@ -9,7 +9,7 @@
 
 def test(dtype):
     rng = np.random.default_rng()
-    mat = np.array(np.ones((rows, cols)).astype(np.int32), dtype=dtype)
+    mat = np.ones((rows, cols), dtype=dtype)
     mat = rng.random((rows, cols)).astype(dtype)
     mat_copy = mat.copy()
     assert (mat == mat_copy).all()
@@ -44,6 +44,11 @@ def test(dtype):
         mat2 = np.matmul(mat, mat.T)
         assert np.isclose(mat2.astype(np.double), mat2_ref).all()
 
+    vec = np.ones((rows,), dtype=dtype)
+    norm = np.linalg.norm(vec)
+    norm_ref = np.linalg.norm(vec.astype(np.double))
+    assert norm == norm_ref
+
 
 def test_cast(from_dtype, to_dtype):
     np.can_cast(from_dtype, to_dtype)
@@ -63,8 +68,17 @@ def test_cast(from_dtype, to_dtype):
 test_cast(user_type.CustomDouble, np.int32)
 test_cast(np.int32, user_type.CustomDouble)
 
-test(user_type.CustomFloat)
-
 v = user_type.CustomDouble(1)
 a = np.array(v)
 assert type(v) is a.dtype.type
+
+test(user_type.CustomFloat)
+
+test_cast(user_type.CustomFloat, np.float32)
+test_cast(np.double, user_type.CustomFloat)
+
+test_cast(user_type.CustomFloat, np.int64)
+test_cast(np.int64, user_type.CustomFloat)
+
+test_cast(user_type.CustomFloat, np.int32)
+test_cast(np.int32, user_type.CustomFloat)

unittest/user_type.cpp

@@ -201,14 +201,19 @@ BOOST_PYTHON_MODULE(user_type) {
 
   eigenpy::registerCast<DoubleType, double>(true);
   eigenpy::registerCast<double, DoubleType>(true);
+  eigenpy::registerCast<DoubleType, float>(false);
+  eigenpy::registerCast<float, DoubleType>(true);
   eigenpy::registerCast<DoubleType, int>(false);
   eigenpy::registerCast<int, DoubleType>(true);
   eigenpy::registerCast<DoubleType, long long>(false);
   eigenpy::registerCast<long long, DoubleType>(true);
   eigenpy::registerCast<DoubleType, long>(false);
   eigenpy::registerCast<long, DoubleType>(true);
+
   eigenpy::registerCast<FloatType, double>(true);
   eigenpy::registerCast<double, FloatType>(false);
+  eigenpy::registerCast<FloatType, float>(true);
+  eigenpy::registerCast<float, FloatType>(true);
   eigenpy::registerCast<FloatType, long long>(false);
   eigenpy::registerCast<long long, FloatType>(true);
   eigenpy::registerCast<FloatType, int>(false);