ortools.math_opt.python.sparse_containers

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