support NCF encodings and grid mappings

Author: mjreno

flopy/discretization/grid.py

@@ -1288,14 +1288,15 @@ def write_shapefile(self, filename="grid.shp", crs=None, prjfile=None, **kwargs)
         )
         return
 
-    def dataset(self, modeltime=None, mesh=None):
+    def dataset(self, modeltime=None, mesh=None, encoding=None):
         """
         Method to generate baseline xarray dataset
 
         Parameters
         ----------
         modeltime : FloPy ModelTime object
         mesh : mesh type
+        encoding : variable encoding dictionary
 
         Returns
         -------

flopy/discretization/structuredgrid.py

@@ -1770,11 +1770,12 @@ def get_plottable_layer_array(self, a, layer):
         assert plotarray.shape == required_shape, msg
         return plotarray
 
-    def dataset(self, modeltime=None, mesh=None):
+    def dataset(self, modeltime=None, mesh=None, encoding=None):
         """
         modeltime : FloPy ModelTime object
         mesh : mesh type
                valid mesh types are "layered" or None
+        encoding : variable encoding dictionary
         """
         from ..utils import import_optional_dependency
 
@@ -1787,11 +1788,11 @@ def dataset(self, modeltime=None, mesh=None):
         ds.attrs["modflow_grid"] = "STRUCTURED"
 
         if mesh and mesh.upper() == "LAYERED":
-            return self._layered_mesh_dataset(ds, modeltime)
+            return self._layered_mesh_dataset(ds, modeltime, encoding)
         elif mesh is None:
-            return self._structured_dataset(ds, modeltime)
+            return self._structured_dataset(ds, modeltime, encoding)
 
-    def _layered_mesh_dataset(self, ds, modeltime=None):
+    def _layered_mesh_dataset(self, ds, modeltime=None, encoding=None):
         FILLNA_INT32 = np.int32(-2147483647)
         FILLNA_DBL = 9.96920996838687e36
         lenunits = {0: "m", 1: "ft", 2: "m", 3: "m"}
@@ -1894,9 +1895,25 @@ def _layered_mesh_dataset(self, ds, modeltime=None):
         ds["mesh_face_nodes"].attrs["_FillValue"] = FILLNA_INT32
         ds["mesh_face_nodes"].attrs["start_index"] = np.int32(1)
 
+        if encoding is not None and "wkt" in encoding and encoding["wkt"] is not None:
+            ds = ds.assign({"projection": ([], np.int32(1))})
+            # wkt override to existing crs
+            ds["projection"].attrs["wkt"] = encoding["wkt"]
+            ds["mesh_node_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_node_y"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_y"].attrs["grid_mapping"] = "projection"
+        elif self.crs is not None:
+            ds = ds.assign({"projection": ([], np.int32(1))})
+            ds["projection"].attrs["wkt"] = self.crs.to_wkt()
+            ds["mesh_node_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_node_y"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_y"].attrs["grid_mapping"] = "projection"
+
         return ds
 
-    def _structured_dataset(self, ds, modeltime=None):
+    def _structured_dataset(self, ds, modeltime=None, encoding=None):
         lenunits = {0: "m", 1: "ft", 2: "m", 3: "m"}
 
         x = self.xoffset + self.xycenters[0]
@@ -1953,6 +1970,18 @@ def _structured_dataset(self, ds, modeltime=None):
         ds["x"].attrs["long_name"] = "Easting"
         ds["x"].attrs["bounds"] = "x_bnds"
 
+        if encoding is not None and "wkt" in encoding and encoding["wkt"] is not None:
+            ds = ds.assign({"projection": ([], np.int32(1))})
+            # wkt override to existing crs
+            ds["projection"].attrs["crs_wkt"] = encoding["wkt"]
+            ds["x"].attrs["grid_mapping"] = "projection"
+            ds["y"].attrs["grid_mapping"] = "projection"
+        elif self.crs is not None:
+            ds = ds.assign({"projection": ([], np.int32(1))})
+            ds["projection"].attrs["crs_wkt"] = self.crs.to_wkt()
+            ds["x"].attrs["grid_mapping"] = "projection"
+            ds["y"].attrs["grid_mapping"] = "projection"
+
         return ds
 
     def _set_structured_iverts(self):

flopy/discretization/vertexgrid.py

@@ -600,12 +600,13 @@ def get_plottable_layer_array(self, a, layer):
         assert plotarray.shape == required_shape, msg
         return plotarray
 
-    def dataset(self, modeltime=None, mesh=None):
+    def dataset(self, modeltime=None, mesh=None, encoding=None):
         """
         modeltime : FloPy ModelTime object
         mesh : mesh type
                valid mesh types are "layered" or None
                VertexGrid objects only support layered mesh
+        encoding : variable encoding dictionary
         """
         from ..utils import import_optional_dependency
 
@@ -717,6 +718,22 @@ def dataset(self, modeltime=None, mesh=None):
         ds["mesh_face_nodes"].attrs["_FillValue"] = FILLNA_INT32
         ds["mesh_face_nodes"].attrs["start_index"] = np.int32(1)
 
+        if encoding is not None and "wkt" in encoding and encoding["wkt"] is not None:
+            ds = ds.assign({"projection": ([], np.int32(1))})
+            # wkt override to existing crs
+            ds["projection"].attrs["wkt"] = encoding["wkt"]
+            ds["mesh_node_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_node_y"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_y"].attrs["grid_mapping"] = "projection"
+        elif self.crs is not None:
+            ds = ds.assign({"projection": ([], np.int32(1))})
+            ds["projection"].attrs["wkt"] = self.crs.to_wkt()
+            ds["mesh_node_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_node_y"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_x"].attrs["grid_mapping"] = "projection"
+            ds["mesh_face_y"].attrs["grid_mapping"] = "projection"
+
         return ds
 
     # initialize grid from a grb file

flopy/mf6/mfmodel.py

@@ -1313,7 +1313,6 @@ def write(
         ext_file_action=ExtFileAction.copy_relative_paths,
         netcdf=None,
     ):
-        from ..version import __version__
         """
         Writes out model's package files.
 
@@ -1359,7 +1358,10 @@ def write(
             if write_netcdf:
                 # set data storage to write ascii for netcdf
                 pp._set_netcdf_storage()
-
+            if pp.package_type.startswith("dis"):
+                crs = pp.crs.get_data()
+                if crs is not None and self.modelgrid.crs is None:
+                    self.modelgrid.crs = crs[0][1]
             if (
                 self.simulation_data.verbosity_level.value
                 >= VerbosityLevel.normal.value
@@ -1372,35 +1374,11 @@ def write(
 
         # write netcdf file
         if write_netcdf and netcdf.lower() != "nofile":
-            mesh = netcdf
-            if mesh.upper() == "STRUCTURED":
-                mesh = None
-
-            ds = self.modelgrid.dataset(
-                modeltime=self.modeltime,
-                mesh=mesh,
-            )
-
-            nc_info = self.netcdf_info(mesh=mesh)
-            nc_info["attrs"]["title"] = f"{self.name.upper()} input"
-            nc_info["attrs"]["source"] = f"flopy {__version__}"
-            # :history = "first created 2025/8/21 9:46:2.909" ;
-            # :Conventions = "CF-1.11 UGRID-1.0" ;
-            ds = self.update_dataset(ds, netcdf_info=nc_info, mesh=mesh)
-
-            # write dataset to netcdf
-            fname = self.name_file.nc_filerecord.get_data()[0][0]
-            ds.to_netcdf(
-                os.path.join(self.model_ws, fname),
-                format="NETCDF4",
-                engine="netcdf4"
-            )
-
+            self._write_netcdf(mesh=netcdf)
             if nc_fname is not None:
                 self.name_file.nc_filerecord = None
 
 
-
     def get_grid_type(self):
         """
         Return the type of grid used by model 'model_name' in simulation
@@ -2316,12 +2294,142 @@ def update_dataset(self, dataset, netcdf_info=None, mesh=None, update_data=True)
         else:
             nc_info = netcdf_info
 
+        if (
+            self.simulation.simulation_data.verbosity_level.value
+            >= VerbosityLevel.normal.value
+        ):
+            print(f"    updating model dataset...")
+
         for a in nc_info["attrs"]:
             dataset.attrs[a] = nc_info["attrs"][a]
 
         # add all packages and update data
         for p in self.packagelist:
             # add package var to dataset
+            if (
+                self.simulation.simulation_data.verbosity_level.value
+                >= VerbosityLevel.normal.value
+            ):
+                print(f"    updating dataset for package {p._get_pname()}...")
             dataset = p.update_dataset(dataset, mesh=mesh, update_data=update_data)
 
         return dataset
+
+    def _write_netcdf(self, mesh=None):
+        import datetime
+
+        from ..version import __version__
+        if mesh is not None and mesh.upper() == "STRUCTURED":
+            mesh = None
+
+        encode = {}
+        for pp in self.packagelist:
+            if pp.package_type == "ncf":
+                encode["shuffle"] = pp.shuffle.get_data()
+                encode["deflate"] = pp.deflate.get_data()
+                encode["chunk_time"] = pp.chunk_time.get_data()
+                encode["chunk_face"] = pp.chunk_face.get_data()
+                encode["chunk_x"] = pp.chunk_x.get_data()
+                encode["chunk_y"] = pp.chunk_y.get_data()
+                encode["chunk_z"] = pp.chunk_z.get_data()
+                wkt = pp.wkt.get_data()
+                if wkt is not None:
+                    wkt = wkt[0][1]
+                encode["wkt"] = wkt
+
+        if (
+            self.simulation.simulation_data.verbosity_level.value
+            >= VerbosityLevel.normal.value
+        ):
+            print(f"    creating model dataset...")
+
+        ds = self.modelgrid.dataset(
+            modeltime=self.modeltime,
+            mesh=mesh,
+            encoding=encode,
+        )
+
+        dt = datetime.datetime.now()
+        timestamp = dt.strftime("%m/%d/%Y %H:%M:%S")
+
+        nc_info = self.netcdf_info(mesh=mesh)
+        nc_info["attrs"]["title"] = f"{self.name.upper()} input"
+        nc_info["attrs"]["source"] = f"flopy {__version__}"
+        nc_info["attrs"]["history"] = f"first created {timestamp}"
+        if mesh is None:
+            nc_info["attrs"]["Conventions"] = "CF-1.11"
+        elif mesh.upper() is "LAYERED":
+            nc_info["attrs"]["Conventions"] = "CF-1.11 UGRID-1.0"
+
+        ds = self.update_dataset(
+            ds,
+            netcdf_info=nc_info,
+            mesh=mesh,
+        )
+
+        chunk = False
+        chunk_t = False
+        if mesh is None:
+            if (
+                "chunk_x" in encode
+                and encode["chunk_x"] is not None
+                and "chunk_y" in encode
+                and encode["chunk_y"] is not None
+                and "chunk_z" in encode
+                and encode["chunk_z"] is not None
+            ):
+                chunk = True
+        elif mesh.upper() == "LAYERED":
+            if "chunk_face" in encode and encode["chunk_face"] is not None:
+                chunk = True
+        if "chunk_time" in encode and encode["chunk_time"] is not None:
+            chunk_t = True
+
+        base_encode = {}
+        if "deflate" in encode and encode["deflate"] is not None:
+            base_encode["zlib"] = True
+            base_encode["complevel"] = encode["deflate"]
+        if "shuffle" in encode and encode["deflate"] is not None:
+            base_encode["shuffle"] = True
+
+        encoding = {}
+        chunk_dims = {'time', 'nmesh_face', 'z', 'y', 'x'}
+        for varname, da in ds.data_vars.items():
+            dims = ds.data_vars[varname].dims
+            codes = dict(base_encode)
+            if (
+                not set(dims).issubset(chunk_dims)
+                or not chunk or not chunk_t
+            ):
+                encoding[varname] = codes
+                continue
+            chunksizes = []
+            if "time" in dims:
+                chunksizes.append(encode["chunk_time"])
+            if mesh is None:
+                if "z" in dims:
+                    chunksizes.append(encode["chunk_z"])
+                if "y" in dims:
+                    chunksizes.append(encode["chunk_y"])
+                if "x" in dims:
+                    chunksizes.append(encode["chunk_x"])
+            elif mesh.upper() == "LAYERED" and "nmesh_face" in dims:
+                chunksizes.append(encode["chunk_face"])
+            if len(chunksizes) > 0:
+                codes["chunksizes"] = chunksizes
+            encoding[varname] = codes
+
+        fname = self.name_file.nc_filerecord.get_data()[0][0]
+
+        if (
+            self.simulation.simulation_data.verbosity_level.value
+            >= VerbosityLevel.normal.value
+        ):
+            print(f"    writing NetCDF file {fname}...")
+        # write dataset to netcdf
+        ds.to_netcdf(
+            os.path.join(self.model_ws, fname),
+            format="NETCDF4",
+            engine="netcdf4",
+            encoding=encoding,
+        )

flopy/mf6/mfpackage.py

@@ -3760,6 +3760,8 @@ def _data_shape(shape):
 
             return dims_l
 
+        projection = "projection" in dataset.data_vars
+
         last_path = ''
         pitem = None
         pdata = None
@@ -3774,6 +3776,15 @@ def _data_shape(shape):
             dataset = dataset.assign(var_d)
             for a in nc_info[v]["attrs"]:
                 dataset[varname].attrs[a] = nc_info[v]["attrs"][a]
+            if projection:
+                dims = dataset[varname].dims
+                if "nmesh_face" in dims or "nmesh_node" in dims:
+                    dataset[varname].attrs["grid_mapping"] = "projection"
+                    dataset[varname].attrs["coordinates"] = "mesh_face_x mesh_face_y"
+                elif mesh is None and len(dims) > 1:
+                    # TODO don't set if lon / lat?
+                    dataset[varname].attrs["grid_mapping"] = "projection"
+                    dataset[varname].attrs["coordinates"] = "x y"
 
             if update_data:
                 path = nc_info[v]["attrs"]["modflow_input"]

flopy/utils/datautil.py

@@ -305,6 +305,7 @@ def reset_delimiter_used():
 
     @staticmethod
     def split_data_line(line, external_file=False, delimiter_conf_length=15):
+        no_split_keys = ['crs', 'wkt']
         if PyListUtil.line_num > delimiter_conf_length and PyListUtil.consistent_delim:
             # consistent delimiter has been found.  continue using that
             # delimiter without doing further checks
@@ -358,7 +359,11 @@ def split_data_line(line, external_file=False, delimiter_conf_length=15):
                     max_split_size = len(max_split_list)
                     max_split_type = "combo"
 
-            if max_split_type is not None and max_split_size > 1:
+            if (
+                max_split_type is not None
+                and max_split_size > 1
+                and clean_line[0].lower() not in no_split_keys
+            ):
                 clean_line = max_split_list
                 if PyListUtil.line_num == 0:
                     PyListUtil.delimiter_used = max_split_type