congrads 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

congrads/{constraints.py → constraints/registry.py}

@@ -24,10 +24,6 @@ Usage:
 
 """
 
-import random
-import string
-import warnings
-from abc import ABC, abstractmethod
 from collections.abc import Callable
 from numbers import Number
 from typing import Literal
@@ -54,156 +50,10 @@ from torch import (
 )
 from torch.nn.functional import normalize
 
-from .descriptor import Descriptor
-from .transformations import IdentityTransformation, Transformation
-from .utils import validate_comparator_pytorch, validate_iterable, validate_type
-
-
-class Constraint(ABC):
-    """Abstract base class for defining constraints applied to neural networks.
-
-    A `Constraint` specifies conditions that the neural network outputs
-    should satisfy. It supports monitoring constraint satisfaction
-    during training and can adjust loss to enforce constraints. Subclasses
-    must implement the `check_constraint` and `calculate_direction` methods.
-
-    Args:
-        tags (set[str]): Tags referencing parts of the network where this constraint applies to.
-        name (str, optional): A unique name for the constraint. If not provided,
-            a name is generated based on the class name and a random suffix.
-        enforce (bool, optional): If False, only monitor the constraint
-            without adjusting the loss. Defaults to True.
-        rescale_factor (Number, optional): Factor to scale the
-            constraint-adjusted loss. Defaults to 1.5. Should be greater
-            than 1 to give weight to the constraint.
-
-    Raises:
-        TypeError: If a provided attribute has an incompatible type.
-        ValueError: If any tag in `tags` is not
-            defined in the `descriptor`.
-
-    Note:
-        - If `rescale_factor <= 1`, a warning is issued.
-        - If `name` is not provided, a name is auto-generated,
-          and a warning is logged.
-
-    """
-
-    descriptor: Descriptor = None
-    device = None
-
-    def __init__(
-        self, tags: set[str], name: str = None, enforce: bool = True, rescale_factor: Number = 1.5
-    ) -> None:
-        """Initializes a new Constraint instance.
-
-        Args:
-            tags (set[str]): Tags referencing parts of the network where this constraint applies to.
-            name (str, optional): A unique name for the constraint. If not
-                provided, a name is generated based on the class name and a
-                random suffix.
-            enforce (bool, optional): If False, only monitor the constraint
-                without adjusting the loss. Defaults to True.
-            rescale_factor (Number, optional): Factor to scale the
-                constraint-adjusted loss. Defaults to 1.5. Should be greater
-                than 1 to give weight to the constraint.
-
-        Raises:
-            TypeError: If a provided attribute has an incompatible type.
-            ValueError: If any tag in `tags` is not defined in the `descriptor`.
-
-        Note:
-            - If `rescale_factor <= 1`, a warning is issued.
-            - If `name` is not provided, a name is auto-generated, and a
-            warning is logged.
-        """
-        # Init parent class
-        super().__init__()
-
-        # Type checking
-        validate_iterable("tags", tags, str)
-        validate_type("name", name, str, allow_none=True)
-        validate_type("enforce", enforce, bool)
-        validate_type("rescale_factor", rescale_factor, Number)
-
-        # Init object variables
-        self.tags = tags
-        self.rescale_factor = rescale_factor
-        self.initial_rescale_factor = rescale_factor
-        self.enforce = enforce
-
-        # Perform checks
-        if rescale_factor <= 1:
-            warnings.warn(
-                f"Rescale factor for constraint {name} is <= 1. The network "
-                "will favor general loss over the constraint-adjusted loss. "
-                "Is this intended behavior? Normally, the rescale factor "
-                "should always be larger than 1.",
-                stacklevel=2,
-            )
-
-        # If no constraint_name is set, generate one based
-        # on the class name and a random suffix
-        if name:
-            self.name = name
-        else:
-            random_suffix = "".join(random.choices(string.ascii_uppercase + string.digits, k=6))
-            self.name = f"{self.__class__.__name__}_{random_suffix}"
-            warnings.warn(f"Name for constraint is not set. Using {self.name}.", stacklevel=2)
-
-        # Infer layers from descriptor and tags
-        self.layers = set()
-        for tag in self.tags:
-            if not self.descriptor.exists(tag):
-                raise ValueError(
-                    f"The tag {tag} used with constraint "
-                    f"{self.name} is not defined in the descriptor. Please "
-                    "add it to the correct layer using "
-                    "descriptor.add('layer', ...)."
-                )
-
-            layer, _ = self.descriptor.location(tag)
-            self.layers.add(layer)
-
-    @abstractmethod
-    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
-        """Evaluates whether the given model predictions satisfy the constraint.
-
-        1 IS SATISFIED, 0 IS NOT SATISFIED
-
-        Args:
-            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
-
-        Returns:
-            tuple[Tensor, Tensor]: A tuple where the first element is a tensor of floats
-            indicating whether the constraint is satisfied (with value 1.0
-            for satisfaction, and 0.0 for non-satisfaction, and the second element is a tensor
-            mask that indicates the relevance of each sample (`True` for relevant
-            samples and `False` for irrelevant ones).
-
-        Raises:
-            NotImplementedError: If not implemented in a subclass.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
-        """Compute adjustment directions to better satisfy the constraint.
-
-        Given the model predictions, input batch, and context, this method calculates the direction
-        in which the predictions referenced by a tag should be adjusted to satisfy the constraint.
-
-        Args:
-            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
-
-        Returns:
-            dict[str, Tensor]: Dictionary mapping network layers to tensors that
-                specify the adjustment direction for each tag.
-
-        Raises:
-            NotImplementedError: Must be implemented by subclasses.
-        """
-        raise NotImplementedError
+from ..transformations.base import Transformation
+from ..transformations.registry import IdentityTransformation
+from ..utils.validation import validate_comparator_pytorch, validate_iterable, validate_type
+from .base import Constraint
 
 
 class ImplicationConstraint(Constraint):
@@ -918,11 +768,23 @@ class MonotonicityConstraint(Constraint):
         return {layer: self.compared_rankings}
 
 
-class GroupedMonotonicityConstraint(MonotonicityConstraint):
-    """Constraint that enforces a monotonic relationship between two tags.
+class PerGroupMonotonicityConstraint(MonotonicityConstraint):
+    """Group-wise monotonicity constraint enforced independently per group.
 
-    This constraint ensures that the activations of a prediction tag (`tag_prediction`)
-    are monotonically ascending or descending with respect to a target tag (`tag_reference`).
+    This constraint enforces a monotonic relationship between a prediction tag
+    (`tag_prediction`) and a reference tag (`tag_reference`) **within each group**
+    identified by `tag_group_identifier`.
+
+    For each unique group identifier, the base `MonotonicityConstraint` is applied
+    independently to the corresponding subset of samples. This makes the behavior
+    semantically explicit and easy to reason about, as no interaction or ordering
+    is introduced across different groups.
+
+    Notes:
+        - Groups are treated as fully independent constraint instances.
+        - This implementation prioritizes correctness and interpretability.
+        - It may be less efficient for large numbers of groups due to explicit
+          iteration and repeated constraint evaluation.
     """
 
     def __init__(
@@ -1010,6 +872,106 @@ class GroupedMonotonicityConstraint(MonotonicityConstraint):
         return {layer: self.directions}
 
 
+class EncodedGroupedMonotonicityConstraint(MonotonicityConstraint):
+    """Group-wise monotonicity constraint enforced via rank encoding.
+
+    This constraint enforces a monotonic relationship between a prediction tag
+    (`tag_prediction`) and a reference tag (`tag_reference`) within each group
+    identified by `tag_group_identifier`, using a fully vectorized approach.
+
+    Group independence is achieved by encoding the group identifiers into the
+    prediction and reference values via large offsets, effectively separating
+    the rank spaces of different groups. This allows the base
+    `MonotonicityConstraint` to be applied once to the entire batch without
+    explicit per-group iteration.
+
+    Notes:
+        - Groups are isolated implicitly through rank-space separation.
+        - The logic is less explicit than per-group evaluation and relies on
+            offset-based rank encoding for correctness.
+        - This constraint might cause floating point errors if
+            the prediction or target range is very large.
+    """
+
+    def __init__(
+        self,
+        tag_prediction: str,
+        tag_reference: str,
+        tag_group_identifier: str,
+        rescale_factor_lower: float = 1.5,
+        rescale_factor_upper: float = 1.75,
+        stable: bool = True,
+        direction: Literal["ascending", "descending"] = "ascending",
+        name: str = None,
+        enforce: bool = True,
+    ):
+        """Constraint that enforces monotonicity on a predicted output.
+
+        This constraint ensures that the activations of a prediction tag (`tag_prediction`)
+        are monotonically ascending or descending with respect to a target tag (`tag_reference`).
+
+        Args:
+            tag_prediction (str): Name of the tag whose activations should follow the monotonic relationship.
+            tag_reference (str): Name of the tag that acts as the monotonic reference.
+            tag_group_identifier (str): Name of the tag that identifies groups for separate monotonicity enforcement.
+            rescale_factor_lower (float, optional): Lower bound for rescaling rank differences. Defaults to 1.5.
+            rescale_factor_upper (float, optional): Upper bound for rescaling rank differences. Defaults to 1.75.
+            stable (bool, optional): Whether to use stable sorting when ranking. Defaults to True.
+            direction (str, optional): Direction of monotonicity to enforce, either 'ascending' or 'descending'. Defaults to 'ascending'.
+            name (str, optional): Custom name for the constraint. If None, a descriptive name is auto-generated.
+            enforce (bool, optional): If False, the constraint is only monitored (not enforced). Defaults to True.
+        """
+        # Compose constraint name
+        if name is None:
+            name = f"{tag_prediction} for each {tag_group_identifier} monotonically {direction} by {tag_reference}"
+
+        # Init parent class
+        super().__init__(
+            tag_prediction=tag_prediction,
+            tag_reference=tag_reference,
+            rescale_factor_lower=rescale_factor_lower,
+            rescale_factor_upper=rescale_factor_upper,
+            stable=stable,
+            direction=direction,
+            name=name,
+            enforce=enforce,
+        )
+
+        # Init variables
+        self.tag_prediction = tag_prediction
+        self.tag_reference = tag_reference
+        self.tag_group_identifier = tag_group_identifier
+
+        # Initialize negation factor based on direction
+        self.negation = 1 if direction == "ascending" else -1
+
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluate whether the monotonicity constraint is satisfied."""
+        # Get data and keys
+        ids = self.descriptor.select(self.tag_group_identifier, data)
+        preds = self.descriptor.select(self.tag_prediction, data)
+        targets = self.descriptor.select(self.tag_reference, data)
+        preds_key, _ = self.descriptor.location(self.tag_prediction)
+        targets_key, _ = self.descriptor.location(self.tag_reference)
+
+        new_preds = preds + ids * (preds.max() - preds.min() + 1)
+        new_targets = targets + self.negation * ids * (targets.max() - targets.min() + 1)
+
+        # Create new batch for child constraint
+        new_data = {preds_key: new_preds, targets_key: new_targets}
+
+        # Call super on the adjusted batch
+        checks, _ = super().check_constraint(new_data)
+        self.directions = self.compared_rankings
+
+        return checks, ones_like(checks)
+
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Calculates ranking adjustments for monotonicity enforcement."""
+        layer, _ = self.descriptor.location(self.tag_prediction)
+        return {layer: self.directions}
+
+
 class ANDConstraint(Constraint):
     """A composite constraint that enforces the logical AND of multiple constraints.
 

congrads/core/batch_runner.py

@@ -0,0 +1,200 @@
+"""Defines the BatchRunner, which executes individual batches for training, validation, and testing.
+
+Responsibilities:
+- Move batch data to the appropriate device
+- Run forward passes through the network
+- Compute base and constraint-adjusted losses
+- Perform backpropagation during training
+- Accumulate metrics for loss and other monitored quantities
+- Trigger callbacks at key points in the batch lifecycle
+"""
+
+import torch
+from torch import Tensor
+from torch.nn import Module
+from torch.optim import Optimizer
+
+from ..callbacks.base import CallbackManager
+from ..core.constraint_engine import ConstraintEngine
+from ..metrics import MetricManager
+
+
+class BatchRunner:
+    """Executes a single batch for training, validation, or testing.
+
+    The BatchRunner handles moving data to the correct device, running the network
+    forward, computing base and constraint-adjusted losses, performing backpropagation
+    during training, accumulating metrics, and dispatching callbacks at key points
+    in the batch lifecycle.
+    """
+
+    def __init__(
+        self,
+        network: Module,
+        criterion,
+        optimizer: Optimizer,
+        constraint_engine: ConstraintEngine,
+        metric_manager: MetricManager | None,
+        callback_manager: CallbackManager | None,
+        device: torch.device,
+    ):
+        """Initialize the BatchRunner.
+
+        Args:
+            network: The neural network module to execute.
+            criterion: Loss function callable accepting (output, target, data=batch).
+            optimizer: Optimizer for updating network parameters.
+            constraint_engine: ConstraintEngine instance for evaluating and enforcing constraints.
+            metric_manager: Optional MetricManager for logging batch metrics.
+            callback_manager: Optional CallbackManager for triggering hooks during batch processing.
+            device: Torch device on which to place data and network.
+        """
+        self.network = network
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.constraint_engine = constraint_engine
+        self.metric_manager = metric_manager
+        self.callback_manager = callback_manager
+        self.device = device
+
+    def _to_device(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Move all tensors in the batch to the BatchRunner's device.
+
+        Args:
+            batch: Dictionary of tensors for a single batch.
+
+        Returns:
+            Dictionary of tensors moved to the target device.
+        """
+        return {k: v.to(self.device) for k, v in batch.items()}
+
+    def _run_callbacks(self, hook: str, data: dict) -> dict:
+        """Run the specified callback hook on the batch data.
+
+        Args:
+            hook: Name of the callback hook to run.
+            data: Dictionary containing batch data.
+
+        Returns:
+            Potentially modified batch data after callback execution.
+        """
+        if self.callback_manager is None:
+            return data
+        return self.callback_manager.run(hook, data)
+
+    def train_batch(self, batch: dict[str, Tensor]) -> Tensor:
+        """Run a single training batch.
+
+        Steps performed:
+        1. Move batch to device and run "on_train_batch_start" callback.
+        2. Forward pass through the network.
+        3. Compute base loss using the criterion and accumulate metric.
+        4. Apply constraint-based adjustments to the loss.
+        5. Perform backward pass and optimizer step.
+        6. Run "on_train_batch_end" callback.
+
+        Args:
+            batch: Dictionary of input and target tensors for the batch.
+
+        Returns:
+            Tensor: The base loss computed before constraint adjustments.
+        """
+        batch = self._to_device(batch)
+        batch = self._run_callbacks("on_train_batch_start", batch)
+
+        # Forward
+        batch = self.network(batch)
+        batch = self._run_callbacks("after_train_forward", batch)
+
+        # Base loss
+        loss: Tensor = self.criterion(
+            batch["output"],
+            batch["target"],
+            data=batch,
+        )
+
+        if self.metric_manager is not None:
+            self.metric_manager.accumulate("Loss/train", loss.unsqueeze(0))
+
+        # Constraint-adjusted loss
+        combined_loss = self.constraint_engine.train(batch, loss)
+
+        # Backward
+        self.optimizer.zero_grad()
+        combined_loss.backward()
+        self.optimizer.step()
+
+        batch = self._run_callbacks("on_train_batch_end", batch)
+        return loss
+
+    def valid_batch(self, batch: dict[str, Tensor]) -> Tensor:
+        """Run a single validation batch.
+
+        Steps performed:
+        1. Move batch to device and run "on_valid_batch_start" callback.
+        2. Forward pass through the network.
+        3. Compute base loss using the criterion and accumulate metric.
+        4. Evaluate constraints via the ConstraintEngine (does not modify loss).
+        5. Run "on_valid_batch_end" callback.
+
+        Args:
+            batch: Dictionary of input and target tensors for the batch.
+
+        Returns:
+            Tensor: The base loss computed for the batch.
+        """
+        batch = self._to_device(batch)
+        batch = self._run_callbacks("on_valid_batch_start", batch)
+
+        batch = self.network(batch)
+        batch = self._run_callbacks("after_valid_forward", batch)
+
+        loss: Tensor = self.criterion(
+            batch["output"],
+            batch["target"],
+            data=batch,
+        )
+
+        if self.metric_manager is not None:
+            self.metric_manager.accumulate("Loss/valid", loss.unsqueeze(0))
+
+        self.constraint_engine.validate(batch, loss)
+
+        batch = self._run_callbacks("on_valid_batch_end", batch)
+        return loss
+
+    def test_batch(self, batch: dict[str, Tensor]) -> Tensor:
+        """Run a single test batch.
+
+        Steps performed:
+        1. Move batch to device and run "on_test_batch_start" callback.
+        2. Forward pass through the network.
+        3. Compute base loss using the criterion and accumulate metric.
+        4. Evaluate constraints via the ConstraintEngine (does not modify loss).
+        5. Run "on_test_batch_end" callback.
+
+        Args:
+            batch: Dictionary of input and target tensors for the batch.
+
+        Returns:
+            Tensor: The base loss computed for the batch.
+        """
+        batch = self._to_device(batch)
+        batch = self._run_callbacks("on_test_batch_start", batch)
+
+        batch = self.network(batch)
+        batch = self._run_callbacks("after_test_forward", batch)
+
+        loss: Tensor = self.criterion(
+            batch["output"],
+            batch["target"],
+            data=batch,
+        )
+
+        if self.metric_manager is not None:
+            self.metric_manager.accumulate("Loss/test", loss.unsqueeze(0))
+
+        self.constraint_engine.test(batch, loss)
+
+        batch = self._run_callbacks("on_test_batch_end", batch)
+        return loss