congrads 1.0.6__py3-none-any.whl → 1.1.0__py3-none-any.whl

congrads/descriptor.py

@@ -1,130 +1,166 @@
-"""
-This module defines the `Descriptor` class, which is designed to manage
-the mapping between neuron names, their corresponding layers, and additional
-properties such as constant or variable status. It provides a way to easily 
-place constraints on parts of your network, by referencing the neuron names
+"""This module defines the `Descriptor` class, which allows assigning tags to parts in the network.
+
+It is designed to manage the mapping between tags, their corresponding data dictionary keys and indices,
+and additional properties such as constant or variable status. It provides a way to easily
+place constraints on parts of your network, by referencing the tags
 instead of indices.
 
 The `Descriptor` class allows for easy constraint definitions on parts of
-your neural network. It supports registering neurons with associated layers, 
-indices, and optional attributes, such as whether the layer is constant 
-or variable.
-
-Key Methods:
-
-    - `__init__`: Initializes the `Descriptor` object with empty mappings
-      and sets for managing neurons and layers.
-    - `add`: Registers a neuron with its associated layer, index, and
-      optional constant status.
-
+your neural network. It supports registering tags with associated data dictionary keys,
+indices, and optional attributes, such as whether the data is constant or variable.
 """
 
+from torch import Tensor
+
 from .utils import validate_type
 
 
 class Descriptor:
-    """
-    A class to manage the mapping between neuron names, their corresponding
-    layers, and additional properties (such as min/max values, output,
-    and constant variables).
+    """A class to manage the mapping between tags.
 
-    This class is designed to track the relationship between neurons and
-    layers in a neural network. It allows for the assignment of properties
-    (like minimum and maximum values, and whether a layer is an output,
-    constant, or variable) to each neuron. The data is stored in
-    dictionaries and sets for efficient lookups.
+    It represents data locations in the data dictionary and holds the dictionary keys, indices,
+    and additional properties (such as min/max values, output, and constant variables).
+
+    This class is designed to manage the relationships between the assigned tags and the
+    data dictionary keys in a neural network model. It allows for the assignment of properties
+    (like minimum and maximum values, and whether data is an output, constant, or variable) to
+    each tag. The data is stored in dictionaries and sets for efficient lookups.
 
     Attributes:
-        neuron_to_layer (dict): A dictionary mapping neuron names to
-            their corresponding layer names.
-        neuron_to_index (dict): A dictionary mapping neuron names to
-            their corresponding indices in the layers.
-        constant_layers (set): A set of layer names that represent
-            constant layers.
-        variable_layers (set): A set of layer names that represent
-            variable layers.
+        constant_keys (set): A set of keys that represent constant data in the data dictionary.
+        variable_keys (set): A set of keys that represent variable data in the data dictionary.
+        affects_loss_keys (set): A set of keys that represent data affecting the loss computation.
     """
 
     def __init__(
         self,
     ):
-        """
-        Initializes the Descriptor object.
-        """
-
-        # Define dictionaries that will translate neuron
-        # names to layer and index
-        self.neuron_to_layer: dict[str, str] = {}
-        self.neuron_to_index: dict[str, int] = {}
+        """Initializes the Descriptor object."""
+        # Define dictionaries that will translate tags to keys and indices
+        self._tag_to_key: dict[str, str] = {}
+        self._tag_to_index: dict[str, int] = {}
 
-        # Define sets that will hold the layers based on which type
-        self.constant_layers: set[str] = set()
-        self.variable_layers: set[str] = set()
+        # Define sets that will hold the keys based on which type
+        self.constant_keys: set[str] = set()
+        self.variable_keys: set[str] = set()
+        self.affects_loss_keys: set[str] = set()
 
     def add(
         self,
-        layer_name: str,
-        index: int,
-        neuron_name: str,
+        key: str,
+        tag: str,
+        index: int = None,
         constant: bool = False,
+        affects_loss: bool = True,
     ):
-        """
-        Adds a neuron to the descriptor with its associated layer,
-        index, and properties.
+        """Adds a tag to the descriptor with its associated key, index, and properties.
 
-        This method registers a neuron name and associates it with a
-        layer, its index, and optional properties such as whether
-        the layer is an output or constant layer.
+        This method registers a tag name and associates it with a
+        data dictionary key, its index, and optional properties such as whether
+        the key hold output or constant data.
 
         Args:
-            layer_name (str): The name of the layer where the neuron is located.
-            index (int): The index of the neuron within the layer.
-            neuron_name (str): The name of the neuron.
-            constant (bool, optional): Whether the layer is a constant layer.
-                Defaults to False.
+            key (str): The key on which the tagged data is located in the data dictionary.
+            tag (str): The identifier of the tag.
+            index (int): The index were the data is present. Defaults to None.
+            constant (bool, optional): Whether the data is constant and is not learned. Defaults to False.
+            affects_loss (bool, optional): Whether the data affects the loss computation. Defaults to True.
 
         Raises:
             TypeError: If a provided attribute has an incompatible type.
-            ValueError: If a layer or index is already assigned for a neuron
-                or a duplicate index is used within a layer.
-
+            ValueError: If a key or index is already assigned for a tag or a duplicate index is used within a key.
         """
-
         # Type checking
-        validate_type("layer_name", layer_name, str)
-        validate_type("index", index, int)
-        validate_type("neuron_name", neuron_name, str)
+        validate_type("key", key, str)
+        validate_type("tag", tag, str)
+        validate_type("index", index, int, allow_none=True)
         validate_type("constant", constant, bool)
+        validate_type("affects_loss", affects_loss, bool)
 
         # Other validations
-        if neuron_name in self.neuron_to_layer:
+        if tag in self._tag_to_key:
             raise ValueError(
-                "There already is a layer registered for the neuron with name "
-                f"'{neuron_name}'. Please use a unique name for each neuron."
+                f"There already is a key registered for the tag '{tag}'. "
+                "Please use a unique key name for each tag."
             )
 
-        if neuron_name in self.neuron_to_index:
+        if tag in self._tag_to_index:
             raise ValueError(
-                "There already is an index registered for the neuron with name "
-                f"'{neuron_name}'. Please use a unique name for each neuron."
+                f"There already is an index registered for the tag '{tag}'. "
+                "Please use a unique name for each tag."
             )
 
-        for existing_neuron, assigned_index in self.neuron_to_index.items():
-            if (
-                assigned_index == index
-                and self.neuron_to_layer[existing_neuron] == layer_name
-            ):
+        for existing_tag, assigned_index in self._tag_to_index.items():
+            if assigned_index == index and self._tag_to_key[existing_tag] == key:
                 raise ValueError(
-                    f"The index {index} in layer {layer_name} is already "
-                    "assigned. Every neuron must be assigned a different "
+                    f"The index {index} on key {key} is already "
+                    "assigned. Every tag must be assigned a different "
                     "index that matches the network's output."
                 )
 
         # Add to dictionaries and sets
+        # TODO this now happens on key level, can this also be done on tag level?
         if constant:
-            self.constant_layers.add(layer_name)
+            self.constant_keys.add(key)
         else:
-            self.variable_layers.add(layer_name)
+            self.variable_keys.add(key)
+
+        if affects_loss:
+            self.affects_loss_keys.add(key)
+
+        self._tag_to_key[tag] = key
+        self._tag_to_index[tag] = index
+
+    def exists(self, tag: str) -> bool:
+        """Check if a tag is registered in the descriptor.
+
+        Args:
+            tag (str): The tag identifier to check.
+
+        Returns:
+            bool: True if the tag is registered, False otherwise.
+        """
+        return tag in self._tag_to_key and tag in self._tag_to_index
+
+    def location(self, tag: str) -> tuple[str, int]:
+        """Get the key and index for a given tag.
+
+        Looks up the mapping for a registered tag and returns the associated
+        dictionary key and the index.
+
+        Args:
+            tag (str): The tag identifier. Must be registered.
+
+        Returns:
+            tuple ((str, int)): A tuple containing:
+                - The key in the data dictionary which holds the data (str).
+                - The tensor index where the data is present (int).
+
+        Raises:
+            ValueError: If the tag is not registered in the descriptor.
+        """
+        key = self._tag_to_key.get(tag)
+        index = self._tag_to_index.get(tag)
+        if key is None or index is None:
+            raise ValueError(f"Tag '{tag}' is not registered in descriptor.")
+        return key, index
+
+    def select(self, tag: str, data: dict[str, Tensor]) -> Tensor:
+        """Extract prediction values for a specific tag.
 
-        self.neuron_to_layer[neuron_name] = layer_name
-        self.neuron_to_index[neuron_name] = index
+        Retrieves the key and index associated with a tag and selects
+        the corresponding slice from the given prediction tensor.
+
+        Args:
+            tag (str): The tag identifier. Must be registered.
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
+
+        Returns:
+            Tensor: A tensor slice of shape ``(batch_size, 1)`` containing
+            the predictions for the specified tag.
+
+        Raises:
+            ValueError: If the tag is not registered in the descriptor.
+        """
+        key, index = self.location(tag)
+        return data[key][:, index : index + 1]

congrads/metrics.py

@@ -1,164 +1,92 @@
-"""
-This module defines the `Metric` and `MetricManager` classes, which are 
-used to track and aggregate performance metrics during model training or 
-evaluation in machine learning. These classes support the accumulation of 
-metric values, aggregation using customizable functions (such as mean), 
-and resetting of the metrics.
-
-Classes:
-
-    - Metric: A class that tracks and aggregates a specific metric over 
-      multiple samples, allowing for accumulation, aggregation, and 
-      resetting of values.
-    - MetricManager: A class that manages and tracks multiple metrics 
-      during model training or evaluation, supporting registration, 
-      accumulation, aggregation, and resetting of metrics.
-
-Key Methods:
-
-    - `Metric.__init__`: Initializes a metric with a specified name and 
-      optional accumulator function (defaults to `nanmean`).
-    - `Metric.accumulate`: Accumulates a new value for the metric, 
-      typically a tensor of model output or performance.
-    - `Metric.aggregate`: Aggregates the accumulated values using the 
-      specified accumulator function.
-    - `Metric.reset`: Resets the accumulated values and sample count for 
-      the metric.
-    - `MetricManager.__init__`: Initializes a manager for multiple metrics.
-    - `MetricManager.register`: Registers a new metric with a name, group, 
-      and optional accumulator function.
-    - `MetricManager.accumulate`: Accumulates a new value for the specified 
-      metric.
-    - `MetricManager.aggregate`: Aggregates all metrics in a specified group.
-    - `MetricManager.reset`: Resets all registered metrics in a specified 
-      group.
-
-Each class provides functionality to efficiently track, aggregate, and reset 
-metrics during the training and evaluation phases of machine learning tasks, 
-supporting flexible aggregation strategies and group-based management of 
-metrics.
+"""Module for managing metrics during training.
+
+Provides the `Metric` and `MetricManager` classes for accumulating,
+aggregating, and resetting metrics over training batches. Supports
+grouping metrics and using custom accumulation functions.
 """
 
-from typing import Callable
+from collections.abc import Callable
 
-from torch import Tensor, cat, nanmean
+from torch import Tensor, cat, nanmean, tensor
 
 from .utils import validate_callable, validate_type
 
 
 class Metric:
-    """
-    A class that tracks and aggregates a specific metric over multiple samples.
-
-    This class allows the accumulation of values, their aggregation using a
-    specified function (e.g., mean), and the ability to reset the metrics.
-    It is typically used to track performance metrics during training or
-    evaluation processes in machine learning.
-
-    Args:
-        name (str): The name of the metric.
-        accumulator (Callable[..., Tensor], optional): A function used to
-        aggregate values (defaults to `nanmean`).
-
-    Attributes:
-        name (str): The name of the metric.
-        accumulator (Callable[..., Tensor]): The function used to aggregate
-            values.
-        values (list): A list to store accumulated values.
-        sample_count (int): The count of accumulated samples.
+    """Represents a single metric to be accumulated and aggregated.
 
+    Stores metric values over multiple batches and computes an aggregated
+    result using a specified accumulation function.
     """
 
-    def __init__(
-        self,
-        name: str,
-        accumulator: Callable[..., Tensor] = nanmean,
-    ) -> None:
-        """
-        Constructor method
-        """
+    def __init__(self, name: str, accumulator: Callable[..., Tensor] = nanmean) -> None:
+        """Initialize a Metric instance.
 
+        Args:
+            name (str): Name of the metric.
+            accumulator (Callable[..., Tensor], optional): Function to aggregate
+                accumulated values. Defaults to `torch.nanmean`.
+        """
         # Type checking
         validate_type("name", name, str)
         validate_callable("accumulator", accumulator)
 
         self.name = name
         self.accumulator = accumulator
-
-        self.values = []
+        self.values: list[Tensor] = []
         self.sample_count = 0
 
     def accumulate(self, value: Tensor) -> None:
-        """
-        Accumulates a new value for the metric.
+        """Accumulate a new value for the metric.
 
         Args:
-            value (Tensor): The new value to accumulate, typically a
-                tensor of model output or performance.
+            value (Tensor): Metric values for the current batch.
         """
-
-        self.values.append(value)
+        self.values.append(value.detach().clone())
         self.sample_count += value.size(0)
 
     def aggregate(self) -> Tensor:
-        """
-        Aggregates the accumulated values using the specified
-        accumulator function.
+        """Compute the aggregated value of the metric.
 
         Returns:
-            Tensor: The aggregated result of the accumulated values.
+            Tensor: The aggregated metric value. Returns NaN if no values
+                have been accumulated.
         """
+        if not self.values:
+            return tensor(float("nan"))
 
         combined = cat(self.values)
         return self.accumulator(combined)
 
     def reset(self) -> None:
-        """
-        Resets the accumulated values and sample count for the metric.
-        """
-
+        """Reset the accumulated values and sample count for the metric."""
         self.values = []
         self.sample_count = 0
 
 
 class MetricManager:
-    """
-    A class to manage and track multiple metrics during model
-    training or evaluation.
-
-    This class allows registering metrics, accumulating values for each metric,
-    and recording the aggregated values. It also supports the reset of metrics
-    after each epoch or training step.
+    """Manages multiple metrics and groups for training or evaluation.
 
-    Attributes:
-        metrics (dict[str, Metric]): A dictionary of registered metrics.
-        groups (dict[str, str]): A dictionary mapping metric names to groups.
+    Supports registering metrics, accumulating values by name, aggregating
+    metrics by group, and resetting metrics by group.
     """
 
     def __init__(self) -> None:
-        """
-        Constructor method
-        """
-
+        """Initialize a MetricManager instance."""
         self.metrics: dict[str, Metric] = {}
         self.groups: dict[str, str] = {}
 
     def register(
-        self,
-        name: str,
-        group: str,
-        accumulator: Callable[..., Tensor] = nanmean,
+        self, name: str, group: str = "default", accumulator: Callable[..., Tensor] = nanmean
     ) -> None:
-        """
-        Registers a new metric with the specified name and accumulator function.
+        """Register a new metric under a specified group.
 
         Args:
-            name (str): The name of the metric to register.
-            group (str): The name of the group to assign the metric to.
-            accumulator (Callable[..., Tensor], optional): The function used
-                to aggregate values for the metric (defaults to `nanmean`).
+            name (str): Name of the metric.
+            group (str, optional): Group name for the metric. Defaults to "default".
+            accumulator (Callable[..., Tensor], optional): Function to aggregate
+                accumulated values. Defaults to `torch.nanmean`.
         """
-
         # Type checking
         validate_type("name", name, str)
         validate_type("group", group, str)
@@ -168,44 +96,44 @@ class MetricManager:
         self.groups[name] = group
 
     def accumulate(self, name: str, value: Tensor) -> None:
-        """
-        Accumulates a new value for the specified metric.
+        """Accumulate a value for a specific metric by name.
 
         Args:
-            name (str): The name of the metric.
-            value (Tensor): The new value to accumulate.
+            name (str): Name of the metric.
+            value (Tensor): Metric values for the current batch.
         """
+        if name not in self.metrics:
+            raise KeyError(f"Metric '{name}' is not registered.")
 
         self.metrics[name].accumulate(value)
 
-    def aggregate(self, group: str) -> dict[str, Tensor]:
-        """
-        Aggregates all metrics in a group using the accumulators
-        specified during registration.
+    def aggregate(self, group: str = "default") -> dict[str, Tensor]:
+        """Aggregate all metrics in a specified group.
 
         Args:
-            group (str): The name of the group.
+            group (str, optional): The group of metrics to aggregate. Defaults to "default".
 
         Returns:
-            dict[str, Tensor]: A dictionary with the metric names and the
-                corresponding aggregated values of the selected group.
+            dict[str, Tensor]: Dictionary mapping metric names to their
+                aggregated values.
         """
-
         return {
             name: metric.aggregate()
             for name, metric in self.metrics.items()
             if self.groups[name] == group
         }
 
-    def reset(self, group: str) -> None:
-        """
-        Resets all registered metrics in a group.
+    def reset(self, group: str = "default") -> None:
+        """Reset all metrics in a specified group.
 
         Args:
-            group (str): The name of the group.
+            group (str, optional): The group of metrics to reset. Defaults to "default".
         """
-
         for name, metric in self.metrics.items():
             if self.groups[name] == group:
                 metric.reset()
-                metric.reset()
+
+    def reset_all(self) -> None:
+        """Reset all metrics across all groups."""
+        for metric in self.metrics.values():
+            metric.reset()

congrads/networks.py

@@ -1,52 +1,11 @@
-"""
-This module defines the `MLPNetwork` class, which constructs and
-operates a multi-layer perceptron (MLP) neural network model. The MLP
-network consists of an input layer, multiple hidden layers, and an
-output layer. It allows for configurable hyperparameters such as the
-number of input features, output features, number of hidden layers,
-and the dimensionality of the hidden layers.
-
-Classes:
-
-    - MLPNetwork: A neural network model that implements a multi-layer
-      perceptron with customizable layers and dimensionalities.
-
-Key Methods:
-
-    - `__init__`: Initializes the MLP network with specified input size,
-      output size, number of hidden layers, and hidden layer dimensionality.
-    - `forward`: Performs a forward pass through the network, returning
-      both the input and output of the model.
-    - `linear`: Creates a basic linear block consisting of a Linear layer
-      followed by a ReLU activation function.
-
-The `MLPNetwork` class constructs a fully connected neural network with 
-multiple hidden layers, providing flexibility in designing the network 
-architecture. It can be used for regression, classification, or other 
-machine learning tasks that require a feedforward neural network structure.
-"""
+"""Module defining the network architectures and components."""
 
+from torch import Tensor
 from torch.nn import Linear, Module, ReLU, Sequential
 
 
 class MLPNetwork(Module):
-    """
-    A multi-layer perceptron (MLP) neural network model consisting of
-    an input layer, multiple hidden layers, and an output layer.
-
-    This class constructs an MLP with configurable hyperparameters such as the
-    number of input features, output features, number of hidden layers, and
-    the dimensionality of hidden layers. It provides methods for both
-    building the model and performing a forward pass through the network.
-
-    Args:
-        n_inputs (int, optional): The number of input features. Defaults to 25.
-        n_outputs (int, optional): The number of output features. Defaults to 2.
-        n_hidden_layers (int, optional): The number of hidden layers.
-            Defaults to 2.
-        hidden_dim (int, optional): The dimensionality of the hidden layers.
-            Defaults to 35.
-    """
+    """A multi-layer perceptron (MLP) neural network with configurable hidden layers."""
 
     def __init__(
         self,
@@ -54,11 +13,18 @@ class MLPNetwork(Module):
         n_outputs,
         n_hidden_layers=3,
         hidden_dim=35,
+        activation=None,
     ):
-        """
-        Initializes the MLPNetwork.
-        """
+        """Initialize the MLPNetwork.
 
+        Args:
+            n_inputs (int, optional): Number of input features. Defaults to 25.
+            n_outputs (int, optional): Number of output features. Defaults to 2.
+            n_hidden_layers (int, optional): Number of hidden layers. Defaults to 3.
+            hidden_dim (int, optional): Dimensionality of hidden layers. Defaults to 35.
+            activation (nn.Module, optional): Activation function module (e.g.,
+                `ReLU()`, `Tanh()`, `LeakyReLU(0.1)`). Defaults to `ReLU()`.
+        """
         super().__init__()
 
         # Init object variables
@@ -67,48 +33,36 @@ class MLPNetwork(Module):
         self.n_hidden_layers = n_hidden_layers
         self.hidden_dim = hidden_dim
 
-        # Set up the components of our model
-        self.input = Linear(self.n_inputs, self.hidden_dim)
-        self.hidden = Sequential(
-            *(
-                self.linear(self.hidden_dim, self.hidden_dim)
-                for _ in range(n_hidden_layers)
-            )
-        )
-        self.out = Linear(self.hidden_dim, self.n_outputs)
+        # Default activation function
+        if activation is None:
+            activation = ReLU()
+        self.activation = activation
 
-    def forward(self, data):
-        """
-        Performs a forward pass through the network.
+        # Build network layers
+        layers = []
 
-        Args:
-            data (Tensor): The input tensor to be passed through the network.
+        # Input layer with activation
+        layers.append(Linear(n_inputs, hidden_dim))
+        layers.append(self.activation)
 
-        Returns:
-            dict: A dictionary containing the 'input' (original input) and
-            'output' (predicted output) of the network.
-        """
+        # Hidden layers (with activation after each)
+        for _ in range(n_hidden_layers - 1):
+            layers.append(Linear(hidden_dim, hidden_dim))
+            layers.append(self.activation)
 
-        output = self.out(self.hidden(self.input(data)))
+        # Output layer (no activation by default)
+        layers.append(Linear(hidden_dim, n_outputs))
 
-        return {"input": data, "output": output}
+        self.network = Sequential(*layers)
 
-    @staticmethod
-    def linear(in_features, out_features):
-        """
-        Creates a basic linear block with a linear transformation followed
-        by a ReLU activation function.
+    def forward(self, data: dict[str, Tensor]):
+        """Run a forward pass through the network.
 
         Args:
-            in_features (int): The number of input features.
-            out_features (int): The number of output features.
+            data (dict[str, Tensor]): Input data to be processed by the network.
 
         Returns:
-            nn.Module: A sequential module consisting of a Linear layer
-            and ReLU activation.
+            dict: The original data tensor augmented with the network's output (having key "output").
         """
-
-        return Sequential(
-            Linear(in_features, out_features),
-            ReLU(),
-        )
+        data["output"] = self.network(data["input"])
+        return data