BorgwardtLab / proteinshake

Protein structure datasets for machine learning.
https://proteinshake.ai
BSD 3-Clause "New" or "Revised" License
99 stars 8 forks source link

Transforms subpackage + surface computation #128

Closed cgoliver closed 1 year ago

cgoliver commented 1 year ago

Propose a new subpackage proteinshake.transforms. Subclasses of this package take a protein dict as input, apply some processing and return a new protein dictionary. This follows the template of pytorch transforms classes.

We currently have two transforms:

Tranform objects are passed to the representation conversion methods to_point(), to_graph(), to_voxel().

Test it out:

from proteinshake.datasets import TMAlignDataset
from proteinshake.transforms import SurfaceTransform
import tempfile
with tempfile.TemporaryDirectory() as tf:
da = TMAlignDataset(root=tf)
da_surf = da.to_point(
                                      resolution='atom',
                                      transform=SurfaceTransform()
                                      ).torch()

TODO: implement a compose() transform