source: trunk/LMDZ.GENERIC/utilities/photochemistry/photochem_postproc.py @ 3431

Last change on this file since 3431 was 3431, checked in by mmaurice, 2 months ago

Generic PCM:

Add photochemistry postprocessing and visualization python routines
along with an introduction notebook to utilities.

MM
File size: 24.0 KB
Line 
1"""
2Generic PCM Photochemistry post-processing library
3Written by Maxime Maurice in 2024 anno domini
4"""
5
6import numpy as np
7import matplotlib.pyplot as plt
8import matplotlib.ticker as tk
9import xarray as xr
10from scipy.constants import R, N_A
11
12import warnings
13warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 't'")
14warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 'lon'")
15warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 'lat'")
16
17M          = {'co2':44,  # Molar masses
18              'o':16,    # TODO: automatic parser
19              'o1d':16,
20              'o2':32,
21              'o3':48,
22              'h':1,
23              'h2':2,
24              'oh':17,
25              'h2o_vap':18,
26              'ho2':33,
27              'h2o2':34,
28              'co':28,
29              'cho':29,
30              'ch2o':30}
31
32background = 'co2' # background gas of the atmosphere
33
34class GPCM_simu:
35    """ Generic PCM Simulation
36
37    Stores the netCDF file and path to the simulation
38
39    Attributes
40    ----------
41    data : xr.Dataset
42        NetCDF file of the simulation
43    path : str
44        Path to simulation directory
45
46    Methods
47    -------
48    get_profile(field,**kw)
49        Get profile of a field (either local or averaged)
50    plot_meridional_slice(field,logcb,labelcb,**kw)
51        Plot a meridional slice of a field
52    plot_time_evolution(field,logcb,labelcb,**kw)
53        Plot time evolution of a field (as a time vs altitude contour plot)
54    plot_profile(field,**kw)
55        Plot a profile of a field
56
57    """
58    def __init__(self,path,datafilename='diagfi',verbose=False):
59        """
60        Parameters
61        ----------
62        path : str
63            Path to simulation directory
64        datafilename : str (optional)
65            Name of the netCDF file (by default: diagfi)
66            works with start, startfi, concat etc.
67            Do not add .nc type suffix
68        """
69        self.path = path
70        try:
71            self.data = xr.open_dataset(path+'/'+datafilename+'.nc',decode_times=False)
72            print(path+'/'+datafilename,'loaded, simulations lasts',self.data['Time'].values[-1],'sols')
73        except:
74            raise Exception('Data not found')
75
76    def __getitem__(self,field):
77        return self.data[field]
78
79    def __setitem__(self,field,value):
80        self.data[field] = value
81
82    def get_profile(self,field,**kw):
83        """ Get profile of a field (either local or averaged)
84
85        Parameters
86        ----------
87        field : str
88            Field name
89        t : float (optional)
90            Time at which to select (if nothing specified use time-average)
91        lat : float (optional)
92            Latitude at which to select (if nothing specified use area-weighted meridional average)
93        lon : float (optional)
94            Longitude at which to select (if nothing specified use zonal average)
95
96        """
97       
98        if self['latitude'].size == 1:
99            # 1D
100            return self[field][:,0,0]
101        else:
102            # 3D
103            if 'lat' in kw:
104                if 'lon' in kw:
105                    return self[field].sel(latitude=kw['lat'],method='nearest').sel(longitude=kw['lon'],method='nearest')
106                else:
107                    return self[field].sel(latitude=kw['lat'],method='nearest').mean(dim='longitude')
108            else:
109                # Latitude-averaged profile: need to weight by the grid cell surface area
110                if 'lon' in kw:
111                    return (self['aire']*self[field]/self['aire'].mean(dim='latitude')).mean(dim='latitude').sel(longitude=kw['lon'],method='nearest')
112                else:
113                    return (self['aire']*self[field]/self['aire'].mean(dim='latitude').mean(dim='longitude')).mean(dim='latitude').mean(dim='longitude')
114
115    def get_subset(self,field='all',**kw):
116        """ Get a subset at fixed given coordinate of the data
117
118        Parameters
119        ----------
120        field : str (optional, default = 'all')
121            Field name. If nothing or 'all'
122            specified, return all fields.
123        t : float (optional)
124            Time of the slice. If nothing
125            specified, use time average
126        lon : float (optional)
127            Longitude of the slice. If nothing
128            specified, use meridional average
129        lat : float (optional)
130            Latitude of the slice. If nothing
131            specified, use area-weighted zonal average
132        alt : float (optional)
133            Altitude of the slice. If nothing
134            specified, use time-average
135
136        Raise
137        -----
138        Slice direction not provided
139        """
140        if len(kw) == 0:
141            raise Exception('Slice direction not provided')
142
143        if field == 'all':
144            data_subset = self.data
145        else:
146            data_subset = self[field]
147           
148        if 't' in kw:
149            data_subset = data_subset.sel(Time=kw['t'],method='nearest')
150        if 'lon' in kw:
151            data_subset = data_subset.sel(longitude=kw['lon'],method='nearest')
152        if 'lat' in kw:
153            data_subset = data_subset.sel(latitude=kw['lat'],method='nearest')
154        if 'alt' in kw:
155            data_subset = data_subset.sel(altitude=kw['alt'],method='nearest')
156
157        return data_subset
158
159    def plot_meridional_slice(self,field,t='avg',lon='avg',logcb=False,labelcb=None,**plt_kw):
160        """ Plot a meridional slice of a field
161
162        Parameters
163        ----------
164        field : str
165            Field name to plot
166        logcb : bool (optional)
167            Use logarithmic colorscale
168        labelcb : str (optional)
169            Use custom colorbar label
170        t : float (keyword)
171            Time at which to plot (if nothing specified use time average)
172        lon : float (keyword)
173            Longitude at which to plot (if nothing specified use zonal average)
174        """
175
176        if self['latitude'].size == 1:
177            # safety check
178            raise Exception('Trying to plot a meridional slice of a 1D simulation')
179
180        meridional_slice = self[field]
181       
182        if t == 'avg':
183            meridional_slice = meridional_slice.mean(dim='Time')
184        else:
185            meridional_slice = meridional_slice.sel(Time=t,method='nearest')
186        if lon == 'avg':
187            meridional_slice = meridional_slice.mean(dim='longitude')
188        else:
189           meridional_slice = meridional_slice.sel(longitude=lon,method='nearest') 
190           
191        if logcb:
192            plt.contourf(self['latitude'],self['altitude'],meridional_slice,locator=tk.LogLocator(),**plt_kw)
193        else:
194            plt.contourf(self['latitude'],self['altitude'],meridional_slice,**plt_kw)
195           
196        plt.colorbar(label=field if labelcb==None else labelcb)
197        plt.xlabel('latitude [°]')
198        plt.ylabel('altitude [km]')
199
200    def plot_time_evolution(self,field,lat='avg',lon='avg',logcb=False,labelcb=None,**plt_kw):
201        """ Plot time evolution of a field (as a time vs altitude contour plot)
202
203        Parameters
204        ----------
205        field : str
206            Field name to plot
207        lat : float (optional)
208            Latitude at which to plot (if nothing specified use area-weighted meridional average)
209        lon : float (optional)
210            Longitude at which to plot (if nothing specified use zonal average)
211        logcb : bool (optional)
212            Use logarithmic colorscale
213        labelcb : str (optional)
214            Use custom colorbar label
215        matplotlib contourf keyword arguments
216        """
217
218        time_evolution = self[field]
219       
220        if lat == 'avg':
221            time_evolution = (self['aire']*time_evolution/self['aire'].mean(dim='latitude')).mean(dim='latitude')
222        else:
223            time_evolution = time_evolution.sel(latitude=lat,method='nearest')
224        if lon == 'avg':
225            time_evolution = time_evolution.mean(dim='longitude')
226        else:
227            time_evolution = time_evolution.sel(longitude=lon,method='nearest')
228           
229        if logcb:
230            plt.contourf(self['Time'],self['altitude'],time_evolution.T,locator=tk.LogLocator(),**plt_kw)
231        else:
232            plt.contourf(self['Time'],self['altitude'],time_evolution.T,**plt_kw)
233           
234        plt.colorbar(label=field if labelcb==None else labelcb)
235        plt.xlabel('time [day]')
236        plt.ylabel('altitude [km]')
237
238    def plot_atlas(self,field,t='avg',alt='avg',logcb=False,labelcb=None,**plt_kw):
239        """ Plot atlas of a field
240
241        Parameters
242        ----------
243        field : str
244            Field name to plot
245        t : float (optional)
246            Time at which to pot (if nothing specified, use time average)
247        alt : float (optional)
248            Altitude at which to plot (if nothing specified use vertical average)
249        logcb : bool (optional)
250            Use logarithmic colorscale
251        labelcb : str (optional)
252            Use custom colorbar label
253        matplotlib contourf keyword arguments
254        """
255
256        atlas = self[field]
257       
258        if t == 'avg':
259            atlas = atlas.mean(dim='Time')
260        else:
261            atlas = atlas.sel(Time=t,method='nearest')
262           
263        if  'altitude' in atlas.coords:
264            if alt == 'avg':
265                atlas = atlas.mean(dim='altitude')
266            else:
267                atlas = atlas.sel(altitude=alt,method='nearest')
268           
269        if logcb:
270            plt.contourf(self['longitude'],self['latitude'],atlas,locator=tk.LogLocator(),**plt_kw)
271        else:
272            plt.contourf(self['longitude'],self['latitude'],atlas,**plt_kw)
273           
274        plt.colorbar(label=field if labelcb==None else labelcb)
275        plt.xlabel('longitude [°]')
276        plt.ylabel('matitude [°]')
277
278    def plot_profile(self,field,t='avg',lon='avg',lat='avg',logx=False,**plt_kw):
279        """ Plot a profile of a field
280
281        Parameters
282        ----------
283        field : str
284            Field name to plot
285        logx : bool (optional)
286            Use logarithmic x axis
287        t : float (optional)
288            Time at which to select (if nothing specified use time-average)
289        lat : float (optional)
290            Latitude at which to plot (if nothing specified use area-weighted meridional average)
291        lon : float (optional)
292            Longitude at which to plot (if nothing specified use zonal average)
293        matplotlib's plot / semilogx keyword arguments
294        """
295
296        profile = self[field]
297       
298        if t == 'avg':
299            profile = profile.mean(dim='Time')
300        else:
301            profile = profile.sel(Time=t,method='nearest')
302           
303        if self['latitude'].size > 1:
304           
305            if lat == 'avg':
306                profile = (self['aire']*profile/self['aire'].mean(dim='latitude')).mean(dim='latitude')
307            else:
308                profile = profile.sel(latitude=lat,method='nearest')
309            if lon == 'avg':
310                profile = profile.mean(dim='longitude')
311            else:
312                profile = profile.sel(longitude=lon,method='nearest')
313           
314        if logx:
315            plt.semilogx(profile,self['altitude'],**plt_kw)
316        else:
317            plt.plot(profile,self['altitude'],**plt_kw)
318           
319        plt.xlabel(field+' ['+self[field].units+']')
320        plt.ylabel('altitude [km]')
321           
322class reaction:
323    """ Instantiates a basic two-body reaction
324   
325    Attributes
326    ----------
327    formula : str
328        Reaction formula (e.g. "A + B -> C + D")
329    reactants : list
330        Reactanting molecules formulae (e.g. ["A", "B"])
331    products : list
332        Produced molecules formulae (e.g. ["C", "D"])
333    constant : fun
334        Reaction rate constant, function of potentially temperature and densities
335
336    Methods
337    -------
338    rate(T,densities)
339        Reaction rate for given temperature and densities
340    """
341   
342    def __init__(self,reactants,products,constant):
343        """
344        Parameters
345        ----------
346        reactants : list
347            Reacting molecules formulae (e.g. ["A", "B"])
348        products : list
349            Produced molecules formulae (e.g. ["C", "D"])
350        constant : fun
351            Reaction rate constant, function of potentially temperature and densities
352        """
353       
354        self.formula   = ''.join([r_+' + ' for r_ in reactants[:-1]])+reactants[-1]+' -> '+''.join([r_+' + ' for r_ in products[:-1]])+products[-1]
355        self.products  = products
356        self.reactants = reactants
357        self.constant  = constant
358
359    def rate(self,T,densities):
360        """ Computes reaction rate
361
362        Parameters
363        ----------
364        T : float
365            Temperature [K]
366        densities : dict
367            Molecular densities [cm^-3]
368
369        Returns
370        -------
371        float
372            Value of the reaction rate [cm^-3.s^-1]
373        """
374       
375        return self.constant(T,densities[background])*densities[self.reactants[0]]*densities[self.reactants[1]]
376
377class termolecular_reaction(reaction):
378
379    def __init__(self,reactants,products,constant):
380
381        self.formula   = ''.join([r_+' + ' for r_ in reactants[:-1]])+reactants[-1]+' + M -> '+''.join([r_+' + ' for r_ in products[:-1]])+products[-1]+' + M'
382        self.products  = products
383        self.reactants = reactants
384        self.constant  = constant
385
386class photolysis(reaction):
387
388    def __init__(self,reactants,products):
389
390        self.formula   = ''.join([r_+' + ' for r_ in reactants[:-1]])+reactants[-1]+' + hv -> '+''.join([r_+' + ' for r_ in products[:-1]])+products[-1]
391        self.products  = products
392        self.reactants = reactants
393
394    def rate(self,j,densities):
395        """ Computes reaction rate
396
397        Parameters
398        ----------
399        j : float
400            Photolysis rate [s^-1]
401        densities : dict
402            Molecular densities [cm^-3]
403
404        Returns
405        -------
406        float
407            Value of the reaction rate [cm^-3.s^-1]
408        """
409
410        return j*densities[self.reactants[0]]
411
412class reaction_constant:
413    """ Basic (Arrhenius) reaction rate constant
414   
415    Instantiates type 1 rate constant for a particular reaction
416    (https://lmdz-forge.lmd.jussieu.fr/mediawiki/Planets/index.php/Photochemistry#Reaction_rate_formulae)
417   
418    Attributes
419    ----------
420    params : dict
421        Reaction-specific set of parameters for the rate constant: a,T0,c,b,d
422
423    Methods
424    -------
425    call(T,density)
426        Compute the reaction rate for given temperature and density
427    """
428
429    def __init__(self,params):
430        """
431        Parameters
432        ----------
433        params : dict
434            Reaction-specific set of parameters for the rate constant: a,T0,c,b,d
435        """
436        self.params = params
437
438    def __call__(self,T,density):
439        """
440        Parameters
441        ----------
442        T : float
443            Temperature [K]
444        density : float
445            Background gas density [cm^-3]
446
447        Returns
448        -------
449        float
450            Value of the reaction rate constant [cm^3.s^-1]
451        """
452        return self.params['a']*(T/self.params['T0'])**self.params['c']*np.exp(-self.params['b']/T)*density**self.params['d']
453
454class reaction_constant_type2(reaction_constant):
455    """ Type 2 reaction rate constant
456   
457    Instantiates type 2 rate constant for a particular reaction
458    (https://lmdz-forge.lmd.jussieu.fr/mediawiki/Planets/index.php/Photochemistry#Reaction_rate_formulae)
459   
460    Attributes
461    ----------
462    params : dict
463        Reaction-specific set of parameters for the rate constant: k0,T0,n,a0,kinf,m,b0,g,h,dup,ddown,fc
464
465    Methods
466    -------
467    call(T,density)
468        Computes the reaction rate for given temperature and density
469    """
470    def __call__(self,T,density):
471        """
472        Parameters
473        ----------
474        T : float
475            Temperature [K]
476        density : float
477            Background gas density [cm^-3]
478
479        Returns
480        -------
481        float
482            The value of the reaction rate constant [cm^3.s^-1]
483        """
484        num = self.params['k0']*(T/self.params['T0'])**self.params['n']*np.exp(-self.params['a0']/T)
485        den = self.params['kinf']*(T/self.params['T0'])**self.params['m']*np.exp(-self.params['b0']/T)
486   
487        return self.params['g']*np.exp(-self.params['h']/T)+num*density**self.params['dup']/(1+num/den*density**self.params['ddown'])*self.params['fc']**(1/(1+np.log10(num/den*density)**2))
488
489# TODO: implement type 3 reaction constant
490
491def read_reactfile(path):
492    """ Reads the reactfile formatted for simulations with the Generic PCM
493   
494    Parameters
495    ----------
496    path : str
497        Path to the reactfile (to become reaction.def)
498
499    Returns
500    -------
501    dict
502        Keys are reactions formulae, items are reactions instances
503    """
504
505    reactions = {}
506    with open(path+'/chemnetwork/reactfile') as reactfile:
507        for line in reactfile:
508            # Commented line
509            if line[0] == '!':
510                # Hard-coded reaction
511                if 'hard' in line and 'coded' in line:
512                    hard_coded_reaction = reaction(line[1:51].split(),line[51:101].split(),None)
513                    print('reaction ',hard_coded_reaction.formula,'seems to be hard-coded. Add it manually if needed.')
514                continue
515            else:
516                reactants = line[:50].split()
517                products  = line[50:100].split()
518               
519                # Photolysis
520                if 'hv' in line:
521                    reactants.pop(reactants.index('hv'))
522                    new_reaction = photolysis(reactants,products)
523                   
524                # Other reactions
525                else:
526                    cst_params = [float(val) for val in line[101:].split()]
527                   
528                    # three-body reaction
529                    if 'M' in line:
530                       
531                        # if third body is not the background gas
532                        if line[line.index('M')+2] != ' ':
533                            third_body = reactants[reactants.index('M')+1]
534                        else:
535                            third_body = 'background'
536                        reactants.pop(reactants.index('M'))
537                        try:
538                            products.pop(reactants.index('M'))
539                        except:
540                            pass
541                        if int(line[100]) == 1:
542                            rate_constant = reaction_constant({param_key:cst_params[i] for i,param_key in enumerate(['a','b','c','T0','d'])})
543                            # if the third body is not the background gas, we treat it as a two-body reaction
544                            if third_body != 'background':
545                                products.append(third_body)
546                                rate_constant.params['d'] = 0
547                        elif int(line[100]) == 2:
548                            rate_constant = reaction_constant_type2({param_key:cst_params[i] for i,param_key in enumerate(['k0','n','a0','kinf','m','b0','T0','fc','g','h','dup','ddown'])})
549                            if third_body != 'background':
550                                raise Exception('Dont know how to handle three body reaction with type 2 constant when third body is not the background gas.')
551                        else:
552                            raise Exception('rate constant parameterization type ',line[100],' not recognized.')
553                        new_reaction = termolecular_reaction(reactants,products,rate_constant)
554
555                    # two-body reaction
556                    else:
557                        rate_constant = reaction_constant({param_key:cst_params[i] for i,param_key in enumerate(['a','b','c','T0','d'])})
558                        new_reaction = reaction(reactants,products,rate_constant)
559               
560                reactions[new_reaction.formula] = new_reaction
561
562    return reactions
563               
564
565def density(P,T,VMR=1):
566    """ Computes molecular density using the perfect gas law
567
568    Parameters
569    ----------
570    P : float
571        Pressure [Pa]
572    T : float
573        Temperature [K]
574    VMR : float (optional)
575        Volume mixing ratio [cm^3/cm^3]
576
577    Returns
578    -------
579    float
580        Molecular density [cm^-3]
581    """
582   
583    m3_to_cm3 = 1e6
584    return VMR * P / R / T / m3_to_cm3 * N_A
585
586def compute_rates(s,reactions='read'):
587    """ Computes reaction rates for a simulation
588
589    Parameters
590    ----------
591    s : GPCM_simu
592        Simulation object
593    reactions : dict or 'read' (optional)
594        Dictionnary of reactions whose rate to compute as returned by read_reactfile
595        If nothing passed, call read_reactfile to identify reactions
596
597    Returns
598    -------
599    GPCM_simu
600        Simulation object with reactions rates, rates constants, species vmr and densities as new fields
601    """
602
603    if reactions == 'read':
604        reactions      = read_reactfile(s.path)
605        s.species   = []
606        s.reactions = {} # reactions dict will be merged at the end
607
608    # Register new species
609    for r in reactions:
610        for sp in reactions[r].reactants:
611            if not sp in s.species:
612                s.species.append(sp)
613        for sp in reactions[r].products:
614            if not sp in s.species:
615                s.species.append(sp)
616       
617    densities = {}
618
619    # Background density
620    s['total density']     = density(s['p'],s['temp']).assign_attrs({'units':'cm^-3.s^-1'})
621
622    for sp in s.species:
623        # volume mixing ratios
624        s[sp+' vmr'] = s[sp] * M['co2'] / M[sp]
625        s[sp+' vmr'] = s[sp+' vmr'].assign_attrs({'units':'m^3/m^3'})
626        # molecular densities
627        s[sp+' density'] = density(s['p'],s['temp'],VMR=s[sp+' vmr'])
628        s[sp+' density'] = s[sp+' density'].assign_attrs({'units':'cm^-3'})
629        densities[sp] = s[sp+' density']
630
631    for r in reactions:
632
633        # Photolysis
634        if type(reactions[r]) == photolysis:
635
636            # Cases with branching ratios
637            if reactions[r].reactants[0] == 'co2':
638                if 'o1d' in reactions[r].products:
639                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jco2_o1d'],densities)
640                else:
641                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jco2_o'],densities)
642            elif reactions[r].reactants[0] == 'o2':
643                if 'o1d' in reactions[r].products:
644                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo2_o1d'],densities)
645                else:
646                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo2_o'],densities)
647            elif reactions[r].reactants[0] == 'o3':
648                if 'o1d' in reactions[r].products:
649                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo3_o1d'],densities)
650                else:
651                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo3_o'],densities)
652            elif reactions[r].reactants[0] == 'ch2o':
653                if 'cho' in reactions[r].products:
654                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jch2o_cho'],densities)
655                else:
656                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jch2o_co'],densities)
657            elif reactions[r].reactants[0] == 'h2o_vap':
658                s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jh2o'],densities)
659            else:
660                # General case
661                s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['j'+reactions[r].reactants[0]],densities)
662        else:
663            s['k ('+reactions[r].formula+')'] = reactions[r].constant(s['temp'],densities[background])
664            s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['temp'],densities)
665
666            # Termolecular reaction
667            if type(reactions[r]) == termolecular_reaction:
668                s['k ('+reactions[r].formula+')'] = s['k ('+reactions[r].formula+')'].assign_attrs({'units':'cm^6.s^-1'})
669
670            # Bimolecular reaction
671            else:
672                s['k ('+reactions[r].formula+')'] = s['k ('+reactions[r].formula+')'].assign_attrs({'units':'cm^3.s^-1'})
673               
674        s['rate ('+reactions[r].formula+')'] = s['rate ('+reactions[r].formula+')'].assign_attrs({'units':'cm^-3.s^-1'})
675               
676    s.reactions = s.reactions | reactions
677   
678    return s
679
Note: See TracBrowser for help on using the repository browser.