Merge pull request #115 from aai-institute/feature/dvc

Feature: dvc

benchmarks/.gitattributes

@@ -0,0 +1 @@
+*.png filter=lfs diff=lfs merge=lfs -text

benchmarks/html/style.css

@@ -28,7 +28,13 @@
     line-height: 1.2;
 }
 
-.benchmark-table td, table th {
+.benchmark-table td {
+    text-align: left;
+    vertical-align: middle;
+    min-width: 90px;
+}
+
+.benchmark-table th {
     text-align: left;
     vertical-align: middle;
 }

benchmarks/navierstokes/plot_navierstokes.py

@@ -0,0 +1,89 @@
+import torch
+import matplotlib.pyplot as plt
+from continuity.benchmarks import NavierStokes
+from continuity.operators import FourierNeuralOperator
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+ns = NavierStokes()
+
+operator = FourierNeuralOperator(
+    ns.train_dataset.shapes,
+    grid_shape=(64, 64, 10),
+    width=32,
+    depth=4,
+    device=device,
+)
+
+operator.load(
+    "mlruns/271016623891034109/8755b17d3af9494db843e3a8d0c42ad6/artifacts/final.pt"
+)
+operator.eval()
+
+# Compute train loss
+loss_fn = ns.losses[0]
+
+
+def compute_loss(dataset):
+    train_loader = torch.utils.data.DataLoader(dataset, batch_size=1)
+    avg_loss = 0
+    max_loss, min_loss = 0, 1e10
+    max_i, min_i = 0, 0
+    for i, xuyv in enumerate(train_loader):
+        x, u, y, v = [t.to(device) for t in xuyv]
+        loss = loss_fn(operator, x, u, y, v)
+        avg_loss += loss.detach()
+        if loss > max_loss:
+            max_loss = loss
+            max_i = i
+        if loss < min_loss:
+            min_loss = loss
+            min_i = i
+    avg_loss = avg_loss / len(train_loader)
+    return avg_loss, max_loss, max_i, min_loss, min_i
+
+
+loss_train, max_loss, max_i_train, min_loss, min_i_train = compute_loss(
+    ns.train_dataset
+)
+print(f"rel. error train = {loss_train:.4e}")
+print(f"min loss = {min_loss:.4e} at index {min_i_train}")
+print(f"max loss = {max_loss:.4e} at index {max_i_train}")
+
+# Compute test loss
+loss_test, max_loss, max_i_test, min_loss, min_i_test = compute_loss(ns.test_dataset)
+print(f"rel. error test =  {loss_test:.4e}")
+print(f"min loss = {min_loss:.4e} at index {min_i_test}")
+print(f"max loss = {max_loss:.4e} at index {max_i_test}")
+
+
+# Plot
+def plot_sample(split, sample):
+    dataset = ns.train_dataset if split == "train" else ns.test_dataset
+    x, u, y, v = [t.to(device) for t in dataset[sample : sample + 1]]
+    v_pred = operator(x, u, y)
+    v = v.reshape(1, 64, 64, 10, 1).cpu()
+    v_pred = v_pred.reshape(1, 64, 64, 10, 1).detach().cpu()
+
+    fig, axs = plt.subplots(10, 3, figsize=(4, 16))
+
+    axs[0][0].set_title("Truth")
+    axs[0][1].set_title("Prediction")
+    axs[0][2].set_title("Error")
+    for t in range(10):
+        axs[t][0].imshow(v[0, :, :, t, 0], cmap="jet")
+        axs[t][1].imshow(v_pred[0, :, :, t, 0], cmap="jet")
+        im = axs[t][2].imshow((v - v_pred)[0, :, :, t, 0], cmap="jet")
+        fig.colorbar(im, ax=axs[t][2])
+        axs[t][0].axis("off")
+        axs[t][1].axis("off")
+        axs[t][2].axis("off")
+
+    plt.tight_layout()
+    plt.savefig(f"navierstokes/ns_{split}_{sample}.png", dpi=500)
+
+
+plot_sample("train", min_i_train)
+plot_sample("train", max_i_train)
+plot_sample("test", min_i_test)
+plot_sample("test", max_i_test)

benchmarks/navierstokes/results.txt

@@ -0,0 +1,16 @@
+rel. error train = 1.8508e-02
+min loss = 8.8748e-03 at index 237
+max loss = 3.1433e-02 at index 420
+rel. error test =  1.8408e-01
+min loss = 1.0220e-01 at index 144
+max loss = 4.4655e-01 at index 179
+
+--
+
+Reference: _Li, Zongyi, et al. "Fourier neural operator for parametric partial
+differential equations." arXiv preprint arXiv:2010.08895 (2020)_
+
+Table 1: nu=1e−5  T=20  N=1000
+FNO-3D: 0.1893
+
+(Ours: 1.8408e-01)

benchmarks/navierstokes/run_navierstokes.py

@@ -0,0 +1,20 @@
+from functools import partial
+from continuity.benchmarks.run import BenchmarkRunner, RunConfig
+from continuity.benchmarks import NavierStokes
+from continuity.operators import FourierNeuralOperator
+
+config = RunConfig(
+    benchmark_factory=NavierStokes,
+    operator_factory=partial(
+        FourierNeuralOperator,
+        grid_shape=(64, 64, 10),
+        width=32,
+        depth=4,
+    ),
+    lr=1e-3,
+    max_epochs=100,
+    batch_size=10,
+)
+
+if __name__ == "__main__":
+    BenchmarkRunner.run(config)

build_scripts/copy_benchmarks.py

@@ -44,3 +44,5 @@ def on_shutdown():
         file.unlink()
     for file in docs_benchmarks_dir.glob("*.png"):
         file.unlink()
+    for file in docs_benchmarks_dir.glob("*.svg"):
+        file.unlink()

data/.dvc/.gitignore

@@ -0,0 +1,3 @@
+/config.local
+/tmp
+/cache

data/.dvc/config

@@ -0,0 +1,4 @@
+[core]
+    remote = gdrive
+['remote "gdrive"']
+    url = gdrive://1Mts9tmPjKzqw-as_j1XTRwa9ulNssy5a

data/.dvcignore

@@ -0,0 +1,3 @@
+# Add patterns of files dvc should ignore, which could improve
+# the performance. Learn more at
+# https://dvc.org/doc/user-guide/dvcignore

data/.gitignore

@@ -0,0 +1,2 @@
+/flame
+/navierstokes

data/README.md

@@ -0,0 +1,37 @@
+# Data
+
+## Prerequisites
+
+We use `dvc` to manage the data. You can install the required packages by
+installing the benchmarks requirements.
+
+```
+pip install -e .[benchmarks]
+```
+
+## Downloading the data
+
+The data is stored in a remote storage on GDrive.
+To download the data, you can run:
+
+```
+cd data
+dvc pull <NAME>
+```
+
+where `<NAME>` is the name of the data set you want to download,
+e.g., `flame` or `navierstokes`, or empty.
+
+
+## Data sets
+
+### FLAME
+
+`data/flame` contains the dataset from [2023 FLAME AI
+Challenge](https://www.kaggle.com/competitions/2023-flame-ai-challenge/data).
+
+### Navier-Stokes
+
+`data/navierstokes` contains a part of the dataset linked in
+[neuraloperator/graph-pde](https://github.com/neuraloperator/graph-pde)
+(Zongyi Li et al. 2020).

data/flame.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 2e61c8311b09a4fdf29d3ec3527cf629.dir
+  size: 415040265
+  nfiles: 13138
+  path: flame

data/navierstokes.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: 0bb228674e976bab14e9493606e14a27.dir
+  size: 412877192
+  nfiles: 1
+  hash: md5
+  path: navierstokes

docs/benchmarks/.gitignore

@@ -1,3 +1,5 @@
 table.html
 style.css
 img
+*.png
+*.svg

docs/benchmarks/index.md

@@ -2,6 +2,41 @@ This is an overview of some benchmark results to compare the performance
 of different operator architectures on various problems.
 
 The benchmarks are implemented in the `benchmarks` directory and we refer to
-this directory for more information on how the benchmarks are run.
+this directory for detailed information on how the benchmarks are run.
+
+## [NavierStokes](../api/continuity/benchmarks/#continuity.benchmarks.NavierStokes)
+
+Reference: _Li, Zongyi, et al. "Fourier neural operator for parametric partial
+differential equations." arXiv preprint arXiv:2010.08895 (2020)_ _Table 1 ($\nu$ = 1e−5  T=20  N=1000)_
+
+_reported for_ FNO-3D: __0.1893__ (rel. test error)
+
+[FourierNeuralOperator](../api/continuity/operators/#continuity.operators.FourierNeuralOperator):
+0.0185 (rel. train error)  __0.1841__ (rel. test error)
+
+<table>
+<tr>
+<td>
+Best training sample<br>
+<img src="img/ns_train_237.png" alt="Best training sample"/>
+rel. error = 8.8748e-03
+</td>
+<td>
+Worst training sample<br>
+<img src="img/ns_train_420.png" alt="Worst training sample"/>
+rel. error = 3.1433e-02
+</td>
+<td>
+Best test sample<br>
+<img src="img/ns_test_144.png" alt="Best test sample"/>
+rel. error = 1.0220e-01
+</td>
+<td>
+Worst test sample<br>
+<img src="img/ns_test_179.png" alt="Worst test sample"/>
+rel. error = 4.4655e-01
+</td>
+</tr>
+</table>
 
 {% include 'benchmarks/table.html' %}

pyproject.toml

@@ -80,8 +80,11 @@ docs = [
     "pygments",
 ]
 benchmark = [
+    "dvc",
+    "dvc-gdrive",
     "mlflow",
     "optuna>=3.5.0,<4.0.0",
+    "scipy",
 ]
 
 [tool.setuptools.dynamic]

src/continuity/benchmarks/__init__.py

@@ -7,5 +7,6 @@
 from .benchmark import Benchmark
 from .sine import SineRegular, SineUniform
 from .flame import Flame
+from .navierstokes import NavierStokes
 
-__all__ = ["Benchmark", "SineRegular", "SineUniform", "Flame"]
+__all__ = ["Benchmark", "SineRegular", "SineUniform", "Flame", "NavierStokes"]