Rigid body implementation ouline with Cabana software framework

rb_cabana.cpp

/****************************************************************************
 * Copyright (c) 2018-2022 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                                    *
 ****************************************************************************/

#include <Cabana_Core.hpp>

#include <iostream>

using DataTypes = Cabana::MemberTypes< double>;
const int VectorLength = 8;
using MemorySpace = Kokkos::HostSpace;
using ExecutionSpace = Kokkos::DefaultHostExecutionSpace;
using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
using AoSoAType = Cabana::AoSoA<DataTypes, DeviceType, VectorLength>;

// rigid body data type
// cm, vcm, mass, force, total no of bodies, indices limits
using RigidBodyDataType = Cabana::MemberTypes<double, int[2]>;
using RBAoSoA = Cabana::AoSoA<RigidBodyDataType, DeviceType, VectorLength>;

void setup_rigid_body_properties(auto aosoa, auto rb){
  auto rb_mass = Cabana::slice<0>( rb, "mass" );
  auto rb_limits = Cabana::slice<1>( rb, "limits" );

  auto aosoa_m = Cabana::slice<0>( aosoa, "mass" );

  auto total_mass_reduce = KOKKOS_LAMBDA( const int i, double& total_m )
    {
      total_m += aosoa_m( i );
    };

  for ( std::size_t i = 0; i < rb_mass.size(); ++i )
    {
      double total_mass = 0.0;
      Kokkos::RangePolicy<ExecutionSpace> policy( rb_limits(i, 0), rb_limits(i, 1) );
      Kokkos::parallel_reduce(
                              "total_mass_computation", policy,
                              total_mass_reduce, total_mass);
      rb_mass(i) = total_mass;
    }

  // print total mass of the rigid body after compuation
  for ( std::size_t i = 0; i < rb_mass.size(); ++i )
    {
      std::cout << "total mass: " << rb_mass(i) << "\n";
    }
}


//---------------------------------------------------------------------------//
// TODO: explain this function in short
//---------------------------------------------------------------------------//
void run()
{
  int comm_rank = -1;
  MPI_Comm_rank( MPI_COMM_WORLD, &comm_rank );

  if ( comm_rank == 0 )
    std::cout << "Two AoSoAs example\n" << std::endl;

  auto num_particles = 16;

  std::vector<double> m_array = { -1, -1, -1, -1, -1, -1, -1, -1, 0.5, 0.5, 0.5, 0.5, 1, 1, 1, 1};
  std::vector<int> rb_limits = {8, 16};

  // create aosoa
  AoSoAType aosoa( "particles", num_particles );

  auto m = Cabana::slice<0>( aosoa, "mass" );

  // initialize the particle properties
  for ( std::size_t i = 0; i < aosoa.size(); ++i )
    {
      m( i ) = m_array[i];
    }

  /*
    Create the rigid body data type.
  */
  int num_bodies = 2;
  RBAoSoA rb( "rb", num_bodies );
  // get the rigid body limits
  {
  auto limits = Cabana::slice<1>( rb, "total_no_bodies" );
  // limits of the first body
  limits(0, 0) = 8;
  limits(0, 1) = 12;

  limits(1, 0) = 12;
  limits(1, 1) = 16;

  setup_rigid_body_properties(aosoa, rb);
  }
}

int main( int argc, char* argv[] )
{

  MPI_Init( &argc, &argv );
  Kokkos::initialize( argc, argv );

  run();

  Kokkos::finalize();

  MPI_Finalize();
  return 0;
}