Fix TensorFlow CUDA poisson solver
@@ -5,8 +5,7 @@ import tensorflow as tf
from numbers import Number
from phi import math
-from phi.physics.pressuresolver.solver_api import PressureSolver
-
+from phi.physics.pressuresolver.solver_api import PoissonSolver
if tf.__version__[0] == '2':
logging.info('Adjusting for tensorflow 2.0')
@@ -14,19 +13,18 @@ if tf.__version__[0] == '2':
tf.disable_eager_execution()
-# Load Custom Ops
+# --- Load Custom Ops ---
current_dir = os.path.dirname(os.path.realpath(__file__))
kernel_path = os.path.join(current_dir, 'cuda/build/pressure_solve_op.so')
assert os.path.isfile(kernel_path), 'CUDA binaries not found at %s. Run "python setup.py cuda" to compile them' % kernel_path
pressure_op = tf.load_op_library(kernel_path)
-class CUDASolver(PressureSolver):
+class CUDASolver(PoissonSolver):
def __init__(self, accuracy=1e-5, gradient_accuracy='same',
max_iterations=2000, max_gradient_iterations='same'):
- PressureSolver.__init__(self, 'CUDA Conjugate Gradient', supported_devices=('GPU',),
- supports_loop_counter=True, supports_guess=False, supports_continuous_masks=False)
+ PoissonSolver.__init__(self, 'CUDA Conjugate Gradient', supported_devices=('GPU',), supports_loop_counter=True, supports_guess=False, supports_continuous_masks=False)
assert isinstance(accuracy, Number), 'invalid accuracy: %s' % accuracy
assert gradient_accuracy == 'same' or isinstance(gradient_accuracy, Number), 'invalid gradient_accuracy: %s' % gradient_accuracy
assert max_gradient_iterations in ['same', 'mirror'] or isinstance(max_gradient_iterations, Number), 'invalid max_gradient_iterations: %s' % max_gradient_iterations
@@ -40,8 +38,10 @@ class CUDASolver(PressureSolver):
else:
self.max_gradient_iterations = max_gradient_iterations
- def solve(self, divergence, domain, pressure_guess):
- # pressure_guess: not used in this implementation, Kernel takes the last pressure value for initial_guess
+ def solve(self, divergence, domain, guess):
+ """
+ :param guess: not used in this implementation, Kernel takes the last pressure value for initial_guess
+ """
active, accessible = domain.active_tensor(extend=1), domain.accessible_tensor(extend=1)
def pressure_gradient(op, grad):
@@ -59,7 +59,7 @@ class CUDASolver(PressureSolver):
def cuda_solve_forward(divergence, active_mask, fluid_mask, accuracy, max_iterations):
# Setup
- dimensions = divergence.get_shape()[1:-1]
+ dimensions = math.staticshape(divergence)[1:-1]
dimensions = dimensions[::-1] # the custom op needs it in the x,y,z order
dim_array = np.array(dimensions)
dim_product = np.prod(dimensions)
