def placeholder(shape,
dtype=np.float32,
basename=None,
item_condition=struct.VARIABLES):
if struct.isstruct(dtype):
def placeholder_map(trace):
shape, dtype = trace.value
return tf.placeholder(dtype, shape, _tf_name(trace, basename))
zipped = struct.zip([shape, dtype],
leaf_condition=is_static_shape,
item_condition=item_condition)
return struct.map(placeholder_map,
zipped,
leaf_condition=is_static_shape,
trace=True,
item_condition=item_condition)
else:
f = lambda trace: tf.placeholder(dtype, trace.value,
_tf_name(trace, basename))
return struct.map(f,
shape,
leaf_condition=is_static_shape,
trace=True,
item_condition=item_condition)
def placeholder(shape, dtype=np.float32, basename='Placeholder'):
if struct.isstruct(dtype):
def placeholder_map(trace):
shape, dtype = trace.value
return tf.placeholder(dtype, shape, _tf_name(trace, basename))
zipped = struct.zip([shape, dtype], leaf_condition=is_static_shape)
return struct.map(placeholder_map, zipped, leaf_condition=is_static_shape, trace=True)
else:
def f(trace): return tf.placeholder(dtype, trace.value, _tf_name(trace, basename))
return struct.map(f, shape, leaf_condition=is_static_shape, trace=True)
def _add_default_fields(self):
def add_default_field(trace):
field = trace.value
if isinstance(field, (CenteredGrid, StaggeredGrid)):
def field_generator():
world_state = self.world.state
return trace.find_in(world_state)
self.add_field(field.name[0].upper() + field.name[1:], field_generator)
return None
struct.map(add_default_field, self.world.state, leaf_condition=lambda x: isinstance(x, (CenteredGrid, StaggeredGrid)), trace=True, content_type=struct.INVALID)
def test_identity(self):
for obj in generate_test_structs():
with struct.unsafe():
obj2 = struct.map(lambda s: s, obj, recursive=False)
self.assertEqual(obj, obj2)
obj3 = struct.map(lambda t: t, obj, recursive=True)
self.assertEqual(obj, obj3)
obj4 = struct.map(lambda t: t,
obj,
item_condition=struct.ALL_ITEMS)
self.assertEqual(obj, obj4)
def dataset_handle(shape, dtype, frames=None):
"""
Creates a single virtual TensorFlow dataset (iterator_handle) for the given struct.
The dataset is expected to hold contain all fields required for loading the obj given the current context item condition.
From the dataset, graph input tensors are derived and arranged into a struct of the same shape as obj.
If an integer is passed to frames, a list of such structs is created by unstacking the second-outer-most dimension of the dataset.
:param shape: tensor shape or struct of tensor shapes
:param dtype: data type of struct of data types matching shape
:param frames: Number of frames contained in each example of the dataset. Expects shape (batch_size, frames, ...)
:type frames: int or None
:return: list of struct and placeholder.
1. If frames=None: valid struct corresponding to obj. If frames>1: list thereof
2. placeholder for a TensorFlow dataset iterator handle (dtype=string)
:rtype: tuple
"""
shapes = tuple(struct.flatten(shape, leaf_condition=is_static_shape))
if struct.isstruct(dtype):
dtypes = tuple(struct.flatten(dtype))
assert len(dtypes) == len(shapes)
else:
dtypes = [dtype] * len(shapes)
if frames is not None:
shapes = tuple(
[shape[0:1] + (frames, ) + shape[1:] for shape in shapes])
# --- TF Dataset handle from string ---
iterator_handle = tf.placeholder(tf.string,
shape=[],
name='dataset_iterator_handle')
iterator = tf.data.Iterator.from_string_handle(iterator_handle,
output_types=dtypes,
output_shapes=shapes)
next_element = iterator.get_next()
# --- Create resulting struct by splitting `next_element`s ---
if frames is None:
next_element_list = list(next_element)
next_struct = struct.map(lambda _: next_element_list.pop(0),
shape,
leaf_condition=is_static_shape)
else:
# --- Remap structures -> to `frames` long list of structs ---
next_struct = []
for frame_idx in range(frames):
next_element_list = list(next_element)
frame_struct = struct.map(
lambda _: next_element_list.pop(0)[:, frame_idx, ...],
shape,
leaf_condition=is_static_shape)
next_struct.append(frame_struct)
return next_struct, iterator_handle
def downsample2x(tensor, interpolation='linear', axes=None):
if struct.isstruct(tensor):
return struct.map(lambda s: downsample2x(s, interpolation, axes),
tensor,
recursive=False)
if interpolation.lower() != 'linear':
raise ValueError('Only linear interpolation supported')
rank = spatial_rank(tensor)
if axes is None:
axes = range(rank)
tensor = math.pad(
tensor, [[0, 0]] +
[([0, 1] if
(dim % 2) != 0 and _contains_axis(axes, ax, rank) else [0, 0])
for ax, dim in enumerate(tensor.shape[1:-1])] + [[0, 0]], 'replicate')
for axis in axes:
upper_slices = tuple([(slice(1, None, 2) if i == axis else slice(None))
for i in range(rank)])
lower_slices = tuple([(slice(0, None, 2) if i == axis else slice(None))
for i in range(rank)])
tensor_sum = tensor[(slice(None), ) + upper_slices +
(slice(None), )] + tensor[(slice(None), ) +
lower_slices +
(slice(None), )]
tensor = tensor_sum / 2
return tensor
def diffuse(field, amount, substeps=1):
u"""
Simulate a finite-time diffusion process of the form dF/dt = α · ΔF on a given `Field` F with diffusion coefficient α.
If `field` is periodic (set via `extrapolation='periodic'`), diffusion may be simulated in Fourier space.
Otherwise, finite differencing is used to approximate the
:param field: CenteredGrid, StaggeredGrid or ConstantField
:param amount: number of Field, typically α · dt
:param substeps: number of iterations to use
:return: Field of same type as `field`
:rtype: Field
"""
if isinstance(field, ConstantField):
return field
if isinstance(field, StaggeredGrid):
return struct.map(
lambda grid: diffuse(grid, amount, substeps=substeps),
field,
leaf_condition=lambda x: isinstance(x, CenteredGrid))
assert isinstance(
field, CenteredGrid), "Cannot diffuse field of type '%s'" % type(field)
if field.extrapolation == 'periodic' and not isinstance(amount, Field):
fft_laplace = -(2 * pi)**2 * field.squared_frequencies
diffuse_kernel = math.exp(fft_laplace * amount)
return math.real(
math.ifft(field.fft() * math.to_complex(diffuse_kernel)))
else:
data = field.data
if isinstance(amount, Field):
amount = amount.at(field).data
else:
amount = math.batch_align(amount, 0, data)
for i in range(substeps):
data += amount / substeps * field.laplace().data
return field.with_data(data)
def test_state_collection(self):
fluid = Fluid(Domain([1, 1]))
fluid2 = Fluid(Domain([2, 2]))
c1 = StateCollection([fluid])
assert c1.fluid is fluid
assert fluid in c1
assert c1[fluid] is fluid
assert isinstance(repr(c1), six.string_types)
assert len(c1) == len(c1.shape) == len(c1.staticshape) == len(c1.dtype)
assert c1.shape.fluid.density.data == (1, 1, 1, 1)
self.assertIsInstance(c1.dtype.fluid.density.data, numpy.dtype)
c2 = StateCollection()
assert len(c2) == 0
c2 = c2.state_added(fluid)
assert c2 == c1
assert hash(c2) == hash(c1)
c3 = c2.state_replaced(fluid2)
assert c3 != c2
assert c3.fluid is fluid2
c4 = c3.state_removed(fluid2)
assert len(c4) == 0
c5 = struct.map(lambda x: x, c1)
assert isinstance(c5, StateCollection)
assert c5 == c1
def downsample2x(tensor, interpolation='linear'):
if struct.isstruct(tensor):
return struct.map(lambda s: downsample2x(s, interpolation),
tensor,
recursive=False)
if interpolation.lower() != 'linear':
raise ValueError('Only linear interpolation supported')
dims = range(spatial_rank(tensor))
tensor = math.pad(tensor,
[[0, 0]] + [([0, 1] if (dim % 2) != 0 else [0, 0])
for dim in tensor.shape[1:-1]] + [[0, 0]],
'SYMMETRIC')
for dimension in dims:
upper_slices = tuple([
(slice(1, None, 2) if i == dimension else slice(None))
for i in dims
])
lower_slices = tuple([
(slice(0, None, 2) if i == dimension else slice(None))
for i in dims
])
sum = tensor[(slice(None), ) + upper_slices +
(slice(None), )] + tensor[(slice(None), ) + lower_slices +
(slice(None), )]
tensor = sum / 2
return tensor
def build_graph_input(obj, input_type='placeholder', frames=None):
"""
Create placeholders for tensors in the supplied state.
:param obj: struct or StateProxy
:param input_type: 'placeholder' or 'dataset_handle'
:param frames: Number of input frames. If not None, returns a list of input structs.
:return:
1. Valid state containing or derived from created placeholders or dataset handle
2. dict mapping from placeholders to their default names (using struct.names)
"""
if isinstance(obj, StateProxy):
obj = obj.state
assert struct.isstruct(obj)
# --- Shapes and names ---
writable_obj = _transform_for_writing(obj)
shape = _writing_staticshape(obj)
names = struct.names(writable_obj)
if input_type == 'placeholder':
if frames is not None: raise NotImplementedError()
with _unsafe():
placeholders = placeholder(shape)
graph_in = struct.map(
lambda x: x,
placeholders) # validates fields, splits staggered tensors
return graph_in, {placeholders: names}
elif input_type == 'dataset_handle':
with _unsafe():
dtypes = struct.dtype(writable_obj)
dataset_nodes, iterator_handle = dataset_handle(shape,
dtypes,
frames=frames)
graph_in = struct.map(
lambda x: x,
dataset_nodes) # validates fields, splits staggered tensors
shapes = struct.flatten(struct.staticshape(dataset_nodes),
leaf_condition=is_static_shape)
dtypes = struct.flatten(struct.dtype(dataset_nodes))
return graph_in, {
'names': struct.flatten(names),
'iterator_handle': iterator_handle,
'shapes': shapes,
'dtypes': dtypes,
'frames': frames
}
else:
raise ValueError(input_type)
def test_paths(self):
obj = {
'Vels': [CenteredGrid(numpy.zeros([1, 4, 1]), box[0:1], name='v')]
}
with struct.unsafe():
names = struct.flatten(
struct.map(lambda attr: attr.path(), obj, trace=True))
self.assertEqual(names[0], 'Vels.0.data')
def test_names(self):
for obj in generate_test_structs():
with struct.unsafe():
names = struct.flatten(
struct.map(lambda attr: attr.name, obj, trace=True))
self.assertGreater(len(names), 0)
for name in names:
self.assertIsInstance(name, str)
def constant(value, dtype=None, basename='const'):
def f(trace):
return tf.constant(
trace.value,
dtype=TF_BACKEND.precision_dtype if dtype is None else dtype,
name=_tf_name(trace, basename))
return struct.map(f, value, trace=True)
def torch_to_numpy(obj, item_condition=struct.ALL_ITEMS):
def to_numpy(obj):
if isinstance(obj, torch.Tensor):
return obj.numpy()
else:
return obj
return struct.map(to_numpy, obj, item_condition=item_condition)
def placeholder_like(obj, basename=None):
warnings.warn(
"placeholder_like may not respect the batch dimension. "
"For State objects, use placeholder(state.shape) instead.",
DeprecationWarning,
stacklevel=2)
f = lambda attr: tf.placeholder(attr.value.dtype, attr.value.shape,
_tf_name(attr, basename))
return struct.map(f, obj, leaf_condition=is_static_shape, trace=True)
def variable(initial_value, dtype=None, basename='Variable', trainable=True):
def f(attr):
return tf.Variable(
attr.value,
name=_tf_name(attr, basename),
dtype=TF_BACKEND.precision_dtype if dtype is None else dtype,
trainable=trainable)
return struct.map(f, initial_value, trace=True)
def _transform_for_writing(obj):
def f(value):
if isinstance(value, field.StaggeredGrid):
return value.staggered_tensor()
if isinstance(value, field.CenteredGrid):
return value.data
else:
return value
data = struct.map(f, obj, lambda x: isinstance(x, (field.StaggeredGrid, field.CenteredGrid)), content_type='format')
return data
def read(self, obj, frame=0):
if struct.isstruct(obj):
obj = _transform_for_writing(obj)
names = struct.flatten(obj)
if not np.all([isinstance(n, six.string_types) for n in names]):
names = struct.names(obj)
data = struct.map(lambda name: self.read_array(self._filename(name), frame), names)
return data
else:
return self.read_array('unnamed', frame)
def randfreq(shape, dtype=np.float32, power=8):
def genarray(shape):
fft = randn(shape, dtype) + 1j * randn(shape, dtype)
k = fftfreq(shape[1:-1], mode='absolute')
shape_fac = math.sqrt(math.mean(shape[1:-1])) # 16: 4, 64: 8, 256: 24,
fft *= (1 / (k + 1))**power * power * shape_fac
array = math.real(math.ifft(fft)).astype(dtype)
return array
return struct.map(genarray, shape, leaf_condition=is_static_shape)
def _get_batch(self, indices):
data_list = self._cache.get(indices, self._load, add_to_cache=True)
data = list_swap_axes(data_list)
data_map = {
self.streams[i]: data[i]
for i in range(len(self._streams))
}
return struct.map(lambda x, is_stream: data_map[x] if is_stream else x,
struct.zip([self._fields, self.stream_mask]),
content_type=struct.INVALID)
def read_sim_frame(directory: math.Tensor,
names: str or tuple or list or dict or struct.Struct,
frame: int,
convert_to_backend=True):
def single_read(name):
name = _slugify_filename(name)
files = math.map(lambda dir_: _filename(dir_, name, frame), directory)
return read(files, convert_to_backend=convert_to_backend)
return struct.map(single_read, names)
def broadcast_function(backend, func, args, kwargs):
backend_func = getattr(backend, func)
obj, build_arguments = argument_assembler(args, kwargs)
def f(*values):
args, kwargs = build_arguments(values)
result = backend_func(*args, **kwargs)
return result
with struct.unsafe():
return struct.map(f, obj)
def _writing_staticshape(obj):
def f(value):
if isinstance(value, field.StaggeredGrid):
shape = math.staticshape(value.staggered_tensor())
return (value._batch_size,) + shape[1:]
if isinstance(value, field.CenteredGrid):
return value.staticshape.data
else:
return value
data = struct.map(f, obj, lambda x: isinstance(x, (field.StaggeredGrid, field.CenteredGrid)), content_type='format_staticshape')
return data
def _transform_for_writing(obj):
def f(value):
if isinstance(value, field.StaggeredGrid):
return value.staggered_tensor()
if isinstance(value, field.CenteredGrid):
return value.data
else:
return value
with struct.unsafe():
data = struct.map(f, obj, lambda x: isinstance(x, (field.StaggeredGrid, field.CenteredGrid)))
return data
def variable(initial_value,
dtype=np.float32,
basename='Variable',
trainable=True):
def f(attr):
return tf.Variable(attr.value,
name=_tf_name(attr, basename),
dtype=dtype,
trainable=trainable)
return struct.map(f, initial_value, trace=True)
def accessible_tensor(self, extend=0):
"""
Scalar channel encoding cells that are accessible, i.e. not solid, as ones and obstacles as zero.
:param extend: Extend the grid in all directions beyond the grid size specified by the domain
"""
pad_values = struct.map(lambda solid: int(not solid), Material.solid(self.domain.boundaries))
if isinstance(pad_values, (list, tuple)):
pad_values = [0] + list(pad_values) + [0]
result = math.pad(self.accessible.data, [[0,0]] + [[extend, extend]] * self.rank + [[0,0]], constant_values=pad_values)
return result
def placeholder(shape, dtype=None, basename='Placeholder'):
if struct.isstruct(dtype):
def placeholder_map(trace):
shape, dtype = trace.value
return tf.placeholder(dtype, shape, _tf_name(trace, basename))
zipped = struct.zip([shape, dtype], leaf_condition=is_static_shape)
return struct.map(placeholder_map,
zipped,
leaf_condition=is_static_shape,
trace=True)
else:
def f(trace):
return tf.placeholder(
TF_BACKEND.precision_dtype if dtype is None else dtype,
trace.value, _tf_name(trace, basename))
return struct.map(f, shape, leaf_condition=is_static_shape, trace=True)
def load_state(state):
if isinstance(state, StateProxy):
state = state.state
assert isinstance(state, State)
state = _transform_for_writing(state)
names = struct.names(state)
with _unsafe():
placeholders = placeholder(state.shape)
state_in = struct.map(
lambda x: x,
placeholders) # validates fields, splits staggered tensors
return state_in, {placeholders: names}
def _add_default_fields(self):
def add_default_field(trace):
field = trace.value
if isinstance(field, (CenteredGrid, StaggeredGrid)):
def field_generator():
world_state = self.world.state
return trace.find_in(world_state)
self.add_field(field.name[0].upper() + field.name[1:],
field_generator)
return None
old_worldstate = world.state
with struct.unsafe():
struct.map(add_default_field,
world.state,
leaf_condition=lambda x: isinstance(
x, (CenteredGrid, StaggeredGrid)),
trace=True)
is_same = old_worldstate is world.state
print()
def shape(self):
"""
Similar to phi.math.shape(self) but respects unknown dimensions.
"""
def tensorshape(tensor):
if tensor is None: return None
default_batched_shape = staticshape(tensor)
if len(default_batched_shape) >= 2:
return (self._batch_size,) + default_batched_shape[1:]
else:
return default_batched_shape
with struct.unsafe():
return struct.map(tensorshape, self, item_condition=struct.VARIABLES)
