Refactor Khalil2016 static feature with strong typed features

libecole/src/observation/khalil-2016.cpp

@@ -1,12 +1,14 @@
 #include <algorithm>
-#include <array>
 #include <cmath>
 #include <limits>
 #include <set>
+#include <tuple>
 #include <type_traits>
 #include <utility>
 
 #include <nonstd/span.hpp>
+#include <xtensor/xfixed.hpp>
+#include <xtensor/xview.hpp>
 
 #include "ecole/observation/khalil-2016.hpp"
 #include "ecole/scip/model.hpp"
@@ -17,6 +19,9 @@ namespace observation {
 
 namespace {
 
+using Feature = Khalil2016::Feature;
+using Static = Khalil2016::Feature::Static;
+using Dynamic = Khalil2016::Feature::Dynamic;
 using value_type = Khalil2016Obs::value_type;
 
 /*************************
@@ -205,45 +210,117 @@ auto scip_row_get_vals(scip::Row* const row) noexcept -> nonstd::span<scip::real
 	return {SCIProwGetVals(row), static_cast<std::size_t>(n_cols)};
 }
 
-/***********************************
- *  Features extraction functions  *
- ***********************************/
+/****************************
+ *  Feature data structure  *
+ ****************************/
+
+/**
+ * Strongly type feature value.
+ *
+ * The value of a feature is wrapped in a type that encode the name of the feature.
+ *
+ * This make it possible to have strong guarantees that features are not accidentally read/written
+ * in the wrong column while minimizing overhead (because the name is known at compile time);
+ */
+template <typename F, F feature_name> struct FeatureValue {
+	constexpr static auto name = feature_name;
+	value_type value;
+};
+
+template <Feature::Static feature_name> auto feature(value_type val) noexcept {
+	return FeatureValue<Feature::Static, feature_name>{val};
+}
+template <Feature::Dynamic feature_name> auto feature(value_type val) noexcept {
+	return FeatureValue<Feature::Dynamic, feature_name>{val};
+}
+
+/************************************
+ *  Feature compile time functions  *
+ ************************************/
+
+template <typename... Features, std::size_t... I>
+constexpr auto
+features_tuple_to_tensor_impl(std::tuple<Features...> const& features, std::index_sequence<I...>) {
+	constexpr auto n_features = sizeof...(Features);
+	return xt::xtensor_fixed<value_type, xt::xshape<n_features>>{std::get<I>(features).value...};
+}
+
+/**
+ * Compile time convertion of a tuple of Feature to a xt::xfixed_tensor.
+ *
+ * The type of features is a @ref FeatureValue that wraps a value with its associated name.
+ */
+template <typename... Features>
+constexpr auto features_to_tensor(std::tuple<Features...> const& features) {
+	constexpr auto n_features = sizeof...(Features);
+	return features_tuple_to_tensor_impl(features, std::make_index_sequence<n_features>{});
+}
+
+/**
+ * Check that a tuple of feature is contiguous.
+ *
+ * With this compile time check other functions can bypass individually setting each feature by its
+ * index and directy assign a whole tensor.
+ */
+template <typename... Features> constexpr auto is_contiguous(std::tuple<Features...>) {
+	constexpr auto n_features = sizeof...(Features);
+	constexpr std::array<std::size_t, n_features> arr = {static_cast<std::size_t>(Features::name)...};
+
+	for (std::size_t i = 0; i < arr.size(); ++i) {
+		if ((i < 0) && (arr[i] != arr[i - 1] + 1)) {
+			return false;
+		}
+	}
+	return true;
+}
+
+/**
+ * Return at compile time the feature index of the first feature in a tuple.
+ */
+template <typename... Features, typename Tuple = std::tuple<Features...>>
+constexpr auto first_index(Tuple) {
+	return static_cast<std::size_t>(std::tuple_element_t<0, Tuple>::name);
+}
+
+/**
+ * Return at compile time the feature index of the last feature in a tuple.
+ */
+template <typename... Features, typename Tuple = std::tuple<Features...>>
+constexpr auto last_index(Tuple) {
+	constexpr auto last_index = std::tuple_size<Tuple>::value - 1;
+	return static_cast<std::size_t>(std::tuple_element_t<last_index, Tuple>::name);
+}
+
+/******************************************
+ *  Static features extraction functions  *
+ ******************************************/
 
 /* Feature as defined and split in table 1 of the paper Khalil et al. "Learning to Branch in Mixed
  * Integer Programming" Thirtieth AAAI Conference on Artificial Intelligence. 2016.
  *
  * https://web.archive.org/web/20200812151256/https://www.cc.gatech.edu/~lsong/papers/KhaLebSonNemDil16.pdf
  */
 
-template <std::size_t N> using Features = std::array<value_type, N>;
-
 /**
  * Objective function coeffs.
  *
  * Value of the coefficient (raw, positive only, negative only).
  */
-auto objective_function_coefficient(scip::Col* const col) noexcept -> Features<3> {
+auto objective_function_coefficient(scip::Col* const col) noexcept {
 	auto const obj = SCIPcolGetObj(col);
-	return {obj, std::max(obj, 0.), std::min(obj, 0.)};
+	return std::make_tuple(
+		feature<Static::obj_coef>(obj),
+		feature<Static::obj_coef_pos_part>(std::max(obj, 0.)),
+		feature<Static::obj_coef_neg_part>(std::min(obj, 0.)));
 }
 
 /**
  * Num. constraints.
  *
  * Number of constraints that the variable participates in (with a non-zero coefficient).
  */
-auto number_constraints(scip::Col* const col) noexcept -> Features<1> {
-	return {static_cast<value_type>(SCIPcolGetNNonz(col))};
-}
-
-/**
- * Stats. for constraint degrees helper function for caching.
- */
-auto static_stats_for_constraint_degree_stats(nonstd::span<scip::Row*> const rows) noexcept
-	-> StatsFeatures {
-	auto transform = [](auto const row) { return static_cast<std::size_t>(SCIProwGetNNonz(row)); };
-	auto filter = [](auto const /* degree */) { return true; };
-	return compute_stats(rows, transform, filter);
+auto number_constraints(scip::Col* const col) noexcept {
+	return std::make_tuple(feature<Static::n_rows>(static_cast<value_type>(SCIPcolGetNNonz(col))));
 }
 
 /**
@@ -253,12 +330,16 @@ auto static_stats_for_constraint_degree_stats(nonstd::span<scip::Row*> const row
  * A variable may participate in multiple constraints, and statistics over those constraints'
  * degrees are used.
  * The constraint degree is computed on the root LP (mean, stdev., min, max)
- *
- * @param root_stats the output of @ref static_stats_for_constraint_degree_stats.
- *        It is extracted in another function because the result is reused elsewhere.
  */
-auto static_stats_for_constraint_degree(StatsFeatures const& root_stats) noexcept -> Features<4> {
-	return {root_stats.mean, root_stats.variance, root_stats.min, root_stats.max};
+auto static_stats_for_constraint_degree(nonstd::span<scip::Row*> const rows) noexcept {
+	auto transform = [](auto const row) { return static_cast<std::size_t>(SCIProwGetNNonz(row)); };
+	auto filter = [](auto const /* degree */) { return true; };
+	auto const stats = compute_stats(rows, transform, filter);
+	return std::make_tuple(
+		feature<Static::rows_deg_mean>(stats.mean),
+		feature<Static::rows_deg_var>(stats.variance),
+		feature<Static::rows_deg_min>(stats.min),
+		feature<Static::rows_deg_max>(stats.max));
 }
 
 /**
@@ -268,9 +349,14 @@ auto static_stats_for_constraint_degree(StatsFeatures const& root_stats) noexcep
  * (count, mean, stdev., min, max).
  */
 auto stats_for_constraint_positive_coefficients(
-	nonstd::span<scip::real> const coefficients) noexcept -> Features<5> {
+	nonstd::span<scip::real> const coefficients) noexcept {
 	auto const stats = compute_stats(coefficients, identity, greater_than(0.));
-	return {stats.count, stats.mean, stats.variance, stats.min, stats.max};
+	return std::make_tuple(
+		feature<Static::rows_pos_coefs_count>(stats.count),
+		feature<Static::rows_pos_coefs_mean>(stats.mean),
+		feature<Static::rows_pos_coefs_var>(stats.variance),
+		feature<Static::rows_pos_coefs_min>(stats.min),
+		feature<Static::rows_pos_coefs_max>(stats.max));
 }
 
 /**
@@ -280,11 +366,70 @@ auto stats_for_constraint_positive_coefficients(
  * (count, mean, stdev., min, max).
  */
 auto stats_for_constraint_negative_coefficients(
-	nonstd::span<scip::real> const coefficients) noexcept -> Features<5> {
+	nonstd::span<scip::real> const coefficients) noexcept {
 	auto const stats = compute_stats(coefficients, identity, lesser_than(0.));
-	return {stats.count, stats.mean, stats.variance, stats.min, stats.max};
+	return std::make_tuple(
+		feature<Static::rows_neg_coefs_count>(stats.count),
+		feature<Static::rows_neg_coefs_mean>(stats.mean),
+		feature<Static::rows_neg_coefs_var>(stats.variance),
+		feature<Static::rows_neg_coefs_min>(stats.min),
+		feature<Static::rows_neg_coefs_max>(stats.max));
+}
+
+/**
+ * Extract the static features for a single columns.
+ */
+auto extract_static_features(scip::Col* const col) {
+	auto const rows = scip_col_get_rows(col);
+	auto const coefficients = scip_col_get_vals(col);
+
+	auto features = std::tuple_cat(
+		objective_function_coefficient(col),
+		number_constraints(col),
+		static_stats_for_constraint_degree(rows),
+		stats_for_constraint_positive_coefficients(coefficients),
+		stats_for_constraint_negative_coefficients(coefficients)  //
+	);
+
+	// Make sure at compile time that feature are retuned in correct order
+	using Tuple = decltype(features);
+	static_assert(is_contiguous(Tuple{}), "Static features are permuted");
+	static_assert(first_index(Tuple{}) == 0, "Static features must start at 0");
+	static_assert(last_index(Tuple{}) == Feature::n_static - 1, "Missing static features");
+
+	return features_to_tensor(features);
 }
 
+/**
+ * Extract the static features for all LP columns in a Model.
+ */
+auto extract_static_features(scip::Model const& model) {
+	auto const columns = model.lp_columns();
+	xt::xtensor<value_type, 2> static_features{{columns.size(), Feature::n_static}, 0.};
+
+	// Similar to the following but slice iteration not working on xt::xtensor
+	// std::transform(columns, xt::axis_slice_begin(features, 1), return extract_static_features)
+	// https://github.com/xtensor-stack/xtensor/issues/2116
+	auto const n_columns = columns.size();
+	for (std::size_t i = 0; i < n_columns; ++i) {
+		xt::row(static_features, static_cast<std::ptrdiff_t>(i)) = extract_static_features(columns[i]);
+	}
+
+	return static_features;
+}
+
+/*******************************************
+ *  Dynamic features extraction functions  *
+ *******************************************/
+
+/* Feature as defined and split in table 1 of the paper Khalil et al. "Learning to Branch in Mixed
+ * Integer Programming" Thirtieth AAAI Conference on Artificial Intelligence. 2016.
+ *
+ * https://web.archive.org/web/20200812151256/https://www.cc.gatech.edu/~lsong/papers/KhaLebSonNemDil16.pdf
+ */
+
+template <std::size_t N> using Features = std::array<value_type, N>;
+
 /**
  * Slack and ceil distances.
  *
@@ -407,8 +552,8 @@ auto min_max_for_ratios_constraint_coeffs_rhs(
  *
  * The statistics are over the ratios of a variable's coefficient, to the sum over all other
  * variables' coefficients, for a given constraint.
- * Four versions of these ratios are considered: positive (negative) coefficient to sum of positive
- * (negative) coefficients.
+ * Four versions of these ratios are considered: positive (negative) coefficient to sum of
+ * positive (negative) coefficients.
  */
 auto min_max_for_one_to_all_coefficient_ratios(
 	nonstd::span<scip::Row*> const rows,
@@ -526,13 +671,12 @@ auto stats_for_active_constraint_coefficients_weights(scip::Model const& model)
  * Stats. for active constraint coefficients.
  *
  * An active constraint at a node LP is one which is binding with equality at the optimum.
- * We consider 4 weighting schemes for an active constraint: unit weight, inverse of the sum of the
- * coefficients of all variables in constraint, inverse of the sum of the coefficients of only
+ * We consider 4 weighting schemes for an active constraint: unit weight, inverse of the sum of
+ * the coefficients of all variables in constraint, inverse of the sum of the coefficients of only
  * candidate variables in constraint, dual cost of the constraint.
- * Given the absolute value of the coefficients of xj in the active constraints, we compute the sum,
- * mean, stdev., max. and min. of those values, for each of the weighting schemes.
- * We also compute the weighted number of active constraints that xj is in, with the same 4
- * weightings.
+ * Given the absolute value of the coefficients of xj in the active constraints, we compute the
+ * sum, mean, stdev., max. and min. of those values, for each of the weighting schemes. We also
+ * compute the weighted number of active constraints that xj is in, with the same 4 weightings.
  */
 auto stats_for_active_constraint_coefficients(
 	Scip* const scip,
@@ -610,11 +754,14 @@ auto stats_for_active_constraint_coefficients(
  *  Main extraction function  *
  ******************************/
 
-auto extract_feat(scip::Model const& model) {
-	std::size_t constexpr n_features = 72UL;
-	std::size_t const n_candidate_vars = model.pseudo_branch_cands().size();
+auto extract_features(scip::Model const& model, xt::xtensor<value_type, 2> const& static_features) {
+	using namespace xt::placeholders;
+	using Feature = Khalil2016::Feature;
 
-	xt::xtensor<value_type, 2> observation{{n_candidate_vars, n_features}, std::nan("")};
+	xt::xtensor<value_type, 2> observation{
+		{model.pseudo_branch_cands().size(), Feature::n_static + Feature::n_dynamic},
+		std::nan(""),
+	};
 
 	auto const scip = model.get_scip_ptr();
 	auto const active_rows_weights = stats_for_active_constraint_coefficients_weights(model);
@@ -626,11 +773,8 @@ auto extract_feat(scip::Model const& model) {
 		auto const coefficients = scip_col_get_vals(col);
 		auto const root_stats = static_stats_for_constraint_degree_stats(rows);
 		// Static features
-		iter = copy(objective_function_coefficient(col), iter);
-		iter = copy(number_constraints(col), iter);
-		iter = copy(static_stats_for_constraint_degree(root_stats), iter);
-		iter = copy(stats_for_constraint_positive_coefficients(coefficients), iter);
-		iter = copy(stats_for_constraint_negative_coefficients(coefficients), iter);
+		// FIXME
+		iter += Khalil2016::Feature::n_static;
 		// Dynamic features
 		iter = copy(slack_ceil_distances(scip, col), iter);
 		iter = copy(pseudocosts(scip, var), iter);
@@ -644,7 +788,7 @@ auto extract_feat(scip::Model const& model) {
 	}
 
 	// Make sure we iterated over as many element as there are in the tensor
-	assert(static_cast<std::size_t>(iter - observation.begin()) == observation.size());
+	assert(iter == observation.end());
 
 	return observation;
 }
@@ -655,12 +799,16 @@ auto extract_feat(scip::Model const& model) {
  *  Observation extracting function  *
  *************************************/
 
+void Khalil2016::reset(scip::Model& model) {
+	static_features = extract_static_features(model);
+}
+
 auto Khalil2016::obtain_observation(scip::Model& model) -> nonstd::optional<Khalil2016Obs> {
 	if (model.get_stage() == SCIP_STAGE_SOLVING) {
-		return extract_feat(model);
+		return extract_features(model, static_features);
 	}
 	return {};
 }
 
 }  // namespace observation
-}  // namespace ecole
+};  // namespace ecole