Restructure project.

data/.gitignore

@@ -0,0 +1,2 @@
+*.csv
+*.dat

data/flame/README.md

@@ -0,0 +1,2 @@
+Dataset from [2023 FLAME AI
+Challenge](https://www.kaggle.com/competitions/2023-flame-ai-challenge/data).

data.py → src/continuity/data/dataset.py

@@ -4,7 +4,7 @@
 from typing import List, Optional
 import pandas as pd
 import torch
-from model import device
+from continuity.model.model import device
 
 def tensor(x):
     return torch.tensor(x, dtype=torch.float32)

src/continuity/data/flame.py

@@ -0,0 +1,55 @@
+import numpy as np
+import pandas as pd
+from continuity.data.dataset import DataSet, Sensor, Observation
+
+
+class Flame(DataSet):
+    """Turbulent flow samples from flame dataset"""
+
+    def __init__(self, size, batch_size):
+        self.num_sensors = 16 * 16
+        self.size = size
+
+        self.coordinate_dim = 2
+        self.num_channels = 1
+
+        # Generate observations
+        observations = [
+            self._generate_observation(i)
+            for i in range(self.size)
+        ]
+
+        super().__init__(observations, batch_size)
+
+    def _generate_observation(self, i: int):
+        # Load data
+        input_path = 'flame/'
+        res = 'LR'
+
+        df = pd.read_csv(input_path + 'train.csv')
+        data_path = input_path + f"flowfields/{res}/train"
+
+        filename = df[f'ux_filename'][i+1]
+        u = np.fromfile(
+            data_path + "/" + filename, dtype="<f4"
+        ).reshape(16, 16, 1)
+
+        # Normalize
+        u = (u - u.mean()) / u.std()
+
+        # Positions
+        x = np.stack(
+            np.meshgrid(
+                np.linspace(-1, 1, 16),
+                np.linspace(-1, 1, 16),
+            ),
+            axis=2
+        )
+
+        # Sensors
+        sensors = [
+            Sensor(x[i][j], u[i][j])
+            for i in range(16) for j in range(16)
+        ]
+        return Observation(sensors)
+

src/continuity/data/sine.py

@@ -0,0 +1,32 @@
+import numpy as np
+import pandas as pd
+from continuity.data.dataset import DataSet, Sensor, Observation
+
+
+class SineWaves(DataSet):
+    """Sine waves"""
+
+    def __init__(self, num_sensors, size, batch_size):
+        self.num_sensors = num_sensors
+        self.size = size
+
+        self.coordinate_dim = 1
+        self.num_channels = 1
+
+        # Generate observations
+        observations = [
+            self._generate_observation(i / max(1, self.size-1))
+            for i in range(self.size)
+        ]
+
+        super().__init__(observations, batch_size)
+
+    def _generate_observation(self, i: int):
+        x = np.linspace(-1, 1, self.num_sensors)
+        u = np.sin((1+i) * np.pi * x)
+        sensors = [
+            Sensor(np.array([x]), np.array([u]))
+            for x, u in zip(x, u)
+        ]
+        return Observation(sensors)
+

plotting.py → src/continuity/plotting/plotting.py

@@ -2,7 +2,7 @@
 import io
 import numpy as np
 import matplotlib.pyplot as plt
-from model import device
+from continuity.model.model import device
 
 
 def plot_observation(observation, ax=plt.gca()):

tests/.gitignore

@@ -0,0 +1 @@
+runs/

main.py → tests/sine.py

@@ -1,9 +1,8 @@
 import torch
 import matplotlib.pyplot as plt
-from data import SineWaves, Flame
-from plotting import *
-from model import FullyConnected, DeepONet, NeuralOperator
-from dadaptation import DAdaptSGD
+from continuity.data.sine import SineWaves
+from continuity.plotting.plotting import *
+from continuity.model.model import NeuralOperator
 from torch.utils.tensorboard import SummaryWriter
 
 # Set random seed
@@ -14,33 +13,18 @@ def main():
     size = 4
 
     # Create data set
-    dataset = SineWaves(32, size, batch_size=32)
-    # dataset = Flame(size, batch_size=1)
+    dataset = SineWaves(
+        num_sensors=32,
+        size=size,
+        batch_size=32,
+    )
 
-    # Create model
-    # model = FullyConnected(
-    #     coordinate_dim=dataset.coordinate_dim,
-    #     num_channels=dataset.num_channels,
-    #     num_sensors=dataset.num_sensors,
-    #     width=128,
-    #     depth=128,
-    # )
-    # model = DeepONet(
-    #     coordinate_dim=dataset.coordinate_dim,
-    #     num_channels=dataset.num_channels,
-    #     num_sensors=dataset.num_sensors,
-    #     branch_width=32,
-    #     branch_depth=4,
-    #     trunk_width=128,
-    #     trunk_depth=32,
-    #     basis_functions=4,
-    # )
     model = NeuralOperator(
         coordinate_dim=dataset.coordinate_dim,
         num_channels=dataset.num_channels,
         depth=1,
-        kernel_width=128,
-        kernel_depth=32,
+        kernel_width=32,
+        kernel_depth=8,
     )
 
     # Load model
@@ -50,13 +34,12 @@ def main():
         model.load_state_dict(torch.load(f'{log_dir}/model_weights.pth'))
 
     # Train model
-    epochs = 10000
+    epochs = 100
     writer = SummaryWriter()
 
     train = True
     if train:
         # Setup optimizer and fit model
-        # optimizer = DAdaptSGD(model.parameters(), lr=0.1)
         optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
         criterion = torch.nn.MSELoss()