Merge pull request #44 from praksharma/development

2D heat conduction and refactoring

DeepINN/config.py

@@ -11,6 +11,10 @@ def __init__(self, float_type=torch.float32, random_seed=42, device = 'cuda'):
         self.device = device
         # Add more configuration parameters as needed
 
+        self.apply_seeds()
+        self.apply_float_type()
+        self.default_device()
+
     def apply_seeds(self):
         torch.manual_seed(self.random_seed)
 

DeepINN/domain/standard.py

@@ -16,10 +16,9 @@ def __init__(self, pde_equation, collocation_object, bc_object) -> None:
         super().__init__()
 
         self.pde = pde_equation
-        self.pde_sampler = collocation_object
+        self.pde_sampler = collocation_object # class of the PDE
         # if the bcs is not list, then make it a list
-        self.bc_sampler = bc_object if isinstance(bc_object, (list)) else [bc_object]
-
+        self.bc_sampler = bc_object if isinstance(bc_object, (list)) else [bc_object] # class of the boundary condition
         self.bc_list_len = len(self.bc_sampler)
 
     def sample_collocation_points(self):

DeepINN/model.py

@@ -1,7 +1,8 @@
 import torch
 import sys
 from .backend import loss_metric, choose_optimiser
-from  .config import Config
+from .config import Config
+from .utils import timer
 
 class Model():
     """
@@ -27,10 +28,10 @@ def compile(self, optimiser_string : str, lr : float, metrics_string : str, devi
 
     def compile_domain(self):
         # sample collocation points
-        self.collocation_point_sample, self.collocation_point_labels = self.domain.sample_collocation_labels()
+        self.collocation_point_sample, self.collocation_point_labels = self.domain.sample_collocation_labels() # list of collocation points and their labels both as tensors
 
         # sample boundary points
-        self.boundary_point_sample, self.boundary_point_labels = self.domain.sample_boundary_labels()
+        self.boundary_point_sample, self.boundary_point_labels = self.domain.sample_boundary_labels() # list of boundary points and their labels both as tensors
         print("Domain compiled", file=sys.stderr, flush=True)
 
     def compile_network(self):
@@ -46,25 +47,39 @@ def compile_network(self):
                                                  )
         print("Network compiled", file=sys.stderr, flush=True)
 
-    def train(self, iterations : int = None, display_every : int = None):
-
+    def initialise_training(self, iterations : int = None):
         if self.iter == 0: # We are running a fresh training
             self.training_history = []  # Initialize an empty list for storing loss values
             self.iterations = iterations
             # Load all the seeds, data types, devices etc.
-            self.config.apply_seeds()
-            self.config.apply_float_type()
-            self.config.default_device()
+            # self.config.apply_seeds()
+            # self.config.apply_float_type()
+            # self.config.default_device()
 
             # In 1D problem we need to combine the BCs as there is only one point for each BC, which returns an undefined feature scaling because the ub and lb are same in the denominator, so we get infinity
             # For problem with multiple points on each boundary, we don't need to combine them.
             if self.boundary_point_sample[0].size()[0] == 1: # if row is 1 in the particular boundary tensor
                 self.boundary_point_sample = torch.cat(self.boundary_point_sample, dim=0)
                 self.boundary_point_labels = torch.cat(self.boundary_point_labels, dim=0)
+            else:
+                self.boundary_point_sample = torch.cat(self.boundary_point_sample, dim=0)
+                self.boundary_point_labels = torch.cat(self.boundary_point_labels, dim=0)
 
             # Set requires_grad=True for self.collocation_point_sample
             self.collocation_point_sample.requires_grad = True
 
+    def train(self, iterations : int = None, display_every : int = 1):
+        """_summary_
+
+        Args:
+            iterations (int): _description_. Number of iterations.
+            display_every (int, optional): _description_. Display the loss every display_every iterations. Defaults to 1.           
+        """
+        self.initialise_training(iterations)
+        self.trainer()
+
+    @timer
+    def trainer(self):
         # implement training loop
         while self.iter <= self.iterations:
 
@@ -93,6 +108,4 @@ def train(self, iterations : int = None, display_every : int = None):
             self.iter = self.iter + 1
         else:
             print('Training finished')
-            #elapsed = time.time() - start_time                
-            #print('Training time: %.2f' % (elapsed))
-            #print(f"Final loss: {total_loss}")
+

DeepINN/nn/FCNN.py

@@ -1,7 +1,7 @@
 import torch
 from .base import BaseNetwork
 
-class FullyConnected(BaseNetwork, torch.nn.Module):
+class FullyConnected(BaseNetwork):
     """
     Implementation of Fully Connected neural network
     """
@@ -28,10 +28,10 @@ def __init__(self,
         self.linears = torch.nn.ModuleList([torch.nn.Linear(self.layer_size[i], self.layer_size[i+1]) for i in range(0,len(self.layer_size)-1)])
 
         # initialise the weights
-        self.init()
+        self.weight_init()
 
 
-    def init(self):
+    def weight_init(self):
         # weight initialisation
         for i in range(len(self.layer_size)-1):
 

DeepINN/nn/base.py

@@ -1,12 +1,13 @@
 from abc import ABC, abstractmethod
 from .utils import activation, initialiser
+import torch
 
-class BaseNetwork():
+class BaseNetwork(torch.nn.Module):
     """
     Base class for all neural networks
     """
     def __init__(self) -> None:
-        super().__init__() # intialise all methods from nn.Module
+        super().__init__() # initialise all methods from nn.Module
 
         self.activation_function = activation
         self.initialiser_function = initialiser        

DeepINN/utils/__init__.py

@@ -1,12 +1,3 @@
-"""Useful helper methods for the definition and evaluation of a problem.
-
-For the creation of conditions, some differential operators are implemented under
-torchphysics.utils.differentialoperators.
-
-For the evaluation of the trained model, some plot and animation functionalities are provided.
-They can give you a rough overview of the determined solution. These lay under
-torchphysics.utils.plotting
-"""
 from .differentialoperators import (laplacian,
                                     grad,
                                     div,
@@ -20,7 +11,7 @@
 
 from .data import PointsDataset, PointsDataLoader, DeepONetDataLoader
 
-from .user_fun import UserFunction, tensor2numpy
+from .user_fun import UserFunction, tensor2numpy, timer
 from .plotting import plot, animate, scatter
 from .evaluation import compute_min_and_max
 

DeepINN/utils/user_fun.py

@@ -5,6 +5,8 @@
 import inspect
 import copy 
 import torch
+import functools
+import time
 
 from ..geometry.spaces.points import Points
 
@@ -317,3 +319,14 @@ def tensor2numpy(tensor_list):
     Converts a list of torch.tensors to numpy arrays.
     """
     return [tensor.detach().cpu().numpy() for tensor in tensor_list]
+
+def timer(func):
+    """Print the runtime of the decorated function"""
+    @functools.wraps(func)
+    def wrapper_timer(*args, **kwargs):
+        start_time = time.perf_counter()
+        func(*args, **kwargs) # execute the decorated function
+        end_time = time.perf_counter()
+        run_time = end_time - start_time
+        print(f"Time taken: {func.__name__!r} in {run_time:.4f} secs")
+    return wrapper_timer

docs/requirements.txt

@@ -1,3 +1,4 @@
 jupyter-book
 matplotlib
 numpy
+ghp-import