def multi_advect(self, fields, interpolation="LINEAR", dt=1):
assert isinstance(
fields,
(list, tuple)), "first parameter must be either a tuple or list"
inputs_lists = []
coords_lists = []
value_generators = []
for field in fields:
if isinstance(field, StaggeredGrid):
i, c, v = self._mac_block_advection(field.staggered, dt)
else:
i, c, v = self._centered_block_advection(field, dt)
inputs_lists.append(i)
coords_lists.append(c)
value_generators.append(v)
inputs = math.concat(sum(inputs_lists, []), 0)
coords = math.concat(sum(coords_lists, []), 0)
all_advected = math.resample(inputs,
coords,
interpolation=interpolation,
boundary="REPLICATE")
all_advected = math.reshape(all_advected, [self.spatial_rank, -1] +
list(all_advected.shape[1:]))
all_advected = math.unstack(all_advected)
results = []
abs_i = 0
for i in range(len(inputs_lists)):
n = len(inputs_lists[0])
assigned_advected = all_advected[abs_i:abs_i + n]
results.append(value_generators[i](assigned_advected))
abs_i += n
return results
def _padded_resample(self, points):
data = self.padded([[1, 1]] * self.rank).data
local_points = self.box.global_to_local(points)
local_points = local_points * math.to_float(
self.resolution) + 0.5 # depends on amount of padding
resampled = math.resample(data,
local_points,
boundary='replicate',
interpolation=self.interpolation)
return resampled
def _advect_centered_field(self, field, dt, interpolation):
idx = indices_tensor(field)
centered_velocity = self.at_centers()[
..., ::-1] # assume right number of components
sample_coords = idx - centered_velocity * dt
result = math.resample(field,
sample_coords,
interpolation=interpolation,
boundary="REPLICATE")
return result
def sample_at(self, points, collapse_dimensions=True):
if not isinstance(self.extrapolation, six.string_types):
return self._padded_resample(points)
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
resampled = math.resample(self.data,
local_points,
boundary=_pad_mode(self.extrapolation),
interpolation=self.interpolation)
return resampled
def sample_at(self, points):
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
resampled = math.resample(self.data,
local_points,
boundary=_pad_mode(self.extrapolation),
interpolation=self.interpolation,
constant_values=_pad_value(
self.extrapolation_value))
return resampled
def sample_at(self, points, collapse_dimensions=True):
if not isinstance(self.extrapolation, six.string_types):
return self._padded_resample(points)
local_points = self.box.global_to_local(points)
local_points = local_points * math.to_float(self.resolution) - 0.5
if self.extrapolation == 'periodic':
data = math.pad(self.data,
[[0, 0]] + [[0, 1]] * self.rank + [[0, 0]],
mode='wrap')
local_points = local_points % math.to_float(
math.staticshape(self.data)[1:-1])
resampled = math.resample(data,
local_points,
interpolation=self.interpolation)
else:
boundary = 'replicate' if self.extrapolation == 'boundary' else 'zero'
resampled = math.resample(self.data,
local_points,
boundary=boundary,
interpolation=self.interpolation)
return resampled
def sample_at(self, points, collapse_dimensions=True):
if not isinstance(self.extrapolation, six.string_types):
return self._padded_resample(points)
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
boundary = {
'periodic': 'circular',
'boundary': 'replicate',
'constant': 'constant'
}[self.extrapolation]
resampled = math.resample(self.data,
local_points,
boundary=boundary,
interpolation=self.interpolation)
return resampled
def _advect_mac(self, field_mac, dt, interpolation):
# resample in each dimension
idx = indices_tensor(self.staggered)
advected_component_fields = []
dims = range(len(self.staggered.shape) - 2)
for d in dims: # z,y,x
velocity_at_staggered_points = self.at_faces(len(dims) - d -
1)[..., ::-1]
sample_coords = idx - velocity_at_staggered_points * dt
d_comp = len(dims) - d - 1
advected = math.resample(field_mac[..., d_comp:d_comp + 1],
sample_coords,
interpolation=interpolation,
boundary="REPLICATE")
advected_component_fields.append(advected)
all_advected = math.concat(advected_component_fields[::-1], axis=-1)
return all_advected