Merge pull request #4 from ComputationalDesignLab/abhi-add-rbf-implementation

Add RBF model implementation

Author: Abhijnan03

scimlstudio/models/__init__.py

@@ -1 +1,2 @@
-from .polynomial import Polynomial
+from .polynomial import Polynomial
+from .rbf import RBF

scimlstudio/models/rbf.py

@@ -0,0 +1,187 @@
+import torch
+from ..base_model import BaseModel
+from ..utils.transformations import Standardize, Normalize
+
+class RBF(BaseModel):
+
+    def __init__(self, x_train: torch.Tensor, y_train: torch.Tensor, sigma: float,
+                 input_transform: Normalize | Standardize | None = None, output_transform: Normalize | Standardize | None = None, 
+                 basis: str = "gaussian"):
+
+        """
+            Class definition for parametric radial basis function models
+
+            f(x) = w^T psi(x), psi_i(x) = psi(r), where r = ||x - c_i||
+
+            NOTE: Only the Gaussian and multiquadric basis functions are implemented
+
+            Parameters
+            ----------
+            xtrain: torch.Tensor
+                Input training data for the RBF model with size (N, d)
+
+            ytrain: torch.Tensor
+                Output training data for the RBF model with size (N, 1)
+
+            sigma: float
+                Value of sigma for calculating the basis function
+
+            input_transform:  Normalize | Standardize | None
+                Object of a transformation class for applying a transform to the inputs of the model
+                The default is None which means the application of no transform to inputs
+
+            output_transform:  Normalize | Standardize | None
+                Object of a transformation for applying a transform to the outputs of the model
+                The default is None which means the application of no transform to outputs
+
+            basis: str
+                String that specifies which parametric basis function should be used
+                User has a choice between "gaussian", "multiquadric" and "inverse"
+
+                gaussian --> exp(-r^2/(2 * sigma^2))
+                multiquadric --> (r^2 + sigma^2)^0.5
+                inverse --> (r^2 + sigma^2)^(-0.5)        
+        """
+
+        super().__init__()
+
+        # Checking inputs
+        assert isinstance(x_train, torch.Tensor) and x_train.ndim == 2, "xtrain must be a 2D tensor array"
+        assert isinstance(y_train, torch.Tensor) and y_train.ndim == 2, "ytrain must be a 2D tensor array"
+        assert x_train.shape[0] == y_train.shape[0], "number of samples in input and output training data must be the same"
+        assert x_train.device == y_train.device, "input and output training data must be on the same device"
+        assert isinstance(sigma, float) and sigma > 0, "sigma value must be a positive floating point value"
+        assert isinstance(basis, str), "basis choice must be a string"
+        assert basis in ["gaussian", "multiquadric", "inverse"], "basis choice is not a valid choice. choice must be one of [gaussian, multiquadric, inverse]"
+
+        self.sigma = sigma
+        self.basis = basis
+        self.input_transform = input_transform
+        self.output_transform = output_transform
+        self.basis_function = self.set_basis_function(self.basis)
+        
+        # Set training data for the model
+        self.xtrain = self.transform_values(x_train, input_transform)
+        self.ytrain = self.transform_values(y_train, output_transform)
+
+    def fit(self) -> torch.Tensor:
+
+        """
+            Method for fitting the model to the provided training data
+
+            Returns
+            -------
+            rbf_weights: torch.Tensor
+                Tensor array with the fitted weights of the model
+
+            basis_matrix: torch.Tensor
+                Tensor array containing the basis matrix for the model
+        """
+        
+        # Basis matrix definition
+        ns = self.xtrain.shape[0]
+        basis_matrix = torch.zeros((ns, ns)).to(self.xtrain)
+
+        # Assigning values to the basis matrix
+        for i in range(ns):
+            for j in range(ns):
+                # Calculating the value of r
+                r = torch.linalg.norm(self.xtrain[i] - self.xtrain[j], ord=2)
+
+                # Putting value in basis matrix 
+                basis_matrix[i, j] = self.basis_function(r)
+        
+        basis_matrix_inverse = torch.linalg.pinv(basis_matrix)
+        self.rbf_weights = torch.matmul(basis_matrix_inverse, self.ytrain)
+
+        return self.rbf_weights.clone(), basis_matrix
+    
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+
+        """
+            Method for to predict values from the model for specified points
+
+            Parameters
+            ----------
+            x: torch.Tensor
+                Data for which predictions must be made. Must be a 2D tensor array
+
+            Returns
+            -------
+            ypred: torch.Tensor
+                Predictions of the model on the specified points
+        """
+        # Checks for training and inputs
+        assert hasattr(self, "rbf_weights"), "model has not been trained. call fit method before predict"
+        assert isinstance(x, torch.Tensor) and x.ndim == 2, "input data provided to the model must be a 2D tensor array"
+        assert x.device == self.xtrain.device, "provided input data must be on the same device as the training data"
+
+        x = self.transform_values(x, self.input_transform)
+
+        # Basis matrix definition
+        basis_matrix = torch.zeros((x.shape[0], self.xtrain.shape[0])).to(self.xtrain)
+
+        # Assigning values to the basis matrix
+        for i in range(x.shape[0]):
+            for j in range(self.xtrain.shape[0]):
+                r = torch.linalg.norm(x[i] - self.xtrain[j], ord=2)
+                basis_matrix[i, j] = self.basis_function(r)
+
+        ypred = torch.matmul(basis_matrix, self.rbf_weights)
+        if self.output_transform is not None:
+            ypred = self.output_transform.inverse_transform(ypred)
+
+        return ypred
+        
+    def set_basis_function(self, basis: str): 
+
+        """
+            Method for setting the basis function of the RBF model based on user input
+
+            Parameters
+            ----------
+            basis: str
+                String that specifies the user choice for the basis function
+
+            Returns
+            -------
+            basis_function: 
+                Python function that calculates the values of the basis function based on the value of r
+        """
+        if basis == "gaussian":
+            basis_function = lambda r: torch.exp(-r**2/(2 * self.sigma ** 2))
+        elif basis == "multiquadric":
+            basis_function = lambda r: torch.sqrt(r**2 + self.sigma ** 2)
+        elif basis == "inverse":
+            basis_function = lambda r: 1 / torch.sqrt(r**2 + self.sigma ** 2)
+
+        return basis_function
+
+    def transform_values(self, data: torch.Tensor, transform: Normalize | Standardize | None) -> torch.Tensor:
+
+        """
+            Method for transforming values based on given transform
+
+            Parameters
+            ----------
+            data: torch.Tensor
+                Data that must be transformed
+
+            transform: Normalize | Standardize | None
+                Object of a transformation class that will be used to transform the data
+
+            Returns
+            -------
+            transformed_data: torch.Tensor
+                Data after the application of the transform
+                If transform is None, then the original data is returned
+        """
+
+        # Checking if transform is None and applying the transform accordingly
+        if transform is None:
+            transformed_data = data
+        else:
+            assert isinstance(transform, Normalize) or isinstance(transform, Standardize), "transform must be an instance of Normalize or Standardize class"
+            transformed_data = transform.transform(data)
+
+        return transformed_data

tests/test_rbf.py

@@ -0,0 +1,118 @@
+import unittest, torch, math
+from scimlstudio.utils import evaluate_scalar
+from scimlstudio.models import RBF
+from scimlstudio.utils import Normalize, Standardize
+from pyDOE3 import lhs, halton_sequence
+
+# Defining the device and data type
+tkwargs = {"device": torch.device("cuda" if torch.cuda.is_available() else "cpu"), "dtype": torch.float64}
+
+class TestRBF(unittest.TestCase):
+
+    """
+        Class defining the test cases for the RBF model
+    """
+    
+    def test_rbf(self):
+
+        # Generating some data for the test cases
+        # Function used for the test cases is a sinusoidal function (https://www.sfu.ca/~ssurjano/curretal88sin.html)
+        x_train = torch.linspace(0, 1, 10, **tkwargs)
+        y_train = torch.sin(2*math.pi*(x_train - 0.1))
+
+        x_test = torch.linspace(0, 1, 100, **tkwargs)
+        y_test = torch.sin(2*math.pi*(x_test - 0.1))
+
+        # Creating and fitting RBF models
+        for basis in ["gaussian", "multiquadric", "inverse"]:
+            rbf = RBF(x_train.reshape(-1,1), y_train.reshape(-1,1), sigma = 0.1, basis = basis)
+            _, basis_matrix = rbf.fit()
+            train_pred = rbf.predict(x_train.reshape(-1,1))
+            test_pred = rbf.predict(x_test.reshape(-1,1))
+            r2_value = evaluate_scalar(test_pred.reshape(-1,), y_test, "r2")
+
+            torch.testing.assert_close(train_pred.reshape(-1,), y_train, rtol=0, atol=1e-6, check_device=True, check_dtype=True) # interpolation check
+            assert basis_matrix.ndim == 2 # basis matrix must be 2D
+            if basis == "gaussian":
+                assert (torch.diag(basis_matrix) == 1).all() # diagonal of basis matrix must be all 1s if basis is gaussian
+            assert round(r2_value, 6) < 1 
+
+        # Creating and fitting RBF models with normalize transformations
+        input_transform = Normalize(x_train.reshape(-1,1))
+        output_transform = Normalize(y_train.reshape(-1,1))
+        for basis in ["gaussian", "multiquadric", "inverse"]:
+            rbf = RBF(x_train.reshape(-1,1), y_train.reshape(-1,1), sigma = 0.1, basis = basis, input_transform=input_transform, output_transform=output_transform)
+            _, basis_matrix = rbf.fit()
+            train_pred = rbf.predict(x_train.reshape(-1,1))
+            test_pred = rbf.predict(x_test.reshape(-1,1))
+            r2_value = evaluate_scalar(test_pred.reshape(-1,), y_test, "r2")
+
+            torch.testing.assert_close(train_pred.reshape(-1,), y_train, rtol=0, atol=1e-6, check_device=True, check_dtype=True) # interpolation check
+            assert basis_matrix.ndim == 2 # basis matrix must be 2D
+            if basis == "gaussian":
+                assert (torch.diag(basis_matrix) == 1).all() # diagonal of basis matrix must be all 1s if basis is gaussian
+            assert round(r2_value, 6) < 1 
+
+    def test_rbf_5D(self):
+
+        # Generating some data for the test cases
+        # Function used for the test cases is the Friedman function (https://www.sfu.ca/~ssurjano/fried.html)
+        x_train = torch.tensor(halton_sequence(num_points=15, dimension=5), **tkwargs)
+        y_train = 10*torch.sin(math.pi*x_train[:,0]*x_train[:,1]) + 20 * (x_train[:,2] - 0.5) ** 2 + 10 * x_train[:,3] + 5 * x_train[:,4]
+
+        x_test = torch.tensor(lhs(n=5, samples=100, criterion='cm', iterations=100, seed=10), **tkwargs)
+        y_test = 10*torch.sin(math.pi*x_test[:,0]*x_test[:,1]) + 20 * (x_test[:,2] - 0.5) ** 2 + 10 * x_test[:,3] + 5 * x_test[:,4]
+
+        # Creating and fitting RBF models
+        for basis in ["gaussian", "multiquadric", "inverse"]:
+            rbf = RBF(x_train, y_train.reshape(-1,1), sigma = 0.1, basis = basis)
+            _, basis_matrix = rbf.fit()
+            train_pred = rbf.predict(x_train)
+            test_pred = rbf.predict(x_test)
+            r2_value = evaluate_scalar(test_pred.reshape(-1,), y_test, "r2")
+
+            torch.testing.assert_close(train_pred.reshape(-1,), y_train, rtol=0, atol=1e-6, check_device=True, check_dtype=True) # interpolation check
+            assert basis_matrix.ndim == 2 # basis matrix must be 2D
+            if basis == "gaussian":
+                assert (torch.diag(basis_matrix) == 1).all() # diagonal of basis matrix must be all 1s if basis is gaussian
+            assert round(r2_value, 6) < 1 
+
+        # Creating and fitting RBF models with standardize transformations
+        input_transform = Standardize(x_train)
+        output_transform = Standardize(y_train.reshape(-1,1))
+        for basis in ["gaussian", "multiquadric", "inverse"]:
+            rbf = RBF(x_train, y_train.reshape(-1,1), sigma = 0.1, basis = basis, input_transform=input_transform, output_transform=output_transform)
+            _, basis_matrix = rbf.fit()
+            train_pred = rbf.predict(x_train)
+            test_pred = rbf.predict(x_test)
+            r2_value = evaluate_scalar(test_pred.reshape(-1,), y_test, "r2")
+
+            torch.testing.assert_close(train_pred.reshape(-1,), y_train, rtol=0, atol=1e-6, check_device=True, check_dtype=True) # interpolation check
+            assert basis_matrix.ndim == 2 # basis matrix must be 2D
+            if basis == "gaussian":
+                assert (torch.diag(basis_matrix) == 1).all() # diagonal of basis matrix must be all 1s if basis is gaussian
+            assert round(r2_value, 6) < 1
+
+    def test_inputs(self):
+
+        # Generate dummy data
+        x_random = torch.rand(15)
+        y_random = torch.rand(15)
+
+        with self.assertRaises(Exception):
+            _ = RBF(x_random.reshape(-1,1), y_random, sigma = 0.1)
+
+        with self.assertRaises(Exception):
+            _ = RBF(x_random, y_random.reshape(-1,1), sigma = 0.1)
+
+        with self.assertRaises(Exception):
+            _ = RBF(x_random.reshape(-1,1), y_random.reshape(-1,1), sigma = -0.1)
+
+        with self.assertRaises(Exception):
+            _ = RBF(x_random.reshape(-1,1), y_random.reshape(-1,1), sigma = 0.1, basis="linear")
+
+        _ = RBF(x_random.reshape(-1,1), y_random.reshape(-1,1), sigma = 0.1, basis="inverse")
+
+if __name__ == '__main__':
+    unittest.main()
+