Coverage for datesat / symbolic_int / alpha_beta_int.py: 82.6%
192 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-10 23:47 +0000
1"""
2Alpha-beta DateSAT implementation using epoch-based conversion.
4This module implements the alpha-beta approach where dates are represented
5as (months, days) since an epoch.
6"""
8from typing import List, Tuple, Union
10from z3 import (
11 And,
12 ArithRef,
13 BoolRef,
14 CheckSatResult,
15 If,
16 Int,
17 IntVal,
18 ModelRef,
19 Not,
20 Optimize,
21 Or,
22 Solver,
23 sat,
24 unsat,
25)
26from ..core import Date, Period
27from .naive_int import (
28 eom_clamp,
29 days_in_month
30)
31from .epoch_days_int import (
32 ymd_from_days_since_epoch,
33 days_since_epoch_from_ymd
34)
36# -------------------------------
37# Alpha (months-since-epoch) helpers
38# Epoch month: 2000-03 (alpha = 0)
39# alpha = 12*y + m - (12*2000 + 3)
40# Inverse: let k = alpha + (12*2000 + 3), then
41# y = (k - 1) / 12
42# m = k - 12*y
43# -------------------------------
44# Python int epoch constants (for arithmetic outside Z3)
45_EPOCH_YEAR = 2000
46_EPOCH_MONTH = 3
47# Z3 epoch constants
48_EPOCH_LINEAR = _EPOCH_YEAR * 12 + _EPOCH_MONTH # 12*2000 + 3
49# Alpha bounds constants (months since epoch)
50_ALPHA_MIN = (1900 - _EPOCH_YEAR) * 12 + (3 - _EPOCH_MONTH) # -1200
51_ALPHA_MAX = (2100 - _EPOCH_YEAR) * 12 + (2 - _EPOCH_MONTH) # 1199
54def months_since_epoch_from_ym(y, m) -> ArithRef:
55 """Z3-pure: compute months-since-epoch (alpha) from year/month."""
56 return (y * 12 + m) - _EPOCH_LINEAR
59def ym_from_months_since_epoch(alpha) -> Tuple[ArithRef, ArithRef]:
60 """Z3-pure inverse: decode (year, month) from alpha months-since-epoch."""
61 k = alpha + _EPOCH_LINEAR
62 y = (k - IntVal(1)) / 12
63 m = k - y * 12
64 return y, m
67class DateVar:
68 """Symbolic date variable using alpha-beta representation.
70 alpha (months_var): months since epoch month 2000-03 (March 2000 = 0)
71 beta (beta_var): extra days within that month (0-based), so DOM = 1+beta
72 """
74 def __init__(self, name: str):
75 """Create a symbolic date variable."""
76 self.name = name
77 # Alpha: Z3 integer variable for months since epoch-month
78 self.months_var = Int(f"{name}_months")
79 # Beta: Z3 integer variable for extra days (0-based) within month
80 self.beta_var = Int(f"{name}_beta")
81 # Solver reference for adding bounds to intermediate dates (set after creation if needed)
82 self._solver = None
84 def __str__(self) -> str:
85 return f"DateVar({self.name})"
87 @property
88 def year(self) -> ArithRef:
89 """Get symbolic year component (decodes from months_var)."""
90 y, _ = ym_from_months_since_epoch(self.months_var)
91 return y
93 @property
94 def month(self) -> ArithRef:
95 """Get symbolic month component (decodes from months_var)."""
96 _, m = ym_from_months_since_epoch(self.months_var)
97 return m
99 @property
100 def day(self) -> ArithRef:
101 """Get symbolic day component (beta_var + 1, since beta is 0-based)."""
102 return self.beta_var + IntVal(1)
104 def to_concrete_date(self, model: ModelRef) -> Date:
105 """Convert Z3 model to concrete Date using (alpha, beta)."""
106 alpha_val = model.evaluate(self.months_var, model_completion=True).as_long()
107 beta_val = model.evaluate(self.beta_var, model_completion=True).as_long()
108 k = alpha_val + (_EPOCH_YEAR * 12 + _EPOCH_MONTH)
109 year = (k - 1) // 12
110 month = k - year * 12
111 day = beta_val + 1
112 try:
113 return Date(year, month, day)
114 except ValueError:
115 # Intermediate result went out of bounds - use unbounded date
116 return Date(year, month, day, bounded=False)
118 def __ge__(self, other) -> BoolRef:
119 """Support x >= date comparison."""
120 if isinstance(other, Date):
121 alpha_o = months_since_epoch_from_ym(
122 IntVal(other.year), IntVal(other.month)
123 )
124 beta_o = IntVal(other.day - 1)
126 return Or(
127 self.months_var > alpha_o,
128 And(self.months_var == alpha_o, self.beta_var >= beta_o),
129 )
130 elif isinstance(other, DateVar): 130 ↛ 138line 130 didn't jump to line 138 because the condition on line 130 was always true
131 return Or(
132 self.months_var > other.months_var,
133 And(
134 self.months_var == other.months_var, self.beta_var >= other.beta_var
135 ),
136 )
137 else:
138 raise TypeError(f"Cannot compare DateVar with {type(other)}")
140 def __le__(self, other) -> BoolRef:
141 """Support x <= date comparison."""
142 if isinstance(other, Date):
143 alpha_o = months_since_epoch_from_ym(
144 IntVal(other.year), IntVal(other.month)
145 )
146 beta_o = IntVal(other.day - 1)
148 return Or(
149 self.months_var < alpha_o,
150 And(self.months_var == alpha_o, self.beta_var <= beta_o),
151 )
152 elif isinstance(other, DateVar): 152 ↛ 160line 152 didn't jump to line 160 because the condition on line 152 was always true
153 return Or(
154 self.months_var < other.months_var,
155 And(
156 self.months_var == other.months_var, self.beta_var <= other.beta_var
157 ),
158 )
159 else:
160 raise TypeError(f"Cannot compare DateVar with {type(other)}")
162 def __lt__(self, other) -> BoolRef:
163 """Support x < date comparison."""
164 if isinstance(other, (Date, DateVar)): 164 ↛ 167line 164 didn't jump to line 167 because the condition on line 164 was always true
165 return Not(self.__ge__(other))
166 else:
167 raise TypeError(f"Cannot compare DateVar with {type(other)}")
169 def __gt__(self, other) -> BoolRef:
170 """Support x > date comparison."""
171 if isinstance(other, (Date, DateVar)): 171 ↛ 174line 171 didn't jump to line 174 because the condition on line 171 was always true
172 return Not(self.__le__(other))
173 else:
174 raise TypeError(f"Cannot compare DateVar with {type(other)}")
176 def __eq__(self, other) -> BoolRef:
177 """Support x == date comparison."""
178 if isinstance(other, Date):
179 alpha_o = months_since_epoch_from_ym(
180 IntVal(other.year), IntVal(other.month)
181 )
182 beta_o = IntVal(other.day - 1)
184 return And(self.months_var == alpha_o, self.beta_var == beta_o)
185 elif isinstance(other, DateVar): 185 ↛ 190line 185 didn't jump to line 190 because the condition on line 185 was always true
186 return And(
187 self.months_var == other.months_var, self.beta_var == other.beta_var
188 )
189 else:
190 raise TypeError(f"Cannot compare DateVar with {type(other)}")
192 def __ne__(self, other) -> BoolRef:
193 """Support x != date comparison using ordinal arithmetic."""
194 if isinstance(other, (Date, DateVar)): 194 ↛ 197line 194 didn't jump to line 197 because the condition on line 194 was always true
195 return Not(self.__eq__(other))
196 else:
197 raise TypeError(f"Cannot compare DateVar with {type(other)}")
199 def _add_bounds(self) -> None:
200 """Add date validation bounds to this DateVar if solver is available."""
201 if self._solver is None: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 return
204 # Alpha bounds: months since 2000-03
205 # 1900-03 => -1200, 2100-02 => 1199
206 self._solver.add(self.months_var >= IntVal(_ALPHA_MIN))
207 self._solver.add(self.months_var <= IntVal(_ALPHA_MAX))
209 # Beta bounds depend on month length: 0 <= beta < days_in_month(y,m)
210 y, m = ym_from_months_since_epoch(self.months_var)
211 self._solver.add(self.beta_var >= 0)
212 self._solver.add(self.beta_var < days_in_month(y, m))
214 def __add__(self, other) -> "DateVar":
215 """DateVar + Period using alpha for Y/M and beta for D.
216 Steps:
217 - Fast path 1: If days-only period and result stays within month, simple addition.
218 - Fast path 2: If days-only period but crosses month boundary, normalize via ordinal.
219 - Full path: Add months to alpha, EOM clamp beta, add days to beta.
220 If result stays within month, simple addition; otherwise normalize via ordinal.
221 """
222 if isinstance(other, Period): 222 ↛ 309line 222 didn't jump to line 309 because the condition on line 222 was always true
223 result = DateVar(f"{self.name}_plus")
224 months_delta = IntVal(other.years * 12 + other.months)
225 days_delta = IntVal(other.days)
227 # Fast path: days-only period
228 if other.years == 0 and other.months == 0:
229 # Get current month info for within-month check
230 y0, m0 = ym_from_months_since_epoch(self.months_var)
231 dim0 = days_in_month(y0, m0)
233 # Add days directly to beta
234 new_beta = self.beta_var + days_delta
236 # Check if result stays within same month
237 stays_in_month = And(new_beta >= IntVal(0), new_beta < dim0)
239 # Within-month fast path: simple addition
240 alpha_within = self.months_var
241 beta_within = new_beta
243 # Fallback: normalize via epoch conversion (handles month overflow)
244 d0 = new_beta + IntVal(1) # Convert 0-based beta to 1-based day
245 # Convert Y/M/D to epoch days, then back to Y/M/D (normalizes across boundaries)
246 epoch_days = days_since_epoch_from_ymd(y0, m0, d0)
247 y2, m2, d2 = ymd_from_days_since_epoch(epoch_days)
248 months_ordinal = months_since_epoch_from_ym(y2, m2)
249 beta_ordinal = d2 - IntVal(1)
251 # Select result based on within-month condition
252 result.months_var = If(stays_in_month, alpha_within, months_ordinal)
253 result.beta_var = If(stays_in_month, beta_within, beta_ordinal)
255 # Add bounds to intermediate result
256 result._solver = self._solver
257 result._add_bounds()
258 return result
260 # Full path: Add to alpha and beta directly, then normalize
261 # Step 1: Add months to alpha directly
262 alpha1 = self.months_var + months_delta
263 y1, m1 = ym_from_months_since_epoch(alpha1)
264 dim1 = days_in_month(y1, m1)
266 # Step 2: EOM clamp beta (needed when adding months - e.g., Jan 31 + 1 month = Feb 28/29)
267 d1 = eom_clamp(y1, m1, self.beta_var + IntVal(1))
268 beta1 = d1 - IntVal(1) # Convert back to 0-based beta
270 # Fast path: years/months-only period (no days)
271 if other.days == 0:
272 # No day addition needed - we're done!
273 result.months_var = alpha1
274 result.beta_var = beta1
276 result._solver = self._solver
277 result._add_bounds()
278 return result
280 else:
281 # Step 3: Add days to beta directly
282 new_beta = beta1 + days_delta
284 # Check if result stays within same month
285 stays_in_month = And(new_beta >= IntVal(0), new_beta < dim1)
287 # Within-month fast path: simple addition
288 alpha_within = alpha1
289 beta_within = new_beta
291 # Fallback: Normalize via epoch conversion (handles all overflow cases)
292 d_temp = new_beta + IntVal(1) # Convert 0-based beta to 1-based day
293 # Convert Y/M/D to epoch days, then back to Y/M/D (normalizes across boundaries)
294 epoch_days = days_since_epoch_from_ymd(y1, m1, d_temp)
295 y2, m2, d2 = ymd_from_days_since_epoch(epoch_days)
297 months_ordinal = months_since_epoch_from_ym(y2, m2)
298 beta_ordinal = d2 - IntVal(1)
300 # Select result based on within-month condition
301 result.months_var = If(stays_in_month, alpha_within, months_ordinal)
302 result.beta_var = If(stays_in_month, beta_within, beta_ordinal)
304 # Add bounds to intermediate result
305 result._solver = self._solver
306 result._add_bounds()
307 return result
308 else:
309 raise TypeError(f"Cannot add {type(other)} to DateVar")
311 def __sub__(self, other) -> "DateVar":
312 """DateVar - Period implemented as DateVar + (-Period)."""
313 if isinstance(other, Period): 313 ↛ 317line 313 didn't jump to line 317 because the condition on line 313 was always true
314 neg = Period(-other.years, -other.months, -other.days)
315 return self.__add__(neg)
316 else:
317 raise TypeError(f"Cannot subtract {type(other)} from DateVar")
320class AlphaBetaSolver:
321 """Alpha-beta date constraint solver using epoch-based conversion."""
323 def __init__(self, timeout_ms=600000, use_maxsat=False):
324 """Initialize the solver with timeout.
326 Args:
327 timeout_ms: Timeout in milliseconds (default: 60 seconds)
328 use_maxsat: If True, use MaxSAT optimization with soft constraints
329 """
330 self.use_maxsat = use_maxsat
331 if use_maxsat: 331 ↛ 332line 331 didn't jump to line 332 because the condition on line 331 was never true
332 self.solver = Optimize()
333 else:
334 self.solver = Solver()
335 self.solver.set("timeout", timeout_ms)
336 self.date_vars = {}
337 self.constraints = []
338 self.timeout_ms = timeout_ms
340 def add_date_var(self, name: str) -> DateVar:
341 """Add a symbolic date variable with basic constraints."""
342 date_var = DateVar(name)
343 date_var._solver = self.solver
344 self.date_vars[name] = date_var
346 # Add bounds using _add_bounds method
347 date_var._add_bounds()
348 return date_var
350 def add_constraint(self, constraint: BoolRef) -> None:
351 """Add a constraint to the solver."""
352 self.constraints.append(constraint)
353 self.solver.add(constraint)
355 def check(self) -> CheckSatResult:
356 """Check if constraints are satisfiable."""
357 return self.solver.check()
359 def model(self) -> ModelRef:
360 """Get the model if satisfiable."""
361 return self.solver.model()
363 def get_concrete_dates(self, model: ModelRef) -> dict:
364 """Get concrete dates from the model."""
365 return {
366 name: var.to_concrete_date(model) for name, var in self.date_vars.items()
367 }
369 def solve(self) -> Union[bool, dict]:
370 """Solve the constraints."""
371 # Add MaxSAT soft constraints if enabled
372 if self.use_maxsat: 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true
373 from datetime import date
375 today = date.today()
376 # Calculate months since epoch for today
377 today_months = (today.year - _EPOCH_YEAR) * 12 + (
378 today.month - _EPOCH_MONTH
379 )
381 # Convert years to months
382 months_50_years = 50 * 12 # 600 months
383 months_10_years = 10 * 12 # 120 months
385 # Add soft constraints for each date variable
386 for name, date_var in self.date_vars.items():
387 # High weight: today ± 50 years
388 within_50_years = And(
389 date_var.months_var >= IntVal(today_months - months_50_years),
390 date_var.months_var <= IntVal(today_months + months_50_years),
391 )
392 self.solver.add_soft(within_50_years, weight=100)
394 # Low weight: today ± 10 years
395 within_10_years = And(
396 date_var.months_var >= IntVal(today_months - months_10_years),
397 date_var.months_var <= IntVal(today_months + months_10_years),
398 )
399 self.solver.add_soft(within_10_years, weight=10)
401 result = self.check()
402 if result == sat:
403 model = self.model()
404 return {
405 "status": "sat",
406 "dates": self.get_concrete_dates(model),
407 }
408 elif result == unsat: 408 ↛ 412line 408 didn't jump to line 412 because the condition on line 408 was always true
409 return {"status": "unsat", "dates": {}}
410 else:
411 # result == unknown (timeout or resource limit)
412 return {"status": "timeout", "dates": {}}
414 def to_smt2(self) -> str:
415 """Return the current problem in SMT-LIB v2 format."""
416 return self.solver.to_smt2()
418 def get_assertions(self) -> List[BoolRef]:
419 """Return the list of current Z3 assertions (BoolRef)."""
420 return list(self.solver.assertions())