BeamMonitor.H

Go to the documentation of this file.

/* Copyright 2022-2026 The Regents of the University of California, through Lawrence
 *           Berkeley National Laboratory (subject to receipt of any required
 *           approvals from the U.S. Dept. of Energy). All rights reserved.
 *
 * This file is part of ImpactX.
 *
 * Authors: Axel Huebl, Chad Mitchell
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_ELEMENTS_DIAGS_BEAMMONITOR_H
#define IMPACTX_ELEMENTS_DIAGS_BEAMMONITOR_H
 
#include "elements/mixin/thin.H"
#include "particles/CovarianceMatrix.H"
#include "particles/ImpactXParticleContainer.H"
#include "elements/mixin/lineartransport.H"
 
#include <AMReX_Extension.H>
#include <AMReX_REAL.H>
 
#include <any>
#include <string>
#include <unordered_map>
#include <vector>
 
 
namespace impactx::elements::diagnostics
{
namespace detail
{
    class ImpactXParticleCounter {
    public:
        using ParticleContainer = typename ImpactXParticleContainer::ContainerLike<amrex::PolymorphicArenaAllocator>;
        using ParticleIter = typename ParticleContainer::ParIterType;
 
        ImpactXParticleCounter (ParticleContainer & pc);
 
        unsigned long long GetTotalNumParticles () { return m_Total; }
 
        std::vector<unsigned long long> m_ParticleOffsetAtRank;
        std::vector<unsigned long long> m_ParticleSizeAtRank;
    private:
        void GetParticleOffsetOfProcessor (amrex::Long const &numParticles,
                                           unsigned long long &offset,
                                           unsigned long long &sum) const;
 
        int m_MPIRank = 0;
        int m_MPISize = 1;
 
        unsigned long long m_Total = 0;
 
        std::vector<unsigned long long> m_ParticleCounterByLevel;
    };
} // namespace detail

    struct BeamMonitor
    : public mixin::LinearTransport<BeamMonitor>,
      public mixin::Thin
    {
        static constexpr auto type = "BeamMonitor";
        using PType = typename ImpactXParticleContainer::ParticleType;
        using PinnedContainer = typename ImpactXParticleContainer::ContainerLike<amrex::PolymorphicArenaAllocator>;
 
        static constexpr char const * DEFAULT_backend = "default";
        static constexpr char const * DEFAULT_encoding = "g";
        static constexpr int DEFAULT_period_sample_intervals = 1;

        BeamMonitor (
            std::string series_name,
            std::string backend = DEFAULT_backend,
            std::string encoding = DEFAULT_encoding,
            int period_sample_intervals = DEFAULT_period_sample_intervals
        );
 
        BeamMonitor (BeamMonitor const & other) = default;
        BeamMonitor (BeamMonitor && other) = default;
        BeamMonitor& operator= (BeamMonitor const & other) = default;
        BeamMonitor& operator= (BeamMonitor && other) = default;

        void reverse () { }

        void prepare (
            PinnedContainer & pc,
            std::vector<std::string> const & real_soa_names,
            std::vector<std::string> const & int_soa_names,
            RefPart const & ref_part,
            int step
        );

        void operator() (
            ImpactXParticleContainer & pc,
            int step,
            int period
        );

        void operator() (
            PinnedContainer::ParIterType & pti,
            std::vector<std::string> const & real_soa_names,
            std::vector<std::string> const & int_soa_names,
            RefPart const & ref_part
        );

        using Thin::operator();

        using LinearTransport::operator();

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
        {
            // nothing to do
            return Map6x6::Identity();
        }

        AMREX_FORCE_INLINE
        std::string name () const { return m_series_name; }

        AMREX_FORCE_INLINE
        bool has_name () const { return true; }

        AMREX_FORCE_INLINE
        std::string backend () const { return m_OpenPMDFileType; }

        AMREX_FORCE_INLINE
        std::string encoding () const { return m_encoding; }

        AMREX_FORCE_INLINE
        int period_sample_intervals () const { return m_period_sample_intervals; }

        static inline std::map<std::string, std::any> m_unique_series = {};

        void
        finalize ();
 
    private:
        void open ();
 
        std::string m_series_name; 
        std::string m_OpenPMDFileType; 
        std::string m_encoding; 
        std::any m_series; 
        std::string m_series_key; 
        int m_step = 0; 
 
        int m_file_min_digits = 6; 
 
        int m_period_sample_intervals = 1; 

        std::vector<uint64_t> m_offset;

        std::unordered_map<std::string, amrex::ParticleReal> m_rbc;
 
    };

    void
    add_optional_properties (
        std::string const & element_name,
        ImpactXParticleContainer & pc
    );
 
} // namespace impactx::diagnostics
 
#endif // IMPACTX_ELEMENTS_DIAGS_BEAMMONITOR_H