Repository: willbakst/pytorch-lattice (language: python)

Issue:
Implement Random Tiny Lattices Layer
---

## About:
Add a module that implements the optimized random tiny lattices layer.

---

**Is your feature request related to a problem? Please describe.**
Lattices can only handle so many features, so to combat this one technique is to use an ensemble of lattices that each operate on a subset of the features. The RTL layer optimizes the construction and calling of this ensemble by using the multi-unit functionality of the Lattice layer to parallelize the computation of input subsets with the same constraint structure. 

**Describe the solution you'd like**
A `torch.Module` that implements the random tiny lattices functionality as a layer that can be used to build a `CalibratedLatticeEnsemble`.

**Describe alternatives you've considered**
N/A

**Additional context**
N/A


Discussion / hints:

Fix patch:
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 4de96d1..45c2e19 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,6 +1,8 @@
 {
   "files.autoSave": "onFocusChange",
-  "editor.rulers": [88],
+  "editor.rulers": [
+    88
+  ],
   "editor.formatOnSave": true,
   "editor.formatOnSaveMode": "file",
   "files.insertFinalNewline": true,
diff --git a/pyproject.toml b/pyproject.toml
index 0872a4f..cb9fb42 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -72,7 +72,7 @@ skip-magic-trailing-comma = false
 line-ending = "auto"
 
 [tool.mypy]
-exclude = ["examples"]
+exclude = ["examples", "venv"]
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/pytorch_lattice/layers/__init__.py b/pytorch_lattice/layers/__init__.py
index 15c21e9..b97b156 100644
--- a/pytorch_lattice/layers/__init__.py
+++ b/pytorch_lattice/layers/__init__.py
@@ -3,3 +3,4 @@
 from .lattice import Lattice
 from .linear import Linear
 from .numerical_calibrator import NumericalCalibrator
+from .rtl import RTL
diff --git a/pytorch_lattice/layers/rtl.py b/pytorch_lattice/layers/rtl.py
new file mode 100644
index 0000000..cdda156
--- /dev/null
+++ b/pytorch_lattice/layers/rtl.py
@@ -0,0 +1,267 @@
+"""A PyTorch module implementing a calibrated modeling layer for Random Tiny Lattices.
+
+This module implements an ensemble of tiny lattices that each operate on a subset of the
+inputs. It utilizes the multi-unit functionality of the Lattice module to better
+optimize speed performance by putting feature subsets that have the same constraint
+structure into the same Lattice module as multiple units.
+"""
+import logging
+from typing import Optional, Union
+
+import numpy as np
+import torch
+
+from ..enums import Interpolation, LatticeInit, Monotonicity
+from .lattice import Lattice
+
+
+class RTL(torch.nn.Module):
+    """A module that efficiently implements Random Tiny Lattices.
+
+    This module creates an ensemble of lattices where each lattice in the ensemble takes
+    as input a subset of the input features. For further efficiency, input subsets with
+    the same constraint structure all go through the same lattice as multiple units in
+    parallel. When creating the ensemble structure, features are shuffled and uniformly
+    repeated if there are more available slots in the ensemble structure than there are
+    features.
+
+    Attributes:
+      - All `__init__` arguments.
+
+    Example:
+    ```python
+    inputs=torch.tensor(...)  # shape: (batch_size, D)
+    monotonicities = List[Monotonicity...]  # len: D
+    random_tiny_lattices = RTL(
+      monotonicities,
+      num_lattices=5
+      lattice_rank=3,  # num_lattices * lattice_rank must be greater than D
+    )
+    output1 = random_tiny_lattices(inputs)
+
+    # You can stack RTL modules based on the previous RTL's output monotonicities.
+    rtl2 = RTL(random_tiny_lattices.output_monotonicities(), ...)
+    outputs2 = rtl2(outputs)
+    ```
+    """
+
+    def __init__(
+        self,
+        monotonicities: list[Monotonicity],
+        num_lattices: int,
+        lattice_rank: int,
+        lattice_size: int = 2,
+        output_min: Optional[float] = None,
+        output_max: Optional[float] = None,
+        kernel_init: LatticeInit = LatticeInit.LINEAR,
+        clip_inputs: bool = True,
+        interpolation: Interpolation = Interpolation.HYPERCUBE,
+        average_outputs: bool = False,
+        random_seed: int = 42,
+    ) -> None:
+        """Initializes an instance of 'RTL'.
+
+        Args:
+            monotonicities: List of `Monotonicity.INCREASING` or `None`
+              indicating monotonicities of input features, ordered respectively.
+            num_lattices: number of lattices in RTL structure.
+            lattice_rank: number of inputs for each lattice in RTL structure.
+            output_min: Minimum output of each lattice in RTL.
+            output_max: Maximum output of each lattice in RTL.
+            kernel_init: Initialization scheme to use for lattices.
+            clip_inputs: Whether input should be clipped to the range of each lattice.
+            interpolation: Interpolation scheme for each lattice in RTL.
+            average_outputs: Whether to average the outputs of every lattice RTL.
+            random_seed: seed used for shuffling.
+
+        Raises:
+            ValueError: If size of RTL, determined by `num_lattices * lattice_rank`, is
+             too small to support the number of input features.
+        """
+        super().__init__()
+
+        if len(monotonicities) > num_lattices * lattice_rank:
+            raise ValueError(
+                f"RTL with {num_lattices}x{lattice_rank}D structure cannot support "
+                + f"{len(monotonicities)} input features."
+            )
+        self.monotonicities = monotonicities
+        self.num_lattices = num_lattices
+        self.lattice_rank = lattice_rank
+        self.lattice_size = lattice_size
+        self.output_min = output_min
+        self.output_max = output_max
+        self.kernel_init = kernel_init
+        self.clip_inputs = clip_inputs
+        self.interpolation = interpolation
+        self.average_outputs = average_outputs
+        self.random_seed = random_seed
+
+        rtl_indices = np.array(
+            [i % len(self.monotonicities) for i in range(num_lattices * lattice_rank)]
+        )
+        np.random.seed(self.random_seed)
+        np.random.shuffle(rtl_indices)
+        split_rtl_indices = [list(arr) for arr in np.split(rtl_indices, num_lattices)]
+        swapped_rtl_indices = self._ensure_unique_sublattices(split_rtl_indices)
+        monotonicity_groupings = {}
+        for lattice_indices in swapped_rtl_indices:
+            monotonic_count = sum(
+                1
+                for idx in lattice_indices
+                if self.monotonicities[idx] == Monotonicity.INCREASING
+            )
+            if monotonic_count not in monotonicity_groupings:
+                monotonicity_groupings[monotonic_count] = [lattice_indices]
+            else:
+                monotonicity_groupings[monotonic_count].append(lattice_indices)
+        for monotonic_count, groups in monotonicity_groupings.items():
+            for i, lattice_indices in enumerate(groups):
+                sorted_indices = sorted(
+                    lattice_indices,
+                    key=lambda x: (self.monotonicities[x] is None),
+                    reverse=False,
+                )
+                groups[i] = sorted_indices
+
+        self._lattice_layers = {}
+        for monotonic_count, groups in monotonicity_groupings.items():
+            self._lattice_layers[monotonic_count] = (
+                Lattice(
+                    lattice_sizes=[self.lattice_size] * self.lattice_rank,
+                    output_min=self.output_min,
+                    output_max=self.output_max,
+                    kernel_init=self.kernel_init,
+                    monotonicities=[Monotonicity.INCREASING] * monotonic_count
+                    + [None] * (lattice_rank - monotonic_count),
+                    clip_inputs=self.clip_inputs,
+                    interpolation=self.interpolation,
+                    units=len(groups),
+                ),
+                groups,
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward method computed by using forward methods of each lattice in ensemble.
+
+        Args:
+            x: input tensor of feature values with shape `(batch_size, num_features)`.
+
+        Returns:
+            `torch.Tensor` containing the outputs of each lattice within RTL structure.
+            If `average_outputs == True`, then all outputs are averaged into a tensor of
+            shape `(batch_size, 1)`. If `average_outputs == False`, shape of tensor is
+            `(batch_size, num_lattices)`.
+        """
+        forward_results = []
+        for _, (lattice, group) in sorted(self._lattice_layers.items()):
+            if len(group) > 1:
+                lattice_input = torch.stack([x[:, idx] for idx in group], dim=-2)
+            else:
+                lattice_input = x[:, group[0]]
+            forward_results.append(lattice.forward(lattice_input))
+        result = torch.cat(forward_results, dim=-1)
+        if not self.average_outputs:
+            return result
+        result = torch.mean(result, dim=-1, keepdim=True)
+
+        return result
+
+    @torch.no_grad()
+    def output_monotonicities(self) -> list[Union[Monotonicity, None]]:
+        """Gives the monotonicities of the outputs of RTL.
+
+        Returns:
+            List of `Monotonicity` corresponding to each output of the RTL layer, in the
+            same order as outputs.
+        """
+        monotonicities = []
+        for monotonic_count, (lattice, _) in sorted(self._lattice_layers.items()):
+            if monotonic_count:
+                monotonicity = Monotonicity.INCREASING
+            else:
+                monotonicity = None
+            for _ in range(lattice.units):
+                monotonicities.append(monotonicity)
+
+        return monotonicities
+
+    @torch.no_grad()
+    def apply_constraints(self) -> None:
+        """Enforces constraints for each lattice in RTL."""
+        for lattice, _ in self._lattice_layers.values():
+            lattice.apply_constraints()
+
+    @torch.no_grad()
+    def assert_constraints(self, eps: float = 1e-6) -> list[list[str]]:
+        """Asserts that each Lattice in RTL satisfies all constraints.
+
+        Args:
+          eps: allowed constraints violations.
+
+        Returns:
+          List of lists, each with constraints violations for an individual Lattice.
+        """
+        return list(
+            lattice.assert_constraints(eps=eps)
+            for lattice, _ in self._lattice_layers.values()
+        )
+
+    @staticmethod
+    def _ensure_unique_sublattices(
+        rtl_indices: list[list[int]],
+        max_swaps: int = 10000,
+    ) -> list[list[int]]:
+        """Attempts to ensure every lattice in RTL structure contains unique features.
+
+        Args:
+            rtl_indices: list of lists where inner lists are groupings of
+                indices of input features to RTL layer.
+            max_swaps: maximum number of swaps to perform before giving up.
+
+        Returns:
+            List of lists where elements between inner lists have been swapped in
+            an attempt to remove any duplicates from every grouping.
+        """
+        swaps = 0
+        num_sublattices = len(rtl_indices)
+
+        def find_swap_candidate(current_index, element):
+            """Helper function to find the next sublattice not containing element."""
+            for offset in range(1, num_sublattices):
+                candidate_index = (current_index + offset) % num_sublattices
+                if element not in rtl_indices[candidate_index]:
+                    return candidate_index
+            return None
+
+        for i, sublattice in enumerate(rtl_indices):
+            unique_elements = set()
+            for element in sublattice:
+                if element in unique_elements:
+                    swap_with = find_swap_candidate(i, element)
+                    if swap_with is not None:
+                        for swap_element in rtl_indices[swap_with]:
+                            if swap_element not in sublattice:
+                                # Perform the swap
+                                idx_to_swap = rtl_indices[swap_with].index(swap_element)
+                                idx_duplicate = sublattice.index(element)
+                                (
+                                    rtl_indices[swap_with][idx_to_swap],
+                                    sublattice[idx_duplicate],
+                                ) = element, swap_element
+                                swaps += 1
+                                break
+                    else:
+                        logging.info(
+                            "Some lattices in RTL may use the same feature multiple "
+                            "times."
+                        )
+                        return rtl_indices
+                else:
+                    unique_elements.add(element)
+                if swaps >= max_swaps:
+                    logging.info(
+                        "Some lattices in RTL may use the same feature multiple times."
+                    )
+                    return rtl_indices
+        return rtl_indices


Test patch:
diff --git a/tests/layers/test_rtl.py b/tests/layers/test_rtl.py
new file mode 100644
index 0000000..aeb06c2
--- /dev/null
+++ b/tests/layers/test_rtl.py
@@ -0,0 +1,370 @@
+"""Tests for RTL layer."""
+from itertools import cycle
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from pytorch_lattice import Interpolation, LatticeInit, Monotonicity
+from pytorch_lattice.layers import RTL, Lattice
+
+
+@pytest.mark.parametrize(
+    "monotonicities, num_lattices, lattice_rank, output_min, output_max, kernel_init,"
+    "clip_inputs, interpolation, average_outputs",
+    [
+        (
+            [
+                None,
+                None,
+                None,
+                None,
+            ],
+            3,
+            3,
+            None,
+            2.0,
+            LatticeInit.LINEAR,
+            True,
+            Interpolation.HYPERCUBE,
+            True,
+        ),
+        (
+            [
+                Monotonicity.INCREASING,
+                Monotonicity.INCREASING,
+                None,
+                None,
+            ],
+            3,
+            3,
+            -1.0,
+            4.0,
+            LatticeInit.LINEAR,
+            False,
+            Interpolation.SIMPLEX,
+            False,
+        ),
+        (
+            [Monotonicity.INCREASING, None] * 25,
+            20,
+            5,
+            None,
+            None,
+            LatticeInit.LINEAR,
+            True,
+            Interpolation.HYPERCUBE,
+            True,
+        ),
+    ],
+)
+def test_initialization(
+    monotonicities,
+    num_lattices,
+    lattice_rank,
+    output_min,
+    output_max,
+    kernel_init,
+    clip_inputs,
+    interpolation,
+    average_outputs,
+):
+    """Tests that RTL Initialization works properly."""
+    rtl = RTL(
+        monotonicities=monotonicities,
+        num_lattices=num_lattices,
+        lattice_rank=lattice_rank,
+        output_min=output_min,
+        output_max=output_max,
+        kernel_init=kernel_init,
+        clip_inputs=clip_inputs,
+        interpolation=interpolation,
+        average_outputs=average_outputs,
+    )
+    assert rtl.monotonicities == monotonicities
+    assert rtl.num_lattices == num_lattices
+    assert rtl.lattice_rank == lattice_rank
+    assert rtl.output_min == output_min
+    assert rtl.output_max == output_max
+    assert rtl.kernel_init == kernel_init
+    assert rtl.interpolation == interpolation
+    assert rtl.average_outputs == average_outputs
+
+    total_lattices = 0
+    for monotonic_count, (lattice, group) in rtl._lattice_layers.items():
+        # test monotonic features have been sorted to front of list for lattice indices
+        for single_lattice_indices in group:
+            for i in range(lattice_rank):
+                if i < monotonic_count:
+                    assert (
+                        rtl.monotonicities[single_lattice_indices[i]]
+                        == Monotonicity.INCREASING
+                    )
+                else:
+                    assert rtl.monotonicities[single_lattice_indices[i]] is None
+
+        assert len(lattice.monotonicities) == len(lattice.lattice_sizes)
+        assert (
+            sum(1 for _ in lattice.monotonicities if _ == Monotonicity.INCREASING)
+            == monotonic_count
+        )
+        assert lattice.output_min == rtl.output_min
+        assert lattice.output_max == rtl.output_max
+        assert lattice.kernel_init == rtl.kernel_init
+        assert lattice.clip_inputs == rtl.clip_inputs
+        assert lattice.interpolation == rtl.interpolation
+
+        # test number of lattices created is correct
+        total_lattices += lattice.units
+
+    assert total_lattices == num_lattices
+
+
+@pytest.mark.parametrize(
+    "monotonicities, num_lattices, lattice_rank",
+    [
+        ([None] * 9, 2, 2),
+        ([Monotonicity.INCREASING] * 10, 3, 3),
+    ],
+)
+def test_initialization_invalid(
+    monotonicities,
+    num_lattices,
+    lattice_rank,
+):
+    """Tests that RTL Initialization raises error when RTL is too small."""
+    with pytest.raises(ValueError) as exc_info:
+        RTL(
+            monotonicities=monotonicities,
+            num_lattices=num_lattices,
+            lattice_rank=lattice_rank,
+        )
+    assert (
+        str(exc_info.value)
+        == f"RTL with {num_lattices}x{lattice_rank}D structure cannot support "
+        + f"{len(monotonicities)} input features."
+    )
+
+
+@pytest.mark.parametrize(
+    "num_features, num_lattices, lattice_rank, units, expected_lattice_args,"
+    "expected_result, expected_avg",
+    [
+        (
+            6,
+            6,
+            3,
+            [3, 2, 1],
+            [
+                torch.tensor([[[0.0, 0.1, 0.2], [0.3, 0.4, 0.5], [0.0, 0.1, 0.2]]]),
+                torch.tensor([[[0.3, 0.4, 0.5], [0.0, 0.1, 0.2]]]),
+                torch.tensor([[0.3, 0.4, 0.5]]),
+            ],
+            torch.tensor([[0.0, 0.0, 0.0, 1.0, 1.0, 2.0]]),
+            torch.tensor([[2 / 3]]),
+        ),
+        (
+            3,
+            3,
+            2,
+            [1, 1, 1],
+            [
+                torch.tensor([[0.0, 0.1]]),
+                torch.tensor([[0.2, 0.0]]),
+                torch.tensor([[0.1, 0.2]]),
+            ],
+            torch.tensor([[0.0, 1.0, 2.0]]),
+            torch.tensor([[1.0]]),
+        ),
+        (
+            6,
+            7,
+            5,
+            [2, 3, 2],
+            [
+                torch.tensor([[[0.0, 0.1, 0.2, 0.3, 0.4], [0.5, 0.0, 0.1, 0.2, 0.3]]]),
+                torch.tensor(
+                    [
+                        [
+                            [0.4, 0.5, 0.0, 0.1, 0.2],
+                            [0.3, 0.4, 0.5, 0.0, 0.1],
+                            [0.2, 0.3, 0.4, 0.5, 0.0],
+                        ]
+                    ]
+                ),
+                torch.tensor([[[0.1, 0.2, 0.3, 0.4, 0.5], [0.0, 0.1, 0.2, 0.3, 0.4]]]),
+            ],
+            torch.tensor([[0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 2.0]]),
+            torch.tensor([[1.0]]),
+        ),
+    ],
+)
+def test_forward(
+    num_features,
+    num_lattices,
+    lattice_rank,
+    units,
+    expected_lattice_args,
+    expected_result,
+    expected_avg,
+):
+    """Tests forward function of RTL Lattice."""
+    rtl = RTL(
+        monotonicities=[None, Monotonicity.INCREASING] * (num_features // 2),
+        num_lattices=num_lattices,
+        lattice_rank=lattice_rank,
+    )
+    # generate indices for each lattice in cyclic fashion based off units
+    groups = []
+    feature_indices = cycle(range(num_features))
+    for lattice_units in units:
+        group = [
+            [next(feature_indices) for _ in range(lattice_rank)]
+            for _ in range(lattice_units)
+        ]
+        groups.append(group)
+    lattice_indices = {i: groups[i % len(groups)] for i in range(len(units))}
+    rtl._lattice_layers = {
+        i: (Lattice(lattice_sizes=[2] * lattice_rank, units=unit), lattice_indices[i])
+        for i, unit in enumerate(units)
+    }
+
+    mock_forwards = []
+    for monotonic_count, (lattice, _) in rtl._lattice_layers.items():
+        mock_forward = Mock()
+        lattice.forward = mock_forward
+        mock_forward.return_value = torch.full(
+            (1, units[monotonic_count]),
+            float(monotonic_count),
+            dtype=torch.float32,
+        )
+        mock_forwards.append(mock_forward)
+
+    x = torch.arange(0, num_features * 0.1, 0.1).unsqueeze(0)
+    result = rtl.forward(x)
+
+    # Assert the calls and results for each mock_forward based on expected_outs
+    for i, mock_forward in enumerate(mock_forwards):
+        mock_forward.assert_called_once()
+        assert torch.allclose(
+            mock_forward.call_args[0][0], expected_lattice_args[i], atol=1e-6
+        )
+    assert torch.allclose(result, expected_result)
+    rtl.average_outputs = True
+    result = rtl.forward(x)
+    assert torch.allclose(result, expected_avg)
+
+
+@pytest.mark.parametrize(
+    "monotonic_counts, units, expected_out",
+    [
+        (
+            [0, 1, 2, 3],
+            [2, 1, 1, 1],
+            [None] * 2 + [Monotonicity.INCREASING] * 3,
+        ),
+        (
+            [0, 4, 5, 7],
+            [1, 2, 3, 4],
+            [None] + [Monotonicity.INCREASING] * 9,
+        ),
+        ([0], [3], [None] * 3),
+        ([1, 2, 3], [1, 1, 1], [Monotonicity.INCREASING] * 3),
+    ],
+)
+def test_output_monotonicities(
+    monotonic_counts,
+    units,
+    expected_out,
+):
+    """Tests output_monotonicities function."""
+    rtl = RTL(
+        monotonicities=[None, Monotonicity.INCREASING],
+        num_lattices=3,
+        lattice_rank=3,
+    )
+    rtl._lattice_layers = {
+        monotonic_count: (Lattice(lattice_sizes=[2, 2], units=units[i]), [])
+        for i, monotonic_count in enumerate(monotonic_counts)
+    }
+    assert rtl.output_monotonicities() == expected_out
+
+
+def test_apply_constraints():
+    """Tests RTL apply_constraints function."""
+    rtl = RTL(
+        monotonicities=[None, Monotonicity.INCREASING],
+        num_lattices=3,
+        lattice_rank=3,
+    )
+    mock_constrains = []
+    for lattice, _ in rtl._lattice_layers.values():
+        mock_constrain = Mock()
+        lattice.apply_constraints = mock_constrain
+        mock_constrains.append(mock_constrain)
+
+    rtl.apply_constraints()
+    for mock_constrain in mock_constrains:
+        mock_constrain.assert_called_once()
+
+
+def test_assert_constraints():
+    """Tests RTL assert_constraints function."""
+    rtl = RTL(
+        monotonicities=[None, Monotonicity.INCREASING],
+        num_lattices=3,
+        lattice_rank=3,
+    )
+    mock_asserts = []
+    for lattice, _ in rtl._lattice_layers.values():
+        mock_assert = Mock()
+        lattice.assert_constraints = mock_assert
+        mock_assert.return_value = "violation"
+        mock_asserts.append(mock_assert)
+
+    violations = rtl.assert_constraints()
+    for mock_assert in mock_asserts:
+        mock_assert.assert_called_once()
+
+    assert violations == ["violation"] * len(rtl._lattice_layers)
+
+
+@pytest.mark.parametrize(
+    "rtl_indices",
+    [
+        [[1, 1], [1, 2], [1, 3], [1, 4], [1, 5], [6, 6]],
+        [[1, 1, 1], [2, 2, 2], [3, 3, 3]],
+        [[1, 1, 1], [2, 2, 2], [1, 2, 3], [3, 3, 3]],
+        [
+            [1, 1, 2],
+            [2, 3, 4],
+            [1, 5, 5],
+            [4, 6, 7],
+            [1, 3, 4],
+            [2, 3, 3],
+            [4, 5, 6],
+            [6, 6, 6],
+        ],
+    ],
+)
+def test_ensure_unique_sublattices_possible(rtl_indices):
+    """Tests _ensure_unique_sublattices removes duplicates from groups when possible."""
+    swapped_indices = RTL._ensure_unique_sublattices(rtl_indices)
+    for group in swapped_indices:
+        assert len(set(group)) == len(group)
+
+
+@pytest.mark.parametrize(
+    "rtl_indices, max_swaps",
+    [
+        ([[1, 1], [1, 2], [1, 3]], 100),
+        ([[1, 1], [2, 2], [3, 3], [4, 4]], 2),
+    ],
+)
+def test_ensure_unique_sublattices_impossible(rtl_indices, max_swaps):
+    """Tests _ensure_unique_sublattices logs when it can't remove duplicates."""
+    with patch("logging.info") as mock_logging_info:
+        RTL._ensure_unique_sublattices(rtl_indices, max_swaps)
+        mock_logging_info.assert_called_with(
+            "Some lattices in RTL may use the same feature multiple times."
+        )