first changes to add state_type

webgeodyn/inout/pygeodyn_hdf5.py

@@ -6,19 +6,21 @@ import numpy as np
 from .default import giveMeasureTypeUnits
 
 
-def load(dataDirectory, dataModel, keepRealisations):
+def load(dataDirectory, dataModel, keepRealisations, state_type='analysed'):
     """ Loading function for pygeodyn files of hdf5 format. Also adds the data to the dataModel.
 
     :param dataDirectory: Location of the pygeodyn files
     :type dataDirectory: os.path
     :param dataModel: Model in which to add the loaded measures
     :type dataModel: Model
-    :param keepRealisations: If True, all realisationsr are kept in the data. Else, the data is averaged over the realisations
+    :param keepRealisations: If True, all realisations are kept in the data. Else, the data is averaged over the realisations
     :type keepRealisations: bool
     :return: 0 if everything went well, -1 otherwise
     :rtype: int
+    :param state_type: Either forecast, computed or analysed depending on the type of states needed
     """
     firstpoint = 3
+    assert state_type in ('computed', 'analysed', 'forecast')
 
     measures_to_load = ['MF', 'SV', 'ER', 'U']
 
@@ -30,7 +32,7 @@ def load(dataDirectory, dataModel, keepRealisations):
     print('Reading:', hdf_filename)
 
     with h5py.File(hdf_filename) as hdf_file:
-        computed_data = hdf_file['computed']
+        computed_data = hdf_file[state_type]
 
         times = np.array(computed_data['times'])[firstpoint:]
 

webgeodyn/models/model.py

@@ -18,7 +18,7 @@ class Model:
     # CONSTRUCTING METHODS
 
     def __init__(self, dataDirectory=None, dataFormat='default', nth=90, nph=180, normalisation=semiNormalisedSchmidt,
-                 keepRealisations=True):
+                 keepRealisations=True, state_type='analysed'):
         """
         :param dataDirectory: directory where the data is located
         :type dataDirectory: str (path)
@@ -56,7 +56,7 @@ class Model:
         self.notSH_measures = {}
 
         if dataDirectory is not None:
-            self.loadDataFromFile(keepRealisations)
+            self.loadDataFromFile(keepRealisations, state_type=state_type)
 
     def copy(self):
         """
@@ -238,7 +238,7 @@ class Model:
                     )
         return
 
-    def loadDataFromFile(self, keepRealisations):
+    def loadDataFromFile(self, keepRealisations, state_type='analysed'):
         """
         Loads the data from files according to the dataDirectory and dataFormat of the model. Uses dedicated load functions of the inout module.
 
@@ -258,7 +258,7 @@ class Model:
         if self.dataFormat is not None:
             print("Reading model using predefined format", self.dataFormat)
             load_fct = importlib.import_module("webgeodyn.inout." + self.dataFormat).load
-            load_fct(self.dataDirectory, self, keepRealisations)
+            load_fct(self.dataDirectory, self, keepRealisations, state_type=state_type)
 
         # Else raise an error
         else: