Program Listing for File multivariate_uniform_distribution.hpp

↰ Return to documentation for file (include/beluga/random/multivariate_uniform_distribution.hpp)

// Copyright 2024 Ekumen, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef BELUGA_RANDOM_MULTIVARIATE_UNIFORM_DISTRIBUTION_HPP
#define BELUGA_RANDOM_MULTIVARIATE_UNIFORM_DISTRIBUTION_HPP

#include <random>

#include <sophus/se2.hpp>
#include <sophus/se3.hpp>

#include <Eigen/Geometry>

#include <beluga/sensor/data/occupancy_grid.hpp>

namespace beluga {


template <class T, class Constraint>
class MultivariateUniformDistribution;

template <>
class MultivariateUniformDistribution<Sophus::SE2d, Eigen::AlignedBox2d> {
 public:

  explicit MultivariateUniformDistribution(const Eigen::AlignedBox2d& box)
      : x_distribution_{box.min().x(), box.max().x()}, y_distribution_{box.min().y(), box.max().y()} {}


  template <class URNG>
  [[nodiscard]] Sophus::SE2d operator()(URNG& engine) {
    return Sophus::SE2d{
        Sophus::SO2d::sampleUniform(engine),
        Sophus::Vector2d{
            x_distribution_(engine),
            y_distribution_(engine),
        },
    };
  }

 private:
  std::uniform_real_distribution<double> x_distribution_;
  std::uniform_real_distribution<double> y_distribution_;
};

MultivariateUniformDistribution(const Eigen::AlignedBox2d&)
    -> MultivariateUniformDistribution<Sophus::SE2d, Eigen::AlignedBox2d>;

template <>
class MultivariateUniformDistribution<Sophus::SE3d, Eigen::AlignedBox3d> {
 public:

  explicit MultivariateUniformDistribution(const Eigen::AlignedBox3d& box)
      : x_distribution_{box.min().x(), box.max().x()},
        y_distribution_{box.min().y(), box.max().y()},
        z_distribution_{box.min().z(), box.max().z()} {}


  template <class URNG>
  [[nodiscard]] Sophus::SE3d operator()(URNG& engine) {
    return Sophus::SE3d{
        Sophus::SO3d::sampleUniform(engine),
        Sophus::Vector3d{
            x_distribution_(engine),
            y_distribution_(engine),
            z_distribution_(engine),
        },
    };
  }

 private:
  std::uniform_real_distribution<double> x_distribution_;
  std::uniform_real_distribution<double> y_distribution_;
  std::uniform_real_distribution<double> z_distribution_;
};

MultivariateUniformDistribution(const Eigen::AlignedBox3d&)
    -> MultivariateUniformDistribution<Sophus::SE3d, Eigen::AlignedBox3d>;


template <class OccupancyGrid>
class MultivariateUniformDistribution<Sophus::SE2d, OccupancyGrid> {
 public:

  constexpr explicit MultivariateUniformDistribution(const OccupancyGrid& grid)
      : free_states_{compute_free_states(grid)}, distribution_{0, free_states_.size() - 1} {
    assert(!free_states_.empty());
  }


  template <class URNG>
  [[nodiscard]] Sophus::SE2d operator()(URNG& engine) {
    return {Sophus::SO2d::sampleUniform(engine), free_states_[distribution_(engine)]};
  }

 private:
  std::vector<Eigen::Vector2d> free_states_;
  std::uniform_int_distribution<std::size_t> distribution_;

  static std::vector<Eigen::Vector2d> compute_free_states(const OccupancyGrid& grid) {
    return grid.coordinates_for(grid.free_cells(), OccupancyGrid::Frame::kGlobal) | ranges::to<std::vector>;
  }
};

template <class Derived>
MultivariateUniformDistribution(const BaseOccupancyGrid2<Derived>&)
    -> MultivariateUniformDistribution<Sophus::SE2d, Derived>;

}  // namespace beluga

#endif