From 518b90fa90357fa132f39cc680ac34f9572ba4d9 Mon Sep 17 00:00:00 2001
From: "Nicholas R. Jankowski" <jankowski.nicholas@gmail.com>
Date: Fri, 5 Aug 2022 16:46:00 -0400
Subject: [PATCH] hygexxx function vectorization (bug #34363)

---
 NEWS                         |  8 ++++
 inst/distributions/hygecdf.m | 18 ++++++--
 inst/distributions/hygeinv.m | 37 +++++++++++----
 inst/distributions/hygepdf.m | 90 +++++++++++++++++++++++++++++-------
 inst/distributions/hygernd.m | 22 ++++++---
 5 files changed, 140 insertions(+), 35 deletions(-)

diff --git a/NEWS b/NEWS
index d7d0154..6aef084 100644
--- a/NEWS
+++ b/NEWS
@@ -36,6 +36,14 @@ Summary of important user-visible changes for statistics 1.5.0:
  
  ** harmmean.m: fully Matlab compatible.
  
+ ** hygepdf.m: added optional parameter "vectorexpand" to facilitate
+    vectorization of other hyge functions. Allows different inputs lengths 
+    for x and t,m,n parameters, with broadcast expanded output (bug #34363)
+    
+ ** hygecdf.m: improved vectorization for non-scalar inputs.
+    hygeinv.m
+    hygernd.m
+ 
  ** ismissing.m: corrects handling of n-D arrays, NaN indicators, and improves
     matlab compatibility for different data types. (patch #10102)
  
diff --git a/inst/distributions/hygecdf.m b/inst/distributions/hygecdf.m
index e207660..d65074e 100644
--- a/inst/distributions/hygecdf.m
+++ b/inst/distributions/hygecdf.m
@@ -1,3 +1,4 @@
+## Copyright (C) 2022 Nicholas R. Jankowski
 ## Copyright (C) 2012 Rik Wehbring
 ## Copyright (C) 1997-2016 Kurt Hornik
 ##
@@ -63,10 +64,19 @@ function cdf = hygecdf (x, t, m, n)
       cdf = discrete_cdf (x, 0 : n, hygepdf (0 : n, t, m, n));
     endif
   else
-    for i = find (ok(:)')  # Must be row vector arg to for loop
-      v = 0 : n(i);
-      cdf(i) = discrete_cdf (x(i), v, hygepdf (v, t(i), m(i), n(i)));
-    endfor
+    ok &= ! isnan (x);
+    n = n(ok);
+    num_n = numel (n);
+    v = 0 : max (n(:));
+
+    ## manually perform discrete_cdf to enable vectorizing with array input
+    p = cumsum (hygepdf (v, t(ok), m(ok), n, "vectorexpand"), 2);
+    sz_p = size (p);
+    end_locs = sub2ind (sz_p, [1 : num_n]', n(:) + 1);
+    p ./= (p(end_locs));
+    v_col_idx = min (lookup (v, x(ok)(:)) + 1, n(:) + 2);
+    cdf(ok) = [zeros(num_n, 1), p](sub2ind (sz_p + [0,1], ...
+                                                    [1 : num_n]', v_col_idx));
   endif
 
 endfunction
diff --git a/inst/distributions/hygeinv.m b/inst/distributions/hygeinv.m
index 01fd211..0f35494 100644
--- a/inst/distributions/hygeinv.m
+++ b/inst/distributions/hygeinv.m
@@ -1,3 +1,4 @@
+## Copyright (C) 2022 Nicholas R. Jankowski
 ## Copyright (C) 2012 Rik Wehbring
 ## Copyright (C) 1997-2016 Kurt Hornik
 ##
@@ -65,16 +66,34 @@ function inv = hygeinv (x, t, m, n)
       inv(x == 0) = 0;  # Hack to return correct value for start of distribution
     endif
   else
-    for i = find (ok(:)')  # Must be row vector arg to for loop
-      v = 0 : n(i);
-      if (x(i) == 0)
-        inv(i) = 0;  # Hack to return correct value for start of distribution
-      else
-        inv(i) = discrete_inv (x(i), v, hygepdf (v, t(i), m(i), n(i)));
-      endif
-    endfor
-  endif
+    k = (x == 0);
+    inv (ok & k) = 0; # set any x=0 to 0 if not already set to output NaN
+    k = (x == 1);
+    inv (ok & k) = n(ok & k);
+    ok &= (x>0 & x<1); #remove 0's and x's outside (0,1), leave unfilled as NaN
+
+    if any(ok(:))
+      n = n(ok);
+      v = 0 : max (n(:));
+
+      ## manually perform discrete_inv to enable vectorizing with array input
+      p = cumsum (hygepdf (v, t(ok), m(ok), n, "vectorexpand"), 2);
+      sz_p = size (p);
+      end_locs = sub2ind (sz_p, [1 : numel(n)]', n(:) + 1);
+
+      ## manual row-wise vectorization of lookup, which returns index of element
+      ## less than or equal to test value, zero if test value less than lowest
+      ## number, and max index if greater than highest number. operated on
+      ## flipped p, adjusting for different vector lengths in array rows.
+      p = (p ./ p(end_locs))(:, end:-1:1) - x(ok)(:);
+      p(p>=0) = NaN;
+      [p_match, p_match_idx] = max (p, [], 2);
+      p_match_idx(isnan(p_match)) = v(end) + 2;
 
+      inv(ok) = v(v(end) - p_match_idx + 3);
+
+    endif
+  endif
 endfunction
 
 
diff --git a/inst/distributions/hygepdf.m b/inst/distributions/hygepdf.m
index 5ce8bf6..2268bc9 100644
--- a/inst/distributions/hygepdf.m
+++ b/inst/distributions/hygepdf.m
@@ -1,3 +1,4 @@
+## Copyright (C) 2022 Nicholas R. Jankowski
 ## Copyright (C) 2012 Rik Wehbring
 ## Copyright (C) 1996-2016 Kurt Hornik
 ##
@@ -16,7 +17,8 @@
 ## <http://www.gnu.org/licenses/>.
 
 ## -*- texinfo -*-
-## @deftypefn {} {} hygepdf (@var{x}, @var{t}, @var{m}, @var{n})
+## @deftypefn {} {@var{pdf} =} hygepdf (@var{x}, @var{t}, @var{m}, @var{n})
+## @deftypefnx {} {@var{pdf} =} hygepdf (@dots{}, "vectorexpand")
 ## Compute the probability density function (PDF) at @var{x} of the
 ## hypergeometric distribution with parameters @var{t}, @var{m}, and @var{n}.
 ##
@@ -25,34 +27,78 @@
 ## total size @var{t} containing @var{m} marked items.
 ##
 ## The parameters @var{t}, @var{m}, and @var{n} must be positive integers
-## with @var{m} and @var{n} not greater than @var{t}.
+## with @var{m} and @var{n} not greater than @var{t}.  They and @var{x} may be
+## scalars or arrays, but all non-scalars must be the same size.
+##
+## The output @var{pdf} will be the same size as the input array.
+##
+## If the optional parameter @code{vectorexpand} is provided, @var{x} may be an
+## array with size different from parameters @var{t}, @var{m}, and @var{n}
+## (which must still be of a common size or scalar).  Each element of @var{x}
+## will be evaluated against each set of parameters @var{t}, @var{m}, and
+## @var{n} in columnwise order. The output @var{pdf} will be an array of size
+## @code{@var{r} x @var{s}}, where @code{@var{r} = numel (@var{t})}, and
+## @code{@var{s} = numel (@var{x})}.
+##
 ## @end deftypefn
 
 ## Author: KH <Kurt.Hornik@wu-wien.ac.at>
 ## Description: PDF of the hypergeometric distribution
 
-function pdf = hygepdf (x, t, m, n)
+function pdf = hygepdf (x, t, m, n, vect_expand = "")
 
-  if (nargin != 4)
+  if (!any (nargin == [4,5]))
     print_usage ();
   endif
 
-  if (! isscalar (t) || ! isscalar (m) || ! isscalar (n))
-    [retval, x, t, m, n] = common_size (x, t, m, n);
-    if (retval > 0)
-      error ("hygepdf: X, T, M, and N must be of common size or scalars");
-    endif
-  endif
-
   if (iscomplex (x) || iscomplex (t) || iscomplex (m) || iscomplex (n))
     error ("hygepdf: X, T, M, and N must not be complex");
   endif
 
+  if strcmpi (vect_expand, "vectorexpand")
+    ## Expansion to improve vectorization of hyge calling functions.
+    ## Project inputs over a 2D array with x(:) as a row vector and t,m,n as
+    ## a column vector. each pdf(i,j) is hygepdf(x(j), t(i), m(i), n(i))
+    ## Following expansion, remainder of algorithm processes as normal.
+
+    if (! isscalar (t) || ! isscalar (m) || ! isscalar (n))
+      [retval, t, m, n] = common_size (t, m, n);
+      if (retval > 0)
+        error ("hygepdf: T, M, and N must be of common size or scalars");
+      endif
+      ## ensure col vectors before expansion
+      t = t(:);
+      m = m(:);
+      n = n(:);
+    endif
+
+    ## expand x,t,m,n to arrays of size numel(t) x numel(x)
+    sz = [numel(t), numel(x)];
+    x = x(:)'; # ensure row vector before expansion
+
+    x = x(ones (sz(1), 1), :);
+    t = t(:, ones (sz(2), 1));
+    m = m(:, ones (sz(2), 1));
+    n = n(:, ones (sz(2), 1));
+
+  else
+
+    if (! isscalar (t) || ! isscalar (m) || ! isscalar (n))
+      [retval, x, t, m, n] = common_size (x, t, m, n);
+      if (retval > 0)
+        error ("hygepdf: X, T, M, and N must be of common size or scalars");
+      endif
+    endif
+
+    sz = size (x);
+
+  endif
+
   if (isa (x, "single") || isa (t, "single")
       || isa (m, "single") || isa (n, "single"))
-    pdf = zeros (size (x), "single");
+    pdf = zeros (sz, "single");
   else
-    pdf = zeros (size (x));
+    pdf = zeros (sz);
   endif
 
   ## everything in nel gives NaN
@@ -65,7 +111,7 @@ function pdf = hygepdf (x, t, m, n)
 
   k = ! nel & ! zel;
   if (any (k(:)))
-    if (isscalar (t) && isscalar (m) && isscalar (n))
+    if (isscalar (t))
       pdf(k) = exp (gammaln (m+1) - gammaln (m-x(k)+1) - gammaln (x(k)+1) +...
                  gammaln (t-m+1) - gammaln (t-m-n+x(k)+1) - ...
                  gammaln (n-x(k)+1) - gammaln (t+1) + gammaln (t-n+1) + ...
@@ -78,6 +124,7 @@ function pdf = hygepdf (x, t, m, n)
     endif
   endif
 
+
 endfunction
 
 
@@ -94,18 +141,29 @@ endfunction
 %!assert (hygepdf (x, 4, 2, 5), [NaN NaN NaN NaN NaN], eps)
 %!assert (hygepdf ([x, NaN], 4, 2, 2), [y, NaN], eps)
 
-#### Test class of input preserved
+## Test class of input preserved
 %!assert (hygepdf (single ([x, NaN]), 4, 2, 2), single ([y, NaN]), eps("single"))
 %!assert (hygepdf ([x, NaN], single (4), 2, 2), single ([y, NaN]), eps("single"))
 %!assert (hygepdf ([x, NaN], 4, single (2), 2), single ([y, NaN]), eps("single"))
 %!assert (hygepdf ([x, NaN], 4, 2, single (2)), single ([y, NaN]), eps("single"))
 
+## Test vector expansion
+%!test
+%! z = zeros(3,5);
+%! z([4,5,6,8,9,12]) = [1, 0.5, 1/6, 0.5, 2/3, 1/6];
+%! assert (hygepdf (x, 4, [0 1 2], 2,"vectorexpand"), z, eps);
+%! assert (hygepdf (x, 4, [0 1 2]', 2,"vectorexpand"), z, eps);
+%! assert (hygepdf (x', 4, [0 1 2], 2,"vectorexpand"), z, eps);
+%! assert (hygepdf (2, 4, [0 1 2], 2,"vectorexpand"), z(:,4), eps);
+%! assert (hygepdf (x, 4, 1, 2,"vectorexpand"), z(2,:), eps);
+%! assert (hygepdf ([NaN,x], 4, [0 1 2]', 2,"vectorexpand"),[NaN(3,1), z], eps);
+
 ## Test input validation
 %!error hygepdf ()
 %!error hygepdf (1)
 %!error hygepdf (1,2)
 %!error hygepdf (1,2,3)
-%!error hygepdf (1,2,3,4,5)
+%!error hygepdf (1,2,3,4,5,6)
 %!error hygepdf (1, ones (3), ones (2), ones (2))
 %!error hygepdf (1, ones (2), ones (3), ones (2))
 %!error hygepdf (1, ones (2), ones (2), ones (3))
diff --git a/inst/distributions/hygernd.m b/inst/distributions/hygernd.m
index be10206..fd09043 100644
--- a/inst/distributions/hygernd.m
+++ b/inst/distributions/hygernd.m
@@ -1,3 +1,4 @@
+## Copyright (C) 2022 Nicholas R. Jankowski
 ## Copyright (C) 2012 Rik Wehbring
 ## Copyright (C) 1997-2016 Kurt Hornik
 ##
@@ -97,12 +98,21 @@ function rnd = hygernd (t, m, n, varargin)
     endif
   else
     rnd = NaN (sz, cls);
-    rn = rand (sz);
-    for i = find (ok(:)')  # Must be row vector arg to for loop
-      v = 0 : n(i);
-      p = hygepdf (v, t(i), m(i), n(i));
-      rnd(i) = v(lookup (cumsum (p(1 : end-1)) / sum (p), rn(i)) + 1);
-    endfor
+    n = n(ok);
+    num_n = numel (n);
+    v = 0 : max (n(:));
+    p = cumsum (hygepdf (v, t(ok), m(ok), n, "vectorexpand"), 2);
+
+    ## manual row-wise vectorization of lookup, which returns index of element
+    ## less than or equal to test value, zero if test value is less than lowest
+    ## number, and max index if greater than highest number.
+
+    end_locs = sub2ind (size (p), [1 : num_n]', n(:) + 1);
+    p = (p ./ p(end_locs)) - rand (num_n, 1);
+    p(p>=0) = NaN;  #NaN values ignored by max
+    [p_match, p_match_idx] = max (p, [], 2);
+    p_match_idx(isnan(p_match)) = 0; #rand < min(p) gives NaN, reset to 0
+    rnd(ok) = v(p_match_idx + 1);
   endif
 
 endfunction
-- 
2.33.0.windows.2

