Skip to content

Commit 568babc

Browse files
github-actions[bot]github-actions[bot]
committed
Github action: auto-update.
1 parent c30e783 commit 568babc

File tree

266 files changed

+56510
-0
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

266 files changed

+56510
-0
lines changed

‎stable/_downloads/01111202d22a556abc06831baa4765ac/plot_darcy_flow.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
"""
2+
A simple Darcy-Flow dataset
3+
===========================
4+
In this example, we demonstrate how to use the small Darcy-Flow example we ship with the package
5+
"""
6+
7+
# %%
8+
# Import the library
9+
# ------------------
10+
# We first import our `neuralop` library and required dependencies.
11+
12+
import matplotlib.pyplot as plt
13+
from neuralop.data.datasets import load_darcy_flow_small
14+
from neuralop.layers.embeddings import GridEmbedding2D
15+
16+
# %%
17+
# Load the dataset
18+
# ----------------
19+
# Training samples are 16x16 and we load testing samples at both
20+
# 16x16 and 32x32 (to test resolution invariance).
21+
22+
train_loader, test_loaders, data_processor = load_darcy_flow_small(
23+
n_train=100, batch_size=4,
24+
test_resolutions=[16, 32], n_tests=[50, 50], test_batch_sizes=[4, 2],
25+
)
26+
27+
train_dataset = train_loader.dataset
28+
29+
# %%
30+
# Visualizing the data
31+
# --------------------
32+
33+
for res, test_loader in test_loaders.items():
34+
print(res)
35+
# Get first batch
36+
batch = next(iter(test_loader))
37+
x = batch['x']
38+
y = batch['y']
39+
40+
print(f'Testing samples for res {res} have shape {x.shape[1:]}')
41+
42+
43+
data = train_dataset[0]
44+
x = data['x']
45+
y = data['y']
46+
47+
print(f'Training samples have shape {x.shape[1:]}')
48+
49+
50+
# Which sample to view
51+
index = 0
52+
53+
data = train_dataset[index]
54+
data = data_processor.preprocess(data, batched=False)
55+
56+
# The first step of the default FNO model is a grid-based
57+
# positional embedding. We will add it manually here to
58+
# visualize the channels appended by this embedding.
59+
positional_embedding = GridEmbedding2D(in_channels=1)
60+
# at train time, data will be collated with a batch dim.
61+
# we create a batch dim to pass into the embedding, then re-squeeze
62+
x = positional_embedding(data['x'].unsqueeze(0)).squeeze(0)
63+
y = data['y']
64+
fig = plt.figure(figsize=(7, 7))
65+
ax = fig.add_subplot(2, 2, 1)
66+
ax.imshow(x[0], cmap='gray')
67+
ax.set_title('input x')
68+
ax = fig.add_subplot(2, 2, 2)
69+
ax.imshow(y.squeeze())
70+
ax.set_title('input y')
71+
ax = fig.add_subplot(2, 2, 3)
72+
ax.imshow(x[1])
73+
ax.set_title('x: 1st pos embedding')
74+
ax = fig.add_subplot(2, 2, 4)
75+
ax.imshow(x[2])
76+
ax.set_title('x: 2nd pos embedding')
77+
fig.suptitle('Visualizing one input sample', y=0.98)
78+
plt.tight_layout()
79+
fig.show()

‎stable/_downloads/07fcc19ba03226cd3d83d4e40ec44385/auto_examples_python.zip‎

38.5 KB
Binary file not shown.

‎stable/_downloads/082e73328a5caf8c1fe9ad7fe05cf68f/plot_incremental_FNO_darcy.ipynb‎

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"\n# Training an FNO with incremental meta-learning\nIn this example, we demonstrate how to use the small Darcy-Flow \nexample we ship with the package to demonstrate the Incremental FNO\nmeta-learning algorithm\n"
8+
]
9+
},
10+
{
11+
"cell_type": "code",
12+
"execution_count": null,
13+
"metadata": {
14+
"collapsed": false
15+
},
16+
"outputs": [],
17+
"source": [
18+
"import torch\nimport matplotlib.pyplot as plt\nimport sys\nfrom neuralop.models import FNO\nfrom neuralop.data.datasets import load_darcy_flow_small\nfrom neuralop.utils import count_model_params\nfrom neuralop.training import AdamW\nfrom neuralop.training.incremental import IncrementalFNOTrainer\nfrom neuralop.data.transforms.data_processors import IncrementalDataProcessor\nfrom neuralop import LpLoss, H1Loss"
19+
]
20+
},
21+
{
22+
"cell_type": "markdown",
23+
"metadata": {},
24+
"source": [
25+
"Loading the Darcy flow dataset\n\n"
26+
]
27+
},
28+
{
29+
"cell_type": "code",
30+
"execution_count": null,
31+
"metadata": {
32+
"collapsed": false
33+
},
34+
"outputs": [],
35+
"source": [
36+
"train_loader, test_loaders, output_encoder = load_darcy_flow_small(\n n_train=100,\n batch_size=16,\n test_resolutions=[16, 32],\n n_tests=[100, 50],\n test_batch_sizes=[32, 32],\n)"
37+
]
38+
},
39+
{
40+
"cell_type": "markdown",
41+
"metadata": {},
42+
"source": [
43+
"Choose device\n\n"
44+
]
45+
},
46+
{
47+
"cell_type": "code",
48+
"execution_count": null,
49+
"metadata": {
50+
"collapsed": false
51+
},
52+
"outputs": [],
53+
"source": [
54+
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
55+
]
56+
},
57+
{
58+
"cell_type": "markdown",
59+
"metadata": {},
60+
"source": [
61+
"Set up the incremental FNO model\nWe start with 2 modes in each dimension\nWe choose to update the modes by the incremental gradient explained algorithm\n\n"
62+
]
63+
},
64+
{
65+
"cell_type": "code",
66+
"execution_count": null,
67+
"metadata": {
68+
"collapsed": false
69+
},
70+
"outputs": [],
71+
"source": [
72+
"incremental = True\nif incremental:\n starting_modes = (2, 2)\nelse:\n starting_modes = (16, 16)"
73+
]
74+
},
75+
{
76+
"cell_type": "markdown",
77+
"metadata": {},
78+
"source": [
79+
"set up model\n\n"
80+
]
81+
},
82+
{
83+
"cell_type": "code",
84+
"execution_count": null,
85+
"metadata": {
86+
"collapsed": false
87+
},
88+
"outputs": [],
89+
"source": [
90+
"model = FNO(\n max_n_modes=(16, 16),\n n_modes=starting_modes,\n hidden_channels=32,\n in_channels=1,\n out_channels=1,\n)\nmodel = model.to(device)\nn_params = count_model_params(model)"
91+
]
92+
},
93+
{
94+
"cell_type": "markdown",
95+
"metadata": {},
96+
"source": [
97+
"Set up the optimizer and scheduler\n\n"
98+
]
99+
},
100+
{
101+
"cell_type": "code",
102+
"execution_count": null,
103+
"metadata": {
104+
"collapsed": false
105+
},
106+
"outputs": [],
107+
"source": [
108+
"optimizer = AdamW(model.parameters(), lr=8e-3, weight_decay=1e-4)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30)\n\n\n# If one wants to use Incremental Resolution, one should use the IncrementalDataProcessor - When passed to the trainer, the trainer will automatically update the resolution\n# Incremental_resolution : bool, default is False\n# if True, increase the resolution of the input incrementally\n# uses the incremental_res_gap parameter\n# uses the subsampling_rates parameter - a list of resolutions to use\n# uses the dataset_indices parameter - a list of indices of the dataset to slice to regularize the input resolution\n# uses the dataset_resolution parameter - the resolution of the input\n# uses the epoch_gap parameter - the number of epochs to wait before increasing the resolution\n# uses the verbose parameter - if True, print the resolution and the number of modes\ndata_transform = IncrementalDataProcessor(\n in_normalizer=None,\n out_normalizer=None,\n device=device,\n subsampling_rates=[2, 1],\n dataset_resolution=16,\n dataset_indices=[2, 3],\n epoch_gap=10,\n verbose=True,\n)\n\ndata_transform = data_transform.to(device)"
109+
]
110+
},
111+
{
112+
"cell_type": "markdown",
113+
"metadata": {},
114+
"source": [
115+
"Set up the losses\n\n"
116+
]
117+
},
118+
{
119+
"cell_type": "code",
120+
"execution_count": null,
121+
"metadata": {
122+
"collapsed": false
123+
},
124+
"outputs": [],
125+
"source": [
126+
"l2loss = LpLoss(d=2, p=2)\nh1loss = H1Loss(d=2)\ntrain_loss = h1loss\neval_losses = {\"h1\": h1loss, \"l2\": l2loss}\nprint(\"\\n### N PARAMS ###\\n\", n_params)\nprint(\"\\n### OPTIMIZER ###\\n\", optimizer)\nprint(\"\\n### SCHEDULER ###\\n\", scheduler)\nprint(\"\\n### LOSSES ###\")\nprint(\"\\n### INCREMENTAL RESOLUTION + GRADIENT EXPLAINED ###\")\nprint(f\"\\n * Train: {train_loss}\")\nprint(f\"\\n * Test: {eval_losses}\")\nsys.stdout.flush()"
127+
]
128+
},
129+
{
130+
"cell_type": "markdown",
131+
"metadata": {},
132+
"source": [
133+
"Set up the IncrementalTrainer\nother options include setting incremental_loss_gap = True\nIf one wants to use incremental resolution set it to True\nIn this example we only update the modes and not the resolution\nWhen using the incremental resolution one should keep in mind that the numnber of modes initially set should be strictly less than the resolution\nAgain these are the various paramaters for the various incremental settings\nincremental_grad : bool, default is False\n if True, use the base incremental algorithm which is based on gradient variance\n uses the incremental_grad_eps parameter - set the threshold for gradient variance\n uses the incremental_buffer paramater - sets the number of buffer modes to calculate the gradient variance\n uses the incremental_max_iter parameter - sets the initial number of iterations\n uses the incremental_grad_max_iter parameter - sets the maximum number of iterations to accumulate the gradients\nincremental_loss_gap : bool, default is False\n if True, use the incremental algorithm based on loss gap\n uses the incremental_loss_eps parameter\n\n"
134+
]
135+
},
136+
{
137+
"cell_type": "code",
138+
"execution_count": null,
139+
"metadata": {
140+
"collapsed": false
141+
},
142+
"outputs": [],
143+
"source": [
144+
"# Finally pass all of these to the Trainer\ntrainer = IncrementalFNOTrainer(\n model=model,\n n_epochs=20,\n data_processor=data_transform,\n device=device,\n verbose=True,\n incremental_loss_gap=False,\n incremental_grad=True,\n incremental_grad_eps=0.9999,\n incremental_loss_eps = 0.001,\n incremental_buffer=5,\n incremental_max_iter=1,\n incremental_grad_max_iter=2,\n)"
145+
]
146+
},
147+
{
148+
"cell_type": "markdown",
149+
"metadata": {},
150+
"source": [
151+
"Train the model\n\n"
152+
]
153+
},
154+
{
155+
"cell_type": "code",
156+
"execution_count": null,
157+
"metadata": {
158+
"collapsed": false
159+
},
160+
"outputs": [],
161+
"source": [
162+
"trainer.train(\n train_loader,\n test_loaders,\n optimizer,\n scheduler,\n regularizer=False,\n training_loss=train_loss,\n eval_losses=eval_losses,\n)"
163+
]
164+
},
165+
{
166+
"cell_type": "markdown",
167+
"metadata": {},
168+
"source": [
169+
"Plot the prediction, and compare with the ground-truth\nNote that we trained on a very small resolution for\na very small number of epochs\nIn practice, we would train at larger resolution, on many more samples.\n\nHowever, for practicity, we created a minimal example that\ni) fits in just a few Mb of memory\nii) can be trained quickly on CPU\n\nIn practice we would train a Neural Operator on one or multiple GPUs\n\n"
170+
]
171+
},
172+
{
173+
"cell_type": "code",
174+
"execution_count": null,
175+
"metadata": {
176+
"collapsed": false
177+
},
178+
"outputs": [],
179+
"source": [
180+
"test_samples = test_loaders[32].dataset\n\nfig = plt.figure(figsize=(7, 7))\nfor index in range(3):\n data = test_samples[index]\n # Input x\n x = data[\"x\"].to(device)\n # Ground-truth\n y = data[\"y\"].to(device)\n # Model prediction\n out = model(x.unsqueeze(0))\n ax = fig.add_subplot(3, 3, index * 3 + 1)\n x = x.cpu().squeeze().detach().numpy()\n y = y.cpu().squeeze().detach().numpy()\n ax.imshow(x, cmap=\"gray\")\n if index == 0:\n ax.set_title(\"Input x\")\n plt.xticks([], [])\n plt.yticks([], [])\n\n ax = fig.add_subplot(3, 3, index * 3 + 2)\n ax.imshow(y.squeeze())\n if index == 0:\n ax.set_title(\"Ground-truth y\")\n plt.xticks([], [])\n plt.yticks([], [])\n\n ax = fig.add_subplot(3, 3, index * 3 + 3)\n ax.imshow(out.cpu().squeeze().detach().numpy())\n if index == 0:\n ax.set_title(\"Model prediction\")\n plt.xticks([], [])\n plt.yticks([], [])\n\nfig.suptitle(\"Inputs, ground-truth output and prediction.\", y=0.98)\nplt.tight_layout()\nfig.show()"
181+
]
182+
}
183+
],
184+
"metadata": {
185+
"kernelspec": {
186+
"display_name": "Python 3",
187+
"language": "python",
188+
"name": "python3"
189+
},
190+
"language_info": {
191+
"codemirror_mode": {
192+
"name": "ipython",
193+
"version": 3
194+
},
195+
"file_extension": ".py",
196+
"mimetype": "text/x-python",
197+
"name": "python",
198+
"nbconvert_exporter": "python",
199+
"pygments_lexer": "ipython3",
200+
"version": "3.13.1"
201+
}
202+
},
203+
"nbformat": 4,
204+
"nbformat_minor": 0
205+
}

‎stable/_downloads/0ac9f102e4f1903984e3b4a7d517c88f/plot_darcy_flow.ipynb‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"\n# A simple Darcy-Flow dataset\nIn this example, we demonstrate how to use the small Darcy-Flow example we ship with the package\n"
8+
]
9+
},
10+
{
11+
"cell_type": "markdown",
12+
"metadata": {},
13+
"source": [
14+
"## Import the library\nWe first import our `neuralop` library and required dependencies.\n\n"
15+
]
16+
},
17+
{
18+
"cell_type": "code",
19+
"execution_count": null,
20+
"metadata": {
21+
"collapsed": false
22+
},
23+
"outputs": [],
24+
"source": [
25+
"import matplotlib.pyplot as plt\nfrom neuralop.data.datasets import load_darcy_flow_small\nfrom neuralop.layers.embeddings import GridEmbedding2D"
26+
]
27+
},
28+
{
29+
"cell_type": "markdown",
30+
"metadata": {},
31+
"source": [
32+
"## Load the dataset\nTraining samples are 16x16 and we load testing samples at both \n16x16 and 32x32 (to test resolution invariance).\n\n"
33+
]
34+
},
35+
{
36+
"cell_type": "code",
37+
"execution_count": null,
38+
"metadata": {
39+
"collapsed": false
40+
},
41+
"outputs": [],
42+
"source": [
43+
"train_loader, test_loaders, data_processor = load_darcy_flow_small(\n n_train=100, batch_size=4, \n test_resolutions=[16, 32], n_tests=[50, 50], test_batch_sizes=[4, 2],\n )\n\ntrain_dataset = train_loader.dataset"
44+
]
45+
},
46+
{
47+
"cell_type": "markdown",
48+
"metadata": {},
49+
"source": [
50+
"## Visualizing the data\n\n"
51+
]
52+
},
53+
{
54+
"cell_type": "code",
55+
"execution_count": null,
56+
"metadata": {
57+
"collapsed": false
58+
},
59+
"outputs": [],
60+
"source": [
61+
"for res, test_loader in test_loaders.items():\n print(res)\n # Get first batch\n batch = next(iter(test_loader))\n x = batch['x']\n y = batch['y']\n\n print(f'Testing samples for res {res} have shape {x.shape[1:]}')\n\n\ndata = train_dataset[0]\nx = data['x']\ny = data['y']\n\nprint(f'Training samples have shape {x.shape[1:]}')\n\n\n# Which sample to view\nindex = 0\n\ndata = train_dataset[index]\ndata = data_processor.preprocess(data, batched=False)\n\n# The first step of the default FNO model is a grid-based\n# positional embedding. We will add it manually here to\n# visualize the channels appended by this embedding.\npositional_embedding = GridEmbedding2D(in_channels=1)\n# at train time, data will be collated with a batch dim.\n# we create a batch dim to pass into the embedding, then re-squeeze\nx = positional_embedding(data['x'].unsqueeze(0)).squeeze(0)\ny = data['y']\nfig = plt.figure(figsize=(7, 7))\nax = fig.add_subplot(2, 2, 1)\nax.imshow(x[0], cmap='gray')\nax.set_title('input x')\nax = fig.add_subplot(2, 2, 2)\nax.imshow(y.squeeze())\nax.set_title('input y')\nax = fig.add_subplot(2, 2, 3)\nax.imshow(x[1])\nax.set_title('x: 1st pos embedding')\nax = fig.add_subplot(2, 2, 4)\nax.imshow(x[2])\nax.set_title('x: 2nd pos embedding')\nfig.suptitle('Visualizing one input sample', y=0.98)\nplt.tight_layout()\nfig.show()"
62+
]
63+
}
64+
],
65+
"metadata": {
66+
"kernelspec": {
67+
"display_name": "Python 3",
68+
"language": "python",
69+
"name": "python3"
70+
},
71+
"language_info": {
72+
"codemirror_mode": {
73+
"name": "ipython",
74+
"version": 3
75+
},
76+
"file_extension": ".py",
77+
"mimetype": "text/x-python",
78+
"name": "python",
79+
"nbconvert_exporter": "python",
80+
"pygments_lexer": "ipython3",
81+
"version": "3.13.1"
82+
}
83+
},
84+
"nbformat": 4,
85+
"nbformat_minor": 0
86+
}

‎stable/_downloads/0d78e075dd52a34e158d7f5f710dfe89/plot_incremental_FNO_darcy.zip‎

16 KB
Binary file not shown.

0 commit comments

Comments
 (0)