kd_tree_environment.h

Go to the documentation of this file.

// -----------------------------------------------------------------------------
//
// Copyright (C) 2021 CERN & University of Surrey for the benefit of the
// BioDynaMo collaboration. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
//
// See the LICENSE file distributed with this work for details.
// See the NOTICE file distributed with this work for additional information
// regarding copyright ownership.
//
// -----------------------------------------------------------------------------
 
#ifndef CORE_ENVIRONMENT_KD_TREE_ENVIRONMENT_
#define CORE_ENVIRONMENT_KD_TREE_ENVIRONMENT_
 
#include "core/container/agent_flat_idx_map.h"
#include "core/container/math_array.h"
#include "core/environment/environment.h"
#include "core/simulation.h"
 
namespace bdm {
 
struct NanoFlannAdapter {
  using coord_t = Real3;
  using idx_t = uint64_t;
 
  NanoFlannAdapter() { rm_ = Simulation::GetActive()->GetResourceManager(); }
 
  inline size_t kdtree_get_point_count() const {
    return Simulation::GetActive()->GetResourceManager()->GetNumAgents();
  }
 
  inline real_t kdtree_distance(const coord_t& p1, const idx_t idx_p2,
                                size_t /*size*/) const {
    AgentHandle ah = flat_idx_map_.GetAgentHandle(idx_p2);
    auto& p2 = rm_->GetAgent(ah)->GetPosition();
    return p1 * p2;
  }
 
  inline real_t kdtree_get_pt(const idx_t idx, int dim) const {
    AgentHandle ah = flat_idx_map_.GetAgentHandle(idx);
    auto& pos = rm_->GetAgent(ah)->GetPosition();
    return pos[dim];
  }
 
  template <class BBOX>
  bool kdtree_get_bbox(BBOX&) const {
    return false;
  }
 
  AgentFlatIdxMap flat_idx_map_;
  ResourceManager* rm_ = nullptr;
};
 
class KDTreeEnvironment : public Environment {
 public:
  struct NanoflannImpl;
 
  KDTreeEnvironment();
 
  ~KDTreeEnvironment() override;
 
  std::array<int32_t, 6> GetDimensions() const override;
 
  std::array<int32_t, 2> GetDimensionThresholds() const override;
 
  LoadBalanceInfo* GetLoadBalanceInfo() override;
 
  NeighborMutexBuilder* GetNeighborMutexBuilder() override;
 
  void Clear() override;
 
  void ForEachNeighbor(Functor<void, Agent*, real_t>& lambda,
                       const Agent& query, real_t squared_radius) override;
 
  void ForEachNeighbor(Functor<void, Agent*>& lambda, const Agent& query,
                       void* criteria) override;
 
  void ForEachNeighbor(Functor<void, Agent*, real_t>& lambda,
                       const Real3& query_position, real_t squared_radius,
                       const Agent* query_agent = nullptr) override;
 
 protected:
  void UpdateImplementation() override;
 
 private:
  // Hide nanoflann-specific types from header (pimpl idiom)
  std::unique_ptr<NanoflannImpl> impl_;
  std::array<int32_t, 6> grid_dimensions_;
  std::array<int32_t, 2> threshold_dimensions_;
  NanoFlannAdapter* nf_adapter_ = nullptr;
 
  void RoundOffGridDimensions(const std::array<real_t, 6>& grid_dimensions);
 
  void CheckGridGrowth();
};
 
}  // namespace bdm
 
#endif  // CORE_ENVIRONMENT_KD_TREE_ENVIRONMENT_