Practical Exam WS22/23#

Introduction to Scientific Machine Learning for Engineers#

Task: PDE Regression

We will use one of the recent PDE benchmarking projects PDEArena to do regression of flow fields. The task is to autoregressively predict the next flow fields of a Computational Fluid Dynamics simulation (in this case incompressible Navier-Stokes) from the current flow fields. The flow fields comprises of:

  • scalar field \(u\),

  • 2D velocity field represented by its components in x and y directions \(\mathbf{v}_x\) and \(\mathbf{v}_y\) respectively.

For a more detailed problem definition, you might want to check the experiments section of the reference paper Towards Multi-spatiotemporal-scale Generalized PDE Modeling (Section 4, specifically Velocity function formulation of Navier-Stokes equations) accompanying PDEArena.


Your tasks will be as follows:

  1. [28 Pts] Dimensionality Reduction and Regression
    In this task, we ask you to transform the data to a lower dimensional space and do simple regression on this low-dimensional space.

    1. [4 Pts] Argue which dimensionality reduction and regression approaches are suitable for this problem.

    2. [12 Pts] Implement the dimensionality reduction method, fit its parameters, and demonstrate its performance qualitatively + quantitatively.

    3. [12 Pts] Implement the regression model, fit its parameters, and demonstrate its performance qualitatively + quantitatively.

  2. [22 Pts] CNN Regression
    CNNsco have proven their effectiveness on image classification and segmentation tasks. Here, we ask you to apply a CNN to the same regression problem from the previous task.

    1. [4 Pts] Argue which properties a CNN needs to have for this task. Manually compute the shape of the latent feature maps in each hidden layer of the CNN you chose.

    2. [8 Pts] Implement the CNN-based regression model.

    3. [10 Pts] Apply 5 tricks for performance improvement of your model. For each of them, train the model and demonstrate its performance qualitatively and quantitatively. If the performance drops after some of these tricks, explain why this could happen.

  3. [10 Pts] Benchmarking
    For comparison, train on our data one of the state-of-the-art (SOTA) models DilResNet-128 (see Towards Multi-spatiotemporal-scale Generalized PDE Modeling (Section 3. PDE Surrogates, specifically (Dilated) ResNet) provided in the benchmarking repository.

    1. [2 Pts] Train the DilResNet-128 model on our data. Try a few different hyperparameter configurations.

    2. [8 Pts] Compare the performance of both regression approaches you implemented against the SOTA model in a table (including number of learnable parameters and one-step MSE). Discuss the performance and characteristics of all three methods. Then, make a suggestion for further improvements on each of the approaches.

Note: Tuning hyperparameters is not the core of these tasks. However, your should demonstrate somewhat plausible results.

0 Setup#

Everything needed before you can start with the actual tasks.

0.1 Installation#

Choose one of the two installation options below and stick to it, i.e. don’t mix the two precedures.

0.1.1 Conda Installation#

If you want to install the project locally, then please use Conda and follow the steps from the official installation instructions here: https://microsoft.github.io/pdearena/install/.

0.1.2 Colab Installation#

# install pdearena

# clone the PDEArena repository
!rm -r *
!git clone https://github.com/microsoft/pdearena && cd pdearena &&  git checkout c8d1ea8880872ad57e1f0fdf014c2d3283628a0f

# move everything to root
!mv pdearena pdearena_temp && mv pdearena_temp/* . && rm -r pdearena_temp
# data generation capable installation
!pip install -e ".[datagen]"
Cloning into 'pdearena'...
remote: Enumerating objects: 300, done.
remote: Counting objects: 100% (152/152), done.
remote: Compressing objects: 100% (91/91), done.
remote: Total 300 (delta 55), reused 95 (delta 42), pack-reused 148
Receiving objects: 100% (300/300), 3.43 MiB | 16.72 MiB/s, done.
Resolving deltas: 100% (71/71), done.
Note: switching to 'c8d1ea8880872ad57e1f0fdf014c2d3283628a0f'.

You are in 'detached HEAD' state. You can look around, make experimental
changes and commit them, and you can discard any commits you make in this
state without impacting any branches by switching back to a branch.

If you want to create a new branch to retain commits you create, you may
do so (now or later) by using -c with the switch command. Example:

  git switch -c <new-branch-name>

Or undo this operation with:

  git switch -

Turn off this advice by setting config variable advice.detachedHead to false

HEAD is now at c8d1ea8 scaling should be other way
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Obtaining file:///content
  Preparing metadata (setup.py) ... ?25l?25hdone
Collecting pytorch-lightning>=1.7
  Downloading pytorch_lightning-2.0.0-py3-none-any.whl (715 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 715.6/715.6 KB 7.1 MB/s eta 0:00:00
?25hRequirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (1.22.4)
Requirement already satisfied: xarray in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (2022.12.0)
Requirement already satisfied: h5py in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (3.8.0)
Requirement already satisfied: click in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (8.1.3)
Requirement already satisfied: torch>=1.12 in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (1.13.1+cu116)
Collecting torchdata
  Downloading torchdata-0.6.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.6 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.6/4.6 MB 67.3 MB/s eta 0:00:00
?25hRequirement already satisfied: matplotlib in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (3.7.1)
Collecting jsonargparse
  Downloading jsonargparse-4.20.0-py3-none-any.whl (183 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 183.6/183.6 KB 16.8 MB/s eta 0:00:00
?25hCollecting omegaconf
  Downloading omegaconf-2.3.0-py3-none-any.whl (79 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 79.5/79.5 KB 7.7 MB/s eta 0:00:00
?25hRequirement already satisfied: tensorboard in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (2.11.2)
Requirement already satisfied: pytest in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (3.6.4)
Collecting pytest-mock
  Downloading pytest_mock-3.10.0-py3-none-any.whl (9.3 kB)
Collecting pytest-explicit
  Downloading pytest_explicit-1.0.1-py3-none-any.whl (4.4 kB)
Collecting phiflow==2.1
  Downloading phiflow-2.1.0.tar.gz (260 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 260.2/260.2 KB 22.3 MB/s eta 0:00:00
?25h  Preparing metadata (setup.py) ... ?25l?25hdone
Collecting fdtd==0.2.5
  Downloading fdtd-0.2.5-py3-none-any.whl (45 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 45.1/45.1 KB 4.2 MB/s eta 0:00:00
?25hRequirement already satisfied: joblib in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (1.1.1)
Collecting juliapkg
  Downloading juliapkg-0.1.10-py3-none-any.whl (15 kB)
Requirement already satisfied: tqdm in /usr/local/lib/python3.9/dist-packages (from pdearena==0.1.1) (4.65.0)
Requirement already satisfied: scipy in /usr/local/lib/python3.9/dist-packages (from fdtd==0.2.5->pdearena==0.1.1) (1.10.1)
Requirement already satisfied: typing-extensions>=4.0.0 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning>=1.7->pdearena==0.1.1) (4.5.0)
Requirement already satisfied: packaging>=17.1 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning>=1.7->pdearena==0.1.1) (23.0)
Collecting lightning-utilities>=0.7.0
  Downloading lightning_utilities-0.8.0-py3-none-any.whl (20 kB)
Collecting torchmetrics>=0.7.0
  Downloading torchmetrics-0.11.4-py3-none-any.whl (519 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 519.2/519.2 KB 17.8 MB/s eta 0:00:00
?25hRequirement already satisfied: fsspec[http]>2021.06.0 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning>=1.7->pdearena==0.1.1) (2023.3.0)
Requirement already satisfied: PyYAML>=5.4 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning>=1.7->pdearena==0.1.1) (6.0)
Collecting semantic-version~=2.9
  Downloading semantic_version-2.10.0-py2.py3-none-any.whl (15 kB)
Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (2.8.2)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (0.11.0)
Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (1.4.4)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (3.0.9)
Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (1.0.7)
Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (8.4.0)
Requirement already satisfied: importlib-resources>=3.2.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (5.12.0)
Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.9/dist-packages (from matplotlib->pdearena==0.1.1) (4.39.0)
Collecting antlr4-python3-runtime==4.9.*
  Downloading antlr4-python3-runtime-4.9.3.tar.gz (117 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 117.0/117.0 KB 8.4 MB/s eta 0:00:00
?25h  Preparing metadata (setup.py) ... ?25l?25hdone
Requirement already satisfied: py>=1.5.0 in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (1.11.0)
Requirement already satisfied: setuptools in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (63.4.3)
Requirement already satisfied: attrs>=17.4.0 in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (22.2.0)
Requirement already satisfied: more-itertools>=4.0.0 in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (9.1.0)
Requirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (1.15.0)
Requirement already satisfied: pluggy<0.8,>=0.5 in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (0.7.1)
Requirement already satisfied: atomicwrites>=1.0 in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (1.4.1)
Collecting pytest
  Downloading pytest-7.2.2-py3-none-any.whl (317 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 317.2/317.2 KB 21.1 MB/s eta 0:00:00
?25hCollecting iniconfig
  Downloading iniconfig-2.0.0-py3-none-any.whl (5.9 kB)
Collecting pytest
  Downloading pytest-7.2.1-py3-none-any.whl (317 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 317.1/317.1 KB 12.1 MB/s eta 0:00:00
?25h  Downloading pytest-7.2.0-py3-none-any.whl (316 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 316.8/316.8 KB 12.5 MB/s eta 0:00:00
?25h  Downloading pytest-7.1.3-py3-none-any.whl (298 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 298.2/298.2 KB 18.2 MB/s eta 0:00:00
?25h  Downloading pytest-7.1.2-py3-none-any.whl (297 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 297.0/297.0 KB 27.2 MB/s eta 0:00:00
?25h  Downloading pytest-7.1.1-py3-none-any.whl (297 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 297.0/297.0 KB 12.9 MB/s eta 0:00:00
?25h  Downloading pytest-7.1.0-py3-none-any.whl (297 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 297.0/297.0 KB 15.1 MB/s eta 0:00:00
?25h  Downloading pytest-7.0.1-py3-none-any.whl (296 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 297.0/297.0 KB 20.3 MB/s eta 0:00:00
?25h  Downloading pytest-7.0.0-py3-none-any.whl (296 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 296.5/296.5 KB 8.6 MB/s eta 0:00:00
?25h  Downloading pytest-6.2.5-py3-none-any.whl (280 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 280.7/280.7 KB 11.3 MB/s eta 0:00:00
?25hRequirement already satisfied: toml in /usr/local/lib/python3.9/dist-packages (from pytest->pdearena==0.1.1) (0.10.2)
  Downloading pytest-6.2.4-py3-none-any.whl (280 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 280.6/280.6 KB 23.3 MB/s eta 0:00:00
?25h  Downloading pytest-6.2.3-py3-none-any.whl (280 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 280.5/280.5 KB 16.2 MB/s eta 0:00:00
?25h  Downloading pytest-6.2.2-py3-none-any.whl (280 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 280.1/280.1 KB 15.7 MB/s eta 0:00:00
?25h  Downloading pytest-6.2.1-py3-none-any.whl (279 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 279.9/279.9 KB 10.0 MB/s eta 0:00:00
?25h  Downloading pytest-6.2.0-py3-none-any.whl (279 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 279.6/279.6 KB 11.7 MB/s eta 0:00:00
?25h  Downloading pytest-6.1.2-py3-none-any.whl (272 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 272.4/272.4 KB 14.8 MB/s eta 0:00:00
?25h  Downloading pytest-6.1.1-py3-none-any.whl (272 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 272.2/272.2 KB 13.0 MB/s eta 0:00:00
?25h  Downloading pytest-6.1.0-py3-none-any.whl (272 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 272.3/272.3 KB 18.0 MB/s eta 0:00:00
?25h  Downloading pytest-6.0.2-py3-none-any.whl (270 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 270.8/270.8 KB 15.9 MB/s eta 0:00:00
?25h  Downloading pytest-6.0.1-py3-none-any.whl (270 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 270.6/270.6 KB 4.6 MB/s eta 0:00:00
?25h  Downloading pytest-6.0.0-py3-none-any.whl (270 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 270.6/270.6 KB 23.8 MB/s eta 0:00:00
?25h  Downloading pytest-5.4.3-py3-none-any.whl (248 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 248.1/248.1 KB 15.1 MB/s eta 0:00:00
?25h  Downloading pytest-5.4.2-py3-none-any.whl (247 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 247.4/247.4 KB 12.6 MB/s eta 0:00:00
?25h  Downloading pytest-5.4.1-py3-none-any.whl (246 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 246.7/246.7 KB 9.8 MB/s eta 0:00:00
?25h  Downloading pytest-5.4.0-py3-none-any.whl (247 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 247.0/247.0 KB 21.6 MB/s eta 0:00:00
?25h  Downloading pytest-5.3.5-py3-none-any.whl (235 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 235.2/235.2 KB 6.3 MB/s eta 0:00:00
?25h  Downloading pytest-5.3.4-py3-none-any.whl (235 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 235.3/235.3 KB 3.2 MB/s eta 0:00:00
?25h  Downloading pytest-5.3.3-py3-none-any.whl (235 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 235.5/235.5 KB 13.4 MB/s eta 0:00:00
?25h  Downloading pytest-5.3.2-py3-none-any.whl (234 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 234.5/234.5 KB 17.3 MB/s eta 0:00:00
?25h  Downloading pytest-5.3.1-py3-none-any.whl (233 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 233.9/233.9 KB 14.5 MB/s eta 0:00:00
?25h  Downloading pytest-5.3.0-py3-none-any.whl (233 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 233.7/233.7 KB 10.4 MB/s eta 0:00:00
?25h  Downloading pytest-5.2.4-py3-none-any.whl (227 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 227.4/227.4 KB 6.2 MB/s eta 0:00:00
?25h  Downloading pytest-5.2.3-py3-none-any.whl (227 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 227.4/227.4 KB 8.4 MB/s eta 0:00:00
?25h  Downloading pytest-5.2.2-py3-none-any.whl (227 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 227.3/227.3 KB 7.7 MB/s eta 0:00:00
?25h  Downloading pytest-5.2.1-py3-none-any.whl (226 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 226.8/226.8 KB 10.7 MB/s eta 0:00:00
?25h  Downloading pytest-5.2.0-py3-none-any.whl (226 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 226.7/226.7 KB 9.2 MB/s eta 0:00:00
?25h  Downloading pytest-5.1.3-py3-none-any.whl (224 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 224.1/224.1 KB 16.5 MB/s eta 0:00:00
?25h  Downloading pytest-5.1.2-py3-none-any.whl (224 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 224.2/224.2 KB 21.5 MB/s eta 0:00:00
?25h  Downloading pytest-5.1.1-py3-none-any.whl (223 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 223.6/223.6 KB 9.1 MB/s eta 0:00:00
?25h  Downloading pytest-5.1.0-py3-none-any.whl (223 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 223.6/223.6 KB 14.0 MB/s eta 0:00:00
?25h  Downloading pytest-5.0.1-py3-none-any.whl (221 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 221.8/221.8 KB 10.0 MB/s eta 0:00:00
?25h  Downloading pytest-5.0.0-py3-none-any.whl (221 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 221.8/221.8 KB 21.1 MB/s eta 0:00:00
?25hINFO: pip is looking at multiple versions of pytest-mock to determine which version is compatible with other requirements. This could take a while.
Collecting pytest-mock
  Downloading pytest_mock-3.9.0-py3-none-any.whl (9.1 kB)
  Downloading pytest_mock-3.8.2-py3-none-any.whl (9.1 kB)
  Downloading pytest_mock-3.8.1-py3-none-any.whl (9.1 kB)
  Downloading pytest_mock-3.8.0-py3-none-any.whl (9.1 kB)
  Downloading pytest_mock-3.7.0-py3-none-any.whl (12 kB)
  Downloading pytest_mock-3.6.1-py3-none-any.whl (12 kB)
  Downloading pytest_mock-3.6.0-py3-none-any.whl (12 kB)
INFO: pip is looking at multiple versions of pytest-mock to determine which version is compatible with other requirements. This could take a while.
  Downloading pytest_mock-3.5.1-py3-none-any.whl (12 kB)
  Downloading pytest_mock-3.5.0-py3-none-any.whl (12 kB)
  Downloading pytest_mock-3.4.0-py3-none-any.whl (11 kB)
  Downloading pytest_mock-3.3.1-py3-none-any.whl (11 kB)
  Downloading pytest_mock-3.3.0-py3-none-any.whl (11 kB)
INFO: This is taking longer than usual. You might need to provide the dependency resolver with stricter constraints to reduce runtime. See https://pip.pypa.io/warnings/backtracking for guidance. If you want to abort this run, press Ctrl + C.
  Downloading pytest_mock-3.2.0-py3-none-any.whl (10 kB)
Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (3.4.1)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (0.6.1)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (1.8.1)
Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (1.4.0)
Requirement already satisfied: protobuf<4,>=3.9.2 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (3.19.6)
Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (2.2.3)
Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (2.27.1)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (0.4.6)
Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (2.16.2)
Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (0.40.0)
Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.9/dist-packages (from tensorboard->pdearena==0.1.1) (1.51.3)
Requirement already satisfied: urllib3>=1.25 in /usr/local/lib/python3.9/dist-packages (from torchdata->pdearena==0.1.1) (1.26.15)
Collecting torch>=1.12
  Downloading torch-2.0.0-cp39-cp39-manylinux1_x86_64.whl (619.9 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 619.9/619.9 MB 2.6 MB/s eta 0:00:00
?25hCollecting nvidia-nccl-cu11==2.14.3
  Downloading nvidia_nccl_cu11-2.14.3-py3-none-manylinux1_x86_64.whl (177.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 177.1/177.1 MB 7.6 MB/s eta 0:00:00
?25hRequirement already satisfied: networkx in /usr/local/lib/python3.9/dist-packages (from torch>=1.12->pdearena==0.1.1) (3.0)
Collecting nvidia-cufft-cu11==10.9.0.58
  Downloading nvidia_cufft_cu11-10.9.0.58-py3-none-manylinux1_x86_64.whl (168.4 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 168.4/168.4 MB 7.7 MB/s eta 0:00:00
?25hRequirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from torch>=1.12->pdearena==0.1.1) (3.10.0)
Collecting triton==2.0.0
  Downloading triton-2.0.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (63.3 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 63.3/63.3 MB 12.2 MB/s eta 0:00:00
?25hCollecting nvidia-cusparse-cu11==11.7.4.91
  Downloading nvidia_cusparse_cu11-11.7.4.91-py3-none-manylinux1_x86_64.whl (173.2 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 173.2/173.2 MB 7.5 MB/s eta 0:00:00
?25hCollecting nvidia-cuda-runtime-cu11==11.7.99
  Downloading nvidia_cuda_runtime_cu11-11.7.99-py3-none-manylinux1_x86_64.whl (849 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 849.3/849.3 KB 67.8 MB/s eta 0:00:00
?25hCollecting nvidia-cusolver-cu11==11.4.0.1
  Downloading nvidia_cusolver_cu11-11.4.0.1-2-py3-none-manylinux1_x86_64.whl (102.6 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 102.6/102.6 MB 10.5 MB/s eta 0:00:00
?25hCollecting nvidia-cuda-cupti-cu11==11.7.101
  Downloading nvidia_cuda_cupti_cu11-11.7.101-py3-none-manylinux1_x86_64.whl (11.8 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 11.8/11.8 MB 96.3 MB/s eta 0:00:00
?25hCollecting nvidia-cudnn-cu11==8.5.0.96
  Downloading nvidia_cudnn_cu11-8.5.0.96-2-py3-none-manylinux1_x86_64.whl (557.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 557.1/557.1 MB 3.0 MB/s eta 0:00:00
?25hCollecting nvidia-cuda-nvrtc-cu11==11.7.99
  Downloading nvidia_cuda_nvrtc_cu11-11.7.99-2-py3-none-manylinux1_x86_64.whl (21.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 21.0/21.0 MB 72.8 MB/s eta 0:00:00
?25hCollecting nvidia-nvtx-cu11==11.7.91
  Downloading nvidia_nvtx_cu11-11.7.91-py3-none-manylinux1_x86_64.whl (98 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 98.6/98.6 KB 12.9 MB/s eta 0:00:00
?25hRequirement already satisfied: jinja2 in /usr/local/lib/python3.9/dist-packages (from torch>=1.12->pdearena==0.1.1) (3.1.2)
Collecting nvidia-cublas-cu11==11.10.3.66
  Downloading nvidia_cublas_cu11-11.10.3.66-py3-none-manylinux1_x86_64.whl (317.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 317.1/317.1 MB 4.7 MB/s eta 0:00:00
?25hCollecting nvidia-curand-cu11==10.2.10.91
  Downloading nvidia_curand_cu11-10.2.10.91-py3-none-manylinux1_x86_64.whl (54.6 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 54.6/54.6 MB 17.9 MB/s eta 0:00:00
?25hRequirement already satisfied: sympy in /usr/local/lib/python3.9/dist-packages (from torch>=1.12->pdearena==0.1.1) (1.11.1)
Collecting lit
  Downloading lit-15.0.7.tar.gz (132 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 132.3/132.3 KB 18.8 MB/s eta 0:00:00
?25h  Preparing metadata (setup.py) ... ?25l?25hdone
Requirement already satisfied: cmake in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch>=1.12->pdearena==0.1.1) (3.25.2)
Requirement already satisfied: pandas>=1.3 in /usr/local/lib/python3.9/dist-packages (from xarray->pdearena==0.1.1) (1.4.4)
Collecting aiohttp!=4.0.0a0,!=4.0.0a1
  Downloading aiohttp-3.8.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.0/1.0 MB 60.9 MB/s eta 0:00:00
?25hRequirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard->pdearena==0.1.1) (5.3.0)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard->pdearena==0.1.1) (0.2.8)
Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard->pdearena==0.1.1) (4.9)
Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.9/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard->pdearena==0.1.1) (1.3.1)
Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.9/dist-packages (from importlib-resources>=3.2.0->matplotlib->pdearena==0.1.1) (3.15.0)
Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.9/dist-packages (from markdown>=2.6.8->tensorboard->pdearena==0.1.1) (6.0.0)
Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.9/dist-packages (from pandas>=1.3->xarray->pdearena==0.1.1) (2022.7.1)
Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard->pdearena==0.1.1) (2.0.12)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard->pdearena==0.1.1) (3.4)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard->pdearena==0.1.1) (2022.12.7)
Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.9/dist-packages (from werkzeug>=1.0.1->tensorboard->pdearena==0.1.1) (2.1.2)
Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.9/dist-packages (from sympy->torch>=1.12->pdearena==0.1.1) (1.3.0)
Collecting multidict<7.0,>=4.5
  Downloading multidict-6.0.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (114 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 114.2/114.2 KB 12.5 MB/s eta 0:00:00
?25hCollecting async-timeout<5.0,>=4.0.0a3
  Downloading async_timeout-4.0.2-py3-none-any.whl (5.8 kB)
Collecting aiosignal>=1.1.2
  Downloading aiosignal-1.3.1-py3-none-any.whl (7.6 kB)
Collecting frozenlist>=1.1.1
  Downloading frozenlist-1.3.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (158 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 158.8/158.8 KB 19.7 MB/s eta 0:00:00
?25hCollecting yarl<2.0,>=1.0
  Downloading yarl-1.8.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (264 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 264.6/264.6 KB 11.7 MB/s eta 0:00:00
?25hRequirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.9/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard->pdearena==0.1.1) (0.4.8)
Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.9/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard->pdearena==0.1.1) (3.2.2)
Building wheels for collected packages: phiflow, antlr4-python3-runtime, lit
  Building wheel for phiflow (setup.py) ... ?25l?25hdone
  Created wheel for phiflow: filename=phiflow-2.1.0-py3-none-any.whl size=291457 sha256=42bba1d43c9c1128eab34ff980c1a5590b2a3eb6c0cde205800af21861a15e78
  Stored in directory: /root/.cache/pip/wheels/97/ca/31/5b4619811513bea63af90681a473a46737973602650b236e54
  Building wheel for antlr4-python3-runtime (setup.py) ... ?25l?25hdone
  Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.9.3-py3-none-any.whl size=144573 sha256=2b7d7736fb0ebd1eff6106e0c2c4c715841568d6c85212b1ed394f68e4a341b5
  Stored in directory: /root/.cache/pip/wheels/23/cf/80/f3efa822e6ab23277902ee9165fe772eeb1dfb8014f359020a
  Building wheel for lit (setup.py) ... ?25l?25hdone
  Created wheel for lit: filename=lit-15.0.7-py3-none-any.whl size=90003 sha256=8fe91bcf781f8ceab3bc76796ba577d7cd17ad773b351c6f230bfd864256a346
  Stored in directory: /root/.cache/pip/wheels/b7/68/18/2ad49b416abb9139c8217c349fd9df0674da8f0d1952db2ea5
Successfully built phiflow antlr4-python3-runtime lit
Installing collected packages: lit, antlr4-python3-runtime, semantic-version, omegaconf, nvidia-nvtx-cu11, nvidia-nccl-cu11, nvidia-cusparse-cu11, nvidia-curand-cu11, nvidia-cufft-cu11, nvidia-cuda-runtime-cu11, nvidia-cuda-nvrtc-cu11, nvidia-cuda-cupti-cu11, nvidia-cublas-cu11, multidict, lightning-utilities, jsonargparse, frozenlist, async-timeout, yarl, pytest-mock, pytest-explicit, nvidia-cusolver-cu11, nvidia-cudnn-cu11, juliapkg, aiosignal, phiflow, fdtd, aiohttp, triton, torch, torchmetrics, torchdata, pytorch-lightning, pdearena
  Attempting uninstall: torch
    Found existing installation: torch 1.13.1+cu116
    Uninstalling torch-1.13.1+cu116:
      Successfully uninstalled torch-1.13.1+cu116
  Running setup.py develop for pdearena
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
torchvision 0.14.1+cu116 requires torch==1.13.1, but you have torch 2.0.0 which is incompatible.
torchtext 0.14.1 requires torch==1.13.1, but you have torch 2.0.0 which is incompatible.
torchaudio 0.13.1+cu116 requires torch==1.13.1, but you have torch 2.0.0 which is incompatible.
fastai 2.7.11 requires torch<1.14,>=1.7, but you have torch 2.0.0 which is incompatible.
Successfully installed aiohttp-3.8.4 aiosignal-1.3.1 antlr4-python3-runtime-4.9.3 async-timeout-4.0.2 fdtd-0.2.5 frozenlist-1.3.3 jsonargparse-4.20.0 juliapkg-0.1.10 lightning-utilities-0.8.0 lit-15.0.7 multidict-6.0.4 nvidia-cublas-cu11-11.10.3.66 nvidia-cuda-cupti-cu11-11.7.101 nvidia-cuda-nvrtc-cu11-11.7.99 nvidia-cuda-runtime-cu11-11.7.99 nvidia-cudnn-cu11-8.5.0.96 nvidia-cufft-cu11-10.9.0.58 nvidia-curand-cu11-10.2.10.91 nvidia-cusolver-cu11-11.4.0.1 nvidia-cusparse-cu11-11.7.4.91 nvidia-nccl-cu11-2.14.3 nvidia-nvtx-cu11-11.7.91 omegaconf-2.3.0 pdearena-0.1.1 phiflow-2.1.0 pytest-explicit-1.0.1 pytest-mock-3.2.0 pytorch-lightning-2.0.0 semantic-version-2.10.0 torch-2.0.0 torchdata-0.6.0 torchmetrics-0.11.4 triton-2.0.0 yarl-1.8.2
# Restart the runtime as prompted above to get the most recent changes.
import os
os.kill(os.getpid(), 9)
# The memory we are given in a Colab session
!df -H
Filesystem      Size  Used Avail Use% Mounted on
overlay          84G   34G   50G  41% /
tmpfs            68M     0   68M   0% /dev
shm             6.2G     0  6.2G   0% /dev/shm
/dev/root       2.1G  1.2G  882M  58% /usr/sbin/docker-init
/dev/sda1        84G   51G   33G  61% /opt/bin/.nvidia
tmpfs           6.9G   37k  6.9G   1% /var/colab
tmpfs           6.9G     0  6.9G   0% /proc/acpi
tmpfs           6.9G     0  6.9G   0% /proc/scsi
tmpfs           6.9G     0  6.9G   0% /sys/firmware
# change your runtime type to GPU. Then upon execution of the following command, 
# you should see which GPU and CUDA Version you have been assigned to.
!nvidia-smi
Sat Mar 18 21:18:31 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   64C    P0    32W /  70W |      0MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

0.2 Imports#

# some imports
import h5py
import time
import os
import copy

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib import rc
rc('animation', html='jshtml')
plt.rcParams.update({'font.size': 14})

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import datasets, models, transforms
from torch.utils.tensorboard import SummaryWriter

from pdearena.visualization import plot_scalar, plot_2dvec, plot_scalar_sequence_comparison

# Load the TensorBoard notebook extension for logging training runs
%load_ext tensorboard
/usr/local/lib/python3.9/dist-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: libtorch_cuda_cu.so: cannot open shared object file: No such file or directory
  warn(f"Failed to load image Python extension: {e}")

0.3 Dataset#

In the end, you should have the following data:

data/NSEsmoke/NavierStokes2D_train_142_0.50000_5.h5
data/NSEsmoke/NavierStokes2D_valid_242_0.50000.h5
data/NSEsmoke/pde_train_seed_42.yaml
data/NSEsmoke/pde_valid_seed_42.yaml
# Generate or load 20 training and one validation trajectories).
# Note that generating the dataset anew (takes 7-8 min) might be faster than uploading from
# Moodle. Downloading `from_dropbox` (takes 5 sec) seems to be the fastest approach.

how_to_get_data = 'from_dropbox'

# clean data dir in case previous download failed
!rm -r data/NSEsmoke

if how_to_get_data == 'create':
  # libraries for data generation
  !pip install phiflow==2.1 fdtd==0.2.5 joblib juliapkg tqdm==4.64.1

  !(python scripts/generate_data.py \
  base=pdedatagen/configs/navierstokes2dsmoke.yaml \
  experiment=smoke mode=train samples=20 seed=42 \
  pdeconfig.sample_rate=10 dirname=data/NSEsmoke \
  pdeconfig.init_args.nx=64 pdeconfig.init_args.ny=64)

  !(python scripts/generate_data.py \
  base=pdedatagen/configs/navierstokes2dsmoke.yaml \
  experiment=smoke mode=valid samples=1 seed=42 \
  pdeconfig.sample_rate=10 dirname=data/NSEsmoke \
  pdeconfig.init_args.nx=64 pdeconfig.init_args.ny=64)

elif how_to_get_data == 'from_moodle':
  # 1. download the data from Moodle 

  # 2. create a `data/NSEsmoke/` folder
  !mkdir -p data/NSEsmoke

  # 3. upload all 4 files from the NSEsmoke dataset on Moodle to `data/NSEsmoke`. 

elif how_to_get_data == 'from_dropbox':
  !mkdir data
  !wget https://www.dropbox.com/s/5gmqfimharsc5af/nse_data.tar.gz
  !tar xf nse_data.tar.gz --directory=data .
  !rm nse_data.tar.gz
rm: cannot remove 'data/NSEsmoke': No such file or directory
--2023-03-18 21:18:50--  https://www.dropbox.com/s/5gmqfimharsc5af/nse_data.tar.gz
Resolving www.dropbox.com (www.dropbox.com)... 162.125.3.18, 2620:100:6018:18::a27d:312
Connecting to www.dropbox.com (www.dropbox.com)|162.125.3.18|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: /s/raw/5gmqfimharsc5af/nse_data.tar.gz [following]
--2023-03-18 21:18:50--  https://www.dropbox.com/s/raw/5gmqfimharsc5af/nse_data.tar.gz
Reusing existing connection to www.dropbox.com:443.
HTTP request sent, awaiting response... 302 Found
Location: https://ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com/cd/0/inline/B4h8MBEKiZ9ofjQs5kQoYCkoSwSsJlsLcfBeDUh_Bu2zNAsHNNRcqsI_D5esMFRXL5kpA7l82-zdPL6FyiuLDldmhqSAeQuuublaGg_iQ8ck9T4rqOwW2dhsRRwTGpRgiiIE2nOg604Sb0usen3yfapDuShAtl82r37Joc_F4BX5Og/file# [following]
--2023-03-18 21:18:50--  https://ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com/cd/0/inline/B4h8MBEKiZ9ofjQs5kQoYCkoSwSsJlsLcfBeDUh_Bu2zNAsHNNRcqsI_D5esMFRXL5kpA7l82-zdPL6FyiuLDldmhqSAeQuuublaGg_iQ8ck9T4rqOwW2dhsRRwTGpRgiiIE2nOg604Sb0usen3yfapDuShAtl82r37Joc_F4BX5Og/file
Resolving ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com (ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com)... 162.125.3.15, 2620:100:6018:15::a27d:30f
Connecting to ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com (ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com)|162.125.3.15|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: /cd/0/inline2/B4h3J2c__dwBab4TiwZbJUjyEFH-6Y0vHOp1hjzOIdlyxaT6Ew34QwtsGBItScUS-lsU-iA6PnCdG1JDQy6qoUg1AbJbwzryHQYNQZUogYjuk6Q_nA-ZisTKVzOABkew_NHneYkN0L5SZ79X8FS24L_V4zdKjOzJmgvhf067g1m2WC3epO-Zc8abfD0N-NPQMMwoBiUCxz73k3egNwPAGx01Z-T5u8JLarn1IVIfPJfc3stidz2erYme1uTFwU5yoAdZiBMru6K8lLMAedBXklcq8lIq7eVY3cKnMKt66gCiraGKQpTBl_HKctev_nDpy8a0aui9dY0ByPQAjbLDWefMJcaVHbtGNLKwzjPNtSVdVLrk5ZF9s47KJ4GmVkQmHLjofsw7-bf5_X-hAa2kStFx3E4rWgkliFedQtqjaPmJBw/file [following]
--2023-03-18 21:18:51--  https://ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com/cd/0/inline2/B4h3J2c__dwBab4TiwZbJUjyEFH-6Y0vHOp1hjzOIdlyxaT6Ew34QwtsGBItScUS-lsU-iA6PnCdG1JDQy6qoUg1AbJbwzryHQYNQZUogYjuk6Q_nA-ZisTKVzOABkew_NHneYkN0L5SZ79X8FS24L_V4zdKjOzJmgvhf067g1m2WC3epO-Zc8abfD0N-NPQMMwoBiUCxz73k3egNwPAGx01Z-T5u8JLarn1IVIfPJfc3stidz2erYme1uTFwU5yoAdZiBMru6K8lLMAedBXklcq8lIq7eVY3cKnMKt66gCiraGKQpTBl_HKctev_nDpy8a0aui9dY0ByPQAjbLDWefMJcaVHbtGNLKwzjPNtSVdVLrk5ZF9s47KJ4GmVkQmHLjofsw7-bf5_X-hAa2kStFx3E4rWgkliFedQtqjaPmJBw/file
Reusing existing connection to ucb503934a4c955079a96dc1de82.dl.dropboxusercontent.com:443.
HTTP request sent, awaiting response... 200 OK
Length: 59686123 (57M) [application/octet-stream]
Saving to: ‘nse_data.tar.gz’

nse_data.tar.gz     100%[===================>]  56.92M   101MB/s    in 0.6s    

2023-03-18 21:18:52 (101 MB/s) - ‘nse_data.tar.gz’ saved [59686123/59686123]
# see the data files
!ls -lh data/NSEsmoke/
total 111M
-rw-rw-r-- 1 1000 1000 106M Jan 28 01:58 NavierStokes2D_train_142_0.50000_20.h5
-rw-rw-r-- 1 1000 1000 5.3M Jan 28 01:49 NavierStokes2D_valid_242_0.50000.h5
-rw-rw-r-- 1 1000 1000  338 Jan 28 01:48 pde_train_seed_42.yaml
-rw-rw-r-- 1 1000 1000  338 Jan 28 01:48 pde_valid_seed_42.yaml

0.4 Visualize the Data#

Let’s have a look at the training data.

train_path = 'data/NSEsmoke/NavierStokes2D_train_142_0.50000_20.h5'
valid_path = 'data/NSEsmoke/NavierStokes2D_valid_242_0.50000.h5'

# Let's see what is contained in the dataset

with h5py.File(train_path, 'r') as f:
  temp = f['train']

  # get the number of training trajectories, length of each trajectory, and 
  # number of pixels in x and y
  num_traj_train, len_traj, nx, ny = temp['u'].shape

  # print the names of variables contained in the dataset and their shapes
  for k, v in temp.items():
    print(f'property: {k: <10} shape: {v.shape}')
property: buo_y      shape: (20,)
property: dt         shape: (20,)
property: dx         shape: (20,)
property: dy         shape: (20,)
property: t          shape: (20, 56)
property: u          shape: (20, 56, 64, 64)
property: vx         shape: (20, 56, 64, 64)
property: vy         shape: (20, 56, 64, 64)
property: x          shape: (20, 64)
property: y          shape: (20, 64)

Out of the many properties the dataset contains, we are interested in the 2D velocity field described by vx and vy, and the scalar field u.

with h5py.File(train_path, 'r') as f:
  u = f['train']["u"]
  vx = f['train']["vx"]
  vy = f['train']["vy"]

  print("u_train shape: ", u.shape)
  assert u.shape == (num_traj_train, 56, nx, nx), "Wrong data shapes"

  train_dict_input = {'u': u[:, :-1], 'vx': vx[:, :-1], 'vy': vy[:, :-1]}
  train_dict_target = {'u': u[:, 1:], 'vx': vx[:, 1:], 'vy': vy[:, 1:]}

with h5py.File(valid_path, 'r') as f:
  u = f['valid']["u"]
  vx = f['valid']["vx"]
  vy = f['valid']["vy"]
 
  print("u_valid shape: ", u.shape)
  assert u.shape == (1, 56, nx, nx), "Wrong data shapes"
  valid_dict_input = {'u': u[:, 1:], 'vx': vx[:, 1:], 'vy': vy[:, 1:]}
  valid_dict_target = {'u': u[:, :-1], 'vx': vx[:, :-1], 'vy': vy[:, :-1]}
u_train shape:  (20, 56, 64, 64)
u_valid shape:  (1, 56, 64, 64)

We first visualize the initial state of each of the training and validation simulations

def plot_with_colorbar(fig, ax, x, title):
  im = ax.imshow(x, cmap='binary')
  divider = make_axes_locatable(ax)
  cax = divider.append_axes('bottom', size='5%', pad=0.5)
  fig.colorbar(im, cax=cax, orientation='horizontal')
  ax.set_title(title)

for name, field in train_dict_input.items():

  fig, axs = plt.subplots(1, 6, figsize=(20,5))

  # training
  for i in range(5):
    plot_with_colorbar(fig, axs[i], field[i][0], f'train {i}')

  # validation
  plot_with_colorbar(fig, axs[5], valid_dict_input[name][0][0], 'valid')

  plt.suptitle(f'Initial {name}')
  plt.show()
_images/ea2e6f5306cefd0e3936561e218f9e309a248a651c682df2468f544272eae690.png _images/5203289f661b8fd02df1fe828cffbc27c145dc68ea48584c7f4140bcbb08cc76.png _images/0bd8c26daa6a3786913814d7f04a51f0ecd59988c99a92229273b21a8fa451d8.png

The pairs we train on look something line this.

def plot_without_colorbar(ax, x, title):
  im = ax.imshow(x, cmap='binary')
  ax.set_title(title)

fig, axs = plt.subplots(3, 2, figsize=(12,12))

# inputs
for i, (name, field) in enumerate(train_dict_input.items()):
  plot_without_colorbar(axs[i, 0], field[i][0], f'input {name}')

# targets
for i, (name, field) in enumerate(train_dict_target.items()):
  plot_without_colorbar(axs[i, 1], field[i][0], f'target {name}')

plt.subplots_adjust(left=0.1,
                    bottom=0.1,
                    right=0.9,
                    top=0.9,
                    wspace=1.0,
                    hspace=0.3)
plt.show()
_images/fc7c02c154efc6ac52fd22b4fe14c601251a5cc3004185f12fdf3a2468f476d5.png

Now, let’s visualize the full first training trajectory of the scalar field u.

# animate one full input trajectory

fig, ax = plt.subplots()

def animate(frame_num):
    im = ax.imshow(train_dict_input['u'][0, frame_num], cmap='binary')
    return im

anim = FuncAnimation(fig, animate, frames=55, interval=100)
anim
_images/fb7faf2873a42744e2784731b21a8f7f0c3717ed561e6566ae0d771a08143611.png

0.5 Dataloader#

For some of the tasks you might want to apply batched gradient descent. For this purpose, we provide a dataloader similar to those from the lecture.

class NavierStokesDataset(Dataset):
    """Dataset Class for Navier-Stokes dataset.

    Args:
        file_path (str): path to .h5 file.
        transform (torchvision.transforms.transforms): input transform.
        transform_target (torchvision.transforms.transforms): target transform.

    Returns:
        (torch.Tensor, torch.Tensor): Tuple containing input and target tensor.
            Each tensor has the shape (3, nx, nx) which corresponds to the
            concatenated particle scalar field and 2D velocity vector field. The
            spatial grid consists of (nx, nx) points.
    """

    def __init__(self, file_path: str, transform=None, target_transform=None):
        super().__init__()

        mode = 'train'  if 'train' in file_path else 'valid'

        with h5py.File(file_path, "r") as f:
            scalar_field = f[mode]['u'][:] 
            velocity_x = f[mode]['vx'][:]
            velocity_y = f[mode]['vy'][:]
            # each with shape (..., 56, nx, nx)
            # (num_trajectories, num_time_steps, num_grid_points, num_grid_points)

        full_data = torch.tensor(
            [scalar_field, velocity_x, velocity_y], 
            dtype=torch.float32
        )
        full_data = full_data.permute(1, 2, 0, 3, 4)

        self.inputs = full_data[:, :-1].reshape(-1, 3, nx, nx)
        self.targets = full_data[:, 1:].reshape(-1, 3, nx, nx)

        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.inputs)
    
    def __getitem__(self, idx):
        input = self.inputs[idx]
        target = self.targets[idx]

        if self.transform:
            input = self.transform(input)
        if self.target_transform:
            target = self.target_transform(target)

        return input, target
# This is an example of how to use the data loader:

batch_size = 4

data_train = NavierStokesDataset(file_path=train_path)
loader_train = DataLoader(data_train, batch_size=batch_size, shuffle=True)

data_valid = NavierStokesDataset(file_path=valid_path)
loader_valid = DataLoader(data_valid, batch_size=batch_size, shuffle=False)

for i, sample in enumerate(loader_train):
  print(f'sample {i+1} with shapes: {sample[0].shape} and {sample[1].shape}')
  break
sample 1 with shapes: torch.Size([4, 3, 64, 64]) and torch.Size([4, 3, 64, 64])
<ipython-input-11-bc031b970b20>:28: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ../torch/csrc/utils/tensor_new.cpp:245.)
  full_data = torch.tensor(

1 Dimensionality Reduction and Regression#

1.1 [4 Pts] Model Selection#

1.2 [12 Pts] Dimensionality Reduction#

1.3 [12 Pts] Regression#

2 CNN Regression#

2.1 [4 Pts] Model Selection#

2.2 [8 Pts] Implementation#

2.3 [10 Pts] Performance Improvement#

3 Benchmarking#

3.1 [2 Pts] Training the Reference#

For training the PDEArena models, we need specific versions of some of the libraries. In the next cell we install them. This might take 2-3 minutes.

# install dependencies

!pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
!pip install scipy==1.10.0 flask==2.2.0 torchtext==0.13.1 pytorch-lightning==1.7.7 xarray==0.20.1 h5py==3.7.0 click==8.0.4 torchdata==0.4.1 "jsonargparse[signatures]==4.15.1" rich==12.6.0 omegaconf==2.2.3 
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/, https://download.pytorch.org/whl/cu113
Collecting torch==1.12.1+cu113
  Downloading https://download.pytorch.org/whl/cu113/torch-1.12.1%2Bcu113-cp39-cp39-linux_x86_64.whl (1837.7 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.8/1.8 GB 825.6 kB/s eta 0:00:00
?25hCollecting torchvision==0.13.1+cu113
  Downloading https://download.pytorch.org/whl/cu113/torchvision-0.13.1%2Bcu113-cp39-cp39-linux_x86_64.whl (23.4 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 23.4/23.4 MB 48.6 MB/s eta 0:00:00
?25hCollecting torchaudio==0.12.1+cu113
  Downloading https://download.pytorch.org/whl/cu113/torchaudio-0.12.1%2Bcu113-cp39-cp39-linux_x86_64.whl (3.8 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.8/3.8 MB 22.0 MB/s eta 0:00:00
?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch==1.12.1+cu113) (4.5.0)
Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.9/dist-packages (from torchvision==0.13.1+cu113) (8.4.0)
Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from torchvision==0.13.1+cu113) (2.27.1)
Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from torchvision==0.13.1+cu113) (1.22.4)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests->torchvision==0.13.1+cu113) (3.4)
Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.9/dist-packages (from requests->torchvision==0.13.1+cu113) (2.0.12)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests->torchvision==0.13.1+cu113) (2022.12.7)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->torchvision==0.13.1+cu113) (1.26.15)
Installing collected packages: torch, torchvision, torchaudio
  Attempting uninstall: torch
    Found existing installation: torch 2.0.0
    Uninstalling torch-2.0.0:
      Successfully uninstalled torch-2.0.0
  Attempting uninstall: torchvision
    Found existing installation: torchvision 0.14.1+cu116
    Uninstalling torchvision-0.14.1+cu116:
      Successfully uninstalled torchvision-0.14.1+cu116
  Attempting uninstall: torchaudio
    Found existing installation: torchaudio 0.13.1+cu116
    Uninstalling torchaudio-0.13.1+cu116:
      Successfully uninstalled torchaudio-0.13.1+cu116
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
torchtext 0.14.1 requires torch==1.13.1, but you have torch 1.12.1+cu113 which is incompatible.
torchdata 0.6.0 requires torch==2.0.0, but you have torch 1.12.1+cu113 which is incompatible.
Successfully installed torch-1.12.1+cu113 torchaudio-0.12.1+cu113 torchvision-0.13.1+cu113
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting scipy==1.10.0
  Downloading scipy-1.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (34.4 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 34.4/34.4 MB 15.2 MB/s eta 0:00:00
?25hCollecting flask==2.2.0
  Downloading Flask-2.2.0-py3-none-any.whl (101 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 101.1/101.1 KB 12.8 MB/s eta 0:00:00
?25hCollecting torchtext==0.13.1
  Downloading torchtext-0.13.1-cp39-cp39-manylinux1_x86_64.whl (1.9 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.9/1.9 MB 81.3 MB/s eta 0:00:00
?25hCollecting pytorch-lightning==1.7.7
  Downloading pytorch_lightning-1.7.7-py3-none-any.whl (708 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 708.1/708.1 KB 52.7 MB/s eta 0:00:00
?25hCollecting xarray==0.20.1
  Downloading xarray-0.20.1-py3-none-any.whl (835 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 835.7/835.7 KB 63.9 MB/s eta 0:00:00
?25hCollecting h5py==3.7.0
  Downloading h5py-3.7.0-cp39-cp39-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (4.5 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.5/4.5 MB 95.5 MB/s eta 0:00:00
?25hCollecting click==8.0.4
  Downloading click-8.0.4-py3-none-any.whl (97 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 97.5/97.5 KB 14.3 MB/s eta 0:00:00
?25hCollecting torchdata==0.4.1
  Downloading torchdata-0.4.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.4 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.4/4.4 MB 72.5 MB/s eta 0:00:00
?25hCollecting jsonargparse[signatures]==4.15.1
  Downloading jsonargparse-4.15.1-py3-none-any.whl (158 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 159.0/159.0 KB 17.9 MB/s eta 0:00:00
?25hCollecting rich==12.6.0
  Downloading rich-12.6.0-py3-none-any.whl (237 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 237.5/237.5 KB 29.0 MB/s eta 0:00:00
?25hCollecting omegaconf==2.2.3
  Downloading omegaconf-2.2.3-py3-none-any.whl (79 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 79.3/79.3 KB 12.1 MB/s eta 0:00:00
?25hRequirement already satisfied: numpy<1.27.0,>=1.19.5 in /usr/local/lib/python3.9/dist-packages (from scipy==1.10.0) (1.22.4)
Requirement already satisfied: importlib-metadata>=3.6.0 in /usr/local/lib/python3.9/dist-packages (from flask==2.2.0) (6.0.0)
Requirement already satisfied: Jinja2>=3.0 in /usr/local/lib/python3.9/dist-packages (from flask==2.2.0) (3.1.2)
Requirement already satisfied: itsdangerous>=2.0 in /usr/local/lib/python3.9/dist-packages (from flask==2.2.0) (2.1.2)
Requirement already satisfied: Werkzeug>=2.2.0 in /usr/local/lib/python3.9/dist-packages (from flask==2.2.0) (2.2.3)
Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from torchtext==0.13.1) (2.27.1)
Requirement already satisfied: torch==1.12.1 in /usr/local/lib/python3.9/dist-packages (from torchtext==0.13.1) (1.12.1+cu113)
Requirement already satisfied: tqdm in /usr/local/lib/python3.9/dist-packages (from torchtext==0.13.1) (4.65.0)
Requirement already satisfied: tensorboard>=2.9.1 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning==1.7.7) (2.11.2)
Requirement already satisfied: fsspec[http]!=2021.06.0,>=2021.05.0 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning==1.7.7) (2023.3.0)
Collecting pyDeprecate>=0.3.1
  Downloading pyDeprecate-0.3.2-py3-none-any.whl (10 kB)
Requirement already satisfied: torchmetrics>=0.7.0 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning==1.7.7) (0.11.4)
Requirement already satisfied: typing-extensions>=4.0.0 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning==1.7.7) (4.5.0)
Requirement already satisfied: PyYAML>=5.4 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning==1.7.7) (6.0)
Requirement already satisfied: packaging>=17.0 in /usr/local/lib/python3.9/dist-packages (from pytorch-lightning==1.7.7) (23.0)
Requirement already satisfied: pandas>=1.1 in /usr/local/lib/python3.9/dist-packages (from xarray==0.20.1) (1.4.4)
Requirement already satisfied: urllib3>=1.25 in /usr/local/lib/python3.9/dist-packages (from torchdata==0.4.1) (1.26.15)
Collecting portalocker>=2.0.0
  Downloading portalocker-2.7.0-py2.py3-none-any.whl (15 kB)
Collecting docstring-parser>=0.15
  Downloading docstring_parser-0.15-py3-none-any.whl (36 kB)
Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /usr/local/lib/python3.9/dist-packages (from rich==12.6.0) (2.6.1)
Collecting commonmark<0.10.0,>=0.9.0
  Downloading commonmark-0.9.1-py2.py3-none-any.whl (51 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 51.1/51.1 KB 7.0 MB/s eta 0:00:00
?25hRequirement already satisfied: antlr4-python3-runtime==4.9.* in /usr/local/lib/python3.9/dist-packages (from omegaconf==2.2.3) (4.9.3)
Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.9/dist-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (3.8.4)
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata>=3.6.0->flask==2.2.0) (3.15.0)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.9/dist-packages (from Jinja2>=3.0->flask==2.2.0) (2.1.2)
Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.9/dist-packages (from pandas>=1.1->xarray==0.20.1) (2022.7.1)
Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.9/dist-packages (from pandas>=1.1->xarray==0.20.1) (2.8.2)
Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (1.51.3)
Requirement already satisfied: protobuf<4,>=3.9.2 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (3.19.6)
Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (0.40.0)
Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (3.4.1)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (1.8.1)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (0.6.1)
Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (63.4.3)
Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (1.4.0)
Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (2.16.2)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.9/dist-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.7) (0.4.6)
Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.9/dist-packages (from requests->torchtext==0.13.1) (2.0.12)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests->torchtext==0.13.1) (2022.12.7)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests->torchtext==0.13.1) (3.4)
Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.9/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (6.0.4)
Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.9/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (1.8.2)
Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.9/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (1.3.1)
Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.9/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (22.2.0)
Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.9/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (4.0.2)
Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.9/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.7) (1.3.3)
Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (4.9)
Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (1.15.0)
Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (5.3.0)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (0.2.8)
Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.9/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (1.3.1)
Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.9/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (0.4.8)
Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.9/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->pytorch-lightning==1.7.7) (3.2.2)
Installing collected packages: commonmark, scipy, rich, pyDeprecate, portalocker, omegaconf, jsonargparse, h5py, docstring-parser, click, torchtext, torchdata, flask, xarray, pytorch-lightning
  Attempting uninstall: scipy
    Found existing installation: scipy 1.10.1
    Uninstalling scipy-1.10.1:
      Successfully uninstalled scipy-1.10.1
  Attempting uninstall: omegaconf
    Found existing installation: omegaconf 2.3.0
    Uninstalling omegaconf-2.3.0:
      Successfully uninstalled omegaconf-2.3.0
  Attempting uninstall: jsonargparse
    Found existing installation: jsonargparse 4.20.0
    Uninstalling jsonargparse-4.20.0:
      Successfully uninstalled jsonargparse-4.20.0
  Attempting uninstall: h5py
    Found existing installation: h5py 3.8.0
    Uninstalling h5py-3.8.0:
      Successfully uninstalled h5py-3.8.0
  Attempting uninstall: click
    Found existing installation: click 8.1.3
    Uninstalling click-8.1.3:
      Successfully uninstalled click-8.1.3
  Attempting uninstall: torchtext
    Found existing installation: torchtext 0.14.1
    Uninstalling torchtext-0.14.1:
      Successfully uninstalled torchtext-0.14.1
  Attempting uninstall: torchdata
    Found existing installation: torchdata 0.6.0
    Uninstalling torchdata-0.6.0:
      Successfully uninstalled torchdata-0.6.0
  Attempting uninstall: flask
    Found existing installation: Flask 2.2.3
    Uninstalling Flask-2.2.3:
      Successfully uninstalled Flask-2.2.3
  Attempting uninstall: xarray
    Found existing installation: xarray 2022.12.0
    Uninstalling xarray-2022.12.0:
      Successfully uninstalled xarray-2022.12.0
  Attempting uninstall: pytorch-lightning
    Found existing installation: pytorch-lightning 2.0.0
    Uninstalling pytorch-lightning-2.0.0:
      Successfully uninstalled pytorch-lightning-2.0.0
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
xarray-einstats 0.5.1 requires xarray>=2022.09.0, but you have xarray 0.20.1 which is incompatible.
arviz 0.15.1 requires xarray>=0.21.0, but you have xarray 0.20.1 which is incompatible.
Successfully installed click-8.0.4 commonmark-0.9.1 docstring-parser-0.15 flask-2.2.0 h5py-3.7.0 jsonargparse-4.15.1 omegaconf-2.2.3 portalocker-2.7.0 pyDeprecate-0.3.2 pytorch-lightning-1.7.7 rich-12.6.0 scipy-1.10.0 torchdata-0.4.1 torchtext-0.13.1 xarray-0.20.1

Basic model training script. You are expected to change some of the parameters.

# adapted from here: https://microsoft.github.io/pdearena/train/

!(python scripts/train.py -c configs/navierstokes2d.yaml \
--data.data_dir=data/NSEsmoke \
--trainer.max_epochs=10 \
--data.batch_size=4 \
--data.time_gap=0 --data.time_history=4 --data.time_future=1 \
--model.name=DilResNet-128 \
--model.lr=2e-4 \
--optimizer=AdamW --optimizer.lr=2e-4 --optimizer.weight_decay=1e-5 \
--lr_scheduler=LinearWarmupCosineAnnealingLR \
--lr_scheduler.warmup_epochs=1 \
--lr_scheduler.max_epochs=10 --lr_scheduler.eta_min=1e-7)
Global seed set to 42
/usr/local/lib/python3.9/dist-packages/pytorch_lightning/cli.py:601: UserWarning: `PDEModel.configure_optimizers` will be overridden by `PDECLI.configure_optimizers`.
  _warn(
Trainer already configured with model summary callbacks: [<class 'pytorch_lightning.callbacks.rich_model_summary.RichModelSummary'>]. Skipping setting a default `ModelSummary` callback.
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓
┃     Name                                Type               Params ┃
┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩
│ 0  │ model                              │ ResNet            │  4.2 M │
│ 1  │ model.conv_in1                     │ Conv2d            │  1.7 K │
│ 2  │ model.conv_in2                     │ Conv2d            │ 16.5 K │
│ 3  │ model.conv_out1                    │ Conv2d            │ 16.5 K │
│ 4  │ model.conv_out2                    │ Conv2d            │    387 │
│ 5  │ model.layers                       │ ModuleList        │  4.1 M │
│ 6  │ model.layers.0                     │ Sequential        │  1.0 M │
│ 7  │ model.layers.0.0                   │ DilatedBasicBlock │  1.0 M │
│ 8  │ model.layers.0.0.dilation_layers   │ ModuleList        │  1.0 M │
│ 9  │ model.layers.0.0.dilation_layers.0 │ Conv2d            │  147 K │
│ 10 │ model.layers.0.0.dilation_layers.1 │ Conv2d            │  147 K │
│ 11 │ model.layers.0.0.dilation_layers.2 │ Conv2d            │  147 K │
│ 12 │ model.layers.0.0.dilation_layers.3 │ Conv2d            │  147 K │
│ 13 │ model.layers.0.0.dilation_layers.4 │ Conv2d            │  147 K │
│ 14 │ model.layers.0.0.dilation_layers.5 │ Conv2d            │  147 K │
│ 15 │ model.layers.0.0.dilation_layers.6 │ Conv2d            │  147 K │
│ 16 │ model.layers.0.0.norm_layers       │ ModuleList        │      0 │
│ 17 │ model.layers.0.0.norm_layers.0     │ Identity          │      0 │
│ 18 │ model.layers.0.0.norm_layers.1     │ Identity          │      0 │
│ 19 │ model.layers.0.0.norm_layers.2     │ Identity          │      0 │
│ 20 │ model.layers.0.0.norm_layers.3     │ Identity          │      0 │
│ 21 │ model.layers.0.0.norm_layers.4     │ Identity          │      0 │
│ 22 │ model.layers.0.0.norm_layers.5     │ Identity          │      0 │
│ 23 │ model.layers.0.0.norm_layers.6     │ Identity          │      0 │
│ 24 │ model.layers.0.0.activation        │ GELU              │      0 │
│ 25 │ model.layers.1                     │ Sequential        │  1.0 M │
│ 26 │ model.layers.1.0                   │ DilatedBasicBlock │  1.0 M │
│ 27 │ model.layers.1.0.dilation_layers   │ ModuleList        │  1.0 M │
│ 28 │ model.layers.1.0.dilation_layers.0 │ Conv2d            │  147 K │
│ 29 │ model.layers.1.0.dilation_layers.1 │ Conv2d            │  147 K │
│ 30 │ model.layers.1.0.dilation_layers.2 │ Conv2d            │  147 K │
│ 31 │ model.layers.1.0.dilation_layers.3 │ Conv2d            │  147 K │
│ 32 │ model.layers.1.0.dilation_layers.4 │ Conv2d            │  147 K │
│ 33 │ model.layers.1.0.dilation_layers.5 │ Conv2d            │  147 K │
│ 34 │ model.layers.1.0.dilation_layers.6 │ Conv2d            │  147 K │
│ 35 │ model.layers.1.0.norm_layers       │ ModuleList        │      0 │
│ 36 │ model.layers.1.0.norm_layers.0     │ Identity          │      0 │
│ 37 │ model.layers.1.0.norm_layers.1     │ Identity          │      0 │
│ 38 │ model.layers.1.0.norm_layers.2     │ Identity          │      0 │
│ 39 │ model.layers.1.0.norm_layers.3     │ Identity          │      0 │
│ 40 │ model.layers.1.0.norm_layers.4     │ Identity          │      0 │
│ 41 │ model.layers.1.0.norm_layers.5     │ Identity          │      0 │
│ 42 │ model.layers.1.0.norm_layers.6     │ Identity          │      0 │
│ 43 │ model.layers.2                     │ Sequential        │  1.0 M │
│ 44 │ model.layers.2.0                   │ DilatedBasicBlock │  1.0 M │
│ 45 │ model.layers.2.0.dilation_layers   │ ModuleList        │  1.0 M │
│ 46 │ model.layers.2.0.dilation_layers.0 │ Conv2d            │  147 K │
│ 47 │ model.layers.2.0.dilation_layers.1 │ Conv2d            │  147 K │
│ 48 │ model.layers.2.0.dilation_layers.2 │ Conv2d            │  147 K │
│ 49 │ model.layers.2.0.dilation_layers.3 │ Conv2d            │  147 K │
│ 50 │ model.layers.2.0.dilation_layers.4 │ Conv2d            │  147 K │
│ 51 │ model.layers.2.0.dilation_layers.5 │ Conv2d            │  147 K │
│ 52 │ model.layers.2.0.dilation_layers.6 │ Conv2d            │  147 K │
│ 53 │ model.layers.2.0.norm_layers       │ ModuleList        │      0 │
│ 54 │ model.layers.2.0.norm_layers.0     │ Identity          │      0 │
│ 55 │ model.layers.2.0.norm_layers.1     │ Identity          │      0 │
│ 56 │ model.layers.2.0.norm_layers.2     │ Identity          │      0 │
│ 57 │ model.layers.2.0.norm_layers.3     │ Identity          │      0 │
│ 58 │ model.layers.2.0.norm_layers.4     │ Identity          │      0 │
│ 59 │ model.layers.2.0.norm_layers.5     │ Identity          │      0 │
│ 60 │ model.layers.2.0.norm_layers.6     │ Identity          │      0 │
│ 61 │ model.layers.3                     │ Sequential        │  1.0 M │
│ 62 │ model.layers.3.0                   │ DilatedBasicBlock │  1.0 M │
│ 63 │ model.layers.3.0.dilation_layers   │ ModuleList        │  1.0 M │
│ 64 │ model.layers.3.0.dilation_layers.0 │ Conv2d            │  147 K │
│ 65 │ model.layers.3.0.dilation_layers.1 │ Conv2d            │  147 K │
│ 66 │ model.layers.3.0.dilation_layers.2 │ Conv2d            │  147 K │
│ 67 │ model.layers.3.0.dilation_layers.3 │ Conv2d            │  147 K │
│ 68 │ model.layers.3.0.dilation_layers.4 │ Conv2d            │  147 K │
│ 69 │ model.layers.3.0.dilation_layers.5 │ Conv2d            │  147 K │
│ 70 │ model.layers.3.0.dilation_layers.6 │ Conv2d            │  147 K │
│ 71 │ model.layers.3.0.norm_layers       │ ModuleList        │      0 │
│ 72 │ model.layers.3.0.norm_layers.0     │ Identity          │      0 │
│ 73 │ model.layers.3.0.norm_layers.1     │ Identity          │      0 │
│ 74 │ model.layers.3.0.norm_layers.2     │ Identity          │      0 │
│ 75 │ model.layers.3.0.norm_layers.3     │ Identity          │      0 │
│ 76 │ model.layers.3.0.norm_layers.4     │ Identity          │      0 │
│ 77 │ model.layers.3.0.norm_layers.5     │ Identity          │      0 │
│ 78 │ model.layers.3.0.norm_layers.6     │ Identity          │      0 │
│ 79 │ train_criterion                    │ CustomMSELoss     │      0 │
│ 80 │ rollout_criterion                  │ MSELoss           │      0 │
└────┴────────────────────────────────────┴───────────────────┴────────┘
Trainable params: 4.2 M                                                         
Non-trainable params: 0                                                         
Total params: 4.2 M                                                             
Total estimated model params size (MB): 16                                      
2023-03-18 21:23:35.597722: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/lib64-nvidia
2023-03-18 21:23:35.598264: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/lib64-nvidia
2023-03-18 21:23:35.598294: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Sanity Checking DataLoader 1:   0% 0/2 [00:00<?, ?it/s]/usr/local/lib/python3.9/dist-packages/pytorch_lightning/utilities/data.py:98: UserWarning: Trying to infer the `batch_size` from an ambiguous collection. The batch size we found is 1. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`.
  warning_cache.warn(
Epoch 0: : 60it [00:09,  6.22it/s, loss=1.51, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 0: : 60it [00:11,  5.04it/s, loss=1.51, v_num=0]
Epoch 1: : 60it [00:09,  6.33it/s, loss=0.765, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 1: : 60it [00:12,  5.00it/s, loss=0.618, v_num=0]
Epoch 2: : 60it [00:09,  6.14it/s, loss=0.332, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 2: : 60it [00:12,  4.88it/s, loss=0.303, v_num=0]
Epoch 3: : 60it [00:09,  6.01it/s, loss=0.24, v_num=0] 
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 3: : 60it [00:12,  4.84it/s, loss=0.227, v_num=0]
Epoch 4: : 60it [00:09,  6.08it/s, loss=0.22, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 4: : 60it [00:12,  4.87it/s, loss=0.204, v_num=0]
Epoch 5: : 60it [00:09,  6.18it/s, loss=0.179, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 5: : 60it [00:12,  4.99it/s, loss=0.168, v_num=0]
Epoch 6: : 60it [00:09,  6.24it/s, loss=0.167, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 6: : 60it [00:12,  4.93it/s, loss=0.175, v_num=0]
Epoch 7: : 60it [00:09,  6.21it/s, loss=0.155, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 7: : 60it [00:12,  4.93it/s, loss=0.142, v_num=0]
Epoch 8: : 60it [00:09,  6.18it/s, loss=0.148, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 8: : 60it [00:12,  4.99it/s, loss=0.148, v_num=0]
Epoch 9: : 60it [00:09,  6.15it/s, loss=0.148, v_num=0]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Validation DataLoader 0: : 0it [00:00, ?it/s]
Epoch 9: : 60it [00:12,  4.92it/s, loss=0.145, v_num=0]
Epoch 9: : 60it [00:12,  4.92it/s, loss=0.145, v_num=0]`Trainer.fit` stopped: `max_epochs=10` reached.
Epoch 9: : 60it [00:12,  4.81it/s, loss=0.145, v_num=0]
Restoring states from the checkpoint path at /content/outputs/ckpts/epoch_005.ckpt
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Loaded model weights from checkpoint at /content/outputs/ckpts/epoch_005.ckpt
Testing: 0it [00:00, ?it/s]
%tensorboard --logdir "outputs/tb/version_0"

3.2 [8 Pts] Performance Analysis#