diff --git a/phi/vis/_vis_base.py b/phi/vis/_vis_base.py index 892c7bc1a..a665ba4da 100644 --- 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)