smashbox.init.smashbox

  1from smashbox.init import param
  2from smashbox.model import model
  3from smashbox.stats import mystats
  4from smashbox.plot import myplot
  5from smashbox.init import multimodel_statistics
  6import smashbox
  7
  8import os
  9import shutil
 10import copy
 11import pyhdf5_handler
 12import smash
 13import glob
 14import pickle
 15
 16
 17class SmashBox:
 18    """
 19    Main class which store main parameters 'myparam' and functions to create or delete
 20     models.
 21    """
 22
 23    def __init__(self):
 24
 25        if not os.path.exists(
 26            os.path.join(os.path.expanduser("~"), ".smashbox", "asset")
 27        ):
 28            shutil.copytree(
 29                os.path.join(smashbox.__path__[0], "asset"),
 30                os.path.join(os.path.expanduser("~"), ".smashbox", "asset"),
 31            )
 32
 33        self.myparam = param.param()
 34        if os.path.exists(
 35            os.path.join(os.path.expanduser("~"), ".smashbox", "param.yaml")
 36        ):
 37            self.myparam.load_param(
 38                os.path.join(
 39                    os.path.expanduser("~"), ".smashbox", "default_param.yaml"
 40                )
 41            )
 42        else:
 43            self.myparam.write_param(
 44                os.path.join(
 45                    os.path.expanduser("~"), ".smashbox", "default_param.yaml"
 46                )
 47            )
 48
 49    def help(self):
 50        """Display help..."""
 51
 52        print("! Welcome to SmashBox !")
 53        print(
 54            "First of all, you must configure main parameters stored in 'myparam' attribute."
 55        )
 56        print(
 57            " - self.myparam.list_param() : list all parameters to configure in self.myparam.param."
 58        )
 59        print(
 60            " - self.myparam.set_param(attr, value) : set value to attribute attr."
 61        )
 62        print(
 63            " - self.myparam.get_param(attr) : get the value of attribute attr."
 64        )
 65        print(
 66            " - self.myparam.list_asset_files() : list data available with SmashBox."
 67        )
 68        print(
 69            " - self.myparam.param : attribute which contains all parameters."
 70        )
 71        print("")
 72        print("When ready, you will able to create a new model attribute:")
 73        print(" - self.newmodel('mymodel')")
 74        print("")
 75        print(
 76            "After the model 'mymodel' has been created, you can adjust the model setup 'self.mymodel.mysetup':"
 77        )
 78        print(
 79            " - self.mymodel.mysetup.setup: dictionnary which contain the setup"
 80        )
 81        print(
 82            " - self.mymodel.mysetup.list_available_setup(): list preconfigured and ready to use setup in SmashBox"
 83        )
 84        print(
 85            " - self.mymodel.mysetup.update_setup(dict): update the smash setup with the input dictionary dict."
 86        )
 87        print("")
 88        print(
 89            "When ready, you will be able to create the mesh, the Smash model and run a simulation:"
 90        )
 91        print(" - self.mymodel.make_mesh()")
 92        print(" - self.mymodel.make_model()")
 93
 94    def newmodel(self, name="mymodel"):
 95        """
 96        Create a new model attribute. This attribute will handle meshing, hydrological modelisation, plot and stattistics.
 97
 98        Parameters
 99        ----------
100
101        name : str
102            The name of the attribute.
103
104        Examples
105        --------
106
107        >>> es=smashbox.SmashBox()
108        >>> sb.newmodel("RealCollobrier")
109
110        """
111
112        name = name.replace(" ", "_")
113
114        if not hasattr(self, name):
115
116            setattr(self, name, model.model(name, self.myparam))
117
118        else:
119            raise ValueError(
120                f"Error: Cannot create model '{name}', model attribute '{name}' already exist"
121            )
122
123        if not hasattr(self, "multimodel_statistics"):
124
125            setattr(
126                self,
127                "multimodel_statistics",
128                multimodel_statistics.multimodel_statistics(self),
129            )
130
131    def import_smash_model(
132        self, name="mymodel", model=None, load_catchment_polygon=False
133    ):
134        """Import a smash model inside a smashbox container to use plot and stats capabilities.
135
136        :param name: name of the smabox model, defaults to "mymodel"
137        :type name: str
138        :param model: A smash model from smash.Model(), defaults to None
139        :type model: smash.Model()
140        :param load_catchment_polygon: load polygon data, defaults to False
141        :type load_catchment_polygon: bool
142
143        :rtype: None
144        """
145        self.newmodel(name)
146        sb_model = getattr(self, name)
147        sb_model.mysetup.update_setup(
148            pyhdf5_handler.read_object_as_dict(model.setup)
149        )
150        sb_model.mymesh.mesh = pyhdf5_handler.read_object_as_dict(model.mesh)
151        if load_catchment_polygon:
152            sb_model.mymesh.load_catchment_polygon(self.myparam.param)
153        sb_model.mysmashmodel.smash = model
154
155    def delmodel(self, modelname):
156        """
157        Delete a model attribute.
158
159        Parameters
160        ----------
161
162        modelname : str
163            The name of the attribute to delete.
164
165        Examples
166        --------
167
168        >>> es=smashbox.SmashBox()
169        >>> sb.newmodel("RealCollobrier")
170        >>> sb.delmodel("RealCollobrier")
171
172        """
173        if hasattr(self, modelname) and isinstance(
174            getattr(self, modelname), model.model
175        ):
176            delattr(self, modelname)
177
178        else:
179
180            raise ValueError(f"Error: {modelname} is not a valid model name.")
181
182    def copymodel(
183        self, modelsource, modelname, copy_smash_model=True, only=None
184    ):
185        """
186        Copy a model container
187
188        Parameters
189        ----------
190
191        modelsource : str
192            The name of the source model attribute
193        modelname : str
194            The name of the new model attribute.
195
196        Examples
197        --------
198
199        >>> es=smashbox.SmashBox()
200        >>> sb.newmodel("RealCollobrier")
201        >>> sb.copymodel("RealCollobrier", RealCollobrier2")
202        """
203
204        if only is None:
205            only = [
206                "mysmashmodel",
207                "warmup_model",
208                "optimize_model",
209                "validation_model",
210            ]
211
212        if not hasattr(self, modelsource):
213            raise ValueError(
214                f"Error: Source model '{modelsource}' does not exist"
215            )
216
217        if not hasattr(self, modelname):
218
219            self.newmodel(modelname)
220
221            source = getattr(self, modelsource)
222            destination = getattr(self, modelname)
223
224            setattr(
225                destination,
226                "_myparam",
227                copy.deepcopy(getattr(source, "_myparam")),
228            )
229            setattr(
230                destination, "mymesh", copy.deepcopy(getattr(source, "mymesh"))
231            )
232            setattr(
233                destination,
234                "mysetup",
235                copy.deepcopy(getattr(source, "mysetup")),
236            )
237
238            if copy_smash_model:
239
240                setup_update = {
241                    "read_prcp": False,
242                    "read_pet": False,
243                    "read_snow": False,
244                    "read_qobs": False,
245                    "read_temp": False,
246                }
247                destination.mysetup.update_setup(setup_update)
248
249                if (
250                    "mysmashmodel" in only
251                    and source.mysmashmodel.smash is not None
252                ):
253                    setattr(
254                        destination.mysmashmodel,
255                        "smash",
256                        source.mysmashmodel.smash.copy(),
257                    )
258                    for key, value in setup_update.items():
259                        setattr(
260                            destination.mysmashmodel.smash.setup, key, value
261                        )
262
263                if (
264                    "warmup_model" in only
265                    and source.warmup_model.smash is not None
266                ):
267                    setattr(
268                        destination.warmup_model,
269                        "smash",
270                        source.warmup_model.smash.copy(),
271                    )
272                    for key, value in setup_update.items():
273                        setattr(
274                            destination.warmup_model.smash.setup, key, value
275                        )
276
277                if (
278                    "optimize_model" in only
279                    and source.optimize_model.smash is not None
280                ):
281                    setattr(
282                        destination.optimize_model,
283                        "smash",
284                        source.optimize_model.smash.copy(),
285                    )
286                    for key, value in setup_update.items():
287                        setattr(
288                            destination.optimize_model.smash.setup, key, value
289                        )
290
291                if (
292                    "validation_model" in only
293                    and source.validation_model.smash is not None
294                ):
295                    setattr(
296                        destination,
297                        "validation_model",
298                        source.validation_model.smash.copy(),
299                    )
300                    for key, value in setup_update.items():
301                        setattr(
302                            destination.validation_model.smash.setup,
303                            key,
304                            value,
305                        )
306
307                if hasattr(source, "_fstates"):
308                    setattr(
309                        destination,
310                        "_fstates",
311                        copy.deepcopy(getattr(source, "_fstates")),
312                    )
313                if hasattr(source, "_istates"):
314                    setattr(
315                        destination,
316                        "_istates",
317                        copy.deepcopy(getattr(source, "_istates")),
318                    )
319                if hasattr(source, "_myatmos_data_connector"):
320                    setattr(
321                        destination,
322                        "_myatmos_data_connector",
323                        copy.deepcopy(
324                            getattr(source, "_myatmos_data_connector")
325                        ),
326                    )
327
328        else:
329            raise ValueError(
330                f"Error: Cannot copy model '{modelsource}', model attribute '{modelname}' already exist"
331            )
332
333        # get new model object
334        # mynewmodel = getattr(self, modelname)
335
336        # reset attributes
337        # mynewmodel.mysmashmodel = None
338        # mynewmodel.extra_smash_results = None
339        # mynewmodel._fstates = None
340        # mynewmodel._istates = None
341        # mynewmodel._myatmos_data_connector = None
342        # mynewmodel.warmup_model = None
343        # mynewmodel.mystats = mystats.mystats(self)
344        # mynewmodel.myplot = myplot.myplot(self)
345
346    def copymodeldata(self, modelsource, modelname):
347        """
348        Copy a model data
349
350        Parameters
351        ----------
352
353        modelsource : str
354            The name of the source model attribute
355        modelname : str
356            The name of the new model attribute.
357
358        Examples
359        --------
360
361        >>> es=smashbox.SmashBox()
362        >>> sb.newmodel("RealCollobrier")
363        >>> sb.copymodeldata("RealCollobrier", RealCollobrier2")
364        """
365        if not hasattr(self, modelsource):
366            raise ValueError(
367                f"Error: Source model '{modelsource}' does not exist"
368            )
369
370        if not hasattr(self, modelname):
371            raise ValueError(
372                f"Error: destination model '{modelname}' does not exist"
373            )
374
375        source = getattr(self, modelsource)
376        destination = getattr(self, modelname)
377
378        if (
379            source.mysmashmodel is not None
380            and destination.mysmashmodel is not None
381        ):
382            destination.mysmashmodel.smash.atmos_data = (
383                source.mysmashmodel.smash.atmos_data.copy()
384            )
385            destination.mysmashmodel.smash.response_data = (
386                source.mysmashmodel.smash.response_data.copy()
387            )
388        if (
389            source.warmup_model is not None
390            and destination.warmup_model is not None
391        ):
392            destination.warmup_model.smash.atmos_data = (
393                source.warmup_model.smash.atmos_data.copy()
394            )
395            destination.warmup_model.smash.response_data = (
396                source.warmup_model.smash.response_data.copy()
397            )
398        if (
399            source.optimize_model is not None
400            and destination.optimize_model is not None
401        ):
402            destination.optimize_model.smash.atmos_data = (
403                source.optimize_model.smash.atmos_data.copy()
404            )
405            destination.optimize_model.smash.response_data = (
406                source.optimize_model.smash.response_data.copy()
407            )
408        if (
409            source.validation_model is not None
410            and destination.validation_model is not None
411        ):
412            destination.validation_model.smash.atmos_data = (
413                source.validation_model.smash.atmos_data.copy()
414            )
415            destination.validation_model.smash.response_data = (
416                source.validation_model.smash.response_data.copy()
417            )
418
419    def load_containers(
420        self, path, only=None, exclude=None, load_catchment_polygon=True
421    ):
422
423        subdir_list = os.listdir(path)
424
425        if only is not None:
426            for d in only:
427                if not d in subdir_list:
428                    raise ValueError(f"{d} does not exist in {path}")
429
430            subdir_list = only
431
432        if exclude is not None:
433            for d in exclude:
434                if d in subdir_list:
435                    subdir_list.remove(d)
436
437        for subdir in subdir_list:
438            box = pyhdf5_handler.read_hdf5file_as_dict(
439                os.path.join(path, subdir, "smashbox.hdf5")
440            )
441
442            mysmashmodel = None
443            warmup_model = None
444            optimize_model = None
445            validation_model = None
446
447            if os.path.exists(os.path.join(path, subdir, "mysmashmodel.hdf5")):
448                mysmashmodel = smash.io.read_model(
449                    os.path.join(path, subdir, "mysmashmodel.hdf5")
450                )
451            if os.path.exists(os.path.join(path, subdir, "warmup_model.hdf5")):
452                warmup_model = smash.io.read_model(
453                    os.path.join(path, subdir, "warmup_model.hdf5")
454                )
455            if os.path.exists(
456                os.path.join(path, subdir, "optimize_model.hdf5")
457            ):
458                optimize_model = smash.io.read_model(
459                    os.path.join(path, subdir, "optimize_model.hdf5")
460                )
461            if os.path.exists(
462                os.path.join(path, subdir, "validation_model.hdf5")
463            ):
464                validation_model = smash.io.read_model(
465                    os.path.join(path, subdir, "validation_model.hdf5")
466                )
467
468            for key, value in box["param"].items():
469                if getattr(self.myparam.param, key) is None:
470                    self.myparam.set_param(key, value)
471                else:
472                    dtype = type(getattr(self.myparam.param, key))
473                    try:
474                        converted_value = dtype(value)
475                        self.myparam.set_param(key, converted_value)
476                    except (ValueError, TypeError) as e:
477                        try:
478                            self.myparam.set_param(key, value)
479                        except (ValueError, TypeError) as ee:
480                            print(
481                                f"Type error for {key}: {dtype}/{type(value)} - {e}"
482                            )
483                            print(
484                                f"Can't import {key}: {dtype}/{type(value)} - {ee}"
485                            )
486
487            self.newmodel(subdir)
488
489            container = getattr(self, subdir)
490
491            for key, value in box["mysetup"]["setup"].items():
492                container.mysetup.set_setup(key, value)
493
494            container.mymesh.mesh = box["mymesh"]["mesh"]
495
496            # getting models extra result and stats
497            for model in [
498                "mysmashmodel",
499                "warmup_model",
500                "optimize_model",
501                "validation_model",
502            ]:
503
504                if model in box:
505                    structure = pyhdf5_handler.src.object_handler.generate_dict_structure(
506                        box[model], include_method=False
507                    )
508                    pyhdf5_handler.src.object_handler.map_dict_to_object(
509                        structure,
510                        box[model],
511                        getattr(container, model),
512                    )
513
514            setattr(container.mysmashmodel, "smash", mysmashmodel)
515            setattr(container.warmup_model, "smash", warmup_model)
516            setattr(container.optimize_model, "smash", optimize_model)
517            setattr(container.validation_model, "smash", validation_model)
518
519            container._myparam = copy.deepcopy(self.myparam)
520            if load_catchment_polygon:
521                container.mymesh.load_catchment_polygon(
522                    param=container._myparam.param
523                )
524
525    # def load_pickle_containers(self, path, only=None, exclude=None):
526
527    #     subdir_list = os.listdir(path)
528
529    #     if only is not None:
530    #         for d in only:
531    #             if not os.path.isdir(d):
532    #                 raise ValueError()
533
534    #         subdir_list = only
535    #     if exclude is not None:
536    #         for d in exclude:
537    #             if os.path.isdir(d):
538    #                 subdir_list.remove(d)
539
540    #     for subdir in subdir_list:
541
542    #         mysmashmodel = None
543    #         warmup_model = None
544    #         optimize_model = None
545
546    #         if os.path.exists(os.path.join(path, subdir, "mysmashmodel.hdf5")):
547    #             mysmashmodel = smash.io.read_model(
548    #                 os.path.join(path, subdir, "mysmashmodel.hdf5")
549    #             )
550    #         if os.path.exists(os.path.join(path, subdir, "warmup_model.hdf5")):
551    #             warmup_model = smash.io.read_model(
552    #                 os.path.join(path, subdir, "warmup_model.hdf5")
553    #             )
554    #         if os.path.exists(os.path.join(path, subdir, "optimize_model.hdf5")):
555    #             optimize_model = smash.io.read_model(
556    #                 os.path.join(path, subdir, "optimize_model.hdf5")
557    #             )
558
559    #         self.newmodel(subdir)
560    #         container = getattr(self, subdir)
561
562    #         for file in glob.glob("*.pkl"):
563    #             key = os.path.basename(file).split(".")[0]
564    #             obj = pickle.load(file)
565    #             setattr(container, key, obj)
566
567    #         container.mysmashmodel = mysmashmodel
568    #         container.warmup_model = warmup_model
569    #         container.optimize_model = optimize_model