diff --git a/src/gstools/field/base.py b/src/gstools/field/base.py index 2006b8587..6174965b6 100755 --- a/src/gstools/field/base.py +++ b/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 diff --git a/tests/test_export.py b/tests/test_export.py index b32898f45..1f60737a8 100644 --- a/tests/test_export.py +++ b/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()