relationalai 0.12.7__py3-none-any.whl → 0.12.9__py3-none-any.whl

relationalai/semantics/reasoners/optimization/solvers_dev.py

@@ -6,7 +6,7 @@ import time
 
 from relationalai.semantics.snowflake import Table
 from relationalai.semantics import std
-from relationalai.semantics.internal import internal as b # TODO(coey) change b name or remove b.?
+from relationalai.semantics.internal import internal as b
 from relationalai.semantics.rel.executor import RelExecutor
 from relationalai.semantics.lqp.executor import LQPExecutor
 from relationalai.tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
@@ -57,18 +57,29 @@ class SolverModelDev:
         # self.objective_values = model.Relationship(f"point {{i:int}} has objective value {{val:{data_type}}}", short_name=_name("objective_values"))
         # self.primal_statuses = model.Relationship("point {i:int} has primal status {status:str}", short_name=_name("primal_statuses"))
 
+        self._model_info = {
+            "num_variables": self.variables,
+            "num_constraints": self.constraints,
+            "num_min_objectives": self.min_objectives,
+            "num_max_objectives": self.max_objectives,
+        }
+
     # TODO(coey) assert that it is a property? not just a relationship.
-    def solve_for(self, expr: b.Relationship | b.Fragment, populate: bool = True, **kwargs):
-        where = []
+    def solve_for(self, expr, where: list = [], populate: bool = True, **kwargs):
         if isinstance(expr, b.Fragment):
-            assert expr._select and len(expr._select) == 1 and expr._where, "Fragment input for `solve_for` must have exactly one select and a where clause"
-            rel = expr._select[0]
-            where.extend(expr._where)
+            # TODO(coey) remove in future
+            raise ValueError("select fragment argument to `solve_for` is deprecated; instead use `where = [conditions...]` kwarg to specify optional grounding conditions")
+        elif isinstance(expr, b.Expression):
+            # must be of the form rel(a, ..., x) where the last element is the decision variable
+            rel = expr._op
+            assert isinstance(rel, b.Relationship)
+            params = expr._params
         elif isinstance(expr, b.Relationship):
             rel = expr
+            params = [b.field_to_type(self._model, f) for f in rel._fields]
         else:
-            raise ValueError(f"Invalid expression type {type(expr)} for `solve_for`; must be a Relationship or Fragment")
-        assert rel._parent and rel._short_name, "Relationship for `solve_for` must have a parent and a short name"
+            raise ValueError(f"Invalid expression type {type(expr)} for `solve_for`")
+        assert len(params) == len(rel._fields)
         assert rel not in self._variable_relationships
 
         self._variable_relationships.add(rel)
@@ -81,7 +92,7 @@ class SolverModelDev:
             new_kwargs["type"] = "cont" if self._data_type == "float" else "int"
         for (key, val) in new_kwargs.items():
             if key == "name":
-                assert isinstance(val, (_Any, list)), f"Expected {key} to be a value or list, got {type(val)}"
+                assert isinstance(val, _Any) or isinstance(val, list), f"Expected {key} to be a value or list, got {type(val)}"
                 defs.append(self.variable_name(node, make_name(val)))
             elif key == "type":
                 assert val in ("cont", "int", "bin"), f"Unsupported variable type {val} for `solve_for`; must be cont, int, or bin"
@@ -96,9 +107,7 @@ class SolverModelDev:
                     self.satisfy(b.require(b.Expression(b.Relationship.builtins[op], rel, val)).where(*where))
             else:
                 raise ValueError(f"Invalid keyword argument {key} for `solve_for`")
-
-        where.append(_make_hash((rel._short_name, rel._parent), node))
-        b.define(*defs).where(*where)
+        b.define(*defs).where(*where, _make_hash((rel._short_name, rel._parent or 0), node))
 
         if populate:
             # get variable values from the result point (populated by the solver)
@@ -108,12 +117,10 @@ class SolverModelDev:
 
         return None
 
-    def minimize(self, expr: _Number, name: _String | list | None = None):
-        assert isinstance(expr, _Number)
+    def minimize(self, expr, name: _String | list | None = None):
         return self._handle_expr(self.min_objectives, expr, name)
 
-    def maximize(self, expr: _Number, name: _String | list | None = None):
-        assert isinstance(expr, _Number)
+    def maximize(self, expr, name: _String | list | None = None):
         return self._handle_expr(self.max_objectives, expr, name)
 
     def satisfy(self, expr: b.Fragment, check: bool = False, name: _String | list | None = None):
@@ -351,18 +358,10 @@ class SolverModelDev:
     # get scalar information
     def __getattr__(self, name: str):
         df = None
-        # model info
-        if name in {"num_variables", "num_constraints", "num_min_objectives", "num_max_objectives"}:
-            map = {
-                "num_variables": self.variables,
-                "num_constraints": self.constraints,
-                "num_min_objectives": self.min_objectives,
-                "num_max_objectives": self.max_objectives,
-            }
+        if name in self._model_info:
             node = b.Hash.ref()
-            df = b.select(b.count(node).where(map[name](node)) | 0).to_df()
-        # result info
-        if name in {"error", "termination_status", "solver_version", "printed_model", "solve_time_sec", "objective_value", "result_count"}:
+            df = b.select(b.count(node).where(self._model_info[name](node)) | 0).to_df()
+        elif name in {"error", "termination_status", "solver_version", "printed_model", "solve_time_sec", "objective_value", "result_count"}:
             val = b.String.ref()
             df = b.select(val).where(self.result_info(name, val)).to_df()
         if df is not None:
@@ -404,7 +403,7 @@ def _rewrite(expr: b.Producer | b.Fragment, ctx: ExprContext):
     elif isinstance(expr, (b.Relationship, b.RelationshipRef, b.RelationshipFieldRef)):
         rel = expr if isinstance(expr, b.Relationship) else expr._relationship
         if rel in ctx.solver_model._variable_relationships:
-            return std.hash(rel._short_name, expr._parent)
+            return std.hash(rel._short_name, expr._parent or 0)
         return None
 
     elif isinstance(expr, b.Union):
@@ -465,11 +464,18 @@ def _rewrite(expr: b.Producer | b.Fragment, ctx: ExprContext):
         subctx = ExprContext(sm)
         ctx.subcontext.append(subctx)
         subctx.where.extend(expr._where._where)
-        arg_hash = b.Hash.ref()
-        subctx.where.append(_make_hash((sm._expr_id, *pre_args), arg_hash)) # TODO also add sym_arg here?
         sm._expr_id += 1
         subctx.define.append(sm.operator(node, op))
-        subctx.define.append(sm.unordered_args_hash(node, arg_hash, sym_arg)) # TODO what if some values are data not hashes?
+
+        # special_ordered_set_type_2 has two ordered arguments: rank and variables
+        if op == "special_ordered_set_type_2":
+            assert len(pre_args) == 1, "special_ordered_set_type_2 expects exactly 2 arguments (rank, variables)"
+            subctx.define.append(sm.ordered_args_hash(node, pre_args[0], sym_arg))
+        else:
+            # other aggregate operators use unordered args
+            arg_hash = b.Hash.ref()
+            subctx.where.append(_make_hash((sm._expr_id, *pre_args), arg_hash))
+            subctx.define.append(sm.unordered_args_hash(node, arg_hash, sym_arg))
         return node
 
     elif isinstance(expr, b.Fragment):
@@ -502,7 +508,7 @@ def _expr_strings_rec(x, names_dict, ops_dict, args_dict):
             s = f"({s})"
         arg_strs.append(s)
 
-    if op in agg_ops:
+    if op in agg_ops and not op == "special_ordered_set_type_2":
         # sort unordered args to improve determinism
         arg_strs.sort()
 
@@ -522,7 +528,7 @@ infix_ops = set(["+", "-", "*", "/", "^"])
 infix_comps = set(["=", "!=", "<", "<=", ">", ">=", "implies"])
 infixs = infix_ops.union(infix_comps)
 prefix_ops = set(["abs", "log", "exp"])
-agg_ops = set(["sum", "count", "min", "max", "all_different"])
+agg_ops = set(["sum", "count", "min", "max", "all_different", "special_ordered_set_type_2"])
 
 # _variable_types = {
 #     "cont": 40,