def test_functional_gradient(self):
def f(x: math.Tensor, y: math.Tensor):
assert isinstance(x, math.Tensor)
assert isinstance(y, math.Tensor)
pred = x
loss = math.l2_loss(pred - y)
return loss, pred
for backend in BACKENDS:
if backend.supports(Backend.functional_gradient):
with backend:
x_data = math.tensor(2.)
y_data = math.tensor(1.)
dx, = math.functional_gradient(f,
wrt=[0],
get_output=False)(x_data,
y_data)
math.assert_close(dx, 1, msg=backend.name)
dx, dy = math.functional_gradient(f, [0, 1],
get_output=False)(x_data,
y_data)
math.assert_close(dx, 1, msg=backend.name)
math.assert_close(dy, -1, msg=backend.name)
(loss, pred), (dx, dy) = math.functional_gradient(
f, [0, 1], get_output=True)(x_data, y_data)
math.assert_close(loss, 0.5, msg=backend.name)
math.assert_close(pred, x_data, msg=backend.name)
math.assert_close(dx, 1, msg=backend.name)
math.assert_close(dy, -1, msg=backend.name)
def test_batch_independence(self):
def simulate(centers):
world = World()
fluid = world.add(Fluid(Domain(x=5,
y=4,
boundaries=CLOSED,
bounds=Box(0, [40, 32])),
buoyancy_factor=0.1,
batch_size=centers.shape[0]),
physics=IncompressibleFlow())
world.add(Inflow(Sphere(center=centers, radius=3), rate=0.2))
world.add(
Fan(Sphere(center=centers, radius=5), acceleration=[1.0, 0]))
world.step(dt=1.5)
world.step(dt=1.5)
world.step(dt=1.5)
assert not math.close(fluid.density.values, 0)
print()
return fluid.density.values.batch[0], fluid.velocity.values.batch[
0]
d1, v1 = simulate(math.tensor([[5, 16], [5, 4]], names='batch,vector'))
d2, v2 = simulate(math.tensor([[5, 16], [5, 16]],
names='batch,vector'))
math.assert_close(d1, d2)
math.assert_close(v1, v2)
def test_np_speed_sum(self):
print()
np1, np2 = rnpv(64), rnpv(256)
t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector'))
t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector'))
_assert_equally_fast(lambda: np.sum(np1), lambda: math.sum(t1, dim=t1.shape), n=10000)
_assert_equally_fast(lambda: np.sum(np2), lambda: math.sum(t2, dim=t1.shape), n=10000)
def points(self,
points: Tensor or Number or tuple or list,
values: Tensor or Number = None,
radius: Tensor or float or int or None = None,
extrapolation: math.Extrapolation = math.extrapolation.ZERO,
color: str or Tensor or tuple or list or None = None) -> PointCloud:
"""
Create a `phi.field.PointCloud` from the given `points`.
The created field has no channel dimensions and all points carry the value `1`.
Args:
points: point locations in physical units
values: (optional) values of the particles, defaults to 1.
radius: (optional) size of the particles
extrapolation: (optional) extrapolation to use, defaults to extrapolation.ZERO
color: (optional) color used when plotting the points
Returns:
`phi.field.PointCloud` object
"""
extrapolation = extrapolation if isinstance(extrapolation, math.Extrapolation) else self.boundaries[extrapolation]
if radius is None:
radius = math.mean(self.bounds.size) * 0.005
# --- Parse points: tuple / list ---
if isinstance(points, (tuple, list)):
if len(points) == 0: # no points
points = math.zeros(instance(points=0), channel(vector=1))
elif isinstance(points[0], Number): # single point
points = math.tensor([points], instance('points'), channel('vector'))
else:
points = math.tensor(points, instance('points'), channel('vector'))
elements = Sphere(points, radius)
if values is None:
values = math.tensor(1.)
return PointCloud(elements, values, extrapolation, add_overlapping=False, bounds=self.bounds, color=color)
def test_multi_dim_tensor_from_numpy(self):
v = math.tensor(np.ones([1, 4, 3, 2]), batch('batch'), spatial('x,y'),
channel('vector'))
self.assertEqual((1, 4, 3, 2), v.shape.sizes)
v = math.tensor(np.ones([10, 4, 3, 2]), batch('batch'), spatial('x,y'),
channel('vector'))
self.assertEqual((10, 4, 3, 2), v.shape.sizes)
def grid_sample(self,
resolution: math.Shape,
size,
shape: math.Shape = None):
shape = (self._shape if shape is None else shape) & resolution
for dim in channel(self._shape):
if dim.item_names[0] is None:
warnings.warn(
f"Please provide item names for Noise dim {dim} using {dim}='x,y,z'",
FutureWarning)
shape &= channel(**{dim.name: resolution.names})
rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(
random_normal(shape)) # Note: there is no complex32
with math.NUMPY:
k = math.fftfreq(resolution) * resolution / math.tensor(
size) * math.tensor(self.scale) # in physical units
k = math.vec_squared(k)
lowest_frequency = 0.1
weight_mask = math.to_float(k > lowest_frequency)
# --- Compute 1/k ---
k._native[(0, ) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k._native[(0, ) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k**self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array, dim=array.shape.non_batch)
array -= math.mean(array, dim=array.shape.non_batch)
array = math.to_float(array)
return array
def test_extrapolate_valid_3D_3x3x3_2(self):
valid = tensor([[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[0, 0, 1],
[0, 0, 0],
[1, 0, 0]],
[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]]], spatial('x, y, z'))
values = tensor([[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[1, 0, 4],
[0, 0, 0],
[2, 0, 0]],
[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]]], spatial('x, y, z'))
expected_valid = math.ones(spatial(x=3, y=3, z=3))
expected_values = tensor([[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]],
[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]],
[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]]], spatial('x, y, z'))
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid, 2)
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def test_cos(self):
for backend in BACKENDS:
with backend:
math.assert_close(math.cos(math.zeros(spatial(x=4))), 1, abs_tolerance=1e-6, msg=backend.name)
math.assert_close(math.cos(math.tensor(math.PI / 2)), 0, abs_tolerance=1e-6, msg=backend.name)
math.assert_close(math.cos(math.tensor(math.PI)), -1, abs_tolerance=1e-6, msg=backend.name)
math.assert_close(math.cos(math.tensor(math.PI * 3 / 2)), 0, abs_tolerance=1e-6, msg=backend.name)
def test_hessian(self):
def f(x, y):
return math.l1_loss(x**2 * y), x, y
eval_hessian = math.hessian(f,
wrt=[
0,
],
get_output=True,
get_gradient=True,
dim_suffixes=('1', '2'))
for backend in BACKENDS:
if backend.supports(Backend.hessian):
with backend:
x = math.tensor([(0.01, 1, 2)], channel('vector', 'v'))
y = math.tensor([1., 2.], batch('batch'))
(L, x, y), g, H, = eval_hessian(x, y)
math.assert_close(L, (5.0001, 10.0002), msg=backend.name)
math.assert_close(g.batch[0].vector[0], (0.02, 2, 4),
msg=backend.name)
math.assert_close(g.batch[1].vector[0], (0.04, 4, 8),
msg=backend.name)
math.assert_close(2,
H.v1[0].v2[0].batch[0],
H.v1[1].v2[1].batch[0],
H.v1[2].v2[2].batch[0],
msg=backend.name)
math.assert_close(4,
H.v1[0].v2[0].batch[1],
H.v1[1].v2[1].batch[1],
H.v1[2].v2[2].batch[1],
msg=backend.name)
def test_closest_grid_values_1d(self):
grid = math.tensor([0, 1, 2, 3], names='x')
coords = math.tensor([[0.1], [1.9], [0.5], [3.1]], names='x,vector')
closest = math.closest_grid_values(grid, coords, extrapolation.ZERO)
math.assert_close(
closest,
math.tensor([(0, 1), (1, 2), (0, 1), (3, 0)], names='x,closest_x'))
def test_cumulative_sum(self):
t = math.tensor([(0, 1, 2, 3), (-1, -2, -3, -4)], spatial('y,x'))
for backend in BACKENDS:
with backend:
t_ = math.tensor(t)
x_ = math.cumulative_sum(t_, 'x')
math.assert_close(x_, [(0, 1, 3, 6), (-1, -3, -6, -10)], msg=backend.name)
y_ = math.cumulative_sum(t_, t.shape[0])
math.assert_close(y_, [(0, 1, 2, 3), (-1, -1, -1, -1)], msg=backend.name)
def test_np_speed_op2(self):
print()
np1, np2 = rnpv(64), rnpv(64)
t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector'))
t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector'))
_assert_equally_fast(lambda: np1 + np2, lambda: t1 + t2, n=10000)
np1, np2 = rnpv(256), rnpv(256)
t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector'))
t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector'))
_assert_equally_fast(lambda: np1 + np2, lambda: t1 + t2, n=1000)
def test_tensor_from_constant(self):
for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
with backend:
for const in (1, 1.5, True, 1+1j):
tens = math.tensor(const, convert=False)
self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens))
math.assert_close(tens, const)
tens = math.tensor(const)
self.assertEqual(backend, math.choose_backend(tens), f'{const} was not converted to the specified backend')
math.assert_close(tens, const)
def test_tensor_from_tuple_of_numbers(self):
data_tuple = (1, 2, 3)
for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
with backend:
tens = math.tensor(data_tuple, convert=False)
self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens))
math.assert_close(tens, data_tuple)
tens = math.tensor(data_tuple)
self.assertEqual(backend, math.choose_backend(tens))
math.assert_close(tens, data_tuple)
def test_tensor_from_tuple_of_tensor_like(self):
native = ((1, 2, 3), math.zeros(vector=3))
for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
with backend:
tens = math.tensor(native, names=['stack', 'vector'], convert=False)
self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens))
self.assertEqual(shape(stack=2, vector=3), tens.shape)
tens = math.tensor(native, names=['stack', 'vector'])
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(shape(stack=2, vector=3), tens.shape)
def test_tensor_from_native(self):
for creation_backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
native = creation_backend.ones((4,))
for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
with backend:
tens = math.tensor(native, convert=False)
self.assertEqual(creation_backend, math.choose_backend(tens))
math.assert_close(tens, native)
tens = math.tensor(native)
self.assertEqual(backend, math.choose_backend(tens), f'Conversion failed from {creation_backend} to {backend}')
math.assert_close(tens, native)
def test_boolean_mask_batched(self):
for backend in BACKENDS:
with backend:
x = math.expand(math.range(spatial('x'), 4), batch(batch=2)) * math.tensor([1, -1])
mask = math.tensor([[True, False, True, False], [False, True, False, False]], batch('batch'), spatial('x'))
selected = math.boolean_mask(x, 'x', mask)
expected_0 = math.tensor([(0, -0), (2, -2)], spatial('x'), channel('vector'))
expected_1 = math.tensor([(1, -1)], spatial('x'), channel('vector'))
math.assert_close(selected.batch[0], expected_0, msg=backend.name)
math.assert_close(selected.batch[1], expected_1, msg=backend.name)
math.assert_close(selected, x.x[mask], msg=backend.name)
def test_extrapolate_valid(self):
valid = tensor([[0, 0, 0], [0, 1, 1], [1, 0, 0]], 'x, y')
values = tensor([[1, 0, 0], [0, 4, 0], [2, 0, 0]], 'x, y')
expected_valid = tensor([[0, 1, 1], [1, 1, 1], [1, 1, 1]], 'x, y')
expected_values = tensor([[1, 4, 0], [3, 4, 0], [2, 3, 0]], 'x, y')
new_values, new_valid = math.extrapolate_valid_values(values, valid, 1)
self.assertTrue(new_values == expected_values)
self.assertTrue(new_valid == expected_valid)
def test_grid_sample_gradient_1d(self):
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
grid = math.tensor([0., 1, 2, 3], spatial('x'))
coords = math.tensor([0.5, 1.5], instance('points'))
with math.record_gradients(grid, coords):
sampled = math.grid_sample(grid, coords, extrapolation.ZERO)
loss = math.mean(math.l2_loss(sampled)) / 2
grad_grid, grad_coords = math.gradients(loss, grid, coords)
math.assert_close(grad_grid, math.tensor([0.125, 0.5, 0.375, 0], spatial('x')), msg=backend)
math.assert_close(grad_coords, math.tensor([0.25, 0.75], instance('points')), msg=backend)
def test__periodic_2d_arakawa_poisson_bracket(self):
"""test _periodic_2d_arakawa_poisson_bracket implementation"""
with math.precision(64):
# Define parameters to test
test_params = {
'grid_size': [(4, 4), (32, 32)],
'dx': [0.1, 1],
'gen_func': [
lambda grid_size: np.random.rand(*grid_size).reshape(
grid_size)
]
}
# Generate test cases as the product
test_cases = [
dict(zip(test_params, v))
for v in product(*test_params.values())
]
for params in test_cases:
grid_size = params['grid_size']
d1 = params['gen_func'](grid_size)
d2 = params['gen_func'](grid_size)
dx = params['dx']
padding = extrapolation.PERIODIC
ref = self.arakawa_reference_implementation(
np.pad(d1.copy(), 1, mode='wrap'),
np.pad(d2.copy(), 1, mode='wrap'), dx)[1:-1, 1:-1]
d1_tensor = field.CenteredGrid(
values=math.tensor(d1, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
d2_tensor = field.CenteredGrid(
values=math.tensor(d2, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
val = math._nd._periodic_2d_arakawa_poisson_bracket(
d1_tensor.values, d2_tensor.values, dx)
try:
math.assert_close(ref,
val,
rel_tolerance=1e-14,
abs_tolerance=1e-14)
except BaseException as e:
abs_error = math.abs(val - ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(ref)
print(val)
raise AssertionError(e, params, variation_str)
def _op2(self, other, operator) -> 'SampledField':
if isinstance(other, Field):
other_values = reduce_sample(other, self._elements)
values = operator(self._values, other_values)
extrapolation_ = operator(self._extrapolation, other.extrapolation)
return self.with_values(values).with_extrapolation(extrapolation_)
else:
if isinstance(other, (tuple, list)) and len(other) == self.spatial_rank:
other = math.tensor(other, self.points.shape['vector'])
else:
other = math.tensor(other)
values = operator(self._values, other)
return self.with_values(values)
def test_tensor_from_tensor(self):
ref = math.batch_stack([math.zeros(x=5), math.zeros(x=4)], 'stack')
for backend in (NUMPY_BACKEND, TORCH_BACKEND, TF_BACKEND):
with backend:
tens = math.tensor(ref, convert=False)
self.assertEqual(math.NUMPY_BACKEND, math.choose_backend(tens))
self.assertEqual(2, tens.shape.get_size('stack'))
self.assertEqual(('stack', 'x'), tens.shape.names)
tens = math.tensor(ref)
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(backend, math.choose_backend(tens.stack[0]))
self.assertEqual(backend, math.choose_backend(tens.stack[1]))
tens = math.tensor(ref, names=('n1', 'n2'))
self.assertEqual(backend, math.choose_backend(tens))
def test_tensor_from_native(self):
for creation_backend in BACKENDS:
native = creation_backend.ones((4, ))
for backend in BACKENDS:
with backend:
tens = math.tensor(native, convert=False)
self.assertEqual(creation_backend, tens.default_backend)
math.assert_close(tens, native)
tens = math.tensor(native)
self.assertEqual(
backend, tens.default_backend,
f'Conversion failed from {creation_backend} to {backend}'
)
math.assert_close(tens, native)
def test_tensor_from_tensor(self):
ref = math.stack([math.zeros(spatial(x=5)),
math.zeros(spatial(x=4))], batch('stack'))
for backend in BACKENDS:
with backend:
tens = math.tensor(ref, convert=False)
self.assertEqual(math.NUMPY, math.choose_backend(tens))
self.assertEqual(2, tens.shape.get_size('stack'))
self.assertEqual(('stack', 'x'), tens.shape.names)
tens = math.tensor(ref)
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(backend, math.choose_backend(tens.stack[0]))
self.assertEqual(backend, math.choose_backend(tens.stack[1]))
tens = math.tensor(ref, batch('n1', 'n2'))
self.assertEqual(backend, math.choose_backend(tens))
def test_slice_by_item_name(self):
t = math.tensor(spatial(x=4, y=3))
math.assert_close(t.dims['x'], 4)
math.assert_close(t.dims['y'], 3)
math.assert_close(t.dims['y,x'], (3, 4))
math.assert_close(t.dims[('y', 'x')], (3, 4))
math.assert_close(t.dims[spatial('x,y')], (4, 3))
def test_boolean_mask_dim_missing(self):
for backend in BACKENDS:
with backend:
x = math.random_uniform(spatial(x=2))
mask = math.tensor([True, False, True, True], batch('selection'))
selected = math.boolean_mask(x, 'selection', mask)
self.assertEqual(set(spatial(x=2) & batch(selection=3)), set(selected.shape), msg=backend.name)
def test_grid_sample(self):
for backend in BACKENDS:
with backend:
grid = math.sum(math.meshgrid(x=[1, 2, 3], y=[0, 3]), 'vector') # 1 2 3 | 4 5 6
coords = math.tensor([(0, 0), (0.5, 0), (0, 0.5), (-2, -1)], instance('list'), channel('vector'))
interp = math.grid_sample(grid, coords, extrapolation.ZERO)
math.assert_close(interp, [1, 1.5, 2.5, 0], msg=backend.name)
def read_single_field(file: str, convert_to_backend=True) -> SampledField:
stored = np.load(file, allow_pickle=True)
ftype = stored['field_type']
implemented_types = ('CenteredGrid', 'StaggeredGrid')
if ftype in implemented_types:
data = stored['data']
dim_item_names = stored.get('dim_item_names',
(None, ) * len(data.shape))
data = NativeTensor(
data,
math.Shape(data.shape, tuple(stored['dim_names']),
tuple(stored['dim_types']), tuple(dim_item_names)))
if convert_to_backend:
data = math.tensor(data, convert=convert_to_backend)
bounds_item_names = stored.get('bounds_item_names', (None, ) *
len(stored['lower'] + stored['upper']))
lower = math.wrap(
stored['lower'], math.channel(vector=tuple(bounds_item_names))
) if stored['lower'].ndim > 0 else math.wrap(stored['lower'])
upper = math.wrap(stored['upper'],
math.channel(vector=tuple(bounds_item_names)))
extrapolation = math.extrapolation.from_dict(
stored['extrapolation'][()])
if ftype == 'CenteredGrid':
return CenteredGrid(data,
bounds=geom.Box(lower, upper),
extrapolation=extrapolation)
elif ftype == 'StaggeredGrid':
data_ = unstack_staggered_tensor(data, extrapolation)
return StaggeredGrid(data_,
bounds=geom.Box(lower, upper),
extrapolation=extrapolation)
raise NotImplementedError(f"{ftype} not implemented ({implemented_types})")
def test_extrapolate_valid_3x3(self):
valid = tensor([[0, 0, 0],
[0, 0, 1],
[1, 0, 0]], spatial('x, y'))
values = tensor([[1, 0, 2],
[0, 0, 4],
[2, 0, 0]], spatial('x, y'))
expected_valid = tensor([[0, 1, 1],
[1, 1, 1],
[1, 1, 1]], spatial('x, y'))
expected_values = tensor([[1, 4, 4],
[2, 3, 4],
[2, 3, 4]], spatial('x, y'))
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid)
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def native_call(f,
*inputs,
channels_last=None,
channel_dim='vector',
extrapolation=None) -> SampledField or Tensor:
"""
Similar to `phi.math.native_call()`.
Args:
f: Function to be called on native tensors of `inputs.values`.
The function output must have the same dimension layout as the inputs and the batch size must be identical.
*inputs: `SampledField` or `phi.Tensor` instances.
extrapolation: (Optional) Extrapolation of the output field. If `None`, uses the extrapolation of the first input field.
Returns:
`SampledField` matching the first `SampledField` in `inputs`.
"""
input_tensors = [
i.values if isinstance(i, SampledField) else tensor(i) for i in inputs
]
values = math.native_call(f,
*input_tensors,
channels_last=channels_last,
channel_dim=channel_dim)
for i in inputs:
if isinstance(i, SampledField):
result = i.with_values(values=values)
if extrapolation is not None:
result = result.with_extrapolation(extrapolation)
return result
else:
raise AssertionError("At least one input must be a SampledField.")
