def _get_voxel_center(self, df, mut_chain, mut_res):
if self.center_at_mut:
# Use CA position of the mutated residue as grid center
sel = ((df.chain == mut_chain) & (df.residue == mut_res) &
(df.name == 'CA'))
pos = df[sel][['x', 'y', 'z']].astype(np.float32)
else:
pos = df[['x', 'y', 'z']].astype(np.float32)
return get_center(pos)
def _voxelize(self, atoms):
# Use center of protein as subgrid center
pos = atoms[['x', 'y', 'z']].astype(np.float32)
center = get_center(pos)
# Generate random rotation matrix
rot_mat = gen_rot_matrix(self.grid_config,
random_seed=self.random_seed)
# Transform protein/ligand into voxel grids and rotate
grid = get_grid(atoms,
center,
config=self.grid_config,
rot_mat=rot_mat)
# Last dimension is atom channel, so we need to move it to the front
# per pytroch style
grid = np.moveaxis(grid, -1, 0)
return grid
def _voxelize(self, atoms, is_active):
# Use center of ligand as subgrid center
ligand_pos = atoms[atoms.chain == 'L'][['x', 'y',
'z']].astype(np.float32)
ligand_center = get_center(ligand_pos)
# Generate random rotation matrix
rot_mat = gen_rot_matrix(self.grid_config,
random_seed=self.random_seed)
# Transform protein/ligand into voxel grids and rotate
grid = get_grid(atoms,
ligand_center,
config=self.grid_config,
rot_mat=rot_mat)
if self.add_flag:
# Add inactive (0) or active (1) flag
flag = np.full(grid.shape[:-1] + (1, ), is_active)
grid = np.concatenate([grid, flag], axis=3)
# Last dimension is atom channel, so we need to move it to the front
# per pytroch style
grid = np.moveaxis(grid, -1, 0)
return grid