Coverage for HARK/parser.py: 81%

1from HARK.distributions import Bernoulli, Lognormal, MeanOneLogNormal 

2from sympy.utilities.lambdify import lambdify 

3from sympy.parsing.sympy_parser import parse_expr 

4import yaml 

5 

6 

7class ControlToken: 

8 """ 

9 Represents a parsed Control variable. 

10 """ 

11 

12 def __init__(self, args): 

13 pass 

14 

15 

16class Expression: 

17 def __init__(self, text): 

18 self.txt 

19 self.expr = parse_expr(text) 

20 self.npf = self.func() 

21 

22 # first derivatives. 

23 self.grad = { 

24 sym.__str__(): self.expr.diff(sym) for sym in list(self.expr.free_symbols) 

25 } 

26 

27 def func(self): 

28 return lambdify(list(self.expr.free_symbols), self.expr, "numpy") 

29 

30 

31def tuple_constructor_from_class(cls): 

32 def constructor(loader, node): 

33 value = loader.construct_mapping(node) 

34 return (cls, value) 

35 

36 return constructor 

37 

38 

39def math_text_to_lambda(text): 

40 """ 

41 Returns a function represented by the given mathematical text. 

42 """ 

43 expr = parse_expr(text) 

44 func = lambdify(list(expr.free_symbols), expr, "numpy") 

45 return func 

46 

47 

48def harklang_loader(): 

49 """ 

50 A PyYAML loader that supports tags for HARKLang, 

51 such as random variables and model tags. 

52 """ 

53 loader = yaml.SafeLoader 

54 yaml.SafeLoader.add_constructor( 

55 "!Bernoulli", tuple_constructor_from_class(Bernoulli) 

56 ) 

57 yaml.SafeLoader.add_constructor( 

58 "!MeanOneLogNormal", tuple_constructor_from_class(MeanOneLogNormal) 

59 ) 

60 yaml.SafeLoader.add_constructor( 

61 "!Lognormal", tuple_constructor_from_class(Lognormal) 

62 ) 

63 yaml.SafeLoader.add_constructor( 

64 "!Control", tuple_constructor_from_class(ControlToken) 

65 ) 

66 

67 return loader