source: CMIP6dreq/trunk/dreqPy/scope.py @ 489

Subversion URL: http://proj.badc.rl.ac.uk/svn/exarch/CMIP6dreq/trunk/dreqPy/scope.py@489
Revision 489, 17.1 KB checked in by mjuckes, 5 years ago (diff)

test

Line 
1"""Date Request Scoping module
2---------------------------
3The scope.py module contains the dreqQuery class and a set of ancilliary functions. The dreqQuery class contains methods for analysing the data request.
4"""
5import dreq
6from utilities import cmvFilter
7import collections, string, operator
8import sys
9
10python2 = True
11if sys.version_info[0] == 3:
12  python2 = False
13  from functools import reduce
14
15class baseException(Exception):
16  """Basic exception for general use in code."""
17
18  def __init__(self,msg):
19    self.msg = 'scope:: %s' % msg
20
21  def __str__(self):
22    return repr( self.msg )
23
24  def __repr__(self):
25    return self.msg
26
27nt_mcfg = collections.namedtuple( 'mcfg', ['nho','nlo','nha','nla','nlas','nls','nh1'] )
28class cmpd(object):
29  def __init__(self,dct):
30    self.d = dct
31  def cmp(self,x,y,):
32    return cmp( self.d[x], self.d[y] )
33
34    self.default_mcfg = nt_mcfg._make( [259200,60,64800,40,20,5,100] )
35
36def filter1( a, b ):
37  if b < 0:
38    return a
39  else:
40    return min( [a,b] )
41
42npy = {'daily':365, u'Annual':1, u'fx':0.01, u'1hr':24*365, u'3hr':8*365, u'monClim':12, u'Timestep':100, u'6hr':4*365, u'day':365, u'1day':365, u'mon':12, u'yr':1, u'1mon':12, 'month':12, 'year':1, 'monthly':12, 'hr':24*365, 'other':24*365, 'subhr':24*365, 'Day':365, '6h':4*365,
43'3 hourly':8*365, '':1 }
44## There are 4 cmor variables with blank frequency ....
45
46def vol01( sz, v, npy, freq, inx ):
47  n1 = npy[freq]
48  s = sz[inx.uid[v].stid]
49  assert type(s) == type(1), 'Non-integer size found for %s' % v
50  assert type(n1) in (type(1),type(0.)), 'Non-number "npy" found for %s, [%s]' % (v,freq)
51  return s*n1
52
53class col_list(object):
54  def __init__(self):
55    self.a = collections.defaultdict(list)
56
57class col_count(object):
58  def __init__(self):
59    self.a = collections.defaultdict(int)
60
61class dreqQuery(object):
62  __doc__ = """Methods to analyse the data request, including data volume estimates"""
63  def __init__(self,dq=None,tierMax=-1):
64    if dq == None:
65      self.dq = dreq.loadDreq()
66    else:
67      self.dq=dq
68    self.rlu = {}
69    for i in self.dq.coll['objective'].items:
70      k = '%s.%s' % (i.mip,i.label)
71      assert not k in self.rlu, 'Duplicate label in objectives: %s' % k
72      self.rlu[k] = i.uid
73
74    self.cmvFilter = cmvFilter( self )
75    self.tierMax = tierMax
76
77    self.mips = set()
78    for i in  self.dq.coll['requestItem'].items:
79       self.mips.add(i.mip)
80    self.mipls = sorted( list( self.mips ) )
81
82    self.default_mcfg = nt_mcfg._make( [259200,60,64800,40,20,5,100] )
83    self.mcfg = self.default_mcfg._asdict()
84    ##for k in self.default_mcfg.__dict__.keys():
85      ##self.mcfg[k] = self.default_mcfg.__dict__[k]
86    self.szcfg()
87    self.requestItemExpAll(  )
88
89  def szcfg(self):
90    self.szss = {}
91    self.sz = {}
92    for i in self.dq.coll['spatialShape'].items:
93      type = 'a'
94      if i.levelFlag == False:
95        ds =  i.dimensions.split( '|' )
96        if ds[-1] in ['site', 'basin']:
97          vd = ds[-2]
98        else:
99          vd = ds[-1]
100 
101        if vd[:4] == 'olev' or vd == 'rho':
102          type = 'o'
103          nz = self.mcfg['nlo']
104        elif vd[:4] == 'alev':
105          nz = self.mcfg['nla']
106        elif vd in ['slevel','sdepth']:
107          nz = self.mcfg['nls']
108        elif vd == 'aslevel':
109          nz = self.mcfg['nlas']
110        else:
111          print ( 'Failed to parse dimensions %s' % i.dimensions )
112          raise
113      else:
114        nz = i.levels
115
116      dims = set( i.dimensions.split( '|' ) )
117      if 'latitude' in dims and 'longitude' in dims:
118        if type == 'o':
119          nh = self.mcfg['nho']
120        else:
121          nh = self.mcfg['nha']
122      else:
123        nh = 10
124
125      self.szss[i.uid] = nh*nz
126    for i in self.dq.coll['structure'].items:
127      s = self.szss[i.spid]
128      if i.odims != '':
129        s = s*5
130      self.sz[i.uid] = s
131
132  def getRequestLinkByMip( self, mipSel ):
133    """Return the set of request links which are associated with specified MIP"""
134    if type(mipSel) == type(''):
135      t1 = lambda x: x == mipSel
136    elif type(mipSel) == type(set()):
137      t1 = lambda x: x in mipSel
138
139    s = set()
140    for i in self.dq.coll['objectiveLink'].items:
141      if t1(i.label):
142        s.add( self.dq.inx.uid[i.rid] )
143    ##self.rqs = list({self.dq.inx.uid[i.rid] for i in self.dq.coll['objectiveLink'].items if t1(i.label) })
144    self.rqs = list( s )
145    return self.rqs
146
147  def getRequestLinkByObjective( self, objSel ):
148    """Return the set of request links which are associated with specified objectives"""
149    if type(objSel) == type(''):
150      t1 = lambda x: x == self.rlu[objSel]
151    elif type(objSel) == type(set()):
152      t1 = lambda x: x in [self.rlu[i] for i in objSel]
153
154    s = set()
155    for i in self.dq.coll['objectiveLink'].items:
156      if t1(i.label):
157        s.add( self.dq.inx.uid[i.oid] )
158##
159    self.rqs = list( s )
160    ##self.rqs = list({self.dq.inx.uid[i.rid] for i in self.dq.coll['objectiveLink'].items if t1(i.oid) })
161    return self.rqs
162
163  def varGroupXexpt(self, rqList ):
164    """For a list of request links, return a list of variable group IDs for each experiment"""
165    self.cc = collections.defaultdict( list )
166    ## dummy = {self.cc[i.expt].append(i.rlid) for i in self.dq.coll['requestItem'].items if i.rlid in {j.uid for j in rqList} }
167    return self.cc
168
169  def yearsInRequest(self, rql ):
170    self.ntot = sum( [i.ny for i in self.dq.coll['requestItem'].items if i.rlid == rql.uid] )
171    return self.ntot
172
173  def volByExpt( self, l1, ex, exptList, pmax=2, cc=None, retainRedundantRank=False, intersection=False ):
174    """volByExpt: calculates the total data volume associated with an experiment/experiment group and a list of request items.
175          The calculation has some approximations concerning the number of years in each experiment group."""
176##
177## cc: an optional collector, to accumulate indexed volumes
178##
179    inx = self.dq.inx
180    imips = set()
181    for i in l1:
182      imips.add(i.mip)
183    ##imips = {i.mip for i in l1}
184##
185## rql is the set of all request links which are associated with a request item for this experiment set
186##
187    l1p = set()
188    for i in l1:
189      if (i.esid == ex) and (i.preset < 0 or i.preset <= pmax):
190        l1p.add(i)
191
192    rql0 = set()
193    for i in l1p:
194       rql0.add(i.rlid)
195
196    rqlInv = set()
197    for u in rql0:
198      if inx.uid[u]._h.label == 'remarks':
199        rqlInv.add( u )
200    if len(rqlInv) != 0:
201      print ( 'WARNING.001.00002: %s invalid request links from request items ...' % len(rqlInv) )
202    rql = set()
203    for u in rql0:
204       if inx.uid[u]._h.label != 'remarks':
205         rql0.add( u ) 
206
207    if len( rql ) == 0:
208      self.vars = set()
209      return (0,{},{} )
210
211## The complete set of variables associated with these requests:
212    tm = 1
213    if tm == 0:
214      s = set()
215      for i in rql:
216        s.add( inx.uid[i].refid )
217      rqvg = list( s )
218    else:
219      cc1 = collections.defaultdict( set )
220      for i in rql:
221        cc1[inx.uid[i].mip].add( inx.uid[i].refid )
222
223      if intersection:
224        ccv = {}
225#
226# set of request variables for each MIP
227##
228        for k in cc1:
229          thisc = reduce( operator.or_, [set( inx.iref_by_sect[vg].a['requestVar'] ) for vg in cc1[k] ] )
230          s = set()
231          for l in list(thisc):
232             if inx.uid[l].priority <= pmax:
233               s.add( inx.uid[l].vid )
234          ccv[k] = s
235
236        if len( ccv.keys() ) < len( list(imips) ):
237          vars = set()
238        else:
239          vars =  reduce( operator.and_, [ccv[k] for k in ccv] )
240      else:
241        rqvg = reduce( operator.or_, [cc1[k] for k in cc1] )
242
243###To obtain a set of variables associated with this collection of variable groups:
244
245        col1 = reduce( operator.or_, [set( inx.iref_by_sect[vg].a['requestVar'] ) for vg in rqvg ] )
246
247###The collector col1 here accumulates all the record uids, resulting in a single collection. These are request variables, to get a set of CMOR variables at priority <= pmax:
248        vars = set()
249        for l in list(col1):
250           if inx.uid[l].priority <= pmax:
251             vars.add(inx.uid[l].vid)
252##
253## if looking for the union, would have to do a filter here ... after looking up which vars are requested by each MIP ...
254##
255## possibly some code re-arrangement would help.
256## e.g. create a set for each MIP a couple of lines back ....
257
258### filter out cases where the request does not point to a CMOR variable.
259    ##vars = {vid for vid in vars if inx.uid[vid][0] == u'CMORvar'}
260    thisvars = set()
261    for vid in vars:
262       if inx.uid[vid]._h.label == u'CMORvar':
263             thisvars.add(vid)
264    vars = thisvars
265##
266## filter by configuration option and rank
267##
268    if not retainRedundantRank:
269      len1 = len(vars)
270      cmv = self.cmvFilter.filterByChoiceRank(cmv=vars)
271
272      vars = cmv
273   
274    self.vars = vars
275
276    e = {}
277    for u in rql:
278### for request variables which reference the variable group attached to the link, add the associate CMOR variables, subject to priority
279      i = inx.uid[u]
280      e[i.uid] = set()
281      for x in inx.iref_by_sect[i.refid].a['requestVar']:
282           if inx.uid[x].priority <= pmax:
283              e[i.uid].add( inx.uid[x].vid )
284#
285# for each variable, calculate the maximum number of years across all the request links which reference that variable.
286##
287## for each request item we have nymax, nenmax, nexmax.
288##
289    nym = {}
290    for v in vars:
291      s = set()
292      for i in l1p:
293        if i.esid == ex and v in e[i.rlid]:
294          s.add( self.rqiExp[i.uid][2] )
295      ##nym[v] = max( {self.rqiExp[i.uid][2] for i in l1p if i.esid == ex and v in e[i.rlid]} )
296      nym[v] = max( s )
297
298    szv = {}
299    ov = []
300    for v in vars:
301      szv[v] = self.sz[inx.uid[v].stid]*npy[inx.uid[v].frequency]
302      ov.append( self.dq.inx.uid[v] )
303    ee = self.listIndexDual( ov, 'frequency', 'label', acount=None, alist=None, cdict=szv, cc=cc )
304    ff = {}
305    for v in vars:
306      ff[v] = self.sz[ inx.uid[v].stid ] * npy[inx.uid[v].frequency] * nym[v]
307    self.ngptot = sum( [  ff[v]  for v in vars] )
308    return (self.ngptot, ee, ff )
309
310  def esid_to_exptList(self,esid,deref=False):
311    if not esid in self.dq.inx.uid:
312      print ( 'Attempt to dereferece invalid uid: %s' % esid )
313      raise
314
315    if self.dq.inx.uid[esid]._h.label == 'experiment':
316      expts = [esid,]
317    elif self.dq.inx.uid[esid]._h.label != 'remarks':
318      if esid in self.dq.inx.iref_by_sect and 'experiment' in self.dq.inx.iref_by_sect[esid].a:
319        expts = self.dq.inx.iref_by_sect[esid].a['experiment']
320      else:
321        expts = []
322    else:
323      ##print ( 'WARNING: request link not associated with valid experiment group' )
324      ##raise
325      return None
326
327    if self.tierMax > 0:
328      expts = [i for i in expts if self.dq.inx.uid[i].tier <= self.tierMax]
329
330    if deref:
331      return [self.dq.inx.uid[e] for e in expts]
332    else:
333      return expts
334 
335##
336## need to call this on load
337## then use instead of i.ny etc below
338##
339  def requestItemExpAll( self ):
340    self.rqiExp = {}
341    for rqi in self.dq.coll['requestItem'].items:
342      a,b,c = self.requestItemExp( rqi )
343      if a != None:
344        self.rqiExp[rqi.uid] = (a,b,c)
345
346  def requestItemExp( self, rqi ):
347    assert rqi._h.label == "requestItem", 'Argument to requestItemExp must be a requestItem'
348    u = rqi.esid
349    if self.dq.inx.uid[u]._h.label == 'experiment':
350      expts = [u,]
351    elif self.dq.inx.uid[u]._h.label != 'remarks':
352      if u in self.dq.inx.iref_by_sect and 'experiment' in self.dq.inx.iref_by_sect[u].a:
353        expts = self.dq.inx.iref_by_sect[u].a['experiment']
354      else:
355        expts = []
356    else:
357      # print ( 'WARNING: request link not associated with valid experiment group'  )
358      ##rqi.__info__()
359      ##raise
360      return (None, None, None)
361
362    if self.tierMax > 0:
363      expts = [i for i in expts if self.dq.inx.uid[i].tier <= self.tierMax]
364
365    if len(expts) > 0:
366      e = [self.dq.inx.uid[i] for i in expts]
367      dat = [ (i.ntot, i.yps, i.ensz, i.nstart, filter1(i.yps,rqi.nymax), filter1(i.ensz,rqi.nenmax) ) for i in e]
368      nytot = sum( [x[-2]*x[-1] for x in dat ] )
369    else:
370      dat = [ (0,0,0,0,0) ]
371      nytot = 0
372   
373    return (expts, dat, nytot )
374   
375
376  def setTierMax( self, tierMax ):
377    """Set the maxium tier and recompute request sizes"""
378    if tierMax != self.tierMax:
379      self.tierMax = tierMax
380      self.requestItemExpAll(  )
381
382  def summaryByMip( self, pmax=1 ):
383    bytesPerFloat = 2.
384    for m in self.mipls:
385      v = self.volByMip( m, pmax=pmax )
386      print ( '%12.12s: %6.2fTb' % (m,v*bytesPerFloat*1.e-12) )
387
388  def volByMip( self, mip, pmax=2, retainRedundantRank=False):
389
390    if type(mip) in [type( '' ),type( u'') ]:
391      if mip not in self.mips:
392        print ( self.mips )
393        raise baseException( 'volByMip: Name of mip not recognised: %s' % mip )
394      l1 = [i for i in  self.dq.coll['requestItem'].items if i.mip == mip]
395    elif type(mip) == type( set()):
396      nf = [ m for m in mip if m not in self.mips]
397      if len(nf) > 0:
398        raise baseException( 'volByMip: Name of mip(s) not recognised: %s' % str(nf) )
399      l1 = [i for i in  self.dq.coll['requestItem'].items if i.mip in mip]
400    else:
401      raise baseException( 'volByMip: "mip" (1st explicit argument) should be type string or set: %s -- %s' % (mip, type(mip))   )
402     
403    #### The set of experiments/experiment groups:
404    exps = set()
405    for i in l1:
406      exps.add( i.esid )
407    self.volByE = {}
408    vtot = 0
409    cc = collections.defaultdict( col_count )
410    self.allVars = set()
411    for e in exps:
412      expts = self.esid_to_exptList(e,deref=True)
413      if expts != None:
414        self.volByE[e] = self.volByExpt( l1, e, expts, pmax=pmax, cc=cc, retainRedundantRank=retainRedundantRank )
415        vtot += self.volByE[e][0]
416        self.allVars = self.allVars.union( self.vars )
417    self.indexedVol = cc
418
419    return vtot
420
421  def listIndexDual(self, ll, a1, a2, acount=None, alist=None, cdict=None, cc=None ):
422    do_count = acount != None
423    do_list = alist != None
424    assert not (do_count and do_list), 'It is an error to request both list and count'
425    if not (do_count or do_list):
426      acount = '__number__'
427      do_count = True
428
429    if cc == None:
430      if do_count:
431        cc = collections.defaultdict( col_count )
432      elif do_list:
433        cc = collections.defaultdict( col_list )
434
435    if do_count:
436      for l in ll:
437        if cdict != None:
438          v = cdict[l.uid]
439        elif acount == '__number__':
440          v = 1
441        else:
442          v = l.__dict__[acount]
443
444        cc[ l.__dict__[a1] ].a[ l.__dict__[a2] ] += v
445    elif do_list:
446      for l in ll:
447        if cdict != None:
448          v = cdict[l.uid]
449        elif alist == '__item__':
450          v = l
451        else:
452          v = l.__dict__[alist]
453        cc[ l.__dict__[a1] ].a[ l.__dict__[a2] ].append( v )
454
455    od = {}
456    for k in cc.keys():
457      d2 = {}
458      for k2 in cc[k].a.keys():
459        d2[k2] = cc[k].a[k2]
460      od[k] = d2
461    return od
462
463class dreqUI(object):
464  """Data Request Command line.
465-------------------------
466      -v : print version and exit;
467      --unitTest : run some simple tests;
468      -m <mip>:  MIP of list of MIPs (comma separated);
469      -h :       help: print help text;
470      -t <tier> maxmum tier;
471      -p <priority>  maximum priority;
472      --printLinesMax <n>: Maximum number of lines to be printed
473      --printVars  : If present, a summary of the variables fitting the selection options will be printed
474"""
475  def __init__(self,args):
476    self.adict = {}
477    self.knownargs = {'-m':('m',True), '-p':('p',True), '-t':('t',True), '-h':('h',False), '--printLinesMax':('plm',True), '--printVars':('vars',False)} 
478    aa = args[:]
479    while len(aa) > 0:
480      a = aa.pop(0)
481      if a in self.knownargs:
482        b = self.knownargs[a][0]
483        if self.knownargs[a][1]:
484          v = aa.pop(0)
485          self.adict[b] = v
486        else:
487          self.adict[b] = True
488
489    if 'm' in self.adict:
490      self.adict['m'] = set(self.adict['m'].split(',') )
491
492    integerArgs = set( ['p','t','plm'] )
493    for i in integerArgs.intersection( self.adict ):
494      self.adict[i] = int( self.adict[i] )
495
496  def run(self, dq=None):
497    if 'h' in self.adict:
498      print ( self.__doc__ )
499      return
500
501    if not 'm' in self.adict:
502      print ( 'Current version requires -m argument'  )
503      print ( self.__doc__ )
504      sys.exit(0)
505
506    if dq == None:
507      self.dq = dreq.loadDreq()
508    else:
509      self.dq = None
510
511    sc = dreqQuery( dq=self.dq )
512
513    ok = True
514    for i in self.adict['m']:
515        if i not in sc.mips:
516          ok = False
517          print ( 'NOT FOUND: ',i )
518    assert ok,'Available MIPs: %s' % str(sc.mips)
519
520    tierMax = self.adict.get( 't', 2 )
521    sc.setTierMax(  tierMax )
522    pmax = self.adict.get( 'p', 2 )
523    v0 = sc.volByMip( self.adict['m'], pmax=pmax )
524    print ( '%7.2fTb' % (v0*2.*1.e-12) )
525    cc = collections.defaultdict( int )
526    for e in sc.volByE:
527      for v in sc.volByE[e][2]:
528          cc[v] += sc.volByE[e][2][v]
529    x = 0
530    for v in cc:
531      x += cc[v]
532   
533    if python2:
534      vl = sorted( cc.keys(), cmp=cmpd(cc).cmp, reverse=True )
535    else:
536      vl = sorted( cc.keys(), key=lambda x: cc[x], reverse=True )
537    if self.adict.get( 'vars', False ):
538      printLinesMax = self.adict.get( 'plm', 20 )
539      if printLinesMax > 0:
540        mx = min( [printLinesMax,len(vl)] )
541      else:
542        mx = len(vl)
543
544      for v in vl[:mx]:
545        print ( self.dq.inx.uid[v].label, '%7.2fTb' % (cc[v]*2.*1.e-12) )
Note: See TracBrowser for help on using the repository browser.