Φ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-02-28 18:59:53.092046: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] 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:1740769193.106527    2203 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1740769193.110906    2203 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-02-28 18:59:53.127566: 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-02-28 18:59:55.294772: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:152] 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.551 ± 0.314 (3e-02...9e-01)

Slice a Tensor¶

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

Print a Tensor¶

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

(examplesᵇ=10, vectorᶜ=x,y) 0.551 ± 0.314 (3e-02...9e-01)
(examplesᵇ=10, vectorᶜ=x,y)
[[0.9, 0.9],
 [0.5, 0.7],
 [0.1, 0.0],
 [0.3, 0.8],
 [0.5, 0.6],
 [0.7, 0.9],
 [0.9, 0.8],
 [0.4, 0.8],
 [0.0, 0.2],
 [0.8, 0.1]]

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.9/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.7117158 , 0.7819734 , 0.46078986, 0.7632882 , 0.9983441 ,
        0.87600124],
       [0.8105186 , 0.0474633 , 0.42952222, 0.34084898, 0.4319815 ,
        0.8809404 ],
       [0.42782837, 0.8493887 , 0.8199676 , 0.5723619 , 0.8848299 ,
        0.14131516],
       [0.40969414, 0.36407202, 0.86129844, 0.87103117, 0.54137987,
        0.9451648 ],
       [0.94180036, 0.5700126 , 0.32170886, 0.9192251 , 0.36841327,
        0.99123365],
       [0.3243165 , 0.7712629 , 0.6265182 , 0.12021291, 0.85183316,
        0.03610802],
       [0.03077155, 0.1386311 , 0.7911339 , 0.6066651 , 0.17137617,
        0.6743576 ],
       [0.8729119 , 0.21536678, 0.8630131 , 0.35792142, 0.0314253 ,
        0.9223388 ]], dtype=float32)
In [11]:
math.reshaped_native(data, ['extra', data.shape], to_numpy=True)

/tmp/ipykernel_2203/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.7117158 , 0.7819734 , 0.46078986, 0.7632882 , 0.9983441 ,
        0.87600124, 0.8105186 , 0.0474633 , 0.42952222, 0.34084898,
        0.4319815 , 0.8809404 , 0.42782837, 0.8493887 , 0.8199676 ,
        0.5723619 , 0.8848299 , 0.14131516, 0.40969414, 0.36407202,
        0.86129844, 0.87103117, 0.54137987, 0.9451648 , 0.94180036,
        0.5700126 , 0.32170886, 0.9192251 , 0.36841327, 0.99123365,
        0.3243165 , 0.7712629 , 0.6265182 , 0.12021291, 0.85183316,
        0.03610802, 0.03077155, 0.1386311 , 0.7911339 , 0.6066651 ,
        0.17137617, 0.6743576 , 0.8729119 , 0.21536678, 0.8630131 ,
        0.35792142, 0.0314253 , 0.9223388 ]], 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_2203/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.9/x64/lib/python3.12/site-packages/phiml/math/_shape.py:2769: RuntimeWarning: Stacking shapes with incompatible item names will result in item names being lost. For vector Got ('x', 'y') and ('x',)
  warnings.warn(f"Stacking shapes with incompatible item names will result in item names 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]:
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)

/tmp/ipykernel_2203/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.243 ± 0.230 (4e-07...7e-01)
Final loss: (batchᵇ=100) 0.214 ± 0.159 (3e-08...5e-01)

Parameter Count¶

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