carray-netcdf 1.0.13 → 2.0.0

data/lib/carray/netcdf/group.rb

@@ -0,0 +1,517 @@
+module CA
+  module NetCDF
+    # A NetCDF-4 group. The root group is a {File} (a subclass of this).
+    #
+    #   nc.groups             # => {String => Group}
+    #   nc.group("model")     # => Group   (subgroup by name in this group)
+    #   nc.def_group("model") # => Group   (create + return)
+    #   nc["temp"]            # => Var     (in this group)
+    #   nc["/model/temp"]     # => Var     (full path from root)
+    #
+    # Dims defined in an ancestor group are visible to a `def_var` / {Var}
+    # in a descendant (netcdf-4 dimension inheritance), so `def_var`
+    # resolves a dim name by walking up the group tree.
+    class Group
+
+      include CA::NetCDF
+
+      # @param ncid [Integer] group id
+      # @param file [File] root file (for codec + mode + close)
+      # @param parent [Group, nil] parent group (nil for root)
+      def initialize(ncid, file, parent = nil)
+        @ncid   = ncid
+        @file   = file
+        @parent = parent
+        @dims   = {}
+        @vars   = {}
+        @groups = {}
+        load_metadata
+      end
+
+      # @!attribute [r] ncid
+      #   @return [Integer]
+      # @!attribute [r] file
+      #   @return [File] root file
+      # @!attribute [r] parent
+      #   @return [Group, nil]
+      attr_reader :ncid, :file, :parent
+
+      # @return [String] `""` for the root group, `"myGroup"` for a named
+      #   group
+      def name = nc_inq_grpname(@ncid)
+
+      # @return [String] `"/"` for root, `"/model/ocean"` for a nested
+      #   group
+      def path = nc_inq_grpname_full(@ncid)
+
+      # @return [Boolean]
+      def root? = @parent.nil?
+
+      # @!group Introspection (this group only)
+
+      # @return [Hash{String => Dim}]
+      def dims       = @dims
+      # @return [Hash{String => Var}]
+      def vars       = @vars
+      # @return [Hash{String => Group}]
+      def groups     = @groups
+      # @param name [#to_s]
+      # @return [Dim, nil]
+      def dim(name)  = @dims[name.to_s]
+      # @param name [#to_s]
+      # @return [Var, nil]
+      def var(name)  = @vars[name.to_s]
+      # @param name [#to_s]
+      # @return [Group, nil]
+      def group(name) = @groups[name.to_s]
+      # @return [Boolean]
+      def dim?(name)   = @dims.key?(name.to_s)
+      # @return [Boolean]
+      def var?(name)   = @vars.key?(name.to_s)
+      # @return [Boolean]
+      def group?(name) = @groups.key?(name.to_s)
+
+      # Group-scoped global attributes (each group has its own set in
+      # netcdf-4). Cached; invalidated by our own {#def_att}. Other paths
+      # that mutate attributes (`rename_att` / `del_att` / a direct
+      # `nc_put_att`) aren't exposed by layer 2, so we don't have to watch
+      # for them.
+      # @return [Hash{String => Object}] frozen
+      def atts
+        @atts_cache ||= read_group_attributes
+      end
+
+      # @param name [#to_s]
+      # @return [Object, nil]
+      def att(name) = atts[name.to_s]
+
+      # @!endgroup
+
+      # @!group Access
+
+      # Look up a variable by name or path.
+      #
+      #   nc["temp"]         # Var in this group
+      #   nc["/foo/bar/x"]   # Var at that path from the root
+      #   nc["foo/x"]        # Var at that path relative to this group
+      #
+      # @param name_or_path [String, Symbol]
+      # @return [Var]
+      # @raise [KeyError] if the variable / any path segment doesn't exist
+      def [](name_or_path)
+        s = name_or_path.to_s
+        return resolve_path(s) if s.include?("/")
+        @vars[s] or raise KeyError, "no variable #{s.inspect} in group #{path}"
+      end
+
+      # Convenience combined variable-lookup and write.
+      #
+      #   nc[name] = ca                # whole-variable write
+      #   nc[name, *index] = ca        # hyperslab write
+      #   nc["/path/name", ...] = ca   # full-path write
+      #
+      # @return [Object] the value written
+      def []=(*argv)
+        value = argv.pop
+        name  = argv.shift
+        self[name][*argv] = value
+      end
+
+      # @!endgroup
+
+      # @!group Definition (scoped to this group)
+
+      # Define a dimension.
+      # @param name [#to_s]
+      # @param length [Integer, :unlimited]
+      # @return [Dim]
+      def def_dim(name, length)
+        n   = name.to_s
+        len = length == :unlimited ? NC_UNLIMITED : Integer(length)
+        @file.ensure_define_mode
+        @dims[n] = Dim.new(self, nc_def_dim(@ncid, n, len),
+                           unlimited: length == :unlimited)
+      end
+
+      # NetCDF-4 per-variable storage / filter options that {#def_var}
+      # recognizes in its keyword arguments (any other kwarg becomes a
+      # variable attribute). See {#apply_var_options} for accepted values.
+      # @return [Array<Symbol>]
+      VAR_OPTION_KEYS = %i[
+        chunks deflate shuffle fletcher32 endian no_fill
+      ].freeze
+
+      # Define a variable in this group.
+      #
+      # @example Positional
+      #   def_var("temp", :float, ["time","lat","lon"], units: "K")
+      # @example Copy definition
+      #   def_var("temp", template: src["temp"])
+      # @example Look up name in a template group
+      #   def_var("temp", template: src)
+      # @example Rename on copy
+      #   def_var("temp2", template: src["temp"])
+      # @example Override type
+      #   def_var("temp", :double, template: src["temp"])
+      # @example Override / add attributes
+      #   def_var("temp", template: src, units: "F")
+      #
+      # When `template:` is a {Var}, its definition (type, dims, atts) is
+      # used; positional `type` / `dims` and kwargs are treated as
+      # overrides. When `template:` is a {File} / {Group}, its variable of
+      # the same name (first positional arg) is looked up and used as the
+      # template — the rename form must use the {Var} version.
+      #
+      # Kwargs are variable attributes, plus the reserved storage / filter
+      # options (see {VAR_OPTION_KEYS}). Overriding attributes merge into
+      # the template's; a kwarg keyed the same as a template attribute
+      # wins.
+      #
+      # @param name [#to_s] variable name
+      # @param type [Symbol, String, Integer, nil] on-disk type; required
+      #   unless a `template:` is given
+      # @param dims [Array<String, Symbol, Dim>, nil] dimensions
+      # @param template [Var, Group, File, nil] source of a copied
+      #   definition
+      # @param atts [Hash] variable attributes and/or storage options
+      #   (see {VAR_OPTION_KEYS})
+      # @return [Var]
+      def def_var(name, type = nil, dims = nil, template: nil, **atts)
+        n = name.to_s
+
+        # Accept `dims:` as a keyword too — a natural spelling when template:
+        # is in use and the positional slot would otherwise force `nil` fillers.
+        if atts.key?(:dims)
+          raise ArgumentError,
+                "def_var: `dims` given both positionally and as a keyword" if dims
+          dims = atts.delete(:dims)
+        end
+        # Symmetric for type: `type:` as keyword.  Same collision rule.
+        if atts.key?(:type)
+          raise ArgumentError,
+                "def_var: `type` given both positionally and as a keyword" if type
+          type = atts.delete(:type)
+        end
+
+        if template
+          src = case template
+                when Var    then template
+                when Group  then template.var(n) or
+                                   raise KeyError,
+                                         "def_var(#{n.inspect}, template: ...): no variable " \
+                                         "#{n.inspect} in template group #{template.path}"
+                else
+                  raise ArgumentError,
+                        "def_var template: expected Var / Group / File, got #{template.class}"
+                end
+          type ||= src.nc_type
+          dims ||= src.dims.map(&:name)
+          # template.atts (String keys, frozen) merged under any explicit
+          # atts (Symbol / String keys).  Explicit wins on collision.
+          atts = src.atts.merge(atts.transform_keys(&:to_s))
+        end
+
+        raise ArgumentError,
+              "def_var: `type` is required (or pass a template:)" if type.nil?
+        dims ||= []
+
+        # Accept a name (String / Symbol) or a Dim object.  Names are resolved
+        # by walking up the group tree for netcdf-4 dim inheritance.
+        dimids = Array(dims).map { |d| d.is_a?(Dim) ? d.to_i : resolve_dim(d.to_s).to_i }
+        # Split reserved storage/filter kwargs off before treating the rest
+        # as variable attributes.  Per-var kwargs win over the file-wide
+        # defaults (nil / false / 0 are all "explicit"; they override).
+        file_defaults = @file.var_defaults
+        opts = {}
+        VAR_OPTION_KEYS.each do |k|
+          if atts.key?(k)
+            opts[k] = atts.delete(k)
+          elsif file_defaults.key?(k)
+            opts[k] = file_defaults[k]
+          end
+        end
+        @file.ensure_define_mode
+        nct    = nc_type_for(type)
+        varid  = nc_def_var(@ncid, n, nct, dimids)
+        apply_var_options(varid, opts) unless opts.empty?
+        # A CArray value keeps its own on-disk type (so a schema loader's
+        # `_FillValue: {short: -32768}` on a float variable stays short);
+        # a packed-space attribute takes the variable's type (CF §2.5.1).
+        atts.each { |k, v| put_var_att(@ncid, varid, k.to_s, v, nct) }
+        @vars[n] = Var.new(self, varid)
+      end
+
+      # Apply NetCDF-4 storage / filter options on a freshly defined
+      # variable. Kwargs on {#def_var}:
+      #
+      #   chunks: [1, 73, 144]     # NC_CHUNKED with explicit sizes
+      #   chunks: :contiguous      # NC_CONTIGUOUS
+      #   chunks: :compact         # NC_COMPACT (small vars only)
+      #   deflate: 4               # compression level 1..9 (0 / nil = off)
+      #   shuffle: true            # shuffle filter, usually paired with deflate
+      #   fletcher32: true         # per-chunk fletcher32 checksum
+      #   endian: :little / :big / :native
+      #   no_fill: true            # disable pre-fill (perf)
+      #
+      # These all require NetCDF-4; libnetcdf raises on a NetCDF-3 file
+      # and the raise propagates.
+      # @api private
+      private def apply_var_options(varid, opts)
+        if (c = opts[:chunks])
+          case c
+          when :contiguous then nc_def_var_chunking(@ncid, varid, NC_CONTIGUOUS, [])
+          when :compact    then nc_def_var_chunking(@ncid, varid, NC_COMPACT, [])
+          when Array       then nc_def_var_chunking(@ncid, varid, NC_CHUNKED, c)
+          else raise ArgumentError, "chunks: expected Array / :contiguous / :compact, got #{c.inspect}"
+          end
+        end
+        if opts.key?(:deflate) || opts.key?(:shuffle)
+          level   = (opts[:deflate] || 0).to_i
+          on      = level.positive? ? 1 : 0
+          shuffle = opts[:shuffle] ? 1 : 0
+          nc_def_var_deflate(@ncid, varid, shuffle, on, level)
+        end
+        if opts[:fletcher32]
+          nc_def_var_fletcher32(@ncid, varid, 1)
+        end
+        if (e = opts[:endian])
+          endian = case e
+                   when :native then NC_ENDIAN_NATIVE
+                   when :little then NC_ENDIAN_LITTLE
+                   when :big    then NC_ENDIAN_BIG
+                   else raise ArgumentError, "endian: expected :native / :little / :big, got #{e.inspect}"
+                   end
+          nc_def_var_endian(@ncid, varid, endian)
+        end
+        if opts.key?(:no_fill)
+          nc_def_var_fill(@ncid, varid, opts[:no_fill] ? 1 : 0, nil)
+        end
+      end
+
+      # Define a group-level (global) attribute. A CArray value keeps its
+      # own type, as it does on a variable.
+      # @param name [#to_s]
+      # @param value [Object] Ruby value or CArray
+      # @return [Object] the value written
+      def def_att(name, value)
+        @file.ensure_define_mode
+        if value.is_a?(CArray)
+          put_att_typed(@ncid, NC_GLOBAL, name.to_s, value.to_a, TYPE_TO_NC[value.data_type])
+        else
+          nc_put_att(@ncid, NC_GLOBAL, name.to_s, value)
+        end
+        @atts_cache = nil    # invalidate — new attribute is now visible
+        value
+      end
+
+      # Copy a dimension into this group.
+      #
+      # @example
+      #   copy_dim(src.dim("time"))                 # same name
+      #   copy_dim(src.dim("time"), name: "t")      # rename
+      #   copy_dim("time", from: src)               # look up by name in src
+      #
+      # An unlimited source becomes an unlimited destination dim.
+      # Coordinate values (the same-named variable, if any) are NOT
+      # copied — use {#copy_var} separately, or `File.create(template:,
+      # vars: [...])`.
+      #
+      # @param dim_or_name [Dim, String, Symbol]
+      # @param from [Group, File, nil] required when `dim_or_name` is a
+      #   name
+      # @param name [String, Symbol, nil] destination name (defaults to
+      #   the source's)
+      # @return [Dim]
+      def copy_dim(dim_or_name, from: nil, name: nil)
+        src = resolve_dim_source(dim_or_name, from)
+        dst_name = (name || src.name).to_s
+        len = src.unlimited? ? :unlimited : src.length
+        def_dim(dst_name, len)
+      end
+
+      # Copy a variable's definition AND its data into this group.
+      #
+      # @example
+      #   copy_var(src["lat"])                      # same name, values included
+      #   copy_var(src["lat"], name: "latitude")    # rename
+      #   copy_var("lat", from: src)                # look up by name in src
+      #   copy_var("lat", from: src, units: "deg")  # override / add attributes
+      #
+      # The data copy is bit-identical: raw on-disk values are read via
+      # {Var#raw} (no codec) and written via layer 1 (no codec). Users who
+      # need a codec-mediated copy should call `def_var(template:)` and
+      # write `dst[name] = src[name].to_ca` themselves.
+      #
+      # Extra kwargs are forwarded to {#def_var}(template:), so per-var
+      # storage options (`deflate:`, `chunks:`, …) and attribute
+      # overrides both work.
+      #
+      # @param var_or_name [Var, String, Symbol]
+      # @param from [Group, File, nil]
+      # @param name [String, Symbol, nil]
+      # @param overrides [Hash] attribute / storage overrides
+      # @return [Var]
+      def copy_var(var_or_name, from: nil, name: nil, **overrides)
+        src = resolve_var_source(var_or_name, from)
+        dst_name = (name || src.name).to_s
+        def_var(dst_name, template: src, **overrides)
+        dst = @vars[dst_name]
+        ensure_data_mode
+        # Bit-identical raw copy (skip both source's decode and dst's encode).
+        if src.scalar?
+          buf = CArray.new(src.data_type, [1])
+          buf[0] = src.raw
+          nc_put_var(@ncid, dst.varid, buf)
+        else
+          # nc_put_var writes exactly the variable's current shape.  For an
+          # unlimited-dim var, the destination shape starts at length 0 in the
+          # newly-defined var, so nc_put_var would silently write nothing;
+          # nc_put_vara with explicit count grows the unlimited dim to fit.
+          ca = src.raw
+          nc_put_vara(@ncid, dst.varid, ca.shape.map { 0 }, ca.shape, ca)
+        end
+        dst
+      end
+
+      # Define a subgroup.
+      # @param name [#to_s]
+      # @return [Group]
+      def def_group(name)
+        n = name.to_s
+        @file.ensure_define_mode
+        child = nc_def_grp(@ncid, n)
+        @groups[n] = Group.new(child, @file, self)
+      end
+
+      # @!endgroup
+
+      # @!group Forwarding to the root {File} (codec + mode management)
+
+      # @return [Object] the codec used for this file
+      def codec              = @file.codec
+      # @return [void]
+      def ensure_define_mode = @file.ensure_define_mode
+      # @return [void]
+      def ensure_data_mode   = @file.ensure_data_mode
+      # @api private
+      def __transaction__(**opts, &block) = @file.__transaction__(**opts, &block)
+
+      # @!endgroup
+
+      # @return [String]
+      def inspect = "#<#{self.class} #{path.inspect} vars=#{@vars.keys} groups=#{@groups.keys}>"
+
+      # Resolve a dimension name by walking up the group tree — netcdf-4
+      # dimension inheritance.
+      # @param name [#to_s]
+      # @return [Dim]
+      # @raise [KeyError] if not found in the ancestor chain
+      def resolve_dim(name)
+        g = self
+        while g
+          d = g.dim(name)
+          return d if d
+          g = g.parent
+        end
+        raise KeyError, "no dimension #{name.inspect} in scope of #{path}"
+      end
+
+      # The {Dim} with this id, in this group or an ancestor (netcdf-4
+      # dimids are unique across a file, so the id is unambiguous where a
+      # name may be shadowed).
+      # @param dimid [Integer]
+      # @return [Dim, nil]
+      # @api private
+      def dim_by_id(dimid)
+        g = self
+        while g
+          d = g.dims.each_value.find { |x| x.to_i == dimid }
+          return d if d
+          g = g.parent
+        end
+        nil
+      end
+
+      # ---- template / copy source resolution -----------------------------
+      # Common: accept the object directly (Dim / Var) or a name + `from:`
+      # group; raise cleanly on the invalid combinations.
+
+      # @api private
+      private def resolve_dim_source(dim_or_name, from)
+        if dim_or_name.is_a?(Dim)
+          raise ArgumentError,
+                "copy_dim: pass either a Dim or (name, from:), not both" if from
+          dim_or_name
+        else
+          raise ArgumentError,
+                "copy_dim(name, from:): `from:` is required" unless from
+          from.dim(dim_or_name.to_s) or
+            raise KeyError,
+                  "copy_dim: no dimension #{dim_or_name.inspect} in #{from.path}"
+        end
+      end
+
+      # @api private
+      private def resolve_var_source(var_or_name, from)
+        if var_or_name.is_a?(Var)
+          raise ArgumentError,
+                "copy_var: pass either a Var or (name, from:), not both" if from
+          var_or_name
+        else
+          raise ArgumentError,
+                "copy_var(name, from:): `from:` is required" unless from
+          from.var(var_or_name.to_s) or
+            raise KeyError,
+                  "copy_var: no variable #{var_or_name.inspect} in #{from.path}"
+        end
+      end
+
+      # Resolve a "/absolute" or "relative/nested" path down to a Var.
+      # @api private
+      private def resolve_path(str)
+        parts = str.split("/").reject(&:empty?)
+        raise KeyError, "empty path" if parts.empty?
+        leaf = parts.pop
+        start = str.start_with?("/") ? @file : self
+        g = start
+        parts.each do |part|
+          nxt = g.group(part) or
+            raise KeyError, "no group #{part.inspect} in #{g.path}"
+          g = nxt
+        end
+        g[leaf]
+      end
+
+      # @api private
+      private def load_metadata
+        # Dimensions defined IN THIS GROUP (not inherited).  nc_inq_dimids
+        # with include_parents=0 gives exactly this group's dims; falling
+        # back to nc_inq_ndims/nc_inq_dimname for the root of a netcdf-3
+        # file (no groups, no dimids call needed but works anyway).
+        # nc_inq_dimids's include_parents arg is Ruby-truthy in the netcdf gem
+        # (RTEST), so pass `false` — plain `0` is truthy and would include
+        # parent-group dims.  We deliberately want THIS group's dims only.
+        # libnetcdf's nc_inq_dimids/varids/grps work on netcdf-3 too, so no
+        # fallback path is needed; a real error should propagate rather than
+        # be silently swallowed (any rescue here used to also catch
+        # permission errors and pretend "netcdf-3 layout" — misleading).
+        nc_inq_dimids(@ncid, false).each { |d| dm = Dim.new(self, d); @dims[dm.name] = dm }
+        nc_inq_varids(@ncid).each        { |v| vr = Var.new(self, v); @vars[vr.name] = vr }
+        nc_inq_grps(@ncid).each do |gid|
+          g = Group.new(gid, @file, self)
+          @groups[g.name] = g
+        end
+      end
+
+      # @api private
+      private def read_group_attributes
+        atts = {}
+        nc_inq_natts(@ncid).times do |i|
+          name = nc_inq_attname(@ncid, NC_GLOBAL, i)
+          atts[name] = nc_get_att(@ncid, NC_GLOBAL, name)
+        end
+        atts.freeze
+      end
+    end
+  end
+end