# Standard modules.
import gc
import sys
import time
from functools import partial

# CV/ML.
import numpy as np

import torch
import torch.nn.functional as F

print(f"Python:{sys.version}")
print(f"Numpy:{np.__version__}")
print(f"PyTorch:{torch.__version__}")

def get_perf_str(val):
    token_si = ["", "m", "µ", "n", "p"]
    exp_si = [1, 1e3, 1e6, 1e9, 1e12]
    perf_str = f"{val:3g}s"
    si = ""
    sval = val
    for token, exp in zip(token_si, exp_si):
        if val * exp > 1.0:
            si = token
            sval = val * exp
            break
    perf_str = f"{sval:3g}{si}s"
    return perf_str

def print_perf_time(intervals, top_k=None):
    if top_k is not None:
        intervals = np.sort(intervals)[:top_k]
    min = intervals.min()
    max = intervals.max()
    mean = intervals.mean()
    std = intervals.std()

    smin = get_perf_str(min)
    smax = get_perf_str(max)
    mean = get_perf_str(mean)
    std = get_perf_str(std)
    if top_k:
        print(f"Top {top_k} summary: Max {smax}, Min {smin}, Mean +/- Std {mean} +/- {std}")
    else:
        print(f"Overall summary: Max {smax}, Min {smin}, Mean +/- Std {mean} +/- {std}")

class PerfMeasure():
    def __init__(self,
                 trials=100,
                 top_k=10):
        self.trials = trials
        self.top_k = top_k

    def __call__(self, func):
        gc.collect()
        gc.disable()
        intervals = []
        for _ in range(self.trials):
            start = time.perf_counter()
            func()
            end = time.perf_counter()
            intervals.append(end - start)
        intervals = np.array(intervals)
        print_perf_time(intervals)
        if self.top_k:
            print_perf_time(intervals, self.top_k)
        gc.enable()
        gc.collect()

TRIALS = 100
TOPK = 10
pmeasure = PerfMeasure(TRIALS, TOPK)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Target device is {DEVICE}.")
JIT_OPT = False

def matrix_interp_torch(track):
    orig_shape = track.shape
    tlength = orig_shape[0]
    mask = track[:, 0, -1] != 0
    valid = mask.sum()
    if valid == tlength:
        return track

    xs = torch.where(mask != 0)[0]
    ys = track.reshape([tlength, -1])[xs, :]
    x = torch.arange(tlength, device=xs.device)

    # ========================================================================
    # Interpolation.
    # ========================================================================
    xs = xs.to(ys.dtype)
    x = x.to(ys.dtype)
    # Pad control points for extrapolation.
    # Unexpectedly, torch.finfo(torch.float64).min returns -inf.
    # So we use torch.finfo(torch.float32).min alternatively.
    xs = torch.cat([torch.tensor([torch.finfo(torch.float32).min], device=xs.device),
                    xs,
                    torch.tensor([torch.finfo(torch.float32).max], device=xs.device)], dim=0)
    ys = torch.cat([ys[:1], ys, ys[-1:]], dim=0)

    # Compute slopes, pad at the edges to flatten.
    sloops = (ys[1:] - ys[:-1]) / torch.unsqueeze((xs[1:] - xs[:-1]), dim=-1)
    sloops = F.pad(sloops[:-1], (0, 0, 1, 1))

    # Solve for intercepts.
    intercepts = ys - sloops * torch.unsqueeze(xs, dim=-1)

    # Search for the line parameters at each input data point.
    # Create a grid of the inputs and piece breakpoints for thresholding.
    # Rely on argmax stopping on the first true when there are duplicates,
    # which gives us an index into the parameter vectors.
    mask_bk_indicator = torch.unsqueeze(xs, dim=-2) > torch.unsqueeze(x, dim=-1)
    idx = torch.argmax(mask_bk_indicator.to(torch.int32), dim=-1)
    sloop = sloops[idx]
    intercept = intercepts[idx]

    # Apply the linear mapping at each input data point.
    y = sloop * torch.unsqueeze(x, dim=-1) + intercept
    y = y.to(ys.dtype)
    y = y.reshape(orig_shape)
    return y

def partsbased_interp_torch(trackdata, device="cpu"):
    trackdata = torch.from_numpy(trackdata).to(device)
    pose = trackdata[:, :33]
    lhand = trackdata[:, 33:33+21]
    rhand = trackdata[:, 33+21:33+21+21]
    face = trackdata[:, 33+21+21:]

    pose = matrix_interp_torch(pose)
    lhand = matrix_interp_torch(lhand)
    rhand = matrix_interp_torch(rhand)
    face = matrix_interp_torch(face)
    return torch.cat([pose, lhand, rhand, face], dim=1)

【コード解説】
- 引数:
  - trackdata: `[T, J, C]` 形状の追跡点配列．欠損値はゼロ埋めされている必要があります．
    こちらの追跡点には全部位のデータが含まれていることを想定しています．
  - device: torch.Tensorを取り扱うデバイスを指定 [cpu/cuda]．
- 3行目: `trackdata` を `numpy.ndarray` から `torch.Tensor` に変換
- 4-7行目: 追跡点を部位毎に分割
- 9-12行目: 部位毎に線形補間を実行
- 13行目: 部位毎の追跡点を結合して返す

trackdata = np.load("finger_far0_non_static.npy")
reftrack = np.load("finger_far0_non_static_interp.npy")
# Remove person axis.
trackdata = trackdata[0]
reftrack = reftrack[0]

# Torch.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.time()
newtrack = partsbased_interp_torch(trackdata, device=DEVICE)
interval = time.time() - start
print_perf_time(np.array([interval]))

diff = (reftrack - newtrack.detach().cpu().numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_torch, trackdata=trackdata, device=DEVICE)
pmeasure(target_fn)

# Torch.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.time()
newtrack = partsbased_interp_torch(trackdata[:-1], device=DEVICE)
interval = time.time() - start
print_perf_time(np.array([interval]))

diff = (reftrack[:-1] - newtrack.detach().cpu().numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_torch, trackdata=trackdata[:-1], device=DEVICE)
pmeasure(target_fn)

@torch.jit.script
def matrix_interp_torch_jit(track):
    ...

# Torch.
with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.time()
    newtrack = partsbased_interp_torch_jit(trackdata, device=DEVICE)
    interval = time.time() - start
    print_perf_time(np.array([interval]))

    diff = (reftrack - newtrack.detach().cpu().numpy()).sum()
    print(f"Sum of error:{diff}")

    print("Time after second call.")
    target_fn = partial(partsbased_interp_torch_jit, trackdata=trackdata, device=DEVICE)
    pmeasure(target_fn)

# Torch.
with torch.jit.optimized_execution(JIT_OPT):
    # The 1st call may be slow because of the computation graph construction.
    print(f"Time of first call.")
    start = time.time()
    newtrack = partsbased_interp_torch_jit(trackdata[:-1], device=DEVICE)
    interval = time.time() - start
    print_perf_time(np.array([interval]))

    diff = (reftrack[:-1] - newtrack.detach().cpu().numpy()).sum()
    print(f"Sum of error:{diff}")

    print("Time after second call.")
    target_fn = partial(partsbased_interp_torch_jit, trackdata=trackdata[:-1], device=DEVICE)
    pmeasure(target_fn)