ΦFlow Cookbook¶

This notebook lists useful code snippets.

Import for NumPy, TensorFlow, Jax, PyTorch¶

In [1]:
from phi.flow import *
from phi.tf.flow import *
from phi.jax.stax.flow import *
from phi.torch.flow import *
2026-02-06 17:00:14.777256: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2026-02-06 17:00:14.825945: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2026-02-06 17:00:16.813338: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303)

Select GPU or CPU¶

In [2]:
backend.default_backend().list_devices('GPU')
Out[2]:
[]
In [3]:
backend.default_backend().list_devices('CPU')
Out[3]:
[torch device 'CPU' (CPU 'cpu') | 15990 MB | 4 processors | ]
In [4]:
assert backend.default_backend().set_default_device('CPU')

Use 64 bit FP precision¶

In [5]:
math.set_global_precision(32)  # single precision is the default
x32 = math.random_normal(batch(b=4))

with math.precision(64):  ## operations within this context will use 32 bit floats
    x64 = math.to_float(x32)

Sample Random Values¶

In [6]:
data = math.random_normal(batch(examples=10)) * .1  # batch of scalar values
data = math.random_uniform(batch(examples=10), channel(vector='x,y'))  # batch of vectors
data
Out[6]:
(examplesᵇ=10, vectorᶜ=x,y) 0.601 ± 0.237 (6e-02...1e+00)

Slice a Tensor¶

In [7]:
data.examples[0]
Out[7]:
(x=0.758, y=0.398)

Print a Tensor¶

In [8]:
print(data)
print(f"{data:full:shape:dtype:color:.1f}")
(examplesᵇ=10, vectorᶜ=x,y) 0.601 ± 0.237 (6e-02...1e+00)
(examplesᵇ=10, vectorᶜ=x,y)
[[0.8, 0.4],
 [0.7, 0.7],
 [0.5, 0.6],
 [0.7, 0.2],
 [0.8, 0.3],
 [0.6, 0.7],
 [0.2, 1.0],
 [0.1, 0.6],
 [0.9, 0.7],
 [0.7, 0.8]]

Plot a Tensor¶

In [9]:
data = math.random_uniform(spatial(x=8, y=6))
vis.plot(data)  # or vis.show(data)
/opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/vis/_matplotlib/_matplotlib_plots.py:167: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  plt.tight_layout()  # because subplot titles can be added after figure creation
Out[9]:

Convert a Tensor to NumPy¶

In [10]:
data.numpy(order='x,y')
Out[10]:
array([[9.4992673e-01, 3.7592220e-01, 8.7513375e-01, 4.6079612e-01,
        5.3339213e-01, 5.7152241e-01],
       [3.6209714e-01, 8.7821102e-01, 1.2370884e-02, 3.2324076e-01,
        4.5080185e-03, 9.2221463e-01],
       [8.7916315e-02, 2.9954505e-01, 9.5225763e-01, 1.8308771e-01,
        6.7085326e-02, 7.7423012e-01],
       [2.8557402e-01, 6.3253170e-01, 5.0047016e-01, 8.7799120e-01,
        6.4364505e-01, 8.8393360e-01],
       [6.4275575e-01, 4.2908192e-01, 3.8105249e-04, 9.0660989e-02,
        2.9684055e-01, 7.3211509e-01],
       [5.9397995e-02, 6.4171064e-01, 8.6398959e-01, 4.1184765e-01,
        1.3389200e-01, 4.8882216e-01],
       [3.4517133e-01, 8.6708307e-01, 1.6758502e-01, 7.2620261e-01,
        8.8897759e-01, 6.7500055e-01],
       [7.2264993e-01, 1.2240374e-01, 8.2114720e-01, 4.9355567e-02,
        3.6108041e-01, 4.9218041e-01]], dtype=float32)
In [11]:
math.reshaped_native(data, ['extra', data.shape], to_numpy=True)
/tmp/ipykernel_2577/2990683702.py:1: DeprecationWarning: phiml.math.reshaped_native() is deprecated. Use Tensor.native() instead.
  math.reshaped_native(data, ['extra', data.shape], to_numpy=True)
Out[11]:
array([[9.4992673e-01, 3.7592220e-01, 8.7513375e-01, 4.6079612e-01,
        5.3339213e-01, 5.7152241e-01, 3.6209714e-01, 8.7821102e-01,
        1.2370884e-02, 3.2324076e-01, 4.5080185e-03, 9.2221463e-01,
        8.7916315e-02, 2.9954505e-01, 9.5225763e-01, 1.8308771e-01,
        6.7085326e-02, 7.7423012e-01, 2.8557402e-01, 6.3253170e-01,
        5.0047016e-01, 8.7799120e-01, 6.4364505e-01, 8.8393360e-01,
        6.4275575e-01, 4.2908192e-01, 3.8105249e-04, 9.0660989e-02,
        2.9684055e-01, 7.3211509e-01, 5.9397995e-02, 6.4171064e-01,
        8.6398959e-01, 4.1184765e-01, 1.3389200e-01, 4.8882216e-01,
        3.4517133e-01, 8.6708307e-01, 1.6758502e-01, 7.2620261e-01,
        8.8897759e-01, 6.7500055e-01, 7.2264993e-01, 1.2240374e-01,
        8.2114720e-01, 4.9355567e-02, 3.6108041e-01, 4.9218041e-01]],
      dtype=float32)

Compute Pair-wise Distances¶

In [12]:
points = math.tensor([(0, 0), (0, 1), (1, 0)], instance('points'), channel('vector'))
distances = points - math.rename_dims(points, 'points', 'others')
math.print(math.vec_length(distances))
[[0.       , 1.       , 1.       ],
 [1.       , 0.       , 1.4142135],
 [1.       , 1.4142135, 0.       ]]
/tmp/ipykernel_2577/2195475714.py:3: DeprecationWarning: phiml.math.length is deprecated in favor of phiml.math.norm
  math.print(math.vec_length(distances))

Construct a CenteredGrid¶

In [13]:
zero_grid = CenteredGrid(0, 0, x=32, y=32, bounds=Box(x=1, y=1))
y_grid = CenteredGrid((0, 1), extrapolation.BOUNDARY, x=32, y=32)
noise_grid = CenteredGrid(Noise(), extrapolation.PERIODIC, x=32, y=32)
sin_curve = CenteredGrid(lambda x: math.sin(x), extrapolation.PERIODIC, x=100, bounds=Box(x=2 * PI))

vis.plot(zero_grid, y_grid, noise_grid, sin_curve, size=(12, 3))
/opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:1174: RuntimeWarning: invalid value encountered in scalar power
  result = op(n1, n2)
Out[13]:

Construct a StaggeredGrid¶

In [14]:
zero_grid = StaggeredGrid(0, 0, x=32, y=32, bounds=Box(x=1, y=1))
y_grid = StaggeredGrid((0, 1), extrapolation.BOUNDARY, x=32, y=32)
noise_grid = StaggeredGrid(Noise(), extrapolation.PERIODIC, x=32, y=32)
sin_curve = StaggeredGrid(lambda x: math.sin(x), extrapolation.PERIODIC, x=100, bounds=Box(x=2 * PI))

vis.plot(zero_grid, y_grid, noise_grid, sin_curve, size=(12, 3))
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[14], line 1
----> 1 zero_grid = StaggeredGrid(0, 0, x=32, y=32, bounds=Box(x=1, y=1))
      2 y_grid = StaggeredGrid((0, 1), extrapolation.BOUNDARY, x=32, y=32)
      3 noise_grid = StaggeredGrid(Noise(), extrapolation.PERIODIC, x=32, y=32)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_grid.py:176, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_)
    174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}"
    175 assert values.shape.instance_rank == 0, f"Instance dimensions not supported for grids. Got values with shape {values.shape}"
--> 176 return Field(elements, values, extrapolation)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:44, in _FieldType.__call__(cls, geometry, values, boundary, variable_attrs, value_attrs, **sampling_kwargs)
     42         values = expand(wrap(values), non_batch(geometry) - 'vector')
     43 result = cls.__new__(cls, geometry, values, boundary, variable_attrs, value_attrs)
---> 44 result.__init__(geometry, values, boundary, variable_attrs, value_attrs)  # also calls __post_init__()
     45 return result

File <string>:8, in __init__(self, geometry, values, boundary, variable_attrs, value_attrs)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:83, in Field.__post_init__(self)
     82 def __post_init__(self):
---> 83     at = self.sampled_at
     84     if at in {'center', 'face'}:
     85         math.merge_shapes(self.values, non_batch(self.sampled_elements).non_channel)  # shape check

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:385, in Field.sampled_at(self)
    383 for name, s_shape in self.geometry.sets.items():
    384     if s_shape.non_batch in v_shape:  # all necessary dims present in values
--> 385         if v_shape.only(s_shape, reorder=True).sizes == s_shape.sizes:
    386             return name
    387 raise ValueError(f"Could not determine where the values of this Field are sampled. Geometry sets: {self.geometry.sets}, Field values shape: {v_shape}")

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:221, in Tensor.__bool__(self)
    220 def __bool__(self):
--> 221     assert self.rank == 0, f"Cannot convert tensor with non-empty shape {self.shape} to bool. Use tensor.any or tensor.all instead."
    222     from ._ops import all_
    223     if not self.default_backend.supports(Backend.jit_compile):  # NumPy

AssertionError: Cannot convert tensor with non-empty shape (~vectorᵈ=x,y) to bool. Use tensor.any or tensor.all instead.

Construct StaggeredGrid from NumPy Arrays¶

Given matching arrays vx and vy, we can construct a StaggeredGrid. Note that the shapes of the arrays must match the extrapolation!

In [15]:
vx = math.tensor(np.zeros([33, 32]), spatial('x,y'))
vy = math.tensor(np.zeros([32, 33]), spatial('x,y'))
StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.BOUNDARY)

vx = math.tensor(np.zeros([32, 32]), spatial('x,y'))
vy = math.tensor(np.zeros([32, 32]), spatial('x,y'))
StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.PERIODIC)

vx = math.tensor(np.zeros([31, 32]), spatial('x,y'))
vy = math.tensor(np.zeros([32, 31]), spatial('x,y'))
StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), 0)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[15], line 3
      1 vx = math.tensor(np.zeros([33, 32]), spatial('x,y'))
      2 vy = math.tensor(np.zeros([32, 33]), spatial('x,y'))
----> 3 StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.BOUNDARY)
      5 vx = math.tensor(np.zeros([32, 32]), spatial('x,y'))
      6 vy = math.tensor(np.zeros([32, 32]), spatial('x,y'))

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_grid.py:176, in StaggeredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_)
    174 assert values.shape.spatial_rank == elements.bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}"
    175 assert values.shape.instance_rank == 0, f"Instance dimensions not supported for grids. Got values with shape {values.shape}"
--> 176 return Field(elements, values, extrapolation)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:44, in _FieldType.__call__(cls, geometry, values, boundary, variable_attrs, value_attrs, **sampling_kwargs)
     42         values = expand(wrap(values), non_batch(geometry) - 'vector')
     43 result = cls.__new__(cls, geometry, values, boundary, variable_attrs, value_attrs)
---> 44 result.__init__(geometry, values, boundary, variable_attrs, value_attrs)  # also calls __post_init__()
     45 return result

File <string>:8, in __init__(self, geometry, values, boundary, variable_attrs, value_attrs)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:83, in Field.__post_init__(self)
     82 def __post_init__(self):
---> 83     at = self.sampled_at
     84     if at in {'center', 'face'}:
     85         math.merge_shapes(self.values, non_batch(self.sampled_elements).non_channel)  # shape check

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_field.py:385, in Field.sampled_at(self)
    383 for name, s_shape in self.geometry.sets.items():
    384     if s_shape.non_batch in v_shape:  # all necessary dims present in values
--> 385         if v_shape.only(s_shape, reorder=True).sizes == s_shape.sizes:
    386             return name
    387 raise ValueError(f"Could not determine where the values of this Field are sampled. Geometry sets: {self.geometry.sets}, Field values shape: {v_shape}")

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:221, in Tensor.__bool__(self)
    220 def __bool__(self):
--> 221     assert self.rank == 0, f"Cannot convert tensor with non-empty shape {self.shape} to bool. Use tensor.any or tensor.all instead."
    222     from ._ops import all_
    223     if not self.default_backend.supports(Backend.jit_compile):  # NumPy

AssertionError: Cannot convert tensor with non-empty shape (~vectorᵈ=x,y) to bool. Use tensor.any or tensor.all instead.

BFGS Optimization¶

In [16]:
def loss_function(x):
    return math.l2_loss(math.cos(x))

initial_guess = math.tensor([1, -1], math.batch('batch'))
math.minimize(loss_function, Solve('L-BFGS-B', 0, 1e-3, x0=initial_guess))
Out[16]:
(1.574, -1.574) along batchᵇ

Linear Solve¶

In [17]:
def f(x):
    return 2 * x

math.solve_linear(f, 84, Solve('CG', 1e-5, x0=0))
Out[17]:
tensor([42.])

Sparse Matrix Construction¶

In [18]:
from functools import partial

periodic_laplace = partial(math.laplace, padding=extrapolation.PERIODIC)
example_input = math.ones(spatial(x=3))
matrix, bias = math.matrix_from_function(periodic_laplace, example_input)
math.print(matrix)
x=0     -2.   1.   1.  along ~x
x=1      1.  -2.   1.  along ~x
x=2      1.   1.  -2.  along ~x

Sampling a Function¶

In [19]:
def f(x):
    return math.l2_loss(math.sin(x))

f_grid = CenteredGrid(f, x=100, y=100, bounds=Box(x=2*PI, y=2*PI))
vis.plot(f_grid)
Out[19]:

Plot Optimization Trajectories¶

In [20]:
def minimize(x0):
    with math.SolveTape(record_trajectories=True) as solves:
        math.minimize(f, Solve('BFGS', 0, 1e-5, x0=x0))
    return solves[0].x  # shape (trajectory, x, y, vector)

trajectories = CenteredGrid(minimize, x=8, y=8, bounds=Box(x=2*PI, y=2*PI)).values
segments = []
for start, end in zip(trajectories.trajectory[:-1].trajectory, trajectories.trajectory[1:].trajectory):
    segments.append(PointCloud(start, end - start, bounds=Box(x=2*PI, y=2*PI)))
anim_segments = field.stack(segments, batch('time'))
vis.plot(f_grid, anim_segments, overlay='args', animate='time', color='#FFFFFF', frame_time=500)
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[20], line 6
      3         math.minimize(f, Solve('BFGS', 0, 1e-5, x0=x0))
      4     return solves[0].x  # shape (trajectory, x, y, vector)
----> 6 trajectories = CenteredGrid(minimize, x=8, y=8, bounds=Box(x=2*PI, y=2*PI)).values
      7 segments = []
      8 for start, end in zip(trajectories.trajectory[:-1].trajectory, trajectories.trajectory[1:].trajectory):

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/field/_grid.py:75, in CenteredGrid(values, boundary, bounds, resolution, extrapolation, convert, **resolution_)
     73     values = sample(values, elements)
     74 elif callable(values):
---> 75     values = sample_function(values, elements, 'center', extrapolation)
     76 else:
     77     if isinstance(values, (tuple, list)) and len(values) == resolution.rank:

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phi/geom/_geom.py:842, in sample_function(f, elements, at, extrapolation)
    840     values = math.map_s2b(f)(*pos.vector)
    841 else:
--> 842     values = math.map_s2b(f)(pos)
    843 assert isinstance(values, math.Tensor), f"values function must return a Tensor but returned {type(values)}"
    844 return values

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_functional.py:1230, in map_types.<locals>.retyped_f(*args, **kwargs)
   1228     retyped_kwarg, input_types = forward_retype(v, input_types)
   1229     retyped_kwargs[k] = retyped_kwarg
-> 1230 output = f(*retyped_args, **retyped_kwargs)
   1231 restored_output = reverse_retype(output, input_types)
   1232 return restored_output

Cell In[20], line 3, in minimize(x0)
      1 def minimize(x0):
      2     with math.SolveTape(record_trajectories=True) as solves:
----> 3         math.minimize(f, Solve('BFGS', 0, 1e-5, x0=x0))
      4     return solves[0].x

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:465, in minimize(f, solve)
    463     iterations = reshaped_tensor(ret[-1].iterations, [batch_dims])
    464     function_evaluations = stack([reshaped_tensor(r.function_evaluations, [batch_dims]) for r in ret], batch('trajectory'))
--> 465     result = SolveInfo(solve, x_, residual, iterations, function_evaluations, converged, diverged, ret[-1].method, ret[-1].message, t)
    466 for tape in _SOLVE_TAPES:
    467     tape._add(solve, trj, result)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_optimize.py:181, in SolveInfo.__init__(self, solve, x, residual, iterations, function_evaluations, converged, diverged, method, msg, solve_time)
    179     _, res_tensors = disassemble_tree(residual, cache=False)
    180     msg_fun = partial(_default_solve_info_msg, solve=solve)
--> 181     msg = map_(msg_fun, msg, converged.trajectory[-1], diverged.trajectory[-1], iterations.trajectory[-1], method=method, residual=res_tensors[0], dims=converged.shape.without('trajectory'))
    182 self.msg = msg
    183 """ `str`, termination message """

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_functional.py:1431, in map_(function, dims, range, unwrap_scalars, expand_results, simplify, map_name, *args, **kwargs)
   1429     assert all(r is None for r in results), f"map function returned None for some elements, {results}"
   1430     return None
-> 1431 return stack(results, dims_, expand_values=expand_results, simplify=simplify, layout_non_matching=True)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:260, in stack(values, dim, expand_values, simplify, layout_non_matching, **kwargs)
    258 # --- Fallback: multi-level stack ---
    259 for dim_ in reversed(dim):
--> 260     values = [stack(values[i:i + dim_.size], dim_, **kwargs) for i in range(0, len(values), dim_.size)]
    261 return values[0]

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:200, in stack(values, dim, expand_values, simplify, layout_non_matching, **kwargs)
    198 if any(isinstance(v, (tuple, list, dict)) for v in values_):
    199     from ._tensors import wrap, layout
--> 200     if _is_data_array(values_):
    201         tensors = [wrap(v) for v in values_]
    202         return stack(tensors, dim)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_magic_ops.py:865, in _is_data_array(sequence)
    863 def _is_data_array(sequence):
    864     try:
--> 865         all([np.asarray(v).dtype != object for v in sequence])
    866     except ValueError:  # e.g. inhomogeneous
    867         return False

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:122, in Tensor.__array__(self, dtype)
    120 if self.rank > 1:
    121     warnings.warn("Automatic conversion of Φ-ML tensors to NumPy can cause problems because the dimension order is not guaranteed.", SyntaxWarning, stacklevel=3)
--> 122 return self.numpy(self._shape)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:117, in Tensor.numpy(self, order, force_expand)
     92 def numpy(self, order: Union[str, tuple, list, Shape] = None, force_expand=True) -> np.ndarray:
     93     """
     94     Converts this tensor to a `numpy.ndarray` with dimensions ordered according to `order`.
     95     
   (...)    115         ValueError if the tensor cannot be transposed to match target_shape
    116     """
--> 117     return self.backend.numpy(self.native(order, force_expand))

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/phiml/math/_tensors.py:177, in Tensor.backend(self)
    175 @property
    176 def backend(self) -> Backend:
--> 177     raise NotImplementedError(self.__class__)

NotImplementedError: <class 'phiml.math._tree.Layout'>

Neural Network Training¶

In [21]:
net = dense_net(1, 1, layers=[8, 8], activation='ReLU')  # Implemented for PyTorch, TensorFlow, Jax-Stax
optimizer = adam(net, 1e-3)
BATCH = batch(batch=100)

def loss_function(data: Tensor):
    prediction = math.native_call(net, data)
    label = math.sin(data)
    return math.l2_loss(prediction - label), data, label

print(f"Initial loss: {loss_function(math.random_normal(BATCH))[0]}")
for i in range(100):
    loss, _data, _label = update_weights(net, optimizer, loss_function, data=math.random_normal(BATCH))
print(f"Final loss: {loss}")
Initial loss: (batchᵇ=100) 0.173 ± 0.150 (3e-05...5e-01)
Final loss: (batchᵇ=100) 0.125 ± 0.142 (6e-08...7e-01)

Parameter Count¶

In [22]:
parameter_count(net)
Out[22]:
97