use of edu.sdsc.mmtf.spark.ml.ProteinSequenceEncoder in project mm-dev by sbl-sdsc.
the class ShapeTypeDemo method main.
public static void main(String[] args) throws IOException {
String path = MmtfReader.getMmtfReducedPath();
if (args.length != 1) {
System.err.println("Usage: " + ShapeTypeDemo.class.getSimpleName() + " <dataset output file");
System.exit(1);
}
SparkConf conf = new SparkConf().setMaster("local[*]").setAppName(ShapeTypeDemo.class.getSimpleName());
JavaSparkContext sc = new JavaSparkContext(conf);
long start = System.nanoTime();
// load a representative PDB chain from the 40% seq. identity Blast Clusters
int sequenceIdentity = 90;
JavaPairRDD<String, StructureDataInterface> pdb = MmtfReader.readSequenceFile(path, sc).flatMapToPair(// extract polymer chains
new StructureToPolymerChains()).filter(// get representative subset
new Pisces(sequenceIdentity, 2.5));
// get a data set with sequence info
Dataset<Row> seqData = PolymerSequenceExtractor.getDataset(pdb);
// convert to BioJava data structure
JavaPairRDD<String, Structure> structures = pdb.mapValues(new StructureToBioJava());
// calculate shape data and convert to dataset
JavaRDD<Row> rows = structures.map(t -> getShapeData(t));
Dataset<Row> data = JavaRDDToDataset.getDataset(rows, "structureChainId", "shape");
// there are only few symmetric chain, leave them out
data = data.filter("shape != 'EXCLUDE'");
// join calculated data with the sequence data
data = seqData.join(data, "structureChainId").cache();
data.show(10);
// create a Word2Vector representation of the protein sequences
ProteinSequenceEncoder encoder = new ProteinSequenceEncoder(data);
// create 2-grams
int n = 2;
// 25-amino residue window size for Word2Vector
int windowSize = 25;
// dimension of feature vector
int vectorSize = 50;
data = encoder.overlappingNgramWord2VecEncode(n, windowSize, vectorSize).cache();
// save data in .parquet file
data.write().mode("overwrite").format("parquet").save(args[0]);
long end = System.nanoTime();
System.out.println((end - start) / 1E9 + " sec.");
sc.close();
}
Aggregations