carray 3.0.2 → 3.0.4

data/lib/carray/axis_group.rb

@@ -213,6 +213,25 @@ class AxisGroup
     end
   end
 
+  # The keywords a group reduction takes besides axis:.  Checked before the
+  # work, so an unknown keyword is refused rather than dropped.
+  REDUCE_KEYWORDS = [:axis, :min_count, :fill_value].freeze
+
+  # all / any take axis: only, as the core ones do.
+  LOGICAL_REDUCE_KEYWORDS = [:axis].freeze
+
+  # A boolean array carrying value's mask: what a count reads, for a value
+  # whose payload the group kernel cannot read.  Called from C.
+  private_class_method def self.count_stand_in (value)
+    return value if value.boolean? || value.integer? || value.float?
+    stand_in = CArray.boolean(*value.shape)
+    stand_in.mask = value.mask if value.has_mask?
+    stand_in
+  end
+
+  # A count of the mask reads every cell of the group, as the core's does.
+  MASK_COUNTS = [:count, :count_not_masked, :count_masked].freeze
+
   # Parse an axis: spec into [has_group, fused_band_slot_positions].
   def self.parse_axis (axis)
     has_group = false
@@ -332,58 +351,109 @@ class CAGroupIterator
   # -- arrived as "wrong number of arguments", which reads as a method that
   # does not take one rather than one whose engine has no such fold. Both are
   # composed here instead, out of folds the engine does have.
-  alias __count_present__ count
-
-  # @overload count
-  #   Per-group count of present (non-masked) cells.
-  # @overload count(UNDEF)
-  #   Per-group count of masked cells, as {#count_masked}.
-  # @overload count(v)
-  #   Per-group count of cells equal to `v`.
-  #   @return [CArray]
-  def count (*args, **kw)
-    return __count_present__(**kw) if args.empty?
-    if args.size > 1
-      raise ArgumentError, "wrong number of arguments (given #{args.size}, expected 0..1)"
+  # They are in a module in front of the engine's own count, all and any,
+  # which they reach with super; so are the reductions below that choose
+  # between the engine and the core.
+  module ComposedFolds
+
+    # @overload count
+    #   Per-group count of present (non-masked) cells.
+    # @overload count(UNDEF)
+    #   Per-group count of masked cells, as {#count_masked}.
+    # @overload count(v)
+    #   Per-group count of cells equal to `v`.
+    #   @return [CArray]
+    def count (*args, **kw)
+      if args.empty?
+        return core_counts?(kw) ? reduce_by_core(:count, kw) : super(**kw)
+      end
+      if args.size > 1
+        raise ArgumentError, "wrong number of arguments (given #{args.size}, expected 0..1)"
+      end
+      v = args.first
+      return count_masked(**kw) if v.equal?(UNDEF)
+      # eq leaves a masked cell UNDEF, so the sum leaves it out and its
+      # min_count: counts the present cells, as the core count(v) does
+      regroup(value.eq(v).int64, spec).sum(**kw).int64
     end
-    v = args.first
-    return count_masked(**kw) if v.equal?(UNDEF)
-    # a masked cell equals nothing, and eq marks it UNDEF; drop that to false
-    # so it is simply not counted
-    hit = value.eq(v)
-    hit = hit.strip_mask(false) if hit.has_mask?
-    self.class.__build__(hit.int64, spec).sum(**kw).int64
-  end
 
-  alias __all_folded__ all
-  alias __any_folded__ any
+    # @overload all
+    #   Per-group `all` over a boolean payload.
+    #   @return [CArray]
+    def all (**kw)
+      boolean_payload!(:all)
+      core_counts?(kw) ? reduce_by_core(:all, kw) : super(**kw)
+    end
 
-  # @overload all
-  #   Per-group `all` over a boolean payload.
-  #   @return [CArray]
-  def all (**kw)
-    boolean_payload!(:all)
-    __all_folded__(**kw)
-  end
+    # @overload any
+    #   Per-group `any` over a boolean payload.
+    #   @return [CArray]
+    def any (**kw)
+      boolean_payload!(:any)
+      core_counts?(kw) ? reduce_by_core(:any, kw) : super(**kw)
+    end
 
-  # @overload any
-  #   Per-group `any` over a boolean payload.
-  #   @return [CArray]
-  def any (**kw)
-    boolean_payload!(:any)
-    __any_folded__(**kw)
+    # @overload count_not_masked
+    #   Per-group count of present (non-masked) cells.
+    #   @return [CArray]
+    def count_not_masked (**kw)
+      core_counts?(kw) ? reduce_by_core(:count_not_masked, kw) : super(**kw)
+    end
+
+    private
+
+    # Whether a count or all / any of each group is taken from the core, as
+    # the other reductions of a payload the engine does not read are. The
+    # engine answers these for every payload; with the fast paths off (a
+    # development build run with CARRAY_GENERAL_PATH set) the core does --
+    # except for a band folded into the count, which has no route but the
+    # engine.
+    def core_counts? (kw)
+      return false if CArray.__fast_paths_enabled__
+      has_group, fused = AxisGroup.parse_axis(kw[:axis])
+      has_group && fused.empty?
+    end
   end
+  prepend ComposedFolds
 
   # Per-group classified cell count (mask-independent) = count on the
   # mask-stripped value, so every classified cell is counted regardless of the
   # value mask (unlike count / count_not_masked, which count present cells).
   def elements (**kw)
-    self.class.__build__(value.value, spec).count(**kw)
+    regroup(value.value, spec).count(**kw)
   end
 
   # Per-group count of value-masked cells = elements - count_not_masked.
   def count_masked (**kw)
-    elements(**kw) - count_not_masked(**kw)
+    has_group, = AxisGroup.parse_axis(kw[:axis])
+    return value.count_masked(**kw) unless has_group
+    check_reduce_keywords(:count_masked, kw)
+    plain = { axis: kw[:axis] }
+    finish_reduction(:count_masked, kw,
+                                   elements(**plain) - count_not_masked(**plain))
+  end
+
+  # The scatter kernel reads an integer or float payload, and answers there
+  # as the core does. Any other payload -- boolean, complex, object, a Face --
+  # is reduced group by group by the core, through the composite categorical,
+  # so the data type, the Face and a refusal are the core's: a boolean's sum
+  # is uint64, a time column's minimum a time, its sum refused.
+  CArray.send(:define_named_methods, ComposedFolds,
+              [:sum, :accumulate, :prod, :mean, :min, :max,
+               :variance, :stddev, :variancep, :stddevp],
+              :scatter_or_core, super_block: true)
+
+  # The scatter kernel is the block (super).
+  private def scatter_or_core (op, args, kw)
+    unless args.empty?
+      raise ArgumentError, "wrong number of arguments (given #{args.size}, expected 0)"
+    end
+    has_group, fused = AxisGroup.parse_axis(kw[:axis])
+    if ! has_group || kernel_reads_value?(fused)
+      yield(**kw)
+    else
+      reduce_by_core(op, kw)
+    end
   end
 
   # Per-group [min, max] pair (matching CArray#minmax).
@@ -418,7 +488,7 @@ class CAGroupIterator
   # value|weight mask, so masked cells drop out. Empty group -> 0.0 (identity).
   # No weighted kernel needed -- it is a plain group-sum of a derived array.
   def wsum (weights, **kw)
-    self.class.__build__(value * weights, spec).sum(**kw)
+    regroup(value * weights, spec).sum(**kw)
   end
 
   # Weighted mean = Sum(v*w) / Sum(w), both over the combined present-set. The
@@ -427,13 +497,21 @@ class CAGroupIterator
   # (no present pair) is UNDEF (matching CArray#wmean); a present group whose
   # weights sum to zero yields NaN/Inf (core's 0/0).
   def wmean (weights, **kw)
+    has_group, = AxisGroup.parse_axis(kw[:axis])
+    return value.wmean(weights, **kw) unless has_group
+    check_reduce_keywords(:wmean, kw)
+    plain = { axis: kw[:axis] }
     prod = value * weights
-    num  = self.class.__build__(prod, spec).sum(**kw)
-    den  = self.class.__build__(value * 0 + weights, spec).sum(**kw)
-    cnt  = self.class.__build__(prod, spec).count_not_masked(**kw)
-    out  = num / den
-    out[cnt.eq(0)] = UNDEF          # no present pair -> masked (empty group)
-    out
+    num  = regroup(prod, spec).sum(**plain)
+    den  = regroup(value * 0 + weights, spec).sum(**plain)
+    cnt  = regroup(prod, spec).count_not_masked(**plain)
+    # No present pair is UNDEF (an empty group), set on the denominator so
+    # the division never meets an empty group's 0 / 0.
+    den  = den.copy
+    den[cnt.eq(0)] = UNDEF
+    # min_count: counts the present pairs, as the core wmean does.
+    pairs = regroup(prod, spec)
+    pairs.finish_reduction(:wmean, kw, num / den)
   end
 
   # ---- segment scan: within-group running statistics ---------------------
@@ -444,7 +522,7 @@ class CAGroupIterator
   # in row-major position order along the grouped axes (per band). Excluded /
   # source-masked cells are UNDEF (they join no running total). All five route
   # through the shared scan_op dispatcher, driving the fused C kernel
-  # __axis_group_scan__ (ext/ca_axis_group.c) with peak O(k) extra. A scan
+  # __scan_by_axis_groups__ (ext/ca_axis_group.c) with peak O(k) extra. A scan
   # cannot fold a band into the statistic (that would collapse an axis), so it
   # accepts only axis: :group; without :group each delegates to the value's
   # same-named scan.
@@ -493,6 +571,59 @@ class CAGroupIterator
     scan_op(:cumcount, kw)
   end
 
+  # @overload lag(n = 1, axis: :group)
+  #   Per-group value `n` steps earlier, counting a group's members in
+  #   row-major position order along the grouped axes (per band), as the
+  #   scans do; UNDEF where the group has none, so a lag never reaches
+  #   another group. Without :group it is the value's own `lag` along an
+  #   axis. A negative `n` looks ahead.
+  #   @return [CArray] source-shaped
+  def lag (n = 1, axis: nil)
+    has_group, fused = AxisGroup.parse_axis(axis)
+    return value.lag(n, axis: axis) unless has_group
+    unless fused.empty?
+      raise ArgumentError,
+            "axis_group lag: folding a band into a lag (axis: [:group, k]) " \
+            "is not supported; a lag preserves the source shape"
+    end
+    scan_by_core(:lag, n)
+  end
+
+  # @overload strip_mask(method:, axis: :group)
+  #   Per-group fill of the masked cells, counting a group's members in
+  #   row-major position order along the grouped axes (per band), as the
+  #   scans do: `:forward` / `:ffill` or `:backward` / `:bfill`, never across
+  #   into another group. Without :group it is the value's own
+  #   `strip_mask(method:, axis:)`.
+  #   @return [CArray] source-shaped
+  def strip_mask (method:, axis: nil)
+    has_group, fused = AxisGroup.parse_axis(axis)
+    return value.strip_mask(method: method, axis: axis) unless has_group
+    CAIterator.hold_backward?(method)
+    unless fused.empty?
+      raise ArgumentError,
+            "axis_group strip_mask: folding a band into a fill (axis: [:group, k]) " \
+            "is not supported; a fill preserves the source shape"
+    end
+    scan_by_core(:strip_mask, method: method)
+  end
+
+  # @overload order(axis: :group, descending: false, method: :ordinal, kind: :quick)
+  #   Per-group rank of each cell among its group's members (int64, `0` the
+  #   smallest; per band), as `CArray#order` ranks them; a cell in no group is
+  #   UNDEF. Without :group it is the value's own `order`.
+  #   @return [CArray] source-shaped
+  def order (axis: nil, **kw)
+    has_group, fused = AxisGroup.parse_axis(axis)
+    return value.order(axis: axis, **kw) unless has_group
+    unless fused.empty?
+      raise ArgumentError,
+            "axis_group order: folding a band into an order (axis: [:group, k]) " \
+            "is not supported; an order preserves the source shape"
+    end
+    scan_by_core(:order, **kw)
+  end
+
   # ---- materialize path: order statistics + generic iteration ------------
   #
   # These need every member of a group held together, so (unlike the scatter
@@ -534,6 +665,41 @@ class CAGroupIterator
     order_stat(:quantile, [], kw)
   end
 
+  # @overload first(axis: :group)
+  #   Per-group first present value, in row-major position order along the
+  #   grouped axes (per band); UNDEF for a group with none. Without :group it
+  #   is the value's own `first`.
+  #   @return [CArray]
+  # @overload last(axis: :group)
+  #   Per-group last present value; UNDEF for a group with none.
+  #   @return [CArray]
+  # @overload nunique(axis: :group)
+  #   Per-group number of distinct present values (int64); 0 for a group with
+  #   none.
+  #   @return [CArray]
+  #
+  # These take no `min_count:` or `fill_value:`, as the core ones do not. A
+  # group is gathered whole, so folding a band into one (axis: [:group, k]) is
+  # not supported, as for the order statistics.
+  CArray.send(:define_named_methods, self, [:first, :last, :nunique], :group_position)
+
+  private def group_position (op, **kw)
+    has_group, fused = AxisGroup.parse_axis(kw[:axis])
+    return value.public_send(op, **kw) unless has_group
+    unless fused.empty?
+      raise NotImplementedError,
+            "axis_group: folding a band into #{op} (axis: [:group, k]) is not " \
+            "supported; #{op} gathers every member of a group"
+    end
+    unknown = kw.keys - [:axis]
+    unless unknown.empty?
+      raise ArgumentError,
+            "#{op}: unknown keyword#{unknown.size > 1 ? "s" : ""}: " +
+            unknown.map(&:inspect).join(", ")
+    end
+    composite_order(op, [])
+  end
+
   # @overload sort_addr
   #   Per-group sorted flat source addresses (group-major). Any grouping (single
   #   or composite group slots, rank-1 or rank-N, flat or band-preserving) is
@@ -569,12 +735,13 @@ class CAGroupIterator
     raise LocalJumpError, "no block given (yield)" unless block
     ccat, _kd, gslots, bslots, gaxes = composite_layout
     return value.group_by_category(ccat).map(data_type: data_type, &block) if bslots.empty?
-    dt  = data_type || value.data_type
-    out = CArray.new(dt, value.shape)
+    out = nil
     each_band_block(ccat, gslots, bslots, gaxes) do |val_idx, _oi, _co, gi|
-      out[*val_idx] = gi.map(data_type: dt, &block)
+      result = gi.map(data_type: data_type, &block)
+      out ||= blank_of(result, value.shape)
+      out[*val_idx] = result
     end
-    out
+    out || (data_type ? CArray.new(data_type, value.shape) : blank_of(value, value.shape))
   end
 
   # @overload reduce { |members| ... }
@@ -599,8 +766,50 @@ class CAGroupIterator
     out
   end
 
+  protected
+
+  # min_count: masks a group with fewer contributing cells (present cells,
+  # or every cell for a count of the mask); fill_value: then fills the masked
+  # cells -- as a core reduction does with each.  Protected: wmean finishes
+  # through the iterator over its products.
+  def finish_reduction (op, kw, result)
+    need = CArray.__present_cells_needed_per_result_cell__(op, kw[:min_count], 0)
+    if need > 0
+      counts = if AxisGroup::MASK_COUNTS.include?(op.to_sym)
+                 elements(axis: kw[:axis])
+               else
+                 count_not_masked(axis: kw[:axis])
+               end
+      if result.is_a?(CArray)
+        result[counts.lt(need)] = UNDEF
+      elsif counts < need
+        result = UNDEF
+      end
+    end
+    CArray.__fill_undefined_result_cells__(result, kw[:fill_value])
+  end
+
   private
 
+  # An unknown keyword is refused before the work is done.
+  def check_reduce_keywords (op, kw)
+    allowed = [:all, :any].include?(op.to_sym) ? AxisGroup::LOGICAL_REDUCE_KEYWORDS : AxisGroup::REDUCE_KEYWORDS
+    unknown = kw.keys - allowed
+    unless unknown.empty?
+      raise ArgumentError,
+            "#{op}: unknown keyword#{unknown.size > 1 ? "s" : ""}: " +
+            unknown.map(&:inspect).join(", ")
+    end
+    CArray.__present_cells_needed_per_result_cell__(op, kw[:min_count], 0)
+    nil
+  end
+
+  # The same grouping over another value.
+  def regroup (value, spec)
+    self.class.send(:build_for, value, spec)
+  end
+
+
   # Build the composite categorical over the group block (the sub-array spanning
   # every group axis). Each group slot -- rank-1 or rank-N -- contributes its
   # per-axis codes; the slot codes are broadcast into the group-block rank (size
@@ -622,7 +831,8 @@ class CAGroupIterator
   # askable. Asking a one-cell array of the same data type lets the core's own
   # refusal through, unworded by us.
   def boolean_payload! (op)
-    CArray.new(value.data_type, [1]).public_send(op)
+    probe = CAIterator.core_probe(value)
+    probe.public_send(op) if probe
     nil
   end
 
@@ -649,6 +859,13 @@ class CAGroupIterator
     [ccat, kdims, gslots, bslots, gaxes]
   end
 
+  # With no band axis the group block is the whole value, grouped in one
+  # pass. With the fast paths off it goes through the band loop instead,
+  # which then has one position.
+  def whole_block? (bslots)
+    bslots.empty? && CArray.__fast_paths_enabled__
+  end
+
   # Iterate the band positions of a band-preserving grouping. For each band
   # position it yields the group block grouped by the composite categorical,
   # together with two indices: val_idx (source order -- nil across the group
@@ -696,8 +913,74 @@ class CAGroupIterator
             "axis_group scan: folding a band into a scan (axis: [:group, k]) " \
             "is not supported; a scan preserves the source shape"
     end
+    return scan_by_core(op) unless kernel_scans?(op)
     group_axes, bundles, = spec.reduce_plan([])
-    value.__axis_group_scan__(group_axes, bundles, op)
+    value.__scan_by_axis_groups__(group_axes, bundles, op)
+  end
+
+  # Whether the scatter kernels read this payload: an integer or a float,
+  # not a Face. They answer there in the core's data types. With the fast
+  # paths off (a development build run with CARRAY_GENERAL_PATH set) each
+  # group is reduced by the core instead -- except a band folded into the
+  # statistic (+fused+), which has no route but the kernel.
+  def kernel_reads_value? (fused = [])
+    return false if value.face? || ! (value.integer? || value.float?)
+    CArray.__fast_paths_enabled__ || ! fused.empty?
+  end
+
+  # Whether the scan kernel gives the core's answer: it reads an integer,
+  # float, boolean or object payload (not a complex, not a Face), and answers cumsum / cumprod in
+  # float64 (an object in object), cummax / cummin in the payload's data type
+  # and cumcount in int64 -- not the core's answer for a boolean's running sum.
+  def kernel_scans? (op)
+    return false unless CArray.__fast_paths_enabled__
+    return false if value.face? || value.complex?
+    return true if value.object? || op == :cumcount
+    core = CAIterator.core_probe(value).public_send(op).data_type
+    kernel = [:cumsum, :cumprod].include?(op) ? CA_FLOAT64 : value.data_type
+    core == kernel
+  end
+
+  # A scan the kernel does not answer as the core does: each group block is
+  # grouped through the composite categorical, whose scans are the core's.
+  def scan_by_core (op, *args, **kw)
+    ccat, _kd, gslots, bslots, gaxes = composite_layout
+    return value.group_by_category(ccat).public_send(op, *args, **kw) if whole_block?(bslots)
+    out = nil
+    each_band_block(ccat, gslots, bslots, gaxes) do |val_idx, _oi, _co, gi|
+      result = gi.public_send(op, *args, **kw)
+      out ||= blank_of(result, value.shape)
+      out[*val_idx] = result
+    end
+    out || blank_of(value, value.shape)
+  end
+
+  # A reduction of a payload the scatter kernel does not read, group by
+  # group through the composite categorical. Folding a band into the
+  # statistic needs the kernel.
+  def reduce_by_core (op, kw)
+    _, fused = AxisGroup.parse_axis(kw[:axis])
+    unless fused.empty?
+      raise ArgumentError,
+            "axis_group #{op}: folding a band into the statistic " \
+            "(axis: [:group, k]) is not supported for a #{value.class} of " \
+            "data type #{value.data_type_name}"
+    end
+    check_reduce_keywords(op, kw)
+    finish_reduction(op, kw, composite_order(op, []))
+  end
+
+  # An array of +shape+ of the same kind as +model+ -- its data type, or its
+  # Face -- with every cell UNDEF. A read-only Face cannot be written cell by
+  # cell, so its values are held in an object array.
+  def blank_of (model, shape)
+    out = if model.face? && model.read_only?
+            CArray.new(CA_OBJECT, shape)
+          else
+            model.__send__(:blank_like, shape)
+          end
+    out[] = UNDEF
+    out
   end
 
   # Dispatch an order statistic (median / percentile / quantile). No :group ->
@@ -707,13 +990,19 @@ class CAGroupIterator
   # operation, and remains a follow-up.
   def order_stat (op, args, kw)
     has_group, fused = AxisGroup.parse_axis(kw[:axis])
-    return value.public_send(op, *args) unless has_group
+    return value.public_send(op, *args, **kw) unless has_group
     unless fused.empty?
       raise NotImplementedError,
             "axis_group: folding a band into an order statistic (axis: [:group, k]) " \
             "is not supported; order statistics gather every member of a group"
     end
-    composite_order(op, args)
+    check_reduce_keywords(op, kw)
+    res = composite_order(op, args)
+    if op == :quantile
+      res.map { |r| finish_reduction(op, kw, r) }
+    else
+      finish_reduction(op, kw, res)
+    end
   end
 
   # Per-group order statistic over the composite grouping. A flat grouping (no
@@ -727,20 +1016,20 @@ class CAGroupIterator
   def composite_order (op, args)
     ccat, kdims, gslots, bslots, gaxes = composite_layout
     nout = (op == :quantile) ? 5 : 1
-    if bslots.empty?
+    if whole_block?(bslots)
       res = value.group_by_category(ccat).public_send(op, *args)
       return nout == 1 ? res.reshape(*kdims) : res.map { |r| r.reshape(*kdims) }
     end
     out_shape = spec.slot_meta.map { |m| m[:kind] == :group ? m[:k] : m[:len] }
-    outs = Array.new(nout) { CArray.float64(*out_shape) }
+    outs = nil
     each_band_block(ccat, gslots, bslots, gaxes) do |_vi, out_idx, _co, gi|
-      result = gi.public_send(op, *args)
-      if nout == 1
-        outs[0][*out_idx] = result.reshape(*kdims)
-      else
-        result.each_with_index { |r, q| outs[q][*out_idx] = r.reshape(*kdims) }
-      end
+      results = gi.public_send(op, *args)
+      results = [results] unless results.is_a?(Array)    # quantile answers five
+      # each output takes the kind the categorical answered in
+      outs ||= results.map { |r| blank_of(r, out_shape) }
+      results.each_with_index { |r, q| outs[q][*out_idx] = r.reshape(*kdims) }
     end
+    outs ||= Array.new(nout) { blank_of(value, out_shape) }
     nout == 1 ? outs[0] : outs
   end
 
@@ -782,7 +1071,15 @@ class CAGroupIterator
     end
     out
   end
+
+  # lag and strip_mask(method:) give the source's values moved or held in
+  # place, so the result shows the source's attributes, as CArray#lag and
+  # CArray#strip_mask(method:) do.
+  prepend(Module.new do
+    def lag (*args, **keywords) = value.__send__(:carry_attrs, super)
+    def strip_mask (**keywords) = value.__send__(:carry_attrs, super)
+  end)
 end
 
 require "carray/categorical"
-CArray.__register_axis_group_classes__(CACategorical, AxisGroup)
+CArray.send(:register_axis_group_classes, CACategorical, AxisGroup)