Merge pull request #55 from dschwoerer/devel

Release v0.2.1

xemc3/cli/_common.py

@@ -1,6 +1,8 @@
 import os
 from argparse import ArgumentParser
 
+from .. import __version__
+
 
 def commonparser(desc, ext="nc") -> ArgumentParser:
     parser = ArgumentParser(description=desc)
@@ -17,6 +19,9 @@ def commonparser(desc, ext="nc") -> ArgumentParser:
         help=f"Specify the name for the output file. Defaults to `dir.{ext}` "
         f"when not given. Otherwise `name.{ext}` is used.",
     )
+    parser.add_argument(
+        "-V", "--version", action="version", version=f"xemc3 {__version__}"
+    )
     return parser
 
 

xemc3/cli/xdivertor.py

@@ -127,6 +127,9 @@ def parser() -> ArgumentParser:
         # ],
         help="Data to plot",
     )
+    parser.add_argument(
+        "-V", "--version", action="version", version=f"xemc3 {xemc3.__version__}"
+    )
     parser.add_argument("path", nargs=1, help="Path of data")
     return parser
 

xemc3/core/load.py

@@ -344,7 +344,14 @@ def add_metadata(ds: xr.Dataset):
 def write_mappings(da: xr.DataArray, fn: str) -> None:
     """Write the mappings data to fortran"""
     with open(fn, "w") as f:
-        infos = da.attrs["numcells"], da.attrs["plasmacells"], da.attrs["other"]
+        at_keys = "numcells", "plasmacells", "other"
+        for key in at_keys:
+            assert key in da.attrs, (
+                "Writing requires to be the following attributes to be present:"
+                + ", ".join([f'"{k}"' for k in at_keys])
+                + ". Please ensure that they are copied over from the read file."
+            )
+        infos = tuple([da.attrs[k] for k in at_keys])
         f.write("%12d %11d %11d\n" % infos)
         _block_write(f, da.data.flatten(order="F") + 1, " %11d", 6)
 
@@ -846,6 +853,14 @@ def read_plates(dir: str) -> xr.Dataset:
     return ds
 
 
+def _dir_of(fn: str) -> str:
+    if os.path.isdir(fn):
+        return fn
+    if "/" in fn:
+        return fn.rsplit("/", 1)[0]
+    return "."
+
+
 def get_plates(dir: str, cache: bool = True) -> xr.Dataset:
     """
     Read the target fluxes from the EMC3_post processing routine
@@ -863,14 +878,8 @@ def get_plates(dir: str, cache: bool = True) -> xr.Dataset:
     xr.Dataset
         The target profiles
     """
-    if os.path.isdir(dir):
-        fn = "TARGET_PROFILES"
-    else:
-        if "/" in dir:
-            dir, fn = dir.rsplit("/", 1)
-        else:
-            fn = dir
-            dir = "."
+    dir = _dir_of(dir)
+
     if cache:
         try:
             if os.path.getmtime(dir + "/TARGET_PROFILES.nc") > os.path.getmtime(
@@ -1020,6 +1029,7 @@ def read_mapped(
     """
 
     if isinstance(mapping, xr.Dataset):
+        mapping = ensure_mapping(_dir_of(fn), mapping)
         mapping = mapping["_plasma_map"]
     if kinetic:
         max = np.max(mapping.data) + 1
@@ -1268,6 +1278,10 @@ def write_mapped(
     if skip_first is not None:
         for d in datas:
             if skip_first:
+                assert "print_before" in d.attrs, (
+                    'The "print_before" attribute is needed for writing. '
+                    "Please ensure it is preserved or copied from the read data."
+                )
                 assert d.attrs["print_before"] != ""
             else:
                 if "print_before" in d.attrs:
@@ -1687,14 +1701,15 @@ def read_fort_file(ds: xr.Dataset, fn: str, type: str = "mapped", **opts) -> xr.
     if type == "mapping":
         ds["_plasma_map"] = read_mappings(fn, ds.R_bounds.data.shape[:3])
     elif type == "mapped":
+        ds = ensure_mapping(_dir_of(fn), ds)
         datas = read_mapped(fn, ds["_plasma_map"], **opts, squeeze=False)
         opts = {}
     elif type == "full":
         datas = [read_magnetic_field(fn, ds)]
     elif type == "plates_mag":
         datas = [read_plates_mag(fn, ds)]
     elif type == "geom":
-        ds_ = read_locations(fn.rsplit("/", 1)[0])
+        ds_ = read_locations(_dir_of(fn))
         ds = ds.assign_coords(ds_.coords)
         assert opts == {}, "Unexpected arguments: " + ", ".join(
             [f"{k}={v}" for k, v in opts.items()]

xemc3/core/plot_2d.py

@@ -87,10 +87,10 @@ def plot_target(ds, phi, fmt=None, ax=None, figsize=None, aspect=True):
         if k.endswith("_dims") and k.startswith(f"_{plate_prefix}"):
             das[k] = ds[k]
     for p in das.emc3.iter_plates(symmetry=True, segments=5):
-        assert len(p.phi.shape) == 1
-        if not (p.phi.min() <= phi <= p.phi.max()):
+        assert len(p.plate_phi.shape) == 1
+        if not (p.plate_phi.min() <= phi <= p.plate_phi.max()):
             continue
-        p["plate_phi_plus1"] = p.phi
+        p["plate_phi_plus1"] = p.plate_phi
         p = p.interp(
             plate_phi_plus1=phi, assume_sorted=False, kwargs=dict(bounds_error=True)
         )