Φ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 *

2025-08-05 10:59:37.646386: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1754391577.660580    2395 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1754391577.665045    2395 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
W0000 00:00:1754391577.676940    2395 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
W0000 00:00:1754391577.676956    2395 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
W0000 00:00:1754391577.676958    2395 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
W0000 00:00:1754391577.676959    2395 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
2025-08-05 10:59:37.681422: 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 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2025-08-05 10:59:39.816690: 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') | 15996 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.459 ± 0.321 (4e-03...1e+00)

Slice a Tensor¶

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

Print a Tensor¶

In [8]:
print(data)
print(f"{data:full:shape:dtype:color:.1f}")

(examplesᵇ=10, vectorᶜ=x,y) 0.459 ± 0.321 (4e-03...1e+00)
(examplesᵇ=10, vectorᶜ=x,y)
[[0.4, 0.9],
 [0.4, 0.2],
 [0.9, 1.0],
 [0.4, 0.5],
 [0.3, 0.2],
 [0.0, 0.0],
 [0.1, 0.7],
 [0.2, 0.8],
 [0.8, 1.0],
 [0.2, 0.3]]

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.11/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([[0.3053004 , 0.12921482, 0.6108594 , 0.8945456 , 0.11128128,
        0.17524141],
       [0.57702404, 0.8549522 , 0.3089379 , 0.0227446 , 0.51802444,
        0.3481202 ],
       [0.62038374, 0.6008116 , 0.7939808 , 0.8664693 , 0.68513113,
        0.8591997 ],
       [0.5470452 , 0.72687906, 0.44776207, 0.7905265 , 0.09467643,
        0.07597095],
       [0.51995486, 0.47496855, 0.4884587 , 0.88716173, 0.02450246,
        0.924181  ],
       [0.9550246 , 0.6012362 , 0.10125625, 0.6585826 , 0.29133928,
        0.68323684],
       [0.81157273, 0.12852931, 0.42233598, 0.60510194, 0.31528926,
        0.51809126],
       [0.707283  , 0.9658    , 0.72280055, 0.89666796, 0.3610866 ,
        0.9003617 ]], dtype=float32)
In [11]:
math.reshaped_native(data, ['extra', data.shape], to_numpy=True)

/tmp/ipykernel_2395/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([[0.3053004 , 0.12921482, 0.6108594 , 0.8945456 , 0.11128128,
        0.17524141, 0.57702404, 0.8549522 , 0.3089379 , 0.0227446 ,
        0.51802444, 0.3481202 , 0.62038374, 0.6008116 , 0.7939808 ,
        0.8664693 , 0.68513113, 0.8591997 , 0.5470452 , 0.72687906,
        0.44776207, 0.7905265 , 0.09467643, 0.07597095, 0.51995486,
        0.47496855, 0.4884587 , 0.88716173, 0.02450246, 0.924181  ,
        0.9550246 , 0.6012362 , 0.10125625, 0.6585826 , 0.29133928,
        0.68323684, 0.81157273, 0.12852931, 0.42233598, 0.60510194,
        0.31528926, 0.51809126, 0.707283  , 0.9658    , 0.72280055,
        0.89666796, 0.3610866 , 0.9003617 ]], 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_2395/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))
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))

/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2834: RuntimeWarning: Stacking shapes with incompatible labels will result in labels being lost. For vector Got ('x', 'y') and ('x',)
  warnings.warn(f"Stacking shapes with incompatible labels will result in labels being lost. For {dim.name} Got {prev.slice_names} and {dim.slice_names}", RuntimeWarning)
Out[14]:

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)
Out[15]:
Field[(xˢ=32, yˢ=32, vectorᶜ=x,y)]

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]:
array([42.], dtype=float32)

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)

/tmp/ipykernel_2395/1489350838.py:9: UserWarning: bounds argument is deprecated since 2.5 and will be ignored.
  segments.append(PointCloud(start, end - start, bounds=Box(x=2*PI, y=2*PI)))
Out[20]:
Your browser does not support the video tag. 
<Figure size 640x480 with 0 Axes>

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.248 ± 0.244 (1e-05...9e-01)
Final loss: (batchᵇ=100) 0.119 ± 0.138 (6e-04...8e-01)

Parameter Count¶

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