Adding histogram_list and organizing _list code

caph1993 · Jun 7, 2024 · d906104 · d906104
1 parent 96bdcdf
commit d906104
Show file tree

Hide file tree

Showing 8 changed files with 293 additions and 120 deletions.
diff --git a/src/array/_globals.kwargs.ts b/src/array/_globals.kwargs.ts
@@ -1,12 +1,12 @@
 //@ts-check
 import NDArray from "../NDArray";
-import { Arr, ArrJS, kwargs_decorator, frequently_used_parsers } from "../array/kwargs";
+import { Arr, ArrOrAny, kwargs_decorator, frequently_used_parsers } from "../array/kwargs";
 
 
 export namespace Func_clip {
   export type Implementation = (a: Arr, a_min: Arr, a_max: Arr, out: Arr) => Arr;
-  export type Kwargs = { a?: ArrJS, a_min?: ArrJS, a_max?: ArrJS, out?: NDArray<any> | null };
-  export type Wrapper = (a: ArrJS, a_min: ArrJS, a_max: ArrJS, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, a_min?: ArrOrAny, a_max?: ArrOrAny, out?: NDArray<any> | null };
+  export type Wrapper = (a: ArrOrAny, a_min: ArrOrAny, a_max: ArrOrAny, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [
     ["a", undefined],

diff --git a/src/array/_globals.ts b/src/array/_globals.ts
@@ -5,6 +5,17 @@ import { TypedArrayConstructor } from "../dtypes";
 const { np, NDArray: __NDArray } = GLOBALS;
 if (!__NDArray) throw new Error(`Programming error: NDArray not defined`);
 
+
+// Types used everywhere
+export type Arr = NDArray<any>;
+export type ArrOrAny = NDArray<any> | number | boolean | any[];
+export type ArrOrConst = NDArray<any> | number | boolean;
+
+export type AxisArg = number | null;
+export type ArrOrNull = NDArray<any> | null;
+
+
+
 // Functions to avoid importing NDArray (because if I import NDArray, I can't use it as a type annotation in the same file)
 export const _NDArray = __NDArray;
 

diff --git a/src/array/kwargs.ts b/src/array/kwargs.ts
@@ -1,16 +1,12 @@
 import { NDArray } from '../NDArray';
 import { TypedArrayConstructor } from '../dtypes';
-import { asarray, broadcast_n_shapes, isarray } from "./_globals";
+import { broadcast_n_shapes, isarray } from "./_globals";
+import { asarray, Arr, ArrOrAny, ArrOrConst, AxisArg, ArrOrNull, } from './_globals';
+export { asarray, Arr, ArrOrAny, ArrOrConst, AxisArg, ArrOrNull } from './_globals';
 import { new_from } from './basic';
 import { Where } from './indexes';
 
 
-export type Arr = NDArray<any>;
-export type ArrJS = NDArray<any> | number | boolean;
-export type ArrNum = NDArray<any> | number | boolean;
-
-export type AxisArg = number | null;
-export type OutArg = NDArray<any> | null;
 
 
 
@@ -22,8 +18,8 @@ export type OutArg = NDArray<any> | null;
 
 
 export function kwargs_decorator<
-  Wrapper extends (...args: any[]) => ArrJS,
-  Implementation extends (...args: any[]) => NDArray,
+  Wrapper extends (...args: any[]) => ArrOrConst,
+  Implementation extends (...args: any[]) => Arr,
 >({ defaults, implementation, parsers, this_as_first_arg }: {
   defaults: [string, any][],
   implementation: Implementation,
@@ -165,29 +161,11 @@ export const frequently_used_parsers = {
   }
 }
 
-type NDArray_non_0D = NDArray<any> | number[];
-
-// Used in statistics.ts:
-export namespace Func_a_q_axis {
-  export type Implementation = (a: NDArray<any>, q: number, axis: number) => NDArray<any>;
-  export type Kwargs = { a?: NDArray_non_0D, q?: number, axis?: AxisArg };
-  export type Wrapper = (a: NDArray_non_0D | Kwargs, q: number | Kwargs, axis?: AxisArg | Kwargs) => NDArray<any>;
-  export const decorator = kwargs_decorator<Wrapper, Implementation>;
-  export const defaults: [string, any][] = [["a", undefined], ["q", undefined], ["axis", null]];
-  export const parsers = [
-    frequently_used_parsers.a_axis_flatten,
-    (kwargs) => { kwargs.q = asarray(kwargs.q); },
-  ];
-  export const defaultDecorator = (implementation: Implementation) => decorator({
-    defaults, implementation, parsers
-  })
-}
-
 // Used by sort:
 export namespace Func_a_lastAxis {
-  export type Implementation = (a: ArrJS, axis: number) => NDArray<any>;
-  export type Kwargs = { a?: ArrJS, axis?: AxisArg };
-  export type Wrapper = (a: ArrJS | Kwargs, axis?: AxisArg | Kwargs) => NDArray<any>;
+  export type Implementation = (a: ArrOrAny, axis: number) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, axis?: AxisArg };
+  export type Wrapper = (a: ArrOrAny | Kwargs, axis?: AxisArg | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["axis", -1]];
   export const parsers = [
@@ -202,8 +180,8 @@ export namespace Func_a_lastAxis {
 // Used by elementwise operators and methods:
 export namespace Func_a_out {
   export type Implementation = (a: NDArray<any>, out: NDArray<any> | null) => NDArray<any>;
-  export type Kwargs = { a?: ArrJS, out?: NDArray<any> | null };
-  export type Wrapper = (a: ArrJS | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, out?: NDArray<any> | null };
+  export type Wrapper = (a: ArrOrAny | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["out", null]];
   export const parsers = [
@@ -236,8 +214,8 @@ export namespace Method_out {
 // Used by round
 export namespace Func_a_decimals_out {
   export type Implementation = (a: NDArray<any>, decimals: number, out: NDArray<any> | null) => NDArray<any>;
-  export type Kwargs = { a?: ArrJS, decimals?: number, out?: NDArray<any> | null };
-  export type Wrapper = (a: ArrJS | Kwargs, decimals: number | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, decimals?: number, out?: NDArray<any> | null };
+  export type Wrapper = (a: ArrOrAny | Kwargs, decimals: number | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["decimals", 0], ["out", null]];
   export const parsers = [
@@ -270,8 +248,8 @@ export namespace Method_a_decimals_out {
 // Used by binary operators and methods:
 export namespace Func_a_other_out {
   export type Implementation = (a: NDArray<any>, other: NDArray<any>, out: NDArray<any> | null) => NDArray<any>;
-  export type Kwargs = { a: ArrJS, other?: ArrJS, out?: NDArray<any> | null };
-  export type Wrapper = (a: ArrJS | Kwargs, other: ArrJS | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
+  export type Kwargs = { a: ArrOrAny, other?: ArrOrAny, out?: NDArray<any> | null };
+  export type Wrapper = (a: ArrOrAny | Kwargs, other: ArrOrAny | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["other", undefined], ["out", null]];
   export const parsers = [
@@ -285,10 +263,10 @@ export namespace Func_a_other_out {
 }
 export namespace Method_other_out {
   export type Implementation = Func_a_other_out.Implementation;
-  export type Kwargs = { other?: ArrJS, out?: NDArray<any> | null };
+  export type Kwargs = { other?: ArrOrAny, out?: NDArray<any> | null };
   export type Wrapper<
     T extends TypedArrayConstructor = Float64ArrayConstructor,
-  > = (other: ArrJS | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<T>;
+  > = (other: ArrOrAny | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<T>;
   export const defaults: [string, any][] = [["a", undefined], ["other", undefined], ["out", null]];
   export const parsers = [
     frequently_used_parsers.isarray('a'),
@@ -307,8 +285,8 @@ export namespace Method_other_out {
 // Used by assign operators and methods:
 export namespace Func_a_values_where {
   export type Implementation = (a: NDArray<any>, values: NDArray<any>, where: Where) => NDArray<any>;
-  export type Kwargs = { a: NDArray<any>, values?: ArrJS, where: Where };
-  export type Wrapper = (a: NDArray<any> | Kwargs, values: ArrJS | Kwargs, ...where: Where) => NDArray<any>;
+  export type Kwargs = { a: NDArray<any>, values?: ArrOrAny, where: Where };
+  export type Wrapper = (a: NDArray<any> | Kwargs, values: ArrOrAny | Kwargs, ...where: Where) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["values", undefined], ["...where", null]];
   export const parsers = [
@@ -321,8 +299,8 @@ export namespace Func_a_values_where {
 }
 export namespace Method_values_where {
   export type Implementation = Func_a_values_where.Implementation;
-  export type Kwargs = { values?: ArrJS, where: Where };
-  export type Wrapper = (values: ArrJS | Kwargs, ...where: Where) => NDArray<any>;
+  export type Kwargs = { values?: ArrOrAny, where: Where };
+  export type Wrapper = (values: ArrOrAny | Kwargs, ...where: Where) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["values", undefined], ["...where", null]];
   export const parsers = [
@@ -337,10 +315,10 @@ export namespace Method_values_where {
 // Used by reduce functions and methods:
 export namespace Func_a_axis_keepdims {
   export type Implementation = (a: NDArray<any>, axis: number, keepdims: boolean) => NDArray<any>;
-  export type Kwargs = { a?: ArrJS, axis?: AxisArg, keepdims?: boolean };
+  export type Kwargs = { a?: ArrOrAny, axis?: AxisArg, keepdims?: boolean };
   export type Wrapper<
     T extends TypedArrayConstructor = Float64ArrayConstructor,
-  > = (a: ArrJS | Kwargs, axis?: AxisArg | Kwargs, keepdims?: boolean | Kwargs) => NDArray<T>;
+  > = (a: ArrOrAny | Kwargs, axis?: AxisArg | Kwargs, keepdims?: boolean | Kwargs) => NDArray<T>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["axis", null], ["keepdims", false]];
   export const parsers = [
@@ -354,6 +332,7 @@ export namespace Func_a_axis_keepdims {
     });
   }
 }
+
 export namespace Method_a_axis_keepdims {
   export type Implementation = Func_a_axis_keepdims.Implementation;
   export type Kwargs = { axis?: AxisArg, keepdims?: boolean };
@@ -378,8 +357,8 @@ export namespace Method_a_axis_keepdims {
 // For norm:
 export namespace Func_a_ord_axis_keepdims {
   export type Implementation = (a: NDArray<any>, ord: number, axis: number, keepdims: boolean) => NDArray<any>;
-  export type Kwargs = { a?: ArrJS, ord?: number, axis?: AxisArg, keepdims?: boolean };
-  export type Wrapper = (a: ArrJS | Kwargs, ord?: number | Kwargs, axis?: AxisArg | Kwargs, keepdims?: boolean | Kwargs) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, ord?: number, axis?: AxisArg, keepdims?: boolean };
+  export type Wrapper = (a: ArrOrAny | Kwargs, ord?: number | Kwargs, axis?: AxisArg | Kwargs, keepdims?: boolean | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["ord", 2], ["axis", null], ["keepdims", false]];
   export const parsers = [
@@ -408,8 +387,8 @@ export namespace Method_a_ord_axis_keepdims {
 // For std:
 export namespace Func_a_axis_ddof_keepdims {
   export type Implementation = (a: NDArray<any>, axis: number, ddof: number, keepdims: boolean) => NDArray<any>;
-  export type Kwargs = { a?: ArrJS, axis?: AxisArg, ddof?: number, keepdims?: boolean };
-  export type Wrapper = (a: ArrJS | Kwargs, axis?: AxisArg | Kwargs, ddof?: number | Kwargs, keepdims?: boolean | Kwargs) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, axis?: AxisArg, ddof?: number, keepdims?: boolean };
+  export type Wrapper = (a: ArrOrAny | Kwargs, axis?: AxisArg | Kwargs, ddof?: number | Kwargs, keepdims?: boolean | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["a", undefined], ["ddof", 0], ["axis", null], ["keepdims", false]];
   export const parsers = [
@@ -439,9 +418,9 @@ export namespace Method_a_axis_ddof_keepdims {
 
 // Used by atan2:
 export namespace Func_y_x_out {
-  export type Implementation = (y: ArrJS, x: ArrJS, out?: NDArray | null) => NDArray<any>;
-  export type Kwargs = { y?: ArrJS, x?: ArrJS, out?: NDArray<any> | null };
-  export type Wrapper = (y: ArrJS | Kwargs, x: ArrJS | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
+  export type Implementation = (y: ArrOrAny, x: ArrOrAny, out?: NDArray | null) => NDArray<any>;
+  export type Kwargs = { y?: ArrOrAny, x?: ArrOrAny, out?: NDArray<any> | null };
+  export type Wrapper = (y: ArrOrAny | Kwargs, x: ArrOrAny | Kwargs, out?: NDArray<any> | null | Kwargs) => NDArray<any>;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [["x", undefined], ["y", undefined], ["out", null]];
   export const parsers = [

diff --git a/src/index.ts b/src/index.ts
@@ -250,8 +250,8 @@ np.geomspace = np.modules.constructors.geomspace;
 np.take = np.modules.indexing.take;
 np.where = np.modules.indexing.where;
 np.nonzero = np.modules.indexing.nonzero;
-np.quantile = np.modules.statistics.kw_exported.quantile;
-np.nanquantile = np.modules.statistics.kw_exported.nanquantile;
+np.quantile = np.modules.statistics.quantile;
+np.nanquantile = np.modules.statistics.nanquantile;
 // np.percentile = np.modules.statistics.kw_exported.percentile;
 // np.median = np.modules.statistics.kw_exported.median;
 // np.average = np.modules.statistics.kw_exported.average;

diff --git a/src/modules/math-functions.kwargs.ts b/src/modules/math-functions.kwargs.ts
@@ -1,12 +1,11 @@
 //@ts-check
-import NDArray from "../NDArray";
-import { Arr, ArrJS, OutArg, kwargs_decorator, frequently_used_parsers } from "../array/kwargs";
-
+import { Arr, ArrOrAny, ArrOrNull, kwargs_decorator, frequently_used_parsers, ArrOrConst } from "../array/kwargs";
+import * as math from "./math-functions";
 
 export namespace Func_clip {
   export type Implementation = (a: Arr, a_min: Arr, a_max: Arr, out: Arr) => Arr;
-  export type Kwargs = { a?: ArrJS, a_min?: ArrJS, a_max?: ArrJS, out?: OutArg };
-  export type Wrapper = (a: ArrJS, a_min: ArrJS, a_max: ArrJS, out?: OutArg | Kwargs) => NDArray<any>;
+  export type Kwargs = { a?: ArrOrAny, a_min?: ArrOrAny, a_max?: ArrOrAny, out?: ArrOrNull };
+  export type Wrapper = (a: ArrOrAny, a_min: ArrOrAny, a_max: ArrOrAny, out?: ArrOrNull | Kwargs) => ArrOrConst;
   export const decorator = kwargs_decorator<Wrapper, Implementation>;
   export const defaults: [string, any][] = [
     ["a", undefined],
@@ -23,4 +22,6 @@ export namespace Func_clip {
   export const defaultDecorator = (implementation: Implementation) => decorator({
     defaults, implementation, parsers
   });
-}
+}
+
+export const clip = Func_clip.defaultDecorator(math.clip);
diff --git a/src/modules/math-functions.ts b/src/modules/math-functions.ts
@@ -1,6 +1,5 @@
 //@ts-check
-import { Arr } from "../array/kwargs";
-import { Func_clip } from "./math-functions.kwargs";
+import { Arr } from "../array/_globals";
 import { op_binary } from "../array/operators";
 import { apply_along_axis, concatenate } from "../array/transform";
 
@@ -76,9 +75,37 @@ export function cumtrapz(y: Arr, x: Arr | null, dx: number, axis: number, initia
   }, y.dtype);
 }
 
+export function trapz(y: Arr, x: Arr | null, dx: number, axis: number) {
+  return apply_along_axis(y, axis, (arr: any[]) => {
+    let sum = 0;
+    for (let i = 1; i < arr.length; i++) {
+      const dxi = x ? x[i] - x[i - 1] : dx;
+      sum += (arr[i] + arr[i - 1]) / 2 * dxi;
+    }
+    return sum;
+  }, y.dtype);
+}
 
+export function interp(x: Arr, xp: Arr, fp: Arr, left: number | null, right: number | null) {
+  const x_flat = x.flat;
+  const xp_flat: number[] = xp.flat;
+  const fp_flat = fp.flat;
+  // Use binary search instead of findIndex:
+  return x_flat.map((v) => {
+    let lo = 0, hi = xp_flat.length;
+    while (lo < hi) {
+      const mid = lo + hi >> 1;
+      if (xp_flat[mid] < v) lo = mid + 1;
+      else hi = mid;
+    }
+    if (lo == 0) return left;
+    if (lo == xp_flat.length) return right;
+    const x0 = xp_flat[lo - 1];
+    const x1 = xp_flat[lo];
+    const f0 = fp_flat[lo - 1];
+    const f1 = fp_flat[lo];
+    return f0 + (f1 - f0) * (v - x0) / (x1 - x0);
+  });
+}
 
 
-export const kw_exported = {
-  clip: Func_clip.defaultDecorator(clip),
-};
diff --git a/src/modules/statistics.kwargs.ts b/src/modules/statistics.kwargs.ts
@@ -0,0 +1,29 @@
+//@ts-check
+import { Arr, AxisArg, kwargs_decorator, frequently_used_parsers, asarray, ArrOrConst, Func_a_axis_keepdims } from "../array/kwargs";
+import * as stats from "./statistics";
+
+
+type NDArray_non_0D = Arr | number[];
+
+// Used in statistics.ts:
+export namespace Func_a_q_axis {
+  export type Implementation = (a: Arr, q: number, axis: number) => Arr;
+  export type Kwargs = { a?: NDArray_non_0D, q?: number, axis?: AxisArg };
+  export type Wrapper = (a: NDArray_non_0D | Kwargs, q: number | Kwargs, axis?: AxisArg | Kwargs) => ArrOrConst;
+  export const decorator = kwargs_decorator<Wrapper, Implementation>;
+  export const defaults: [string, any][] = [["a", undefined], ["q", undefined], ["axis", null]];
+  export const parsers = [
+    frequently_used_parsers.a_axis_flatten,
+    (kwargs) => { kwargs.q = asarray(kwargs.q); },
+  ];
+  export const defaultDecorator = (implementation: Implementation) => decorator({
+    defaults, implementation, parsers
+  })
+}
+
+export const quantile = Func_a_q_axis.defaultDecorator(stats.quantile);
+export const nanquantile = Func_a_q_axis.defaultDecorator(stats.nanquantile);
+export const percentile = Func_a_q_axis.defaultDecorator(stats.percentile);
+export const nanpercentile = Func_a_q_axis.defaultDecorator(stats.nanpercentile);
+export const median = Func_a_axis_keepdims.defaultDecorator(stats.median);
+export const nanmedian = Func_a_axis_keepdims.defaultDecorator(stats.nanmedian);