From 0c5fa7871e81b18b1a245e4df51be8d77e7bafae Mon Sep 17 00:00:00 2001
From: Samuel Burbulla <s.burbulla@appliedai-institute.de>
Date: Sat, 18 Nov 2023 11:38:16 +0200
Subject: [PATCH] Restructure project.

---
 data/.gitignore                               |  2 +
 data/flame/README.md                          |  2 +
 src/continuity/__init__.py                    |  0
 src/continuity/data/__init__.py               |  0
 data.py => src/continuity/data/dataset.py     |  2 +-
 src/continuity/data/flame.py                  | 55 +++++++++++++++++++
 src/continuity/data/sine.py                   | 32 +++++++++++
 src/continuity/model/__init__.py              |  0
 model.py => src/continuity/model/model.py     |  0
 src/continuity/plotting/__init__.py           |  0
 .../continuity/plotting/plotting.py           |  2 +-
 tests/.gitignore                              |  1 +
 main.py => tests/sine.py                      | 39 ++++---------
 13 files changed, 105 insertions(+), 30 deletions(-)
 create mode 100644 data/.gitignore
 create mode 100644 data/flame/README.md
 create mode 100644 src/continuity/__init__.py
 create mode 100644 src/continuity/data/__init__.py
 rename data.py => src/continuity/data/dataset.py (99%)
 create mode 100644 src/continuity/data/flame.py
 create mode 100644 src/continuity/data/sine.py
 create mode 100644 src/continuity/model/__init__.py
 rename model.py => src/continuity/model/model.py (100%)
 create mode 100644 src/continuity/plotting/__init__.py
 rename plotting.py => src/continuity/plotting/plotting.py (97%)
 create mode 100644 tests/.gitignore
 rename main.py => tests/sine.py (61%)

diff --git a/data/.gitignore b/data/.gitignore
new file mode 100644
index 00000000..47a7e5fa
--- /dev/null
+++ b/data/.gitignore
@@ -0,0 +1,2 @@
+*.csv
+*.dat
\ No newline at end of file
diff --git a/data/flame/README.md b/data/flame/README.md
new file mode 100644
index 00000000..4cc75d39
--- /dev/null
+++ b/data/flame/README.md
@@ -0,0 +1,2 @@
+Dataset from [2023 FLAME AI
+Challenge](https://www.kaggle.com/competitions/2023-flame-ai-challenge/data).
\ No newline at end of file
diff --git a/src/continuity/__init__.py b/src/continuity/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/continuity/data/__init__.py b/src/continuity/data/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/data.py b/src/continuity/data/dataset.py
similarity index 99%
rename from data.py
rename to src/continuity/data/dataset.py
index 9fc977fc..c41afb99 100644
--- a/data.py
+++ b/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)
diff --git a/src/continuity/data/flame.py b/src/continuity/data/flame.py
new file mode 100644
index 00000000..6bc0e2e5
--- /dev/null
+++ b/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)
+
diff --git a/src/continuity/data/sine.py b/src/continuity/data/sine.py
new file mode 100644
index 00000000..a5462e72
--- /dev/null
+++ b/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)
+
diff --git a/src/continuity/model/__init__.py b/src/continuity/model/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/model.py b/src/continuity/model/model.py
similarity index 100%
rename from model.py
rename to src/continuity/model/model.py
diff --git a/src/continuity/plotting/__init__.py b/src/continuity/plotting/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/plotting.py b/src/continuity/plotting/plotting.py
similarity index 97%
rename from plotting.py
rename to src/continuity/plotting/plotting.py
index 36494469..e8560aac 100644
--- a/plotting.py
+++ b/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()):
diff --git a/tests/.gitignore b/tests/.gitignore
new file mode 100644
index 00000000..643cb181
--- /dev/null
+++ b/tests/.gitignore
@@ -0,0 +1 @@
+runs/
\ No newline at end of file
diff --git a/main.py b/tests/sine.py
similarity index 61%
rename from main.py
rename to tests/sine.py
index 363a2039..a009cf99 100644
--- a/main.py
+++ b/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()