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.
  # ==============================================================================
  """End-to-end tests for data generation and model training."""
  from absl import flags
  from absl.testing import flagsaver
  from absl.testing import parameterized
  import apache_beam as beam
  import tree
  from datadrivenpdes.core import models
  from datadrivenpdes.core import readers
  from datadrivenpdes.core import tensor_ops
  from datadrivenpdes.pipelines import create_training_data
  import tensorflow as tf
  
  from absl.testing import absltest
  
  FLAGS = flags.FLAGS
  
  
  
  
  
  class IntegrationTest(parameterized.TestCase):
  
    @classmethod
    def setUpClass(cls):
      # create training data
      output_path = FLAGS.test_tmpdir
      output_name = 'temp'
      with flagsaver.flagsaver(
          dataset_path=output_path,
          dataset_name=output_name,
          dataset_type='time_evolution',
          num_shards=1,
          total_time_steps=5,
          example_num_time_steps=8,
          time_step_interval=1,
          num_seeds=4):
        create_training_data.main([], runner=beam.runners.DirectRunner())
  
      metadata_path = '{}/{}.metadata.json'.format(output_path, output_name)
      cls.metadata = readers.load_metadata(metadata_path)
      super(IntegrationTest, cls).setUpClass()
  
    @parameterized.parameters(
        dict(model_cls=models.FiniteDifferenceModel),
        dict(model_cls=models.LinearModel),
        dict(model_cls=models.PseudoLinearModel),
        dict(model_cls=models.NonlinearModel),
        dict(model_cls=models.DirectModel),
    )
    def test_training(self, model_cls):
      # a basic integration test
      equation = readers.get_equation(self.metadata)
      grid = readers.get_output_grid(self.metadata)
      model = model_cls(equation, grid)
  
      def create_inputs(state):
        inputs = tree.map_structure(lambda x: x[:-1], state)
        labels = state['concentration'][1:]
        return inputs, labels
  
      training_data = (
          model.load_data(self.metadata)
          .repeat()
          .shuffle(10)
          .map(create_inputs)
      )
      model.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                    loss='mean_squared_error')
      model.fit(training_data, epochs=1, steps_per_epoch=5)
      model.evaluate(training_data, steps=2)
  
    @parameterized.parameters(
        dict(model_cls=models.LinearModel),
        dict(model_cls=models.PseudoLinearModel),
        dict(model_cls=models.NonlinearModel),
    )
    def test_training_multiple_times(self, model_cls):
      # a basic integration test
      equation = readers.get_equation(self.metadata)
      grid = readers.get_output_grid(self.metadata)
      model = model_cls(equation, grid, num_time_steps=4)
  
      def create_inputs(state):
        # (batch, x, y)
        inputs = tree.map_structure(lambda x: x[:-model.num_time_steps], state)
        # (batch, time, x, y)
        labels = tensor_ops.stack_all_contiguous_slices(
            state['concentration'][1:], model.num_time_steps, new_axis=1)
        return inputs, labels
  
      training_data = (
          model.load_data(self.metadata)
          .map(create_inputs)
          .apply(tf.data.experimental.unbatch())
          .shuffle(10)
          .repeat()
          .batch(4, drop_remainder=True)
          .prefetch(1)
      )
      model.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                    loss='mean_squared_error')
      model.fit(training_data, epochs=1, steps_per_epoch=5)
      model.evaluate(training_data, steps=2)
  
  
  if __name__ == '__main__':
    absltest.main()
  