def cg_loop(x, dx, dy, residual, iterations):
dx_dy = math.sum(dx * dy, axis=non_batch_dims, keepdims=True)
step_size = math.divide_no_nan(math.sum(dx * residual, axis=non_batch_dims, keepdims=True), dx_dy)
x += step_size * dx
residual -= step_size * dy
dx = residual - math.divide_no_nan(math.sum(residual * dy, axis=non_batch_dims, keepdims=True) * dx, dx_dy)
dy = function(dx)
return [x, dx, dy, residual, iterations + 1]
def l1_loss(tensor, batch_norm=True, reduce_batches=True):
if struct.isstruct(tensor):
all_tensors = struct.flatten(tensor)
return sum(l1_loss(tensor, batch_norm, reduce_batches) for tensor in all_tensors)
if reduce_batches:
total_loss = math.sum(math.abs(tensor))
else:
total_loss = math.sum(math.abs(tensor), axis=list(range(1, len(tensor.shape))))
if batch_norm and reduce_batches:
batch_size = math.shape(tensor)[0]
return math.div(total_loss, math.to_float(batch_size))
else:
return total_loss
def normalize_to(target, source=1, epsilon=1e-5, batch_dims=1):
"""
Multiplies the target so that its total content matches the source.
:param target: a tensor
:param source: a tensor or number
:param epsilon: small number to prevent division by zero or None.
:return: normalized tensor of the same shape as target
"""
target_total = math.sum(target, axis=tuple(range(batch_dims, math.ndims(target))), keepdims=True)
denominator = math.maximum(target_total, epsilon) if epsilon is not None else target_total
source_total = math.sum(source, axis=tuple(range(batch_dims, math.ndims(source))), keepdims=True)
return target * (source_total / denominator)
def blur(field, radius, cutoff=None, kernel="1/1+x"):
"""
Warning: This function can cause NaN in the gradients, reason unknown.
Runs a blur kernel over the given tensor.
:param field: tensor
:param radius: weight function curve scale
:param cutoff: kernel size
:param kernel: Type of blur kernel (str). Must be in ('1/1+x', 'gauss')
:return:
"""
if cutoff is None:
cutoff = min(int(round(radius * 3)), *field.shape[1:-1])
xyz = np.meshgrid(
*[range(-int(cutoff), (cutoff) + 1) for _ in field.shape[1:-1]])
d = math.to_float(np.sqrt(np.sum([x**2 for x in xyz], axis=0)))
if kernel == "1/1+x":
weights = math.to_float(1) / (d / radius + 1)
elif kernel.lower() == "gauss":
weights = math.exp(-d / radius / 2)
else:
raise ValueError("Unknown kernel: %s" % kernel)
weights /= math.sum(weights)
weights = math.reshape(weights, list(weights.shape) + [1, 1])
return math.conv(field, weights)
def fftfreq(resolution, mode='vector', dtype=None):
"""
Returns the discrete Fourier transform sample frequencies.
These are the frequencies corresponding to the components of the result of `math.fft` on a tensor of shape `resolution`.
:param resolution: grid resolution measured in cells
:param mode: one of (None, 'vector', 'absolute', 'square')
:param dtype: data type of the returned tensor
:return: tensor holding the frequencies of the corresponding values computed by math.fft
"""
assert mode in ('vector', 'absolute', 'square')
k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution],
indexing='ij')
k = math.expand_dims(math.stack(k, -1), 0)
if dtype is not None:
k = k.astype(dtype)
else:
k = math.to_float(k)
if mode == 'vector':
return k
k = math.sum(k**2, axis=-1, keepdims=True)
if mode == 'square':
return k
else:
return math.sqrt(k)
def _divergence_nd(tensor, relative_shifts):
rank = spatial_rank(tensor)
tensor = math.pad(tensor, _get_pad_width(rank, (-relative_shifts[0], relative_shifts[1])))
components = []
for dimension in range(rank):
lower, upper = _dim_shifted(tensor, dimension, relative_shifts, diminish_others=(-relative_shifts[0], relative_shifts[1]), components=rank - dimension - 1)
components.append(upper - lower)
return math.sum(components, 0)
def l_n_loss(tensor, n, batch_norm=True):
if struct.isstruct(tensor):
all_tensors = struct.flatten(tensor)
return sum(l_n_loss(tensor, n, batch_norm) for tensor in all_tensors)
total_loss = math.sum(tensor**n) / n
if batch_norm:
batch_size = math.shape(tensor)[0]
return math.div(total_loss, math.to_float(batch_size))
else:
return total_loss
def fftfreq(resolution, mode='vector', dtype=np.float32):
assert mode in ('vector', 'absolute', 'square')
k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution], indexing='ij')
k = math.expand_dims(math.stack(k, -1), 0)
k = k.astype(dtype)
if mode == 'vector':
return k
k = math.sum(k**2, axis=-1, keepdims=True)
if mode == 'square':
return k
else:
return math.sqrt(k)
def _sliced_laplace_nd(tensor, axes=None):
"""
Laplace Stencil for N-Dimensions
aggregated from (c)enter, (u)pper, and (l)ower parts
"""
rank = spatial_rank(tensor)
dims = range(rank)
components = []
for ax in dims:
if _contains_axis(axes, ax, rank):
lower, center, upper = _dim_shifted(tensor, ax, (-1, 0, 1), diminish_others=(1, 1), diminish_other_condition=lambda other_ax: _contains_axis(axes, other_ax, rank))
components.append(upper + lower - 2 * center)
return math.sum(components, 0)
def loop_body(pressure, momentum, A_times_momentum, residual, loop_index):
"""
iteratively solve for:
x : pressure
momentum : momentum
laplace_momentum : A_times_momentum
residual : residual
"""
tmp = math.sum(momentum * A_times_momentum,
axis=non_batch_dims,
keepdims=True) # t = sum(mAm)
a = math.divide_no_nan(
math.sum(momentum * residual, axis=non_batch_dims, keepdims=True),
tmp) # a = sum(mr)/sum(mAm)
pressure += a * momentum # p += am
residual -= a * A_times_momentum # r -= aAm
momentum = residual - math.divide_no_nan(
math.sum(residual * A_times_momentum,
axis=non_batch_dims,
keepdims=True) * momentum,
tmp) # m = r-sum(rAm)*m/t = r-sum(rAm)*m/sum(mAm)
A_times_momentum = apply_A(momentum) # Am = A*m
return [pressure, momentum, A_times_momentum, residual, loop_index + 1]
def spatial_sum(tensor):
summed = math.sum(tensor, axis=math.dimrange(tensor))
for i in math.dimrange(tensor):
summed = math.expand_dims(summed, i)
return summed