Source code for aurora.core.optimization

# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias

"""Optimization algorithms for Aurora-GLM."""

from __future__ import annotations

from .base import Optimizer
from .irls import irls
from .lbfgs import lbfgs
from .newton import modified_newton, newton_raphson
from .result import OptimizationResult
from .sparse_solvers import solve_sparse_penalized_ls

__all__ = [
    "OptimizationResult",
    "Optimizer",
    "newton_raphson",
    "modified_newton",
    "irls",
    "lbfgs",
    "optimize",
    "solve_sparse_penalized_ls",
]




def optimize(
    loss_fn,
    init_params,
    method: str = "lbfgs",
    backend: str = "jax",
    **kwargs,
):
    """Unified optimization interface."""
    from ..backends import get_backend

    backend_obj = get_backend(backend)

    optimizers = {
        "newton": newton_raphson,
        "newton-raphson": newton_raphson,
        "modified-newton": modified_newton,
        "levenberg-marquardt": modified_newton,
        "irls": irls,
        "lbfgs": lbfgs,
        "l-bfgs": lbfgs,
    }

    key = method.lower()
    if key not in optimizers:
        raise ValueError(f"Unknown optimization method: {method}. Available: {list(optimizers)}")

    optimizer_fn = optimizers[key]
    return optimizer_fn(loss_fn, init_params, backend=backend_obj, **kwargs)