def __init__(self,
equation,
grid,
stencil_size=5,
initial_accuracy_order=1,
constrained_accuracy_order=1,
learned_keys=None,
fixed_keys=None,
num_time_steps=1,
target=None,
geometric_transforms=None,
predict_permutations=True,
name='linear_model'):
super().__init__(equation, grid, num_time_steps, target, name)
self.learned_keys, self.fixed_keys = (normalize_learned_and_fixed_keys(
learned_keys, fixed_keys, equation))
self.output_layers = build_output_layers(
equation,
grid,
self.learned_keys,
stencil_size,
initial_accuracy_order,
constrained_accuracy_order,
layer_cls=FixedCoefficientsLayer,
predict_permutations=predict_permutations)
self.fd_model = FiniteDifferenceModel(equation, grid,
initial_accuracy_order)
self.geometric_transforms = geometric_transforms or [
geometry.Identity()
]
def __init__(self,
equation,
grid,
stencil_size=5,
initial_accuracy_order=1,
constrained_accuracy_order=1,
learned_keys=None,
fixed_keys=None,
core_model_func=conv2d_stack,
num_time_steps=1,
geometric_transforms=None,
predict_permutations=True,
target=None,
name='pseudo_linear_model',
**kwargs):
# NOTE(jiaweizhuang): Too many input arguments. Only document important or
# confusing ones for now.
# pylint: disable=g-doc-args
"""Initialize class.
Args:
core_model_func: callable (function or class object). It should return
a Keras model (or layer) instance, which contains trainable weights.
The returned core_model instance should take a dict of tensors as input
(see the call() method in the base TimeStepModel class).
Additional kwargs are passed to this callable to specify hyperparameters
of core_model (such as number of layers and convolutional filters).
"""
# pylint: enable=g-doc-args
super().__init__(equation, grid, num_time_steps, target, name)
self.learned_keys, self.fixed_keys = (normalize_learned_and_fixed_keys(
learned_keys, fixed_keys, equation))
self.output_layers = build_output_layers(
equation,
grid,
self.learned_keys,
stencil_size,
initial_accuracy_order,
constrained_accuracy_order,
layer_cls=VaryingCoefficientsLayer,
predict_permutations=predict_permutations)
self.fd_model = FiniteDifferenceModel(equation, grid,
initial_accuracy_order)
self.geometric_transforms = geometric_transforms or [
geometry.Identity()
]
num_outputs = sum(layer.kernel_size
for layer in self.output_layers.values())
self.core_model = core_model_func(num_outputs, **kwargs)
def __init__(self,
equation,
grid,
stencil_size=5,
initial_accuracy_order=1,
constrained_accuracy_order=1,
learned_keys=None,
fixed_keys=None,
core_model_func=conv2d_stack,
num_time_steps=1,
geometric_transforms=None,
predict_permutations=True,
name='pseudo_linear_model',
**kwargs):
super().__init__(equation, grid, num_time_steps, name)
self.learned_keys, self.fixed_keys = (normalize_learned_and_fixed_keys(
learned_keys, fixed_keys, equation))
self.output_layers = build_output_layers(
equation,
grid,
self.learned_keys,
stencil_size,
initial_accuracy_order,
constrained_accuracy_order,
layer_cls=VaryingCoefficientsLayer)
self.fd_model = FiniteDifferenceModel(equation, grid,
initial_accuracy_order)
if not predict_permutations:
# NOTE(shoyer): this only makes sense if geometric_transforms includes
# permutations. Otherwise you won't be predicting every needed tensor.
modeled = set()
for key in sorted(self.output_layers):
value = equation.key_definitions[key]
swapped = value.swap_xy()
if swapped in modeled:
del self.output_layers[key]
modeled.add(value)
if geometric_transforms is None:
geometric_transforms = [geometry.Identity()]
self.geometric_transforms = geometric_transforms
num_outputs = sum(layer.kernel_size
for layer in self.output_layers.values())
self.core_model = core_model_func(num_outputs, **kwargs)