diff --git a/MindSPONGE/README.md b/MindSPONGE/README.md index 7af776d4fa26d83eb1a276ad225779ed0e8bfe04..8109914c45b3224b244b226a19190a9296a411ba 100644 --- a/MindSPONGE/README.md +++ b/MindSPONGE/README.md @@ -64,9 +64,58 @@ transformed_rot = quat_to_rot(quat) transformed_quat = rot_to_quat(rot) ``` -- 简单体系分子模拟 +- 一个简单的分子动力学模拟案例 ```bash +import numpy as np +from mindspore import context +from mindsponge import Sponge +from mindsponge import Molecule +from mindsponge import ForceFieldBase +from mindsponge import DynamicUpdater +from mindsponge.potential import BondEnergy, AngleEnergy +from mindsponge.function import VelocityGenerator +from mindsponge.control import LeapFrog +from mindsponge.callback import WriteH5MD, RunInfo + +context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + +system = Molecule( + atoms=['O', 'H', 'H'], + coordinate=[[0, 0, 0], [0.1, 0, 0], [-0.0333, 0.0943, 0]], + bond=[[[0, 1], [0, 2]]], +) + +bond_energy = BondEnergy( + index=system.bond, + force_constant=[[345000, 345000]], + bond_length=[[0.1, 0.1]], +) + +angle_energy = AngleEnergy( + index=[[1, 0, 2]], + force_constant=[[383]], + bond_angle=[[109.47 / 180 * np.pi]], +) + +energy = ForceFieldBase(energy=[bond_energy, angle_energy]) + +vgen = VelocityGenerator(300) +velocities = vgen(system.coordinate.shape, system.atom_mass) + +opt = DynamicUpdater( + system, + integrator=LeapFrog(system), + time_step=1e-3, + velocity=velocities +) + +md = Sponge(system, energy, opt) + +run_info = RunInfo(10) +cb_h5md = WriteH5MD(system, 'test.h5md', save_freq=10, write_velocity=True, write_force=True) + +md.run(2000, callbacks=[run_info, cb_h5md]) ``` **更多应用案例请见**： @@ -136,9 +185,9 @@ pip install mindscience_sponge*.whl ### CO-CHAIR -- 深圳湾实验室杨奕 +- 深圳湾实验室[杨奕](https://gitee.com/helloyesterday) -- 北京昌平实验室张骏 +- 北京昌平实验室[张骏](https://gitee.com/jz_90) ### SIG @@ -154,7 +203,7 @@ MindSPONGE SIG小组为广大科研人员，老师和学生提供高效易用的 ### 核心贡献者 -- [高毅勤课题组](https://www.chem.pku.edu.cn/gaoyq/) : [杨奕](https://gitee.com/helloyesterday)，[张骏](https://gitee.com/jz_90)，[刘思睿](https://gitee.com/sirui63)，[夏义杰](https://gitee.com/gao_hyp_xyj_admin)，[陈迪青](https://gitee.com/dechin)，[黄渝鹏](https://gitee.com/gao_hyp_xyj_admin) +- [高毅勤课题组](https://www.chem.pku.edu.cn/gaoyq/): [杨奕](https://gitee.com/helloyesterday)，[张骏](https://gitee.com/jz_90)，[刘思睿](https://gitee.com/sirui63)，[夏义杰](https://gitee.com/gao_hyp_xyj_admin)，[陈迪青](https://gitee.com/dechin)，[黄渝鹏](https://gitee.com/gao_hyp_xyj_admin) ## **许可证** diff --git a/MindSPONGE/mindsponge/python/colvar/atoms.py b/MindSPONGE/mindsponge/python/colvar/atoms.py index e32b4d3b7b9508ad3729333c486566c99956ecc1..7bc2ddccbad58318cba6c8dfdb4653f301ef6fbf 100644 --- a/MindSPONGE/mindsponge/python/colvar/atoms.py +++ b/MindSPONGE/mindsponge/python/colvar/atoms.py @@ -58,7 +58,7 @@ class AtomDistances(Colvar): length_unit=length_unit, ) - # (1,b,2) + # (B,b,2) self.index = index self.identity = ops.Identity() @@ -75,8 +75,9 @@ class AtomDistances(Colvar): """ - # (B,b,2,D) + # (B,b,2) index = self.identity(self.index) + # (B,b,2,D) atoms = func.gather_vectors(coordinate, index) # (B,b,D) @@ -120,8 +121,10 @@ class AtomAngles(Colvar): """ + # (B,a,3) + index = self.identity(self.index) # (B,a,3,D) - atoms = func.gather_vectors(coordinate, self.index) + atoms = func.gather_vectors(coordinate, index) # (B,a,1,D) atom0, atom1, atom2 = self.split(atoms) @@ -163,7 +166,7 @@ class AtomTorsions(Colvar): use_pbc=use_pbc, ) - # (1,d,4) + # (B,d,4) self.index = index self.split = ops.Split(-2, 4) @@ -180,8 +183,10 @@ class AtomTorsions(Colvar): """ + # (B,d,4) + index = self.identity(self.index) # (B,d,4,D) - atoms = func.gather_vectors(coordinate, self.index) + atoms = func.gather_vectors(coordinate, index) # (B,d,1,D) atom_a, atom_b, atom_c, atom_d = self.split(atoms) diff --git a/MindSPONGE/mindsponge/python/colvar/colvar.py b/MindSPONGE/mindsponge/python/colvar/colvar.py index 6859e31bee9289074b38ac44a91f63ed59f4da72..38fcd8cbea775c4e47ccbd3769ee00480b2c5569 100644 --- a/MindSPONGE/mindsponge/python/colvar/colvar.py +++ b/MindSPONGE/mindsponge/python/colvar/colvar.py @@ -25,6 +25,7 @@ Collective variables """ import mindspore as ms +from mindspore import ops from mindspore.ops import functional as F from mindspore.nn import Cell from mindspore.common import Tensor @@ -91,6 +92,8 @@ class Colvar(Cell): self.any_periodic = self.periodic.any() self.all_periodic = self.periodic.all() + self.identity = ops.Identity() + @property def length_unit(self): """length unit""" diff --git a/MindSPONGE/mindsponge/python/core/simulation/simulation.py b/MindSPONGE/mindsponge/python/core/simulation/simulation.py index ffe176ef68c4f136a053d830775ba654446dd0cc..3fdd8531469cd6c9321af87e680824a6399eeadf 100644 --- a/MindSPONGE/mindsponge/python/core/simulation/simulation.py +++ b/MindSPONGE/mindsponge/python/core/simulation/simulation.py @@ -28,6 +28,7 @@ import mindspore as ms import mindspore.numpy as msnp from mindspore import Tensor from mindspore import Parameter +from mindspore import context from mindspore import ops from mindspore.ops import functional as F from mindspore.nn import Cell, CellList @@ -113,7 +114,10 @@ class SimulationCell(Cell): self.neighbour_index = self.neighbour_list.neighbours self.neighbour_mask = self.neighbour_list.neighbour_mask - self.get_neighbour_list = self.neighbour_list.get_neighbour_list + + self.no_mask = False + if context.get_context("mode") == context.PYNATIVE_MODE and self.neighbour_list.no_mask: + self.no_mask = True self.num_neighbours = self.neighbour_list.num_neighbours @@ -163,6 +167,15 @@ class SimulationCell(Cell): coordinate, pbc_box = self.system() return self.neighbour_list(coordinate, pbc_box) + def get_neighbour_list(self): + """get neighbour list""" + if self.no_mask: + (neighbour_index,) = self.neighbour_list.get_neighbour_list() + neighbour_mask = None + else: + neighbour_index, neighbour_mask = self.neighbour_list.get_neighbour_list() + return neighbour_index, neighbour_mask + @property def length_unit(self): return self.units.length_unit diff --git a/MindSPONGE/mindsponge/python/data/forcefield/amber.ff14sb.yaml b/MindSPONGE/mindsponge/python/data/forcefield/amber.ff14sb.yaml index 2501a66ccc59b5679ab201f77f9b3cc16fca80df..9285b012098649c27569914fe25a12f5f467e941 100644 --- a/MindSPONGE/mindsponge/python/data/forcefield/amber.ff14sb.yaml +++ b/MindSPONGE/mindsponge/python/data/forcefield/amber.ff14sb.yaml @@ -2691,6 +2691,15 @@ template: - H1 - H1 - H1 + WAT: + atom_charge: + - -0.834 + - 0.417 + - 0.417 + atom_type: + - OW + - HW + - HW parameters: bond_energy: length_unit: nm @@ -9404,6 +9413,8 @@ parameters: - 3C - C8 - CO + - HW + - OW epsilon: - 0.06568879751861095 - 0.0 @@ -9444,6 +9455,8 @@ parameters: - 0.4577295861840248 - 0.4577295861840248 - 0.35982401221990584 + - 0.0 + - 0.315061 sigma: - 0.10690785042497683 - 0.0 @@ -9484,6 +9497,8 @@ parameters: - 0.3399669494829499 - 0.3399669494829499 - 0.3399669494829499 + - 0.0 + - 0.636386 nb_pair_energy: length_unit: nm energy_unit: kj/mol diff --git a/MindSPONGE/mindsponge/python/data/parameters.py b/MindSPONGE/mindsponge/python/data/parameters.py index 114b069ed9c82ce0f038ff7af5d99fcc2d5d9e4e..778f8aadb7d0bb16f4ecf3f3459bab7c2c7ca3ef 100644 --- a/MindSPONGE/mindsponge/python/data/parameters.py +++ b/MindSPONGE/mindsponge/python/data/parameters.py @@ -635,7 +635,7 @@ class ForceFieldParameters: 'C4','CT','CX','C','N','N3','S','SH','P','MG', 'C0','F','Cl','Br','I','2C','3C','C8','CO'] """ - atom_names_count = np.zeros(39) + atom_names_count = np.zeros(41) for i in range(self.atom_nums): this_id = np.where( np.isin(self.vdw_params["atoms"], atom_names[i]))[0] @@ -706,7 +706,7 @@ class ForceFieldParameters: def __call__(self, bonds): # pylint: disable=unused-argument hbonds, non_hbonds = self.get_hbonds(bonds) - atoms_types = self.atom_types + atoms_types = self.atom_types.copy() self.get_vdw_params(atoms_types) atom_types = np.append(atoms_types, self._wildcard) this_bond_params = self.get_bond_params(bonds, atoms_types) @@ -722,17 +722,21 @@ class ForceFieldParameters: np.where(np.isin(bonds, middle_id).sum(axis=1) == 2)[0] ] dihedrals = self.get_dihedrals(angles, dihedral_middle_id) - self.dihedral_params = self.get_dihedral_params(dihedrals, atom_types) + if dihedrals is not None: + self.dihedral_params = self.get_dihedral_params(dihedrals, atom_types) core_id = np.where(np.bincount(bonds.flatten()) > 2)[0] checked_core_id = self.check_idihedral(bonds, core_id) idihedrals, third_id = self.get_idihedrals(bonds, checked_core_id) self.improper_dihedral_params = self.get_idihedral_params( idihedrals, atom_types, third_id ) - self.pair_index = self.get_pair_index(dihedrals, angles, bonds) - pair_params = self.get_pair_params(self.pair_index, self.vdw_param[:, 0][None, :], - self.vdw_param[:, 1][None, :]) - self.excludes = self.get_excludes(bonds, angles, dihedrals, idihedrals) + if dihedrals is not None: + self.pair_index = self.get_pair_index(dihedrals, angles, bonds) + pair_params = self.get_pair_params(self.pair_index, self.vdw_param[:, 0][None, :], + self.vdw_param[:, 1][None, :]) + self.excludes = self.get_excludes(bonds, angles, dihedrals, idihedrals) + else: + pair_params = None return ForceConstants( self.bond_params, self.angle_params, diff --git a/MindSPONGE/mindsponge/python/data/template/protein0.yaml b/MindSPONGE/mindsponge/python/data/template/protein0.yaml index 6a5febea3322a1c75c75c8671c12e3bb9aa14efb..43c56349cc83a138aeee46a11aa20b2d5e7ad190 100644 --- a/MindSPONGE/mindsponge/python/data/template/protein0.yaml +++ b/MindSPONGE/mindsponge/python/data/template/protein0.yaml @@ -5920,3 +5920,24 @@ NME: - 5 head_atom: 0 tail_atom: null +WAT: + residue: WAT + atom_name: + - OW + - HW1 + - HW2 + atom_mass: + - 16.00000 + - 1.00800 + - 1.00800 + atomic_number: + - 8 + - 1 + - 1 + bond: + - - 0 + - 1 + - - 0 + - 2 + head_atom: null + tail_atom: null diff --git a/MindSPONGE/mindsponge/python/partition/distance.py b/MindSPONGE/mindsponge/python/partition/distance.py index c6ab1106198de4a01e94660a8865ce42b7eec644..eafc9719c3512e0cd78726a9468cccf68448f07a 100644 --- a/MindSPONGE/mindsponge/python/partition/distance.py +++ b/MindSPONGE/mindsponge/python/partition/distance.py @@ -140,6 +140,26 @@ class DistanceNeighbours(Cell): def print_info(self): return self + def check_neighbours_number(self, neighbour_mask: Tensor): + """check number of neighbours in neighbour list""" + max_neighbours = F.cast(msnp.max(F.cast(msnp.sum(neighbour_mask, -1), ms.float32)), ms.int32) + if max_neighbours > self.num_neighbours: + print( + '================================================================================') + print( + 'Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning!') + print( + '--------------------------------------------------------------------------------') + print('The max number of neighbour atoms is larger than that in neighbour list!') + print('The max number of neighbour atoms:') + print(max_neighbours) + print('The number of neighbour atoms in neighbour list:') + print(self.num_neighbours) + print('Please increase the value of grid_num_scale or num_neighbours!') + print( + '================================================================================') + return self + def construct(self, coordinate: Tensor, pbc_box: Tensor = None, @@ -201,24 +221,7 @@ class DistanceNeighbours(Cell): num_neighbours = num_atoms - 1 else: num_neighbours = self.num_neighbours - max_neighbours = F.cast( - msnp.max(F.cast(msnp.sum(neighbour_mask, -1), ms.float32)), ms.int32) - if max_neighbours > num_neighbours: - print( - '================================================================================') - print( - 'Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning!') - print( - '--------------------------------------------------------------------------------') - print('The max number of neighbour atoms ' - 'is larger than that in neighbour list!') - print('The max number of neighbour atoms:') - print(max_neighbours) - print('The number of neighbour atoms in neighbour list:') - print(num_neighbours) - print('Please increase the value of grid_num_scale or num_neighbours!') - print( - '================================================================================') + self.check_neighbours_number(neighbour_mask) distances = distances[..., 1:num_neighbours+1] neighbours = neighbours[..., 1:num_neighbours+1] diff --git a/MindSPONGE/mindsponge/python/partition/grids.py b/MindSPONGE/mindsponge/python/partition/grids.py index d316e0c6ead59cd81964ca74ebaf5af60ce66364..ed0bff76da705c6612abc68fccfaea242d691cd2 100644 --- a/MindSPONGE/mindsponge/python/partition/grids.py +++ b/MindSPONGE/mindsponge/python/partition/grids.py @@ -296,6 +296,31 @@ class GridNeighbours(Cell): self.exclude_index = F.expand_dims(self.exclude_index, 0) return self + def check_neighbours_number(self, grid_neigh_atoms: Tensor, num_neighbours: int = None): + """check number of neighbours in neighbour list""" + if num_neighbours is None: + num_neighbours = self.num_neighbours + max_neighbours = msnp.sum(grid_neigh_atoms != self.num_atoms, axis=-1) + max_neighbours = F.cast( + msnp.max(F.cast(max_neighbours, ms.float32)), ms.int32) + if max_neighbours > num_neighbours: + print( + '================================================================================') + print( + 'Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning!') + print( + '--------------------------------------------------------------------------------') + print('The max number of neighbour atoms ' + 'is larger than that in neighbour list!') + print('The max number of neighbour atoms:') + print(max_neighbours) + print('The number of neighbour atoms in neighbour list:') + print(num_neighbours) + print('Please increase the value of grid_num_scale or num_neighbours!') + print( + '================================================================================') + return self + def print_info(self): """print information of neighbour list""" print('Calculate neighbour list from grids') @@ -355,26 +380,8 @@ class GridNeighbours(Cell): grid_neigh_atoms, _ = self.sort(F.cast(grid_neigh_atoms, ms.float32)) grid_neigh_atoms = F.cast(grid_neigh_atoms, ms.int32) - # max_neighbours = msnp.sum(grid_neigh_atoms != self.num_atoms, axis=-1) - # max_neighbours = F.cast( - # msnp.max(F.cast(max_neighbours, ms.float32)), ms.int32) - # if num_neighbours < max_neighbours: - # print( - # '================================================================================') - # print( - # 'Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning! Warning!') - # print( - # '--------------------------------------------------------------------------------') - # print('The max number of neighbour atoms ' - # 'is larger than that in neighbour list!') - # print('The max number of neighbour atoms:') - # print(max_neighbours) - # print('The number of neighbour atoms in neighbour list:') - # print(num_neighbours) - # print('Please increase the value of grid_num_scale or num_neighbours!') - # print( - # '================================================================================') - # grid_neigh_atoms = grid_neigh_atoms[..., :num_neighbours] + self.check_neigbours_number(grid_neigh_atoms, num_neighbours) + grid_neigh_atoms = grid_neigh_atoms[..., :num_neighbours] # neighbour atoms for each atom # (B,A,N) diff --git a/MindSPONGE/mindsponge/python/partition/neighbourlist.py b/MindSPONGE/mindsponge/python/partition/neighbourlist.py index 3c629b34356b0d4dd6fc620762f4310beb99b0b7..a5cce550c2cb47d1a819b74ed9a744ea05a6b15a 100644 --- a/MindSPONGE/mindsponge/python/partition/neighbourlist.py +++ b/MindSPONGE/mindsponge/python/partition/neighbourlist.py @@ -28,9 +28,9 @@ import mindspore as ms import mindspore.numpy as msnp from mindspore import Tensor, ms_function from mindspore import Parameter -from mindspore.nn import Cell from mindspore import ops from mindspore.ops import functional as F +from mindspore.nn import Cell from . import FullConnectNeighbours, DistanceNeighbours, GridNeighbours from ..system import Molecule diff --git a/MindSPONGE/mindsponge/python/potential/energy/angle.py b/MindSPONGE/mindsponge/python/potential/energy/angle.py index 09c1ae2c83f41ca0d01d274b17a3d2f1f5de89e6..0fb0168fe73c302decdc8f7b2f928d398f32e08e 100644 --- a/MindSPONGE/mindsponge/python/potential/energy/angle.py +++ b/MindSPONGE/mindsponge/python/potential/energy/angle.py @@ -98,8 +98,13 @@ class AngleEnergy(EnergyCell): # (1,a,3) index = Tensor(index, ms.int32) + if index.shape[-1] != 3: + raise ValueError('The last dimension of index in AngleEnergy must be 3 but got: ' + + str(index.shape[-1])) if index.ndim == 2: index = F.expand_dims(index, 0) + if index.ndim != 3: + raise ValueError('The rank of index must be 2 or 3 but got shape: '+str(index.shape)) self.index = Parameter(index, name='angle_index', requires_grad=False) self.num_angles = index.shape[-2] diff --git a/MindSPONGE/mindsponge/python/potential/energy/bond.py b/MindSPONGE/mindsponge/python/potential/energy/bond.py index 9bb25b4509b6d01add3a98e1e521ff258be53075..badeeabccf2f4b1aa90e9a2c1d44f8057497bdb2 100644 --- a/MindSPONGE/mindsponge/python/potential/energy/bond.py +++ b/MindSPONGE/mindsponge/python/potential/energy/bond.py @@ -100,10 +100,15 @@ class BondEnergy(EnergyCell): force_constant = parameters.get('force_constant') bond_length = parameters.get('bond_length') - # (1,b,2) + # (B,b,2) index = Tensor(index, ms.int32) + if index.shape[-1] != 2: + raise ValueError('The last dimension of index in BondEnergy must be 2 but got: ' + + str(index.shape[-1])) if index.ndim == 2: index = F.expand_dims(index, 0) + if index.ndim != 3: + raise ValueError('The rank of index must be 2 or 3 but got shape: '+str(index.shape)) self.index = Parameter(index, name='bond_index', requires_grad=False) # (B,b) @@ -175,16 +180,16 @@ class BondEnergy(EnergyCell): """ - # (B,M) + # (B,b) dis = self.get_bond_length(coordinate, pbc_box) * self.input_unit_scale - # (B,M) = (B,M) - (1,M) + # (B,b) = (B,b) - (B,b) diff = dis - self.bond_length - # (B,M) + # (B,b) diff2 = F.square(diff) - # (B,M) = (1,M) * (B,M) * k + # (B,b) = (1,b) * (B,b) * k energy = 0.5 * self.force_constant * diff2 - # (B,1) <- (B,M) + # (B,1) <- (B,b) return func.keepdim_sum(energy, -1) diff --git a/MindSPONGE/mindsponge/python/potential/energy/dihedral.py b/MindSPONGE/mindsponge/python/potential/energy/dihedral.py index 57c6707d557e6dda007a40947cd5aa7d7cf2e5be..d6f82b4ac61b840d13881615250f0649271254d3 100644 --- a/MindSPONGE/mindsponge/python/potential/energy/dihedral.py +++ b/MindSPONGE/mindsponge/python/potential/energy/dihedral.py @@ -101,8 +101,14 @@ class DihedralEnergy(EnergyCell): # (1,d,4) index = Tensor(index, ms.int32) + if index.shape[-1] != 4: + raise ValueError('The last dimension of index in DihedralEnergy must be 2 but got: ' + + str(index.shape[-1])) if index.ndim == 2: index = F.expand_dims(index, 0) + if index.ndim != 3: + raise ValueError( + 'The rank of index must be 2 or 3 but got shape: '+str(index.shape)) self.index = Parameter(index, name='dihedral_index', requires_grad=False) # (1,d) diff --git a/MindSPONGE/mindsponge/python/potential/energy/paris.py b/MindSPONGE/mindsponge/python/potential/energy/paris.py index 00d4ced622b633b991843b2bbf59f59149933570..3e643d4b7cb49ba35510c3d0f6f0dee814af4272 100644 --- a/MindSPONGE/mindsponge/python/potential/energy/paris.py +++ b/MindSPONGE/mindsponge/python/potential/energy/paris.py @@ -126,8 +126,13 @@ class NonbondPairwiseEnergy(EnergyCell): # (1,p,2) index = Tensor(index, ms.int32) + if index.shape[-1] != 2: + raise ValueError('The last dimension of index in NonbondPairwiseEnergy must be 2 but got: ' + + str(index.shape[-1])) if index.ndim == 2: index = F.expand_dims(index, 0) + if index.ndim != 3: + raise ValueError('The rank of index must be 2 or 3 but got shape: '+str(index.shape)) self.index = Parameter(index, name='pairs_index', requires_grad=False) self.num_pairs = index.shape[-2] diff --git a/MindSPONGE/mindsponge/python/potential/forcefield.py b/MindSPONGE/mindsponge/python/potential/forcefield.py index 185eae86ef76a1ae7e55b28ca2d8bdb13a8f1b32..199312684b269d2702ce53e155f1cadb2f41fd1a 100644 --- a/MindSPONGE/mindsponge/python/potential/forcefield.py +++ b/MindSPONGE/mindsponge/python/potential/forcefield.py @@ -259,15 +259,17 @@ class ForceField(ForceFieldBase): # Generate Forcefield Parameters ff_dict = get_forcefield_parameters(parameters) + template = ff_dict.get('template') + parameters = ff_dict.get('parameters') for residue in system.residue: - residue.build_atom_type(ff_dict['template']) - residue.build_atom_charge(ff_dict['template']) + residue.build_atom_type(template.get(residue.name)) + residue.build_atom_charge(template.get(residue.name)) system.build_system() ff_params = ForceFieldParameters( - system.atom_type, copy.deepcopy(ff_dict), atom_names=system.atom_name[0], - atom_charges=self.identity(system.atom_charge).asnumpy()[None, :]) + system.atom_type[0], copy.deepcopy(ff_dict), atom_names=system.atom_name[0], + atom_charges=self.identity(system.atom_charge).asnumpy()) if isinstance(system.bond, np.ndarray): force_params = ff_params(system.bond[0]) @@ -282,8 +284,9 @@ class ForceField(ForceFieldBase): bond_force_constant = Tensor(force_params.bond_params[:, 2][None, :], ms.float32) bond_length = Tensor(force_params.bond_params[:, 3][None, :], ms.float32) - length_unit = ff_dict['parameters']['bond_energy']['length_unit'] - energy_unit = ff_dict['parameters']['bond_energy']['energy_unit'] + bond_params: dict = parameters.get('bond_energy') + length_unit = bond_params.get('length_unit') + energy_unit = bond_params.get('energy_unit') bond_energy = BondEnergy(bond_index, force_constant=bond_force_constant, bond_length=bond_length, use_pbc=use_pbc, length_unit=length_unit, energy_unit=energy_unit) @@ -295,7 +298,8 @@ class ForceField(ForceFieldBase): angle_force_constant = Tensor(force_params.angle_params[:, 3][None, :], ms.float32) bond_angle = Tensor(force_params.angle_params[:, 4][None, :], ms.float32) - energy_unit = ff_dict['parameters']['angle_energy']['energy_unit'] + angle_params: dict = parameters.get('angle_energy') + energy_unit = angle_params.get('energy_unit') angle_energy = AngleEnergy(angle_index, force_constant=angle_force_constant, bond_angle=bond_angle, use_pbc=use_pbc, energy_unit=energy_unit) energy.append(angle_energy) @@ -322,7 +326,8 @@ class ForceField(ForceFieldBase): phase = msnp.append(phase, Tensor( force_params.improper_dihedral_params[:, 6][None, :], ms.float32), axis=1) - energy_unit = ff_dict['parameters']['dihedral_energy']['energy_unit'] + dihedral_params: dict = parameters.get('dihedral_energy') + energy_unit = dihedral_params.get('energy_unit') dihedral_energy = DihedralEnergy(dihedral_index, force_constant=dihe_force_constant, periodicity=periodicity, phase=phase, use_pbc=use_pbc, energy_unit=energy_unit) @@ -330,8 +335,9 @@ class ForceField(ForceFieldBase): # Electronic energy if system.atom_charge is not None: - length_unit = ff_dict['parameters']['coulomb_energy']['length_unit'] - energy_unit = ff_dict['parameters']['coulomb_energy']['energy_unit'] + coulomb_params: dict = parameters.get('coulomb_energy') + length_unit = coulomb_params.get('length_unit') + energy_unit = coulomb_params.get('energy_unit') ele_energy = CoulombEnergy(atom_charge=system.atom_charge, use_pbc=use_pbc, length_unit=length_unit, energy_unit=energy_unit) energy.append(ele_energy) @@ -343,8 +349,9 @@ class ForceField(ForceFieldBase): epsilon = Tensor(force_params.vdw_param[:, 0][None, :], ms.float32) sigma = Tensor(force_params.vdw_param[:, 1][None, :], ms.float32) - length_unit = ff_dict['parameters']['vdw_energy']['length_unit'] - energy_unit = ff_dict['parameters']['vdw_energy']['energy_unit'] + vdw_params: dict = parameters.get('vdw_energy') + length_unit = vdw_params.get('length_unit') + energy_unit = vdw_params.get('energy_unit') vdw_energy = LennardJonesEnergy(epsilon=epsilon, sigma=sigma, use_pbc=use_pbc, length_unit=length_unit, energy_unit=energy_unit) energy.append(vdw_energy) @@ -356,11 +363,12 @@ class ForceField(ForceFieldBase): epsilon_ij = Tensor(force_params.pair_params[:, 1][None, :], ms.float32) sigma_ij = Tensor(force_params.pair_params[:, 2][None, :], ms.float32) - r_scale = ff_dict['parameters']['nb_pair_energy']['r_scale'] - r6_scale = ff_dict['parameters']['nb_pair_energy']['r6_scale'] - r12_scale = ff_dict['parameters']['nb_pair_energy']['r12_scale'] - length_unit = ff_dict['parameters']['nb_pair_energy']['length_unit'] - energy_unit = ff_dict['parameters']['nb_pair_energy']['energy_unit'] + pair_params: dict = parameters.get('nb_pair_energy') + length_unit = pair_params.get('length_unit') + energy_unit = pair_params.get('energy_unit') + r_scale = pair_params.get('r_scale') + r6_scale = pair_params.get('r6_scale') + r12_scale = pair_params.get('r12_scale') pair_energy = NonbondPairwiseEnergy(pair_index, qiqj=qiqj, epsilon_ij=epsilon_ij, sigma_ij=sigma_ij, r_scale=r_scale, r6_scale=r6_scale, r12_scale=r12_scale, length_unit=length_unit, energy_unit=energy_unit, use_pbc=use_pbc) diff --git a/MindSPONGE/mindsponge/python/system/modeling/pdb_parser.py b/MindSPONGE/mindsponge/python/system/modeling/pdb_parser.py index 67c5ddf02f3021e8fa7c864b87a9fff75fbc05ea..96d171001501895209c91329fc7da84a0a475af9 100644 --- a/MindSPONGE/mindsponge/python/system/modeling/pdb_parser.py +++ b/MindSPONGE/mindsponge/python/system/modeling/pdb_parser.py @@ -44,7 +44,7 @@ resdict = {'ALA': 0, 'ARG': 1, 'ASN': 2, 'ASP': 3, 'CYS': 4, 'NGLN': 5, 'NGLU': 6, 'NGLY': 7, 'NHIS': 8, 'NILE': 9, 'NLEU': 10, 'NLYS': 11, 'NMET': 12, 'NPHE': 13, 'NPRO': 14, 'NSER': 15, 'NTHR': 16, 'NTRP': 17, 'NTYR': 18, 'NVAL': 19, - 'CHIE': 8, 'HIE': 8, 'NHIE': 8 + 'CHIE': 8, 'HIE': 8, 'NHIE': 8, 'WAT': 22 } atom_types = [ @@ -78,6 +78,7 @@ restype_name_to_atom14_names = { 'TYR': ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'OH', '', ''], 'VAL': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '', '', '', '', '', '', ''], 'UNK': ['', '', '', '', '', '', '', '', '', '', '', '', '', ''], + 'WAT': ['OW', '', '', '', '', '', '', '', '', '', '', '', '', ''], } restype_name_to_atom14_masks = { 'ALA': [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], @@ -101,7 +102,8 @@ restype_name_to_atom14_masks = { 'TRP': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 'TYR': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], 'VAL': [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0], - 'UNK': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]} + 'UNK': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 'WAT': [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]} atom14_order_dict = {'ALA': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4}, 'ARG': {'N': 0, @@ -238,7 +240,8 @@ atom14_order_dict = {'ALA': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4}, 'CZ': 10, 'OH': 11}, 'VAL': {'N': 0, 'CA': 1, 'C': 2, 'O': 3, 'CB': 4, 'CG1': 5, 'CG2': 6}, - 'UNK': {}} + 'UNK': {}, + 'WAT': {'OW': 0}} atom14_to_atom37_dict = {'ALA': [0, 1, 2, 4, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'ARG': [0, 1, 2, 4, 3, 5, 11, 23, 32, 29, 30, 0, 0, 0], @@ -261,7 +264,8 @@ atom14_to_atom37_dict = {'ALA': [0, 1, 2, 4, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'TRP': [0, 1, 2, 4, 3, 5, 12, 13, 24, 21, 22, 33, 34, 28], 'TYR': [0, 1, 2, 4, 3, 5, 12, 13, 20, 21, 32, 31, 0, 0], 'VAL': [0, 1, 2, 4, 3, 6, 7, 0, 0, 0, 0, 0, 0, 0], - 'UNK': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]} + 'UNK': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + 'WAT': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]} def read_pdb(pdb_name, ignoreh=False): diff --git a/MindSPONGE/mindsponge/python/system/molecule/protein.py b/MindSPONGE/mindsponge/python/system/molecule/protein.py index d62780b8831e2d4be76ffa6bae216edf7a2a0dea..5de01bab60533a5c9a57cf5d66c6f104c726df40 100644 --- a/MindSPONGE/mindsponge/python/system/molecule/protein.py +++ b/MindSPONGE/mindsponge/python/system/molecule/protein.py @@ -87,10 +87,11 @@ class Protein(Molecule): _, _, _, _, _ = read_pdb( pdb, ignore_hydrogen) - if residue_name[0] != 'ACE' and residue_name[0] != 'NME': - residue_name[0] = 'N' + residue_name[0] - if residue_name[-1] != 'ACE' and residue_name[-1] != 'NME': - residue_name[-1] = 'C' + residue_name[-1] + if len(residue_name) > 1: + if residue_name[0] != 'ACE' and residue_name[0] != 'NME': + residue_name[0] = 'N' + residue_name[0] + if residue_name[-1] != 'ACE' and residue_name[-1] != 'NME': + residue_name[-1] = 'C' + residue_name[-1] self.init_resname = init_res_names self.init_resid = init_res_ids diff --git a/MindSPONGE/mindsponge/python/system/residue/amino.py b/MindSPONGE/mindsponge/python/system/residue/amino.py index d82dd0600dab88b3596640c3903adf06b309daae..a3b458caa0db8fb386a8c35f7f7172cae22d075f 100644 --- a/MindSPONGE/mindsponge/python/system/residue/amino.py +++ b/MindSPONGE/mindsponge/python/system/residue/amino.py @@ -23,10 +23,9 @@ """ Molecule """ -import numpy as np from mindspore import ms_class from .residue import Residue - +from ...data.template import get_template_index @ms_class class AminoAcid(Residue): @@ -65,5 +64,4 @@ class AminoAcid(Residue): name=(name.replace('HIE', 'HIS') if 'HIE' in name else name), template=template, ) - self.atom_index = np.where( - np.array(list(template[self.name]['atom_name'])) == self.atom_name[0][:, None])[-1] + self.atom_index = get_template_index(template[self.name], self.atom_name) diff --git a/MindSPONGE/mindsponge/python/system/residue/residue.py b/MindSPONGE/mindsponge/python/system/residue/residue.py index cbfc379b4afc35e36d6246aa22ccbb9e60054515..d194a7c889e864a139f4f6c69d8440447e3b5103 100644 --- a/MindSPONGE/mindsponge/python/system/residue/residue.py +++ b/MindSPONGE/mindsponge/python/system/residue/residue.py @@ -458,35 +458,40 @@ class Residue: self.bond = Tensor(self._get_bond(template, self.atom_index), ms.int32) return self - def _get_atom_mass(self, template: dict, atom_index: ndarray = None) -> ndarray: + @classmethod + def _get_atom_mass(cls, template: dict, atom_index: ndarray = None) -> ndarray: """get atom mass from template and atom index""" atom_mass = np.array(template.get('atom_mass'), np.float32) if atom_index is not None: atom_mass = atom_mass[atom_index] return atom_mass - def _get_atomic_number(self, template: dict, atom_index: ndarray = None) -> ndarray: + @classmethod + def _get_atomic_number(cls, template: dict, atom_index: ndarray = None) -> ndarray: """get atomic number from template and atom index""" atomic_number = np.array(template.get('atomic_number'), np.int32) if atom_index is not None: atomic_number = atomic_number[atom_index] return atomic_number - def _get_atom_type(self, template: dict, atom_index: ndarray = None) -> ndarray: + @classmethod + def _get_atom_type(cls, template: dict, atom_index: ndarray = None) -> ndarray: """get atom type from template and atom index""" - atom_type = np.array(template[self.name]['atom_type'], np.str_) + atom_type = np.array(template.get('atom_type'), np.str_) if atom_index is not None: atom_type = atom_type[atom_index] return atom_type - def _get_atom_charge(self, template: dict, atom_index: ndarray = None) -> ndarray: + @classmethod + def _get_atom_charge(cls, template: dict, atom_index: ndarray = None) -> ndarray: """get atom charge from template and atom index""" - atom_charge = np.array(template[self.name]['atom_charge'], np.float32) + atom_charge = np.array(template['atom_charge'], np.float32) if atom_index is not None: atom_charge = atom_charge[atom_index] return atom_charge - def _get_bond(self, template: dict, atom_index: ndarray = None) -> ndarray: + @classmethod + def _get_bond(cls, template: dict, atom_index: ndarray = None) -> ndarray: """get bond from template and atom index""" bond = np.array(template.get('bond')) if atom_index is not None: