|
| 1 | +# SPDX-License-Identifier: LGPL-3.0-or-later |
| 2 | +import unittest |
| 3 | + |
| 4 | +import numpy as np |
| 5 | + |
| 6 | +from deepmd.pt.model.descriptor import ( |
| 7 | + DescrptSeA, |
| 8 | +) |
| 9 | +from deepmd.pt.model.task import ( |
| 10 | + EnergyFittingNet, |
| 11 | +) |
| 12 | +from deepmd.pt.utils.utils import ( |
| 13 | + to_numpy_array, |
| 14 | + to_torch_tensor, |
| 15 | +) |
| 16 | + |
| 17 | + |
| 18 | +def _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds): |
| 19 | + merged_output_stat = [] |
| 20 | + nsys = len(sys_natoms) |
| 21 | + ndof = len(avgs) |
| 22 | + for ii in range(nsys): |
| 23 | + sys_dict = {} |
| 24 | + tmp_data_f = [] |
| 25 | + tmp_data_a = [] |
| 26 | + for jj in range(ndof): |
| 27 | + rng = np.random.default_rng(2025 * ii + 220 * jj) |
| 28 | + tmp_data_f.append( |
| 29 | + rng.normal(loc=avgs[jj], scale=stds[jj], size=(sys_nframes[ii], 1)) |
| 30 | + ) |
| 31 | + rng = np.random.default_rng(220 * ii + 1636 * jj) |
| 32 | + tmp_data_a.append( |
| 33 | + rng.normal( |
| 34 | + loc=avgs[jj], scale=stds[jj], size=(sys_nframes[ii], sys_natoms[ii]) |
| 35 | + ) |
| 36 | + ) |
| 37 | + tmp_data_f = np.transpose(tmp_data_f, (1, 2, 0)) |
| 38 | + tmp_data_a = np.transpose(tmp_data_a, (1, 2, 0)) |
| 39 | + sys_dict["fparam"] = to_torch_tensor(tmp_data_f) |
| 40 | + sys_dict["aparam"] = to_torch_tensor(tmp_data_a) |
| 41 | + merged_output_stat.append(sys_dict) |
| 42 | + return merged_output_stat |
| 43 | + |
| 44 | + |
| 45 | +def _brute_fparam_pt(data, ndim): |
| 46 | + adata = [to_numpy_array(ii["fparam"]) for ii in data] |
| 47 | + all_data = [] |
| 48 | + for ii in adata: |
| 49 | + tmp = np.reshape(ii, [-1, ndim]) |
| 50 | + if len(all_data) == 0: |
| 51 | + all_data = np.array(tmp) |
| 52 | + else: |
| 53 | + all_data = np.concatenate((all_data, tmp), axis=0) |
| 54 | + avg = np.average(all_data, axis=0) |
| 55 | + std = np.std(all_data, axis=0) |
| 56 | + return avg, std |
| 57 | + |
| 58 | + |
| 59 | +def _brute_aparam_pt(data, ndim): |
| 60 | + adata = [to_numpy_array(ii["aparam"]) for ii in data] |
| 61 | + all_data = [] |
| 62 | + for ii in adata: |
| 63 | + tmp = np.reshape(ii, [-1, ndim]) |
| 64 | + if len(all_data) == 0: |
| 65 | + all_data = np.array(tmp) |
| 66 | + else: |
| 67 | + all_data = np.concatenate((all_data, tmp), axis=0) |
| 68 | + avg = np.average(all_data, axis=0) |
| 69 | + std = np.std(all_data, axis=0) |
| 70 | + return avg, std |
| 71 | + |
| 72 | + |
| 73 | +class TestEnerFittingStat(unittest.TestCase): |
| 74 | + def test(self) -> None: |
| 75 | + descrpt = DescrptSeA(6.0, 5.8, [46, 92], neuron=[25, 50, 100], axis_neuron=16) |
| 76 | + fitting = EnergyFittingNet( |
| 77 | + descrpt.get_ntypes(), |
| 78 | + descrpt.get_dim_out(), |
| 79 | + neuron=[240, 240, 240], |
| 80 | + resnet_dt=True, |
| 81 | + numb_fparam=3, |
| 82 | + numb_aparam=3, |
| 83 | + ) |
| 84 | + avgs = [0, 10, 100] |
| 85 | + stds = [2, 0.4, 0.00001] |
| 86 | + sys_natoms = [10, 100] |
| 87 | + sys_nframes = [5, 2] |
| 88 | + all_data = _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds) |
| 89 | + frefa, frefs = _brute_fparam_pt(all_data, len(avgs)) |
| 90 | + arefa, arefs = _brute_aparam_pt(all_data, len(avgs)) |
| 91 | + fitting.compute_input_stats(all_data, protection=1e-2) |
| 92 | + frefs_inv = 1.0 / frefs |
| 93 | + arefs_inv = 1.0 / arefs |
| 94 | + frefs_inv[frefs_inv > 100] = 100 |
| 95 | + arefs_inv[arefs_inv > 100] = 100 |
| 96 | + np.testing.assert_almost_equal(frefa, to_numpy_array(fitting.fparam_avg)) |
| 97 | + np.testing.assert_almost_equal( |
| 98 | + frefs_inv, to_numpy_array(fitting.fparam_inv_std) |
| 99 | + ) |
| 100 | + np.testing.assert_almost_equal(arefa, to_numpy_array(fitting.aparam_avg)) |
| 101 | + np.testing.assert_almost_equal( |
| 102 | + arefs_inv, to_numpy_array(fitting.aparam_inv_std) |
| 103 | + ) |
0 commit comments