bundle_adjustment

Bundle adjustment over plain arrays, rig-free (ADR-0001).

build_state packs per-view rvecs, tvecs, intrs, dists and the observed 2D into a flat state, with a small grammar for fixing and sharing parameters across views, and bundle_adjust solves it; it depends only on deeperfly.geometry.

BASolution

Bases: NamedTuple

Optimized parameters unpacked back into per-view arrays plus points.

Source code in src/deeperfly/bundle_adjustment/core.py

class BASolution(NamedTuple):
    """Optimized parameters unpacked back into per-view arrays plus points."""

    rvecs: np.ndarray
    tvecs: np.ndarray
    intrs: np.ndarray
    dists: np.ndarray
    pts3d: np.ndarray

BAState

Bases: NamedTuple

Packed bundle-adjustment state ready for bundle_adjust.

values is the flat parameter vector and fixed the mask held constant; the *_idx arrays recover the per-view arrays and may alias to share. pts2d is the observations, and mirrored, the apparatus's and never solved, rides beside them.

Source code in src/deeperfly/bundle_adjustment/state.py

class BAState(NamedTuple):
    """Packed bundle-adjustment state ready for `bundle_adjust`.

    `values` is the flat parameter vector and `fixed` the mask held constant; the
    `*_idx` arrays recover the per-view arrays and may alias to share. `pts2d` is the
    observations, and `mirrored`, the apparatus's and never solved, rides beside them.
    """

    values: Float[np.ndarray, "n_params"]
    fixed: Bool[np.ndarray, "n_params"]
    rvecs_idx: Int[np.ndarray, "V 3"]
    tvecs_idx: Int[np.ndarray, "V 3"]
    intrs_idx: Int[np.ndarray, "V P"]
    dists_idx: Int[np.ndarray, "V K"]
    pts3d_idx: Int[np.ndarray, "N 3"]
    pts2d: Float[np.ndarray, "V N 2"]
    mirrored: Bool[np.ndarray, "V"]

build_state

build_state(
    rvecs: Float[ndarray, "V 3"],
    tvecs: Float[ndarray, "V 3"],
    intrs: Float[ndarray, "V P"],
    dists: Float[ndarray, "V K"],
    mirrored: Bool[ndarray, V],
    pts2d: Float[ndarray, "V N 2"],
    names: list[str] | None = None,
    *,
    fixed: Sequence[str] = (),
    shared: Sequence[Sequence[str]] = (),
    pts3d: Float[ndarray, "N 3"] | None = None,
) -> BAState

Build a BAState from per-view arrays and fix/share specs.

Intrinsics and distortion are per view unless shared. names (default "0", "1"...) resolves references and must be one unique name per view. A reference is "<view>.<param>" or "<view>.<param>[i]", <view> a name or * and <param> one of rvec, tvec, intr (alias kmat) or dist. Each shared group ties its members element-wise into one slot, so fixing one fixes all, and a clash of initial values keeps the lowest slot's, with a warning. pts3d defaults to a triangulation.

Raises:

Type	Description
ValueError	If the per-view arrays disagree on the view count, names is not one unique name per view, or a reference is malformed or unknown.

Source code in src/deeperfly/bundle_adjustment/state.py

def build_state(
    rvecs: Float[np.ndarray, "V 3"],
    tvecs: Float[np.ndarray, "V 3"],
    intrs: Float[np.ndarray, "V P"],
    dists: Float[np.ndarray, "V K"],
    mirrored: Bool[np.ndarray, "V"],
    pts2d: Float[np.ndarray, "V N 2"],
    names: list[str] | None = None,
    *,
    fixed: Sequence[str] = (),
    shared: Sequence[Sequence[str]] = (),
    pts3d: Float[np.ndarray, "N 3"] | None = None,
) -> BAState:
    """Build a [BAState][] from per-view arrays and fix/share specs.

    Intrinsics and distortion are per view unless `shared`. `names` (default `"0"`,
    `"1"`...) resolves references and must be one unique name per view. A reference
    is `"<view>.<param>"` or `"<view>.<param>[i]"`, `<view>` a name or `*` and
    `<param>` one of `rvec`, `tvec`, `intr` (alias `kmat`) or `dist`. Each `shared`
    group ties its members element-wise into one slot, so fixing one fixes all, and a
    clash of initial values keeps the lowest slot's, with a warning. `pts3d` defaults
    to a triangulation.

    Raises:
        ValueError: If the per-view arrays disagree on the view count, `names` is not
            one unique name per view, or a reference is malformed or unknown.
    """
    rvecs, tvecs, intrs, dists = (
        np.asarray(a, dtype=float) for a in (rvecs, tvecs, intrs, dists)
    )
    mirrored = np.asarray(mirrored, dtype=bool)
    pts2d = np.asarray(pts2d, dtype=float)
    n_views = len(rvecs)
    for label, arr in (
        ("tvecs", tvecs),
        ("intrs", intrs),
        ("dists", dists),
        ("mirrored", mirrored),
        ("pts2d", pts2d),
    ):
        if len(arr) != n_views:
            raise ValueError(f"{label} has {len(arr)} rows but rvecs has {n_views}")
    if names is None:
        names = [str(i) for i in range(n_views)]
    elif len(names) != n_views:
        # A short list would silently narrow `*`, leaving `"*.intr"` partly free.
        raise ValueError(f"got {len(names)} view names for {n_views} views: {names}")
    elif len(set(names)) != n_views:
        dupes = sorted({n for n in names if names.count(n) > 1})
        raise ValueError(f"duplicate view names {dupes}")
    name_to_row = {name: i for i, name in enumerate(names)}

    if pts3d is None:
        pts3d = np.asarray(
            initialize_pts3d(pts2d, rvecs, tvecs, intrs, mirrored), dtype=float
        )
    else:
        pts3d = np.asarray(pts3d, dtype=float)

    arrs = [rvecs, tvecs, intrs, dists, pts3d]
    values = np.concatenate([a.ravel() for a in arrs])
    offsets = np.cumsum([0, *(a.size for a in arrs)])

    def make_idx(arr, off):
        return np.arange(arr.size).reshape(arr.shape) + off

    rvecs_idx = make_idx(rvecs, offsets[0])
    tvecs_idx = make_idx(tvecs, offsets[1])
    intrs_idx = make_idx(intrs, offsets[2])
    dists_idx = make_idx(dists, offsets[3])
    pts3d_idx = make_idx(pts3d, offsets[4])

    param_idx = {
        "rvec": rvecs_idx,
        "tvec": tvecs_idx,
        "intr": intrs_idx,
        "dist": dists_idx,
    }

    def ref_indices(ref: str) -> list[np.ndarray]:
        """Resolve one reference to a flat-index array per matched view."""
        view, sep, rest = ref.partition(".")
        if not sep or not rest:
            raise ValueError(f"malformed reference {ref!r}; expected 'view.param'")
        if rest.endswith("]"):
            param_name, _, idx_str = rest[:-1].partition("[")
            index = int(idx_str)
        else:
            param_name, index = rest, None
        if param_name not in _PARAM_ALIASES:
            raise ValueError(
                f"unknown parameter {param_name!r} in {ref!r}; "
                f"expected one of {sorted(_PARAM_ALIASES)}"
            )
        group = _PARAM_ALIASES[param_name]
        matched = names if view == "*" else [view]
        out = []
        for c in matched:
            if c not in name_to_row:
                raise ValueError(f"unknown view {c!r} in {ref!r}")
            row = param_idx[group][name_to_row[c]]
            if index is None:
                out.append(row)
            else:
                if not 0 <= index < row.size:
                    raise ValueError(f"index {index} out of range in {ref!r}")
                out.append(row[index : index + 1])
        return out

    fixed_mask = np.zeros(values.size, dtype=bool)
    for ref in fixed:
        for idxs in ref_indices(ref):
            fixed_mask[idxs] = True

    # Union-find over slots; each shared group ties its members element-wise.
    parent = list(range(values.size))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        parent[find(a)] = find(b)

    conflicts: list[tuple[object, list[float], int]] = []
    for group in shared:
        members = [arr for ref in group for arr in ref_indices(ref)]
        if not members:
            continue
        size = members[0].size
        if any(m.size != size for m in members):
            raise ValueError(
                f"shared group {group} mixes references of different lengths"
            )
        for j in range(size):
            idxs = [int(m[j]) for m in members]
            vals = values[idxs]
            if not np.allclose(vals, vals[0]):
                # The surviving value depends on the whole component, which later
                # groups may extend, so the clash is warned about after compaction.
                loc = group if size == 1 else f"{group} element {j}"
                conflicts.append((loc, vals.tolist(), idxs[0]))
            for m in members[1:]:
                union(idxs[0], int(m[j]))

    # Compact: collapse each connected component to a single slot.
    roots = np.array([find(i) for i in range(values.size)])
    uniq, old_to_new = np.unique(roots, return_inverse=True)
    old_to_new = old_to_new.reshape(-1)  # numpy version-proof: keep 1D
    new_values = np.empty(uniq.size)
    new_fixed = np.zeros(uniq.size, dtype=bool)
    filled = np.zeros(uniq.size, dtype=bool)
    for i in range(values.size):
        k = old_to_new[i]
        if not filled[k]:  # smallest original index wins the initial value
            new_values[k] = values[i]
            filled[k] = True
        new_fixed[k] |= fixed_mask[i]

    for loc, vals, member in conflicts:
        warnings.warn(
            f"shared {loc} have differing initial values {vals}; "
            f"keeping {new_values[old_to_new[member]]}",
            stacklevel=2,
        )

    return BAState(
        new_values,
        new_fixed,
        old_to_new[rvecs_idx],
        old_to_new[tvecs_idx],
        old_to_new[intrs_idx],
        old_to_new[dists_idx],
        old_to_new[pts3d_idx],
        pts2d,
        mirrored,
    )

bundle_adjust

bundle_adjust(
    values: Float[ndarray, n_params],
    fixed: Bool[ndarray, n_params],
    rvecs_idx: Int[ndarray, "V 3"],
    tvecs_idx: Int[ndarray, "V 3"],
    intrs_idx: Int[ndarray, "V P"],
    dists_idx: Int[ndarray, "V K"],
    pts3d_idx: Int[ndarray, "N 3"],
    pts2d: Float[ndarray, "V N 2"],
    mirrored: Bool[ndarray, V],
    loss: str = "linear",
    f_scale: float = 1.0,
    max_nfev: int = 1000,
    weights: Float[ndarray, "V N"] | None = None,
    bone_pairs: Int[ndarray, "B 2"] | None = None,
    bone_targets: Float[ndarray, B] | None = None,
    bone_weight: float = 1.0,
    **kwargs: Any,
) -> tuple[OptimizeResult, BASolution]

Bundle adjustment with a JAX-computed analytic Jacobian.

The low-level solver over the packed state of build_state (see its BAState). loss and f_scale go to scipy.optimize.least_squares, with huber, cauchy and arctan translated to IRLS form (_IRLS_LOSSES). weights (V, N) scale each residual by sqrt(weight); a negative or non-finite one drops the observation, and all zero falls back to uniform with a warning. bone_pairs, bone_targets and bone_weight append a soft bone-length prior bone_weight * (||p_i - p_j|| - target). kwargs go to scipy; x_scale defaults to "jac" (ADR-0057).

Returns:

Type	Description
OptimizeResult	(result, solution): scipy's OptimizeResult and a BASolution unpacked
BASolution	into per-view arrays and points.

Source code in src/deeperfly/bundle_adjustment/core.py

def bundle_adjust(
    values: Float[np.ndarray, "n_params"],
    fixed: Bool[np.ndarray, "n_params"],
    rvecs_idx: Int[np.ndarray, "V 3"],
    tvecs_idx: Int[np.ndarray, "V 3"],
    intrs_idx: Int[np.ndarray, "V P"],
    dists_idx: Int[np.ndarray, "V K"],
    pts3d_idx: Int[np.ndarray, "N 3"],
    pts2d: Float[np.ndarray, "V N 2"],
    mirrored: Bool[np.ndarray, "V"],
    loss: str = "linear",
    f_scale: float = 1.0,
    max_nfev: int = 1000,
    weights: Float[np.ndarray, "V N"] | None = None,
    bone_pairs: Int[np.ndarray, "B 2"] | None = None,
    bone_targets: Float[np.ndarray, "B"] | None = None,
    bone_weight: float = 1.0,
    **kwargs: Any,
) -> tuple[OptimizeResult, BASolution]:
    """Bundle adjustment with a JAX-computed analytic Jacobian.

    The low-level solver over the packed state of
    [build_state][deeperfly.bundle_adjustment.state.build_state] (see its
    [BAState][]). `loss` and `f_scale` go to `scipy.optimize.least_squares`, with
    huber, cauchy and arctan translated to IRLS form (`_IRLS_LOSSES`). `weights`
    `(V, N)` scale each residual by `sqrt(weight)`; a negative or non-finite one drops
    the observation, and all zero falls back to uniform with a warning. `bone_pairs`,
    `bone_targets` and `bone_weight` append a soft bone-length prior
    `bone_weight * (||p_i - p_j|| - target)`. `kwargs` go to scipy; `x_scale`
    defaults to `"jac"` (ADR-0057).

    Returns:
        `(result, solution)`: scipy's `OptimizeResult` and a [BASolution][] unpacked
        into per-view arrays and points.
    """
    obs_view, obs_pt = np.where(np.isfinite(pts2d).all(axis=-1))
    pts2d_observed = pts2d[obs_view, obs_pt]
    n_obs = len(obs_view)
    n_free = int((~fixed).sum())

    if weights is None:
        sqrt_w = np.ones(n_obs)
    else:
        w = np.asarray(weights, dtype=float)[obs_view, obs_pt]
        # A junk weight just drops its observation: zero residual and Jacobian rows.
        w = np.where(np.isfinite(w), np.maximum(w, 0.0), 0.0)
        if w.size and not w.any():
            # Every observation weighs zero, which would leave the solve ill-posed,
            # so confidence cannot disable it silently: fall back to uniform weights.
            warnings.warn(
                "all confidence weights are zero (or non-positive); "
                "falling back to uniform weights for bundle adjustment",
                stacklevel=2,
            )
            w = np.ones(n_obs)
        sqrt_w = np.sqrt(w)

    use_bones = bone_pairs is not None and len(bone_pairs) > 0
    if use_bones:
        if bone_targets is None:
            raise ValueError("bone_targets is required when bone_pairs is given")
        bone_pairs = np.asarray(bone_pairs)
        bone_targets = np.asarray(bone_targets, dtype=float)
        bone_i, bone_j = bone_pairs[:, 0], bone_pairs[:, 1]
        n_bones = len(bone_pairs)
    else:
        bone_i = bone_j = np.zeros(0, dtype=int)
        n_bones = 0

    x0 = values[~fixed]
    rvecs = values[rvecs_idx]
    tvecs = values[tvecs_idx]
    intrs = values[intrs_idx]
    dists = values[dists_idx]
    pts3d = values[pts3d_idx]

    free_cumsum = np.cumsum(~fixed)

    def free_assign(idx):
        assign_mask = ~fixed[idx]
        source_idx = free_cumsum[idx][assign_mask] - 1
        return assign_mask, source_idx

    rvecs_assign = free_assign(rvecs_idx)
    tvecs_assign = free_assign(tvecs_idx)
    intrs_assign = free_assign(intrs_idx)
    dists_assign = free_assign(dists_idx)
    pts3d_assign = free_assign(pts3d_idx)

    def unpack(x):
        rvecs[rvecs_assign[0]] = x[rvecs_assign[1]]
        tvecs[tvecs_assign[0]] = x[tvecs_assign[1]]
        intrs[intrs_assign[0]] = x[intrs_assign[1]]
        dists[dists_assign[0]] = x[dists_assign[1]]
        pts3d[pts3d_assign[0]] = x[pts3d_assign[1]]
        return rvecs, tvecs, intrs, dists, pts3d

    obs_mirrored = np.asarray(mirrored, dtype=bool)[obs_view]

    def _project_obs(rvecs, tvecs, intrs, dists, pts3d):
        # Pass NumPy straight in and let the jit kernel convert once, rather than
        # pre-staging intermediate JAX arrays.
        return _project_per_obs(
            pts3d[obs_pt],
            rvecs[obs_view],
            tvecs[obs_view],
            intrs[obs_view],
            dists[obs_view],
            obs_mirrored,
        )

    n_resid = 2 * n_obs + n_bones

    def residuals(x):
        rvecs, tvecs, intrs, dists, pts3d = unpack(x)
        pts2d_predicted = _project_obs(rvecs, tvecs, intrs, dists, pts3d)
        reproj = (
            (np.asarray(pts2d_predicted) - pts2d_observed) * sqrt_w[:, None]
        ).ravel()
        if not use_bones:
            return reproj
        lengths = np.asarray(_bone_len_per(pts3d[bone_i], pts3d[bone_j]))
        bone_resid = bone_weight * (lengths - bone_targets)
        return np.concatenate([reproj, bone_resid])

    # Sparsity: each observation's two rows depend on its 3D point and its view's rvec,
    # tvec, intr and dist slots, in `_jac_per_obs`'s order.
    free_idx_map = np.full(values.size, -1, dtype=np.int64)
    free_idx_map[~fixed] = np.arange(n_free)
    cols_per_obs = np.concatenate(
        [
            pts3d_idx[obs_pt],
            rvecs_idx[obs_view],
            tvecs_idx[obs_view],
            intrs_idx[obs_view],
            dists_idx[obs_view],
        ],
        axis=1,
    )
    free_cols_per_obs = free_idx_map[cols_per_obs]
    valid = free_cols_per_obs >= 0
    cols_nz = free_cols_per_obs[valid]
    obs_idx = np.broadcast_to(np.arange(n_obs)[:, None], cols_per_obs.shape)
    rows_x = 2 * obs_idx[valid]
    rows = np.concatenate([rows_x, rows_x + 1])
    cols = np.concatenate([cols_nz, cols_nz])

    # Bone-length sparsity. Each bone row b depends only on the 3 coordinate
    # slots of its two endpoints; its row index sits after the reprojection rows.
    if use_bones:
        bone_cols = np.concatenate(
            [pts3d_idx[bone_i], pts3d_idx[bone_j]], axis=1
        )  # (B,6)
        free_bone_cols = free_idx_map[bone_cols]
        bone_valid = free_bone_cols >= 0
        bone_rows_all = np.broadcast_to(
            (2 * n_obs + np.arange(n_bones))[:, None], bone_cols.shape
        )
        bone_rows = bone_rows_all[bone_valid]
        bone_cols_nz = free_bone_cols[bone_valid]
        rows = np.concatenate([rows, bone_rows])
        cols = np.concatenate([cols, bone_cols_nz])

    # The pattern is constant, so the CSR layout is built once and each iteration's
    # values are scattered into it with one np.bincount (summing duplicate entries).
    coo_key = rows.astype(np.int64) * n_free + cols
    uniq_key, coo_to_csr = np.unique(coo_key, return_inverse=True)
    coo_to_csr = coo_to_csr.reshape(-1)  # numpy version-proof: keep 1D
    nnz_csr = uniq_key.size
    csr_indices = (uniq_key % n_free).astype(np.int64)
    csr_rows = uniq_key // n_free
    csr_indptr = np.concatenate(
        [[0], np.cumsum(np.bincount(csr_rows, minlength=n_resid))]
    ).astype(np.int64)

    def jac(x):
        rvecs, tvecs, intrs, dists, pts3d = unpack(x)
        jac_blocks = _jac_per_obs(
            pts3d[obs_pt],
            rvecs[obs_view],
            tvecs[obs_view],
            intrs[obs_view],
            dists[obs_view],
            obs_mirrored,
        )
        # Each block is (n_obs, 2, *param_dim); concatenate along the last
        # axis to align with cols_per_obs above. Weight rows by sqrt(weight).
        jac_full = np.concatenate([np.asarray(j) for j in jac_blocks], axis=-1)
        jac_full = jac_full * sqrt_w[:, None, None]
        data = np.concatenate([jac_full[:, 0, :][valid], jac_full[:, 1, :][valid]])
        if use_bones:
            dpi, dpj = _bone_jac_per(pts3d[bone_i], pts3d[bone_j])
            bone_jac = bone_weight * np.concatenate(
                [np.asarray(dpi), np.asarray(dpj)], axis=1
            )  # (B, 6)
            data = np.concatenate([data, bone_jac[bone_valid]])
        csr_data = np.bincount(coo_to_csr, weights=data, minlength=nnz_csr)
        return csr_matrix((csr_data, csr_indices, csr_indptr), shape=(n_resid, n_free))

    kwargs.setdefault("x_scale", "jac")
    result = least_squares(
        residuals,
        x0,
        jac=jac,
        method="trf",
        tr_solver="lsmr",
        loss=_IRLS_LOSSES.get(loss, loss),
        f_scale=f_scale,
        max_nfev=max_nfev,
        **kwargs,
    )

    solution = BASolution(*unpack(result.x))
    return result, solution

initialize_pts3d

initialize_pts3d(
    pts2d: Float[ndarray, "V *pts 2"],
    rvecs: Float[ndarray, "V 3"],
    tvecs: Float[ndarray, "V 3"],
    intrs: Float[ndarray, "V P"] | Float[ndarray, P],
    mirrored: Bool[ndarray, V],
) -> Float[Array, "*pts 3"]

Triangulate initial (*pts, 3) points from the 2D and the view poses.

Source code in src/deeperfly/bundle_adjustment/state.py

def initialize_pts3d(
    pts2d: Float[np.ndarray, "V *pts 2"],
    rvecs: Float[np.ndarray, "V 3"],
    tvecs: Float[np.ndarray, "V 3"],
    intrs: Float[np.ndarray, "V P"] | Float[np.ndarray, "P"],
    mirrored: Bool[np.ndarray, "V"],
) -> Float[Array, "*pts 3"]:
    """Triangulate initial `(*pts, 3)` points from the 2D and the view poses."""
    return triangulate_dlt(pts2d, pmats(rvecs, tvecs, intrs, mirrored))

Packed parameter state for bundle adjustment.

The state concatenates [rvecs, tvecs, intrs, dists, pts3d], index arrays recover the shapes, and a fixed mask holds parameters constant. Two index entries may point at one slot, which is how parameters are shared; build_state exposes it through a reference grammar ("f.tvec[2]", "*.intr").

BAState

Bases: NamedTuple

Packed bundle-adjustment state ready for bundle_adjust.

values is the flat parameter vector and fixed the mask held constant; the *_idx arrays recover the per-view arrays and may alias to share. pts2d is the observations, and mirrored, the apparatus's and never solved, rides beside them.

Source code in src/deeperfly/bundle_adjustment/state.py

class BAState(NamedTuple):
    """Packed bundle-adjustment state ready for `bundle_adjust`.

    `values` is the flat parameter vector and `fixed` the mask held constant; the
    `*_idx` arrays recover the per-view arrays and may alias to share. `pts2d` is the
    observations, and `mirrored`, the apparatus's and never solved, rides beside them.
    """

    values: Float[np.ndarray, "n_params"]
    fixed: Bool[np.ndarray, "n_params"]
    rvecs_idx: Int[np.ndarray, "V 3"]
    tvecs_idx: Int[np.ndarray, "V 3"]
    intrs_idx: Int[np.ndarray, "V P"]
    dists_idx: Int[np.ndarray, "V K"]
    pts3d_idx: Int[np.ndarray, "N 3"]
    pts2d: Float[np.ndarray, "V N 2"]
    mirrored: Bool[np.ndarray, "V"]

build_state

build_state(
    rvecs: Float[ndarray, "V 3"],
    tvecs: Float[ndarray, "V 3"],
    intrs: Float[ndarray, "V P"],
    dists: Float[ndarray, "V K"],
    mirrored: Bool[ndarray, V],
    pts2d: Float[ndarray, "V N 2"],
    names: list[str] | None = None,
    *,
    fixed: Sequence[str] = (),
    shared: Sequence[Sequence[str]] = (),
    pts3d: Float[ndarray, "N 3"] | None = None,
) -> BAState

Build a BAState from per-view arrays and fix/share specs.

Intrinsics and distortion are per view unless shared. names (default "0", "1"...) resolves references and must be one unique name per view. A reference is "<view>.<param>" or "<view>.<param>[i]", <view> a name or * and <param> one of rvec, tvec, intr (alias kmat) or dist. Each shared group ties its members element-wise into one slot, so fixing one fixes all, and a clash of initial values keeps the lowest slot's, with a warning. pts3d defaults to a triangulation.

Raises:

Type	Description
ValueError	If the per-view arrays disagree on the view count, names is not one unique name per view, or a reference is malformed or unknown.

Source code in src/deeperfly/bundle_adjustment/state.py

def build_state(
    rvecs: Float[np.ndarray, "V 3"],
    tvecs: Float[np.ndarray, "V 3"],
    intrs: Float[np.ndarray, "V P"],
    dists: Float[np.ndarray, "V K"],
    mirrored: Bool[np.ndarray, "V"],
    pts2d: Float[np.ndarray, "V N 2"],
    names: list[str] | None = None,
    *,
    fixed: Sequence[str] = (),
    shared: Sequence[Sequence[str]] = (),
    pts3d: Float[np.ndarray, "N 3"] | None = None,
) -> BAState:
    """Build a [BAState][] from per-view arrays and fix/share specs.

    Intrinsics and distortion are per view unless `shared`. `names` (default `"0"`,
    `"1"`...) resolves references and must be one unique name per view. A reference
    is `"<view>.<param>"` or `"<view>.<param>[i]"`, `<view>` a name or `*` and
    `<param>` one of `rvec`, `tvec`, `intr` (alias `kmat`) or `dist`. Each `shared`
    group ties its members element-wise into one slot, so fixing one fixes all, and a
    clash of initial values keeps the lowest slot's, with a warning. `pts3d` defaults
    to a triangulation.

    Raises:
        ValueError: If the per-view arrays disagree on the view count, `names` is not
            one unique name per view, or a reference is malformed or unknown.
    """
    rvecs, tvecs, intrs, dists = (
        np.asarray(a, dtype=float) for a in (rvecs, tvecs, intrs, dists)
    )
    mirrored = np.asarray(mirrored, dtype=bool)
    pts2d = np.asarray(pts2d, dtype=float)
    n_views = len(rvecs)
    for label, arr in (
        ("tvecs", tvecs),
        ("intrs", intrs),
        ("dists", dists),
        ("mirrored", mirrored),
        ("pts2d", pts2d),
    ):
        if len(arr) != n_views:
            raise ValueError(f"{label} has {len(arr)} rows but rvecs has {n_views}")
    if names is None:
        names = [str(i) for i in range(n_views)]
    elif len(names) != n_views:
        # A short list would silently narrow `*`, leaving `"*.intr"` partly free.
        raise ValueError(f"got {len(names)} view names for {n_views} views: {names}")
    elif len(set(names)) != n_views:
        dupes = sorted({n for n in names if names.count(n) > 1})
        raise ValueError(f"duplicate view names {dupes}")
    name_to_row = {name: i for i, name in enumerate(names)}

    if pts3d is None:
        pts3d = np.asarray(
            initialize_pts3d(pts2d, rvecs, tvecs, intrs, mirrored), dtype=float
        )
    else:
        pts3d = np.asarray(pts3d, dtype=float)

    arrs = [rvecs, tvecs, intrs, dists, pts3d]
    values = np.concatenate([a.ravel() for a in arrs])
    offsets = np.cumsum([0, *(a.size for a in arrs)])

    def make_idx(arr, off):
        return np.arange(arr.size).reshape(arr.shape) + off

    rvecs_idx = make_idx(rvecs, offsets[0])
    tvecs_idx = make_idx(tvecs, offsets[1])
    intrs_idx = make_idx(intrs, offsets[2])
    dists_idx = make_idx(dists, offsets[3])
    pts3d_idx = make_idx(pts3d, offsets[4])

    param_idx = {
        "rvec": rvecs_idx,
        "tvec": tvecs_idx,
        "intr": intrs_idx,
        "dist": dists_idx,
    }

    def ref_indices(ref: str) -> list[np.ndarray]:
        """Resolve one reference to a flat-index array per matched view."""
        view, sep, rest = ref.partition(".")
        if not sep or not rest:
            raise ValueError(f"malformed reference {ref!r}; expected 'view.param'")
        if rest.endswith("]"):
            param_name, _, idx_str = rest[:-1].partition("[")
            index = int(idx_str)
        else:
            param_name, index = rest, None
        if param_name not in _PARAM_ALIASES:
            raise ValueError(
                f"unknown parameter {param_name!r} in {ref!r}; "
                f"expected one of {sorted(_PARAM_ALIASES)}"
            )
        group = _PARAM_ALIASES[param_name]
        matched = names if view == "*" else [view]
        out = []
        for c in matched:
            if c not in name_to_row:
                raise ValueError(f"unknown view {c!r} in {ref!r}")
            row = param_idx[group][name_to_row[c]]
            if index is None:
                out.append(row)
            else:
                if not 0 <= index < row.size:
                    raise ValueError(f"index {index} out of range in {ref!r}")
                out.append(row[index : index + 1])
        return out

    fixed_mask = np.zeros(values.size, dtype=bool)
    for ref in fixed:
        for idxs in ref_indices(ref):
            fixed_mask[idxs] = True

    # Union-find over slots; each shared group ties its members element-wise.
    parent = list(range(values.size))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        parent[find(a)] = find(b)

    conflicts: list[tuple[object, list[float], int]] = []
    for group in shared:
        members = [arr for ref in group for arr in ref_indices(ref)]
        if not members:
            continue
        size = members[0].size
        if any(m.size != size for m in members):
            raise ValueError(
                f"shared group {group} mixes references of different lengths"
            )
        for j in range(size):
            idxs = [int(m[j]) for m in members]
            vals = values[idxs]
            if not np.allclose(vals, vals[0]):
                # The surviving value depends on the whole component, which later
                # groups may extend, so the clash is warned about after compaction.
                loc = group if size == 1 else f"{group} element {j}"
                conflicts.append((loc, vals.tolist(), idxs[0]))
            for m in members[1:]:
                union(idxs[0], int(m[j]))

    # Compact: collapse each connected component to a single slot.
    roots = np.array([find(i) for i in range(values.size)])
    uniq, old_to_new = np.unique(roots, return_inverse=True)
    old_to_new = old_to_new.reshape(-1)  # numpy version-proof: keep 1D
    new_values = np.empty(uniq.size)
    new_fixed = np.zeros(uniq.size, dtype=bool)
    filled = np.zeros(uniq.size, dtype=bool)
    for i in range(values.size):
        k = old_to_new[i]
        if not filled[k]:  # smallest original index wins the initial value
            new_values[k] = values[i]
            filled[k] = True
        new_fixed[k] |= fixed_mask[i]

    for loc, vals, member in conflicts:
        warnings.warn(
            f"shared {loc} have differing initial values {vals}; "
            f"keeping {new_values[old_to_new[member]]}",
            stacklevel=2,
        )

    return BAState(
        new_values,
        new_fixed,
        old_to_new[rvecs_idx],
        old_to_new[tvecs_idx],
        old_to_new[intrs_idx],
        old_to_new[dists_idx],
        old_to_new[pts3d_idx],
        pts2d,
        mirrored,
    )

initialize_pts3d

initialize_pts3d(
    pts2d: Float[ndarray, "V *pts 2"],
    rvecs: Float[ndarray, "V 3"],
    tvecs: Float[ndarray, "V 3"],
    intrs: Float[ndarray, "V P"] | Float[ndarray, P],
    mirrored: Bool[ndarray, V],
) -> Float[Array, "*pts 3"]

Triangulate initial (*pts, 3) points from the 2D and the view poses.

Source code in src/deeperfly/bundle_adjustment/state.py

def initialize_pts3d(
    pts2d: Float[np.ndarray, "V *pts 2"],
    rvecs: Float[np.ndarray, "V 3"],
    tvecs: Float[np.ndarray, "V 3"],
    intrs: Float[np.ndarray, "V P"] | Float[np.ndarray, "P"],
    mirrored: Bool[np.ndarray, "V"],
) -> Float[Array, "*pts 3"]:
    """Triangulate initial `(*pts, 3)` points from the 2D and the view poses."""
    return triangulate_dlt(pts2d, pmats(rvecs, tvecs, intrs, mirrored))

Bundle adjustment with JAX-accelerated residuals and Jacobians.

bundle_adjust wraps scipy.optimize.least_squares (TRF and LSMR); the per-observation residual and Jacobian come from jax.vmap and jax.jacrev on project_full_one, assembled into a sparse matrix through a precomputed pattern. The packed state is deeperfly.bundle_adjustment.state's.

BASolution

Bases: NamedTuple

Optimized parameters unpacked back into per-view arrays plus points.

Source code in src/deeperfly/bundle_adjustment/core.py

class BASolution(NamedTuple):
    """Optimized parameters unpacked back into per-view arrays plus points."""

    rvecs: np.ndarray
    tvecs: np.ndarray
    intrs: np.ndarray
    dists: np.ndarray
    pts3d: np.ndarray

bundle_adjust

bundle_adjust(
    values: Float[ndarray, n_params],
    fixed: Bool[ndarray, n_params],
    rvecs_idx: Int[ndarray, "V 3"],
    tvecs_idx: Int[ndarray, "V 3"],
    intrs_idx: Int[ndarray, "V P"],
    dists_idx: Int[ndarray, "V K"],
    pts3d_idx: Int[ndarray, "N 3"],
    pts2d: Float[ndarray, "V N 2"],
    mirrored: Bool[ndarray, V],
    loss: str = "linear",
    f_scale: float = 1.0,
    max_nfev: int = 1000,
    weights: Float[ndarray, "V N"] | None = None,
    bone_pairs: Int[ndarray, "B 2"] | None = None,
    bone_targets: Float[ndarray, B] | None = None,
    bone_weight: float = 1.0,
    **kwargs: Any,
) -> tuple[OptimizeResult, BASolution]

Bundle adjustment with a JAX-computed analytic Jacobian.

The low-level solver over the packed state of build_state (see its BAState). loss and f_scale go to scipy.optimize.least_squares, with huber, cauchy and arctan translated to IRLS form (_IRLS_LOSSES). weights (V, N) scale each residual by sqrt(weight); a negative or non-finite one drops the observation, and all zero falls back to uniform with a warning. bone_pairs, bone_targets and bone_weight append a soft bone-length prior bone_weight * (||p_i - p_j|| - target). kwargs go to scipy; x_scale defaults to "jac" (ADR-0057).

Returns:

Type	Description
OptimizeResult	(result, solution): scipy's OptimizeResult and a BASolution unpacked
BASolution	into per-view arrays and points.

Source code in src/deeperfly/bundle_adjustment/core.py

def bundle_adjust(
    values: Float[np.ndarray, "n_params"],
    fixed: Bool[np.ndarray, "n_params"],
    rvecs_idx: Int[np.ndarray, "V 3"],
    tvecs_idx: Int[np.ndarray, "V 3"],
    intrs_idx: Int[np.ndarray, "V P"],
    dists_idx: Int[np.ndarray, "V K"],
    pts3d_idx: Int[np.ndarray, "N 3"],
    pts2d: Float[np.ndarray, "V N 2"],
    mirrored: Bool[np.ndarray, "V"],
    loss: str = "linear",
    f_scale: float = 1.0,
    max_nfev: int = 1000,
    weights: Float[np.ndarray, "V N"] | None = None,
    bone_pairs: Int[np.ndarray, "B 2"] | None = None,
    bone_targets: Float[np.ndarray, "B"] | None = None,
    bone_weight: float = 1.0,
    **kwargs: Any,
) -> tuple[OptimizeResult, BASolution]:
    """Bundle adjustment with a JAX-computed analytic Jacobian.

    The low-level solver over the packed state of
    [build_state][deeperfly.bundle_adjustment.state.build_state] (see its
    [BAState][]). `loss` and `f_scale` go to `scipy.optimize.least_squares`, with
    huber, cauchy and arctan translated to IRLS form (`_IRLS_LOSSES`). `weights`
    `(V, N)` scale each residual by `sqrt(weight)`; a negative or non-finite one drops
    the observation, and all zero falls back to uniform with a warning. `bone_pairs`,
    `bone_targets` and `bone_weight` append a soft bone-length prior
    `bone_weight * (||p_i - p_j|| - target)`. `kwargs` go to scipy; `x_scale`
    defaults to `"jac"` (ADR-0057).

    Returns:
        `(result, solution)`: scipy's `OptimizeResult` and a [BASolution][] unpacked
        into per-view arrays and points.
    """
    obs_view, obs_pt = np.where(np.isfinite(pts2d).all(axis=-1))
    pts2d_observed = pts2d[obs_view, obs_pt]
    n_obs = len(obs_view)
    n_free = int((~fixed).sum())

    if weights is None:
        sqrt_w = np.ones(n_obs)
    else:
        w = np.asarray(weights, dtype=float)[obs_view, obs_pt]
        # A junk weight just drops its observation: zero residual and Jacobian rows.
        w = np.where(np.isfinite(w), np.maximum(w, 0.0), 0.0)
        if w.size and not w.any():
            # Every observation weighs zero, which would leave the solve ill-posed,
            # so confidence cannot disable it silently: fall back to uniform weights.
            warnings.warn(
                "all confidence weights are zero (or non-positive); "
                "falling back to uniform weights for bundle adjustment",
                stacklevel=2,
            )
            w = np.ones(n_obs)
        sqrt_w = np.sqrt(w)

    use_bones = bone_pairs is not None and len(bone_pairs) > 0
    if use_bones:
        if bone_targets is None:
            raise ValueError("bone_targets is required when bone_pairs is given")
        bone_pairs = np.asarray(bone_pairs)
        bone_targets = np.asarray(bone_targets, dtype=float)
        bone_i, bone_j = bone_pairs[:, 0], bone_pairs[:, 1]
        n_bones = len(bone_pairs)
    else:
        bone_i = bone_j = np.zeros(0, dtype=int)
        n_bones = 0

    x0 = values[~fixed]
    rvecs = values[rvecs_idx]
    tvecs = values[tvecs_idx]
    intrs = values[intrs_idx]
    dists = values[dists_idx]
    pts3d = values[pts3d_idx]

    free_cumsum = np.cumsum(~fixed)

    def free_assign(idx):
        assign_mask = ~fixed[idx]
        source_idx = free_cumsum[idx][assign_mask] - 1
        return assign_mask, source_idx

    rvecs_assign = free_assign(rvecs_idx)
    tvecs_assign = free_assign(tvecs_idx)
    intrs_assign = free_assign(intrs_idx)
    dists_assign = free_assign(dists_idx)
    pts3d_assign = free_assign(pts3d_idx)

    def unpack(x):
        rvecs[rvecs_assign[0]] = x[rvecs_assign[1]]
        tvecs[tvecs_assign[0]] = x[tvecs_assign[1]]
        intrs[intrs_assign[0]] = x[intrs_assign[1]]
        dists[dists_assign[0]] = x[dists_assign[1]]
        pts3d[pts3d_assign[0]] = x[pts3d_assign[1]]
        return rvecs, tvecs, intrs, dists, pts3d

    obs_mirrored = np.asarray(mirrored, dtype=bool)[obs_view]

    def _project_obs(rvecs, tvecs, intrs, dists, pts3d):
        # Pass NumPy straight in and let the jit kernel convert once, rather than
        # pre-staging intermediate JAX arrays.
        return _project_per_obs(
            pts3d[obs_pt],
            rvecs[obs_view],
            tvecs[obs_view],
            intrs[obs_view],
            dists[obs_view],
            obs_mirrored,
        )

    n_resid = 2 * n_obs + n_bones

    def residuals(x):
        rvecs, tvecs, intrs, dists, pts3d = unpack(x)
        pts2d_predicted = _project_obs(rvecs, tvecs, intrs, dists, pts3d)
        reproj = (
            (np.asarray(pts2d_predicted) - pts2d_observed) * sqrt_w[:, None]
        ).ravel()
        if not use_bones:
            return reproj
        lengths = np.asarray(_bone_len_per(pts3d[bone_i], pts3d[bone_j]))
        bone_resid = bone_weight * (lengths - bone_targets)
        return np.concatenate([reproj, bone_resid])

    # Sparsity: each observation's two rows depend on its 3D point and its view's rvec,
    # tvec, intr and dist slots, in `_jac_per_obs`'s order.
    free_idx_map = np.full(values.size, -1, dtype=np.int64)
    free_idx_map[~fixed] = np.arange(n_free)
    cols_per_obs = np.concatenate(
        [
            pts3d_idx[obs_pt],
            rvecs_idx[obs_view],
            tvecs_idx[obs_view],
            intrs_idx[obs_view],
            dists_idx[obs_view],
        ],
        axis=1,
    )
    free_cols_per_obs = free_idx_map[cols_per_obs]
    valid = free_cols_per_obs >= 0
    cols_nz = free_cols_per_obs[valid]
    obs_idx = np.broadcast_to(np.arange(n_obs)[:, None], cols_per_obs.shape)
    rows_x = 2 * obs_idx[valid]
    rows = np.concatenate([rows_x, rows_x + 1])
    cols = np.concatenate([cols_nz, cols_nz])

    # Bone-length sparsity. Each bone row b depends only on the 3 coordinate
    # slots of its two endpoints; its row index sits after the reprojection rows.
    if use_bones:
        bone_cols = np.concatenate(
            [pts3d_idx[bone_i], pts3d_idx[bone_j]], axis=1
        )  # (B,6)
        free_bone_cols = free_idx_map[bone_cols]
        bone_valid = free_bone_cols >= 0
        bone_rows_all = np.broadcast_to(
            (2 * n_obs + np.arange(n_bones))[:, None], bone_cols.shape
        )
        bone_rows = bone_rows_all[bone_valid]
        bone_cols_nz = free_bone_cols[bone_valid]
        rows = np.concatenate([rows, bone_rows])
        cols = np.concatenate([cols, bone_cols_nz])

    # The pattern is constant, so the CSR layout is built once and each iteration's
    # values are scattered into it with one np.bincount (summing duplicate entries).
    coo_key = rows.astype(np.int64) * n_free + cols
    uniq_key, coo_to_csr = np.unique(coo_key, return_inverse=True)
    coo_to_csr = coo_to_csr.reshape(-1)  # numpy version-proof: keep 1D
    nnz_csr = uniq_key.size
    csr_indices = (uniq_key % n_free).astype(np.int64)
    csr_rows = uniq_key // n_free
    csr_indptr = np.concatenate(
        [[0], np.cumsum(np.bincount(csr_rows, minlength=n_resid))]
    ).astype(np.int64)

    def jac(x):
        rvecs, tvecs, intrs, dists, pts3d = unpack(x)
        jac_blocks = _jac_per_obs(
            pts3d[obs_pt],
            rvecs[obs_view],
            tvecs[obs_view],
            intrs[obs_view],
            dists[obs_view],
            obs_mirrored,
        )
        # Each block is (n_obs, 2, *param_dim); concatenate along the last
        # axis to align with cols_per_obs above. Weight rows by sqrt(weight).
        jac_full = np.concatenate([np.asarray(j) for j in jac_blocks], axis=-1)
        jac_full = jac_full * sqrt_w[:, None, None]
        data = np.concatenate([jac_full[:, 0, :][valid], jac_full[:, 1, :][valid]])
        if use_bones:
            dpi, dpj = _bone_jac_per(pts3d[bone_i], pts3d[bone_j])
            bone_jac = bone_weight * np.concatenate(
                [np.asarray(dpi), np.asarray(dpj)], axis=1
            )  # (B, 6)
            data = np.concatenate([data, bone_jac[bone_valid]])
        csr_data = np.bincount(coo_to_csr, weights=data, minlength=nnz_csr)
        return csr_matrix((csr_data, csr_indices, csr_indptr), shape=(n_resid, n_free))

    kwargs.setdefault("x_scale", "jac")
    result = least_squares(
        residuals,
        x0,
        jac=jac,
        method="trf",
        tr_solver="lsmr",
        loss=_IRLS_LOSSES.get(loss, loss),
        f_scale=f_scale,
        max_nfev=max_nfev,
        **kwargs,
    )

    solution = BASolution(*unpack(result.x))
    return result, solution