update to the comments in the observationoperator

da/stilt/observationoperator.py

@@ -83,8 +83,9 @@ class STILTObservationOperator(ObservationOperator):
         self.dacycle = dacycle
         self.outputdir = dacycle['dir.output']
         self.rc = dacycle['da.obsoperator.rc']
- 
+        # Read the stilt.rc file for settings to the observationoperator 
         rc_options = rc.read(self.rc)
+        # Also read in the tracer information
         rc_options['tracers'] = rc_options['tracers'].split(',')
         self.rc_options = rc_options
         self.tracer_info = {tracer: None for tracer in rc_options['tracers']}
@@ -92,9 +93,11 @@ class STILTObservationOperator(ObservationOperator):
             self.tracer_info[tracer] = rc.read(self.rc_options['tracerdir']+'/'+tracer)
         
     def parse_samplefile(self):
+        """Function that reads the sample files created in obspack.py"""
         infile = nc.Dataset(self.dacycle['ObsOperator.inputfile'])
         with warnings.catch_warnings():
-            warnings.simplefilter("ignore") 
+            warnings.simplefilter("ignore")
+            # Add the observation ids, tracers and sites
             obs_ids = infile['obs_id'][:]
             obs_ids = [b''.join(obs_id).decode() for obs_id in obs_ids]
             self.tracer, self.site, self.obs_id = [], [], []
@@ -103,10 +106,12 @@ class STILTObservationOperator(ObservationOperator):
                 self.tracer.append(tracer)
                 self.site.append(site)
                 self.obs_id.append(obs_id.split('~')[-1])
+            # Times and latitude and longitudes are static; only need once
             self.times = infile['date_components'][:]
             
             self.lat = infile['latitude'][:]
             self.lon = infile['longitude'][:]
+        # Set the start time
         starttime = self.dacycle['time.sample.start']
         self.starttime = starttime
         self.year = starttime.year
@@ -115,7 +120,10 @@ class STILTObservationOperator(ObservationOperator):
         self.hour = starttime.hour
     
     def prepare_run(self,proc):
-        """ Prepare the running of the actual forecast model, for example compile code """
+        """ Prepare the running of the actual forecast model.
+        - Initialise the file that will hold the simulated values
+        - Initialise the number of ensemble members
+        - Update the rc file (stilt_X.rc)"""
         import os
 
         # Define the name of the file that will contain the modeled output of each observation
@@ -124,14 +132,16 @@ class STILTObservationOperator(ObservationOperator):
         self.update_rc(self.rc,proc)
 
     def update_rc(self,name,proc):
+        """Function that updates the .rc files.
+        The outputdir and time of the new rc file are adjusted"""
         if 'stilt.rc' in name:
             path,dummy= name.split('stilt.rc')
             shutil.copyfile(name,path+'stilt_%s.rc'%proc)
             name=os.path.join(path,'stilt_%s.rc'%proc)
         self.rc_filename = name
 
-        logging.info('RC name %s'%name)
         new_rc = rc.read(name)
+        # update the time and the outputdir
         new_rc['cycstadate'] = self.starttime
         new_rc['outdir'] = self.rc_options['outdir']
         rc.write(name, new_rc)
@@ -139,26 +149,41 @@ class STILTObservationOperator(ObservationOperator):
         logging.debug('STILT rc-file updated successfully')
 
     
-    def get_time_index_nc(self):
-        cycledate = self.dacycle['time.start']
+    def get_time_index_nc(self, time=None):
+        """Function that gets the time index from the flux files
+        based on the cycletime and the first time in all the files (hardcoded in stilt.rc)
+        """
+        if time == None:
+            # Get the time of the current cycle
+            cycledate = self.dacycle['time.start']
+        # Get the start date of all cycles
         startdate = self.rc_options['files_startdate']
         startdate = datetime.datetime.strptime(startdate, '%Y-%m-%d %H:%M:%S')
+        # Get the difference between the current and the start
+        # Note that this is in hours, and thus assumes that the flux files are hourly as well
         timediff = cycledate - startdate
         timediff_hours = int(timediff.total_seconds()/3600)
         time_index = int(timediff_hours)
         return time_index
     
     def get_time_indices(self):
+        """Function that gets the time indices in the flux files
+        Because if the footprint is for 24 hours back, we need the fluxes 24 hours back"""
+
         time_index = self.get_time_index_nc()
         
         return slice(time_index-self.numtimes, time_index)
 
     
     def get_latlon_index_nc(self, ncfile):
-        
+        """Function that gets the indices of a lat/lon point in a .nc file.
+        This can be used for e.g. the background concentrations"""
+        # Initialise some names that could indicate latitude. Expand at will
         lat_opts = ['lat', 'latitude', 'Lat', 'LAT', 'LATITUDE', 'Latitude']
+        # Same for longitude
         lon_opts = ['lon', 'longigude', 'Lon', 'LON', 'LONGITUDE', 'Longitude']
         
+        # Get the index via the minimum of the difference.
         for latopt in lat_opts:
             try:
                 lat_ind = np.argmin(abs(ncfile[latopt][:] - self.lat))
@@ -174,6 +199,8 @@ class STILTObservationOperator(ObservationOperator):
         return lat_ind, lon_ind
     
     def get_foot(self, site):
+        """Function that gets the footprint for the current time and site.
+        Returns a 3D np.array with dims (time, lat, lon)"""
         path = self.rc_options['footdir'] + '/' + site.upper() + '24'
         fname = path + '/footprint_{0}_{1}x{2:02d}x{3:02d}x{4:02d}*.nc'.format(site.upper(),
                                                               self.year, self.month, self.day, self.hour)
@@ -185,21 +212,18 @@ class STILTObservationOperator(ObservationOperator):
         self.numtimes = len(footprint)
         return np.flip(np.flipud(footprint), axis=1)
     
-    def get_center_of_mass(self):
+    def get_center_of_mass(fp):
         from scipy import ndimage
+        i = -1
+        total_influence = 0
+        while total_influence < 0.0000001:
+            i+=1
+            total_influence = fp[i].sum()
         
-        foot = self.get_foot(site)
-        center_of_mass = []
-        
-        for footi in foot:
-            # Get the center of mass (tuple of floats)
-            center_of_mass_i = ndimage.measurements.center_of_mass(footi)
-            # Round the center of mass to nearest int to make it usable as index
-            center_of_mass_i = np.array(np.rint(center_of_mass_i), dtype=int)
-            # Append the rounded center of mass to a list
-            center_of_mass.append(center_of_mass_i)
-        
-        return(center_of_mass)
+        center_of_mass = ndimage.measurements.center_of_mass(fp[i])
+        center_of_mass = np.array(np.rint(center_of_mass), dtype=int)
+        # Note that we could also calculate the index in the file with this i and the 'get_time_indices' func. 
+        return center_of_mass, i
     
     def get_background(self, tracer_info):
         
@@ -221,14 +245,20 @@ class STILTObservationOperator(ObservationOperator):
         return 0
 
     def get_atm_increase(self, site, tracer_info):
-
+        """Function that calculates the atmospheric increase due to the exchange of fluxes over the footprint"""
+        # First, get the footprint
         foot = self.get_foot(site)
+        # Get the tracer fluxes
         file = tracer_info['path']
         data = nc.Dataset(file, 'r')
-        indices = self.get_time_indices()
         fluxmap = data[tracer_info['fluxname']]
+        # Get the time indices for the tracer fluxes
+        # It is assumed that the files all start at the same time
+        # And therefore this is not tracer-dependent
+        indices = self.get_time_indices()
         fluxes = fluxmap[indices]
         
+        # Multiply the flux with the footprint. Take the half-life into account
         half_life = eval(tracer_info['half_life'])
         if half_life:
             atm_increase = 0 # Initialise the atmospheric increase
@@ -242,12 +272,15 @@ class STILTObservationOperator(ObservationOperator):
         return atm_increase
     
     def get_concentration(self, site,  tracer_info):
+        """Function that calculates the simulated concentration for a given site and tracer  as
+        [flux * foot] + background"""
         return self.get_atm_increase(site, tracer_info) + self.get_background(tracer_info)
 
     def run_forecast_model(self,proc,out_q):
+        """Function that runs the forecast model parallel
+        First prepares run, then runs"""
         self.parse_samplefile()
         self.prepare_run(proc)
-#        self.parse_rc(self.rc_filename)
         self.run(proc)
         outdict = {}
         outdict[proc] = self.simulated_file
@@ -255,8 +288,8 @@ class STILTObservationOperator(ObservationOperator):
 
 
     def run(self,proc):
-        """
-        Call stilt R executable : //give input file lists
+        """Function that calculates the concentration for each site and tracer.
+        Calculates the concentration based on 'get_concentration' and then saves the data
         """
         concentrations = []
         for tracer, site, id in zip(self.tracer, self.site, self.obs_id):
@@ -274,8 +307,8 @@ class STILTObservationOperator(ObservationOperator):
          # Add a dimension
          dimid = f.add_dim('obs_num', len(self.concentrations))
          dim_site = f.add_dim('charstring', 10)
-#         dimid = ('obs_num',)
-         # Add the observation numbers
+         
+        # Add the observation numbers
          savedict = io.std_savedict.copy()
          savedict['name'] = "obs_num"
          savedict['dtype'] = "int"
@@ -319,9 +352,6 @@ class STILTObservationOperator(ObservationOperator):
          savedict['comment'] = "Tracer"
          f.add_data(savedict)
 
-################### End Class STILT ###################
-
-
 if __name__ == "__main__":
     pass