Merge pull request #66 from boutproject/fix-squashoutput-global-attributes

Copy file attributes in squashoutput()

Author: johnomotani

boutdata/data.py

@@ -1102,6 +1102,12 @@ def __init__(
                 for i in range(self.grid_info["npes"])
             ]
 
+        if self._DataFileCaching or self._parallel:
+            # Keep reference to 0'th file, for reading attributes
+            self._file0 = DataFile(self._file_list[0])
+        else:
+            self._file0 = None
+
         if self._info:
             print(
                 "mxsub = {} mysub = {} mz = {}\n".format(
@@ -1145,6 +1151,8 @@ def __del__(self):
                 connection.send(None)
                 worker.join()
                 connection.close()
+        if self._file0 is not None:
+            self._file0.close()
 
     def _init_caching(self):
         """
@@ -1224,6 +1232,64 @@ def evolvingVariables(self):
         """Return a list of names of time-evolving variables"""
         return self.grid_info["evolvingVariableNames"]
 
+    def get_attribute(self, variable, attrname):
+        """Get an attribute of a variable
+
+        Parameters
+        ----------
+        variable : str
+            Name of variable to get attribute from
+        attrname : str
+            Name of attribute
+
+        Returns
+        -------
+        Value of attribute
+        """
+        if self._file0 is None:
+            with DataFile(self._file_list[0]) as f:
+                return f.attributes(variable)[attrname]
+        else:
+            return self._file0.attributes(variable)[attrname]
+
+    def get_file_attribute(self, attrname):
+        """Get an attribute of the output files.
+
+        Attribute is taken from the rank-0 file. No checking is done that the attribute
+        is consistent between all the output files.
+
+        Parameters
+        ----------
+        attrname : str
+            Name of attribute
+
+        Returns
+        -------
+        Value of attribute
+        """
+        if self._file0 is None:
+            with DataFile(self._file_list[0]) as f:
+                return f.read_file_attribute(attrname)
+        else:
+            return self._file0.read_file_attribute(attrname)
+
+    def list_file_attributes(self):
+        """List all file attributes of output files
+
+        List is taken from the rank-0 file. No checking is done that the file attributes
+        are consistent between all the output files.
+
+        Returns
+        -------
+        List of str
+            Names of the file attributes
+        """
+        if self._file0 is None:
+            with DataFile(self._file_list[0]) as f:
+                return f.list_file_attributes()
+        else:
+            return self._file0.list_file_attributes()
+
     def redistribute(self, npes, nxpe=None, mxg=2, myg=2, include_restarts=True):
         """Create a new set of dump files for npes processors.
 

boutdata/squashoutput.py

@@ -226,6 +226,12 @@ def squashoutput(
         var = None
         gc.collect()
 
+    # Copy file attributes
+    for attrname in outputs.list_file_attributes():
+        attrval = outputs.get_file_attribute(attrname)
+        for f in files:
+            f.write_file_attribute(attrname, attrval)
+
     for f in files:
         f.close()
 

boutdata/tests/make_test_data.py

@@ -75,6 +75,12 @@
     },
 }
 
+expected_file_attributes = {
+    "global_str_attribute": "foobar",
+    "global_int_attribute": 42,
+    "global_float_attribute": 7.0,
+}
+
 
 def make_grid_info(
     *, mxg=2, myg=2, nxpe=1, nype=1, ixseps1=None, ixseps2=None, xpoints=0
@@ -321,6 +327,9 @@ def createScalar(name, value):
         createScalar("PE_YIND", i // nxpe)
         createScalar("MYPE", i)
 
+        for attrname, attr in expected_file_attributes.items():
+            setattr(outputfile, attrname, attr)
+
     return result
 
 

boutdata/tests/test_collect.py

@@ -6,12 +6,14 @@
 
 from boutdata.collect import collect
 from boutdata.squashoutput import squashoutput
+from boututils.datafile import DataFile
 
 from boutdata.tests.make_test_data import (
     apply_slices,
     create_dump_file,
     concatenate_data,
     expected_attributes,
+    expected_file_attributes,
     make_grid_info,
     remove_xboundaries,
     remove_yboundaries,
@@ -57,7 +59,7 @@ def check_collected_data(
         Arrays should be global (not per-process).
     fieldperp_global_yind : int
         Global y-index where FieldPerps are expected to be defined.
-    path : pathlib.Path or str
+    path : pathlib.Path
         Path to collect data from.
     squash : bool
         If True, call `squashoutput()` and delete the `BOUT.dmp.*.nc` files (so that we
@@ -104,6 +106,14 @@ def check_collected_data(
             fieldperp_global_yind,
         )
 
+    if squash:
+        filename = path.joinpath("boutdata.nc")
+    else:
+        filename = path.joinpath("BOUT.dmp.0.nc")
+    with DataFile(str(filename)) as f:
+        for attrname, attr in expected_file_attributes.items():
+            assert f.read_file_attribute(attrname) == attr
+
 
 def check_variable(
     varname, actual, expected_data, expected_attributes, fieldperp_global_yind

setup.py

@@ -45,7 +45,7 @@
         "numpy",
         "matplotlib",
         "scipy",
-        "boututils",
+        "boututils>=0.1.9",
         "importlib-metadata ; python_version<'3.8'",
     ],
     classifiers=[