# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Secondary Structure Blosum62 Encoder Demo
#
# This demo creates a dataset of sequence segments dericed from a non-redundant set. The dataset contains the sequence segment, the DSSP Q8 and DSSP Q3 code of the center residue in a seuqnce segment, and a Blosum62 encoding of the sequence segment.
#
# ## Imports

from pyspark import SparkConf, SparkContext, SQLContext
from mmtfPyspark.ml import ProteinSequenceEncoder
from mmtfPyspark.mappers import StructureToPolymerChains
from mmtfPyspark.filters import ContainsLProteinChain
from mmtfPyspark.datasets import secondaryStructureSegmentExtractor
from mmtfPyspark.webfilters import Pisces
from mmtfPyspark.io import mmtfReader

# ## Configure Spark Context

# +
conf = SparkConf() \
            .setMaster("local[*]") \
            .setAppName("SecondaryStructureBlosumEncoderDemo")

sc = SparkContext(conf = conf)
# -

#  ## Read MMTF Hadoop sequence file and 
#  
#  Create a non-redundant set(<=20% seq. identity) of L-protein chains

# +
path = "../../resources/mmtf_reduced_sample/"
sequenceIdentity = 20
resolution = 2.0
fraction = 0.1
seed = 123

pdb = mmtfReader \
        .read_sequence_file(path, sc) \
        .flatMap(StructureToPolymerChains()) \
        .filter(Pisces(sequenceIdentity, resolution)) \
        .filter(ContainsLProteinChain()) \
        .sample(False, fraction, seed)
# -

# ## Get content

segmentLength = 11
data = secondaryStructureSegmentExtractor.get_dataset(pdb, segmentLength).cache()
print(f"original data   : {data.count()}")

# ## Drop Q3 and sequence duplicates

data = data.dropDuplicates(["labelQ3", "sequence"]).cache()
print(f"- duplicate Q3/seq  : {data.count()}")

# ## Drop sequence duplicates

data = data.dropDuplicates(["sequence"])
print(f"- duplicate seq  : {data.count()}")

# ## Blosum62 Encoding

# +
encoder = ProteinSequenceEncoder(data)
data = encoder.blosum62_encode()

data.printSchema()
data.show(5, False)
# -

# ## Terminate Spark Context

sc.stop()

StarcoderdataJupyterScriptsDedupFiltered