Use normalisation from paper and add PyTorch notebook example (#9)

- The convergence is very dependent on the data scaling and corresponding hyperparameters, so we are again using the (-2, 2) data scaling for the branch net
- Created a notebook that should give a good introduction to DeepONets and in particular the implementation of the wave propagation. Plotting methods done for 2D data.

Author: nikolasborrel

.vscode/launch.json

@@ -12,26 +12,26 @@
             "console": "integratedTerminal"
         },
         {
-            "name": "Python: Train 3D",
+            "name": "Python: Train",
             "type": "debugpy",
             "request": "launch",
             "module": "deeponet_acoustics.main_train",
             "args": [
                 "--path_settings",
-                "json_setups/threeD/settings.json",
+                "json_setups/threeD/choras_cube.json",
             ],
             "console": "integratedTerminal"
-        },
+        },        
         {
-            "name": "Python: Inference 3D",
+            "name": "Python: Inference",
             "type": "debugpy",
             "request": "launch",
             "module": "deeponet_acoustics.main_inference",
             "args": [
                 "--path_settings",
-                "json_setups/threeD/settings.json",
+                "json_setups/twoD/rect2x2.json",
                 "--path_eval_settings",
-                "json_setups/threeD/settings_eval.json",
+                "json_setups/twoD/rect2x2_eval.json",
             ],
             "console": "integratedTerminal"
         },

deeponet_acoustics/datahandlers/datagenerators.py

@@ -16,6 +16,7 @@
 import h5py
 import jax.numpy as jnp
 import numpy as np
+import torch
 from torch.utils.data import Dataset
 
 import deeponet_acoustics.datahandlers.io as IO
@@ -160,6 +161,7 @@ def __init__(
         data_prune=1,
         norm_data=False,
         MAXNUM_DATASETS=sys.maxsize,
+        u_p_range=None,
     ):
         filenames_xdmf = IO.pathsToFileType(data_path, ".xdmf", exclude="rectilinear")
         self.normalize_data = norm_data
@@ -202,9 +204,15 @@ def __init__(
         self.N = len(self.datasets)
 
         # Calculate min and max pressure values from sampled datasets
-        self._u_p_min, self._u_p_max = _calculate_u_pressure_minmax(
-            self.datasets, self.tag_ufield
-        )
+        if u_p_range is not None:
+            self._u_p_min, self._u_p_max = u_p_range
+            print(
+                f"Using specified u pressure range: [{self._u_p_min:.4f}, {self._u_p_max:.4f}]"
+            )
+        else:
+            self._u_p_min, self._u_p_max = _calculate_u_pressure_minmax(
+                self.datasets, self.tag_ufield
+            )
 
     # --- required abstract properties implemented ---
     @property
@@ -264,6 +272,7 @@ def __init__(
         data_prune=1,
         norm_data=False,
         MAXNUM_DATASETS=sys.maxsize,
+        u_p_range=None,
     ):
         filenamesH5 = IO.pathsToFileType(data_path, ".h5", exclude="rectilinear")
         self.data_prune = data_prune
@@ -295,10 +304,13 @@ def __init__(
             )
             self.tsteps = r[self.tags_field[0]].attrs["time_steps"]
             self.tsteps = jnp.array([t for t in self.tsteps if t <= tmax / t_norm])
-            self.tsteps = self.tsteps * t_norm
+            self.tsteps = (
+                self.tsteps * t_norm
+            )  # corresponding to c = 1 for same spatial / temporal resolution
 
             if self.normalize_data:
                 self.mesh = self.normalize_spatial(self.mesh)
+                # normalize relative to spatial dimension to keep ratio
                 self.tsteps = self.normalize_temporal(self.tsteps)
 
         self.tt = np.repeat(self.tsteps, self.mesh.shape[0])
@@ -308,16 +320,21 @@ def __init__(
         for i in range(0, min(MAXNUM_DATASETS, len(filenamesH5))):
             filename = filenamesH5[i]
             if Path(filename).exists():
-                self.datasets.append(
-                    h5py.File(filename, "r")
-                )  # add file handles and keeps open
+                # add file handles and keeps open
+                self.datasets.append(h5py.File(filename, "r"))
             else:
                 print(f"Could not be found (ignoring): {filename}")
 
         # Calculate min and max pressure values from sampled datasets
-        self._u_p_min, self._u_p_max = _calculate_u_pressure_minmax(
-            self.datasets, self.tag_ufield
-        )
+        if u_p_range is not None:
+            self._u_p_min, self._u_p_max = u_p_range
+            print(
+                f"Using specified u pressure range: [{self._u_p_min:.4f}, {self._u_p_max:.4f}]"
+            )
+        else:
+            self._u_p_min, self._u_p_max = _calculate_u_pressure_minmax(
+                self.datasets, self.tag_ufield
+            )
 
     # --- required abstract properties implemented ---
     @property
@@ -332,8 +349,11 @@ def P(self):
 
     @property
     def xxyyzztt(self):
-        xxyyzz = np.tile(self.mesh, (len(self.tsteps), 1))
-        return np.hstack((xxyyzz, self.tt.reshape(-1, 1)))
+        return np.hstack((self.xxyyzz, self.tt.reshape(-1, 1)))
+
+    @property
+    def xxyyzz(self):
+        return np.tile(self.mesh, (len(self.tsteps), 1))
 
     def normalize_spatial(self, data):
         return _normalize_spatial(data, self.xmin, self.xmax)
@@ -412,6 +432,7 @@ def __init__(
         flatten_ic: bool = True,
         data_prune: int = 1,
         norm_data: bool = False,
+        u_p_range=None,
     ) -> None:
         self.data_prune = data_prune
         self._normalize_data = norm_data
@@ -478,9 +499,15 @@ def __init__(
                 )
 
         # Calculate min and max pressure values from generated datasets
-        self._u_p_min, self._u_p_max = _calculate_u_pressure_minmax(
-            self.datasets, self.tag_ufield, max_samples=self.N
-        )
+        if u_p_range is not None:
+            self._u_p_min, self._u_p_max = u_p_range
+            print(
+                f"Using specified u pressure range: [{self._u_p_min:.4f}, {self._u_p_max:.4f}]"
+            )
+        else:
+            self._u_p_min, self._u_p_max = _calculate_u_pressure_minmax(
+                self.datasets, self.tag_ufield, max_samples=self.N
+            )
 
         self.tt = np.repeat(self.tsteps, self.mesh.shape[0])
 
@@ -509,8 +536,11 @@ def normalize_data(self) -> bool:
     @property
     def xxyyzztt(self) -> np.ndarray[float]:
         """Spatio-temporal coordinates stacked as [x, y, z, t]."""
-        xxyyzz = np.tile(self.mesh, (len(self.tsteps), 1))
-        return np.hstack((xxyyzz, self.tt.reshape(-1, 1)))
+        return np.hstack((self.xxyyzz, self.tt.reshape(-1, 1)))
+
+    @property
+    def xxyyzz(self) -> np.ndarray[float]:
+        return np.tile(self.mesh, (len(self.tsteps), 1))
 
     def normalize_spatial(self, data: np.ndarray[float]) -> np.ndarray[float]:
         return _normalize_spatial(data, self._xmin, self._xmax)
@@ -606,7 +636,7 @@ def __getitem__(self, idx):
                 0:num_tsteps, :: self.data.data_prune
             ].flatten()[indxs_coord]
         elif self.data.simulationDataType == SimulationDataType.XDMF:
-            s = np.empty((self.P), dtype=jnp.float32)
+            s = np.empty((self.P), dtype=np.float32)
             for j in range(num_tsteps):
                 s[j * self.data.P_mesh : (j + 1) * self.P_mesh] = dataset[
                     self.data.tags_field[j]
@@ -615,7 +645,9 @@ def __getitem__(self, idx):
         elif self.data.simulationDataType == SimulationDataType.SOURCE_ONLY:
             s = []
         else:
-            raise Exception("Data format unknown: should be H5COMPACT or XDMF")
+            raise Exception(
+                "Data format unknown: should be H5COMPACT, XDMF or SOURCE_ONLY"
+            )
 
         # normalize
         x0 = (
@@ -624,19 +656,38 @@ def __getitem__(self, idx):
             else []
         )
 
-        inputs = jnp.asarray(u), jnp.asarray(y)
-        return inputs, jnp.asarray(s), indxs_coord, x0
+        inputs = np.asarray(u), np.asarray(y)
+        return inputs, np.asarray(s), indxs_coord, x0
 
 
 def get_number_of_sources(data_path: str):
     return len(IO.pathsToFileType(data_path, ".h5", exclude="rectilinear"))
 
 
 def numpy_collate(batch):
+    """Collate function for JAX - converts batches to JAX arrays."""
     if isinstance(batch[0], np.ndarray):
         return jnp.stack(batch)
     elif isinstance(batch[0], (tuple, list)):
         transposed = zip(*batch)
         return [numpy_collate(samples) for samples in transposed]
     else:
         return np.array(batch)
+
+
+def pytorch_collate(batch):
+    """Collate function for PyTorch - converts batches to PyTorch tensors.
+
+    Use this collator with PyTorch DataLoader for educational notebooks:
+        DataLoader(dataset, batch_size=2, collate_fn=pytorch_collate)
+    """
+
+    if isinstance(batch[0], np.ndarray):
+        return torch.from_numpy(np.stack(batch)).float()
+    elif isinstance(batch[0], jnp.ndarray):
+        return torch.from_numpy(np.array(batch)).float()
+    elif isinstance(batch[0], (tuple, list)):
+        transposed = zip(*batch)
+        return [pytorch_collate(samples) for samples in transposed]
+    else:
+        return torch.tensor(batch).float()

deeponet_acoustics/end2end/inference.py

@@ -133,10 +133,9 @@ def inference(
     model.plotLosses(settings.dirs.figs_dir)
 
     tdim = len(metadata.tsteps)
-    xxyyzztt = metadata.xxyyzztt
-    y_in = y_feat_fn(xxyyzztt)
+    y_in = y_feat_fn(metadata.xxyyzztt)
 
-    xxyyzz_phys = metadata.denormalize_spatial(xxyyzztt[:, 0:3])
+    xxyyzz_phys = metadata.denormalize_spatial(metadata.xxyyzz)
     mesh_phys = metadata.denormalize_spatial(metadata.mesh)
     tsteps_phys = metadata.denormalize_temporal(metadata.tsteps / c_phys)
 
@@ -196,9 +195,15 @@ def inference(
         r0_list_norm = metadata.normalize_spatial(r0_list[i_src])
 
         (u_test_i, *_), s_test_i, _, x0 = dataset[i_src]
-
-        x0 = metadata.denormalize_spatial(x0)
-        x0_srcs.append(x0)
+        if len(x0) > 0:
+            x0 = metadata.denormalize_spatial(x0)
+            x0_srcs.append(x0)
+        else:
+            print(
+                "Warning: test data does not have source position data - setting index as coordinate"
+            )
+            x0 = i_src
+            x0_srcs.append([x0])
 
         y_rcvs = np.repeat(np.array(r0_list_norm), len(metadata.tsteps), axis=0)
         tsteps_rcvs = np.tile(metadata.tsteps, len(r0_list_norm))

deeponet_acoustics/end2end/train.py

@@ -65,6 +65,7 @@ def train(settings_dict: dict[str, Any]):
         t_norm=c_phys,
         norm_data=settings.normalize_data,
         flatten_ic=flatten_ic,
+        u_p_range=(-2.0, 2.0),
     )
     dataset = DatasetStreamer(
         metadata, training.batch_size_coord, y_feat_extract_fn=y_feat_fn
@@ -75,6 +76,7 @@ def train(settings_dict: dict[str, Any]):
         t_norm=c_phys,
         norm_data=settings.normalize_data,
         flatten_ic=flatten_ic,
+        u_p_range=(-2.0, 2.0),
     )
     dataset_val = DatasetStreamer(
         metadata_val, training.batch_size_coord, y_feat_extract_fn=y_feat_fn

deeponet_acoustics/models/deeponet.py

@@ -219,13 +219,13 @@ def train(
                 i += 1
 
                 if do_timings:
-                    timer.startTiming("backprop") if do_timings else None
+                    timer.startTiming("backprop")
                     self.params, self.opt_state, _ = self.step(
                         self.params, self.opt_state, data_batch
                     )
-                    jax.block_until_ready(self.params) if do_timings else None
-                    jax.block_until_ready(self.opt_state) if do_timings else None
-                    timer.endTiming("backprop") if do_timings else None
+                    jax.block_until_ready(self.params)
+                    jax.block_until_ready(self.opt_state)
+                    timer.endTiming("backprop")
 
                     timer.writeTimings(
                         {
@@ -235,8 +235,8 @@ def train(
                         }
                     )
                     timer.resetTimings()
-                    timer.startTiming("total_iter") if do_timings else None
-                    timer.startTiming("dataloader") if do_timings else None
+                    timer.startTiming("total_iter")
+                    timer.startTiming("dataloader")
                 else:
                     self.params, self.opt_state, _ = self.step(
                         self.params, self.opt_state, data_batch

deeponet_acoustics/scripts/convertH5/split_2D_by_source.py

@@ -88,6 +88,19 @@ def split_2d_by_source_position(input_file: Path, output_path: Path):
             x0_srcs = None
 
         print(f"Converting {num_sources} source positions from 2D to 3D format...")
+
+        # Compute denormalized values once for all sources
+        c_phys = float(pressure_attrs["c_phys"])
+        fmax_normalized = float(pressure_attrs["fmax"])
+        fmax_phys = round(fmax_normalized * c_phys, 1)
+        dt_normalized = float(pressure_attrs["dt"])
+        dt_phys = dt_normalized / c_phys
+
+        # Update pressure_attrs with denormalized values for use in root metadata
+        pressure_attrs_denormalized = pressure_attrs.copy()
+        pressure_attrs_denormalized["fmax"] = fmax_phys
+        pressure_attrs_denormalized["dt"] = dt_phys
+
         src_pos_2d = None
 
         for src_idx in range(num_sources):
@@ -120,10 +133,12 @@ def split_2d_by_source_position(input_file: Path, output_path: Path):
                     "pressures", data=src_pressures
                 )
 
-                # Add time_steps attribute to pressures (from t dataset)
-                pressures_dataset.attrs["time_steps"] = t.astype(np.float32)
+                # Denormalize time_steps by dividing by physical speed of sound
+                time_steps_phys = (t / c_phys).astype(np.float32)
+                pressures_dataset.attrs["time_steps"] = time_steps_phys
+
                 # Note: Adding all pressure attributes that weren't in original 3D data
-                for attr_name, attr_value in pressure_attrs.items():
+                for attr_name, attr_value in pressure_attrs_denormalized.items():
                     if attr_name != "time_steps":  # Avoid duplicate
                         pressures_dataset.attrs[attr_name] = attr_value
 
@@ -147,13 +162,15 @@ def split_2d_by_source_position(input_file: Path, output_path: Path):
                 for attr_name, attr_value in upressure_attrs.items():
                     upressures_dataset.attrs[attr_name] = attr_value
 
-            # Create metadata JSON file in source folder
-            create_metadata_json(src_folder, pressure_attrs, src_pos_2d)
+                # Create metadata JSON file in source folder using corrected attributes
+                create_metadata_json(
+                    src_folder, dict(pressures_dataset.attrs), src_pos_2d
+                )
 
             print(f"Created {output_file}")
 
         # Create metadata JSON file in root directory as well
-        create_metadata_json(output_path, pressure_attrs, None)
+        create_metadata_json(output_path, pressure_attrs_denormalized, None)
 
         print(f"Conversion complete! Created {num_sources} files in {output_path}")