Merge pull request #40 from andrrizzi/neural

Neural spline transformer improvements

tfep/app/mixedmaf.py

@@ -811,11 +811,30 @@ def _get_transformer(
             maf_periodic_dof_indices = maf_periodic_dof_indices - torch.searchsorted(
                 maf_conditioning_dof_indices, maf_periodic_dof_indices)
 
-        return tfep.nn.transformers.NeuralSplineTransformer(
-            x0=x0.detach(),
-            xf=xf.detach(),
+        # Find all non-periodic DOFs (after filtering the conditioning ones).
+        mask = torch.full(x0.shape, fill_value=True)
+        mask[maf_periodic_dof_indices] = False
+        maf_nonperiodic_dof_indices = torch.tensor(range(len(x0)))[mask]
+
+        # Standard splines for non-periodic DOFs.
+        spline = tfep.nn.transformers.NeuralSplineTransformer(
+            x0=x0[maf_nonperiodic_dof_indices].detach(),
+            xf=xf[maf_nonperiodic_dof_indices].detach(),
             n_bins=5,
-            circular=maf_periodic_dof_indices.detach(),
+            circular=False,
+        )
+
+        # Circular splines for periodic DOFs.
+        circular_spline = tfep.nn.transformers.NeuralSplineTransformer(
+            x0=x0[maf_periodic_dof_indices].detach(),
+            xf=xf[maf_periodic_dof_indices].detach(),
+            n_bins=5,
+            circular=True,
+        )
+
+        return tfep.nn.transformers.MixedTransformer(
+            transformers=[spline, circular_spline],
+            indices=[maf_nonperiodic_dof_indices, maf_periodic_dof_indices],
         )
 
 

tfep/nn/transformers/affine.py

@@ -42,7 +42,7 @@ class AffineTransformer(MAFTransformer):
 
     """
     # Number of parameters needed by the transformer for each input dimension.
-    n_parameters_per_input = 2
+    n_parameters_per_feature = 2
 
     def forward(self, x: torch.Tensor, parameters: torch.Tensor) -> tuple[torch.Tensor]:
         """Apply the affine transformation to the input.
@@ -112,7 +112,7 @@ def get_identity_parameters(self, n_features: int) -> torch.Tensor:
             Shape ``(2*n_features)``. The parameters for the identity.
 
         """
-        return torch.zeros(size=(self.n_parameters_per_input*n_features,))
+        return torch.zeros(size=(self.n_parameters_per_feature*n_features,))
 
     def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
         """Returns the degrees associated to the conditioner's output.
@@ -131,13 +131,13 @@ def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
             transformer as parameters.
 
         """
-        return degrees_in.tile((self.n_parameters_per_input,))
+        return degrees_in.tile((self.n_parameters_per_feature,))
 
     def _split_parameters(self, parameters):
         """Divide shift from log scale."""
         # From (batch, 2*n_features) to (batch, 2, n_features).
         batch_size = parameters.shape[0]
-        parameters = parameters.reshape(batch_size, self.n_parameters_per_input, -1)
+        parameters = parameters.reshape(batch_size, self.n_parameters_per_feature, -1)
         return parameters[:, 0], parameters[:, 1]
 
 
@@ -162,7 +162,7 @@ class VolumePreservingShiftTransformer(MAFTransformer):
 
     """
     # Number of parameters needed by the transformer for each input dimension.
-    n_parameters_per_input = 1
+    n_parameters_per_feature = 1
 
     def __init__(
             self,
@@ -252,7 +252,7 @@ def get_identity_parameters(self, n_features: int) -> torch.Tensor:
             Shape ``(n_features)``. The parameters for the identity.
 
         """
-        return torch.zeros(size=(self.n_parameters_per_input*n_features,))
+        return torch.zeros(size=(self.n_parameters_per_feature*n_features,))
 
     def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
         """Returns the degrees associated to the conditioner's output.
@@ -271,7 +271,7 @@ def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
             transformer as parameters.
 
         """
-        return degrees_in.tile((self.n_parameters_per_input,))
+        return degrees_in.tile((self.n_parameters_per_feature,))
 
 
 # =============================================================================

tfep/nn/transformers/moebius.py

@@ -54,7 +54,7 @@ class MoebiusTransformer(MAFTransformer):
 
     """
     # Number of parameters needed by the transformer for each input dimension.
-    n_parameters_per_input = 1
+    n_parameters_per_feature = 1
 
     def __init__(self, dimension: int, max_radius: float = 0.99, unit_sphere: bool = False):
         """Constructor.
@@ -170,7 +170,7 @@ def get_identity_parameters(self, n_features: int) -> torch.Tensor:
             vector to perform the identity function with a Moebius transformer.
 
         """
-        return torch.zeros(size=(self.n_parameters_per_input*n_features,))
+        return torch.zeros(size=(self.n_parameters_per_feature*n_features,))
 
     def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
         """Returns the degrees associated to the conditioner's output.
@@ -189,7 +189,7 @@ def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
             transformer as parameters.
 
         """
-        return degrees_in.tile((self.n_parameters_per_input,))
+        return degrees_in.tile((self.n_parameters_per_feature,))
 
 
 class SymmetrizedMoebiusTransformer(MAFTransformer):
@@ -220,7 +220,7 @@ class SymmetrizedMoebiusTransformer(MAFTransformer):
 
     """
     # Number of parameters needed by the transformer for each input dimension.
-    n_parameters_per_input = 1
+    n_parameters_per_feature = 1
 
     def __init__(self, dimension: int, max_radius: float = 0.99):
         """Constructor.
@@ -330,7 +330,7 @@ def get_identity_parameters(self, n_features: int) -> torch.Tensor:
             vector to perform the identity function with a Moebius transformer.
 
         """
-        return torch.zeros(size=(self.n_parameters_per_input*n_features,))
+        return torch.zeros(size=(self.n_parameters_per_feature*n_features,))
 
     def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
         """Returns the degrees associated to the conditioner's output.
@@ -349,7 +349,7 @@ def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
             transformer as parameters.
 
         """
-        return degrees_in.tile((self.n_parameters_per_input,))
+        return degrees_in.tile((self.n_parameters_per_feature,))
 
 
 # =============================================================================

tfep/nn/transformers/sos.py

@@ -77,7 +77,7 @@ def parameters_per_polynomial(self):
         return self.degree_polynomials + 1
 
     @property
-    def n_parameters_per_input(self):
+    def n_parameters_per_feature(self):
         """Number of parameters needed by the transformer for each input dimension."""
         return self.parameters_per_polynomial * self.n_polynomials + 1
 
@@ -105,7 +105,7 @@ def forward(self, x: torch.Tensor, parameters: torch.Tensor) -> tuple[torch.Tens
         """
         # From (batch, n_parameters*n_features) to (batch, n_parameters, n_features).
         batch_size = parameters.shape[0]
-        parameters = parameters.reshape(batch_size, self.n_parameters_per_input, -1)
+        parameters = parameters.reshape(batch_size, self.n_parameters_per_feature, -1)
         return sos_polynomial_transformer(x, parameters)
 
     def inverse(self, y: torch.Tensor, parameters: torch.Tensor) -> tuple[torch.Tensor]:
@@ -131,7 +131,7 @@ def get_identity_parameters(self, n_features: int) -> torch.Tensor:
             and degree of the polynomials.
 
         """
-        id_conditioner = torch.zeros(size=(self.n_parameters_per_input, n_features))
+        id_conditioner = torch.zeros(size=(self.n_parameters_per_feature, n_features))
         # The sum of the squared linear parameters must be 1.
         id_conditioner[1::self.parameters_per_polynomial].fill_(np.sqrt(1 / self.n_polynomials))
         return id_conditioner.flatten()
@@ -153,7 +153,7 @@ def get_degrees_out(self, degrees_in: torch.Tensor) -> torch.Tensor:
             transformer as parameters.
 
         """
-        return degrees_in.tile((self.n_parameters_per_input,))
+        return degrees_in.tile((self.n_parameters_per_feature,))
 
 
 # =============================================================================