add whichND_both

Author: mohawk2

Changes

@@ -27,6 +27,7 @@
 - badflag now not propagated back to inputs if set on outputs (#517)
 - add new_around_pointer (#505) - thanks @chrisarg for idea
 - add do_print to perldl to match pdl2
+- add whichND_both - thanks @guillepo for inspiration
 
 2.098 2025-01-03
 - fix Windows build problems

lib/PDL/Primitive.pd

@@ -3388,6 +3388,8 @@ C<$c_i>.
  [5 6 7 8 9]
  [0 1 2 3 4]
 
+See also L</whichND_both> for the n-dimensional version.
+
 =cut
 
 EOD
@@ -3727,7 +3729,7 @@ sub PDL::whereND :lvalue {
 
 EOD
 
-pp_add_exported("", 'whichND');
+pp_add_exported("", 'whichND whichND_both');
 pp_addpm(<<'EOD');
 =head2 whichND
 
@@ -3825,6 +3827,53 @@ sub PDL::whichND {
   return $mask ? pdl(indx,0) : PDL->new_from_specification(indx,0) if !$ndims;
   _one2nd($mask, $mask->flat->which->dummy(0,$ndims)->sever);
 }
+
+=head2 whichND_both
+
+=for ref
+
+Return the coordinates of non-zero values in a mask.
+
+=for usage
+
+Like L</which_both>, but returns the N-dimensional coordinates (like
+L</whichND>) of the nonzero, zero values in the mask PDL. The
+returned values arrive as an array-of-vectors suitable for use in
+L<indexND|PDL::Slices/indexND> or L<range|PDL::Slices/range>.
+Added in 2.099.
+
+ ($nonzero_coords, $zero_coords) = whichND_both($mask);
+
+SEE ALSO:
+
+L</which> finds coordinates of nonzero values in a 1-D mask.
+
+L</where> extracts values from a data PDL that are associated
+with nonzero values in a mask PDL.
+
+L<PDL::Slices/indexND> can be fed the coordinates to return the values.
+
+=for example
+
+ pdl> $s=sequence(10,10,3,4)
+ pdl> ($x, $y, $z, $w)=whichND($s == 203); p $x, $y, $z, $w
+ [3] [0] [2] [0]
+ pdl> print $s->at(list(cat($x,$y,$z,$w)))
+ 203
+
+=cut
+
+*whichND_both = \&PDL::whichND_both;
+sub PDL::whichND_both {
+  my $mask = PDL->topdl(shift);
+  return ((PDL->new_from_specification(indx,$mask->ndims,0))x2) if !$mask->nelem;
+  my $ndims = $mask->getndims;
+  if (!$ndims) {
+    my ($t, $f) = (pdl(indx,0), PDL->new_from_specification(indx,0));
+    return $mask ? ($t,$f) : ($f,$t);
+  }
+  map _one2nd($mask, $_->dummy(0,$ndims)->sever), $mask->flat->which_both;
+}
 EOD
 
 #

t/primitive-selector.t

@@ -106,6 +106,12 @@ subtest 'which' => sub {
     is_pdl $zero, pdl(indx, 3, 5), 'zero indices';
   };
 
+  subtest 'whichND_both' => sub {
+    my ( $nonzero, $zero ) = whichND_both(PDL::MatrixOps::identity(2));
+    is_pdl $nonzero, indx('0 0; 1 1'), 'nonzero indices';
+    is_pdl $zero, indx('1 0; 0 1'), 'zero indices';
+  };
+
   subtest 'whichover' => sub {
     my $a = pdl q[3 4 6 3 2 3 5 6 1 7];
     my $b = $a->uniq;