diff --git a/CHANGELOG.md b/CHANGELOG.md index 8338f33..1d671ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,23 @@ Project Changelog ================= +Release 1.2.0 (19 Sept 2020) +---------------------------- + +* Added support for Raysect 0.7. +* Replaced unsafe SOLPSFunction3D and SOLPSVectorFunction3D with safe SOLPSFunction2D and SOLPSVectorFunction2D (use AxisymmetricMapper(SOLPSFunction2D) and VectorAxisymmetricMapper(SOLPSVectorFunction2D) for 3D). +* Added correct initialisation of properties in SOLPSMesh and SOLPSSimulation. +* Added new attributes to SOLPSMesh for basis vectors, cell connection areas, indices of neighbouring cells and new methods to_poloidal() and to_cartesian() for converting vectors defined on a grid from/to (poloidal, radial)/(R, Z). +* Fixed incorrect calculation of cell basis vectors. +* Inverted the indexing of data arrays and made all arrays row-major. +* Added electron_velocities and neutral_listproperties to SOLPSSimulation. +* Added 2D and 3D interpolators for plasma parameters to SOLPSSimulation. +* Added parsing of additional quantities in load_solps_from_mdsplus(), load_solps_from_raw_output() and load_solps_from_balance(). +* Added support for B2 stand-alone simulations. +* Added support for arbitrary plasma chemical composition. +* Fixed incorrect calculation of velocities. +* Other small fixes and improvements. + Release 1.1.0 (30 July 2020) ---------------------------- diff --git a/cherab/solps/__init__.py b/cherab/solps/__init__.py index 03d35fc..e53b2db 100644 --- a/cherab/solps/__init__.py +++ b/cherab/solps/__init__.py @@ -17,7 +17,7 @@ # See the Licence for the specific language governing permissions and limitations # under the Licence. -from .solps_3d_functions import SOLPSFunction3D, SOLPSVectorFunction3D +from .solps_2d_functions import SOLPSFunction2D, SOLPSVectorFunction2D from .mesh_geometry import SOLPSMesh from .solps_plasma import SOLPSSimulation from .formats import load_solps_from_raw_output, load_solps_from_mdsplus, load_solps_from_balance diff --git a/cherab/solps/b2/parse_b2_block_file.py b/cherab/solps/b2/parse_b2_block_file.py index 03a51c5..7c25c13 100755 --- a/cherab/solps/b2/parse_b2_block_file.py +++ b/cherab/solps/b2/parse_b2_block_file.py @@ -47,16 +47,16 @@ def _make_solps_data_object(_data): # Multiple 2D data field (e.g. na) if number > nxyg: - _data = np.array(_data).reshape((nxg, nyg, int(number / nxyg)), order='F') + _data = np.array(_data).reshape((number // nxyg, nyg, nxg)) if debug: - print('Mesh data field {} with dimensions: {:d} x {:d} x {:d}'.format(name, nxg, nyg, int(number/nxyg))) + print('Mesh data field {} with dimensions: {:d} x {:d} x {:d}'.format(name, number // nxyg, nyg, nxg)) return MESH_DATA, _data # 2D data field (e.g. ne) elif number == nxyg: - _data = np.array(_data).reshape((nxg, nyg), order='F') + _data = np.array(_data).reshape((nyg, nxg)) if debug: - print('Mesh data field {} with dimensions: {:d} x {:d}'.format(name, nxg, nyg)) + print('Mesh data field {} with dimensions: {:d} x {:d}'.format(name, nyg, nxg)) return MESH_DATA, _data # Additional information field (e.g. zamin) diff --git a/cherab/solps/eirene/eirene.py b/cherab/solps/eirene/eirene.py index 7f8fcc2..3eeb80e 100755 --- a/cherab/solps/eirene/eirene.py +++ b/cherab/solps/eirene/eirene.py @@ -16,8 +16,6 @@ # See the Licence for the specific language governing permissions and limitations # under the Licence. -import numpy as np - # Code based on script by Felix Reimold (2016) class Eirene: @@ -91,7 +89,7 @@ def __init__(self, nx, ny, na, nm, ni, ns, species_labels, version=None, @property def nx(self): """ - Number of grid cells in the x direction + Number of grid cells in the poloidal direction :rtype: int """ @@ -100,7 +98,7 @@ def nx(self): @property def ny(self): """ - Number of grid cells in the y direction + Number of grid cells in the radial direction :rtype: int """ @@ -395,7 +393,7 @@ def elosm(self): @elosm.setter def elosm(self, value): - self._check_dimensions(value, 1) + self._check_dimensions(value, self._nm) self._elosm = value @property @@ -409,7 +407,7 @@ def edism(self): @edism.setter def edism(self, value): - self._check_dimensions(value, 1) + self._check_dimensions(value, self._nm) self._edism = value @property @@ -426,13 +424,13 @@ def eradt(self, value): self._check_dimensions(value, 1) self._eradt = value - def _check_dimensions(self, data, dim2): + def _check_dimensions(self, data, dim0): """ Checks compatibility of the data array dimension with the species number and grid size. - :param dim2: size of the 2nd dimenion + :param dim0: size of the 1st dimenion :return: """ - if not data.shape == (self._nx, self._ny, dim2): + if not data.shape == (dim0, self._ny, self._nx): raise ValueError("Array with shape {0} obtained, but {1} expected".format(data.shape, - (self._nx, self._ny, dim2))) + (dim0, self._ny, self._nx))) diff --git a/cherab/solps/eirene/parser/fort44.py b/cherab/solps/eirene/parser/fort44.py index 230e4ce..ce36d60 100644 --- a/cherab/solps/eirene/parser/fort44.py +++ b/cherab/solps/eirene/parser/fort44.py @@ -48,7 +48,7 @@ def load_fort44_file(file_path, debug=False): # Look up file parsing function and call it to obtain for44 block and update data dictionary parser = assign_fort44_parser(data["version"]) - return parser(file_path) + return parser(file_path, debug) def assign_fort44_parser(file_version): diff --git a/cherab/solps/eirene/parser/fort44_2013.py b/cherab/solps/eirene/parser/fort44_2013.py index 13ade5d..9836db7 100644 --- a/cherab/solps/eirene/parser/fort44_2013.py +++ b/cherab/solps/eirene/parser/fort44_2013.py @@ -98,8 +98,8 @@ def load_fort44_2013(file_path, debug=False): eirene.emism = read_block44(file_handle, 1, nx, ny) # Molecular Halpha Emission # Radiated power (elosm, edism, eradt) - eirene.elosm = read_block44(file_handle, 1, nx, ny) # Power loss due to molecules (including dissociation) - eirene.edism = read_block44(file_handle, 1, nx, ny) # Power loss due to molecule dissociation + eirene.elosm = read_block44(file_handle, nm, nx, ny) # Power loss due to molecules (including dissociation) + eirene.edism = read_block44(file_handle, nm, nx, ny) # Power loss due to molecule dissociation eirene.eradt = read_block44(file_handle, 1, nx, ny) # Neutral radiated power return eirene diff --git a/cherab/solps/eirene/parser/fort44_2017.py b/cherab/solps/eirene/parser/fort44_2017.py index b261c19..965da3e 100644 --- a/cherab/solps/eirene/parser/fort44_2017.py +++ b/cherab/solps/eirene/parser/fort44_2017.py @@ -105,7 +105,7 @@ def load_fort44_2017(file_path, debug=False): _ = read_block44(file_handle, eirene.nm, eirene.nx, eirene.ny) # Molecule particle source # Radiated power (elosm, edism, eradt) - eirene.edism = read_block44(file_handle, 1, eirene.nx, eirene.ny) # Power loss due to molecule dissociation + eirene.edism = read_block44(file_handle, eirene.nm, eirene.nx, eirene.ny) # Power loss due to molecule dissociation # Consume lines until eradt is reached while True: diff --git a/cherab/solps/eirene/parser/utility.py b/cherab/solps/eirene/parser/utility.py index 72e842d..f3d0944 100644 --- a/cherab/solps/eirene/parser/utility.py +++ b/cherab/solps/eirene/parser/utility.py @@ -25,9 +25,9 @@ def read_block44(file_handle, ns, nx, ny): :param file_handle: A python core file handle object as a result of a call to open('./fort.44'). :param int ns: total number of species - :param int nx: number of grid x cells - :param int ny: number of grid y cells - :return: ndarray of data with shape [nx, ny, ns] + :param int nx: number of grid poloidal cells + :param int ny: number of grid radial cells + :return: ndarray of data with shape [ns, ny, nx] """ data = [] npoints = ns * nx * ny @@ -37,5 +37,5 @@ def read_block44(file_handle, ns, nx, ny): # This is a comment line. Ignore continue data.extend(line) - data = np.asarray(data, dtype=float).reshape((nx, ny, ns), order='F') + data = np.asarray(data, dtype=float).reshape((ns, ny, nx)) return data diff --git a/cherab/solps/formats/balance.py b/cherab/solps/formats/balance.py index 90e4ed9..c0b8e36 100644 --- a/cherab/solps/formats/balance.py +++ b/cherab/solps/formats/balance.py @@ -19,139 +19,186 @@ import numpy as np from scipy.io import netcdf -from scipy.constants import elementary_charge -from raysect.core.math.function.float import Discrete2DMesh +from scipy import constants -from cherab.core.math.mappers import AxisymmetricMapper -from cherab.core.atomic.elements import hydrogen, deuterium, helium, beryllium, carbon, nitrogen, oxygen, neon, \ - argon, krypton, xenon +from cherab.core.atomic.elements import lookup_isotope from cherab.solps.mesh_geometry import SOLPSMesh -from cherab.solps.solps_plasma import SOLPSSimulation - -Q = elementary_charge - -# key is nuclear charge Z and atomic mass AMU -_popular_species = { - (1, 2): deuterium, - (6, 12.0): carbon, - (2, 4.003): helium, - (7, 14.0): nitrogen, - (10, 20.180): neon, - (18, 39.948): argon, - (18, 40.0): argon, - (36, 83.798): krypton, - (54, 131.293): xenon -} +from cherab.solps.solps_plasma import SOLPSSimulation, prefer_element, eirene_flux_to_velocity, b2_flux_to_velocity + def load_solps_from_balance(balance_filename): """ Load a SOLPS simulation from SOLPS balance.nc output files. """ - # Open the file - fhandle = netcdf.netcdf_file(balance_filename, 'r') - - # Load SOLPS mesh geometry - cr_x = fhandle.variables['crx'].data.copy() - cr_z = fhandle.variables['cry'].data.copy() - vol = fhandle.variables['vol'].data.copy() - - # Re-arrange the array dimensions in the way CHERAB expects... - cr_x = np.moveaxis(cr_x, 0, -1) - cr_z = np.moveaxis(cr_z, 0, -1) - - # Create the SOLPS mesh - mesh = SOLPSMesh(cr_x, cr_z, vol) - - sim = SOLPSSimulation(mesh) - - # TODO: add code to load SOLPS velocities and magnetic field from files - - # Load electron species - sim._electron_temperature = fhandle.variables['te'].data.copy() / Q - sim._electron_density = fhandle.variables['ne'].data.copy() - - ########################################## - # Load each plasma species in simulation # - ########################################## - - sim._species_list = [] - n_species = len(fhandle.variables['am'].data) + el_charge = constants.elementary_charge + rydberg_energy = constants.value('Rydberg constant times hc in eV') - for i in range(n_species): - - # Extract the nuclear charge - if fhandle.variables['species'].data[i, 1] == b'D': - zn = 1 - if fhandle.variables['species'].data[i, 1] == b'C': - zn = 6 - if fhandle.variables['species'].data[i, 1] == b'N': - zn = 7 - if fhandle.variables['species'].data[i, 1] == b'N' and fhandle.variables['species'].data[i, 2] == b'e': - zn = 10 - if fhandle.variables['species'].data[i, 1] == b'A' and fhandle.variables['species'].data[i, 2] == b'r': - zn = 18 - - am = float(fhandle.variables['am'].data[i]) # Atomic mass - charge = int(fhandle.variables['za'].data[i]) # Ionisation/charge - species = _popular_species[(zn, am)] - - # If we only need to populate species_list, there is probably a faster way to do this... - sim.species_list.append(species.symbol + str(charge)) - - tmp = fhandle.variables['na'].data.copy() - tmp = np.moveaxis(tmp, 0, -1) - sim._species_density = tmp - - # Make Mesh Interpolator function for inside/outside mesh test. - inside_outside_data = np.ones(mesh.num_tris) - inside_outside = AxisymmetricMapper(Discrete2DMesh(mesh.vertex_coords, mesh.triangles, inside_outside_data, limit=False)) - sim._inside_mesh = inside_outside - - # Load the neutrals data - try: - D0_indx = sim.species_list.index('D0') - except ValueError: - D0_indx = None - - # Replace the deuterium neutrals density (from the fluid neutrals model by default) with - # the values calculated by EIRENE - do the same for other neutrals? - if 'dab2' in fhandle.variables.keys(): - if D0_indx is not None: - b2_len = np.shape(sim.species_density[:, :, D0_indx])[-1] - eirene_len = np.shape(fhandle.variables['dab2'].data)[-1] - sim.species_density[:, :, D0_indx] = fhandle.variables['dab2'].data[0, :, 0:b2_len-eirene_len] + # Open the file + with netcdf.netcdf_file(balance_filename, 'r') as fhandle: + + # Load SOLPS mesh geometry + mesh = load_mesh_from_netcdf(fhandle) + + # Load each plasma species in simulation + + species_list = [] + neutral_indx = [] + am = np.round(fhandle.variables['am'].data).astype(np.int) # Atomic mass number + charge = fhandle.variables['za'].data.astype(np.int) # Ionisation/charge + species_names = fhandle.variables['species'].data.copy() + ns = am.size + for i in range(ns): + symbol = ''.join([b.decode('utf-8').strip(' 0123456789+-') for b in species_names[i]]) # also strips isotope number + if symbol not in ('D', 'T'): + isotope = lookup_isotope(symbol, number=am[i]) # will throw an error for D or T + species = prefer_element(isotope) # Prefer Element over Isotope if the mass number is the same + else: + species = lookup_isotope(symbol) + + # If we only need to populate species_list, there is probably a faster way to do this... + species_list.append((species.name, charge[i])) + if charge[i] == 0: + neutral_indx.append(i) + + sim = SOLPSSimulation(mesh, species_list) + nx = mesh.nx + ny = mesh.ny + + ########################## + # Magnetic field vectors # + sim.b_field = fhandle.variables['bb'].data.copy()[:3] + # sim.b_field_cylindrical is created automatically + + # Load electron species + sim.electron_temperature = fhandle.variables['te'].data.copy() / el_charge + sim.electron_density = fhandle.variables['ne'].data.copy() + + # Load ion temperature + sim.ion_temperature = fhandle.variables['ti'].data.copy() / el_charge + + # Load species density + sim.species_density = fhandle.variables['na'].data.copy() + + # Load parallel velocity + parallel_velocity = fhandle.variables['ua'].data.copy() + + # Load poloidal and radial particle fluxes for velocity calculation + if 'fna_tot' in fhandle.variables: + fna = fhandle.variables['fna_tot'].data.copy() + # Obtaining velocity from B2 flux + sim.velocities_cylindrical = b2_flux_to_velocity(sim, fna[:, 0], fna[:, 1], parallel_velocity) + else: # trying to obtain particle flux from components + fna = 0 + for key in fhandle.variables.keys(): + if 'fna_' in key: + fna += fhandle.variables[key].data.copy() + if fna is not 0: + # Obtaining velocity from B2 flux + sim.velocities_cylindrical = b2_flux_to_velocity(sim, fna[:, 0], fna[:, 1], parallel_velocity) + + # Obtaining additional data from EIRENE and replacing data for neutrals + if 'dab2' in fhandle.variables: + sim.species_density[neutral_indx] = fhandle.variables['dab2'].data.copy()[:, :ny, :nx] # in case of large grid size + b2_standalone = False + else: + b2_standalone = True + + if not b2_standalone: + + # Obtaining neutral atom velocity from EIRENE flux + # Note that if the output for fluxes was turned off, pfluxa and rfluxa' are all zeros + if 'pfluxa' in fhandle.variables and 'rfluxa' in fhandle.variables: + neutral_poloidal_flux = fhandle.variables['pfluxa'].data.copy()[:, :ny, :nx] + neutral_radial_flux = fhandle.variables['rfluxa'].data.copy()[:, :ny, :nx] + + if np.any(neutral_poloidal_flux) or np.any(neutral_radial_flux): + neutral_velocities = eirene_flux_to_velocity(sim, neutral_poloidal_flux, neutral_radial_flux, + parallel_velocity[neutral_indx]) + if sim.velocities_cylindrical is not None: + sim.velocities_cylindrical[neutral_indx] = neutral_velocities + sim.velocities_cylindrical = sim.velocities_cylindrical # Updating sim.velocities + else: + # No 'fna_*' keys in balance.nc and b2 species velocities are not set + velocities_cylindrical = np.zeros((len(sim.species_list, 3, ny, nx))) + velocities_cylindrical[neutral_indx] = neutral_velocities + sim.velocities_cylindrical = velocities_cylindrical + + # Obtaining neutral temperatures + if 'tab2' in fhandle.variables: + sim.neutral_temperature = fhandle.variables['tab2'].data.copy()[:, :ny, :nx] + + # Calculate the total radiated power + b2_ploss = 0 + eirene_ecoolrate = 0 + eirene_potential_loss = 0 + + # Total radiated power from B2, not including neutrals + if 'b2stel_she_bal' in fhandle.variables: + b2_ploss = np.sum(fhandle.variables['b2stel_she_bal'].data, axis=0) / mesh.vol + + # Electron energy loss due to interactions with neutrals + if 'eirene_mc_eael_she_bal' in fhandle.variables: + eirene_ecoolrate = np.sum(fhandle.variables['eirene_mc_eael_she_bal'].data, axis=0) / mesh.vol + + # Ionisation rate from EIRENE, needed to calculate the energy loss to overcome the ionisation potential of atoms + if 'eirene_mc_papl_sna_bal' in fhandle.variables: + tmp = np.sum(fhandle.variables['eirene_mc_papl_sna_bal'].data, axis=0)[1] + eirene_potential_loss = rydberg_energy * tmp * el_charge / mesh.vol + + # This will be negative (energy sink); multiply by -1 + total_rad = -1.0 * (b2_ploss + (eirene_ecoolrate - eirene_potential_loss)) + + else: + # Total radiated power from B2, not including neutrals + b2_ploss = 0 + potential_loss = 0 + + if 'b2stel_she_bal' in fhandle.variables: + b2_ploss = np.sum(fhandle.variables['b2stel_she_bal'].data, axis=0) / mesh.vol + + if 'b2stel_sna_ion_bal' in fhandle.variables: + potential_loss = np.sum(fhandle.variables['b2stel_sna_ion_bal'].data, axis=0) / mesh.vol + + # Save total radiated power to the simulation object + total_rad = rydberg_energy * el_charge * potential_loss - b2_ploss + + if total_rad is not 0: + sim.total_radiation = total_rad - eirene_run = True - else: - eirene_run = False + return sim - # Calculate the total radiated power - if eirene_run: - # Total radiated power from B2, not including neutrals - b2_ploss = np.sum(fhandle.variables['b2stel_she_bal'].data, axis=0) / vol - # Electron energy loss due to interactions with neutrals - if 'eirene_mc_eael_she_bal' in fhandle.variables.keys(): - eirene_ecoolrate = np.sum(fhandle.variables['eirene_mc_eael_she_bal'].data, axis=0) / vol +def load_mesh_from_netcdf(fhandle): - # Ionisation rate from EIRENE, needed to calculate the energy loss to overcome the ionisation potential of atoms - if 'eirene_mc_papl_sna_bal' in fhandle.variables.keys(): - eirene_potential_loss = 13.6 * np.sum(fhandle.variables['eirene_mc_papl_sna_bal'].data, axis=(0))[1, :, :] * Q / vol + # Load SOLPS mesh geometry + # Re-arrange the array dimensions in the way CHERAB expects... + r = fhandle.variables['crx'].data.copy() + z = fhandle.variables['cry'].data.copy() + vol = fhandle.variables['vol'].data.copy() - # This will be negative (energy sink); multiply by -1 - sim._total_rad = -1.0 * (b2_ploss + (eirene_ecoolrate-eirene_potential_loss)) + # Loading neighbouring cell indices + neighbix = np.zeros(r.shape, dtype=np.int) + neighbiy = np.zeros(r.shape, dtype=np.int) - else: - # Total radiated power from B2, not including neutrals - b2_ploss = np.sum(fhandle.variables['b2stel_she_bal'].data, axis=0)/vol + neighbix[0] = fhandle.variables['leftix'].data.astype(np.int) # poloidal prev. + neighbix[1] = fhandle.variables['bottomix'].data.astype(np.int) # radial prev. + neighbix[2] = fhandle.variables['rightix'].data.astype(np.int) # poloidal next + neighbix[3] = fhandle.variables['topix'].data.astype(np.int) # radial next - potential_loss = np.sum(fhandle.variables['b2stel_sna_ion_bal'].data, axis=0)/vol + neighbiy[0] = fhandle.variables['leftiy'].data.astype(np.int) + neighbiy[1] = fhandle.variables['bottomiy'].data.astype(np.int) + neighbiy[2] = fhandle.variables['rightiy'].data.astype(np.int) + neighbiy[3] = fhandle.variables['topiy'].data.astype(np.int) - # Save total radiated power to the simulation object - sim._total_rad = 13.6 * Q * potential_loss - b2_ploss + # In SOLPS cell indexing starts with -1 (guarding cell), but in SOLPSMesh -1 means no neighbour. + neighbix += 1 + neighbiy += 1 + neighbix[neighbix == r.shape[2]] = -1 + neighbiy[neighbiy == r.shape[1]] = -1 - fhandle.close() + # Create the SOLPS mesh + mesh = SOLPSMesh(r, z, vol, neighbix, neighbiy) - return sim + return mesh diff --git a/cherab/solps/formats/mdsplus.py b/cherab/solps/formats/mdsplus.py index bc9e3ae..09b2c5e 100644 --- a/cherab/solps/formats/mdsplus.py +++ b/cherab/solps/formats/mdsplus.py @@ -17,20 +17,13 @@ # See the Licence for the specific language governing permissions and limitations # under the Licence. -import re import numpy as np -from math import sqrt -from raysect.core import Point2D -from raysect.core.math.function.float import Discrete2DMesh -from cherab.core.math.mappers import AxisymmetricMapper +from cherab.core.atomic.elements import lookup_isotope from cherab.solps.mesh_geometry import SOLPSMesh -from cherab.solps.solps_plasma import SOLPSSimulation +from cherab.solps.solps_plasma import SOLPSSimulation, prefer_element, eirene_flux_to_velocity, b2_flux_to_velocity -_SPECIES_REGEX = '([a-zA-z]+)\+?([0-9]+)' - -# TODO: violates interface of SOLPSSimulation.... puts numpy arrays in the object where they should be function2D def load_solps_from_mdsplus(mds_server, ref_number): """ Load a SOLPS simulation from a MDSplus server. @@ -40,119 +33,100 @@ def load_solps_from_mdsplus(mds_server, ref_number): :rtype: SOLPSSimulation """ - from MDSplus import Connection as MDSConnection + from MDSplus import Connection as MDSConnection, MdsException # Setup connection to server conn = MDSConnection(mds_server) conn.openTree('solps', ref_number) # Load SOLPS mesh geometry and lookup arrays - mesh = load_mesh_from_mdsplus(conn) - sim = SOLPSSimulation(mesh) - ni = mesh.nx - nj = mesh.ny + mesh = load_mesh_from_mdsplus(conn, MdsException) + + # Load each plasma species in simulation + ns = conn.get(r'\SOLPS::TOP.IDENT.NS').data() # Number of species + zn = conn.get(r'\SOLPS::TOP.SNAPSHOT.GRID.ZN').data().astype(np.int) # Nuclear charge + am = np.round(conn.get(r'\SOLPS::TOP.SNAPSHOT.GRID.AM').data()).astype(np.int) # Atomic mass number + charge = conn.get(r'\SOLPS::TOP.SNAPSHOT.GRID.ZA').data().astype(np.int) # Ionisation/charge + + species_list = [] + neutral_indx = [] + for i in range(ns): + isotope = lookup_isotope(zn[i], number=am[i]) + species = prefer_element(isotope) # Prefer Element over Isotope if the mass number is the same + species_list.append((species.name, charge[i])) + if charge[i] == 0: + neutral_indx.append(i) + + sim = SOLPSSimulation(mesh, species_list) + nx = mesh.nx + ny = mesh.ny ########################## # Magnetic field vectors # - raw_b_field = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.B').data(), 0, 2) - b_field_vectors_cartesian = np.zeros((ni, nj, 3)) - b_field_vectors = np.zeros((ni, nj, 3)) - for i in range(ni): - for j in range(nj): - bparallel = raw_b_field[i, j, 0] - bradial = raw_b_field[i, j, 1] - btoroidal = raw_b_field[i, j, 2] - b_field_vectors[i, j] = (bparallel, bradial, btoroidal) - - pv = mesh.poloidal_grid_basis[i, j, 0] # parallel basis vector - rv = mesh.poloidal_grid_basis[i, j, 1] # radial basis vector - - bx = pv.x * bparallel + rv.x * bradial # component of B along poloidal x - by = pv.y * bparallel + rv.y * bradial # component of B along poloidal y - b_field_vectors_cartesian[i, j] = (bx, btoroidal, by) - sim._b_field_vectors = b_field_vectors - sim._b_field_vectors_cartesian = b_field_vectors_cartesian - - # Load electron species - sim._electron_temperature = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.TE').data(), 0, 1) # (32, 98) => (98, 32) - sim._electron_density = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.NE').data(), 0, 1) # (32, 98) => (98, 32) - - ############################ - # Load each plasma species # - ############################ - - # Master list of species, e.g. ['D0', 'D+1', 'C0', 'C+1', ... - species_list = conn.get('\SOLPS::TOP.IDENT.SPECIES').data() + sim.b_field = conn.get(r'\SOLPS::TOP.SNAPSHOT.B').data()[:3] + # sim.b_field_cylindrical is created automatically + + # Load electron temperature and density + sim.electron_temperature = conn.get(r'\SOLPS::TOP.SNAPSHOT.TE').data() + sim.electron_density = conn.get(r'\SOLPS::TOP.SNAPSHOT.NE').data() + + # Load ion temperature + sim.ion_temperature = conn.get(r'\SOLPS::TOP.SNAPSHOT.TI').data() + + # Load species density + sim.species_density = conn.get(r'\SOLPS::TOP.SNAPSHOT.NA').data() + + # Load parallel velocity + parallel_velocity = conn.get(r'\SOLPS::TOP.SNAPSHOT.UA').data() + + # Load poloidal and radial particle fluxes for velocity calculation + poloidal_flux = conn.get(r'\SOLPS::TOP.SNAPSHOT.FNAX').data() + radial_flux = conn.get(r'\SOLPS::TOP.SNAPSHOT.FNAY').data() + + # B2 fluxes are defined between cells, so correcting array shapes if needed + if poloidal_flux.shape[2] == nx - 1: + poloidal_flux = np.concatenate((np.zeros((ns, ny, 1)), poloidal_flux), axis=2) + + if radial_flux.shape[1] == ny - 1: + radial_flux = np.concatenate((np.zeros((ns, 1, nx)), radial_flux), axis=1) + + # Setting velocities from B2 flux + sim.velocities_cylindrical = b2_flux_to_velocity(sim, poloidal_flux, radial_flux, parallel_velocity) + + # Obtaining additional data from EIRENE and replacing data for neutrals + try: - species_list = species_list.decode('UTF-8') - except AttributeError: # Already a string - pass - sim._species_list = species_list.split() - sim._species_density = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.NA').data(), 0, 2) - sim._rad_par_flux = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.FNAY').data(), 0, 2) # radial particle flux - sim._radial_area = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.SY').data(), 0, 1) # radial contact area - - # Load the neutral atom density from B2 - dab2 = conn.get('\SOLPS::TOP.SNAPSHOT.DAB2').data() - if isinstance(dab2, np.ndarray): - sim._b2_neutral_densities = np.swapaxes(dab2, 0, 2) - - sim._velocities_parallel = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.UA').data(), 0, 2) - sim._velocities_radial = np.zeros((ni, nj, len(sim.species_list))) - sim._velocities_toroidal = np.zeros((ni, nj, len(sim.species_list))) - sim._velocities_cartesian = np.zeros((ni, nj, len(sim.species_list), 3), dtype=np.float64) - - ################################################ - # Calculate the species' velocity distribution # - b2_neutral_i = 0 # counter for B2 neutrals - for k, sp in enumerate(sim.species_list): - - # Identify the species based on its symbol - symbol, charge = re.match(_SPECIES_REGEX, sp).groups() - charge = int(charge) - - # If neutral and B" atomic density available, use B2 density, otherwise use fluid species density. - if charge == 0 and (sim._b2_neutral_densities is not None): - species_dens_data = sim.b2_neutral_densities[:, :, b2_neutral_i] - b2_neutral_i += 1 - else: - species_dens_data = sim.species_density[:, :, k] - - for i in range(ni): - for j in range(nj): - # Load grid basis vectors - pv = mesh.poloidal_grid_basis[i, j, 0] # parallel basis vector - rv = mesh.poloidal_grid_basis[i, j, 1] # radial basis vector - - # calculate field component ratios for velocity conversion - bparallel = b_field_vectors[i, j, 0] - btoroidal = b_field_vectors[i, j, 2] - bplane = sqrt(bparallel**2 + btoroidal**2) - parallel_to_toroidal_ratio = bparallel * btoroidal / (bplane**2) - - # Calculate toroidal and radial velocity components - v_parallel = sim.velocities_parallel[i, j, k] # straight from SOLPS 'UA' variable - v_toroidal = v_parallel * parallel_to_toroidal_ratio - sim.velocities_toroidal[i, j, k] = v_toroidal - # Special case for edge of mesh, no radial velocity expected. - try: - if species_dens_data[i, j] == 0: - v_radial = 0.0 - else: - v_radial = sim.radial_particle_flux[i, j, k] / sim.radial_area[i, j] / species_dens_data[i, j] - except IndexError: - v_radial = 0.0 - sim.velocities_radial[i, j, k] = v_radial - - # Convert velocities to cartesian coordinates - vx = pv.x * v_parallel + rv.x * v_radial # component of v along poloidal x - vy = pv.y * v_parallel + rv.y * v_radial # component of v along poloidal y - sim.velocities_cartesian[i, j, k, :] = (vx, v_toroidal, vy) - - # Make Mesh Interpolator function for inside/outside mesh test. - inside_outside_data = np.ones(mesh.num_tris) - inside_outside = AxisymmetricMapper(Discrete2DMesh(mesh.vertex_coords, mesh.triangles, inside_outside_data, limit=False)) - sim._inside_mesh = inside_outside + # Replace the species densities + neutral_density = conn.get(r'\SOLPS::TOP.SNAPSHOT.DAB2').data() # this will throw a TypeError is neutral_density is not an array + # We can update the data without re-initialising interpolators because they use pointers + sim.species_density[neutral_indx] = neutral_density[:] + + except (MdsException, TypeError): + print("Warning! This is B2 stand-alone simulation.") + b2_standalone = True + else: + b2_standalone = False + + if not b2_standalone: + # Obtaining neutral atom velocity from EIRENE flux + # Note that if the output for fluxes was turned off, PFLA and RFLA' are all zeros + try: + neutral_poloidal_flux = conn.get(r'\SOLPS::TOP.SNAPSHOT.PFLA').data()[:] + neutral_radial_flux = conn.get(r'\SOLPS::TOP.SNAPSHOT.RFLA').data()[:] + + if np.any(neutral_poloidal_flux) or np.any(neutral_radial_flux): + sim.velocities_cylindrical[neutral_indx] = eirene_flux_to_velocity(sim, neutral_poloidal_flux, neutral_radial_flux, + parallel_velocity[neutral_indx]) + sim.velocities_cylindrical = sim.velocities_cylindrical # Updating sim.velocities + + except (MdsException, TypeError): + pass + + # Obtaining neutral temperatures + try: + sim.neutral_temperature = conn.get(r'\SOLPS::TOP.SNAPSHOT.TAB2').data()[:] + except (MdsException, TypeError): + pass ############################### # Load extra data from server # @@ -160,95 +134,76 @@ def load_solps_from_mdsplus(mds_server, ref_number): #################### # Integrated power # - vol = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.VOL').data(), 0, 1) # TODO - this should be a mesh property - linerad = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.RQRAD').data(), 0, 2) - linerad = np.sum(linerad, axis=2) - brmrad = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.RQBRM').data(), 0, 2) - brmrad = np.sum(brmrad, axis=2) - neurad = conn.get('\SOLPS::TOP.SNAPSHOT.ENEUTRAD').data() - if neurad is not None: # need to cope with fact that neurad may not be present!!! - if len(neurad.shape) == 3: - neurad = np.swapaxes(np.abs(np.sum(neurad, axis=0)), 0, 1) + try: + linerad = np.sum(conn.get(r'\SOLPS::TOP.SNAPSHOT.RQRAD').data()[:], axis=0) + except (MdsException, TypeError): + linerad = 0 + + try: + brmrad = np.sum(conn.get(r'\SOLPS::TOP.SNAPSHOT.RQBRM').data()[:], axis=0) + except (MdsException, TypeError): + brmrad = 0 + + try: + eneutrad = conn.get(r'\SOLPS::TOP.SNAPSHOT.ENEUTRAD').data()[:] + if np.ndim(eneutrad) == 3: + neurad = np.abs(np.sum(eneutrad, axis=0)) else: - neurad = np.swapaxes(np.abs(neurad), 0, 1) - else: - neurad = np.zeros(brmrad.shape) + neurad = np.abs(eneutrad) + except (MdsException, TypeError): + neurad = 0 + + total_rad = linerad + brmrad + neurad - total_rad_data = np.zeros(vol.shape) - ni, nj = vol.shape - for i in range(ni): - for j in range(nj): - total_rad_data[i, j] = (linerad[i, j] + brmrad[i, j] + neurad[i, j]) / vol[i, j] - sim._total_rad = total_rad_data + if total_rad is not 0: + sim.total_radiation = total_rad / mesh.vol return sim -def load_mesh_from_mdsplus(mds_connection): +def load_mesh_from_mdsplus(mds_connection, MdsException): """ Load the SOLPS mesh geometry for a given MDSplus connection. :param mds_connection: MDSplus connection object. Already set to the SOLPS tree with pulse #ID. + :param mdsExceptions: MDSplus mdsExceptions module for error handling. """ # Load the R, Z coordinates of the cell vertices, original coordinates are (4, 38, 98) - # re-arrange axes (4, 38, 98) => (98, 38, 4) - x = np.swapaxes(mds_connection.get('\TOP.SNAPSHOT.GRID:R').data(), 0, 2) - z = np.swapaxes(mds_connection.get('\TOP.SNAPSHOT.GRID:Z').data(), 0, 2) + r = mds_connection.get(r'\TOP.SNAPSHOT.GRID:R').data() + z = mds_connection.get(r'\TOP.SNAPSHOT.GRID:Z').data() + + vol = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.VOL').data() + + # Loading neighbouring cell indices + neighbix = np.zeros(r.shape, dtype=np.int) + neighbiy = np.zeros(r.shape, dtype=np.int) - vol = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.VOL').data(), 0, 1) + neighbix[0] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:LEFTIX').data().astype(np.int) + neighbix[1] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:BOTTOMIX').data().astype(np.int) + neighbix[2] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:RIGHTIX').data().astype(np.int) + neighbix[3] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:TOPIX').data().astype(np.int) + + neighbiy[0] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:LEFTIY').data().astype(np.int) + neighbiy[1] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:BOTTOMIY').data().astype(np.int) + neighbiy[2] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:RIGHTIY').data().astype(np.int) + neighbiy[3] = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:TOPIY').data().astype(np.int) + + neighbix[neighbix == r.shape[2]] = -1 + neighbiy[neighbiy == r.shape[1]] = -1 # build mesh object - mesh = SOLPSMesh(x, z, vol) + mesh = SOLPSMesh(r, z, vol, neighbix, neighbiy) ############################# # Add additional parameters # ############################# # add the vessel geometry - mesh.vessel = mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:VESSEL').data() - - # Load the centre points of the grid cells. - cr = np.swapaxes(mds_connection.get('\TOP.SNAPSHOT.GRID:CR').data(), 0, 1) - cz = np.swapaxes(mds_connection.get('\TOP.SNAPSHOT.GRID:CZ').data(), 0, 1) - mesh._cr = cr - mesh._cz = cz - - # Load cell basis vectors - nx = mesh.nx - ny = mesh.ny - - cell_poloidal_basis = np.empty((nx, ny, 2), dtype=object) - for i in range(nx): - for j in range(ny): - - # Work out cell's 2D parallel vector in the poloidal plane - if i == nx - 1: - # Special case for end of array, repeat previous calculation. - # This is because I don't have access to the gaurd cells. - xp_x = cr[i, j] - cr[i-1, j] - xp_y = cz[i, j] - cz[i-1, j] - norm = np.sqrt(xp_x**2 + xp_y**2) - cell_poloidal_basis[i, j, 0] = Point2D(xp_x/norm, xp_y/norm) - else: - xp_x = cr[i+1, j] - cr[i, j] - xp_y = cz[i+1, j] - cz[i, j] - norm = np.sqrt(xp_x**2 + xp_y**2) - cell_poloidal_basis[i, j, 0] = Point2D(xp_x/norm, xp_y/norm) - - # Work out cell's 2D radial vector in the poloidal plane - if j == ny - 1: - # Special case for end of array, repeat previous calculation. - yr_x = cr[i, j] - cr[i, j-1] - yr_y = cz[i, j] - cz[i, j-1] - norm = np.sqrt(yr_x**2 + yr_y**2) - cell_poloidal_basis[i, j, 1] = Point2D(yr_x/norm, yr_y/norm) - else: - yr_x = cr[i, j+1] - cr[i, j] - yr_y = cz[i, j+1] - cz[i, j] - norm = np.sqrt(yr_x**2 + yr_y**2) - cell_poloidal_basis[i, j, 1] = Point2D(yr_x/norm, yr_y/norm) - - mesh._poloidal_grid_basis = cell_poloidal_basis + try: + vessel = mds_connection.get(r'\SOLPS::TOP.SNAPSHOT.GRID:VESSEL').data()[:] + mesh.vessel = vessel + except (MdsException, TypeError): + pass return mesh diff --git a/cherab/solps/formats/raw_pickle.py b/cherab/solps/formats/raw_pickle.py index b4a99aa..c1e0e86 100644 --- a/cherab/solps/formats/raw_pickle.py +++ b/cherab/solps/formats/raw_pickle.py @@ -26,8 +26,8 @@ def load_solps_from_pickle(filename): file_handle = open(filename, 'rb') state = pickle.load(file_handle) - mesh = SOLPSMesh(state['mesh']['cr_r'], state['mesh']['cr_z'], state['mesh']['vol']) - simulation = SOLPSSimulation(mesh) + mesh = SOLPSMesh(state['mesh']['r'], state['mesh']['z'], state['mesh']['vol'], state['mesh']['neighbix'], state['mesh']['neighbiy']) + simulation = SOLPSSimulation(mesh, state['species_list']) simulation.__setstate__(state) file_handle.close() diff --git a/cherab/solps/formats/raw_simulation_files.py b/cherab/solps/formats/raw_simulation_files.py index ea0354d..244338c 100644 --- a/cherab/solps/formats/raw_simulation_files.py +++ b/cherab/solps/formats/raw_simulation_files.py @@ -19,33 +19,14 @@ import os import numpy as np -from raysect.core.math.function.float import Discrete2DMesh +from scipy.constants import elementary_charge -from cherab.core.math.mappers import AxisymmetricMapper -from cherab.core.atomic.elements import hydrogen, deuterium, helium, beryllium, carbon, nitrogen, oxygen, neon, \ - argon, krypton, xenon +from cherab.core.atomic.elements import lookup_isotope from cherab.solps.eirene import load_fort44_file from cherab.solps.b2.parse_b2_block_file import load_b2f_file from cherab.solps.mesh_geometry import SOLPSMesh -from cherab.solps.solps_plasma import SOLPSSimulation - -Q = 1.602E-19 - -# key is nuclear charge Z and atomic mass AMU -_popular_species = { - (1, 2): deuterium, - (2, 4.003): helium, - (2, 4.0): helium, - (6, 12.011): carbon, - (6, 12.0): carbon, - (7, 14.007): nitrogen, - (10, 20.0): neon, - (10, 20.1797): neon, - (18, 39.948): argon, - (36, 83.798): krypton, - (54, 131.293): xenon -} +from cherab.solps.solps_plasma import SOLPSSimulation, prefer_element, eirene_flux_to_velocity, b2_flux_to_velocity # Code based on script by Felix Reimold (2016) @@ -56,93 +37,156 @@ def load_solps_from_raw_output(simulation_path, debug=False): Required files include: * mesh description file (b2fgmtry) * B2 plasma state (b2fstate) - * Eirene output file (fort.44) + * Eirene output file (fort.44), optional :param str simulation_path: String path to simulation directory. :rtype: SOLPSSimulation """ if not os.path.isdir(simulation_path): - RuntimeError("simulation_path must be a valid directory") + raise RuntimeError("simulation_path must be a valid directory.") mesh_file_path = os.path.join(simulation_path, 'b2fgmtry') b2_state_file = os.path.join(simulation_path, 'b2fstate') eirene_fort44_file = os.path.join(simulation_path, "fort.44") if not os.path.isfile(mesh_file_path): - raise RuntimeError("No B2 b2fgmtry file found in SOLPS output directory") + raise RuntimeError("No B2 b2fgmtry file found in SOLPS output directory.") - if not(os.path.isfile(b2_state_file)): - RuntimeError("No B2 b2fstate file found in SOLPS output directory") + if not os.path.isfile(b2_state_file): + raise RuntimeError("No B2 b2fstate file found in SOLPS output directory.") - if not(os.path.isfile(eirene_fort44_file)): - RuntimeError("No EIRENE fort.44 file found in SOLPS output directory") + if not os.path.isfile(eirene_fort44_file): + print("Warning! No EIRENE fort.44 file found in SOLPS output directory. Assuming B2 stand-alone simulation.") + b2_standalone = True + else: + # Load data for neutral species from EIRENE output file + eirene = load_fort44_file(eirene_fort44_file, debug=debug) + b2_standalone = False # Load SOLPS mesh geometry - mesh = load_mesh_from_files(mesh_file_path=mesh_file_path, debug=debug) + _, _, geom_data_dict = load_b2f_file(mesh_file_path, debug=debug) # geom_data_dict is needed also for magnetic field - header_dict, sim_info_dict, mesh_data_dict = load_b2f_file(b2_state_file, debug=debug) - - sim = SOLPSSimulation(mesh) - ni = mesh.nx - nj = mesh.ny - - # TODO: add code to load SOLPS velocities and magnetic field from files + mesh = create_mesh_from_geom_data(geom_data_dict) - # Load electron species - sim._electron_temperature = mesh_data_dict['te']/Q - sim._electron_density = mesh_data_dict['ne'] + ny = mesh.ny # radial + nx = mesh.nx # poloidal - ########################################## - # Load each plasma species in simulation # - ########################################## + header_dict, sim_info_dict, mesh_data_dict = load_b2f_file(b2_state_file, debug=debug) - sim._species_list = [] + # Load each plasma species in simulation + species_list = [] + neutral_indx = [] for i in range(len(sim_info_dict['zn'])): zn = int(sim_info_dict['zn'][i]) # Nuclear charge - am = float(sim_info_dict['am'][i]) # Atomic mass + am = int(round(float(sim_info_dict['am'][i]))) # Atomic mass number charge = int(sim_info_dict['zamax'][i]) # Ionisation/charge - species = _popular_species[(zn, am)] - sim.species_list.append(species.symbol + str(charge)) - - sim._species_density = mesh_data_dict['na'] + isotope = lookup_isotope(zn, number=am) + species = prefer_element(isotope) # Prefer Element over Isotope if the mass number is the same + species_list.append((species.name, charge)) + if charge == 0: # updating neutral index + neutral_indx.append(i) - # Make Mesh Interpolator function for inside/outside mesh test. - inside_outside_data = np.ones(mesh.num_tris) - inside_outside = AxisymmetricMapper(Discrete2DMesh(mesh.vertex_coords, mesh.triangles, inside_outside_data, limit=False)) - sim._inside_mesh = inside_outside + sim = SOLPSSimulation(mesh, species_list) - # Load total radiated power from EIRENE output file - eirene = load_fort44_file(eirene_fort44_file, debug=debug) - sim._eirene = eirene + # Load magnetic field + sim.b_field = geom_data_dict['bb'][:3] + # sim.b_field_cylindrical is created automatically - # Note EIRENE data grid is slightly smaller than SOLPS grid, for example (98, 38) => (96, 36) - # Need to pad EIRENE data to fit inside larger B2 array - nx = mesh.nx - ny = mesh.ny - eradt_raw_data = eirene.eradt.sum(2) - eradt_data = np.zeros((nx, ny)) - eradt_data[1:nx-1, 1:ny-1] = eradt_raw_data - sim._total_rad = eradt_data + # Load electron species + sim.electron_temperature = mesh_data_dict['te'] / elementary_charge + sim.electron_density = mesh_data_dict['ne'] + + # Load ion temperature + sim.ion_temperature = mesh_data_dict['ti'] / elementary_charge + + # Load species density + sim.species_density = mesh_data_dict['na'] + + # Load parallel velocity + parallel_velocity = mesh_data_dict['ua'] + + # Load poloidal and radial particle fluxes for velocity calculation + poloidal_flux = mesh_data_dict['fna'][::2] + radial_flux = mesh_data_dict['fna'][1::2] + + # Obtaining velocity from B2 flux + sim.velocities_cylindrical = b2_flux_to_velocity(sim, poloidal_flux, radial_flux, parallel_velocity) + + if not b2_standalone: + # Obtaining additional data from EIRENE and replacing data for neutrals + # Note EIRENE data grid is slightly smaller than SOLPS grid, for example (98, 38) => (96, 36) + # Need to pad EIRENE data to fit inside larger B2 array + + neutral_density = np.zeros((len(neutral_indx), ny, nx)) + neutral_density[:, 1:-1, 1:-1] = eirene.da + sim.species_density[neutral_indx] = neutral_density + + # Obtaining neutral atom velocity from EIRENE flux + # Note that if the output for fluxes was turned off, eirene.ppa and eirene.rpa are all zeros + if np.any(eirene.ppa) or np.any(eirene.rpa): + neutral_poloidal_flux = np.zeros((len(neutral_indx), ny, nx)) + neutral_poloidal_flux[:, 1:-1, 1:-1] = eirene.ppa + + neutral_radial_flux = np.zeros((len(neutral_indx), ny, nx)) + neutral_radial_flux[:, 1:-1, 1:-1] = eirene.rpa + + neutral_parallel_velocity = np.zeros((len(neutral_indx), ny, nx)) # must be zero outside EIRENE grid + neutral_parallel_velocity[:, 1:-1, 1:-1] = parallel_velocity[neutral_indx, 1:-1, 1:-1] + + sim.velocities_cylindrical[neutral_indx] = eirene_flux_to_velocity(sim, neutral_poloidal_flux, neutral_radial_flux, + neutral_parallel_velocity) + sim.velocities_cylindrical = sim.velocities_cylindrical # Updating sim.velocities + + # Obtaining neutral temperatures + ta = np.zeros((eirene.ta.shape[0], ny, nx)) + ta[:, 1:-1, 1:-1] = eirene.ta + # extrapolating + for i in (0, -1): + ta[:, i, 1:-1] = eirene.ta[:, i, :] + ta[:, 1:-1, i] = eirene.ta[:, :, i] + for i, j in ((0, 0), (0, -1), (-1, 0), (-1, -1)): + ta[:, i, j] = eirene.ta[:, i, j] + sim.neutral_temperature = ta / elementary_charge + + # Obtaining total radiation + if eirene.eradt is not None: + eradt_raw_data = eirene.eradt.sum(0) + total_radiation = np.zeros((ny, nx)) + total_radiation[1:-1, 1:-1] = eradt_raw_data + sim.total_radiation = total_radiation + + sim.eirene_simulation = eirene return sim +def create_mesh_from_geom_data(geom_data): -def load_mesh_from_files(mesh_file_path, debug=False): - """ - Load SOLPS grid description from B2 Eirene output file. + r = geom_data['crx'] + z = geom_data['cry'] + vol = geom_data['vol'] - :param str filepath: full path for B2 eirene mesh description file - :param bool debug: flag for displaying textual debugging information. - :return: tuple of dictionaries. First is the header information such as the version, label, grid size, etc. - Second dictionary has a ndarray for each piece of data found in the file. - """ - _, _, geom_data_dict = load_b2f_file(mesh_file_path, debug=debug) + # Loading neighbouring cell indices + neighbix = np.zeros(r.shape, dtype=np.int) + neighbiy = np.zeros(r.shape, dtype=np.int) + + neighbix[0] = geom_data['leftix'].astype(np.int) # poloidal prev. + neighbix[1] = geom_data['bottomix'].astype(np.int) # radial prev. + neighbix[2] = geom_data['rightix'].astype(np.int) # poloidal next + neighbix[3] = geom_data['topix'].astype(np.int) # radial next + + neighbiy[0] = geom_data['leftiy'].astype(np.int) + neighbiy[1] = geom_data['bottomiy'].astype(np.int) + neighbiy[2] = geom_data['rightiy'].astype(np.int) + neighbiy[3] = geom_data['topiy'].astype(np.int) + + # In SOLPS cell indexing starts with -1 (guarding cell), but in SOLPSMesh -1 means no neighbour. + neighbix += 1 + neighbiy += 1 + neighbix[neighbix == r.shape[2]] = -1 + neighbiy[neighbiy == r.shape[1]] = -1 - cr_x = geom_data_dict['crx'] - cr_z = geom_data_dict['cry'] - vol = geom_data_dict['vol'] + mesh = SOLPSMesh(r, z, vol, neighbix, neighbiy) - # build mesh object - return SOLPSMesh(cr_x, cr_z, vol) + return mesh diff --git a/cherab/solps/mesh_geometry.py b/cherab/solps/mesh_geometry.py index 77d117a..5cab148 100755 --- a/cherab/solps/mesh_geometry.py +++ b/cherab/solps/mesh_geometry.py @@ -17,16 +17,10 @@ # See the Licence for the specific language governing permissions and limitations # under the Licence. -# External imports -from collections import namedtuple - import numpy as np import matplotlib.pyplot as plt from matplotlib.patches import Polygon from matplotlib.collections import PatchCollection -from raysect.core.math.function.float import Discrete2DMesh - -INFINITY = 1E99 class SOLPSMesh: @@ -40,89 +34,114 @@ class SOLPSMesh: triangle vertices. Therefore, each SOLPS rectangular cell is split into two triangular cells. The data points are later interpolated onto the vertex points. - :param ndarray cr_r: Array of cell vertex r coordinates, must be 3 dimensional. Example shape is (98 x 32 x 4). - :param ndarray cr_z: Array of cell vertex z coordinates, must be 3 dimensional. Example shape is (98 x 32 x 4). - :param ndarray vol: Array of cell volumes. Example shape is (98 x 32). + :param ndarray r: Array of cell vertex r coordinates, must be 3 dimensional. Example shape is (4 x 32 x 98). + :param ndarray z: Array of cell vertex z coordinates, must be 3 dimensional. Example shape is (4 x 32 x 98). + :param ndarray vol: Array of cell volumes. Example shape is (32 x 98). + :param ndarray neighbix: Array of poloidal indeces of neighbouring cells in order: left, bottom, right, top, + must be 3 dimensional. Example shape is (4 x 32 x 98). + In SOLPS notation: left/right - poloidal prev./next, bottom/top - radial prev./next. + Cell indexing starts with 0 and -1 means no neighbour. + :param ndarray neighbiy: Array of radial indeces of neighbouring cells in order: left, bottom, right, top, + must be 3 dimensional. Example shape is (4 x 32 x 98). """ - def __init__(self, cr_r, cr_z, vol): + # TODO Make neighbix and neighbix optional in the future, as they can be reconstructed with _triangle_to_grid_map + + def __init__(self, r, z, vol, neighbix, neighbiy): + + if r.shape != z.shape: + raise ValueError('Shape of r array: {0} mismatch the shape of z array: {1}.'.format(r.shape, z.shape)) - self._cr = None - self._cz = None - self._poloidal_grid_basis = None + if vol.shape != r.shape[1:]: + raise ValueError('Shape of vol array: {0} mismatch the grid dimentions: {1}.'.format(vol.shape, r.shape[1:])) - nx = cr_r.shape[0] - ny = cr_r.shape[1] - self._nx = nx - self._ny = ny + if neighbix.shape != r.shape: + raise ValueError('Shape of neighbix array must be {0}, but it is {1}.'.format(r.shape, neighbix.shape)) - self._r = cr_r - self._z = cr_z + if neighbiy.shape != r.shape: + raise ValueError('Shape of neighbix array must be {0}, but it is {}.'.format(r.shape, neighbiy.shape)) + + self._cr = r.sum(0) / 4. + self._cz = z.sum(0) / 4. + + self._nx = r.shape[2] # poloidal + self._ny = r.shape[1] # radial + + self._r = r + self._z = z self._vol = vol - # Iterate through the arrays from MDS plus to pull out unique vertices - unique_vertices = {} - vertex_id = 0 - for i in range(nx): - for j in range(ny): - for k in range(4): - vertex = (cr_r[i, j, k], cr_z[i, j, k]) - try: - unique_vertices[vertex] - except KeyError: - unique_vertices[vertex] = vertex_id - vertex_id += 1 - - # Load these unique vertices into a numpy array for later use in Raysect's mesh interpolator object. - self.num_vertices = len(unique_vertices) - self.vertex_coords = np.zeros((self.num_vertices, 2), dtype=np.float64) - for vertex, vertex_id in unique_vertices.items(): - self.vertex_coords[vertex_id, :] = vertex + self._neighbix = neighbix.astype(np.int) + self._neighbiy = neighbiy.astype(np.int) + + self.vessel = None + + # Calculating poloidal basis vector + self._poloidal_basis_vector = np.zeros((2, self._ny, self._nx)) + vec_r = r[1] - r[0] + vec_z = z[1] - z[0] + vec_magn = np.sqrt(vec_r**2 + vec_z**2) + self._poloidal_basis_vector[0] = np.divide(vec_r, vec_magn, out=np.zeros_like(vec_magn), where=(vec_magn > 0)) + self._poloidal_basis_vector[1] = np.divide(vec_z, vec_magn, out=np.zeros_like(vec_magn), where=(vec_magn > 0)) + + # Calculating radial contact areas + self._radial_area = np.pi * (r[1] + r[0]) * vec_magn + + # Calculating radial basis vector + self._radial_basis_vector = np.zeros((2, self._ny, self._nx)) + vec_r = r[2] - r[0] + vec_z = z[2] - z[0] + vec_magn = np.sqrt(vec_r**2 + vec_z**2) + self._radial_basis_vector[0] = np.divide(vec_r, vec_magn, out=np.zeros_like(vec_magn), where=(vec_magn > 0)) + self._radial_basis_vector[1] = np.divide(vec_z, vec_magn, out=np.zeros_like(vec_magn), where=(vec_magn > 0)) + + # Calculating poloidal contact areas + self._poloidal_area = np.pi * (r[2] + r[0]) * vec_magn + + # For convertion from Cartesian to poloidal + # TODO Make it work with triangle cells + self._inv_det = 1. / (self._poloidal_basis_vector[0] * self._radial_basis_vector[1] - + self._poloidal_basis_vector[1] * self._radial_basis_vector[0]) + + # Finding unique vertices + vertices = np.array([r.flatten(), z.flatten()]).T + self._vertex_coords, unique_vertices = np.unique(vertices, axis=0, return_inverse=True) + self._num_vertices = self._vertex_coords.shape[0] # Work out the extent of the mesh. - rmin = cr_r.flatten().min() - rmax = cr_r.flatten().max() - zmin = cr_z.flatten().min() - zmax = cr_z.flatten().max() - self.mesh_extent = {"minr": rmin, "maxr": rmax, "minz": zmin, "maxz": zmax} + self._mesh_extent = {"minr": r.min(), "maxr": r.max(), "minz": z.min(), "maxz": z.max()} # Number of triangles must be equal to number of rectangle centre points times 2. - self.num_tris = nx * ny * 2 - self.triangles = np.zeros((self.num_tris, 3), dtype=np.int32) - - self._triangle_to_grid_map = np.zeros((nx*ny*2, 2), dtype=np.int32) - tri_index = 0 - for i in range(nx): - for j in range(ny): - # Pull out the index number for each unique vertex in this rectangular cell. - # Unusual vertex indexing is based on SOLPS output, see Matlab code extract from David Moulton. - # cell_r = [r(i,j,1),r(i,j,3),r(i,j,4),r(i,j,2)]; - v1_id = unique_vertices[(cr_r[i, j, 0], cr_z[i, j, 0])] - v2_id = unique_vertices[(cr_r[i, j, 2], cr_z[i, j, 2])] - v3_id = unique_vertices[(cr_r[i, j, 3], cr_z[i, j, 3])] - v4_id = unique_vertices[(cr_r[i, j, 1], cr_z[i, j, 1])] - - # Split the quad cell into two triangular cells. - # Each triangle cell is mapped to the tuple ID (ix, iy) of its parent mesh cell. - self.triangles[tri_index, :] = (v1_id, v2_id, v3_id) - self._triangle_to_grid_map[tri_index, :] = (i, j) - tri_index += 1 - self.triangles[tri_index, :] = (v3_id, v4_id, v1_id) - self._triangle_to_grid_map[tri_index, :] = (i, j) - tri_index += 1 - - tri_indices = np.arange(self.num_tris, dtype=np.int32) - self._tri_index_loopup = Discrete2DMesh(self.vertex_coords, self.triangles, tri_indices) + self._num_tris = self._nx * self._ny * 2 + self._triangles = np.zeros((self._num_tris, 3), dtype=np.int32) + self._triangle_to_grid_map = np.zeros((self._num_tris, 2), dtype=np.int32) + + # Pull out the index number for each unique vertex in this rectangular cell. + # Unusual vertex indexing is based on SOLPS output, see Matlab code extract from David Moulton. + ng = self._nx * self._ny # grid size + self._triangles[0::2, 0] = unique_vertices[0:ng] + self._triangles[0::2, 1] = unique_vertices[2 * ng: 3 * ng] + self._triangles[0::2, 2] = unique_vertices[3 * ng: 4 * ng] + # Split the quad cell into two triangular cells. + self._triangles[1::2, 0] = unique_vertices[3 * ng: 4 * ng] + self._triangles[1::2, 1] = unique_vertices[ng: 2 * ng] + self._triangles[1::2, 2] = unique_vertices[0:ng] + + # Each triangle cell is mapped to the tuple ID (ix, iy) of its parent mesh cell. + ym, xm = np.meshgrid(np.arange(self._ny, dtype=np.int32), np.arange(self._nx, dtype=np.int32), indexing='ij') + self._triangle_to_grid_map[::2, 0] = ym.flatten() + self._triangle_to_grid_map[::2, 1] = xm.flatten() + self._triangle_to_grid_map[1::2] = self._triangle_to_grid_map[::2] @property def nx(self): - """Number of grid cells in the x direction.""" + """Number of grid cells in the poloidal direction.""" return self._nx @property def ny(self): - """Number of grid cells in the y direction.""" + """Number of grid cells in the radial direction.""" return self._ny @property @@ -135,54 +154,143 @@ def cz(self): """Z-coordinate of the cell centres.""" return self._cz + @property + def r(self): + """R-coordinates of the cell vertices.""" + return self._r + + @property + def z(self): + """Z-coordinate of the cell vertices.""" + return self._z + @property def vol(self): """Volume/area of each grid cell.""" return self._vol @property - def poloidal_grid_basis(self): + def neighbix(self): + """Poloidal indeces of neighbouring cells in order: left, bottom, right, top.""" + return self._neighbix + + @property + def neighbiy(self): + """Radial indeces of neighbouring cells in order: left, bottom, right, top.""" + return self._neighbiy + + @property + def radial_area(self): + """Radial contact area.""" + return self._radial_area + + @property + def poloidal_area(self): + """Poloidal contact area.""" + return self._poloidal_area + + @property + def vertex_coordinates(self): + """RZ-coordinates of unique vertices.""" + return self._vertex_coords + + @property + def num_vertices(self): + """Total number of unique vertices.""" + return self._num_vertices + + @property + def mesh_extent(self): + """Extent of the mesh. A dictionary with minr, maxr, minz and maxz keys.""" + return self._mesh_extent + + @property + def num_triangles(self): + """Total number of triangles (the number of cells doubled).""" + return self._num_tris + + @property + def triangles(self): + """Array of triangle vertex indices with (num_thiangles, 3) shape.""" + return self._triangles + + @property + def poloidal_basis_vector(self): """ - Array of 2D basis vectors for grid cells. + Array of 2D poloidal basis vectors for grid cells. - For each cell there is a parallel and radial basis vector. + For each cell there is a poloidal and radial basis vector. Any vector on the poloidal grid can be converted to cartesian with the following transformation. bx = (p_x r_x) ( b_p ) by (p_y r_y) ( b_r ) - :return: ndarray with shape (nx, ny, 2) where the two basis vectors are [parallel, radial] respectively. + :return: ndarray with shape (2, ny, nx). """ - return self._poloidal_grid_basis + return self._poloidal_basis_vector @property - def triangle_to_grid_map(self): + def radial_basis_vector(self): """ - Array mapping every triangle index to a tuple grid cell ID, i.e. (ix, iy). + Array of 2D radial basis vectors for grid cells. - :return: ndarray with shape (nx*ny*2, 2) + For each cell there is a poloidal and radial basis vector. + + Any vector on the poloidal grid can be converted to cartesian with the following transformation. + bx = (p_x r_x) ( b_p ) + by (p_y r_y) ( b_r ) + + :return: ndarray with shape (2, ny, nx). """ - return self._triangle_to_grid_map + return self._radial_basis_vector @property - def triangle_index_lookup(self): + def triangle_to_grid_map(self): """ - Discrete2DMesh instance that looks up a triangle index at any 2D point. + Array mapping every triangle index to a tuple grid cell ID, i.e. (iy, ix). - Useful for mapping from a 2D point -> triangle cell -> parent SOLPS mesh cell - - :return: Discrete2DMesh instance + :return: ndarray with shape (nx*ny*2, 2) """ - return self._tri_index_loopup + return self._triangle_to_grid_map def __getstate__(self): state = { - 'cr_r': self._r, - 'cr_z': self._z, + 'r': self._r, + 'z': self._z, 'vol': self._vol, + 'neighbix': self._neighbix, + 'neighbiy': self._neighbiy } return state + def to_cartesian(self, vec_pol): + """ + Converts the 2D vector defined on mesh from poloidal to cartesian coordinates. + :param ndarray vec_pol: Array of 2D vector with with shape (2, ny, nx). + [0, :, :] - poloidal component, [1, :, :] - radial component + + :return: ndarray with shape (2, ny, nx) + """ + vec_cart = np.zeros((2, self._ny, self._nx)) + vec_cart[0] = self._poloidal_basis_vector[0] * vec_pol[0] + self._radial_basis_vector[0] * vec_pol[1] + vec_cart[1] = self._poloidal_basis_vector[1] * vec_pol[0] + self._radial_basis_vector[1] * vec_pol[1] + + return vec_cart + + def to_poloidal(self, vec_cart): + """ + Converts the 2D vector defined on mesh from cartesian to poloidal coordinates. + :param ndarray vector_on_mesh: Array of 2D vector with with shape (2, ny, nx). + [0, :, :] - R component, [1, :, :] - Z component + + :return: ndarray with shape (2, ny, nx) + """ + vec_pol = np.zeros((2, self._ny, self._nx)) + vec_pol[0] = self._inv_det * (self._radial_basis_vector[1] * vec_cart[0] - self._radial_basis_vector[0] * vec_cart[1]) + vec_pol[1] = self._inv_det * (self._poloidal_basis_vector[0] * vec_cart[1] - self._poloidal_basis_vector[1] * vec_cart[0]) + + return vec_pol + def plot_mesh(self): """ Plot the mesh grid geometry to a matplotlib figure. @@ -190,15 +298,10 @@ def plot_mesh(self): fig, ax = plt.subplots() patches = [] for triangle in self.triangles: - vertices = self.vertex_coords[triangle] + vertices = self.vertex_coordinates[triangle] patches.append(Polygon(vertices, closed=True)) p = PatchCollection(patches, facecolors='none', edgecolors='b') ax.add_collection(p) ax.axis('equal') - return ax - # Code for plotting vessel geometry if available - # for i in range(vessel.shape[0]): - # plt.plot([vessel[i, 0], vessel[i, 2]], [vessel[i, 1], vessel[i, 3]], 'k') - # for i in range(vessel.shape[0]): - # plt.plot([vessel[i, 0], vessel[i, 2]], [vessel[i, 1], vessel[i, 3]], 'or') + return ax diff --git a/cherab/solps/models/radiated_power.pxd b/cherab/solps/models/radiated_power.pxd index 89cae88..0dd7293 100755 --- a/cherab/solps/models/radiated_power.pxd +++ b/cherab/solps/models/radiated_power.pxd @@ -29,7 +29,7 @@ cdef class SOLPSTotalRadiatedPower(InhomogeneousVolumeEmitter): cdef: public double vertical_offset - Function3D total_rad, inside_simulation + Function3D total_rad cpdef Spectrum emission_function(self, Point3D point, Vector3D direction, Spectrum spectrum, World world, Ray ray, Primitive primitive, diff --git a/cherab/solps/models/radiated_power.pyx b/cherab/solps/models/radiated_power.pyx index fa4f0d3..f7f5566 100755 --- a/cherab/solps/models/radiated_power.pyx +++ b/cherab/solps/models/radiated_power.pyx @@ -30,19 +30,15 @@ cimport cython cdef class SOLPSTotalRadiatedPower(InhomogeneousVolumeEmitter): - def __init__(self, object solps_simulation, double vertical_offset=0.0, step=0.01): + def __init__(self, Function3D total_radiation, double vertical_offset=0.0, step=0.01): super().__init__(NumericalIntegrator(step=step)) self.vertical_offset = vertical_offset - self.inside_simulation = solps_simulation.inside_volume_mesh - self.total_rad = solps_simulation.total_radiation_volume + self.total_rad = total_radiation - def __call__(self, x, y , z): + def __call__(self, x, y, z): - if self.inside_simulation.evaluate(x, y, z) < 1.0: - return 0.0 - - return self.total_rad.evaluate(x, y, z) + return self.total_rad.evaluate(x, y, z) # this returns 0 if outside the mesh @cython.boundscheck(False) @cython.wraparound(False) @@ -54,16 +50,18 @@ cdef class SOLPSTotalRadiatedPower(InhomogeneousVolumeEmitter): cdef: double offset_z, wvl_range - offset_z = point.z + self.vertical_offset - if self.inside_simulation.evaluate(point.x, point.y, offset_z) < 1.0: - return spectrum + offset_z = point.z + self.vertical_offset wvl_range = ray.max_wavelength - ray.min_wavelength spectrum.samples_mv[:] = self.total_rad.evaluate(point.x, point.y, offset_z) / (4 * pi * wvl_range * spectrum.bins) return spectrum + def solps_total_radiated_power(world, solps_simulation, step=0.01): + if solps_simulation.total_radiation_f3d is None: + raise RuntimeError('Total radiation is not available for this simulation.') + mesh = solps_simulation.mesh outer_radius = mesh.mesh_extent['maxr'] inner_radius = mesh.mesh_extent['minr'] - 0.001 @@ -71,9 +69,7 @@ def solps_total_radiated_power(world, solps_simulation, step=0.01): lower_z = mesh.mesh_extent['minz'] main_plasma_cylinder = Cylinder(outer_radius, plasma_height, parent=world, - material=SOLPSTotalRadiatedPower(solps_simulation, vertical_offset=lower_z, step=step), + material=SOLPSTotalRadiatedPower(solps_simulation.total_radiation_f3d, vertical_offset=lower_z, step=step), transform=translate(0, 0, lower_z)) return main_plasma_cylinder - - diff --git a/cherab/solps/solps_2d_functions.pyx b/cherab/solps/solps_2d_functions.pyx new file mode 100755 index 0000000..6b40f78 --- /dev/null +++ b/cherab/solps/solps_2d_functions.pyx @@ -0,0 +1,216 @@ +# cython: language_level=3 + +# Copyright 2016-2018 Euratom +# Copyright 2016-2018 United Kingdom Atomic Energy Authority +# Copyright 2016-2018 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas +# +# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the +# European Commission - subsequent versions of the EUPL (the "Licence"); +# You may not use this work except in compliance with the Licence. +# You may obtain a copy of the Licence at: +# +# https://joinup.ec.europa.eu/software/page/eupl5 +# +# Unless required by applicable law or agreed to in writing, software distributed +# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR +# CONDITIONS OF ANY KIND, either express or implied. +# +# See the Licence for the specific language governing permissions and limitations +# under the Licence. + +import numpy as np +cimport numpy as np + +from raysect.core.math.function.float.function2d.interpolate.common cimport MeshKDTree2D +from raysect.core.math.vector cimport Vector3D, new_vector3d +from raysect.core.math.point cimport new_point2d + +from cherab.core.math.function cimport Function2D, VectorFunction2D +cimport cython + + +cdef class SOLPSFunction2D(Function2D): + + cdef: + MeshKDTree2D _kdtree + np.ndarray _grid_data, _triangle_to_grid_map + np.int32_t[:,::1] _triangle_to_grid_map_mv + double[:,::1] _grid_data_mv + + def __init__(self, object vertex_coords not None, object triangles not None, object triangle_to_grid_map not None, object grid_data not None): + + # use numpy arrays to store data internally + vertex_coords = np.array(vertex_coords, dtype=np.float64) + triangles = np.array(triangles, dtype=np.int32) + triangle_to_grid_map = np.array(triangle_to_grid_map, dtype=np.int32) + + # Attention!!! Do not copy grid_data! Attribute self._grid_data must point to the original data array, + # so as not to re-initialize the interpolator if the user changes data values. + + # build kdtree + self._kdtree = MeshKDTree2D(vertex_coords, triangles) + + # populate internal attributes + self._grid_data = grid_data + self._triangle_to_grid_map = triangle_to_grid_map + + self._grid_data_mv = self._grid_data + self._triangle_to_grid_map_mv = self._triangle_to_grid_map + + def __getstate__(self): + return self._grid_data, self._triangle_to_grid_map, self._kdtree + + def __setstate__(self, state): + self._grid_data, self._triangle_to_grid_map, self._kdtree = state + self._triangle_to_grid_map_mv = self._triangle_to_grid_map + self._grid_data_mv = self._grid_data + + def __reduce__(self): + return self.__new__, (self.__class__, ), self.__getstate__() + + @classmethod + def instance(cls, SOLPSFunction2D instance not None, object grid_data=None): + """ + Creates a new interpolator instance from an existing interpolator instance. + The new interpolator instance will share the same internal acceleration + data as the original interpolator. The grid_data of the new instance can + be redefined. + This method should be used if the user has multiple sets of grid_data + that lie on the same mesh geometry. Using this methods avoids the + repeated rebuilding of the mesh acceleration structures by sharing the + geometry data between multiple interpolator objects. + :param SOLPSFunction2D instance: SOLPSFunction2D object. + :param ndarray grid_data: An array containing data on SOLPS grid. + :return: A SOLPSFunction2D object. + :rtype: SOLPSFunction2D + """ + + cdef SOLPSFunction2D m + + # copy source data + m = SOLPSFunction2D.__new__(SOLPSFunction2D) + m._kdtree = instance._kdtree + m._triangle_to_grid_map = instance._triangle_to_grid_map + + # do we have replacement triangle data? + if grid_data is None: + m._grid_data = instance._grid_data + else: + m._grid_data = grid_data + + m._triangle_to_grid_map_mv = m._triangle_to_grid_map + m._grid_data_mv = m._grid_data + + return m + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.initializedcheck(False) + cdef double evaluate(self, double x, double y) except? -1e999: + + cdef: + np.int32_t triangle_id, ix, iy + + if self._kdtree.is_contained(new_point2d(x, y)): + + triangle_id = self._kdtree.triangle_id + ix = self._triangle_to_grid_map_mv[triangle_id, 0] + iy = self._triangle_to_grid_map_mv[triangle_id, 1] + return self._grid_data_mv[ix, iy] + + return 0.0 + +cdef class SOLPSVectorFunction2D(VectorFunction2D): + + cdef: + MeshKDTree2D _kdtree + np.ndarray _grid_vectors, _triangle_to_grid_map + np.int32_t[:,::1] _triangle_to_grid_map_mv + double[:,:,::1] _grid_vectors_mv + + def __init__(self, object vertex_coords not None, object triangles not None, object triangle_to_grid_map not None, object grid_vectors not None): + + # use numpy arrays to store data internally + vertex_coords = np.array(vertex_coords, dtype=np.float64) + triangles = np.array(triangles, dtype=np.int32) + triangle_to_grid_map = np.array(triangle_to_grid_map, dtype=np.int32) + + # Attention!!! Do not copy grid_vectors! Attribute self._grid_vectors must point to the original data array, + # so as not to re-initialize the interpolator if the user changes data values. + + # build kdtree + self._kdtree = MeshKDTree2D(vertex_coords, triangles) + + # populate internal attributes + self._grid_vectors = grid_vectors + self._triangle_to_grid_map = triangle_to_grid_map + self._grid_vectors_mv = self._grid_vectors + self._triangle_to_grid_map_mv = self._triangle_to_grid_map + + def __getstate__(self): + return self._grid_vectors, self._triangle_to_grid_map, self._kdtree + + def __setstate__(self, state): + self._grid_vectors, self._triangle_to_grid_map, self._kdtree = state + self._grid_vectors_mv = self._grid_vectors + self._triangle_to_grid_map_mv = self._triangle_to_grid_map + + def __reduce__(self): + return self.__new__, (self.__class__, ), self.__getstate__() + + @classmethod + def instance(cls, SOLPSVectorFunction2D instance not None, object grid_vectors=None): + """ + Creates a new interpolator instance from an existing interpolator instance. + The new interpolator instance will share the same internal acceleration + data as the original interpolator. The grid_data of the new instance can + be redefined. + This method should be used if the user has multiple sets of grid_data + that lie on the same mesh geometry. Using this methods avoids the + repeated rebuilding of the mesh acceleration structures by sharing the + geometry data between multiple interpolator objects. + :param SOLPSVectorFunction2D instance: SOLPSVectorFunction2D object. + :param ndarray grid_vectors: An array containing vector data on SOLPS grid. + :return: A SOLPSVectorFunction2D object. + :rtype: SOLPSVectorFunction2D + """ + + cdef SOLPSVectorFunction2D m + + # copy source data + m = SOLPSVectorFunction2D.__new__(SOLPSVectorFunction2D) + m._kdtree = instance._kdtree + m._triangle_to_grid_map = instance._triangle_to_grid_map + + # do we have replacement triangle data? + if grid_vectors is None: + m._grid_vectors = instance._grid_vectors + else: + m._grid_vectors = grid_vectors + + m._triangle_to_grid_map_mv = m._triangle_to_grid_map + m._grid_vectors_mv = m._grid_vectors + + return m + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.initializedcheck(False) + cdef Vector3D evaluate(self, double x, double y): + + cdef: + np.int32_t triangle_id, ix, iy + double vx, vy, vz + + if self._kdtree.is_contained(new_point2d(x, y)): + + triangle_id = self._kdtree.triangle_id + ix = self._triangle_to_grid_map_mv[triangle_id, 0] + iy = self._triangle_to_grid_map_mv[triangle_id, 1] + vx = self._grid_vectors_mv[0, ix, iy] + vy = self._grid_vectors_mv[1, ix, iy] + vz = self._grid_vectors_mv[2, ix, iy] + + return new_vector3d(vx, vy, vz) + + return new_vector3d(0, 0, 0) diff --git a/cherab/solps/solps_3d_functions.pyx b/cherab/solps/solps_3d_functions.pyx deleted file mode 100755 index c828ffc..0000000 --- a/cherab/solps/solps_3d_functions.pyx +++ /dev/null @@ -1,108 +0,0 @@ -# cython: language_level=3 - -# Copyright 2016-2018 Euratom -# Copyright 2016-2018 United Kingdom Atomic Energy Authority -# Copyright 2016-2018 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas -# -# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the -# European Commission - subsequent versions of the EUPL (the "Licence"); -# You may not use this work except in compliance with the Licence. -# You may obtain a copy of the Licence at: -# -# https://joinup.ec.europa.eu/software/page/eupl5 -# -# Unless required by applicable law or agreed to in writing, software distributed -# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR -# CONDITIONS OF ANY KIND, either express or implied. -# -# See the Licence for the specific language governing permissions and limitations -# under the Licence. - -from libc.math cimport atan2, M_PI -import numpy as np -from numpy cimport ndarray - -from raysect.core.math.vector cimport Vector3D, new_vector3d -from raysect.core.math.transform cimport rotate_z - -from cherab.core.math.mappers cimport AxisymmetricMapper -from cherab.core.math.function cimport Function3D, Discrete2DMesh, VectorFunction3D -cimport cython - -cdef double RAD_TO_DEG = 360 / (2*M_PI) - - -cdef class SOLPSFunction3D(Function3D): - - cdef: - AxisymmetricMapper _triangle_index_lookup - int[:,:] _triangle_to_grid_map - double[:,:] _grid_values - - def __init__(self, Discrete2DMesh triangle_index_lookup, int[:,:] triangle_to_grid_map, double[:,:] grid_values): - - # todo: this is unsafe - converting an int to a double and performing operations - rounding errors - self._triangle_index_lookup = AxisymmetricMapper(triangle_index_lookup) - self._triangle_to_grid_map = triangle_to_grid_map - self._grid_values = grid_values - - @cython.boundscheck(False) - @cython.wraparound(False) - @cython.initializedcheck(False) - cdef double evaluate(self, double x, double y, double z) except? -1e999: - - cdef: - int tri_index # Index of the underlying mesh triangle - int ix # SOLPS grid x coordinate - int iy # SOLPS grid y coordinate - - try: - tri_index = self._triangle_index_lookup.evaluate(x, y, z) - except ValueError: - return 0.0 # Return zero if outside of mesh bounds - - ix, iy = self._triangle_to_grid_map[tri_index, :] - - return self._grid_values[ix, iy] - - -cdef class SOLPSVectorFunction3D(VectorFunction3D): - - cdef: - AxisymmetricMapper _triangle_index_lookup - int[:,:] _triangle_to_grid_map - double[:,:,:] _grid_vectors - - def __init__(self, Discrete2DMesh triangle_index_lookup, object triangle_to_grid_map, object grid_vectors): - - # todo: this is unsafe - converting an int to a double and performing operations - rounding errors - self._triangle_index_lookup = AxisymmetricMapper(triangle_index_lookup) - self._triangle_to_grid_map = triangle_to_grid_map - self._grid_vectors = grid_vectors - - @cython.boundscheck(False) - @cython.wraparound(False) - @cython.initializedcheck(False) - cdef Vector3D evaluate(self, double x, double y, double z): - - cdef: - int tri_index # Index of the underlying mesh triangle - int ix # SOLPS grid x coordinate - int iy # SOLPS grid y coordinate - double vx, vy, vz - Vector3D v - - try: - tri_index = self._triangle_index_lookup.evaluate(x, y, z) - except ValueError: - return Vector3D(0.0, 0.0, 0.0) # Return zero vector if outside of mesh bounds. - - ix, iy = self._triangle_to_grid_map[tri_index, :] - - # print(self._grid_vectors.shape) - # Lookup vector for this grid cell. - vx, vy, vz = self._grid_vectors[ix, iy, :] - v = new_vector3d(vx, vy, vz) - - # Rotate vector field around the z-axis. - return v.transform(rotate_z(atan2(y, x) * RAD_TO_DEG)) diff --git a/cherab/solps/solps_plasma.py b/cherab/solps/solps_plasma.py index 4b21488..f2bac0a 100755 --- a/cherab/solps/solps_plasma.py +++ b/cherab/solps/solps_plasma.py @@ -17,228 +17,570 @@ # See the Licence for the specific language governing permissions and limitations # under the Licence. -import re import pickle import numpy as np -import matplotlib.pyplot as plt from scipy.constants import atomic_mass, electron_mass # Raysect imports -from raysect.core.math.function.float import Discrete2DMesh -from raysect.core import translate, Point3D, Vector3D, Node, AffineMatrix3D +from raysect.core import translate, Vector3D from raysect.primitive import Cylinder -from raysect.optical import Spectrum # CHERAB core imports from cherab.core import Plasma, Species, Maxwellian -from cherab.core.math.mappers import AxisymmetricMapper -from cherab.core.atomic.elements import ( - hydrogen, deuterium, helium, beryllium, carbon, boron, nitrogen, oxygen, neon, - argon, krypton, xenon -) +from cherab.core.math.function import ConstantVector3D +from cherab.core.math.mappers import AxisymmetricMapper, VectorAxisymmetricMapper +from cherab.core.atomic.elements import lookup_isotope, lookup_element # This SOLPS package imports -from .solps_3d_functions import SOLPSFunction3D, SOLPSVectorFunction3D +from cherab.solps.eirene import Eirene +from .solps_2d_functions import SOLPSFunction2D, SOLPSVectorFunction2D from .mesh_geometry import SOLPSMesh -Q = 1.602E-19 +class SOLPSSimulation: -_species_symbol_map = { - 'D': deuterium, - 'C': carbon, - 'He': helium, - 'B': boron, - 'N': nitrogen, - 'Ne': neon, - 'Ar': argon, - 'Kr': krypton, - 'Xe': xenon, -} + def __init__(self, mesh, species_list): -_SPECIES_REGEX = '([a-zA-z]+)\+?([0-9]+)' + # Mesh and species_list cannot be changed after initialisation + if isinstance(mesh, SOLPSMesh): + self._mesh = mesh + else: + raise ValueError('Argument "mesh" must be a SOLPSMesh instance.') -# TODO: this interface is half broken - some routines expect internal data as arrays, others as function 3d -class SOLPSSimulation: + # Make Mesh Interpolator function for inside/outside mesh test. + inside_outside_data = np.ones((self._mesh.ny, self._mesh.nx)) + self._inside_mesh = SOLPSFunction2D(mesh.vertex_coordinates, mesh.triangles, mesh.triangle_to_grid_map, inside_outside_data) - def __init__(self, mesh): + # Creating a sample SOLPSVectorFunction2D for KDtree to use later + sample_vector = np.ones((3, self._mesh.ny, self._mesh.nx)) + self._sample_vector_f2d = SOLPSVectorFunction2D(mesh.vertex_coordinates, mesh.triangles, mesh.triangle_to_grid_map, sample_vector) - self.mesh = mesh + if not len(species_list): + raise ValueError('Argument "species_list" must contain at least one species.') + self._species_list = tuple(species_list) # adding additional species is not allowed + self._neutral_list = tuple([sp for sp in self._species_list if sp[1] == 0]) self._electron_temperature = None + self._electron_temperature_f2d = None + self._electron_temperature_f3d = None self._electron_density = None - self._species_list = None + self._electron_density_f2d = None + self._electron_density_f3d = None + self._electron_velocities = None + self._electron_velocities_cylindrical = None + self._electron_velocities_cylindrical_f2d = None + self._electron_velocities_cartesian = None + self._ion_temperature = None + self._ion_temperature_f2d = None + self._ion_temperature_f3d = None + self._neutral_temperature = None + self._neutral_temperature_f2d = None + self._neutral_temperature_f3d = None self._species_density = None - self._rad_par_flux = None - self._radial_area = None - self._b2_neutral_densities = None - self._velocities_parallel = None - self._velocities_radial = None - self._velocities_toroidal = None + self._species_density_f2d = None + self._species_density_f3d = None + self._velocities = None + self._velocities_cylindrical = None + self._velocities_cylindrical_f2d = None self._velocities_cartesian = None - self._inside_mesh = None - self._total_rad = None - self._b_field_vectors = None - self._b_field_vectors_cartesian = None - self._parallel_velocities = None - self._radial_velocities = None + self._total_radiation = None + self._total_radiation_f2d = None + self._total_radiation_f3d = None + self._b_field = None + self._b_field_cylindrical = None + self._b_field_cylindrical_f2d = None + self._b_field_cartesian = None self._eirene_model = None self._b2_model = None self._eirene = None + @property + def mesh(self): + """ + SOLPSMesh instance. + :return: + """ + return self._mesh + + @property + def species_list(self): + """ + Tuple of species elements in the form (species name, charge). + :return: + """ + return self._species_list + + @property + def neutral_list(self): + """ + Tuple of species elements in the form (species name, charge). + :return: + """ + return self._neutral_list + @property def electron_temperature(self): """ Simulated electron temperatures at each mesh cell. + Array of shape (ny, nx). :return: """ return self._electron_temperature + @property + def electron_temperature_f2d(self): + """ + Function2D interpolator for electron temperature. + Returns electron temperature at a given point (R, Z). + :return: + """ + return self._electron_temperature_f2d + + @property + def electron_temperature_f3d(self): + """ + Function3D interpolator for electron temperature. + Returns electron temperature at a given point (x, y, z). + :return: + """ + return self._electron_temperature_f3d + + @electron_temperature.setter + def electron_temperature(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("electron_temperature", value, (self.mesh.ny, self.mesh.nx)) + self._electron_temperature = value + self._electron_temperature_f2d = SOLPSFunction2D.instance(self._inside_mesh, value) + self._electron_temperature_f3d = AxisymmetricMapper(self._electron_temperature_f2d) + + @property + def ion_temperature(self): + """ + Simulated ion temperatures at each mesh cell. + Array of shape (ny, nx). + :return: + """ + return self._ion_temperature + + @property + def ion_temperature_f2d(self): + """ + Function2D interpolator for ion temperature. + Returns ion temperature at a given point (R, Z). + :return: + """ + return self._ion_temperature_f2d + + @property + def ion_temperature_f3d(self): + """ + Function3D interpolator for ion temperature. + Returns ion temperature at a given point (x, y, z). + :return: + """ + return self._ion_temperature_f3d + + @ion_temperature.setter + def ion_temperature(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("ion_temperature", value, (self.mesh.ny, self.mesh.nx)) + self._ion_temperature = value + self._ion_temperature_f2d = SOLPSFunction2D.instance(self._inside_mesh, value) + self._ion_temperature_f3d = AxisymmetricMapper(self._ion_temperature_f2d) + + @property + def neutral_temperature(self): + """ + Simulated neutral atom (effective) temperatures at each mesh cell. + Array of shape (na, ny, nx). + :return: + """ + return self._neutral_temperature + + @property + def neutral_temperature_f2d(self): + """ + Dictionary of Function2D interpolators for neutral atom (effective) temperatures. + Accessed by neutral atom index or neutral_list elements. + E.g., neutral_temperature_f2d[0] or neutral_temperature_f2d[('deuterium', 0))]. + Each entry returns neutral atom (effective) temperature at a given point (R, Z). + :return: + """ + return self._neutral_temperature_f2d + + @property + def neutral_temperature_f3d(self): + """ + Dictionary of Function3D interpolators for neutral atom (effective) temperatures. + Accessed by neutral atom index or neutral_list elements. + E.g., neutral_temperature_f3d[0] or neutral_temperature_f3d[('deuterium', 0))]. + Each entry returns neutral atom (effective) temperature at a given point (x, y, z). + :return: + """ + return self._neutral_temperature_f3d + + @neutral_temperature.setter + def neutral_temperature(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("neutral_temperature", value, (len(self._neutral_list), self.mesh.ny, self.mesh.nx)) + self._neutral_temperature = value + self._neutral_temperature_f2d = {} + self._neutral_temperature_f3d = {} + for k, sp in enumerate(self._neutral_list): + self._neutral_temperature_f2d[k] = SOLPSFunction2D.instance(self._inside_mesh, value[k]) + self._neutral_temperature_f2d[sp] = self._neutral_temperature_f2d[k] + self._neutral_temperature_f3d[k] = AxisymmetricMapper(self._neutral_temperature_f2d[k]) + self._neutral_temperature_f3d[sp] = self._neutral_temperature_f3d[k] + @property def electron_density(self): """ Simulated electron densities at each mesh cell. + Array of shape (ny, nx) :return: """ return self._electron_density @property - def species_list(self): + def electron_density_f2d(self): """ - Text list of species elements present in the simulation. + Function2D interpolator for electron density. + Returns electron density at a given point (R, Z). :return: """ - return self._species_list + return self._electron_density_f2d @property - def species_density(self): + def electron_density_f3d(self): """ - Array of species densities at each mesh cell. + Function3D interpolator for electron density. + Returns electron density at a given point (x, y, z). :return: """ - return self._species_density + return self._electron_density_f3d + + @electron_density.setter + def electron_density(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("electron_density", value, (self.mesh.ny, self.mesh.nx)) + self._electron_density = value + self._electron_density_f2d = SOLPSFunction2D.instance(self._inside_mesh, value) + self._electron_density_f3d = AxisymmetricMapper(self._electron_density_f2d) + + @property + def electron_velocities(self): + """ + Electron velocities in poloidal coordinates (e_pol, e_rad, e_tor) at each mesh cell. + Array of shape (3, ny, nx): + [0, :, :] - poloidal, [1, :, :] - radial, [2, :, :] - toroidal. + :return: + """ + return self._electron_velocities + + @electron_velocities.setter + def electron_velocities(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("electron_velocities", value, (3, self.mesh.ny, self.mesh.nx)) + + # Converting to cylindrical coordinates + velocities_cylindrical = np.zeros(value.shape) + velocities_cylindrical[1] = -value[2] + velocities_cylindrical[[0, 2]] = self.mesh.to_cartesian(value[:2]) + + self._electron_velocities = value + self._electron_velocities_cylindrical = velocities_cylindrical + self._electron_velocities_cylindrical_f2d = SOLPSVectorFunction2D.instance(self._sample_vector_f2d, velocities_cylindrical) + self._electron_velocities_cartesian = VectorAxisymmetricMapper(self._electron_velocities_cylindrical_f2d) + + @property + def electron_velocities_cylindrical(self): + """ + Electron velocities in cylindrical coordinates (R, phi, Z) at each mesh cell. + Array of shape (3, ny, nx): [0, :, :] - R, [1, :, :] - phi, [2, :, :] - Z. + :return: + """ + return self._electron_velocities_cylindrical + + @electron_velocities_cylindrical.setter + def electron_velocities_cylindrical(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("electron_velocities_cylindrical", value, (3, self.mesh.ny, self.mesh.nx)) + + # Converting to poloidal coordinates + velocities = np.zeros(value.shape) + velocities[:, 2] = -value[:, 1] + velocities[:2] = self.mesh.to_poloidal(value[[0, 2]]) + + self._electron_velocities_cylindrical = value + self._electron_velocities = velocities + self._electron_velocities_cylindrical_f2d = SOLPSVectorFunction2D.instance(self._sample_vector_f2d, value) + self._electron_velocities_cartesian = VectorAxisymmetricMapper(self._electron_velocities_cylindrical_f2d) + + @property + def electron_velocities_cylindrical_f2d(self): + """ + VectorFunction2D interpolator for electron velocities in cylindrical coordinates. + Returns a vector of electron velocity at a given point (R, Z). + :return: + """ + return self._electron_velocities_cylindrical_f2d @property - def radial_particle_flux(self): + def electron_velocities_cartesian(self): """ - Blah + VectorFunction3D interpolator for electron velocities in Cartesian coordinates. + Returns a vector of electron velocity at a given point (x, y, z). :return: """ - return self._rad_par_flux + return self._electron_velocities_cartesian @property - def radial_area(self): + def species_density(self): """ - Blah + Simulated species densities at each mesh cell. + Array of shape (ns, ny, nx). :return: """ - return self._radial_area + return self._species_density @property - def b2_neutral_densities(self): + def species_density_f2d(self): """ - Neutral atom densities from B2 + Dictionary of Function2D interpolators for species densities. + Accessed by species index or species_list elements. + E.g., species_density_f2d[1] or species_density_f2d[('deuterium', 1))]. + Each entry returns species density at a given point (R, Z). :return: """ - return self._b2_neutral_densities + return self._species_density_f2d @property - def velocities_parallel(self): - return self._velocities_parallel + def species_density_f3d(self): + """ + Dictionary of Function3D interpolators for species densities. + Accessed by species index or species_list elements. + E.g., species_density_f3d[1] or species_density_f3d[('deuterium', 1))]. + Each entry returns species density at a given point (x, y, z). + :return: + """ + return self._species_density_f3d + + @species_density.setter + def species_density(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("species_density", value, (len(self._species_list), self.mesh.ny, self.mesh.nx)) + self._species_density = value + self._species_density_f2d = {} + self._species_density_f3d = {} + for k, sp in enumerate(self._species_list): + self._species_density_f2d[k] = SOLPSFunction2D.instance(self._inside_mesh, value[k]) + self._species_density_f2d[sp] = self._species_density_f2d[k] + self._species_density_f3d[k] = AxisymmetricMapper(self._species_density_f2d[k]) + self._species_density_f3d[sp] = self._species_density_f3d[k] @property - def velocities_radial(self): - return self._velocities_radial + def velocities(self): + """ + Species velocities in poloidal coordinates (e_pol, e_rad, e_tor) at each mesh cell. + Array of shape (ns, 3, ny, nx): + [:, 0, :, :] - poloidal, [:, 1, :, :] - radial, [:, 2, :, :] - toroidal. + :return: + """ + return self._velocities + + @velocities.setter + def velocities(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("velocities", value, (len(self.species_list), 3, self.mesh.ny, self.mesh.nx)) + + # Converting to cylindrical coordinates + velocities_cylindrical = np.zeros(value.shape) + velocities_cylindrical[:, 1] = -value[:, 2] + for k in range(value.shape[0]): + velocities_cylindrical[k, [0, 2]] = self.mesh.to_cartesian(value[k, :2]) + + self._velocities = value + self._velocities_cylindrical = velocities_cylindrical + self._velocities_cylindrical_f2d = {} + self._velocities_cartesian = {} + for k, sp in enumerate(self._species_list): + self._velocities_cylindrical_f2d[k] = SOLPSVectorFunction2D.instance(self._sample_vector_f2d, velocities_cylindrical[k]) + self._velocities_cylindrical_f2d[sp] = self._velocities_cylindrical_f2d[k] + self._velocities_cartesian[k] = VectorAxisymmetricMapper(self._velocities_cylindrical_f2d[k]) + self._velocities_cartesian[sp] = self._velocities_cartesian[k] + + @property + def velocities_cylindrical(self): + """ + Species velocities in cylindrical coordinates (R, phi, Z) at each mesh cell. + Array of shape (ns, 3, ny, nx): [:, 0, :, :] - R, [:, 1, :, :] - phi, [:, 2, :, :] - Z. + :return: + """ + return self._velocities_cylindrical + + @velocities_cylindrical.setter + def velocities_cylindrical(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("velocities_cylindrical", value, (len(self.species_list), 3, self.mesh.ny, self.mesh.nx)) + + # Converting to poloidal coordinates + velocities = np.zeros(value.shape) + velocities[:, 2] = -value[:, 1] + for k in range(value.shape[0]): + velocities[k, :2] = self.mesh.to_poloidal(value[k, [0, 2]]) + + self._velocities_cylindrical = value + self._velocities = velocities + self._velocities_cylindrical_f2d = {} + self._velocities_cartesian = {} + for k, sp in enumerate(self._species_list): + self._velocities_cylindrical_f2d[k] = SOLPSVectorFunction2D.instance(self._sample_vector_f2d, value[k]) + self._velocities_cylindrical_f2d[sp] = self._velocities_cylindrical_f2d[k] + self._velocities_cartesian[k] = VectorAxisymmetricMapper(self._velocities_cylindrical_f2d[k]) + self._velocities_cartesian[sp] = self._velocities_cartesian[k] @property - def velocities_toroidal(self): - return self._velocities_toroidal + def velocities_cylindrical_f2d(self): + """ + Dictionary of VectorFunction2D interpolators for species velocities + in cylindrical coordinates. + Accessed by species index or species_list elements. + E.g., velocities_cylindrical_f2d[1] or velocities_cylindrical_f2d[('deuterium', 1))]. + Each entry returns a vector of species velocity at a given point (R, Z). + :return: + """ + return self._velocities_cylindrical_f2d @property def velocities_cartesian(self): + """ + Dictionary of VectorFunction3D interpolators for species velocities + in Cartesian coordinates. + Accessed by species index or species_list elements. + E.g., velocities_cartesian[1] or velocities_cartesian[('deuterium', 1))]. + Each entry returns a vector of species velocity at a given point (x, y, z). + :return: + """ return self._velocities_cartesian + @property + def inside_mesh(self): + """ + Function2D for testing if point p is inside the simulation mesh. + """ + return self._inside_mesh + @property def inside_volume_mesh(self): """ Function3D for testing if point p is inside the simulation mesh. """ - if self._inside_mesh is None: - raise RuntimeError("Inside mesh test not available for this simulation") - else: - return self._inside_mesh + return AxisymmetricMapper(self._inside_mesh) @property def total_radiation(self): """ - Total radiation + Total radiation at each mesh cell. + Array of shape (ny, nx). This is not calculated from the CHERAB emission models, instead it comes from the SOLPS output data. - Is calculated from the sum of all integrated line emission and all Bremmstrahlung. The signals used are 'RQRAD' - and 'RQBRM'. Final output is in W/str? + Is calculated from the sum of all integrated line emission and all Bremmstrahlung. + The signals used are 'RQRAD' and 'RQBRM'. Final output is in W m-3. """ - if self._total_rad is None: - raise RuntimeError("Total radiation not available for this simulation") - else: - return self._total_rad - # TODO: decide is this a 2D or 3D interface? + return self._total_radiation + @property - def total_radiation_volume(self): + def total_radiation_f2d(self): + """ + Function2D interpolator for total radiation. + Returns total radiation at a given point (R, Z). """ - Total radiation volume. - This is not calculated from the CHERAB emission models, instead it comes from the SOLPS output data. - Is calculated from the sum of all integrated line emission and all Bremmstrahlung. The signals used are 'RQRAD' - and 'RQBRM'. Final output is in W/str? + return self._total_radiation_f2d - :returns: Function3D + @property + def total_radiation_f3d(self): + """ + Function3D interpolator for total radiation. + Returns total radiation at a given point (x, y, z). """ - mapped_radiation_data = _map_data_onto_triangles(self._total_rad) - radiation_mesh_2d = Discrete2DMesh(self.mesh.vertex_coords, self.mesh.triangles, mapped_radiation_data, limit=False) - # return AxisymmetricMapper(radiation_mesh_2d) - return radiation_mesh_2d + return self._total_radiation_f3d + + @total_radiation.setter + def total_radiation(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("total_radiation", value, (self.mesh.ny, self.mesh.nx)) + self._total_radiation = value + self._total_radiation_f2d = SOLPSFunction2D.instance(self._inside_mesh, value) + self._total_radiation_f3d = AxisymmetricMapper(self._total_radiation_f2d) @property - def parallel_velocities(self): + def b_field(self): """ - Plasma velocity field at each mesh cell. Equivalent to 'UA' in SOLPS. + Magnetic B field in poloidal coordinates (e_pol, e_rad, e_tor) at each mesh cell. + Array of shape (3, ny, nx): [0, :, :] - poloidal, [1, :, :] - radial, [2, :, :] - toroidal. """ - if self._parallel_velocities is None: - raise RuntimeError("Parallel velocities not available for this simulation") - else: - return self._parallel_velocities + + return self._b_field + + @b_field.setter + def b_field(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("b_field", value, (3, self.mesh.ny, self.mesh.nx)) + + # Converting to cylindrical system + b_field_cylindrical = np.zeros(value.shape) + b_field_cylindrical[1] = -value[2] + b_field_cylindrical[[0, 2]] = self.mesh.to_cartesian(value[:2]) + + self._b_field_cylindrical = b_field_cylindrical + self._b_field = value + self._b_field_cylindrical_f2d = SOLPSVectorFunction2D.instance(self._sample_vector_f2d, b_field_cylindrical) + self._b_field_cartesian = VectorAxisymmetricMapper(self._b_field_cylindrical_f2d) @property - def radial_velocities(self): + def b_field_cylindrical(self): """ - Calculated radial velocity components for each species. + Magnetic B field in poloidal coordinates (R, phi, Z) at each mesh cell. + Array of shape (3, ny, nx): [0, :, :] - R, [1, :, :] - phi, [2, :, :] - Z. """ - if self._parallel_velocities is None: - raise RuntimeError("Radial velocities not available for this simulation") - else: - return self._radial_velocities + + return self._b_field_cylindrical @property - def b_field(self): + def b_field_cylindrical_f2d(self): """ - Magnetic B field at each mesh cell in mesh cell coordinates (b_parallel, b_radial b_toroidal). + VectorFunction2D interpolator for magnetic B field in cylindrical coordinates. + Returns a vector of magnetic field at a given point (R, Z). """ - if self._b_field_vectors is None: - raise RuntimeError("Magnetic field not available for this simulation") - else: - return self._b_field_vectors + + return self._b_field_cylindrical_f2d @property def b_field_cartesian(self): """ - Magnetic B field at each mesh cell in cartesian coordinates (Bx, By, Bz). + VectorFunction3D interpolator for magnetic B field in Cartesian coordinates. + Returns a vector of magnetic field at a given point (x, y, z). """ - if self._b_field_vectors_cartesian is None: - raise RuntimeError("Magnetic field not available for this simulation") - else: - return self._b_field_vectors_cartesian + + return self._b_field_cartesian + + @b_field_cylindrical.setter + def b_field_cylindrical(self, value): + value = np.array(value, dtype=np.float64, copy=False) + _check_shape("b_field_cylindrical", value, (3, self.mesh.ny, self.mesh.nx)) + + # Converting to poloidal system + b_field = np.zeros(value.shape) + b_field[2] = -value[1] + b_field[:2] = self.mesh.to_poloidal(value[[0, 2]]) + + self._b_field = b_field + self._b_field_cylindrical = value + self._b_field_cylindrical_f2d = SOLPSVectorFunction2D.instance(self._sample_vector_f2d, value) + self._b_field_cartesian = VectorAxisymmetricMapper(self._b_field_cylindrical_f2d) @property def eirene_simulation(self): @@ -247,31 +589,29 @@ def eirene_simulation(self): :rtype: Eirene """ - if self._eirene is None: - raise RuntimeError("EIRENE simulation data not available for this SOLPS simulation") - else: - return self._eirene + + return self._eirene + + @eirene_simulation.setter + def eirene_simulation(self, value): + if not isinstance(value, Eirene): + raise ValueError('Attribute "eirene_simulation" must be an Eirene instance.') + + self._eirene = value def __getstate__(self): state = { - 'mesh': self.mesh.__getstate__(), + 'mesh': self._mesh.__getstate__(), + 'species_list': self._species_list, 'electron_temperature': self._electron_temperature, + 'ion_temperature': self._ion_temperature, + 'neutral_temperature': self._neutral_temperature, 'electron_density': self._electron_density, - 'species_list': self._species_list, 'species_density': self._species_density, - 'rad_par_flux': self._rad_par_flux, - 'radial_area': self._radial_area, - 'b2_neutral_densities': self._b2_neutral_densities, - 'velocities_parallel': self._velocities_parallel, - 'velocities_radial': self._velocities_radial, - 'velocities_toroidal': self._velocities_toroidal, - 'velocities_cartesian': self._velocities_cartesian, - 'inside_mesh': self._inside_mesh, - 'total_rad': self._total_rad, - 'b_field_vectors': self._b_field_vectors, - 'b_field_vectors_cartesian': self._b_field_vectors_cartesian, - 'parallel_velocities': self._parallel_velocities, - 'radial_velocities': self._radial_velocities, + 'electron_velocities_cylindrical': self._electron_velocities_cylindrical, + 'velocities_cylindrical': self._velocities_cylindrical, + 'total_radiation': self._total_radiation, + 'b_field_cylindrical': self._b_field_cylindrical, 'eirene_model': self._eirene_model, 'b2_model': self._b2_model, 'eirene': self._eirene @@ -279,24 +619,24 @@ def __getstate__(self): return state def __setstate__(self, state): - self.mesh = SOLPSMesh(**state['mesh']) - self._electron_temperature = state['electron_temperature'] - self._electron_density = state['electron_density'] - self._species_list = state['species_list'] - self._species_density = state['species_density'] - self._rad_par_flux = state['rad_par_flux'] - self._radial_area = state['radial_area'] - self._b2_neutral_densities = state['b2_neutral_densities'] - self._velocities_parallel = state['velocities_parallel'] - self._velocities_radial = state['velocities_radial'] - self._velocities_toroidal = state['velocities_toroidal'] - self._velocities_cartesian = state['velocities_cartesian'] - self._inside_mesh = state['inside_mesh'] - self._total_rad = state['total_rad'] - self._b_field_vectors = state['b_field_vectors'] - self._b_field_vectors_cartesian = state['b_field_vectors_cartesian'] - self._parallel_velocities = state['parallel_velocities'] - self._radial_velocities = state['radial_velocities'] + if state['electron_temperature'] is not None: + self.electron_temperature = state['electron_temperature'] # will create _f2d() and _f3d() + if state['ion_temperature'] is not None: + self.ion_temperature = state['ion_temperature'] + if state['neutral_temperature'] is not None: + self.neutral_temperature = state['neutral_temperature'] + if state['electron_density'] is not None: + self.electron_density = state['electron_density'] + if state['species_density'] is not None: + self.species_density = state['species_density'] + if state['electron_velocities_cylindrical'] is not None: + self.velocities_cylindrical = state['electron_velocities_cylindrical'] + if state['velocities_cylindrical'] is not None: + self.velocities_cylindrical = state['velocities_cylindrical'] + if state['total_radiation'] is not None: + self.total_radiation = state['total_radiation'] + if state['b_field_cylindrical'] is not None: + self.b_field_cylindrical = state['b_field_cylindrical'] self._eirene_model = state['eirene_model'] self._b2_model = state['b2_model'] self._eirene = state['eirene'] @@ -307,140 +647,6 @@ def save(self, filename): pickle.dump(self.__getstate__(), file_handle) file_handle.close() - # def plot_electrons(self): - # """ Make a plot of the electron temperature and density in the SOLPS mesh plane. """ - # - # me = self.mesh.mesh_extent - # xl, xu = (me['minr'], me['maxr']) - # yl, yu = (me['minz'], me['maxz']) - # - # te_samples = np.zeros((500, 500)) - # ne_samples = np.zeros((500, 500)) - # - # xrange = np.linspace(xl, xu, 500) - # yrange = np.linspace(yl, yu, 500) - # - # plasma = self.plasma - # for i, x in enumerate(xrange): - # for j, y in enumerate(yrange): - # ne_samples[j, i] = plasma.electron_distribution.density(x, 0.0, y) - # te_samples[j, i] = plasma.electron_distribution.effective_temperature(x, 0.0, y) - # - # plt.figure() - # plt.imshow(ne_samples, extent=[xl, xu, yl, yu], origin='lower') - # plt.colorbar() - # plt.xlim(xl, xu) - # plt.ylim(yl, yu) - # plt.title("electron density") - # plt.figure() - # plt.imshow(te_samples, extent=[xl, xu, yl, yu], origin='lower') - # plt.colorbar() - # plt.xlim(xl, xu) - # plt.ylim(yl, yu) - # plt.title("electron temperature") - # - # def plot_species_density(self, species, ionisation): - # """ - # Make a plot of the requested species density in the SOLPS mesh plane. - # - # :param Element species: The species to plot. - # :param int ionisation: The charge state of the species to plot. - # """ - # - # species_dist = self.plasma.get_species(species, ionisation) - # - # me = self.mesh.mesh_extent - # xl, xu = (me['minr'], me['maxr']) - # yl, yu = (me['minz'], me['maxz']) - # species_samples = np.zeros((500, 500)) - # - # xrange = np.linspace(xl, xu, 500) - # yrange = np.linspace(yl, yu, 500) - # - # for i, x in enumerate(xrange): - # for j, y in enumerate(yrange): - # species_samples[j, i] = species_dist.distribution.density(x, 0.0, y) - # - # plt.figure() - # plt.imshow(species_samples, extent=[xl, xu, yl, yu], origin='lower') - # plt.colorbar() - # plt.xlim(xl, xu) - # plt.ylim(yl, yu) - # plt.title("Species {} - stage {} - density".format(species.name, ionisation)) - # - # def plot_pec_emission_lines(self, emission_lines, title="", vmin=None, vmax=None, log=False): - # """ - # Make a plot of the given PEC emission lines - # - # :param list emission_lines: List of PEC emission lines. - # :param str title: The title of the plot. - # :param float vmin: The minimum value for clipping the plots (default=None). - # :param float vmax: The maximum value for clipping the plots (default=None). - # :param bool log: Toggle a log plot for the data (default=False). - # """ - # - # me = self.mesh.mesh_extent - # xl, xu = (me['minr'], me['maxr']) - # yl, yu = (me['minz'], me['maxz']) - # emission_samples = np.zeros((500, 500)) - # - # xrange = np.linspace(xl, xu, 500) - # yrange = np.linspace(yl, yu, 500) - # - # for i, x in enumerate(xrange): - # for j, y in enumerate(yrange): - # for emitter in emission_lines: - # emission_samples[j, i] += emitter.emission(Point3D(x, 0.0, y), Vector3D(0, 0, 0), Spectrum(350, 700, 800)).total() - # - # if log: - # emission_samples = np.log(emission_samples) - # plt.figure() - # plt.imshow(emission_samples, extent=[xl, xu, yl, yu], origin='lower', vmin=vmin, vmax=vmax) - # plt.colorbar() - # plt.xlim(xl, xu) - # plt.ylim(yl, yu) - # plt.title(title) - # - # def plot_radiated_power(self): - # """ - # Make a plot of the given PEC emission lines - # - # :param list emission_lines: List of PEC emission lines. - # :param str title: The title of the plot. - # """ - # - # mesh = self.mesh - # me = mesh.mesh_extent - # total_rad = self.total_radiation - # - # xl, xu = (me['minr'], me['maxr']) - # yl, yu = (me['minz'], me['maxz']) - # - # # tri_index_lookup = mesh.triangle_index_lookup - # - # emission_samples = np.zeros((500, 500)) - # - # xrange = np.linspace(xl, xu, 500) - # yrange = np.linspace(yl, yu, 500) - # - # for i, x in enumerate(xrange): - # for j, y in enumerate(yrange): - # - # try: - # # k, l = mesh.triangle_to_grid_map[int(tri_index_lookup(x, y))] - # # emission_samples[i, j] = total_rad[k, l] - # emission_samples[j, i] = total_rad(x, 0, y) - # - # except ValueError: - # continue - # - # plt.figure() - # plt.imshow(emission_samples, extent=[xl, xu, yl, yu], origin='lower') - # plt.colorbar() - # plt.xlim(xl, xu) - # plt.ylim(yl, yu) - # plt.title("Radiated Power (W/m^3)") - def create_plasma(self, parent=None, transform=None, name=None): """ Make a CHERAB plasma object from this SOLPS simulation. @@ -452,6 +658,16 @@ def create_plasma(self, parent=None, transform=None, name=None): :rtype: Plasma """ + # Checking if the minimal required data is available to create a plasma object + if self.electron_density_f3d is None: + raise RuntimeError("Unable to create plasma object: electron density is not set.") + if self.electron_temperature_f3d is None: + raise RuntimeError("Unable to create plasma object: electron temperature is not set.") + if self.species_density_f3d is None: + raise RuntimeError("Unable to create plasma object: species density is not set.") + if self.ion_temperature_f3d is None: + raise RuntimeError("Unable to create plasma object: ion temperature is not set.") + mesh = self.mesh name = name or "SOLPS Plasma" plasma = Plasma(parent=parent, transform=transform, name=name) @@ -460,77 +676,216 @@ def create_plasma(self, parent=None, transform=None, name=None): plasma.geometry = Cylinder(radius, height) plasma.geometry_transform = translate(0, 0, mesh.mesh_extent['minz']) - tri_index_lookup = self.mesh.triangle_index_lookup - tri_to_grid = self.mesh.triangle_to_grid_map - - if isinstance(self._b_field_vectors, np.ndarray): - plasma.b_field = SOLPSVectorFunction3D(tri_index_lookup, tri_to_grid, self._b_field_vectors_cartesian) - else: + if self.b_field_cartesian is None: print('Warning! No magnetic field data available for this simulation.') + else: + plasma.b_field = self.b_field_cartesian # Create electron species - triangle_data = _map_data_onto_triangles(self._electron_temperature) - electron_te_interp = Discrete2DMesh(mesh.vertex_coords, mesh.triangles, triangle_data, limit=False) - electron_temp = AxisymmetricMapper(electron_te_interp) - triangle_data = _map_data_onto_triangles(self._electron_density) - electron_ne_interp = Discrete2DMesh.instance(electron_te_interp, triangle_data) - electron_dens = AxisymmetricMapper(electron_ne_interp) - electron_velocity = lambda x, y, z: Vector3D(0, 0, 0) - plasma.electron_distribution = Maxwellian(electron_dens, electron_temp, electron_velocity, electron_mass) - - if not isinstance(self.velocities_cartesian, np.ndarray): - print('Warning! No velocity field data available for this simulation.') - - b2_neutral_i = 0 # counter for B2 neutrals - for k, sp in enumerate(self.species_list): + if self.electron_velocities_cartesian is None: + print('Warning! No electron velocity data available for this simulation.') + electron_velocity = ConstantVector3D(Vector3D(0, 0, 0)) + else: + electron_velocity = self.electron_velocities_cartesian + plasma.electron_distribution = Maxwellian(self.electron_density_f3d, self.electron_temperature_f3d, electron_velocity, electron_mass) - # Identify the species based on its symbol - symbol, charge = re.match(_SPECIES_REGEX, sp).groups() - charge = int(charge) - species_type = _species_symbol_map[symbol] + if self.velocities_cartesian is None: + print('Warning! No species velocities data available for this simulation.') - # If neutral and B" atomic density available, use B2 density, otherwise use fluid species density. - if isinstance(self.b2_neutral_densities, np.ndarray) and charge == 0: - species_dens_data = self.b2_neutral_densities[:, :, b2_neutral_i] - b2_neutral_i += 1 - else: - species_dens_data = self.species_density[:, :, k] + if self.neutral_temperature_f3d is None: + print('Warning! No neutral atom temperature data available for this simulation.') + + neutral_i = 0 # neutrals count + for k, sp in enumerate(self.species_list): + + try: + species_type = lookup_element(sp[0]) + except ValueError: + species_type = lookup_isotope(sp[0]) - triangle_data = _map_data_onto_triangles(species_dens_data) - dens = AxisymmetricMapper(Discrete2DMesh.instance(electron_te_interp, triangle_data)) - # dens = SOLPSFunction3D(tri_index_lookup, tri_to_grid, species_dens_data) + charge = sp[1] # Create the velocity vector lookup function - if isinstance(self.velocities_cartesian, np.ndarray): - velocity = SOLPSVectorFunction3D(tri_index_lookup, tri_to_grid, self.velocities_cartesian[:, :, k, :]) + if self.velocities_cartesian is not None: + velocity = self.velocities_cartesian[k] else: - velocity = lambda x, y, z: Vector3D(0, 0, 0) + velocity = ConstantVector3D(Vector3D(0, 0, 0)) + + if charge or self.neutral_temperature is None: # ions or neutral atoms (neutral temperature is not available) + distribution = Maxwellian(self.species_density_f3d[k], self.ion_temperature_f3d, velocity, + species_type.atomic_weight * atomic_mass) + + else: # neutral atoms with neutral temperature + distribution = Maxwellian(self.species_density_f3d[k], self._neutral_temperature_f3d[neutral_i], velocity, + species_type.atomic_weight * atomic_mass) + neutral_i += 1 - distribution = Maxwellian(dens, electron_temp, velocity, species_type.atomic_weight * atomic_mass) plasma.composition.add(Species(species_type, charge, distribution)) return plasma -def _map_data_onto_triangles(solps_dataset): +def _check_shape(name, value, shape): + if value.shape != shape: + raise ValueError('Shape of "{0}": {1} mismatch the shape of SOLPS grid: {2}.'.format(name, value.shape, shape)) + + +def prefer_element(isotope): """ - Reshape a SOLPS data array so that it matches the triangles in the SOLPS mesh. + Return Element instance, if the element of this isotope has the same mass number. + """ + el_mass_number = int(round(isotope.element.atomic_weight)) + if el_mass_number == isotope.mass_number: + return isotope.element + + return isotope - :param ndarray solps_dataset: Given SOLPS dataset, typically of shape (98 x 32). - :return: New 1D ndarray with shape (98*32*2) + +def b2_flux_to_velocity(sim, poloidal_flux, radial_flux, parallel_velocity): + """ + Calculates velocities of all species at cell centres using B2 particle fluxes defined + at cell faces. + + :param SOLPSSimulation sim: SOLPSSimulation instance. + :param ndarray poloidal_flux: Poloidal flux of species in s-1. Must be a 3 dimensional array of + shape (num_species, mesh.ny, mesh.nx). + :param ndarray radial_flux: Radial flux of species in s-1. Must be a 3 dimensional array of + shape (num_species, mesh.ny, mesh.nx). + :param ndarray parallel_velocity: Parallel velocity of species in m/s. Must be a 3 dimensional + array of shape (num_species, mesh.ny, mesh.nx). + Parallel velocity is a velocity projection on magnetic + field direction. + + :return ndarray: Velocities of all species in cylindrical coordinates. + Array of shape (num_species, 3, mesh.ny, mesh.nx). + """ + if sim.b_field_cylindrical is None: + raise RuntimeError('Attribute "b_field_cylindrical" is not set.') + if sim.species_density is None: + raise RuntimeError('Attribute "species_density" is not set.') + + mesh = sim.mesh + b = sim.b_field_cylindrical + + poloidal_flux = np.array(poloidal_flux, dtype=np.float64, copy=False) + radial_flux = np.array(radial_flux, dtype=np.float64, copy=False) + parallel_velocity = np.array(parallel_velocity, dtype=np.float64, copy=False) + + nx = mesh.nx # poloidal + ny = mesh.ny # radial + ns = len(sim.species_list) # number of species + + _check_shape('poloidal_flux', poloidal_flux, (ns, ny, nx)) + _check_shape('radial_flux', radial_flux, (ns, ny, nx)) + _check_shape('parallel_velocity', parallel_velocity, (ns, ny, nx)) + + poloidal_area = mesh.poloidal_area + radial_area = mesh.radial_area + leftix = mesh.neighbix[0] # poloidal prev. + leftiy = mesh.neighbiy[0] + bottomix = mesh.neighbix[1] # radial prev. + bottomiy = mesh.neighbiy[1] + rightix = mesh.neighbix[2] # poloidal next. + rightiy = mesh.neighbiy[2] + topix = mesh.neighbix[3] # radial next. + topiy = mesh.neighbiy[3] + + # Converting s-1 --> m-2 s-1 + poloidal_flux = np.divide(poloidal_flux, poloidal_area, out=np.zeros_like(poloidal_flux), where=poloidal_area > 0) + radial_flux = np.divide(radial_flux, radial_area, out=np.zeros_like(radial_flux), where=radial_area > 0) + + # Obtaining left velocity + dens_neighb = sim.species_density[:, leftiy, leftix] # density in the left neighbouring cell + has_neighbour = ((leftix > -1) * (leftiy > -1)) # check if has left neighbour + neg_flux = (poloidal_flux < 0) * (sim.species_density > 0) # will use density in this cell if flux is negative + pos_flux = (poloidal_flux > 0) * (dens_neighb > 0) * has_neighbour # will use density in neighbouring cell if flux is positive + velocity_left = np.divide(poloidal_flux, sim.species_density, out=np.zeros((ns, ny, nx)), where=neg_flux) + velocity_left = np.divide(poloidal_flux, dens_neighb, out=velocity_left, where=pos_flux) + poloidal_face_normal = mesh.radial_basis_vector[[1, 0]] + poloidal_face_normal[1] *= -1 + velocity_left = velocity_left[:, None] * poloidal_face_normal # to vector in RZ + + # Obtaining bottom velocity + dens_neighb = sim.species_density[:, bottomiy, bottomix] + has_neighbour = ((bottomix > -1) * (bottomiy > -1)) + neg_flux = (radial_flux < 0) * (sim.species_density > 0) + pos_flux = (poloidal_flux > 0) * (dens_neighb > 0) * has_neighbour + velocity_bottom = np.divide(radial_flux, sim.species_density, out=np.zeros((ns, ny, nx)), where=neg_flux) + velocity_bottom = np.divide(radial_flux, dens_neighb, out=velocity_bottom, where=pos_flux) + radial_face_normal = mesh.poloidal_basis_vector[[1, 0]] + radial_face_normal[0] *= -1 + velocity_bottom = velocity_bottom[:, None] * radial_face_normal # to RZ + + # Obtaining right and top velocities + velocity_right = velocity_left[:, :, rightiy, rightix] + velocity_right[:, :, (rightix < 0) + (rightiy < 0)] = 0 + + velocity_top = velocity_bottom[:, :, topiy, topix] + velocity_top[:, :, (topix < 0) + (topiy < 0)] = 0 + + vcyl = np.zeros((ns, 3, ny, nx)) # velocities in cylindrical coordinates + + # Projection of velocity on RZ-plane + vcyl[:, [0, 2]] = 0.25 * (velocity_bottom + velocity_left + velocity_top + velocity_right) + + # Obtaining toroidal velocity + bmagn = np.sqrt((b * b).sum(0)) + vcyl[:, 1] = (parallel_velocity * bmagn - vcyl[:, 0] * b[0] - vcyl[:, 2] * b[2]) / b[1] + + return vcyl + + +def eirene_flux_to_velocity(sim, poloidal_flux, radial_flux, parallel_velocity): """ + Calculates velocities of neutral atoms using Eirene particle fluxes defined at cell centre. + + :param SOLPSSimulation sim: SOLPSSimulation instance. + :param ndarray poloidal_flux: Poloidal flux of atoms in m-2 s-1. Must be a 3 dimensional array of + shape (num_atoms, mesh.ny, mesh.nx). + :param ndarray radial_flux: Radial flux of atoms in m-2 s-1. Must be a 3 dimensional array of + shape (num_atoms, mesh.ny, mesh.nx). + :param ndarray parallel_velocity: Parallel velocity of atoms in m/s. Must be a 3 dimensional + array of shape (num_atoms, mesh.ny, mesh.nx). + Parallel velocity is a velocity projection on magnetic + field direction. + + :return ndarray: Velocities of neutral atoms in cylindrical coordinates. + Array of shape (num_atoms, 3, mesh.ny, mesh.nx). + """ + if sim.b_field_cylindrical is None: + raise RuntimeError('Attribute "b_field_cylindrical" is not set.') + if sim.species_density is None: + raise RuntimeError('Attribute "species_density" is not set.') + + mesh = sim.mesh + b = sim.b_field_cylindrical + + poloidal_flux = np.array(poloidal_flux, dtype=np.float64, copy=False) + radial_flux = np.array(radial_flux, dtype=np.float64, copy=False) + parallel_velocity = np.array(parallel_velocity, dtype=np.float64, copy=False) + + nx = mesh.nx # poloidal + ny = mesh.ny # radial + ns = len(sim.neutral_list) # number of neutral atoms + + _check_shape('poloidal_flux', poloidal_flux, (ns, ny, nx)) + _check_shape('radial_flux', radial_flux, (ns, ny, nx)) + _check_shape('parallel_velocity', parallel_velocity, (ns, ny, nx)) + + neutral_indx = [k for k, sp in enumerate(sim.species_list) if sp[1] == 0] + density = sim.species_density[neutral_indx, :] + + # Obtaining velocity + poloidal_velocity = np.divide(poloidal_flux, density, out=np.zeros_like(density), where=(density > 0)) + radial_velocity = np.divide(radial_flux, density, out=np.zeros_like(density), where=(density > 0)) - solps_mesh_shape = solps_dataset.shape - triangle_data = np.zeros(solps_mesh_shape[0] * solps_mesh_shape[1] * 2, dtype=np.float64) + vcyl = np.zeros((ns, 3, ny, nx)) # velocities in cylindrical coordinates - tri_index = 0 - for i in range(solps_mesh_shape[0]): - for j in range(solps_mesh_shape[1]): + # Projection of velocity on RZ-plane + vcyl[:, [0, 2]] = (poloidal_velocity[:, None] * mesh.poloidal_basis_vector + radial_velocity[:, None] * mesh.radial_basis_vector) - # Same data - triangle_data[tri_index] = solps_dataset[i, j] - tri_index += 1 - triangle_data[tri_index] = solps_dataset[i, j] - tri_index += 1 + # Obtaining toroidal velocity + bmagn = np.sqrt((b * b).sum(0)) + vcyl[:, 1] = (parallel_velocity * bmagn - vcyl[:, 0] * b[0] - vcyl[:, 2] * b[2]) / b[1] - return triangle_data + return vcyl