Program Listing for File vdb_likelihood_field_model.hpp

↰ Return to documentation for file (include/beluga_vdb/sensor/vdb_likelihood_field_model.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_VDB_SENSOR_VDB_LIKELIHOOD_FIELD_MODEL_HPP
#define BELUGA_VDB_SENSOR_VDB_LIKELIHOOD_FIELD_MODEL_HPP

#include <algorithm>
#include <cmath>
#include <random>
#include <vector>

#include <openvdb/openvdb.h>
#include <beluga/3d_embedding.hpp>

#include <Eigen/Core>

#include <range/v3/numeric/accumulate.hpp>
#include <range/v3/range/conversion.hpp>
#include <range/v3/view/all.hpp>
#include <range/v3/view/transform.hpp>
#include <sophus/se2.hpp>
#include <sophus/se3.hpp>
#include <sophus/so2.hpp>
#include <sophus/so3.hpp>

namespace beluga_vdb {


struct VDBLikelihoodFieldModelParam {

  double max_obstacle_distance = 100.0;
  double max_laser_distance = 2.0;
  double z_hit = 0.5;
  double z_random = 0.5;

  double sigma_hit = 0.2;
};


template <typename GridT, typename PointCloud, class StateType = Sophus::SE3d>
class VDBLikelihoodFieldModel {
  static_assert(
      std::is_same_v<StateType, Sophus::SE2d> or std::is_same_v<StateType, Sophus::SE3d>,
      "VDB likelihood field sensor model only supports SE2 and SE3 state types.");

 public:
  using state_type = StateType;

  using weight_type = double;

  using measurement_type = PointCloud;

  using map_type = GridT;

  using param_type = VDBLikelihoodFieldModelParam;


  explicit VDBLikelihoodFieldModel(const param_type& params, const map_type& grid)
      : params_{params},
        grid_{openvdb::gridPtrCast<map_type>(grid.deepCopyGrid())},
        accessor_{grid_->getConstAccessor()},
        transform_{grid_->transform()},
        background_{grid_->background()},
        two_squared_sigma_{2 * params.sigma_hit * params.sigma_hit},
        amplitude_{params.z_hit / (params.sigma_hit * std::sqrt(2 * Sophus::Constants<double>::pi()))},
        offset_{params.z_random / params.max_laser_distance} {
    openvdb::initialize();
    assert(two_squared_sigma_ > 0.0);
    assert(amplitude_ > 0.0);
  }


  [[nodiscard]] auto operator()(measurement_type&& measurement) const {
    // Transform each point from the sensor frame to the origin frame
    auto transformed_points = measurement.points() | ranges::views::transform([&](const auto& point) {
                                return measurement.origin() * point.template cast<double>();
                              }) |
                              ranges::to<std::vector>();
    return [this, points = std::move(transformed_points)](const state_type& state) -> weight_type {
      return ranges::accumulate(
          points |  //
              ranges::views::transform([this, state = beluga::To3d(state)](const auto& point) {
                const Eigen::Vector3d point_in_state_frame = state * point;
                const openvdb::math::Coord ijk = transform_.worldToIndexCellCentered(
                    openvdb::math::Vec3d(point_in_state_frame.x(), point_in_state_frame.y(), point_in_state_frame.z()));
                const auto distance = accessor_.isValueOn(ijk) ? accessor_.getValue(ijk) : background_;
                return amplitude_ * std::exp(-(distance * distance) / two_squared_sigma_) + offset_;
              }),
          1.0, std::plus{});
    };
  }

 private:
  param_type params_;
  const typename map_type::Ptr grid_;
  const typename map_type::ConstAccessor accessor_;
  const openvdb::math::Transform transform_;
  const typename map_type::ValueType background_;
  double two_squared_sigma_;
  double amplitude_;
  double offset_;
};

template <typename GridT, typename PointCloud>
using VDBLikelihoodFieldModel2 = VDBLikelihoodFieldModel<GridT, PointCloud, Sophus::SE2d>;

template <typename GridT, typename PointCloud>
using VDBLikelihoodFieldModel3 = VDBLikelihoodFieldModel<GridT, PointCloud, Sophus::SE3d>;

}  // namespace beluga_vdb

#endif