Coverage for ase / optimize / gpmin / prior.py: 73.53%

1# fmt: off 

2 

3import numpy as np 

4 

5 

6class Prior(): 

7 """Base class for all priors for the bayesian optimizer. 

8 

9 The __init__ method and the prior method are implemented here. 

10 Each child class should implement its own potential method, that will be 

11 called by the prior method implemented here. 

12 

13 When used, the prior should be initialized outside the optimizer and the 

14 Prior object should be passed as a function to the optimizer. 

15 """ 

16 

17 def __init__(self): 

18 """Basic prior implementation.""" 

19 

20 def prior(self, x): 

21 """Actual prior function, common to all Priors""" 

22 if len(x.shape) > 1: 

23 n = x.shape[0] 

24 return np.hstack([self.potential(x[i, :]) for i in range(n)]) 

25 else: 

26 return self.potential(x) 

27 

28 

29class ZeroPrior(Prior): 

30 """ZeroPrior object, consisting on a constant prior with 0eV energy.""" 

31 

32 def __init__(self): 

33 Prior.__init__(self) 

34 

35 def potential(self, x): 

36 return np.zeros(x.shape[0] + 1) 

37 

38 

39class ConstantPrior(Prior): 

40 """Constant prior, with energy = constant and zero forces 

41 

42 Parameters 

43 ---------- 

44 

45 constant: energy value for the constant. 

46 

47 Example: 

48 

49 >>> from ase.optimize import GPMin 

50 >>> from ase.optimize.gpmin.prior import ConstantPrior 

51 >>> op = GPMin(atoms, Prior = ConstantPrior(10) 

52 """ 

53 

54 def __init__(self, constant): 

55 self.constant = constant 

56 Prior.__init__(self) 

57 

58 def potential(self, x): 

59 d = x.shape[0] 

60 output = np.zeros(d + 1) 

61 output[0] = self.constant 

62 return output 

63 

64 def set_constant(self, constant): 

65 self.constant = constant 

66 

67 

68class CalculatorPrior(Prior): 

69 """CalculatorPrior object, allows the user to 

70 use another calculator as prior function instead of the 

71 default constant. 

72 

73 Parameters 

74 ---------- 

75 

76 atoms: the Atoms object 

77 calculator: one of ASE's calculators 

78 """ 

79 

80 def __init__(self, atoms, calculator): 

81 Prior.__init__(self) 

82 self.atoms = atoms.copy() 

83 self.atoms.calc = calculator 

84 

85 def potential(self, x): 

86 self.atoms.set_positions(x.reshape(-1, 3)) 

87 V = self.atoms.get_potential_energy(force_consistent=True) 

88 gradV = -self.atoms.get_forces().reshape(-1) 

89 return np.append(np.array(V).reshape(-1), gradV)