Writing n-dimensional Code with ΦML¶

Colab   •   🌐 ΦML   •   📖 Documentation   •   🔗 API   •   ▶ Videos   •   Examples

ΦML's dimension types allow you to write abstract code that scales with the number of spatial dimensions.

In [1]:
%%capture
!pip install phiml

from phiml import math
from phiml.math import spatial, channel, instance

Grids¶

Grids are a popular data structure that in n dimensions. In ΦML, each axis of the grid is represented by a spatial dimension.

In [2]:
grid_1d = math.random_uniform(spatial(x=5))
grid_2d = math.random_uniform(spatial(x=3, y=3))
grid_3d = math.random_uniform(spatial(x=16, y=16, z=16))

Note that the dimension names are arbitrary. We chose x, y, z for readability.

Now, let's write a function that outputs the mean of the direct neighbors of each cell. In 1D, this would be the stencil (.5, 0, .5) and in 2D (0, .25, 0; .25, 0, .25; 0, .25, 0).

In [3]:
def neighbor_mean(grid):
    left, right = math.shift(grid, (-1, 1), padding=math.extrapolation.PERIODIC)
    return math.mean([left, right], math.non_spatial)

This function uses math.shift() to access the left and right neighbor in each direction. By default, shift shifts in all spatial dimensions and lists the result along a new channel dimension. Then we can take the mean of the right and the left values to compute the mean of all neighbors.

We can now evaluate the function in 1D, 2D, 3D, etc. and it will automatically derive the correct stencil.

In [4]:
neighbor_mean(grid_1d)
Out[4]:
(0.391, 0.144, 0.181, 0.513, 0.172) along xˢ
In [5]:
neighbor_mean(grid_2d)
Out[5]:
(xˢ=3, yˢ=3) 0.605 ± 0.072 (4e-01...7e-01)
In [6]:
neighbor_mean(grid_3d)
Out[6]:
(xˢ=16, yˢ=16, zˢ=16) 0.503 ± 0.117 (1e-01...9e-01)

To make sure that the stencil is correct, we can look at the matrix representation of our function.

In [7]:
math.print(math.matrix_from_function(neighbor_mean, grid_1d)[0])

x=0     0.   0.5  0.   0.   0.5  along ~x
x=1     0.5  0.   0.5  0.   0.   along ~x
x=2     0.   0.5  0.   0.5  0.   along ~x
x=3     0.   0.   0.5  0.   0.5  along ~x
x=4     0.5  0.   0.   0.5  0.   along ~x
In [8]:
math.print(math.matrix_from_function(neighbor_mean, grid_2d)[0])

x&y=0     0.    0.25  0.25  0.25  0.    0.    0.25  0.    0.    along ~x&~y
x&y=1     0.25  0.    0.25  0.    0.25  0.    0.    0.25  0.    along ~x&~y
x&y=2     0.25  0.25  0.    0.    0.    0.25  0.    0.    0.25  along ~x&~y
x&y=3     0.25  0.    0.    0.    0.25  0.25  0.25  0.    0.    along ~x&~y
x&y=4     0.    0.25  0.    0.25  0.    0.25  0.    0.25  0.    along ~x&~y
x&y=5     0.    0.    0.25  0.25  0.25  0.    0.    0.    0.25  along ~x&~y
x&y=6     0.25  0.    0.    0.25  0.    0.    0.    0.25  0.25  along ~x&~y
x&y=7     0.    0.25  0.    0.    0.25  0.    0.25  0.    0.25  along ~x&~y
x&y=8     0.    0.    0.25  0.    0.    0.25  0.25  0.25  0.    along ~x&~y

The same principle holds for all grid functions in the phiml.math library. For example, if we perform a Fourier transform, the algorithm will be selected based on the number of spatial dimensions. A 1D FFT will always be performed on our 1D grid, even if we add additional non-spatial dimensions.

In [9]:
math.fft(grid_1d)  # 1D FFT
Out[9]:
((1.4010619+0j), (-0.23663189+0.7188296j), (-0.43180037+0.36949766j), (-0.43180037-0.36949766j), (-0.23663189-0.7188296j)) along xˢ complex64
In [10]:
math.fft(grid_2d)  # 2D FFT
Out[10]:
(xˢ=3, yˢ=3) complex64 |...| < 5.443326950073242
In [11]:
math.fft(grid_3d)  # 3D FFT
Out[11]:
(xˢ=16, yˢ=16, zˢ=16) complex64 |...| < 2061.266845703125

Dimensions as Components¶

Not all applications involving physical space use grids to represent data. Take point clouds or particles for instance. In these cases, we would represent the dimensionality not by the number of spatial dimensions but by the number of vector components.

In [12]:
points_1d = math.random_uniform(instance(points=4), channel(vector='x'))
points_2d = math.random_uniform(instance(points=4), channel(vector='x,y'))
points_3d = math.random_uniform(instance(points=4), channel(vector='x,y,z'))

In these cases, the generalization to n dimensions is usually trivial. Take the following function that computes the pairwise distances.

In [13]:
def pairwise_distances(x):
    return math.vec_length(math.rename_dims(x, 'points', 'others') - x)

Here, we compute the distances between each pair of particles on a matrix with dimensions points and others. The intermediate matrix of position distances inherits the vector dimension from x and math.vec_length() sums all components. Consequently, this function computes the correct distances in 1D, 2D and 3D.

In [14]:
pairwise_distances(points_1d)

/tmp/ipykernel_2894/2630121943.py:2: DeprecationWarning: phiml.math.length is deprecated in favor of phiml.math.norm
  return math.vec_length(math.rename_dims(x, 'points', 'others') - x)
Out[14]:
(othersⁱ=4, pointsⁱ=4) 0.355 ± 0.284 (0e+00...8e-01)
In [15]:
pairwise_distances(points_2d)

/tmp/ipykernel_2894/2630121943.py:2: DeprecationWarning: phiml.math.length is deprecated in favor of phiml.math.norm
  return math.vec_length(math.rename_dims(x, 'points', 'others') - x)
Out[15]:
(othersⁱ=4, pointsⁱ=4) 0.250 ± 0.179 (0e+00...5e-01)
In [16]:
pairwise_distances(points_3d)

/tmp/ipykernel_2894/2630121943.py:2: DeprecationWarning: phiml.math.length is deprecated in favor of phiml.math.norm
  return math.vec_length(math.rename_dims(x, 'points', 'others') - x)
Out[16]:
(othersⁱ=4, pointsⁱ=4) 0.463 ± 0.346 (0e+00...9e-01)

Further Reading¶

Here, we focussed on spatial dimensions, but each dimension type plays a unique role in ΦML.

The library ΦFlow uses ΦML to implement an n-dimensional incompressible fluid solver.

🌐 ΦML   •   📖 Documentation   •   🔗 API   •   ▶ Videos   •   Examples