feat(network): implement a MonomialNN
a MonomialNN expands every input to a specified list of powers
import torch
import torch.nn as nn
-from .function_basis import RealSphericalHarmonics
+from warnings import warn
class FCNN(nn.Module):
@@ -38,6 +38,38 @@ class FCNN(nn.Module):
return x
+class MonomialNN(nn.Module):
+ def __init__(self, n_input_units, degrees=1):
+ super(MonomialNN, self).__init__()
+
+ self.n_input_units = n_input_units
+ if isinstance(degrees, int):
+ degrees = [d for d in range(1, degrees + 1)]
+ self.degrees = tuple(degrees)
+
+ if len(self.degrees) == 0:
+ raise ValueError(f"No degrees used, check `degrees` argument again")
+ if 0 in degrees:
+ warn("One of the degrees is 0 which might introduce redundant features")
+ if len(set(self.degrees)) < len(self.degrees):
+ warn(f"Duplicate degrees found: {self.degrees}")
+
+ def forward(self, x):
+ if len(x.shape) != 2 or x.shape[1] != self.n_input_units:
+ raise ValueError(f"Expected input shape (-1, {self.n_input_units}), got {x.shape} instead")
+ return torch.cat([x ** d for d in self.degrees], dim=1)
+
+ def __repr__(self):
+ return f"{self.__class__.__name__}(" \
+ f"n_input_units={self.n_input_units}" \
+ f", " \
+ f"degrees={self.degrees}" \
+ f")"
+
+ def __str__(self):
+ return self.__repr__()
+
+
class SinActv(nn.Module):
"""The sin activation function.
"""
