update comments

da/stilt/observationoperator.py

@@ -79,7 +79,8 @@ class STILTObservationOperator(ObservationOperator):
         self.startdate=None
 
     def setup(self, dacycle):
-        """ Execute all steps needed to prepare the ObsOperator for use inside CTDAS, only done at the very first cycle normally """
+        """ Execute all steps needed to prepare the ObsOperator for use inside CTDAS, only done at the very first cycle normally
+        Input: dacycle: dict with information on the data assimilatin cycle """
         self.dacycle = dacycle
         self.outputdir = dacycle['dir.output']
         self.rc = dacycle['da.obsoperator.rc']
@@ -123,7 +124,9 @@ class STILTObservationOperator(ObservationOperator):
         """ 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)"""
+        - Update the rc file (stilt_X.rc)
+        Input: 
+            proc: int: number of process"""
         import os
 
         # Define the name of the file that will contain the modeled output of each observation
@@ -133,7 +136,10 @@ class STILTObservationOperator(ObservationOperator):
 
     def update_rc(self,name,proc):
         """Function that updates the .rc files.
-        The outputdir and time of the new rc file are adjusted"""
+        The outputdir and time of the new rc file are adjusted
+        Input: 
+            name: str: rc filename
+            proc: number of process"""
         if 'stilt.rc' in name:
             path,dummy= name.split('stilt.rc')
             shutil.copyfile(name,path+'stilt_%s.rc'%proc)
@@ -152,16 +158,18 @@ class STILTObservationOperator(ObservationOperator):
     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)
+        Input:
+            time: datetime.datetime: The time for which the index needs to be found. Default: current time cycle datetime
         """
         if time == None:
             # Get the time of the current cycle
-            cycledate = self.dacycle['time.start']
+            time = 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 = time - startdate
         timediff_hours = int(timediff.total_seconds()/3600)
         time_index = int(timediff_hours)
         return time_index
@@ -177,7 +185,9 @@ class STILTObservationOperator(ObservationOperator):
     
     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"""
+        This can be used for e.g. the background concentrations
+        Input: 
+            ncfile: netCDF4.Dataset with latitude and longitude dimensions"""
         # Initialise some names that could indicate latitude. Expand at will
         lat_opts = ['lat', 'latitude', 'Lat', 'LAT', 'LATITUDE', 'Latitude']
         # Same for longitude
@@ -200,7 +210,9 @@ class STILTObservationOperator(ObservationOperator):
     
     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)"""
+        Returns a 3D np.array with dims (time, lat, lon)
+        Input:
+            site: str of 3 letter sitename. Is converted to uppercase. Example: 'hei'"""
         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)
@@ -213,6 +225,10 @@ class STILTObservationOperator(ObservationOperator):
         return np.flip(np.flipud(footprint), axis=1)
     
     def get_center_of_mass(fp):
+        """Function that finds the center of mass of the first footprint and the time corresponding to it.
+        Input: 
+            fp: 3d np.array of the footprint
+        time is taken from the observationoperator self."""
         from scipy import ndimage
         i = -1
         total_influence = 0
@@ -226,7 +242,8 @@ class STILTObservationOperator(ObservationOperator):
         return center_of_mass, i
     
     def get_background(self, tracer_info):
-        
+        """Function that finds the background concentration, based on the tracer info
+        Input: tracer info: Dict with infromation on the background concentration"""
 #         data = nc.Dataset(file, 'r+')
 #         # Set the data in 'days since 2016-1-1 00:00:00', similar to the footprints.
 #         # Note that this should, preferably, not be hardcoded.
@@ -245,7 +262,13 @@ 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"""
+        """Function that calculates the atmospheric increase due to the exchange of fluxes over the footprint
+        Input: 
+            site: str of 3 chars with the location. Example: 'hei'
+            tracer_info: dict with information on the tracer:
+                path to fluxes,
+                half-life,
+                recalculation factor"""
         # First, get the footprint
         foot = self.get_foot(site)
         # Get the tracer fluxes
@@ -273,12 +296,21 @@ class STILTObservationOperator(ObservationOperator):
     
     def get_concentration(self, site,  tracer_info):
         """Function that calculates the simulated concentration for a given site and tracer  as
-        [flux * foot] + background"""
+        [flux * foot] + background
+        Input:
+            site: str of 3 chars with the location. Example: 'hei'
+            tracer_info: dict with information on the tracer:
+                path to fluxes,
+                half-life,
+                recalculation factor"""
         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"""
+        First prepares run, then runs
+        Input: 
+            proc: int of the current process
+            out_q: instance of multiprocessing.Queue()"""
         self.parse_samplefile()
         self.prepare_run(proc)
         self.run(proc)
@@ -290,12 +322,15 @@ class STILTObservationOperator(ObservationOperator):
     def run(self,proc):
         """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):
-            logging.info('making concentration for tracer {}, site {} and obs_id {}'.format(tracer, site, id))
-            conc = self.get_concentration(site, self.tracer_info[tracer])
-            concentrations.append(conc)
+        Input:
+            proc: int of process"""
+        for mem in range(self.dacycle['da.optimizer.nmembers']):
+            concentrations_member_i = []
+            for tracer, site, id in zip(self.tracer, self.site, self.obs_id):
+                logging.info('making concentration for tracer {}, site {} and obs_id {}'.format(tracer, site, id))
+                conc = self.get_concentration(site, self.tracer_info[tracer])
+                concentrations_member_i.append(conc)
+            concentrations.append(concentrations_member_i)
         self.concentrations = concentrations
         self.save_data()
 
@@ -320,13 +355,14 @@ class STILTObservationOperator(ObservationOperator):
          f.add_data(savedict)
 
          # Add the simulated concentrations
+         dimmember = f.createDimension('nmembers', size=self.dacycle['da.optimizer.nmembers'])
          savedict = io.std_savedict.copy()
          savedict['name'] = "simulated"
          savedict['dtype'] = "float"
          savedict['long_name'] = "Simulated_concentration"
          savedict['units'] = "mol mol-1"
          savedict['values'] = self.concentrations
-         savedict['dims'] = dimid
+         savedict['dims'] = dimid + dimmember
          savedict['comment'] = "Simulated value by STILTObservationOperator"
          f.add_data(savedict)
 

templates/template.py

@@ -29,7 +29,7 @@ from da.tools.initexit import start_logger, validate_opts_args, parse_options, C
 from da.stilt.pipeline import forward_pipeline, header, footer
 from da.platform.cartesius import CartesiusPlatform
 from da.baseclasses.dasystem import DaSystem
-from da.co2gridded.statevector import CO2GriddedStateVector 
+from da.stilt.statevector import CO2GriddedStateVector 
 from da.baseclasses.obs import Observations
 from da.baseclasses.optimizer import Optimizer
 from da.stilt.observationoperator import STILTObservationOperator