Using Data Dictionary metadata¶

import imas.training

# 1. Load the training data for the core_profiles IDS:
entry = imas.training.get_training_db_entry()
core_profiles = entry.get("core_profiles")

# 1a. Print the coordinate of profiles_1d[0].electrons.temperature
print(core_profiles.profiles_1d[0].electrons.temperature.coordinates[0])
# Do you recognize the coordinate? Yes, as shown in the first line of the output, this
# is "profiles_1d[0]/grid/rho_tor_norm".

# 1b. Print the coordinate of profiles_1d:
print(core_profiles.profiles_1d.coordinates[0])
# What do you notice? This prints the core_profiles.time array:
#   <IDSNumericArray (IDS:core_profiles, time, FLT_1D)>
#   numpy.ndarray([  3.98722186, 432.93759781, 792.        ])

# 1c. Change the time mode and print again
core_profiles.ids_properties.homogeneous_time = \
    imas.ids_defs.IDS_TIME_MODE_HETEROGENEOUS
print(core_profiles.profiles_1d.coordinates[0])
# What has changed? Now we get a numpy array with values -9e+40:
#   [-9.e+40 -9.e+40 -9.e+40]
#
# In heterogeneous time, the coordinate of profiles_1d is profiles_1d/time, which is a
# scalar. IMAS-Python will construct a numpy array for you where
#   array[i] := profiles_1d[i]/time
# Since we didn't set these values, they are set to the default EMPTY_FLOAT, which is
# -9e+40.

# 2. Load the training data for the equilibrium IDS:
equilibrium = entry.get("equilibrium")

# 2a. What is the coordinate of time_slice/profiles_2d?
slice0 = equilibrium.time_slice[0]
print(slice0.profiles_2d.metadata.coordinates)
# This will output:
#   (IDSCoordinate('1...N'),)
# The coordinate of profiles_2d is an index. When requesting the coordinate values,
# IMAS-Python will generate an index array for you:
print(slice0.profiles_2d.coordinates[0])
# -> array([0])

# 2b. What are the coordinates of ``time_slice/profiles_2d/b_field_r``?
print(slice0.profiles_2d[0].b_field_r.metadata.coordinates)
# This is a 2D array and therefore there are two coordinates:
#   (IDSCoordinate('time_slice(itime)/profiles_2d(i1)/grid/dim1'),
#    IDSCoordinate('time_slice(itime)/profiles_2d(i1)/grid/dim2'))

import imas

# 1. Create an empty distributions IDS
distributions = imas.IDSFactory().distributions()

# 2. Use the metadata attribute to find the coordinates of
#    distribution/profiles_2d/density
print(distributions.metadata["distribution/profiles_2d/density"].coordinates)
# Alternative, by resizing the Arrays of Structures:
distributions.distribution.resize(1)
distributions.distribution[0].profiles_2d.resize(1)
p2d = distributions.distribution[0].profiles_2d[0]
print(p2d.density.metadata.coordinates)
# This outputs (newlines added for clarity):
#  (IDSCoordinate('distribution(i1)/profiles_2d(itime)/grid/r
#                  OR distribution(i1)/profiles_2d(itime)/grid/rho_tor_norm'),
#   IDSCoordinate('distribution(i1)/profiles_2d(itime)/grid/z
#                  OR distribution(i1)/profiles_2d(itime)/grid/theta_geometric
#                  OR distribution(i1)/profiles_2d(itime)/grid/theta_straight'))
#
# What do you notice: in both dimensions there are multiple options for the coordinate.

# 3. Retrieve the coordinate values through the ``coordinates`` attribute.
# This will raise a coordinate lookup error because IMAS-Python cannot choose which of the
# coordinates to use:
try:
    print(p2d.density.coordinates[0])
except Exception as exc:
    print(exc)

# 4a. Use the IDSCoordinate.references attribute:
# Example for the first dimension:
coordinate_options = p2d.density.metadata.coordinates[0].references
# 4b. Use IDSPath.goto:
for option in coordinate_options:
    coordinate_node = option.goto(p2d.density)
    print(coordinate_node)
# This will print:
#   <IDSNumericArray (IDS:distributions, distribution[0]/profiles_2d[0]/grid/r, empty FLT_1D)>
#   <IDSNumericArray (IDS:distributions, distribution[0]/profiles_2d[0]/grid/rho_tor_norm, empty FLT_1D)>

import itertools  # python standard library iteration tools

import imas
import imas.training
import pint

# 1. Load core_profiles IDS from training DBEntry
entry = imas.training.get_training_db_entry()
cp = entry.get("core_profiles")

# 2. Select the first time slice of profiles_1d
p1d = cp.profiles_1d[0]

# 3.
# Create pint UnitRegistry
ureg = pint.UnitRegistry()

# Convert DD units to Pint Units
_dd_to_pint = {
    "-": ureg("dimensionless"),
    "Atomic Mass Unit": ureg("unified_atomic_mass_unit"),
    "Elementary Charge Unit": ureg("elementary_charge"),
}
def dd_to_pint(dd_unit):
    if dd_unit in _dd_to_pint:
        return _dd_to_pint[dd_unit]
    return ureg(dd_unit)
# End of translation

# 4. Calculate mass density:
# 4a. Create mass_density variable with units:
mass_density = ureg("0 kg.m^-3")
# 4b. Loop over all ion and neutral species
for species in itertools.chain(p1d.ion, p1d.neutral):
    mass = sum(
        element.a * dd_to_pint(element.a.metadata.units)
        for element in species.element
    )
    density = species.density * dd_to_pint(species.density.metadata.units)
    mass_density += mass * density

# 4c. Print the total mass density
print(mass_density)
# Note that the species mass is given in Atomic Mass Units, but pint
# automatically converted this to kilograms for us, because we defined
# mass_density in kg/m^3!