<s> """<STR_LIT>""" <EOL> from __future__ import division <EOL> from __future__ import print_function <EOL> from __future__ import unicode_literals <EOL> import fnmatch <EOL> import json <EOL> import os <EOL> import ga4gh . datamodel as datamodel <EOL> import ga4gh . datamodel . reads as reads <EOL> import ga4gh . datamodel . sequenceAnnotations as sequenceAnnotations <EOL> import ga4gh . datamodel . variants as variants <EOL> import ga4gh . exceptions as exceptions <EOL> import ga4gh . protocol as protocol <EOL> class AbstractDataset ( datamodel . DatamodelObject ) : <EOL> """<STR_LIT>""" <EOL> compoundIdClass = datamodel . DatasetCompoundId <EOL> def __init__ ( self , localId ) : <EOL> super ( AbstractDataset , self ) . __init__ ( None , localId ) <EOL> self . _variantSetIds = [ ] <EOL> self . _variantSetIdMap = { } <EOL> self . _featureSetIds = [ ] <EOL> self . _featureSetIdMap = { } <EOL> self . _readGroupSetIds = [ ] <EOL> self . _readGroupSetIdMap = { } <EOL> self . _readGroupSetNameMap = { } <EOL> self . _variantAnnotationSetIds = [ ] <EOL> self . _variantAnnotationSetIdMap = { } <EOL> self . _description = None <EOL> def addVariantSet ( self , variantSet ) : <EOL> """<STR_LIT>""" <EOL> id_ = variantSet . getId ( ) <EOL> self . _variantSetIdMap [ id_ ] = variantSet <EOL> self . _variantSetIds . append ( id_ ) <EOL> def addVariantAnnotationSet ( self , variantAnnotationSet ) : <EOL> """<STR_LIT>""" <EOL> id_ = variantAnnotationSet . getId ( ) <EOL> self . _variantAnnotationSetIdMap [ id_ ] = variantAnnotationSet <EOL> self . _variantAnnotationSetIds . append ( id_ ) <EOL> def addFeatureSet ( self , featureSet ) : <EOL> """<STR_LIT>""" <EOL> id_ = featureSet . getId ( ) <EOL> self . _featureSetIdMap [ id_ ] = featureSet <EOL> self . _featureSetIds . append ( id_ ) <EOL> def addReadGroupSet ( self , readGroupSet ) : <EOL> """<STR_LIT>""" <EOL> id_ = readGroupSet . getId ( ) <EOL> self . _readGroupSetIdMap [ id_ ] = readGroupSet <EOL> self . _readGroupSetNameMap [ readGroupSet . getLocalId ( ) ] = readGroupSet <EOL> self . _readGroupSetIds . append ( id_ ) <EOL> def toProtocolElement ( self ) : <EOL> dataset = protocol . Dataset ( ) <EOL> dataset . id = self . getId ( ) <EOL> dataset . name = self . getLocalId ( ) <EOL> dataset . description = self . getDescription ( ) <EOL> return dataset <EOL> def getVariantSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ self . _variantSetIdMap [ id_ ] for id_ in self . _variantSetIds ] <EOL> def getNumVariantSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . _variantSetIds ) <EOL> def getVariantAnnotationSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ self . _variantAnnotationSetIdMap [ id_ ] for id_ in <EOL> self . _variantAnnotationSetIds ] <EOL> def getVariantAnnotationSet ( self , id_ ) : <EOL> """<STR_LIT>""" <EOL> if id_ not in self . _variantAnnotationSetIdMap : <EOL> raise exceptions . AnnotationSetNotFoundException ( id_ ) <EOL> return self . _variantAnnotationSetIdMap [ id_ ] <EOL> def getNumVariantAnnotationSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . _variantAnnotationSetIds ) <EOL> def getVariantSet ( self , id_ ) : <EOL> """<STR_LIT>""" <EOL> if id_ not in self . _variantSetIdMap : <EOL> raise exceptions . VariantSetNotFoundException ( id_ ) <EOL> return self . _variantSetIdMap [ id_ ] <EOL> def getVariantSetByIndex ( self , index ) : <EOL> """<STR_LIT>""" <EOL> return self . _variantSetIdMap [ self . _variantSetIds [ index ] ] <EOL> def getVariantAnnotationSetByIndex ( self , index ) : <EOL> """<STR_LIT>""" <EOL> return self . _variantAnnotationSetIdMap [ <EOL> self . _variantAnnotationSetIds [ index ] ] <EOL> def getFeatureSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ self . _featureSetIdMap [ id_ ] for id_ in self . _featureSetIds ] <EOL> def getNumFeatureSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . _featureSetIds ) <EOL> def getFeatureSet ( self , id_ ) : <EOL> """<STR_LIT>""" <EOL> if id_ not in self . _featureSetIdMap : <EOL> raise exceptions . FeatureSetNotFoundException ( id_ ) <EOL> return self . _featureSetIdMap [ id_ ] <EOL> def getFeatureSetByIndex ( self , index ) : <EOL> """<STR_LIT>""" <EOL> return self . _featureSetIdMap [ self . _featureSetIds [ index ] ] <EOL> def getNumReadGroupSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . _readGroupSetIds ) <EOL> def getReadGroupSets ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ self . _readGroupSetIdMap [ id_ ] for id_ in self . _readGroupSetIds ] <EOL> def getReadGroupSetByName ( self , name ) : <EOL> """<STR_LIT>""" <EOL> if name not in self . _readGroupSetNameMap : <EOL> raise exceptions . ReadGroupSetNameNotFoundException ( name ) <EOL> return self . _readGroupSetNameMap [ name ] <EOL> def getReadGroupSetByIndex ( self , index ) : <EOL> """<STR_LIT>""" <EOL> return self . _readGroupSetIdMap [ self . _readGroupSetIds [ index ] ] <EOL> def getReadGroupSet ( self , id_ ) : <EOL> """<STR_LIT>""" <EOL> if id_ not in self . _readGroupSetIdMap : <EOL> raise exceptions . ReadGroupNotFoundException ( id_ ) <EOL> return self . _readGroupSetIdMap [ id_ ] <EOL> def getDescription ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _description <EOL> class SimulatedDataset ( AbstractDataset ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( <EOL> self , localId , referenceSet , randomSeed = <NUM_LIT:0> , <EOL> numVariantSets = <NUM_LIT:1> , numCalls = <NUM_LIT:1> , variantDensity = <NUM_LIT:0.5> , <EOL> numReadGroupSets = <NUM_LIT:1> , numReadGroupsPerReadGroupSet = <NUM_LIT:1> , <EOL> numAlignments = <NUM_LIT:1> , numFeatureSets = <NUM_LIT:1> ) : <EOL> super ( SimulatedDataset , self ) . __init__ ( localId ) <EOL> self . _description = "<STR_LIT>" . format ( localId ) <EOL> for i in range ( numVariantSets ) : <EOL> localId = "<STR_LIT>" . format ( i ) <EOL> seed = randomSeed + i <EOL> variantSet = variants . SimulatedVariantSet ( <EOL> self , localId , seed , numCalls , variantDensity ) <EOL> self . addVariantSet ( variantSet ) <EOL> variantAnnotationSet = variants . SimulatedVariantAnnotationSet ( <EOL> self , "<STR_LIT>" . format ( i ) , variantSet ) <EOL> self . addVariantAnnotationSet ( variantAnnotationSet ) <EOL> for i in range ( numReadGroupSets ) : <EOL> localId = '<STR_LIT>' . format ( i ) <EOL> seed = randomSeed + i <EOL> readGroupSet = reads . SimulatedReadGroupSet ( <EOL> self , localId , referenceSet , seed , <EOL> numReadGroupsPerReadGroupSet , numAlignments ) <EOL> self . addReadGroupSet ( readGroupSet ) <EOL> for i in range ( numFeatureSets ) : <EOL> localId = "<STR_LIT>" . format ( i ) <EOL> seed = randomSeed + i <EOL> featureSet = sequenceAnnotations . SimulatedFeatureSet ( <EOL> self , localId , seed ) <EOL> self . addFeatureSet ( featureSet ) <EOL> class FileSystemDataset ( AbstractDataset ) : <EOL> """<STR_LIT>""" <EOL> variantsDirName = "<STR_LIT>" <EOL> readsDirName = "<STR_LIT>" <EOL> featuresDirName = "<STR_LIT>" <EOL> def __init__ ( self , localId , dataDir , dataRepository ) : <EOL> super ( FileSystemDataset , self ) . __init__ ( localId ) <EOL> self . _dataDir = dataDir <EOL> self . _setMetadata ( ) <EOL> variantSetDir = os . path . join ( dataDir , self . variantsDirName ) <EOL> for localId in os . listdir ( variantSetDir ) : <EOL> relativePath = os . path . join ( variantSetDir , localId ) <EOL> if os . path . isdir ( relativePath ) : <EOL> variantSet = variants . HtslibVariantSet ( <EOL> self , localId , relativePath , dataRepository ) <EOL> self . addVariantSet ( variantSet ) <EOL> if variantSet . isAnnotated ( relativePath ) : <EOL> variantAnnotationSet = variants . HtslibVariantAnnotationSet ( <EOL> self , localId , relativePath , dataRepository , <EOL> variantSet ) <EOL> self . addVariantAnnotationSet ( variantAnnotationSet ) <EOL> readGroupSetDir = os . path . join ( dataDir , self . readsDirName ) <EOL> for filename in os . listdir ( readGroupSetDir ) : <EOL> if fnmatch . fnmatch ( filename , '<STR_LIT>' ) : <EOL> localId , _ = os . path . splitext ( filename ) <EOL> bamPath = os . path . join ( readGroupSetDir , filename ) <EOL> readGroupSet = reads . HtslibReadGroupSet ( <EOL> self , localId , bamPath , dataRepository ) <EOL> self . addReadGroupSet ( readGroupSet ) <EOL> featureSetDir = os . path . join ( dataDir , self . featuresDirName ) <EOL> for filename in os . listdir ( featureSetDir ) : <EOL> if fnmatch . fnmatch ( filename , '<STR_LIT>' ) : <EOL> localId , _ = os . path . splitext ( filename ) <EOL> fullPath = os . path . join ( featureSetDir , filename ) <EOL> featureSet = sequenceAnnotations . Gff3DbFeatureSet ( <EOL> self , localId , fullPath , dataRepository ) <EOL> self . addFeatureSet ( featureSet ) <EOL> def _setMetadata ( self ) : <EOL> metadataFileName = '<STR_LIT>' . format ( self . _dataDir ) <EOL> if os . path . isfile ( metadataFileName ) : <EOL> with open ( metadataFileName ) as metadataFile : <EOL> metadata = json . load ( metadataFile ) <EOL> try : <EOL> self . _description = metadata [ '<STR_LIT:description>' ] <EOL> except KeyError as err : <EOL> raise exceptions . MissingDatasetMetadataException ( <EOL> metadataFileName , str ( err ) ) </s>
