diff --git a/mindspore/nn/probability/toolbox/uncertainty_evaluation.py b/mindspore/nn/probability/toolbox/uncertainty_evaluation.py index eb6e24c1d5..00644d1fef 100644 --- a/mindspore/nn/probability/toolbox/uncertainty_evaluation.py +++ b/mindspore/nn/probability/toolbox/uncertainty_evaluation.py @@ -87,7 +87,7 @@ class UncertaintyEvaluation: self.epi_model = deepcopy(model) self.ale_model = deepcopy(model) self.epi_train_dataset = train_dataset - self.ale_train_dataset = train_dataset + self.ale_train_dataset = deepcopy(train_dataset) self.task_type = task_type self.epochs = Validator.check_positive_int(epochs) self.epi_uncer_model_path = epi_uncer_model_path @@ -101,7 +101,8 @@ class UncertaintyEvaluation: if not isinstance(model, Cell): raise TypeError('The model should be Cell type.') if task_type not in ('regression', 'classification'): - raise ValueError('The task should be regression or classification.') + raise ValueError( + 'The task should be regression or classification.') if task_type == 'classification': self.num_classes = Validator.check_positive_int(num_classes) else: @@ -119,15 +120,19 @@ class UncertaintyEvaluation: self.epi_uncer_model = EpistemicUncertaintyModel(self.epi_model) if self.epi_uncer_model.drop_count == 0 and self.epi_train_dataset is not None: if self.task_type == 'classification': - net_loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") + net_loss = SoftmaxCrossEntropyWithLogits( + sparse=True, reduction="mean") net_opt = Adam(self.epi_uncer_model.trainable_params()) - model = Model(self.epi_uncer_model, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) + model = Model(self.epi_uncer_model, net_loss, + net_opt, metrics={"Accuracy": Accuracy()}) else: net_loss = MSELoss() net_opt = Adam(self.epi_uncer_model.trainable_params()) - model = Model(self.epi_uncer_model, net_loss, net_opt, metrics={"MSE": MSE()}) + model = Model(self.epi_uncer_model, net_loss, + net_opt, metrics={"MSE": MSE()}) if self.save_model: - config_ck = CheckpointConfig(keep_checkpoint_max=self.epochs) + config_ck = CheckpointConfig( + keep_checkpoint_max=self.epochs) ckpoint_cb = ModelCheckpoint(prefix='checkpoint_epi_uncer_model', directory=self.epi_uncer_model_path, config=config_ck) @@ -137,7 +142,8 @@ class UncertaintyEvaluation: model.train(self.epochs, self.epi_train_dataset, dataset_sink_mode=False, callbacks=[LossMonitor()]) else: - uncer_param_dict = load_checkpoint(self.epi_uncer_model_path) + uncer_param_dict = load_checkpoint( + self.epi_uncer_model_path) load_param_into_net(self.epi_uncer_model, uncer_param_dict) def _eval_epistemic_uncertainty(self, eval_data, mc=10): @@ -164,15 +170,19 @@ class UncertaintyEvaluation: Get the model which can obtain the aleatoric uncertainty. """ if self.ale_train_dataset is None: - raise ValueError('The train dataset should not be None when evaluating aleatoric uncertainty.') + raise ValueError( + 'The train dataset should not be None when evaluating aleatoric uncertainty.') if self.ale_uncer_model is None: - self.ale_uncer_model = AleatoricUncertaintyModel(self.ale_model, self.num_classes, self.task_type) + self.ale_uncer_model = AleatoricUncertaintyModel( + self.ale_model, self.num_classes, self.task_type) net_loss = AleatoricLoss(self.task_type) net_opt = Adam(self.ale_uncer_model.trainable_params()) if self.task_type == 'classification': - model = Model(self.ale_uncer_model, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) + model = Model(self.ale_uncer_model, net_loss, + net_opt, metrics={"Accuracy": Accuracy()}) else: - model = Model(self.ale_uncer_model, net_loss, net_opt, metrics={"MSE": MSE()}) + model = Model(self.ale_uncer_model, net_loss, + net_opt, metrics={"MSE": MSE()}) if self.save_model: config_ck = CheckpointConfig(keep_checkpoint_max=self.epochs) ckpoint_cb = ModelCheckpoint(prefix='checkpoint_ale_uncer_model', @@ -284,7 +294,8 @@ class AleatoricUncertaintyModel(Cell): self.ale_model = ale_model self.var_layer = Dense(num_classes, num_classes) else: - self.ale_model, self.var_layer, self.pred_layer = self._make_aleatoric(ale_model) + self.ale_model, self.var_layer, self.pred_layer = self._make_aleatoric( + ale_model) def construct(self, x): if self.task == 'classification': @@ -327,7 +338,8 @@ class AleatoricLoss(Cell): self.exp = P.Exp() self.normal = C.normal self.to_tensor = P.ScalarToArray() - self.entropy = SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") + self.entropy = SoftmaxCrossEntropyWithLogits( + sparse=True, reduction="mean") else: self.mean = P.ReduceMean() self.exp = P.Exp() @@ -337,7 +349,8 @@ class AleatoricLoss(Cell): y_pred, var = data_pred if self.task == 'classification': sample_times = 10 - epsilon = self.normal((1, sample_times), self.to_tensor(0.0), self.to_tensor(1.0), 0) + epsilon = self.normal((1, sample_times), self.to_tensor( + 0.0), self.to_tensor(1.0), 0) total_loss = 0 for i in range(sample_times): y_pred_i = y_pred + epsilon[0][i] * var @@ -345,5 +358,6 @@ class AleatoricLoss(Cell): total_loss += loss avg_loss = total_loss / sample_times return avg_loss - loss = self.mean(0.5 * self.exp(-var) * self.pow(y - y_pred, 2) + 0.5 * var) + loss = self.mean(0.5 * self.exp(-var) * + self.pow(y - y_pred, 2) + 0.5 * var) return loss diff --git a/mindspore/ops/operations/sponge_ops.py b/mindspore/ops/operations/sponge_ops.py index d25e74553b..bc436e8384 100644 --- a/mindspore/ops/operations/sponge_ops.py +++ b/mindspore/ops/operations/sponge_ops.py @@ -34,13 +34,14 @@ class BondForce(PrimitiveWithInfer): dr = (x_1-x_2, y_1-y_2, z_1-z_2) F = (F_x, F_y, F_z) = 2*k*(1 - r_0/|dr|)*dr + Args: + atom_numbers(int32): the number of atoms N. + bond_numbers(int32): the number of harmonic bonds M. + Inputs: - - **atom_numbers** (int32) - the number of atoms N. - - **bond_numbers** (int32) - the number of harmonic bonds M. - - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate - value of each atom. + - **uint_crd_f** (Tensor, uint32 ) - [N, 3], the unsigned int coordinate value of each atom. - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor (x, y, z), - between the real space float coordinates and the unsigned int coordinates. + between the real space float coordinates and the unsigned int coordinates. - **atom_a** (Tensor, int32) - [M,], the first atom index of each bond. - **atom_b** (Tensor, int32) - [M,], the second atom index of each bond. - **bond_k** (Tensor, float32) - [M,], the force constant of each bond. @@ -51,11 +52,12 @@ class BondForce(PrimitiveWithInfer): Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, bond_numbers, atom_numbers): + assert isinstance(bond_numbers, int) + assert isinstance(atom_numbers, int) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -67,28 +69,28 @@ class BondForce(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0", cls_name) return uint_crd_f_shape def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) return bond_r0_type class BondEnergy(PrimitiveWithInfer): """ - BondEnergyCuda: + BondEnergy: Calculate the harmonic potential energy between each bonded atom pair. Assume our system has N atoms and M harmonic bonds. @@ -98,25 +100,28 @@ class BondEnergy(PrimitiveWithInfer): dr = (x_1-x_2, y_1-y_2, z_1-z_2) E = k*(|dr| - r_0)^2 + Args: + Same as operator BondForce(). + Inputs: Same as operator BondForce(). - .. math:: + .. math:: - dr = (x_1-x_2, y_1-y_2, z_1-z_2) - E = k*(|dr| - r_0)^2 + dr = (x_1-x_2, y_1-y_2, z_1-z_2) + E = k*(|dr| - r_0)^2 Outputs: - - **bond_ene** (Tensor, float32) - [M,], the harmonic potential energy - for each bond. + - **bond_ene** (Tensor, float32) - [M,], the harmonic potential energy for each bond. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, bond_numbers, atom_numbers): + assert isinstance(bond_numbers, int) + assert isinstance(atom_numbers, int) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -128,49 +133,52 @@ class BondEnergy(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0", cls_name) return bond_k_shape def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) return bond_r0_type class BondAtomEnergy(PrimitiveWithInfer): """ - BondAtomEnergyCuda: + BondAtomEnergy: Add the potential energy caused by simple harmonic bonds to the total potential energy of each atom. The calculation formula is the same as operator BondEnergy(). + Args: + Same as operator BondForce(). + Inputs: Same as operator BondForce(). Outputs: - - **atom_ene** (Tensor, float32) - [N,], the accumulated potential - energy for each atom. + - **atom_ene** (Tensor, float32) - [N,], the accumulated potential energy for each atom. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, bond_numbers, atom_numbers): + assert isinstance(bond_numbers, int) + assert isinstance(atom_numbers, int) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -182,23 +190,22 @@ class BondAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) - + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0", cls_name) return [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) return bond_r0_type @@ -210,6 +217,9 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator BondForce() and BondEnergy(). + Args: + Same as operator BondForce(). + Inputs: Same as operator BondForce(). @@ -219,11 +229,12 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer): Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, bond_numbers, atom_numbers): + assert isinstance(bond_numbers, int) + assert isinstance(atom_numbers, int) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -235,25 +246,25 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0", cls_name) return uint_crd_f_shape, [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) return bond_r0_type, bond_r0_type @@ -272,21 +283,24 @@ class BondForceWithAtomVirial(PrimitiveWithInfer): dr = (x_1-x_2, y_1-y_2, z_1-z_2) virial = |dr|*(|dr| - r_0)*k + Args: + Same as operator BondForce(). + Inputs: Same as operator BondForce() Outputs: - **frc_f** (Tensor, float32) - [N, 3], same as operator BondForce(). - - **atom_v** (Tensor, float32) - [N,], the accumulated virial coefficient - for each atom. + - **atom_v** (Tensor, float32) - [N,], the accumulated virial coefficient for each atom. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, bond_numbers, atom_numbers): + assert isinstance(bond_numbers, int) + assert isinstance(atom_numbers, int) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -298,25 +312,25 @@ class BondForceWithAtomVirial(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) - validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(bond_k_shape[0], M, Rel.EQ, "bond_k", cls_name) + validator.check_int(bond_r0_shape[0], M, Rel.EQ, "bond_r0", cls_name) return uint_crd_f_shape, [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name) return bond_r0_type, bond_r0_type @@ -352,34 +366,36 @@ class DihedralForce(PrimitiveWithInfer): F_c = - dEdrl - dEdrjpart F_d = dEdrl + Args: + dihedral_numbers(int32): the number of dihedral terms M. + Inputs: - **dihedral_numbers** (int32) - the number of dihedral terms M. - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value of each atom. - - **scalar_f** (Tensor, float32) - [3, ], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M, ], the 1st atom index of each dihedral. - - **atom_b** (Tensor, int32) - [M, ], the 2nd atom index of each dihedral. - - **atom_c** (Tensor, int32) - [M, ], the 3rd atom index of each dihedral. - - **atom_d** (Tensor, int32) - [M, ], the 4th atom index of each dihedral. - 4 atoms are connected in the form a-b-c-d. - - **ipn** (Tensor, int32) - [M, ], the period of dihedral angle of each dihedral. - - **pk** (Tensor, float32) - [M, ], the force constant of each dihedral. - - **gamc** (Tensor, float32) - [M, ], k*cos(phi_0) of each dihedral. - - **gams** (Tensor, float32) - [M, ], k*sin(phi_0) of each dihedral. - - **pn** (Tensor, float32) - [M, ], the floating point form of ipn. + value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each dihedral. + - **atom_b** (Tensor, int32) - [M,], the 2nd atom index of each dihedral. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each dihedral. + - **atom_d** (Tensor, int32) - [M,], the 4th atom index of each dihedral. + 4 atoms are connected in the form a-b-c-d. + - **ipn** (Tensor, int32) - [M,], the period of dihedral angle of each dihedral. + - **pk** (Tensor, float32) - [M,], the force constant of each dihedral. + - **gamc** (Tensor, float32) - [M,], k*cos(phi_0) of each dihedral. + - **gams** (Tensor, float32) - [M,], k*sin(phi_0) of each dihedral. + - **pn** (Tensor, float32) - [M,], the floating point form of ipn. Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, dihedral_numbers): + assert isinstance(dihedral_numbers, int) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -390,32 +406,32 @@ class DihedralForce(PrimitiveWithInfer): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): cls_name = self.name M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn", cls_name) return uint_crd_f_shape def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d_type', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn_type', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk_type', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) return pn_type @@ -431,21 +447,23 @@ class DihedralEnergy(PrimitiveWithInfer): E = k(1 + cos(n*phi - phi_0)) + Args: + Same as operator DihedralForce(). + Inputs: Same as operator DihedralForce(). Outputs: - - **ene** (Tensor, float32) - [M, ], the potential energy for each - dihedral term. + - **ene** (Tensor, float32) - [M,], the potential energy for each + dihedral term. Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, dihedral_numbers): + assert isinstance(dihedral_numbers, int) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -456,32 +474,32 @@ class DihedralEnergy(PrimitiveWithInfer): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): cls_name = self.name M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn", cls_name) return [M,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d_type', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn_type', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk_type', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) return pn_type @@ -495,21 +513,23 @@ class DihedralAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator DihedralEnergy(). + Args: + Same as operator DihedralEnergy(). + Inputs: Same as operator DihedralEnergy(). Outputs: - - **ene** (Tensor, float32) - [N, ], the accumulated potential - energy for each atom. + - **ene** (Tensor, float32) - [N,], the accumulated potential + energy for each atom. Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, dihedral_numbers): + assert isinstance(dihedral_numbers, int) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -521,32 +541,32 @@ class DihedralAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = uint_crd_f_shape[0] M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn", cls_name) return [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d_type', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn_type', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk_type', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) return pn_type @@ -559,21 +579,23 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator DihedralForce() and DihedralEnergy(). + Args: + Same as operator DihedralForce(). + Inputs: Same as operator DihedralForce(). Outputs: - **frc_f** (Tensor, float32) - [N, 3], same as operator DihedralForce(). - - **ene** (Tensor, float32) - [N, ], same as operator DihedralAtomEnergy(). + - **ene** (Tensor, float32) - [N,], same as operator DihedralAtomEnergy(). Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, dihedral_numbers): + assert isinstance(dihedral_numbers, int) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -585,32 +607,32 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = uint_crd_f_shape[0] M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) - validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) - validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) - validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) - validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) - validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(atom_d_shape[0], M, Rel.EQ, "atom_d", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn", cls_name) return uint_crd_f_shape, [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_d_type', atom_d_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('ipn_type', ipn_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('pk_type', pk_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name) return pn_type, pn_type @@ -634,31 +656,29 @@ class AngleForce(PrimitiveWithInfer): - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}] F_b = -F_a - F_c + Args: + angle_numbers(int32): the number of angles M. + Inputs: - - **angle_numbers** (int32) - the number of angles M. - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate - value of each atom. - - **scaler_f** (Tensor, float32) - [3, ], the 3-D scale factor between - the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M, ], the 1st atom index of each angle. - - **atom_b** (Tensor, int32) - [M, ], the 2nd and the central atom index - of each angle. - - **atom_c** (Tensor, int32) - [M, ], the 3rd atom index of each angle. - - **angle_k** (Tensor, float32) - [M, ], the force constant for each angle. - - **angle_theta0** (Tensor, float32) - [M, ], the equilibrium position value - for each angle. + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **scaler_f** (Tensor, float32) - [3,], the 3-D scale factor between + the real space float coordinates and the unsigned int coordinates. + - **atom_a** (Tensor, int32) - [M,], the 1st atom index of each angle. + - **atom_b** (Tensor, int32) - [M,], the 2nd and the central atom index of each angle. + - **atom_c** (Tensor, int32) - [M,], the 3rd atom index of each angle. + - **angle_k** (Tensor, float32) - [M,], the force constant for each angle. + - **angle_theta0** (Tensor, float32) - [M,], the equilibrium position value for each angle. Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, angle_numbers): + assert isinstance(angle_numbers, int) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -669,24 +689,24 @@ class AngleForce(PrimitiveWithInfer): angle_theta0_shape): cls_name = self.name M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0", cls_name) return uint_crd_f_shape def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k_type', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0_type', angle_theta0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) return angle_k_type @@ -703,21 +723,22 @@ class AngleEnergy(PrimitiveWithInfer): theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) E = k*(theta - theta_0)^2 + Args: + Same as operator AngleForce(). + Inputs: Same as operator AngleForce(). Outputs: - - **ene** (Tensor, float32) - [M, ], the potential energy for - each angle term. + - **ene** (Tensor, float32) - [M,], the potential energy for each angle term. Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, angle_numbers): + assert isinstance(angle_numbers, int) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -728,24 +749,24 @@ class AngleEnergy(PrimitiveWithInfer): angle_theta0_shape): cls_name = self.name M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0", cls_name) return [M,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k_type', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0_type', angle_theta0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) return angle_k_type @@ -758,21 +779,22 @@ class AngleAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator AngleEnergy(). + Args: + Same as operator AngleForce(). + Inputs: Same as operator AngleForce(). Outputs: - - **ene** (Tensor, float32) - [N, ], the accumulated potential energy - for each atom. + - **ene** (Tensor, float32) - [N,], the accumulated potential energy for each atom. Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, angle_numbers): + assert isinstance(angle_numbers, int) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -784,24 +806,24 @@ class AngleAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = uint_crd_f_shape[0] M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0", cls_name) return [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k_type', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0_type', angle_theta0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) return angle_k_type @@ -813,21 +835,23 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator AngleForce() and AngleEnergy(). + Args: + Same as operator AngleForce(). + Inputs: Same as operator AngleForce(). Outputs: - **frc_f** (Tensor, float32) - [N, 3], same as operator AngleForce(). - - **ene** (Tensor, float) - [N, ], same as operator AngleAtomEnergy(). + - **ene** (Tensor, float) - [N,], same as operator AngleAtomEnergy(). Supported Platforms: ``GPU`` - - Examples: """ @prim_attr_register def __init__(self, angle_numbers): + assert isinstance(angle_numbers, int) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -839,24 +863,24 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = uint_crd_f_shape[0] M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) - validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) - validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) - validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) - validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) - validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) - validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) + validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) + validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) + validator.check_int(atom_c_shape[0], M, Rel.EQ, "atom_c", cls_name) + validator.check_int(angle_k_shape[0], M, Rel.EQ, "angle_k", cls_name) + validator.check_int(angle_theta0_shape[0], M, Rel.EQ, "angle_theta0", cls_name) return uint_crd_f_shape, [N,] def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_c_type', atom_c_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('angle_k_type', angle_k_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('angle_theta0_type', angle_theta0_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name) return angle_k_type, angle_k_type @@ -876,38 +900,34 @@ class Dihedral14LJForce(PrimitiveWithInfer): dr = (x_a-x_b, y_a-y_b, z_a-z_b) F = k*(-12*A/|dr|^{14} + 6*B/|dr|^{8})*dr + Args: + dihedral_14_numbers(int32): the number of necessary dihedral 1,4 terms M. + atom_numbers(int32): the number of atoms N. + Inputs: - - **dihedral_14_numbers** (int32) - the number of necessary dihedral - 1,4 terms M. - - **atom_numbers** (int32) - the number of atoms N. - - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate - value of each atom. - - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each - atom. + + - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJ_type** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. - **charge** (Tensor, float32) - [N,], the charge of each atom. - - **boxlength_f** (Tensor, float32) - [3,], the length of molecular - simulation box in 3 dimensions. - - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral - 1,4 term. - - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral - 1,4 term. + - **boxlength_f** (Tensor, float32) - [3,], the length of molecular simulation box in 3 dimensions. + - **a_14** (Tensor, int32) - [M,], the first atom index of each dihedral 1,4 term. + - **b_14** (Tensor, int32) - [M,], the second atom index of each dihedral 1,4 term. - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the - Lennard-Jones part of force correction of each dihedral 1,4 term. - - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones - scheme of each atom pair type. - - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones - shceme of each atom pair type. + Lennard-Jones part of force correction of each dihedral 1,4 term. + - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. + - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers self.init_prim_io_names( @@ -918,23 +938,22 @@ class Dihedral14LJForce(PrimitiveWithInfer): self.add_prim_attr('atom_numbers', self.atom_numbers) def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, - lj_scale_factor_shape, - LJ_type_A_shape, LJ_type_B_shape): + lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): return uint_crd_f_shape def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor_type', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A_type', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B_type', LJ_type_B_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) return LJ_type_B_type @@ -950,20 +969,24 @@ class Dihedral14LJEnergy(PrimitiveWithInfer): dr = (x_a-x_b, y_a-y_b, z_a-z-b) E = k*(A/|dr|^{12} - B/|dr|^{6}) + Args: + Same as operator Dihedral14LJForce(). + Inputs: Same as operator Dihedral14LJForce(). Outputs: - **ene** (Tensor, float32) - [M,], the Lennard-Jones potential - energy correction for each necessary dihedral 1,4 term. + energy correction for each necessary dihedral 1,4 term. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -980,15 +1003,15 @@ class Dihedral14LJEnergy(PrimitiveWithInfer): def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor_type', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A_type', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B_type', LJ_type_B_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) return LJ_type_A_type @@ -1008,22 +1031,26 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): dr = (x_a-x_b, y_a-y_b, z_a-z_b) F = -k*q_a*q_b/|r|^3*dr + Args: + Same as operator Dihedral14LJForce(). + Inputs: - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the - Coulomb part of force correction for each dihedral 1,4 terms. + Coulomb part of force correction for each dihedral 1,4 terms. - The rest of the inputs is the same as operator Dihedral14LJForce(). + The rest of the inputs is the same as operator Dihedral14LJForce(). Outputs: - **frc_f** (Tensor, float) - [N, 3], the force felt by each atom. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1041,16 +1068,16 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, lj_scale_factor_type, cf_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor_type', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('cf_scale_factor_type', cf_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A_type', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B_type', LJ_type_B_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) return LJ_type_A_type @@ -1067,22 +1094,24 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): Dihedral14LJForceWithDirectCF(), and the energy correction part is the same as operator Dihedral14LJEnergy() and Dihedral14CFEnergy(). + Args: + Same as operator Dihedral14LJForce(). + Inputs: Same as operator Dihedral14LJForceWithdirectCF(). Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. - - **atom_energy** (Tensor, float32) - [N,], the accumulated potential - energy for each atom. + - **atom_energy** (Tensor, float32) - [N,], the accumulated potential energy for each atom. Supported Platforms: ``GPU`` - Examples: """ - @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1100,16 +1129,16 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, lj_scale_factor_type, cf_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor_type', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('cf_scale_factor_type', cf_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A_type', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B_type', LJ_type_B_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) return charge_dtype, charge_dtype @@ -1123,21 +1152,23 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator Dihedral14LJEnergy(). + Args: + Same as operator Dihedral14LJForce(). + Inputs: Same as operator Dihedral14LJForce(). Outputs: - - **ene** (Tensor, float32) - [N,], the accumulated potential energy of - each atom. + - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. Supported Platforms: ``GPU`` - Examples: """ - @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1154,16 +1185,16 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer): def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, lj_scale_factor_type, LJ_type_A_type, LJ_type_B_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor_type', lj_scale_factor_type, [mstype.float32], + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_A_type', LJ_type_A_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('LJ_type_B_type', LJ_type_B_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_A', LJ_type_A_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('LJ_type_B', LJ_type_B_type, [mstype.float32], self.name) return LJ_type_A_type @@ -1180,21 +1211,23 @@ class Dihedral14CFEnergy(PrimitiveWithInfer): dr = (x_a-x_b, y_a-y_b, z_a-z_b) E = k*q_a*q_b/|dr| + Args: + Same as operator Dihedral14LJForce(). + Inputs: - The meaning and type of each input is the same as that of operator - Dihedral14LJForceWithDirectCF(). + The meaning and type of each input is the same as that of operator Dihedral14LJForceWithDirectCF(). Outputs: - - **ene** (Tensor, float32) - [M,], the accumulated potential energy - of each atom. + - **ene** (Tensor, float32) - [M,], the accumulated potential energy of each atom. Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1210,13 +1243,13 @@ class Dihedral14CFEnergy(PrimitiveWithInfer): def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, cf_scale_factor_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('lj_scale_factor_type', cf_scale_factor_type, [mstype.float32], + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('lj_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) return charge_dtype @@ -1231,23 +1264,24 @@ class Dihedral14CFAtomEnergy(PrimitiveWithInfer): The calculation formula is the same as operator Dihedral14CFEnergy(). + Args: + Same as operator Dihedral14LJForce(). + Inputs: - The meaning and type of each input is the same as that of operator - Dihedral14LJForceWithDirectCF(). + The meaning and type of each input is the same as that of operator Dihedral14LJForceWithDirectCF(). Outputs: - - **ene** (Tensor, float32) - [N,], the accumulated potential energy - of each atom. + - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. Supported Platforms: ``GPU`` - Examples: """ - @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): + assert isinstance(nb14_numbers, int) + assert isinstance(atom_numbers, int) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1263,13 +1297,13 @@ class Dihedral14CFAtomEnergy(PrimitiveWithInfer): def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, cf_scale_factor_type): - validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('LJtype_dtype', LJtype_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('charge_dtype', charge_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('boxlength_f_type', boxlength_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('a_14_type', a_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('b_14_type', b_14_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('cf_scale_factor_type', cf_scale_factor_type, [mstype.float32], + validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('LJtype', LJtype_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name) return charge_dtype @@ -1288,23 +1322,21 @@ class MDIterationLeapFrog(PrimitiveWithInfer): scheme for efficient configurational sampling for classical/quantum canonical ensembles via molecular dynamics. DOI: 10.1063/1.4991621. - Inputs: - - **float4_numbers** (int32) - total length to store random numbers. - - **atom_numbers** (int32) - the number of atoms N. - - **dt** (float32) - time step for finite difference. - - **half_dt** (float32) - half of time step for finite difference. - - **exp_gamma** (float32) - parameter in Liu's dynamic, equals - exp(-gamma_ln * dt), where gamma_ln is the firction factor in Langvin - dynamics. - - **max_velocity** (float32) - the upper limit of velocity, when the - veclocity overflows, scale it to the upper limit. - - **is_max_velocity** (int32) - whether the max velocity control is - open or not. + Args: + float4_numbers(int32): total length to store random numbers. + atom_numbers(int32): the number of atoms N. + dt(float32): time step for finite difference. + half_dt(float32): half of time step for finite difference. + exp_gamma(float32): parameter in Liu's dynamic, equals exp(-gamma_ln * dt), + where gamma_ln is the firction factor in Langvin dynamics. + max_velocity(float32): the upper limit of velocity, when the veclocity overflows, + scale it to the upper limit. + is_max_velocity(int32): whether the max velocity control is open or not. - - **mass_inverse** (Tensor, float32) - [N,], the inverse value of - mass of each atom. + Inputs: + - **mass_inverse** (Tensor, float32) - [N,], the inverse value of mass of each atom. - **sqrt_mass** (Tensor, float32) - [N,], the inverse square root value - of effect mass in Liu's dynamics of each atom. + of effect mass in Liu's dynamics of each atom. Outputs: - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. @@ -1314,11 +1346,17 @@ class MDIterationLeapFrog(PrimitiveWithInfer): Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): + assert isinstance(float4_numbers, int) + assert isinstance(atom_numbers, int) + assert isinstance(half_dt, float) + assert isinstance(dt, float) + assert isinstance(exp_gamma, float) + assert isinstance(is_max_velocity, int) + assert isinstance(max_velocity, float) self.float4_numbers = float4_numbers self.atom_numbers = atom_numbers self.half_dt = half_dt @@ -1342,8 +1380,8 @@ class MDIterationLeapFrog(PrimitiveWithInfer): return [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3] def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype): - validator.check_tensor_dtype_valid('mass_inverse_dtype', mass_inverse_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('sqrt_mass_dtype', sqrt_mass_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name) return mass_inverse_dtype, mass_inverse_dtype, mass_inverse_dtype, mass_inverse_dtype @@ -1353,42 +1391,39 @@ class PMEReciprocalForce(PrimitiveWithInfer): PMEReciprocalForce: Calculate the reciprocal part of long-range Coulumb force using - PME(Particle Meshed Ewald) method. Assume the number of atoms is - N. + PME(Particle Meshed Ewald) method. Assume the number of atoms is N. The detailed calculation formula of PME(Particle Meshed Ewald) method can be found in this paper: A Smooth Particle Mesh Ewald Method. DOI: 10.1063/1.470117. + Args: + atom_numbers(int32): the number of atoms, N. + beta(float32): the PME beta parameter, determined by the + non-bond cutoff value and simulation precision tolerance. + fftx(int32): the number of points for Fourier transform in dimension X. + ffty(int32): the number of points for Fourier transform in dimension Y. + fftz(int32): the number of points for Fourier transform in dimension Z. + Inputs: - - **atom_numbers** (int32) - the number of atoms, N. - - **beta** (float32) - the PME beta parameter, determined by the - non-bond cutoff value and simulation precision tolerance. - - **fftx** (int32) - the number of points for Fourier transform - in dimension X. - - **ffty** (int32) - the number of points for Fourier transform - in dimension Y. - - **fftz** (int32) - the number of points for Fourier transform - in dimension Z. - - - **boxlength** (Tensor, float32) - [3,], the length of simulation - box in 3 dimensions. - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value of each atom. - - **charge** (Tensor, float32) - [N,], the charge carried by each - atom. + - **boxlength** (Tensor, float32) - [3,], the length of simulation box in 3 dimensions. + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. Outputs: - **force** (Tensor, float32) - [N, 3], the force felt by each atom. Supported Platforms: ```GPU``` - - Examples: """ @prim_attr_register def __init__(self, atom_numbers, beta, fftx, ffty, fftz): + assert isinstance(atom_numbers, int) + assert isinstance(beta, float) + assert isinstance(fftx, int) + assert isinstance(ffty, int) + assert isinstance(fftz, int) self.atom_numbers = atom_numbers self.beta = beta self.fftx = fftx @@ -1406,9 +1441,9 @@ class PMEReciprocalForce(PrimitiveWithInfer): return uint_crd_shape def infer_dtype(self, boxlength_type, uint_crd_type, charge_type): - validator.check_tensor_dtype_valid('boxlength_type', boxlength_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('uint_crd_type', uint_crd_type, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('charge_type', charge_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('boxlength', boxlength_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name) return charge_type @@ -1420,18 +1455,20 @@ class PMEExcludedForce(PrimitiveWithInfer): PME(Particle Meshed Ewald) method. Assume the number of atoms is N, and the length of excluded list is E. + Args: + atom_numbers(int32): the number of atoms, N. + beta(float32): the PME beta parameter, determined by the + non-bond cutoff value and simulation precision tolerance. + Inputs: - - **atom_numbers** (int32) - the number of atoms, N. - - **beta** (float32) - the PME beta parameter, determined by the - non-bond cutoff value and simulation precision tolerance. - **scaler** (Tensor, float32) - [3,], the scale factor between real space - coordinates and its unsigned int value. + coordinates and its unsigned int value. - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index - in excluded list for each atom. + in excluded list for each atom. - **excluded_numbers** (Tensor, int32) - [N,], the number of atom excluded - in excluded list for each atom. + in excluded list for each atom. - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - list of each atom. + list of each atom. The rest of the input is the same as that of operator PMEReciprocalForce(). @@ -1440,12 +1477,12 @@ class PMEExcludedForce(PrimitiveWithInfer): Supported Platforms: ```GPU``` - - Examples: """ @prim_attr_register def __init__(self, atom_numbers, beta): + assert isinstance(atom_numbers, int) + assert isinstance(beta, float) self.atom_numbers = atom_numbers self.beta = beta self.init_prim_io_names( @@ -1460,14 +1497,14 @@ class PMEExcludedForce(PrimitiveWithInfer): def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type, excluded_atom_numbers_type): - validator.check_tensor_dtype_valid('sacler_type', sacler_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('uint_crd_type', uint_crd_type, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('charge_type', charge_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('excluded_list_start_type', excluded_list_start_type, [mstype.int32], + validator.check_tensor_dtype_valid('sacler', sacler_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('excluded_list_type', excluded_list_type, [mstype.int32], + validator.check_tensor_dtype_valid('excluded_list', excluded_list_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('excluded_atom_numbers_type', excluded_atom_numbers_type, [mstype.int32], + validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers_type, [mstype.int32], self.name) return charge_type @@ -1482,8 +1519,11 @@ class PMEEnergy(PrimitiveWithInfer): E = sum_{ij} q_iq_j/r_{ij} + Args: + same as operator PMEReciprocalForce(). + Inputs: - Same as that of operator PMEReciprocalForce(), PMEExcludedForce() + Same as operator PMEReciprocalForce(), PMEExcludedForce() and PMEDirectAtomEnergy(). Outputs: @@ -1494,12 +1534,15 @@ class PMEEnergy(PrimitiveWithInfer): Supported Platforms: ```GPU`` - - Examples: """ @prim_attr_register def __init__(self, atom_numbers, beta, fftx, ffty, fftz): + assert isinstance(atom_numbers, int) + assert isinstance(beta, float) + assert isinstance(fftx, int) + assert isinstance(ffty, int) + assert isinstance(fftz, int) self.atom_numbers = atom_numbers self.beta = beta self.fftx = fftx @@ -1552,39 +1595,33 @@ class LJEnergy(PrimitiveWithInfer): dr = (x_a-x_b, y_a-y_b, z_a-z_b) E = A/|dr|^{12} - B/|dr|^{6} + Agrs: + atom_numbers(int32): the number of atoms, N. + cutoff_square(float32): the square value of cutoff. + Inputs: - - **atom_numbers** (int32) - the number of atoms, N. - - **cutoff_square** (float32) - the square value of cutoff. - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate - value of each atom. - - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each - atom. - - **charge** (Tensor, float32) - [N,], the charge carried by each - atom. + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. + - **LJtype** (Tensor, int32) - [N,], the Lennard-Jones type of each atom. + - **charge** (Tensor, float32) - [N,], the charge carried by each atom. - **scaler** (Tensor, float32) - [3,], the scale factor between real - space coordinate and its unsigned int value. + space coordinate and its unsigned int value. - **nl_numbers** - (Tensor, int32) - [N,], the each atom. - - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, - the max number is 800. - - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient - of each kind of atom pair. - - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient - of each kind of atom pair. + - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. + - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. + - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. Outputs: - - **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones - potential energy of each atom. - - **d_LJ_energy_sum** (float32), the sum of Lennard-Jones potential - energy of each atom. + - **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones potential energy of each atom. + - **d_LJ_energy_sum** (float32), the sum of Lennard-Jones potential energy of each atom. Supported Platforms: ```GPU``` - - Examples: """ @prim_attr_register def __init__(self, atom_numbers, cutoff_square): + assert isinstance(atom_numbers, int) + assert isinstance(cutoff_square, float) self.atom_numbers = atom_numbers self.cutoff_square = cutoff_square self.init_prim_io_names( @@ -1620,6 +1657,9 @@ class LJForce(PrimitiveWithInfer): dr = (x_a-x_b, y_a-y_b, z_a-z_b) F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr + Args: + Same as operator LJEnergy(). + Inputs: Same as operator LJEnergy(). @@ -1628,11 +1668,12 @@ class LJForce(PrimitiveWithInfer): Supported Platforms: ```GPU``` - Examples: """ @prim_attr_register def __init__(self, atom_numbers, cutoff_square): + assert isinstance(atom_numbers, int) + assert isinstance(cutoff_square, float) self.atom_numbers = atom_numbers self.cutoff_square = cutoff_square self.init_prim_io_names( @@ -1665,23 +1706,25 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer): The calculation formula of Lennard-Jones part is the same as operator LJForce(), and the PME direct part is within PME method. - Inputs: - - **pme_beta** (float32), PME beta parameter, same as operator - PMEReciprocalForce(). + Args: + pme_beta(float32): PME beta parameter, same as operator PMEReciprocalForce(). + The rest is same as operator LJEnergy(). - the rest of the input is the same as operator LJForce(). + Inputs: + same as operator LJForce(). Outputs: - **frc** (Tensor, float32), [N, 3], the force felt by each atom. Supported Platforms: ```GPU``` - - Examples: """ @prim_attr_register def __init__(self, atom_numbers, cutoff, pme_beta): + assert isinstance(atom_numbers, int) + assert isinstance(cutoff, float) + assert isinstance(pme_beta, float) self.atom_numbers = atom_numbers self.cutoff = cutoff self.pme_beta = pme_beta @@ -1713,7 +1756,6 @@ class GetCenterOfGeometry(PrimitiveWithInfer): Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register @@ -1729,14 +1771,14 @@ class GetCenterOfGeometry(PrimitiveWithInfer): def infer_shape(self, center_atoms_shape, crd_f_shape): cls_name = self.name N = self.center_numbers - validator.check_int(center_atoms_shape[0], N, Rel.EQ, "center_atoms_shape", cls_name) - validator.check_int(crd_f_shape[0], N, Rel.EQ, "crd_f_shape", cls_name) - validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f_shape", cls_name) + validator.check_int(center_atoms_shape[0], N, Rel.EQ, "center_atoms", cls_name) + validator.check_int(crd_f_shape[0], N, Rel.EQ, "crd_f", cls_name) + validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f", cls_name) return [3,] def infer_dtype(self, center_atoms_dtype, crd_f_dtype): - validator.check_tensor_dtype_valid('center_atoms_dtype', center_atoms_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('crd_f_dtype', crd_f_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('center_atoms', center_atoms_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('crd_f', crd_f_dtype, [mstype.float32], self.name) return crd_f_dtype @@ -1745,15 +1787,20 @@ class MDTemperature(PrimitiveWithInfer): """ MDTemperature: + Calculate the temperature. + Supported Platforms: ``GPU`` - Examples: """ @prim_attr_register - def __init__(self, residue_numbers): + def __init__(self, residue_numbers, atom_numbers): + assert isinstance(residue_numbers, int) + assert isinstance(atom_numbers, int) self.residue_numbers = residue_numbers + self.atom_numbers = atom_numbers self.add_prim_attr('residue_numbers', self.residue_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) self.init_prim_io_names( inputs=['start', 'end', 'atom_vel_f', 'atom_mass'], outputs=['ek']) @@ -1761,18 +1808,22 @@ class MDTemperature(PrimitiveWithInfer): def infer_shape(self, start_shape, end_shape, atom_vel_f_shape, atom_mass_shape): cls_name = self.name N = self.residue_numbers - validator.check_int(start_shape[0], N, Rel.EQ, "start_shape", cls_name) - validator.check_int(end_shape[0], N, Rel.EQ, "end_shape", cls_name) - validator.check_int(atom_vel_f_shape[0], N, Rel.EQ, "atom_vel_f_shape", cls_name) - validator.check_int(atom_vel_f_shape[1], 3, Rel.EQ, "atom_vel_f_shape", cls_name) - validator.check_int(atom_mass_shape[0], N, Rel.EQ, "atom_mass_shape", cls_name) + M = self.atom_numbers + validator.check_int(len(start_shape), 1, Rel.EQ, "start", cls_name) + validator.check_int(start_shape[0], N, Rel.EQ, "end", cls_name) + validator.check_int(len(end_shape), 1, Rel.EQ, "start", cls_name) + validator.check_int(end_shape[0], N, Rel.EQ, "end", cls_name) + validator.check_int(atom_vel_f_shape[0], M, Rel.EQ, "atom_vel_f", cls_name) + validator.check_int(atom_vel_f_shape[1], 3, Rel.EQ, "atom_vel_f", cls_name) + validator.check_int(len(atom_mass_shape), 1, Rel.EQ, "atom_mass", cls_name) + validator.check_int(atom_mass_shape[0], M, Rel.EQ, "atom_mass", cls_name) return [N,] def infer_dtype(self, start_dtype, end_dtype, atom_vel_f_dtype, atom_mass_dtype): - validator.check_tensor_dtype_valid('start_dtype', start_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('end_dtype', end_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('atom_vel_f_dtype', atom_vel_f_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_mass_dtype', atom_mass_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('start', start_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('end', end_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('atom_vel_f', atom_vel_f_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('atom_mass', atom_mass_dtype, [mstype.float32], self.name) return atom_mass_dtype @@ -1786,74 +1837,53 @@ class NeighborListUpdate(PrimitiveWithInfer): grid is M, the maximum number of atoms in single atom's neighbor list is L, and the number of total atom in excluded list is E. + Args: + grid_numbers(int32): the total number of grids divided. + refresh_count(int32): the counter which counts how many + iteration steps have passed since last update. + not_first_time(int32): whether to construct the neighbor + list first time or not. + Nxy(int32): the total number of grids divided in xy plane. + excluded_atom_numbers(int32): the total atom numbers in the excluded list. + cutoff(float32): the cutoff distance for short-range force calculation. + skin(float32): the overflow value of cutoff to maintain a neighbor list. + cutoff_square(float32): the suqare value of cutoff. + half_skin_square(float32): skin*skin/4, indicates the maximum + square value of the distance atom allowed to move between two updates. + cutoff_with_skin(float32): cutoff + skin, indicates the + radius of the neighbor list for each atom. + half_cutoff_with_skin(float32): cutoff_with_skin/2. + cutoff_with_skin_square(float32): the square value of cutoff_with_skin. + refresh_interval(int32): the number of iteration steps between two updates of neighbor list. + max_atom_in_grid_numbers(int32): the maximum number of atoms in one grid. + Inputs: - - **grid_numbers** (int32) - the total number of grids divided. - - **refresh_count** (int32) - the counter which counts how many - iteration steps have passed since last update. - - **not_first_time** (int32) - whether to construct the neighbor - list first time or not. - - **Nxy** (int32) - the total number of grids divided in xy plane. - - **excluded_atom_numbers** (int32) - the total atom numbers in - the excluded list. - - **cutoff** (float32) - the cutoff distance for short-range force - calculation. - - **skin** (float32) - the overflow value of cutoff to maintain a - neighbor list. - - **cutoff_square** (float32) - the suqare value of cutoff. - - **half_skin_square** (float32) - skin*skin/4, indicates the maximum - square value of the distance atom allowed to move between two updates. - - **cutoff_with_skin** (float32) - cutoff + skin, indicates the - radius of the neighbor list for each atom. - - **half_cutoff_with_skin** (float32) - cutoff_with_skin/2. - - **cutoff_with_skin_square** (float32) - the square value of - cutoff_with_skin. - - **refresh_interval** (int32) - the number of iteration steps - between two updates of neighbor list. - - **max_atom_in_grid_numbers** (int32) - the maximum number of atoms - in one grid. - - - **atom_numbers_in_grid_bucket** (Tensor, int32) - [G,], the number - of atoms in each grid bucket. - - **bucket** (Tensor, int32) - (Tensor,int32) - [G, M], the atom - indices in each grid bucket. + - **atom_numbers_in_grid_bucket** (Tensor, int32) - [G,], the number of atoms in each grid bucket. + - **bucket** (Tensor, int32) - (Tensor,int32) - [G, M], the atom indices in each grid bucket. - **crd** (Tensor, float32) - [N,], the coordinates of each atom. - - **box_length** (Tensor, float32) - [3,], the length of 3 dimensions - of the simulation box. - - **grid_N** (Tensor, int32) - [3,], the number of grids divided of - 3 dimensions of the simulation box. + - **box_length** (Tensor, float32) - [3,], the length of 3 dimensions of the simulation box. + - **grid_N** (Tensor, int32) - [3,], the number of grids divided of 3 dimensions of the simulation box. - **grid_length_inverse** (float32) - the inverse value of grid length. - - **atom_in_grid_serial** (Tensor, int32) - [N,], the grid index for - each atom. - - **old_crd** (Tensor, float32) - [N, 3], the coordinates before update - of each atom. + - **atom_in_grid_serial** (Tensor, int32) - [N,], the grid index for each atom. + - **old_crd** (Tensor, float32) - [N, 3], the coordinates before update of each atom. - **crd_to_uint_crd_cof** (Tensor, float32) - [3,], the scale factor - between the unsigned int value and the real space coordinates. - - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates - value fo each atom. - - **gpointer** (Tensor, int32) - [G, 125], the 125 nearest neighbor grids - (including self) of each grid. - - **nl_atom_numbers** (Tensor, int32) - [N,], the number of atoms in - neighbor list of each atom. - - **nl_atom_serial** (Tensor, int32) - [N, L], the indices of atoms in - neighbor list of each atom. + between the unsigned int value and the real space coordinates. + - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value fo each atom. + - **gpointer** (Tensor, int32) - [G, 125], the 125 nearest neighbor grids (including self) of each grid. + - **nl_atom_numbers** (Tensor, int32) - [N,], the number of atoms in neighbor list of each atom. + - **nl_atom_serial** (Tensor, int32) - [N, L], the indices of atoms in neighbor list of each atom. - **uint_dr_to_dr_cof** (Tensor, float32) - [3,], the scale factor between - the real space coordinates and the unsigned int value. - - **excluded_list_start** (Tensor, int32) - [N,], the start excluded - index in excluded list for each atom. - - **excluded_numbers** (Tensor, int32) - [N,], the number of atom excluded - in excluded list for each atom. - - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - list of each atom. - - **need_refresh_flag** (Tensor, int32) - [N,], whether the neighbor list - of each atom need update or not. + the real space coordinates and the unsigned int value. + - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index in excluded list for each atom. + - **excluded_numbers** (Tensor, int32) - [N,], the number of atom excluded in excluded list for each atom. + - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded list of each atom. + - **need_refresh_flag** (Tensor, int32) - [N,], whether the neighbor list of each atom need update or not. Outputs: - **res** (float32) Supported Platforms: ```GPU``` - - Examples: """ @prim_attr_register @@ -1923,34 +1953,34 @@ class NeighborListUpdate(PrimitiveWithInfer): assert len(need_refresh_flag_shape) == 1 validator.check_int(atom_numbers_in_grid_bucket_shape[0], self.grid_numbers, Rel.EQ, - "atom_numbers_in_grid_bucket_shape", self.name) - validator.check_int(bucket_shape[0], self.grid_numbers, Rel.EQ, "bucket_shape", self.name) - validator.check_int(bucket_shape[1], self.max_atom_in_grid_numbers, Rel.EQ, "bucket_shape", self.name) - validator.check_int(crd_shape[0], self.atom_numbers, Rel.EQ, "crd_shape", self.name) - validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape", self.name) - validator.check_int(box_length_shape[0], 3, Rel.EQ, "box_length_shape", self.name) - validator.check_int(grid_N_shape[0], 3, Rel.EQ, "grid_N_shape", self.name) - validator.check_int(grid_length_inverse_shape[0], 3, Rel.EQ, "grid_length_inverse_shape", self.name) - validator.check_int(atom_in_grid_serial_shape[0], self.atom_numbers, Rel.EQ, "atom_in_grid_serial_shape", + "atom_numbers_in_grid_bucket", self.name) + validator.check_int(bucket_shape[0], self.grid_numbers, Rel.EQ, "bucket", self.name) + validator.check_int(bucket_shape[1], self.max_atom_in_grid_numbers, Rel.EQ, "bucket", self.name) + validator.check_int(crd_shape[0], self.atom_numbers, Rel.EQ, "crd", self.name) + validator.check_int(crd_shape[1], 3, Rel.EQ, "crd", self.name) + validator.check_int(box_length_shape[0], 3, Rel.EQ, "box_length", self.name) + validator.check_int(grid_N_shape[0], 3, Rel.EQ, "grid_N", self.name) + validator.check_int(grid_length_inverse_shape[0], 3, Rel.EQ, "grid_length_inverse", self.name) + validator.check_int(atom_in_grid_serial_shape[0], self.atom_numbers, Rel.EQ, "atom_in_grid_serial", self.name) - validator.check_int(old_crd_shape[0], self.atom_numbers, Rel.EQ, "old_crd_shape", self.name) - validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd_shape", self.name) - validator.check_int(crd_to_uint_crd_cof_shape[0], 3, Rel.EQ, "crd_to_uint_crd_cof_shape", self.name) - validator.check_int(uint_crd_shape[0], self.atom_numbers, Rel.EQ, "uint_crd_shape", self.name) - validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape", self.name) - validator.check_int(gpointer_shape[0], self.grid_numbers, Rel.EQ, "gpointer_shape", self.name) - validator.check_int(gpointer_shape[1], 125, Rel.EQ, "gpointer_shape", self.name) - validator.check_int(nl_atom_numbers_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_numbers_shape", self.name) - validator.check_int(nl_atom_serial_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_serial_shape", self.name) - validator.check_int(nl_atom_serial_shape[1], self.max_neighbor_numbers, Rel.EQ, "nl_atom_serial_shape", + validator.check_int(old_crd_shape[0], self.atom_numbers, Rel.EQ, "old_crd", self.name) + validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd", self.name) + validator.check_int(crd_to_uint_crd_cof_shape[0], 3, Rel.EQ, "crd_to_uint_crd_cof", self.name) + validator.check_int(uint_crd_shape[0], self.atom_numbers, Rel.EQ, "uint_crd", self.name) + validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd", self.name) + validator.check_int(gpointer_shape[0], self.grid_numbers, Rel.EQ, "gpointer", self.name) + validator.check_int(gpointer_shape[1], 125, Rel.EQ, "gpointer", self.name) + validator.check_int(nl_atom_numbers_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_numbers", self.name) + validator.check_int(nl_atom_serial_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_serial", self.name) + validator.check_int(nl_atom_serial_shape[1], self.max_neighbor_numbers, Rel.EQ, "nl_atom_serial", self.name) - validator.check_int(uint_dr_to_dr_cof_shape[0], 3, Rel.EQ, "uint_dr_to_dr_cof_shape", self.name) - validator.check_int(excluded_list_start_shape[0], self.atom_numbers, Rel.EQ, "excluded_list_start_shape", + validator.check_int(uint_dr_to_dr_cof_shape[0], 3, Rel.EQ, "uint_dr_to_dr_cof", self.name) + validator.check_int(excluded_list_start_shape[0], self.atom_numbers, Rel.EQ, "excluded_list_start", self.name) - validator.check_int(excluded_list_shape[0], self.excluded_atom_numbers, Rel.EQ, "excluded_list_shape", + validator.check_int(excluded_list_shape[0], self.excluded_atom_numbers, Rel.EQ, "excluded_list", self.name) - validator.check_int(excluded_numbers_shape[0], self.atom_numbers, Rel.EQ, "excluded_numbers_shape", self.name) - validator.check_int(need_refresh_flag_shape[0], 1, Rel.EQ, "need_refresh_flag_shape", self.name) + validator.check_int(excluded_numbers_shape[0], self.atom_numbers, Rel.EQ, "excluded_numbers", self.name) + validator.check_int(need_refresh_flag_shape[0], 1, Rel.EQ, "need_refresh_flag", self.name) return [1,] @@ -1959,30 +1989,134 @@ class NeighborListUpdate(PrimitiveWithInfer): uint_crd_dtype, gpointer_dtype, nl_atom_numbers_dtype, nl_atom_serial_dtype, uint_dr_to_dr_cof_dtype, excluded_list_start_dtype, excluded_list_dtype, excluded_numbers_dtype, need_refresh_flag_dtype): - validator.check_tensor_dtype_valid('atom_numbers_in_grid_bucket_dtype', atom_numbers_in_grid_bucket_dtype, + validator.check_tensor_dtype_valid('atom_numbers_in_grid_bucket', atom_numbers_in_grid_bucket_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bucket_dtype', bucket_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('crd_dtype', crd_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('box_length_dtype', box_length_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('grid_N_dtype', grid_N_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('grid_length_inverse_dtype', grid_length_inverse_dtype, [mstype.float32], + validator.check_tensor_dtype_valid('bucket', bucket_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('box_length', box_length_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('grid_N', grid_N_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('grid_length_inverse', grid_length_inverse_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_in_grid_serial_dtype', atom_in_grid_serial_dtype, [mstype.int32], + validator.check_tensor_dtype_valid('atom_in_grid_serial', atom_in_grid_serial_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('old_crd_dtype', old_crd_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('crd_to_uint_crd_cof_dtype', crd_to_uint_crd_cof_dtype, [mstype.float32], + validator.check_tensor_dtype_valid('old_crd', old_crd_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('crd_to_uint_crd_cof', crd_to_uint_crd_cof_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('uint_crd_dtype', uint_crd_dtype, [mstype.uint32], self.name) - validator.check_tensor_dtype_valid('gpointer_dtype', gpointer_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_atom_numbers_dtype', nl_atom_numbers_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('nl_atom_serial_dtype', nl_atom_serial_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('uint_dr_to_dr_cof_dtype', uint_dr_to_dr_cof_dtype, [mstype.float32], + validator.check_tensor_dtype_valid('uint_crd', uint_crd_dtype, [mstype.uint32], self.name) + validator.check_tensor_dtype_valid('gpointer', gpointer_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_atom_numbers', nl_atom_numbers_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('nl_atom_serial', nl_atom_serial_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('uint_dr_to_dr_cof', uint_dr_to_dr_cof_dtype, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('excluded_list_start_dtype', excluded_list_start_dtype, [mstype.int32], + validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('excluded_list_dtype', excluded_list_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('excluded_numbers_dtype', excluded_numbers_dtype, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('need_refresh_flag_dtype', need_refresh_flag_dtype, [mstype.int32], + validator.check_tensor_dtype_valid('excluded_list', excluded_list_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('excluded_numbers', excluded_numbers_dtype, [mstype.int32], self.name) + validator.check_tensor_dtype_valid('need_refresh_flag', need_refresh_flag_dtype, [mstype.int32], self.name) return mstype.float32 + + +class MDIterationLeapFrogWithRF(PrimitiveWithInfer): + """ + MDIterationLeapFrog: + + One step of classical leap frog algorithm to solve the finite difference + Hamiltonian equations of motion for certain system, using Langevin dynamics + with Liu's thermostat scheme. Assume the number of atoms is N and the target + control temperature is T. + + Detailed iteration formula can be found in this paper: A unified thermostat + scheme for efficient configurational sampling for classical/quantum canonical + ensembles via molecular dynamics. DOI: 10.1063/1.4991621. + + Inputs: + - **float4_numbers** (int32) - total length to store random numbers. + - **atom_numbers** (int32) - the number of atoms N. + - **dt** (float32) - time step for finite difference. + - **half_dt** (float32) - half of time step for finite difference. + - **exp_gamma** (float32) - parameter in Liu's dynamic, equals + exp(-gamma_ln * dt), where gamma_ln is the firction factor in Langvin + dynamics. + - **max_velocity** (float32) - the upper limit of velocity, when the + veclocity overflows, scale it to the upper limit. + - **is_max_velocity** (int32) - whether the max velocity control is + open or not. + + - **mass_inverse** (Tensor, float32) - [N,], the inverse value of + mass of each atom. + - **sqrt_mass** (Tensor, float32) - [N,], the inverse square root value + of effect mass in Liu's dynamics of each atom. + - **vel** (Tensor, float32) - [N, 3], the velocity of each atom. + - **crd** (Tensor, float32) - [N, 3], the coordinate of each atom. + - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. + - **acc** (Tensor, float32) - [N, 3], the acceleration of each atom. + - **random force** (Tensor, float32) - [N, 3], the random forces. + + Outputs: + - **res** (float32) + + Supported Platforms: + ``GPU`` + Examples: + """ + + @prim_attr_register + def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): + assert isinstance(float4_numbers, int) + assert isinstance(atom_numbers, int) + assert isinstance(half_dt, float) + assert isinstance(dt, float) + assert isinstance(exp_gamma, float) + assert isinstance(is_max_velocity, int) + assert isinstance(max_velocity, float) + self.float4_numbers = float4_numbers + self.atom_numbers = atom_numbers + self.half_dt = half_dt + self.dt = dt + self.exp_gamma = exp_gamma + self.is_max_velocity = is_max_velocity + self.max_velocity = max_velocity + + self.init_prim_io_names( + inputs=['mass_inverse', 'sqrt_mass', 'vel_in', 'crd_in', 'frc_in', 'acc_in', 'random_force'], + outputs=['res']) + self.add_prim_attr('float4_numbers', self.float4_numbers) + self.add_prim_attr('atom_numbers', self.atom_numbers) + self.add_prim_attr('half_dt', self.half_dt) + self.add_prim_attr('dt', self.dt) + self.add_prim_attr('exp_gamma', self.exp_gamma) + self.add_prim_attr('is_max_velocity', self.is_max_velocity) + self.add_prim_attr('max_velocity', self.max_velocity) + + def infer_shape(self, mass_inverse_shape, sqrt_mass_shape, vel_in_shape, crd_in_shape, frc_in_shape, acc_in_shape, + random_force_shape): + N = self.atom_numbers + validator.check_int(len(mass_inverse_shape), 1, Rel.EQ, "mass_inverse", self.name) + validator.check_int(len(sqrt_mass_shape), 1, Rel.EQ, "mass_inverse", self.name) + validator.check_int(mass_inverse_shape[0], N, Rel.EQ, "mass_inverse", self.name) + validator.check_int(sqrt_mass_shape[0], N, Rel.EQ, "mass_inverse", self.name) + validator.check_int(vel_in_shape[0], N, Rel.EQ, "vel_in", self.name) + validator.check_int(vel_in_shape[1], 3, Rel.EQ, "vel_in", self.name) + validator.check_int(crd_in_shape[0], N, Rel.EQ, "crd_in", self.name) + validator.check_int(crd_in_shape[1], 3, Rel.EQ, "crd_in", self.name) + validator.check_int(frc_in_shape[0], N, Rel.EQ, "frc_in", self.name) + validator.check_int(frc_in_shape[1], 3, Rel.EQ, "frc_in", self.name) + validator.check_int(acc_in_shape[0], N, Rel.EQ, "acc_in", self.name) + validator.check_int(acc_in_shape[1], 3, Rel.EQ, "acc_in", self.name) + validator.check_int(random_force_shape[0], N, Rel.EQ, "random_force", self.name) + validator.check_int(random_force_shape[1], 3, Rel.EQ, "random_force", self.name) + + return [1,] + + def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype, vel_in_dtype, crd_in_dtype, frc_in_dtype, acc_in_dtype, + rf_dtype): + validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('vel_in', vel_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('crd_in', crd_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('frc_in', frc_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('acc_in', acc_in_dtype, [mstype.float32], self.name) + validator.check_tensor_dtype_valid('rf', rf_dtype, [mstype.float32], self.name) + return mstype.float32 diff --git a/model_zoo/research/hpc/sponge/main.py b/model_zoo/research/hpc/sponge/main.py new file mode 100644 index 0000000000..ae9532c57d --- /dev/null +++ b/model_zoo/research/hpc/sponge/main.py @@ -0,0 +1,55 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""main""" + +import time +import argparse +from mindspore import context +from src.simulation_initial import Simulation + +parser = argparse.ArgumentParser(description='Sponge Controller') +parser.add_argument('--i', type=str, default=None, help='input file') +parser.add_argument('--amber_parm', type=str, default=None, + help='paramter file in AMBER type') +parser.add_argument('--c', type=str, default=None, + help='initial coordinates file') +parser.add_argument('--r', type=str, default="restrt", help='') +parser.add_argument('--x', type=str, default="mdcrd", help='') +parser.add_argument('--o', type=str, default="mdout", help="") +parser.add_argument('--box', type=str, default="mdbox", help='') +args_opt = parser.parse_args() + +context.set_context(mode=context.PYNATIVE_MODE, + device_target="GPU", device_id=0, save_graphs=True) + +if __name__ == "__main__": + start = time.time() + simulation = Simulation(args_opt) + simulation.Main_Initial() + res = simulation.Initial_Neighbor_List_Update(not_first_time=0) + md_info = simulation.md_info + md_info.step_limit = 1 + for i in range(1, md_info.step_limit + 1): + print("steps: ", i) + md_info.steps = i + simulation.Main_Before_Calculate_Force() + simulation.Main_Calculate_Force() + simulation.Main_Calculate_Energy() + simulation.Main_After_Calculate_Energy() + temperature = simulation.Main_Print() + simulation.Main_Iteration_2() + end = time.time() + print("Main time(s):", end - start) + simulation.Main_Destroy() diff --git a/model_zoo/research/hpc/sponge/src/Langevin_Liujian_md.py b/model_zoo/research/hpc/sponge/src/Langevin_Liujian_md.py new file mode 100644 index 0000000000..222b6db13c --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/Langevin_Liujian_md.py @@ -0,0 +1,109 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Langevin Liujian MD class""" + +import math + +import numpy as np +from mindspore.common.tensor import Tensor +import mindspore.common.dtype as mstype + + +class Langevin_Liujian: + """Langevin_Liujian class""" + + def __init__(self, controller, atom_numbers): + self.atom_numbers = atom_numbers + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + + self.CONSTANT_TIME_CONVERTION = 20.455 + self.CONSTANT_kB = 0.00198716 + + self.target_temperature = 300.0 if "target_temperature" not in controller.Command_Set else float( + controller.Command_Set["target_temperature"]) + self.gamma_ln = 1.0 if "langevin_gamma" not in controller.Command_Set else float( + controller.Command_Set["langevin_gamma"]) + self.rand_seed = 0 if "langevin_seed" not in controller.Command_Set else float( + controller.Command_Set["langevin_seed"]) # jiahong0315 + self.max_velocity = 10000.0 if "velocity_max" not in controller.Command_Set else float( + controller.Command_Set["velocity_max"]) + assert self.max_velocity > 0 + self.is_max_velocity = 0 if "velocity_max" not in controller.Command_Set else 1 + print("target temperature is ", self.target_temperature) + print("friction coefficient is ", self.gamma_ln, "ps^-1") + print("random seed is ", self.rand_seed) + self.dt = float(controller.Command_Set["dt"]) + self.dt *= self.CONSTANT_TIME_CONVERTION + self.half_dt = 0.5 * self.dt + self.float4_numbers = math.ceil(3.0 * self.atom_numbers / 4.0) + self.gamma_ln = self.gamma_ln / self.CONSTANT_TIME_CONVERTION + self.exp_gamma = math.exp(-1 * self.gamma_ln * self.dt) + self.sqrt_gamma = math.sqrt((1. - self.exp_gamma * self.exp_gamma) * self.target_temperature * self.CONSTANT_kB) + self.h_sqrt_mass = [0] * self.atom_numbers + for i in range(self.atom_numbers): + self.h_sqrt_mass[i] = self.sqrt_gamma * math.sqrt(1. / self.h_mass[i]) + self.d_sqrt_mass = Tensor(self.h_sqrt_mass, mstype.float32) + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + self.h_mass = [0] * self.atom_numbers + for idx, val in enumerate(context): + if "%FLAG MASS" in val: + count = 0 + start_idx = idx + information = [] + while count < self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + + for i in range(self.atom_numbers): + self.h_mass[i] = information[i] + break + + def MDIterationLeapFrog_Liujian(self, atom_numbers, half_dt, dt, exp_gamma, inverse_mass, sqrt_mass_inverse, vel, + crd, frc, random_frc): + """compute MDIterationLeapFrog Liujian""" + inverse_mass = inverse_mass.reshape((-1, 1)) + sqrt_mass_inverse = sqrt_mass_inverse.reshape((-1, 1)) + acc = inverse_mass * frc + vel = vel + dt * acc + crd = crd + half_dt * vel + vel = exp_gamma * vel + sqrt_mass_inverse * random_frc + crd = crd + half_dt * vel + frc = Tensor(np.zeros((atom_numbers, 3)), mstype.float32) + return vel, crd, frc, acc + + def MD_Iteration_Leap_Frog(self, d_mass_inverse, vel_in, crd_in, frc_in): + """MD_Iteration_Leap_Frog""" + np.random.seed(int(self.rand_seed)) + self.rand_force = Tensor(np.zeros((self.atom_numbers, 3)), mstype.float32) + # self.rand_force = Tensor(np.random.randn(self.atom_numbers, 3), mstype.float32) + vel, crd, frc, acc = self.MDIterationLeapFrog_Liujian(atom_numbers=self.atom_numbers, half_dt=self.half_dt, + dt=self.dt, exp_gamma=self.exp_gamma, + inverse_mass=d_mass_inverse, + sqrt_mass_inverse=self.d_sqrt_mass, + vel=vel_in, crd=crd_in, + frc=frc_in, random_frc=self.rand_force) + return vel, crd, frc, acc diff --git a/model_zoo/research/hpc/sponge/src/angle.py b/model_zoo/research/hpc/sponge/src/angle.py new file mode 100644 index 0000000000..f6d00edbdc --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/angle.py @@ -0,0 +1,175 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""angle class""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class Angle(nn.Cell): + """Angle class""" + + def __init__(self, controller): + super(Angle, self).__init__() + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + + self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) + self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) + self.atom_c = Tensor(np.asarray(self.h_atom_c, np.int32), mstype.int32) + self.angle_k = Tensor(np.asarray(self.h_angle_k, np.float32), mstype.float32) + self.angle_theta0 = Tensor(np.asarray(self.h_angle_theta0, np.float32), mstype.float32) + + def read_process1(self, context): + """read_information_from_amberfile process1""" + for idx, val in enumerate(context): + if idx < len(context) - 1: + if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: + start_idx = idx + 2 + count = 0 + value = list(map(int, context[start_idx].strip().split())) + self.angle_with_H_numbers = value[4] + self.angle_without_H_numbers = value[5] + self.angle_numbers = self.angle_with_H_numbers + self.angle_without_H_numbers + information = [] + information.extend(value) + while count < 15: + start_idx += 1 + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.angle_type_numbers = information[16] + print("angle type numbers ", self.angle_type_numbers) + break + + def read_process2(self, context): + """read_information_from_amberfile process2""" + angle_count = 0 + for idx, val in enumerate(context): + if "%FLAG ANGLES_INC_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 4 * self.angle_with_H_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for _ in range(self.angle_with_H_numbers): + self.h_atom_a[angle_count] = information[angle_count * 4 + 0] / 3 + self.h_atom_b[angle_count] = information[angle_count * 4 + 1] / 3 + self.h_atom_c[angle_count] = information[angle_count * 4 + 2] / 3 + self.h_type[angle_count] = information[angle_count * 4 + 3] - 1 + angle_count += 1 + + break + return angle_count + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + self.read_process1(context) + + self.h_atom_a = [0] * self.angle_numbers + self.h_atom_b = [0] * self.angle_numbers + self.h_atom_c = [0] * self.angle_numbers + self.h_type = [0] * self.angle_numbers + angle_count = self.read_process2(context) + + for idx, val in enumerate(context): + if "%FLAG ANGLES_WITHOUT_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 4 * self.angle_without_H_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.angle_without_H_numbers): + self.h_atom_a[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 0] / 3 + self.h_atom_b[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 1] / 3 + self.h_atom_c[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 2] / 3 + self.h_type[angle_count] = information[(angle_count - self.angle_with_H_numbers) * 4 + 3] - 1 + angle_count += 1 + break + + self.type_k = [0] * self.angle_type_numbers + for idx, val in enumerate(context): + if "%FLAG ANGLE_FORCE_CONSTANT" in val: + count = 0 + start_idx = idx + information = [] + while count < self.angle_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + # print(start_idx) + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.type_k = information[:self.angle_type_numbers] + break + + self.type_theta0 = [0] * self.angle_type_numbers + for idx, val in enumerate(context): + if "%FLAG ANGLE_EQUIL_VALUE" in val: + count = 0 + start_idx = idx + information = [] + while count < self.angle_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.type_theta0 = information[:self.angle_type_numbers] + break + if self.angle_numbers != angle_count: + print("angle count %d != angle_number %d ", angle_count, self.angle_numbers) + + self.h_angle_k = [] + self.h_angle_theta0 = [] + for i in range(self.angle_numbers): + self.h_angle_k.append(self.type_k[self.h_type[i]]) + self.h_angle_theta0.append(self.type_theta0[self.h_type[i]]) + + def Angle_Energy(self, uint_crd, uint_dr_to_dr_cof): + """compute angle energy""" + self.angle_energy = P.AngleEnergy(self.angle_numbers)(uint_crd, uint_dr_to_dr_cof, self.atom_a, self.atom_b, + self.atom_c, self.angle_k, self.angle_theta0) + self.sigma_of_angle_ene = P.ReduceSum()(self.angle_energy) + return self.sigma_of_angle_ene + + def Angle_Force_With_Atom_Energy(self, uint_crd, scaler): + """compute angle force with atom energy""" + print("angele angle numbers:", self.angle_numbers) + self.afae = P.AngleForceWithAtomEnergy(angle_numbers=self.angle_numbers) + frc, ene = self.afae(uint_crd, scaler, self.atom_a, self.atom_b, self.atom_c, self.angle_k, self.angle_theta0) + return frc, ene diff --git a/model_zoo/research/hpc/sponge/src/bond.py b/model_zoo/research/hpc/sponge/src/bond.py new file mode 100644 index 0000000000..9fc526d6f0 --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/bond.py @@ -0,0 +1,163 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""bond class""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class Bond(nn.Cell): + """bond class""" + + def __init__(self, controller, md_info): + super(Bond, self).__init__() + + self.atom_numbers = md_info.atom_numbers + + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + + self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) + self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) + self.bond_k = Tensor(np.asarray(self.h_k, np.float32), mstype.float32) + self.bond_r0 = Tensor(np.asarray(self.h_r0, np.float32), mstype.float32) + + def process1(self, context): + """process1: read information from amberfile""" + for idx, val in enumerate(context): + if idx < len(context) - 1: + if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: + start_idx = idx + 2 + count = 0 + value = list(map(int, context[start_idx].strip().split())) + self.bond_with_hydrogen = value[2] + self.bond_numbers = value[3] + self.bond_numbers += self.bond_with_hydrogen + print(self.bond_numbers) + information = [] + information.extend(value) + while count < 16: + start_idx += 1 + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.bond_type_numbers = information[15] + print("bond type numbers ", self.bond_type_numbers) + break + + for idx, val in enumerate(context): + if "%FLAG BOND_FORCE_CONSTANT" in val: + count = 0 + start_idx = idx + information = [] + while count < self.bond_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.bond_type_k = information[:self.bond_type_numbers] + break + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + self.process1(context) + + for idx, val in enumerate(context): + if "%FLAG BOND_EQUIL_VALUE" in val: + count = 0 + start_idx = idx + information = [] + while count < self.bond_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.bond_type_r = information[:self.bond_type_numbers] + break + + for idx, val in enumerate(context): + if "%FLAG BONDS_INC_HYDROGEN" in val: + self.h_atom_a = [0] * self.bond_numbers + self.h_atom_b = [0] * self.bond_numbers + self.h_k = [0] * self.bond_numbers + self.h_r0 = [0] * self.bond_numbers + + count = 0 + start_idx = idx + information = [] + while count < 3 * self.bond_with_hydrogen: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + + for i in range(self.bond_with_hydrogen): + self.h_atom_a[i] = information[3 * i + 0] / 3 + self.h_atom_b[i] = information[3 * i + 1] / 3 + tmpi = information[3 * i + 2] - 1 + self.h_k[i] = self.bond_type_k[tmpi] + self.h_r0[i] = self.bond_type_r[tmpi] + break + + for idx, val in enumerate(context): + if "%FLAG BONDS_WITHOUT_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 3 * (self.bond_numbers - self.bond_with_hydrogen): + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + + for i in range(self.bond_with_hydrogen, self.bond_numbers): + self.h_atom_a[i] = information[3 * (i - self.bond_with_hydrogen) + 0] / 3 + self.h_atom_b[i] = information[3 * (i - self.bond_with_hydrogen) + 1] / 3 + tmpi = information[3 * (i - self.bond_with_hydrogen) + 2] - 1 + self.h_k[i] = self.bond_type_k[tmpi] + self.h_r0[i] = self.bond_type_r[tmpi] + break + + def Bond_Energy(self, uint_crd, uint_dr_to_dr_cof): + """compute bond energy""" + self.bond_energy = P.BondEnergy(self.bond_numbers, self.atom_numbers)(uint_crd, uint_dr_to_dr_cof, self.atom_a, + self.atom_b, self.bond_k, self.bond_r0) + self.sigma_of_bond_ene = P.ReduceSum()(self.bond_energy) + return self.sigma_of_bond_ene + + def Bond_Force_With_Atom_Energy(self, uint_crd, scaler): + """compute bond force with atom energy""" + self.bfatomenergy = P.BondForceWithAtomEnergy(bond_numbers=self.bond_numbers, + atom_numbers=self.atom_numbers) + frc, atom_energy = self.bfatomenergy(uint_crd, scaler, self.atom_a, self.atom_b, self.bond_k, self.bond_r0) + return frc, atom_energy diff --git a/model_zoo/research/hpc/sponge/src/dihedral.py b/model_zoo/research/hpc/sponge/src/dihedral.py new file mode 100644 index 0000000000..5529a5ad29 --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/dihedral.py @@ -0,0 +1,221 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""dihedral class""" + +import math +import numpy as np + +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class Dihedral(nn.Cell): + """dihedral class""" + + def __init__(self, controller): + super(Dihedral, self).__init__() + self.CONSTANT_Pi = 3.1415926535897932 + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) + self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) + self.atom_c = Tensor(np.asarray(self.h_atom_c, np.int32), mstype.int32) + self.atom_d = Tensor(np.asarray(self.h_atom_d, np.int32), mstype.int32) + + self.pk = Tensor(np.asarray(self.pk, np.float32), mstype.float32) + self.gamc = Tensor(np.asarray(self.gamc, np.float32), mstype.float32) + self.gams = Tensor(np.asarray(self.gams, np.float32), mstype.float32) + self.pn = Tensor(np.asarray(self.pn, np.float32), mstype.float32) + self.ipn = Tensor(np.asarray(self.ipn, np.int32), mstype.int32) + + def process1(self, context): + """process1: read information from amberfile""" + for idx, val in enumerate(context): + if idx < len(context) - 1: + if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: + start_idx = idx + 2 + count = 0 + value = list(map(int, context[start_idx].strip().split())) + self.dihedral_with_hydrogen = value[6] + self.dihedral_numbers = value[7] + self.dihedral_numbers += self.dihedral_with_hydrogen + information = [] + information.extend(value) + while count < 15: + start_idx += 1 + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.dihedral_type_numbers = information[17] + print("dihedral type numbers ", self.dihedral_type_numbers) + break + + self.phase_type = [0] * self.dihedral_type_numbers + self.pk_type = [0] * self.dihedral_type_numbers + self.pn_type = [0] * self.dihedral_type_numbers + + for idx, val in enumerate(context): + if "%FLAG DIHEDRAL_FORCE_CONSTANT" in val: + count = 0 + start_idx = idx + information = [] + while count < self.dihedral_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.pk_type = information[:self.dihedral_type_numbers] + break + + for idx, val in enumerate(context): + if "%FLAG DIHEDRAL_PHASE" in val: + count = 0 + start_idx = idx + information = [] + while count < self.dihedral_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.phase_type = information[:self.dihedral_type_numbers] + break + + for idx, val in enumerate(context): + if "%FLAG DIHEDRAL_PERIODICITY" in val: + count = 0 + start_idx = idx + information = [] + while count < self.dihedral_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.pn_type = information[:self.dihedral_type_numbers] + break + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + + self.process1(context) + + self.h_atom_a = [0] * self.dihedral_numbers + self.h_atom_b = [0] * self.dihedral_numbers + self.h_atom_c = [0] * self.dihedral_numbers + self.h_atom_d = [0] * self.dihedral_numbers + self.pk = [] + self.gamc = [] + self.gams = [] + self.pn = [] + self.ipn = [] + for idx, val in enumerate(context): + if "%FLAG DIHEDRALS_INC_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 5 * self.dihedral_with_hydrogen: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.dihedral_with_hydrogen): + self.h_atom_a[i] = information[i * 5 + 0] / 3 + self.h_atom_b[i] = information[i * 5 + 1] / 3 + self.h_atom_c[i] = information[i * 5 + 2] / 3 + self.h_atom_d[i] = abs(information[i * 5 + 3] / 3) + tmpi = information[i * 5 + 4] - 1 + self.pk.append(self.pk_type[tmpi]) + tmpf = self.phase_type[tmpi] + if abs(tmpf - self.CONSTANT_Pi) <= 0.001: + tmpf = self.CONSTANT_Pi + tmpf2 = math.cos(tmpf) + if abs(tmpf2) < 1e-6: + tmpf2 = 0 + self.gamc.append(tmpf2 * self.pk[i]) + tmpf2 = math.sin(tmpf) + if abs(tmpf2) < 1e-6: + tmpf2 = 0 + self.gams.append(tmpf2 * self.pk[i]) + self.pn.append(abs(self.pn_type[tmpi])) + self.ipn.append(int(self.pn[i] + 0.001)) + break + for idx, val in enumerate(context): + if "%FLAG DIHEDRALS_WITHOUT_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 5 * (self.dihedral_numbers - self.dihedral_with_hydrogen): + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.dihedral_with_hydrogen, self.dihedral_numbers): + self.h_atom_a[i] = information[(i - self.dihedral_with_hydrogen) * 5 + 0] / 3 + self.h_atom_b[i] = information[(i - self.dihedral_with_hydrogen) * 5 + 1] / 3 + self.h_atom_c[i] = information[(i - self.dihedral_with_hydrogen) * 5 + 2] / 3 + self.h_atom_d[i] = abs(information[(i - self.dihedral_with_hydrogen) * 5 + 3] / 3) + tmpi = information[(i - self.dihedral_with_hydrogen) * 5 + 4] - 1 + self.pk.append(self.pk_type[tmpi]) + tmpf = self.phase_type[tmpi] + if abs(tmpf - self.CONSTANT_Pi) <= 0.001: + tmpf = self.CONSTANT_Pi + tmpf2 = math.cos(tmpf) + if abs(tmpf2) < 1e-6: + tmpf2 = 0 + self.gamc.append(tmpf2 * self.pk[i]) + tmpf2 = math.sin(tmpf) + if abs(tmpf2) < 1e-6: + tmpf2 = 0 + self.gams.append(tmpf2 * self.pk[i]) + self.pn.append(abs(self.pn_type[tmpi])) + self.ipn.append(int(self.pn[i] + 0.001)) + break + for i in range(self.dihedral_numbers): + if self.h_atom_c[i] < 0: + self.h_atom_c[i] *= -1 + + def Dihedral_Engergy(self, uint_crd, uint_dr_to_dr_cof): + """compute dihedral energy""" + self.dihedral_energy = P.DihedralEnergy(self.dihedral_numbers)(uint_crd, uint_dr_to_dr_cof, self.atom_a, + self.atom_b, self.atom_c, self.atom_d, self.ipn, + self.pk, self.gamc, self.gams, self.pn) + self.sigma_of_dihedral_ene = P.ReduceSum()(self.dihedral_energy) + return self.sigma_of_dihedral_ene + + def Dihedral_Force_With_Atom_Energy(self, uint_crd, scaler): + """compute dihedral force and atom energy""" + self.dfae = P.DihedralForceWithAtomEnergy(dihedral_numbers=self.dihedral_numbers) + self.frc, self.ene = self.dfae(uint_crd, scaler, self.atom_a, self.atom_b, self.atom_c, self.atom_d, + self.ipn, self.pk, self.gamc, self.gams, self.pn) + return self.frc, self.ene diff --git a/model_zoo/research/hpc/sponge/src/lennard_jones.py b/model_zoo/research/hpc/sponge/src/lennard_jones.py new file mode 100644 index 0000000000..4be20eeccd --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/lennard_jones.py @@ -0,0 +1,122 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""lennard jones""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class Lennard_Jones_Information(nn.Cell): + """class Lennard Jones Information""" + + def __init__(self, controller): + super(Lennard_Jones_Information, self).__init__() + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + self.atom_LJ_type = Tensor(np.asarray(self.atom_LJ_type, dtype=np.int32), mstype.int32) + self.LJ_A = Tensor(np.asarray(self.LJ_A, dtype=np.float32), mstype.float32) + self.LJ_B = Tensor(np.asarray(self.LJ_B, dtype=np.float32), mstype.float32) + self.LJ_energy_sum = 0 + self.LJ_energy = 0 + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + + for idx, val in enumerate(context): + if idx < len(context) - 1: + if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: + start_idx = idx + 2 + count = 0 + value = list(map(int, context[start_idx].strip().split())) + self.atom_numbers = value[0] + self.atom_type_numbers = value[1] + self.pair_type_numbers = int(self.atom_type_numbers * (self.atom_type_numbers + 1) / 2) + print(self.pair_type_numbers) + break + self.atom_LJ_type = [0] * self.atom_numbers + for idx, val in enumerate(context): + if "%FLAG ATOM_TYPE_INDEX" in val: + count = 0 + start_idx = idx + information = [] + while count < self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.atom_numbers): + self.atom_LJ_type[i] = information[i] - 1 + break + self.LJ_A = [0] * self.pair_type_numbers + for idx, val in enumerate(context): + if "%FLAG LENNARD_JONES_ACOEF" in val: + count = 0 + start_idx = idx + information = [] + while count < self.pair_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.pair_type_numbers): + self.LJ_A[i] = 12.0 * information[i] + break + self.LJ_B = [0] * self.pair_type_numbers + for idx, val in enumerate(context): + if "%FLAG LENNARD_JONES_BCOEF" in val: + count = 0 + start_idx = idx + information = [] + while count < self.pair_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.pair_type_numbers): + self.LJ_B[i] = 6.0 * information[i] + break + + def LJ_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof, nl_atom_numbers, nl_atom_serial, cutoff_square): + """compute LJ energy""" + uint_crd, LJtype, charge = uint_crd_with_LJ + self.LJ_energy = P.LJEnergy(self.atom_numbers, cutoff_square) \ + (uint_crd, LJtype, charge, uint_dr_to_dr_cof, nl_atom_numbers, nl_atom_serial, self.LJ_A, self.LJ_B) + self.LJ_energy_sum = P.ReduceSum()(self.LJ_energy) + return self.LJ_energy_sum + + def LJ_Force_With_PME_Direct_Force(self, atom_numbers, uint_crd_with_LJ, uint_dr_to_dr_cof, nl_number, nl_serial, + cutoff, beta): + """compute LJ force with PME direct force""" + assert atom_numbers == self.atom_numbers + assert isinstance(uint_crd_with_LJ, tuple) + uint_crd_f, LJtype, charge = uint_crd_with_LJ + self.ljfd = P.LJForceWithPMEDirectForce(atom_numbers, cutoff, beta) + frc = self.ljfd(uint_crd_f, LJtype, charge, uint_dr_to_dr_cof, nl_number, nl_serial, self.LJ_A, self.LJ_B) + return frc diff --git a/model_zoo/research/hpc/sponge/src/md_information.py b/model_zoo/research/hpc/sponge/src/md_information.py new file mode 100644 index 0000000000..d9fccbac60 --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/md_information.py @@ -0,0 +1,212 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""md information""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class md_information(nn.Cell): + """class md information""" + + def __init__(self, controller): + super(md_information, self).__init__() + CONSTANT_TIME_CONVERTION = 20.455 + CONSTANT_UINT_MAX_FLOAT = 4294967296.0 + self.md_task = controller.md_task + self.mode = 0 if "mode" not in controller.Command_Set else int(controller.Command_Set["mode"]) + self.dt = 0.001 * CONSTANT_TIME_CONVERTION if "dt" not in controller.Command_Set \ + else float(controller.Command_Set["dt"]) * CONSTANT_TIME_CONVERTION + self.skin = 2.0 if "skin" not in controller.Command_Set \ + else float(controller.Command_Set["skin"]) + self.trans_vec = [self.skin, self.skin, self.skin] + self.trans_vec_minus = -1 * self.trans_vec + self.step_limit = 1000 if "step_limit" not in controller.Command_Set else int( + controller.Command_Set["step_limit"]) + self.netfrc = 0 if "net_force" not in controller.Command_Set else int(controller.Command_Set["net_force"]) + self.ntwx = 1000 if "write_information_interval" not in controller.Command_Set else \ + int(controller.Command_Set["write_information_interval"]) + self.ntce = self.step_limit + 1 if "calculate_energy_interval" not in controller.Command_Set else \ + int(controller.Command_Set["calculate_energy_interval"]) + self.atom_numbers = 0 + self.residue_numbers = 0 + self.density = 0.0 + self.lin_serial = [] + self.h_res_start = [] + self.h_res_end = [] + self.h_mass = [] + self.h_mass_inverse = [] + self.h_charge = [] + self.steps = 0 + + if controller.amber_parm is not None: + self.read_basic_system_information_from_amber_file(controller.amber_parm) + + if "amber_irest" in controller.Command_Set: + amber_irest = int(controller.Command_Set["amber_irest"]) + if controller.initial_coordinates_file is not None: + self.read_basic_system_information_from_rst7(controller.initial_coordinates_file, amber_irest) + + self.crd_to_uint_crd_cof = [CONSTANT_UINT_MAX_FLOAT / self.box_length[0], + CONSTANT_UINT_MAX_FLOAT / self.box_length[1], + CONSTANT_UINT_MAX_FLOAT / self.box_length[2]] + self.uint_dr_to_dr_cof = [1.0 / self.crd_to_uint_crd_cof[0], 1.0 / self.crd_to_uint_crd_cof[1], + 1.0 / self.crd_to_uint_crd_cof[2]] + self.density *= 1e24 / 6.023e23 / (self.box_length[0] * self.box_length[1] * self.box_length[2]) + self.frc = Tensor(np.zeros((self.atom_numbers, 3)), mstype.float32) + self.crd = Tensor(np.array(self.coordinate, dtype=np.float32).reshape((self.atom_numbers, 3)), mstype.float32) + self.crd_n = np.array(self.coordinate).reshape([self.atom_numbers, 3]) + self.crd_old = Tensor(np.zeros([self.atom_numbers, 3], dtype=np.float32), mstype.float32) + self.uint_crd = Tensor(np.zeros([self.atom_numbers, 3], dtype=np.uint32), mstype.uint32) + self.charge = Tensor(self.h_charge, mstype.float32) + self.crd_to_uint_crd_cof_n = np.array(self.crd_to_uint_crd_cof) + self.crd_to_uint_crd_cof = Tensor(self.crd_to_uint_crd_cof, mstype.float32) + self.uint_dr_to_dr_cof = Tensor(self.uint_dr_to_dr_cof, mstype.float32) + self.uint_crd_with_LJ = None + self.d_mass_inverse = Tensor(self.h_mass_inverse, mstype.float32) + self.d_res_start = Tensor(self.h_res_start, mstype.int32) + self.d_res_end = Tensor(self.h_res_end, mstype.int32) + self.d_mass = Tensor(self.h_mass, mstype.float32) + + def process1(self, context): + """process1: read basic system information from amber file""" + for idx, val in enumerate(context): + if idx < len(context) - 1: + if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: + start_idx = idx + 2 + value = list(map(int, context[start_idx].strip().split())) + self.atom_numbers = value[0] + count = len(value) - 1 + while count < 10: + start_idx += 1 + value = list(map(int, context[start_idx].strip().split())) + count += len(value) + self.residue_numbers = list(map(int, context[start_idx].strip().split()))[10 - (count - 10)] + break + + def read_basic_system_information_from_amber_file(self, path): + """read basic system information from amber file""" + file = open(path, 'r') + context = file.readlines() + file.close() + self.process1(context) + + if self.residue_numbers != 0 and self.atom_numbers != 0: + for idx, val in enumerate(context): + if "%FLAG RESIDUE_POINTER" in val: + count = 0 + start_idx = idx + while count != self.residue_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + self.lin_serial.extend(value) + count += len(value) + for i in range(self.residue_numbers - 1): + self.h_res_start.append(self.lin_serial[i] - 1) + self.h_res_end.append(self.lin_serial[i + 1] - 1) + self.h_res_start.append(self.lin_serial[-1] - 1) + self.h_res_end.append(self.atom_numbers + 1 - 1) + break + + for idx, val in enumerate(context): + if "%FLAG MASS" in val: + count = 0 + start_idx = idx + while count != self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + self.h_mass.extend(value) + count += len(value) + for i in range(self.atom_numbers): + if self.h_mass[i] == 0: + self.h_mass_inverse.append(0.0) + else: + self.h_mass_inverse.append(1.0 / self.h_mass[i]) + self.density += self.h_mass[i] + break + for idx, val in enumerate(context): + if "%FLAG CHARGE" in val: + count = 0 + start_idx = idx + while count != self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + self.h_charge.extend(value) + count += len(value) + break + + def read_basic_system_information_from_rst7(self, path, irest): + """read basic system information from rst7""" + file = open(path, 'r') + context = file.readlines() + file.close() + atom_numbers = int(context[1].strip().split()[0]) + if atom_numbers != self.atom_numbers: + print("ERROR") + else: + print("check atom_numbers") + information = [] + count = 0 + start_idx = 1 + if irest == 1: + self.simulation_start_time = float(context[1].strip().split()[1]) + while count <= 6 * self.atom_numbers + 3: + start_idx += 1 + # print(start_idx) + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.coordinate = information[: 3 * self.atom_numbers] + self.velocity = information[3 * self.atom_numbers: 6 * self.atom_numbers] + self.box_length = information[6 * self.atom_numbers:6 * self.atom_numbers + 3] + else: + while count <= 3 * self.atom_numbers + 3: + start_idx += 1 + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.coordinate = information[: 3 * self.atom_numbers] + self.velocity = [0.0] * (3 * self.atom_numbers) + self.box_length = information[3 * self.atom_numbers:3 * self.atom_numbers + 3] + self.vel = Tensor(self.velocity, mstype.float32).reshape((self.atom_numbers, 3)) + self.acc = Tensor(np.zeros((self.atom_numbers, 3), dtype=np.float32), mstype.float32) + + def MD_Information_Crd_To_Uint_Crd(self): + """transform the crd to uint crd""" + uint_crd = self.crd.asnumpy() * (0.5 * self.crd_to_uint_crd_cof.asnumpy()) * 2 + self.uint_crd = Tensor(uint_crd, mstype.uint32) + + return self.uint_crd + + def Centerize(self): + return + + def MD_Information_Temperature(self): + """compute temperature""" + self.mdtemp = P.MDTemperature(self.residue_numbers, self.atom_numbers) + self.res_ek_energy = self.mdtemp(self.d_res_start, self.d_res_end, self.vel, self.d_mass) + self.d_temperature = P.ReduceSum()(self.res_ek_energy) + return self.d_temperature diff --git a/model_zoo/research/hpc/sponge/src/nb14.py b/model_zoo/research/hpc/sponge/src/nb14.py new file mode 100644 index 0000000000..01f8f9213e --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/nb14.py @@ -0,0 +1,195 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""nb14""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class NON_BOND_14(nn.Cell): + """class Non bond 14""" + + def __init__(self, controller, dihedral, atom_numbers): + super(NON_BOND_14, self).__init__() + + self.dihedral_with_hydrogen = dihedral.dihedral_with_hydrogen + self.dihedral_numbers = dihedral.dihedral_numbers + self.dihedral_type_numbers = dihedral.dihedral_type_numbers + self.atom_numbers = atom_numbers + + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + + self.atom_a = Tensor(np.asarray(self.h_atom_a, np.int32), mstype.int32) + self.atom_b = Tensor(np.asarray(self.h_atom_b, np.int32), mstype.int32) + self.lj_scale_factor = Tensor(np.asarray(self.h_lj_scale_factor, np.float32), mstype.float32) + self.cf_scale_factor = Tensor(np.asarray(self.h_cf_scale_factor, np.float32), mstype.float32) + + def process1(self, context): + """process1: read information from amberfile""" + for idx, val in enumerate(context): + if "%FLAG SCEE_SCALE_FACTOR" in val: + count = 0 + start_idx = idx + information = [] + while count < self.dihedral_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.cf_scale_type = information[:self.dihedral_type_numbers] + break + + for idx, val in enumerate(context): + if "%FLAG SCNB_SCALE_FACTOR" in val: + count = 0 + start_idx = idx + information = [] + while count < self.dihedral_type_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(float, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.lj_scale_type = information[:self.dihedral_type_numbers] + break + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + + self.cf_scale_type = [0] * self.dihedral_type_numbers + self.lj_scale_type = [0] * self.dihedral_type_numbers + + self.h_atom_a = [0] * self.dihedral_numbers + self.h_atom_b = [0] * self.dihedral_numbers + self.h_lj_scale_factor = [0] * self.dihedral_numbers + self.h_cf_scale_factor = [0] * self.dihedral_numbers + nb14_numbers = 0 + for idx, val in enumerate(context): + if "%FLAG DIHEDRALS_INC_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 5 * self.dihedral_with_hydrogen: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.dihedral_with_hydrogen): + tempa = information[i * 5 + 0] + tempi = information[i * 5 + 1] + tempi2 = information[i * 5 + 2] + tempb = information[i * 5 + 3] + tempi = information[i * 5 + 4] + + tempi -= 1 + if tempi2 > 0: + self.h_atom_a[nb14_numbers] = tempa / 3 + self.h_atom_b[nb14_numbers] = abs(tempb / 3) + self.h_lj_scale_factor[nb14_numbers] = self.lj_scale_type[tempi] + if self.h_lj_scale_factor[nb14_numbers] != 0: + self.h_lj_scale_factor[nb14_numbers] = 1.0 / self.h_lj_scale_factor[nb14_numbers] + self.h_cf_scale_factor[nb14_numbers] = self.cf_scale_type[tempi] + if self.h_cf_scale_factor[nb14_numbers] != 0: + self.h_cf_scale_factor[nb14_numbers] = 1.0 / self.h_cf_scale_factor[nb14_numbers] + nb14_numbers += 1 + break + for idx, val in enumerate(context): + if "%FLAG DIHEDRALS_WITHOUT_HYDROGEN" in val: + count = 0 + start_idx = idx + information = [] + while count < 5 * (self.dihedral_numbers - self.dihedral_with_hydrogen): + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + for i in range(self.dihedral_with_hydrogen, self.dihedral_numbers): + tempa = information[(i - self.dihedral_with_hydrogen) * 5 + 0] + tempi = information[(i - self.dihedral_with_hydrogen) * 5 + 1] + tempi2 = information[(i - self.dihedral_with_hydrogen) * 5 + 2] + tempb = information[(i - self.dihedral_with_hydrogen) * 5 + 3] + tempi = information[(i - self.dihedral_with_hydrogen) * 5 + 4] + + tempi -= 1 + if tempi2 > 0: + self.h_atom_a[nb14_numbers] = tempa / 3 + self.h_atom_b[nb14_numbers] = abs(tempb / 3) + self.h_lj_scale_factor[nb14_numbers] = self.lj_scale_type[tempi] + if self.h_lj_scale_factor[nb14_numbers] != 0: + self.h_lj_scale_factor[nb14_numbers] = 1.0 / self.h_lj_scale_factor[nb14_numbers] + self.h_cf_scale_factor[nb14_numbers] = self.cf_scale_type[tempi] + if self.h_cf_scale_factor[nb14_numbers] != 0: + self.h_cf_scale_factor[nb14_numbers] = 1.0 / self.h_cf_scale_factor[nb14_numbers] + nb14_numbers += 1 + break + + self.nb14_numbers = nb14_numbers + + def Non_Bond_14_LJ_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof, LJ_A, LJ_B): + """compute Non bond 14 LJ energy""" + assert isinstance(uint_crd_with_LJ, tuple) + uint_crd, LJtype, charge = uint_crd_with_LJ + self.LJ_energy = P.Dihedral14LJEnergy(self.nb14_numbers, self.atom_numbers)(uint_crd, LJtype, charge, + uint_dr_to_dr_cof, self.atom_a, + self.atom_b, self.lj_scale_factor, + LJ_A, LJ_B) + self.nb14_lj_energy_sum = P.ReduceSum()(self.LJ_energy) + + return self.nb14_lj_energy_sum + + def Non_Bond_14_CF_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof): + """compute Non bond 14 CF energy""" + assert isinstance(uint_crd_with_LJ, tuple) + uint_crd, LJtype, charge = uint_crd_with_LJ + self.CF_energy = P.Dihedral14CFEnergy(self.nb14_numbers, self.atom_numbers)(uint_crd, LJtype, charge, + uint_dr_to_dr_cof, self.atom_a, + self.atom_b, self.cf_scale_factor) + self.nb14_cf_energy_sum = P.ReduceSum()(self.CF_energy) + return self.nb14_cf_energy_sum + + def Non_Bond_14_LJ_CF_Energy(self, uint_crd_with_LJ, uint_dr_to_dr_cof, LJ_A, LJ_B): + """compute Non bond 14 LJ and CF energy""" + assert isinstance(uint_crd_with_LJ, tuple) + self.nb14_lj_energy_sum = self.Non_Bond_14_LJ_Energy(uint_crd_with_LJ, uint_dr_to_dr_cof, LJ_A, LJ_B) + self.nb14_cf_energy_sum = self.Non_Bond_14_CF_Energy(uint_crd_with_LJ, uint_dr_to_dr_cof) + + return self.nb14_lj_energy_sum, self.nb14_cf_energy_sum + + def Non_Bond_14_LJ_CF_Force_With_Atom_Energy(self, uint_crd_with_LJ, boxlength, LJ_A, LJ_B): + """compute Non bond 14 LJ CF force and atom energy""" + self.d14lj = P.Dihedral14LJCFForceWithAtomEnergy(nb14_numbers=self.nb14_numbers, atom_numbers=self.atom_numbers) + assert isinstance(uint_crd_with_LJ, tuple) + uint_crd_f, LJtype, charge = uint_crd_with_LJ + self.frc, self.atom_ene = self.d14lj(uint_crd_f, LJtype, charge, boxlength, self.atom_a, self.atom_b, + self.lj_scale_factor, self.cf_scale_factor, LJ_A, LJ_B) + return self.frc, self.atom_ene diff --git a/model_zoo/research/hpc/sponge/src/neighbor_list.py b/model_zoo/research/hpc/sponge/src/neighbor_list.py new file mode 100644 index 0000000000..d9991fa849 --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/neighbor_list.py @@ -0,0 +1,207 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""neighbour list""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class nb_infomation(nn.Cell): + """neighbour list""" + + def __init__(self, controller, atom_numbers, box_length): + super(nb_infomation, self).__init__() + self.refresh_interval = 20 if "neighbor_list_refresh_interval" not in controller.Command_Set else \ + int(controller.Command_Set["neighbor_list_refresh_interval"]) + self.max_atom_in_grid_numbers = 64 if "max_atom_in_grid_numbers" not in controller.Command_Set else \ + int(controller.Command_Set["max_atom_in_grid_numbers"]) + self.max_neighbor_numbers = 800 if "max_neighbor_numbers" not in controller.Command_Set else \ + int(controller.Command_Set["max_neighbor_numbers"]) + self.skin = 2.0 if "skin" not in controller.Command_Set else float(controller.Command_Set["skin"]) + self.cutoff = 10.0 if "cut" not in controller.Command_Set else float(controller.Command_Set["cut"]) + self.cutoff_square = self.cutoff * self.cutoff + self.cutoff_with_skin = self.cutoff + self.skin + self.half_cutoff_with_skin = 0.5 * self.cutoff_with_skin + self.cutoff_with_skin_square = self.cutoff_with_skin * self.cutoff_with_skin + self.half_skin_square = 0.25 * self.skin * self.skin + self.atom_numbers = atom_numbers + self.box_length = box_length + + if controller.amber_parm is not None: + file_path = controller.amber_parm + self.read_information_from_amberfile(file_path) + + self.Initial_Neighbor_Grid() + self.not_first_time = 0 + self.refresh_count = 0 + + self.atom_numbers_in_grid_bucket = Tensor(np.asarray(self.atom_numbers_in_grid_bucket, np.int32), mstype.int32) + self.bucket = Tensor( + np.asarray(self.bucket, np.int32).reshape([self.grid_numbers, self.max_atom_in_grid_numbers]), mstype.int32) + self.grid_N = Tensor(np.asarray(self.grid_N, np.int32), mstype.int32) + self.grid_length_inverse = Tensor(np.asarray(self.grid_length_inverse, np.float32), mstype.float32) + self.atom_in_grid_serial = Tensor(np.zeros(self.atom_numbers, np.int32), mstype.int32) + self.pointer = Tensor(np.asarray(self.pointer, np.int32).reshape([self.grid_numbers, 125]), mstype.int32) + self.nl_atom_numbers = Tensor(np.zeros(self.atom_numbers, np.int32), mstype.int32) + self.nl_atom_serial = Tensor(np.zeros([self.atom_numbers, self.max_neighbor_numbers], np.int32), mstype.int32) + self.excluded_list_start = Tensor(np.asarray(self.excluded_list_start, np.int32), mstype.int32) + self.excluded_list = Tensor(np.asarray(self.excluded_list, np.int32), mstype.int32) + self.excluded_numbers = Tensor(np.asarray(self.excluded_numbers, np.int32), mstype.int32) + self.need_refresh_flag = Tensor(np.asarray([0], np.int32), mstype.int32) + + def read_information_from_amberfile(self, file_path): + """read information from amberfile""" + file = open(file_path, 'r') + context = file.readlines() + file.close() + self.excluded_list_start = [0] * self.atom_numbers + self.excluded_numbers = [0] * self.atom_numbers + + for idx, val in enumerate(context): + if idx < len(context) - 1: + if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]: + start_idx = idx + 2 + count = 0 + value = list(map(int, context[start_idx].strip().split())) + information = [] + information.extend(value) + while count < 11: + start_idx += 1 + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + self.excluded_atom_numbers = information[10] + break + for idx, val in enumerate(context): + if "%FLAG NUMBER_EXCLUDED_ATOMS" in val: + count = 0 + start_idx = idx + information = [] + while count < self.atom_numbers: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + count = 0 + for i in range(self.atom_numbers): + self.excluded_numbers[i] = information[i] + self.excluded_list_start[i] = count + count += information[i] + break + + total_count = sum(self.excluded_numbers) + self.excluded_list = [] + for idx, val in enumerate(context): + if "%FLAG EXCLUDED_ATOMS_LIST" in val: + count = 0 + start_idx = idx + information = [] + while count < total_count: + start_idx += 1 + if "%FORMAT" in context[start_idx]: + continue + else: + value = list(map(int, context[start_idx].strip().split())) + information.extend(value) + count += len(value) + + count = 0 + for i in range(self.atom_numbers): + tmp_list = [] + for _ in range(self.excluded_numbers[i]): + tmp_list.append(information[count] - 1) + count += 1 + tmp_list = sorted(tmp_list) + self.excluded_list.extend(tmp_list) + break + + def fun(self, Nx, Ny, Nz, l, m, temp_grid_serial, count): + """fun to replace the for""" + for n in range(-2, 3): + xx = Nx + l + if xx < 0: + xx = xx + self.Nx + elif xx >= self.Nx: + xx = xx - self.Nx + yy = Ny + m + if yy < 0: + yy = yy + self.Ny + elif yy >= self.Ny: + yy = yy - self.Ny + zz = Nz + n + if zz < 0: + zz = zz + self.Nz + elif zz >= self.Nz: + zz = zz - self.Nz + temp_grid_serial[count] = zz * self.Nxy + yy * self.Nx + xx + count += 1 + return temp_grid_serial, count + + def Initial_Neighbor_Grid(self): + """initial neighbour grid""" + half_cutoff = self.half_cutoff_with_skin + self.Nx = int(self.box_length[0] / half_cutoff) + self.Ny = int(self.box_length[1] / half_cutoff) + self.Nz = int(self.box_length[2] / half_cutoff) + self.grid_N = [self.Nx, self.Ny, self.Nz] + self.grid_length = [self.box_length[0] / self.Nx, self.box_length[1] / self.Ny, self.box_length[2] / self.Nz] + self.grid_length_inverse = [1.0 / self.grid_length[0], 1.0 / self.grid_length[1], 1.0 / self.grid_length[2]] + self.Nxy = self.Nx * self.Ny + self.grid_numbers = self.Nz * self.Nxy + + self.atom_numbers_in_grid_bucket = [0] * self.grid_numbers + self.bucket = [-1] * (self.grid_numbers * self.max_atom_in_grid_numbers) + self.pointer = [] + temp_grid_serial = [0] * 125 + for i in range(self.grid_numbers): + Nz = int(i / self.Nxy) + Ny = int((i - self.Nxy * Nz) / self.Nx) + Nx = i - self.Nxy * Nz - self.Nx * Ny + count = 0 + for l in range(-2, 3): + for m in range(-2, 3): + temp_grid_serial, count = self.fun(Nx, Ny, Nz, l, m, temp_grid_serial, count) + temp_grid_serial = sorted(temp_grid_serial) + self.pointer.extend(temp_grid_serial) + + def NeighborListUpdate(self, crd, old_crd, uint_crd, crd_to_uint_crd_cof, uint_dr_to_dr_cof, box_length, + not_first_time=0): + """NeighborList Update""" + self.not_first_time = not_first_time + self.neighbor_list_update = P.NeighborListUpdate(grid_numbers=self.grid_numbers, atom_numbers=self.atom_numbers, + refresh_count=self.refresh_count, + not_first_time=self.not_first_time, + Nxy=self.Nxy, excluded_atom_numbers=self.excluded_atom_numbers, + cutoff_square=self.cutoff_square, + half_skin_square=self.half_skin_square, + cutoff_with_skin=self.cutoff_with_skin, + half_cutoff_with_skin=self.half_cutoff_with_skin, + cutoff_with_skin_square=self.cutoff_with_skin_square, + refresh_interval=self.refresh_interval, cutoff=self.cutoff, + skin=self.skin, + max_atom_in_grid_numbers=self.max_atom_in_grid_numbers, + max_neighbor_numbers=self.max_neighbor_numbers) + + res = self.neighbor_list_update(self.atom_numbers_in_grid_bucket, self.bucket, crd, box_length, self.grid_N, + self.grid_length_inverse, self.atom_in_grid_serial, old_crd, + crd_to_uint_crd_cof, uint_crd, self.pointer, self.nl_atom_numbers, + self.nl_atom_serial, uint_dr_to_dr_cof, self.excluded_list_start, + self.excluded_list, self.excluded_numbers, self.need_refresh_flag) + return res diff --git a/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py b/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py new file mode 100644 index 0000000000..80130940a6 --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/particle_mesh_ewald.py @@ -0,0 +1,146 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""PME""" + +import math +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn +from mindspore.ops import operations as P + + +class Particle_Mesh_Ewald(nn.Cell): + """class Particle_Mesh_Ewald""" + + def __init__(self, controller, md_info): + super(Particle_Mesh_Ewald, self).__init__() + self.cutoff = 10.0 if "cut" not in controller.Command_Set else float(controller.Command_Set["cut"]) + self.tolerance = 0.00001 if "PME_Direct_Tolerance" not in controller.Command_Set else float( + controller.Command_Set["PME_Direct_Tolerance"]) + self.fftx = -1 if "fftx" not in controller.Command_Set else int(controller.Command_Set["fftx"]) + self.ffty = -1 if "ffty" not in controller.Command_Set else int(controller.Command_Set["ffty"]) + self.fftz = -1 if "fftz" not in controller.Command_Set else int(controller.Command_Set["fftz"]) + self.atom_numbers = md_info.atom_numbers + self.box_length = md_info.box_length + + if self.fftx < 0: + self.fftx = self.Get_Fft_Patameter(self.box_length[0]) + if self.ffty < 0: + self.ffty = self.Get_Fft_Patameter(self.box_length[1]) + if self.fftz < 0: + self.fftz = self.Get_Fft_Patameter(self.box_length[2]) + + self.beta = self.Get_Beta(self.cutoff, self.tolerance) + self.box_length = Tensor(np.asarray(self.box_length, np.float32), mstype.float32) + + print("========== ", self.fftx, self.ffty, self.fftz, self.tolerance, self.beta) + + def Get_Beta(self, cutoff, tolerance): + """get beta""" + high = 1.0 + ihigh = 1 + while 1: + tempf = math.erfc(high * cutoff) / cutoff + if tempf <= tolerance: + break + high *= 2 + ihigh += 1 + ihigh += 50 + low = 0.0 + for _ in range(1, ihigh): + beta = (low + high) / 2 + tempf = math.erfc(beta * cutoff) / cutoff + if tempf >= tolerance: + low = beta + else: + high = beta + return beta + + def Check_2357_Factor(self, number): + """check 2357 factor""" + while number > 0: + if number == 1: + return 1 + tempn = number / 2 + if tempn * 2 != number: + break + number = tempn + while number > 0: + if number == 1: + return 1 + tempn = number / 3 + if tempn * 3 != number: + break + number = tempn + while number > 0: + if number == 1: + return 1 + tempn = number / 5 + if tempn * 5 != number: + break + number = tempn + while number > 0: + if number == 1: + return 1 + tempn = number / 7 + if tempn * 7 != number: + break + number = tempn + return 0 + + def Get_Fft_Patameter(self, length): + """get fft parameter""" + tempi = math.ceil(length + 3) >> 2 << 2 + if 60 <= tempi <= 68: + tempi = 64 + elif 120 <= tempi <= 136: + tempi = 128 + elif 240 <= tempi <= 272: + tempi = 256 + elif 480 <= tempi <= 544: + tempi = 512 + elif 960 <= tempi <= 1088: + tempi = 1024 + while 1: + if self.Check_2357_Factor(tempi): + return tempi + tempi += 4 + + def PME_Energy(self, uint_crd, charge, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof, excluded_list_start, + excluded_list, excluded_numbers): + """PME_Energy""" + self.pmee = P.PMEEnergy(self.atom_numbers, self.beta, self.fftx, self.ffty, self.fftz) + self.reciprocal_energy, self.self_energy, self.direct_energy, self.correction_energy = \ + self.pmee(self.box_length, uint_crd, charge, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof, + excluded_list_start, excluded_list, excluded_numbers) + return self.reciprocal_energy, self.self_energy, self.direct_energy, self.correction_energy + + def PME_Excluded_Force(self, uint_crd, scaler, charge, excluded_list_start, excluded_list, + excluded_numbers): + """PME Excluded Force""" + self.pmeef = P.PMEExcludedForce(atom_numbers=self.atom_numbers, beta=self.beta) + self.frc = self.pmeef(uint_crd, scaler, charge, excluded_list_start, excluded_list, excluded_numbers) + return self.frc + + def PME_Reciprocal_Force(self, uint_crd, charge): + """PME reciprocal force""" + self.pmerf = P.PMEReciprocalForce(self.atom_numbers, self.beta, self.fftx, self.ffty, self.fftz) + self.frc = self.pmerf(self.box_length, uint_crd, charge) + return self.frc + + def Energy_Device_To_Host(self): + """Energy_Device_To_Host""" + self.ee_ene = self.reciprocal_energy + self.self_energy + self.direct_energy + self.correction_energy + return self.ee_ene diff --git a/model_zoo/research/hpc/sponge/src/simulation_initial.py b/model_zoo/research/hpc/sponge/src/simulation_initial.py new file mode 100644 index 0000000000..5eed6cea9b --- /dev/null +++ b/model_zoo/research/hpc/sponge/src/simulation_initial.py @@ -0,0 +1,243 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""simulation""" + +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import Tensor, nn + +from Langevin_Liujian_md import Langevin_Liujian +from angle import Angle +from bond import Bond +from dihedral import Dihedral +from lennard_jones import Lennard_Jones_Information +from md_information import md_information +from nb14 import NON_BOND_14 +from neighbor_list import nb_infomation +from particle_mesh_ewald import Particle_Mesh_Ewald + + +class controller: + """class controller""" + def __init__(self, args_opt): + self.input_file = args_opt.i + self.initial_coordinates_file = args_opt.c + self.amber_parm = args_opt.amber_parm + self.restrt = args_opt.r + self.mdcrd = args_opt.x + self.mdout = args_opt.o + self.mdbox = args_opt.box + + self.Command_Set = {} + self.md_task = None + self.commands_from_in_file() + + def commands_from_in_file(self): + """commands from in file""" + file = open(self.input_file, 'r') + context = file.readlines() + file.close() + self.md_task = context[0].strip() + for val in context: + if "=" in val: + assert len(val.strip().split("=")) == 2 + flag, value = val.strip().split("=") + value = value.replace(",", '') + flag = flag.replace(" ", "") + if flag not in self.Command_Set: + self.Command_Set[flag] = value + else: + print("ERROR COMMAND FILE") + + +class Simulation(nn.Cell): + """class simulation""" + + def __init__(self, args_opt): + super(Simulation, self).__init__() + self.control = controller(args_opt) + self.md_info = md_information(self.control) + self.bond = Bond(self.control, self.md_info) + self.angle = Angle(self.control) + self.dihedral = Dihedral(self.control) + self.nb14 = NON_BOND_14(self.control, self.dihedral, self.md_info.atom_numbers) + self.nb_info = nb_infomation(self.control, self.md_info.atom_numbers, self.md_info.box_length) + self.LJ_info = Lennard_Jones_Information(self.control) + self.liujian_info = Langevin_Liujian(self.control, self.md_info.atom_numbers) + self.pme_method = Particle_Mesh_Ewald(self.control, self.md_info) + self.box_length = Tensor(np.asarray(self.md_info.box_length, np.float32), mstype.float32) + self.file = None + + def Main_Before_Calculate_Force(self): + """Main Before Calculate Force""" + _ = self.md_info.MD_Information_Crd_To_Uint_Crd() + self.md_info.uint_crd_with_LJ = (self.md_info.uint_crd, self.LJ_info.atom_LJ_type, self.md_info.charge) + return self.md_info.uint_crd, self.md_info.uint_crd_with_LJ + + def Initial_Neighbor_List_Update(self, not_first_time): + """Initial Neighbor List Update""" + res = self.nb_info.NeighborListUpdate(self.md_info.crd, self.md_info.crd_old, self.md_info.uint_crd, + self.md_info.crd_to_uint_crd_cof, self.md_info.uint_dr_to_dr_cof, + self.box_length, not_first_time) + + return res + + def Main_Calculate_Force(self): + """main calculate force""" + self.bond.atom_numbers = self.md_info.atom_numbers + md_info = self.md_info + LJ_info = self.LJ_info + nb_info = self.nb_info + pme_method = self.pme_method + bond_frc, _ = self.bond.Bond_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) + frc_t = bond_frc.asnumpy() + + angle_frc, _ = self.angle.Angle_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) + frc_t += angle_frc.asnumpy() + + dihedral_frc, _ = self.dihedral.Dihedral_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) + frc_t += dihedral_frc.asnumpy() + + nb14_frc, _ = self.nb14.Non_Bond_14_LJ_CF_Force_With_Atom_Energy(md_info.uint_crd_with_LJ, + md_info.uint_dr_to_dr_cof, LJ_info.LJ_A, + LJ_info.LJ_B) + frc_t += nb14_frc.asnumpy() + + lj_frc = LJ_info.LJ_Force_With_PME_Direct_Force( + md_info.atom_numbers, md_info.uint_crd_with_LJ, md_info.uint_dr_to_dr_cof, nb_info.nl_atom_numbers, + nb_info.nl_atom_serial, nb_info.cutoff, pme_method.beta) + frc_t += lj_frc.asnumpy() + + pme_excluded_frc = pme_method.PME_Excluded_Force( + md_info.uint_crd, md_info.uint_dr_to_dr_cof, md_info.charge, + nb_info.excluded_list_start, nb_info.excluded_list, + nb_info.excluded_numbers) + frc_t += pme_excluded_frc.asnumpy() + + pme_reciprocal_frc = pme_method.PME_Reciprocal_Force(md_info.uint_crd, md_info.charge) + frc_t += pme_reciprocal_frc.asnumpy() + + self.md_info.frc = Tensor(frc_t, mstype.float32) + + return self.md_info.frc + + def Main_Calculate_Energy(self): + """main calculate energy""" + _ = self.bond.Bond_Energy(self.md_info.uint_crd, self.md_info.uint_dr_to_dr_cof) + _ = self.angle.Angle_Energy(self.md_info.uint_crd, self.md_info.uint_dr_to_dr_cof) + _ = self.dihedral.Dihedral_Engergy(self.md_info.uint_crd, self.md_info.uint_dr_to_dr_cof) + _ = self.nb14.Non_Bond_14_LJ_CF_Energy(self.md_info.uint_crd_with_LJ, self.md_info.uint_dr_to_dr_cof, + self.LJ_info.LJ_A, + self.LJ_info.LJ_B) + + _ = self.LJ_info.LJ_Energy(self.md_info.uint_crd_with_LJ, self.md_info.uint_dr_to_dr_cof, + self.nb_info.nl_atom_numbers, self.nb_info.nl_atom_serial, + self.nb_info.cutoff_square) + _ = self.pme_method.PME_Energy( + self.md_info.uint_crd, self.md_info.charge, self.nb_info.nl_atom_numbers, self.nb_info.nl_atom_serial, + self.md_info.uint_dr_to_dr_cof, self.nb_info.excluded_list_start, self.nb_info.excluded_list, + self.nb_info.excluded_numbers) + _ = self.pme_method.Energy_Device_To_Host() + + def Main_After_Calculate_Energy(self): + """main after calculate energy""" + md_info = self.md_info + LJ_info = self.LJ_info + bond = self.bond + angle = self.angle + dihedral = self.dihedral + nb14 = self.nb14 + pme_method = self.pme_method + + md_info.total_potential_energy = 0 + md_info.total_potential_energy += bond.sigma_of_bond_ene + md_info.total_potential_energy += angle.sigma_of_angle_ene + md_info.total_potential_energy += dihedral.sigma_of_dihedral_ene + md_info.total_potential_energy += nb14.nb14_lj_energy_sum + nb14.nb14_cf_energy_sum + md_info.total_potential_energy += LJ_info.LJ_energy_sum + pme_method.Energy_Device_To_Host() + md_info.total_potential_energy += pme_method.ee_ene + print("md_info.total_potential_energy", md_info.total_potential_energy) + + def Main_Iteration_2(self): + """main iteration2""" + md_info = self.md_info + control = self.control + liujian_info = self.liujian_info + + if md_info.mode > 0 and int(control.Command_Set["thermostat"]) == 1: + md_info.vel, md_info.crd, md_info.frc, md_info.acc = liujian_info.MD_Iteration_Leap_Frog( + md_info.d_mass_inverse, md_info.vel, md_info.crd, md_info.frc) + + def Main_After_Iteration(self): + """main after iteration""" + md_info = self.md_info + nb_info = self.nb_info + md_info.Centerize() + _ = nb_info.NeighborListUpdate(md_info.crd, md_info.crd_old, md_info.uint_crd, + md_info.crd_to_uint_crd_cof, + md_info.uint_dr_to_dr_cof, self.box_length, not_first_time=1) + + def Main_Print(self): + """compute the temperature""" + md_info = self.md_info + temperature = md_info.MD_Information_Temperature() + md_info.h_temperature = temperature + steps = md_info.steps + temperature = temperature.asnumpy() + total_potential_energy = md_info.total_potential_energy.asnumpy() + sigma_of_bond_ene = self.bond.sigma_of_bond_ene.asnumpy() + sigma_of_angle_ene = self.angle.sigma_of_angle_ene.asnumpy() + sigma_of_dihedral_ene = self.dihedral.sigma_of_dihedral_ene.asnumpy() + nb14_lj_energy_sum = self.nb14.nb14_lj_energy_sum.asnumpy() + nb14_cf_energy_sum = self.nb14.nb14_cf_energy_sum.asnumpy() + LJ_energy_sum = self.LJ_info.LJ_energy_sum.asnumpy() + ee_ene = self.pme_method.ee_ene.asnumpy() + print("_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ " + "_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_") + print("{:>7.0f} {:>7.3f} {:>11.3f}".format(steps, float(temperature), float(total_potential_energy)), end=" ") + if self.bond.bond_numbers > 0: + print("{:>10.3f}".format(float(sigma_of_bond_ene)), end=" ") + if self.angle.angle_numbers > 0: + print("{:>11.3f}".format(float(sigma_of_angle_ene)), end=" ") + if self.dihedral.dihedral_numbers > 0: + print("{:>14.3f}".format(float(sigma_of_dihedral_ene)), end=" ") + if self.nb14.nb14_numbers > 0: + print("{:>10.3f} {:>10.3f}".format(float(nb14_lj_energy_sum), float(nb14_cf_energy_sum)), end=" ") + + print("{:>7.3f}".format(float(LJ_energy_sum)), end=" ") + print("{:>12.3f}".format(float(ee_ene))) + + if self.file is not None: + self.file.write("{:>7.0f} {:>7.3f} {:>11.3f} {:>10.3f} {:>11.3f} {:>14.3f} {:>10.3f} {:>10.3f} {:>7.3f}" + " {:>12.3f}\n".format(steps, float(temperature), float(total_potential_energy), + float(sigma_of_bond_ene), float(sigma_of_angle_ene), + float(sigma_of_dihedral_ene), float(nb14_lj_energy_sum), + float(nb14_cf_energy_sum), float(LJ_energy_sum), float(ee_ene))) + + return temperature + + def Main_Initial(self): + """main initial""" + if self.control.mdout: + self.file = open(self.control.mdout, 'w') + self.file.write("_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ " + "_ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_\n") + + def Main_Destroy(self): + """main destroy""" + if self.file is not None: + self.file.close() + print("Save successfully!")