finhjb 0.1.4

This is a Python library for solving Hamilton-Jacobi-Bellman (HJB) equations using JAX. It provides a framework for numerical solutions to dynamic optimization problems in finance and economics.

FinHJB

简体中文 README

FinHJB is a Python library for solving one-dimensional Hamilton-Jacobi-Bellman (HJB) equations with JAX.

It provides typed interfaces for model construction and a high-level solver API for:

• policy iteration,
• boundary update/search,
• parameter continuation (sensitivity analysis),
• result serialization and reloading.

Installation

• Python: >=3.10
• Recommended: uv

uv sync

For editable installation with pip:

pip install -e .

Quick Start

from dataclasses import dataclass

import jax.numpy as jnp
from jaxtyping import Array

import finhjb as fjb


class Parameter(fjb.AbstractParameter):
    r: float = 0.03
    sigma: float = 0.15


class PolicyDict(fjb.AbstractPolicyDict):
    investment: Array


@dataclass
class Boundary(fjb.AbstractBoundary[Parameter]):
    @staticmethod
    def compute_v_left(p: Parameter) -> float:
        return 0.0

    @staticmethod
    def compute_v_right(p: Parameter, s_max: float) -> float:
        return 1.0 + 0.1 * s_max


@dataclass
class Policy(fjb.AbstractPolicy[Parameter, PolicyDict]):
    @staticmethod
    def initialize(grid: fjb.Grid, p: Parameter) -> PolicyDict:
        return PolicyDict(investment=jnp.full_like(grid.s, 0.1))

    @staticmethod
    @fjb.explicit_policy(order=1)
    def update_investment(grid: fjb.Grid) -> fjb.Grid:
        grid.policy["investment"] = jnp.maximum(1e-6, 0.5 * grid.v / grid.dv)
        return grid


@dataclass
class Model(fjb.AbstractModel[Parameter, PolicyDict]):
    @staticmethod
    def hjb_residual(v, dv, d2v, s, policy, jump, boundary, p):
        inv = policy["investment"]
        return -p.r * v + (s - inv) * dv + 0.5 * p.sigma**2 * d2v


solver = fjb.Solver(
    boundary=Boundary(p=Parameter(), s_min=0.0, s_max=0.2),
    model=Model(policy=Policy()),
    policy_guess=True,
    number=400,
    config=fjb.Config(pi_method="scan", derivative_method="central"),
)

state, history = solver.solve()
print(state.converged, state.best_error)

Main APIs

Top-level exports include:

• Config
• Solver
• Grid, Grids, ImmutableBoundary
• AbstractBoundary, BoundaryConditionTarget
• AbstractModel, AbstractParameter
• AbstractPolicy, AbstractPolicyDict
• AbstractValueGuess, LinearInitialValue, QuadraticInitialValue
• explicit_policy, implicit_policy
• load_grid, load_grids, load_sensitivity_result

Solver Workflows

• Solve: state, history = solver.solve()
• Boundary update (model must implement update_boundary(grid)): state, history = solver.boundary_update()
• Boundary search: state = solver.boundary_search(method="hybr")
• Sensitivity analysis: result = solver.sensitivity_analysis(method="hybr", param_name="sigma", param_values=...)

Save / Load

state.grid.save("solution_grid")
loaded_grid = fjb.load_grid("solution_grid")

Similarly:

• grids.save(path) + fjb.load_grids(path)
• result.save(path) + fjb.load_sensitivity_result(path)

Configuration Highlights

Config controls derivative rules and convergence behavior:

• derivative_method: central | forward | backward
• pi_method: scan | anderson
• pe_*, pi_*, bs_* tolerances and iteration limits
• aa_* settings for Anderson acceleration

Testing

uv run pytest

Coverage gate (configured in project settings):

uv run pytest --cov=src/finhjb --cov-fail-under=85

Documentation

• English docs: docs/en/index.md
• Chinese docs: docs/zh/index.md