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 6
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: PetscBool map_sphere;
131: PetscReal sphere_inner_radius_90degree[LANDAU_MAX_GRIDS];
132: PetscReal sphere_inner_radius_45degree[LANDAU_MAX_GRIDS];
133: PetscInt cells0[3];
134: /* AMR */
135: PetscBool use_p4est;
136: PetscInt numRERefine; /* RE: refinement along v_perp=0, z > 0 */
137: PetscInt nZRefine1; /* RE: origin refinement after v_perp=0 refinement */
138: PetscInt nZRefine2; /* RE: origin refinement after origin AMR refinement */
139: PetscInt numAMRRefine[LANDAU_MAX_GRIDS]; /* normal AMR - refine from origin */
140: PetscInt postAMRRefine[LANDAU_MAX_GRIDS]; /* uniform refinement of AMR */
141: PetscBool simplex;
142: char filename[PETSC_MAX_PATH_LEN];
143: PetscReal thermal_speed[LANDAU_MAX_GRIDS];
144: PetscBool sphere_uniform_normal;
145: /* relativistic */
146: PetscBool use_energy_tensor_trick;
147: PetscBool use_relativistic_corrections;
148: /* physics */
149: PetscReal thermal_temps[LANDAU_MAX_SPECIES];
150: PetscReal masses[LANDAU_MAX_SPECIES]; /* mass of each species */
151: PetscReal charges[LANDAU_MAX_SPECIES]; /* charge of each species */
152: PetscReal n[LANDAU_MAX_SPECIES]; /* number density of each species */
153: PetscReal m_0; /* reference mass */
154: PetscReal v_0; /* reference velocity */
155: PetscReal n_0; /* reference number density */
156: PetscReal t_0; /* reference time */
157: PetscReal Ez;
158: PetscReal epsilon0;
159: PetscReal k;
160: PetscReal lambdas[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS];
161: PetscReal electronShift;
162: PetscInt num_species;
163: PetscInt num_grids;
164: PetscInt species_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices
165: PetscInt mat_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices
166: // batching
167: PetscBool jacobian_field_major_order; // this could be a type but lets not get pedantic
168: VecScatter plex_batch;
169: Vec work_vec;
170: IS batch_is;
171: PetscErrorCode (*seqaij_mult)(Mat, Vec, Vec);
172: PetscErrorCode (*seqaij_multtranspose)(Mat, Vec, Vec);
173: PetscErrorCode (*seqaij_solve)(Mat, Vec, Vec);
174: PetscErrorCode (*seqaij_getdiagonal)(Mat, Vec);
175: /* COO */
176: Mat J;
177: Mat M;
178: Vec X;
179: /* derived type */
180: void *data;
181: /* computing */
182: LandauDeviceType deviceType;
183: DM pack;
184: DM plex[LANDAU_MAX_GRIDS];
185: LandauStaticData SData_d; /* static geometric data on device */
186: /* diagnostics */
187: PetscInt verbose;
188: PetscLogEvent events[20];
189: PetscLogStage stage;
190: PetscObjectState norm_state;
191: PetscInt batch_sz;
192: PetscInt batch_view_idx;
193: } LandauCtx;
195: #define LANDAU_SPECIES_MAJOR
196: #if !defined(LANDAU_SPECIES_MAJOR)
197: #define LAND_PACK_IDX(_b, _g) (_b * ctx->num_grids + _g)
198: #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_b * _mat_off[_ngrid] + _mat_off[_g])
199: #else
200: #define LAND_PACK_IDX(_b, _g) (_g * ctx->batch_sz + _b)
201: #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_nbch * _mat_off[_g] + _b * (_mat_off[_g + 1] - _mat_off[_g]))
202: #endif
204: typedef struct {
205: PetscReal scale;
206: LandauIdx gid; // Landau matrix index (<10,000)
207: } pointInterpolationP4est;
208: typedef struct _lP4estVertexMaps {
209: LandauIdx (*gIdx)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND]; // #elems * LANDAU_MAX_NQND
210: LandauIdx num_elements;
211: LandauIdx num_reduced;
212: LandauIdx num_face; // (Q or Q^2 for 3D)
213: LandauDeviceType deviceType;
214: PetscInt Nf;
215: pointInterpolationP4est (*c_maps)[LANDAU_MAX_Q_FACE];
216: struct _lP4estVertexMaps *d_self;
217: void *vp1, *vp2, *vp3;
218: PetscInt numgrids;
219: } P4estVertexMaps;
221: #if defined(PETSC_HAVE_KOKKOS)
222: 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);
223: PETSC_EXTERN PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt);
224: PETSC_EXTERN PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *, PetscInt);
225: 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 *);
226: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *);
227: #endif