Allow general discrete grids

[timmens]

In this PR, I add the functionality to allow general discrete grids.

This entails two major updates:

1. Previously, only index grids were allowed. Now an updating step is added, which takes a model containing an arbitrary discrete grid, overwrites the grid with a corresponding index grid, and adds a function to the model that transforms the index values to the correct label. This means that the user can now use arbitrary integer grids, while internally, we can assume that all grids are integer grids.
2. We now force users to define discrete grids by instantiating DiscreteGrid with a dataclass that provides a mapping between a category name and its code. See the new syntax below.

Closes #82, #87.

New Syntax

Note

For most use cases, the codes (the field values of the dataclass) should be indices; however, with this PR it is now possible to use non-index codes; we therefore showcase this feature below

from dataclasses import dataclass
from lcm import DiscreteGrid

@dataclass
class HealthState:
   bad: int = -1
   neutral: int = 0
   good: int = 1
   
grid = DiscreteGrid(HealthState)

grid.categories  # = ["bad", "neutral", "good"]
grid.codes  # = [-1, 0, 1]
grid.to_jax()  # jnp.array([-1, 0, 1])

[hmgaudecker]

Great start, path forward as discussed in person!

[timmens]

Should provide only the following option to create discrete grids:

@dataclass
class RetirementStatus:
    working: int = 0
    retired: int = 1
    
DiscreteGrid(RetirementStatus)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[hmgaudecker]

As discussed, really nice!

[hmgaudecker]

Just had a brief look -- very nice! Need to stop now.

Apologies for yet another round of renaming suggestions, but looking at it I wasn't certain that a newcomer would immediately understand what is expected where.

[ChristianZimpelmann]

Very nice! I really see how the OOP features are beneficial.

Just a few questions or comments.

[hmgaudecker]

Excellent, thanks!

[hmgaudecker]

Just to check whether you really need the two functions here? I.e., with and without underscore? Can't we just decorate the function itself or does that lead to issues because it takes an argument?

[timmens]

Yes, the argument becomes a problem. While pytest supports fixtures that have arguments, it seems like you have to parametrize it via the parametrize decorator, and that would not work in our case, because we need to use different versions of it in the same test.