def sample(value, domain, batch_size=None, name=None):
assert isinstance(domain, Domain)
if isinstance(value, Field):
assert_same_rank(
value.rank, domain.rank,
'rank of value (%s) does not match domain (%s)' %
(value.rank, domain.rank))
if isinstance(value,
CenteredGrid) and value.box == domain.box and np.all(
value.resolution == domain.resolution):
data = value.data
else:
point_field = CenteredGrid.getpoints(domain.box,
domain.resolution)
point_field._batch_size = batch_size
data = value.at(point_field).data
else: # value is constant
if callable(value):
x = CenteredGrid.getpoints(
domain.box, domain.resolution).copied_with(
extrapolation=Material.extrapolation_mode(
domain.boundaries),
name=name)
value = value(x)
return value
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.add(
math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )), value)
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries),
name=name)
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 accessible(self, accessible):
if accessible is not None:
assert isinstance(accessible, CenteredGrid)
assert accessible.rank == self.domain.rank
return accessible
else:
return self.domain.centered_grid(1, extrapolation=Material.extrapolation_mode(self.domain.boundaries))
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 accessible_tensor(self, extend=0):
"""
Scalar channel encoding cells that are accessible, i.e. not solid, as ones and obstacles as zero.
:param extend: Extend the grid in all directions beyond the grid size specified by the domain
"""
pad_values = struct.map(lambda solid: int(not solid), Material.solid(self.domain.boundaries))
if isinstance(pad_values, (list, tuple)):
pad_values = [0] + list(pad_values) + [0]
result = math.pad(self.accessible.data, [[0,0]] + [[extend, extend]] * self.rank + [[0,0]], constant_values=pad_values)
return result
def active(self, active):
extrapolation = _active_extrapolation(
Material.extrapolation_mode(self.domain.boundaries))
if active is not None:
assert isinstance(active, CenteredGrid)
assert active.rank == self.domain.rank
if active.extrapolation != extrapolation:
active = active.copied_with(extrapolation=extrapolation)
return active
else:
return self.domain.centered_grid(1, extrapolation=extrapolation)
def solve(self, divergence, domain, guess, enable_backprop):
assert isinstance(domain, PoissonDomain)
fluid_mask = domain.accessible_tensor(extend=1)
extrapolation = Material.extrapolation_mode(domain.domain.boundaries)
def apply_A(pressure):
pressure = CenteredGrid(pressure, extrapolation=extrapolation)
pressure_padded = pressure.padded([[1, 1]] * pressure.rank)
return _weighted_sliced_laplace_nd(pressure_padded.data,
weights=fluid_mask)
return conjugate_gradient(divergence,
apply_A,
guess,
self.accuracy,
self.max_iterations,
back_prop=enable_backprop)
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
def solve_pressure_forward(divergence,
fluid_mask,
max_iterations,
guess,
accuracy,
domain,
back_prop=False):
from phi.physics.material import Material
extrapolation = Material.extrapolation_mode(domain.domain.boundaries)
def apply_A(pressure):
pressure = CenteredGrid(pressure, extrapolation=extrapolation)
pressure_padded = pressure.padded([[1, 1]] * pressure.rank)
return _weighted_sliced_laplace_nd(pressure_padded.data,
weights=fluid_mask)
return conjugate_gradient(divergence,
apply_A,
guess,
accuracy,
max_iterations,
back_prop=back_prop)
def sample(value, domain, batch_size=None):
assert isinstance(domain, Domain)
if isinstance(value, Field):
assert_same_rank(
value.rank, domain.rank,
'rank of value (%s) does not match domain (%s)' %
(value.rank, domain.rank))
if isinstance(value,
CenteredGrid) and value.box == domain.box and np.all(
value.resolution == domain.resolution):
data = value.data
else:
data = value.sample_at(
CenteredGrid.getpoints(domain.box, domain.resolution).data)
else: # value is constant
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )) + value
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries))
def solve(self, field, domain, guess, enable_backprop):
assert isinstance(domain, FluidDomain)
dimensions = list(field.shape[1:-1])
A = sparse_pressure_matrix(dimensions, domain.active_tensor(extend=1),
domain.accessible_tensor(extend=1),
Material.periodic(domain.domain.boundaries))
def np_solve_p(div):
div_vec = div.reshape([-1, A.shape[0]])
pressure = [
scipy.sparse.linalg.spsolve(A, div_vec[i, ...])
for i in range(div_vec.shape[0])
]
return np.array(pressure).reshape(div.shape).astype(np.float32)
def np_solve_p_gradient(op, grad_in):
return math.py_func(np_solve_p, [grad_in], np.float32, field.shape)
pressure = math.py_func(np_solve_p, [field],
np.float32,
field.shape,
grad=np_solve_p_gradient)
return pressure, None
def apply_A(pressure):
from phi.physics.material import Material
mode = 'replicate' if Material.solid(domain.domain.boundaries) else 'constant'
padded = math.pad(pressure, [[0,0]] + [[1,1]]*(math.ndims(pressure)-2) + [[0,0]], mode=mode)
return _weighted_sliced_laplace_nd(padded, weights=fluid_mask)