<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Train a model for time-integration
This notebook shows an example of training a simple model (and model evaluation) for time-integration with Burgers' equation.
We trained many more models in our paper, but hopefully this gives an idea of what this is like.
Copyright 2019 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
Step1:  Run training
To help ensure this example runs quickly, we train for only 100000 steps, even though in our paper we used 40000 steps.
Step2:  Run evaluation
One key parameter here is the "warmup" time cutoff, which we use to ensure that we are only asking the neural network to make predictions on fully developed solutions, after all transiants are removed. We used warmup=10 for Burgers, warmup=100 for KS and warmup=50 for KdV.
Step3:  See also ks_spectral.nc and kdv_spectral.nc in the same directory for reference simulations with KS and KdV equations.
Step4:  Very simple evaluation
Evaluations have been saved to burgers-checkpoints/results.nc, but we'll download them from cloud storage instead
Step5:  An example solution from our reference model, at high resolution
Step6:  Coarse-grained simulation with our neural network
Step7:  Difference between the neural network results and coarse-grained reference results
<ASSISTANT_TASK:>
Python Code:
! pip install -q -U xarray matplotlib
! rm -rf data-driven-discretization-1d
! git clone https://github.com/google/data-driven-discretization-1d.git
! pip install -q -e data-driven-discretization-1d
# install the seaborn bug-fix from https://github.com/mwaskom/seaborn/pull/1602
! pip install -U -q git+git://github.com/stfnrpplngr/seaborn.git@309a9de383fac4db1c66dbf87815c4ba0c439c59
# Ensure we're using Tensorflow 1.x in Colab. If not using Colab, remove this magic.
%tensorflow_version 1.x
import tensorflow as tf
assert tf.__version__[:2] == '1.'
Explanation: Train a model for time-integration
This notebook shows an example of training a simple model (and model evaluation) for time-integration with Burgers' equation.
We trained many more models in our paper, but hopefully this gives an idea of what this is like.
Copyright 2019 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.
Setup environment
End of explanation
! gsutil cp gs://data-driven-discretization-public/training-data/burgers.h5 .
%%time
! python data-driven-discretization-1d/pde_superresolution/scripts/run_training.py \
--checkpoint_dir burgers-checkpoints \
--equation burgers \
--hparams resample_factor=16,learning_stops=[5000,10000] \
--input_path burgers.h5
Explanation: Run training
To help ensure this example runs quickly, we train for only 100000 steps, even though in our paper we used 40000 steps.
End of explanation
# Use pre-computed "exact" solution from WENO.
# You could also run this yourself using scripts/create_exact_data.py
! gsutil cp gs://data-driven-discretization-public/time-evolution/exact/burgers_weno.nc .
Explanation: Run evaluation
One key parameter here is the "warmup" time cutoff, which we use to ensure that we are only asking the neural network to make predictions on fully developed solutions, after all transiants are removed. We used warmup=10 for Burgers, warmup=100 for KS and warmup=50 for KdV.
End of explanation
import xarray
# remove extra samples, so evaluation runs faster
reference = xarray.open_dataset('burgers_weno.nc').isel(sample=slice(10)).load()
reference.to_netcdf('burgers_weno_10samples.nc')
%%time
! python data-driven-discretization-1d/pde_superresolution/scripts/run_evaluation.py \
--checkpoint_dir burgers-checkpoints \
--exact_solution_path burgers_weno_10samples.nc \
--equation_name burgers \
--stop_times "[10]" \
--num_samples 10 \
--warmup 10 \
--time_delta 0.1 \
--time_max 50 \
--logtostderr
Explanation: See also ks_spectral.nc and kdv_spectral.nc in the same directory for reference simulations with KS and KdV equations.
End of explanation
! gsutil cp gs://data-driven-discretization-public/time-evolution/model/burgers_16x_samples.nc .
import xarray
results = xarray.open_dataset('burgers_16x_samples.nc').load()
results
reference
Explanation: Very simple evaluation
Evaluations have been saved to burgers-checkpoints/results.nc, but we'll download them from cloud storage instead:
End of explanation
reference.y[0].sel(time=slice(10, 60)).plot.imshow()
Explanation: An example solution from our reference model, at high resolution:
End of explanation
results.y[0].plot.imshow()
Explanation: Coarse-grained simulation with our neural network:
End of explanation
(results.y.sel(sample=0)
 - reference.y.sel(sample=0, time=slice(10, 60)).coarsen(x=16).mean()
   .assign_coords(x=results.x)).plot.imshow()
Explanation: Difference between the neural network results and coarse-grained reference results:
End of explanation
<END_TASK>