Ability to create sub classes

[staadecker]

Summary

I would love to be able to create subclasses of Expression, Set, Param, and Var. Specifically I'd love to be able to override __init__() to handle some custom arguments. Unfortunately, the way many of the classes are setup at the moment make this impossible. For example, __new__() in Set prevents properly creating an instance of a subclass of set.

Rationale

Being able to extend Pyomo would allow greater flexibility for users.

Example

import pyomo.environ as pyo


class CustomSet(pyo.Set):
    def __init__(self, *args, custom_param=None, **kwargs):
        self.custom_param = custom_param
        super().__init__(*args, **kwargs)


# Works fine
m1 = pyo.AbstractModel()
m1.FirstSet = pyo.Set(dimen=2)
m1.DependantSet = pyo.Set(m1.FirstSet)

# Doesn't work, raises "can't apply a Set operator to an indexed Set component"
m2 = pyo.AbstractModel()
m2.FirstSet = CustomSet(dimen=2, custom_param="some_value")
m2.DependantSet = pyo.Set(m2.FirstSet)

[ghackebeil]

The pyomo.kernel modeling layer was designed for this purpose. It might be worth trying out:

https://pyomo.readthedocs.io/en/stable/library_reference/kernel/index.html

[jsiirola]

Can you elaborate on what is preventing you from creating subclasses? There are numerous examples of subclassing Pyomo components (in particular subclassing Block - but all components follow a similar design pattern)

[staadecker]

@ghackebeil that looks really great! It would still be nice for the pyomo.core to support subclassing. I'm working with a very large existing code base and switching the entire code base to kernel isn't possible.

@jsiirola I've added an example to the issue.

[blnicho]

@staadecker I'm cleaning up old issues. Can this be closed?

[staadecker]

I mean the issue is still valid afaik, but if the decision is Pyomo doesn't have the interest/time to implement a fix then it would make sense to close it. I know this is not something I'm working on atm.

[jsiirola]

I think we should leave this open: at the very least, we need to better document the AML class hierarchy design, as well as tools like the declare_custom_block decorator (we even have core developers occasionally try inheriting from core classes in ways other than we originally intended). Longer term, generalizing the approach for managing Scalar / Indexed components (possibly avoiding the current diamond hierarchies) would be a big win.

[goroderickgo]

@jsiirola, following up on this thread, I'm in similar pickle to @staadecker (building on an existing codebase that uses the AML syntax). It's also hard for me to make the case to my team that we should switch to the kernel library, since the docs say that it is "experimental", so it could be helpful if the Pyomo devs could provide more guidance on the stability of the library (or how to help contribute!) 🙏🏼.

I added the line self._ctype = self.ctype to my subclass's __init__ method, which seemed like one important piece to getting indexed variables (and expressions) to work; however, I'm finding that I can't initialize scalar variables (or expressions) properly, which I discovered by trying to use any of the methods or attributes that I would expect on a pyo.Var (e.g., .pprint(), .value.

class NewVar(pyo.Var):
    def __init__(self, *args, a=0, **kwargs):
        super().__init__(*args, **kwargs)
        self._ctype = self.ctype

        self.a = a


m = pyo.ConcreteModel()

m.i = pyo.RangeSet(1, 100)

# Indexed variables have all the behaviors of pyo.Var I expect (e.g., m.x.pprint() works)
m.x = NewVar(m.i, within=pyo.NonNegativeReals, bounds=(3, 18)) 

# Scalar variables don't work (e.g., m.y.pprint() doesn't work, and I get a TypeError if I try to build expressions using this sub-class)
m.y = NewVar(within=pyo.NonNegativeReals, bounds=(0, 102))

I tried to look at how __new__ works, but this is stretching my level of understanding (and I may be doing things I shouldn't be doing 😄). Would appreciate any guidance you can provide on how to sub-class to get this working!

---

Update:

Actually, after some more fiddling with __new__, I think I figured something out, but would be good if someone had ideas/suggestions if I'm doing something bad. Thanks again!

import pyomo.environ as pyo

class NewVar(pyo.Var):
    def __new__(cls, *args, a=0, **kwargs):
        obj = pyo.Var(*args, **kwargs)
        obj.a = a
        return obj


class NewExpression(pyo.Expression):
    def __new__(cls, *args, a=0, **kwargs):
        obj = pyo.Expression(*args, **kwargs)
        obj.a = a
        return obj