implemented sampline via autoregressive inverse transform

Author: LeanderK

analysis/plot_helper.py

@@ -165,14 +165,6 @@ def plot_in_diverging_theme(
     # plt.figure()
     # f = plt.gcf()
     # the_ax = plt.gca()
-    if plot_max_value is None:
-        the_max = data.max()
-    else:
-        the_max = plot_max_value
-
-    if norm is not None:
-        the_max = norm.vmax
-
     if with_contour and plot_borders_below:
         if hatched:
             contour = the_ax.contourf(
@@ -197,22 +189,32 @@ def plot_in_diverging_theme(
             origin="upper",
             zorder=0,
         )
-    the_data = data.copy()
-
-    # negative values are valid, positve values are invalid
-    # I know this is stupid, but it is what it is
-    the_data[valid.T] = (-1) * the_data[valid.T]
-    # the_data[~valid] = (-1) * the_data[~valid]
-
-    the_ax.imshow(
-        the_data,
-        extent=img_extent,
-        origin="upper",
-        cmap=the_cmap,
-        vmin=-the_max,
-        vmax=the_max,
-        interpolation="nearest",
-    )
+
+    if data is not None:
+        if plot_max_value is None:
+            the_max = data.max()
+        else:
+            the_max = plot_max_value
+
+        if norm is not None:
+            the_max = norm.vmax
+
+        the_data = data.copy()
+
+        # negative values are valid, positve values are invalid
+        # I know this is stupid, but it is what it is
+        the_data[valid.T] = (-1) * the_data[valid.T]
+        # the_data[~valid] = (-1) * the_data[~valid]
+
+        the_ax.imshow(
+            the_data,
+            extent=img_extent,
+            origin="upper",
+            cmap=the_cmap,
+            vmin=-the_max,
+            vmax=the_max,
+            interpolation="nearest",
+        )
 
     if with_valid:
         # plot valid tensor with true being pink and false transparent

analysis/plot_samples.ipynb

@@ -27,7 +27,7 @@ def bisection_search(
             upper = mid
         count += 1
 
-    print(f"bisection_search took {count} iterations")
+    # print(f"bisection_search took {count} iterations")
     return (upper + lower) / 2
 
 
@@ -53,7 +53,7 @@ def conditioned_function(
     """
 
     def wrapped_function(x: torch.Tensor) -> torch.Tensor:
-        x_full = torch.cat([x, constant.unsqueeze(0)], dim=-1)
+        x_full = torch.vmap(lambda x: torch.cat([constant, x]))(x)
         return f(x_full)
 
     return wrapped_function
@@ -122,7 +122,7 @@ def condition_linear_ineq(linear_ineq: lra.LinearInequality):
             return linear_ineq_res
 
         the_constraints = constraints.map_constraints(
-            f=condition_linear_ineq, drop_vars=list(sampled_dimensions.keys())
+            f=condition_linear_ineq, #, drop_vars=list(sampled_dimensions.keys())
         )
 
         degree = get_sub_total_degree(poly, i)
@@ -137,7 +137,7 @@ def condition_linear_ineq(linear_ineq: lra.LinearInequality):
         lower, upper = limits[var]
 
         constrained_var_map_dict = {
-            var: (i - idx) for var, idx in variable_map.items() if i >= idx
+            var: (i - idx) for var, idx in variable_map.items() if idx >= i
         }
 
         integration_kwargs = {
@@ -183,7 +183,7 @@ def sample_spline_distribution(
     Sample a single point from the SplineSQ2D distribution.
     """
     assert not dist.is_batched()
-    mixture_weights = dist.mixture_weights.unsqueeze(0)
+    mixture_weights = dist.mixture_weights.squeeze(0)
     # sample the component index from the mixture weights
     component_index = torch.multinomial(mixture_weights, 1, replacement=True).item()
 
@@ -209,11 +209,19 @@ def sample_spline_distribution(
     f.to(device)
     f = f.to(precision)
 
-    return inverse_transform_sampling_gasp(
+    p = inverse_transform_sampling_gasp(
         boxed_problem,
         f,
         device=device,
         precision=precision,
         gasp_kwargs=gasp_kwargs,
         wmi_pa_mode=wmi_pa_mode,
     )
+
+    # assert that in box
+    variable_map = dist.var_positions
+    for var, idx in variable_map.items():
+        lb, ub = box.constraints[var]
+        assert lb <= p[idx] <= ub, f"Sampled point {p[idx]} is out of bounds for {var} ({lb}, {ub})"
+    
+    return p

pal/distribution/sampling/inverse_transform.py

@@ -420,7 +420,7 @@ class SplineSQ2D(ConstrainedDistribution, torch.nn.Module):
     differences: torch.Tensor  # (num_knots-1, 2) because of memory layout
     integrals_2dgrid: (
         torch.Tensor
-    )  # (b, num_mixtures, num_knots, num_knots), b=1 gets broadcasted
+    )  # (b, num_mixtures, num_knots-1, num_knots-1), b=1 gets broadcasted
     poly_params: (
         torch.Tensor
     )  # (b, num_mixtures, 2, num_knots, 4), b=1 gets broadcasted
@@ -552,15 +552,18 @@ def enumerate_pieces(
     ) -> list[tuple[lra.Box, torch.Tensor, Callable[[torch.Tensor], torch.Tensor]]]:
         results = []
 
+        y_pos_dict = {i: name for name, i in self.var_positions.items()}
+        y_pos_dict = y_pos_dict
+
         for i in range(self.knots.shape[0] - 1):
             for j in range(self.knots.shape[0] - 1):
                 lower_x0 = self.knots[i, 0]
                 upper_x0 = self.knots[i + 1, 0]
                 lower_x1 = self.knots[j, 1]
                 upper_x1 = self.knots[j + 1, 1]
 
-                varname0 = self.y_pos_dict[0]
-                varname1 = self.y_pos_dict[1]
+                varname0 = y_pos_dict[0]
+                varname1 = y_pos_dict[1]
 
                 box = lra.Box(
                     id=(i, j),
@@ -570,7 +573,7 @@ def enumerate_pieces(
                     },
                 )
 
-                integrals = self.integrals_2dgrid[:, i, j]
+                integrals = self.integrals_2dgrid[:, :, i, j]
                 p1 = self.poly_params[:, :, 0, i]
                 p2 = self.poly_params[:, :, 1, j]
 

pal/distribution/spline_distribution.py

@@ -279,6 +279,8 @@ def __init__(
         VariableMapMixin.__init__(self, variable_map_dict)
         self.register_buffer("coeffs", coeffs)
         self.register_buffer("powers", powers)
+        if params.shape[0] == 1:
+            params = params.squeeze(0)
         self.register_buffer("params", params)
         assert coeffs.shape[0] == powers.shape[0]
 

pal/distribution/torch_polynomial.py

@@ -126,7 +126,7 @@ def __and__(self, other: "LRA"):
         if isinstance(other, And):
             return And(*(self.children + other.children))
         else:
-            return And(*(self.children + [other]))
+            return And(*(list(self.children) + [other]))
 
     def __str__(self) -> str:
         return "(" + " & ".join([str(child) for child in self.children]) + ")"
@@ -259,7 +259,8 @@ def to_logic(var: str, lb: float, ub: float):
             return self._expression
 
     def map_constraints(
-        self, f: Callable[[LinearInequality], LinearInequality]
+        self,
+        f: Callable[[LinearInequality], LinearInequality]
     ) -> "LRAProblem":
         """
         Applies a function `f` to each `LinearConstraint` in the expression tree while keeping
@@ -268,6 +269,8 @@ def map_constraints(
         Args:
             f (Callable[[LinearConstraint], LinearConstraint]): The function to apply to each
             `LinearConstraint`.
+            drop_vars (list[str]): A list of variables to to drop as a result of the mapping.
+                This is used for removing variables that are no longer needed after the mapping.
 
         Returns:
             LinearIneqLogicTree: A new expression tree with the modified `LinearConstraint` objects.
@@ -285,15 +288,15 @@ def recurse_expression(expr: LRA) -> LRA:
                 return Or(*mapped_children)
 
         if self.expression is None:
-            return LRAProblem(None, self._variables)
+            return LRAProblem(None, self._variables, self._name)
         else:
             expr = recurse_expression(self._expression)
             variables = gather_variables(expr)
             if isinstance(self._variables, dict):
                 sub_vars = {var: self._variables[var] for var in variables}
             else:
                 sub_vars = [var for var in self._variables if var in variables]
-            return LRAProblem(expr, sub_vars)
+            return LRAProblem(expr, sub_vars, self._name)
 
     def get_global_limits(self) -> dict[str, tuple[float, float]]:
         """
@@ -309,9 +312,9 @@ def get_global_limits(self) -> dict[str, tuple[float, float]]:
                 "Global limits are not available when variables are provided as a list."
             )
 
-    def __and__(self, other: LRA | Box):
+    def __and__(self, other: LRA | Box) -> "LRAProblem":
         if isinstance(other, LRA):
-            return LRAProblem(self.expression & other, self._variables)
+            return LRAProblem(self.expression & other, self._variables, self._name)
         elif isinstance(other, Box):
             # via global bounds
             if isinstance(self._variables, dict):

pal/logic/lra.py

@@ -1,5 +1,5 @@
 import pal.logic.lra as lra
-from pysmt.shortcuts import get_env, Symbol, Real, Plus, And, Or
+from pysmt.shortcuts import get_env, Symbol, Real, Plus, And, Or, Bool
 import pysmt.typing as stypes
 from pysmt.fnode import FNode
 
@@ -25,19 +25,23 @@ def recursive_translate(
         node: lra.LRA,
     ) -> FNode:
         if isinstance(node, lra.LinearInequality):
-            # canonical is ax + by + cz + ... < d
-            left_sum = Plus(
-                *[
-                    Real(lhs_coeff) * get_symb(lhs_var)
-                    for lhs_var, lhs_coeff in node.lhs.items()
-                ]
-            )
-            if node.symbol == "<=":
-                return left_sum <= Real(node.rhs)
-            elif node.symbol == ">=":
-                return left_sum >= Real(node.rhs)
+            if not node.is_empty():
+                # canonical is ax + by + cz + ... < d
+                left_sum = Plus(
+                    *[
+                        Real(lhs_coeff) * get_symb(lhs_var)
+                        for lhs_var, lhs_coeff in node.lhs.items()
+                    ]
+                )
+                if node.symbol == "<=":
+                    return left_sum <= Real(node.rhs)
+                elif node.symbol == ">=":
+                    return left_sum >= Real(node.rhs)
+                else:
+                    raise NotImplementedError()
             else:
-                raise NotImplementedError()
+                # empty constraints are always true
+                return Bool(True)
         elif isinstance(node, lra.And):
             recursive_translate_children = [
                 recursive_translate(c) for c in node.children

pal/logic/lra_pysmt.py

@@ -4,11 +4,11 @@
 from pysmt.fnode import FNode
 import torch
 
+from pal.logic.lra import box_to_lra
 from pal.logic.lra_pysmt import translate_to_pysmt
 from pal.distribution.constrained_distribution import (
     ConditionalConstraintedDistribution,
     ConstrainedDistributionBuilder,
-    box_to_lra,
 )
 from pal.wmi.gasp.gasp.torch.wmipa.numerical_symb_integrator_pa import (
     FunctionMode,

pal/wmi/compute_integral.py

@@ -316,8 +316,8 @@ def convert_to_problem_wrapper(problem):
                 if len(results) == 0:
                     return torch.zeros([1, len(problems)]), 0
                 else:
-                    if any(r.shape[0] == 1 for r in results):
-                        print("hi")
+                    # if any(r.shape[0] == 1 for r in results):
+                    #     print("hi")
                     results = torch.concatenate(results, dim=-1)
                     return results, 0
             case WeightedFormulaMode():