congrads 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

congrads/descriptor.py

@@ -1,65 +1,169 @@
+"""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 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.validation import validate_type
+
+
 class Descriptor:
-    """
-    A class to manage the mapping of neurons to layers and their properties
-    (e.g., output, constant, or variable) in a neural network.
+    """A class to manage the mapping between tags.
 
-    This class enables the organization and description of network elements,
-    such as associating neurons with specific layers and categorizing layers
-    as outputs, constants, or variables.
+    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 allows users to easily place constraints on parts of the network by
-    referencing the name that is configured in this class.
+    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:
+        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,
     ):
-        """
-        Initialize the Descriptor class with empty mappings for neurons and layers.
-
-        This includes:
-            - `neuron_to_layer`: A dictionary mapping neuron names to their corresponding layer names.
-            - `neuron_to_index`: A dictionary mapping neuron names to their corresponding index within a layer.
-            - `output_layers`: A set that holds the names of layers marked as output layers.
-            - `constant_layers`: A set that holds the names of layers marked as constant layers.
-            - `variable_layers`: A set that holds the names of layers marked as variable layers.
-        """
-
-        # 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.output_layers: set[str] = set()
-        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,
-        neuron_names: list[str],
-        output: bool = False,
+        key: str,
+        tag: str,
+        index: int = None,
         constant: bool = False,
+        affects_loss: bool = True,
     ):
-        """
-        Add a layer to the descriptor, associating it with neurons and marking it
-        as an output or constant layer.
+        """Adds a tag to the descriptor with its associated key, index, and properties.
+
+        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 to be added.
-            neuron_names (list[str]): A list of neuron names that belong to the layer.
-            output (bool, optional): If True, mark this layer as an output layer. Defaults to False.
-            constant (bool, optional): If True, mark this layer as 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 key or index is already assigned for a tag or a duplicate index is used within a key.
         """
+        # Type checking
+        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)
 
-        if output:
-            self.output_layers.add(layer_name)
+        # Other validations
+        if tag in self._tag_to_key:
+            raise ValueError(
+                f"There already is a key registered for the tag '{tag}'. "
+                "Please use a unique key name for each tag."
+            )
 
+        if tag in self._tag_to_index:
+            raise ValueError(
+                f"There already is an index registered for the tag '{tag}'. "
+                "Please use a unique name for each tag."
+            )
+
+        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} 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
 
-        for index, neuron_name in enumerate(neuron_names):
-            self.neuron_to_layer[neuron_name] = layer_name
-            self.neuron_to_index[neuron_name] = 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 | None]:
+        """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 | None)): A tuple containing:
+                - The key in the data dictionary which holds the data (str).
+                - The tensor index where the data is present or None (int | None).
+
+        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:
+            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.
+
+        Retrieves the key and index associated with a tag and selects
+        the corresponding slice from the given prediction tensor.
+        Returns the full tensor if no index was specified when registering the tag.
+
+        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, or the full tensor if no index was specified when registering the tag.
+
+        Raises:
+            ValueError: If the tag is not registered in the descriptor.
+        """
+        key, index = self.location(tag)
+        if index is None:
+            return data[key]
+        return data[key][:, index : index + 1]

congrads/metrics.py

@@ -1,64 +1,139 @@
-from torch import Tensor, tensor, sum, numel
-from torchmetrics import Metric
+"""Module for managing metrics during training.
 
-# NOTE
+Provides the `Metric` and `MetricManager` classes for accumulating,
+aggregating, and resetting metrics over training batches. Supports
+grouping metrics and using custom accumulation functions.
+"""
 
+from collections.abc import Callable
 
-class ConstraintSatisfactionRatio(Metric):
-    """
-    A custom metric to calculate the ratio of satisfied constraints in a neural network model.
-    It computes the proportion of constraints that have been satisfied,
-    where satisfaction is determined based on the provided constraint results.
+from torch import Tensor, cat, nanmean, tensor
+
+from .utils.validation import validate_callable, validate_type
 
-    This metric tracks the number of unsatisfied constraints and the total number of constraints
-    during the training process, and computes the ratio of satisfied constraints once all updates
-    have been made.
 
-    Attributes:
-        unsatisfied (Tensor): Tracks the number of unsatisfied constraints.
-        total (Tensor): Tracks the total number of constraints processed.
+class Metric:
+    """Represents a single metric to be accumulated and aggregated.
 
-    Note:
-        For more information about custom metrics, we refer to the Pytorch Lightning documentation
-        at https://lightning.ai/docs/torchmetrics/stable/pages/implement.html
+    Stores metric values over multiple batches and computes an aggregated
+    result using a specified accumulation function.
     """
 
-    def __init__(self, **kwargs):
-        """
-        Initializes the ConstraintSatisfactionRatio metric by setting up the
-        state variables to track the number of unsatisfied and total constraints.
+    def __init__(self, name: str, accumulator: Callable[..., Tensor] = nanmean) -> None:
+        """Initialize a Metric instance.
 
         Args:
-            **kwargs: Additional arguments to pass to the base Metric class constructor.
+            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: list[Tensor] = []
+        self.sample_count = 0
 
-        # Init parent class
-        super().__init__(**kwargs)
+    def accumulate(self, value: Tensor) -> None:
+        """Accumulate a new value for the metric.
+
+        Args:
+            value (Tensor): Metric values for the current batch.
+        """
+        self.values.append(value.detach().clone())
+        self.sample_count += value.size(0)
 
-        # Init scalar tensors that will hold metric values
-        self.add_state("unsatisfied", default=tensor(0), dist_reduce_fx="sum")
-        self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
+    def aggregate(self) -> Tensor:
+        """Compute the aggregated value of the metric.
 
-    def update(self, constraint_result: Tensor) -> None:
+        Returns:
+            Tensor: The aggregated metric value. Returns NaN if no values
+                have been accumulated.
         """
-        Updates the state of the metric with the latest constraint results.
+        if not self.values:
+            return tensor(float("nan"))
+
+        combined = cat(self.values)
+        return self.accumulator(combined)
+
+    def reset(self) -> None:
+        """Reset the accumulated values and sample count for the metric."""
+        self.values = []
+        self.sample_count = 0
+
+
+class MetricManager:
+    """Manages multiple metrics and groups for training or evaluation.
+
+    Supports registering metrics, accumulating values by name, aggregating
+    metrics by group, and resetting metrics by group.
+    """
+
+    def __init__(self) -> None:
+        """Initialize a MetricManager instance."""
+        self.metrics: dict[str, Metric] = {}
+        self.groups: dict[str, str] = {}
+
+    def register(
+        self, name: str, group: str = "default", accumulator: Callable[..., Tensor] = nanmean
+    ) -> None:
+        """Register a new metric under a specified group.
 
         Args:
-            constraint_result (Tensor): A tensor representing the result of
-                                         the constraint checks, where each
-                                         element indicates whether a constraint
-                                         is satisfied (e.g., 0 for satisfied,
-                                         1 for unsatisfied).
+            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`.
         """
-        self.unsatisfied += sum(constraint_result)
-        self.total += numel(constraint_result)
+        # Type checking
+        validate_type("name", name, str)
+        validate_type("group", group, str)
+        validate_callable("accumulator", accumulator)
+
+        self.metrics[name] = Metric(name, accumulator)
+        self.groups[name] = group
+
+    def accumulate(self, name: str, value: Tensor) -> None:
+        """Accumulate a value for a specific metric by name.
 
-    def compute(self) -> Tensor:
+        Args:
+            name (str): Name of the metric.
+            value (Tensor): Metric values for the current batch.
         """
-        Computes the constraint satisfaction ratio, defined as:
-        1 - (number of unsatisfied constraints / total constraints).
+        if name not in self.metrics:
+            raise KeyError(f"Metric '{name}' is not registered.")
+
+        self.metrics[name].accumulate(value)
+
+    def aggregate(self, group: str = "default") -> dict[str, Tensor]:
+        """Aggregate all metrics in a specified group.
+
+        Args:
+            group (str, optional): The group of metrics to aggregate. Defaults to "default".
 
         Returns:
-            Tensor: The satisfaction ratio as a scalar tensor.
+            dict[str, Tensor]: Dictionary mapping metric names to their
+                aggregated values.
         """
-        return 1 - (self.unsatisfied.float() / self.total)
+        return {
+            name: metric.aggregate()
+            for name, metric in self.metrics.items()
+            if self.groups[name] == group
+        }
+
+    def reset(self, group: str = "default") -> None:
+        """Reset all metrics in a specified group.
+
+        Args:
+            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()
+
+    def reset_all(self) -> None:
+        """Reset all metrics across all groups."""
+        for metric in self.metrics.values():
+            metric.reset()

congrads/networks/registry.py

@@ -0,0 +1,68 @@
+"""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 with configurable hidden layers."""
+
+    def __init__(
+        self,
+        n_inputs,
+        n_outputs,
+        n_hidden_layers=3,
+        hidden_dim=35,
+        activation=None,
+    ):
+        """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
+        self.n_inputs = n_inputs
+        self.n_outputs = n_outputs
+        self.n_hidden_layers = n_hidden_layers
+        self.hidden_dim = hidden_dim
+
+        # Default activation function
+        if activation is None:
+            activation = ReLU()
+        self.activation = activation
+
+        # Build network layers
+        layers = []
+
+        # Input layer with activation
+        layers.append(Linear(n_inputs, hidden_dim))
+        layers.append(self.activation)
+
+        # 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 layer (no activation by default)
+        layers.append(Linear(hidden_dim, n_outputs))
+
+        self.network = Sequential(*layers)
+
+    def forward(self, data: dict[str, Tensor]):
+        """Run a forward pass through the network.
+
+        Args:
+            data (dict[str, Tensor]): Input data to be processed by the network.
+
+        Returns:
+            dict: The original data tensor augmented with the network's output (having key "output").
+        """
+        data["output"] = self.network(data["input"])
+        return data

congrads/transformations/base.py

@@ -0,0 +1,37 @@
+"""Module defining transformations and components."""
+
+from abc import ABC, abstractmethod
+
+from torch import Tensor
+
+from ..utils.validation import validate_type
+
+
+class Transformation(ABC):
+    """Abstract base class for tag data transformations."""
+
+    def __init__(self, tag: str):
+        """Initialize a Transformation.
+
+        Args:
+            tag (str): Tag this transformation applies to.
+        """
+        validate_type("tag", tag, str)
+
+        super().__init__()
+        self.tag = tag
+
+    @abstractmethod
+    def __call__(self, data: Tensor) -> Tensor:
+        """Apply the transformation to the input tensor.
+
+        Args:
+            data (Tensor): Input tensor representing network data.
+
+        Returns:
+            Tensor: Transformed tensor.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError

congrads/transformations/registry.py

@@ -0,0 +1,86 @@
+"""Module holding specific transformation implementations."""
+
+from numbers import Number
+
+from torch import Tensor
+
+from ..utils.validation import validate_callable, validate_type
+from .base import Transformation
+
+
+class IdentityTransformation(Transformation):
+    """A transformation that returns the input unchanged."""
+
+    def __call__(self, data: Tensor) -> Tensor:
+        """Return the input tensor without any modification.
+
+        Args:
+            data (Tensor): Input tensor.
+
+        Returns:
+            Tensor: The same input tensor.
+        """
+        return data
+
+
+class DenormalizeMinMax(Transformation):
+    """A transformation that denormalizes data using min-max scaling."""
+
+    def __init__(self, tag: str, min: Number, max: Number):
+        """Initialize a min-max denormalization transformation.
+
+        Args:
+            tag (str): Tag this transformation applies to.
+            min (Number): Minimum value used for denormalization.
+            max (Number): Maximum value used for denormalization.
+        """
+        validate_type("min", min, Number)
+        validate_type("max", max, Number)
+
+        super().__init__(tag)
+
+        self.min = min
+        self.max = max
+
+    def __call__(self, data: Tensor) -> Tensor:
+        """Denormalize the input tensor using the min-max range.
+
+        Args:
+            data (Tensor): Normalized input tensor (typically in range [0, 1]).
+
+        Returns:
+            Tensor: Denormalized tensor in the range [min, max].
+        """
+        return data * (self.max - self.min) + self.min
+
+
+class ApplyOperator(Transformation):
+    """A transformation that applies a binary operator to the input tensor."""
+
+    def __init__(self, tag: str, operator: callable, value: Number):
+        """Initialize an operator-based transformation.
+
+        Args:
+            tag (str): Tag this transformation applies to.
+            operator (callable): A callable that takes two arguments (tensor, value)
+                and returns a tensor.
+            value (Number): The value to use as the second argument in the operator.
+        """
+        validate_callable("operator", operator)
+        validate_type("value", value, Number)
+
+        super().__init__(tag)
+
+        self.operator = operator
+        self.value = value
+
+    def __call__(self, data: Tensor) -> Tensor:
+        """Apply the operator to the input tensor and the specified value.
+
+        Args:
+            data (Tensor): Input tensor.
+
+        Returns:
+            Tensor: Result of applying `operator(data, value)`.
+        """
+        return self.operator(data, self.value)