Merge pull request #87 from dschwoerer/read_mode2

v0.2.6

xemc3/core/load.py

@@ -329,11 +329,27 @@ def read_mappings(fn: str, dims: typing.Sequence[int]) -> xr.DataArray:
     return da
 
 
-def add_metadata(ds: xr.Dataset):
+def ensure_metadata(ds: xr.Dataset) -> xr.Dataset:
+    if not ds:
+        return ds
+    meta = get_default_metadata()
+    for k in meta:
+        if k not in ds.attrs:
+            ds.attrs.update(meta)
+            return ds
+    return ds
+
+
+def add_metadata(ds: xr.Dataset) -> xr.Dataset:
+    ds.attrs.update(get_default_metadata())
+    return ds
+
+
+def get_default_metadata() -> dict:
     # Delay import to avoid circular dependency
     from .. import __version__
 
-    ds.attrs.update(
+    return dict(
         title="EMC3-EIRENE Simulation data",
         software_name="xemc3",
         software_version=__version__,
@@ -401,8 +417,7 @@ def read_locations(path: str, ds: typing.Optional[xr.Dataset] = None) -> xr.Data
     ds.emc3.unit("z_bounds", "m")
     ds.emc3.unit("phi_bounds", "radian")
     assert isinstance(ds, xr.Dataset)
-    add_metadata(ds)
-    return ds
+    return add_metadata(ds)
 
 
 def scrape(f: typing.TextIO, *, ignore="!", verbose=False) -> str:
@@ -634,7 +649,9 @@ def read_plates_raw(cwd: str, fn: str) -> typing.Sequence[xr.Dataset]:
                 f"Unexpected string `{s}` while reading {plate}." + raise_issue
             )
             _, geom = s
-            r, z, phi = read_plate(cwd + geom)
+            if not geom.startswith("/"):
+                geom = cwd + geom
+            r, z, phi = read_plate(geom)
             nx, ny = r.shape
             nx -= 1
             ny -= 1
@@ -725,10 +742,12 @@ def read_plates_raw(cwd: str, fn: str) -> typing.Sequence[xr.Dataset]:
                         newpos[ix::xref, iy::yref, :] = pos @ A
 
                     corrs[i] = newpos
+                dims = [plate_prefix + x for x in ("phi", "x")]
+                dims += ["delta_" + x for x in dims]
                 corrs_da = [
                     xr.DataArray(
                         data=a.reshape(nxr, nyr, 2, 2),
-                        dims=("phi", "x", "delta_phi", "delta_x"),
+                        dims=dims,
                     )
                     for a in corrs
                 ]
@@ -1435,7 +1454,10 @@ def read_fort_file_pub(
         if not isinstance(ds, xr.Dataset):
             ds = xr.Dataset()
         return read_fort_file(ds, fn, **defaults)
-    ds = ensure_mapping("/".join(fn.split("/")[:-1]), ds, type == "mapped", fn=fn)
+    if type == "target_flux":
+        ds = ds or xr.Dataset()
+    else:
+        ds = ensure_mapping("/".join(fn.split("/")[:-1]), ds, type == "mapped", fn=fn)
     assert isinstance(ds, xr.Dataset)
     return read_fort_file(ds, fn, **defaults)
 
@@ -1497,7 +1519,7 @@ def read_fort_file(ds: xr.Dataset, fn: str, type: str = "mapped", **opts) -> xr.
         [f"{k}={v}" for k, v in opts.items()]
     )
     if datas is None:
-        return ds
+        return ensure_metadata(ds)
     vars = vars.copy()
     assert opts == {}
     keys = [k for k in vars]
@@ -1526,7 +1548,7 @@ def read_fort_file(ds: xr.Dataset, fn: str, type: str = "mapped", **opts) -> xr.
         assert (
             varopts == {}
         ), f"variable {var} has options {varopts} but didn't expect anything"
-    return ds
+    return ensure_metadata(ds)
 
 
 def guess_type(ds: xr.Dataset, key: typing.Hashable) -> str: