Marimo notebook: All elementary functions from a single binary operator

eml.py

# Run this notebook with:
#  marimo edit --sandbox eml.py
#
# /// script
# dependencies = [
#     "marimo",
#     "sympy",
# ]
# requires-python = ">=3.12"
# ///

import marimo

__generated_with = "0.22.4"
app = marimo.App(width="medium", app_title="eml(x, y)")


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Build operations using only EML and 1

    This is a puzzle inspired by the recent paper [All elementary functions from a single binary operator](https://arxiv.org/abs/2603.21852)

    The goal of this notebook is to build up a variety of functions using only the function EML and the constant 1. **You are allowed to build on functions that you wrote from earlier steps**.

    $$
    \begin{align}
    eml(a, b) = e^a - ln(b)
    \end{align}
    $$

    ## Examples

    These come directly from the abstract of the paper. Just a demo of how the notebook works.
    """)
    return


@app.cell
def _(check_function, eml, sp):
    def eml_exp(a):
        "Implement $e^a$ using only eml, a, and 1"
        return eml(a, 1)


    check_function(eml_exp, sp.exp)
    return


@app.cell
def _(check_function, eml, sp):
    def eml_ln(a):
        "Implement $ln(a)$ using only eml, a, and 1"
        return eml(1, eml(eml(1, a), 1))


    check_function(eml_ln, sp.log)
    return (eml_ln,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Puzzle 1: Basics

    See page 22 of the paper for a diagram showing the relative difficulty of the operations. Start from the inside and work your way out.
    """)
    return


@app.cell
def _(check_function, eml, sp):
    def eml_e():
        "Find constant $e$ using only eml and 1"
        raise NotImplementedError


    check_function(eml_e, lambda: sp.exp(1))
    return (eml_e,)


@app.cell
def _(check_function, eml, eml_ln):
    def eml_sub(a, b):
        "Implement $a - b$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_sub, lambda a, b: a - b)
    return (eml_sub,)


@app.cell
def _(check_function, eml, eml_e, eml_ln, eml_sub):
    def eml_minus_1():
        "Find constant $-1$ using only eml and 1"
        raise NotImplementedError


    check_function(eml_minus_1, lambda: -1)
    return (eml_minus_1,)


@app.cell
def _(check_function, eml, eml_e, eml_ln, eml_minus_1, eml_sub):
    def eml_2():
        "Find constant $2$ using only eml and 1"
        raise NotImplementedError


    check_function(eml_2, lambda: 2)
    return (eml_2,)


@app.cell
def _(check_function, eml, eml_e, eml_ln, eml_sub):
    def eml_unary_minus(a):
        "Implement unary minus ($-a$) using only eml, a, and 1"
        raise NotImplementedError


    check_function(eml_unary_minus, lambda a: -a)
    return (eml_unary_minus,)


@app.cell
def _(check_function, eml_sub, eml_unary_minus):
    def eml_plus(a, b):
        "Implement $a + b$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_plus, lambda a, b: a + b)
    return


@app.cell
def _(check_function, eml):
    def eml_inverse(a):
        "Implement multiplicative inverse $1 / a$ using only eml, a, and 1"
        raise NotImplementedError


    check_function(eml_inverse, lambda a: 1 / a)
    return (eml_inverse,)


@app.cell
def _(check_function, eml):
    def eml_times(a, b):
        "Implement multiplication $a \\times b$ using only eml, a, b, and 1"
        left = eml(1, eml(eml(1, eml(1, a)), 1))
        return eml(eml(1, eml(eml(left, b), 1)), 1)


    check_function(eml_times, lambda a, b: a * b)
    return (eml_times,)


@app.cell
def _(check_function, eml_times):
    def eml_square(a):
        "Implement the square function $a^2$ using only eml, a, and 1"
        raise NotImplementedError


    check_function(eml_square, lambda a: a**2)
    return


@app.cell
def _(check_function, eml_inverse, eml_times):
    def eml_divide(a, b):
        "Implement division $a \\div b$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_divide, lambda a, b: a / b)
    return (eml_divide,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Puzzle 2: More complex combinations

    You finally implemented division! You are ready for secondary school!
    """)
    return


@app.cell
def _(check_function, eml_2, eml_divide):
    def eml_half(a):
        "Implement division by 2 $a \\div 2$ using only eml, a, and 1"
        # Remember you can do this!
        return eml_divide(a, eml_2())


    check_function(eml_half, lambda a: a / 2)
    return


@app.cell
def _(check_function):
    def eml_x_plus_y_over_2(a, b):
        "Implement the function $(a + b) \\div 2$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_x_plus_y_over_2, lambda a, b: (a + b) / 2)
    return


@app.cell
def _(check_function, sp):
    def eml_sqrt(a):
        "Implement the square root function $\\sqrt{a}$ using only eml, a, and 1"
        raise NotImplementedError


    check_function(eml_sqrt, lambda a: sp.sqrt(a))
    return


@app.cell
def _(check_function):
    def eml_pow(a, b):
        "Implement the exponentiation function $a^b$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_pow, lambda a, b: a**b)
    return


@app.cell
def _(check_function, sp):
    def eml_log(a, b):
        "Implement the logarithm function $log_{a} b$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_log, lambda a, b: sp.log(b, a))
    return


@app.cell
def _(check_function, sp):
    def eml_pi():
        "Find constant $\\pi$ using only eml and 1"
        raise NotImplementedError


    check_function(eml_pi, lambda: sp.S.Pi)
    return


@app.cell
def _(check_function, sp):
    def eml_pythag(a, b):
        "Implement the pythagorean formula $\\sqrt{a^2+b^2}$ using only eml, a, b, and 1"
        raise NotImplementedError


    check_function(eml_pythag, lambda a, b: sp.sqrt(a**2 + b**2))
    return


@app.cell
def _(check_function, sp):
    def eml_sigma(a):
        "Implement the function $1 \\div (1 + e^-a)$ using only eml, a, and 1"
        raise NotImplementedError


    check_function(eml_sigma, lambda a: 1 / (1 + sp.exp(-a)))
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Puzzle 3: Trigonometry

    The original author used computer-aided search to derive this system. But surely you can do this by hand!
    """)
    return


@app.cell
def _(check_function, sp):
    def eml_cosh(a):
        "Implement the hyperolic cosine $\\cosh(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_cosh, lambda a: sp.cosh(a))
    return


@app.cell
def _(check_function, sp):
    def eml_sinh(a):
        "Implement the hyperbolic sine $\\sinh(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_sinh, lambda a: sp.sinh(a))
    return


@app.cell
def _(check_function, sp):
    def eml_tanh(a):
        "Implement the hyperbolic tangent $\\tanh(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_tanh, lambda a: sp.tanh(a))
    return


@app.cell
def _(check_function, sp):
    def eml_cos(a):
        "Implement the cosine function $\\cos(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_cos, lambda a: sp.cos(a))
    return


@app.cell
def _(check_function, sp):
    def eml_sin(a):
        "Implement the sine function $\\sin(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_sin, lambda a: sp.sin(a))
    return


@app.cell
def _(check_function, sp):
    def eml_tan(a):
        "Implement the tangent function $\\tan(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_tan, lambda a: sp.tan(a))
    return


@app.cell
def _(check_function, sp):
    def eml_arcsinh(a):
        "Implement the hyperbolic arcsine $\\arcsinh(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_arcsinh, lambda a: sp.asinh(a))
    return


@app.cell
def _(check_function, sp):
    def eml_arcosh(a):
        "Implement the hyperbolic arccosine $\\mathrm{arcosh}(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_arcosh, lambda a: sp.acosh(a))
    return


@app.cell
def _(check_function, sp):
    def eml_arccos(a):
        "Implement the arccosine function $\\arccos(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_arccos, lambda a: sp.acos(a))
    return


@app.cell
def _(check_function, sp):
    def eml_artanh(a):
        "Implement the hyperbolic arctangent $\\mathrm{artanh}(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_artanh, lambda a: sp.atanh(a))
    return


@app.cell
def _(check_function, sp):
    def eml_arcsin(a):
        "Implement the arcsine function $\\arcsin(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_arcsin, lambda a: sp.asin(a))
    return


@app.cell
def _(check_function, sp):
    def eml_arctan(a):
        "Implement the arctangent function $\\arctan(a)$ using only eml, a, and 1"
        raise NotImplementedError

    check_function(eml_arctan, lambda a: sp.atan(a))
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Finish!

    Congratulations if you've made it this far!
    """)
    return


@app.cell
def _():
    import marimo as mo
    import sympy as sp
    import types

    return mo, sp


@app.cell
def _(sp):
    class eml(sp.Function):
        "This is the core EML function."

        def doit(self, deep=False, **hints):
            a, b = self.args
            # Recursively call doit() on the args whenever deep=True.
            # Be sure to pass deep=True and **hints through here.
            if deep:
                a, b = a.doit(deep=deep, **hints), b.doit(deep=deep, **hints)
            return sp.exp(a) - sp.log(b)

        def rewrite(self, target, *args, deep=True, **hints):
            a, b = self.args
            if deep:
                a = (
                    a.rewrite(target, *args, deep=deep, **hints)
                    if isinstance(a, sp.Basic)
                    else a
                )
                b = (
                    b.rewrite(target, *args, deep=deep, **hints)
                    if isinstance(b, sp.Basic)
                    else b
                )
            if target == sp.exp:
                if b == 1:
                    return sp.exp(a)
            if target == sp.log:
                # original: eml(1, eml(eml(1, a), 1))
                if (
                    a == 1
                    and isinstance(b, eml)
                    and b.args[1] == 1
                    and isinstance(b.args[0], eml)
                    and b.args[0].args[0] == 1
                ):
                    return sp.log(b.args[0].args[1])
                # after rewrite(sp.exp): eml(1, exp(x)) -> log(x)
                if (
                    a == 1
                    and isinstance(b, sp.exp)
                    and isinstance(b.args[0], eml)
                    and b.args[0].args[0] == 1
                ):
                    return sp.log(b.args[0].args[1])
            return eml(a, b)

    return (eml,)


@app.cell
def _(eml, mo, sp):
    def check_function(test, expect):
        try:
            a, b = sp.symbols("a b", positive=True)
            if test.__code__.co_argcount == 0:
                have_exp = test()
                want = expect()
            elif test.__code__.co_argcount == 1:
                have_exp = test(a)
                want = expect(a)
            else:
                have_exp = test(a, b)
                want = expect(a, b)

            steps = ""
            have = have_exp
            for _ in range(20):
                steps += f"\n\n$$ {sp.latex(have)} $$"
                next = (
                    have.rewrite(sp.exp, deep=False)
                    .rewrite(sp.log, deep=False)
                    .doit()
                )
                if have == next or have == want:
                    break
                have = next

            # Validate that only eml, 1, a, b are used
            for node in sp.preorder_traversal(have_exp):
                if isinstance(node, eml):
                    continue
                if isinstance(node, sp.Integer) and node == 1:
                    continue
                if node in (a, b):
                    continue
                return mo.md(
                    f"### ❌ {test.__name__}: Incorrect\n\n"
                    f"Uses disallowed term: `{node}`. Only `eml`, `1`, `a`, and `b` are allowed.\n\n"
                    f"Your result so far: ${sp.latex(sp.simplify(have))}$"
                )

            if sp.simplify(have - want) == 0:
                ops = sum(
                    1
                    for _ in sp.preorder_traversal(have_exp)
                    if isinstance(_, eml)
                )
                return mo.md(
                    f"### ✅ {test.__name__}: Correct!\n$$ {sp.latex(sp.simplify(have))} = {sp.latex(want)} $$\n\n**EML operations used: {ops}**"
                )
            else:
                return mo.md(f"### ❌ {test.__name__}: Incorrect\n{steps}")
        except NotImplementedError as ex:
            return mo.md(
                f"### ⏳ {test.__name__}: Not started yet\n\n**Next task**: {test.__doc__}"
            )

    return (check_function,)


if __name__ == "__main__":
    app.run()

very fun puzzlel! I integrated it to emlvm under emlvm puzzle