text: # python3
  # Copyright 2018 Google LLC
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Functions for evaluation of various state quantities and saving data.
  
  This module provides functions for evaluation of common quantities and
  IO functions for writing dataset and metadata.
  """
  import json
  import operator
  from typing import Any, Dict, List, Mapping, Tuple, TypeVar
  
  import numpy as np
  from datadrivenpdes.core import equations
  from datadrivenpdes.core import grids
  from datadrivenpdes.core import integrate
  from datadrivenpdes.core import models
  from datadrivenpdes.core import states
  from datadrivenpdes.core import tensor_ops
  from datadrivenpdes.core import utils
  import tensorflow as tf
  from tensorflow.io import gfile
  
  nest = tf.nest
  
  
  KeyedTensors = Dict[states.StateDefinition, tf.Tensor]
  StateGridTensors = Dict[Tuple[states.StateDefinition, grids.Grid], tf.Tensor]
  StatisticsDict = Dict[
      Tuple[states.StateDefinition, grids.Grid], Tuple[float, float]]
  
  
  def convert_to_tf_example(state: StateGridTensors) -> bytes:
    """Generates serialized tensorflow example holding dataset_states.
  
    Generates tf.Example from the states in dataset_states. The key is generated
    from the StateDefinition and corresponding grid. Example is then serialized.
  
    Args:
      state: A dict of states to be serialized.
  
    Returns:
      tf.Example serialized to a string.
  
    Raises:
      ValueError: The number of states doesn't equal to the number of grids.
    """
    def _floats_feature(value: np.ndarray) -> tf.train.Feature:
      """Helper function to convert list of floats to a tf.train.Feature."""
      return tf.train.Feature(
          float_list=tf.train.FloatList(value=np.ravel(value)))
  
    feature = {}
    for (state_key, grid), tensor in state.items():
      feature[utils.component_name(state_key, grid)] = _floats_feature(tensor)  # pytype: disable=wrong-arg-types
    example = tf.train.Example(features=tf.train.Features(feature=feature))
    return example.SerializeToString()
  
  
  T = TypeVar('T')
  
  
  def unstack_dict(
      state: Mapping[T, tf.Tensor],
      num: int,
  ) -> List[Dict[T, tf.Tensor]]:
    """Unstack keyed tensors into a list."""
    unstacked = tf.nest.map_structure(
        lambda x: tf.unstack(x, num), state)
    result = [{} for _ in range(num)]
    for k, tensor_list in unstacked.items():
      for i, tensor in enumerate(tensor_list):
        result[i][k] = tensor
    return result
  
  
  K = TypeVar('K')
  V = TypeVar('V')
  
  
  def merge(entries: List[Tuple[K, V]]) -> Dict[K, V]:
    """Merge a list of key-value pairs into a dict, raising if keys are reused."""
    result = {}
    for k, v in entries:
      if k in result:
        raise ValueError('duplicate entries with key {}'.format(k))
      result[k] = v
    return result
  
  
  class Builder:
    """Defines interface for building datasets.
  
    Defines an interface that datasets must implement to be compatible with the
    dataset generation and pipeline.
  
    We use a two stage process for generating data:
    1. Create random states, and integrate them for a long number of time steps.
    2. Sample short movies from the long time trajectory at regular intervals.
  
    For the sake of simplicity, we do time integration again in step (2), but this
    is not strictly necessary.
    """
  
    def __init__(
        self,
        equation: equations.Equation,
        simulation_grid: grids.Grid,
        output_grid: grids.Grid,
        initial_condition_steps: np.ndarray,
        example_num_time_steps: int,
    ):
      """Builder constructor."""
      self.equation = equation
      self.simulation_grid = simulation_grid
      self.output_grid = output_grid
      self.initial_condition_steps = initial_condition_steps
      self.example_num_time_steps = example_num_time_steps
  
      time_step_ratio = (equation.get_time_step(output_grid)
                         / equation.get_time_step(simulation_grid))
      self.time_step_ratio = round(time_step_ratio)
      if self.time_step_ratio != time_step_ratio:
        raise ValueError(
            f'time step ratio must be an integer: {time_step_ratio}')
  
    @property
    def model(self):
      return models.FiniteDifferenceModel(self.equation, self.simulation_grid)
  
    @property
    def coarse_model(self):
      return models.FiniteDifferenceModel(self.equation, self.output_grid)
  
    def integrate_for_initial_conditions(
        self, state: Dict[str, tf.Tensor]) -> List[Dict[str, tf.Tensor]]:
      """Integrate from random seeds for initial conditions."""
      assert tf.executing_eagerly()
      integrated = integrate.integrate_steps(
          self.model, state, self.initial_condition_steps, axis=0)
      return unstack_dict(integrated, len(self.initial_condition_steps))
  
    def integrate_each_example(
        self, state: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      """Integrate each initial condition to produce examples for training."""
      assert tf.executing_eagerly()
      steps = self.time_step_ratio * np.arange(self.example_num_time_steps)
      return integrate.integrate_steps(self.model, state, steps, axis=0)
  
    def postprocess(
        self, state: Dict[str, tf.Tensor]) -> StateGridTensors:
      """Post-process integrated states."""
      raise NotImplementedError
  
    def convert_to_tf_example(self, state: StateGridTensors) -> bytes:
      """Converts input_data to serialized tf.Example protos."""
      assert tf.executing_eagerly()
      return convert_to_tf_example(state)
  
    def save_metadata(
        self,
        statistics_dict: StatisticsDict,
        records_path: str,
        metadata_path: str,
        num_shards: int,
        flags: Mapping[str, Any],
    ) -> None:
      """Saves dataset metadata to `path`."""
      data_components = []
      for (key, grid), (mean, variance) in statistics_dict.items():
        data_components.append(
            dict(
                grid=grid.to_config(),
                state_definition=key.to_config(),
                statistics=dict(mean=mean, variance=variance),
            )
        )
  
      # TODO(shoyer): use official APIs for this
      file_names = [
          records_path +
          '-{0:0{width}}-of-{1:0{width}}'.format(shard, num_shards, width=5)
          for shard in range(num_shards)]
  
      config = dict(
          components=data_components,
          file_names=file_names,
          equation=self.equation.to_config(),
          model=self.model.to_config(),
          output_grid=self.output_grid.to_config(),
          simulation_grid=self.simulation_grid.to_config(),
          initial_condition_steps=self.initial_condition_steps.tolist(),
          example_num_time_steps=self.example_num_time_steps,
          flags=flags,
      )
      with gfile.GFile(metadata_path, 'w') as f:
        f.write(json.dumps(config, indent=4, sort_keys=True))
  
  
  class TimeDerivatives(Builder):
    """Provides functions to build dataset of states and time derivatives."""
  
    def postprocess(self, state: Dict[str, tf.Tensor]) -> StateGridTensors:
      """Resample states to low resolution."""
      result = []
      key_defs = self.equation.key_definitions
  
      # solution
      coarse_state = self.equation.regrid(
          state, self.simulation_grid, self.output_grid)
      for k, v in coarse_state.items():
        result.append(((key_defs[k].exact(), self.output_grid), v))
  
      # time derivatives
      time_derivative = self.model.time_derivative(state)
  
      coarse_time_derivative = self.equation.regrid(
          time_derivative, self.simulation_grid, self.output_grid)
      for k, v in coarse_time_derivative.items():
        key = (key_defs[k].time_derivative().exact(), self.output_grid)
        result.append((key, v))
  
      baseline_time_derivative = self.coarse_model.time_derivative(coarse_state)
      for k, v in baseline_time_derivative.items():
        result.append(((key_defs[k].baseline(), self.output_grid), v))
  
      return merge(result)
  
  
  class AllDerivatives(Builder):
    """Provides functions to build dataset of states and all derivatives."""
  
    def postprocess(self, state: Dict[str, tf.Tensor]) -> StateGridTensors:
      """Post-process integrated states."""
      result = []
      key_defs = self.equation.key_definitions
  
      # solution
      coarse_state = self.equation.regrid(
          state, self.simulation_grid, self.output_grid)
      for k, v in coarse_state.items():
        result.append(((key_defs[k].exact(), self.output_grid), v))
  
      # spatial derivatives
      spatial_derivatives = self.model.spatial_derivatives(state)
      coarse_spatial_derivatives = self.equation.regrid(
          spatial_derivatives, self.simulation_grid, self.output_grid)
      for k, v in coarse_spatial_derivatives.items():
        if k not in state:
          result.append(((key_defs[k].exact(), self.output_grid), v))
  
      baseline_spatial_derivatives = self.coarse_model.spatial_derivatives(
          coarse_state)
      for k, v in baseline_spatial_derivatives.items():
        if k not in state:
          result.append(((key_defs[k].baseline(), self.output_grid), v))
  
      # time derivatives
      time_derivative = self.model.time_derivative(state)
  
      coarse_time_derivative = self.equation.regrid(
          time_derivative, self.simulation_grid, self.output_grid)
      for k, v in coarse_time_derivative.items():
        key = (key_defs[k].time_derivative().exact(), self.output_grid)
        result.append((key, v))
  
      baseline_time_derivative = self.coarse_model.time_derivative(coarse_state)
      for k, v in baseline_time_derivative.items():
        key = (key_defs[k].time_derivative().baseline(), self.output_grid)
        result.append((key, v))
  
      return merge(result)
  
  
  class HighResolution(Builder):
    """Provides functions to build dataset of states at high resolution."""
  
    def postprocess(self, state: Dict[str, tf.Tensor]) -> StateGridTensors:
      key_defs = self.equation.key_definitions
      return {(key_defs[k].exact(), self.simulation_grid): v
              for k, v in state.items()}
  
  
  class TimeEvolution(Builder):
    """Save the results of time-evolution."""
  
    def postprocess(self, state: Dict[str, tf.Tensor]) -> StateGridTensors:
      """Resample states to low resolution."""
      result = []
      key_defs = self.equation.key_definitions
  
      # exact solution
      coarse_state = self.equation.regrid(
          state, self.simulation_grid, self.output_grid)
      for k, v in coarse_state.items():
        result.append(((key_defs[k].exact(), self.output_grid), v))
  
      # baseline solution
      initial_coarse_state = nest.map_structure(
          operator.itemgetter(0), coarse_state)
      integrated_baseline = integrate.integrate_steps(
          self.coarse_model, initial_coarse_state,
          np.arange(self.example_num_time_steps))
      for k, v in integrated_baseline.items():
        result.append(((key_defs[k].baseline(), self.output_grid), v))
  
      return merge(result)
  
  
  DATASET_TYPES = {
      'time_derivatives': TimeDerivatives,
      'all_derivatives': AllDerivatives,
      'high_resolution': HighResolution,
      'time_evolution': TimeEvolution,
  }
  