More loopy 4

namedisl/__init__.py

@@ -36,6 +36,7 @@
     PwMultiAff,
     PwQPolynomial,
     QPolynomial,
+    Term,
     affs_from_domain_space,
     make_aff,
     make_constraint,
@@ -74,6 +75,7 @@
     "QPolynomial",
     "Set",
     "Space",
+    "Term",
     "affs_from_domain_space",
     "align_two",
     "make_aff",

namedisl/expression_like.py

@@ -18,6 +18,10 @@
 .. autofunction:: make_pw_aff
 .. autofunction:: pw_affs_from_domain_space
 
+Quasipolynomial term
+--------------------
+.. autoclass:: Term
+
 Quasipolynomial
 ---------------
 .. autoclass:: QPolynomial
@@ -67,7 +71,7 @@
 """
 
 import operator
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from functools import cached_property
 from typing import (
@@ -86,6 +90,7 @@
 from .core import (
     DimType,
     IslAffLikeT_co,
+    IslExpressionLikeT,
     IslExpressionLikeT_co,
     IslHasCoefficientsT_co,
     IslObject,
@@ -110,14 +115,17 @@
     bound="_NamedExpressionLike[IslScalarExpressionLike]")
 
 
-def _unparam_expr_domain(obj: NamedExpressionLikeT) -> NamedExpressionLikeT:
+def _unparam_expr_domain(obj: IslExpressionLikeT) -> IslExpressionLikeT:
     """If the domain of obj is a param domain, this will make it not a param domain."""
 
     # Oh isl. There *has* to be a better way.
 
+    if isinstance(obj, (isl.PwMultiAff, isl.Constraint)):
+        raise NotImplementedError(f"not supported for {type(obj)}")
+
     dt = isl.dim_type.in_
-    d = obj._obj.dim(dt)
-    return type(obj)(obj._obj.add_dims(dt, 1).drop_dims(dt, d, 1), obj.space)
+    d = obj.dim(dt)
+    return cast("IslExpressionLikeT", obj.add_dims(dt, 1).drop_dims(dt, d, 1))
 
 
 def _align_two_expr_likes(
@@ -130,9 +138,9 @@ def _align_two_expr_likes(
     lhs, rhs = align_two(lhs, rhs)
 
     if lhs._obj.get_domain_space().is_params():
-        lhs = _unparam_expr_domain(lhs)
+        lhs = type(lhs)(_unparam_expr_domain(lhs._obj), lhs.space)
     if rhs._obj.get_domain_space().is_params():
-        rhs = _unparam_expr_domain(rhs)
+        rhs = type(rhs)(_unparam_expr_domain(rhs._obj), rhs.space)
 
     return lhs, rhs
 
@@ -324,7 +332,7 @@ class _NamedHasCoefficients(NamedIslObject[IslHasCoefficientsT_co]):
     .. automethod:: get_div
     .. automethod:: get_div_coefficient
 
-    .. automethod:: get_constant
+    .. autoattribute:: constant
     """
 
     @property
@@ -341,7 +349,8 @@ def get_coefficient(self, name: str) -> isl.Val:
         dt, idx = self.space.name_to_dim[name]
         return self._obj.get_coefficient_val(dt.as_isl(), idx)
 
-    def get_constant(self) -> isl.Val:
+    @property
+    def constant(self) -> isl.Val:
         return self._obj.get_constant_val()
 
 
@@ -350,7 +359,7 @@ class Aff(_NamedAffLike[isl.Aff], _NamedHasCoefficients[isl.Aff]):
     .. automethod:: zero_on_domain
     .. automethod:: as_pw_aff
     .. automethod:: set_coefficient
-    .. automethod:: get_denominator
+    .. automethod:: denominator
 
     .. autoattribute:: var_affs
 
@@ -375,7 +384,7 @@ def set_coefficient(self, name: str, value: int) -> Aff:
         dt, idx = self.space.name_to_dim[name]
         return Aff(self._obj.set_coefficient_val(dt.as_isl(), idx, value), self.space)
 
-    def get_denominator(self) -> isl.Val:
+    def denominator(self) -> isl.Val:
         return self._obj.get_denominator_val()
 
     @cached_property
@@ -404,6 +413,8 @@ def make_aff(src: isl.Aff) -> Aff:
 
 def make_aff(src: str | isl.Aff, ctx: isl.Context | None = None) -> Aff:
     obj = isl.Aff(src, ctx) if isinstance(src, str) else src
+    if obj.get_domain_space().is_params():
+        obj = _unparam_expr_domain(obj)
     return Aff(obj, Space.from_isl(obj, Aff.active_dim_types))
 
 
@@ -505,7 +516,7 @@ class PwAff(_NamedAffLike[isl.PwAff]):
     .. automethod:: zero_like_me
     .. autoattribute:: var_pw_affs
     .. automethod:: where
-    .. automethod:: get_pieces
+    .. autoattribute:: pieces
     .. automethod:: coalesce
     .. automethod:: eq_set
     .. automethod:: ne_set
@@ -515,7 +526,7 @@ class PwAff(_NamedAffLike[isl.PwAff]):
     .. automethod:: lt_set
     .. automethod:: max
     .. automethod:: min
-    .. automethod:: get_aggregate_domain
+    .. automethod:: aggregate_domain
     .. automethod:: union_max
     .. automethod:: union_min
     .. automethod:: union_add
@@ -613,7 +624,7 @@ def min(self, other: PwAff) -> PwAff:
         self_a, other_a = _align_two_expr_likes(self, other)
         return PwAff(self_a._obj.min(other_a._obj), self_a.space)
 
-    def get_pieces(self) -> list[tuple[Set, Aff]]:
+    def pieces(self) -> list[tuple[Set, Aff]]:
         set_space = self.space.as_set_space()
         from .set_like import Set
         return [
@@ -626,7 +637,7 @@ def get_pieces(self) -> list[tuple[Set, Aff]]:
     def coalesce(self) -> PwAff:
         return PwAff(self._obj.coalesce(), self.space)
 
-    def get_aggregate_domain(self) -> Set:
+    def aggregate_domain(self) -> Set:
         from .set_like import Set
         agg_domain = self._obj.get_aggregate_domain()
         return Set(
@@ -661,6 +672,8 @@ def make_pw_aff(src: str | isl.PwAff, ctx: isl.Context | None = None) -> PwAff:
     Create a :class:`PwAff` from isl syntax or an :class:`islpy.PwAff`.
     """
     obj = isl.PwAff(src, ctx) if isinstance(src, str) else src
+    if obj.get_domain_space().is_params():
+        obj = _unparam_expr_domain(obj)
     return PwAff(obj, Space.from_isl(obj, PwAff. active_dim_types))
 
 
@@ -709,14 +722,51 @@ def __pow__(self, other: int) -> Self:
         return type(self)(cast("IslPolynomialLikeT_co", self._obj ** other), self.space)
 
 
+@dataclass(frozen=True, eq=False)
+class Term:
+    """
+    .. autoattribute:: space
+    .. autoattribute:: coefficient
+    .. automethod:: get_exp
+    .. autoattribute:: num_divs
+    .. automethod:: get_div
+    .. automethod:: get_div_exp
+    """
+    # Term is super-rudimentary
+    _obj: isl.Term
+    space: Space
+
+    @property
+    def coefficient(self) -> isl.Val:
+        return self._obj.get_coefficient_val()
+
+    def get_exp(self, name: str) -> int:
+        dt, idx = self.space.name_to_dim[name]
+        return self._obj.get_exp(dt.as_isl(), idx)
+
+    @property
+    def num_divs(self) -> int:
+        return self._obj.dim(isl.dim_type.div)
+
+    def get_div(self, index: int) -> Aff:
+        return Aff(self._obj.get_div(index), self.space)
+
+    def get_div_exp(self, index: int) -> int:
+        return self._obj.get_exp(isl.dim_type.div, index)
+
+
 class QPolynomial(_NamedPolynomialLike[isl.QPolynomial]):
     __doc__ = f"""
+    .. automethod:: terms
     {_NamedPolynomialLike.__doc__}
     {_NamedExpressionLike.__doc__}
     {NamedIslObject.__doc__}
     """
     _isl_type: ClassVar[type[IslObject]] = isl.QPolynomial
 
+    def terms(self) -> Sequence[Term]:
+        return [Term(trm, self.space) for trm in self._obj.get_terms()]
+
 
 @overload
 def make_qpolynomial(src: str, ctx: isl.Context | None = None) -> QPolynomial:
@@ -741,20 +791,23 @@ def make_qpolynomial(
     else:
         obj = src
 
+    if obj.get_domain_space().is_params():
+        obj = _unparam_expr_domain(obj)
+
     return QPolynomial(obj, Space.from_isl(obj, QPolynomial.active_dim_types))
 
 
 class PwQPolynomial(_NamedPolynomialLike[isl.PwQPolynomial]):
     __doc__ = f"""
-    .. automethod:: get_pieces
+    .. automethod:: pieces
     {_NamedPolynomialLike.__doc__}
     {_NamedExpressionLike.__doc__}
     {NamedIslObject.__doc__}
     """
 
     _isl_type: ClassVar[type[IslObject]] = isl.PwQPolynomial
 
-    def get_pieces(self) -> list[tuple[Set, QPolynomial]]:
+    def pieces(self) -> list[tuple[Set, QPolynomial]]:
         set_space = self.space.as_set_space()
         from .set_like import Set
         return [
@@ -782,6 +835,8 @@ def make_pw_qpolynomial(
     Create a :class:`PwQPolynomial` from isl syntax or an isl object.
     """
     obj = isl.PwQPolynomial(src, ctx) if isinstance(src, str) else src
+    if obj.get_domain_space().is_params():
+        obj = _unparam_expr_domain(obj)
     return PwQPolynomial(
         obj,
         Space.from_isl(obj, PwQPolynomial.active_dim_types))

namedisl/set_like.py

@@ -320,7 +320,7 @@ class BasicSet(_NamedIslSetLike[isl.BasicSet], _NamedIslBasic[isl.BasicSet]):
     .. automethod:: add_constraint
     .. autoattribute:: var_affs
     .. automethod:: as_set
-    .. automethod:: get_constraints
+    .. automethod:: constraints
     {_NamedIslSetLike.__doc__}
     {_NamedIslBasic.__doc__}
     {_NamedIslSetOrMapLike.__doc__}
@@ -336,7 +336,7 @@ def add_constraint(self, cns: Constraint, /) -> BasicSet:
                 raise ValueError("spaces don't match")
         return BasicSet(self._obj.add_constraint(cns._obj), self.space)
 
-    def get_constraints(self):
+    def constraints(self):
         from .expression_like import Constraint
         return [
             Constraint(cns, self.space.with_empty_dim_type(DimType.in_))
@@ -382,7 +382,7 @@ class Set(_NamedIslSetLike[isl.Set], _NamedIslUnbasic[isl.Set]):
     .. automethod:: complement
     .. automethod:: simple_hull
     .. automethod:: convex_hull
-    .. automethod:: get_basic_sets
+    .. automethod:: basic_sets
     .. automethod:: dim_max
     .. automethod:: dim_min
     .. autoattribute:: var_affs
@@ -405,23 +405,27 @@ def simple_hull(self):
     def convex_hull(self) -> BasicSet:
         return BasicSet(self._obj.convex_hull(), self.space)
 
-    def get_basic_sets(self) -> list[BasicSet]:
+    def basic_sets(self) -> list[BasicSet]:
         return [BasicSet(bs, self.space) for bs in self._obj.get_basic_sets()]
 
     def dim_max(self, name: str, *, cache: Cache | None = None) -> PwAff:
         dt, idx = self.space.name_to_dim[name]
         if dt != DimType.out:
             raise ValueError("can only take max with respect to set dimensions")
-        from .expression_like import PwAff
-        return PwAff(with_cache(cache, isl.Set.dim_max, self._obj, idx),
+        from .expression_like import PwAff, _unparam_expr_domain
+        isl_result = _unparam_expr_domain(
+            with_cache(cache, isl.Set.dim_max, self._obj, idx))
+        return PwAff(isl_result,
             self.space.drop_dim_type(DimType.out).with_empty_dim_type(DimType.in_))
 
     def dim_min(self, name: str, *, cache: Cache | None = None) -> PwAff:
         dt, idx = self.space.name_to_dim[name]
         if dt != DimType.out:
             raise ValueError("can only take min with respect to set dimensions")
-        from .expression_like import PwAff
-        return PwAff(with_cache(cache, isl.Set.dim_min, self._obj, idx),
+        from .expression_like import PwAff, _unparam_expr_domain
+        isl_result = _unparam_expr_domain(
+            with_cache(cache, isl.Set.dim_min, self._obj, idx))
+        return PwAff(isl_result,
             self.space.drop_dim_type(DimType.out).with_empty_dim_type(DimType.in_))
 
     @cached_property
@@ -508,7 +512,7 @@ class BasicMap(_NamedIslMapLike[isl.BasicMap], _NamedIslBasic[isl.BasicMap]):
     .. automethod:: range
     .. automethod:: intersect_domain
     .. automethod:: intersect_range
-    .. automethod:: get_constraints
+    .. automethod:: constraints
     {_NamedIslMapLike.__doc__}
     {_NamedIslBasic.__doc__}
     {_NamedIslSetOrMapLike.__doc__}
@@ -522,7 +526,7 @@ def add_constraint(self, cns: Constraint, /) -> BasicMap:
                 raise ValueError("spaces don't match")
         return BasicMap(self._obj.add_constraint(cns._obj), self.space)
 
-    def get_constraints(self):
+    def constraints(self):
         from .expression_like import Constraint
         return [
             Constraint(cns, self.space) for cns in self._obj.get_constraints()]
@@ -595,7 +599,7 @@ class Map(_NamedIslMapLike[isl.Map], _NamedIslUnbasic[isl.Map]):
     .. automethod:: complement
     .. automethod:: simple_hull
     .. automethod:: convex_hull
-    .. automethod:: get_basic_maps
+    .. automethod:: basic_maps
     .. automethod:: domain
     .. automethod:: range
     .. automethod:: intersect_domain
@@ -625,7 +629,7 @@ def simple_hull(self):
     def convex_hull(self) -> BasicMap:
         return BasicMap(self._obj.convex_hull(), self.space)
 
-    def get_basic_maps(self) -> list[BasicMap]:
+    def basic_maps(self) -> list[BasicMap]:
         return [BasicMap(bs, self.space) for bs in self._obj.get_basic_maps()]
 
     def domain(self) -> Set:

namedisl/test/test_set_like.py

@@ -311,8 +311,8 @@ def test_map_coalesce() -> None:
     map_ = nisl.make_map(
         "{ [i] -> [j = i] : 0 <= i < 5 or 5 <= i < 10 }"
     )
-    assert len(map_.get_basic_maps()) == 2
-    assert len(map_.coalesce().get_basic_maps()) == 1
+    assert len(map_.basic_maps()) == 2
+    assert len(map_.coalesce().basic_maps()) == 1
 
 
 @pytest.mark.parametrize("ndims_domain", [2, 3, 4, 5])