def solve(self, field, domain, guess):
assert isinstance(domain, FluidDomain)
active_mask = domain.active_tensor(extend=1)
fluid_mask = domain.accessible_tensor(extend=1)
dimensions = math.staticshape(field)[1:-1]
N = int(np.prod(dimensions))
periodic = Material.periodic(domain.domain.boundaries)
if math.choose_backend([field, active_mask, fluid_mask]).matches_name('SciPy'):
A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask, periodic)
else:
sidx, sorting = sparse_indices(dimensions, periodic)
sval_data = sparse_values(dimensions, active_mask, fluid_mask, sorting, periodic)
backend = math.choose_backend(field)
sval_data = backend.cast(sval_data, field.dtype)
A = backend.sparse_tensor(indices=sidx, values=sval_data, shape=[N, N])
if self.autodiff:
return sparse_cg(field, A, self.max_iterations, guess, self.accuracy, back_prop=True)
else:
def pressure_gradient(op, grad):
return sparse_cg(grad, A, max_gradient_iterations, None, self.gradient_accuracy)[0]
pressure, iteration = math.with_custom_gradient(sparse_cg,
[field, A, self.max_iterations, guess, self.accuracy],
pressure_gradient, input_index=0, output_index=0,
name_base='scg_pressure_solve')
max_gradient_iterations = iteration if self.max_gradient_iterations == 'mirror' else self.max_gradient_iterations
return pressure, iteration
def solve(self, field, domain, guess, enable_backprop):
assert isinstance(domain, FluidDomain)
active_mask = domain.active_tensor(extend=1)
fluid_mask = domain.accessible_tensor(extend=1)
dimensions = math.staticshape(field)[1:-1]
N = int(np.prod(dimensions))
periodic = Material.periodic(domain.domain.boundaries)
if math.choose_backend([field, active_mask,
fluid_mask]).matches_name('SciPy'):
A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask,
periodic)
else:
sidx, sorting = sparse_indices(dimensions, periodic)
sval_data = sparse_values(dimensions, active_mask, fluid_mask,
sorting, periodic)
backend = math.choose_backend(field)
sval_data = backend.cast(sval_data, field.dtype)
A = backend.sparse_tensor(indices=sidx,
values=sval_data,
shape=[N, N])
div_vec = math.reshape(field, [-1, int(np.prod(field.shape[1:]))])
if guess is not None:
guess = math.reshape(guess, [-1, int(np.prod(field.shape[1:]))])
def apply_A(pressure):
return math.matmul(A, pressure)
result_vec, iterations = conjugate_gradient(div_vec, apply_A, guess,
self.accuracy,
self.max_iterations,
enable_backprop)
return math.reshape(result_vec, math.shape(field)), iterations
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_tensor_from_tuple_of_tensor_like(self):
native = ([1, 2, 3], math.zeros(channel(vector=3)))
for backend in BACKENDS:
with backend:
tens = wrap(native, batch(stack=2), channel(vector=3))
self.assertEqual(math.NUMPY, math.choose_backend(tens))
self.assertEqual(
batch(stack=2) & channel(vector=3), tens.shape)
tens = tensor(native, batch(stack=2), channel(vector=3))
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(
batch(stack=2) & channel(vector=3), tens.shape)
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_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 poisson_solve(input_field, poisson_domain, solver=None, guess=None):
"""
Solves the Poisson equation Δp = input_field for p.
:param input_field: CenteredGrid
:param poisson_domain: PoissonDomain instance
:param solver: PoissonSolver to use, None for default
:param guess: CenteredGrid with same size and resolution as input_field
:return: p as CenteredGrid, iteration count as int or None if not available
:rtype: CenteredGrid, int
"""
assert isinstance(input_field, CenteredGrid)
if guess is not None:
assert isinstance(guess, CenteredGrid)
assert guess.compatible(input_field)
guess = guess.data
if isinstance(poisson_domain, Domain):
poisson_domain = PoissonDomain(poisson_domain)
if solver is None:
from .sparse import SparseSciPy, SparseCG
if math.choose_backend([
input_field.data, poisson_domain.active.data,
poisson_domain.accessible.data
]).matches_name('SciPy'):
solver = SparseSciPy()
else:
from phi.physics.pressuresolver.fourier import FourierSolver
solver = FourierSolver() & SparseCG()
pressure, iteration = solver.solve(input_field.data,
poisson_domain,
guess=guess)
pressure = CenteredGrid(pressure,
input_field.box,
extrapolation=input_field.extrapolation,
name='pressure')
return pressure, iteration
def poisson_solve(input_field, poisson_domain, solver=None):
"""
Solves the Poisson equation Δp = input_field for p.
:param input_field: CenteredGrid
:param poisson_domain: PoissonDomain instance
:param solver: PoissonSolver to use, None for default
:return: p as CenteredGrid, iteration count as int or None if not available
:rtype: CenteredGrid, int
"""
from .sparse import SparseSciPy, SparseCG
assert isinstance(input_field, CenteredGrid)
if isinstance(poisson_domain, Domain):
poisson_domain = PoissonDomain(poisson_domain)
if solver is None:
if math.choose_backend([
input_field.data, poisson_domain.active.data,
poisson_domain.accessible.data
]).matches_name('SciPy'):
solver = SparseSciPy()
else:
solver = SparseCG()
pressure, iteration = solver.solve(input_field.data,
poisson_domain,
guess=None)
pressure = CenteredGrid(pressure, input_field.box, name='pressure')
return pressure, iteration
def solve(self, divergence, domain, pressure_guess):
assert isinstance(domain, FluidDomain)
active_mask = domain.active_tensor(extend=1)
fluid_mask = domain.accessible_tensor(extend=1)
dimensions = list(divergence.shape[1:-1])
N = int(np.prod(dimensions))
if math.choose_backend(divergence).matches_name('SciPy'):
A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask)
else:
sidx, sorting = sparse_indices(dimensions)
sval_data = sparse_values(dimensions, active_mask, fluid_mask,
sorting)
A = math.choose_backend(divergence).sparse_tensor(indices=sidx,
values=sval_data,
shape=[N, N])
if self.autodiff:
return sparse_cg(divergence,
A,
self.max_iterations,
pressure_guess,
self.accuracy,
back_prop=True)
else:
def pressure_gradient(op, grad):
return sparse_cg(grad, A, max_gradient_iterations, None,
self.gradient_accuracy)[0]
pressure, iteration = math.with_custom_gradient(
sparse_cg, [
divergence, A, self.max_iterations, pressure_guess,
self.accuracy
],
pressure_gradient,
input_index=0,
output_index=0,
name_base='scg_pressure_solve')
max_gradient_iterations = iteration if self.max_gradient_iterations == 'mirror' else self.max_gradient_iterations
return pressure, iteration
def general_sample_at(self, points, reduce):
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
result = general_grid_sample_nd(
self.data,
local_points,
boundary=_pad_mode(self.extrapolation),
constant_values=_pad_value(self.extrapolation_value),
math=math.choose_backend([self.data, points]),
reduce=reduce)
return result
def __dataop__(self, other, linear_if_scalar, data_operator):
if isinstance(other, Field):
assert self.compatible(other), 'Fields are not compatible: %s and %s' % (self, other)
flags = propagate_flags_operation(self.flags+other.flags, False, self.rank, self.component_count)
self_data = self.data if self.has_points else self.at(other).data
other_data = other.data if other.has_points else other.at(self).data
backend = math.choose_backend([self_data, other_data])
self_data_tensor = backend.as_tensor(self_data)
other_data_tensor = backend.as_tensor(other_data)
data = data_operator(self_data_tensor, other_data_tensor)
else:
flags = propagate_flags_operation(self.flags, linear_if_scalar, self.rank, self.component_count)
data = data_operator(self.data, other)
return self.copied_with(data=data, flags=flags)
def test_tensor_from_constant(self):
for backend in BACKENDS:
with backend:
for const in (1, 1.5, True, 1 + 1j):
tens = math.wrap(const)
self.assertEqual(math.NUMPY, tens.default_backend)
self.assertTrue(isinstance(tens.native(),
(int, float, bool, complex)),
msg=backend)
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 solve(self, divergence, domain, pressure_guess):
assert isinstance(domain, FluidDomain)
active_mask = domain.active_tensor(extend=1)
fluid_mask = domain.accessible_tensor(extend=1)
dimensions = list(divergence.shape[1:-1])
N = int(np.prod(dimensions))
if math.choose_backend(divergence).matches_name('TensorFlow'):
import tensorflow as tf
if tf.__version__[0] == '2':
print('Adjusting for tensorflow 2.0')
tf = tf.compat.v1
tf.disable_eager_execution()
sidx, sorting = sparse_indices(dimensions)
sval_data = sparse_values(dimensions, active_mask, fluid_mask,
sorting)
A = tf.SparseTensor(indices=sidx,
values=sval_data,
dense_shape=[N, N])
else:
A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask)
if self.autodiff:
return sparse_cg(divergence,
A,
self.max_iterations,
pressure_guess,
self.accuracy,
back_prop=True)
else:
def pressure_gradient(op, grad):
return sparse_cg(grad, A, max_gradient_iterations, None,
self.gradient_accuracy)[0]
pressure, iteration = math.with_custom_gradient(
sparse_cg, [
divergence, A, self.max_iterations, pressure_guess,
self.accuracy
],
pressure_gradient,
input_index=0,
output_index=0,
name_base='scg_pressure_solve')
max_gradient_iterations = iteration if self.max_gradient_iterations == 'mirror' else self.max_gradient_iterations
return pressure, iteration
def poisson_solve(input_field, poisson_domain, solver=None, guess=None, gradient='implicit'):
"""
Solves the Poisson equation Δp = input_field for p.
:param gradient: one of ('implicit', 'autodiff', 'inverse')
If 'autodiff', use the built-in autodiff for backpropagation.
The intermediate results of each loop iteration will be permanently stored if backpropagation is used.
If 'implicit', computes a forward pressure solve in reverse accumulation backpropagation.
This requires less memory but is only accurate if the solution is fully converged.
:param input_field: CenteredGrid
:param poisson_domain: PoissonDomain instance
:param solver: PoissonSolver to use, None for default
:param guess: CenteredGrid with same size and resolution as input_field
:return: p as CenteredGrid, iteration count as int or None if not available
:rtype: CenteredGrid, int
"""
assert isinstance(input_field, CenteredGrid)
if guess is not None:
assert isinstance(guess, CenteredGrid)
assert guess.compatible(input_field)
guess = guess.data
if isinstance(poisson_domain, Domain):
poisson_domain = PoissonDomain(poisson_domain)
if solver is None:
solver = _choose_solver(input_field.resolution, math.choose_backend([input_field.data, poisson_domain.active.data, poisson_domain.accessible.data]))
if not struct.any(Material.open(poisson_domain.domain.boundaries)): # has no open boundary
input_field = input_field - math.mean(input_field.data, axis=tuple(range(1, 1 + input_field.rank)), keepdims=True) # Subtract mean divergence
assert gradient in ('autodiff', 'implicit', 'inverse')
if gradient == 'autodiff':
pressure, iteration = solver.solve(input_field.data, poisson_domain, guess, enable_backprop=True)
else:
if gradient == 'implicit':
def poisson_gradient(_op, grad):
return poisson_solve(CenteredGrid.sample(grad, poisson_domain.domain), poisson_domain, solver, None, gradient=gradient)[0].data
else: # gradient = 'inverse'
def poisson_gradient(_op, grad):
return CenteredGrid.sample(grad, poisson_domain.domain).laplace(physical_units=False).data
pressure, iteration = math.with_custom_gradient(solver.solve, [input_field.data, poisson_domain, guess, False], poisson_gradient, input_index=0, output_index=0, name_base='poisson_solve')
pressure = CenteredGrid(pressure, input_field.box, extrapolation=input_field.extrapolation, name='pressure')
return pressure, iteration
