def __add__(self, other):
if other is 0:
return self
assert isinstance(other, Gravity), type(other)
if self._batch_size is not None:
assert self._batch_size == other._batch_size
# Add gravity
if math.is_scalar(self.gravity) and math.is_scalar(other.gravity):
return Gravity(self.gravity + other.gravity)
else:
rank = math.staticshape(other.gravity)[-1] if math.is_scalar(self.gravity)\
else math.staticshape(self.gravity)[-1]
sum_tensor = gravity_tensor(self, rank) + gravity_tensor(
other, rank)
return Gravity(sum_tensor)
def gravity_tensor(gravity, rank):
if isinstance(gravity, Gravity):
gravity = gravity.gravity
if math.is_scalar(gravity):
gravity = gravity * GLOBAL_AXIS_ORDER.up_vector(rank)
assert math.staticshape(gravity)[-1] == rank
return math.to_float(
math.expand_dims(gravity, 0,
rank + 2 - len(math.staticshape(gravity))))
def __init__(self,
accuracy=1e-5,
gradient_accuracy='same',
max_iterations=2000,
max_gradient_iterations='same',
autodiff=False):
"""
Conjugate gradient solver using sparse matrix multiplications.
:param accuracy: the maximally allowed error for each cell, measured in terms of field values.
:param gradient_accuracy: accuracy applied during backpropagation, number of 'same' to use forward accuracy
:param max_iterations: integer specifying maximum conjugent gradient loop iterations or None for no limit
:param max_gradient_iterations: maximum loop iterations during backpropagation,
'same' uses the number from max_iterations,
'mirror' sets the maximum to the number of iterations that were actually performed in the forward pass
:param autodiff: If autodiff=True, use the built-in autodiff for backpropagation.
The intermediate results of each loop iteration will be permanently stored if backpropagation is used.
If False, replaces autodiff by a forward pressure solve in reverse accumulation backpropagation.
This requires less memory but is only accurate if the solution is fully converged.
"""
PoissonSolver.__init__(self,
'Sparse Conjugate Gradient',
supported_devices=('CPU', 'GPU'),
supports_guess=True,
supports_loop_counter=True,
supports_continuous_masks=True)
assert math.is_scalar(accuracy), 'invalid accuracy: %s' % accuracy
assert gradient_accuracy == 'same' or math.is_scalar(
gradient_accuracy
), 'invalid gradient_accuracy: %s' % gradient_accuracy
assert max_gradient_iterations in ['same', 'mirror'] or isinstance(
max_gradient_iterations, int
), 'invalid max_gradient_iterations: %s' % max_gradient_iterations
self.accuracy = accuracy
self.gradient_accuracy = accuracy if gradient_accuracy == 'same' else gradient_accuracy
self.max_iterations = max_iterations
if max_gradient_iterations == 'same':
self.max_gradient_iterations = max_iterations
elif max_gradient_iterations == 'mirror':
self.max_gradient_iterations = 'mirror'
else:
self.max_gradient_iterations = max_gradient_iterations
assert not autodiff, 'Cannot specify max_gradient_iterations when autodiff=True'
self.autodiff = autodiff
def gravity_tensor(gravity, rank):
if isinstance(gravity, Gravity):
gravity = gravity.gravity
if math.is_scalar(gravity):
return math.to_float(
math.expand_dims([gravity] + [0] * (rank - 1), 0, rank + 1))
else:
assert math.staticshape(gravity)[-1] == rank
return math.to_float(
math.expand_dims(gravity, 0,
rank + 2 - len(math.staticshape(gravity))))
def __init__(self, accuracy=1e-5, max_iterations=2000):
"""
Conjugate gradient solver that geometrically calculates laplace pressure in each iteration.
Unlike most other solvers, this algorithm is TPU compatible but usually performs worse than SparseCG.
Obstacles are allowed to vary between examples but the same number of iterations is performed for each example in one batch.
:param accuracy: the maximally allowed error on the divergence channel for each cell
:param max_iterations: integer specifying maximum conjugent gradient loop iterations or None for no limit
:param autodiff:
"""
PoissonSolver.__init__(self,
'Single-Phase Conjugate Gradient',
supported_devices=('CPU', 'GPU', 'TPU'),
supports_guess=True,
supports_loop_counter=True,
supports_continuous_masks=True)
assert math.is_scalar(accuracy), 'invalid accuracy: %s' % accuracy
self.accuracy = accuracy
self.max_iterations = max_iterations
