ortools.math_opt.python.sparse_containers

  1# Copyright 2010-2024 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 model
 22
 23
 24VarOrConstraintType = TypeVar(
 25    "VarOrConstraintType", model.Variable, model.LinearConstraint
 26)
 27
 28
 29def to_sparse_double_vector_proto(
 30    terms: Mapping[VarOrConstraintType, float]
 31) -> sparse_containers_pb2.SparseDoubleVectorProto:
 32    """Converts a sparse vector from proto to dict representation."""
 33    result = sparse_containers_pb2.SparseDoubleVectorProto()
 34    if terms:
 35        id_and_values = [(key.id, value) for (key, value) in terms.items()]
 36        id_and_values.sort()
 37        ids, values = zip(*id_and_values)
 38        result.ids[:] = ids
 39        result.values[:] = values
 40    return result
 41
 42
 43def to_sparse_int32_vector_proto(
 44    terms: Mapping[VarOrConstraintType, int]
 45) -> sparse_containers_pb2.SparseInt32VectorProto:
 46    """Converts a sparse vector from proto to dict representation."""
 47    result = sparse_containers_pb2.SparseInt32VectorProto()
 48    if terms:
 49        id_and_values = [(key.id, value) for (key, value) in terms.items()]
 50        id_and_values.sort()
 51        ids, values = zip(*id_and_values)
 52        result.ids[:] = ids
 53        result.values[:] = values
 54    return result
 55
 56
 57def parse_variable_map(
 58    proto: sparse_containers_pb2.SparseDoubleVectorProto, mod: model.Model
 59) -> Dict[model.Variable, float]:
 60    """Converts a sparse vector of variables from proto to dict representation."""
 61    result = {}
 62    for index, var_id in enumerate(proto.ids):
 63        result[mod.get_variable(var_id)] = proto.values[index]
 64    return result
 65
 66
 67def parse_linear_constraint_map(
 68    proto: sparse_containers_pb2.SparseDoubleVectorProto, mod: model.Model
 69) -> Dict[model.LinearConstraint, float]:
 70    """Converts a sparse vector of linear constraints from proto to dict representation."""
 71    result = {}
 72    for index, lin_con_id in enumerate(proto.ids):
 73        result[mod.get_linear_constraint(lin_con_id)] = proto.values[index]
 74    return result
 75
 76
 77class SparseVectorFilter(Generic[VarOrConstraintType]):
 78    """Restricts the variables or constraints returned in a sparse vector.
 79
 80    The default behavior is to return entries for all variables/constraints.
 81
 82    E.g. when requesting the solution to an optimization problem, use this class
 83    to restrict the variables that values are returned for.
 84
 85    Attributes:
 86      skip_zero_values: Do not include key value pairs with value zero.
 87      filtered_items: If not None, include only key value pairs these keys. Note
 88        that the empty set is different (don't return any keys) from None (return
 89        all keys).
 90    """
 91
 92    def __init__(
 93        self,
 94        *,
 95        skip_zero_values: bool = False,
 96        filtered_items: Optional[Iterable[VarOrConstraintType]] = None,
 97    ):
 98        self._skip_zero_values: bool = skip_zero_values
 99        self._filtered_items: Optional[Set[VarOrConstraintType]] = (
100            None if filtered_items is None else frozenset(filtered_items)
101        )  # pytype: disable=annotation-type-mismatch  # attribute-variable-annotations
102
103    @property
104    def skip_zero_values(self) -> bool:
105        return self._skip_zero_values
106
107    @property
108    def filtered_items(self) -> Optional[FrozenSet[VarOrConstraintType]]:
109        return (
110            self._filtered_items
111        )  # pytype: disable=bad-return-type  # attribute-variable-annotations
112
113    def to_proto(self):
114        """Returns an equivalent proto representation."""
115        result = sparse_containers_pb2.SparseVectorFilterProto()
116        result.skip_zero_values = self._skip_zero_values
117        if self._filtered_items is not None:
118            result.filter_by_ids = True
119            result.filtered_ids[:] = sorted(t.id for t in self._filtered_items)
120        return result
121
122
123VariableFilter = SparseVectorFilter[model.Variable]
124LinearConstraintFilter = SparseVectorFilter[model.LinearConstraint]