ENH: Single action sets and state spaces

[hmgaudecker]

We currently have constructs like:

@dataclass
class LaborSupply:
    work: int = 0
    retire: int = 1

in order to be able to write LaborSupply.work.

Similarly, we have (in #191/ #160):

@dataclass
class RegimeID:
    working: int = 0
    dead: int = 1

and

def next_regime(period: int, n_periods: int) -> ScalarInt:
    certain_retirement = period >= n_periods - 2
    return jnp.where(certain_retirement, RegimeID.retirement, RegimeID.working)

The dataclasses really span the space one of these three types of variables can do. That we use that to not write out integers is one thing, but it would be great if we could also use it for type hinting like so:

def next_regime(period: int, n_periods: int) -> RegimeID:

We may be able to do that via a custom combination of Enums and classes, syntax would ideally look like:

@lcm.regime_space
class RegimeID:
    working: int
    dead: int

@lcm.one_action_set
class LaborSupply:
    work: int
    retire: int


@lcm.one_state_space
class LaggedLaborSupply:
    working: int
    retired: int

Whether we can use that to further restrict the type hints for state-related next-functions is a different story, but probably not crucial