import numpy as np
from adios2 import Stream
[docs]
class DiffusionCoefficients:
"""
Stores the profiles of diffusion model coefficients defined as functions
of the poloidal magnetic flux ψ.
The four diffusion coefficients are:
- ptl_diffusivity (m^2/s)
- momentum_diffusivity (m^2/s)
- heat_conductivity (m^2/s)
- ptl_pinch_velocity (m/s)
The data layout of each coefficient is (nspecies, npsi) where the second dimension
(the ψ-grid) is contiguous in memory. Also the ψ–grid is stored as a 1D NumPy array of length npsi.
Parameters:
nspecies: Number of particle species (default: 2).
npsi: Number of points in the ψ–grid.
"""
def __init__(self, adios, nspecies=2, npsi=100):
self.adios = adios
self.nspecies = int(nspecies)
self.npsi = npsi
# Allocate the 4 coefficient arrays and the psi–grid.
self.ptl_diffusivity = np.zeros((nspecies, npsi), order='C', dtype=np.float64)
self.momentum_diffusivity = np.zeros((nspecies, npsi), order='C', dtype=np.float64)
self.heat_conductivity = np.zeros((nspecies, npsi), order='C', dtype=np.float64)
self.ptl_pinch_velocity = np.zeros((nspecies, npsi), order='C', dtype=np.float64)
self.psi = np.zeros(npsi, dtype=np.float64)
[docs]
def write_to_file(self, filename="xgc.diffusion_coeff.bp", gstep=-1):
"""
Writes the diffusion coefficient profiles and the ψ–grid to an Adios BP file as a new step.
For each diffusion coefficient, the data for each species is split into an individual variable
using a species suffix. The species suffixes (for species indices 0 through 6) are:
0: _elec
1: _ion
2: _imp1
3: _imp2
4: _imp3
5: _imp4
6: _imp5
In addition, the file will store:
- n_species: a scalar integer (int64).
- psi: a 1D array (of length npsi).
The method opens the specified file (filename) using the Adios2 FileWriter in append mode,
writes all the individual diffusion profiles, and then ends the step.
"""
# Define species suffixes.
species_suffixes = ["_elec", "_ion", "_imp1", "_imp2", "_imp3", "_imp4", "_imp5"]
IOName = "input.diffusion_coefficients"
io = self.adios.declare_io(IOName)
if io.engine_type().upper() == "FILE":
print("Setting default engine FileStream for {0}".format(filename))
io.set_engine("FileStream")
# Open the file for writing (appending a new step)
with Stream(io, filename, "a") as writer:
writer.begin_step()
# Define a list of tuples: (coefficient base name, corresponding array)
coeffs = [
("ptl_diffusivity", self.ptl_diffusivity),
("momentum_diffusivity", self.momentum_diffusivity),
("heat_conductivity", self.heat_conductivity),
("ptl_pinch_velocity", self.ptl_pinch_velocity)
]
# Loop over coefficients and species.
for coeff_name, coeff_array in coeffs:
for i in range(self.nspecies):
# Form the variable name using the base name and species suffix.
var_name = coeff_name + species_suffixes[i]
# Extract the data for this species (a 1D array of length npsi).
data = coeff_array[i, :]
writer.write(var_name, data, shape=data.shape, start=(0,), count=(self.npsi,))
# Write out the number of species as an int64 scalar.
writer.write("n_species", self.nspecies)
# Write out the psi–grid as a 1D double array.
writer.write("psi", self.psi, shape=data.shape, start=(0,), count=(self.npsi,))
# Write the time step index to which this data belongs
writer.write("gstep", gstep)
writer.end_step()
self.adios.remove_io(IOName)