[field] Add Field.as_spheres()

tum-pbs · May 11, 2024 · ed8b668 · ed8b668
1 parent b6a9884
commit ed8b668
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/phi/field/_field.py b/phi/field/_field.py
@@ -273,6 +273,9 @@ def as_points(self, list_dim: Optional[Shape] = instance('elements')) -> 'Field'
         Returns this field as a PointCloud.
         This replaces the `Field.geometry` with a `phi.geom.Point` instance while leaving the sample points unchanged.
 
+        See Also:
+            `Field.as_spheres()`.
+
         Args:
             list_dim: If not `None`, packs spatial, instance and dual dims.
                 Defaults to `instance('elements')`.
@@ -288,6 +291,31 @@ def as_points(self, list_dim: Optional[Shape] = instance('elements')) -> 'Field'
             values = pack_dims(values, dims, list_dim)
         return Field(Point(points), values, self._boundary)
 
+    def as_spheres(self, list_dim: Optional[Shape] = instance('elements')) -> 'Field':
+        """
+        Returns this field as a PointCloud with spherical / circular elements, preserving element volumes.
+        This replaces the `Field.geometry` with a `phi.geom.Sphere` instance while leaving the sample points unchanged.
+
+        See Also:
+            `Field.as_points()`.
+
+        Args:
+            list_dim: If not `None`, packs spatial, instance and dual dims.
+                Defaults to `instance('elements')`.
+
+        Returns:
+            `Field` with same values and boundaries but `Sphere` geometry.
+        """
+        points = self.sampled_elements.center
+        volumes = self.sampled_elements.volume
+        values = self._values
+        if list_dim:
+            dims = non_batch(points).non_channel & non_batch(points).non_channel
+            points = pack_dims(points, dims, list_dim)
+            values = pack_dims(values, dims, list_dim)
+            volumes = pack_dims(volumes, dims, list_dim)
+        return Field(Sphere(points, volume=volumes), values, self._boundary)
+
     def at_centers(self, **kwargs) -> 'Field':
         """
         Interpolates the values to the cell centers.

diff --git a/tests/commit/field/test__field.py b/tests/commit/field/test__field.py
@@ -80,3 +80,13 @@ def test_as_points(self):
         math.assert_close([1, 2, 3, 4], points.values)
         self.assertIsInstance(points.geometry, Point)
         math.assert_close(math.flatten(grid.points), math.flatten(points.points))
+
+    def test_as_spheres(self):
+        values = math.wrap([[1, 2], [3, 4]], spatial('x,y'))
+        grid = CenteredGrid(values, 1)
+        spheres = grid.as_spheres()
+        math.assert_close([1, 2, 3, 4], spheres.values)
+        self.assertIsInstance(spheres.geometry, Sphere)
+        math.assert_close(1, spheres.geometry.volume)
+        self.assertEqual(math.EMPTY_SHAPE, spheres.geometry.volume.shape)
+        math.assert_close(math.flatten(grid.points), math.flatten(spheres.points))