Add CSV export method to Field class

src/gstools/field/base.py

@@ -434,6 +434,74 @@ def vtk_export(
             fieldname=fieldname,
         )
 
+    def csv_export(self, filename, fields=None):  # pragma: no cover
+        """Export the stored field(s) to csv with unstructured mesh type.
+
+        First columns with be the components of the pos tuple followed by
+        the selected fields.
+
+        Parameters
+        ----------
+        filename : :class:`str`
+            Filename of the file to be saved, including the path.
+        fields : :class:`list` of :class:`str`, optional
+            Fields that should be stored. All by default.
+        """
+        fields = self.field_names if fields is None else fields
+        if not fields:
+            msg = "csv_export: no fields selected"
+            raise ValueError(msg)
+        if not all(field in self for field in fields):
+            unknown = set(fields) - set(self.field_names)
+            msg = f"csv_export: some fields are unknown: {unknown}"
+            raise ValueError(msg)
+        # generate axis names
+        s_dim = self.dim - int(self.temporal)
+        if self.latlon:
+            p_names = ["lat", "lon"]
+        elif self.dim < 4:
+            p_names = ["x", "y", "z"][:s_dim]
+        else:
+            p_names = [f"x{d}" for d in range(s_dim)]
+        if self.temporal:
+            p_names.append("t")
+        # generate fields
+        if self.mesh_type != "unstructured":
+            pos = generate_grid(self.pos)
+            if self.value_type == "vector":
+                out_f = {
+                    f"{f}_{p_names[i]}": self[f][i].reshape(-1)
+                    for f in fields
+                    for i in range(self.dim)
+                }
+            else:
+                out_f = {f: self[f].reshape(-1) for f in fields}
+        else:
+            pos = self.pos
+            if self.value_type == "vector":
+                out_f = {
+                    f"{f}_{p_names[i]}": self[f][i]
+                    for f in fields
+                    for i in range(self.dim)
+                }
+            else:
+                out_f = {f: self[f] for f in fields}
+        # generate output matrix
+        p_names += out_f.keys()
+        data = np.empty((len(pos[0]), len(p_names)), dtype=float)
+        for i, p in enumerate(pos):
+            data[:, i] = p
+        for i, f in enumerate(out_f.values()):
+            data[:, i + self.dim] = f
+        np.savetxt(
+            filename,
+            data,
+            fmt="%s",
+            delimiter=",",
+            header=",".join(p_names),
+            comments="",
+        )
+
     def plot(
         self, field="field", fig=None, ax=None, **kwargs
     ):  # pragma: no cover

tests/test_export.py

@@ -8,6 +8,7 @@
 import numpy as np
 
 from gstools import SRF, Exponential, Gaussian
+from gstools.field import Field
 from gstools.random import MasterRNG
 
 HAS_PYVISTA = False
@@ -19,19 +20,39 @@
     pass
 
 
+def get_first_line(file):
+    with open(file) as f:
+        first_line = f.readline().strip("\n")
+    return first_line
+
+
 class TestExport(unittest.TestCase):
     def setUp(self):
         self.test_dir = tempfile.mkdtemp()
         # structured field with a size 100x100x100 and a grid-size of 1x1x1
-        x = y = z = range(50)
+        x = y = z = range(10)
         model = Gaussian(dim=3, var=0.6, len_scale=20)
-        self.srf_structured = SRF(model)
+        self.srf_structured = SRF(model, seed=20170519)
         self.srf_structured((x, y, z), mesh_type="structured")
+        # vector
+        model = Gaussian(dim=2, var=0.6, len_scale=10)
+        self.srf_vector = SRF(model, generator="VectorField", seed=19841203)
+        self.srf_vector((x, y), mesh_type="structured")
+        # latlon temporal
+        model = Gaussian(latlon=True, temporal=True, var=0.6, len_scale=20)
+        self.srf_latlon_temp = SRF(model, seed=20170519)
+        self.srf_latlon_temp((x, y, z), mesh_type="structured")
+        self.srf_latlon_temp((x, y, z), mesh_type="structured", store="other")
+        # 4d
+        x = y = z = v = range(3)
+        model = Gaussian(dim=4, var=0.6, len_scale=1)
+        self.srf_4d = SRF(model, seed=20170519)
+        self.srf_4d((x, y, z, v), mesh_type="structured")
         # unstrucutred field
         seed = MasterRNG(19970221)
         rng = np.random.RandomState(seed())
-        x = rng.randint(0, 100, size=1000)
-        y = rng.randint(0, 100, size=1000)
+        x = rng.randint(0, 100, size=100)
+        y = rng.randint(0, 100, size=100)
         model = Exponential(
             dim=2, var=1, len_scale=[12.0, 3.0], angles=np.pi / 8.0
         )
@@ -59,6 +80,41 @@ def test_pyevtk_export(self):
         self.srf_unstructured.vtk_export(ufilename)
         self.assertTrue(os.path.isfile(ufilename + ".vtu"))
 
+    def test_csv_export(self):
+        # Structured
+        sfilename = os.path.join(self.test_dir, "structured.csv")
+        self.srf_structured.csv_export(sfilename)
+        self.assertTrue(os.path.isfile(sfilename))
+        self.assertTrue(get_first_line(sfilename) == "x,y,z,field")
+        # Unstructured
+        ufilename = os.path.join(self.test_dir, "unstructured.csv")
+        self.srf_unstructured.csv_export(ufilename)
+        self.assertTrue(os.path.isfile(ufilename))
+        self.assertTrue(get_first_line(ufilename) == "x,y,field")
+        # latlon temp
+        lfilename = os.path.join(self.test_dir, "latlon.csv")
+        self.srf_latlon_temp.csv_export(lfilename)
+        self.assertTrue(os.path.isfile(lfilename))
+        self.assertTrue(get_first_line(lfilename) == "lat,lon,t,field,other")
+        # vector
+        vfilename = os.path.join(self.test_dir, "vector.csv")
+        self.srf_vector.csv_export(vfilename)
+        self.assertTrue(os.path.isfile(vfilename))
+        self.assertTrue(get_first_line(vfilename) == "x,y,field_x,field_y")
+        # 4D
+        dfilename = os.path.join(self.test_dir, "4D.csv")
+        self.srf_4d.csv_export(dfilename)
+        self.assertTrue(os.path.isfile(dfilename))
+        self.assertTrue(get_first_line(dfilename) == "x0,x1,x2,x3,field")
+        # check errors
+        field = Field(dim=4)
+        with self.assertRaises(ValueError):
+            field.csv_export("test.csv")
+        with self.assertRaises(ValueError):
+            field.csv_export("test.csv", fields=[])
+        with self.assertRaises(ValueError):
+            field.csv_export("test.csv", fields=["wow"])
+
 
 if __name__ == "__main__":
     unittest.main()