Unit handling and round-trip .nc conversion fixes (#13)

* Fix roundtrip conversion of named_options

* Fix unit-conversion error ("" as ureg.dimensionless, not skip conversion)

* Pre-commit fixes

* Remove unnecessary string catch in fileio

* Add test to check conversion to dimensionless

* Add netcdf4 to environment

* Add back in fix for str in netcdf

* Mypy fixes

* Add default value for netcdf engine

* Update cfspopcon/file_io.py

Co-authored-by: Christoph Hasse <[email protected]>

* fix pre-commit error

---------

Co-authored-by: Christoph Hasse <[email protected]>

cfspopcon/atomic_data/read_radas_data.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import xarray as xr
-from scipy.interpolate import RectBivariateSpline, RegularGridInterpolator  # type: ignore[import]
+from scipy.interpolate import RectBivariateSpline, RegularGridInterpolator  # type: ignore[import-untyped]
 
 from ..named_options import Impurity
 from ..unit_handling import convert_units, magnitude, ureg

cfspopcon/cli.py

@@ -5,9 +5,9 @@
 from pathlib import Path
 
 import click
-import matplotlib.pyplot as plt  # type:ignore[import]
+import matplotlib.pyplot as plt
 import xarray as xr
-from ipdb import launch_ipdb_on_exception  # type:ignore[import]
+from ipdb import launch_ipdb_on_exception  # type:ignore[import-untyped]
 
 from cfspopcon import file_io
 from cfspopcon.input_file_handling import read_case

cfspopcon/file_io.py

@@ -1,8 +1,9 @@
 """Functions for saving results to file and loading those files."""
 import json
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
+import numpy as np
 import xarray as xr
 
 from .helpers import convert_named_options
@@ -26,7 +27,9 @@ def sanitize_variable(val: xr.DataArray, key: str) -> xr.DataArray:
     return val
 
 
-def write_dataset_to_netcdf(dataset: xr.Dataset, filepath: Path) -> None:
+def write_dataset_to_netcdf(
+    dataset: xr.Dataset, filepath: Path, netcdf_writer: Literal["netcdf4", "scipy", "h5netcdf"] = "netcdf4"
+) -> None:
     """Write a dataset to a NetCDF file."""
     serialized_dataset = dataset.copy()
     for key in dataset.keys():
@@ -37,7 +40,7 @@ def write_dataset_to_netcdf(dataset: xr.Dataset, filepath: Path) -> None:
         assert isinstance(key, str)  # because hashable type of key is broader str but we  know it's str
         serialized_dataset[key] = sanitize_variable(dataset[key], key)
 
-    serialized_dataset.to_netcdf(filepath)
+    serialized_dataset.to_netcdf(filepath, engine=netcdf_writer)
 
 
 def promote_variable(val: xr.DataArray, key: str) -> Any:
@@ -47,7 +50,8 @@ def promote_variable(val: xr.DataArray, key: str) -> Any:
     except KeyError:
         pass
 
-    if val.dtype == object:
+    # scipy i/o yields objects for strings while netcdf4 yields np.str_
+    if val.dtype == object or val.dtype.type == np.str_:
         if val.size == 1:
             return convert_named_options(key, val.item())
         else:

cfspopcon/formulas/divertor_metrics.py

@@ -1,6 +1,6 @@
 """Divertor loading and functions to calculate OMP pitch (for q_parallel calculation)."""
 import numpy as np
-from scipy import constants  # type: ignore[import]
+from scipy import constants  # type: ignore[import-untyped]
 
 from ..unit_handling import ureg, wraps_ufunc
 

cfspopcon/formulas/fusion_reaction_data/reaction_energies.py

@@ -4,7 +4,7 @@
 
 from numpy import float64
 from numpy.typing import NDArray
-from scipy import constants  # type: ignore[import]
+from scipy import constants  # type: ignore[import-untyped]
 
 
 def reaction_energy_DT(

cfspopcon/formulas/plasma_profile_data/density_and_temperature_profile_fits.py

@@ -54,7 +54,7 @@
 import pandas as pd
 import yaml
 from numpy.typing import NDArray
-from scipy.interpolate import RectBivariateSpline  # type: ignore[import]
+from scipy.interpolate import RectBivariateSpline  # type: ignore[import-untyped]
 
 plasma_profiles_directory = Path(__file__).parent
 

cfspopcon/transform.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import xarray as xr
-from scipy.interpolate import griddata  # type:ignore[import]
+from scipy.interpolate import griddata  # type:ignore[import-untyped]
 
 from cfspopcon.unit_handling import Unit, magnitude
 

cfspopcon/unit_handling/default_units.py

@@ -252,7 +252,7 @@ def convert_to_default_units(value: Quantity, key: str) -> Quantity:
 def convert_to_default_units(value: Union[float, Quantity, xr.DataArray], key: str) -> Union[float, Quantity, xr.DataArray]:
     """Convert an array or scalar to default units."""
     unit = DEFAULT_UNITS[key]
-    if unit is None or unit == "":
+    if unit is None:
         return value
     elif isinstance(value, (xr.DataArray, Quantity)):
         return convert_units(value, unit)

cfspopcon/unit_handling/setup_unit_handling.py

@@ -7,7 +7,7 @@
 import numpy as np
 import numpy.typing as npt
 import pint
-import pint_xarray  # type:ignore[import]
+import pint_xarray  # type:ignore[import-untyped]
 import xarray as xr
 from pint.errors import DimensionalityError
 from typing_extensions import ParamSpec

pyproject.toml

@@ -38,6 +38,7 @@ xarray = "^2023.4.1"
 pint-xarray = "^0.3"
 ipdb = "^0.13.13"
 click = "^8.1.0"
+netcdf4 = "^1.6.5"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^2.20.0"

tests/test_unit_handling/test_conversions.py

@@ -0,0 +1,15 @@
+import warnings
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from cfspopcon.unit_handling import dimensionless_magnitude, ureg, Quantity, ureg, convert_to_default_units
+
+
+def test_conversion_of_dimensionless():
+
+    val = Quantity(2.0, ureg.percent)
+
+    assert np.isclose(dimensionless_magnitude(val), 2e-2)
+    assert np.isclose(convert_to_default_units(val, key="beta"), 2e-2)