From 87b57104a5cd14cd7a1c7374d3d04b4652727300 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Mon, 11 Apr 2022 16:24:14 -0700
Subject: [PATCH] Add tests

---
 .github/workflows/pytest.yml         |  9 ++--
 environment/requirements.txt         |  2 -
 mattport/nerf/modules/encoding.py    | 42 ++++++++++++------
 tests/nerf/modules/test_encodings.py | 65 ++++++++++++++++++++++++++++
 tests/nerf/test_encodings.py         |  9 ----
 5 files changed, 97 insertions(+), 30 deletions(-)
 create mode 100644 tests/nerf/modules/test_encodings.py
 delete mode 100644 tests/nerf/test_encodings.py

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 320012ddc..53b264d9d 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -16,16 +16,15 @@ jobs:
 
     steps:
     - uses: actions/checkout@v3
-    - name: Set up Python 3.10
+    - name: Set up Python 3.8.13
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "3.8.13"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install flake8 pytest
-        if [ -f environment/requirements.txt ]; then pip install -r environment/requirements.txt; fi
+        pip install -r environment/requirements.txt
         python setup.py develop
     - name: Test with pytest
       run: |
-        pytest tests/nerf/test_graph.py
+        pytest tests
diff --git a/environment/requirements.txt b/environment/requirements.txt
index ef64b7b0c..0c43402bc 100644
--- a/environment/requirements.txt
+++ b/environment/requirements.txt
@@ -4,6 +4,4 @@ pytest
 torch==1.11.0
 torchtyping
 torchvision==0.12.0
-torchaudio==0.11.0
 black
-flake8
\ No newline at end of file
diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/modules/encoding.py
index 0146d04cb..53ce14854 100644
--- a/mattport/nerf/modules/encoding.py
+++ b/mattport/nerf/modules/encoding.py
@@ -1,7 +1,6 @@
 """
 Encoding functions
 """
-from base64 import encode
 import torch
 from torch import nn
 from torchtyping import TensorType
@@ -16,6 +15,8 @@ class Encoding(nn.Module):
             in_dim (int): Input dimension of tensor
         """
         super().__init__()
+        if in_dim <= 0:
+            raise ValueError("Input dimension should be greater than zero")
         self.in_dim = in_dim
 
     def get_out_dim(self) -> int:
@@ -55,12 +56,11 @@ class ScalingAndOffset(Encoding):
         """
         super().__init__(in_dim)
 
-        self.in_dim = in_dim
         self.scaling = scaling
         self.offset = offset
 
     def get_out_dim(self) -> int:
-        return self.in_dim * self.num_frequencies * 2
+        return self.in_dim
 
     def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
         return self.scaling * in_tensor + self.offset
@@ -105,16 +105,36 @@ class NeRFEncoding(Encoding):
 class RFFEncoding(Encoding):
     """Random Fourier Feature encoding"""
 
-    def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
-        """Encodes an input tensor.
+    def __init__(self, in_dim: int, num_frequencies: int, scale: float) -> None:
+        """
 
         Args:
-            x (TensorType[..., "input_dim"]): Input tensor to be encoded
+            in_dim (int): Input dimension of tensor
+            num_frequencies (int): Number of encoding frequencies
+            scale (float): Std of Gaussian to sample frequencies. Must be greater than zero
+        """
+        super().__init__(in_dim)
+
+        self.num_frequencies = num_frequencies
+        if not scale > 0:
+            raise ValueError("RFF encoding scale should be greater than zero")
+        self.scale = scale
+
+    def get_out_dim(self) -> int:
+        return self.num_frequencies * 2
+
+    def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+        """
+        Args:
+            in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1.
 
         Returns:
-            TensorType[..., "output_dim"]: A encoded input tensor
+            TensorType[..., "output_dim"]: Output values will be between -1 and 1
         """
-        raise NotImplementedError
+        B = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies))
+        scaled_inputs = in_tensor @ B  # [..., "num_frequencies"]
+        encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1)
+        return encoded_inputs
 
 
 class HashEncoding(Encoding):
@@ -130,9 +150,3 @@ class HashEncoding(Encoding):
             TensorType[..., "output_dim"]: A encoded input tensor
         """
         raise NotImplementedError
-
-
-encodings_dict = {
-    "RFF": RFFEncoding,
-    "hash": HashEncoding,
-}
diff --git a/tests/nerf/modules/test_encodings.py b/tests/nerf/modules/test_encodings.py
new file mode 100644
index 000000000..47bd74668
--- /dev/null
+++ b/tests/nerf/modules/test_encodings.py
@@ -0,0 +1,65 @@
+"""
+Encoding Tests
+"""
+import pytest
+import torch
+from mattport.nerf.modules import encoding
+
+
+def test_scaling_and_offset():
+    """Test scaling and offset encoder"""
+    in_dim = 4
+    in_tensor = torch.ones((2, 3, in_dim))
+
+    scaling = 2.0
+    offset = 4.5
+    encoder = encoding.ScalingAndOffset(in_dim=in_dim, scaling=scaling, offset=offset)
+
+    assert encoder.get_out_dim() == in_dim
+    encoded = encoder.encode(in_tensor)
+    assert encoded.shape[-1] == in_dim
+    assert in_tensor * 6.5 == pytest.approx(encoded)
+
+    with pytest.raises(ValueError):
+        encoding.ScalingAndOffset(in_dim=-1)
+
+
+def test_nerf_encoder():
+    """Test NeRF encoder"""
+    in_dim = 4
+    out_dim = 24
+    in_tensor = torch.ones((2, 3, in_dim))
+
+    num_frequencies = 3
+    min_freq_exp = 0
+    max_freq_exp = 3
+    encoder = encoding.NeRFEncoding(
+        in_dim=in_dim, num_frequencies=num_frequencies, min_freq_exp=min_freq_exp, max_freq_exp=max_freq_exp
+    )
+    assert encoder.get_out_dim() == out_dim
+
+    in_tensor = torch.ones((2, 3, in_dim))
+    encoded = encoder.encode(in_tensor)
+    assert encoded.shape[-1] == out_dim
+    assert torch.max(encoded) == 1
+
+    in_tensor = torch.zeros((2, 3, in_dim))
+    encoded = encoder.encode(in_tensor)
+    assert encoded.shape[-1] == out_dim
+    assert torch.min(encoded) == 0
+
+
+def test_rff_encoder():
+    """Test RFF encoder"""
+    in_dim = 3
+    out_dim = 24
+    in_tensor = torch.ones((2, 3, in_dim))
+
+    num_frequencies = 12
+    scale = 5
+    encoder = encoding.RFFEncoding(in_dim=in_dim, num_frequencies=num_frequencies, scale=scale)
+    assert encoder.get_out_dim() == out_dim
+
+    in_tensor = torch.ones((2, 3, in_dim))
+    encoded = encoder.encode(in_tensor)
+    assert encoded.shape[-1] == out_dim
diff --git a/tests/nerf/test_encodings.py b/tests/nerf/test_encodings.py
deleted file mode 100644
index 99f26096c..000000000
--- a/tests/nerf/test_encodings.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import pytest
-
-
-def test_positional_encoding():
-    assert True
-
-
-def test_lower_case():
-    assert "MATTPORT".lower() == "mattport"