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 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)
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
