Date: 2025-03-29T05:58:17+00:00

Actual source code: petsclandau.h

      1: #pragma once

      3: #include <petscdmplex.h>
      4: #include <petscts.h>

      6: /* MANSEC = TS */
      7: /* SUBMANSEC = LANDAU */

      9: PETSC_EXTERN PetscErrorCode DMPlexLandauPrintNorms(Vec, PetscInt);
     10: PETSC_EXTERN PetscErrorCode DMPlexLandauCreateVelocitySpace(MPI_Comm, PetscInt, const char[], Vec *, Mat *, DM *);
     11: PETSC_EXTERN PetscErrorCode DMPlexLandauDestroyVelocitySpace(DM *);
     12: PETSC_EXTERN PetscErrorCode DMPlexLandauAccess(DM, Vec, PetscErrorCode (*)(DM, Vec, PetscInt, PetscInt, PetscInt, void *), void *);
     13: PETSC_EXTERN PetscErrorCode DMPlexLandauAddMaxwellians(DM, Vec, PetscReal, PetscReal[], PetscReal[], PetscInt, PetscInt, PetscInt, void *);
     14: PETSC_EXTERN PetscErrorCode DMPlexLandauCreateMassMatrix(DM dm, Mat *Amat);
     15: PETSC_EXTERN PetscErrorCode DMPlexLandauIFunction(TS, PetscReal, Vec, Vec, Vec, void *);
     16: PETSC_EXTERN PetscErrorCode DMPlexLandauIJacobian(TS, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *);

     18: typedef PetscInt LandauIdx;

     20: /* the Fokker-Planck-Landau context */
     21: #if !defined(LANDAU_MAX_SPECIES)
     22:   #if defined(PETSC_USE_DMLANDAU_2D)
     23:     #define LANDAU_MAX_SPECIES 10
     24:     #define LANDAU_MAX_GRIDS   3
     25:   #else
     26:     #define LANDAU_MAX_SPECIES 10
     27:     #define LANDAU_MAX_GRIDS   3
     28:   #endif
     29: #else
     30:   #define LANDAU_MAX_GRIDS 3
     31: #endif

     33: #if !defined(LANDAU_MAX_Q)
     34:   #if defined(LANDAU_MAX_NQND)
     35:     #error "LANDAU_MAX_NQND but not LANDAU_MAX_Q. Use -DLANDAU_MAX_Q=4 for Q3 elements"
     36:   #endif
     37:   #if defined(PETSC_USE_DMLANDAU_2D)
     38:     #define LANDAU_MAX_Q 6
     39:   #else
     40:     #define LANDAU_MAX_Q 5
     41:   #endif
     42: #else
     43:   #undef LANDAU_MAX_NQND
     44: #endif

     46: #if defined(PETSC_USE_DMLANDAU_2D)
     47:   #define LANDAU_MAX_Q_FACE   LANDAU_MAX_Q
     48:   #define LANDAU_MAX_NQND     (LANDAU_MAX_Q * LANDAU_MAX_Q)
     49:   #define LANDAU_MAX_BATCH_SZ 1024
     50:   #define LANDAU_DIM          2
     51: #else
     52:   #define LANDAU_MAX_Q_FACE   (LANDAU_MAX_Q * LANDAU_MAX_Q)
     53:   #define LANDAU_MAX_NQND     (LANDAU_MAX_Q * LANDAU_MAX_Q * LANDAU_MAX_Q)
     54:   #define LANDAU_MAX_BATCH_SZ 64
     55:   #define LANDAU_DIM          3
     56: #endif

     58: typedef enum {
     59:   LANDAU_KOKKOS,
     60:   LANDAU_CPU
     61: } LandauDeviceType;

     63: // static data - will be "device" data
     64: typedef struct {
     65:   void *invJ;    // nip*dim*dim
     66:   void *D;       // nq*nb*dim
     67:   void *B;       // nq*nb
     68:   void *alpha;   // ns
     69:   void *beta;    // ns
     70:   void *invMass; // ns
     71:   void *w;       // nip
     72:   void *x;       // nip
     73:   void *y;       // nip
     74:   void *z;       // nip
     75:   void *Eq_m;    // ns - dynamic
     76:   void *f;       //  nip*Nf - dynamic (IP)
     77:   void *dfdx;    // nip*Nf - dynamic (IP)
     78:   void *dfdy;    // nip*Nf - dynamic (IP)
     79:   void *dfdz;    // nip*Nf - dynamic (IP)
     80:   int   dim_, ns_, nip_, nq_, nb_;
     81:   void *NCells;         // remove and use elem_offset - TODO
     82:   void *species_offset; // for each grid, but same for all batched vertices
     83:   void *mat_offset;     // for each grid, but same for all batched vertices
     84:   void *elem_offset;    // for each grid, but same for all batched vertices
     85:   void *ip_offset;      // for each grid, but same for all batched vertices
     86:   void *ipf_offset;     // for each grid, but same for all batched vertices
     87:   void *ipfdf_data;     // for each grid, but same for all batched vertices
     88:   void *maps;           // for each grid, but same for all batched vertices
     89:   // COO
     90:   void     *coo_elem_offsets;
     91:   void     *coo_elem_point_offsets;
     92:   void     *coo_elem_fullNb;
     93:   void     *coo_vals;
     94:   void     *lambdas;
     95:   LandauIdx coo_n_cellsTot;
     96:   LandauIdx coo_size;
     97:   LandauIdx coo_max_fullnb;
     98: } LandauStaticData;

    100: typedef enum {
    101:   LANDAU_EX2_TSSOLVE,
    102:   LANDAU_MATRIX_TOTAL,
    103:   LANDAU_OPERATOR,
    104:   LANDAU_JACOBIAN_COUNT,
    105:   LANDAU_JACOBIAN,
    106:   LANDAU_MASS,
    107:   LANDAU_F_DF,
    108:   LANDAU_KERNEL,
    109:   KSP_FACTOR,
    110:   KSP_SOLVE,
    111:   LANDAU_NUM_TIMERS
    112: } LandauOMPTimers;

    114: typedef struct {
    115:   PetscBool interpolate; /* Generate intermediate mesh elements */
    116:   PetscBool gpu_assembly;
    117:   MPI_Comm  comm; /* global communicator to use for errors and diagnostics */
    118:   double    times[LANDAU_NUM_TIMERS];
    119:   PetscBool use_matrix_mass;
    120:   /* FE */
    121:   PetscFE fe[LANDAU_MAX_SPECIES];
    122:   /* geometry  */
    123:   PetscReal radius[LANDAU_MAX_GRIDS];
    124:   PetscReal radius_par[LANDAU_MAX_GRIDS];
    125:   PetscReal radius_perp[LANDAU_MAX_GRIDS];
    126:   PetscReal re_radius;      /* RE: radius of refinement along v_perp=0, z>0 */
    127:   PetscReal vperp0_radius1; /* RE: radius of refinement along v_perp=0 */
    128:   PetscReal vperp0_radius2; /* RE: radius of refinement along v_perp=0 after origin AMR refinement */
    129:   PetscBool sphere;
    130:   PetscReal sphere_inner_radius_90degree[LANDAU_MAX_GRIDS];
    131:   PetscReal sphere_inner_radius_45degree[LANDAU_MAX_GRIDS];
    132:   PetscInt  cells0[3];
    133:   /* AMR */
    134:   PetscBool use_p4est;
    135:   PetscInt  numRERefine;                     /* RE: refinement along v_perp=0, z > 0 */
    136:   PetscInt  nZRefine1;                       /* RE: origin refinement after v_perp=0 refinement */
    137:   PetscInt  nZRefine2;                       /* RE: origin refinement after origin AMR refinement */
    138:   PetscInt  numAMRRefine[LANDAU_MAX_GRIDS];  /* normal AMR - refine from origin */
    139:   PetscInt  postAMRRefine[LANDAU_MAX_GRIDS]; /* uniform refinement of AMR */
    140:   PetscBool simplex;
    141:   char      filename[PETSC_MAX_PATH_LEN];
    142:   PetscReal thermal_speed[LANDAU_MAX_GRIDS];
    143:   PetscBool sphere_uniform_normal;
    144:   /* relativistic */
    145:   PetscBool use_energy_tensor_trick;
    146:   PetscBool use_relativistic_corrections;
    147:   /* physics */
    148:   PetscReal thermal_temps[LANDAU_MAX_SPECIES];
    149:   PetscReal masses[LANDAU_MAX_SPECIES];  /* mass of each species  */
    150:   PetscReal charges[LANDAU_MAX_SPECIES]; /* charge of each species  */
    151:   PetscReal n[LANDAU_MAX_SPECIES];       /* number density of each species  */
    152:   PetscReal m_0;                         /* reference mass */
    153:   PetscReal v_0;                         /* reference velocity */
    154:   PetscReal n_0;                         /* reference number density */
    155:   PetscReal t_0;                         /* reference time */
    156:   PetscReal Ez;
    157:   PetscReal epsilon0;
    158:   PetscReal k;
    159:   PetscReal lambdas[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS];
    160:   PetscReal electronShift;
    161:   PetscInt  num_species;
    162:   PetscInt  num_grids;
    163:   PetscInt  species_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices
    164:   PetscInt  mat_offset[LANDAU_MAX_GRIDS + 1];     // for each grid, but same for all batched vertices
    165:   // batching
    166:   PetscBool  jacobian_field_major_order; // this could be a type but lets not get pedantic
    167:   VecScatter plex_batch;
    168:   Vec        work_vec;
    169:   IS         batch_is;
    170:   PetscErrorCode (*seqaij_mult)(Mat, Vec, Vec);
    171:   PetscErrorCode (*seqaij_multtranspose)(Mat, Vec, Vec);
    172:   PetscErrorCode (*seqaij_solve)(Mat, Vec, Vec);
    173:   PetscErrorCode (*seqaij_getdiagonal)(Mat, Vec);
    174:   /* COO */
    175:   Mat J;
    176:   Mat M;
    177:   Vec X;
    178:   /* derived type */
    179:   void *data;
    180:   /* computing */
    181:   LandauDeviceType deviceType;
    182:   DM               pack;
    183:   DM               plex[LANDAU_MAX_GRIDS];
    184:   LandauStaticData SData_d; /* static geometric data on device */
    185:   /* diagnostics */
    186:   PetscInt         verbose;
    187:   PetscLogEvent    events[20];
    188:   PetscLogStage    stage;
    189:   PetscObjectState norm_state;
    190:   PetscInt         batch_sz;
    191:   PetscInt         batch_view_idx;
    192: } LandauCtx;

    194: #define LANDAU_SPECIES_MAJOR
    195: #if !defined(LANDAU_SPECIES_MAJOR)
    196:   #define LAND_PACK_IDX(_b, _g)                         (_b * ctx->num_grids + _g)
    197:   #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_b * _mat_off[_ngrid] + _mat_off[_g])
    198: #else
    199:   #define LAND_PACK_IDX(_b, _g)                         (_g * ctx->batch_sz + _b)
    200:   #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_nbch * _mat_off[_g] + _b * (_mat_off[_g + 1] - _mat_off[_g]))
    201: #endif

    203: typedef struct {
    204:   PetscReal scale;
    205:   LandauIdx gid; // Landau matrix index (<10,000)
    206: } pointInterpolationP4est;
    207: typedef struct _lP4estVertexMaps {
    208:   LandauIdx (*gIdx)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND]; // #elems *  LANDAU_MAX_NQND
    209:   LandauIdx        num_elements;
    210:   LandauIdx        num_reduced;
    211:   LandauIdx        num_face; // (Q or Q^2 for 3D)
    212:   LandauDeviceType deviceType;
    213:   PetscInt         Nf;
    214:   pointInterpolationP4est (*c_maps)[LANDAU_MAX_Q_FACE];
    215:   struct _lP4estVertexMaps *d_self;
    216:   void                     *vp1, *vp2, *vp3;
    217:   PetscInt                  numgrids;
    218: } P4estVertexMaps;

    220: #if defined(PETSC_HAVE_KOKKOS)
    221: PETSC_EXTERN PetscErrorCode LandauKokkosJacobian(DM[], const PetscInt, const PetscInt, const PetscInt, const PetscInt, const PetscInt[], PetscReal[], PetscScalar[], const PetscScalar[], const LandauStaticData *, const PetscReal, const PetscLogEvent[], const PetscInt[], const PetscInt[], Mat[], Mat);
    222: PETSC_EXTERN PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt);
    223: PETSC_EXTERN PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *, PetscInt);
    224: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataSet(DM, const PetscInt, const PetscInt, const PetscInt, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], LandauStaticData *);
    225: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *);
    226: #endif