#!/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 *
[docs]class Voxels(torch.nn.Module):
[docs] 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()
'''
[docs] 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
'''
[docs]class PointCloudSurface(torch.nn.Module):
[docs] 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
[docs] 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 = (maxpoints // padding_mask.sum(-1).min() + 1) * 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
[docs]class PointCloudVolume(torch.nn.Module):
[docs] 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
[docs] def forward(self, coords, radius, maxpoints=5000):
"""
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