Coverage for /builds/ase/ase/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 constant: energy value for the constant. 

45 

46 Example: 

47 

48 >>> from ase.optimize import GPMin 

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

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

51 """ 

52 

53 def __init__(self, constant): 

54 self.constant = constant 

55 Prior.__init__(self) 

56 

57 def potential(self, x): 

58 d = x.shape[0] 

59 output = np.zeros(d + 1) 

60 output[0] = self.constant 

61 return output 

62 

63 def set_constant(self, constant): 

64 self.constant = constant 

65 

66 

67class CalculatorPrior(Prior): 

68 """CalculatorPrior object, allows the user to 

69 use another calculator as prior function instead of the 

70 default constant. 

71 

72 Parameters: 

73 

74 atoms: the Atoms object 

75 calculator: one of ASE's calculators 

76 """ 

77 

78 def __init__(self, atoms, calculator): 

79 Prior.__init__(self) 

80 self.atoms = atoms.copy() 

81 self.atoms.calc = calculator 

82 

83 def potential(self, x): 

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

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

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

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