def solve(self, field, domain, guess):
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 np_solve_p_gradient(op, grad_in):
return math.py_func(np_solve_p, [grad_in], np.float32, field.shape)
def np_solve_p_gradient(op, grad_in):
return math.py_func(np_solve_p, [grad_in], np.float32,
divergence.shape)