The Function Representation

import jax.numpy as jnp
import plotly.graph_objects as go

from lcm import AgeGrid, LinSpacedGrid, Model, Regime, categorical
from lcm.typing import ContinuousAction, ContinuousState, FloatND

blue, orange, green = "#4C78A8", "#F58518", "#54A24B"


def utility(consumption: ContinuousAction, risk_aversion: float) -> FloatND:
    return consumption ** (1 - risk_aversion) / (1 - risk_aversion)


def next_wealth(
    wealth: ContinuousState,
    consumption: ContinuousAction,
    interest_rate: float,
) -> ContinuousState:
    return (1 + interest_rate) * (wealth - consumption)


def borrowing_constraint(
    consumption: ContinuousAction, wealth: ContinuousState
) -> FloatND:
    return consumption <= wealth


@categorical(ordered=False)
class RegimeId:
    working_life: int
    retirement: int


retirement_regime = Regime(
    transition=None,
    functions={"utility": utility},
    constraints={"borrowing_constraint": borrowing_constraint},
    actions={"consumption": LinSpacedGrid(start=1, stop=400, n_points=50)},
    states={"wealth": LinSpacedGrid(start=1, stop=400, n_points=10)},
)

working_life_regime = Regime(
    transition=lambda: RegimeId.retirement,
    functions={"utility": utility},
    constraints={"borrowing_constraint": borrowing_constraint},
    actions={"consumption": LinSpacedGrid(start=1, stop=400, n_points=50)},
    states={
        "wealth": LinSpacedGrid(start=1, stop=400, n_points=10),
    },
    state_transitions={
        "wealth": next_wealth,
    },
)

model = Model(
    description="Minimal consumption-savings model",
    ages=AgeGrid(start=25, stop=65, step="20Y"),
    regimes={"working_life": working_life_regime, "retirement": retirement_regime},
    regime_id_class=RegimeId,
)

params = {
    "discount_factor": 0.95,
    "risk_aversion": 1.5,
    "interest_rate": 0.04,
}

from lcm.regime_building.Q_and_F import _get_U_and_F

internal_regime = model.internal_regimes["retirement"]

u_and_f = _get_U_and_F(
    functions=internal_regime.solve_functions.functions,
    constraints=internal_regime.solve_functions.constraints,
)
u_and_f.__signature__

_u, _f = u_and_f(consumption=100.0, wealth=50.0, utility__risk_aversion=1.5)
print(f"Utility: {_u}, feasible: {_f}")

from lcm.utils.dispatchers import productmap

_variables = ("wealth", "consumption")
u_and_f_mapped = productmap(
    func=u_and_f, variables=_variables, batch_sizes=dict.fromkeys(_variables, 0)
)
grid_arrays = {name: g.to_jax() for name, g in internal_regime.grids.items()}
u, f = u_and_f_mapped(**grid_arrays, utility__risk_aversion=1.5)

V_arr = jnp.max(u, axis=1, where=f, initial=-jnp.inf)
wealth_grid = internal_regime.grids["wealth"].to_jax()

print(f"V_arr shape: {V_arr.shape} ({len(wealth_grid)} wealth grid points)")

fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=wealth_grid,
        y=V_arr,
        mode="markers",
        marker={"color": blue, "size": 8},
        name="Pre-calculated values",
    )
)
fig.update_layout(
    xaxis_title="Wealth (x)",
    yaxis_title="V(x)",
    width=600,
    height=400,
)
fig.show()

from lcm.regime_building.V import create_v_interpolation_info, get_V_interpolator

v_interpolation_info = create_v_interpolation_info(retirement_regime)

scalar_value_function = get_V_interpolator(
    v_interpolation_info=v_interpolation_info,
    state_prefix="next_",
    V_arr_name="V_arr",
)
scalar_value_function.__signature__

value_function = productmap(
    func=scalar_value_function,
    variables=("next_wealth",),
    batch_sizes={"next_wealth": 0},
)

wealth_points_new = jnp.array([10.0, 25.0, 75.0, 210.0, 300.0])
wealth_all = jnp.concatenate([wealth_grid, wealth_points_new])

V_via_func = value_function(next_wealth=wealth_all, V_arr=V_arr)

fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=wealth_grid,
        y=V_arr,
        mode="lines+markers",
        marker={"color": blue, "size": 8},
        line={"color": blue},
        name="Pre-calculated values (linear interpolation)",
    )
)
fig.add_trace(
    go.Scatter(
        x=wealth_all,
        y=V_via_func,
        mode="markers",
        marker={"color": orange, "size": 6},
        name="Function representation output",
    )
)
fig.update_layout(
    xaxis_title="Wealth (x)",
    yaxis_title="V(x)",
    width=700,
    height=400,
)
fig.show()

from lcm.regime_building.V import _get_lookup_function

lookup = _get_lookup_function(array_name="V_arr", axis_names=["wealth_index"])
print(f"Signature: {lookup.__signature__}")

# Look up values at indices 0, 2, 5
lookup(wealth_index=jnp.array([0, 2, 5]), V_arr=V_arr)

from lcm.regime_building.V import _get_coordinate_finder

wealth_grid = LinSpacedGrid(start=1, stop=400, n_points=10)

wealth_coordinate_finder = _get_coordinate_finder(
    in_name="wealth",
    grid=wealth_grid,
)
print(f"Signature: {wealth_coordinate_finder.__signature__}")

wealth_values = jnp.array([1.0, (1 + 45.333336) / 2, 390.0])
coords = wealth_coordinate_finder(wealth=wealth_values)

for w, c in zip(wealth_values, coords, strict=True):
    print(f"  wealth = {w:8.2f}  →  coordinate = {float(c):.4f}")

from lcm.regime_building.V import _get_interpolator

value_function_interpolator = _get_interpolator(
    name_of_values_on_grid="V_arr",
    axis_names=["wealth_index"],
)
print(f"Signature: {value_function_interpolator.__signature__}")

wealth_indices = wealth_coordinate_finder(wealth=wealth_values)
V_interpolations = value_function_interpolator(wealth_index=wealth_indices, V_arr=V_arr)

fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=wealth_grid.to_jax(),
        y=V_arr,
        mode="markers",
        marker={"color": blue, "size": 8},
        name="Pre-calculated values",
    )
)
fig.add_trace(
    go.Scatter(
        x=wealth_values,
        y=V_interpolations,
        mode="markers",
        marker={"color": orange, "size": 6},
        name="Interpolated values",
    )
)
fig.update_layout(
    xaxis_title="Wealth (x)",
    yaxis_title="V(x)",
    width=600,
    height=400,
)
fig.show()

space_info = create_v_interpolation_info(retirement_regime)

funcs = {}

print(f"Discrete states: {space_info.discrete_states}")

# Step 1: Discrete lookup — identity (no discrete states to index by)
def discrete_lookup(V_arr):
    return V_arr


funcs["__interpolation_data__"] = discrete_lookup

# Step 2: Coordinate finder for wealth
from lcm.grids.coordinates import get_linspace_coordinate


def wealth_coordinate_finder(wealth):
    return get_linspace_coordinate(value=wealth, start=1, stop=400, n_points=10)


funcs["__wealth_coord__"] = wealth_coordinate_finder

# Step 3: Interpolator using map_coordinates
from lcm.regime_building.ndimage import map_coordinates


def interpolator(__interpolation_data__, __wealth_coord__):
    coordinates = jnp.array([__wealth_coord__])
    return map_coordinates(input=__interpolation_data__, coordinates=coordinates)


funcs["__fval__"] = interpolator

# Compose with dags
from dags import concatenate_functions

value_function = concatenate_functions(functions=funcs, targets="__fval__")
print(f"Composed signature: {value_function.__signature__}")

V_evaluated = value_function(wealth=wealth_grid.to_jax(), V_arr=V_arr)

fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=wealth_grid.to_jax(),
        y=V_arr,
        mode="markers",
        marker={"color": blue, "size": 8},
        name="Pre-calculated values",
    )
)
fig.add_trace(
    go.Scatter(
        x=wealth_grid.to_jax(),
        y=V_evaluated,
        mode="markers",
        marker={"color": orange, "size": 6},
        name="Re-implemented function representation",
    )
)
fig.update_layout(
    xaxis_title="Wealth (x)",
    yaxis_title="V(x)",
    width=600,
    height=400,
)
fig.show()