AlgorithmStepLibrary.py

Go to the documentation of this file.

from exaparty.api.particle.AlgorithmStep import AlgorithmStep
from exaparty.api.particle.SPHParticle import SPHParticle
 
 
DefaultIncludes = """
#include "exaparty/CellMetaData.h"
"""
 
 
class ParticleVertexIteratorTemplate(str):
    """!
    A mini-wrapper around a string for particle iterator.
    Exists for convenience: adds self.insert() method,
    which inserts the function which the iterators are supposed to call
    on a particle or particle pair.
 
    The function insertion works via string replacement. It will
    look for the string '*FUNCTION*' to be replaced.
    """
 
    def __init__(self, template):
        """!
        template: str
            the string containing the particle interacion iterator template
        """
        self.template = template
        return
 
    def insert(self, function):
        """!
        insert a function into the template string.
 
        Returns the particle iterator template with the function inserted.
        """
        return self.template.replace("*FUNCTION*", function)
 
 
class ParticleCellIteratorTemplate(str):
    """!
 
    Counterpart of ParticleVertexIteratorTemplate for cells.
 
    """
 
    def __init__(self, template):
        """!
        template: str
            the string containing the particle interacion iterator template
        """
        self.template = template
        return
 
    def insert(self, function):
        """!
        insert a function into the template string.
 
        Returns the particle iterator template with the function inserted.
        """
 
        return """
        ::exaparty::CellMetaData cellMetaData(
          marker,
          numberOfCoalescedLocalParticlesPerVertex,
          numberOfCoalescedActiveParticlesPerVertex,
          localParticleIsContainedInCell
        );        
""" + self.template.replace(
            "*FUNCTION*", function
        )
 
 
def _get_particle_iterator_templates(particle_name):
    """!
    Construct particle interaction iterator call templates given the
    desired particle interaction kernel realisation. These signature
    templates are strings containing c++ code snippets of the correct
    iterator call signature.
 
    In total, there are four types of iterators. The essential two are:
 
    - The one that's used to run over all particle assigned to a vertex. It
      is called iterate_over_particles_assigned_to_vertex.
    - There is one to iterate over all particle pairs associated with a
      cell. This one is called iterate_over_particles_assigned_to_cell.
 
    From hereon, we can define some "derived" iterators:
 
    - One special iterator runs over all particles associated with a vertex
      and allows them to move. This one is titled
      iterate_over_particles_assigned_to_vertex_for_move and equals
      iterate_over_particles_assigned_to_vertex. However, it embeds the
      latter into some bookkeeping for the validation of the
      vertex-particle assocation.
    - Another special iterator runs over all particles but allows them to
      reduce into some global space. If you use vectorisation, this one
      cannot be vectorised by definition.
 
    Both special iterators are constructed to default values unless you
    explicitly set them beforehand.
 
    The iterators require a function to be passed as an argument to
    act on the particles. In the template snippets being generated here,
    the functions shall be replaced by a placeholder string '*FUNCTION*'.
 
    Finally, the code snippets being returned are stored as a
    ParticleVertexIteratorTemplate object.
 
    @param particle_name: str
        the particle class name.
 
 
    4 ParticleVertexIteratorTemplate objects containing particle iterator code snippets.
 
    @return vertex_for_move: snippet (in a ParticleVertexIteratorTemplate object) to
        use for particle interactons which will move assigned to a vertex
    @return cell: snippet (in a ParticleVertexIteratorTemplate object) to use for
        particle-particle interactions assigned to a cell
    @return vertex: snippet (in a ParticleVertexIteratorTemplate object) to use for
        particle interactions assigned to a vertex
    @return vertex_with_global_reduction: snippet (in a ParticleVertexIteratorTemplate
        object) to use for particle interactions assigned to a vertex which also
        contain a global reduction
    """
 
    PARTICLE = particle_name
 
    iterate_over_particles_assigned_to_vertex_for_move = None
    iterate_over_particles_assigned_to_vertex = f"""
        // Template in src/exaparty/api/particle/AlgorithmStepLibrary.py
        // iterate_over_particles_assigned_to_vertex
 
        ::exaparty::kernels::forAllParticles(*FUNCTION*<globaldata::{PARTICLE}>,
                                           numberOfCoalescedAssignedParticles,
                                           *FUNCTION*UpdateParticlePredicate<globaldata::{PARTICLE}>
                                           );
        // end template
        """
    iterate_over_particles_assigned_to_cell = f"""
        // Template in src/exaparty/api/particle/AlgorithmStepLibrary.py
        // iterate_over_particles_assigned_to_cell
        
        ::exaparty::kernels::forAllParticlePairs<*FUNCTION*<globaldata::{PARTICLE}>, *FUNCTION*UpdateParticlePredicate<globaldata::{PARTICLE}>, *FUNCTION*EvaluatePairInteractionPredicate<globaldata::{PARTICLE}>>(
                                               cellMetaData,
                                               localParticles,
                                               activeParticles,
                                               ::exaparty::kernels::kernelrealisation::Generic
                                               );
        // end template
"""
    iterate_over_particles_assigned_to_vertex_with_global_reduction = f"""
        // Template in src/exaparty/api/particle/AlgorithmStepLibrary.py
        // iterate_over_particles_assigned_to_vertex_with_global_reduction
 
        ::exaparty::kernels::forAllParticles<*FUNCTION*<globaldata::{PARTICLE}>, *FUNCTION*UpdateParticlePredicate<globaldata::{PARTICLE}>>(marker,
                                           assignedParticles,
                                           numberOfCoalescedAssignedParticles,
                                           ::exaparty::kernels::kernelrealisation::Generic
                                           );
        // end template
        """
 
    #  Set defaults or raise error if there shouldn't be defaults and we missed something.
 
    if iterate_over_particles_assigned_to_vertex_for_move is None:
        iterate_over_particles_assigned_to_vertex_for_move = f"""
            // Template in src/exaparty/api/particle/AlgorithmStepLibrary.py
            // iterate_over_particles_assigned_to_vertex_for_move
 
            auto oldParticlePositions = ::toolbox::particles::assignmentchecks::recordParticlePositions(assignedParticles);
            ::exaparty::kernels::forAllParticles<*FUNCTION*<globaldata::{PARTICLE}>, *FUNCTION*ParticleUpdatePredicate<globaldata::{PARTICLE}>>(marker,
                                               assignedParticles,
                                               numberOfCoalescedAssignedParticles,
                                               ::exaparty::kernels::kernelrealisation::Generic
                                               );
            ::toolbox::particles::assignmentchecks::traceParticleMovements(assignedParticles,
                                                                           oldParticlePositions,
                                                                           marker.x(),
                                                                           marker.h(),
                                                                           _spacetreeId);
            // end template
            """
        
 
    if iterate_over_particles_assigned_to_cell is None:
        raise ValueError("We shouldn't be here.")
 
    if iterate_over_particles_assigned_to_vertex is None:
        raise ValueError("We shouldn't be here.")
 
    if iterate_over_particles_assigned_to_vertex_with_global_reduction is None:
        iterate_over_particles_assigned_to_vertex_with_global_reduction = f"""
        // Template in src/exaparty/api/particle/AlgorithmStepLibrary.py,
        // iterate_over_particles_assigned_to_vertex_with_global_reduction
        ::exaparty::kernels::forAllParticles<*FUNCTION*<globaldata::{PARTICLE}>,
                                           ::exaparty::kernels::alwaysUpdateInVertexKernel<globaldata::{PARTICLE}>>(marker,
                                           assignedParticles,
                                           ::exaparty::kernels::kernelrealisation::Generic
                                           );
        // end template
        """
 
    #  finally, wrap the templates into ParticleVertexIteratorTemplate objects
    vertex_for_move = ParticleVertexIteratorTemplate(
        iterate_over_particles_assigned_to_vertex_for_move
    )
    cell = ParticleCellIteratorTemplate(iterate_over_particles_assigned_to_cell)
    vertex = ParticleVertexIteratorTemplate(
        iterate_over_particles_assigned_to_vertex_with_global_reduction
    )
    vertex_with_global_reduction = ParticleVertexIteratorTemplate(
        iterate_over_particles_assigned_to_vertex_with_global_reduction
    )
 
    return vertex_for_move, cell, vertex, vertex_with_global_reduction
 
 
def get_algorithm_step_dict(particle):
    """!
 
    Sets up a dict of the algorithm steps for this scheme so they can be accessed
    at various points and by various models.
 
    """
 
    algorithm_steps_dict = {}
 
    project_namespace = particle._exaparty_project_namespace
 
    PARTICLE = particle.name
    NAMESPACE = "::".join(project_namespace)
 
    (
        v_iter_for_move,
        c_iter,
        v_iter,
        v_iter_with_global_reduction,
    ) = _get_particle_iterator_templates(PARTICLE)
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    sph_kick1 = AlgorithmStep(
        name="SPH_Kick1",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.ALTER_LOCAL_STATE,
        touch_vertex_last_time_kernel=v_iter.insert(
            "::exaparty::kernels::legacy::leapfrogKickWithGlobalTimeStepSize"
        ),
        includes=DefaultIncludes
        + """
            #include "exaparty/kernels/legacy/Leapfrog.h"
        """,
    )
    algorithm_steps_dict["SPH_Kick1"] = sph_kick1
 
    # --------------------------------------------------------------------------
    #
    #
    # Correctness: Whenever we drift a particle, we might subsequently have to
    # resort it. That is, we assign it to another vertex. Therefore, it can
    # happen that another cell checks a particle again and that we notably call
    # touchVertexLastTime() with one particle twice. Therefore, we only know
    # that we touch a particle at least once in touchVertexLastTime, that cell
    # updates can be done multiple times (subsequent cell and vertex updates
    # will mask out any changes to the particle as the moved flag is set, but
    # they check), and that a cell update can follow a touchVertexLastTime()
    # (on a finer level if we move particles from fine to coarse).
    # --------------------------------------------------------------------------
    sph_drift = AlgorithmStep(
        name="SPH_Drift",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.CHANGE_POSITION_OR_INTERACTION_RADIUS,
        touch_vertex_first_time_kernel=v_iter.insert(
            "::exaparty::timestepping::resetMovedParticleMarker"
        ),
        touch_vertex_last_time_kernel=v_iter_for_move.insert(
            "::exaparty::kernels::legacy::leapfrogDriftWithGlobalTimeStepSize"
        ),
        includes=DefaultIncludes
        + """
            #include "Constants.h"
            #include "exaparty/kernels/legacy/Leapfrog.h"
            #include "exaparty/timestepping/GlobalTimeStepping.h"
            #include "exaparty/timestepping/TimeStepping.h"
        """,
        cell_kernel_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce_MayOverwritePreviousCellUpdatesFromSameSweep",
        touch_vertex_first_time_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce",
        touch_vertex_last_time_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce",
    )
    algorithm_steps_dict["SPH_Drift"] = sph_drift
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # From the SPH side of things, this doesn't need to be an AlgorithmStep
    # on its own. We could merge it together with the drifts. However, in order
    # to ensure that the particles which have just been moved in the preceeding
    # drift sweep are sorted correctly into the cells and vertices, we need an
    # extra sweep.
    predictHydro = AlgorithmStep(
        name="PredictHydro",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.ALTER_LOCAL_STATE,
        touch_vertex_first_time_kernel=v_iter.insert(
            "::exaparty::timestepping::resetMovedParticleMarker"
        )
        + v_iter.insert("::exaparty::kernels::legacy::hydroPredictExtra"),
        includes=DefaultIncludes
        + """
            #include "Constants.h"
            #include "exaparty/kernels/legacy/Leapfrog.h"
            #include "exaparty/timestepping/TimeStepping.h"
           """,
        touch_vertex_first_time_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce_MayOverwritePreviousCellUpdatesFromSameSweep",
        cell_kernel_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce_MayOverwritePreviousCellUpdatesFromSameSweep",
    )
    algorithm_steps_dict["PredictHydro"] = predictHydro
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    sph_densityWithConstantSearchRadius = AlgorithmStep(
        name="SPH_DensityLoopWithConstantSearchRadius",
        dependencies=AlgorithmStep.Dependencies.NEIGHBOURS,
        effect=AlgorithmStep.Effect.ALTER_LOCAL_STATE_AND_MIGHT_RERUN,
        prepare_traversal_kernel=f"""globaldata::{PARTICLE}::getSpecies().clearRerunPreviousGridSweepFlag();""",
        touch_vertex_first_time_kernel=v_iter.insert(
            "::exaparty::kernels::legacy::prepareDensity"
        ),
        cell_kernel=c_iter.insert("::exaparty::kernels::legacy::densityKernel"),
        touch_vertex_last_time_kernel=(
            v_iter.insert("::exaparty::kernels::legacy::endDensityCalculation")
            + v_iter_with_global_reduction.insert(
                "::exaparty::kernels::legacy::updateSmoothingLengthAndRerunIfRequired"
            )
        ),
        includes=DefaultIncludes
        + """
            #include "exaparty/kernels/legacy/SmoothingLength.h"
            #include "exaparty/kernels/legacy/Density.h"
        """,
    )
    algorithm_steps_dict["SPH_Density"] = sph_densityWithConstantSearchRadius
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    sph_force = AlgorithmStep(
        name="SPH_ForceLoop",
        dependencies=AlgorithmStep.Dependencies.NEIGHBOURS,
        effect=AlgorithmStep.Effect.ALTER_LOCAL_STATE,
        touch_vertex_first_time_kernel=(
            v_iter.insert(
                "::exaparty::kernels::legacy::resetSmoothingLengthIterationCounter"
            )
            + v_iter.insert("::exaparty::kernels::legacy::prepareHydroForce")
            + v_iter.insert("::exaparty::kernels::legacy::resetAcceleration")
        ),
        cell_kernel=c_iter.insert("::exaparty::kernels::legacy::forceKernel"),
        touch_vertex_last_time_kernel=v_iter.insert(
            "::exaparty::kernels::legacy::endHydroForceCalculation"
        ),
        includes=DefaultIncludes
        + """
          #include "exaparty/kernels/legacy/SmoothingLength.h"
          #include "exaparty/kernels/legacy/HydroForce.h"
        """,
    )
    algorithm_steps_dict["SPH_Force"] = sph_force
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    sph_kick2 = AlgorithmStep(
        name="SPH_Kick2",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.ALTER_LOCAL_STATE,
        touch_vertex_last_time_kernel=(
            v_iter.insert(
                "::exaparty::kernels::legacy::leapfrogKickWithGlobalTimeStepSize"
            )
            + v_iter.insert("::exaparty::kernels::legacy::resetPredictedValues")
        ),
        includes=DefaultIncludes
        + """
            #include "Constants.h"
            #include "exaparty/kernels/legacy/Leapfrog.h"
        """,
    )
    algorithm_steps_dict["SPH_Kick2"] = sph_kick2
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    timestep = AlgorithmStep(
        name="SPH_Timestep",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.ALTER_GLOBAL_STATE,
        touch_vertex_last_time_kernel="::exaparty::kernels::legacy::computeNextTimestep(marker, assignedParticles);",
        includes=DefaultIncludes
        + """
            #include "Constants.h"
            #include "exaparty/kernels/legacy/Timestep.h"
        """,
    )
    algorithm_steps_dict["SPH_Timestep"] = timestep
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    sph_reduceGlobals = AlgorithmStep(
        name="SPH_ReduceGlobalQuantities",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.ALTER_GLOBAL_STATE,
        touch_vertex_last_time_kernel=v_iter_with_global_reduction.insert(
            "::exaparty::statistics::reduceVelocityAndSearchRadius"
        ),
        prepare_traversal_kernel=f"""
                                   globaldata::{PARTICLE}::getSpecies().clearSearchRadius();
                                   globaldata::{PARTICLE}::getSpecies().clearVelocity();
        """,
        unprepare_traversal_kernel=f"""
                                   globaldata::{PARTICLE}::getSpecies().setTimeStamp( globaldata::{PARTICLE}::getSpecies().getMinTimeStamp() + globaldata::{PARTICLE}::getSpecies().getMinTimeStepSize(), false );
                                   ::exaparty::statistics::reportSearchRadiusVTDt<globaldata::{PARTICLE}>( "{PARTICLE}" );
        """,
        includes=DefaultIncludes
        + """
                                   #include "Constants.h"
                                   #include "exaparty/statistics/Reports.h"
                                   #include "exaparty/statistics/Statistics.h"
        """,
    )
    algorithm_steps_dict["SPH_ReduceGlobalQuantities"] = sph_reduceGlobals
 
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # The firstInitParticle step  may adapt a particle's search radius to fit
    # globally set limits. If the radius changes, and since this is the first
    # step in the run, a particle may be passed further down the tree, in which
    # case the dependency checks fail unless you permit the
    # `_MayOverwritePreviousCellUpdatesFromSameSweep` dependency policy.
    firstInit = AlgorithmStep(
        name="SPH_FirstInitParticle",
        dependencies=AlgorithmStep.Dependencies.SELF,
        effect=AlgorithmStep.Effect.CHANGE_POSITION_OR_INTERACTION_RADIUS,
        touch_vertex_last_time_kernel=v_iter_for_move.insert(
            "::exaparty::kernels::legacy::firstInitParticle"
        ),
        includes=DefaultIncludes
        + """
            #include "Constants.h"
            #include "exaparty/kernels/legacy/Swift.h"
        """,
        cell_kernel_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce_MayOverwritePreviousCellUpdatesFromSameSweep",
        touch_vertex_first_time_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce",
        touch_vertex_last_time_dependency_policy="TouchAtLeastOnce_AllPreviousStepsUpdateAtLeastOnce",
    )
    algorithm_steps_dict["SPH_FirstInitParticle"] = firstInit
 
    return algorithm_steps_dict