VolumeMaker.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#   Copyright 2021 Gabriele Orlando
#
#   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.

import torch,math
from pyuul.sources.globalVariables import *

import numpy as np
import random

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
setup_seed(100)

class Voxels(torch.nn.Module):
	
	def __init__(self,  device=torch.device("cpu"),sparse=True):
		"""
		Constructor for the Voxels class, which builds the main PyUUL object.

		Parameters
		----------

		device : torch.device
			The device on which the model should run. E.g. torch.device("cuda") or torch.device("cpu:0")
		sparse : bool
			Use sparse tensors calculation when possible

		Returns
		-------
		"""
		super(Voxels, self).__init__()

		self.sparse=sparse
		self.boxsize = None
		self.dev = device

	def __transform_coordinates(self,coords,radius=None):
		"""
		Private function that transform the coordinates to fit them in the 3d box. It also takes care of the resolution.

		Parameters
		----------
		coords : torch.Tensor
			Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 )
		radius : torch.Tensor or None
			Radius of the atoms. Shape ( batch, numberOfAtoms )

		Returns
		-------
		coords : torch.Tensor
			transformed coordinates

		"""
		coords = (coords*self.dilatation)- self.translation
		if not radius is None:
			radius = radius*self.dilatation
			return coords,radius
		else:
			return coords
	'''
	def get_coords_voxel(self, voxel_indices, resolution):
		"""
		returns the coordinates of the center of the voxel provided its indices.

		Parameters
		----------
		voxel_indices : torch.Tensor
			Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 ) 
		resolution : torch.Tensor or None
			Radius of the atoms. Shape ( batch, numberOfAtoms )

		Returns
		-------
		"""
		#voxel_indices is a n,3 long tensor
		centersCoords = voxel_indices + 0.5*resolution
		return (centersCoords + self.translation)/self.dilatation
	'''
	def __define_spatial_conformation(self,mincoords,cubes_around_atoms_dim,resolution):
		"""
		Private function that defines the space of the volume. Takes resolution and margins into consideration.

		Parameters
		----------
		mincoords : torch.Tensor
			minimum coordinates of each macromolecule of the batch. Shape ( batch, 3 )
		cubes_around_atoms_dim : int
			maximum distance in number of voxels to check for atom contribution to occupancy of a voxel
		resolution : float
			side in A of a voxel. The lower this value is the higher the resolution of the final representation will be
		Returns
		-------
		"""
		self.translation=(mincoords-(cubes_around_atoms_dim)).unsqueeze(1)
		self.dilatation = 1.0/resolution

	'''
	def find_cubes_indices(self,coords):
		coords_scaled = self.transform_coordinates(coords)
		return torch.trunc(coords_scaled.data).long()
	'''

	def forward( self,coords, radius,channels,numberchannels=None,resolution=1, cubes_around_atoms_dim=5, steepness=10,function="sigmoid"):
		"""
		Voxels representation of the macromolecules

		Parameters
		----------
		coords : torch.Tensor
			Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 ). Can be calculated from a PDB file using utils.parsePDB
		radius : torch.Tensor
			Radius of the atoms. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToRadius
		channels: torch.LongTensor
			channels of the atoms. Atoms of the same type shold belong to the same channel. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToChannels
		numberchannels : int or None
			maximum number of channels. if None, max(atNameHashing) + 1 is used

		cubes_around_atoms_dim : int
			maximum distance in number of voxels for which the contribution to occupancy is taken into consideration. Every atom that is farer than cubes_around_atoms_dim voxels from the center of a voxel does no give any contribution to the relative voxel occupancy
		resolution : float
			side in A of a voxel. The lower this value is the higher the resolution of the final representation will be

		steepness : float or int
			steepness of the sigmoid occupancy function.

		function : "sigmoid" or "gaussian"
			occupancy function to use. Can be sigmoid (every atom has a sigmoid shaped occupancy function) or gaussian (based on Li et al. 2014)
		Returns
		-------
		volume : torch.Tensor
			voxel representation of the macromolecules in the batch. Shape ( batch, channels, x,y,z), where x,y,z are the size of the 3D volume in which the macromolecules have been represented

		"""
		padding_mask = ~channels.eq(PADDING_INDEX)
		if numberchannels is None:
			numberchannels = int(channels[padding_mask].max().cpu().data+1)
		self.featureVectorSize = numberchannels
		self.function = function

		arange_type = torch.int16

		gx = torch.arange(-cubes_around_atoms_dim, cubes_around_atoms_dim + 1, device=self.dev, dtype=arange_type)
		gy = torch.arange(-cubes_around_atoms_dim, cubes_around_atoms_dim + 1, device=self.dev, dtype=arange_type)
		gz = torch.arange(-cubes_around_atoms_dim, cubes_around_atoms_dim + 1, device=self.dev, dtype=arange_type)
		self.lato = gx.shape[0]

		x1 = gx.unsqueeze(1).expand(self.lato, self.lato).unsqueeze(-1)
		x2 = gy.unsqueeze(0).expand(self.lato, self.lato).unsqueeze(-1)

		xy = torch.cat([x1, x2], dim=-1).unsqueeze(2).expand(self.lato, self.lato, self.lato, 2)
		x3 = gz.unsqueeze(0).unsqueeze(1).expand(self.lato, self.lato, self.lato).unsqueeze(-1)

		del gx, gy, gz, x1, x2

		self.standard_cube = torch.cat([xy, x3], dim=-1).unsqueeze(0).unsqueeze(0)



		### definition of the box ###
		# you take the maximum and min coord on each dimension (every prot in the batch shares the same box. In the future we can pack, but I think this is not the bottleneck)
		# I scale by resolution
		# I add the cubes in which I define the gradient. One in the beginning and one at the end --> 2*



		mincoords = torch.min(coords[:, :, :], dim=1)[0]
		mincoords = torch.trunc(mincoords / resolution)

			
		box_size_x = (math.ceil(torch.max(coords[padding_mask][:,0])/resolution)-mincoords[:,0].min())+(2*cubes_around_atoms_dim+1)
		box_size_y = (math.ceil(torch.max(coords[padding_mask][:,1])/resolution)-mincoords[:,1].min())+(2*cubes_around_atoms_dim+1)
		box_size_z = (math.ceil(torch.max(coords[padding_mask][:,2])/resolution)-mincoords[:,2].min())+(2*cubes_around_atoms_dim+1)
		#############################

		self.__define_spatial_conformation(mincoords,cubes_around_atoms_dim,resolution)	#define the spatial transforms to coordinates
		coords,radius = self.__transform_coordinates(coords,radius)

		boxsize = (int(box_size_x),int(box_size_y),int(box_size_z))
		self.boxsize=boxsize

		#selecting best types for indexing
		if max(boxsize)<256: # i can use byte tensor
			self.dtype_indices=torch.uint8
		else:
			self.dtype_indices = torch.int16

		if self.function=="sigmoid":
			volume = self.__forward_actual_calculation(coords, boxsize, radius, channels,padding_mask,steepness,resolution)
		elif self.function=="gaussian":
			volume = self.__forward_actual_calculationGaussian(coords, boxsize, radius, channels, padding_mask,resolution)
		return volume

	def __forward_actual_calculationGaussian(self, coords_scaled, boxsize, radius, atNameHashing, padding_mask,resolution):
		"""
		private function for the calculation of the gaussian voxel occupancy

		Parameters
		----------
		coords_scaled : torch.LongTensor
			Discrete Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 )
		boxsize : torch.LongTensor
			The size of the box in which the macromolecules are represented
		radius : torch.Tensor
			Radius of the atoms. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToRadius
		atNameHashing: torch.LongTensor
			channels of the atoms. Atoms of the same type shold belong to the same channel. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToChannels
		resolution : float
			side in A of a voxel. The lower this value is the higher the resolution of the final representation will be
		padding_mask : torch.BoolTensor
			tensor to mask the padding. Shape (batch, numberOfAtoms)
		Returns
		-------
		volume : torch.Tensor
			voxel representation of the macromolecules in the batch with Gaussian occupancy function. Shape ( batch, channels, x,y,z), where x,y,z are the size of the 3D volume in which the macromolecules have been represented

		"""
		batch = coords_scaled.shape[0]
		dev = self.dev
		L = coords_scaled.shape[1]

		discrete_coordinates = torch.trunc(coords_scaled.data).to(self.dtype_indices)

		#### making everything in the volume shape

		# implicit_cube_formation
		radius = radius.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		atNameHashing = atNameHashing.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		coords_scaled = coords_scaled.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		discrete_coordinates = discrete_coordinates.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		distmat_standard_cube = torch.norm(
			coords_scaled - ((discrete_coordinates + self.standard_cube + 1) + 0.5 * resolution), dim=-1).to(
			coords_scaled.dtype)

		atNameHashing = atNameHashing.long()
		#### old sigmoid stuff
		'''
		exponent = self.steepness*(distmat_standard_cube-radius)

		exp_mask = exponent.ge(10)
		exponent = torch.masked_fill(exponent,exp_mask, 10)

		volume_cubes = 1.0/(1.0+torch.exp(exponent))
		'''
		### from doi: 10.1142/S0219633614400021 eq 1
		sigma = 0.93
		exponent = -distmat_standard_cube[padding_mask] ** 2 / (sigma ** 2 * radius[padding_mask] ** 2)
		exp_mask = exponent.ge(10)
		exponent = torch.masked_fill(exponent, exp_mask, 10)
		volume_cubes = torch.exp(exponent)

		#### index_put everything ###
		batch_list = torch.arange(batch,device=dev).unsqueeze(1).unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(batch,L,self.lato,self.lato,self.lato)

		cubes_coords = (discrete_coordinates[padding_mask] + self.standard_cube.squeeze(0) + 1)[~exp_mask]
		atNameHashing = atNameHashing[padding_mask].expand(-1,self.lato,self.lato,self.lato)
		if self.sparse:

			index_tens = torch.cat(
					[batch_list[padding_mask][~exp_mask].view(-1).unsqueeze(0),
					 atNameHashing[~exp_mask].unsqueeze(0),
					 cubes_coords[:,0].unsqueeze(0),
					 cubes_coords[:,1].unsqueeze(0),
					 cubes_coords[:,2].unsqueeze(0),
					])
			#index_tens = torch.cat(index)

			volume_cubes = volume_cubes[~exp_mask].view(-1)
			volume_cubes = torch.log(1 - volume_cubes.contiguous())
			#powOrExpIsNotImplementedInSparse
			volume = torch.sparse_coo_tensor(indices=index_tens, values=volume_cubes.exp(), size=[batch, self.featureVectorSize, boxsize[0] , boxsize[1] , boxsize[2] ]).coalesce()
			volume = torch.sparse_coo_tensor(volume.indices(),1 - volume.values(), volume.shape)

		else:
			volume = torch.zeros(batch,boxsize[0]+1,boxsize[1]+1,boxsize[2]+1,self.featureVectorSize,device=dev,dtype=torch.float)
			#index = (batch_list[padding_mask].view(-1),cubes_coords[padding_mask][:,:,:,:,0].view(-1), cubes_coords[padding_mask][:,:,:,:,1].view(-1), cubes_coords[padding_mask][:,:,:,:,2].view(-1), atNameHashing[padding_mask].view(-1) )
			index = (batch_list[padding_mask][~exp_mask].view(-1).long(),
					 cubes_coords[:,0].long(),
					 cubes_coords[:,1].long(),
					 cubes_coords[:,2].long(),
					 atNameHashing[~exp_mask])
			volume_cubes=volume_cubes[~exp_mask].view(-1)

			volume_cubes = torch.log(1 - volume_cubes.contiguous())
			volume = 1- torch.exp(volume.index_put(index,volume_cubes,accumulate=True))
			#volume = 1 - torch.exp(volume.index_put(index, torch.log(1 - volume_cubes.contiguous().view(-1)), accumulate=True))
			volume=volume.permute(0,4,1,2,3)
			#volume = -torch.nn.functional.threshold(-volume,-1,-1)

		return volume



		return volume

	def __sparseClamp(self,volume, minv, maxv):
		vals = volume.values()
		ind = volume.indices()

		vals = vals.clamp(minv, maxv)
		volume = torch.sparse_coo_tensor(indices=ind, values=vals, size=volume.shape).coalesce()
		return volume

	def __forward_actual_calculation(self, coords_scaled, boxsize, radius,atNameHashing,padding_mask,steepness,resolution):
		"""
		private function for the calculation of the gaussian voxel occupancy

		Parameters
		----------
		coords_scaled : torch.LongTensor
			Discrete Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 )
		boxsize : torch.LongTensor
			The size of the box in which the macromolecules are represented
		radius : torch.Tensor
			Radius of the atoms. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToRadius
		atNameHashing: torch.LongTensor
			channels of the atoms. Atoms of the same type shold belong to the same channel. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToChannels
		resolution : float
			side in A of a voxel. The lower this value is the higher the resolution of the final representation will be
		padding_mask : torch.BoolTensor
			tensor to mask the padding. Shape (batch, numberOfAtoms)
		steepness : float
			steepness of the sigmoid function (coefficient of the exponent)

		Returns
		-------
		volume : torch.Tensor
			voxel representation of the macromolecules in the batch with Sigmoid occupancy function. Shape ( batch, channels, x,y,z), where x,y,z are the size of the 3D volume in which the macromolecules have been represented

		"""
		batch = coords_scaled.shape[0]
		dev=self.dev
		L = coords_scaled.shape[1]
		
		discrete_coordinates = torch.trunc(coords_scaled.data).to(self.dtype_indices)

		#### making everything in the volume shape

		#implicit_cube_formation
		radius = radius.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		atNameHashing = atNameHashing.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		coords_scaled = coords_scaled.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		discrete_coordinates = discrete_coordinates.unsqueeze(2).unsqueeze(3).unsqueeze(4)
		distmat_standard_cube = torch.norm(coords_scaled-((discrete_coordinates + self.standard_cube + 1) + 0.5 * resolution), dim=-1).to(coords_scaled.dtype)

		atNameHashing = atNameHashing.long()

		exponent = steepness*(distmat_standard_cube[padding_mask]-radius[padding_mask])
		del distmat_standard_cube
		exp_mask = exponent.ge(10)
		exponent = torch.masked_fill(exponent,exp_mask, 10)

		volume_cubes = 1.0/(1.0+torch.exp(exponent))
		
		#### index_put everything ### 
		batch_list = torch.arange(batch,device=dev).unsqueeze(1).unsqueeze(1).unsqueeze(1).unsqueeze(1).expand(batch,L,self.lato,self.lato,self.lato)

		#cubes_coords = coords_scaled + self.standard_cube + 1
		cubes_coords = (discrete_coordinates[padding_mask] + self.standard_cube.squeeze(0) + 1)[~exp_mask]
		atNameHashing = atNameHashing[padding_mask].expand(-1,self.lato,self.lato,self.lato)
		if self.sparse:

			index_tens = torch.cat(
					[batch_list[padding_mask][~exp_mask].view(-1).unsqueeze(0),
					 atNameHashing[~exp_mask].unsqueeze(0),
					 cubes_coords[:,0].unsqueeze(0),
					 cubes_coords[:,1].unsqueeze(0),
					 cubes_coords[:,2].unsqueeze(0),
					 ])
			#index_tens = torch.cat(index)
			volume = torch.sparse_coo_tensor(indices=index_tens, values=volume_cubes[~exp_mask].view(-1), size=[batch,  self.featureVectorSize, boxsize[0] , boxsize[1] , boxsize[2] ]).coalesce()
			volume = self.__sparseClamp(volume,0,1)

		else:
			volume = torch.zeros(batch,boxsize[0]+1,boxsize[1]+1,boxsize[2]+1,self.featureVectorSize,device=dev,dtype=torch.float)
			#index = (batch_list[padding_mask].view(-1),cubes_coords[padding_mask][:,:,:,:,0].view(-1), cubes_coords[padding_mask][:,:,:,:,1].view(-1), cubes_coords[padding_mask][:,:,:,:,2].view(-1), atNameHashing[padding_mask].view(-1) )
			index = (batch_list[padding_mask][~exp_mask].view(-1).long(),
					 cubes_coords[:,0].long(),
					 cubes_coords[:,1].long(),
					 cubes_coords[:,2].long(),
					 atNameHashing[~exp_mask])
			volume_cubes=volume_cubes[~exp_mask].view(-1)

			volume = volume.index_put(index,volume_cubes.view(-1),accumulate=True)

			volume = -torch.nn.functional.threshold(-volume,-1,-1)
			volume = volume.permute(0,4,1,2,3)

		return volume
	'''
	mesh will be added as soon as pytorch3d becomes a little more stable
	def mesh(self,coords, radius,threshSurface = 0.01):

		atNameHashing= torch.zeros(radius.shape).to(self.dev)
		mask  = radius.eq(PADDING_INDEX)
		atNameHashing = atNameHashing.masked_fill_(mask,PADDING_INDEX)
		vol = self(coords,radius,atNameHashing).to_dense()
		mesh = cubifyNOALIGN(vol.sum(-1),thresh=threshSurface)# creates pytorch 3d mesh from cubes. It uses a MODIFIED version of pytorch3d with no align
		return mesh
	'''

class PointCloudSurface(torch.nn.Module):
	def __init__(self,device="cpu"):
		"""
		Constructor for the CloudPointSurface class, which builds the main PyUUL object for cloud surface.

		Parameters
		----------
		device : torch.device
			The device on which the model should run. E.g. torch.device("cuda") or torch.device("cpu:0")


		Returns
		-------
		"""
		super(PointCloudSurface, self).__init__()

		self.device=device

	def __buildStandardSphere(self,npoints=50): # Fibonacci lattice

		goldenRatio = (1 + 5 ** 0.5) / 2
		i = torch.arange(0, npoints,device=self.device)
		theta = 2 * math.pi * i / goldenRatio
		phi = torch.acos(1 - 2 * (i + 0.5) / npoints)

		x, y, z = torch.cos(theta) * torch.sin(phi), torch.sin(theta) * torch.sin(phi), torch.cos(phi)

		coords=torch.cat([x.unsqueeze(-1),y.unsqueeze(-1),z.unsqueeze(-1)],dim=-1)
		#plot_volume(False,20*coords.unsqueeze(0))

		return coords

	def forward(self, coords, radius, maxpoints=5000,external_radius_factor=1.4):
		"""
		Function to calculate the surface cloud point representation of macromolecules

		Parameters
		----------
		coords : torch.Tensor
			Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 ). Can be calculated from a PDB file using utils.parsePDB
		radius : torch.Tensor
			Radius of the atoms. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToRadius
		maxpoints : int
			number of points per macromolecule in the batch
		external_radius_factor=1.4
			multiplicative factor of the radius in order ot define the place to sample the points around each atom. The higher this value is, the smoother the surface will be
		Returns
		-------
		surfacePointCloud : torch.Tensor
			surface point cloud representation of the macromolecules in the batch. Shape ( batch, channels, numberOfAtoms, 3)

		"""
		padding_mask = ~radius.eq(PADDING_INDEX)

		batch = coords.shape[0]
		npoints = torch.div(maxpoints,(padding_mask.sum(-1).min() + 1), rounding_mode="floor") * 2  # we ensure that the smallest protein has at least maxpoints points

		sphere = self.__buildStandardSphere(npoints)
		finalPoints=[]

		for b in range(batch):

			distmat = torch.cdist(coords[b][padding_mask[b]].unsqueeze(0), coords[b][padding_mask[b]].unsqueeze(0))
			L=distmat.shape[1]
			AtomSelfContributionMask = torch.eye(L, dtype=torch.bool, device=self.device).unsqueeze(0)
			triangular_mask = ~torch.tril(torch.ones((L, L), dtype=torch.bool, device=self.device), diagonal=-1).unsqueeze(0)

			#todoMask = (distmat[b].le(5) & (~AtomSelfContributionMask) & triangular_mask).squeeze(0)
			external_radius = radius * external_radius_factor
			todoMask = (distmat[0].le(5) & (~AtomSelfContributionMask)).squeeze(0)
			points = coords[b][padding_mask[b]].unsqueeze(0).unsqueeze(-2) - sphere.unsqueeze(0).unsqueeze(1) * external_radius[b][padding_mask[b]].unsqueeze(0).unsqueeze(-1).unsqueeze(-1)

			p = points.expand( L, L, npoints, 3)[todoMask]
			c = coords[b][padding_mask[b]].unsqueeze(1).unsqueeze(-2).expand( L, L, points.shape[2], 3)[todoMask]
			r = radius[b][padding_mask[b]].unsqueeze(1).unsqueeze(-2).expand( L, L, points.shape[2])[todoMask]
			occupancy = self.__occupancy(p, c, r)

			point_index = torch.arange(0,L*npoints,device=self.device).view(L,npoints).unsqueeze(0).expand(L,L,npoints)[todoMask]
			point_occupancy =torch.zeros((L*npoints),dtype=torch.float,device=self.device)
			point_occupancy = point_occupancy.index_put_([point_index.view(-1)], occupancy.view(-1), accumulate=True)
			point_occupancy = (1- torch.exp(point_occupancy))

			points_on_surfaceMask = point_occupancy.le(0.5)

			points=points.permute(0,3,1,2).view(3,-1).transpose(0,1)[points_on_surfaceMask]
			random_indices = torch.randint(0, points.shape[0], [maxpoints], device=self.device)
			sampled_points = points[random_indices,:]

			finalPoints +=[sampled_points]

		return torch.cat(finalPoints,dim=0)

	def __occupancy(self, points, coords, radius):

		dist = torch.norm(points-coords,dim=-1)


		sigma=0.93
		exponent = -dist**2/(sigma**2 * radius**2)
		exp_mask = exponent.ge(10)
		exponent = torch.masked_fill(exponent, exp_mask, 10)

		occupancy_on_points = torch.exp(exponent)
		return torch.log(1-occupancy_on_points)
		return occupancy_on_points
		del exponent

		AtomSelfContributionMask = torch.eye(L,dtype=torch.bool,device=self.device).unsqueeze(0).expand(batch,L,L)
		occupancy_on_points[AtomSelfContributionMask]=0.0

		occupancy = (1-torch.exp(torch.log(1-occupancy_on_points).sum(2)))#.sum(dim=-1)/npoints
		#if log_correction:
		#	occupancy = -torch.log(occupancy + 1) # log scaling
		return occupancy

class PointCloudVolume(torch.nn.Module):
	def __init__(self, device="cpu"):
		"""
		Constructor for the CloudPointSurface class, which builds the main PyUUL object for volumetric point cloud.

		Parameters
		----------
		device : torch.device
			The device on which the model should run. E.g. torch.device("cuda") or torch.device("cpu:0")


		Returns
		-------
		"""
		super(PointCloudVolume, self).__init__()

		self.device = device

	def forward(self, coords, radius, maxpoints=500):

		"""
		Function to calculate the volumetric cloud point representation of macromolecules

		Parameters
		----------
		coords : torch.Tensor
			Coordinates of the atoms. Shape ( batch, numberOfAtoms, 3 ). Can be calculated from a PDB file using utils.parsePDB
		radius : torch.Tensor
			Radius of the atoms. Shape ( batch, numberOfAtoms ). Can be calculated from a PDB file using utils.parsePDB and utils.atomlistToRadius
		maxpoints : int
			number of points per macromolecule in the batch

		Returns
		-------
		PointCloudVolume : torch.Tensor
			volume point cloud representation of the macromolecules in the batch. Shape ( batch, channels, numberOfAtoms, 3)

		"""

		padding_mask = ~radius.eq(PADDING_INDEX)

		#npoints = torch.div(maxpoints, padding_mask.sum(-1).min()) + 1 # we ensure that the smallest protein has at least 5000 points

		batch = coords.shape[0]
		L =  coords.shape[1]

		batched = []
		for i in range(batch):
			mean = coords[i][padding_mask[i]]

			sampled = radius[i][padding_mask[i]].sqrt().unsqueeze(-1) * torch.randn((mean.size()), device=self.device) + mean
			p = sampled.view(-1,3)
			random_indices = torch.randint(0, p.shape[0], [maxpoints], device=self.device)
			batched+=[p[random_indices].unsqueeze(0)]

		batched = torch.cat(batched,dim=0)
		return batched