Skip to content

Commit

Permalink
BBNBI can be run via libascot
Browse files Browse the repository at this point in the history
  • Loading branch information
miekkasarki committed Jul 12, 2024
1 parent f8c12b6 commit 2a09dca
Show file tree
Hide file tree
Showing 8 changed files with 415 additions and 154 deletions.
8 changes: 5 additions & 3 deletions a5py/ascotpy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ def _initmpi(self, mpirank, mpisize, mpiroot=0):

def _init(self, data, bfield=None, efield=None, plasma=None,
wall=None, neutral=None, boozer=None, mhd=None, asigma=None,
switch=False):
nbi=None, switch=False):
"""Read, offload, and initialize input data so it can be accessed
by libascot.
Expand Down Expand Up @@ -168,6 +168,8 @@ def _init(self, data, bfield=None, efield=None, plasma=None,
QID of the MHD data to be initialized or the data as a dictionary.
asigma : str or dict
QID of the atomicdata to be initialized or the data as a dictionary.
nbi : str or dict
QID of the NBI data to be initialized or the data as a dictionary.
switch : bool
If ``True``, free input that has been
"""
Expand All @@ -179,7 +181,7 @@ def _init(self, data, bfield=None, efield=None, plasma=None,
args = locals() # Contains function arguments and values in a dictionary
to_be_provided = [] # Inputs to be directly injected (provided)
for inp in ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
"mhd", "asigma"]:
"mhd", "asigma", "nbi"]:
if args[inp] is None:
# This input is not going to be initialized
continue
Expand Down Expand Up @@ -225,7 +227,7 @@ def _init(self, data, bfield=None, efield=None, plasma=None,
# an exception, in which case sim.qid_* would point to data which is not
# initialized.
for inp in ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
"mhd", "asigma"]:
"mhd", "asigma", "nbi"]:
if inputs2read.value & getattr(ascot2py, "hdf5_input_" + inp):
setattr(self._sim, "qid_" + inp, args[inp])

Expand Down
36 changes: 23 additions & 13 deletions a5py/ascotpy/ascot2py.py
Original file line number Diff line number Diff line change
Expand Up @@ -2615,6 +2615,9 @@ class struct_c__SA_asigma_loc_data(Structure):
asigma_loc_eval_sigmav = _libraries['libascot.so'].asigma_loc_eval_sigmav
asigma_loc_eval_sigmav.restype = a5err
asigma_loc_eval_sigmav.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, real, ctypes.c_int32, ctypes.c_int32, real, real, real, real, ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(struct_c__SA_asigma_loc_data)]
asigma_loc_eval_cx = _libraries['libascot.so'].asigma_loc_eval_cx
asigma_loc_eval_cx.restype = a5err
asigma_loc_eval_cx.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, real, real, ctypes.c_int32, ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), real, ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.POINTER(struct_c__SA_asigma_loc_data)]
asigma_loc_eval_bms = _libraries['libascot.so'].asigma_loc_eval_bms
asigma_loc_eval_bms.restype = a5err
asigma_loc_eval_bms.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, real, real, ctypes.c_int32, ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), real, ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.POINTER(struct_c__SA_asigma_loc_data)]
Expand Down Expand Up @@ -2690,6 +2693,9 @@ class struct_c__SA_asigma_data(Structure):
asigma_eval_sigmav = _libraries['libascot.so'].asigma_eval_sigmav
asigma_eval_sigmav.restype = a5err
asigma_eval_sigmav.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, real, ctypes.c_int32, ctypes.c_int32, real, real, real, real, asigma_reac_type, ctypes.POINTER(struct_c__SA_asigma_data)]
asigma_eval_cx = _libraries['libascot.so'].asigma_eval_cx
asigma_eval_cx.restype = a5err
asigma_eval_cx.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, real, real, ctypes.c_int32, ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), real, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(struct_c__SA_asigma_data)]
asigma_eval_bms = _libraries['libascot.so'].asigma_eval_bms
asigma_eval_bms.restype = a5err
asigma_eval_bms.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, real, real, ctypes.c_int32, ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), real, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(struct_c__SA_asigma_data)]
Expand Down Expand Up @@ -3115,6 +3121,9 @@ class struct_c__SA_afsi_data(Structure):
biosaw_calc_B = _libraries['libascot.so'].biosaw_calc_B
biosaw_calc_B.restype = None
biosaw_calc_B.argtypes = [ctypes.c_int32, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), 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)]
bbnbi_simulate = _libraries['libascot.so'].bbnbi_simulate
bbnbi_simulate.restype = None
bbnbi_simulate.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.c_int32, real, real, 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_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.POINTER(struct_c__SA_particle_state)), ctypes.POINTER(ctypes.c_double)]
__all__ = \
['B_2DS_data', 'B_2DS_eval_B', 'B_2DS_eval_B_dB',
'B_2DS_eval_psi', 'B_2DS_eval_psi_dpsi', 'B_2DS_eval_rho_drho',
Expand Down Expand Up @@ -3156,19 +3165,20 @@ class struct_c__SA_afsi_data(Structure):
'N0_3D_init', 'N0_3D_init_offload', 'N0_3D_offload_data',
'Reaction', 'SIMULATION_MODE', 'a5err', 'afsi_data', 'afsi_run',
'afsi_test_dist', 'afsi_test_thermal', 'afsi_thermal_data',
'asigma_data', 'asigma_eval_bms', 'asigma_eval_sigma',
'asigma_eval_sigmav', 'asigma_extrapolate', 'asigma_free_offload',
'asigma_init', 'asigma_init_offload', 'asigma_loc_data',
'asigma_loc_eval_bms', 'asigma_loc_eval_sigma',
'asigma_loc_eval_sigmav', 'asigma_loc_free_offload',
'asigma_loc_init', 'asigma_loc_init_offload',
'asigma_loc_offload_data', 'asigma_offload_data',
'asigma_reac_type', 'asigma_type', 'asigma_type_loc',
'biosaw_calc_B', 'boozer_data', 'boozer_eval_psithetazeta',
'boozer_free_offload', 'boozer_init', 'boozer_init_offload',
'boozer_offload_data', 'boschhale_reaction', 'boschhale_sigma',
'boschhale_sigmav', 'diag_data', 'diag_free', 'diag_free_offload',
'diag_init', 'diag_init_offload', 'diag_offload_data',
'asigma_data', 'asigma_eval_bms', 'asigma_eval_cx',
'asigma_eval_sigma', 'asigma_eval_sigmav', 'asigma_extrapolate',
'asigma_free_offload', 'asigma_init', 'asigma_init_offload',
'asigma_loc_data', 'asigma_loc_eval_bms', 'asigma_loc_eval_cx',
'asigma_loc_eval_sigma', 'asigma_loc_eval_sigmav',
'asigma_loc_free_offload', 'asigma_loc_init',
'asigma_loc_init_offload', 'asigma_loc_offload_data',
'asigma_offload_data', 'asigma_reac_type', 'asigma_type',
'asigma_type_loc', 'bbnbi_simulate', 'biosaw_calc_B',
'boozer_data', 'boozer_eval_psithetazeta', 'boozer_free_offload',
'boozer_init', 'boozer_init_offload', 'boozer_offload_data',
'boschhale_reaction', 'boschhale_sigma', 'boschhale_sigmav',
'diag_data', 'diag_free', 'diag_free_offload', 'diag_init',
'diag_init_offload', 'diag_offload_data',
'diag_orb_check_plane_crossing', 'diag_orb_check_radial_crossing',
'diag_orb_data', 'diag_orb_free', 'diag_orb_init',
'diag_orb_offload_data', 'diag_orb_update_fo',
Expand Down
123 changes: 122 additions & 1 deletion a5py/ascotpy/libsimulate.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from numpy.ctypeslib import ndpointer

from a5py import physlib
from a5py.routines.virtualrun import VirtualRun
from a5py.routines.virtualrun import VirtualRun, VirtualBBNBIRun
from a5py.exceptions import *

from .libascot import _LIBASCOT
Expand Down Expand Up @@ -199,6 +199,23 @@ def simulation_initinputs(self, bfield=True, efield=True, plasma=True,
wall=wall, boozer=boozer, mhd=mhd, asigma=asigma, switch=switch)
self._pack()

def simulation_initbbnbi(
self,
bfield=True,
plasma=True,
neutral=True,
wall=True,
asigma=True,
nbi=True,
switch=True):
"""Initialize inputs for BBNBI simulation.
This method must be called before running BBNBI simulation.
"""
self.input_init(
bfield=bfield, plasma=plasma, neutral=neutral,
wall=wall, asigma=asigma, switch=switch)
self._init(self.data, nbi=getattr(self.data, "nbi").active.get_qid())

def simulation_initmarkers(self, **mrk):
"""Create markers for the interactive simulations.
Expand Down Expand Up @@ -454,6 +471,110 @@ def read(self):
self._diag_offload_array,
diagorb=diagorb, dist5d=dist5d)

def simulation_bbnbi(self, nprt, t1=0, t2=0, printsummary=True):
"""Run BBNBI simulation.
Parameters
----------
nprt : Number of markers to be launched in total.
This number is the combined total for all injectors which are then
distributed among the injectors in proportion to the injector power.
For example, suppose there are 3 injectors with P_1 = 10 MW and
P_2 = 5 MW and P_3 = 5 MW. If `nprt` = 10000, then 5000 markers are
generated for injector 1 and 2500 markers for both injectors 2
and 3.
The number of ionized markers is either equal or smaller than `nprt`
as some of the markers are lost as shinethrough.
t1 : float, optional
The time instant at which the injector is turned on.
In the usual case where the background is time independent,
the parameters `t1` and `t2` only affect the marker initial time
which will be uniformly distributed in `[t1, t2]`.
printsummary : bool, optional
If True, summary of marker endstates is printed after the simulation
completes.
Returns
-------
run : :class:`VirtualBBNBI`
An object that acts almost exactly as :class:`BBNBIGroup` except that
the data is read from C arrays in the memory instead of HDF5 file.
The run can be used until the output is freed with
:meth:`simulation_free`. Previous data must be freed before
rerunning the simulation.
Raises
------
AscotInitException
If inputs are not packed, markers are not initialized or previous
results have not been freed.
"""
if not _LIBASCOT:
raise AscotInitException(
"Python interface disabled as libascot.so is not found")
if self._nmrk.value != 0:
raise AscotInitException(
"Free markers before rerunning the simulation")
if self._diag_occupied:
raise AscotInitException(
"Free previous results before running the simulation")
# Initialize diagnostics array and endstate
self._endstate = ctypes.pointer(ascot2py.struct_c__SA_particle_state())
ascot2py.diag_init_offload(ctypes.byref(self._sim.diag_offload_data),
ctypes.byref(self._diag_offload_array),
nprt)
self._diag_occupied = True

# Simulate and print stdout/stderr if requested
def runsim():
ascot2py.bbnbi_simulate(
ctypes.byref(self._sim), nprt, t1, t2,
self._bfield_offload_array,
self._plasma_offload_array,
self._neutral_offload_array,
self._wall_offload_array, self._wall_int_offload_array,
self._asigma_offload_array,
self._nbi_offload_array,
ctypes.byref(self._endstate),
self._diag_offload_array)

if self._mute == "no":
runsim()
else:
with wurlitzer.pipes() as (out, err):
runsim()
err = err.read()
if self._mute == "err" and len(err) > 1: print(err)

# Print summary
self._nmrk.value = nprt
if self._sim.mpi_rank == self._sim.mpi_root and printsummary:
ascot2py.print_marker_summary(self._endstate, self._nmrk)

class VirtualInput():
"""Wrapper for marker and options inputs.
"""

def __init__(self, inp):
self.inp = inp

def read(self):
return self.inp

diagorb = None
if self._sim.diag_offload_data.diagorb_collect:
diagorb = self._sim.diag_offload_data.diagorb
dist5d = None
if self._sim.diag_offload_data.dist5D_collect:
dist5d = self._sim.diag_offload_data.dist5D
return VirtualBBNBIRun(
self, nprt, self._endstate, VirtualInput(self._virtualoptions),
self._diag_offload_array, dist5d=dist5d)

def simulation_free(self, inputs=False, markers=False, diagnostics=False):
"""Free resources used by the interactive simulation.
Expand Down
56 changes: 50 additions & 6 deletions a5py/routines/virtualrun.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from a5py.ascot5io import State, Orbits, Dist
from a5py.ascot5io.dist import DistData
from .runmixin import RunMixin
from .bbnbi5 import BBNBIMixin

from a5py.ascotpy.libascot import _LIBASCOT
if _LIBASCOT:
Expand All @@ -30,16 +31,10 @@ def __init__(self, ascot, nmrk, inistate, endstate, options, markers,
Ascot instance for input interpolation.
nmrk : int
Number of markers in the simulation.
npoint : int
Number of data points reserved per marker.
inistate :
Pointer to marker inistate array.
endstate :
Pointer to marker endstate array.
mode : int
What simulation mode was active.
orbmode : int
What orbit collection mode was active.
orbits : array_like, optional
The diagnostics array containing the orbit data, if present.
dist5d : array_like, optional
Expand Down Expand Up @@ -79,6 +74,55 @@ def pointer_increment(idx):

self._dist5d = VirtualDist("5d", dist5d, diag_offload_array)

class VirtualBBNBIRun(BBNBIMixin):
"""Virtual :class:`BBNBIGroup` whose data exists solely in the memory.
"""

def __init__(self, ascot, nmrk, state, options, diag_offload_array,
dist5d=None, dist5drho=None):
"""Initialize fields that allow this instance to replicate
:class:`BBNBIGroup` behavior.
Sets state and orbit attributes which are used by the methods in
:class:`BBNBIMixin`.
Parameters
----------
ascot : :class:`Ascot`
Ascot instance for input interpolation.
nmrk : int
Number of markers in the simulation.
state :
Pointer to marker state array.
dist5d : array_like, optional
The 5d dist data struct if present.
dist5drho : array_like, optional
The diagnostics array containing the 5D rhodist data, if present.
"""
self.options = options
# There's a need for better solution here in case inputs are provided
#for inp in ["bfield", "efield", "plasma", "neutral", "wall", "boozer",
# "mhd", "asigma"]:
# grp = ascot.data[inp].active
# setattr(self, inp, grp)
self._state = VirtualState(ascot, ascot._nmrk, state)

def pointer_increment(idx):
"""Returns pointer to the diagnostics array on index idx.
"""
array_at_idx = ctypes.cast(diag_offload_array,
ctypes.POINTER(ctypes.c_double))
ptr = ctypes.cast(ctypes.pointer(array_at_idx),
ctypes.POINTER(ctypes.c_void_p))
ptr.contents.value += idx*ctypes.sizeof(ctypes.c_double)
return array_at_idx

if dist5d is not None:
data = pointer_increment(
ascot._sim.diag_offload_data.offload_dist5D_index)

self._dist5d = VirtualDist("5d", dist5d, diag_offload_array)

class VirtualState():
"""Like :class:`State` but the data is in C array.
"""
Expand Down
9 changes: 5 additions & 4 deletions src/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ OBJS= math.o list.o octree.o error.c \
$(SPLINEOBJS) \
neutral.o plasma.o particle.o endcond.o B_field.o gctransform.o \
E_field.o wall.o simulate.o diag.o offload.o boozer.o mhd.o \
random.o print.c hdf5_interface.o suzuki.o nbi.o biosaw.o \
random.o print.o hdf5_interface.o suzuki.o nbi.o biosaw.o \
asigma.o mpi_interface.o boschhale.o

BINS=test_math test_nbi test_bsearch \
Expand All @@ -165,13 +165,13 @@ ifeq ($(CC),h5pcc)
libascot.so: CFLAGS+=-shlib
endif

libascot.so: libascot.o ascot5_main.o libascot_mem.o afsi.o $(OBJS)
libascot.so: libascot.o ascot5_main.o libascot_mem.o afsi.o bbnbi5.o $(OBJS)
$(CC) $(CFLAGS) -o $@ $^ $(LFLAGS)

ascot5_main: ascot5_main.o $(OBJS)
$(CC) -o $@ $^ $(CFLAGS)

bbnbi5: bbnbi5.o $(OBJS)
bbnbi5: bbnbi5_main.o $(OBJS)
$(CC) -o $@ $^ $(CFLAGS)

doc:
Expand Down Expand Up @@ -235,7 +235,8 @@ ASCOT2PY_HEADERFILES=ascot5.h particle.h mpi_interface.h endcond.h simulate.h \
$(DIAGHEADERS) $(BFHEADERS) $(EFHEADERS) $(WALLHEADERS) $(PLSHEADERS) \
$(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
boozer.h mhd.h nbi.h asigma.h afsi.h ascot5_main.h biosaw.h boschhale.h\
bbnbi5.h

ascot2py.py : libascot.so
$(eval CLANGOMP=$(shell clang --print-resource-dir)) \
Expand Down
Loading

0 comments on commit 2a09dca

Please sign in to comment.