ortools.math_opt.python.sparse_containers

  1# Copyright 2010-2025 Google LLC
  2# Licensed under the Apache License, Version 2.0 (the "License");
  3# you may not use this file except in compliance with the License.
  4# You may obtain a copy of the License at
  5#
  6#     http://www.apache.org/licenses/LICENSE-2.0
  7#
  8# Unless required by applicable law or agreed to in writing, software
  9# distributed under the License is distributed on an "AS IS" BASIS,
 10# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 11# See the License for the specific language governing permissions and
 12# limitations under the License.
 13
 14"""Sparse vectors and matrices using variables and constraints from Model.
 15
 16Analogous to sparse_containers.proto, with bidirectional conversion.
 17"""
 18from typing import Dict, FrozenSet, Generic, Iterable, Mapping, Optional, Set, TypeVar
 19
 20from ortools.math_opt import sparse_containers_pb2
 21from ortools.math_opt.python import linear_constraints
 22from ortools.math_opt.python import model
 23from ortools.math_opt.python import quadratic_constraints
 24from ortools.math_opt.python import variables
 25
 26
 27VarOrConstraintType = TypeVar(
 28    "VarOrConstraintType",
 29    variables.Variable,
 30    linear_constraints.LinearConstraint,
 31    quadratic_constraints.QuadraticConstraint,
 32)
 33
 34
 35def to_sparse_double_vector_proto(
 36    terms: Mapping[VarOrConstraintType, float],
 37) -> sparse_containers_pb2.SparseDoubleVectorProto:
 38    """Converts a sparse vector from proto to dict representation."""
 39    result = sparse_containers_pb2.SparseDoubleVectorProto()
 40    if terms:
 41        id_and_values = [(key.id, value) for (key, value) in terms.items()]
 42        id_and_values.sort()
 43        ids, values = zip(*id_and_values)
 44        result.ids[:] = ids
 45        result.values[:] = values
 46    return result
 47
 48
 49def to_sparse_int32_vector_proto(
 50    terms: Mapping[VarOrConstraintType, int],
 51) -> sparse_containers_pb2.SparseInt32VectorProto:
 52    """Converts a sparse vector from proto to dict representation."""
 53    result = sparse_containers_pb2.SparseInt32VectorProto()
 54    if terms:
 55        id_and_values = [(key.id, value) for (key, value) in terms.items()]
 56        id_and_values.sort()
 57        ids, values = zip(*id_and_values)
 58        result.ids[:] = ids
 59        result.values[:] = values
 60    return result
 61
 62
 63def parse_variable_map(
 64    proto: sparse_containers_pb2.SparseDoubleVectorProto,
 65    mod: model.Model,
 66    validate: bool = True,
 67) -> Dict[variables.Variable, float]:
 68    """Converts a sparse vector of variables from proto to dict representation."""
 69    result = {}
 70    for index, var_id in enumerate(proto.ids):
 71        result[mod.get_variable(var_id, validate=validate)] = proto.values[index]
 72    return result
 73
 74
 75def parse_linear_constraint_map(
 76    proto: sparse_containers_pb2.SparseDoubleVectorProto,
 77    mod: model.Model,
 78    validate: bool = True,
 79) -> Dict[linear_constraints.LinearConstraint, float]:
 80    """Converts a sparse vector of linear constraints from proto to dict representation."""
 81    result = {}
 82    for index, lin_con_id in enumerate(proto.ids):
 83        result[mod.get_linear_constraint(lin_con_id, validate=validate)] = proto.values[
 84            index
 85        ]
 86    return result
 87
 88
 89def parse_quadratic_constraint_map(
 90    proto: sparse_containers_pb2.SparseDoubleVectorProto,
 91    mod: model.Model,
 92    validate: bool = True,
 93) -> Dict[quadratic_constraints.QuadraticConstraint, float]:
 94    """Converts a sparse vector of quadratic constraints from proto to dict representation."""
 95    result = {}
 96    for index, quad_con_id in enumerate(proto.ids):
 97        result[mod.get_quadratic_constraint(quad_con_id, validate=validate)] = (
 98            proto.values[index]
 99        )
100    return result
101
102
103class SparseVectorFilter(Generic[VarOrConstraintType]):
104    """Restricts the variables or constraints returned in a sparse vector.
105
106    The default behavior is to return entries for all variables/constraints.
107
108    E.g. when requesting the solution to an optimization problem, use this class
109    to restrict the variables that values are returned for.
110
111    Attributes:
112      skip_zero_values: Do not include key value pairs with value zero.
113      filtered_items: If not None, include only key value pairs these keys. Note
114        that the empty set is different (don't return any keys) from None (return
115        all keys).
116    """
117
118    def __init__(
119        self,
120        *,
121        skip_zero_values: bool = False,
122        filtered_items: Optional[Iterable[VarOrConstraintType]] = None,
123    ):
124        self._skip_zero_values: bool = skip_zero_values
125        self._filtered_items: Optional[Set[VarOrConstraintType]] = (
126            None if filtered_items is None else frozenset(filtered_items)
127        )  # pytype: disable=annotation-type-mismatch  # attribute-variable-annotations
128
129    @property
130    def skip_zero_values(self) -> bool:
131        return self._skip_zero_values
132
133    @property
134    def filtered_items(self) -> Optional[FrozenSet[VarOrConstraintType]]:
135        return (
136            self._filtered_items
137        )  # pytype: disable=bad-return-type  # attribute-variable-annotations
138
139    def to_proto(self) -> sparse_containers_pb2.SparseVectorFilterProto:
140        """Returns an equivalent proto representation."""
141        result = sparse_containers_pb2.SparseVectorFilterProto()
142        result.skip_zero_values = self._skip_zero_values
143        if self._filtered_items is not None:
144            result.filter_by_ids = True
145            result.filtered_ids[:] = sorted(t.id for t in self._filtered_items)
146        return result
147
148
149VariableFilter = SparseVectorFilter[variables.Variable]
150LinearConstraintFilter = SparseVectorFilter[linear_constraints.LinearConstraint]
151QuadraticConstraintFilter = SparseVectorFilter[
152    quadratic_constraints.QuadraticConstraint
153]