Coverage for datesat / symbolic_int / alpha_beta_table_int.py: 86.4%
251 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-10 23:47 +0000
1"""
2Alpha-beta-table DateSAT using a 4-year (48-month) table.
4Representation:
5- alpha (months_var): months since epoch month 2000-03 (March 2000 = 0)
6- beta (beta_var): 0-based day index within that month (DOM = beta + 1)
8We avoid full ordinal decode by using a 48-month DIM/DBM table.
9"""
11from typing import List, Tuple, Union
13from z3 import (
14 And,
15 ArithRef,
16 BoolRef,
17 CheckSatResult,
18 If,
19 Int,
20 IntSort,
21 IntVal,
22 K,
23 ModelRef,
24 Not,
25 Optimize,
26 Or,
27 Select,
28 Solver,
29 Store,
30 sat,
31 unsat,
32)
33from ..core import Date, Period
35_EPOCH_YEAR = 2000
36_EPOCH_MONTH = 3
37# Linearized epoch month as a Z3 Int numeral
38_EPOCH_LINEAR = IntVal(_EPOCH_YEAR * 12 + _EPOCH_MONTH)
39_FOUR_YEAR_MONTHS = 48
40_FOUR_YEAR_DAYS = 1461
41# Alpha bounds constants (months since epoch)
42_ALPHA_MIN = (1900 - _EPOCH_YEAR) * 12 + (3 - _EPOCH_MONTH) # -1200
43_ALPHA_MAX = (2100 - _EPOCH_YEAR) * 12 + (2 - _EPOCH_MONTH) # 1199
46def eom_clamp(dim, beta) -> ArithRef:
47 return If(
48 beta < IntVal(0), IntVal(0), If(beta > dim - IntVal(1), dim - IntVal(1), beta)
49 )
52def _is_leap_py(y: int) -> bool:
53 return (y % 4 == 0 and y % 100 != 0) or (y % 400 == 0)
56def _days_in_month_py(y: int, m: int) -> int:
57 if m == 2:
58 return 29 if _is_leap_py(y) else 28
59 if m in (4, 6, 9, 11):
60 return 30
61 return 31
64def _add_months(y: int, m: int, delta: int) -> tuple[int, int]:
65 total = (y * 12 + m) + delta
66 y2 = (total - 1) // 12
67 m2 = total - y2 * 12
68 return y2, m2
71def build_dim_dbm_48_from_epoch() -> tuple[list[int], list[int]]:
72 dim: list[int] = [0] * _FOUR_YEAR_MONTHS
73 dbm: list[int] = [0] * _FOUR_YEAR_MONTHS
74 y, m = _EPOCH_YEAR, _EPOCH_MONTH
75 cum = 0
76 for i in range(_FOUR_YEAR_MONTHS):
77 dbm[i] = cum
78 d = _days_in_month_py(y, m)
79 dim[i] = d
80 cum += d
81 y, m = _add_months(y, m, 1)
82 assert cum == _FOUR_YEAR_DAYS
83 return dim, dbm
86def const_array(values: list[int]):
87 a = K(IntSort(), IntVal(0))
88 for i, v in enumerate(values):
89 a = Store(a, IntVal(i), IntVal(v))
90 return a
93_DIM48_LIST_PY, _DBM48_LIST_PY = build_dim_dbm_48_from_epoch()
94_DIM48_LIST = const_array(_DIM48_LIST_PY)
95_DBM48_LIST = const_array(_DBM48_LIST_PY)
98def mod48(x):
99 return x % IntVal(_FOUR_YEAR_MONTHS)
102def alpha_to_abs_month(alpha):
103 return alpha + _EPOCH_LINEAR
106def months_since_epoch_from_ym(y, m):
107 return (y * IntVal(12) + m) - _EPOCH_LINEAR
110class DateVar:
111 """Symbolic date variable using alpha-beta representation.
113 alpha (months_var): months since epoch month 2000-03 (March 2000 = 0)
114 beta (beta_var): extra days within that month (0-based), so DOM = 1+beta
115 """
117 def __init__(self, name: str):
118 """Create a symbolic date variable."""
119 self.name = name
120 # Alpha: Z3 integer variable for months since epoch-month
121 self.months_var = Int(f"{name}_months")
122 # Beta: Z3 integer variable for extra days (0-based) within month
123 self.beta_var = Int(f"{name}_beta")
124 # Solver reference for adding bounds to intermediate dates (set after creation if needed)
125 self._solver = None
127 def __str__(self) -> str:
128 return f"DateVar({self.name})"
130 @property
131 def year(self) -> ArithRef:
132 """Get symbolic year component (decodes from months_var)."""
133 k = self.months_var + _EPOCH_LINEAR
134 y = (k - IntVal(1)) / IntVal(12)
135 return y
137 @property
138 def month(self) -> ArithRef:
139 """Get symbolic month component (decodes from months_var)."""
140 k = self.months_var + _EPOCH_LINEAR
141 y = (k - IntVal(1)) / IntVal(12)
142 m = k - y * IntVal(12)
143 return m
145 @property
146 def day(self) -> ArithRef:
147 """Get symbolic day component (beta_var + 1, since beta is 0-based)."""
148 return self.beta_var + IntVal(1)
150 def to_concrete_date(self, model: ModelRef) -> Date:
151 """Convert Z3 model to concrete Date using (alpha, beta)."""
152 alpha_val = model.evaluate(self.months_var, model_completion=True).as_long()
153 beta_val = model.evaluate(self.beta_var, model_completion=True).as_long()
154 k = alpha_val + (_EPOCH_YEAR * 12 + _EPOCH_MONTH)
155 year = (k - 1) // 12
156 month = k - year * 12
157 day = beta_val + 1
158 try:
159 return Date(year, month, day)
160 except ValueError:
161 return Date(year, month, day, bounded=False)
163 def _add_bounds(self) -> None:
164 """Add date validation bounds to this DateVar if solver is available."""
165 if self._solver is None: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true
166 return
168 # Alpha bounds: months since 2000-03
169 # 1900-03 => -1200, 2100-02 => 1199
170 self._solver.add(self.months_var >= IntVal(_ALPHA_MIN))
171 self._solver.add(self.months_var <= IntVal(_ALPHA_MAX))
173 # Beta bounds: 0 <= beta < DIM
174 idx = mod48(self.months_var)
175 dim = Select(_DIM48_LIST, idx)
176 self._solver.add(And(self.beta_var >= IntVal(0), self.beta_var < dim))
178 def __ge__(self, other) -> BoolRef:
179 """Support x >= date comparison."""
180 if isinstance(other, Date):
181 alpha_o = months_since_epoch_from_ym(
182 IntVal(other.year), IntVal(other.month)
183 )
184 beta_o = IntVal(other.day - 1)
186 return Or(
187 self.months_var > alpha_o,
188 And(self.months_var == alpha_o, self.beta_var >= beta_o),
189 )
190 elif isinstance(other, DateVar): 190 ↛ 198line 190 didn't jump to line 198 because the condition on line 190 was always true
191 return Or(
192 self.months_var > other.months_var,
193 And(
194 self.months_var == other.months_var, self.beta_var >= other.beta_var
195 ),
196 )
197 else:
198 raise TypeError(f"Cannot compare DateVar with {type(other)}")
200 def __le__(self, other) -> BoolRef:
201 """Support x <= date comparison."""
202 if isinstance(other, Date):
203 alpha_o = months_since_epoch_from_ym(
204 IntVal(other.year), IntVal(other.month)
205 )
206 beta_o = IntVal(other.day - 1)
208 return Or(
209 self.months_var < alpha_o,
210 And(self.months_var == alpha_o, self.beta_var <= beta_o),
211 )
212 elif isinstance(other, DateVar): 212 ↛ 220line 212 didn't jump to line 220 because the condition on line 212 was always true
213 return Or(
214 self.months_var < other.months_var,
215 And(
216 self.months_var == other.months_var, self.beta_var <= other.beta_var
217 ),
218 )
219 else:
220 raise TypeError(f"Cannot compare DateVar with {type(other)}")
222 def __lt__(self, other) -> BoolRef:
223 """Support x < date comparison."""
224 if isinstance(other, (Date, DateVar)): 224 ↛ 227line 224 didn't jump to line 227 because the condition on line 224 was always true
225 return Not(self.__ge__(other))
226 else:
227 raise TypeError(f"Cannot compare DateVar with {type(other)}")
229 def __gt__(self, other) -> BoolRef:
230 """Support x > date comparison."""
231 if isinstance(other, (Date, DateVar)): 231 ↛ 234line 231 didn't jump to line 234 because the condition on line 231 was always true
232 return Not(self.__le__(other))
233 else:
234 raise TypeError(f"Cannot compare DateVar with {type(other)}")
236 def __eq__(self, other) -> BoolRef:
237 """Support x == date comparison."""
238 if isinstance(other, Date):
239 alpha_o = months_since_epoch_from_ym(
240 IntVal(other.year), IntVal(other.month)
241 )
242 beta_o = IntVal(other.day - 1)
244 return And(self.months_var == alpha_o, self.beta_var == beta_o)
245 elif isinstance(other, DateVar): 245 ↛ 250line 245 didn't jump to line 250 because the condition on line 245 was always true
246 return And(
247 self.months_var == other.months_var, self.beta_var == other.beta_var
248 )
249 else:
250 raise TypeError(f"Cannot compare DateVar with {type(other)}")
252 def __ne__(self, other) -> BoolRef:
253 """Support x != date comparison using ordinal arithmetic."""
254 if isinstance(other, (Date, DateVar)): 254 ↛ 257line 254 didn't jump to line 257 because the condition on line 254 was always true
255 return Not(self.__eq__(other))
256 else:
257 raise TypeError(f"Cannot compare DateVar with {type(other)}")
259 def __add__(self, other) -> "DateVar":
260 if isinstance(other, Period): 260 ↛ 369line 260 didn't jump to line 369 because the condition on line 260 was always true
261 result = DateVar(f"{self.name}_plus")
262 months_delta = IntVal(other.years * 12 + other.months)
263 days_delta = IntVal(other.days)
265 # Fast path: days-only period (skip month shift)
266 if other.years == 0 and other.months == 0:
267 # Check if result stays within same month
268 alpha1 = self.months_var
269 idx1 = mod48(alpha1)
270 dim1 = Select(_DIM48_LIST, idx1)
271 beta1 = eom_clamp(dim1, self.beta_var)
273 # Within-month fast path: if beta1 + days_delta stays in [0, dim1)
274 new_beta = beta1 + days_delta
275 stays_in_month = And(new_beta >= IntVal(0), new_beta < dim1)
277 # Within-month: simple addition
278 alpha_within = alpha1
279 beta_within = new_beta
281 # Fallback: use full table lookup (when days cross month boundary)
282 base48 = Select(_DBM48_LIST, idx1) + beta1
283 total = base48 + days_delta
285 q0 = total / IntVal(_FOUR_YEAR_DAYS)
286 r0 = total % IntVal(_FOUR_YEAR_DAYS)
288 # Compute idx2 by scanning all 48 months with century correction at target
289 best = IntVal(0)
290 for i in range(1, _FOUR_YEAR_MONTHS):
291 dbm_i_corr = Select(_DBM48_LIST, IntVal(i))
292 best = If(r0 >= dbm_i_corr, IntVal(i), best)
294 idx2 = best
295 diff2 = idx2 - idx1
296 beta2 = r0 - (Select(_DBM48_LIST, idx2))
298 dim2 = Select(_DIM48_LIST, idx2)
299 carry = If(beta2 >= dim2, IntVal(1), IntVal(0))
301 alpha_ordinal = alpha1 + q0 * IntVal(_FOUR_YEAR_MONTHS) + diff2 + carry
302 beta_ordinal = If(carry == IntVal(1), beta2 - dim2, beta2)
304 # Select result based on within-month condition
305 result.months_var = If(stays_in_month, alpha_within, alpha_ordinal)
306 result.beta_var = If(stays_in_month, beta_within, beta_ordinal)
307 # Add bounds to intermediate result
308 result._solver = self._solver
309 result._add_bounds()
310 return result
312 # Full path: with month shift
313 alpha1 = self.months_var + months_delta
314 idx1 = mod48(alpha1)
315 dim1 = Select(_DIM48_LIST, idx1)
316 beta1 = eom_clamp(dim1, self.beta_var)
318 # Fast path: years/months-only period (no days)
319 if other.days == 0:
320 # No day addition needed - we're done!
321 result.months_var = alpha1
322 result.beta_var = beta1
324 result._solver = self._solver
325 result._add_bounds()
326 return result
327 else:
328 # Within-month fast path: if adding days stays in same month
329 new_beta = beta1 + days_delta
330 stays_in_month = And(new_beta >= IntVal(0), new_beta < dim1)
332 # Within-month: simple addition
333 alpha_within = alpha1
334 beta_within = new_beta
336 # Full table lookup path
337 base48 = Select(_DBM48_LIST, idx1) + beta1
338 total = base48 + days_delta
340 q0 = total / IntVal(_FOUR_YEAR_DAYS)
341 r0 = total % IntVal(_FOUR_YEAR_DAYS)
343 # Compute idx2 by scanning all 48 months with century correction at target
344 best = IntVal(0)
345 for i in range(1, _FOUR_YEAR_MONTHS):
346 dbm_i_corr = Select(_DBM48_LIST, IntVal(i))
347 best = If(r0 >= dbm_i_corr, IntVal(i), best)
349 idx2 = best
350 diff2 = idx2 - idx1
351 beta2 = r0 - (Select(_DBM48_LIST, idx2))
353 # End-of-month overflow carry: if beta2 equals/exceeds the month length,
354 # advance one month and wrap beta into the next month.
355 dim2 = Select(_DIM48_LIST, idx2)
356 carry = If(beta2 >= dim2, IntVal(1), IntVal(0))
358 alpha_ordinal = alpha1 + q0 * IntVal(_FOUR_YEAR_MONTHS) + diff2 + carry
359 beta_ordinal = If(carry == IntVal(1), beta2 - dim2, beta2)
361 # Select result based on within-month condition
362 result.months_var = If(stays_in_month, alpha_within, alpha_ordinal)
363 result.beta_var = If(stays_in_month, beta_within, beta_ordinal)
364 # Add bounds to intermediate result
365 result._solver = self._solver
366 result._add_bounds()
367 return result
368 else:
369 raise TypeError(f"Cannot add {type(other)} to DateVar")
371 def __sub__(self, other) -> "DateVar":
372 """DateVar - Period implemented as DateVar + (-Period)."""
373 if isinstance(other, Period): 373 ↛ 377line 373 didn't jump to line 377 because the condition on line 373 was always true
374 neg = Period(-other.years, -other.months, -other.days)
375 return self.__add__(neg)
376 else:
377 raise TypeError(f"Cannot subtract {type(other)} from DateVar")
380class AlphaBetaTableSolver:
381 """Alpha-beta date constraint solver using epoch-based conversion."""
383 def __init__(self, timeout_ms=600000, use_maxsat=False):
384 """Initialize the solver with timeout.
386 Args:
387 timeout_ms: Timeout in milliseconds (default: 60 seconds)
388 use_maxsat: If True, use MaxSAT optimization with soft constraints
389 """
390 self.use_maxsat = use_maxsat
391 if use_maxsat: 391 ↛ 392line 391 didn't jump to line 392 because the condition on line 391 was never true
392 self.solver = Optimize()
393 else:
394 self.solver = Solver()
395 self.solver.set("timeout", timeout_ms)
396 self.date_vars = {}
397 self.constraints = []
398 self.timeout_ms = timeout_ms
400 def add_date_var(self, name: str) -> DateVar:
401 """Add a symbolic date variable with basic constraints."""
402 date_var = DateVar(name)
403 date_var._solver = self.solver
404 self.date_vars[name] = date_var
406 # Add bounds using _add_bounds method
407 date_var._add_bounds()
408 return date_var
410 def add_constraint(self, constraint: BoolRef) -> None:
411 """Add a constraint to the solver."""
412 self.constraints.append(constraint)
413 self.solver.add(constraint)
415 def check(self) -> CheckSatResult:
416 """Check if constraints are satisfiable."""
417 return self.solver.check()
419 def model(self) -> ModelRef:
420 """Get the model if satisfiable."""
421 return self.solver.model()
423 def get_concrete_dates(self, model: ModelRef) -> dict:
424 """Get concrete dates from the model."""
425 return {
426 name: var.to_concrete_date(model) for name, var in self.date_vars.items()
427 }
429 def solve(self) -> Union[bool, dict]:
430 """Solve the constraints."""
431 # Add MaxSAT soft constraints if enabled
432 if self.use_maxsat: 432 ↛ 433line 432 didn't jump to line 433 because the condition on line 432 was never true
433 from datetime import date
435 today = date.today()
436 # Calculate months since epoch for today
437 today_months = (today.year - _EPOCH_YEAR) * 12 + (
438 today.month - _EPOCH_MONTH
439 )
441 # Convert years to months
442 months_50_years = 50 * 12 # 600 months
443 months_10_years = 10 * 12 # 120 months
445 # Add soft constraints for each date variable
446 for name, date_var in self.date_vars.items():
447 # High weight: today ± 50 years
448 within_50_years = And(
449 date_var.months_var >= IntVal(today_months - months_50_years),
450 date_var.months_var <= IntVal(today_months + months_50_years),
451 )
452 self.solver.add_soft(within_50_years, weight=100)
454 # Low weight: today ± 10 years
455 within_10_years = And(
456 date_var.months_var >= IntVal(today_months - months_10_years),
457 date_var.months_var <= IntVal(today_months + months_10_years),
458 )
459 self.solver.add_soft(within_10_years, weight=10)
461 result = self.check()
462 if result == sat:
463 model = self.model()
464 return {
465 "status": "sat",
466 "dates": self.get_concrete_dates(model),
467 }
468 elif result == unsat: 468 ↛ 472line 468 didn't jump to line 472 because the condition on line 468 was always true
469 return {"status": "unsat", "dates": {}}
470 else:
471 # result == unknown (timeout or resource limit)
472 return {"status": "timeout", "dates": {}}
474 def to_smt2(self) -> str:
475 """Return the current problem in SMT-LIB v2 format."""
476 return self.solver.to_smt2()
478 def get_assertions(self) -> List[BoolRef]:
479 """Return the list of current Z3 assertions (BoolRef)."""
480 return list(self.solver.assertions())