Merge pull request #527 from ubermag/521-to_vtk

521 to_vtk

discretisedfield/field.py

@@ -3378,7 +3378,10 @@ def to_vtk(self):
 
         This method convers at `discretisedfield.Field` into a
         `vtk.vtkRectilinearGrid`. The field data (``field.array``) is stored as
-        ``CELL_DATA`` of the ``RECTILINEAR_GRID``. Scalar fields (``nvdim=1``)
+        ``CELL_DATA`` of the ``RECTILINEAR_GRID``. Only fields with ``ndim=3``
+        can be converted. The ``x``, ``y``, and ``z``, spatial coordinates of
+        the VTK array will be in the same order as ``field.mesh.region.dims``.
+        Scalar fields (``nvdim=1``)
         contain one VTK array called ``field``. Vector fields (``nvdim>1``)
         contain one VTK array called ``field`` containing vector data and
         scalar VTK arrays for each field component (called
@@ -3420,15 +3423,13 @@ def to_vtk(self):
         rgrid = vtkRectilinearGrid()
         rgrid.SetDimensions(*(n + 1 for n in self.mesh.n))
 
-        rgrid.SetXCoordinates(
-            vns.numpy_to_vtk(np.fromiter(self.mesh.vertices.x, float))
-        )
-        rgrid.SetYCoordinates(
-            vns.numpy_to_vtk(np.fromiter(self.mesh.vertices.y, float))
-        )
-        rgrid.SetZCoordinates(
-            vns.numpy_to_vtk(np.fromiter(self.mesh.vertices.z, float))
-        )
+        for dim, setter in zip(
+            self.mesh.region.dims,
+            [rgrid.SetXCoordinates, rgrid.SetYCoordinates, rgrid.SetZCoordinates],
+        ):
+            setter(
+                vns.numpy_to_vtk(np.fromiter(getattr(self.mesh.vertices, dim), float))
+            )
 
         cell_data = rgrid.GetCellData()
         field_norm = vns.numpy_to_vtk(

discretisedfield/tests/test_field.py

@@ -13,6 +13,7 @@
 import pytest
 import pyvista as pv
 import scipy.fft as spfft
+import vtk
 import xarray as xr
 
 import discretisedfield as df
@@ -2791,6 +2792,49 @@ def test_write_read_ovf(tmp_path):
         assert f_read.array.nbytes > 0
 
 
+def test_to_vtk():
+    p1 = (0, 0, 0)
+    p2 = (5e-9, 2e-9, 1e-9)
+    cell = (1e-9, 1e-9, 1e-9)
+
+    def value_fun(point):
+        a, b, c = point
+        return (3 + a * b, a - 2 * b, a * b * c)
+
+    region = df.Region(p1=p1, p2=p2)
+    mesh = df.Mesh(region=region, cell=cell)
+
+    f = df.Field(mesh, nvdim=3, value=value_fun)
+    f_vtk = f.to_vtk()
+    assert isinstance(f_vtk, vtk.vtkRectilinearGrid)
+    cell_data = f_vtk.GetCellData()
+    assert cell_data.GetArrayName(0) == "norm"
+    assert cell_data.GetArrayName(1) == "x"
+    assert cell_data.GetArrayName(2) == "y"
+    assert cell_data.GetArrayName(3) == "z"
+    assert cell_data.GetArrayName(4) == "field"
+
+    # Test with different dims
+    dims = ["a", "b", "c"]
+    vdims = ["va", "vb", "vc"]
+    vdim_mapping = dict(zip(vdims, dims))
+
+    region = df.Region(p1=p1, p2=p2, dims=dims)
+    mesh = df.Mesh(region=region, cell=cell)
+
+    f = df.Field(mesh, nvdim=3, value=value_fun, vdims=vdims, vdim_mapping=vdim_mapping)
+    f.to_vtk()
+
+    f_vtk = f.to_vtk()
+    assert isinstance(f_vtk, vtk.vtkRectilinearGrid)
+    cell_data = f_vtk.GetCellData()
+    assert cell_data.GetArrayName(0) == "norm"
+    assert cell_data.GetArrayName(1) == "va"
+    assert cell_data.GetArrayName(2) == "vb"
+    assert cell_data.GetArrayName(3) == "vc"
+    assert cell_data.GetArrayName(4) == "field"
+
+
 def test_write_read_vtk(tmp_path):
     filename = "testfile.vtk"