RFOF functions should now be callable via libascot

ascot4fusion · Jul 15, 2024 · 8800240 · 8800240
1 parent 5cdecc9
commit 8800240
Show file tree

Hide file tree

Showing 11 changed files with 291 additions and 170 deletions.
diff --git a/a5py/ascotpy/__init__.py b/a5py/ascotpy/__init__.py
@@ -12,7 +12,7 @@
 from a5py.routines.plotting import openfigureifnoaxes, plt
 from a5py.exceptions import AscotInitException
 
-from .libascot    import LibAscot, _LIBASCOT
+from .libascot    import LibAscot, _LIBASCOT, PTR_RFOF
 from .libsimulate import LibSimulate
 from .libproviders import LibProviders
 
@@ -104,6 +104,7 @@ def __init__(self):
         self._asigma_offload_array  = ctypes.POINTER(ctypes.c_double)()
         self._nbi_offload_array     = ctypes.POINTER(ctypes.c_double)()
         self._diag_offload_array    = ctypes.POINTER(ctypes.c_double)()
+        self._rfof_initialized      = False
 
         self._wall_int_offload_array = ctypes.POINTER(ctypes.c_int)()
 
@@ -424,6 +425,25 @@ def _requireinit(self, *inputs):
             if qid == Ascotpy.DUMMY_QID:
                 raise AscotInitException(inp + " is not initialized")
 
+    def input_init_rfof(self):
+        """
+        """
+        if self._rfof_initialized:
+            raise AscotInitException("RFOF input is already initialized")
+        fun = _LIBASCOT.rfof_interface_initev_excl_marker_stuff
+        fun.restype  = None
+        fun.argtypes = [PTR_RFOF]
+        fun(ctypes.byref(self._sim.rfof_data))
+        self._rfof_initialized = True
+
+    def input_free_rfof(self):
+        """
+        """
+        if not self._rfof_initialized:
+            raise AscotInitException("RFOF input is not initialized")
+
+        self._rfof_initialized = False
+
     def input_initialized(self):
         """Get inputs that are currently initialized.
 

diff --git a/a5py/ascotpy/ascot2py.py b/a5py/ascotpy/ascot2py.py
@@ -2779,6 +2779,76 @@ class struct_c__SA_nbi_data(Structure):
 nbi_inject = _libraries['libascot.so'].nbi_inject
 nbi_inject.restype = None
 nbi_inject.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(struct_c__SA_nbi_injector), ctypes.POINTER(ctypes.POINTER(None))]
+class struct_prt_rfof(Structure):
+    pass
+
+struct_prt_rfof._pack_ = 1 # source:False
+struct_prt_rfof._fields_ = [
+    ('dmu', ctypes.c_double),
+    ('dvpar', ctypes.c_double),
+    ('de', ctypes.c_double),
+    ('deCumulative', ctypes.c_double),
+    ('dpitch', ctypes.c_double),
+    ('maxAcc', ctypes.c_double),
+    ('RFdt', ctypes.c_double),
+]
+
+prt_rfof = struct_prt_rfof
+class struct_c__SA_rfof_data(Structure):
+    pass
+
+struct_c__SA_rfof_data._pack_ = 1 # source:False
+struct_c__SA_rfof_data._fields_ = [
+    ('cptr_rfof_input_params', ctypes.POINTER(None)),
+    ('cptr_rfglobal', ctypes.POINTER(None)),
+]
+
+rfof_data = struct_c__SA_rfof_data
+rfof_interface_initev_excl_marker_stuff = _libraries['libascot.so'].rfof_interface_initev_excl_marker_stuff
+rfof_interface_initev_excl_marker_stuff.restype = None
+rfof_interface_initev_excl_marker_stuff.argtypes = [ctypes.POINTER(struct_c__SA_rfof_data)]
+rfof_init = _libraries['libascot.so'].rfof_init
+rfof_init.restype = None
+rfof_init.argtypes = [ctypes.POINTER(struct_c__SA_rfof_data), ctypes.POINTER(struct_c__SA_rfof_data)]
+rfof_interface_initialise_res_mem = _libraries['libascot.so'].rfof_interface_initialise_res_mem
+rfof_interface_initialise_res_mem.restype = None
+rfof_interface_initialise_res_mem.argtypes = [ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_initialise_diagnostics = _libraries['libascot.so'].rfof_interface_initialise_diagnostics
+rfof_interface_initialise_diagnostics.restype = None
+rfof_interface_initialise_diagnostics.argtypes = [ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_allocate_rfof_marker = _libraries['libascot.so'].rfof_interface_allocate_rfof_marker
+rfof_interface_allocate_rfof_marker.restype = None
+rfof_interface_allocate_rfof_marker.argtypes = [ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_set_marker_pointers = _libraries['libascot.so'].rfof_interface_set_marker_pointers
+rfof_interface_set_marker_pointers.restype = None
+rfof_interface_set_marker_pointers.argtypes = [ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
+rfof_interface_do_rfof_stuff_gc = _libraries['libascot.so'].rfof_interface_do_rfof_stuff_gc
+rfof_interface_do_rfof_stuff_gc.restype = None
+rfof_interface_do_rfof_stuff_gc.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), rfof_data, ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
+rfof_interface_reset_icrh_mem = _libraries['libascot.so'].rfof_interface_reset_icrh_mem
+rfof_interface_reset_icrh_mem.restype = None
+rfof_interface_reset_icrh_mem.argtypes = [ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
+rfof_interface_deallocate_rfof_input_param = _libraries['libascot.so'].rfof_interface_deallocate_rfof_input_param
+rfof_interface_deallocate_rfof_input_param.restype = None
+rfof_interface_deallocate_rfof_input_param.argtypes = [ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_deallocate_rfglobal = _libraries['libascot.so'].rfof_interface_deallocate_rfglobal
+rfof_interface_deallocate_rfglobal.restype = None
+rfof_interface_deallocate_rfglobal.argtypes = [ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_deallocate_res_mem = _libraries['libascot.so'].rfof_interface_deallocate_res_mem
+rfof_interface_deallocate_res_mem.restype = None
+rfof_interface_deallocate_res_mem.argtypes = [ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
+rfof_interface_deallocate_diagnostics = _libraries['libascot.so'].rfof_interface_deallocate_diagnostics
+rfof_interface_deallocate_diagnostics.restype = None
+rfof_interface_deallocate_diagnostics.argtypes = [ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_deallocate_marker = _libraries['libascot.so'].rfof_interface_deallocate_marker
+rfof_interface_deallocate_marker.restype = None
+rfof_interface_deallocate_marker.argtypes = [ctypes.POINTER(ctypes.POINTER(None))]
+rfof_interface_get_rf_wave_local = _libraries['libascot.so'].rfof_interface_get_rf_wave_local
+rfof_interface_get_rf_wave_local.restype = None
+rfof_interface_get_rf_wave_local.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]
+rfof_interface_eval_resonance_function = _libraries['libascot.so'].rfof_interface_eval_resonance_function
+rfof_interface_eval_resonance_function.restype = None
+rfof_interface_eval_resonance_function.argtypes = [ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.POINTER(None)), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_int32)]
 
 # values for enumeration 'SIMULATION_MODE'
 SIMULATION_MODE__enumvalues = {
@@ -2807,6 +2877,7 @@ class struct_c__SA_sim_offload_data(Structure):
     ('asigma_offload_data', asigma_offload_data),
     ('nbi_offload_data', nbi_offload_data),
     ('diag_offload_data', diag_offload_data),
+    ('rfof_data', rfof_data),
     ('sim_mode', ctypes.c_int32),
     ('enable_ada', ctypes.c_int32),
     ('record_mode', ctypes.c_int32),
@@ -2875,17 +2946,6 @@ class struct_c__SA_mccc_data(Structure):
     ('include_gcdiff', ctypes.c_int32),
 ]
 
-class struct_c__SA_rfof_data(Structure):
-    pass
-
-struct_c__SA_rfof_data._pack_ = 1 # source:False
-struct_c__SA_rfof_data._fields_ = [
-    ('cptr_rfof_input_params', ctypes.POINTER(None)),
-    ('cptr_rfglobal', ctypes.POINTER(None)),
-    ('icrh_initialised', ctypes.c_int32),
-    ('PADDING_0', ctypes.c_ubyte * 4),
-]
-
 struct_c__SA_sim_data._pack_ = 1 # source:False
 struct_c__SA_sim_data._fields_ = [
     ('B_data', B_field_data),
@@ -2900,7 +2960,7 @@ class struct_c__SA_rfof_data(Structure):
     ('diag_data', diag_data),
     ('random_data', ctypes.POINTER(None)),
     ('mccc_data', struct_c__SA_mccc_data),
-    ('rfof_data', struct_c__SA_rfof_data),
+    ('rfof_data', rfof_data),
     ('sim_mode', ctypes.c_int32),
     ('enable_ada', ctypes.c_int32),
     ('record_mode', ctypes.c_int32),
@@ -3283,12 +3343,27 @@ class struct_c__SA_afsi_data(Structure):
     'plasma_init', 'plasma_init_offload', 'plasma_offload_data',
     'plasma_type', 'plasma_type_1D', 'plasma_type_1DS',
     'plasma_type_1Dt', 'prepare_markers', 'print_marker_summary',
-    'real', 'sigma_CX', 'sigma_ioniz', 'sigma_recomb', 'sigmav_BMS',
-    'sigmav_CX', 'sigmav_ioniz', 'sigmav_recomb', 'sigmaveff_CX',
-    'sigmaveff_ioniz', 'sigmaveff_recomb', 'sim_data', 'sim_init',
-    'sim_offload_data', 'simulate', 'simulate_init_offload',
-    'simulate_mode_fo', 'simulate_mode_gc', 'simulate_mode_hybrid',
-    'simulate_mode_ml', 'size_t', 'struct_c__SA_B_2DS_data',
+    'prt_rfof', 'real', 'rfof_data', 'rfof_init',
+    'rfof_interface_allocate_rfof_marker',
+    'rfof_interface_deallocate_diagnostics',
+    'rfof_interface_deallocate_marker',
+    'rfof_interface_deallocate_res_mem',
+    'rfof_interface_deallocate_rfglobal',
+    'rfof_interface_deallocate_rfof_input_param',
+    'rfof_interface_do_rfof_stuff_gc',
+    'rfof_interface_eval_resonance_function',
+    'rfof_interface_get_rf_wave_local',
+    'rfof_interface_initev_excl_marker_stuff',
+    'rfof_interface_initialise_diagnostics',
+    'rfof_interface_initialise_res_mem',
+    'rfof_interface_reset_icrh_mem',
+    'rfof_interface_set_marker_pointers', 'sigma_CX', 'sigma_ioniz',
+    'sigma_recomb', 'sigmav_BMS', 'sigmav_CX', 'sigmav_ioniz',
+    'sigmav_recomb', 'sigmaveff_CX', 'sigmaveff_ioniz',
+    'sigmaveff_recomb', 'sim_data', 'sim_init', 'sim_offload_data',
+    'simulate', 'simulate_init_offload', 'simulate_mode_fo',
+    'simulate_mode_gc', 'simulate_mode_hybrid', 'simulate_mode_ml',
+    'size_t', 'struct_c__SA_B_2DS_data',
     'struct_c__SA_B_2DS_offload_data', 'struct_c__SA_B_3DS_data',
     'struct_c__SA_B_3DS_offload_data', 'struct_c__SA_B_GS_data',
     'struct_c__SA_B_GS_offload_data', 'struct_c__SA_B_STS_data',
@@ -3343,7 +3418,7 @@ class struct_c__SA_afsi_data(Structure):
     'struct_c__SA_wall_2d_offload_data', 'struct_c__SA_wall_3d_data',
     'struct_c__SA_wall_3d_offload_data', 'struct_c__SA_wall_data',
     'struct_c__SA_wall_offload_data', 'struct_diag_transcoef_link',
-    'union_c__SA_input_particle_0', 'wall_2d_data',
+    'struct_prt_rfof', 'union_c__SA_input_particle_0', 'wall_2d_data',
     'wall_2d_find_intersection', 'wall_2d_free_offload',
     'wall_2d_hit_wall', 'wall_2d_init', 'wall_2d_init_offload',
     'wall_2d_inside', 'wall_2d_offload_data', 'wall_3d_data',

diff --git a/a5py/ascotpy/libascot.py b/a5py/ascotpy/libascot.py
@@ -39,6 +39,7 @@ def from_param(cls, obj):
     PTR_REAL = _ndpointerwithnull(ctypes.c_double, flags="C_CONTIGUOUS")
     PTR_INT  = _ndpointerwithnull(ctypes.c_int,    flags="C_CONTIGUOUS")
     PTR_SIM  = ctypes.POINTER(ascot2py.struct_c__SA_sim_offload_data)
+    PTR_RFOF = ctypes.POINTER(ascot2py.struct_c__SA_rfof_data)
     PTR_ARR  = ctypes.POINTER(ctypes.c_double)
     STRUCT_DIST5DOFFLOAD = ascot2py.struct_c__SA_dist_5D_offload_data
     STRUCT_DIST5D        = ascot2py.struct_c__SA_dist_5D_data
@@ -52,6 +53,7 @@ def from_param(cls, obj):
     PTR_REAL  = None
     PTR_INT   = None
     PTR_SIM   = None
+    PTR_RFOF  = None
     PTR_ARR   = None
     STRUCT_DIST5DOFFLOAD = None
     STRUCT_DIST5D        = None
@@ -879,3 +881,62 @@ def input_findpsi0(self, psi1):
             raise RuntimeError("Failed to converge.")
 
         return (rz[0], rz[1], psi)
+
+    @parseunits(m="kg", q="C", vpar="m/s", r="m", phi="rad", z="m", t="s")
+    def input_eval_rfof(self, m, q, vpar, r, phi, z, t):
+        """Evaluate Evaluate ICRH electric field and the resonance condition.
+
+        The evaluated electric field consists of left-hand (-) and right-hand (+)
+        circularly polarized components. The resonance condition is given by
+
+        omega_wave - n * omega_gyro - k_parallel * v_parallel
+        - k_perp dot v_drift = 0.
+
+        Parameters
+        ----------
+        m : float
+            Test particle mass (for calculating the resonance).
+        q : float
+            Test particle charge (for calculating the resonance).
+        vpar : float
+            Test particle parallel velocity (for calculating the resonance).
+        r : array_like, (n,)
+            R coordinates where data is evaluated.
+        phi : array_like (n,)
+            phi coordinates where data is evaluated.
+        z : array_like (n,)
+            z coordinates where data is evaluated.
+        t : array_like (n,)
+            Time coordinates where data is evaluated.
+
+        Returns
+        -------
+        eplus : array_like, (n,)
+            Reaction cross-section.
+        eminus : array_like, (n,)
+        rescond : array_like, (n,)
+
+        Raises
+        ------
+        AssertionError
+            If required data has not been initialized.
+        """
+        self._requireinit("bfield")
+        if not self._rfof_initialized:
+            raise RuntimeError("RFOF data not initialized.")
+
+        Neval = r.size
+        eplus    = np.NaN * np.zeros((Neval,), dtype="f8") * unyt.V/unyt.m
+        eminus   = np.NaN * np.zeros((Neval,), dtype="f8") * unyt.V/unyt.m
+        res_cond = np.NaN * np.zeros((Neval,), dtype="f8") * unyt.dimensionless
+
+        fun = _LIBASCOT.libascot_eval_rfof
+        fun.restype  = None
+        fun.argtypes = [PTR_SIM, PTR_ARR,
+                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
+                        ctypes.c_double, ctypes.c_double, ctypes.c_double,
+                        PTR_REAL, PTR_REAL, PTR_REAL]
+        fun(ctypes.byref(self._sim), self._bfield_offload_array,
+            Neval, r, phi, z, t, m, q, vpar, eplus, eminus, res_cond)
+
+        return eplus, eminus, res_cond
diff --git a/src/Makefile b/src/Makefile
@@ -241,7 +241,7 @@ ASCOT2PY_HEADERFILES=ascot5.h particle.h mpi_interface.h endcond.h simulate.h  \
 	$(ASIGMAHEADERS) $(N0HEADERS) $(MHDHEADERS) hdf5_interface.h           \
 	libascot_mem.h diag.h B_field.h E_field.h wall.h plasma.h neutral.h    \
 	boozer.h mhd.h nbi.h asigma.h afsi.h ascot5_main.h biosaw.h boschhale.h\
-	bbnbi5.h
+	bbnbi5.h rfof_interface.h
 
 ascot2py.py : libascot.so
 	$(eval CLANGOMP=$(shell clang --print-resource-dir)) \

diff --git a/src/ascot5.h b/src/ascot5.h
@@ -130,4 +130,7 @@ typedef double real;  /**< Double precision float   */
 /** @brief Default depth of octree struct */
 #define WALL_OCTREE_DEPTH 7
 
+/** @brief Input filename where RFOF parameters are stored */
+#define RFOF_CODEPARAM_XML "rfof_codeparam.xml"
+
 #endif
diff --git a/src/ascot5_main.c b/src/ascot5_main.c
@@ -179,6 +179,10 @@ int main(int argc, char** argv) {
         return 1;
     };
 
+    if(sim.enable_icrh) {
+        rfof_interface_initev_excl_marker_stuff(&(sim.rfof_data));
+    }
+
     /* Initialize marker states array ps and free marker input p */
     int n_proc; /* Number of markers allocated for this MPI process */
     particle_state* ps;
@@ -228,6 +232,12 @@ int main(int argc, char** argv) {
     /* Free input data */
     offload_free_offload(&offload_data, &offload_array, &int_offload_array);
 
+    if(sim.enable_icrh) {
+        rfof_interface_deallocate_rfof_input_param(
+            &(sim.rfof_data.cptr_rfof_input_params));
+        rfof_interface_deallocate_rfglobal(&(sim.rfof_data.cptr_rfglobal));
+    }
+
     /* Write output and clean */
     if( write_output(&sim, pout, n_gathered, diag_offload_array) ) {
         goto CLEANUP_FAILURE;