mcp-solver

""" Cardinality constraint templates for PySAT. This module provides template functions for common PySAT patterns related to cardinality constraints. """ import sys from typing import Any # Import PySAT but protect against failure try: from pysat.card import CardEnc, EncType from pysat.formula import CNF except ImportError: print("PySAT solver not found. Install with: pip install python-sat") sys.exit(1) def at_most_k( variables: list[int], k: int, encoding: EncType = EncType.seqcounter ) -> list[list[int]]: """ Create clauses ensuring at most k variables can be true. This is a robust implementation that handles edge cases better and falls back to direct encoding when needed. Args: variables: List of variable IDs k: Maximum number of variables that can be true encoding: Encoding type to try (will fall back to direct encoding if needed) Returns: List of clauses representing the constraint """ # Handle edge cases if k >= len(variables): return [] # No constraint needed # Use direct encoding for small sets (more reliable across versions) if len(variables) <= 20: clauses = [] import itertools for combo in itertools.combinations(variables, k + 1): clauses.append([-v for v in combo]) return clauses # For larger sets, try PySAT's encoding but handle failures try: return CardEnc.atmost(variables, k, encoding=encoding).clauses except Exception: # Fall back to direct encoding if PySAT's fails # This will be slower but more reliable clauses = [] import itertools for combo in itertools.combinations(variables, k + 1): clauses.append([-v for v in combo]) return clauses def at_least_k( variables: list[int], k: int, encoding: EncType = EncType.seqcounter ) -> list[list[int]]: """ Create clauses ensuring at least k variables must be true. Uses De Morgan's law with the robust at_most_k implementation. Args: variables: List of variable IDs k: Minimum number of variables that must be true encoding: Encoding type to try (will fall back to direct encoding if needed) Returns: List of clauses representing the constraint """ # We use De Morgan's law: at least k of variables = at most (n-k) of negated variables # For small instances, we can directly call at_most_k with negated vars if len(variables) <= 100: # Arbitrary threshold to avoid excessive computation neg_vars = [-v for v in variables] return at_most_k(neg_vars, len(variables) - k, encoding=encoding) # For larger instances, try PySAT's direct implementation try: return CardEnc.atleast(variables, k, encoding=encoding).clauses except Exception: # Fall back to at_most_k with De Morgan's law neg_vars = [-v for v in variables] return at_most_k(neg_vars, len(variables) - k, encoding=encoding) def exactly_k( variables: list[int], k: int, encoding: EncType = EncType.seqcounter ) -> CNF: """ Create CNF formula enforcing exactly k variables are true. Args: variables: List of variable IDs k: Exact number of variables that must be true encoding: Encoding type (default: sequential counter) Returns: CNF formula with the constraint """ return CardEnc.equals(variables, k, encoding=encoding) def one_hot_encoding( variables: list[int], encoding: EncType = EncType.seqcounter ) -> CNF: """ Create CNF formula enforcing exactly one variable is true. Args: variables: List of variable IDs encoding: Encoding type (default: sequential counter) Returns: CNF formula with the constraint """ return CardEnc.equals(variables, 1, encoding=encoding) def add_cardinality_constraint( formula: CNF, constraint_type: str, variables: list[int], k: int, encoding: EncType = EncType.seqcounter, ) -> CNF: """ Add a cardinality constraint to an existing CNF formula. Args: formula: Existing CNF formula constraint_type: Type of constraint ("atmost", "atleast", "exactly") variables: List of variable IDs k: Parameter value (max/min/exact number) encoding: Encoding type (default: sequential counter) Returns: CNF formula with the added constraint """ # Create constraint formula if constraint_type.lower() == "atmost": constraint = CardEnc.atmost(variables, k, encoding=encoding) elif constraint_type.lower() == "atleast": constraint = CardEnc.atleast(variables, k, encoding=encoding) elif constraint_type.lower() in ["exactly", "equals"]: constraint = CardEnc.equals(variables, k, encoding=encoding) else: raise ValueError(f"Unknown constraint type: {constraint_type}") # Add constraint clauses to the formula for clause in constraint.clauses: formula.append(clause) return formula def create_balanced_partitioning( variables: list[int], num_parts: int = 2, encoding: EncType = EncType.seqcounter ) -> dict[str, Any]: """ Create a balanced partitioning of variables. Args: variables: List of variable IDs num_parts: Number of partitions (default: 2) encoding: Encoding type (default: sequential counter) Returns: Dictionary with formulas and information about the partitioning """ if num_parts < 2: raise ValueError("Number of partitions must be at least 2") n = len(variables) # Create result dictionary result = {"formula": CNF(), "part_variables": [], "constraints": []} # Create part indicator variables # p_i_j = 1 means variable i is in part j top_var = max(abs(var) for var in variables) part_vars = [] for i, var in enumerate(variables): part_var_row = [] for j in range(num_parts): part_var = top_var + 1 + i * num_parts + j part_var_row.append(part_var) part_vars.append(part_var_row) # Each variable must be in exactly one part one_hot = CardEnc.equals(part_var_row, 1, encoding=encoding) result["constraints"].append(("one_hot", var, one_hot)) for clause in one_hot.clauses: result["formula"].append(clause) result["part_variables"] = part_vars # Calculate target size for each part target_size = n // num_parts remainder = n % num_parts # Create part size constraints for j in range(num_parts): # Variables indicating if each element is in part j part_j_vars = [part_vars[i][j] for i in range(n)] # Part size constraint part_size = target_size + (1 if j < remainder else 0) part_size_constraint = CardEnc.equals(part_j_vars, part_size, encoding=encoding) result["constraints"].append(("part_size", j, part_size_constraint)) # Add constraints to the formula for clause in part_size_constraint.clauses: result["formula"].append(clause) return result