MpdRdmeSolverExtended.h

Go to the documentation of this file.

/*
 * University of Illinois Open Source License
 * Copyright 2024 Luthey-Schulten Group,
 * All rights reserved.
 *
 * Extended RDME solver with particle size support (ByteLattice compatible)
 */
 
#ifndef LM_MPDRDMESOLVEREXTENDED_H_
#define LM_MPDRDMESOLVEREXTENDED_H_
 
#include "MpdRdmeSolver.h"
#include "CudaByteLattice.h"
#include "CudaByteLatticeExtended.h"
#include "cuda/particle_size_kernels.cuh"
 
using lm::rdme::MpdRdmeSolver;
using lm::rdme::CudaByteLattice;
using lm::rdme::CudaByteLatticeExtended;
using lm::rdme::Lattice;
 
 
namespace lm {

class MpdRdmeSolverExtended : public lm::rdme::MpdRdmeSolver
{
public:
    MpdRdmeSolverExtended();
    virtual ~MpdRdmeSolverExtended();
 
    // Particle size management
    void setParticleSizes(const uint32_t* sizes, uint32_t numTypes);
    void enableParticleSizeChecking(bool enable);
    uint32_t getParticleSize(particle_t particleType) const;
 
    // Validation methods
    uint32_t validateConstraints() const;
    void printConstraintViolations() const;
 
    // Efficient size checking methods
    void initializeFastSizeChecking();
    void useFastSizeAwareDiffusion(bool enable);
    bool isFastSizeCheckingEnabled() const { return useFastSizeKernels_; }
 
    // Override solver methods
    virtual void buildModel(const uint numberSpeciesA,
                           const uint numberReactionsA,
                           const uint* initialSpeciesCountsA,
                           const uint* reactionTypesA,
                           const double* KA,
                           const int* SA,
                           const uint* DA,
                           const uint kCols);
 
    virtual void allocateLattice(lattice_size_t latticeXSize,
                               lattice_size_t latticeYSize,
                               lattice_size_t latticeZSize,
                               site_size_t particlesPerSite,
                               const unsigned int bytes_per_particle,
                               si_dist_t latticeSpacing);
 
    virtual void buildDiffusionModel(const uint numberSiteTypesA,
                                   const double* DFA,
                                   const uint* RLA,
                                   lattice_size_t latticeXSize,
                                   lattice_size_t latticeYSize,
                                   lattice_size_t latticeZSize,
                                   site_size_t particlesPerSite,
                                   const unsigned int bytes_per_particle,
                                   si_dist_t latticeSpacing,
                                   const uint8_t* latticeData,
                                   const uint8_t* latticeSitesData,
                                   bool rowMajorData = true);
 
    virtual void generateTrajectory();
    virtual void runTimestep(CudaByteLattice* lattice, uint32_t timestep);
 
protected:
    // Particle size state
    bool particleSizeEnabled;
    bool particleSizesInitialized;
    uint32_t numParticleTypes;
    uint32_t* hostParticleSizes;
    bool useFastSizeKernels_;  // Flag to use new efficient kernels
 
    // Reaction data for GPU kernels
    int8_t* deviceReactionMatrix;
    uint8_t* deviceReactionLocationMatrix;
 
    // Helper methods
    void initializeParticleSizes();
    void copyParticleSizesToDevice();
    void copyReactionDataToDevice();
    void generateSizeAwareTrajectory();
    void runSizeAwareTimestep(lm::rdme::CudaByteLatticeExtended* lattice, uint32_t timestep);
    void runStandardTimestep(lm::rdme::CudaByteLatticeExtended* lattice, uint32_t timestep);
};
 
} // namespace lm
 
#endif // LM_MPDRDMESOLVEREXTENDED_H_