add OrthogonalBlock to make input orthonormal

pina/model/layers/orthogonal.py

@@ -0,0 +1,50 @@
+"""Module for OrthogonalBlock layer, to make the input orthonormal."""
+
+import torch
+
+
+class OrthogonalBlock(torch.nn.Module):
+    """
+    Module to make the input orthonormal.
+    The module takes a tensor of size [N, M] and returns a tensor of
+    size [N, M] where the columns are orthonormal.
+    """
+    def __init__(self, dim=-1):
+        """
+        Initialize the OrthogonalBlock module.
+
+        :param int dim: The dimension where to orthogonalize.
+        """
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, X):
+        """
+        Forward pass of the OrthogonalBlock module using a Gram-Schmidt
+        algorithm.
+
+        :raises Warning: If the dimension is greater than the other dimensions.
+
+        :param torch.Tensor X: The input tensor to orthogonalize.
+        :return: The orthonormal tensor.
+        """
+        # check dim is less than all the other dimensions
+        if X.shape[self.dim] > min(X.shape):
+            raise Warning("The dimension where to orthogonalize is greater\
+                          than the other dimensions")
+
+        result = torch.zeros_like(X)
+        # normalize first basis
+        X_0 = torch.select(X, self.dim, 0)
+        result_0 = torch.select(result, self.dim, 0)
+        result_0 += X_0/torch.norm(X_0)
+        # iterate over the rest of the basis with Gram-Schmidt
+        for i in range(1, X.shape[self.dim]):
+            v = torch.select(X, self.dim, i)
+            for j in range(i):
+                v -= torch.sum(v * torch.select(result, self.dim, j),
+                               dim=self.dim, keepdim=True) * torch.select(
+                                   result, self.dim, j)
+            result_i = torch.select(result, self.dim, i)
+            result_i += v/torch.norm(v)
+        return result