def predict_result(inputs: tf.Tensor, hparams: tf.contrib.training.HParams) -> tf.Tensor: """Infer predictions from inputs with our forward model. Args: inputs: float32 Tensor with dimensions [batch, x]. hparams: model hyperparameters. Returns: Float32 Tensor with dimensions [batch, x] with inferred time derivatives. """ if hparams.model_target in {'flux', 'time_derivative'}: # use dummy values (all zeros) for space derivatives if hparams.space_derivatives_weight: raise ValueError( 'space derivatives are not predicted by model {}'.format( hparams.model_target)) _, equation = equations.from_hparams(hparams) num_derivatives = len(equation.DERIVATIVE_ORDERS) space_derivatives = tf.zeros( tf.concat([tf.shape(inputs), [num_derivatives]], axis=0)) time_derivative = predict_time_derivative(inputs, hparams) else: space_derivatives = predict_space_derivatives(inputs, hparams) _, equation = equations.from_hparams(hparams) time_derivative = apply_space_derivatives(space_derivatives, inputs, equation) if hparams.num_time_steps: integrated_solution = predict_time_evolution(inputs, hparams) else: integrated_solution = None return result_stack(space_derivatives, time_derivative, integrated_solution)
def integrate_model_from_warm_start( checkpoint_dir: str, y0: np.ndarray, hparams: tf.contrib.training.HParams = None, random_seed: int = 0, times: np.ndarray = _DEFAULT_TIMES, warmup: float = 0, integrate_method: str = 'RK23') -> xarray.Dataset: """Integrate the given PDE with standard and modeled finite differences.""" if hparams is None: hparams = training.load_hparams(checkpoint_dir) logging.info('integrating %s with seed=%s', hparams.equation, random_seed) _, equation_coarse = equations.from_hparams(hparams, random_seed=random_seed) logging.info('solving neural network model at low resolution') checkpoint_path = training.checkpoint_dir_to_path(checkpoint_dir) differentiator = SavedModelDifferentiator( checkpoint_path, equation_coarse, hparams) solution_model, num_evals_model = odeint( y0, differentiator, warmup+times, method=integrate_method) results = xarray.Dataset( data_vars={'y': (('time', 'x'), solution_model)}, coords={'time': warmup+times, 'x': equation_coarse.grid.solution_x, 'num_evals': num_evals_model}) return results
def predict_flux_directly(inputs, hparams, reuse=tf.AUTO_REUSE): """Predict flux directly, without using the equation of motion.""" _, equation = equations.from_hparams(hparams) dx = equation.grid.solution_dx output = _multilayer_conv1d(inputs, hparams, num_targets=1, reuse=reuse) flux = tf.squeeze(output, axis=-1) return equations.staggered_first_derivative(flux, dx)
def integrate_exact_baseline_and_model( checkpoint_dir: str, hparams: tf.contrib.training.HParams = None, random_seed: int = 0, times: np.ndarray = _DEFAULT_TIMES, warmup: float = 0, integrate_method: str = 'RK23', exact_filter_interval: float = None) -> xarray.Dataset: """Integrate the given PDE with standard and modeled finite differences.""" if hparams is None: hparams = training.load_hparams(checkpoint_dir) logging.info('integrating %s with seed=%s', hparams.equation, random_seed) equation_fine, equation_coarse = equations.from_hparams( hparams, random_seed=random_seed) logging.info('solving the "exact" model at high resolution') ds_solution_exact = integrate_exact( equation_fine, times, warmup, integrate_method=integrate_method, filter_interval=exact_filter_interval) solution_exact = ds_solution_exact['y'].data num_evals_exact = ds_solution_exact['num_evals'].item() # resample to the coarse grid y0 = equation_coarse.grid.resample(solution_exact[0, :]) if np.isnan(y0).any(): raise ValueError('solution contains NaNs') logging.info('solving baseline finite differences at low resolution') differentiator = PolynomialDifferentiator(equation_coarse) solution_baseline, num_evals_baseline = odeint( y0, differentiator, warmup+times, method=integrate_method) logging.info('solving neural network model at low resolution') checkpoint_path = training.checkpoint_dir_to_path(checkpoint_dir) differentiator = SavedModelDifferentiator( checkpoint_path, equation_coarse, hparams) solution_model, num_evals_model = odeint( y0, differentiator, warmup+times, method=integrate_method) results = xarray.Dataset({ 'y_exact': (('time', 'x_high'), solution_exact), 'y_baseline': (('time', 'x_low'), solution_baseline), 'y_model': (('time', 'x_low'), solution_model), }, coords={ 'time': warmup+times, 'x_low': equation_coarse.grid.solution_x, 'x_high': equation_fine.grid.solution_x, 'num_evals_exact': num_evals_exact, 'num_evals_baseline': num_evals_baseline, 'num_evals_model': num_evals_model, }) return results
def __init__(self, snapshots: np.ndarray, hparams: tf.contrib.training.HParams, training: bool = False): """Initialize an object for running inference. Args: snapshots: np.ndarray with shape [examples, x] with high-resolution training data. hparams: hyperparameters for training. training: whether to evaluate on training or validation datasets. """ if training: dataset_type = model.Dataset.TRAINING else: dataset_type = model.Dataset.VALIDATION dataset = model.make_dataset(snapshots, hparams, dataset_type=dataset_type, repeat=False, evaluation=True) iterator = dataset.make_initializable_iterator() data = iterator.get_next() _, coarse_equation = equations.from_hparams(hparams) predictions = model.predict_result(data['inputs'], hparams) loss_per_head = model.loss_per_head(predictions, labels=data['labels'], baseline=data['baseline'], hparams=hparams) loss = model.weighted_loss(loss_per_head, hparams) results = dict(data, predictions=predictions) metrics = { k: tf.contrib.metrics.streaming_concat(v) for k, v in results.items() } metrics['loss'] = tf.metrics.mean(loss) space_loss, time_loss, integrated_loss = model.result_unstack( loss_per_head, coarse_equation) metrics['loss/space_derivatives'] = tf.metrics.mean(space_loss) metrics['loss/time_derivative'] = tf.metrics.mean(time_loss) if integrated_loss is not None: metrics['loss/integrated_solution'] = tf.metrics.mean( integrated_loss) initializer = tf.group(iterator.initializer, tf.local_variables_initializer()) self._initializer = initializer self._metrics = metrics
def predict_time_derivative(inputs: tf.Tensor, hparams: tf.contrib.training.HParams, reuse: object = tf.AUTO_REUSE) -> tf.Tensor: """Infer time evolution from inputs with our forward model. Args: inputs: float32 Tensor with dimensions [batch, x]. hparams: model hyperparameters. reuse: whether or not to reuse TensorFlow variables. Returns: Float32 Tensor with dimensions [batch, x] with inferred time derivatives. """ space_derivatives = predict_space_derivatives(inputs, hparams, reuse=reuse) _, equation = equations.from_hparams(hparams) return apply_space_derivatives(space_derivatives, inputs, equation)
def predict_time_evolution(inputs: tf.Tensor, hparams: tf.contrib.training.HParams) -> tf.Tensor: """Infer time evolution from inputs with our neural network model. Args: inputs: float32 Tensor with dimensions [batch, x]. hparams: model hyperparameters. Returns: Float32 Tensor with dimensions [batch, x, num_time_steps+1] with the integrated solution. """ def func(y, t): del t # unused return predict_time_derivative(y, hparams, reuse=True) _, equation = equations.from_hparams(hparams) return integrate_ode(func, inputs, hparams.num_time_steps, equation.time_step)
def _multilayer_conv1d(inputs, hparams, num_targets, reuse=tf.AUTO_REUSE): """Apply multiple conv1d layers with input normalization.""" _, equation = equations.from_hparams(hparams) assert_consistent_solution(equation, inputs) net = inputs[:, :, tf.newaxis] net /= equation.standard_deviation activation = _NONLINEARITIES[hparams.nonlinearity] for _ in range(hparams.num_layers - 1): net = layers.conv1d_periodic_layer(net, filters=hparams.filter_size, kernel_size=hparams.kernel_size, activation=activation, center=True) if hparams.num_layers == 0: raise NotImplementedError('not implemented yet') net = layers.conv1d_periodic_layer(net, filters=num_targets, kernel_size=hparams.kernel_size, activation=None, center=True) return net
def run_integrate( seed_and_initial_condition, checkpoint_dir=FLAGS.checkpoint_dir, times=np.arange(0, FLAGS.time_max + FLAGS.time_delta, FLAGS.time_delta), warmup=FLAGS.warmup, integrate_method=FLAGS.integrate_method, ): random_seed, y0 = seed_and_initial_condition _, equation_coarse = equations.from_hparams( hparams, random_seed=random_seed) checkpoint_path = training.checkpoint_dir_to_path(checkpoint_dir) differentiator = integrate.SavedModelDifferentiator( checkpoint_path, equation_coarse, hparams) solution_model, num_evals_model = integrate.odeint( y0, differentiator, warmup+times, method=integrate_method) results = xarray.Dataset( data_vars={'y': (('time', 'x'), solution_model)}, coords={'time': warmup+times, 'x': equation_coarse.grid.solution_x, 'num_evals': num_evals_model, 'sample': random_seed}) return results
def predict_result(inputs: tf.Tensor, hparams: tf.contrib.training.HParams) -> tf.Tensor: """Infer predictions from inputs with our forward model. Args: inputs: float32 Tensor with dimensions [batch, x]. hparams: model hyperparameters. Returns: Float32 Tensor with dimensions [batch, x] with inferred time derivatives. """ space_derivatives = predict_space_derivatives(inputs, hparams) _, equation = equations.from_hparams(hparams) time_derivative = apply_space_derivatives(space_derivatives, inputs, equation) if hparams.num_time_steps: integrated_solution = predict_time_evolution(inputs, hparams) else: integrated_solution = None return result_stack(space_derivatives, time_derivative, integrated_solution)
def predict_space_derivatives_directly(inputs, hparams, reuse=tf.AUTO_REUSE): """Predict finite difference coefficients directly from a neural net.""" _, equation = equations.from_hparams(hparams) num_targets = len(equation.DERIVATIVE_ORDERS) return _multilayer_conv1d(inputs, hparams, num_targets, reuse=reuse)
def predict_coefficients(inputs: tf.Tensor, hparams: tf.contrib.training.HParams, reuse: object = tf.AUTO_REUSE) -> tf.Tensor: """Predict finite difference coefficients with a neural networks. Args: inputs: float32 Tensor with dimensions [batch, x]. hparams: model hyperparameters. reuse: whether or not to reuse TensorFlow variables. Returns: Float32 Tensor with dimensions [batch, x, derivative, coefficient]. Raises: ValueError: if inputs does not have the expected size for the equation. ValueError: if polynomial accuracy constraints are infeasible. """ # TODO(shoyer): refactor to use layer classes to hold variables, like # tf.keras.layers, instead of relying on reuse. _, equation = equations.from_hparams(hparams) assert_consistent_solution(equation, inputs) with tf.variable_scope('predict_coefficients', reuse=reuse): num_derivatives = len(equation.DERIVATIVE_ORDERS) grid = polynomials.regular_grid( equation.GRID_OFFSET, derivative_order=0, accuracy_order=hparams.coefficient_grid_min_size, dx=equation.grid.solution_dx) net = inputs[:, :, tf.newaxis] net /= equation.standard_deviation activation = _NONLINEARITIES[hparams.nonlinearity] for _ in range(hparams.num_layers - 1): net = layers.conv1d_periodic_layer(net, filters=hparams.filter_size, kernel_size=hparams.kernel_size, activation=activation, center=True) if not hparams.polynomial_accuracy_order: if hparams.num_layers == 0: raise NotImplementedError net = layers.conv1d_periodic_layer(net, filters=num_derivatives * grid.size, kernel_size=hparams.kernel_size, activation=None, center=True) new_dims = [num_derivatives, grid.size] outputs = tf.reshape( net, tf.concat([tf.shape(inputs), new_dims], axis=0)) outputs.set_shape(inputs.shape[:2].concatenate(new_dims)) if hparams.ensure_unbiased_coefficients: if 0 in equation.DERIVATIVE_ORDERS: raise ValueError( 'ensure_unbiased not yet supported for 0th order ' 'spatial derivatives') outputs -= tf.reduce_mean(outputs, axis=-1, keepdims=True) else: poly_accuracy_layers = [] for derivative_order in equation.DERIVATIVE_ORDERS: method = FINITE_VOL if equation.CONSERVATIVE else FINITE_DIFF poly_accuracy_layers.append( polynomials.PolynomialAccuracyLayer( grid=grid, method=method, derivative_order=derivative_order, accuracy_order=hparams.polynomial_accuracy_order, out_scale=hparams.polynomial_accuracy_scale)) input_sizes = [layer.input_size for layer in poly_accuracy_layers] if hparams.num_layers > 0: net = layers.conv1d_periodic_layer( net, filters=sum(input_sizes), kernel_size=hparams.kernel_size, activation=None, center=True) else: initializer = tf.initializers.zeros() coefficients = tf.get_variable('coefficients', (sum(input_sizes), ), initializer=initializer) net = tf.tile(coefficients[tf.newaxis, tf.newaxis, :], [tf.shape(inputs)[0], inputs.shape[1].value, 1]) cum_sizes = np.cumsum(input_sizes) starts = [0] + cum_sizes[:-1].tolist() stops = cum_sizes.tolist() zipped = zip(starts, stops, poly_accuracy_layers) outputs = tf.stack([ layer.apply(net[..., start:stop]) for start, stop, layer in zipped ], axis=-2) assert outputs.shape.as_list()[-1] == grid.size return outputs
def model_inputs(fine_inputs: tf.Tensor, hparams: tf.contrib.training.HParams, evaluation: bool = False) -> Dict[str, tf.Tensor]: """Create coarse model inputs from high resolution simulations. Args: fine_inputs: float32 Tensor with shape [batch, x] with results of high-resolution simulations. hparams: model hyperparameters. evaluation: bool indicating whether to create data for evaluation or model training. Returns: Dict of tensors with entries: - 'labels': float32 Tensor with shape [batch, x//factor, derivative] with finite difference derivatives computed at high resolution. - 'baseline': float32 Tensor with shape [batch, x//factor, derivative] with finite difference derivatives computed from low resolution inputs. - 'inputs': float32 Tensor with shape [batch, x//factor] with low resolution inputs. """ fine_equation, coarse_equation = equations.from_hparams(hparams) assert fine_equation.grid.resample_factor == 1 resample_method = 'mean' if coarse_equation.CONSERVATIVE else 'subsample' resample = duckarray.RESAMPLE_FUNCS[resample_method] if evaluation: ground_truth_order = None else: if hparams.ground_truth_order == -1: ground_truth_order = None else: ground_truth_order = hparams.ground_truth_order fine_derivatives = baseline_result(fine_inputs, fine_equation, hparams.num_time_steps, accuracy_order=ground_truth_order) labels = resample(fine_derivatives, factor=hparams.resample_factor, axis=1) coarse_inputs = resample(fine_inputs, factor=hparams.resample_factor, axis=1) baseline = baseline_result(coarse_inputs, coarse_equation, hparams.num_time_steps, accuracy_order=1) if not evaluation and hparams.noise_probability: if hparams.noise_type == 'white': filtered = False elif hparams.noise_type == 'filtered': filtered = True else: raise ValueError('invalid noise_type: {}'.format( hparams.noise_type)) coarse_inputs = apply_noise(coarse_inputs, hparams.noise_probability, hparams.noise_amplitude, filtered=filtered) return {'labels': labels, 'baseline': baseline, 'inputs': coarse_inputs}