"""INIT (Algorithm 3 of the paper): the cheapest feasible curve over finite per-joint state sets, by a shortest path over the joints. A candidate curve is one state per joint. Piece i is determined by the states at joints i and i+1, the cost is a sum over the pieces, and a curve is feasible exactly when every piece is. INIT keeps, for every state at a joint, the cheapest feasible way to reach it (V) and a back-pointer. A pair (s, s') is skipped without a piece test when V(s) + w(s, s') exceeds the budget U (the cost of a known feasible curve). With the cost-to-go rule the test is V(s) + w(s, s') + togo(s') > U, where togo(s') is a lower bound on the cost of the remaining pieces; it skips more pairs and returns the same curve. The instance supplies the state sets, the piece cost and the exact piece test (an object `inst`): inst.n_pieces M; the joints are 0..M inst.n_states(j) number of states at joint j inst.cost(i, s, V) V(s) + w(s, s') for the states s (indices at joint i, costs so far V) and every state s' at joint i+1, w the piece cost: float (len(s), n_states(i+1)) inst.free(i, s) exact test of piece i: bool (len(s), n_states(i+1)), True where the piece from s to s' is feasible (for tests that are table look-ups); or instead inst.free_pairs(i, s, t) the same test for the listed pairs (s[k], t[k]) only: bool (len(s),) (for tests that cost work per pair; INIT then lists the pairs that pass the budget) inst.togo(j) cost-to-go lower bound of every state at joint j, (n_states(j),) (rule 'togo') inst.n_labels, inst.label(i, s, lab) (optional) a state can carry a label fixed by the pieces before it (e.g. the route class so far): label returns the label of the end state of every pair, int (len(s), n_states(i+1)), or (len(s), 1) on a piece that changes no label. INIT then keeps V per (label, state) and returns the cheapest curve per final label. """ import time import numpy as np def init(inst, U=np.inf, rule='togo', improve_only=True, ties='first', chunk=2e7): """INIT. rule 'togo' (default): budget with cost-to-go; 'budget': the plain budget rule V(s) + w > U (same curve, more piece tests). improve_only: also skip pairs that cannot lower the current V(s') (same curve). Sources are processed cheapest first, `chunk` matrix entries at a time. ties: among ways of equal cost into a state, keep the first one processed ('first', fewest piece tests) or the one from the lowest (label, state) index ('state', independent of the processing order). Returns dict(curves, layers, time): curves = {label: (cost, [state index at joint 0..M])}, the cheapest feasible curve with cost <= U per final label (empty = FAIL); layers = per-piece counts (sources, targets, pairs, budget_pass, tested, free, live).""" M = inst.n_pieces; L = getattr(inst, 'n_labels', 1); label = getattr(inst, 'label', None); t0 = time.time() pairs = hasattr(inst, 'free_pairs'); first = ties == 'first' idx = np.arange(inst.n_states(0)); lab = np.zeros(len(idx), np.int64); V = np.zeros(len(idx)) if rule == 'togo': keep = inst.togo(0) <= U; idx, lab, V = idx[keep], lab[keep], V[keep] back = [(idx, lab, None)]; layers = [] for i in range(M): tl = time.time(); nt = inst.n_states(i + 1); cols = np.arange(nt) tg = inst.togo(i + 1)[None, :] if rule == 'togo' else 0.0 Vn = np.full((L, nt), np.inf); pred = np.full((L, nt), -1, np.int64) st = dict(piece=i, sources=len(V), targets=nt, pairs=len(V) * nt, budget_pass=0, tested=0, free=0) od = np.argsort(V, kind='stable'); cs = max(1, int(chunk // max(1, nt))) for b in range(0, len(od), cs): si = od[b:b + cs] if first else np.sort(od[b:b + cs]); s = idx[si] tot = inst.cost(i, s, V[si]) # V(s) + w(s, s') need = tot + tg <= U if U < np.inf else np.ones(tot.shape, bool) # budget rule st['budget_pass'] += int(np.count_nonzero(need)) lb = label(i, s, lab[si]) if label else None if improve_only: cur = Vn[0][None, :] if lb is None else Vn[lb[:, 0]] if lb.shape[1] == 1 else Vn[lb, cols[None, :]] need &= (tot < cur) if first else (tot <= cur) n_need = int(np.count_nonzero(need)); st['tested'] += n_need if n_need == 0: continue if pairs: # exact piece test, listed pairs a, t = np.nonzero(need); ok = inst.free_pairs(i, s[a], t); a, t = a[ok], t[ok]; st['free'] += len(a) else: # exact piece test, whole rows ok = inst.free(i, s) & need; n_ok = int(np.count_nonzero(ok)); st['free'] += n_ok if n_ok * 20 >= ok.size: # many pairs: reduce over the matrix val = np.where(ok, tot, np.inf) if lb is not None and lb.shape[1] == 1: lb = np.broadcast_to(lb, ok.shape) for l in ([0] if lb is None else np.flatnonzero(np.bincount(lb[ok], minlength=L))): vl = val if lb is None else np.where(lb == l, val, np.inf) a = vl.argmin(0); v = vl[a, cols]; better = v < Vn[l] if not first: better |= (v == Vn[l]) & (si[a] < pred[l]) Vn[l, better] = v[better]; pred[l, better] = si[a[better]] continue a, t = np.nonzero(ok) if len(a) == 0: continue val = tot[a, t]; k = t if lb is None else lb[a, t if lb.shape[1] > 1 else 0] * nt + t # key: (label, s') o = np.lexsort((val, k)); k = k[o]; head = np.r_[True, k[1:] != k[:-1]] # cheapest per key a, k, val = a[o][head], k[head], val[o][head] better = val < Vn.ravel()[k] if not first: better |= (val == Vn.ravel()[k]) & (si[a] < pred.ravel()[k]) Vn.ravel()[k[better]] = val[better]; pred.ravel()[k[better]] = si[a[better]] ll, tt = np.nonzero(np.isfinite(Vn)) st.update(live=len(tt), time=time.time() - tl); layers.append(st) if len(tt) == 0: return dict(curves={}, layers=layers, time=time.time() - t0) idx, lab, V = tt, ll, Vn[ll, tt]; back.append((idx, lab, pred[ll, tt])) curves = {} for l in np.unique(lab): j = int(np.flatnonzero(lab == l)[np.argmin(V[lab == l])]); cost = float(V[j]); path = [0] * (M + 1) for lev in range(M, -1, -1): path[lev] = int(back[lev][0][j]) if lev > 0: j = int(back[lev][2][j]) curves[int(l)] = (cost, path) return dict(curves=curves, layers=layers, time=time.time() - t0)