abstract_particle_set.hpp

/******************************************************************************
*                                                                             *
*  Copyright 2020 Max Planck Institute for Dynamics and Self-Organization     *
*                                                                             *
*  This file is part of TurTLE.                                               *
*                                                                             *
*  TurTLE is free software: you can redistribute it and/or modify             *
*  it under the terms of the GNU General Public License as published          *
*  by the Free Software Foundation, either version 3 of the License,          *
*  or (at your option) any later version.                                     *
*                                                                             *
*  TurTLE is distributed in the hope that it will be useful,                  *
*  but WITHOUT ANY WARRANTY; without even the implied warranty of             *
*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
*  GNU General Public License for more details.                               *
*                                                                             *
*  You should have received a copy of the GNU General Public License          *
*  along with TurTLE.  If not, see <http://www.gnu.org/licenses/>             *
*                                                                             *
* Contact: Cristian.Lalescu@ds.mpg.de                                         *
*                                                                             *
******************************************************************************/



#ifndef ABSTRACT_PARTICLE_SET_HPP
#define ABSTRACT_PARTICLE_SET_HPP

#include <array>
#include <cassert>
#include <vector>
#include <memory>
#include "field.hpp"
#include "particles/p2p/p2p_distr_mpi.hpp"
#include "particles/particles_sampling.hpp"



/** Brings together particle information with interpolation functionality.
 *
 * This is an abstract class that defines the functionality required by the
 * particle solver code.
 * We define methods for:
 *  - accessing particle set data and miscelaneous information.
 *  - updating particle information and redistributing data among MPI processes
 *    accordingly.
 *  - approximating fields at particle locations ("sampling").
 *  - applying a generic particle-to-particle interaction object (P2PKernel)
 *
 */

class abstract_particle_set
{
    public:
        using partsize_t = long long int;
        using particle_rnumber = double;

        // virtual destructor
        virtual ~abstract_particle_set(){}

        // extract particle information
        virtual particle_rnumber* getParticleState() const = 0;
        virtual partsize_t* getParticleIndex() const = 0;
        virtual int setParticleState(particle_rnumber *) = 0;
        virtual int getParticleState(particle_rnumber *) = 0;

        virtual std::unique_ptr<particle_rnumber[]> extractFromParticleState(
                const int firstComponent,
                const int lastComponent) const = 0;

        virtual partsize_t getLocalNumberOfParticles() const = 0;
        virtual partsize_t getTotalNumberOfParticles() const = 0;
        virtual partsize_t* getParticlesPerPartition() const = 0;
        virtual int getStateSize() const = 0;

        virtual std::vector<hsize_t> getParticleFileLayout() = 0;

        // get p2p computer
        virtual p2p_distr_mpi<partsize_t, particle_rnumber>* getP2PComputer() = 0;

        // generic loop over data
        template <class func_type>
        int LOOP(func_type expression)
        {
            TIMEZONE("abstract_particle_set::LOOP");
            for (partsize_t idx = 0; idx < this->getLocalNumberOfParticles()*this->getStateSize(); idx++)
                expression(idx);
            return EXIT_SUCCESS;
        }
        template <class func_type>
        int LOOP_state(func_type expression)
        {
            TIMEZONE("abstract_particle_set::LOOP_state");
            for (partsize_t idx_part = 0; idx_part < this->getLocalNumberOfParticles(); idx_part++)
            for (unsigned int cc = 0; cc < this->getStateSize(); cc++)
                expression(idx_part, cc);
            return EXIT_SUCCESS;
        }

        int copy_state_tofrom(
                particle_rnumber *dst,
                const particle_rnumber *src)
        {
            return this->LOOP(
                    [&](const partsize_t idx)
                    {
                        dst[idx] = src[idx];
                    });
        }

        int copy_state_tofrom(
                std::unique_ptr<particle_rnumber[]> &dst,
                const std::unique_ptr<particle_rnumber[]> &src)
        {
            return this->LOOP(
                    [&](const partsize_t idx)
                    {
                        dst[idx] = src[idx];
                    });
        }

        /** Reorganize particles within MPI domain.
         *
         * Based on current particle locations, redistribute the full state
         * data among the MPI processes.
         * Optional: also redistribute a list of arrays of the same shape.
         */
        virtual int redistribute(
                std::vector<std::unique_ptr<particle_rnumber[]>> &additional_data) = 0;

        // sample field values at particle location
        virtual int sample(
                const field<float,
                            FFTW,
                            ONE>         &field_to_sample,
                particle_rnumber *result) = 0;
        virtual int sample(
                const field<float,
                            FFTW,
                            THREE>       &field_to_sample,
                particle_rnumber *result) = 0;
        virtual int sample(
                const field<float,
                            FFTW,
                            THREExTHREE> &field_to_sample,
                particle_rnumber *result) = 0;
        virtual int sample(
                const field<double,
                            FFTW,
                            ONE>         &field_to_sample,
                particle_rnumber *result) = 0;
        virtual int sample(
                const field<double,
                            FFTW,
                            THREE>       &field_to_sample,
                particle_rnumber *result) = 0;
        virtual int sample(
                const field<double,
                            FFTW,
                            THREExTHREE> &field_to_sample,
                particle_rnumber *result) = 0;

        // apply p2p kernel
        // IMPORTANT: this method shuffles particle arrays.
        // If you need some data to be shuffled according to the same permutation,
        // please place it in `additional_data`.
        // This applies for instance to rhs data during time steps.
        template <int state_size,
                  class p2p_kernel_class>
        int applyP2PKernel(
                p2p_kernel_class &p2p_kernel,
                std::vector<std::unique_ptr<particle_rnumber[]>> &additional_data)
        {
            TIMEZONE("abstract_particle_set::applyP2PKernel");
            // there must be at least one array with additional data,
            // since the p2p kernel expects an array where to store the result of the computation
            assert(int(additional_data.size()) > int(0));
            this->getP2PComputer()->template compute_distr<p2p_kernel_class,
                                                           state_size, // I'd prefer to use this->getStateSize() here. not possible because virtual function is not constexpr...
                                                           state_size>(
                    p2p_kernel,
                    this->getParticlesPerPartition(),
                    this->getParticleState(),
                    additional_data.data(),
                    int(additional_data.size()),
                    this->getParticleIndex());
            return EXIT_SUCCESS;
        }

        template <typename field_rnumber,
                  field_backend be,
                  field_components fc>
        int writeSample(
                field<field_rnumber, be, fc> *field_to_sample,
                particles_output_sampling_hdf5<partsize_t, particle_rnumber, 3> *particle_sample_writer,
                const std::string species_name,
                const std::string field_name,
                const int iteration)
        {
            TIMEZONE("abstract_particle_set::writeSample");
            // set file layout
            particle_sample_writer->setParticleFileLayout(this->getParticleFileLayout());
            // allocate position array
            std::unique_ptr<particle_rnumber[]> xx = this->extractFromParticleState(0, 3);
            // allocate temporary array
            std::unique_ptr<particle_rnumber[]> pdata(new particle_rnumber[ncomp(fc)*this->getLocalNumberOfParticles()]);
            // clean up temporary array
            std::fill_n(pdata.get(), ncomp(fc)*this->getLocalNumberOfParticles(), 0);
            this->sample(*field_to_sample, pdata.get());
            particle_sample_writer->template save_dataset<ncomp(fc)>(
                    species_name,
                    field_name,
                    xx.get(),
                    &pdata,
                    this->getParticleIndex(),
                    this->getLocalNumberOfParticles(),
                    iteration);
            // deallocate temporary array
            delete[] pdata.release();
            // deallocate position array
            delete[] xx.release();
            return EXIT_SUCCESS;
        }

        template <typename field_rnumber,
                  field_backend be,
                  field_components fc>
        int writeSample(
                field<field_rnumber, be, fc> *field_to_sample,
                const std::string file_name,
                const std::string species_name,
                const std::string field_name,
                const int iteration)
        {
            TIMEZONE("abstract_particle_set::writeSample_direct_to_file");
            // allocate temporary array
            std::unique_ptr<particle_rnumber[]> pdata(new particle_rnumber[ncomp(fc)*this->getLocalNumberOfParticles()]);
            // clean up temporary array
            set_particle_data_to_zero<partsize_t, particle_rnumber, ncomp(fc)>(
                    pdata.get(), this->getLocalNumberOfParticles());
            this->sample(*field_to_sample, pdata.get());
            hid_t file_id = H5Fopen(file_name.c_str(), H5F_ACC_RDWR, H5P_DEFAULT);
            hdf5_tools::gather_and_write_with_single_rank(
                    field_to_sample->myrank,
                    0,
                    field_to_sample->comm,
                    pdata.get(),
                    this->getLocalNumberOfParticles(),
                    file_id,
                    species_name + std::string("/") + field_name + std::string("/") + std::to_string(iteration));
            H5Fclose(file_id);
            // deallocate temporary array
            delete[] pdata.release();
            return EXIT_SUCCESS;
        }

        int writeStateChunk(
                const int i0,
                const int contiguous_state_chunk,
                particles_output_sampling_hdf5<partsize_t, particle_rnumber, 3> *particle_sample_writer,
                const std::string species_name,
                const std::string field_name,
                const int iteration)
        {
            TIMEZONE("abstract_particle_set::writeStateChunk");
            assert(i0 >= 0);
            assert(i0 <= this->getStateSize()-contiguous_state_chunk);
            // set file layout
            particle_sample_writer->setParticleFileLayout(this->getParticleFileLayout());
            // allocate position array
            std::unique_ptr<particle_rnumber[]> xx = this->extractFromParticleState(0, 3);
            // allocate temporary array
            std::unique_ptr<particle_rnumber[]> yy = this->extractFromParticleState(i0, i0+contiguous_state_chunk);
            switch(contiguous_state_chunk)
            {
                case 1:
                    particle_sample_writer->template save_dataset<1>(
                            species_name,
                            field_name,
                            xx.get(),
                            &yy,
                            this->getParticleIndex(),
                            this->getLocalNumberOfParticles(),
                            iteration);
                    break;
                case 3:
                    particle_sample_writer->template save_dataset<3>(
                            species_name,
                            field_name,
                            xx.get(),
                            &yy,
                            this->getParticleIndex(),
                            this->getLocalNumberOfParticles(),
                            iteration);
                    break;
                case 9:
                    particle_sample_writer->template save_dataset<9>(
                            species_name,
                            field_name,
                            xx.get(),
                            &yy,
                            this->getParticleIndex(),
                            this->getLocalNumberOfParticles(),
                            iteration);
                    break;
                default:
                    DEBUG_MSG("code not specialized for this value of contiguous_state_chunk in abstract_particle_set::writeStateChunk\n");
                    std::cerr << "error in abstract_particle_set::writeStateChunk. please contact maintainer.\n" << std::endl;
            }
            // deallocate temporary array
            delete[] yy.release();
            // deallocate position array
            delete[] xx.release();
            return EXIT_SUCCESS;
        }

        int writeStateTriplet(
                const int i0,
                particles_output_sampling_hdf5<partsize_t, particle_rnumber, 3> *particle_sample_writer,
                const std::string species_name,
                const std::string field_name,
                const int iteration)
        {
            TIMEZONE("abstract_particle_set::writeStateTriplet");
            return this->writeStateChunk(
                    i0, 3,
                    particle_sample_writer,
                    species_name,
                    field_name,
                    iteration);
        }

        int writeStateComponent(
                const int i0,
                particles_output_sampling_hdf5<partsize_t, particle_rnumber, 3> *particle_sample_writer,
                const std::string species_name,
                const std::string field_name,
                const int iteration)
        {
            TIMEZONE("abstract_particle_set::writeStateComponent");
            return this->writeStateChunk(
                    i0, 1,
                    particle_sample_writer,
                    species_name,
                    field_name,
                    iteration);
        }
};

#endif//ABSTRACT_PARTICLE_SET_HPP