Bump version: 0.5.0 → 0.5.1

* Added PCA initializer and component for OmegaMatrix or LinearMappings

* [QA] Add default configuration for pre commit hooks

* [QA] Add more pre commit checks

* [QA] Add more pre commit checks

* test(githooks): Add gitlint to check commit messages on commit

* docs(githooks): Add usage guide for pre-commit  to readme

* fix(githooks): mypy only checks source now

reverts changes on docs conf.py

* docs(githooks): Fix typo

Co-authored-by: staps@hs-mittweida.de <staps@hs-mittweida.de>
Co-authored-by: Alexander Engelsberger <alexanderengelsberger@gmail.com>

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.5.0
+current_version = 0.5.1
 commit = True
 tag = True
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)

.github/ISSUE_TEMPLATE/bug_report.md

@@ -23,9 +23,9 @@ A clear and concise description of what you expected to happen.
 If applicable, add screenshots to help explain your problem.
 
 **Desktop (please complete the following information):**
- - OS: [e.g. Ubuntu 20.10]
- - Prototorch Version: [e.g. v0.4.0]
- - Python Version: [e.g. 3.9.5]
+- OS: [e.g. Ubuntu 20.10]
+- Prototorch Version: [e.g. v0.4.0]
+- Python Version: [e.g. 3.9.5]
 
 **Additional context**
-Add any other context about the problem here.
+Add any other context about the problem here.

.gitignore

@@ -155,4 +155,4 @@ scratch*
 .vscode/
 
 reports
-artifacts
+artifacts

.pre-commit-config.yaml

@@ -0,0 +1,54 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.0.1
+    hooks:
+    -   id: trailing-whitespace
+    -   id: end-of-file-fixer
+    -   id: check-yaml
+    -   id: check-added-large-files
+    -   id: check-ast
+    -   id: check-case-conflict
+
+
+- repo: https://github.com/myint/autoflake
+  rev: v1.4
+  hooks:
+  -   id: autoflake
+
+- repo: http://github.com/PyCQA/isort
+  rev: 5.8.0
+  hooks:
+  -   id: isort
+
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.902'
+    hooks:
+    -   id: mypy
+        files: prototorch
+        additional_dependencies: [types-pkg_resources]
+
+-   repo: https://github.com/pre-commit/mirrors-yapf
+    rev: 'v0.31.0'  # Use the sha / tag you want to point at
+    hooks:
+    -   id: yapf
+
+-   repo: https://github.com/pre-commit/pygrep-hooks
+    rev: v1.9.0  # Use the ref you want to point at
+    hooks:
+    -   id: python-use-type-annotations
+    -   id: python-no-log-warn
+    -   id: python-check-blanket-noqa
+
+
+-   repo: https://github.com/asottile/pyupgrade
+    rev: v2.19.4
+    hooks:
+    -   id: pyupgrade
+
+-   repo: https://github.com/jorisroovers/gitlint
+    rev: "v0.15.1"
+    hooks:
+    -   id: gitlint
+        args: [--contrib=CT1, --ignore=B6, --msg-filename]

README.md

@@ -48,6 +48,17 @@ pip install -e .[all]
 The documentation is available at <https://www.prototorch.ml/en/latest/>. Should
 that link not work try <https://prototorch.readthedocs.io/en/latest/>.
 
+## Contribution
+
+This repository contains definition for [git hooks](https://githooks.com).
+[Pre-commit](https://pre-commit.com) gets installed as development dependency with prototorch.
+Please install the hooks by running
+```bash
+pre-commit install
+pre-commit install --hook-type commit-msg
+```
+before creating the first commit.
+
 ## Bibtex
 
 If you would like to cite the package, please use this:

RELEASE.md

@@ -2,17 +2,15 @@
 
 ## Release 0.5.0
 
-- Breaking: Removed deprecated `prototorch.modules.Prototypes1D`
-  - Use `prototorch.components.LabeledComponents` instead
+- Breaking: Removed deprecated `prototorch.modules.Prototypes1D`.
+  - Use `prototorch.components.LabeledComponents` instead.
 
 ## Release 0.2.0
 
-### Includes
 - Fixes in example scripts.
 
 ## Release 0.1.1-dev0
 
-### Includes
 - Minor bugfixes.
 - 100% line coverage.
 

docs/source/conf.py

@@ -23,7 +23,7 @@ author = "Jensun Ravichandran"
 
 # The full version, including alpha/beta/rc tags
 #
-release = "0.5.0"
+release = "0.5.1"
 
 # -- General configuration ---------------------------------------------------
 

examples/glvq_iris.py

@@ -1,120 +0,0 @@
-"""ProtoTorch GLVQ example using 2D Iris data."""
-
-import numpy as np
-import torch
-from matplotlib import pyplot as plt
-from prototorch.components import LabeledComponents, StratifiedMeanInitializer
-from prototorch.functions.competitions import wtac
-from prototorch.functions.distances import euclidean_distance
-from prototorch.modules.losses import GLVQLoss
-from sklearn.datasets import load_iris
-from sklearn.preprocessing import StandardScaler
-from torchinfo import summary
-
-# Prepare and preprocess the data
-scaler = StandardScaler()
-x_train, y_train = load_iris(return_X_y=True)
-x_train = x_train[:, [0, 2]]
-scaler.fit(x_train)
-x_train = scaler.transform(x_train)
-
-
-# Define the GLVQ model
-class Model(torch.nn.Module):
-    def __init__(self):
-        """GLVQ model for training on 2D Iris data."""
-        super().__init__()
-        prototype_initializer = StratifiedMeanInitializer([x_train, y_train])
-        prototype_distribution = {"num_classes": 3, "prototypes_per_class": 3}
-        self.proto_layer = LabeledComponents(
-            prototype_distribution,
-            prototype_initializer,
-        )
-
-    def forward(self, x):
-        prototypes, prototype_labels = self.proto_layer()
-        distances = euclidean_distance(x, prototypes)
-        return distances, prototype_labels
-
-
-# Build the GLVQ model
-model = Model()
-
-# Print summary using torchinfo (might be buggy/incorrect)
-print(summary(model))
-
-# Optimize using SGD optimizer from `torch.optim`
-optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
-criterion = GLVQLoss(squashing="sigmoid_beta", beta=10)
-
-x_in = torch.Tensor(x_train)
-y_in = torch.Tensor(y_train)
-
-# Training loop
-TITLE = "Prototype Visualization"
-fig = plt.figure(TITLE)
-for epoch in range(70):
-    # Compute loss
-    distances, prototype_labels = model(x_in)
-    loss = criterion([distances, prototype_labels], y_in)
-
-    # Compute Accuracy
-    with torch.no_grad():
-        predictions = wtac(distances, prototype_labels)
-        correct = predictions.eq(y_in.view_as(predictions)).sum().item()
-    acc = 100.0 * correct / len(x_train)
-
-    print(
-        f"Epoch: {epoch + 1:03d} Loss: {loss.item():05.02f} Acc: {acc:05.02f}%"
-    )
-
-    # Optimizer step
-    optimizer.zero_grad()
-    loss.backward()
-    optimizer.step()
-
-    # Get the prototypes form the model
-    prototypes = model.proto_layer.components.numpy()
-    if np.isnan(np.sum(prototypes)):
-        print("Stopping training because of `nan` in prototypes.")
-        break
-
-    # Visualize the data and the prototypes
-    ax = fig.gca()
-    ax.cla()
-    ax.set_title(TITLE)
-    ax.set_xlabel("Data dimension 1")
-    ax.set_ylabel("Data dimension 2")
-    cmap = "viridis"
-    ax.scatter(x_train[:, 0], x_train[:, 1], c=y_train, edgecolor="k")
-    ax.scatter(
-        prototypes[:, 0],
-        prototypes[:, 1],
-        c=prototype_labels,
-        cmap=cmap,
-        edgecolor="k",
-        marker="D",
-        s=50,
-    )
-
-    # Paint decision regions
-    x = np.vstack((x_train, prototypes))
-    x_min, x_max = x[:, 0].min() - 1, x[:, 0].max() + 1
-    y_min, y_max = x[:, 1].min() - 1, x[:, 1].max() + 1
-    xx, yy = np.meshgrid(np.arange(x_min, x_max, 1 / 50),
-                         np.arange(y_min, y_max, 1 / 50))
-    mesh_input = np.c_[xx.ravel(), yy.ravel()]
-
-    torch_input = torch.Tensor(mesh_input)
-    d = model(torch_input)[0]
-    w_indices = torch.argmin(d, dim=1)
-    y_pred = torch.index_select(prototype_labels, 0, w_indices)
-    y_pred = y_pred.reshape(xx.shape)
-
-    # Plot voronoi regions
-    ax.contourf(xx, yy, y_pred, cmap=cmap, alpha=0.35)
-
-    ax.set_xlim(left=x_min + 0, right=x_max - 0)
-    ax.set_ylim(bottom=y_min + 0, top=y_max - 0)
-
-    plt.pause(0.1)

examples/gmlvq_tecator.py

@@ -1,103 +0,0 @@
-"""ProtoTorch "siamese" GMLVQ example using Tecator."""
-
-import matplotlib.pyplot as plt
-import torch
-from prototorch.components import LabeledComponents, StratifiedMeanInitializer
-from prototorch.datasets.tecator import Tecator
-from prototorch.functions.distances import sed
-from prototorch.modules.losses import GLVQLoss
-from prototorch.utils.colors import get_legend_handles
-from torch.utils.data import DataLoader
-
-# Prepare the dataset and dataloader
-train_data = Tecator(root="./artifacts", train=True)
-train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
-
-
-class Model(torch.nn.Module):
-    def __init__(self, **kwargs):
-        """GMLVQ model as a siamese network."""
-        super().__init__()
-        prototype_initializer = StratifiedMeanInitializer(train_loader)
-        prototype_distribution = {"num_classes": 2, "prototypes_per_class": 2}
-
-        self.proto_layer = LabeledComponents(
-            prototype_distribution,
-            prototype_initializer,
-        )
-
-        self.omega = torch.nn.Linear(in_features=100,
-                                     out_features=100,
-                                     bias=False)
-        torch.nn.init.eye_(self.omega.weight)
-
-    def forward(self, x):
-        protos = self.proto_layer.components
-        plabels = self.proto_layer.component_labels
-
-        # Process `x` and `protos` through `omega`
-        x_map = self.omega(x)
-        protos_map = self.omega(protos)
-
-        # Compute distances and output
-        dis = sed(x_map, protos_map)
-        return dis, plabels
-
-
-# Build the GLVQ model
-model = Model()
-
-# Print a summary of the model
-print(model)
-
-# Optimize using Adam optimizer from `torch.optim`
-optimizer = torch.optim.Adam(model.parameters(), lr=0.001_0)
-scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=75, gamma=0.1)
-criterion = GLVQLoss(squashing="identity", beta=10)
-
-# Training loop
-for epoch in range(150):
-    epoch_loss = 0.0  # zero-out epoch loss
-    optimizer.zero_grad()  # zero-out gradients
-    for xb, yb in train_loader:
-        # Compute loss
-        distances, plabels = model(xb)
-        loss = criterion([distances, plabels], yb)
-        epoch_loss += loss.item()
-        # Backprop
-        loss.backward()
-    # Take a gradient descent step
-    optimizer.step()
-    scheduler.step()
-
-    lr = optimizer.param_groups[0]["lr"]
-    print(f"Epoch: {epoch + 1:03d} Loss: {epoch_loss:06.02f} lr: {lr:07.06f}")
-
-# Get the omega matrix form the model
-omega = model.omega.weight.data.numpy().T
-
-# Visualize the lambda matrix
-title = "Lambda Matrix Visualization"
-fig = plt.figure(title)
-ax = fig.gca()
-ax.set_title(title)
-im = ax.imshow(omega.dot(omega.T), cmap="viridis")
-plt.show()
-
-# Get the prototypes form the model
-protos = model.proto_layer.components.numpy()
-plabels = model.proto_layer.component_labels.numpy()
-
-# Visualize the prototypes
-title = "Tecator Prototypes"
-fig = plt.figure(title)
-ax = fig.gca()
-ax.set_title(title)
-ax.set_xlabel("Spectral frequencies")
-ax.set_ylabel("Absorption")
-clabels = ["Class 0 - Low fat", "Class 1 - High fat"]
-handles, colors = get_legend_handles(clabels, marker="line", zero_indexed=True)
-for x, y in zip(protos, plabels):
-    ax.plot(x, c=colors[int(y)])
-ax.legend(handles, clabels)
-plt.show()

examples/gtlvq_mnist.py

@@ -1,183 +0,0 @@
-"""
-ProtoTorch GTLVQ example using MNIST data.
-The GTLVQ is placed as an classification model on
-top of a CNN, considered as featurer extractor.
-Initialization of subpsace and prototypes in
-Siamnese fashion
-For more info about GTLVQ see:
-DOI:10.1109/IJCNN.2016.7727534
-"""
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torchvision
-from prototorch.functions.helper import calculate_prototype_accuracy
-from prototorch.modules.losses import GLVQLoss
-from prototorch.modules.models import GTLVQ
-from torchvision import transforms
-
-# Parameters and options
-num_epochs = 50
-batch_size_train = 64
-batch_size_test = 1000
-learning_rate = 0.1
-momentum = 0.5
-log_interval = 10
-cuda = "cuda:0"
-random_seed = 1
-device = torch.device(cuda if torch.cuda.is_available() else "cpu")
-
-# Configures reproducability
-torch.manual_seed(random_seed)
-np.random.seed(random_seed)
-
-# Prepare and preprocess the data
-train_loader = torch.utils.data.DataLoader(
-    torchvision.datasets.MNIST(
-        "./files/",
-        train=True,
-        download=True,
-        transform=torchvision.transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize((0.1307, ), (0.3081, ))
-        ]),
-    ),
-    batch_size=batch_size_train,
-    shuffle=True,
-)
-
-test_loader = torch.utils.data.DataLoader(
-    torchvision.datasets.MNIST(
-        "./files/",
-        train=False,
-        download=True,
-        transform=torchvision.transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize((0.1307, ), (0.3081, ))
-        ]),
-    ),
-    batch_size=batch_size_test,
-    shuffle=True,
-)
-
-
-# Define the GLVQ model plus appropriate feature extractor
-class CNNGTLVQ(torch.nn.Module):
-    def __init__(
-        self,
-        num_classes,
-        subspace_data,
-        prototype_data,
-        tangent_projection_type="local",
-        prototypes_per_class=2,
-        bottleneck_dim=128,
-    ):
-        super(CNNGTLVQ, self).__init__()
-
-        # Feature Extractor - Simple CNN
-        self.fe = nn.Sequential(
-            nn.Conv2d(1, 32, 3, 1),
-            nn.ReLU(),
-            nn.Conv2d(32, 64, 3, 1),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Dropout(0.25),
-            nn.Flatten(),
-            nn.Linear(9216, bottleneck_dim),
-            nn.Dropout(0.5),
-            nn.LeakyReLU(),
-            nn.LayerNorm(bottleneck_dim),
-        )
-
-        # Forward pass of subspace and prototype initialization data through feature extractor
-        subspace_data = self.fe(subspace_data)
-        prototype_data[0] = self.fe(prototype_data[0])
-
-        # Initialization of GTLVQ
-        self.gtlvq = GTLVQ(
-            num_classes,
-            subspace_data,
-            prototype_data,
-            tangent_projection_type=tangent_projection_type,
-            feature_dim=bottleneck_dim,
-            prototypes_per_class=prototypes_per_class,
-        )
-
-    def forward(self, x):
-        # Feature Extraction
-        x = self.fe(x)
-
-        # GTLVQ Forward pass
-        dis = self.gtlvq(x)
-        return dis
-
-
-# Get init data
-subspace_data = torch.cat(
-    [next(iter(train_loader))[0],
-     next(iter(test_loader))[0]])
-prototype_data = next(iter(train_loader))
-
-# Build the CNN GTLVQ  model
-model = CNNGTLVQ(
-    10,
-    subspace_data,
-    prototype_data,
-    tangent_projection_type="local",
-    bottleneck_dim=128,
-).to(device)
-
-# Optimize using SGD optimizer from `torch.optim`
-optimizer = torch.optim.Adam(
-    [{
-        "params": model.fe.parameters()
-    }, {
-        "params": model.gtlvq.parameters()
-    }],
-    lr=learning_rate,
-)
-criterion = GLVQLoss(squashing="sigmoid_beta", beta=10)
-
-# Training loop
-for epoch in range(num_epochs):
-    for batch_idx, (x_train, y_train) in enumerate(train_loader):
-        model.train()
-        x_train, y_train = x_train.to(device), y_train.to(device)
-        optimizer.zero_grad()
-
-        distances = model(x_train)
-        plabels = model.gtlvq.cls.component_labels.to(device)
-
-        # Compute loss.
-        loss = criterion([distances, plabels], y_train)
-        loss.backward()
-        optimizer.step()
-
-        # GTLVQ uses projected SGD, which means to orthogonalize the subspaces after every gradient update.
-        model.gtlvq.orthogonalize_subspace()
-
-        if batch_idx % log_interval == 0:
-            acc = calculate_prototype_accuracy(distances, y_train, plabels)
-            print(
-                f"Epoch: {epoch + 1:02d}/{num_epochs:02d} Epoch Progress: {100. * batch_idx / len(train_loader):02.02f} % Loss: {loss.item():02.02f} \
-              Train Acc: {acc.item():02.02f}")
-
-    # Test
-    with torch.no_grad():
-        model.eval()
-        correct = 0
-        total = 0
-        for x_test, y_test in test_loader:
-            x_test, y_test = x_test.to(device), y_test.to(device)
-            test_distances = model(torch.tensor(x_test))
-            test_plabels = model.gtlvq.cls.prototype_labels.to(device)
-            i = torch.argmin(test_distances, 1)
-            correct += torch.sum(y_test == test_plabels[i])
-            total += y_test.size(0)
-        print("Accuracy of the network on the test images: %d %%" %
-              (torch.true_divide(correct, total) * 100))
-
-# Save the model
-PATH = "./glvq_mnist_model.pth"
-torch.save(model.state_dict(), PATH)

examples/lgmlvq_iris.py

@@ -1,108 +0,0 @@
-"""ProtoTorch LGMLVQ example using 2D Iris data."""
-
-import numpy as np
-import torch
-from matplotlib import pyplot as plt
-from prototorch.components import LabeledComponents, StratifiedMeanInitializer
-from prototorch.functions.competitions import stratified_min
-from prototorch.functions.distances import lomega_distance
-from prototorch.modules.losses import GLVQLoss
-from sklearn.datasets import load_iris
-from sklearn.metrics import accuracy_score
-
-# Prepare training data
-x_train, y_train = load_iris(True)
-x_train = x_train[:, [0, 2]]
-
-
-# Define the model
-class Model(torch.nn.Module):
-    def __init__(self):
-        """Local-GMLVQ model."""
-        super().__init__()
-
-        prototype_initializer = StratifiedMeanInitializer([x_train, y_train])
-        prototype_distribution = [1, 2, 2]
-        self.proto_layer = LabeledComponents(
-            prototype_distribution,
-            prototype_initializer,
-        )
-
-        omegas = torch.eye(2, 2).repeat(5, 1, 1)
-        self.omegas = torch.nn.Parameter(omegas)
-
-    def forward(self, x):
-        protos, plabels = self.proto_layer()
-        omegas = self.omegas
-        dis = lomega_distance(x, protos, omegas)
-        return dis, plabels
-
-
-# Build the model
-model = Model()
-
-# Optimize using Adam optimizer from `torch.optim`
-optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
-criterion = GLVQLoss(squashing="sigmoid_beta", beta=10)
-
-x_in = torch.Tensor(x_train)
-y_in = torch.Tensor(y_train)
-
-# Training loop
-title = "Prototype Visualization"
-fig = plt.figure(title)
-for epoch in range(100):
-    # Compute loss
-    dis, plabels = model(x_in)
-    loss = criterion([dis, plabels], y_in)
-    y_pred = np.argmin(stratified_min(dis, plabels).detach().numpy(), axis=1)
-    acc = accuracy_score(y_train, y_pred)
-    log_string = f"Epoch: {epoch + 1:03d} Loss: {loss.item():05.02f} "
-    log_string += f"Acc: {acc * 100:05.02f}%"
-    print(log_string)
-
-    # Take a gradient descent step
-    optimizer.zero_grad()
-    loss.backward()
-    optimizer.step()
-
-    # Get the prototypes form the model
-    protos = model.proto_layer.components.numpy()
-
-    # Visualize the data and the prototypes
-    ax = fig.gca()
-    ax.cla()
-    ax.set_title(title)
-    ax.set_xlabel("Data dimension 1")
-    ax.set_ylabel("Data dimension 2")
-    cmap = "viridis"
-    ax.scatter(x_train[:, 0], x_train[:, 1], c=y_train, edgecolor="k")
-    ax.scatter(
-        protos[:, 0],
-        protos[:, 1],
-        c=plabels,
-        cmap=cmap,
-        edgecolor="k",
-        marker="D",
-        s=50,
-    )
-
-    # Paint decision regions
-    x = np.vstack((x_train, protos))
-    x_min, x_max = x[:, 0].min() - 1, x[:, 0].max() + 1
-    y_min, y_max = x[:, 1].min() - 1, x[:, 1].max() + 1
-    xx, yy = np.meshgrid(np.arange(x_min, x_max, 1 / 50),
-                         np.arange(y_min, y_max, 1 / 50))
-    mesh_input = np.c_[xx.ravel(), yy.ravel()]
-
-    d, plabels = model(torch.Tensor(mesh_input))
-    y_pred = np.argmin(stratified_min(d, plabels).detach().numpy(), axis=1)
-    y_pred = y_pred.reshape(xx.shape)
-
-    # Plot voronoi regions
-    ax.contourf(xx, yy, y_pred, cmap=cmap, alpha=0.35)
-
-    ax.set_xlim(left=x_min + 0, right=x_max - 0)
-    ax.set_ylim(bottom=y_min + 0, top=y_max - 0)
-
-    plt.pause(0.1)

prototorch/__init__.py

@@ -1,6 +1,7 @@
 """ProtoTorch package."""
 
 import pkgutil
+from typing import List
 
 import pkg_resources
 
@@ -8,7 +9,7 @@ from . import components, datasets, functions, modules, utils
 from .datasets import *
 
 # Core Setup
-__version__ = "0.5.0"
+__version__ = "0.5.1"
 
 __all_core__ = [
     "datasets",
@@ -19,7 +20,7 @@ __all_core__ = [
 ]
 
 # Plugin Loader
-__path__ = pkgutil.extend_path(__path__, __name__)
+__path__: List[str] = pkgutil.extend_path(__path__, __name__)
 
 
 def discover_plugins():

prototorch/components/components.py

@@ -3,13 +3,84 @@
 import warnings
 
 import torch
+from torch.nn.parameter import Parameter
+
 from prototorch.components.initializers import (ClassAwareInitializer,
                                                 ComponentsInitializer,
-                                                CustomLabelsInitializer,
                                                 EqualLabelsInitializer,
                                                 UnequalLabelsInitializer,
                                                 ZeroReasoningsInitializer)
-from torch.nn.parameter import Parameter
+
+from .initializers import parse_data_arg
+
+
+def get_labels_object(distribution):
+    if isinstance(distribution, dict):
+        if "num_classes" in distribution.keys():
+            labels = EqualLabelsInitializer(
+                distribution["num_classes"],
+                distribution["prototypes_per_class"])
+        else:
+            clabels = list(distribution.keys())
+            dist = list(distribution.values())
+            labels = UnequalLabelsInitializer(dist, clabels)
+    elif isinstance(distribution, tuple):
+        num_classes, prototypes_per_class = distribution
+        labels = EqualLabelsInitializer(num_classes, prototypes_per_class)
+    elif isinstance(distribution, list):
+        labels = UnequalLabelsInitializer(distribution)
+    else:
+        msg = f"`distribution` not understood." \
+            f"You have provided: {distribution=}."
+        raise ValueError(msg)
+    return labels
+
+
+def _precheck_initializer(initializer):
+    if not isinstance(initializer, ComponentsInitializer):
+        emsg = f"`initializer` has to be some subtype of " \
+            f"{ComponentsInitializer}. " \
+            f"You have provided: {initializer=} instead."
+        raise TypeError(emsg)
+
+
+class LinearMapping(torch.nn.Module):
+    """LinearMapping is a learnable Mapping Matrix."""
+    def __init__(self,
+                 mapping_shape=None,
+                 initializer=None,
+                 *,
+                 initialized_linearmapping=None):
+        super().__init__()
+
+        # Ignore all initialization settings if initialized_components is given.
+        if initialized_linearmapping is not None:
+            self._register_mapping(initialized_linearmapping)
+            if num_components is not None or initializer is not None:
+                wmsg = "Arguments ignored while initializing Components"
+                warnings.warn(wmsg)
+        else:
+            self._initialize_mapping(mapping_shape, initializer)
+
+    @property
+    def mapping_shape(self):
+        return self._omega.shape
+
+    def _register_mapping(self, components):
+        self.register_parameter("_omega", Parameter(components))
+
+    def _initialize_mapping(self, mapping_shape, initializer):
+        _precheck_initializer(initializer)
+        _mapping = initializer.generate(mapping_shape)
+        self._register_mapping(_mapping)
+
+    @property
+    def mapping(self):
+        """Tensor containing the component tensors."""
+        return self._omega.detach()
+
+    def forward(self):
+        return self._omega
 
 
 class Components(torch.nn.Module):
@@ -21,29 +92,46 @@ class Components(torch.nn.Module):
                  initialized_components=None):
         super().__init__()
 
-        self.num_components = num_components
-
         # Ignore all initialization settings if initialized_components is given.
         if initialized_components is not None:
-            self.register_parameter("_components",
-                                    Parameter(initialized_components))
+            self._register_components(initialized_components)
             if num_components is not None or initializer is not None:
                 wmsg = "Arguments ignored while initializing Components"
                 warnings.warn(wmsg)
         else:
-            self._initialize_components(initializer)
+            self._initialize_components(num_components, initializer)
 
-    def _precheck_initializer(self, initializer):
-        if not isinstance(initializer, ComponentsInitializer):
-            emsg = f"`initializer` has to be some subtype of " \
-                f"{ComponentsInitializer}. " \
-                f"You have provided: {initializer=} instead."
-            raise TypeError(emsg)
+    @property
+    def num_components(self):
+        return len(self._components)
 
-    def _initialize_components(self, initializer):
-        self._precheck_initializer(initializer)
-        _components = initializer.generate(self.num_components)
-        self.register_parameter("_components", Parameter(_components))
+    def _register_components(self, components):
+        self.register_parameter("_components", Parameter(components))
+
+    def _initialize_components(self, num_components, initializer):
+        _precheck_initializer(initializer)
+        _components = initializer.generate(num_components)
+        self._register_components(_components)
+
+    def add_components(self,
+                       num=1,
+                       initializer=None,
+                       *,
+                       initialized_components=None):
+        if initialized_components is not None:
+            _components = torch.cat([self._components, initialized_components])
+        else:
+            _precheck_initializer(initializer)
+            _new = initializer.generate(num)
+            _components = torch.cat([self._components, _new])
+        self._register_components(_components)
+
+    def remove_components(self, indices=None):
+        mask = torch.ones(self.num_components, dtype=torch.bool)
+        mask[indices] = False
+        _components = self._components[mask]
+        self._register_components(_components)
+        return mask
 
     @property
     def components(self):
@@ -54,7 +142,7 @@ class Components(torch.nn.Module):
         return self._components
 
     def extra_repr(self):
-        return f"components.shape: {tuple(self._components.shape)}"
+        return f"(components): (shape: {tuple(self._components.shape)})"
 
 
 class LabeledComponents(Components):
@@ -68,39 +156,60 @@ class LabeledComponents(Components):
                  *,
                  initialized_components=None):
         if initialized_components is not None:
-            components, component_labels = initialized_components
+            components, component_labels = parse_data_arg(
+                initialized_components)
             super().__init__(initialized_components=components)
-            self._labels = component_labels
+            self._register_labels(component_labels)
         else:
-            _labels = self._initialize_labels(distribution)
+            labels = get_labels_object(distribution)
+            self.initial_distribution = labels.distribution
+            _labels = labels.generate()
             super().__init__(len(_labels), initializer=initializer)
-            self.register_buffer("_labels", _labels)
+            self._register_labels(_labels)
 
-    def _initialize_components(self, initializer):
+    def _register_labels(self, labels):
+        self.register_buffer("_labels", labels)
+
+    @property
+    def distribution(self):
+        clabels, counts = torch.unique(self._labels,
+                                       sorted=True,
+                                       return_counts=True)
+        return dict(zip(clabels.tolist(), counts.tolist()))
+
+    def _initialize_components(self, num_components, initializer):
         if isinstance(initializer, ClassAwareInitializer):
-            self._precheck_initializer(initializer)
-            _components = initializer.generate(self.num_components,
-                                               self.distribution)
-            self.register_parameter("_components", Parameter(_components))
+            _precheck_initializer(initializer)
+            _components = initializer.generate(num_components,
+                                               self.initial_distribution)
+            self._register_components(_components)
         else:
-            super()._initialize_components(initializer)
+            super()._initialize_components(num_components, initializer)
 
-    def _initialize_labels(self, distribution):
-        if type(distribution) == dict:
-            if "num_classes" in distribution.keys():
-                labels = EqualLabelsInitializer(
-                    distribution["num_classes"],
-                    distribution["prototypes_per_class"])
-            else:
-                labels = CustomLabelsInitializer(distribution)
-        elif type(distribution) == tuple:
-            num_classes, prototypes_per_class = distribution
-            labels = EqualLabelsInitializer(num_classes, prototypes_per_class)
-        elif type(distribution) == list:
-            labels = UnequalLabelsInitializer(distribution)
+    def add_components(self, distribution, initializer):
+        _precheck_initializer(initializer)
 
-        self.distribution = labels.distribution
-        return labels.generate()
+        # Labels
+        labels = get_labels_object(distribution)
+        new_labels = labels.generate()
+        _labels = torch.cat([self._labels, new_labels])
+        self._register_labels(_labels)
+
+        # Components
+        if isinstance(initializer, ClassAwareInitializer):
+            _new = initializer.generate(len(new_labels), distribution)
+        else:
+            _new = initializer.generate(len(new_labels))
+        _components = torch.cat([self._components, _new])
+        self._register_components(_components)
+
+    def remove_components(self, indices=None):
+        # Components
+        mask = super().remove_components(indices)
+
+        # Labels
+        _labels = self._labels[mask]
+        self._register_labels(_labels)
 
     @property
     def component_labels(self):
@@ -112,7 +221,7 @@ class LabeledComponents(Components):
 
 
 class ReasoningComponents(Components):
-    """ReasoningComponents generate a set of components and a set of reasoning matrices.
+    r"""ReasoningComponents generate a set of components and a set of reasoning matrices.
 
     Every Component has a reasoning matrix assigned.
 
@@ -141,7 +250,7 @@ class ReasoningComponents(Components):
             super().__init__(len(self._reasonings), initializer=initializer)
 
     def _initialize_reasonings(self, reasonings):
-        if type(reasonings) == tuple:
+        if isinstance(reasonings, tuple):
             num_classes, num_components = reasonings
             reasonings = ZeroReasoningsInitializer(num_classes, num_components)
 

prototorch/components/initializers.py

@@ -13,21 +13,30 @@ def parse_data_arg(data_arg):
 
     if isinstance(data_arg, DataLoader):
         data = torch.tensor([])
-        labels = torch.tensor([])
+        targets = torch.tensor([])
         for x, y in data_arg:
             data = torch.cat([data, x])
-            labels = torch.cat([labels, y])
+            targets = torch.cat([targets, y])
     else:
-        data, labels = data_arg
+        data, targets = data_arg
         if not isinstance(data, torch.Tensor):
             wmsg = f"Converting data to {torch.Tensor}."
             warnings.warn(wmsg)
             data = torch.Tensor(data)
-        if not isinstance(labels, torch.Tensor):
-            wmsg = f"Converting labels to {torch.Tensor}."
+        if not isinstance(targets, torch.Tensor):
+            wmsg = f"Converting targets to {torch.Tensor}."
             warnings.warn(wmsg)
-            labels = torch.Tensor(labels)
-    return data, labels
+            targets = torch.Tensor(targets)
+    return data, targets
+
+
+def get_subinitializers(data, targets, clabels, subinit_type):
+    initializers = dict()
+    for clabel in clabels:
+        class_data = data[targets == clabel]
+        class_initializer = subinit_type(class_data)
+        initializers[clabel] = (class_initializer)
+    return initializers
 
 
 # Components
@@ -37,18 +46,22 @@ class ComponentsInitializer(object):
 
 
 class DimensionAwareInitializer(ComponentsInitializer):
-    def __init__(self, c_dims):
+    def __init__(self, dims):
         super().__init__()
-        if isinstance(c_dims, Iterable):
-            self.components_dims = tuple(c_dims)
+        if isinstance(dims, Iterable):
+            self.components_dims = tuple(dims)
         else:
-            self.components_dims = (c_dims, )
+            self.components_dims = (dims, )
 
 
 class OnesInitializer(DimensionAwareInitializer):
+    def __init__(self, dims, scale=1.0):
+        super().__init__(dims)
+        self.scale = scale
+
     def generate(self, length):
         gen_dims = (length, ) + self.components_dims
-        return torch.ones(gen_dims)
+        return torch.ones(gen_dims) * self.scale
 
 
 class ZerosInitializer(DimensionAwareInitializer):
@@ -58,78 +71,73 @@ class ZerosInitializer(DimensionAwareInitializer):
 
 
 class UniformInitializer(DimensionAwareInitializer):
-    def __init__(self, c_dims, min=0.0, max=1.0):
-        super().__init__(c_dims)
-
-        self.min = min
-        self.max = max
+    def __init__(self, dims, minimum=0.0, maximum=1.0, scale=1.0):
+        super().__init__(dims)
+        self.minimum = minimum
+        self.maximum = maximum
+        self.scale = scale
 
     def generate(self, length):
         gen_dims = (length, ) + self.components_dims
-        return torch.ones(gen_dims).uniform_(self.min, self.max)
+        return torch.ones(gen_dims).uniform_(self.minimum,
+                                             self.maximum) * self.scale
 
 
 class DataAwareInitializer(ComponentsInitializer):
-    def __init__(self, data):
+    def __init__(self, data, transform=torch.nn.Identity()):
         super().__init__()
         self.data = data
+        self.transform = transform
+
+    def __del__(self):
+        del self.data
 
 
 class SelectionInitializer(DataAwareInitializer):
     def generate(self, length):
         indices = torch.LongTensor(length).random_(0, len(self.data))
-        return self.data[indices]
+        return self.transform(self.data[indices])
 
 
 class MeanInitializer(DataAwareInitializer):
     def generate(self, length):
         mean = torch.mean(self.data, dim=0)
         repeat_dim = [length] + [1] * len(mean.shape)
-        return mean.repeat(repeat_dim)
+        return self.transform(mean.repeat(repeat_dim))
 
 
-class ClassAwareInitializer(ComponentsInitializer):
+class ClassAwareInitializer(DataAwareInitializer):
     def __init__(self, data, transform=torch.nn.Identity()):
-        super().__init__()
-        data, labels = parse_data_arg(data)
-        self.data = data
-        self.labels = labels
-
-        self.transform = transform
-
-        self.clabels = torch.unique(self.labels)
+        data, targets = parse_data_arg(data)
+        super().__init__(data, transform)
+        self.targets = targets
+        self.clabels = torch.unique(self.targets).int().tolist()
         self.num_classes = len(self.clabels)
 
     def _get_samples_from_initializer(self, length, dist):
         if not dist:
             per_class = length // self.num_classes
-            dist = self.num_classes * [per_class]
-        if type(dist) == dict:
-            dist = dist.values()
-        samples_list = [
-            init.generate(n) for init, n in zip(self.initializers, dist)
-        ]
-        out = torch.vstack(samples_list)
+            dist = dict(zip(self.clabels, self.num_classes * [per_class]))
+        if isinstance(dist, list):
+            dist = dict(zip(self.clabels, dist))
+        samples = [self.initializers[k].generate(n) for k, n in dist.items()]
+        out = torch.vstack(samples)
         with torch.no_grad():
             out = self.transform(out)
         return out
 
     def __del__(self):
         del self.data
-        del self.labels
+        del self.targets
 
 
 class StratifiedMeanInitializer(ClassAwareInitializer):
     def __init__(self, data, **kwargs):
         super().__init__(data, **kwargs)
+        self.initializers = get_subinitializers(self.data, self.targets,
+                                                self.clabels, MeanInitializer)
 
-        self.initializers = []
-        for clabel in self.clabels:
-            class_data = self.data[self.labels == clabel]
-            class_initializer = MeanInitializer(class_data)
-            self.initializers.append(class_initializer)
-
-    def generate(self, length, dist=[]):
+    def generate(self, length, dist):
         samples = self._get_samples_from_initializer(length, dist)
         return samples
 
@@ -138,12 +146,9 @@ class StratifiedSelectionInitializer(ClassAwareInitializer):
     def __init__(self, data, noise=None, **kwargs):
         super().__init__(data, **kwargs)
         self.noise = noise
-
-        self.initializers = []
-        for clabel in self.clabels:
-            class_data = self.data[self.labels == clabel]
-            class_initializer = SelectionInitializer(class_data)
-            self.initializers.append(class_initializer)
+        self.initializers = get_subinitializers(self.data, self.targets,
+                                                self.clabels,
+                                                SelectionInitializer)
 
     def add_noise_v1(self, x):
         return x + self.noise
@@ -155,13 +160,21 @@ class StratifiedSelectionInitializer(ClassAwareInitializer):
         mask = torch.bernoulli(n1) - torch.bernoulli(n2)
         return x + (self.noise * mask)
 
-    def generate(self, length, dist=[]):
+    def generate(self, length, dist):
         samples = self._get_samples_from_initializer(length, dist)
         if self.noise is not None:
             samples = self.add_noise_v1(samples)
         return samples
 
 
+# Omega matrix
+class PcaInitializer(DataAwareInitializer):
+    def generate(self, shape):
+        (input_dim, latent_dim) = shape
+        (_, eigVal, eigVec) = torch.pca_lowrank(self.data, q=latent_dim)
+        return eigVec
+
+
 # Labels
 class LabelsInitializer:
     def generate(self):
@@ -169,20 +182,18 @@ class LabelsInitializer:
 
 
 class UnequalLabelsInitializer(LabelsInitializer):
-    def __init__(self, dist):
+    def __init__(self, dist, clabels=None):
         self.dist = dist
+        self.clabels = clabels or range(len(self.dist))
 
     @property
     def distribution(self):
         return self.dist
 
-    def generate(self, clabels=None, dist=None):
-        if not clabels:
-            clabels = range(len(self.dist))
-        if not dist:
-            dist = self.dist
-        labels = list(chain(*[[i] * n for i, n in zip(clabels, dist)]))
-        return torch.LongTensor(labels)
+    def generate(self):
+        targets = list(
+            chain(*[[i] * n for i, n in zip(self.clabels, self.dist)]))
+        return torch.LongTensor(targets)
 
 
 class EqualLabelsInitializer(LabelsInitializer):
@@ -198,13 +209,6 @@ class EqualLabelsInitializer(LabelsInitializer):
         return torch.arange(self.classes).repeat(self.per_class, 1).T.flatten()
 
 
-class CustomLabelsInitializer(UnequalLabelsInitializer):
-    def generate(self):
-        clabels = list(self.dist.keys())
-        dist = list(self.dist.values())
-        return super().generate(clabels, dist)
-
-
 # Reasonings
 class ReasoningsInitializer:
     def generate(self, length):
@@ -226,3 +230,4 @@ SMI = StratifiedMeanInitializer
 Random = RandomInitializer = UniformInitializer
 Zeros = ZerosInitializer
 Ones = OnesInitializer
+PCA = PcaInitializer

prototorch/datasets/__init__.py

@@ -1,8 +1,6 @@
 """ProtoTorch datasets."""
 
 from .abstract import NumpyDataset
-from .iris import Iris
+from .sklearn import Blobs, Circles, Iris, Moons, Random
 from .spiral import Spiral
 from .tecator import Tecator
-
-__all__ = ['Iris', 'Spiral', 'Tecator']

prototorch/datasets/abstract.py

@@ -15,9 +15,9 @@ import torch
 class NumpyDataset(torch.utils.data.TensorDataset):
     """Create a PyTorch TensorDataset from NumPy arrays."""
     def __init__(self, data, targets):
-        self.data = data
-        self.targets = targets
-        tensors = [torch.Tensor(data), torch.Tensor(targets)]
+        self.data = torch.Tensor(data)
+        self.targets = torch.LongTensor(targets)
+        tensors = [self.data, self.targets]
         super().__init__(*tensors)
 
 

prototorch/datasets/iris.py

@@ -1,40 +0,0 @@
-"""Thin wrapper for the Iris classification dataset from sklearn.
-
-URL:
-    https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_iris.html
-
-"""
-
-from typing import Sequence
-
-from prototorch.datasets.abstract import NumpyDataset
-from sklearn.datasets import load_iris
-
-
-class Iris(NumpyDataset):
-    """
-    Iris Dataset by Ronald Fisher introduced in 1936.
-
-    The dataset contains four measurements from flowers of three species of iris.
-
-    .. list-table:: Iris
-        :header-rows: 1
-
-        * - dimensions
-          - classes
-          - training size
-          - validation size
-          - test size
-        * - 4
-          - 3
-          - 150
-          - 0
-          - 0
-
-    :param dims: select a subset of dimensions
-    """
-    def __init__(self, dims: Sequence[int] = None):
-        x, y = load_iris(return_X_y=True)
-        if dims:
-            x = x[:, dims]
-        super().__init__(x, y)

prototorch/datasets/sklearn.py

@@ -0,0 +1,137 @@
+"""Thin wrappers for a few scikit-learn datasets.
+
+URL:
+    https://scikit-learn.org/stable/modules/classes.html#module-sklearn.datasets
+
+"""
+
+import warnings
+from typing import Sequence, Union
+
+from sklearn.datasets import (load_iris, make_blobs, make_circles,
+                              make_classification, make_moons)
+
+from prototorch.datasets.abstract import NumpyDataset
+
+
+class Iris(NumpyDataset):
+    """Iris Dataset by Ronald Fisher introduced in 1936.
+
+    The dataset contains four measurements from flowers of three species of iris.
+
+    .. list-table:: Iris
+        :header-rows: 1
+
+        * - dimensions
+          - classes
+          - training size
+          - validation size
+          - test size
+        * - 4
+          - 3
+          - 150
+          - 0
+          - 0
+
+    :param dims: select a subset of dimensions
+    """
+    def __init__(self, dims: Sequence[int] = None):
+        x, y = load_iris(return_X_y=True)
+        if dims:
+            x = x[:, dims]
+        super().__init__(x, y)
+
+
+class Blobs(NumpyDataset):
+    """Generate isotropic Gaussian blobs for clustering.
+
+    Read more at
+    https://scikit-learn.org/stable/datasets/sample_generators.html#sample-generators.
+
+    """
+    def __init__(self,
+                 num_samples: int = 300,
+                 num_features: int = 2,
+                 seed: Union[None, int] = 0):
+        x, y = make_blobs(num_samples,
+                          num_features,
+                          centers=None,
+                          random_state=seed,
+                          shuffle=False)
+        super().__init__(x, y)
+
+
+class Random(NumpyDataset):
+    """Generate a random n-class classification problem.
+
+    Read more at
+    https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_classification.html.
+
+    Note: n_classes * n_clusters_per_class <= 2**n_informative must satisfy.
+    """
+    def __init__(self,
+                 num_samples: int = 300,
+                 num_features: int = 2,
+                 num_classes: int = 2,
+                 num_clusters: int = 2,
+                 num_informative: Union[None, int] = None,
+                 separation: float = 1.0,
+                 seed: Union[None, int] = 0):
+        if not num_informative:
+            import math
+            num_informative = math.ceil(math.log2(num_classes * num_clusters))
+            if num_features < num_informative:
+                warnings.warn("Generating more features than requested.")
+                num_features = num_informative
+        x, y = make_classification(num_samples,
+                                   num_features,
+                                   n_informative=num_informative,
+                                   n_redundant=0,
+                                   n_classes=num_classes,
+                                   n_clusters_per_class=num_clusters,
+                                   class_sep=separation,
+                                   random_state=seed,
+                                   shuffle=False)
+        super().__init__(x, y)
+
+
+class Circles(NumpyDataset):
+    """Make a large circle containing a smaller circle in 2D.
+
+    A simple toy dataset to visualize clustering and classification algorithms.
+
+    Read more at
+    https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_circles.html
+
+    """
+    def __init__(self,
+                 num_samples: int = 300,
+                 noise: float = 0.3,
+                 factor: float = 0.8,
+                 seed: Union[None, int] = 0):
+        x, y = make_circles(num_samples,
+                            noise=noise,
+                            factor=factor,
+                            random_state=seed,
+                            shuffle=False)
+        super().__init__(x, y)
+
+
+class Moons(NumpyDataset):
+    """Make two interleaving half circles.
+
+    A simple toy dataset to visualize clustering and classification algorithms.
+
+    Read more at
+    https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_moons.html
+
+    """
+    def __init__(self,
+                 num_samples: int = 300,
+                 noise: float = 0.3,
+                 seed: Union[None, int] = 0):
+        x, y = make_moons(num_samples,
+                          noise=noise,
+                          random_state=seed,
+                          shuffle=False)
+        super().__init__(x, y)

prototorch/datasets/spiral.py

@@ -6,7 +6,7 @@ import torch
 
 def make_spiral(num_samples=500, noise=0.3):
     """Generates the Spiral Dataset.
-    
+
     For use in Prototorch use `prototorch.datasets.Spiral` instead.
     """
     def get_samples(n, delta_t):

prototorch/datasets/tecator.py

@@ -40,9 +40,10 @@ import os
 
 import numpy as np
 import torch
-from prototorch.datasets.abstract import ProtoDataset
 from torchvision.datasets.utils import download_file_from_google_drive
 
+from prototorch.datasets.abstract import ProtoDataset
+
 
 class Tecator(ProtoDataset):
     """
@@ -101,12 +102,12 @@ class Tecator(ProtoDataset):
             x_train, y_train = f["x_train"], f["y_train"]
             x_test, y_test = f["x_test"], f["y_test"]
         training_set = [
-            torch.tensor(x_train, dtype=torch.float32),
-            torch.tensor(y_train),
+            torch.Tensor(x_train),
+            torch.LongTensor(y_train),
         ]
         test_set = [
-            torch.tensor(x_test, dtype=torch.float32),
-            torch.tensor(y_test),
+            torch.Tensor(x_test),
+            torch.LongTensor(y_test),
         ]
 
         with open(os.path.join(self.processed_folder, self.training_file),

prototorch/functions/__init__.py

@@ -2,11 +2,4 @@
 
 from .activations import identity, sigmoid_beta, swish_beta
 from .competitions import knnc, wtac
-
-__all__ = [
-    "identity",
-    "sigmoid_beta",
-    "swish_beta",
-    "knnc",
-    "wtac",
-]
+from .pooling import *

prototorch/functions/activations.py

@@ -5,17 +5,14 @@ import torch
 ACTIVATIONS = dict()
 
 
-# def register_activation(scriptf):
-#     ACTIVATIONS[scriptf.name] = scriptf
-#     return scriptf
-def register_activation(function):
+def register_activation(fn):
     """Add the activation function to the registry."""
-    ACTIVATIONS[function.__name__] = function
-    return function
+    name = fn.__name__
+    ACTIVATIONS[name] = fn
+    return fn
 
 
 @register_activation
-# @torch.jit.script
 def identity(x, beta=0.0):
     """Identity activation function.
 
@@ -29,7 +26,6 @@ def identity(x, beta=0.0):
 
 
 @register_activation
-# @torch.jit.script
 def sigmoid_beta(x, beta=10.0):
     r"""Sigmoid activation function with scaling.
 
@@ -44,7 +40,6 @@ def sigmoid_beta(x, beta=10.0):
 
 
 @register_activation
-# @torch.jit.script
 def swish_beta(x, beta=10.0):
     r"""Swish activation function with scaling.
 

prototorch/functions/competitions.py

@@ -3,42 +3,26 @@
 import torch
 
 
-def stratified_min(distances, labels):
-    clabels = torch.unique(labels, dim=0)
-    num_classes = clabels.size()[0]
-    if distances.size()[1] == num_classes:
-        # skip if only one prototype per class
-        return distances
-    batch_size = distances.size()[0]
-    winning_distances = torch.zeros(num_classes, batch_size)
-    inf = torch.full_like(distances.T, fill_value=float("inf"))
-    # distances_to_wpluses = torch.where(matcher, distances, inf)
-    for i, cl in enumerate(clabels):
-        # cdists = distances.T[labels == cl]
-        matcher = torch.eq(labels.unsqueeze(dim=1), cl)
-        if labels.ndim == 2:
-            # if the labels are one-hot vectors
-            matcher = torch.eq(torch.sum(matcher, dim=-1), num_classes)
-        cdists = torch.where(matcher, distances.T, inf).T
-        winning_distances[i] = torch.min(cdists, dim=1,
-                                         keepdim=True).values.squeeze()
-    if labels.ndim == 2:
-        # Transpose to return with `batch_size` first and
-        # reverse the columns to fix the ordering of the classes
-        return torch.flip(winning_distances.T, dims=(1, ))
+def wtac(distances: torch.Tensor,
+         labels: torch.LongTensor) -> (torch.LongTensor):
+    """Winner-Takes-All-Competition.
 
-    return winning_distances.T  # return with `batch_size` first
+    Returns the labels corresponding to the winners.
 
-
-def wtac(distances, labels):
+    """
     winning_indices = torch.min(distances, dim=1).indices
     winning_labels = labels[winning_indices].squeeze()
     return winning_labels
 
 
-def knnc(distances, labels, k=1):
+def knnc(distances: torch.Tensor,
+         labels: torch.LongTensor,
+         k: int = 1) -> (torch.LongTensor):
+    """K-Nearest-Neighbors-Competition.
+
+    Returns the labels corresponding to the winners.
+
+    """
     winning_indices = torch.topk(-distances, k=k, dim=1).indices
-    # winning_labels = torch.mode(labels[winning_indices].squeeze(),
-    #                             dim=1).values
     winning_labels = torch.mode(labels[winning_indices], dim=1).values
     return winning_labels

prototorch/functions/distances.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import torch
+
 from prototorch.functions.helper import (_check_shapes, _int_and_mixed_shape,
                                          equal_int_shape, get_flat)
 

prototorch/functions/helper.py

@@ -89,6 +89,6 @@ def _check_shapes(signal_int_shape, proto_int_shape):
 
 def _int_and_mixed_shape(tensor):
     shape = mixed_shape(tensor)
-    int_shape = tuple([i if isinstance(i, int) else None for i in shape])
+    int_shape = tuple(i if isinstance(i, int) else None for i in shape)
 
     return shape, int_shape

prototorch/functions/losses.py

@@ -57,3 +57,38 @@ def lvq21_loss(distances, target_labels, prototype_labels):
     mu = dp - dm
 
     return mu
+
+
+# Probabilistic
+def _get_class_probabilities(probabilities, targets, prototype_labels):
+    # Create Label Mapping
+    uniques = prototype_labels.unique(sorted=True).tolist()
+    key_val = {key: val for key, val in zip(uniques, range(len(uniques)))}
+
+    target_indices = torch.LongTensor(list(map(key_val.get, targets.tolist())))
+
+    whole = probabilities.sum(dim=1)
+    correct = probabilities[torch.arange(len(probabilities)), target_indices]
+    wrong = whole - correct
+
+    return whole, correct, wrong
+
+
+def nllr_loss(probabilities, targets, prototype_labels):
+    """Compute the Negative Log-Likelihood Ratio loss."""
+    _, correct, wrong = _get_class_probabilities(probabilities, targets,
+                                                 prototype_labels)
+
+    likelihood = correct / wrong
+    log_likelihood = torch.log(likelihood)
+    return -1.0 * log_likelihood
+
+
+def rslvq_loss(probabilities, targets, prototype_labels):
+    """Compute the Robust Soft Learning Vector Quantization (RSLVQ) loss."""
+    whole, correct, _ = _get_class_probabilities(probabilities, targets,
+                                                 prototype_labels)
+
+    likelihood = correct / whole
+    log_likelihood = torch.log(likelihood)
+    return -1.0 * log_likelihood

prototorch/functions/pooling.py

@@ -0,0 +1,80 @@
+"""ProtoTorch pooling functions."""
+
+from typing import Callable
+
+import torch
+
+
+def stratify_with(values: torch.Tensor,
+                  labels: torch.LongTensor,
+                  fn: Callable,
+                  fill_value: float = 0.0) -> (torch.Tensor):
+    """Apply an arbitrary stratification strategy on the columns on `values`.
+
+    The outputs correspond to sorted labels.
+    """
+    clabels = torch.unique(labels, dim=0, sorted=True)
+    num_classes = clabels.size()[0]
+    if values.size()[1] == num_classes:
+        # skip if stratification is trivial
+        return values
+    batch_size = values.size()[0]
+    winning_values = torch.zeros(num_classes, batch_size, device=labels.device)
+    filler = torch.full_like(values.T, fill_value=fill_value)
+    for i, cl in enumerate(clabels):
+        matcher = torch.eq(labels.unsqueeze(dim=1), cl)
+        if labels.ndim == 2:
+            # if the labels are one-hot vectors
+            matcher = torch.eq(torch.sum(matcher, dim=-1), num_classes)
+        cdists = torch.where(matcher, values.T, filler).T
+        winning_values[i] = fn(cdists)
+    if labels.ndim == 2:
+        # Transpose to return with `batch_size` first and
+        # reverse the columns to fix the ordering of the classes
+        return torch.flip(winning_values.T, dims=(1, ))
+
+    return winning_values.T  # return with `batch_size` first
+
+
+def stratified_sum_pooling(values: torch.Tensor,
+                           labels: torch.LongTensor) -> (torch.Tensor):
+    """Group-wise sum."""
+    winning_values = stratify_with(
+        values,
+        labels,
+        fn=lambda x: torch.sum(x, dim=1, keepdim=True).squeeze(),
+        fill_value=0.0)
+    return winning_values
+
+
+def stratified_min_pooling(values: torch.Tensor,
+                           labels: torch.LongTensor) -> (torch.Tensor):
+    """Group-wise minimum."""
+    winning_values = stratify_with(
+        values,
+        labels,
+        fn=lambda x: torch.min(x, dim=1, keepdim=True).values.squeeze(),
+        fill_value=float("inf"))
+    return winning_values
+
+
+def stratified_max_pooling(values: torch.Tensor,
+                           labels: torch.LongTensor) -> (torch.Tensor):
+    """Group-wise maximum."""
+    winning_values = stratify_with(
+        values,
+        labels,
+        fn=lambda x: torch.max(x, dim=1, keepdim=True).values.squeeze(),
+        fill_value=-1.0 * float("inf"))
+    return winning_values
+
+
+def stratified_prod_pooling(values: torch.Tensor,
+                            labels: torch.LongTensor) -> (torch.Tensor):
+    """Group-wise maximum."""
+    winning_values = stratify_with(
+        values,
+        labels,
+        fn=lambda x: torch.prod(x, dim=1, keepdim=True).squeeze(),
+        fill_value=1.0)
+    return winning_values

prototorch/functions/transforms.py

@@ -0,0 +1,32 @@
+import torch
+
+
+# Functions
+def gaussian(distances, variance):
+    return torch.exp(-(distances * distances) / (2 * variance))
+
+
+def rank_scaled_gaussian(distances, lambd):
+    order = torch.argsort(distances, dim=1)
+    ranks = torch.argsort(order, dim=1)
+
+    return torch.exp(-torch.exp(-ranks / lambd) * distances)
+
+
+# Modules
+class GaussianPrior(torch.nn.Module):
+    def __init__(self, variance):
+        super().__init__()
+        self.variance = variance
+
+    def forward(self, distances):
+        return gaussian(distances, self.variance)
+
+
+class RankScaledGaussianPrior(torch.nn.Module):
+    def __init__(self, lambd):
+        super().__init__()
+        self.lambd = lambd
+
+    def forward(self, distances):
+        return rank_scaled_gaussian(distances, self.lambd)

prototorch/modules/__init__.py

@@ -1 +1,5 @@
-"""ProtoTorch modules."""
+"""ProtoTorch modules."""
+
+from .competitions import *
+from .pooling import *
+from .wrappers import LambdaLayer, LossLayer

prototorch/modules/competitions.py

@@ -0,0 +1,42 @@
+"""ProtoTorch Competition Modules."""
+
+import torch
+
+from prototorch.functions.competitions import knnc, wtac
+
+
+class WTAC(torch.nn.Module):
+    """Winner-Takes-All-Competition Layer.
+
+    Thin wrapper over the `wtac` function.
+
+    """
+    def forward(self, distances, labels):
+        return wtac(distances, labels)
+
+
+class LTAC(torch.nn.Module):
+    """Loser-Takes-All-Competition Layer.
+
+    Thin wrapper over the `wtac` function.
+
+    """
+    def forward(self, probs, labels):
+        return wtac(-1.0 * probs, labels)
+
+
+class KNNC(torch.nn.Module):
+    """K-Nearest-Neighbors-Competition.
+
+    Thin wrapper over the `knnc` function.
+
+    """
+    def __init__(self, k=1, **kwargs):
+        super().__init__(**kwargs)
+        self.k = k
+
+    def forward(self, distances, labels):
+        return knnc(distances, labels, k=self.k)
+
+    def extra_repr(self):
+        return f"k: {self.k}"

prototorch/modules/losses.py

@@ -21,8 +21,8 @@ class GLVQLoss(torch.nn.Module):
 
 
 class NeuralGasEnergy(torch.nn.Module):
-    def __init__(self, lm):
-        super().__init__()
+    def __init__(self, lm, **kwargs):
+        super().__init__(**kwargs)
         self.lm = lm
 
     def forward(self, d):
@@ -38,3 +38,22 @@ class NeuralGasEnergy(torch.nn.Module):
     @staticmethod
     def _nghood_fn(rankings, lm):
         return torch.exp(-rankings / lm)
+
+
+class GrowingNeuralGasEnergy(NeuralGasEnergy):
+    def __init__(self, topology_layer, **kwargs):
+        super().__init__(**kwargs)
+        self.topology_layer = topology_layer
+
+    @staticmethod
+    def _nghood_fn(rankings, topology):
+        winner = rankings[:, 0]
+
+        weights = torch.zeros_like(rankings, dtype=torch.float)
+        weights[torch.arange(rankings.shape[0]), winner] = 1.0
+
+        neighbours = topology.get_neighbours(winner)
+
+        weights[neighbours] = 0.1
+
+        return weights

prototorch/modules/models.py

@@ -1,8 +1,9 @@
 import torch
+from torch import nn
+
 from prototorch.components import LabeledComponents, StratifiedMeanInitializer
 from prototorch.functions.distances import euclidean_distance_matrix
 from prototorch.functions.normalization import orthogonalization
-from torch import nn
 
 
 class GTLVQ(nn.Module):

prototorch/modules/pooling.py

@@ -0,0 +1,32 @@
+"""ProtoTorch Pooling Modules."""
+
+import torch
+
+from prototorch.functions.pooling import (stratified_max_pooling,
+                                          stratified_min_pooling,
+                                          stratified_prod_pooling,
+                                          stratified_sum_pooling)
+
+
+class StratifiedSumPooling(torch.nn.Module):
+    """Thin wrapper over the `stratified_sum_pooling` function."""
+    def forward(self, values, labels):
+        return stratified_sum_pooling(values, labels)
+
+
+class StratifiedProdPooling(torch.nn.Module):
+    """Thin wrapper over the `stratified_prod_pooling` function."""
+    def forward(self, values, labels):
+        return stratified_prod_pooling(values, labels)
+
+
+class StratifiedMinPooling(torch.nn.Module):
+    """Thin wrapper over the `stratified_min_pooling` function."""
+    def forward(self, values, labels):
+        return stratified_min_pooling(values, labels)
+
+
+class StratifiedMaxPooling(torch.nn.Module):
+    """Thin wrapper over the `stratified_max_pooling` function."""
+    def forward(self, values, labels):
+        return stratified_max_pooling(values, labels)

prototorch/modules/wrappers.py

@@ -0,0 +1,36 @@
+"""ProtoTorch Wrappers."""
+
+import torch
+
+
+class LambdaLayer(torch.nn.Module):
+    def __init__(self, fn, name=None):
+        super().__init__()
+        self.fn = fn
+        self.name = name or fn.__name__  # lambda fns get <lambda>
+
+    def forward(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+    def extra_repr(self):
+        return self.name
+
+
+class LossLayer(torch.nn.modules.loss._Loss):
+    def __init__(self,
+                 fn,
+                 name=None,
+                 size_average=None,
+                 reduce=None,
+                 reduction: str = "mean") -> None:
+        super().__init__(size_average=size_average,
+                         reduce=reduce,
+                         reduction=reduction)
+        self.fn = fn
+        self.name = name or fn.__name__  # lambda fns get <lambda>
+
+    def forward(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+    def extra_repr(self):
+        return self.name

prototorch/utils/utils.py

@@ -1,243 +0,0 @@
-"""Utilities that provide various small functionalities."""
-
-import os
-import pickle
-import sys
-from time import time
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-
-def progressbar(title, value, end, bar_width=20):
-    percent = float(value) / end
-    arrow = "=" * int(round(percent * bar_width) - 1) + ">"
-    spaces = "." * (bar_width - len(arrow))
-    sys.stdout.write("\r{}: [{}] {}%".format(title, arrow + spaces,
-                                             int(round(percent * 100))))
-    sys.stdout.flush()
-    if percent == 1.0:
-        print()
-
-
-def prettify_string(inputs, start="", sep=" ", end="\n"):
-    outputs = start + " ".join(inputs.split()) + end
-    return outputs
-
-
-def pretty_print(inputs):
-    print(prettify_string(inputs))
-
-
-def writelog(self, *logs, logdir="./logs", logfile="run.txt"):
-    f = os.path.join(logdir, logfile)
-    with open(f, "a+") as fh:
-        for log in logs:
-            fh.write(log)
-            fh.write("\n")
-
-
-def start_tensorboard(self, logdir="./logs"):
-    cmd = f"tensorboard --logdir={logdir} --port=6006"
-    os.system(cmd)
-
-
-def make_directory(save_dir):
-    if not os.path.exists(save_dir):
-        print(f"Making directory {save_dir}.")
-        os.mkdir(save_dir)
-
-
-def make_gif(filenames, duration, output_file=None):
-    try:
-        import imageio
-    except ModuleNotFoundError as e:
-        print("Please install Protoflow with [other] extra requirements.")
-        raise (e)
-
-    images = list()
-    for filename in filenames:
-        images.append(imageio.imread(filename))
-    if not output_file:
-        output_file = f"makegif.gif"
-    if images:
-        imageio.mimwrite(output_file, images, duration=duration)
-
-
-def gif_from_dir(directory,
-                 duration,
-                 prefix="",
-                 output_file=None,
-                 verbose=True):
-    images = os.listdir(directory)
-    if verbose:
-        print(f"Making gif from {len(images)} images under {directory}.")
-    filenames = list()
-    # Sort images
-    images = sorted(
-        images,
-        key=lambda img: int(os.path.splitext(img)[0].replace(prefix, "")))
-    for image in images:
-        fname = os.path.join(directory, image)
-        filenames.append(fname)
-    if not output_file:
-        output_file = os.path.join(directory, "makegif.gif")
-    make_gif(filenames=filenames, duration=duration, output_file=output_file)
-
-
-def accuracy_score(y_true, y_pred):
-    accuracy = np.sum(y_true == y_pred)
-    normalized_acc = accuracy / float(len(y_true))
-    return normalized_acc
-
-
-def predict_and_score(clf,
-                      x_test,
-                      y_test,
-                      verbose=False,
-                      title="Test accuracy"):
-    y_pred = clf.predict(x_test)
-    accuracy = np.sum(y_test == y_pred)
-    normalized_acc = accuracy / float(len(y_test))
-    if verbose:
-        print(f"{title}: {normalized_acc * 100:06.04f}%")
-    return normalized_acc
-
-
-def remove_nan_rows(arr):
-    """Remove all rows with `nan` values in `arr`."""
-    mask = np.isnan(arr).any(axis=1)
-    return arr[~mask]
-
-
-def remove_nan_cols(arr):
-    """Remove all columns with `nan` values in `arr`."""
-    mask = np.isnan(arr).any(axis=0)
-    return arr[~mask]
-
-
-def replace_in(arr, replacement_dict, inplace=False):
-    """Replace the keys found in `arr` with the values from
-    the `replacement_dict`.
-    """
-    if inplace:
-        new_arr = arr
-    else:
-        import copy
-
-        new_arr = copy.deepcopy(arr)
-    for k, v in replacement_dict.items():
-        new_arr[arr == k] = v
-    return new_arr
-
-
-def train_test_split(data, train=0.7, val=0.15, shuffle=None, return_xy=False):
-    """Split a classification dataset in such a way so as to
-    preserve the class distribution in subsamples of the dataset.
-    """
-    if train + val > 1.0:
-        raise ValueError("Invalid split values for train and val.")
-    Y = data[:, -1]
-    labels = set(Y)
-    hist = dict()
-    for l in labels:
-        data_l = data[Y == l]
-        nl = len(data_l)
-        nl_train = int(nl * train)
-        nl_val = int(nl * val)
-        nl_test = nl - (nl_train + nl_val)
-        hist[l] = (nl_train, nl_val, nl_test)
-
-    train_data = list()
-    val_data = list()
-    test_data = list()
-    for l, (nl_train, nl_val, nl_test) in hist.items():
-        data_l = data[Y == l]
-        if shuffle:
-            np.random.shuffle(data_l)
-        train_l = data_l[:nl_train]
-        val_l = data_l[nl_train:nl_train + nl_val]
-        test_l = data_l[nl_train + nl_val:nl_train + nl_val + nl_test]
-        train_data.append(train_l)
-        val_data.append(val_l)
-        test_data.append(test_l)
-
-    def _squash(data_list):
-        data = np.array(data_list[0])
-        for item in data_list[1:]:
-            data = np.vstack((data, np.array(item)))
-        return data
-
-    train_data = _squash(train_data)
-    if val_data:
-        val_data = _squash(val_data)
-    if test_data:
-        test_data = _squash(test_data)
-    if return_xy:
-        x_train = train_data[:, :-1]
-        y_train = train_data[:, -1]
-        x_val = val_data[:, :-1]
-        y_val = val_data[:, -1]
-        x_test = test_data[:, :-1]
-        y_test = test_data[:, -1]
-        return (x_train, y_train), (x_val, y_val), (x_test, y_test)
-    return train_data, val_data, test_data
-
-
-def class_histogram(data, title="Untitled"):
-    plt.figure(title)
-    plt.clf()
-    plt.title(title)
-    dist, counts = np.unique(data[:, -1], return_counts=True)
-    plt.bar(dist, counts)
-    plt.xticks(dist)
-    print("Call matplotlib.pyplot.show() to see the plot.")
-
-
-def ntimer(n=10):
-    """Wraps a function which wraps another function to time it."""
-    if n < 1:
-        raise (Exception(f"Invalid n = {n} given."))
-
-    def timer(func):
-        """Wraps `func` with a timer and returns the wrapped `func`."""
-        def wrapper(*args, **kwargs):
-            rv = None
-            before = time()
-            for _ in range(n):
-                rv = func(*args, **kwargs)
-            after = time()
-            elapsed = after - before
-            print(f"Elapsed: {elapsed*1e3:02.02f} ms")
-            return rv
-
-        return wrapper
-
-    return timer
-
-
-def memoize(verbose=True):
-    """Wraps a function which wraps another function that memoizes."""
-    def memoizer(func):
-        """Memoize (cache) return values of `func`.
-        Wraps `func` and returns the wrapped `func` so that `func`
-        is executed when the results are not available in the cache.
-        """
-        cache = {}
-
-        def wrapper(*args, **kwargs):
-            t = (pickle.dumps(args), pickle.dumps(kwargs))
-            if t not in cache:
-                if verbose:
-                    print(f"Adding NEW rv {func.__name__}{args}{kwargs} "
-                          "to cache.")
-                cache[t] = func(*args, **kwargs)
-            else:
-                if verbose:
-                    print(f"Using OLD rv {func.__name__}{args}{kwargs} "
-                          "from cache.")
-            return cache[t]
-
-        return wrapper
-
-    return memoizer

setup.py

@@ -1,10 +1,12 @@
 """
-  _____           _     _______             _
- |  __ \         | |   |__   __|           | |
- | |__) | __ ___ | |_ ___ | | ___  _ __ ___| |__
- |  ___/ '__/ _ \| __/ _ \| |/ _ \| '__/ __| '_ \
- | |   | | | (_) | || (_) | | (_) | | | (__| | | |
- |_|   |_|  \___/ \__\___/|_|\___/|_|  \___|_| |_|
+
+ ######
+ #     # #####   ####  #####  ####  #####  ####  #####   ####  #    #
+ #     # #    # #    #   #   #    #   #   #    # #    # #    # #    #
+ ######  #    # #    #   #   #    #   #   #    # #    # #      ######
+ #       #####  #    #   #   #    #   #   #    # #####  #      #    #
+ #       #   #  #    #   #   #    #   #   #    # #   #  #    # #    #
+ #       #    #  ####    #    ####    #    ####  #    #  ####  #    #
 
 ProtoTorch Core Package
 """
@@ -18,7 +20,7 @@ with open("README.md", "r") as fh:
 
 INSTALL_REQUIRES = [
     "torch>=1.3.1",
-    "torchvision>=0.5.0",
+    "torchvision>=0.5.1",
     "numpy>=1.9.1",
     "sklearn",
 ]
@@ -26,7 +28,10 @@ DATASETS = [
     "requests",
     "tqdm",
 ]
-DEV = ["bumpversion"]
+DEV = [
+    "bumpversion",
+    "pre-commit",
+]
 DOCS = [
     "recommonmark",
     "sphinx",
@@ -43,7 +48,7 @@ ALL = DATASETS + DEV + DOCS + EXAMPLES + TESTS
 
 setup(
     name="prototorch",
-    version="0.5.0",
+    version="0.5.1",
     description="Highly extensible, GPU-supported "
     "Learning Vector Quantization (LVQ) toolbox "
     "built using PyTorch and its nn API.",

tests/test_components.py

@@ -1,8 +1,9 @@
 """ProtoTorch components test suite."""
 
-import prototorch as pt
 import torch
 
+import prototorch as pt
+
 
 def test_labcomps_zeros_init():
     protos = torch.zeros(3, 2)

tests/test_functions.py

@@ -4,8 +4,9 @@ import unittest
 
 import numpy as np
 import torch
+
 from prototorch.functions import (activations, competitions, distances,
-                                  initializers, losses)
+                                  initializers, losses, pooling)
 
 
 class TestActivations(unittest.TestCase):
@@ -104,10 +105,28 @@ class TestCompetitions(unittest.TestCase):
                                                         decimal=5)
         self.assertIsNone(mismatch)
 
+    def test_knnc_k1(self):
+        d = torch.tensor([[2.0, 3.0, 1.99, 3.01], [2.0, 3.0, 2.01, 3.0]])
+        labels = torch.tensor([0, 1, 2, 3])
+        actual = competitions.knnc(d, labels, k=1)
+        desired = torch.tensor([2, 0])
+        mismatch = np.testing.assert_array_almost_equal(actual,
+                                                        desired,
+                                                        decimal=5)
+        self.assertIsNone(mismatch)
+
+    def tearDown(self):
+        pass
+
+
+class TestPooling(unittest.TestCase):
+    def setUp(self):
+        pass
+
     def test_stratified_min(self):
         d = torch.tensor([[1.0, 0.0, 2.0, 3.0], [9.0, 8.0, 0, 1]])
         labels = torch.tensor([0, 0, 1, 2])
-        actual = competitions.stratified_min(d, labels)
+        actual = pooling.stratified_min_pooling(d, labels)
         desired = torch.tensor([[0.0, 2.0, 3.0], [8.0, 0.0, 1.0]])
         mismatch = np.testing.assert_array_almost_equal(actual,
                                                         desired,
@@ -118,28 +137,70 @@ class TestCompetitions(unittest.TestCase):
         d = torch.tensor([[1.0, 0.0, 2.0, 3.0], [9.0, 8.0, 0, 1]])
         labels = torch.tensor([0, 0, 1, 2])
         labels = torch.eye(3)[labels]
-        actual = competitions.stratified_min(d, labels)
+        actual = pooling.stratified_min_pooling(d, labels)
         desired = torch.tensor([[0.0, 2.0, 3.0], [8.0, 0.0, 1.0]])
         mismatch = np.testing.assert_array_almost_equal(actual,
                                                         desired,
                                                         decimal=5)
         self.assertIsNone(mismatch)
 
-    def test_stratified_min_simple(self):
+    def test_stratified_min_trivial(self):
         d = torch.tensor([[0.0, 2.0, 3.0], [8.0, 0, 1]])
         labels = torch.tensor([0, 1, 2])
-        actual = competitions.stratified_min(d, labels)
+        actual = pooling.stratified_min_pooling(d, labels)
         desired = torch.tensor([[0.0, 2.0, 3.0], [8.0, 0.0, 1.0]])
         mismatch = np.testing.assert_array_almost_equal(actual,
                                                         desired,
                                                         decimal=5)
         self.assertIsNone(mismatch)
 
-    def test_knnc_k1(self):
-        d = torch.tensor([[2.0, 3.0, 1.99, 3.01], [2.0, 3.0, 2.01, 3.0]])
-        labels = torch.tensor([0, 1, 2, 3])
-        actual = competitions.knnc(d, labels, k=1)
-        desired = torch.tensor([2, 0])
+    def test_stratified_max(self):
+        d = torch.tensor([[1.0, 0.0, 2.0, 3.0, 9.0], [9.0, 8.0, 0, 1, 7.0]])
+        labels = torch.tensor([0, 0, 3, 2, 0])
+        actual = pooling.stratified_max_pooling(d, labels)
+        desired = torch.tensor([[9.0, 3.0, 2.0], [9.0, 1.0, 0.0]])
+        mismatch = np.testing.assert_array_almost_equal(actual,
+                                                        desired,
+                                                        decimal=5)
+        self.assertIsNone(mismatch)
+
+    def test_stratified_max_one_hot(self):
+        d = torch.tensor([[1.0, 0.0, 2.0, 3.0, 9.0], [9.0, 8.0, 0, 1, 7.0]])
+        labels = torch.tensor([0, 0, 2, 1, 0])
+        labels = torch.nn.functional.one_hot(labels, num_classes=3)
+        actual = pooling.stratified_max_pooling(d, labels)
+        desired = torch.tensor([[9.0, 3.0, 2.0], [9.0, 1.0, 0.0]])
+        mismatch = np.testing.assert_array_almost_equal(actual,
+                                                        desired,
+                                                        decimal=5)
+        self.assertIsNone(mismatch)
+
+    def test_stratified_sum(self):
+        d = torch.tensor([[1.0, 0.0, 2.0, 3.0], [9.0, 8.0, 0, 1]])
+        labels = torch.LongTensor([0, 0, 1, 2])
+        actual = pooling.stratified_sum_pooling(d, labels)
+        desired = torch.tensor([[1.0, 2.0, 3.0], [17.0, 0.0, 1.0]])
+        mismatch = np.testing.assert_array_almost_equal(actual,
+                                                        desired,
+                                                        decimal=5)
+        self.assertIsNone(mismatch)
+
+    def test_stratified_sum_one_hot(self):
+        d = torch.tensor([[1.0, 0.0, 2.0, 3.0], [9.0, 8.0, 0, 1]])
+        labels = torch.tensor([0, 0, 1, 2])
+        labels = torch.eye(3)[labels]
+        actual = pooling.stratified_sum_pooling(d, labels)
+        desired = torch.tensor([[1.0, 2.0, 3.0], [17.0, 0.0, 1.0]])
+        mismatch = np.testing.assert_array_almost_equal(actual,
+                                                        desired,
+                                                        decimal=5)
+        self.assertIsNone(mismatch)
+
+    def test_stratified_prod(self):
+        d = torch.tensor([[1.0, 0.0, 2.0, 3.0, 9.0], [9.0, 8.0, 0, 1, 7.0]])
+        labels = torch.tensor([0, 0, 3, 2, 0])
+        actual = pooling.stratified_prod_pooling(d, labels)
+        desired = torch.tensor([[0.0, 3.0, 2.0], [504.0, 1.0, 0.0]])
         mismatch = np.testing.assert_array_almost_equal(actual,
                                                         desired,
                                                         decimal=5)