feat: support tensor composition

Author: mtth

src/opvious/modeling/__init__.py

@@ -22,6 +22,7 @@
     total,
 )
 from .definitions import (
+    ComposeFunction,
     Constraint,
     Dimension,
     Image,
@@ -31,7 +32,9 @@
     TensorLike,
     Variable,
     alias,
+    compose,
     constraint,
+    infer_quantifiables,
     interval,
     objective,
 )
@@ -81,5 +84,8 @@
     # Fragments
     "fragments",
     # Utilities
+    "ComposeFunction",
+    "compose",
+    "infer_quantifiables",
     "method_decorator",
 ]

src/opvious/modeling/definitions.py

@@ -214,9 +214,6 @@ def total_product_count(self):
     The number of arguments must match the tensor's quantification.
     """
 
-    # TODO: Add map method, which appends to _mappers array of transformations.
-    # Once implemented, remove the negate arguments to transformations.
-
     def __init__(
         self,
         image: Image,
@@ -387,6 +384,46 @@ def render_statement(self, label: Label) -> str | None:
         return s
 
 
+type ComposeFunction = Callable[[*tuple[Expression, ...]], Expression]
+
+
+def compose(fn: ComposeFunction, *tensors: TensorLike) -> TensorLike:
+    """Maps a function over one or more tensors
+
+    This is particularly useful when creating model fragments which use
+    tensor-likes as arguments. For example it can be used to flip the indicator
+    in an :class:`~opvious.modeling.fragments.ActivatedVariable`.
+    """
+    if not tensors:
+        raise ValueError("No tensors")
+    return _Composed(fn, tensors)
+
+
+@dataclasses.dataclass(frozen=True)
+class _Composed:
+    function: ComposeFunction
+    tensors: Sequence[TensorLike]
+
+    def __call__(self, *subscripts: ExpressionLike) -> Expression:
+        exprs = [tensor(*subscripts) for tensor in self.tensors]
+        return self.function(*exprs)
+
+
+def infer_quantifiables(tensor: TensorLike) -> tuple[Quantifiable, ...]:
+    """Infers a tensor-like's underlying quantification
+
+    Args:
+        tensor: A tensor or composed tensor with only unary functions
+    """
+    match tensor:
+        case Tensor():
+            return tensor.quantifiables()
+        case _Composed(_, tensors) if len(tensors) == 1:
+            return infer_quantifiables(tensors[0])
+        case _:
+            raise Exception(f"Unable to infer quantifiables from {tensor}")
+
+
 def _is_simple_domain(d: Domain) -> bool:
     if d.mask is not None:
         return False

src/opvious/modeling/fragments.py

@@ -32,6 +32,7 @@
     Variable,
     alias,
     constraint,
+    infer_quantifiables,
     interval,
 )
 from .identifiers import Name
@@ -157,9 +158,9 @@ def __init__(
         lower_bound: bool = True,
         upper_bound: bool = True,
     ) -> None:
+        if not quantifiables:
+            quantifiables = infer_quantifiables(tensor)
         if isinstance(tensor, Tensor):
-            if not quantifiables:
-                quantifiables = tensor.quantifiables()
             if not image:
                 image = tensor.image
             if lower_bound and tensor.image.lower_bound == 0:
@@ -246,9 +247,9 @@ class ActivationVariable(ModelFragment):
             used, if `False` no activation constraint will be added.
         lower_bound: Value of the lower bound used in the deactivation
             constraint. If `True` the variable's image's lower bound will be
-            used, if `False` no deactivation constraint will be added.
+            used, if `False` no deactivation constraint will be added. Can only
+            be used if the tensor is non-negative.
         name: Name of the generated activation variable
-        negate: Negate the returned indicator variable.
         projection: Mask used to project the variable's quantification. When
             this is set, the indicator variable will be set to 1 iff at least
             one of the projected tensor values is positive.
@@ -263,11 +264,10 @@ def __new__(
         upper_bound: ExpressionLike | TensorLike | bool = True,
         lower_bound: ExpressionLike | TensorLike | bool = False,
         name: Name | None = None,
-        negate: bool = False,
         projection: Projection = -1,
     ) -> ActivationVariable:
-        if not quantifiables and isinstance(tensor, Tensor):
-            quantifiables = tensor.quantifiables()
+        if not quantifiables:
+            quantifiables = infer_quantifiables(tensor)
         domains = tuple(domain(q) for q in quantifiables)
 
         def quantification(
@@ -299,8 +299,7 @@ def activates(self) -> Quantified:
                         bound = bound(*cp.lifted)
                     elif bound is True:
                         bound = tensor_image().upper_bound
-                    value = 1 - self.value(*cp) if negate else self.value(*cp)
-                    yield bound * value >= tensor(*cp.lifted)
+                    yield bound * self.value(*cp) >= tensor(*cp.lifted)
 
             @constraint(disabled=lower_bound is False)
             def deactivates(self) -> Quantified:
@@ -317,8 +316,7 @@ def deactivates(self) -> Quantified:
                         bound = bound(*cp)
                     elif bound is True:
                         bound = tensor_image().lower_bound
-                    value = 1 - self.value(*cp) if negate else self.value(*cp)
-                    yield bound * value <= term
+                    yield bound * self.value(*cp) <= term
 
         return _Fragment()
 
@@ -362,7 +360,6 @@ def activation_variable(
     upper_bound: ExpressionLike | TensorLike | bool = True,
     lower_bound: ExpressionLike | TensorLike | bool = False,
     name: Name | None = None,
-    negate: bool = False,
     projection: Projection = -1,
 ) -> Callable[[TensorLike], ActivationVariable]:
     """Transforms a method into an :class:`ActivationVariable` fragment
@@ -379,7 +376,6 @@ def wrapper(fn: TensorLike) -> ActivationVariable:
             lower_bound=lower_bound,
             upper_bound=upper_bound,
             name=name,
-            negate=negate,
             projection=projection,
         )
 
@@ -416,8 +412,8 @@ def __init__(
             raise NotImplementedError()  # TODO: Implement.
 
         self._tensor = tensor
-        if not quantifiables and isinstance(tensor, Tensor):
-            quantifiables = tensor.quantifiables()
+        if not quantifiables:
+            quantifiables = infer_quantifiables(tensor)
         self._domains = tuple(domain(q) for q in quantifiables)
 
         self._piece_count = Parameter.discrete(
@@ -528,7 +524,6 @@ class ActivatedVariable(ModelFragment):
             subscripts
         upper_bound: Tensor upper bound, can be omitted if `tensor` is a
             :class:`~opvious.modeling.Tensor` instance
-        negate: Negate the input indicator
         force_activation: Add constraint to ensure that the derived variable is
             at least equal to `tensor` when `indicator` is non-zero. You may
             choose to omit this if the variable is already pushed up via other
@@ -549,13 +544,11 @@ def __init__(
         upper_bound: ExpressionLike | None = None,
         force_activation: bool = True,
         force_deactivation: bool = True,
-        negate: bool = False,
         name: Name | None = None,
     ) -> None:
         self._tensor = tensor
         self._indicator = indicator
         self._indicator_projection = indicator_projection
-        self._negate = negate
         self._force_activation = force_activation
         self._force_deactivation = force_deactivation
 
@@ -593,12 +586,9 @@ def deactivates(self) -> Quantified:
         This constraint will be omitted if `force_deactivation` is false.
         """
         for cp in self._quantification():
-            toggle = (
-                1 - self._indicator(*cp)
-                if self._negate
-                else self._indicator(*cp)
-            )
-            yield self.value(*cp.lifted) <= self._upper_bound * toggle
+            yield self.value(
+                *cp.lifted
+            ) <= self._upper_bound * self._indicator(*cp)
 
     @constraint(lambda init, self: init(disabled=not self._force_activation))
     def activates(self) -> Quantified:
@@ -607,14 +597,10 @@ def activates(self) -> Quantified:
         This constraint will be omitted if `force_activation` is false.
         """
         for cp in self._quantification():
-            toggle = (
-                self._indicator(*cp)
-                if self._negate
-                else 1 - self._indicator(*cp)
-            )
             yield (
                 self.value(*cp.lifted)
-                >= self._tensor(*cp.lifted) - self._upper_bound * toggle
+                >= self._tensor(*cp.lifted)
+                - self._upper_bound * (1 - self._indicator(*cp))
             )
 
 
@@ -624,7 +610,6 @@ def activated_variable(
     indicator: Tensor,
     indicator_projection: Projection = -1,
     upper_bound: ExpressionLike | None = None,
-    negate: bool = False,
     force_activation: bool = True,
     force_deactivation: bool = True,
     name: Name | None = None,
@@ -641,7 +626,6 @@ def wrapper(fn: TensorLike) -> ActivatedVariable:
             indicator=indicator,
             indicator_projection=indicator_projection,
             upper_bound=upper_bound,
-            negate=negate,
             force_activation=force_activation,
             force_deactivation=force_deactivation,
             name=name,

tests/modeling_test.py

@@ -259,29 +259,21 @@ def child_starts_after_parent_ends(self):
     @om.fragments.activation_variable(
         lambda init, self: init(
             self.competing_tasks,
-            negate=True,
             upper_bound=self.duration.total(),
         )
     )
-    def must_start_after(self, t1, t2):
+    # t2 ends after t1 starts => 1
+    # 0 => t2 ends before t1 starts
+    def may_start_before_end(self, t1, t2):
         return self.task_end(t2) - self.task_start(t1)
 
-    @om.fragments.activation_variable(
-        lambda init, self: init(
-            self.competing_tasks,
-            negate=True,
-            upper_bound=self.duration.total(),
-        )
-    )
-    def must_end_before(self, t1, t2):
-        return self.task_end(t1) - self.task_start(t2)
-
     @om.constraint
     def one_active_task_per_machine(self):
         for t1, t2 in self.competing_tasks:
             yield (
-                self.must_end_before(t1, t2) + self.must_start_after(t1, t2)
-                >= 1
+                self.may_start_before_end(t1, t2) +
+                self.may_start_before_end(t2, t1)
+                <= 1
             )
 
     @om.objective