The Function Representation

In dynamic programming, the value function $V_T(x)$ is computed on a discrete grid but must be evaluated at arbitrary points when solving earlier periods. The function representation turns a pre-computed array $V^\text{arr}_T$ into a callable function that:

Accepts named arguments (e.g., wealth=150.0)
Returns exact values at grid points
Linearly interpolates between grid points

This notebook explains how it works, using a minimal terminal-regime model.

The two steps¶

Converting an array into a callable function requires two things:

Coordinate finding — For continuous variables, convert physical values (e.g., wealth = 150) to generalized coordinates (fractional indices into the grid). See the interpolation notebook for details. Discrete variables use integer codes that directly serve as array indices.
Interpolation — Use the generalized coordinates with map_coordinates to linearly interpolate between grid points.

Worked example¶

We set up a minimal model with a single terminal regime: a retiree choosing consumption given wealth, with CRRA utility. The wealth grid is intentionally coarse (10 points) to clearly show the interpolation behavior.

import inspect

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, ScalarInt

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: ScalarInt
    retirement: ScalarInt


retirement_regime = Regime(
    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(
    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,
    },
)

edges = {"working_life": {"retirement": 25}}

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

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

Computing the last-period value function array¶

In the terminal period, the value function is the maximum of utility over feasible actions. We use the internal regime representation to access the compiled functions and grids.

from _lcm.regime_building.Q_and_F import _get_U_and_F

regime = model._regimes["retirement"]  # noqa: SLF001

u_and_f = _get_U_and_F(
    functions=regime.solution.functions,
    constraints=regime.solution.constraints,
)
inspect.signature(u_and_f)

<Signature (consumption: 'ContinuousAction', utility__risk_aversion: 'float', wealth: 'ContinuousState') -> ('FloatND', 'FloatND')>

The function returns (utility, feasibility) for scalar inputs:

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

Utility: -0.2, feasible: False

To evaluate on the full state-action grid, we use productmap:

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 regime.solution.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 = regime.solution.grids["wealth"].to_jax()

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

V_arr shape: (10,) (10 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()

Creating the function representation¶

The function representation turns V_arr into a callable that can be evaluated at any wealth value. The V_arr_name argument sets the name of the array parameter in the resulting function.

from _lcm.regime_building.V import create_v_interpolation_info, get_V_interpolator

v_interpolation_info = create_v_interpolation_info(user_regime=retirement_regime)

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

<Signature (V_arr: 'FloatND | IntND', next_wealth: 'FloatND') -> 'FloatND'>

This scalar function is then wrapped with productmap so it can evaluate on arrays:

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

Visualizing interpolation¶

We evaluate the function representation on the original grid points (which should match exactly) and on additional points between grid points (which are interpolated).

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()

The orange points from the function representation lie exactly on the blue line connecting the grid points. The function representation behaves like an analytical function corresponding to this piecewise linear interpolation.

Technical details¶

The function representation is assembled from three building blocks, each implemented as a small function with a carefully chosen signature. These functions are composed using dags.concatenate_functions.

Lookup function¶

Indexes into the value function array using named axes. This is important because dags.concatenate_functions matches functions by argument names.

from _lcm.regime_building.V import _get_lookup_function

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

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

Signature: (wealth_index: 'FloatND | IntND', V_arr: 'FloatND | IntND') -> 'FloatND'

Array([-2. , -0.22028813, -0.13457806], dtype=float32)

Coordinate finder¶

Converts physical values to generalized coordinates — fractional indices into the grid. For a linearly spaced grid [1, 45.3, 89.7, ...], the value 23.2 might correspond to coordinate 0.5 (halfway between indices 0 and 1).

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: {inspect.signature(wealth_coordinate_finder)}")

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}")

Signature: (wealth: 'FloatND') -> 'FloatND'
  wealth =     1.00  →  coordinate = 0.0000
  wealth =    23.17  →  coordinate = 0.5000
  wealth =   390.00  →  coordinate = 8.7744

Interpolator¶

Uses the generalized coordinates to linearly interpolate on the value function array via map_coordinates.

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: {inspect.signature(value_function_interpolator)}")

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

Signature: (V_arr: 'FloatND', wealth_index: 'FloatND') -> 'FloatND'

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()

Re-implementation from scratch¶

To understand how the pieces fit together, let’s re-implement the function representation manually using dags.concatenate_functions.

The general idea: create functions for array lookup, coordinate finding, and interpolation, each with signatures that declare their dependencies. Then let dags wire them together.

Steps¶

Discrete lookup — index into the array using discrete positions. With no discrete states in our model, this is the identity (returns the array unchanged)
Coordinate finder for each continuous state — maps values to fractional indices
Interpolator — uses coordinates to interpolate on the array

Implementation¶

space_info = create_v_interpolation_info(user_regime=retirement_regime)

funcs = {}

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

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=jnp.asarray(1.0),
        stop=jnp.asarray(400.0),
        n_points=jnp.asarray(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: {inspect.signature(value_function)}")

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

Composed signature: (V_arr, wealth)

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()

The orange points coincide perfectly with the blue grid points — our manual re-implementation matches pylcm’s built-in function representation.