分段函数无穷级数的Sympy图，符号为索引问题

Question

I'm having trouble creating this Function to plot in python using Sympy, where {r1, r2, r3, ...} are an enumeration of the rationals.

I've tried the following to define each function separately, but the main issue is trying to use a sympy symbol as an index for my rationals list:

import numpy as np
from sympy.abc import x, n
from sympy import Piecewise, piecewise_fold, Sum, IndexedBase, oo, Function

rationals = 10*np.random.rand(10000)

u = Function('u')(x, n)
class u(Function):
    
    '''A function u_n(x) which returns 1/2^n if x > r_n, and 0 otherwise, 
    where r_n taken from an enumeration of the rationals.'''
    
    nargs = 2
    @classmethod
    def eval(cls, x, n):
        if x > rationals[n]:
            return 1/2**n
        else:
            return 0
        
h = Function('h')(x)
class h(Function):
    
    '''Function which evaluates the summation of u_n(x)'''
    
    @classmethod
    def eval(cls, x):
        return Sum(u(x, n), (0, n, oo))

I get the following error when evaluating x > rationals[n]:

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

How do I go around this problem? Are there better ways to write this code if all I want to do is plot h(x), and perhaps make a gif with a slider to see how the function changes as n -> oo?

Answer 1

There are a couple minor issues with the code:

Sum(u(x, n), (0, n, oo)) should be instead Sum(u(x, n), (n, 0, oo))

Also there's no need to write u = Function('u')(x, n) because this is overwritten by the class definition anyway.

But more fundamentally you can't index into a numpy array using a sympy symbol as this will always return an IndexError . In addition, Sum where the limits are infinite is only going to be helpful for well known symbolic summations like Sum(1/n**2, (n, 1, oo)) which you can call the methods .doit() on to obtain pi**2/6 (or .evalf() will give a floating point result).

However, the .doit() method for an infinite Sum won't be able to compute a result for your custom function u as is. It seems to me that you might have better luck trying to write this using purely numerical functions from numpy or scipy.

Sticking with sympy though I think you want something like this?

import numpy as np
from sympy import Function, Sum, symbols, oo, S, plot

rationals = 10*np.random.rand(10000)

class u(Function):
    
    '''A function u_n(x) which returns 1/2^n if x > r_n, and 0 otherwise, 
    where r_n taken from an enumeration of the rationals.'''
    
    nargs = 2
    @classmethod
    def eval(cls, x, n):
        # a guard like this will prevent your IndexError (but you'll be left with an unevaluated `u(x, n)` if you pass non numeric arguments to `u`
        if x.is_Number and x.is_real and n.is_Integer: 
            if x > rationals[n]:
                return 1/2**n
            else:
                return S(0)  # sympy's plot apparently requires a 'sympified' integer

n = symbols("n", integer=True)
        
class h(Function):
    
    '''Function which evaluates the summation of u_n(x)'''
    
    @classmethod
    def eval(cls, x, lim):
        return Sum(u(x, n), (n, 0, lim)).doit()


numeric_result = h(2, 100)
print(numeric_result) # rational
print(numeric_result.evalf()) # float
print(h(2, oo)) # unevaluated

x = symbols("x")
p = plot(h(x,100))
p.show()

Running this code produces eg:

159056405582528143525572386825/633825300114114700748351602688
0.250946760178067
Sum(u(2, n), (n, 0, oo))

so with the infinite upper limit sympy just returns the result unevaluated.

Here's the plot: