Index: trunk/LMDZ.GENERIC/utilities/photochemistry/photochem_postproc.py
===================================================================
--- trunk/LMDZ.GENERIC/utilities/photochemistry/photochem_postproc.py	(revision 3431)
+++ trunk/LMDZ.GENERIC/utilities/photochemistry/photochem_postproc.py	(revision 3511)
@@ -15,19 +15,4 @@
 warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 'lat'")
 
-M          = {'co2':44,  # Molar masses
-              'o':16,    # TODO: automatic parser
-              'o1d':16,
-              'o2':32,
-              'o3':48,
-              'h':1,
-              'h2':2,
-              'oh':17,
-              'h2o_vap':18,
-              'ho2':33,
-              'h2o2':34,
-              'co':28,
-              'cho':29,
-              'ch2o':30}
-
 background = 'co2' # background gas of the atmosphere
 
@@ -46,10 +31,14 @@
     Methods
     -------
-    get_profile(field,**kw)
-        Get profile of a field (either local or averaged)
+    get_subset(field,**kw)
+        Get pa subset at fixed given coordinate of the data
     plot_meridional_slice(field,logcb,labelcb,**kw)
         Plot a meridional slice of a field
     plot_time_evolution(field,logcb,labelcb,**kw)
         Plot time evolution of a field (as a time vs altitude contour plot)
+    plot_time_series(field,lat,lon,alt,logy)
+        Plot time series of a field (at a specific location (lon, lat, alt), averged, or a combination thereof)
+    plot_atlas(field,t,alt,logcb,labelcb)
+        Plot atlas of a field
     plot_profile(field,**kw)
         Plot a profile of a field
@@ -80,39 +69,12 @@
         self.data[field] = value
 
-    def get_profile(self,field,**kw):
-        """ Get profile of a field (either local or averaged)
-
-        Parameters
-        ----------
-        field : str
-            Field name
-        t : float (optional)
-            Time at which to select (if nothing specified use time-average)
-        lat : float (optional)
-            Latitude at which to select (if nothing specified use area-weighted meridional average)
-        lon : float (optional)
-            Longitude at which to select (if nothing specified use zonal average)
-
-        """
-        
-        if self['latitude'].size == 1:
-            # 1D
-            return self[field][:,0,0]
-        else:
-            # 3D
-            if 'lat' in kw:
-                if 'lon' in kw:
-                    return self[field].sel(latitude=kw['lat'],method='nearest').sel(longitude=kw['lon'],method='nearest')
-                else:
-                    return self[field].sel(latitude=kw['lat'],method='nearest').mean(dim='longitude')
-            else:
-                # Latitude-averaged profile: need to weight by the grid cell surface area
-                if 'lon' in kw:
-                    return (self['aire']*self[field]/self['aire'].mean(dim='latitude')).mean(dim='latitude').sel(longitude=kw['lon'],method='nearest')
-                else:
-                    return (self['aire']*self[field]/self['aire'].mean(dim='latitude').mean(dim='longitude')).mean(dim='latitude').mean(dim='longitude')
+    def __area_weight__(self,data_array):
+        return self['aire']*data_array/self['aire'].mean('latitude')
 
     def get_subset(self,field='all',**kw):
         """ Get a subset at fixed given coordinate of the data
+
+        Can also average over a given dimension. In this case,
+        the meridional average is area-weighted.
 
         Parameters
@@ -133,12 +95,5 @@
             Altitude of the slice. If nothing
             specified, use time-average
-
-        Raise
-        -----
-        Slice direction not provided
-        """
-        if len(kw) == 0:
-            raise Exception('Slice direction not provided')
-
+        """
         if field == 'all':
             data_subset = self.data
@@ -147,11 +102,23 @@
             
         if 't' in kw:
-            data_subset = data_subset.sel(Time=kw['t'],method='nearest')
+            if kw['t'] == 'avg':
+                data_subset = data_subset.mean(dim='Time')
+            else:
+                data_subset = data_subset.sel(Time=kw['t'],method='nearest')
+        if 'lat' in kw:
+            if kw['lat'] == 'avg':
+                data_subset = self.__area_weight__(data_subset).mean(dim='latitude')
+            else:
+                data_subset = data_subset.sel(latitude=kw['lat'],method='nearest')
         if 'lon' in kw:
-            data_subset = data_subset.sel(longitude=kw['lon'],method='nearest')
-        if 'lat' in kw:
-            data_subset = data_subset.sel(latitude=kw['lat'],method='nearest')
+            if kw['lon'] == 'avg':
+                data_subset = data_subset.mean(dim='longitude')
+            else:
+                data_subset = data_subset.sel(longitude=kw['lon'],method='nearest')
         if 'alt' in kw:
-            data_subset = data_subset.sel(altitude=kw['alt'],method='nearest')
+            if kw['alt'] == 'avg':
+                data_subset = data_subset.mean(dim='altitude')
+            else:
+                data_subset = data_subset.sel(altitude=kw['alt'],method='nearest')
 
         return data_subset
@@ -178,19 +145,10 @@
             raise Exception('Trying to plot a meridional slice of a 1D simulation')
 
-        meridional_slice = self[field]
-        
-        if t == 'avg':
-            meridional_slice = meridional_slice.mean(dim='Time')
+        meridional_slice = self.get_subset(field,t=t,lon=lon)
+            
+        if logcb:
+            plt.contourf(meridional_slice['latitude'],meridional_slice['altitude'],meridional_slice,locator=tk.LogLocator(),**plt_kw)
         else:
-            meridional_slice = meridional_slice.sel(Time=t,method='nearest')
-        if lon == 'avg':
-            meridional_slice = meridional_slice.mean(dim='longitude')
-        else:
-           meridional_slice = meridional_slice.sel(longitude=lon,method='nearest') 
-            
-        if logcb:
-            plt.contourf(self['latitude'],self['altitude'],meridional_slice,locator=tk.LogLocator(),**plt_kw)
-        else:
-            plt.contourf(self['latitude'],self['altitude'],meridional_slice,**plt_kw)
+            plt.contourf(meridional_slice['latitude'],meridional_slice['altitude'],meridional_slice,**plt_kw)
             
         plt.colorbar(label=field if labelcb==None else labelcb)
@@ -216,23 +174,42 @@
         """
 
-        time_evolution = self[field]
-        
-        if lat == 'avg':
-            time_evolution = (self['aire']*time_evolution/self['aire'].mean(dim='latitude')).mean(dim='latitude')
+        time_evolution = self.get_subset(field,lon=lon,lat=lat)
+            
+        if logcb:
+            plt.contourf(time_evolution['Time'],time_evolution['altitude'],time_evolution.T,locator=tk.LogLocator(),**plt_kw)
         else:
-            time_evolution = time_evolution.sel(latitude=lat,method='nearest')
-        if lon == 'avg':
-            time_evolution = time_evolution.mean(dim='longitude')
-        else:
-            time_evolution = time_evolution.sel(longitude=lon,method='nearest')
-            
-        if logcb:
-            plt.contourf(self['Time'],self['altitude'],time_evolution.T,locator=tk.LogLocator(),**plt_kw)
-        else:
-            plt.contourf(self['Time'],self['altitude'],time_evolution.T,**plt_kw)
+            plt.contourf(time_evolution['Time'],time_evolution['altitude'],time_evolution.T,**plt_kw)
             
         plt.colorbar(label=field if labelcb==None else labelcb)
         plt.xlabel('time [day]')
         plt.ylabel('altitude [km]')
+
+    def plot_time_series(self,field,lat='avg',lon='avg',alt='avg',logy=False,**plt_kw):
+        """ Plot time series of a field (at a specific location (lon, lat, alt), averged, or a combination thereof)
+
+        Parameters
+        ----------
+        field : str
+            Field name to plot
+        lat : float (optional)
+            Latitude at which to plot (if nothing specified use area-weighted meridional average)
+        lon : float (optional)
+            Longitude at which to plot (if nothing specified use zonal average)
+        logy : bool (optional)
+            Use logarithmic y-axis
+        matplotlib plot keyword arguments
+        """
+
+        time_series = self.get_subset(field,lon=lon,lat=lat,alt=alt)
+
+        if not 'label' in plt_kw:
+            plt_kw['label'] = self[field].units
+        if logy:
+            plt.semilogy(time_series['Time'],time_series,**plt_kw)
+        else:
+            plt.plot(time_series['Time'],time_series,**plt_kw)
+            
+        plt.xlabel('time [day]')
+        plt.ylabel(field+' ['+self[field].units+']')
 
     def plot_atlas(self,field,t='avg',alt='avg',logcb=False,labelcb=None,**plt_kw):
@@ -254,21 +231,13 @@
         """
 
-        atlas = self[field]
-        
-        if t == 'avg':
-            atlas = atlas.mean(dim='Time')
+        if 'altitude' in self[field].dims:
+            atlas = self.get_subset(field,t=t,alt=alt)
         else:
-            atlas = atlas.sel(Time=t,method='nearest')
-            
-        if  'altitude' in atlas.coords:
-            if alt == 'avg':
-                atlas = atlas.mean(dim='altitude')
-            else:
-                atlas = atlas.sel(altitude=alt,method='nearest')
+            atlas = self.get_subset(field,t=t)
             
         if logcb:
-            plt.contourf(self['longitude'],self['latitude'],atlas,locator=tk.LogLocator(),**plt_kw)
+            plt.contourf(atlas['longitude'],atlas['latitude'],atlas,locator=tk.LogLocator(),**plt_kw)
         else:
-            plt.contourf(self['longitude'],self['latitude'],atlas,**plt_kw)
+            plt.contourf(atlas['longitude'],atlas['latitude'],atlas,**plt_kw)
             
         plt.colorbar(label=field if labelcb==None else labelcb)
@@ -294,29 +263,68 @@
         """
 
-        profile = self[field]
+        profile = self.get_subset(field,t=t,lon=lon,lat=lat)
         
-        if t == 'avg':
-            profile = profile.mean(dim='Time')
+        if logx:
+            plt.semilogx(profile,profile['altitude'],**plt_kw)
         else:
-            profile = profile.sel(Time=t,method='nearest')
-            
-        if self['latitude'].size > 1:
-            
-            if lat == 'avg':
-                profile = (self['aire']*profile/self['aire'].mean(dim='latitude')).mean(dim='latitude')
-            else:
-                profile = profile.sel(latitude=lat,method='nearest')
-            if lon == 'avg':
-                profile = profile.mean(dim='longitude')
-            else:
-                profile = profile.sel(longitude=lon,method='nearest')
-            
-        if logx:
-            plt.semilogx(profile,self['altitude'],**plt_kw)
-        else:
-            plt.plot(profile,self['altitude'],**plt_kw)
+            plt.plot(profile,profile['altitude'],**plt_kw)
             
         plt.xlabel(field+' ['+self[field].units+']')
         plt.ylabel('altitude [km]')
+
+    def to_chempath(self,t,dt,filename_suffix='_chempath',lon='avg',lat='avg',alt='avg'):
+        """ Create files redable by the chemical path analyzer chempath (DOI: 10.5194/gmd-2024-163)
+
+        Parameters
+        ----------
+        t : float
+            Initial time
+        dt : float
+            Timestep
+        filename_suffix : str (optional)
+            Suffix to append to the files being created (species.txt,
+            reactions.txt, model_time.dat, rates.dat, concentrations.dat)
+        lat : float (optional)
+            Latitude at which to plot (if nothing specified use area-weighted meridional average)
+        lon : float (optional)
+            Longitude at which to plot (if nothing specified use zonal average)
+        alt : float (optional)
+            Altitude of the slice. If nothing
+        matplotlib's plot / semilogx keyword arguments
+        """
+
+        # Save pecies list in chempath format
+        with open(self.path+'/species'+filename_suffix+'.txt', 'w') as fp:
+            for sp in self.species:
+                fp.write("%s\n" % sp)
+
+        # Save reactions list in chempath format
+        with open(self.path+'/reactions'+filename_suffix+'.txt', 'w') as fp:
+            for r in self.reactions:
+                current_formula = self.reactions[r].formula
+                current_formula = current_formula.replace(' ','')
+                current_formula = current_formula.replace('->','=')
+                fp.write("%s\n" % current_formula)
+
+        # Save rates, concentrations and times in chempath format
+        rates   = np.array([],dtype=np.float128)
+        conc    = np.array([],dtype=np.float128)
+        conc_dt = np.array([],dtype=np.float128)
+        times   = np.array([t,t+dt],dtype=np.float128)
+        
+        out    = self.get_subset(t=t,lon=lon,lat=lat,alt=alt)  
+        for r in self.reactions:
+            rates = np.append(rates,out[f'rate ({r})'])
+        for sp in self.species:
+            conc = np.append(conc,out[f'{sp} vmr'])
+        
+        out    = self.get_subset(t=t+dt,lon=lon,lat=lat,alt=alt)  
+        for sp in self.species:
+            conc_dt = np.append(conc_dt,out[f'{sp} vmr'])
+        conc = np.vstack((conc,conc_dt))
+
+        times.tofile(self.path+'/model_time'+filename_suffix+'.dat')
+        rates.tofile(self.path+'/rates'+filename_suffix+'.dat')
+        conc.tofile(self.path+'/concentrations'+filename_suffix+'.dat')
             
 class reaction:
@@ -603,4 +611,5 @@
     if reactions == 'read':
         reactions      = read_reactfile(s.path)
+        s.tracers   = read_traceurs(s.path)
         s.species   = []
         s.reactions = {} # reactions dict will be merged at the end
@@ -609,7 +618,11 @@
     for r in reactions:
         for sp in reactions[r].reactants:
+            if not sp in s.tracers or not s.tracers[sp].is_chim:
+                raise Warning(sp, 'is not recorded as a chimical tracer')
             if not sp in s.species:
                 s.species.append(sp)
         for sp in reactions[r].products:
+            if not sp in s.tracers or not s.tracers[sp].is_chim:
+                raise Warning(sp, 'is not recorded as a chimical tracer')
             if not sp in s.species:
                 s.species.append(sp)
@@ -622,5 +635,5 @@
     for sp in s.species:
         # volume mixing ratios
-        s[sp+' vmr'] = s[sp] * M['co2'] / M[sp]
+        s[sp+' vmr'] = s[sp] * s.tracers['co2'].M / s.tracers[sp].M
         s[sp+' vmr'] = s[sp+' vmr'].assign_attrs({'units':'m^3/m^3'})
         # molecular densities
@@ -678,2 +691,76 @@
     return s
 
+class trac:
+    """ Tracer class
+    
+    A class to store useful informations about simulations tracers
+    
+    Attributes
+    ----------
+    name : string
+        Tracer's name
+    M : float
+        Tracer's molar mass [g/mol]
+    is_chim : bool
+        Is it a photochemical tracer?
+    """
+    def __init__(self,name,M,is_chim):
+        self.name    = name
+        self.M       = M
+        self.is_chim = is_chim
+
+def read_traceurs(path):
+    """ Read the traceurs of a simulation
+
+    Parameters
+    ----------
+    path : string
+        path to simulation
+
+    Returns
+    -------
+    dict
+        dictionnaries of all tracers
+    """
+    tracdict = {}
+    with open(path+'/traceur.def') as tracfile:
+        
+        for iline,line in enumerate(tracfile):
+            
+            # Empty line
+            if len(line.split()) == 0:
+                continue
+                
+            # First line
+            elif iline == 0:
+                if not '#ModernTrac-v1' in line:
+                    raise Exception('Can only read modern traceur.def')
+                continue
+                
+            # Second line (number of tracers)
+            elif iline == 1:
+                ntrac = int(line)
+                continue
+            
+            # Commented line
+            elif line[0] == '!':
+                continue
+
+            # Regular entry
+            else:
+                line    = line.split()
+                name    = line[0]
+                is_chim = False
+                
+                for param in line[1:]:
+                    if param[:4] == 'mmol':
+                        M = float(param[5:])
+                    elif param[:7] == 'is_chim':
+                        is_chim = bool(float(param[8:]))
+                        
+                tracdict[name] = trac(name,M,is_chim)
+                
+    if len(tracdict) != ntrac:
+        raise Exception('Mismatch between announced an read number of tracers')
+        
+    return tracdict
