Github action: auto-update.

Author: github-actions[bot]

dev/_downloads/02a7a82499cf14d2702c9234035c44c4/plot_diffusion_advection_solver.zip

@@ -48,11 +48,11 @@
 
 train_loader, test_loaders = load_spherical_swe(
     n_train=200,
-    batch_size=8,
+    batch_size=32,
     train_resolution=(32, 64),
     test_resolutions=[(32, 64), (64, 128)],
-    n_tests=[50, 50],
-    test_batch_sizes=[10, 10],
+    n_tests=[40, 40],
+    test_batch_sizes=[40, 40],
 )
 
 
@@ -71,6 +71,7 @@
     out_channels=3,
     hidden_channels=64,
     domain_padding=[0.05, 0.05],
+    n_layers=2,
 )
 model = model.to(device)
 
@@ -175,7 +176,7 @@
 #
 # In practice we would train a Neural Operator on one or multiple GPUs
 
-fig = plt.figure(figsize=(7, 7))
+fig = plt.figure(figsize=(14, 7))
 for index, resolution in enumerate([(32, 64), (64, 128)]):
     test_samples = test_loaders[resolution].dataset
     data = test_samples[0]
@@ -195,20 +196,24 @@
     plt.xticks([], [])
     plt.yticks([], [])
 
+    # Compute the min and max to use consistent color mapping
+    vmin = y.min()
+    vmax = y.max()
+
     # Plot ground-truth fields
     ax = fig.add_subplot(2, 3, index * 3 + 2)
-    ax.imshow(y)
+    im_gt = ax.imshow(y, vmin=vmin, vmax=vmax)
     ax.set_title("Ground-truth y")
     plt.xticks([], [])
     plt.yticks([], [])
 
     # Plot model prediction
     ax = fig.add_subplot(2, 3, index * 3 + 3)
-    ax.imshow(out)
+    im_pred = ax.imshow(out, vmin=vmin, vmax=vmax)
     ax.set_title("SFNO prediction")
     plt.xticks([], [])
     plt.yticks([], [])
 
-fig.suptitle("SFNO predictions on spherical shallow water equations", y=0.98)
+fig.suptitle("SFNO predictions on spherical shallow water equations", y=0.98, fontsize=24)
 plt.tight_layout()
 fig.show()

dev/_downloads/07fcc19ba03226cd3d83d4e40ec44385/auto_examples_python.zip

@@ -40,7 +40,7 @@
       },
       "outputs": [],
       "source": [
-        "train_loader, test_loaders = load_spherical_swe(\n    n_train=200,\n    batch_size=8,\n    train_resolution=(32, 64),\n    test_resolutions=[(32, 64), (64, 128)],\n    n_tests=[50, 50],\n    test_batch_sizes=[10, 10],\n)"
+        "train_loader, test_loaders = load_spherical_swe(\n    n_train=200,\n    batch_size=32,\n    train_resolution=(32, 64),\n    test_resolutions=[(32, 64), (64, 128)],\n    n_tests=[40, 40],\n    test_batch_sizes=[40, 40],\n)"
       ]
     },
     {
@@ -58,7 +58,7 @@
       },
       "outputs": [],
       "source": [
-        "model = SFNO(\n    n_modes=(16, 32),\n    in_channels=3,\n    out_channels=3,\n    hidden_channels=64,\n    domain_padding=[0.05, 0.05],\n)\nmodel = model.to(device)\n\n# Count and display the number of parameters\nn_params = count_model_params(model)\nprint(f\"\\nOur model has {n_params} parameters.\")\nsys.stdout.flush()"
+        "model = SFNO(\n    n_modes=(16, 32),\n    in_channels=3,\n    out_channels=3,\n    hidden_channels=64,\n    domain_padding=[0.05, 0.05],\n    n_layers=2,\n)\nmodel = model.to(device)\n\n# Count and display the number of parameters\nn_params = count_model_params(model)\nprint(f\"\\nOur model has {n_params} parameters.\")\nsys.stdout.flush()"
       ]
     },
     {
@@ -159,7 +159,7 @@
       },
       "outputs": [],
       "source": [
-        "fig = plt.figure(figsize=(7, 7))\nfor index, resolution in enumerate([(32, 64), (64, 128)]):\n    test_samples = test_loaders[resolution].dataset\n    data = test_samples[0]\n    # Input x\n    x = data[\"x\"]\n    # Ground-truth\n    y = data[\"y\"][0, ...].numpy()\n    # Model prediction: SFNO output\n    x_in = x.unsqueeze(0).to(device)\n    out = model(x_in).squeeze()[0, ...].detach().cpu().numpy()\n    x = x[0, ...].detach().numpy()\n\n    # Plot input fields\n    ax = fig.add_subplot(2, 3, index * 3 + 1)\n    ax.imshow(x)\n    ax.set_title(f\"Input x {resolution}\")\n    plt.xticks([], [])\n    plt.yticks([], [])\n\n    # Plot ground-truth fields\n    ax = fig.add_subplot(2, 3, index * 3 + 2)\n    ax.imshow(y)\n    ax.set_title(\"Ground-truth y\")\n    plt.xticks([], [])\n    plt.yticks([], [])\n\n    # Plot model prediction\n    ax = fig.add_subplot(2, 3, index * 3 + 3)\n    ax.imshow(out)\n    ax.set_title(\"SFNO prediction\")\n    plt.xticks([], [])\n    plt.yticks([], [])\n\nfig.suptitle(\"SFNO predictions on spherical shallow water equations\", y=0.98)\nplt.tight_layout()\nfig.show()"
+        "fig = plt.figure(figsize=(14, 7))\nfor index, resolution in enumerate([(32, 64), (64, 128)]):\n    test_samples = test_loaders[resolution].dataset\n    data = test_samples[0]\n    # Input x\n    x = data[\"x\"]\n    # Ground-truth\n    y = data[\"y\"][0, ...].numpy()\n    # Model prediction: SFNO output\n    x_in = x.unsqueeze(0).to(device)\n    out = model(x_in).squeeze()[0, ...].detach().cpu().numpy()\n    x = x[0, ...].detach().numpy()\n\n    # Plot input fields\n    ax = fig.add_subplot(2, 3, index * 3 + 1)\n    ax.imshow(x)\n    ax.set_title(f\"Input x {resolution}\")\n    plt.xticks([], [])\n    plt.yticks([], [])\n\n    # Compute the min and max to use consistent color mapping\n    vmin = y.min()\n    vmax = y.max()\n\n    # Plot ground-truth fields\n    ax = fig.add_subplot(2, 3, index * 3 + 2)\n    im_gt = ax.imshow(y, vmin=vmin, vmax=vmax)\n    ax.set_title(\"Ground-truth y\")\n    plt.xticks([], [])\n    plt.yticks([], [])\n\n    # Plot model prediction\n    ax = fig.add_subplot(2, 3, index * 3 + 3)\n    im_pred = ax.imshow(out, vmin=vmin, vmax=vmax)\n    ax.set_title(\"SFNO prediction\")\n    plt.xticks([], [])\n    plt.yticks([], [])\n\nfig.suptitle(\"SFNO predictions on spherical shallow water equations\", y=0.98, fontsize=24)\nplt.tight_layout()\nfig.show()"
       ]
     }
   ],