Commit ad0a6d12 authored by csongor's avatar csongor
Browse files

Move files, rename, tie in with smoothing operator

parent dc2945e5
import numpy as np import numpy as np
import nifty.nifty_utilities as utilities import nifty.nifty_utilities as utilities
from nifty import RGSpace, LMSpace
from nifty.operators.endomorphic_operator import EndomorphicOperator from nifty.operators.endomorphic_operator import EndomorphicOperator
from nifty.operators.fft_operator import FFTOperator from nifty.operators.fft_operator import FFTOperator
import smooth_util as su
from d2o import STRATEGIES
class SmoothingOperator(EndomorphicOperator): class SmoothingOperator(EndomorphicOperator):
# ---Overwritten properties and methods--- # ---Overwritten properties and methods---
def __init__(self, domain=(), field_type=(), sigma=0): def __init__(self, domain=(), field_type=(), sigma=0,
log_distances=False):
self._domain = self._parse_domain(domain) self._domain = self._parse_domain(domain)
self._field_type = self._parse_field_type(field_type) self._field_type = self._parse_field_type(field_type)
...@@ -17,8 +18,8 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -17,8 +18,8 @@ class SmoothingOperator(EndomorphicOperator):
if len(self.domain) != 1: if len(self.domain) != 1:
raise ValueError( raise ValueError(
'ERROR: SmoothOperator accepts only exactly one ' 'ERROR: SmoothOperator accepts only exactly one '
'space as input domain.' 'space as input domain.'
) )
if self.field_type != (): if self.field_type != ():
...@@ -27,13 +28,16 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -27,13 +28,16 @@ class SmoothingOperator(EndomorphicOperator):
'empty tuple.' 'empty tuple.'
) )
self._sigma = sigma self.sigma = sigma
self.log_distances = log_distances
self._direct_smoothing_width = 2.
def _inverse_times(self, x, spaces, types): def _inverse_times(self, x, spaces, types):
return self._smooth_helper(x, spaces, types, inverse=True) return self._smoothing_helper(x, spaces, inverse=True)
def _times(self, x, spaces, types): def _times(self, x, spaces, types):
return self._smooth_helper(x, spaces, types) return self._smoothing_helper(x, spaces, inverse=False)
# ---Mandatory properties and methods--- # ---Mandatory properties and methods---
@property @property
...@@ -50,18 +54,31 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -50,18 +54,31 @@ class SmoothingOperator(EndomorphicOperator):
@property @property
def symmetric(self): def symmetric(self):
return True return False
@property @property
def unitary(self): def unitary(self):
return False return False
# ---Added properties and methods--- # ---Added properties and methods---
@property @property
def sigma(self): def sigma(self):
return self._sigma return self._sigma
def _smooth_helper(self, x, spaces, types, inverse=False): @sigma.setter
def sigma(self, sigma):
self._sigma = np.float(sigma)
@property
def log_distances(self):
return self._log_distances
@log_distances.setter
def log_distances(self, log_distances):
self._log_distances = bool(log_distances)
def _smoothing_helper(self, x, spaces, inverse):
if self.sigma == 0: if self.sigma == 0:
return x.copy() return x.copy()
...@@ -74,6 +91,13 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -74,6 +91,13 @@ class SmoothingOperator(EndomorphicOperator):
else: else:
spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain)) spaces = utilities.cast_axis_to_tuple(spaces, len(x.domain))
try:
result = self._fft_smoothing(x, spaces, inverse)
except ValueError:
result = self._direct_smoothing(x, spaces, inverse)
return result
def _fft_smoothing(self, x, spaces, inverse):
Transformator = FFTOperator(x.domain[spaces[0]]) Transformator = FFTOperator(x.domain[spaces[0]])
# transform to the (global-)default codomain and perform all remaining # transform to the (global-)default codomain and perform all remaining
...@@ -87,9 +111,13 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -87,9 +111,13 @@ class SmoothingOperator(EndomorphicOperator):
transformed_x.val.get_axes_local_distribution_strategy(axes=coaxes) transformed_x.val.get_axes_local_distribution_strategy(axes=coaxes)
kernel = codomain.get_distance_array( kernel = codomain.get_distance_array(
distribution_strategy=axes_local_distribution_strategy) distribution_strategy=axes_local_distribution_strategy)
if self.log_distances:
kernel.apply_scalar_function(np.log, inplace=True)
kernel.apply_scalar_function( kernel.apply_scalar_function(
codomain.get_smoothing_kernel_function(self.sigma), codomain.get_fft_smoothing_kernel_function(self.sigma),
inplace=True) inplace=True)
# now, apply the kernel to transformed_x # now, apply the kernel to transformed_x
...@@ -98,7 +126,7 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -98,7 +126,7 @@ class SmoothingOperator(EndomorphicOperator):
local_transformed_x = transformed_x.val.get_local_data(copy=False) local_transformed_x = transformed_x.val.get_local_data(copy=False)
local_kernel = kernel.get_local_data(copy=False) local_kernel = kernel.get_local_data(copy=False)
reshaper = [local_transformed_x.shape[i] if i in coaxes else 1 reshaper = [transformed_x.shape[i] if i in coaxes else 1
for i in xrange(len(transformed_x.shape))] for i in xrange(len(transformed_x.shape))]
local_kernel = np.reshape(local_kernel, reshaper) local_kernel = np.reshape(local_kernel, reshaper)
...@@ -113,3 +141,111 @@ class SmoothingOperator(EndomorphicOperator): ...@@ -113,3 +141,111 @@ class SmoothingOperator(EndomorphicOperator):
result = Transformator.inverse_times(transformed_x, spaces=spaces) result = Transformator.inverse_times(transformed_x, spaces=spaces)
return result return result
def _direct_smoothing(self, x, spaces, inverse):
# infer affected axes
# we rely on the knowledge, that `spaces` is a tuple with length 1.
affected_axes = x.domain_axes[spaces[0]]
axes_local_distribution_strategy = \
x.val.get_axes_local_distribution_strategy(axes=affected_axes)
distance_array = x.domain[spaces[0]].get_distance_array(
distribution_strategy=axes_local_distribution_strategy)
if self.log_distances:
distance_array.apply_scalar_function(np.log, inplace=True)
# collect the local data + ghost cells
local_data_Q = False
if x.distribution_strategy == 'not':
local_data_Q = True
elif x.distribution_strategy in STRATEGIES['slicing']:
# infer the local start/end based on the slicing information of
# x's d2o. Only gets non-trivial for axis==0.
if 0 not in affected_axes:
local_data_Q = True
else:
# we rely on the fact, that the content of x.domain_axes is
# sorted
true_starts = [x.val.distributor.local_start]
true_starts += [0] * (len(affected_axes) - 1)
true_ends = [x.val.distributor.local_end]
true_ends += [x.shape[i] for i in affected_axes[1:]]
augmented_start = max(0,
true_starts[0] -
self._direct_smoothing_width * self.sigma)
augmented_end = min(x.shape[affected_axes[0]],
true_ends[0] +
self._direct_smoothing_width * self.sigma)
augmented_slice = slice(augmented_start, augmented_end)
augmented_data = x.val.get_data(augmented_slice,
local_keys=True,
copy=False)
augmented_data = augmented_data.get_local_data(copy=False)
augmented_distance_array = distance_array.get_data(
augmented_slice,
local_keys=True,
copy=False)
augmented_distance_array = \
augmented_distance_array.get_local_data(copy=False)
else:
raise ValueError(about._errors.cstring(
"ERROR: Direct smoothing not implemented for given"
"distribution strategy."))
if local_data_Q:
# if the needed data resides on the nodes already, the necessary
# are the same; no matter what the distribution strategy was.
augmented_data = x.val.get_local_data(copy=False)
augmented_distance_array = \
distance_array.get_local_data(copy=False)
true_starts = [0] * len(affected_axes)
true_ends = [x.shape[i] for i in affected_axes]
# perform the convolution along the affected axes
local_result = augmented_data
for index in range(len(affected_axes)):
data_axis = affected_axes[index]
distances_axis = index
true_start = true_starts[index]
true_end = true_ends[index]
local_result = self._direct_smoothing_single_axis(
local_result,
data_axis,
augmented_distance_array,
distances_axis,
true_start,
true_end,
inverse)
result = x.copy_empty()
result.val.set_local_data(local_result, copy=False)
return result
def _direct_smoothing_single_axis(self, data, data_axis, distances,
distances_axis, true_start, true_end,
inverse):
if inverse:
true_sigma = 1 / self.sigma
else:
true_sigma = self.sigma
if (data.dtype == np.dtype('float32')):
smoothed_data = su.apply_along_axis_f(data_axis, data,
startindex=true_start,
endindex=true_end,
distances=distances,
smooth_length=true_sigma)
else:
smoothed_data = su.apply_along_axis(data_axis, data,
startindex=true_start,
endindex=true_end,
distances=distances,
smooth_length=true_sigma)
return smoothed_data
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment