diff --git a/namedisl/__init__.py b/namedisl/__init__.py index 31a2b22..ace697b 100644 --- a/namedisl/__init__.py +++ b/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", diff --git a/namedisl/expression_like.py b/namedisl/expression_like.py index 23d65fe..4c0932d 100644 --- a/namedisl/expression_like.py +++ b/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,12 +791,15 @@ 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__} @@ -754,7 +807,7 @@ class PwQPolynomial(_NamedPolynomialLike[isl.PwQPolynomial]): _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)) diff --git a/namedisl/set_like.py b/namedisl/set_like.py index c17239f..3a69a26 100644 --- a/namedisl/set_like.py +++ b/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: diff --git a/namedisl/test/test_set_like.py b/namedisl/test/test_set_like.py index b66711c..a4cabea 100644 --- a/namedisl/test/test_set_like.py +++ b/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])