Format units with unicode superscript characters

imas_paraview/plugins/profiles_2d.py

@@ -13,7 +13,7 @@
 
 from imas_paraview.ids_util import create_name_recursive
 from imas_paraview.plugins.base_class import GGDVTKPluginBase
-from imas_paraview.util import find_closest_indices
+from imas_paraview.util import find_closest_indices, format_units
 
 logger = logging.getLogger("imas_paraview")
 
@@ -238,8 +238,8 @@ def _create_vtkscalars(self, profile, name):
         """
         values_flat = profile.ravel()
         vtk_scalars = numpy_to_vtk(values_flat.astype(np.float64), deep=True)
-        units = profile.metadata.units
-        vtk_scalars.SetName(f"{name} [{units}]")
+        units = format_units(profile)
+        vtk_scalars.SetName(f"{name} {units}")
         return vtk_scalars
 
     def _create_ugrid(self, vtk_points, vtk_scalars):

imas_paraview/tests/test_util.py

@@ -1,7 +1,8 @@
+import imas
 import numpy as np
 import pytest
 
-from imas_paraview.util import find_closest_indices
+from imas_paraview.util import find_closest_indices, format_units
 
 
 @pytest.mark.parametrize(
@@ -20,3 +21,14 @@ def test_find_closest_indices(request_times, expected_indices):
     time_array = np.array([1.1, 2.2, 3.3, 4.4, 5.5])
     time_indices = find_closest_indices(request_times, time_array)
     assert time_indices == expected_indices
+
+
+def test_format_units():
+    cs = imas.IDSFactory("4.0.0").core_sources()
+    cs.source.resize(1)
+    cs.source[0].profiles_1d.resize(1)
+    p1d = cs.source[0].profiles_1d[0]
+
+    assert format_units(p1d.electrons.particles) == "[m⁻³·s⁻¹]"
+    assert format_units(p1d.electrons.energy) == "[W·m⁻³]"
+    assert format_units(p1d.momentum_phi) == "[kg·m⁻¹·s⁻²]"

imas_paraview/util.py

@@ -8,9 +8,46 @@
 logger = logging.getLogger("imas_paraview")
 
 
+# Translation table to convert numbers and the - sign to their superscript character
+_SUPERSCRIPTS = str.maketrans(
+    {
+        "-": "⁻",
+        "0": "⁰",
+        "1": "¹",
+        "2": "²",
+        "3": "³",
+        "4": "⁴",
+        "5": "⁵",
+        "6": "⁶",
+        "7": "⁷",
+        "8": "⁸",
+        "9": "⁹",
+    }
+)
+
+
 def format_units(node) -> str:
-    """Return the unit of the node surrounded by square brackets."""
-    return f"[{node.metadata.units}]"
+    """Nicely format units
+
+    Use superscript characters for exponentiation, use a center dot instead of a lower
+    dot for composite units, and surround the units with quare brackets.
+
+    Args:
+        node: IDS element
+
+    Example:
+        >>> ids = imas.IDSFactory("4.0.0").core_profiles()
+        >>> ids.profiles_1d.resize(1)
+        >>> ids.profiles_1d[0].ion.resize(1)
+        >>> format_units(ids.profiles_1d[0].ion[0].velocity)
+        '[m·s⁻¹]'
+    """
+    units = node.metadata.units
+    units = "·".join(
+        f"{base}{exponent.translate(_SUPERSCRIPTS)}"
+        for base, _, exponent in (s.partition("^") for s in units.split("."))
+    )
+    return f"[{units}]"
 
 
 def iter_metadata_tree(meta):