basic case funky-trig-substitution integral works! epic

Author: laurgao

src/simpy/expr.py

@@ -228,6 +228,11 @@ def evalf(self, subs: Optional[Dict[str, "Expr"]] = None) -> "Expr":
     def children(self) -> List["Expr"]:
         raise NotImplementedError(f"Cannot get children of {self.__class__.__name__}")
 
+    @property
+    def childless(self) -> bool:
+        """Returns True if it's a basic element like a symbol or a number that doesn't have any subcomponents at all."""
+        return len(self.children()) == 0
+
     def contains(self: "Expr", var: "Symbol") -> bool:
         is_var = isinstance(self, Symbol) and self.name == var.name
         return is_var or any(e.contains(var) for e in self.children())

src/simpy/regex.py

@@ -1,7 +1,11 @@
-"""Custom library for checking what shit exists. Replaces searching the repr with regex.
+"""Custom library for checking what exists within exprs. Replaces searching the repr with regex.
 
 This module is currently still developmental. It does the job often but is not promised to be robust
 outside of the cases it is currently used for. Use with caution.
+
+TODO: write a regex quickstart guide. Maybe spin this off into a subfolder and make the outward facing API more
+intuitive. Perhaps make it similar to the regex library.
++ Unit testing with comprehensively thought-out cases.
 """
 
 from collections import defaultdict
@@ -395,6 +399,28 @@ def count(expr: Expr, query: Expr) -> int:
     return sum(count(e, query) for e in expr.children())
 
 
+def contains(expr: Expr, query: Expr) -> EqResult:
+    """Checks if `query` appears in `expr`. Assumes query contains Any_ objects.
+    Exact any-matches only, no up to factor or sum.
+    Returns a results dictionary like eq() does. with `success` and `matches` keys.
+    """
+    ## base cases ##
+    eq_output = eq(expr, query)
+    if eq_output["success"]:
+        return eq_output
+    if expr.childless:
+        return {"success": False}
+
+    ## recursive cases ##
+    # this isn't super sophisticated for dupes but it's fine for now.
+    for e in expr.children():
+        eq_output_ = contains(e, query)
+        if eq_output_["success"]:
+            return eq_output_
+
+    return {"success": False}
+
+
 def contains_cls(expr: Expr, cls: Type[Expr]) -> bool:
     if isinstance(expr, cls):
         return True
@@ -412,6 +438,10 @@ def general_count(expr: Expr, condition: ExprCondition) -> int:
 
 
 def general_contains(expr: Expr, condition: ExprCondition) -> bool:
+    """contains with a condition function instead of any.
+
+    this is sorta-legacy --- i think we should use contains instead; it's cuter. any-matches are cute.
+    """
     if condition(expr):
         return True
     return any(general_contains(e, condition) for e in expr.children())

src/simpy/transforms.py

@@ -33,7 +33,7 @@
 from .integral_table import check_integral_table
 from .linalg import invert
 from .polynomial import Polynomial, is_polynomial, polynomial_to_expr, rid_ending_zeros, to_const_polynomial
-from .regex import count, general_contains, replace, replace_class, replace_factory
+from .regex import Any_, contains, count, general_contains, replace, replace_class, replace_factory
 from .simplify import pythagorean_simplification
 from .simplify.product_to_sum import product_to_sum_unit
 from .utils import ExprFn, eq_with_var, random_id
@@ -211,6 +211,7 @@ def check(self, node: Node) -> bool:
 class USub(Transform, ABC):
     """Base class for u-substituion transforms."""
 
+    # u (new var) written in terms of x (old var)
     _u: Expr = None
 
     def backward(self, node: Node) -> None:
@@ -377,8 +378,6 @@ def _get_last_heuristic_transform(node: Node, tup=(PullConstant, Additivity)):
     return node.transform
 
 
-# Let's just add all the transforms we've used for now.
-# and we will make this shit good and generalized later.
 class TrigUSub2(USub):
     """
     u-sub of a trig function
@@ -410,8 +409,8 @@ def check(self, node: Node) -> bool:
         if super().check(node) is False:
             return False
 
-        # Since B and C essentially undo each other, we want to make sure that the last
-        # heuristic transform wasn't C.
+        # Since TrigUSub2 and InverseTrigUSub essentially undo each other, we want to make sure that the last
+        # heuristic transform wasn't InverseTrigUSub.
 
         t = _get_last_heuristic_transform(node)
         if isinstance(t, InverseTrigUSub):
@@ -966,6 +965,67 @@ def forward(self, node: Node) -> None:
         self._u = self._u
 
 
+class TrigUSub(USub):
+    """This is the substitution of the form x = 4*cos(theta) (TODO: write better descrip later)"""
+
+    _a: Rat = None  # constant in the square root
+    _exponent: Rat = None  # exponent of the square root (includes the square root, is a fraction w denom = 2)
+    _case: int = None  # 0 for sqrt(a^2 - x^2), 1 for sqrt(a^2 + x^2), 2 for sqrt(-a^2 + x^2)
+
+    def check(self, node: Node):
+        """Check if node.expr contains sqrt(a^2-x^2) or sqrt(a^2+x^2) where a is a constant."""
+        if super().check(node) is False:
+            return False
+
+        # check if any instance of sqrt(a^2 - x^2) or sqrt(a^2 + x^2) or sqrt(-a^2 + x^2) appears.
+        def squared_integer_condition(expr: Expr) -> bool:
+            return isinstance(expr, Rat) and isinstance(sqrt(expr), Rat)
+
+        a_squared = Any_("squared_integer", squared_integer_condition, is_constant=True)
+        any_square_root_exponent = Any_(
+            "square_root_exponent", lambda expr: isinstance(expr, Rat) and expr.denominator == 2, is_constant=True
+        )
+        queries = [
+            (a_squared - node.var**2) ** any_square_root_exponent,
+            (node.var**2 - a_squared) ** any_square_root_exponent,
+            (node.var**2 + a_squared) ** any_square_root_exponent,
+        ]
+        for i, query in enumerate(queries):
+            out = contains(node.expr, query)
+            if out["success"]:
+                self._a = sqrt(out["matches"]["squared_integer"])
+                self._exponent = out["matches"]["square_root_exponent"]
+                self._case = i
+                return True
+
+        return False
+
+    def forward(self, node: Node) -> None:
+        if self._case == 0:
+            # in the case of sqrt(a^2 - x^2):
+            # x = a * sin(theta)
+            # dx = a * cos(theta) d(theta)
+            theta = generate_intermediate_var()
+            dx_dtheta = self._a * cos(theta)
+
+            # so we replace x = a * sin(theta), this effectively leads to
+            # replacing sqrt(a^2 - x^2) = a * cos^2(theta)
+            theta_expr = replace(
+                node.expr,
+                (self._a**2 - node.var**2) ** self._exponent,
+                (self._a * cos(theta)) ** self._exponent.numerator,
+            )
+            if theta_expr.contains(node.var):
+                theta_expr = replace(theta_expr, node.var, self._a * sin(theta))
+
+            new_integrand = theta_expr * dx_dtheta
+            node.add_child(Node(new_integrand, theta, self, node))
+
+            self._u = asin(node.var / self._a)  # theta in terms of x
+
+        return NotImplemented
+
+
 class CompleteTheSquare(Transform):
     """Integration via completing the square"""
 
@@ -1120,6 +1180,7 @@ def backward(self, node: Node) -> None:
     RewriteTrig,
     RewritePythagorean,
     InverseTrigUSub,
+    TrigUSub,
     CompleteTheSquare,
     GenericUSub,
 ]

tests/test_khan_academy_integrals.py

@@ -202,3 +202,11 @@ def test_more_complicated_trig():
     expr = tan(x) ** 5 * sec(x) ** 4
     expected_ans = tan(x) ** 6 / 6 + tan(x) ** 8 / 8
     assert_integral(expr, expected_ans)
+
+
+def test_trigonometric_substitution():
+    # that's this one: https://www.khanacademy.org/math/integral-calculus/ic-integration/ic-trig-substitution/e/integration-using-trigonometric-substitution
+    expr = (4 - x**2) ** Fraction(3, 2)
+    expected_ans = sin(4 * asin(x / 2)) / 2 + 4 * sin(2 * asin(x / 2)) + 6 * asin(x / 2)  # TODO: simplify this
+    ans = integrate(expr)
+    assert_eq_plusc(expected_ans, ans)

tests/test_regex.py

@@ -2,7 +2,7 @@
 from test_utils import x, y
 
 from simpy.expr import *
-from simpy.regex import Any_, any_, any_constant, eq
+from simpy.regex import Any_, any_, any_constant, contains, eq
 
 
 def test_any_basic():
@@ -18,6 +18,14 @@ def test_sort_anys():
     assert eq(sin(x) * sec(x), sec(any_) * sin(any_))
 
 
+def test_eq_with_different_anys():
+    any2 = Any_()
+    expr = sin(x) + cos(y)
+    query = sin(any_) + cos(any2)
+    out = eq(expr, query)
+    assert out["success"]
+
+
 @pytest.mark.parametrize(
     ["sum", "expected"],
     [
@@ -118,3 +126,23 @@ def test_any_constant_fail():
     query = 2 * x + any_constant
     out = eq(expr, query)
     assert not out["success"]
+
+
+def test_any_constant_with_multiple_anys():
+    expr = 2 * x + 3
+    query = 2 * any_ + any_constant
+    out = eq(expr, query)
+    assert out["success"]
+
+
+def test_contains():
+    query = log(sin(any_) ** 2 + cos(any_) ** 2)
+    expr = (log(sin(x) ** 2 + cos(x) ** 2) + 3) ** 2
+    assert contains(expr, query)["success"]
+    assert contains(expr, query)["matches"] == x
+
+
+def test_contains_fail():
+    query = log(sin(any_) ** 2 + cos(any_) ** 2)
+    expr = (log(sin(x) ** 2 + cos(x) ** 3) + 3) ** 2 + 1
+    assert not contains(expr, query)["success"]