Customising the subs() functionality of Function subclasses in sympy

Question

I want to define a class of functions that follow this mathematical definition:

Definition Given a function q(t) and two values, a and b, we define as variation the following function

         /
         | q(t) if s=0 or t=a or t=b
v(t,s) = |
         | v(t,s) else
         \ 

I'm trying to emulate this behaviour definition with a Function subclass:

from sympy import Function
class Variation(Function):
    # Initialize the function with the desired properties
    def __init__(self, path, st, en, name='\\vartheta'):
        self.path = path
        self.st = st
        self.en= en
        self.ends = [self.st, self.en]
        self.name = name

    # here I define the behaviour when called
    def __call__(self, tt, ss):
        if tt in self.ends:
            return self.path(tt)
        elif ss == 0:
            return self.path(tt)
        else:
            return Function(self.name)(tt,ss)  # This is the part that fails to behave 

The function behaves well when called:

from sympy import *
s,t,a,b = symbols('s t a b')
c = Function('c')

Var = Variation(c, a, b)
Var(t,s), Var(a,s), Var(t,0) 

> \vartheta(t,s), q(a), q(t)

but as expected, if we do:

Var(t,s).subs(t,0)

> \vartheta(t,0)

Is there a way to modify the .subs()'s method behaviour? Because as far as I am aware, the integrate() function makes use of subs().

I also tried changing Function(self.name)(tt,ss) to self(tt,ss) but this gave me an infinite loop (also expected).

On the same note, is there a good guide to construct arbitrary mathematical functions on python?

Edit: Tried

def Var(t,s):
    return Piecewise((c(t), s==0), (c(t), t==a), (c(t), t==b), (Function('v')(t,s), True ))

Var(t,s).subs(t,0)

but it had the same problems.

Answer 1

I think you need to reshape how you go about this a bit.

Creating your Var function.

You really want Var in your example to be a Function class. sympy is designed around functions being classes and using the eval() class method to evaluate them. Overriding __call__ of a Function subclass seems to be very non-standard and I haven't seen any sympy built-in functions that use that, so I don't think that's the right approach. One way of doing this would be to create a factory function to create the class for you:

def Variation(path_, st_, en_, v_):

    class Variation(Function):

        nargs = 2

        path = path_
        st = st_
        en = en_
        ends = [st, en]
        v = v_

        @classmethod
        def eval(cls, tt, ss):
            if tt in cls.ends:
                return cls.path(tt)

            if ss == 0:
                return cls.path(tt)

            return cls.v(tt, ss)

    return Variation

Var = Variation(c, a, b, Function(r'\vartheta'))

• I've replaced your 'name' variable with an actual function, which seems more sensible.

Now you can create variations and prevent immediate evaluation (if you want) using the standard flag:

# by default the function is immediately evaluated...
Var(a, t)
>>> c(a)
# ...but that can be overridden
Var(a, t, evaluate=False)
>>> Variation(a, t)

You could also approach this by flattening the Var function and passing st en and path parameters straight into eval(), which removes the extra layer of the factory function:

class Variation(Function):

    @classmethod
    def eval(cls, path, st, en, v, tt, ss):
        if tt in [st, en]:
            return path(tt)
        elif ss == 0:
            return path(tt)
        else:
            return v(tt, ss)

Variation(c, a, b, Function(r'\vartheta'), a, t)
>>> Variation(c, a, b, \vartheta, a, t)

Note that since you can override .eval() you could modify it so that it didn't automatically simplify if you wanted, and just returned a new instance of cls:

class Variation(Function):
    no_eval = True

    @classmethod
    def eval(cls, tt, ss):

        if cls.no_eval:
            return cls(tt, ss, evaluate=False)
        # etc.

Customizing .subs()

By default whenever you do a subs(), sympy will do an eval() as well (as per the .subs() docs). So by default when you do a .subs() with one of your special values, .eval() will be called and the function will be simplified.

However, you can now override ._eval_subs() and do your own thing, if you want:

class Variation(Function):
    ...
    def _eval_subs(self, old, new):
        # return self to do no substitutions at all
        return self
        # return None to continue normally by next calling _subs on the arguments to this Function
        # return some other Expression to return that instead.

Note that anything returned by ._eval_subs() will subsequently be .eval()ed as well. You can override .eval() as explained above if you wanted to get round that.

So I think that answers the question about how to modify the .subs() behaviour...

I don't quite understand what you want with:

is there a good guide to construct arbitrary mathematical functions on python?

I think sympy is pretty good, and has reasonable docs and many many inbuilt examples in its codebase which are easy to crib from. Anyway asking for guides is off-topic (see point 4) on stackoverflow ;-).