pepflow 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

pepflow/scalar.py

@@ -34,28 +34,78 @@ if TYPE_CHECKING:
 
 
 def is_numerical_or_scalar(val: Any) -> bool:
-    return utils.is_numerical(val) or isinstance(val, Scalar)
+    return utils.is_numerical_or_parameter(val) or isinstance(val, Scalar)
 
 
 def is_numerical_or_evaluatedscalar(val: Any) -> bool:
-    return utils.is_numerical(val) or isinstance(val, EvaluatedScalar)
+    return utils.is_numerical_or_parameter(val) or isinstance(val, EvaluatedScalar)
 
 
 @attrs.frozen
-class EvalExpressionScalar:
+class ScalarRepresentation:
     op: utils.Op
     left_scalar: Point | Scalar | float
     right_scalar: Point | Scalar | float
 
 
+@attrs.frozen
+class ZeroScalar:
+    """A special class to represent 0 in scalar."""
+
+    pass
+
+
 @attrs.frozen
 class EvaluatedScalar:
+    """
+    The concrete representation of the abstract :class:`Scalar`.
+
+    Each abstract basis :class:`Scalar` object has a unique concrete
+    representation as a unit vector. The concrete representations of
+    linear combinations of abstract basis :class:`Scalar` objects
+    are linear combinations of the unit vectors. This information is
+    stored in the `vector` attribute.
+
+    Abstract :class:`Scalar` objects can be formed through taking the
+    inner product of two abstract :class:`Point` objects. The
+    concrete representation of an abstract :class:`Scalar` object formed
+    this way is the outer product of the concrete representations of the
+    two abstract :class:`Point` objects, i.e., a matrix. This information
+    is stored in the `matrix` attribute.
+
+    Abstract :class:`Scalar` objects can be added or subtracted with
+    numeric data types. This information is stored in the `constant`
+    attribute.
+
+    :class:`EvaluatedScalar` objects can be constructed as linear combinations
+    of other :class:`EvaluatedScalar` objects. Let `a` and `b` be some numeric
+    data type. Let `u` and `v` be :class:`EvaluatedScalar` objects. Then, we
+    can form a new :class:`EvaluatedScalar` object: `a*u+b*v`.
+
+    Attributes:
+        vector (np.ndarray): The vector component of the concrete
+            representation of the abstract :class:`Scalar`.
+        matrix (np.ndarray): The matrix component of the concrete
+            representation of the abstract :class:`Scalar`.
+        constant (float): The constant component of the concrete
+            representation of the abstract :class:`Scalar`.
+    """
+
     vector: np.ndarray
     matrix: np.ndarray
     constant: float
 
+    @classmethod
+    def zero(cls, num_basis_scalars: int, num_basis_points: int):
+        return EvaluatedScalar(
+            vector=np.zeros(num_basis_scalars),
+            matrix=np.zeros((num_basis_points, num_basis_points)),
+            constant=0.0,
+        )
+
     def __add__(self, other):
-        assert is_numerical_or_evaluatedscalar(other)
+        if not is_numerical_or_evaluatedscalar(other):
+            return NotImplemented
         if utils.is_numerical(other):
             return EvaluatedScalar(
                 vector=self.vector, matrix=self.matrix, constant=self.constant + other
@@ -68,7 +118,8 @@ class EvaluatedScalar:
             )
 
     def __radd__(self, other):
-        assert is_numerical_or_evaluatedscalar(other)
+        if not is_numerical_or_evaluatedscalar(other):
+            return NotImplemented
         if utils.is_numerical(other):
             return EvaluatedScalar(
                 vector=self.vector, matrix=self.matrix, constant=other + self.constant
@@ -81,7 +132,8 @@ class EvaluatedScalar:
             )
 
     def __sub__(self, other):
-        assert is_numerical_or_evaluatedscalar(other)
+        if not is_numerical_or_evaluatedscalar(other):
+            return NotImplemented
         if utils.is_numerical(other):
             return EvaluatedScalar(
                 vector=self.vector, matrix=self.matrix, constant=self.constant - other
@@ -94,7 +146,8 @@ class EvaluatedScalar:
             )
 
     def __rsub__(self, other):
-        assert is_numerical_or_evaluatedscalar(other)
+        if not is_numerical_or_evaluatedscalar(other):
+            return NotImplemented
         if utils.is_numerical(other):
             return EvaluatedScalar(
                 vector=-self.vector, matrix=-self.matrix, constant=other - self.constant
@@ -107,7 +160,8 @@ class EvaluatedScalar:
             )
 
     def __mul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         return EvaluatedScalar(
             vector=self.vector * other,
             matrix=self.matrix * other,
@@ -115,7 +169,8 @@ class EvaluatedScalar:
         )
 
     def __rmul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         return EvaluatedScalar(
             vector=other * self.vector,
             matrix=other * self.matrix,
@@ -126,7 +181,8 @@ class EvaluatedScalar:
         return self.__rmul__(other=-1)
 
     def __truediv__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         return EvaluatedScalar(
             vector=self.vector / other,
             matrix=self.matrix / other,
@@ -136,11 +192,27 @@ class EvaluatedScalar:
 
 @attrs.frozen
 class Scalar:
+    """
+    A :class:`Scalar` object represents linear combination of functions values,
+    inner products of points, and constant scalar values.
+
+    :class:`Scalar` objects can be constructed as linear combinations of
+    other :class:`Scalar` objects. Let `a` and `b` be some numeric data type.
+    Let `x` and `y` be :class:`Scalar` objects. Then, we can form a new
+    :class:`Scalar` object: `a*x+b*y`.
+
+    Attributes:
+        is_basis (bool): True if this point is not formed through a linear
+            combination of other points. False otherwise.
+        tags (list[str]): A list that contains tags that can be used to
+            identify the :class:`Point` object. Tags should be unique.
+    """
+
     # If true, the scalar is the basis for the evaluations of F
     is_basis: bool
 
-    # Not sure on this yet
-    eval_expression: EvalExpressionScalar | None = None
+    # The representation of scalar used for evaluation.
+    eval_expression: ScalarRepresentation | ZeroScalar | None = None
 
     # Human tagged value for the scalar
     tags: list[str] = attrs.field(factory=list)
@@ -159,13 +231,27 @@ class Scalar:
             raise RuntimeError("Did you forget to create a context?")
         pep_context.add_scalar(self)
 
+    @staticmethod
+    def zero() -> Scalar:
+        return Scalar(is_basis=False, eval_expression=ZeroScalar(), tags=["0"])
+
     @property
     def tag(self):
+        """Returns the most recently added tag.
+
+        Returns:
+            str: The most recently added tag of this :class:`Scalar` object.
+        """
         if len(self.tags) == 0:
             raise ValueError("Scalar should have a name.")
         return self.tags[-1]
 
     def add_tag(self, tag: str):
+        """Add a new tag for this :class:`Scalar` object.
+
+        Args:
+            tag (str): The new tag to be added to the `tags` list.
+        """
         self.tags.append(tag)
 
     def __repr__(self):
@@ -174,93 +260,100 @@ class Scalar:
         return super().__repr__()
 
     def _repr_latex_(self):
-        s = repr(self)
-        s = s.replace("star", r"\star")
-        s = s.replace("gradient_", r"\nabla ")
-        return rf"$\\displaystyle {s}$"
+        return utils.str_to_latex(repr(self))
 
     def __add__(self, other):
-        assert is_numerical_or_scalar(other)
-        if utils.is_numerical(other):
-            tag_other = f"{other:.4g}"
+        if not is_numerical_or_scalar(other):
+            return NotImplemented
+        if utils.is_numerical_or_parameter(other):
+            tag_other = utils.numerical_str(other)
         else:
             tag_other = other.tag
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.ADD, self, other),
+            eval_expression=ScalarRepresentation(utils.Op.ADD, self, other),
             tags=[f"{self.tag}+{tag_other}"],
         )
 
     def __radd__(self, other):
-        assert is_numerical_or_scalar(other)
-        if utils.is_numerical(other):
-            tag_other = f"{other:.4g}"
+        if not is_numerical_or_scalar(other):
+            return NotImplemented
+        if utils.is_numerical_or_parameter(other):
+            tag_other = utils.numerical_str(other)
         else:
             tag_other = other.tag
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.ADD, other, self),
+            eval_expression=ScalarRepresentation(utils.Op.ADD, other, self),
             tags=[f"{tag_other}+{self.tag}"],
         )
 
     def __sub__(self, other):
-        assert is_numerical_or_scalar(other)
-        if utils.is_numerical(other):
-            tag_other = f"{other:.4g}"
+        if not is_numerical_or_scalar(other):
+            return NotImplemented
+        if utils.is_numerical_or_parameter(other):
+            tag_other = utils.numerical_str(other)
         else:
             tag_other = utils.parenthesize_tag(other)
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.SUB, self, other),
+            eval_expression=ScalarRepresentation(utils.Op.SUB, self, other),
             tags=[f"{self.tag}-{tag_other}"],
         )
 
     def __rsub__(self, other):
-        assert is_numerical_or_scalar(other)
+        if not is_numerical_or_scalar(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
-        if utils.is_numerical(other):
-            tag_other = f"{other:.4g}"
+        if utils.is_numerical_or_parameter(other):
+            tag_other = utils.numerical_str(other)
         else:
             tag_other = other.tag
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.SUB, other, self),
+            eval_expression=ScalarRepresentation(utils.Op.SUB, other, self),
             tags=[f"{tag_other}-{tag_self}"],
         )
 
     def __mul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical_or_parameter(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
+        tag_other = utils.numerical_str(other)
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.MUL, self, other),
-            tags=[f"{tag_self}*{other:.4g}"],
+            eval_expression=ScalarRepresentation(utils.Op.MUL, self, other),
+            tags=[f"{tag_self}*{tag_other}"],
         )
 
     def __rmul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical_or_parameter(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
+        tag_other = utils.numerical_str(other)
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.MUL, other, self),
-            tags=[f"{other:.4g}*{tag_self}"],
+            eval_expression=ScalarRepresentation(utils.Op.MUL, other, self),
+            tags=[f"{tag_other}*{tag_self}"],
         )
 
     def __neg__(self):
         tag_self = utils.parenthesize_tag(self)
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.MUL, -1, self),
+            eval_expression=ScalarRepresentation(utils.Op.MUL, -1, self),
             tags=[f"-{tag_self}"],
         )
 
     def __truediv__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical_or_parameter(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
+        tag_other = f"1/{utils.numerical_str(other)}"
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.DIV, self, other),
-            tags=[f"1/{other:.4g}*{tag_self}"],
+            eval_expression=ScalarRepresentation(utils.Op.DIV, self, other),
+            tags=[f"{tag_other}*{tag_self}"],
         )
 
     def __hash__(self):
@@ -271,22 +364,101 @@ class Scalar:
             return NotImplemented
         return self.uid == other.uid
 
-    def le(self, other, name: str) -> ctr.Constraint:
+    def le(self, other: Scalar | float | int, name: str) -> ctr.Constraint:
+        """
+        Generate a :class:`Constraint` object that represents the inequality
+        `self` <= `other`.
+
+        Args:
+            other (:class:`Scalar` | float | int): The other side of the
+                relation.
+            name (str): The name of the generated :class:`Constraint` object.
+
+        Returns:
+            :class:`Constraint`: An object that represents the inequality
+            `self` <= `other`.
+        """
         return ctr.Constraint(self - other, comparator=utils.Comparator.LT, name=name)
 
-    def lt(self, other, name: str) -> ctr.Constraint:
+    def lt(self, other: Scalar | float | int, name: str) -> ctr.Constraint:
+        """
+        Generate a :class:`Constraint` object that represents the inequality
+        `self` < `other`.
+
+        Args:
+            other (:class:`Scalar` | float | int): The other side of the
+                relation.
+            name (str): The name of the generated :class:`Constraint` object.
+
+        Returns:
+            :class:`Constraint`: An object that represents the inequality
+            `self` < `other`.
+        """
         return ctr.Constraint(self - other, comparator=utils.Comparator.LT, name=name)
 
-    def ge(self, other, name: str) -> ctr.Constraint:
+    def ge(self, other: Scalar | float | int, name: str) -> ctr.Constraint:
+        """
+        Generate a :class:`Constraint` object that represents the inequality
+        `self` >= `other`.
+
+        Args:
+            other (:class:`Scalar` | float | int): The other side of the
+                relation.
+            name (str): The name of the generated :class:`Constraint` object.
+
+        Returns:
+            :class:`Constraint`: An object that represents the inequality
+            `self` >= `other`.
+        """
         return ctr.Constraint(self - other, comparator=utils.Comparator.GT, name=name)
 
-    def gt(self, other, name: str) -> ctr.Constraint:
+    def gt(self, other: Scalar | float | int, name: str) -> ctr.Constraint:
+        """
+        Generate a :class:`Constraint` object that represents the inequality
+        `self` > `other`.
+
+        Args:
+            other (:class:`Scalar` | float | int): The other side of the
+                relation.
+            name (str): The name of the generated :class:`Constraint` object.
+
+        Returns:
+            :class:`Constraint`: An object that represents the inequality
+            `self` > `other`.
+        """
         return ctr.Constraint(self - other, comparator=utils.Comparator.GT, name=name)
 
-    def eq(self, other, name: str) -> ctr.Constraint:
+    def eq(self, other: Scalar | float | int, name: str) -> ctr.Constraint:
+        """
+        Generate a :class:`Constraint` object that represents the inequality
+        `self` = `other`.
+
+        Args:
+            other (:class:`Scalar` | float | int): The other side of the
+                relation.
+            name (str): The name of the generated :class:`Constraint` object.
+
+        Returns:
+            :class:`Constraint`: An object that represents the inequality
+            `self` = `other`.
+        """
         return ctr.Constraint(self - other, comparator=utils.Comparator.EQ, name=name)
 
     def eval(self, ctx: pc.PEPContext | None = None) -> EvaluatedScalar:
+        """
+        Return the concrete representation of this :class:`Scalar`.
+        Concrete representations of :class:`Scalar` objects are
+        :class:`EvaluatedScalar` objects.
+
+        Args:
+            ctx (:class:`PEPContext` | None): The :class:`PEPContext` object
+                we consider. `None` if we consider the current global
+                :class:`PEPContext` object.
+
+        Returns:
+            :class:`EvaluatedScalar`: The concrete representation of
+            this :class:`Scalar`.
+        """
         from pepflow.expression_manager import ExpressionManager
 
         # Note this can be inefficient.
@@ -296,3 +468,44 @@ class Scalar:
             raise RuntimeError("Did you forget to create a context?")
         em = ExpressionManager(ctx)
         return em.eval_scalar(self)
+
+    def repr_by_basis(
+        self, ctx: pc.PEPContext | None = None, greedy_square: bool = True
+    ) -> str:
+        """Express this :class:`Scalar` object in terms of the basis
+        :class:`Point` and :class:`Scalar` objects of the given
+        :class:`PEPContext`.
+
+        A :class:`Scalar` can be formed by linear combinations of basis
+        :class:`Scalar` objects. A :class:`Scalar` can also be formed through
+        the inner product of two basis :class:`Point` objects. This function
+        returns the representation of this :class:`Scalar` object in terms of
+        the basis :class:`Point` and :class:`Scalar` objects as a `str` where,
+        to refer to the basis :class:`Point` and :class:`Scalar` objects, we
+        use their tags.
+
+        Args:
+            ctx (:class:`PEPContext`): The :class:`PEPContext` object
+                whose basis :class:`Point` and :class:`Scalar` objects we
+                consider. `None` if we consider the current global
+                `PEPContext` object.
+            greedy_square (bool): If `greedy_square` is true, the function will
+                try to return :math:`\\|a-b\\|^2` whenever possible. If not,
+                the function will return
+                :math:`\\|a\\|^2 - 2 * \\langle a, b \\rangle + \\|b\\|^2` instead.
+                `True` by default.
+
+        Returns:
+            str: The representation of this :class:`Scalar` object in terms of
+            the basis :class:`Point` and :class:`Scalar` objects of the given
+            :class:`PEPContext`.
+        """
+        from pepflow.expression_manager import ExpressionManager
+
+        # Note this can be inefficient.
+        if ctx is None:
+            ctx = pc.get_current_context()
+        if ctx is None:
+            raise RuntimeError("Did you forget to create a context?")
+        em = ExpressionManager(ctx)
+        return em.repr_scalar_by_basis(self, greedy_square=greedy_square)