optimise¶

def draw_correspondences(
    img: np.ndarray,
    dstpoints: np.ndarray,
    projpts: np.ndarray,
) -> np.ndarray:
    """Draw matching points (projected vs. desired) on a copy of the image."""
    display = img.copy()
    dstpoints = norm2pix(img.shape, dstpoints, True)
    projpts = norm2pix(img.shape, projpts, True)
    for pts, color in [(projpts, (255, 0, 0)), (dstpoints, (0, 0, 255))]:
        for point in pts:
            circle(display, fltp(point), 3, color, -1, LINE_AA)
    for point_a, point_b in zip(projpts, dstpoints):
        line(display, fltp(point_a), fltp(point_b), (255, 255, 255), 1, LINE_AA)
    return display

def make_objective(
    dstpoints: np.ndarray,
    keypoint_index: np.ndarray,
    shear_cost: float,
    rvec_slice: slice,
) -> Callable[[np.ndarray], float]:
    """Create an objective function for scipy.optimize.minimize.

    Args:
        dstpoints: Target points to match.
        keypoint_index: Index array for keypoint extraction.
        shear_cost: Cost coefficient for shear penalty.
        rvec_slice: Slice object for extracting rotation vector from params.

    Returns:
        Objective function suitable for scipy.optimize.minimize.

    """
    from ..keypoints import project_keypoints

    def objective(pvec: np.ndarray) -> float:
        ppts = project_keypoints(pvec, keypoint_index)
        error = np.sum((dstpoints - ppts) ** 2)
        if shear_cost > 0.0:
            rvec = pvec[rvec_slice]
            rotation_penalty = shear_cost * rvec[0] ** 2
            return error + rotation_penalty
        else:
            return error

    return objective

def optimise_params_jax(
    name: str,
    small: np.ndarray,
    dstpoints: np.ndarray,
    span_counts: list[int],
    params: np.ndarray,
    debug_lvl: int,
) -> np.ndarray:
    """Refine the parameter vector using JAX-accelerated L-BFGS-B optimization.

    Args:
        name: Image name for debug output.
        small: Downscaled image for visualization.
        dstpoints: Target points to match.
        span_counts: Number of keypoints per span.
        params: Initial parameter vector.
        debug_lvl: Debug verbosity level.

    Returns:
        Optimized parameter vector.

    """
    device = get_device(cfg.DEVICE)
    keypoint_index = make_keypoint_index(span_counts)

    # Use the base objective for initial value display
    objective = make_objective(
        dstpoints,
        keypoint_index,
        cfg.SHEAR_COST,
        slice(*cfg.RVEC_IDX),
    )

    print("  initial objective is", objective(params))
    if debug_lvl >= 1:
        projpts = project_keypoints(params, keypoint_index)
        display = draw_correspondences(small, dstpoints, projpts)
        debug_show(name, 4, "keypoints before", display)

    print(f"  optimizing {len(params)} parameters on {device.device_kind.upper()}...")

    start = dt.now()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        with jax.default_device(device):
            result = _run_jax_lbfgsb(dstpoints, keypoint_index, params)
    elapsed = (dt.now() - start).total_seconds()
    print(
        f"  optimization (L-BFGS-B + JAX autodiff) took {elapsed:.2f}s, {result.nfev} evals",
    )

    print(f"  final objective is {result.fun:.6f}")
    params = result.x

    if debug_lvl >= 1:
        _print_diagnostics(params, keypoint_index, small, dstpoints, name)

    return params

def optimise_params_scipy(
    name: str,
    small: np.ndarray,
    dstpoints: np.ndarray,
    span_counts: list[int],
    params: np.ndarray,
    debug_lvl: int,
    method: str,
) -> np.ndarray:
    """Refine the parameter vector using SciPy optimization.

    Args:
        name: Image name for debug output.
        small: Downscaled image for visualization.
        dstpoints: Target points to match.
        span_counts: Number of keypoints per span.
        params: Initial parameter vector.
        debug_lvl: Debug verbosity level.
        method: Optimization method (e.g., 'Powell', 'L-BFGS-B').

    Returns:
        Optimized parameter vector.

    """
    keypoint_index = make_keypoint_index(span_counts)
    objective = make_objective(
        dstpoints,
        keypoint_index,
        cfg.SHEAR_COST,
        slice(*cfg.RVEC_IDX),
    )

    print("  initial objective is", objective(params))
    if debug_lvl >= 1:
        projpts = project_keypoints(params, keypoint_index)
        display = draw_correspondences(small, dstpoints, projpts)
        debug_show(name, 4, "keypoints before", display)

    print(f"  optimizing {len(params)} parameters...")

    # Hint about JAX for L-BFGS-B
    if method == "L-BFGS-B":
        print(
            "OPT_METHOD='L-BFGS-B' can use JAX for fast autodiff. "
            "Install with: pip install page-dewarp[jax]",
            file=sys.stderr,
        )

    start = dt.now()
    result = minimize(
        objective,
        params,
        method=method,
        options={"maxiter": cfg.OPT_MAX_ITER},
    )
    elapsed = (dt.now() - start).total_seconds()
    print(f"  optimization ({method}) took {elapsed:.2f}s, {result.nfev} evals")

    print(f"  final objective is {result.fun:.6f}")
    params = result.x

    if debug_lvl >= 1:
        _print_diagnostics(params, keypoint_index, small, dstpoints, name)

    return params