Added fluxes to velocity conversion to adress cherab#36.

cherab/solps/formats/balance.py

@@ -41,24 +41,15 @@ def load_solps_from_balance(balance_filename):
     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)
+    mesh = load_mesh_from_netcdf(fhandle)
 
     # Load each plasma species in simulation
 
     species_list = []
     n_species = len(fhandle.variables['am'].data)
     for i in range(n_species):
 
-        # Extract the nuclear charge    
+        # Extract the nuclear charge
         if fhandle.variables['species'].data[i, 1] == b'D':
             zn = 1
         if fhandle.variables['species'].data[i, 1] == b'C':
@@ -75,7 +66,7 @@ def load_solps_from_balance(balance_filename):
         isotope = lookup_isotope(zn, number=am)
         species = prefer_element(isotope)  # Prefer Element over Isotope if the mass number is the same
 
-        # If we only need to populate species_list, there is probably a faster way to do this...        
+        # If we only need to populate species_list, there is probably a faster way to do this...
         species_list.append((species, charge))
 
     sim = SOLPSSimulation(mesh, species_list)
@@ -114,28 +105,66 @@ def load_solps_from_balance(balance_filename):
     # 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
+        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.keys():
-            eirene_ecoolrate = np.sum(fhandle.variables['eirene_mc_eael_she_bal'].data, axis=0) / vol
+            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.keys():
-            eirene_potential_loss = rydberg_energy * np.sum(fhandle.variables['eirene_mc_papl_sna_bal'].data, axis=(0))[1, :, :] * el_charge / vol
+            eirene_potential_loss = rydberg_energy * np.sum(fhandle.variables['eirene_mc_papl_sna_bal'].data, axis=(0))[1, :, :] * el_charge / mesh.vol
 
         # This will be negative (energy sink); multiply by -1
         sim.total_radiation = -1.0 * (b2_ploss + (eirene_ecoolrate - eirene_potential_loss))
 
     else:
         # Total radiated power from B2, not including neutrals
-        b2_ploss = np.sum(fhandle.variables['b2stel_she_bal'].data, axis=0) / vol
+        b2_ploss = np.sum(fhandle.variables['b2stel_she_bal'].data, axis=0) / mesh.vol
 
-        potential_loss = np.sum(fhandle.variables['b2stel_sna_ion_bal'].data, axis=0) / vol
+        potential_loss = np.sum(fhandle.variables['b2stel_sna_ion_bal'].data, axis=0) / mesh.vol
 
         # Save total radiated power to the simulation object
         sim.total_radiation = rydberg_energy * el_charge * potential_loss - b2_ploss
 
     fhandle.close()
 
     return sim
+
+
+def load_mesh_from_netcdf(fhandle):
+
+    # Load SOLPS mesh geometry
+    r = fhandle.variables['crx'].data.copy()
+    z = fhandle.variables['cry'].data.copy()
+    vol = fhandle.variables['vol'].data.copy()
+
+    # Re-arrange the array dimensions in the way CHERAB expects...
+    r = np.moveaxis(r, 0, -1)
+    z = np.moveaxis(z, 0, -1)
+
+    # Loading neighbouring cell indices
+    neighbix = np.zeros(r.shape, dtype=np.int)
+    neighbiy = np.zeros(r.shape, dtype=np.int)
+
+    neighbix[:, :, 0] = fhandle.variables['leftix'].data.copy().astype(np.int)  # poloidal prev.
+    neighbix[:, :, 1] = fhandle.variables['bottomix'].data.copy().astype(np.int)  # radial prev.
+    neighbix[:, :, 2] = fhandle.variables['rightix'].data.copy().astype(np.int)  # poloidal next
+    neighbix[:, :, 3] = fhandle.variables['topix'].data.copy().astype(np.int)  # radial next
+
+    neighbiy[:, :, 0] = fhandle.variables['leftiy'].data.copy().astype(np.int)
+    neighbiy[:, :, 1] = fhandle.variables['bottomiy'].data.copy().astype(np.int)
+    neighbiy[:, :, 2] = fhandle.variables['rightiy'].data.copy().astype(np.int)
+    neighbiy[:, :, 3] = fhandle.variables['topiy'].data.copy().astype(np.int)
+
+    # In SOLPS cell indexing starts with -1 (guarding cell), but in SOLPSMesh -1 means no neighbour.
+    if neighbix.min() < -1 or neighbiy.min() < -1:
+        neighbix += 1
+        neighbiy += 1
+    neighbix[neighbix == r.shape[0]] = -1
+    neighbiy[neighbiy == r.shape[1]] = -1
+
+    # Create the SOLPS mesh
+    mesh = SOLPSMesh(r, z, vol, neighbix, neighbiy)
+
+    return mesh

cherab/solps/formats/mdsplus.py

@@ -22,7 +22,7 @@
 
 from cherab.core.atomic.elements import lookup_isotope
 from cherab.solps.mesh_geometry import SOLPSMesh
-from cherab.solps.solps_plasma import SOLPSSimulation, prefer_element
+from cherab.solps.solps_plasma import SOLPSSimulation, prefer_element, b2_flux_to_velocity, eirene_flux_to_velocity
 
 from matplotlib import pyplot as plt
 
@@ -44,19 +44,22 @@ def load_solps_from_mdsplus(mds_server, ref_number):
     conn.openTree('solps', ref_number)
 
     # Load SOLPS mesh geometry and lookup arrays
-    mesh = load_mesh_from_mdsplus(conn)
+    mesh = load_mesh_from_mdsplus(conn, mdsExceptions)
 
-    # Load each plasma species
+    # Load each plasma species in simulation
     ns = conn.get('\SOLPS::TOP.IDENT.NS').data()  # Number of species
     zn = conn.get('\SOLPS::TOP.SNAPSHOT.GRID.ZN').data().astype(np.int)  # Nuclear charge
     am = np.round(conn.get('\SOLPS::TOP.SNAPSHOT.GRID.AM').data()).astype(np.int)  # Atomic mass number
     charge = conn.get('\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, charge[i]))
+        if charge[i] == 0:
+            neutral_indx.append(i)
 
     sim = SOLPSSimulation(mesh, species_list)
     ni = mesh.nx
@@ -75,109 +78,137 @@ def load_solps_from_mdsplus(mds_server, ref_number):
     sim.ion_temperature = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.TI').data(), 0, 1)
 
     # Load species density
-    sim.species_density = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.NA').data(), 0, 2)
+    species_density = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.NA').data(), 0, 2)
 
-    # Load the neutral atom density from Eirene if available
-    try:
-        dab2 = conn.get('\SOLPS::TOP.SNAPSHOT.DAB2').data()
-        if isinstance(dab2, np.ndarray):
-            # Replace the species densities
-            neutral_densities = np.swapaxes(dab2, 0, 2)
-
-            neutral_i = 0  # counter for neutral atoms
-            for k, sp in enumerate(sim.species_list):
-                charge = sp[1]
-                if charge == 0:
-                    sim.species_density[:, :, k] = neutral_densities[:, :, neutral_i]
-                    neutral_i += 1
-    except mdsExceptions.TreeNNF:
-        pass
-
-    # Load the neutral atom temperature from Eirene if available
-    try:
-        tab2 = conn.get('\SOLPS::TOP.SNAPSHOT.TAB2').data()
-        if isinstance(tab2, np.ndarray):
-            sim.neutral_temperature = np.swapaxes(tab2, 0, 2)
-    except mdsExceptions.TreeNNF:
-        pass
+    # Load parallel velocity
+    parallel_velocity = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.UA').data(), 0, 2)
 
-    # TODO: Eirene data (TOP.SNAPSHOT.PFLA, TOP.SNAPSHOT.RFLA) should be used for neutral atoms.
-    velocities = np.zeros((ni, nj, len(sim.species_list), 3))
-    velocities[:, :, :, 0] = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.UA').data(), 0, 2)
+    # Load poloidal and radial particle fluxes for velocity calculation
+    poloidal_flux = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.FNAX').data(), 0, 2)
+    radial_flux = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.FNAY').data(), 0, 2)
 
-    ################################################
-    # Calculate the species' velocity distribution #
+    # B2 fluxes are defined between cells, so correcting array shapes if needed
+    if poloidal_flux.shape[0] == ni - 1:
+        poloidal_flux = np.vstack((np.zeros((1, nj, ns)), poloidal_flux))
 
-    # calculate field component ratios for velocity conversion
-    bplane2 = sim.b_field[:, :, 0]**2 + sim.b_field[:, :, 2]**2
-    parallel_to_toroidal_ratio = sim.b_field[:, :, 0] * sim.b_field[:, :, 2] / bplane2
+    if radial_flux.shape[1] == nj - 1:
+        radial_flux = np.hstack((np.zeros((ni, 1, ns)), radial_flux))
 
-    # Calculate toroidal velocity components
-    velocities[:, :, :, 2] = velocities[:, :, :, 0] * parallel_to_toroidal_ratio[:, :, None]
+    # Obtaining velocity from B2 flux
+    velocities_cartesian = b2_flux_to_velocity(mesh, species_density, poloidal_flux, radial_flux, parallel_velocity, sim.b_field_cartesian)
 
-    # Radial velocity is obtained from radial particle flux
-    radial_particle_flux = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.FNAY').data(), 0, 2)  # radial particle flux
-    radial_area = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.SY').data(), 0, 1)  # radial contact area
-    for k, sp in enumerate(sim.species_list):
-        i, j = np.where(sim.species_density[:, :-1, k] > 0)  # radial_area array corresponds to [:, 1:] in mesh, so maybe [:, 1:, k]
-        velocities[i, j, k, 1] = radial_particle_flux[i, j, k] / radial_area[i, j] / sim.species_density[i, j, k]
+    # Obtaining additional data from EIRENE and replacing data for neutrals
 
-    sim.velocities = velocities
-    # sim.velocities_cartesian is created authomatically
+    b2_standalone = False
+    try:
+        # Replace the species densities
+        neutral_density = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.DAB2').data(), 0, 2)
+        species_density[:, :, neutral_indx] = neutral_density
+    except (mdsExceptions.TreeNNF, np.AxisError):
+        print("Warning! This is B2 stand-alone simulation.")
+        b2_standalone = True
+
+    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 = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.PFLA').data(), 0, 2)
+            neutral_radial_flux = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.RFLA').data(), 0, 2)
+
+            if np.any(neutral_poloidal_flux) or np.any(neutral_radial_flux):
+                neutral_velocities_cartesian = eirene_flux_to_velocity(mesh, neutral_density, neutral_poloidal_flux, neutral_radial_flux,
+                                                                       parallel_velocity[:, :, neutral_indx], sim.b_field_cartesian)
+
+                velocities_cartesian[:, :, neutral_indx, :] = neutral_velocities_cartesian
+        except (mdsExceptions.TreeNNF, np.AxisError):
+            pass
+
+        # Obtaining neutral temperatures
+        try:
+            sim.neutral_temperature = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.TAB2').data(), 0, 2)
+        except (mdsExceptions.TreeNNF, np.AxisError):
+            pass
+
+    sim.species_density = species_density
+    sim.velocities_cartesian = velocities_cartesian  # this also updates sim.velocities
 
     ###############################
     # Load extra data from server #
     ###############################
 
     ####################
     # Integrated power #
-    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)
+    try:
+        linerad = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.RQRAD').data(), 0, 2)
+        linerad = np.sum(linerad, axis=2)
+    except (mdsExceptions.TreeNNF, np.AxisError):
+        linerad = 0
 
-    # need to cope with fact that neurad may not be present!!!
     try:
-        neurad = conn.get('\SOLPS::TOP.SNAPSHOT.ENEUTRAD').data()
-        if isinstance(neurad, np.ndarray):
-            if len(neurad.shape) == 3:
-                neurad = np.swapaxes(np.abs(np.sum(neurad, axis=0)), 0, 1)
-            else:
-                neurad = np.swapaxes(np.abs(neurad), 0, 1)
+        brmrad = np.swapaxes(conn.get('\SOLPS::TOP.SNAPSHOT.RQBRM').data(), 0, 2)
+        brmrad = np.sum(brmrad, axis=2)
+    except (mdsExceptions.TreeNNF, np.AxisError):
+        brmrad = 0
+
+    try:
+        eneutrad = conn.get('\SOLPS::TOP.SNAPSHOT.ENEUTRAD').data()
+        if np.ndim(eneutrad) == 3:  # this will not return error if eneutrad is not np.ndarray
+            neurad = np.swapaxes(np.abs(np.sum(eneutrad, axis=0)), 0, 1)
         else:
-            neurad = 0
-    except mdsExceptions.TreeNNF:
+            neurad = np.swapaxes(np.abs(eneutrad), 0, 1)
+    except (mdsExceptions.TreeNNF, np.AxisError):
         neurad = 0
 
     sim.total_radiation = (linerad + brmrad + neurad) / mesh.vol
 
     return sim
 
 
-def load_mesh_from_mdsplus(mds_connection):
+def load_mesh_from_mdsplus(mds_connection, mdsExceptions):
     """
     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)
+    r = 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)
 
     vol = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.VOL').data(), 0, 1)
 
+    # Loading neighbouring cell indices
+    neighbix = np.zeros(r.shape, dtype=np.int)
+    neighbiy = np.zeros(r.shape, dtype=np.int)
+
+    neighbix[:, :, 0] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:LEFTIX').data().astype(np.int), 0, 1)
+    neighbix[:, :, 1] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:BOTTOMIX').data().astype(np.int), 0, 1)
+    neighbix[:, :, 2] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:RIGHTIX').data().astype(np.int), 0, 1)
+    neighbix[:, :, 3] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:TOPIX').data().astype(np.int), 0, 1)
+
+    neighbiy[:, :, 0] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:LEFTIY').data().astype(np.int), 0, 1)
+    neighbiy[:, :, 1] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:BOTTOMIY').data().astype(np.int), 0, 1)
+    neighbiy[:, :, 2] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:RIGHTIY').data().astype(np.int), 0, 1)
+    neighbiy[:, :, 3] = np.swapaxes(mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:TOPIY').data().astype(np.int), 0, 1)
+
+    neighbix[neighbix == r.shape[0]] = -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
-    vessel = mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:VESSEL').data()
-    if isinstance(vessel, np.ndarray):
-        mesh.vessel = vessel
+    try:
+        vessel = mds_connection.get('\SOLPS::TOP.SNAPSHOT.GRID:VESSEL').data()
+        if isinstance(vessel, np.ndarray):
+            mesh.vessel = vessel
+    except mdsExceptions.TreeNNF:
+        pass
 
     return mesh