source: CONFIG_DEVT/IPSLCM6.5_work_ENSEMBLES/oasis3-mct/util/lucia/pyLucia.py @ 5725

Last change on this file since 5725 was 5725, checked in by aclsce, 3 years ago

Added new oasis3-MCT version to be used to handle ensembles simulations with XIOS.
  • Property svn:executable set to *
File size: 14.3 KB
Line 
1#!/usr/bin/env python3
2
3##########################################################################
4#
5#   pyLucia.py
6#
7#   Visualisation tool
8#
9#   Input:  - timeline files produced by the OASIS coupler
10#             ( load balancing measurement option required )
11#           - json or yaml configuration file (see example below )
12#
13#   Output: - timeline plot of OASIS coupling events
14#             (in graphical format file and visulaisation GUI)
15#
16#   Author : A. Piacentini (2020)
17#            + E. Maisonnave (2020): Color blind friendly default palette
18#
19#
20#   Poem lines :
21#
22#
23#
24##########################################################################
25#
26#  json config file example
27#
28###
29#    {
30#       "Components":[
31#           {"Name":"Atmo",
32#            "File":"timeline_atm.nc"},
33#           {"Name":"Ocean",
34#            "File":"timeline_oce.nc"},
35#           {"Name":"IOserver",
36#            "File":"timeline_ios.nc"}
37#       ],
38#       "Plots":{
39#           "Kind" : true,
40#           "Field": true,
41#           "Component": true
42#       },
43#       "TimeRange":{
44#           "nominFrac":0.25,
45#           "nomaxFrac":0.5,
46#           "nominTime":20,
47#           "nomaxTime":145
48#       },
49#       "Rendering":{
50#           "Display": true,
51#           "File":"Lucia.png",
52#           "EventsBounds": false,
53#           "noPalette":"tab10"
54#       },
55#       "Fields":["Heat","Rain","Love"]
56#   }
57#
58#
59##########################################################################
60#
61#  yaml config file example
62#
63###
64#    ---
65#    Components:
66#    - Name: Atmo
67#      File: timeline_pam_.nc
68#    - Name: Ocean
69#      File: timeline_pim_.nc
70#    - Name: Xios
71#      File: timeline_poum.nc
72#    Plots:
73#      Kind: True
74#      Field: True
75#      Component: True
76#    #TimeRange:
77#    #  minFrac: 0.25
78#    #  maxFrac: 0.5
79#    #  minTime: 10
80#    #  maxTime: 40
81#    Rendering:
82#      Display: True
83#      File: Lucia.png
84#      EventsBounds: False
85#    Fields:
86#    - Heat
87#    - Rain
88#    - Love
89#
90##########################################################################
91 
92import netCDF4
93try:
94    import json
95    has_json = True
96except:
97    has_json = False
98try:
99    import yaml
100    has_yaml = True
101except:
102    has_yaml = False
103import sys, os
104import time
105import numpy as np
106import matplotlib.pyplot as plt
107import matplotlib.collections
108from matplotlib.colors import LinearSegmentedColormap
109from matplotlib.backend_bases import MouseEvent
110import math
111
112if len(sys.argv) == 1:
113    print(">>> Missing configuration file name")
114    exit(1)
115
116config_ok = False
117if has_json:
118    try:
119        jf = open(sys.argv[1])
120        config = json.load(jf)
121        config_ok = True
122    except:
123        pass
124if (not config_ok) and has_yaml:
125    try:
126        jf = open(sys.argv[1])
127        try:
128            config = yaml.load(jf, Loader=yaml.FullLoader)
129            config_ok = True
130        except:
131            try:
132                config = yaml.load(jf)
133                config_ok = True
134            except:
135                pass
136    except:
137        pass       
138if not config_ok:   
139    print(">>> Problem loading configuration file {}".format(sys.argv[1]))
140    exit(1)
141jf.close()
142
143initime = time.time()
144
145nbplots=sum([config["Plots"][i] for i in config["Plots"]])
146if nbplots == 0:
147    print("No plots selected")
148    exit()
149
150
151files = []
152cnam_in = []
153for cp in config["Components"]:
154    files.append(cp["File"])
155    cnam_in.append(cp["Name"])
156
157if "Fields" in config:
158    fieldlb = ["Oasis"]+config["Fields"]
159   
160dofile = "File" in config["Rendering"]
161if dofile:
162    outfile = config["Rendering"]["File"]
163    if outfile:
164        dofile = outfile.lower() != 'none' and outfile.lower() != 'no'
165    else:
166        dofile = False
167
168af_in = [netCDF4.Dataset(fi,'r') for fi in files]
169comp_id = [int(tf.getncattr('component_id'))-1 for tf in af_in]
170cnam = [cnam_in[id] for id in comp_id]
171af = [af_in[id] for id in comp_id]
172
173udpal = "Palette" in config["Rendering"]
174if udpal:
175    udpalette = config["Rendering"]["Palette"]
176
177totprocs = -1
178cprocs = []
179fprocs = []
180for i,tf in enumerate(af):
181    timeini = time.time()
182    nevents = len(tf.dimensions['nx'])
183    nprocs  = len(tf.dimensions['ny'])
184    tstrt   = tf.variables['timer_strt'][:,:].flatten()
185    tstop   = tf.variables['timer_stop'][:,:].flatten()
186    if config["Plots"]["Kind"]:
187        kind    = tf.variables['kind'][:]
188        kindlb  = tf.variables['kind'].getncattr('flag_meanings').split()
189    if config["Plots"]["Field"]:
190        field   = tf.variables['field'][:]
191    if config["Plots"]["Component"]:
192        compo   = tf.variables['component'][:]
193    tf.close()
194
195    if config["Plots"]["Kind"]:
196        lpkind  = np.tile(kind,nprocs)
197    if config["Plots"]["Field"]:
198        lpfield = np.tile(field,nprocs)
199    if config["Plots"]["Component"]:
200        lpcompo = np.tile(compo,nprocs)
201
202    procs = np.arange(nprocs)+totprocs+1
203    procs = np.repeat(procs,nevents)
204    totprocs += nprocs
205    cprocs.append(totprocs+1)
206    fprocs.append(totprocs+1.5)
207
208    polyx = np.array([tstrt,tstop,tstop,tstrt]).T
209    polyy = np.array([procs+0.5,procs+0.5,procs+1.5,procs+1.5]).T
210    lpolyxy = np.dstack((polyx[...,np.newaxis],polyy[...,np.newaxis]))
211
212    if i == 0:
213        if config["Plots"]["Kind"]:
214            pkind = np.copy(lpkind)
215        if config["Plots"]["Field"]:
216            pfield = np.copy(lpfield)
217        if config["Plots"]["Component"]:
218            pcompo = np.copy(lpcompo)
219        polyxy = np.copy(lpolyxy)
220    else:
221        if config["Plots"]["Kind"]:
222            pkind = np.hstack((pkind,lpkind))
223        if config["Plots"]["Field"]:
224            pfield = np.hstack((pfield,lpfield))
225        if config["Plots"]["Component"]:
226            pcompo = np.hstack((pcompo,lpcompo))
227        polyxy = np.concatenate((polyxy,lpolyxy))
228    print('Loaded {} in {} sec.'.format(cnam[i],time.time()-timeini))
229
230# Time range selection
231
232timeini = time.time()
233
234if config["Plots"]["Kind"]:
235    minpkind = np.min(pkind)
236    maxpkind = np.max(pkind)
237if config["Plots"]["Field"]:
238    minpfield = np.min(pfield)
239    maxpfield = np.max(pfield)
240if config["Plots"]["Component"]:
241    minpcompo = np.min(pcompo)
242    maxpcompo = np.max(pcompo)
243
244trange = np.max(polyxy[:,1,0])-np.min(polyxy[:,0,0])
245print("The full trace spans {} sec. and contains {} events".format(trange,polyxy.shape[0]))
246
247uselimits = "TimeRange" in config
248if uselimits:
249    if "minFrac" in config["TimeRange"]:
250        tmin = np.min(polyxy[:,0,0]) + config["TimeRange"]["minFrac"] * trange
251    else:
252        if "minTime" in config["TimeRange"]:
253            tmin = config["TimeRange"]["minTime"]
254        else:
255            tmin = np.min(polyxy[:,0,0])
256    if "maxFrac" in config["TimeRange"]:
257        tmax = np.min(polyxy[:,0,0]) + config["TimeRange"]["maxFrac"] * trange
258    else:
259        if "maxTime" in config["TimeRange"]:
260            tmax = config["TimeRange"]["maxTime"]
261        else:
262            tmax = np.max(polyxy[:,1,0])
263else:
264    tmin = np.min(polyxy[:,0,0])
265    tmax = np.max(polyxy[:,1,0])
266tmin = float(math.floor(tmin))
267tmax = float(math.ceil(tmax))
268if uselimits:
269    ti_msk = np.less_equal(polyxy[:,1,0],tmin)
270    ti_msk = np.logical_or(ti_msk,np.greater_equal(polyxy[:,0,0],tmax))
271   
272    if config["Plots"]["Kind"]:
273        pkind = np.delete(pkind,ti_msk)
274    if config["Plots"]["Field"]:
275        pfield = np.delete(pfield,ti_msk)
276    if config["Plots"]["Component"]:
277        pcompo = np.delete(pcompo,ti_msk)
278    polyxy = np.delete(polyxy,ti_msk,axis=0)
279    print("The selection spans {} sec. between {} and {} and contains {} events".format(tmax-tmin,tmin,tmax,polyxy.shape[0]))
280    print('Time range selection took {} sec'.format(str(time.time()-timeini)))
281
282
283# Plotting
284
285timeini = time.time()
286
287if nbplots == 1:
288    figsz=(11.75,8.25)
289else:
290    figsz=(8.25,11.75)
291
292fig=plt.figure(figsize=figsz, frameon=True)
293
294# Define default color blind friendly palette
295okabe_ito = ['#E69F00', '#56B4E9', '#009E73', '#F0E442', '#0072B2', '#D55E00', '#CC79A7', '#000000']
296cmap_name = "mycolor"
297
298compolb = ['Oasis']+cnam
299compopc = [0]+cprocs
300plotnb = 0
301
302if config["Plots"]["Kind"]:
303    cm = LinearSegmentedColormap.from_list(cmap_name, okabe_ito, N=10)
304    if udpal:
305        palette = udpalette # 'tab10' # or 'Paired'
306    else :
307        palette = cm
308    plotnb += 1
309    kind_ax = plt.subplot(nbplots*100+10+plotnb)
310    kind_ax.set_ylim([0.5,totprocs+1.5])
311    kind_ax.set_xlim([tmin,tmax])
312    kind_pc = matplotlib.collections.PolyCollection(polyxy,pickradius=0.0)
313    kind_pc.set_array(pkind)
314    kind_pc.set_clim([minpkind-.5,maxpkind+.5])
315    cmap = plt.get_cmap(palette, maxpkind-minpkind+1)
316    cbar = plt.colorbar(kind_pc, ticks=np.arange(minpkind,maxpkind+1))
317    kind_pc.set_cmap(cmap)
318   
319    plt.hlines(fprocs[:-1],tmin,tmax,colors='black',linestyles='dashed', linewidth=0.8)
320    kind_ax.add_collection(kind_pc)
321    kind_ax.set_title('KIND')
322    cbar.ax.set_yticklabels(kindlb[minpkind:maxpkind+1])
323    for i,label in enumerate(compolb[1:],start=1):
324        kind_ax.text(tmax*1.005,
325                     compopc[i-1]+(compopc[i]-compopc[i-1])/2,
326                     compolb[i],{'ha': 'left', 'va': 'center'},
327                     rotation = 60)
328
329if config["Plots"]["Field"]:
330    cm = LinearSegmentedColormap.from_list(cmap_name, okabe_ito, N=len(fieldlb))
331    if udpal:
332        palette = udpalette # 'tab10' # or 'Paired'
333    else :
334        palette = cm
335    plotnb += 1
336    field_ax = plt.subplot(nbplots*100+10+plotnb)
337    field_ax.set_ylim([0.5,totprocs+1.5])
338    field_ax.set_xlim([tmin,tmax])
339    field_pc = matplotlib.collections.PolyCollection(polyxy,pickradius=0.0)
340    field_pc.set_array(pfield)
341    field_pc.set_clim([minpfield-.5,maxpfield+.5])
342    cmap = plt.get_cmap(palette, len(fieldlb))
343    cbar = plt.colorbar(field_pc, ticks=np.arange(len(fieldlb)))
344    field_pc.set_cmap(cmap)
345   
346    plt.hlines(fprocs[:-1],tmin,tmax,colors='black',linestyles='dashed', linewidth=0.8)
347    field_ax.add_collection(field_pc)
348    field_ax.set_title('FIELD')
349    if "Fields" in config:
350        cbar.ax.set_yticklabels(fieldlb)
351    for i,label in enumerate(compolb[1:],start=1):
352        field_ax.text(tmax*1.005,
353                      compopc[i-1]+(compopc[i]-compopc[i-1])/2,
354                      compolb[i],{'ha': 'left', 'va': 'center'},
355                      rotation = 60)
356
357if config["Plots"]["Component"]:
358    cm = LinearSegmentedColormap.from_list(cmap_name, okabe_ito, N=len(compolb))
359    if udpal:
360        palette = udpalette # 'tab10' # or 'Paired'
361    else :
362        palette = cm
363    plotnb += 1
364    compo_ax = plt.subplot(nbplots*100+10+plotnb)
365    compo_ax.set_ylim([0.5,totprocs+1.5])
366    compo_ax.set_xlim([tmin,tmax])
367    compo_pc = matplotlib.collections.PolyCollection(polyxy,pickradius=0.0)
368    compo_pc.set_array(pcompo)
369    compo_pc.set_clim([-.5,len(compolb)-.5])
370    cmap = plt.get_cmap(palette, len(compolb))
371    cbar = plt.colorbar(compo_pc, ticks=np.arange(len(compolb)))
372    compo_pc.set_cmap(cmap)
373   
374    plt.hlines(fprocs[:-1],tmin,tmax,colors='black',linestyles='dashed', linewidth=0.8)
375    compo_ax.add_collection(compo_pc)
376    compo_ax.set_title('COMPONENT')
377    cbar.ax.set_yticklabels(compolb)
378    for i,label in enumerate(compolb[1:],start=1):
379        compo_ax.text(tmax*1.005,
380                      compopc[i-1]+(compopc[i]-compopc[i-1])/2,
381                      compolb[i],{'ha': 'left', 'va': 'center'},
382                      rotation = 60)
383
384plt.subplots_adjust(left=0.07,right=1.05,bottom=0.03,top=0.95,wspace=0.0,hspace=0.2)
385
386print('Plot preparation took {} sec'.format(str(time.time()-timeini)))
387
388if dofile:
389    timeini = time.time()
390    plt.savefig(outfile,bbox_inches='tight')
391    print('File output took {} sec'.format(str(time.time()-timeini)))
392print('Overall time {} sec'.format(str(time.time()-initime)))
393
394if config["Rendering"]["Display"]:
395    if "EventsBounds" not in config["Rendering"]:
396        config["Rendering"]["EventsBounds"] = False
397    if config["Rendering"]["EventsBounds"]:
398        if config["Plots"]["Kind"]:
399            kind_pc.set_edgecolor('black')
400            kind_pc.set_linewidth(0.3)
401        if config["Plots"]["Field"]:
402            field_pc.set_edgecolor('black')
403            field_pc.set_linewidth(0.3)
404        if config["Plots"]["Component"]:
405            compo_pc.set_edgecolor('black')
406            compo_pc.set_linewidth(0.3)
407       
408    class pseudo_mouse:
409        def __init__(self,x,y):
410            self.x = x
411            self.y = y
412
413    if config["Plots"]["Kind"]:
414        def format_coord_kind(x, y):
415            comp = compolb[min(np.searchsorted(fprocs,y)+1,len(compolb)-1)]
416            ip = min(max(1,int(round(y))),compopc[-1])
417            proc = ip-1 - compopc[np.searchsorted(fprocs,y)]
418            px, py = kind_pc.get_transform().transform((x,y))
419            cont, ind = kind_pc.contains(pseudo_mouse(px,py))
420            if cont:
421                return 'Time: {:.3f}, Resource: {:.0f}, Component: {}, Rank: {:.0f}, Kind: {}'.format(x, ip, comp, proc, kindlb[pkind[ind['ind'][0]]])
422            else:
423                return 'Time: {:.3f}, Resource: {:.0f}, Component: {}, Rank: {:.0f}'.format(x, ip, comp, proc)
424
425        kind_ax.format_coord = format_coord_kind
426    if config["Plots"]["Field"]:
427        def format_coord_field(x, y):
428            comp = compolb[min(np.searchsorted(fprocs,y)+1,len(compolb)-1)]
429            ip = min(max(1,int(round(y))),compopc[-1])
430            proc = ip-1 - compopc[np.searchsorted(fprocs,y)]
431            px, py = field_pc.get_transform().transform((x,y))
432            cont, ind = field_pc.contains(pseudo_mouse(px,py))
433            if cont:
434                return 'Time: {:.3f}, Resource: {:.0f}, Component: {}, Rank: {:.0f}, Field: {}'.format(x, ip, comp, proc, fieldlb[pfield[ind['ind'][0]]])
435            else:
436                return 'Time: {:.3f}, Resource: {:.0f}, Component: {}, Rank: {:.0f}'.format(x, ip, comp, proc)
437
438        field_ax.format_coord = format_coord_field
439    if config["Plots"]["Component"]:
440        def format_coord_compo(x, y):
441            comp = compolb[min(np.searchsorted(fprocs,y)+1,len(compolb)-1)]
442            ip = min(max(1,int(round(y))),compopc[-1])
443            proc = ip-1 - compopc[np.searchsorted(fprocs,y)]
444            px, py = compo_pc.get_transform().transform((x,y))
445            cont, ind = compo_pc.contains(pseudo_mouse(px,py))
446            if cont:
447                return 'Time: {:.3f}, Resource: {:.0f}, Component: {}, Rank: {:.0f}, Comp: {}'.format(x, ip, comp, proc, compolb[pcompo[ind['ind'][0]]])
448            else:
449                return 'Time: {:.3f}, Resource: {:.0f}, Component: {}, Rank: {:.0f}'.format(x, ip, comp, proc)
450
451        compo_ax.format_coord = format_coord_compo
452    plt.show()
Note: See TracBrowser for help on using the repository browser.