Skip to content

Commit 7312f58

Browse files
author
Han Wang
committed
norm descrpt, ener bias for multicomponent system, add type fitting net for smooth version. continue when nei list length is not enough
1 parent 021ed25 commit 7312f58

File tree

8 files changed

+419
-131
lines changed

8 files changed

+419
-131
lines changed

‎examples/train/water.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
"_comment": " if type < 0, exclude type -(type+1)",
1212
"_comment": " for water (O:0, H:1) it can be",
1313
"_comment": " [0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0]",
14-
"n_neuron": [240, 120, 60, 30, 10],
14+
"fitting_neuron": [240, 120, 60, 30, 10],
1515

1616
"_comment": " traing controls",
1717
"systems": ["../data/water/"],

‎examples/train/water_smth.json‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@
66
"rcut": 6.00,
77
"filter_neuron": [25, 50, 100],
88
"filter_resnet_dt": false,
9-
"n_axis_neuron": 16,
10-
"n_neuron": [240, 240, 240],
11-
"resnet_dt": true,
9+
"axis_neuron": 16,
10+
"fitting_neuron": [240, 240, 240],
11+
"fitting_resnet_dt":true,
12+
"coord_norm": true,
13+
"type_fitting_net": false,
1214

1315
"_comment": " traing controls",
1416
"systems": ["../data/water/"],

‎source/lib/include/ComputeDescriptor.h‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -335,21 +335,25 @@ int format_nlist_fill_a (vector<int > & fmt_nei_idx_a,
335335
sort (sel_nei.begin(), sel_nei.end());
336336

337337
vector<int > nei_iter = sec_a;
338+
int overflowed = -1;
338339
for (unsigned kk = 0; kk < sel_nei.size(); ++kk){
339340
const int & nei_type = sel_nei[kk].type;
340341
if (nei_iter[nei_type] >= sec_a[nei_type+1]) {
341342
int r_idx_iter = (nei_iter[nei_type] ++) - sec_a[nei_type+1] + sec_r[nei_type];
342343
if (r_idx_iter >= sec_r[nei_type+1]) {
343-
return nei_type;
344+
// return nei_type;
345+
overflowed = nei_type;
346+
}
347+
else {
348+
fmt_nei_idx_r[r_idx_iter] = sel_nei[kk].index;
344349
}
345-
fmt_nei_idx_r[r_idx_iter] = sel_nei[kk].index;
346350
}
347351
else {
348352
fmt_nei_idx_a[nei_iter[nei_type] ++] = sel_nei[kk].index;
349353
}
350354
}
351355

352-
return -1;
356+
return overflowed;
353357
}
354358

355359

‎source/op/descrpt.cc‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ class DescrptOp : public OpKernel {
7575
nnei_r = sec_r.back();
7676
nnei = nnei_a + nnei_r;
7777
fill_nei_a = (rcut_a < 0);
78+
count_nei_idx_overflow = 0;
7879
}
7980

8081
void Compute(OpKernelContext* context) override {
@@ -93,24 +94,27 @@ class DescrptOp : public OpKernel {
9394
OP_REQUIRES (context, (natoms_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of natoms should be 1"));
9495
OP_REQUIRES (context, (box_tensor.shape().dims() == 2), errors::InvalidArgument ("Dim of box should be 2"));
9596
OP_REQUIRES (context, (mesh_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of mesh should be 1"));
96-
OP_REQUIRES (context, (avg_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of avg should be 1"));
97-
OP_REQUIRES (context, (std_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of std should be 1"));
97+
OP_REQUIRES (context, (avg_tensor.shape().dims() == 2), errors::InvalidArgument ("Dim of avg should be 2"));
98+
OP_REQUIRES (context, (std_tensor.shape().dims() == 2), errors::InvalidArgument ("Dim of std should be 2"));
9899

99100
OP_REQUIRES (context, (natoms_tensor.shape().dim_size(0) >= 3), errors::InvalidArgument ("number of atoms should be larger than (or equal to) 3"));
100101
auto natoms = natoms_tensor .flat<int>();
101102
int nloc = natoms(0);
102103
int nall = natoms(1);
104+
int ntypes = natoms_tensor.shape().dim_size(0) - 2;
103105
int nsamples = coord_tensor.shape().dim_size(0);
104106

105107
// check the sizes
106108
OP_REQUIRES (context, (nsamples == type_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of samples should match"));
107109
OP_REQUIRES (context, (nsamples == box_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of samples should match"));
108-
OP_REQUIRES (context, (ndescrpt == avg_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of avg should be ndescrpt"));
109-
OP_REQUIRES (context, (ndescrpt == std_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of std should be ndescrpt"));
110+
OP_REQUIRES (context, (ntypes == avg_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of avg should be ntype"));
111+
OP_REQUIRES (context, (ntypes == std_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of std should be ntype"));
110112

111113
OP_REQUIRES (context, (nall * 3 == coord_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of atoms should match"));
112114
OP_REQUIRES (context, (nall == type_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of atoms should match"));
113115
OP_REQUIRES (context, (9 == box_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of box should be 9"));
116+
OP_REQUIRES (context, (ndescrpt == avg_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of avg should be ndescrpt"));
117+
OP_REQUIRES (context, (ndescrpt == std_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of std should be ndescrpt"));
114118

115119
int nei_mode = 0;
116120
if (mesh_tensor.shape().dim_size(0) == 16) {
@@ -156,8 +160,8 @@ class DescrptOp : public OpKernel {
156160
auto type = type_tensor .matrix<int>();
157161
auto box = box_tensor .matrix<VALUETYPE>();
158162
auto mesh = mesh_tensor .flat<int>();
159-
auto avg = avg_tensor .flat<VALUETYPE>();
160-
auto std = std_tensor .flat<VALUETYPE>();
163+
auto avg = avg_tensor .matrix<VALUETYPE>();
164+
auto std = std_tensor .matrix<VALUETYPE>();
161165
auto descrpt = descrpt_tensor ->matrix<VALUETYPE>();
162166
auto descrpt_deriv = descrpt_deriv_tensor ->matrix<VALUETYPE>();
163167
auto rij = rij_tensor ->matrix<VALUETYPE>();
@@ -170,7 +174,6 @@ class DescrptOp : public OpKernel {
170174
// if (type(0, ii) > max_type_v) max_type_v = type(0, ii);
171175
// }
172176
// int ntypes = max_type_v + 1;
173-
int ntypes = natoms_tensor.shape().dim_size(0) - 2;
174177
OP_REQUIRES (context, (ntypes == int(sel_a.size())), errors::InvalidArgument ("number of types should match the length of sel array"));
175178
OP_REQUIRES (context, (ntypes == int(sel_r.size())), errors::InvalidArgument ("number of types should match the length of sel array"));
176179

@@ -270,8 +273,11 @@ class DescrptOp : public OpKernel {
270273
int ret = -1;
271274
if (fill_nei_a){
272275
if ((ret = format_nlist_fill_a (fmt_nlist_a, fmt_nlist_r, d_coord3, ntypes, d_type, region, b_pbc, ii, d_nlist_a[ii], d_nlist_r[ii], rcut_r, sec_a, sec_r)) != -1){
273-
cout << "Radial neighbor list length of type " << ret << " is not enough" << endl;
274-
exit(1);
276+
if (count_nei_idx_overflow == 0) {
277+
cout << "WARNING: Radial neighbor list length of type " << ret << " is not enough" << endl;
278+
flush(cout);
279+
count_nei_idx_overflow ++;
280+
}
275281
}
276282
}
277283

@@ -322,16 +328,16 @@ class DescrptOp : public OpKernel {
322328
assert (int(fmt_nlist_r.size()) == nnei_r);
323329
// record outputs
324330
for (int jj = 0; jj < ndescrpt_a; ++jj) {
325-
descrpt(kk, ii * ndescrpt + jj) = (d_descrpt_a[jj] - avg(jj)) / std(jj);
331+
descrpt(kk, ii * ndescrpt + jj) = (d_descrpt_a[jj] - avg(d_type[ii], jj)) / std(d_type[ii], jj);
326332
}
327333
for (int jj = 0; jj < ndescrpt_r; ++jj) {
328-
descrpt(kk, ii * ndescrpt + ndescrpt_a + jj) = (d_descrpt_r[jj] - avg(ndescrpt_a + jj)) / std(ndescrpt_a + jj);
334+
descrpt(kk, ii * ndescrpt + ndescrpt_a + jj) = (d_descrpt_r[jj] - avg(d_type[ii], ndescrpt_a + jj)) / std(d_type[ii], ndescrpt_a + jj);
329335
}
330336
for (int jj = 0; jj < ndescrpt_a * 12; ++jj) {
331-
descrpt_deriv(kk, ii * ndescrpt * 12 + jj) = d_descrpt_a_deriv[jj] / std(jj/12);
337+
descrpt_deriv(kk, ii * ndescrpt * 12 + jj) = d_descrpt_a_deriv[jj] / std(d_type[ii], jj/12);
332338
}
333339
for (int jj = 0; jj < ndescrpt_r * 12; ++jj) {
334-
descrpt_deriv(kk, ii * ndescrpt * 12 + ndescrpt_a * 12 + jj) = d_descrpt_r_deriv[jj] / std(jj/12 + ndescrpt_a);
340+
descrpt_deriv(kk, ii * ndescrpt * 12 + ndescrpt_a * 12 + jj) = d_descrpt_r_deriv[jj] / std(d_type[ii], jj/12 + ndescrpt_a);
335341
}
336342
for (int jj = 0; jj < nnei_a * 3; ++jj){
337343
rij (kk, ii * nnei * 3 + jj) = d_rij_a[jj];
@@ -371,6 +377,7 @@ class DescrptOp : public OpKernel {
371377
int ndescrpt, ndescrpt_a, ndescrpt_r;
372378
int nnei, nnei_a, nnei_r;
373379
bool fill_nei_a;
380+
int count_nei_idx_overflow;
374381
void
375382
cum_sum (vector<int> & sec,
376383
const vector<int32> & n_sel) const {

‎source/op/descrpt_norot.cc‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ class DescrptNorotOp : public OpKernel {
7373
nnei_r = sec_r.back();
7474
nnei = nnei_a + nnei_r;
7575
fill_nei_a = (rcut_a < 0);
76+
count_nei_idx_overflow = 0;
7677
}
7778

7879
void Compute(OpKernelContext* context) override {
@@ -92,26 +93,29 @@ class DescrptNorotOp : public OpKernel {
9293
OP_REQUIRES (context, (natoms_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of natoms should be 1"));
9394
OP_REQUIRES (context, (box_tensor.shape().dims() == 2), errors::InvalidArgument ("Dim of box should be 2"));
9495
OP_REQUIRES (context, (mesh_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of mesh should be 1"));
95-
OP_REQUIRES (context, (avg_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of avg should be 1"));
96-
OP_REQUIRES (context, (std_tensor.shape().dims() == 1), errors::InvalidArgument ("Dim of std should be 1"));
96+
OP_REQUIRES (context, (avg_tensor.shape().dims() == 2), errors::InvalidArgument ("Dim of avg should be 2"));
97+
OP_REQUIRES (context, (std_tensor.shape().dims() == 2), errors::InvalidArgument ("Dim of std should be 2"));
9798
OP_REQUIRES (context, (fill_nei_a), errors::InvalidArgument ("Rotational free descriptor only support the case rcut_a < 0"));
9899
OP_REQUIRES (context, (sec_r.back() == 0), errors::InvalidArgument ("Rotational free descriptor only support all-angular information: sel_r should be all zero."));
99100

100101
OP_REQUIRES (context, (natoms_tensor.shape().dim_size(0) >= 3), errors::InvalidArgument ("number of atoms should be larger than (or equal to) 3"));
101102
auto natoms = natoms_tensor .flat<int>();
102103
int nloc = natoms(0);
103104
int nall = natoms(1);
105+
int ntypes = natoms_tensor.shape().dim_size(0) - 2;
104106
int nsamples = coord_tensor.shape().dim_size(0);
105107

106108
// check the sizes
107109
OP_REQUIRES (context, (nsamples == type_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of samples should match"));
108110
OP_REQUIRES (context, (nsamples == box_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of samples should match"));
109-
OP_REQUIRES (context, (ndescrpt == avg_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of avg should be ndescrpt"));
110-
OP_REQUIRES (context, (ndescrpt == std_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of std should be ndescrpt"));
111+
OP_REQUIRES (context, (ntypes == avg_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of avg should be ntype"));
112+
OP_REQUIRES (context, (ntypes == std_tensor.shape().dim_size(0)), errors::InvalidArgument ("number of std should be ntype"));
111113

112114
OP_REQUIRES (context, (nall * 3 == coord_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of atoms should match"));
113115
OP_REQUIRES (context, (nall == type_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of atoms should match"));
114116
OP_REQUIRES (context, (9 == box_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of box should be 9"));
117+
OP_REQUIRES (context, (ndescrpt == avg_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of avg should be ndescrpt"));
118+
OP_REQUIRES (context, (ndescrpt == std_tensor.shape().dim_size(1)), errors::InvalidArgument ("number of std should be ndescrpt"));
115119

116120
int nei_mode = 0;
117121
if (mesh_tensor.shape().dim_size(0) == 16) {
@@ -161,8 +165,8 @@ class DescrptNorotOp : public OpKernel {
161165
auto type = type_tensor .matrix<int>();
162166
auto box = box_tensor .matrix<VALUETYPE>();
163167
auto mesh = mesh_tensor .flat<int>();
164-
auto avg = avg_tensor .flat<VALUETYPE>();
165-
auto std = std_tensor .flat<VALUETYPE>();
168+
auto avg = avg_tensor .matrix<VALUETYPE>();
169+
auto std = std_tensor .matrix<VALUETYPE>();
166170
auto descrpt = descrpt_tensor ->matrix<VALUETYPE>();
167171
auto descrpt_deriv = descrpt_deriv_tensor ->matrix<VALUETYPE>();
168172
auto rij = rij_tensor ->matrix<VALUETYPE>();
@@ -174,7 +178,6 @@ class DescrptNorotOp : public OpKernel {
174178
// if (type(0, ii) > max_type_v) max_type_v = type(0, ii);
175179
// }
176180
// int ntypes = max_type_v + 1;
177-
int ntypes = natoms_tensor.shape().dim_size(0) - 2;
178181
OP_REQUIRES (context, (ntypes == int(sel_a.size())), errors::InvalidArgument ("number of types should match the length of sel array"));
179182
OP_REQUIRES (context, (ntypes == int(sel_r.size())), errors::InvalidArgument ("number of types should match the length of sel array"));
180183

@@ -274,8 +277,11 @@ class DescrptNorotOp : public OpKernel {
274277
int ret = -1;
275278
if (fill_nei_a){
276279
if ((ret = format_nlist_fill_a (fmt_nlist_a, fmt_nlist_r, d_coord3, ntypes, d_type, region, b_pbc, ii, d_nlist_a[ii], d_nlist_r[ii], rcut_r, sec_a, sec_r)) != -1){
277-
cout << "Radial neighbor list length of type " << ret << " is not enough" << endl;
278-
exit(1);
280+
if (count_nei_idx_overflow == 0) {
281+
cout << "WARNING: Radial neighbor list length of type " << ret << " is not enough" << endl;
282+
flush(cout);
283+
count_nei_idx_overflow ++;
284+
}
279285
}
280286
}
281287

@@ -306,10 +312,10 @@ class DescrptNorotOp : public OpKernel {
306312
assert (int(fmt_nlist_a.size()) == nnei_a);
307313
// record outputs
308314
for (int jj = 0; jj < ndescrpt_a; ++jj) {
309-
descrpt(kk, ii * ndescrpt + jj) = (d_descrpt_a[jj] - avg(jj)) / std(jj);
315+
descrpt(kk, ii * ndescrpt + jj) = (d_descrpt_a[jj] - avg(d_type[ii], jj)) / std(d_type[ii], jj);
310316
}
311317
for (int jj = 0; jj < ndescrpt_a * 3; ++jj) {
312-
descrpt_deriv(kk, ii * ndescrpt * 3 + jj) = d_descrpt_a_deriv[jj] / std(jj/3);
318+
descrpt_deriv(kk, ii * ndescrpt * 3 + jj) = d_descrpt_a_deriv[jj] / std(d_type[ii], jj/3);
313319
}
314320
for (int jj = 0; jj < nnei_a * 3; ++jj){
315321
rij (kk, ii * nnei * 3 + jj) = d_rij_a[jj];
@@ -335,6 +341,7 @@ class DescrptNorotOp : public OpKernel {
335341
int ndescrpt, ndescrpt_a, ndescrpt_r;
336342
int nnei, nnei_a, nnei_r;
337343
bool fill_nei_a;
344+
int count_nei_idx_overflow;
338345
void
339346
cum_sum (vector<int> & sec,
340347
const vector<int32> & n_sel) const {

‎source/train/Data.py‎

Lines changed: 13 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -10,27 +10,20 @@ class DataSets (object):
1010
def __init__ (self,
1111
sys_path,
1212
set_prefix,
13-
do_norm = False,
1413
seed = None) :
15-
self.dirs = glob.glob (sys_path + "/" + set_prefix + ".*")
14+
self.dirs = glob.glob (os.path.join(sys_path, set_prefix + ".*"))
1615
self.dirs.sort()
17-
if os.path.isfile (sys_path + "/ncopies.raw") :
18-
self.ncopies = np.loadtxt (sys_path + "/ncopies.raw", dtype=np.int32)
19-
else :
20-
self.ncopies = [1]
21-
2216
# load atom type
2317
self.atom_type, self.idx_map, self.idx3_map = self.load_type (sys_path)
24-
25-
self.do_norm = do_norm
26-
self.eavg = self.stats_energy()
27-
# print ("avg is ",self.eavg)
28-
18+
# train dirs
2919
self.test_dir = self.dirs[-1]
3020
if len(self.dirs) == 1 :
3121
self.train_dirs = self.dirs
3222
else :
3323
self.train_dirs = self.dirs[:-1]
24+
# energy norm
25+
self.eavg = self.stats_energy()
26+
# load sets
3427
self.set_count = 0
3528
self.load_batch_set (self.train_dirs[self.set_count % self.get_numb_set()])
3629
self.load_test_set (self.test_dir)
@@ -65,11 +58,11 @@ def get_numb_set (self) :
6558
return len (self.train_dirs)
6659

6760
def stats_energy (self) :
68-
eners = []
69-
for ii in self.dirs:
70-
ei = np.load (ii + "/energy.npy")
71-
eners.append (np.average(ei))
72-
return np.average (eners)
61+
eners = np.array([])
62+
for ii in self.train_dirs:
63+
ei = np.load(os.path.join(ii, "energy.npy"))
64+
eners = np.append(eners, ei)
65+
return np.average(eners)
7366

7467
def cond_load_vec (self,
7568
nframes,
@@ -119,7 +112,6 @@ def load_batch_set (self,
119112
self.box_batch = self.box_batch[idx]
120113
self.type_batch = np.tile (self.atom_type, (nframe, 1))
121114
self.reset_iter ()
122-
if self.do_norm: self.normalization (self.energy_batch)
123115
# sort according to type
124116
self.type_batch = self.type_batch[:, self.idx_map]
125117
self.coord_batch = self.coord_batch[:, self.idx3_map]
@@ -146,7 +138,6 @@ def load_test_set (self,
146138
self.coord_test = self.coord_test[idx]
147139
self.box_test = self.box_test[idx]
148140
self.type_test = np.tile (self.atom_type, (nframe, 1))
149-
if self.do_norm: self.normalization (self.energy_test)
150141
self.type_test = self.type_test[:, self.idx_map]
151142
self.coord_test = self.coord_test[:, self.idx3_map]
152143
self.force_test = self.force_test[:, self.idx3_map]
@@ -174,12 +165,6 @@ def get_batch (self,
174165
self.iterator += batch_size
175166
return self.prop_c_batch.astype(np.float32), self.energy_batch[idx].astype(np.float64), self.force_batch[idx, :].astype(np.float64), self.virial_batch[idx, :].astype(np.float64), self.coord_batch[idx, :].astype(np.float64), self.box_batch[idx, :].astype(np.float64), self.type_batch[idx, :]
176167

177-
def normalization (self,
178-
energy) :
179-
if self.do_norm :
180-
for ii in range (energy.shape[0]) :
181-
energy[ii] -= self.eavg
182-
183168
def get_natoms (self) :
184169
sample_type = self.type_batch[0]
185170
natoms = len(sample_type)
@@ -199,22 +184,18 @@ def get_natoms_vec (self, ntypes) :
199184
tmp = np.append (tmp, natoms_vec)
200185
return tmp.astype(np.int32)
201186

202-
def get_ncopies (self) :
203-
return self.ncopies
204-
205187
def set_numb_batch (self,
206188
batch_size) :
207189
return self.energy_batch.shape[0] // batch_size
208190

209191
def get_sys_numb_batch (self, batch_size) :
210192
return self.set_numb_batch(batch_size) * self.get_numb_set()
211193

212-
def get_bias_atom_e (self) :
213-
natoms = self.get_natoms ()
214-
return self.eavg / natoms
194+
def get_ener (self) :
195+
return self.eavg
215196

216197
if __name__ == '__main__':
217-
data = DataSets (".", "set", do_norm = False)
198+
data = DataSets (".", "set")
218199
energy, force, virial, coord, box, ttype = data.get_batch(1)
219200
print (energy.shape)
220201
print (force.shape)

0 commit comments

Comments
 (0)