#include
#include
#include
// ReSharper disable once CppUnusedIncludeDirective
#include
#include "fast_linear_assignment_sorter.hpp"
#include "metrics.hpp"
namespace py = pybind11;
std::tuple > flas(
const py::array_t &features,
const py::array_t &ids,
const py::array_t &frozen,
const bool wrap, float radius_decay, float weight_swappable, float weight_non_swappable, float weight_hole,
int max_swap_positions, int seed, int optimize_narrow_grids, const std::function& callback
) {
// ids
const py::buffer_info ids_info = ids.request();
if (ids_info.ndim != 2) {
const py::array_t tmp(0);
return std::make_tuple(1, tmp);
}
const int height = static_cast(ids_info.shape[0]);
const int width = static_cast(ids_info.shape[1]);
const int32_t* ids_ptr = static_cast(ids_info.ptr);
// features
const py::buffer_info features_info = features.request();
if (features_info.ndim != 2) {
const py::array_t tmp(0);
return std::make_tuple(2, tmp);
}
const int dim = static_cast(features_info.shape[1]);
const float* features_ptr = static_cast(features_info.ptr);
// frozen
const py::buffer_info frozen_info = frozen.request();
if (frozen_info.ndim != 2) {
const py::array_t result_indices(0);
return std::make_tuple(3, result_indices);
}
if (frozen_info.shape[0] != height || frozen_info.shape[1] != width) {
const py::array_t result_indices(0);
return std::make_tuple(4, result_indices);
}
const bool* frozen_ptr = static_cast(frozen_info.ptr);
// settings
const FlasSettings settings(
wrap, 0.5f, radius_decay, 1, 1.0f, weight_swappable, weight_non_swappable, weight_hole, 1.0f, max_swap_positions,
optimize_narrow_grids
);
// random
RandomEngine rng;
if (seed == -1)
rng.seed(std::random_device()());
else
rng.seed(seed);
// create map fields
auto map_fields = static_cast(malloc(height * width * sizeof(MapField)));
if (map_fields == nullptr) {
std::cerr << "Failed to allocate map_fields.\n" << std::endl;
exit(1);
}
for (int y = 0; y < height; y++) {
for (int x = 0; x < width; x++) {
const int index = x + y * static_cast(width);
int id = ids_ptr[index];
MapField *current_field = &map_fields[index];
const bool field_frozen = frozen_ptr[index];
if (id != -1) {
const float *const float_feature = features_ptr + (id * dim);
init_map_field(current_field, id, float_feature, !field_frozen);
} else {
init_invalid_map_field(current_field, !field_frozen);
}
}
}
// ----------------------------------------
do_sorting_full(
map_fields, static_cast(dim), static_cast(width), static_cast(height), &settings, &rng, callback
);
const py::array_t result_indices({height, width});
const py::buffer_info result_indices_info = result_indices.request();
const auto res_ptr = static_cast(result_indices_info.ptr);
// convert back --------------------------
for (int y = 0; y < height; y++) {
for (int x = 0; x < width; x++) {
const MapField map_field = map_fields[x + y * width];
res_ptr[y * width + x] = map_field.id;
}
}
free(map_fields);
// ---------------------------------------
return std::make_tuple(0, result_indices);
}
std::tuple > flas_no_callback(
const py::array_t &features,
const py::array_t &ids,
const py::array_t &frozen,
const bool wrap, float radius_decay, float weight_swappable, float weight_non_swappable, float weight_hole,
int max_swap_positions, int seed, int optimize_narrow_grids
) {
return flas(
features, ids, frozen, wrap, radius_decay, weight_swappable, weight_non_swappable, weight_hole, max_swap_positions,
seed, optimize_narrow_grids, [](float) { return false; }
);
}
std::tuple get_size(const size_t n_features, const float aspect_ratio) {
uint32_t height = static_cast(std::sqrt(static_cast(n_features) / aspect_ratio));
uint32_t width = static_cast(static_cast(height) * aspect_ratio);
// Adjust height and width to cover all features
while (height * width < n_features) {
if (static_cast(width) / static_cast(height) < aspect_ratio) {
width++;
} else {
height++;
}
}
return std::make_tuple(height, width);
}
std::tuple get_optimal_grid_size(
uint32_t total_num_features, float aspect_ratio, uint32_t min_height, uint32_t min_width
) {
while (total_num_features > min_height * min_width) {
float asp = static_cast(min_width) / static_cast(min_height);
if (asp == aspect_ratio) {
if (min_width < min_height)
min_width++;
else
min_height++;
} else if (asp < aspect_ratio) {
min_width++;
} else {
min_height++;
}
}
return std::make_tuple(min_height, min_width);
}
/**
* @param features The features array with shape [h, w, d].
* @param valid Boolean array defining valid (not holes) fields with shape [h, w]
* @param wrap Whether features can be wrapped around the plane.
* @return A tuple [error_code, num_distances, sum_distance], where
* - error_code indices an error, if != 0
* - num_distances, is the number of distances that where calculated to substitute holes
* - sum_distance, is the sum of all distances that where calculated to substitute holes
*/
std::tuple calc_hole_substitution_distance(
const py::array_t &features,
const py::array_t &valid,
const bool wrap
) {
// features
const py::buffer_info features_info = features.request();
if (features_info.ndim != 3) {
return std::make_tuple(1, 0, 0.0);
}
const int height = static_cast(features_info.shape[0]);
const int width = static_cast(features_info.shape[1]);
const int dim = static_cast(features_info.shape[2]);
const double* features_ptr = static_cast(features_info.ptr);
// valid
const py::buffer_info valid_info = valid.request();
if (valid_info.ndim != 2) {
return std::make_tuple(2, 0, 0.0);
}
if (valid_info.shape[0] != height || valid_info.shape[1] != width) {
return std::make_tuple(3, 0, 0.0); // wrong shape
}
const bool* valid_ptr = static_cast(valid_info.ptr);
// calculation
auto [num_dists, sum_dists] = calc_substitution_distance(height, width, dim, wrap, features_ptr, valid_ptr);
return std::make_tuple(0, num_dists, sum_dists);
}
/**
* @param features The features array with shape [h, w, d].
* @param wrap Whether features can be wrapped around the plane.
* @return A tuple [error_code, num_distances, sum_distance], where
* - error_code indices an error, if != 0
* - num_distances, is the number of distances that where calculated to substitute holes
* - sum_distance, is the sum of all distances that where calculated to substitute holes
*/
std::tuple calc_hole_substitution_distance_all_valid(
const py::array_t &features,
const bool wrap
) {
// features
const py::buffer_info features_info = features.request();
if (features_info.ndim != 3) {
return std::make_tuple(1, 0, 0.0);
}
const int height = static_cast(features_info.shape[0]);
const int width = static_cast(features_info.shape[1]);
const int dim = static_cast(features_info.shape[2]);
const double* features_ptr = static_cast(features_info.ptr);
// calculation
auto [num_dists, sum_dists] = calc_substitution_distance(height, width, dim, wrap, features_ptr, nullptr);
return std::make_tuple(0, num_dists, sum_dists);
}
PYBIND11_MODULE(flas_cpp, m) {
m.def("flas", &flas);
m.def("flas_no_callback", &flas_no_callback);
m.def("get_optimal_grid_size", &get_optimal_grid_size);
m.def("calc_hole_substitution_distance", &calc_hole_substitution_distance);
m.def("calc_hole_substitution_distance_all_valid", &calc_hole_substitution_distance_all_valid);
}