Refactor variable iteration in Khalil2016

Author: AntoinePrv

libecole/include/ecole/tweak/range.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <utility>
+
+#include <nonstd/span.hpp>
+#include <range/v3/range_fwd.hpp>
+
+/**
+ * Tell the range library that `nonstd::span` is a view type.
+ *
+ * See `Rvalue Ranges and Views in C++20 <https://tristanbrindle.com/posts/rvalue-ranges-and-views>`_
+ * FIXME no longer needed when switching to C++20 ``std::span``.
+ * */
+namespace ranges {
+template <typename T, std::size_t Extent> inline constexpr bool enable_borrowed_range<nonstd::span<T, Extent>> = true;
+}  // namespace ranges

libecole/src/observation/khalil-2016.cpp

@@ -551,15 +551,10 @@ void set_dynamic_features(
  * The static features have been computed for all LP columns and stored in the order of `LPcolumns`.
  * We need to find the one associated with the given variable.
  */
-template <typename Tensor>
-void set_precomputed_static_features(
-	Tensor&& out,
-	std::size_t var_idx,
-	xt::xtensor<value_type, 2> const& static_features) {
-
+template <typename TensorOut, typename TensorIn>
+void set_precomputed_static_features(TensorOut&& var_features, TensorIn const& var_static_features) {
 	using namespace xt::placeholders;
-	xt::view(out, xt::range(_, Khalil2016Obs::n_static_features)) =
-		xt::row(static_features, static_cast<std::ptrdiff_t>(var_idx));
+	xt::view(var_features, xt::range(_, Khalil2016Obs::n_static_features)) = var_static_features;
 }
 
 /******************************
@@ -568,18 +563,17 @@ void set_precomputed_static_features(
 
 auto extract_all_features(scip::Model& model, bool pseudo, xt::xtensor<value_type, 2> const& static_features) {
 	auto const branch_cands = pseudo ? model.pseudo_branch_cands() : model.lp_branch_cands();
-	auto const n_branch_cands = branch_cands.size();
-
-	auto observation = xt::xtensor<value_type, 2>{{n_branch_cands, Khalil2016Obs::n_features}, std::nan("")};
+	auto observation = xt::xtensor<value_type, 2>{{model.variables().size(), Khalil2016Obs::n_features}, std::nan("")};
 
 	auto* const scip = model.get_scip_ptr();
 	auto const active_rows_weights = stats_for_active_constraint_coefficients_weights(model);
 
-	for (std::size_t var_idx = 0; var_idx < n_branch_cands; ++var_idx) {
-		auto* const var = branch_cands[var_idx];
-		auto features = xt::row(observation, static_cast<std::ptrdiff_t>(var_idx));
-		set_precomputed_static_features(features, var_idx, static_features);
-		set_dynamic_features(features, scip, var, active_rows_weights);
+	for (auto* var : branch_cands) {
+		auto const var_idx = SCIPvarGetProbindex(var);
+		auto var_features = xt::row(observation, var_idx);
+		auto var_static_features = xt::row(static_features, var_idx);
+		set_precomputed_static_features(var_features, var_static_features);
+		set_dynamic_features(var_features, scip, var, active_rows_weights);
 	}
 
 	return observation;

libecole/tests/src/observation/test-khalil-2016.cpp

@@ -1,12 +1,20 @@
 #include <catch2/catch.hpp>
+#include <range/v3/view/enumerate.hpp>
+#include <range/v3/view/transform.hpp>
+#include <scip/scip.h>
+#include <xtensor/xindex_view.hpp>
 #include <xtensor/xmath.hpp>
+#include <xtensor/xtensor.hpp>
 #include <xtensor/xview.hpp>
 
 #include "ecole/observation/khalil-2016.hpp"
+#include "ecole/tweak/range.hpp"
 
 #include "conftest.hpp"
 #include "observation/unit-tests.hpp"
 
+namespace views = ranges::views;
+
 using namespace ecole;
 
 TEST_CASE("Khalil2016 unit tests", "[unit][obs]") {
@@ -20,6 +28,16 @@ auto in_interval(Tensor const& tensor, T const& lower, T const& upper) {
 	return (lower <= tensor) && (tensor <= upper);
 }
 
+/** Get the features of the pseudo candidate only. */
+template <typename Tensor, typename Range>
+auto obs_pseudo_cands(Tensor const& obs_features, Range const& pseudo_cands_idx) -> Tensor {
+	auto filtered_features = Tensor::from_shape({pseudo_cands_idx.size(), obs_features.shape()[1]});
+	for (auto const [idx, var_idx] : views::enumerate(pseudo_cands_idx)) {
+		xt::row(filtered_features, static_cast<std::ptrdiff_t>(idx)) = xt::row(obs_features, var_idx);
+	}
+	return filtered_features;
+}
+
 TEST_CASE("Khalil2016 return correct observation", "[obs]") {
 	using Features = observation::Khalil2016Obs::Features;
 
@@ -34,20 +52,23 @@ TEST_CASE("Khalil2016 return correct observation", "[obs]") {
 
 	SECTION("Observation features has correct shape") {
 		auto const& obs = optional_obs.value();
-		auto const branch_cands = pseudo ? model.pseudo_branch_cands() : model.lp_branch_cands();
-		REQUIRE(obs.features.shape(0) == branch_cands.size());
+		REQUIRE(obs.features.shape(0) == model.variables().size());
 		REQUIRE(obs.features.shape(1) == observation::Khalil2016Obs::n_features);
 	}
 
-	SECTION("No features are NaN or infinite") {
-		auto const& obs = optional_obs.value();
-		REQUIRE_FALSE(xt::any(xt::isnan(obs.features)));
-		REQUIRE_FALSE(xt::any(xt::isinf(obs.features)));
-	}
-
 	SECTION("Observation has correct values") {
 		auto const& obs = optional_obs.value();
-		auto col = [&obs](auto feat) { return xt::col(obs.features, static_cast<std::ptrdiff_t>(feat)); };
+		auto const branch_cands = pseudo ? model.pseudo_branch_cands() : model.lp_branch_cands();
+		auto obs_pseudo = obs_pseudo_cands(obs.features, views::transform(branch_cands, SCIPvarGetProbindex));
+		auto col = [&obs_pseudo](auto feat) { return xt::col(obs_pseudo, static_cast<std::ptrdiff_t>(feat)); };
+
+		SECTION("No pseudo_candidate features are NaN or infinite") {
+			for (auto* var : branch_cands) {
+				auto const var_idx = SCIPvarGetProbindex(var);
+				REQUIRE_FALSE(xt::any(xt::isnan(xt::row(obs.features, var_idx))));
+				REQUIRE_FALSE(xt::any(xt::isinf(xt::row(obs.features, var_idx))));
+			}
+		}
 
 		SECTION("Objective function coefficients") {
 			REQUIRE(xt::all(col(Features::obj_coef_pos_part) >= 0));

python/src/ecole/core/observation.cpp

@@ -263,8 +263,10 @@ void bind_submodule(py::module_ const& m) {
 		auto_class<Khalil2016Obs>(m, "Khalil2016Obs", R"(
 		Branching candidates features from Khalil et al. (2016).
 
-		The observation is a matrix where rows represent pseudo branching candidates and columns
-		represent features related to these variables.
+		The observation is a matrix where rows represent all variables and columns represent features related
+		to these variables.
+		Only rows representing pseudo branching candidate contain meaningful observation, other rows are filled with
+		``NaN``.
 		See [Khalil2016]_ for a complete reference on this observation function.
 
 		The first :py:attr:`Khalil2016Obs.n_static_features` are static (they do not change through the solving