Self-supervised

This example shows how to train a neural operator on sine functions in a self-supervised manner.

Setup

import torch
import matplotlib.pyplot as plt
from continuiti.benchmarks.sine import SineBenchmark
from continuiti.data.selfsupervised import SelfSupervisedOperatorDataset
from continuiti.operators.integralkernel import NaiveIntegralKernel, NeuralNetworkKernel
from continuiti.trainer import Trainer

Dataset

Create a data set of sine waves: The SineBenchmark generates \(N\) sine waves $$ f(x) = \sin(w_k x), \quad w_k = 1 + \frac{k}{N-1}, $$ $$ \quad k = 0, \dots, N-1. $$ We wrap the dataset by a SelfSupervisedDataset that exports samples for self-supervised training, namely $$ \left(\mathbf{x}, f(\mathbf{x}), x_j, f(x_j)\right), \quad \text{for } j = 1, \dots, M, $$ where \(\mathbf{x} = (x_i)_{i=1 \dots M}\) are the \(M\) equidistantly distributed sensor positions.

benchmark = SineBenchmark(n_sensors=32, n_train=4, n_evaluations=3)
sine = benchmark.train_dataset
dataset = SelfSupervisedOperatorDataset(sine.x, sine.u)

This dataset contains 128 samples. Let's plot a random one!

Operator

In this example, we use a NaiveIntegralKernel as neural operator with a NeuralNetworkKernel as kernel function.

kernel = NeuralNetworkKernel(dataset.shapes, kernel_width=128, kernel_depth=8)
operator = NaiveIntegralKernel(kernel)

Training

Train the neural operator.

Trainer(operator).fit(dataset, tol=1e-3, batch_size=128)

Plotting

Plot model predictions for training data.

Generalization

Plot prediction on a test sample which was not part of the training set.

---

Last update: 2024-04-16
Created: 2023-12-14