Coverage for model_plots / __init__.py: 100%

1import logging 

2 

3logging.basicConfig( 

4 format="%(levelname)-8s %(module)-14s %(message)s", level=logging.INFO 

5) 

6log = logging.getLogger(__name__) 

7 

8# Set matplotlib backend globally in order to avoid potential issues from 

9# people forgetting to set this 

10import matplotlib 

11 

12matplotlib.use("Agg") 

13 

14# --- Custom DAG operations to register -------------------------------------------------------------------------------- 

15import numpy as np 

16import pandas as pd 

17import xarray as xr 

18from utopya.eval import register_operation 

19 

20register_operation(name="pd.Index", func=pd.Index) 

21register_operation(name="np.maximum", func=np.maximum) 

22register_operation(name="np.subtract", func=np.subtract) 

23register_operation(name="np.exp", func=np.exp) 

24register_operation(name="np.sin", func=np.sin) 

25register_operation(name="np.random.randint", func=np.random.randint) 

26register_operation(name="xr.where", func=xr.where) 

27register_operation(name=".idxmax", func=lambda d, *a, **k: d.idxmax(*a, **k)) 

28register_operation(name=".idxmin", func=lambda d, *a, **k: d.idxmin(*a, **k)) 

29register_operation(name=".stack", func=lambda d, *a, **k: d.stack(*a, **k)) 

30register_operation(name=".to_xarray", func=lambda d, *a, **k: d.to_xarray(*a, **k)) 

31register_operation( 

32 name="pd.to_datetime", func=lambda d, *a, **k: pd.to_datetime(d, *a, **k) 

33) 

34register_operation(name="pd.date_range", func=lambda *a, **k: pd.date_range(*a, **k)) 

35register_operation(name=".index", func=lambda d: d.index) 

36register_operation(name=".reset_index", func=lambda d, *a, **k: d.reset_index(*a, **k)) 

37register_operation(name="replace", func=lambda s, *a, **k: s.replace(*a, **k)) 

38register_operation(name="np.nansum", func=np.nansum) 

39register_operation(name="np.histogramdd", func=np.histogramdd) 

40register_operation(name="replace", func=lambda s, *a, **k: s.replace(*a, **k)) 

41register_operation(name="list_of", func=lambda s: [s]) 

42register_operation(name="zip", func=zip) 

43register_operation(name=".capitalize", func=lambda s: s.capitalize()) 

44register_operation(name="np.randint", func=np.random.randint) 

45register_operation(name=".to_xarray", func=lambda d, *a, **k: d.to_xarray(*a, **k)) 

46register_operation(name="np.nansum", func=np.nansum) 

47register_operation(name="np.histogramdd", func=np.histogramdd) 

48register_operation(name="np.ones", func=np.ones) 

49register_operation(name=".start", func=lambda d: d.start) 

50register_operation(name=".stop", func=lambda d: d.stop) 

51register_operation(name=".step", func=lambda d: d.step) 

52 

53from .bar import bar, hist 

54from .data_ops import * 

55from .nw_ops import * 

56from .plots import * 

57 

58# Model operations 

59from .Covid import * 

60from .SIR import *