Merge pull request #179 from bluescarni/pr/fix_wrong_check

Fix wrong input size check

CMakeLists.txt

@@ -11,7 +11,7 @@ if(NOT CMAKE_BUILD_TYPE)
 	FORCE)
 endif()
 
-project(heyoka.py VERSION 5.0.0 LANGUAGES CXX C)
+project(heyoka.py VERSION 5.0.1 LANGUAGES CXX C)
 
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake" "${CMAKE_CURRENT_SOURCE_DIR}/cmake/yacma")
 

doc/changelog.rst

@@ -3,6 +3,15 @@
 Changelog
 =========
 
+5.0.1 (2024-06-14)
+------------------
+
+Fix
+~~~
+
+- Fix an input size check that would wrongly throw on valid code
+  (`#179 <https://github.com/bluescarni/heyoka.py/pull/179>`__).
+
 5.0.0 (2024-06-13)
 ------------------
 

heyoka/_test_var_integrator.py

@@ -234,3 +234,54 @@ def test_batch(self):
             "The array of inputs provided for the evaluation of a Taylor map has 1 column(s), but it must have 2 column(s) instead"
             in str(cm.exception)
         )
+
+    def test_size_check_bug(self):
+        # BUG: wrong size check on the input for a Taylor map (the size is checked against
+        # the number of original state variables instead of the number of variational arguments).
+
+        from . import (
+            make_vars,
+            var_ode_sys,
+            var_args,
+            cos,
+            sin,
+            par,
+            time,
+            taylor_adaptive_batch,
+            taylor_adaptive,
+            core,
+        )
+        import numpy as np
+
+        x, v = make_vars("x", "v")
+
+        orig_sys = [(x, v), (v, cos(time) - par[0] * v - sin(x))]
+
+        vsys = var_ode_sys(orig_sys, var_args.vars | var_args.params, order=2)
+
+        ta = taylor_adaptive(
+            vsys,
+            [0.2, 0.3],
+            pars=[0.4],
+            time=0.5,
+            compact_mode=True,
+        )
+
+        # This would throw.
+        self.assertTrue(np.all(ta.state[:2] == ta.eval_taylor_map([0.0, 0.0, 0.0])))
+
+        # Test the batch case too.
+        ta = taylor_adaptive_batch(
+            vsys,
+            [[0.2, 0.21], [0.3, 0.31]],
+            pars=[[0.4, 0.41]],
+            time=[0.5, 0.51],
+            compact_mode=True,
+        )
+
+        # This would throw.
+        self.assertTrue(
+            np.all(
+                ta.state[:2] == ta.eval_taylor_map([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
+            )
+        )

heyoka/expose_batch_integrators.cpp

@@ -614,11 +614,11 @@ void expose_batch_integrator_impl(py::module_ &m, const std::string &suffix)
                  }
 
                  // Validate the shape for the inputs.
-                 if (boost::numeric_cast<std::uint32_t>(inputs.shape(0)) != ta->get_n_orig_sv()) [[unlikely]] {
+                 if (boost::numeric_cast<std::uint32_t>(inputs.shape(0)) != ta->get_vargs().size()) [[unlikely]] {
                      py_throw(PyExc_ValueError, fmt::format("The array of inputs provided for the evaluation "
                                                             "of a Taylor map has {} row(s), "
                                                             "but it must have {} row(s) instead",
-                                                            inputs.shape(0), ta->get_n_orig_sv())
+                                                            inputs.shape(0), ta->get_vargs().size())
                                                     .c_str());
                  }
 

heyoka/taylor_expose_integrator.cpp

@@ -448,11 +448,11 @@ void expose_taylor_integrator_impl(py::module &m, const std::string &suffix)
                  }
 
                  // Validate the shape for the inputs.
-                 if (boost::numeric_cast<std::uint32_t>(inputs.shape(0)) != ta->get_n_orig_sv()) [[unlikely]] {
+                 if (boost::numeric_cast<std::uint32_t>(inputs.shape(0)) != ta->get_vargs().size()) [[unlikely]] {
                      py_throw(PyExc_ValueError, fmt::format("The array of inputs provided for the evaluation "
                                                             "of a Taylor map has {} elements, "
                                                             "but it must have {} elements instead",
-                                                            inputs.shape(0), ta->get_n_orig_sv())
+                                                            inputs.shape(0), ta->get_vargs().size())
                                                     .c_str());
                  }