add whereND_both

Author: mohawk2

Changes

@@ -27,7 +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
+- add {which,where}ND_both - thanks @guillepo for inspiration
 
 2.098 2025-01-03
 - fix Windows build problems

lib/PDL/Primitive.pd

@@ -3644,7 +3644,7 @@ sub PDL::where_both {
 *where_both = \&PDL::where_both;
 EOD
 
-pp_add_exported("", 'whereND');
+pp_add_exported("", 'whereND whereND_both');
 pp_addpm(<<'EOD');
 =head2 whereND
 
@@ -3670,7 +3670,7 @@ to the original data PDLs, for the purpose of dataflow.
   $sdata = whereND $data, $mask
   ($s1, $s2, ..., $sn) = whereND $d1, $d2, ..., $dn, $mask
 
-  where
+where
 
     $data is M dimensional
     $mask is N < M dimensional
@@ -3727,6 +3727,35 @@ sub PDL::whereND :lvalue {
 }
 *whereND = \&PDL::whereND;
 
+=head2 whereND_both
+
+=for ref
+
+C<where_both> with support for ND masks and broadcasting
+
+This works like L</whichND_both>, but data-flowing slices
+rather than index-sets are returned.
+
+C<whereND_both> differs from C<where_both> in that the mask
+dimensionality is preserved which allows for
+proper broadcasting of the selection operation over
+higher dimensions.
+
+As with C<where_both> the output PDLs are still connected
+to the original data PDLs, for the purpose of dataflow.
+
+=for usage
+
+ ($match_vals, $non_match_vals) = whereND_both($pdl, $mask);
+
+=cut
+
+sub PDL::whereND_both :lvalue {
+  barf "Usage: whereND_both(\$pdl, \$mask)\n" if @_ != 2;
+  my ($arr, $mask) = @_;  # $mask has 0==false, 1==true
+  map $arr->indexND($_), PDL::whichND_both($mask);
+}
+*whereND_both = \&PDL::whereND_both;
 EOD
 
 pp_add_exported("", 'whichND whichND_both');

t/primitive-selector.t

@@ -23,6 +23,12 @@ subtest 'where' => sub {
           'dataflow affected orig';
     };
 
+    subtest 'whereND_both' => sub {
+      my ( $t, $f ) = whereND_both(sequence(2,2,2), pdl(0,1));
+      is_pdl $t, pdl('[1;3] [5;7]'), 'nonzero vals';
+      is_pdl $f, pdl('[0;2] [4;6]'), 'zero vals';
+    };
+
     subtest 'whereND' => sub {
         is_deeply( [ zeroes( 2, 3, 1 )->whereND( pdl '0 0' )->dims ], [ 0, 3, 1 ] );