ruby-grads 1.0.0

data/lib/grads/gridded.rb

@@ -0,0 +1,709 @@
+#
+#
+#
+# A library to read GrADS gridded data into CArray.
+#
+# TODO
+# * options 
+#
+# NOTE
+# * 'dset' should be a file
+#
+# --- supported -------
+# dset
+# dtype
+# title
+# undef
+# options
+# pdef
+# xdef
+# ydef
+# zdef
+# tdef
+# vectorpairs
+# * comment
+#
+# --- not supported ---
+# chsub
+# fileheader
+# theader
+# tailerbytes
+# xyheader
+#
+
+require "carray"
+require "strscan"
+
+module GrADS
+end
+
+class GrADS::Gridded
+end
+
+class GrADS::Gridded::Variable
+
+  def initialize (name, ctl)
+    @name      = name
+    @dim       = ctl.vardims[name]
+    @tsize     = ctl.tsize
+    @dset      = ctl.dset
+    @chunksize = ctl.chunksize
+    @offset    = ctl.varoffsets[name]
+    @undef     = ctl.undef
+  end
+
+  attr_reader :name, :dim, :size, :dset, :offset
+
+  def each_trange (trng=nil)
+    case trng
+    when NilClass, FalseClass
+      (0...@tsize).each do |t|
+        yield(t)
+      end
+    when Range
+      trng.each do |t|
+        yield(t)
+      end
+    else
+      t = trng
+      yield(t)
+    end
+  end
+
+  def tsize (trng=nil)
+    case trng
+    when NilClass, FalseClass
+      return @tsize
+    when Range
+      return trng.to_a.size
+    else
+      return 1
+    end
+  end
+
+  def [] (*argv)
+    trng = argv.shift
+    buf = CArray.float(*@dim)      
+    out = nil
+    open(@dset) { |io|
+      if tsize(trng) == 1
+        each_trange(trng) { |t|
+          io.seek(t * @chunksize + offset)        
+          buf.load_binary(io)
+          out = buf[*argv]
+        }
+      else
+        outlist = []
+        i = 0
+        bind_ok = false
+        each_trange(trng) { |t|
+          t * @chunksize + offset
+          io.seek(t * @chunksize + offset)
+          buf.load_binary(io)      
+          obj = buf[*argv].to_ca
+          if obj.is_a?(CArray)
+            outlist << obj
+          else
+            outlist << CA_FLOAT(obj)
+            bind_ok = true
+          end
+          i += 1
+        }
+        if bind_ok
+          out = CArray.bind(:float, outlist, 0)
+        else
+          out = CArray.merge(:float, outlist, 0)
+        end
+      end
+    }
+    out[:eq, @undef] = UNDEF
+    return out
+  end
+
+  def to_ca
+    return self[]
+  end
+
+end
+
+class GrADS::Gridded
+
+  def initialize (ctl_file)
+    @ctl_file = ctl_file
+    @entries = scan_ctl_file(ctl_file)
+    parse_entries(@entries)
+  end
+
+  attr_reader :ctl_file, :entries
+  attr_reader :dset, :title, :undef
+  attr_reader :pdef, :proj
+  attr_reader :isize, :jsize
+  attr_reader :xsize, :ysize, :zsize, :tsize
+  attr_reader :xaxis, :yaxis, :zaxis, :taxis
+  attr_reader :vectorpair
+  attr_reader :varnames, :vars, :vardims, :varoffsets
+  attr_reader :chunksize
+
+  private
+
+  def scan_ctl_file (ctl_file)
+    io = open(ctl_file, "r:ASCII-8BIT")
+    edefs = ""
+    vdefs = ""
+    headers = ""
+    while line = io.gets
+      case line
+      when /\A\s*vars/i
+        vdefs << line
+        while line = io.gets
+          vdefs << line
+          if line =~ /\A\s*endvars/i
+            break
+          end
+        end
+      when /\A\s*edef/i
+        raise "edef entry is not supported, sorry!"
+        edefs << line
+        while line = io.gets
+          ddefs << line
+          if line =~ /\A\s*endedef/i
+            break
+          end
+        end
+      else
+        headers << line
+      end
+    end
+    entries = scan_headers(headers)
+    entries["vars"] = scan_vars(vdefs)
+    return entries
+  end
+
+  def scan_headers (headers)
+    items = {}
+    buffer = ""
+    headers.each_line do |line|
+      case line
+      when /\A\s*\*/, /\A\s*\Z/
+        next
+      when /A\s*(dtype)\s*(.*)\Z/i
+        raise "data file is not GrADS gridded binary"
+      when /\A\s*(dset|title|undef|options|pdef|xdef|ydef|zdef|tdef|vectorpairs)\s*(.*)\Z/i
+        name = $1.downcase
+        if items[name]
+          items[name] = [items[name]]
+          items[name] << $2.rstrip
+          buffer = items[name].last
+        else
+          items[name] = $2.rstrip
+          buffer = items[name]
+        end
+      when /\A\s*(chsub|fileheader|theader|tailerbyts|xyheader)\s*(.*)\Z/i
+        raise "#{$1} entry is not supported, sorry"
+      else
+        line = line.chomp.sub(/\\\z/,'').rstrip
+        buffer.replace(buffer + " " + line)
+      end
+    end
+    return items
+  end
+
+  def scan_vars (vars)
+    items = {:namelist=>[]}
+    nvars = nil
+    vars.each_line do |line|
+      case line
+      when /\A\s*\*/, /\A\s*\Z/
+        next
+      when /\A\s*vars\s+(\d+)/i
+        nvars = $1.to_i
+      when /\A\s*endvars/i
+        break
+      when /\A\s*(\w+)\s*(.*)\Z/
+        name = $1
+        list = $2
+        items[:namelist].push(name)
+        items[name] = list
+      end
+    end
+    if items[:namelist].size != nvars
+      raise "invalid vars number"
+    end
+    return items
+  end
+
+  def parse_entries (entries)
+    parse_dset(entries["dset"])
+    if entries.has_key?("title")
+      parse_title(entries["title"])
+    end
+    parse_undef(entries["undef"])
+    parse_options(*entries["options"])
+    if entries.has_key?("pdef")
+      parse_pdef(entries["pdef"])
+      parse_xdef(entries["xdef"], with_pdef = true)
+      parse_ydef(entries["ydef"], with_pdef = true)
+    else
+      parse_xdef(entries["xdef"], with_pdef = false)
+      parse_ydef(entries["ydef"], with_pdef = false)
+    end
+    parse_zdef(entries["zdef"])
+    parse_tdef(entries["tdef"])
+    if entries.has_key?("vectorpairs")
+      parse_vectorpairs(entries["vectorpairs"])
+    end
+    parse_vars(entries["vars"][:namelist], entries["vars"])
+  end
+
+  def parse_dset (line)
+    line.strip!
+    if line =~ /\A\^/
+      cwd = File.dirname(File.expand_path(@ctl_file))
+      @dset = File.join(cwd, line[1..-1])
+    else
+      @dset = line
+    end
+  end
+
+  def parse_title (line)
+    @title = line
+  end
+
+  def parse_undef (line)
+    @undef = Float(line)
+  end
+
+  def parse_options (*lines)
+    # yrev             o
+    # zrev             o
+    # template         x
+    # squential        o
+    # 365_day_calendar ?
+    # byteswapped      o
+    # big_endian       o
+    # little_endian    o
+    # cray_32bit_ieee  ?
+  end
+
+  def parse_pdef (line)
+    list = line.split(/\s+/)
+    @pdef  = list
+    @isize = list[0].to_i
+    @jsize = list[1].to_i
+    case list[2].downcase
+    when "nps", "sps"
+      parse_pdef_xps(list[2..-1])
+    when "lcc", "lccr"
+      parse_pdef_lcc(list[2..-1])
+    when "eta.u"
+      parse_pdef_eta_u(list[2..-1])
+    when "pse"
+      parse_pdef_pse(list[2..-1])
+    when "ops"
+      parse_pdef_ops(list[2..-1])
+    when "rotll", "rotllr"
+      parse_pdef_rotll(list[2..-1])
+    when "bilin"      
+      parse_pdef_bilin(list[2..-1])      
+    when "general"
+      parse_pdef_general(list[2..-1])      
+    when "file"      
+      parse_pdef_file(list[2..-1])      
+    else
+      raise "pdef #{list[2].downcase} is not supported"
+    end
+  end
+
+  def parse_pdef_xps (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_pdef_lcc (args)
+    require "proj4r"
+    rlat  = args[1].to_f
+    rlon  = args[2].to_f
+    xi    = args[3].to_f
+    yi    = args[4].to_f
+    slat1 = args[5].to_f
+    slat2 = args[6].to_f
+    slon  = args[7].to_f
+    dx    = args[8].to_f
+    dy    = args[9].to_f
+    slat0 = slat1 > 0 ? 90 : -90
+    proj = PROJ4::Proj.new %{
+      +ellps=sphere +proj=lcc \
+      +lat_1=#{slat1} +lat_2=#{slat2} +lat_0=90 +lon_0=#{slon}
+    }
+    rx, ry = proj.forward(rlon, rlat)
+    rx0 = rx - (xi-1)*dx
+    ry0 = ry - (yi-1)*dy
+    @forward = lambda { |lon, lat|
+      mx, my = proj.forward(lon, lat)
+      [(mx - rx0)/dx, (my - ry0)/dy]
+    }
+    @inverse = lambda { |fi, fj|
+      mx = fi*dx + rx0
+      my = fj*dy + ry0
+      proj.inverse(mx, my)
+    }
+  end
+
+  def parse_pdef_eta_u (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_pdef_pse (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_pdef_ops (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_pdef_rotll (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_pdef_general (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_pdef_file (args)
+    raise NotImplementedError, "not implemented"
+  end
+
+  def parse_xdef (line, with_pdef=false)
+    list = line.split(/\s+/)
+    if with_pdef
+      @xsize = list[0].to_i
+    else
+      @isize = list[0].to_i      
+      @xsize = @isize
+    end
+    case list[1]
+    when /linear/i
+      @xaxis = []
+      @xsize.times do |i|
+        @xaxis << list[2].to_f + i*list[3].to_f
+      end
+    when /levels/i
+      @xaxis = list[2..-1].map{|x| x.to_f}
+    end
+  end
+
+  def parse_ydef (line, with_pdef=false)
+    list = line.split(/\s+/)
+    if with_pdef
+      @ysize = list[0].to_i
+    else
+      @jsize = list[0].to_i
+      @ysize = @jsize
+    end
+    case list[1]
+    when /linear/i
+      @yaxis = []
+      @ysize.times do |i|
+        @yaxis << list[2].to_f + i*list[3].to_f
+      end
+    when /levels/i
+      @yaxis = list[2..-1].map{|y| y.to_f}
+    end
+  end
+
+  def parse_zdef (line)
+    list = line.split(/\s+/)
+    @zsize = list[0].to_i
+    case list[1]
+    when /linear/i
+      @zaxis = []
+      @zsize.times do |i|
+        @zaxis << list[2].to_f + i*list[3].to_f
+      end
+    when /levels/i
+      @zaxis = list[2..-1].map{|z| z.to_f}
+    end
+  end
+
+  def parse_tdef (line)
+    list = line.split(/\s+/)
+    @tsize = list[0].to_i
+    @inittime = DateTime.parse(list[2].sub(/z/i,' '))
+    @taxis = [@inittime]
+    case list[3]
+    when /\A(\d+)(mn|hr|dy|mo|yr)/
+      case $2
+      when "mn"
+        (1..@tsize-1).each do |i|
+          @taxis[i] = @taxis[i-1] + ($1.to_i).quo(1440)
+        end
+      when "hr"
+        (1..@tsize-1).each do |i|
+          @taxis[i] = @taxis[i-1] + ($1.to_i).quo(24)
+        end
+      when "dy"
+        (1..@tsize-1).each do |i|
+          @taxis[i] = @taxis[i-1] + $1.to_i
+        end
+      when "mn"
+        (1..@tsize-1).each do |i|
+          @taxis[i] = @taxis[i-1] >> $1.to_i
+        end
+      when "yr"
+        (1..@tsize-1).each do |i|
+          @taxis[i] = @taxis[i-1] >> (12*$1.to_i)
+        end
+      end
+    else
+      raise "invalid time increment"
+    end
+  end
+
+  def parse_vectorpairs (line)
+    list = line.split(/\s+/)
+    @vectorpair = list.map{|pair| pair.downcase.split(/\s*,\s*/)}
+  end
+
+  def parse_vars (namelist, hash)
+    @varnames = namelist
+    @vars = {}
+    @vardims  = {}
+    @varoffsets = {}
+    @chunksize = 0
+    namelist.each do |name|
+      line = hash[name]
+      ss = StringScanner.new(line)
+      znum  = ss.scan(/\A.*?\s+/).rstrip.to_i
+      units = ss.scan(/\A.*?\s+/).rstrip
+      desc  = ss.rest
+      @vars[name] = [znum, units, desc]
+      if znum == 0
+        @vardims[name] = [@jsize,@isize]
+      else
+        @vardims[name] = [znum, @jsize,@isize]
+      end
+      @varoffsets[name] = @chunksize
+      @chunksize += 4 * @vardims[name].inject(1) {|s,d| s * d}
+    end
+  end
+
+  public
+
+  def inverse (fi, fj)
+    return @inverse[fi,fj]
+  end
+
+  def forward (rlon, rlat)
+    return @forward[rlon, rlat]
+  end
+
+  def var (name)
+    unless @varnames.include?(name)
+      raise "invalid variable name '#{name}'"
+    end
+    return Variable.new(name, self)
+  end
+
+  def template (ctl_file, &block)
+    asm = Writer.new(self)
+    asm.define(&block)
+    asm.write(ctl_file)
+  end
+
+end
+
+class GrADS::Gridded::Writer
+
+  def initialize (ctl = nil)
+    @dset  = nil
+    @title = nil
+    @options = []
+    @vectorpairs = nil
+    @vars = []
+    @nt = 1
+    if ctl
+      e = ctl.entries
+      @undef = e["undef"]
+      @pdef  = e["pdef"]
+      @xdef  = e["xdef"]
+      @ydef  = e["ydef"]
+      @zdef  = e["zdef"]
+      @tdef  = e["tdef"]
+    else
+      @undef = nil
+      @pdef  = nil
+      @xdef  = nil
+      @ydef  = nil
+      @zdef  = nil
+      @tdef  = "1 linear 0z1jan2000"
+    end
+  end
+  
+  private
+  
+  def dset (arg)
+    @dset = arg
+  end
+  
+  def title (arg)
+    @title = arg
+  end
+  
+  def undef! (arg)
+    @undef = arg
+  end
+
+  def options (*args)
+    @options.push(*args)
+  end
+  
+  def pdef (*args)
+    @pdef = args.join(" ")
+  end
+  
+  def xdef (*args)
+    if args.size == 1 and 
+        ( args.first.is_a?(CArray) or args.first.is_a?(Array) )
+      ary = args.first
+      @xdef = [ary.length, "levels", ary.to_a.join("\n")].join(" ")
+    else
+      @xdef = args.join(" ")
+    end
+  end
+
+  def ydef (*args)
+    if args.size == 1 and 
+        ( args.first.is_a?(CArray) or args.first.is_a?(Array) )
+      ary = args.first
+      @ydef = [ary.length, "levels", ary.to_a.join("\n")].join(" ")
+    else
+      @ydef = args.join(" ")
+    end
+  end
+
+  def zdef (*args)
+    if args.size == 1 and 
+        ( args.first.is_a?(CArray) or args.first.is_a?(Array) )
+      ary = args.first
+      @zdef = [ary.length, "levels", ary.to_a.join("\n")].join(" ")
+    else
+      @zdef = args.join(" ")
+    end
+  end
+
+  def tdef (*args)
+    if args.size == 1 and 
+        ( args.first.is_a?(CArray) or args.first.is_a?(Array) )
+      ary = args.first
+      @tdef = [ary.length, "levels", ary.to_a.join("\n")].join(" ")
+    else
+      @tdef = args.join(" ")
+    end
+  end
+
+  def vectorpairs (args)
+    @vectorpairs = args.join(" ")
+  end
+
+  def var2d (name, obj, desc = "x")
+    @vars.push([name, 2, obj, desc])
+    if obj.rank == 3 and obj.dim0 > @nt
+      @nt = obj.dim0
+    end
+  end
+
+  def var3d (name, obj, desc = "x")
+    @vars.push([name, 3, obj, desc])
+    if obj.rank == 4 and obj.dim0 > @nt
+      @nt = obj.dim0
+    end
+  end
+
+  public
+
+  def define (&block)
+    instance_eval(&block)
+  end
+
+  def write (ctl_file)
+    open(ctl_file, "w") { |io|
+      if @dset
+        io.puts "dset #{@dset}"
+      else
+        raise "dset is required"
+      end
+      if @title
+        io.puts "title #{@title}"
+      end
+      io.puts "undef #{@undef}"
+      @options.each do |opt|
+        io.puts %{options #{opt}}
+      end
+      if @pdef
+        io.puts %{pdef #{@pdef}}
+      end
+      io.puts %{xdef #{@xdef}}
+      io.puts %{ydef #{@ydef}}
+      io.puts %{zdef #{@zdef}}
+      io.puts %{tdef #{@tdef}}
+      if @vectorpairs
+        io.puts %{vectorpairs #{@vectorpairs}}
+      end
+
+      if @dset =~ /\A\^/
+        cwd = File.dirname(File.expand_path(ctl_file))
+        dset = File.join(cwd, @dset[1..-1])
+      else
+        dset = @dset
+      end
+
+      io.puts %{vars #{@vars.size}}
+      @vars.each do |name, d, obj, desc|
+        case d
+        when 2
+          io.puts %{#{name} 1 99 #{desc}}
+        when 3
+          case obj.rank 
+          when 3
+            io.puts %{#{name} #{obj.dim0} 99 #{desc} }
+          when 4
+            io.puts %{#{name} #{obj.dim1} 99 #{desc} }
+          else
+            raise "invalid size"
+          end
+        end
+      end
+      io.puts "endvars"
+
+      open(dset, "w") { |dat|
+        @nt.times do |i|
+          @vars.each do |name, d, obj, desc|
+            case d
+            when 2
+              case obj.rank
+              when 2
+                obj.float.dump_binary(dat)
+              when 3
+                obj[i,false].float.dump_binary(dat)
+              end
+            when 3
+              case obj.rank
+              when 3
+                obj.float.dump_binary(dat)
+              when 4
+                obj[i,false].float.dump_binary(dat)
+              end
+            end
+          end
+        end
+      }
+    }
+
+
+  end
+
+  def template (ctlfile, &block)
+    define(&block)
+    write(ctlfile)
+  end
+
+end