Doubt regarding DeePMD Workflow

[AnuragKr]

Hello,

DeePMD Workflow Diagram--

Ref: - https://arxiv.org/abs/1805.09003

First of all I would like to describe my brief rough summary of my understanding of DeePMD-kit working without type embedding in terms of dimension transformation --
Input --
Coord = Nframes * Natoms * 3
Ener = Nframes
Key Steps --

1. Environment matrix Creation -
   -- Input -> Generalized Coordinates
   Nframes * Natoms * 3 -> Nframes * Natoms * Nei_atoms * 4
2. Descriptor Formulation
   -- Generalized Coordinates -> Gi1
   Nframes * Natoms * Nei_atoms * 4 -> Nframes * Natoms * Nei_atoms * M1
   --Gi1 -> Gi2
   Nframes * Natoms * Nei_atoms * M1 -> Nframes * Natoms * Nei_atoms * M2
   Descriptor = (Gi1^TxR)x(R^TxGi2)
   = (M1 * Nei_atoms x Nei_atoms * 4) x (4 * Nei_atoms x Nei_atoms * M2)
   = (M1 * M2)
   = Nframes * Natoms * M1 * M2
3. Fitting Network
   Descriptor -> Energy
   Nframes * Natoms * M1 * M2 -> Nframes * Natoms * 1 -> Nframes * 1

Summary -- First local environment of every atom in a frame get converted into environment matrix and then into descriptor. After that energy is calculated for every atom in a frame and then sum it up to get total energy of a frame.

Doubt -- 1. Does R Correspond to one Frame atom coordinates in terms of real data ?
2. In a codebase where local energy is being converted into total frame energy ? If I am correct in
EnerFitting class it returns local energy for every atoms.
3. Why we are calculating energy based on atoms type ? Is it because of memory constraint ?

Hoping for positive response.

[AnuragKr]

My Findings ---

1. Yes R corresponds to one frame atom coordinates.
2. In the fitting network itself every atom contribution is summed up in a final layer.
3. No it is not because of memory constraints if we have not selected the Atom type embedding approach then each atom type will have its own Multilayer Perceptron (MLP) which means weight can't be shared across atom types and then in the end we do a summation of energy calculated from atom types to derive total energy.

I am closing this discussion with my findings if someone has other reasons, Please share your comment.