Updated the OpenVDB.py with new API structure

Author: spyke7

gridData/OpenVDB.py

@@ -69,6 +69,7 @@
 
 import numpy
 import warnings
+from dataclasses import dataclass
 
 try:
     import openvdb as vdb
@@ -77,6 +78,11 @@
     vdb = None
 
 
+@dataclass
+class DownCastTo:
+    gridType: str
+
+
 class OpenVDBField(object):
     """OpenVDB field object for writing volumetric data.
 
@@ -153,10 +159,12 @@ def __init__(
 
         if grid is not None:
             self._populate(grid, origin, delta)
+            self.vdb_grid = self._create_openvdb_grid()
         else:
             self.grid = None
             self.origin = None
             self.delta = None
+            self.vdb_grid = None
 
     def _populate(self, grid, origin, delta):
         """Populate the field with grid data.
@@ -206,31 +214,70 @@ def _populate(self, grid, origin, delta):
                 "delta must be either a length-3 vector or a 3x3 diagonal matrix"
             )
 
-    def write(self, filename):
-        """Write the field to an OpenVDB file.
-
-        Parameters
-        ----------
-        filename : str
-            Output filename (should end in .vdb)
+    def _get_best_grid_type(self):
+        """Selects the suitable OpenVDB grid type
 
-        Notes
-        -----
-        Limitations in OpenVDB can lead to loss of precision. If the input
-        data is not of type float32, it will be converted to FloatGrid which is float32.
+        Returns
+        -------
+        openvdb.GridBase
 
+        Raises
+        ------
+        TypeError
+            If dtype is not supported or no suitable grid type is available
         """
+        datatypes = {
+            numpy.dtype("bool"): ["BoolGrid"],
+            numpy.dtype("int8"): [DownCastTo("Int32Grid"), DownCastTo("FloatGrid")],
+            numpy.dtype("uint8"): [DownCastTo("Int32Grid"), DownCastTo("FloatGrid")],
+            numpy.dtype("int16"): [DownCastTo("Int32Grid"), DownCastTo("FloatGrid")],
+            numpy.dtype("uint16"): [DownCastTo("Int32Grid"), DownCastTo("FloatGrid")],
+            numpy.dtype("int32"): ["Int32Grid", DownCastTo("FloatGrid")],
+            numpy.dtype("uint32"): ["Int32Grid", DownCastTo("FloatGrid")],
+            numpy.dtype("int64"): ["Int64Grid", DownCastTo("FloatGrid")],
+            numpy.dtype("uint64"): ["Int64Grid", DownCastTo("FloatGrid")],
+            numpy.dtype("float16"): [DownCastTo("HalfGrid"), DownCastTo("FloatGrid")],
+            numpy.dtype("float32"): ["FloatGrid"],
+            numpy.dtype("float64"): ["DoubleGrid", DownCastTo("FloatGrid")],
+        }
+
+        try:
+            vdb_gridtypes = datatypes[self.grid.dtype]
+        except KeyError:
+            raise TypeError(f"Data type {self.grid.dtype} not supported for VDB")
+
+        for gridtype_downcast in vdb_gridtypes:
+            if isinstance(gridtype_downcast, DownCastTo):
+                gridtype = gridtype_downcast.gridType
+            else:
+                gridtype = gridtype_downcast
 
-        if self.grid.dtype == numpy.bool_ or self.grid.dtype == bool:
-            vdb_grid = vdb.BoolGrid()
-            use_tolerance = False
-
-        else:
-            vdb_grid = vdb.FloatGrid()
-            if self.tolerance == None or self.tolerance == 0:
-                use_tolerance = False
+            try:
+                VDB_Grid = getattr(vdb, gridtype)
+            except AttributeError:
+                continue
             else:
-                use_tolerance = True
+                if isinstance(gridtype_downcast, DownCastTo):
+                    warnings.warn(
+                        f"Grid type {vdb_gridtypes[0]} not available. Using {gridtype} instead. Data may lose precision.",
+                        UserWarning,
+                    )
+                return VDB_Grid()
+
+        raise TypeError(
+            f"Could not find any VDB grid type for numpy dtype {self.grid.dtype}"
+        )
+
+    def _create_openvdb_grid(self):
+        """Create and populate an OpenVDB grid
+
+        Returns
+        -------
+        openvdb.GridBase
+
+        """
+
+        vdb_grid = self._get_best_grid_type()
 
         vdb_grid.name = self.name
 
@@ -245,11 +292,27 @@ def write(self, filename):
                 except (TypeError, ValueError) as e:
                     warnings.warn(f"Could not set metadata '{key}': {e}", UserWarning)
 
-        if use_tolerance:
-            vdb_grid.copyFromArray(self.grid, tolerance=self.tolerance)
-            vdb_grid.prune()
-        else:
+        if isinstance(vdb_grid, vdb.BoolGrid) and (
+            self.tolerance is None or self.tolerance == 0
+        ):
             vdb_grid.copyFromArray(self.grid)
             vdb_grid.prune(tolerance=False)
+        else:
+            if self.tolerance is None:
+                vdb_grid.copyFromArray(self.grid)
+            else:
+                vdb_grid.copyFromArray(self.grid, tolerance=self.tolerance)
+
+            vdb_grid.prune()
+
+        return vdb_grid
 
-        vdb.write(filename, grids=[vdb_grid])
+    def write(self, filename):
+        """Write the field to an OpenVDB file.
+
+        Parameters
+        ----------
+        filename : str
+            Output filename (should end in .vdb)
+        """
+        vdb.write(filename, grids=[self.vdb_grid])

gridData/tests/test_vdb.py

@@ -272,6 +272,38 @@ def test_write_vdb_with_tolerance_parameter(self, tmpdir, grid345):
         assert acc.getValue((0, 0, 0)) == pytest.approx(float(data[0, 0, 0]))
         assert acc.getValue((1, 2, 3)) == pytest.approx(float(data[1, 2, 3]))
 
+    def test_write_vdb_float64_conversion_warning(self, tmpdir):
+        data = np.random.random((5, 5, 5))
+
+        g = Grid(data, origin=[0, 0, 0], delta=[1, 1, 1])
+        outfile = str(tmpdir / "float64.vdb")
+
+        with pytest.warns(
+            UserWarning, match="Grid type DoubleGrid not available.*FloatGrid"
+        ):
+            g.export(outfile)
+
+        grids, _ = vdb.readAll(outfile)
+        assert isinstance(grids[0], vdb.FloatGrid)
+
+    def test_write_vdb_int32_conversion_warning(self, tmpdir):
+        data = np.arange(27, dtype=np.int32).reshape((3, 3, 3))
+
+        g = Grid(data, origin=[0, 0, 0], delta=[1, 1, 1])
+        outfile = str(tmpdir / "int32.vdb")
+
+        with pytest.warns(
+            UserWarning, match="Grid type Int32Grid not available.*FloatGrid"
+        ):
+            g.export(outfile)
+
+        grids, _ = vdb.readAll(outfile)
+        assert isinstance(grids[0], vdb.FloatGrid)
+
+        acc = grids[0].getAccessor()
+        assert acc.getValue((0, 0, 0)) == pytest.approx(float(data[0, 0, 0]))
+        assert acc.getValue((1, 1, 1)) == pytest.approx(float(data[1, 1, 1]))
+
 
 @pytest.mark.skipif(
     not HAS_OPENVDB, reason="Need openvdb to test import error handling"