entitysearch2.hh

// -*- tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
// vi: set et ts=4 sw=2 sts=2:
#ifndef DUNE_MESHDIST_ENTITYSEARCH2_HH
#define DUNE_MESHDIST_ENTITYSEARCH2_HH

#include <limits>
#include <map>
#include <vector>

#include <dune/common/fvector.hh>
#include <dune/common/parametertree.hh>
#include <dune/common/timer.hh>
#include <dune/geometry/mappedgeometry.hh>
#include <dune/geometry/quadraturerules.hh>
#include <dune/geometry/test/referenceelementgeometry.hh>
#include <dune/meshdist/kdtree.hh>
#include <dune/meshdist/chainediterator.hh>
#include <dune/grid/common/rangegenerators.hh>

namespace Dune::MeshDist {

template <class GridView, int dimension = GridView::dimensionworld>
class EntitySearch2
{
  using F = typename GridView::ctype;

  using GlobalCoordinate = Dune::FieldVector<F,dimension>;
  using EntitySeed = typename GridView::template Codim<0>::Entity::EntitySeed;
  using SeedRange = std::map<typename GridView::IndexSet::IndexType, EntitySeed>;

  template <class Parametrization>
  static auto boundingBox (const GridView& gridView, const Parametrization& X)
  {
    Timer t;
    auto lower = Dune::FieldVector<F,dimension>(std::numeric_limits<F>::max());
    auto upper = Dune::FieldVector<F,dimension>(std::numeric_limits<F>::min());

    auto localX = localFunction(X);
    for (auto const& e : Dune::elements(gridView))
    {
      localX.bind(e);
      auto refElem = referenceElement(e);
      int nVertices = refElem.size(GridView::dimension);
      for (int i = 0; i < nVertices; ++i)
      {
        auto v = localX(refElem.position(i,GridView::dimension));
        for (int j = 0; j < dimension; ++j)
        {
          lower[j] = std::min<F>(lower[j], v[j]);
          upper[j] = std::max<F>(upper[j], v[j]);
        }
      }
    }

    // TODO: Find better way to increase box size
    auto extra = (upper-lower)/10;
#ifndef NDEBUG
    std::cout << "time(build boundingBox): " << t.elapsed() << std::endl;
#endif
    return std::pair{lower-extra, upper+extra};
  }

  template <class Parametrization>
  static auto centerList (const GridView& gridView, const Parametrization& X)
  {
    Dune::Timer t;
    std::vector<Dune::FieldVector<F,dimension>> centers(gridView.size(0));

    auto localX = localFunction(X);
    const auto& indexSet = gridView.indexSet();
    for (auto const& e : Dune::elements(gridView))
    {
      localX.bind(e);
      auto refElem = referenceElement(e);
      centers[indexSet.index(e)] = localX(refElem.position(0,0));
    }
#ifndef NDEBUG
    std::cout << "time(build centerList): " << t.elapsed() << std::endl;
#endif
    return centers;
  }

  static auto seedList (const GridView& gridView)
  {
    Dune::Timer t;
    std::vector<EntitySeed> seeds(gridView.size(0));

    const auto& indexSet = gridView.indexSet();
    for (auto const& e : Dune::elements(gridView))
    {
      seeds[indexSet.index(e)] = e.seed();
    }
#ifndef NDEBUG
    std::cout << "time(build seedlist): " << t.elapsed() << std::endl;
#endif
    return seeds;
  }


public:
  template <class Parametrization>
  EntitySearch2 (const GridView& gridView, const Parametrization& X)
    : gridView_(gridView)
    , points_(centerList(gridView_,X))
    , entitySeeds_(seedList(gridView))
    , kdtree_(points_, boundingBox(gridView_,X))
  {}

  template <class K>
  auto elements (const Dune::FieldVector<K,dimension>& x, const Dune::ParameterTree& pt = {}) const
  {
    int numElements = pt.get<int>("num_elements", 16);
    std::vector<std::size_t> outIndices(numElements);
    outDistancesSq_.resize(numElements);
    kdtree_.query(x, numElements, outIndices, outDistancesSq_);

    return Dune::transformedRangeView(std::move(outIndices), [&](const auto& index)
    {
      return gridView_.grid().entity(entitySeeds_[index]);
    });
  }

private:
  GridView gridView_;
  std::vector<GlobalCoordinate> points_;
  std::vector<EntitySeed> entitySeeds_;
  KDTree<F,dimension> kdtree_;

  mutable std::vector<F> outDistancesSq_;
};

} // end namespace Dune::MeshDist

#endif // DUNE_MESHDIST_ENTITYSEARCH2_HH