Add deepsets and basic test (#47)

* Add deepsets and basic test

* Update requirements

* Removed PL tests and callbacks

Author: jlperla

econ_layers/callbacks.py

@@ -153,3 +153,102 @@ def forward(self, input):
             return self.OutputRescalingLayer(input, out)
         else:
             return out
+
+class DeepSet(nn.Module):
+    def __init__(
+        self,
+        n_in: int,
+        n_out: int,
+        L: int,        
+        phi_layers: int,
+        rho_layers: int,
+        phi_hidden_dim: int = 128,
+        rho_hidden_dim: int = 128,
+        phi_activator: Optional[nn.Module] = lazy_instance(nn.ReLU),
+        phi_hidden_bias: bool = True,
+        phi_last_activator: Optional[nn.Module] = lazy_instance(nn.Identity),
+        phi_last_bias=True,
+        rho_activator: Optional[nn.Module] = lazy_instance(nn.ReLU),
+        rho_hidden_bias: bool = True,
+        rho_last_activator: Optional[nn.Module] = lazy_instance(nn.Identity),
+        rho_last_bias=True,        
+        OutputRescalingLayer: Optional[nn.Module] = None,
+        InputRescalingLayer: Optional[nn.Module] = None,
+    ):
+        """
+        Init method.
+        """
+        assert n_in == 1  # only supporting univariate states for now
+        super().__init__()  # init the base class
+        self.rho = FlexibleSequential(
+            L,
+            n_out,
+            rho_layers,
+            rho_hidden_dim,
+            rho_activator,
+            rho_hidden_bias,
+            rho_last_activator,
+            rho_last_bias,
+            OutputRescalingLayer = OutputRescalingLayer,
+        )
+
+        self.phi = FlexibleSequential(
+            n_in,
+            L,
+            phi_layers,
+            phi_hidden_dim,
+            phi_activator,
+            phi_hidden_bias,
+            phi_last_activator,
+            phi_last_bias,
+            InputRescalingLayer = InputRescalingLayer,
+        )
+
+    def forward(self, X):
+        num_batches, N = X.shape
+        phi_X = torch.stack(
+            [torch.mean(self.phi(X[i, :].reshape([N, 1])), 0) for i in range(num_batches)]
+        )
+        return self.rho(phi_X)
+
+
+
+class DeepSetMoments(nn.Module):
+    def __init__(
+        self,
+        n_in: int,
+        n_out: int,
+        L: int,        
+        rho_layers: int,
+        rho_hidden_dim: int = 128,
+        rho_activator: Optional[nn.Module] = lazy_instance(nn.ReLU),
+        rho_hidden_bias: bool = True,
+        rho_last_activator: Optional[nn.Module] = lazy_instance(nn.Identity),
+        rho_last_bias=True,        
+        OutputRescalingLayer: Optional[nn.Module] = None,
+    ):
+        """
+        Init method.
+        """
+        assert n_in == 1  # only supporting univariate states for now
+        super().__init__()  # init the base class
+        self.rho = FlexibleSequential(
+            L,
+            n_out,
+            rho_layers,
+            rho_hidden_dim,
+            rho_activator,
+            rho_hidden_bias,
+            rho_last_activator,
+            rho_last_bias,
+            OutputRescalingLayer = OutputRescalingLayer,
+        )
+
+        self.phi = Moments(L)
+
+    def forward(self, X):
+        num_batches, N = X.shape
+        phi_X = torch.stack(
+            [torch.mean(self.phi(X[i, :].reshape([N, 1])), 0) for i in range(num_batches)]
+        )
+        return self.rho(phi_X)

econ_layers/layers.py

@@ -1,6 +1,5 @@
-pytorch-lightning >= 1.5.2
 torch >= 1.10
 pytest
 black
 numpy
-jsonargparse >= 4.0.4
+jsonargparse[signatures] >= 4.19

requirements.txt

@@ -57,6 +57,6 @@
     test_suite="tests",
     tests_require=test_requirements,
     url="https://github.com/HighDimensionalEconLab/econ_layers",
-    version="0.0.27",
+    version="0.0.28",
     zip_safe=False,
 )

setup.py

@@ -0,0 +1,42 @@
+"""Tests for Exponential layer"""
+
+import pytest
+import torch
+import numpy as np
+import torch
+from torch import nn
+import numpy.testing
+import torch.autograd.gradcheck
+from torch.autograd import Variable
+
+from econ_layers.layers import (
+    FlexibleSequential,
+    RescaleOutputsByInput,
+    ScalarExponentialRescaling,
+    DeepSet,
+    DeepSetMoments,
+)
+
+torch.set_printoptions(16)  # to be able to see what is going on
+
+
+# Unit testing of the autodiff.  Easy here, but need to be more advanced later
+def test_deep_set_derivative():
+    n_in = 1 # one dimensional state per "agent".  Only n_in = 1 is supported right now
+    N = 5 # number of "agents"
+    n_out = 2
+    mod = DeepSet(n_in, n_out, L = 2, phi_layers=2, phi_hidden_dim=32, rho_layers=2, rho_hidden_dim=32).double()
+    batches = 10
+    input = (Variable(torch.randn(N, batches).double(), requires_grad=True),)
+    assert torch.autograd.gradcheck(mod, input)
+
+
+def test_deep_set_moments_derivative():
+    n_in = 1 # one dimensional state per "agent".  Only n_in = 1 is supported right now
+    N = 5 # number of "agents"
+    n_out = 2
+    num_moments = 3
+    mod = DeepSetMoments(n_in, n_out, L = num_moments, rho_layers=2, rho_hidden_dim=32).double()
+    batches = 10
+    input = (Variable(torch.randn(N, batches).double(), requires_grad=True),)
+    assert torch.autograd.gradcheck(mod, input)