-
Notifications
You must be signed in to change notification settings - Fork 10
/
NNBasisPotential.jl
173 lines (134 loc) · 4.91 KB
/
NNBasisPotential.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
# TODO: Add to InteratomicPotentials.jl/InteratomicBasisPotentials.jl
"""
NNBasisPotential
Definition of the neural network basis potential composed type.
"""
mutable struct NNBasisPotential <: AbstractPotential
nn
nn_params
ibp_params
end
"""
potential_energy(A::AbstractSystem, p::NNBasisPotential)
`A`: atomic system or atomic configuration.
`p`: neural network basis potential.
Returns the potential energy of a system using a neural network basis potential.
"""
function potential_energy(A::AbstractSystem, p::NNBasisPotential)
b = evaluate_basis(A, p.ibp_params)
return p.nn(b)
end
"""
force(A::AbstractSystem, p::NNBasisPotential)
`A`: atomic system or atomic configuration.
`p`: neural network basis potential.
Returns the force of a system using a neural network basis potential.
"""
function force(A::AbstractSystem, p::NNBasisPotential)
b = evaluate_basis(A, p.ibp_params)
dnndb = first(gradient(p.nn, b))
dbdr = evaluate_basis_d(A, p.ibp_params)
return [[-dnndb ⋅ dbdr[atom][coor,:] for coor in 1:3]
for atom in 1:length(dbdr)]
end
"""
potential_energy(b::Vector, p::NNBasisPotential)
`b`: vector of energy descriptors.
`p`: neural network basis potential.
Returns the potential energy of a system using a neural network basis potential.
"""
function potential_energy(b::Vector, p::NNBasisPotential)
return sum(p.nn(b))
end
"""
potential_energy(b::Vector, ps::Vector, re)
`b`: energy descriptors of a system.
`ps`: neural network parameters. See Flux.destructure.
`re`: neural network restructure. See Flux.destructure.
Returns the potential energy of a system using a neural network basis potential.
"""
function potential_energy(b::Vector, ps::Vector, re)
return sum(re(ps)(b))
end
# TODO: create issue.
#Note: calculating the gradient of the loss function requires in turn
#calculating the gradient of the energy. That is, calculating the gradient of
#a function that calculates another gradient.
#So far I have not found a clean way to do this using the abstractions
#provided by Flux, which in turn is integrated with Zygote.
#Links related to this issue:
#https://discourse.julialang.org/t/how-to-add-norm-of-gradient-to-a-loss-function/69873/16
#https://discourse.julialang.org/t/issue-with-zygote-over-forwarddiff-derivative/70824
#https://github.com/FluxML/Zygote.jl/issues/953#issuecomment-841882071
#
#To solve this for the moment I am calculating one of the gradients analytically.
#To do this I had to use Flux.destructure, which I think makes the code slower
#because of the copies it creates.
"""
grad_mlp(nn_params, x0)
`nn_params`: neural network parameters.
`x0`: first layer input (energy descriptors).
Returns the analytical derivative of a feedforward neural network.
"""
function grad_mlp(nn_params, x0)
dsdy(x) = x>0 ? 1 : 0 # Flux.σ(x) * (1 - Flux.σ(x))
prod = 1; x = x0
n_layers = length(nn_params) ÷ 2
for i in 1:2:2(n_layers-1)-1 # i : 1, 3
y = nn_params[i] * x + nn_params[i+1]
x = Flux.relu.(y) # Flux.σ.(y)
prod = dsdy.(y) .* nn_params[i] * prod
end
i = 2(n_layers)-1
prod = nn_params[i] * prod
return prod
end
"""
force(b::Vector, dbdr::Vector, p::NNBasisPotential)
`b`: energy descriptors of a system.
`dbdr`: force descriptors of a system.
`p`: neural network basis potential.
Returns the force of a system using a neural network basis potential.
"""
function force(b::Vector, dbdr::Vector, p::NNBasisPotential)
dnndb = grad_mlp(p.nn_params, b)
return dnndb ⋅ dbdr
end
"""
force(b::Vector, dbdr::Vector, ps::Vector, re)
`b`: energy descriptors of a system.
`dbdr`: force descriptors of a system.
`ps`: neural network parameters. See Flux.destructure.
`re`: neural network restructure. See Flux.destructure.
Returns the force of a system using a neural network basis potential.
"""
function force(b::Vector, dbdr::Vector, ps::Vector, re)
nn_params = Flux.params(re(ps))
dnndb = grad_mlp(nn_params, b)
return dnndb ⋅ dbdr
end
# Computing the force using ForwardDiff
#function force(b::Vector, dbdr::Vector, p::NNBasisPotential)
# dnndb = ForwardDiff.gradient(x -> sum(p.nn(x)), b)
# return dnndb ⋅ dbdr
#end
#function force(b::Vector, dbdr::Vector, p::NNBasisPotential)
# dnndb = gradient(x -> sum(p.nn(x)), b)[1]
# return dnndb ⋅ dbdr
#end
# Computing the force using ForwardDiff and destructure
#function force(b::Vector, dbdr::Vector, ps::Vector, re)
# dnndb = ForwardDiff.gradient(x -> sum(re(ps)(x)), b)
# return dnndb ⋅ dbdr
#end
# Computing the force using pullback
#function force(b::Vector, dbdr::Vector, p::NNBasisPotential)
# y, pullback = Zygote.pullback(p.nn, b)
# dnndb = pullback(ones(size(y)))[1]
# return dnndb ⋅ dbdr
#end
#function force(b::Vector, dbdr::Vector, ps::Vector, re)
# y, pullback = Zygote.pullback(re(ps), b)
# dnndb = pullback(ones(size(y)))[1]
# return dnndb ⋅ dbdr
#end