1 //
2 // CDDL HEADER START
3 //
4 // The contents of this file are subject to the terms of the Common Development
5 // and Distribution License Version 1.0 (the "License").
6 //
7 // You can obtain a copy of the license at
8 // http://www.opensource.org/licenses/CDDL-1.0.  See the License for the
9 // specific language governing permissions and limitations under the License.
10 //
11 // When distributing Covered Code, include this CDDL HEADER in each file and
12 // include the License file in a prominent location with the name LICENSE.CDDL.
13 // If applicable, add the following below this CDDL HEADER, with the fields
14 // enclosed by brackets "[]" replaced with your own identifying information:
15 //
16 // Portions Copyright (c) [yyyy] [name of copyright owner]. All rights reserved.
17 //
18 // CDDL HEADER END
19 //
20 
21 //
22 // Copyright (c) 2019, Regents of the University of Minnesota.
23 // All rights reserved.
24 //
25 // Contributors:
26 //    Mingjian Wen
27 //
28 
29 #ifndef ANN_IMPLEMENTATION_HPP_
30 #define ANN_IMPLEMENTATION_HPP_
31 
32 #include "ANN.hpp"
33 #include "KIM_LogMacros.hpp"
34 #include "descriptor.h"
35 #include "helper.hpp"
36 #include "network.h"
37 #include <vector>
38 
39 #define DIM 3
40 #define ONE 1.0
41 
42 #define MAX_PARAMETER_FILES 3
43 
44 //==============================================================================
45 //
46 // Declaration of ANNImplementation class
47 //
48 //==============================================================================
49 
50 //******************************************************************************
51 class ANNImplementation
52 {
53  public:
54   ANNImplementation(KIM::ModelDriverCreate * const modelDriverCreate,
55                     KIM::LengthUnit const requestedLengthUnit,
56                     KIM::EnergyUnit const requestedEnergyUnit,
57                     KIM::ChargeUnit const requestedChargeUnit,
58                     KIM::TemperatureUnit const requestedTemperatureUnit,
59                     KIM::TimeUnit const requestedTimeUnit,
60                     int * const ier);
61   ~ANNImplementation();  // no explicit Destroy() needed here
62 
63   int Refresh(KIM::ModelRefresh * const modelRefresh);
64   int Compute(KIM::ModelCompute const * const modelCompute,
65               KIM::ModelComputeArguments const * const modelComputeArguments);
66   int ComputeArgumentsCreate(KIM::ModelComputeArgumentsCreate * const
67                                  modelComputeArgumentsCreate) const;
68   int ComputeArgumentsDestroy(KIM::ModelComputeArgumentsDestroy * const
69                                   modelComputeArgumentsDestroy) const;
70 
71  private:
72   // Constant values that never change
73   //   Set in constructor (via SetConstantValues)
74   //
75   //
76   // ANNImplementation: constants
77 
78   double energyScale_;
79 
80   // Constant values that are read from the input files and never change
81   //   Set in constructor (via functions listed below)
82   // none
83   //
84   // Private Model Parameters
85   //   Memory allocated in AllocatePrivateParameterMemory() (from constructor)
86   //   Memory deallocated in destructor
87   //   Data set in ReadParameterFile routines
88   int ensemble_size_;
89   int last_ensemble_size_;
90   //
91   // KIM API: Model Parameters whose (pointer) values never change
92   //   Memory allocated in AllocateParameterMemory() (from constructor)
93   //   Memory deallocated in destructor
94   //   Data set in ReadParameterFile routines OR by KIM Simulator
95   int active_member_id_;
96 
97   // Mutable values that only change when Refresh() executes
98   //   Set in Refresh (via SetRefreshMutableValues)
99   //
100   //
101   // KIM API: Model Parameters (can be changed directly by KIM Simulator)
102   // none
103   //
104   // ANNImplementation: values (changed only by Refresh())
105   int last_active_member_id_;
106   double influenceDistance_;
107   int modelWillNotRequestNeighborsOfNoncontributingParticles_;
108 
109   // Mutable values that can change with each call to Refresh() and Compute()
110   //   Memory may be reallocated on each call
111   //
112   //
113   // ANNImplementation: values that change
114   int cachedNumberOfParticles_;
115 
116   // descriptor and network
117   Descriptor * descriptor_;
118   NeuralNetwork * network_;
119 
120   // Helper methods
121   //
122   //
123   // Related to constructor
124   void AllocatePrivateParameterMemory();
125   void AllocateParameterMemory();
126 
127   static int
128   OpenParameterFiles(KIM::ModelDriverCreate * const modelDriverCreate,
129                      int const numberParameterFiles,
130                      FILE * parameterFilePointers[MAX_PARAMETER_FILES]);
131   int ProcessParameterFiles(
132       KIM::ModelDriverCreate * const modelDriverCreate,
133       int const numberParameterFiles,
134       FILE * const parameterFilePointers[MAX_PARAMETER_FILES]);
135   static void
136   CloseParameterFiles(int const numberParameterFiles,
137                       FILE * const parameterFilePointers[MAX_PARAMETER_FILES]);
138   int ConvertUnits(KIM::ModelDriverCreate * const modelDriverCreate,
139                    KIM::LengthUnit const requestedLengthUnit,
140                    KIM::EnergyUnit const requestedEnergyUnit,
141                    KIM::ChargeUnit const requestedChargeUnit,
142                    KIM::TemperatureUnit const requestedTemperatureUnit,
143                    KIM::TimeUnit const requestedTimeUnit);
144   int RegisterKIMModelSettings(
145       KIM::ModelDriverCreate * const modelDriverCreate) const;
146   int RegisterKIMComputeArgumentsSettings(
147       KIM::ModelComputeArgumentsCreate * const modelComputeArgumentsCreate)
148       const;
149   int RegisterKIMParameters(KIM::ModelDriverCreate * const modelDriverCreate);
150   int RegisterKIMFunctions(
151       KIM::ModelDriverCreate * const modelDriverCreate) const;
152 
153   //
154   // Related to Refresh()
155   template<class ModelObj>
156   int SetRefreshMutableValues(ModelObj * const modelObj);
157 
158   //
159   // Related to Compute()
160   int SetComputeMutableValues(
161       KIM::ModelComputeArguments const * const modelComputeArguments,
162       bool & isComputeProcess_dEdr,
163       bool & isComputeProcess_d2Edr2,
164       bool & isComputeEnergy,
165       bool & isComputeForces,
166       bool & isComputeParticleEnergy,
167       bool & isComputeVirial,
168       bool & isComputeParticleVirial,
169       int const *& particleSpeciesCodes,
170       int const *& particleContributing,
171       VectorOfSizeDIM const *& coordinates,
172       double *& energy,
173       VectorOfSizeDIM *& forces,
174       double *& particleEnergy,
175       VectorOfSizeSix *& virial,
176       VectorOfSizeSix *& particleVirial);
177   int CheckParticleSpeciesCodes(KIM::ModelCompute const * const modelCompute,
178                                 int const * const particleSpeciesCodes) const;
179   int GetComputeIndex(const bool & isComputeProcess_dEdr,
180                       const bool & isComputeProcess_d2Edr2,
181                       const bool & isComputeEnergy,
182                       const bool & isComputeForces,
183                       const bool & isComputeParticleEnergy,
184                       const bool & isComputeVirial,
185                       const bool & isComputeParticleVirial) const;
186 
187   // compute functions
188   template<bool isComputeProcess_dEdr,
189            bool isComputeProcess_d2Edr2,
190            bool isComputeEnergy,
191            bool isComputeForces,
192            bool isComputeParticleEnergy,
193            bool isComputeVirial,
194            bool isComputeParticleVirial>
195   int Compute(KIM::ModelCompute const * const modelCompute,
196               KIM::ModelComputeArguments const * const modelComputeArguments,
197               const int * const particleSpeciesCodes,
198               const int * const particleContributing,
199               const VectorOfSizeDIM * const coordinates,
200               double * const energy,
201               VectorOfSizeDIM * const forces,
202               double * const particleEnergy,
203               VectorOfSizeSix virial,
204               VectorOfSizeSix * const particleVirial) const;
205 };
206 
207 //==============================================================================
208 //
209 // Definition of ANNImplementation::Compute functions
210 //
211 // NOTE: Here we rely on the compiler optimizations to prune dead code
212 //       after the template expansions.  This provides high efficiency
213 //       and easy maintenance.
214 //
215 //==============================================================================
216 
217 #undef KIM_LOGGER_OBJECT_NAME
218 #define KIM_LOGGER_OBJECT_NAME modelCompute
219 template<bool isComputeProcess_dEdr,
220          bool isComputeProcess_d2Edr2,
221          bool isComputeEnergy,
222          bool isComputeForces,
223          bool isComputeParticleEnergy,
224          bool isComputeVirial,
225          bool isComputeParticleVirial>
Compute(KIM::ModelCompute const * const modelCompute,KIM::ModelComputeArguments const * const modelComputeArguments,const int * const particleSpeciesCodes,const int * const particleContributing,const VectorOfSizeDIM * const coordinates,double * const energy,VectorOfSizeDIM * const forces,double * const particleEnergy,VectorOfSizeSix virial,VectorOfSizeSix * const particleVirial) const226 int ANNImplementation::Compute(
227     KIM::ModelCompute const * const modelCompute,
228     KIM::ModelComputeArguments const * const modelComputeArguments,
229     const int * const particleSpeciesCodes,
230     const int * const particleContributing,
231     const VectorOfSizeDIM * const coordinates,
232     double * const energy,
233     VectorOfSizeDIM * const forces,
234     double * const particleEnergy,
235     VectorOfSizeSix virial,
236     VectorOfSizeSix * const particleVirial) const
237 {
238   int ier = false;
239 
240   if ((isComputeEnergy == false) && (isComputeParticleEnergy == false)
241       && (isComputeForces == false) && (isComputeProcess_dEdr == false)
242       && (isComputeProcess_d2Edr2 == false) && (isComputeVirial == false)
243       && (isComputeParticleVirial == false))
244   { return ier; }
245 
246   ier = true;
247 
248   if (isComputeProcess_dEdr == true)
249   {
250     LOG_ERROR("process_dEdr not supported by this driver");
251     return ier;
252   }
253 
254   if (isComputeProcess_d2Edr2 == true)
255   {
256     LOG_ERROR("process_d2Edr2 not supported by this driver");
257     return ier;
258   }
259 
260   bool need_dE = ((isComputeForces == true) || (isComputeVirial == true)
261                   || (isComputeParticleVirial == true));
262 
263   // ANNImplementation: values that does not change
264   int const Nparticles = cachedNumberOfParticles_;
265 
266   // initialize energy and forces
267   if (isComputeEnergy == true) { *energy = 0.0; }
268 
269   if (isComputeParticleEnergy == true)
270   {
271     for (int i = 0; i < Nparticles; ++i) { particleEnergy[i] = 0.0; }
272   }
273 
274   if (isComputeForces == true)
275   {
276     for (int i = 0; i < Nparticles; ++i)
277     {
278       for (int j = 0; j < DIM; ++j) { forces[i][j] = 0.0; }
279     }
280   }
281 
282   if (isComputeVirial == true)
283   {
284     for (int i = 0; i < 6; ++i) { virial[i] = 0.0; }
285   }
286 
287   if (isComputeParticleVirial == true)
288   {
289     for (int i = 0; i < Nparticles; ++i)
290     {
291       for (int j = 0; j < 6; ++j) { particleVirial[i][j] = 0.0; }
292     }
293   }
294 
295   // calculate generalized coordinates
296   //
297   // Setup loop over contributing particles
298   for (int i = 0; i < Nparticles; i++)
299   {
300     if (!particleContributing[i]) { continue; }
301 
302     // get neighbors of atom i
303     int numnei = 0;
304     int const * n1atom = 0;
305     modelComputeArguments->GetNeighborList(0, i, &numnei, &n1atom);
306 
307     double * GC = nullptr;
308     double ** dGCdr = nullptr;
309     int const Ndescriptors = descriptor_->get_num_descriptors();
310     AllocateAndInitialize1DArray<double>(GC, Ndescriptors);
311     if (need_dE)
312     {
313       AllocateAndInitialize2DArray<double>(
314           dGCdr, Ndescriptors, (numnei + 1) * DIM);
315     }
316 
317     descriptor_->generate_one_atom(
318         i,
319         reinterpret_cast<double const *>(coordinates),
320         particleSpeciesCodes,
321         n1atom,
322         numnei,
323         GC,
324         dGCdr[0],
325         need_dE);
326 
327     // centering and normalization
328     if (descriptor_->need_normalize())
329     {
330       for (int j = 0; j < Ndescriptors; j++)
331       {
332         double mean;
333         double std;
334         descriptor_->get_feature_mean_and_std(j, mean, std);
335         GC[j] = (GC[j] - mean) / std;
336 
337         if (need_dE)
338         {
339           for (int k = 0; k < (numnei + 1) * DIM; k++) { dGCdr[j][k] /= std; }
340         }
341       }
342     }
343 
344     double E = 0;
345     double * dEdGC = nullptr;
346 
347     // select a specific running mode
348     if (ensemble_size_ == 0 || active_member_id_ == 0)
349     {
350       // fully-connected NN
351 
352       network_->set_fully_connected(true);
353       int ensemble_index = 0;  // ignored by in fully-connected mode
354 
355       // NN forward
356       network_->forward(GC, 1, Ndescriptors, ensemble_index);
357       E = network_->get_sum_output();
358 
359       // NN backprop
360       if (need_dE)
361       {
362         network_->backward();
363         dEdGC = network_->get_grad_input();
364       }
365     }
366     else if (active_member_id_ > 0 && active_member_id_ <= ensemble_size_)
367     {
368       // a specific member of the ensemble
369 
370       network_->set_fully_connected(false);
371       int ensemble_index = active_member_id_ - 1;  // internally starts from 0
372 
373       // NN forward
374       network_->forward(GC, 1, Ndescriptors, ensemble_index);
375       E = network_->get_sum_output();
376 
377       // NN backprop
378       if (need_dE)
379       {
380         network_->backward();
381         dEdGC = network_->get_grad_input();
382       }
383     }
384     else if (active_member_id_ == -1)
385     {
386       // average over ensemble
387 
388       network_->set_fully_connected(false);
389       if (need_dE)
390       { AllocateAndInitialize1DArray<double>(dEdGC, Ndescriptors); }
391 
392       for (int iev = 0; iev < ensemble_size_; iev++)
393       {
394         // NN forward
395         network_->forward(GC, 1, Ndescriptors, iev);
396         double eng = network_->get_sum_output();
397         E += eng / ensemble_size_;
398 
399         // NN backprop
400         if (need_dE)
401         {
402           network_->backward();
403           double * deng = network_->get_grad_input();
404           for (int j = 0; j < Ndescriptors; j++)
405           { dEdGC[j] += deng[j] / ensemble_size_; }
406         }
407       }
408     }
409     else
410     {
411       char message[MAXLINE];
412       sprintf(message,
413               "`active_member_id=%d` out of range. Should be [-1, %d]",
414               active_member_id_,
415               ensemble_size_);
416       LOG_ERROR(message);
417       return ier;
418     }
419 
420     // Contribution to energy
421     if (isComputeEnergy == true) { *energy += energyScale_ * E; }
422 
423     // Contribution to particle energy
424     if (isComputeParticleEnergy == true)
425     { particleEnergy[i] += energyScale_ * E; }
426 
427     // Contribution to forces, particle virial, and virial
428     if (need_dE == true)
429     {
430       for (int j = 0; j < Ndescriptors; j++)
431       {
432         for (int k = 0; k < numnei + 1; k++)
433         {
434           int idx;
435           if (k == numnei)
436           {
437             idx = i;  // targeting atom itself
438           }
439           else
440           {
441             idx = n1atom[k];  // neighbors
442           }
443           VectorOfSizeDIM f;
444           f[0] = -dEdGC[j] * dGCdr[j][k * DIM + 0];
445           f[1] = -dEdGC[j] * dGCdr[j][k * DIM + 1];
446           f[2] = -dEdGC[j] * dGCdr[j][k * DIM + 2];
447 
448           if (isComputeForces == true)
449           {
450             forces[idx][0] += energyScale_ * f[0];
451             forces[idx][1] += energyScale_ * f[1];
452             forces[idx][2] += energyScale_ * f[2];
453           }
454 
455           if (isComputeParticleVirial == true || isComputeVirial == true)
456           {
457             VectorOfSizeSix v;
458             v[0] = -energyScale_ * f[0] * coordinates[idx][0];
459             v[1] = -energyScale_ * f[1] * coordinates[idx][1];
460             v[2] = -energyScale_ * f[2] * coordinates[idx][2];
461             v[3] = -energyScale_ * f[1] * coordinates[idx][2];
462             v[4] = -energyScale_ * f[2] * coordinates[idx][0];
463             v[5] = -energyScale_ * f[0] * coordinates[idx][1];
464             if (isComputeParticleVirial == true)
465             {
466               particleVirial[idx][0] += v[0];
467               particleVirial[idx][1] += v[1];
468               particleVirial[idx][2] += v[2];
469               particleVirial[idx][3] += v[3];
470               particleVirial[idx][4] += v[4];
471               particleVirial[idx][5] += v[5];
472             }
473             if (isComputeVirial == true)
474             {
475               virial[0] += v[0];
476               virial[1] += v[1];
477               virial[2] += v[2];
478               virial[3] += v[3];
479               virial[4] += v[4];
480               virial[5] += v[5];
481             }
482           }
483         }
484       }
485     }
486 
487     // deallocate memory
488     Deallocate1DArray(GC);
489     if (need_dE) { Deallocate2DArray(dGCdr); }
490     if (need_dE && (ensemble_size_ != 0) && (active_member_id_ == -1)) {
491       Deallocate1DArray(dEdGC);
492     }
493 
494   }  // loop over i
495 
496   // everything is good
497   ier = false;
498   return ier;
499 }
500 
501 #endif  // ANN_IMPLEMENTATION_HPP_
502