Merge pull request #57 from Jeyhun1/sampling

Adaptive sampling

Author: StillerPatrick

PINNFramework/Adaptive_Sampler.py

@@ -0,0 +1,77 @@
+import torch
+import numpy as np
+from .Sampler import Sampler
+
+
+class AdaptiveSampler(Sampler):
+    def __init__(self, n_seed, model, pde, device = torch.device("cuda")):
+        """
+        Constructor of the AdaptiveSampler class
+
+        Args:            
+            n_seed (int): the number of seed points for adaptive sampling.
+            model: is the model which is trained to represent the underlying PDE.
+            pde (function): function that represents residual of the PDE.
+            device (torch.device): "cuda" or "cpu".
+        """
+        self.n_seed = n_seed
+        self.model = model 
+        self.pde = pde 
+        self.device = device 
+        super(AdaptiveSampler, self).__init__()
+
+    def sample(self, lb, ub, n):
+        """
+        Generate a tuple of 'n' sampled points in [lb,ub] and corresponding weights.
+        
+        Args:
+            lb (numpy.ndarray): lower bound of the domain.
+            ub (numpy.ndarray): upper bound of the domain.
+            n (int): the number of sampled points.
+        """
+        
+        torch.manual_seed(42)
+        np.random.seed(42)
+        
+        lb =  lb.reshape(1,-1)
+        ub =  ub.reshape(1,-1)
+
+        dimension = lb.shape[1]
+        xs = np.random.uniform(lb, ub, size=(self.n_seed, dimension))
+
+        # collocation points
+        xf = np.random.uniform(lb, ub, size=(n, dimension))
+
+        # make the points into tensors
+        xf = torch.tensor(xf).float().to(self.device)
+        xs = torch.tensor(xs).float().to(self.device)
+
+        # prediction with seed points
+        xs.requires_grad = True
+        prediction_seed = self.model(xs)
+
+        # pde residual with seed points
+        loss_seed = self.pde(xs, prediction_seed)
+        losses_xf = torch.zeros_like(xf)
+
+        # Compute the 2-norm distance between seed points and collocation points
+        dist = torch.cdist(xf, xs, p=2)
+
+        # obtain the smallest element of the given tensor
+        knn = dist.topk(1, largest=False)
+
+        # assign the seed loss to the loss of the closest collocation points
+        losses_xf = loss_seed[knn.indices[:, 0]]
+
+        # apply softmax function
+        q_model = torch.softmax(losses_xf, dim=0)
+
+        # obtain 'n' indices sampled from the multinomial distribution
+        indicies_new = torch.multinomial(q_model[:, 0], n, replacement=True)
+
+        # collocation points and corresponding weights
+        xf = xf[indicies_new]
+        weight = q_model[indicies_new].detach()
+        weight = torch.mean(weight, 1, True)
+
+        return xf, weight

PINNFramework/Geometry.py

@@ -0,0 +1,29 @@
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+
+class Geometry(Dataset):
+    def __init__(self, lb, ub, n_points, batch_size, sampler):
+        """
+        Constructor of the Geometry class
+
+        Args:
+            lb (numpy.ndarray): lower bound of the domain.
+            ub (numpy.ndarray): upper bound of the domain.
+            n_points (int): the number of sampled points.
+            batch_size (int): batch size
+            sampler: instance of the Sampler class.
+        """
+        self.lb = lb
+        self.ub = ub
+        self.n_points = n_points
+        self.batch_size = batch_size
+        self.sampler = sampler
+
+
+    def __getitem__(self, idx):
+        raise NotImplementedError("Subclasses should implement '__getitem__' method")
+
+    def __len__(self):
+        raise NotImplementedError("Subclasses should implement '__len__' method")
+

PINNFramework/HPMLoss.py

@@ -1,18 +1,19 @@
 from .PDELoss import PDELoss
 
 class HPMLoss(PDELoss):
-    def __init__(self, dataset, name, hpm_input, hpm_model, norm='L2', weight=1.):
+    def __init__(self, geometry, name, hpm_input, hpm_model, norm='L2', weight=1.):
         """
         Constructor of the HPM loss
+        
         Args:
-            dataset (torch.utils.Dataset): dataset that provides the residual points
+            geometry: instance of the geometry class that defines the domain
             hpm_input(function): function that calculates the needed input for the HPM model. The hpm_input function
             should return a list of tensors, where the last entry is the time_derivative
             hpm_model (torch.nn.Module): model for the HPM, represents the underlying PDE
             norm: Norm used for calculation PDE loss
             weight: Weighting for the loss term
         """
-        super(HPMLoss, self).__init__(dataset, None, name, norm='L2', weight=1.)
+        super(HPMLoss, self).__init__(geometry, None, name, norm='L2', weight=1.)
         self.hpm_input = hpm_input
         self.hpm_model = hpm_model
 

PINNFramework/LHS_Sampler.py

@@ -0,0 +1,31 @@
+import torch
+import numpy as np
+from pyDOE import lhs
+from .Sampler import Sampler
+
+
+class LHSSampler(Sampler):
+    def __init__(self):
+        """
+        Constructor of the LHSSampler class
+        """
+        super(LHSSampler, self).__init__()
+
+    def sample(self, lb, ub, n):        
+        """Generate 'n' number of sample points in [lb,ub]
+        
+        Args:
+            lb (numpy.ndarray): lower bound of the domain.
+            ub (numpy.ndarray): upper bound of the domain.
+            n (int): the number of sampled points.
+        """
+            
+        torch.manual_seed(42)
+        np.random.seed(42)
+        
+        lb =  lb.reshape(1,-1)
+        ub =  ub.reshape(1,-1)
+        
+        dimension = lb.shape[1]
+        xf = lb + (ub - lb) * lhs(dimension, n)
+        return torch.tensor(xf).float()

PINNFramework/ND_Cube.py

@@ -0,0 +1,36 @@
+import numpy as np
+import torch
+from .Geometry import Geometry
+
+class NDCube(Geometry):
+    def __init__(self, lb, ub, n_points, batch_size, sampler):
+        """
+        Constructor of the NDCube class
+
+        Args:
+            lb (numpy.ndarray): lower bound of the domain.
+            ub (numpy.ndarray): upper bound of the domain.
+            n_points (int): the number of sampled points.
+            batch_size (int): batch size
+            sampler: instance of the Sampler class.
+        """
+        super(NDCube, self).__init__(lb, ub, n_points, batch_size, sampler)
+        
+        
+    def __getitem__(self, idx):
+        """
+        Returns data at given index
+        Args:
+            idx (int)
+        """
+        self.x = self.sampler.sample(self.lb,self.ub, self.batch_size) 
+        
+        if type(self.x) is tuple:
+            x, w = self.x
+            return torch.cat((x, w), 1)
+        else:
+            return self.x
+
+    def __len__(self):
+        """Length of the dataset"""
+        return self.n_points // self.batch_size

PINNFramework/PDELoss.py

@@ -2,21 +2,22 @@
 from torch import Tensor as Tensor
 from torch.nn import Module as Module
 from .LossTerm import LossTerm
+from .Adaptive_Sampler import AdaptiveSampler
 
 
 class PDELoss(LossTerm):
-    def __init__(self, dataset, pde, name, norm='L2', weight=1.):
+    def __init__(self, geometry, pde, name, norm='L2', weight=1.):
         """
         Constructor of the PDE Loss
 
         Args:
-            dataset (torch.utils.Dataset): dataset that provides the residual points
+            geometry: instance of the geometry class that defines the domain
             pde (function): function that represents residual of the PDE
             norm: Norm used for calculation PDE loss
             weight: Weighting for the loss term
         """
-        super(PDELoss, self).__init__(dataset, name, norm, weight)
-        self.dataset = dataset
+        super(PDELoss, self).__init__(geometry, name, norm, weight)
+        self.geometry = geometry
         self.pde = pde
 
     def __call__(self, x: Tensor, model: Module, **kwargs):
@@ -26,8 +27,17 @@ def __call__(self, x: Tensor, model: Module, **kwargs):
         x: residual points
         model: model that predicts the solution of the PDE
         """
+
+        if isinstance(self.geometry.sampler, AdaptiveSampler):
+            w = x[:,-1:]
+            x = x[:,:-1]
+        
         x.requires_grad = True  # setting requires grad to true in order to calculate
         u = model.forward(x)
         pde_residual = self.pde(x, u, **kwargs)
-        zeros = torch.zeros(pde_residual.shape, device=pde_residual.device)
-        return self.norm(pde_residual, zeros)
+        
+        if isinstance(self.geometry.sampler, AdaptiveSampler):
+            return 1 / self.geometry.batch_size * torch.mean(1 / w * pde_residual ** 2)
+        else:
+            zeros = torch.zeros(pde_residual.shape, device=pde_residual.device)
+            return self.norm(pde_residual, zeros)

PINNFramework/PINN.py

@@ -10,6 +10,7 @@
 from .PDELoss import PDELoss
 from .JoinedDataset import JoinedDataset
 from .HPMLoss import HPMLoss
+from .Adaptive_Sampler import AdaptiveSampler
 from torch.autograd import grad as grad
 from PINNFramework.callbacks import CallbackList
 
@@ -96,11 +97,14 @@ def __init__(self, model: torch.nn.Module, input_dimension: int, output_dimensio
             self.is_hpm = False
         else:
             raise TypeError("PDE loss has to be an instance of a PDE Loss class")
-
+            
         if isinstance(pde_loss, HPMLoss):
             self.is_hpm = True
             if self.use_gpu:
                 self.pde_loss.hpm_model.cuda()
+        
+        if isinstance(pde_loss.geometry.sampler, AdaptiveSampler):
+            self.pde_loss.geometry.sampler.device = torch.device("cuda" if self.use_gpu else "cpu")
 
         if isinstance(initial_condition, InitialCondition):
             self.initial_condition = initial_condition
@@ -110,12 +114,12 @@ def __init__(self, model: torch.nn.Module, input_dimension: int, output_dimensio
         if not len(initial_condition.dataset):
             raise ValueError("Initial condition dataset is empty")
 
-        if not len(pde_loss.dataset):
-            raise ValueError("PDE dataset is empty")
-
+        if not len(pde_loss.geometry):
+            raise ValueError("Geometry is empty")
+                             
         joined_datasets = {
             initial_condition.name: initial_condition.dataset,
-            pde_loss.name: pde_loss.dataset
+            pde_loss.name: pde_loss.geometry
         }
         if self.rank == 0:
             self.loss_log[initial_condition.name] = float(0.0)  # adding initial condition to the loss_log
@@ -318,14 +322,14 @@ def pinn_loss(self, training_data, track_gradient=False, annealing=False):
         # unpack training data
         # ============== PDE LOSS ============== "
         if type(training_data[self.pde_loss.name]) is not list:
-            pde_loss = self.pde_loss(training_data[self.pde_loss.name][0].type(self.dtype), self.model)
+            pde_loss = self.pde_loss(training_data[self.pde_loss.name][0].type(self.dtype), self.model)             
             if annealing or track_gradient:
                 self.loss_gradients_storage[self.pde_loss.name] = self.loss_gradients(pde_loss)
             pinn_loss = pinn_loss + self.pde_loss.weight * pde_loss
             if self.rank == 0:
                 self.loss_log[self.pde_loss.name] = pde_loss + self.loss_log[self.pde_loss.name]
         else:
-            raise ValueError("Training Data for PDE data is a single tensor consists of residual points ")
+            raise ValueError("Training Data for PDE data is either a single tensor consisting of residual points or a concatenation of residual points and corresponding weights ")
 
         # ============== INITIAL CONDITION ============== "
         if type(training_data[self.initial_condition.name]) is list:
@@ -673,4 +677,4 @@ def take_snapshot(model, file_path, device, n_points):
         # convert all tensors to numpy arrays and save as VTK data
         grid = [x.numpy(), y.numpy(), z.numpy()]
         output = output.view(n_points).to('cpu').numpy()
-        imageToVTK(file_path, grid, pointData={"model output": output})
+        imageToVTK(file_path, grid, pointData={"model output": output})

PINNFramework/Random_Sampler.py

@@ -0,0 +1,30 @@
+import torch
+import numpy as np
+from .Sampler import Sampler
+
+
+class RandomSampler(Sampler):
+    def __init__(self):
+        """
+        Constructor of the RandomSampler (pseudo random sampler) class       
+        """
+        super(RandomSampler, self).__init__()        
+
+    def sample(self, lb, ub, n):
+        """Generate 'n' number of sample points in [lb,ub]
+        
+        Args:
+            lb (numpy.ndarray): lower bound of the domain.
+            ub (numpy.ndarray): upper bound of the domain.
+            n (int): the number of sampled points.
+        """
+            
+        torch.manual_seed(42)
+        np.random.seed(42)
+        
+        lb =  lb.reshape(1,-1)
+        ub =  ub.reshape(1,-1)
+        
+        dimension = lb.shape[1]
+        xf = np.random.uniform(lb,ub,size=(n, dimension))
+        return torch.tensor(xf).float()

PINNFramework/Sampler.py

@@ -0,0 +1,20 @@
+import torch
+import numpy as np
+from abc import ABC, abstractmethod
+
+class Sampler(ABC):
+    def __init__(self):
+        """
+        Constructor of the Sampler class
+        """
+
+    @abstractmethod 
+    def sample(self, lb, ub, n):
+        """Generate 'n' number of sample points in [lb,ub]
+        
+        Args:
+            lb (numpy.ndarray): lower bound of the domain.
+            ub (numpy.ndarray): upper bound of the domain.
+            n (int): the number of sampled points.
+        """
+        

PINNFramework/__init__.py

@@ -11,6 +11,10 @@
 from .WandB_Logger import WandbLogger
 from .TensorBoard_Logger import TensorBoardLogger
 from .PINN import PINN
+from .Random_Sampler import RandomSampler
+from .LHS_Sampler import LHSSampler
+from .Adaptive_Sampler import AdaptiveSampler
+from .ND_Cube import NDCube
 
 import PINNFramework.models
 import PINNFramework.callbacks
@@ -24,6 +28,10 @@
     'NeumannBC',
     'TimeDerivativeBC',
     'PDELoss',
+    'RandomSampler',
+    'LHSSampler',
+    'AdaptiveSampler',
+    'NDCube'
     'HPMLoss',
     'PINN',
     'models',