diff --git a/src/boutdata/collect.py b/src/boutdata/collect.py index f6c9624..deb6cb3 100644 --- a/src/boutdata/collect.py +++ b/src/boutdata/collect.py @@ -953,15 +953,16 @@ def _check_fieldperp_attributes( # and check they are unique if yindex_global is not None and yindex_global != temp_yindex: raise ValueError( - "Found FieldPerp {} at different global y-indices, {} " "and {}".format( + "Found FieldPerp {} at different global y-indices, {} and {}".format( varname, temp_yindex, yindex_global ) ) yindex_global = temp_yindex if fieldperp_yproc is not None and fieldperp_yproc != pe_yind: raise ValueError( - "Found FieldPerp {} on different y-processor indices, " - "{} and {}".format(varname, fieldperp_yproc, pe_yind) + "Found FieldPerp {} on different y-processor indices, {} and {}".format( + varname, fieldperp_yproc, pe_yind + ) ) fieldperp_yproc = pe_yind var_attributes = temp_f_attributes diff --git a/src/boututils/ask.py b/src/boututils/ask.py index 559e71e..fb167ec 100644 --- a/src/boututils/ask.py +++ b/src/boututils/ask.py @@ -53,4 +53,4 @@ def query_yes_no(question, default="yes"): elif choice in valid: return valid[choice] else: - sys.stdout.write("Please respond with 'yes' or 'no' " "(or 'y' or 'n').\n") + sys.stdout.write("Please respond with 'yes' or 'no' (or 'y' or 'n').\n") diff --git a/src/boututils/datafile.py b/src/boututils/datafile.py index b14d960..becf65b 100644 --- a/src/boututils/datafile.py +++ b/src/boututils/datafile.py @@ -281,7 +281,7 @@ def bout_type(self, varname): """ return self.attributes(varname)["bout_type"] - def write(self, name, data, info=False): + def write(self, name, data, info=False, *, dims=None): """Write a variable to file If the variable is not a :py:obj:`~boututils.boutarray.BoutArray` with @@ -297,13 +297,15 @@ def write(self, name, data, info=False): info : bool, optional If True, print information about what is being written to file + dims : tuple(str) or None, optional + If passed, specifies the dimensions to be used to write the variable. Returns ------- None """ - return self.impl.write(name, data, info) + return self.impl.write(name, data, info, dims=dims) def read_file_attribute(self, name): return self.impl.read_file_attribute(name) @@ -525,12 +527,18 @@ def _bout_dimensions_from_var(self, data): return BoutArray.dims_from_type(bout_type) - def write(self, name, data, info=False): + def write(self, name, data, info=False, *, dims=None): if not self.writeable: raise Exception("File not writeable. Open with write=True keyword") s = np.shape(data) + if dims is not None and len(dims) != data.ndim: + raise ValueError( + f"When dims is passed, it must have an entry for each dimension of " + f"`data`, but dims={dims} and data.ndim={data.ndim}." + ) + # Get the variable type t = type(data).__name__ @@ -568,7 +576,10 @@ def write(self, name, data, info=False): # Not found, so add. # Get dimensions - defdims = self._bout_dimensions_from_var(data) + if dims is None: + defdims = self._bout_dimensions_from_var(data) + else: + defdims = dims def find_dim(dim): # Find a dimension with given name and size @@ -626,10 +637,10 @@ def find_dim(dim): # List of (size, 'name') tuples dlist = list(zip(s, defdims)) # Get new list of variables, and turn into a tuple - dims = tuple(map(find_dim, dlist)) + dims_tuple = tuple(map(find_dim, dlist)) # Create the variable - var = self.handle.createVariable(name, t, dims, **self._kwargs) + var = self.handle.createVariable(name, t, dims_tuple, **self._kwargs) if var is None: raise Exception("Couldn't create variable") @@ -899,7 +910,7 @@ def size(self, varname): return None return var.shape - def write(self, name, data, info=False): + def write(self, name, data, info=False, *, dims=None): if not self.writeable: raise Exception("File not writeable. Open with write=True keyword")