[vis] Fix plotting sparse points

tum-pbs · Nov 24, 2024 · 1640ea5 · 1640ea5
1 parent 046375d
commit 1640ea5
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/phi/vis/_vis_base.py b/phi/vis/_vis_base.py
@@ -588,13 +588,18 @@ def _limits(center: Tensor, half: Tensor, is_log: Union[bool, Tensor]):
 
 
 def only_stored_elements(f: Field) -> Field:
-    if not math.is_sparse(f.points):
+    if not math.is_sparse(f.points) and not math.is_sparse(f.values):
         return f
-    elements = f.sampled_elements.at(f.points._values)
+    if math.is_sparse(f.points):
+        points = math.stored_values(f.points)
+    else:
+        mat = f.values != 0
+        points = math.stored_values(f.points * mat)
+    elements = f.sampled_elements.at(points)
     if math.is_sparse(f.values):
-        values = f.values._values
+        values = math.stored_values(f.values)
     else:
-        values = f.values[f.points._indices]
+        values = f.values[math.stored_indices(f.points)]
     return Field(elements, values, math.extrapolation.NONE)