Taichi BFGS (TIBFGS)

This package implements the Broyden-Fletcher-Goldfarb-Shanno (BFGS) non-linear optimization algorithm in Taichi, enabling massively-parallel solutions to non-convex optimization problems.

Installation

pip install tibfgs

Use

To use tibfgs, define an objective function as a ti.func, for example the Ackley function:

@ti.func
def ackley(x: ti.math.vec2) -> ti.f32:
    return (
        -20 * ti.exp(-0.2 * ti.sqrt(0.5 * x.norm_sqr()))
        - ti.exp(
            0.5 * ti.cos(2 * ti.math.pi * x.x) + 0.5 * ti.cos(2 * ti.math.pi * x.y)
        )
        + ti.math.e
        + 20
    )

Initialize \(n\) initial conditions as an \(n \times m\) Numpy array (where \(m\) is the dimension of the problem, in this case \(m=2\)):

n_particles = int(1e6)
x0 = 4 * np.random.rand(n_particles, 2) - 2

Run the BFGS optimizer in parallel on the GPU with:

df = tibfgs.minimize(
    ackley, x0, gtol=1e-3, eps=1e-5, discard_failures=False
)

Which returns a polars DataFrame of solutions with df.head() like:

shape: (5, 10)
┌───────┬──────────┬───────────────┬───────────────┬───┬───────┬───────┬────────────────────────┬────────────────────────┐
│ index ┆ fun      ┆ xk            ┆ x0            ┆ … ┆ feval ┆ geval ┆ gradient               ┆ hessian_inverse        │
│ ---   ┆ ---      ┆ ---           ┆ ---           ┆   ┆ ---   ┆ ---   ┆ ---                    ┆ ---                    │
│ u32   ┆ f32      ┆ array[f32, 2] ┆ array[f32, 2] ┆   ┆ i32   ┆ i32   ┆ array[f32, 2]          ┆ array[f32, (2, 2)]     │
╞═══════╪══════════╪═══════════════╪═══════════════╪═══╪═══════╪═══════╪════════════════════════╪════════════════════════╡
│ 0     ┆ 5.381865 ┆ [-0.982348,   ┆ [-0.968771,   ┆ … ┆ 39    ┆ 13    ┆ [0.190735, 0.190735]   ┆ [[0.617119,            │
│       ┆          ┆ 1.964832]     ┆ 1.981125]     ┆   ┆       ┆       ┆                        ┆ -0.479458], [-0.47…    │
│ 1     ┆ 0.000011 ┆ [-5.7371e-7,  ┆ [1.487202,    ┆ … ┆ 147   ┆ 49    ┆ [1.716614, 0.38147]    ┆ [[0.000008, 0.000002], │
│       ┆          ┆ -0.000005]    ┆ -1.334406]    ┆   ┆       ┆       ┆                        ┆ [0.0000…               │
│ 2     ┆ 3.574453 ┆ [-0.968429,   ┆ [-0.806571,   ┆ … ┆ 36    ┆ 12    ┆ [0.0, 0.0]             ┆ [[0.016302, -0.00451], │
│       ┆          ┆ -0.968395]    ┆ -0.820057]    ┆   ┆       ┆       ┆                        ┆ [-0.004…               │
│ 3     ┆ 4.884083 ┆ [0.000768,    ┆ [-0.342866,   ┆ … ┆ 33    ┆ 11    ┆ [0.0, 0.0]             ┆ [[0.020406, 0.000647], │
│       ┆          ┆ -1.959207]    ┆ -1.756427]    ┆   ┆       ┆       ┆                        ┆ [0.0006…               │
│ 4     ┆ 0.000011 ┆ [0.000003,    ┆ [1.534462,    ┆ … ┆ 111   ┆ 37    ┆ [2.861023, 2.861023]   ┆ [[0.000299,            │
│       ┆          ┆ 0.000003]     ┆ 0.979372]     ┆   ┆       ┆       ┆                        ┆ -0.000029], [-0.00…    │
└───────┴──────────┴───────────────┴───────────────┴───┴───────┴───────┴────────────────────────┴────────────────────────┘

The full schema of this dataframe is:

┌─────────────────┬──────────────────────────────┐
│ Column          ┆ Type                         │
│ ---             ┆ ---                          │
│ str             ┆ object                       │
╞═════════════════╪══════════════════════════════╡
│ index           ┆ UInt32                       │
│ fun             ┆ Float32                      │
│ xk              ┆ Array(Float32, shape=(2,))   │
│ x0              ┆ Array(Float32, shape=(2,))   │
│ message         ┆ String                       │
│ iterations      ┆ UInt16                       │
│ feval           ┆ Int32                        │
│ geval           ┆ Int32                        │
│ gradient        ┆ Array(Float32, shape=(2,))   │
│ hessian_inverse ┆ Array(Float32, shape=(2, 2)) │
└─────────────────┴──────────────────────────────┘

Where:

• index is the particle initial condition index in x0

• fun is the objective function value at the converged solution

• xk is the converged solution

• x0 is the provided initial guess

• message describes the convergence state

• iterations is the number of state update iterations used to reach convergence

• feval is the number of objective function evaluations (including evaluations during gradient evaluation) over all iterations

• geval is the number of gradient evaluations across all iterations

• gradient is the gradient at the converged xk

• hessian_inverse is the inverse of the Hessian matrix maintained by BFGS internally, at the converged xk