How to create a Sympy IndexedBase using a custom subclass of Symbol?

I want to create a 2-D matrix of binary variables x. For now I do it like this: x = IndexedBase('x', shape=(imax,jmax), integer=True)

For a binary variable, the following identity holds: x_i**n == x_i for n>0, and I want to make use of this identity for simplifying my expressions later on. Therefore, the 'integer' assumption should be replaced with a "stronger" 'binary' assumption.

Creating a single binary variable works (Ref: https://stackoverflow.com/a/73953040/7740977):

from sympy import *

class Binary(Symbol):
    def _eval_power(self, other):
        return self

x0 = Binary('x0')
x0**2 == x0

Output: True

What I'm struggling with right now is creating an IndexedBase object where its entries are instances of my Binary class rather than of Symbol. Looking at https://docs.sympy.org/latest/modules/tensor/indexed.html#sympy.tensor.indexed.IndexedBase I saw that assumptions can be inherited if a Symbol is used to initialize the IndexedBase, which led me to unsuccessfully try the following:

x = symbols('x', cls=Binary)
x = IndexedBase(x)

x[0,0]**2 == x[0,0]

Output: False (expected True)

Is there another way to get IndexedBase to use my custom Binary class?

Many thanks in advance!

Solution

It's a different strategy, but instead of defining custom Symbol properties, you may be able to .replace() away exponentiation instead

custom Symbol can be simplified
define Wilds which are broad enough, but won't match nested instances

class Binary(Symbol):
    pass

a = Wild("a", properties=[lambda a: isinstance(a, Indexed) and a.atoms(Binary)])
b = Wild("b", properties=[lambda b: isinstance(b, Number)])

Now set up the problem and use .replace()

>>> x = Binary("x")
>>> A = IndexedBase(x)
>>> expr = A[0,0]
>>> expr
x[0, 0]
>>> expr = expr**2
>>> expr
x[0, 0]**2
>>> srepr(expr)  # NOTE the Indexed wrapper
"Pow(Indexed(IndexedBase(Binary('x')), Integer(0), Integer(0)), Integer(2))"
>>> expr.replace(Pow(a, b), lambda a, b: a)
x[0, 0]

If you're passing this on to someone else, consider a custom function to make it look cleaner and a docstring, but .replace() + Wild are surprisingly powerful and can reach into all sorts of expressions

>>> expr = Integral(5*A[0,0]**2 + sin(A[0,1]**3), x)  # not meaningful
>>> expr  
Integral(sin(x[0, 1]**3) + 5*x[0, 0]**2, x)
>>> expr.replace(Pow(a, b), lambda a, b: a)
Integral(sin(x[0, 1]) + 5*x[0, 0], x)