@thi.ng/tensors/select.js

0D/1D/2D/3D/4D tensors with extensible polymorphic operations and customizable storage

import { identity } from "@thi.ng/api/fn";
import { top } from "./top.js";
const select1 = (a, xform, pred, initial) => {
  const {
    data,
    offset,
    shape: [sx],
    stride: [tx]
  } = a;
  const arg = [-1];
  let value = initial, v;
  for (let x = 0; x < sx; x++) {
    v = xform(data[offset + x * tx]);
    if (pred(v, value)) {
      value = v;
      arg[0] = x;
    }
  }
  return { arg, value };
};
const select2 = (a, xform, pred, initial) => {
  const {
    data,
    offset,
    shape: [sx, sy],
    stride: [tx, ty]
  } = a;
  const arg = [-1, -1];
  let value = initial, v;
  let ox, x, y;
  for (x = 0; x < sx; x++) {
    ox = offset + x * tx;
    for (y = 0; y < sy; y++) {
      v = xform(data[ox + y * ty]);
      if (pred(v, value)) {
        value = v;
        arg[0] = x;
        arg[1] = y;
      }
    }
  }
  return { arg, value };
};
const select3 = (a, xform, pred, initial) => {
  const {
    data,
    offset,
    shape: [sx, sy, sz],
    stride: [tx, ty, tz]
  } = a;
  const arg = [-1, -1, -1];
  let value = initial, v;
  let ox, oy, x, y, z;
  for (x = 0; x < sx; x++) {
    ox = offset + x * tx;
    for (y = 0; y < sy; y++) {
      oy = ox + y * ty;
      for (z = 0; z < sz; z++) {
        v = xform(data[oy + z * tz]);
        if (pred(v, value)) {
          value = v;
          arg[0] = x;
          arg[1] = y;
          arg[2] = z;
        }
      }
    }
  }
  return { arg, value };
};
const select4 = (a, xform, pred, initial) => {
  const {
    data,
    offset,
    shape: [sx, sy, sz, sw],
    stride: [tx, ty, tz, tw]
  } = a;
  const arg = [-1, -1, -1, -1];
  let value = initial, v;
  let ox, oy, oz, x, y, z, w;
  for (x = 0; x < sx; x++) {
    ox = offset + x * tx;
    for (y = 0; y < sy; y++) {
      oy = ox + y * ty;
      for (z = 0; z < sz; z++) {
        oz = oy + z * tz;
        for (w = 0; w < sw; w++) {
          v = xform(data[oz + w * tw]);
          if (pred(v, value)) {
            value = v;
            arg[0] = x;
            arg[1] = y;
            arg[2] = z;
            arg[3] = w;
          }
        }
      }
    }
  }
  return { arg, value };
};
const __select = top(0, void 0, select1, select2, select3, select4);
const select = (a, xform, pred, initial) => __select(a, xform, pred, initial);
const argMin = (a) => __select(a, identity, (a2, b) => a2 < b, Infinity);
const argMax = (a) => __select(a, identity, (a2, b) => a2 > b, -Infinity);
export {
  argMax,
  argMin,
  select,
  select1,
  select2,
  select3,
  select4
};