Cabana_ParticleInit.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef CABANA_PARTICLEINIT_HPP
#define CABANA_PARTICLEINIT_HPP
 
#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>
 
#include <Cabana_ParticleList.hpp>
#include <Cabana_Slice.hpp>
 
#include <random>
#include <type_traits>
 
namespace Cabana
{
 
namespace Impl
{
template <class ArrayType>
auto copyArray( ArrayType corner )
{
    using value_type = std::remove_reference_t<decltype( corner[0] )>;
    Kokkos::Array<value_type, 3> kokkos_corner;
    for ( std::size_t d = 0; d < 3; ++d )
        kokkos_corner[d] = corner[d];
 
    return kokkos_corner;
}
 
} // namespace Impl
 
//---------------------------------------------------------------------------//
// Initialization type tags.

struct InitUniform
{
};

struct InitRandom
{
};
 
//---------------------------------------------------------------------------//
// Fully random
//---------------------------------------------------------------------------//
 
//---------------------------------------------------------------------------//
template <class ExecutionSpace, class InitFunctor, class ParticleListType,
          class ArrayType>
int createParticles(
    InitRandom, ExecutionSpace exec_space, const InitFunctor& create_functor,
    ParticleListType& particle_list, const std::size_t num_particles,
    const ArrayType box_min, const ArrayType box_max,
    const std::size_t previous_num_particles = 0,
    const bool shrink_to_fit = true, const uint64_t seed = 342343901,
    typename std::enable_if<is_particle_list<ParticleListType>::value,
                            int>::type* = 0 )
{
    // Memory space.
    using memory_space = typename ParticleListType::memory_space;
 
    using PoolType = Kokkos::Random_XorShift64_Pool<ExecutionSpace>;
    using RandomType = Kokkos::Random_XorShift64<ExecutionSpace>;
    PoolType pool( seed );
 
    // Resize the aosoa prior to lambda capture.
    auto& aosoa = particle_list.aosoa();
    aosoa.resize( previous_num_particles + num_particles );
 
    // Creation count.
    auto count = Kokkos::View<int*, memory_space>( "particle_count", 1 );
    Kokkos::deep_copy( count, previous_num_particles );
 
    // Copy corners to device accessible arrays.
    auto kokkos_min = Impl::copyArray( box_min );
    auto kokkos_max = Impl::copyArray( box_max );
 
    auto random_coord_op = KOKKOS_LAMBDA( const int p )
    {
        // Particle coordinate.
        double px[3];
 
        auto gen = pool.get_state();
        auto particle = particle_list.getParticle( p );
        for ( int d = 0; d < 3; ++d )
            px[d] = Kokkos::rand<RandomType, double>::draw( gen, kokkos_min[d],
                                                            kokkos_max[d] );
        pool.free_state( gen );
 
        // No volume information, so pass zero.
        int create = create_functor( count( 0 ), px, 0.0, particle );
 
        // If we created a new particle insert it into the list.
        if ( create )
        {
            auto c = Kokkos::atomic_fetch_add( &count( 0 ), 1 );
            particle_list.setParticle( particle, c );
        }
    };
 
    Kokkos::RangePolicy<ExecutionSpace> exec_policy(
        exec_space, previous_num_particles,
        previous_num_particles + num_particles );
    Kokkos::parallel_for( exec_policy, random_coord_op );
    Kokkos::fence();
 
    auto count_host =
        Kokkos::create_mirror_view_and_copy( Kokkos::HostSpace(), count );
    aosoa.resize( count_host( 0 ) );
    if ( shrink_to_fit )
        aosoa.shrinkToFit();
    return count_host( 0 );
}
 
//---------------------------------------------------------------------------//
template <class InitFunctor, class ParticleListType, class ArrayType>
int createParticles( InitRandom tag, const InitFunctor& create_functor,
                     ParticleListType& particle_list,
                     const std::size_t num_particles, const ArrayType box_min,
                     const ArrayType box_max,
                     const std::size_t previous_num_particles = 0,
                     const bool shrink_to_fit = true,
                     const uint64_t seed = 342343901 )
{
    using exec_space = typename ParticleListType::memory_space::execution_space;
    return createParticles( tag, exec_space{}, create_functor, particle_list,
                            num_particles, box_min, box_max,
                            previous_num_particles, shrink_to_fit, seed );
}

template <class ExecutionSpace, class PositionType, class ArrayType>
void createParticles(
    InitRandom, ExecutionSpace exec_space, PositionType& positions,
    const std::size_t num_particles, const ArrayType box_min,
    const ArrayType box_max, const std::size_t previous_num_particles = 0,
    const uint64_t seed = 342343901,
    typename std::enable_if<( is_slice<PositionType>::value ||
                              Kokkos::is_view<PositionType>::value ),
                            int>::type* = 0 )
{
    // Ensure correct space for the particles (View or Slice).
    checkSize( positions, num_particles + previous_num_particles );
 
    // Copy corners to device accessible arrays.
    auto kokkos_min = Impl::copyArray( box_min );
    auto kokkos_max = Impl::copyArray( box_max );
 
    using PoolType = Kokkos::Random_XorShift64_Pool<ExecutionSpace>;
    using RandomType = Kokkos::Random_XorShift64<ExecutionSpace>;
    PoolType pool( seed );
    auto random_coord_op = KOKKOS_LAMBDA( const int p )
    {
        auto gen = pool.get_state();
        for ( int d = 0; d < 3; ++d )
            positions( p, d ) = Kokkos::rand<RandomType, double>::draw(
                gen, kokkos_min[d], kokkos_max[d] );
        pool.free_state( gen );
    };
 
    Kokkos::RangePolicy<ExecutionSpace> exec_policy(
        exec_space, previous_num_particles,
        previous_num_particles + num_particles );
    Kokkos::parallel_for( exec_policy, random_coord_op );
    Kokkos::fence();
}

template <class PositionType, class ArrayType>
void createParticles(
    InitRandom tag, PositionType& positions, const std::size_t num_particles,
    const ArrayType box_min, const ArrayType box_max,
    const std::size_t previous_num_particles = 0,
    const uint64_t seed = 342343901,
    typename std::enable_if<( is_slice<PositionType>::value ||
                              Kokkos::is_view<PositionType>::value ),
                            int>::type* = 0 )
{
    using exec_space = typename PositionType::execution_space;
    createParticles( tag, exec_space{}, positions, num_particles, box_min,
                     box_max, previous_num_particles, seed );
}
 
//---------------------------------------------------------------------------//
// Random with minimum separation
//---------------------------------------------------------------------------//

template <class ExecutionSpace, class InitFunctor, class ParticleListType,
          class PositionTag, class ArrayType>
int createParticles(
    InitRandom tag, ExecutionSpace exec_space,
    const InitFunctor& create_functor, ParticleListType& particle_list,
    PositionTag position_tag, const std::size_t num_particles,
    const double min_dist, const ArrayType box_min, const ArrayType box_max,
    const std::size_t previous_num_particles = 0,
    const bool shrink_to_fit = true, const uint64_t seed = 342343901,
    typename std::enable_if<is_particle_list<ParticleListType>::value,
                            int>::type* = 0 )
{
    double min_dist_sqr = min_dist * min_dist;
 
    // Resize the aosoa prior to lambda capture.
    auto& aosoa = particle_list.aosoa();
    aosoa.resize( previous_num_particles + num_particles );
 
    auto positions = particle_list.slice( position_tag );
 
    // Create the functor which ignores particles within the radius of another.
    auto min_distance_op =
        KOKKOS_LAMBDA( const int id, const double px[3], const double,
                       typename ParticleListType::particle_type& particle )
    {
        // Ensure this particle is not within the minimum distance of any other
        // existing particle.
        for ( int n = 0; n < id; n++ )
        {
            double dx = positions( n, 0 ) - px[0];
            double dy = positions( n, 1 ) - px[1];
            double dz = positions( n, 2 ) - px[2];
            double dist = dx * dx + dy * dy + dz * dz;
            if ( dist < min_dist_sqr )
                return false;
        }
 
        bool create = create_functor( id, px, 0.0, particle );
        return create;
    };
 
    // Pass the functor to the general case.
    return createParticles( tag, exec_space, min_distance_op, particle_list,
                            num_particles, box_min, box_max,
                            previous_num_particles, shrink_to_fit, seed );
}

template <class InitFunctor, class ParticleListType, class PositionTag,
          class ArrayType>
int createParticles( InitRandom tag, const InitFunctor& create_functor,
                     ParticleListType& particle_list, PositionTag position_tag,
                     const std::size_t num_particles, const double min_dist,
                     const ArrayType box_min, const ArrayType box_max,
                     const std::size_t previous_num_particles = 0,
                     const bool shrink_to_fit = true,
                     const uint64_t seed = 342343901 )
{
    using exec_space = typename ParticleListType::memory_space::execution_space;
    return createParticles( tag, exec_space{}, create_functor, particle_list,
                            position_tag, num_particles, min_dist, box_min,
                            box_max, previous_num_particles, shrink_to_fit,
                            seed );
}
 
} // namespace Cabana
 
#endif