o Uݢgop@sdZddlmZddlZddlZddlmZmZmZm Z m Z ddl m Z ddl mZddlZddlmZmZddlmZddlmZddlmZmZdd lmZdd lmZd Zd d egZ e ddGdddZ!dZ"e"gZ#e#Z$dZ%Gddde&Z'GdddZ(GdddZ)dS)zXsz-TargetFeatureSets.__init__..)target_feature_setsitemsallset)r&r<rZ_setr(r(r)r4Rs zTargetFeatureSets.__init__cCst|tsdS|j|jkS)NF)r?r;rDr&otherr(r(r)__eq__\s  zTargetFeatureSets.__eq__cC t|jSr2)hashrDr8r(r(r)__hash__a zTargetFeatureSets.__hash__ KeysView[str]cC |jSr2)rDkeysr8r(r(r)namesds zTargetFeatureSets.namesItemsView[str, set[int]]cCrPr2rDrEr8r(r(r)rEhrNzTargetFeatureSets.itemscCrP)zReturns a view of the items. .. deprecated::4.0 Use `items` instead. Python 3 doesn't do `iteritems`. rTr8r(r(r) iteritemsks zTargetFeatureSets.iteritemskeyrset[int] | NonecCs |j|Sr2)rDr%)r&rVr(r(r)r%ss zTargetFeatureSets.getrIdict[str, set[int]]cCs|j|dSr2)rDupdaterHr(r(r)rYvszTargetFeatureSets.updatercCttj|jS)zReturns the number of distinct feature indices. This is more efficient than len(get_target_feature_indices() because it doesn't bother to sort the indices. )lenrGunionrDvaluesr8r(r(r)count_target_feature_indicesysz.TargetFeatureSets.count_target_feature_indices list[int]cCrZr2)sortedrGr\rDr]r8r(r(r)get_target_feature_indicessz,TargetFeatureSets.get_target_feature_indicesold_to_new_index_translationdict[int, int]csHttsJi}|jD]\}}fdd|D}|||<qt|S)a{Recreates a list of target features by subsetting it. If the Feature Reference holding this object is ever subset, or the index values are changed, we need to update the target_feature_sets to hold these new values instead. This function recreates a list of target features by subsetting it to only those present in a translation table that maps old index positions to new ones. Args: old_to_new_index_translation: a dictionary that tells what old index position is being converted to a new index position Returns: A new TargetFeatureSets csg|] }|vr|qSr(r()rAZ old_indexrbr(r) s z@TargetFeatureSets.subset_to_reduced_features..)r?dictrDrEr;)r&rbnew_target_featuresrindicesZnew_valsr(rdr)subset_to_reduced_featuress  z,TargetFeatureSets.subset_to_reduced_featuresN)r<r=rr>)rrO)rrS)rVrrrW)rIrXrr>rr)rr_)rbrcrr;)r+r,r-r.r4rJrMpropertyrRrErUr%rYr^rarir(r(r(r)r;Ks         r;c@seZdZdZ dudvd d ZdwddZddZddZddZdxddZ dyddZ dzddZ d{d d!Z d|d#d$Z d}d(d)Zd~d*d+Zd~d,d-Z dudd1d2Zdd6d7Z 8ddd.rmr8r(r(r)+get_feature_ids_excluding_deprecated_probessz.)anyrmr8r(r(r)has_deprecated_probesz&FeatureReference.has_deprecated_probescCrw)z=Return the list of feature types excluding deprecated probes.cSsg|] }|jts|jqSr()rryrrrzr(r(r)res  zRFeatureReference.get_feature_types_excluding_deprecated_probes..r}r8r(r(r)-get_feature_types_excluding_deprecated_probessz>FeatureReference.get_feature_types_excluding_deprecated_probes tuple | NonecCsPd}t|jvr&dd|tjD}t|dksJt|dkr$|dnd}|S)z7Extract antigen control feature from feature reference.Nc Ss6g|]}|jtdkr|j|j|j|jtdfqS)FalseN)rr rrrr%r rzr(r(r)res z8FeatureReference.get_antigen_control..rr)r roselect_features_by_typerrmr[)r&Zcontrolr(r(r)get_antigen_controls  z$FeatureReference.get_antigen_controlfrozenset[list[tuple]]c CsLg}|tjD]}||j|j|j|jt d|jt dfqt |S)zExtract the antigen capture content from feature reference. The conent extracted in this function is what we expect to match between inputs to aggr. N) rrrmr#rrrrr%r r frozenset)r&lstr{r(r(r)get_antigen_capture_contents   z,FeatureReference.get_antigen_capture_contentrIis_pdtuple[bool, str | None]c Cs|}|}t||dkrtt|ddd}tt|ddd}t|t|ks/Jt|D]N\}}||d||dksEJ||d||dksSJ||d||dksaJ||d||dkrpd S||d ||d kr|sd Sq3d S) zChecks if two feature references have the same antigen capture features. returns appropriate error message in case of a mismatch rcS|dSNrr(rBr(r(r)zAFeatureReference.equals_antigen_capture_content..)rVcSrrr(rr(r(r)r rrr )FzThe datasets you are trying to aggregate have incompatible MHC alleles for the same control feature id. Please re-run the original multi pipelines with uniform [antigen-specificity] sections.)FzThe datasets you are trying to aggregate have incompatible control feature ids. Please re-run the original multi pipelines with uniform [antigen-specificity] sections.)TN)rr[ differencer`list enumerate) r&rIrZantigen_content_1Zantigen_content_2Zsorted_antigen_content_1Zsorted_antigen_content_2i_r(r(r)equals_antigen_capture_contents z/FeatureReference.equals_antigen_capture_contentcCst| p | }|p||k}dd|jD}dd|jD}||ko9||ko9||ko9|S)a-Checks if two feature references are equal, ignoring missing targets. Only takes into account equality of the target set if it is present in both feature references. Useful for targeted feature reference compatibility checks in preflights. Excludes deprecated probes. cSg|]}|tkr|qSr( HASHTAG_TAGrAtagr(r(r)reGz=FeatureReference.equals_ignore_target_set..cSrr(rrr(r(r)reHr)has_target_featuresrarorr~)r&rIignore_target_setcompatible_target_setsZ self_all_tagsZother_all_tagsr(r(r)equals_ignore_target_set;s"z)FeatureReference.equals_ignore_target_setcCs,| p | }|p||k}|S)zChecks if two feature references have compatible target sets. Useful for targeted feature reference compatibility checks in preflights. Returns: True if either target set does not exist or if they are identical. rra)r&rIrrr(r(r)has_compatible_target_setSs z*FeatureReference.has_compatible_target_setnew_tags new_labels"Collection[Collection[str]] | Nonecst|dksJ|D] }||jvrtdq |dur'|D]}||iqdSt|jt|kr4td|D] }t|t|ksBJq6t|D]\}fdd|D}ddt|j|D}|||qGdS) aAdd new tags and corresponding labels to self. If new labels are None, empty strings are supplied by default. Args: new_tags (list[str]): a list of new tags new_labels (Optional[Collection[Collection[str]]], optional): per feature list of label values corresponding to the new tags. Defaults to None. Raises: ValueError: Tag already present ValueError: Number of labels is not equal to number of features in feature ref rz(tag {} is already present in feature_refNzAnumber of labels does not match number of features in feature_refc3s|]}|VqdSr2r(r@indr(r)rCsz,FeatureReference.add_tags..cSsi|]\}}|j|qSr(r)rAr{labelr(r(r) s z-FeatureReference.add_tags..)r[ro ValueErroradd_tagrmrzip)r&rrrlabelsZ label_listfeature_id_to_label_dictr(rr)add_tagsas,  zFeatureReference.add_tags tag_to_updater!dict[str, str] | dict[bytes, str]cCsX||jvrtdd|dd|j|jD]}|j|vr)||jd|j|<qdS)a<Updates a tag in the feature reference. The features that need to be updated with their new tag values are passed in. Args: tag_to_update (str): a tag to update. Should be in feature_ref feature_id_to_label_dict (Union[dict[str, str], dict[bytes, str]]): a dict of {feature_id: new_value_of_tag} to update. The features that need to be updated mapped to their new tag values. Raises: ValueError: If tag passed in to be updated does not exist in the feature ref z"Tag to update not in Feature ref. Tag to update  . Tags in matrix r!N)rorrmrr%r)r&rrf_defr(r(r) update_tags     zFeatureReference.update_tagr! tag_to_adddefault_tag_valuercCsZ||jvrtdd|dd|j|jD] }||j||j|<q|j|dS)aPAdds a tag to a Feature Reference. Args: tag_to_add (str): name of the tag to add feature_id_to_label_dict (Union[dict[str, str], dict[bytes, str]]):dict of {feature_id: new_value_of_tag} to update. The features that need to be updated mapped to their new tag values default_tag_value (Optional[str], optional): Value of the tag for features not in feature_id_to_label_dict. Defaults to "". Raises: ValueError: If tag passed in to be added already exists in the feature ref z(Tag to be added already in Feature ref. rrrN)rorrmr%rrr#)r&rrrrr(r(r)rs    zFeatureReference.add_tag feature_ref1 feature_ref2cCsf|j|jksJ|j}|j}|jdur|j}n|jdur |j}n |j}||jt|||j|dS)zKConcatenate two feature references, requires unique ids and identical tags.Nrt)rormr<rYrl)rrZ feature_defs1Z feature_defs2Zcombined_target_featuresr(r(r)joins   zFeatureReference.joincCs|ggddS)Nrtr()clsr(r(r)emptyszFeatureReference.emptycCrKr2)r[rmr8r(r(r)get_num_featuresrNz!FeatureReference.get_num_featurescfdd|jDS)zSReturn a list of feature ids of a particular feature type (e.g. "Gene Expression").cg|] }|jkr|jqSr()rrrzrr(r)resz.r}r&rr(rr)get_feature_ids_by_typerz(FeatureReference.get_feature_ids_by_typer_cr)Ncrr()rr)rArrrr(r)resz9FeatureReference.get_indices_for_type..r}rr(rr)get_indices_for_types z%FeatureReference.get_indices_for_typecsfdd|jD}t|S)zGet sorted list of genomes. Empty string is for reverse compatibility. A specific feature type can optionally be specified. cs,h|]}dus |jkr|jtdqS)Nr!)rrr%GENOME_FEATURE_TAGrzrr(r) s  z/FeatureReference.get_genomes..)rmr`)r&rgenomesr(rr) get_genomess zFeatureReference.get_genomescCs |jduSr2)r<r8r(r(r)rrNz$FeatureReference.has_target_featurescCs|jdur |jSdS)zReturns the number of feature indices. If there are no target features, returns zero. This is more efficient than len(get_target_feature_indices()) because it doesn't bother sorting the indices. Nr)r<r^r8r(r(r)r^s z-FeatureReference.count_target_feature_indiceslist[int] | NonecCs|sdS|jS)z}Gets the indices of on-target features within the FeatureReference. Returns None if there is no target set. N)rr<rar8r(r(r)ra s z+FeatureReference.get_target_feature_indiceslist[bytes] | Nonecs"sdSfddDS)zeGets the feature ids of on-target features. Returns None if there is no target set. Ncsg|]}j|jqSr()rmrrArr8r(r)rerz;FeatureReference.get_target_feature_ids..rr8r(r8r)get_target_feature_idssz'FeatureReference.get_target_feature_idscs"fddt|jD}||S)z_Returns all features that match a given type. Args: feature_type: csg|] \}}|jkr|qSr(rrArfdrr(r)re#r|z.)rrmselect_features)r&rrhr(rr)rs z(FeatureReference.select_features_by_typer0csfddt|jDS)zReturns the indices of features within the FeatureReference with genome. Args: genome: Genome name cs&g|]\}}|jtdkr|qSr!)rr%rrr0r(r)re,s zBFeatureReference.get_feature_indices_by_genome..)rrmr&r0r(rr)get_feature_indices_by_genome&s z.FeatureReference.get_feature_indices_by_genomecCs|||S)zgReturns all features that match a given genome. Args: genome: Genome name )rrrr(r(r)select_features_by_genome2z*FeatureReference.select_features_by_genomecCs$t|jD] }|j|krdSqdS)zDetermines if a feature type is present in the FeatureRef. Args: feature_type: Returns: True | False TF)reversedrmr)r&rfeatr(r(r)has_feature_type:s  z!FeatureReference.has_feature_typecCrw)zjGet a list of feature names. Returns: The names of all of the featture_defs. cSg|]}|jqSr()rr@r(r(r)reNz6FeatureReference.get_feature_names..r}r8r(r(r)get_feature_namesHrz"FeatureReference.get_feature_nameslist[str] | list[bytes]cCrw)zeGet a list of feature IDs. Returns: The IDs of all of the feature_defs. cSrr(rr@r(r(r)reVrz4FeatureReference.get_feature_ids..r}r8r(r(r)get_feature_idsPrz FeatureReference.get_feature_idsrh Iterable[int]c sfdd|D}g}i}t|D]\}}|||j<|t||j|j|j|jdqr@j |}t | dkr?d}nd}t |j|dS)zCreate a new FeatureReference that only contains features with the given indices. Any target sets present are updated, but the target_features attribute in the output is set to None if all target features were removed. csg|]}j|qSr(r}rr8r(r)re^z4FeatureReference.select_features..)rrrrrrNrt)rrr#rrrrrrr<rir[rarlro)r&rhZold_defsZnew_defsZtranslation_maprrrgr(r8r)rXs2  z FeatureReference.select_featuresgroup h5py.Groupr>cstD]fdd|jD}tj||ddq|jD]fdd|jD}tj||ddq|jdurV|tj}|j D]\}}tj||ddt |DddqAt|d|jdS) zWrite to an HDF5 group.csg|]}t|qSr()getattrrzcolr(r)re~rz,FeatureReference.to_hdf5..T) compressioncsg|] }|jdqSr)rr%rzrr(r)resNcSg|]}t|qSr()rr@r(r(r)re _all_tag_keys) r"rmcr_h5Zcreate_hdf5_string_datasetror<Z create_group h5_constantsH5_TARGET_SET_ATTRrEr`)r&rdataZtarget_set_grouprVvalr(rr)to_hdf5zs   zFeatureReference.to_hdf5handler cCs>t|j}t|}|||jD] }|||qdS)zWrite to a csv file (only columns specified in CSV_HEADER). Args: handle: the file-like object the feature ref will be written to N)r"rocsvwriterwriterowrmr*)r&rZ csv_headersrrrr(r(r)to_csvs    zFeatureReference.to_csv h5py.Datasetc sddddd d  dd fdd t|d}ddD}t|dkr0d}g}ttd}t|D]&fdd|D}|ttddt|dq>||||dS)!zLoad from an HDF5 group. Args: group (h5py.Dataset): Group to load from. Returns: feature_ref (FeatureReference): New object. noderrrcSs>|jdurgS|jjtjkr|jdv}t||S|ddS)N)z/features/feature_typez/features/genome)shapedtypecharrZSTR_DTYPE_CHARrZread_hdf5_string_dataset)rmemoizer(r(r) _load_strs    z-FeatureReference.from_hdf5.._load_strpath str | bytesrcSs"|dd}t|tr|S|S)z>Strip off leading slash and convert to str (could be unicode).rN)r?rr$)rr(r(r) _format_paths z0FeatureReference.from_hdf5.._format_pathr!rh5py.Group | h5py.Datasetprefix,Generator[tuple[str, list[str]], None, None]c3sb|D]+}||}|d|}t|tjr ||fVqt|tjr.||EdHqdS)N/)r?h5pyZDatasetZGroup)rrrVitemr)r_h5py_dataset_iteratorrr(r)r s  z:FeatureReference.from_hdf5.._h5py_dataset_iteratorrcSs6i|]\}}tjt|vrtj|dd|DqS)cSrr()rr@r(r(r)rerz9FeatureReference.from_hdf5...)rrr osrbasename)rArVrr(r(r)rs z.FeatureReference.from_hdf5..rNcs.i|]}|rt|t|qSr()r )rAk)rrr(r)rs.rr)rrrrr)ror<)rrrr)rrrrr)rrrrrr) rfrEr[r"ranger#rr FEATURE_TYPE)rrror<rmZ num_featuresrr()rr rrrr) from_hdf5s2       zFeatureReference.from_hdf5r2)rmrnrorr<r=)rrrr)rrv)rr)rr)rr)rr)rIrlrrrr)rIrlrr)rrrr)rrrrr)rrrrrr)rrlrrlrrl)rrlrj)rrrrv)rrrr_)rrrr)rr)rr)rrrrl)r0rrr_)r0rrrl)rrrr)rr)rhrrrl)rrrr>)rr rr>)rrrrl)*r+r,r-r.r4rsrJrMrur~rrrrrrrrrr staticmethodr classmethodrrrrrrr^rarrrrrrrrrrrr(r(r(r)rlsX %       &  ("              "  rl)*r. __future__rrr collections.abcrrrrr dataclassesrtypingr rsixr r cellranger.h5_constantsrZcellranger.hdf5Zhdf5rZ&cellranger.feature.antigen.specificityr r cellranger.rna.libraryr&cellranger.targeted.targeted_constantsrrr"rrDEFAULT_FEATURE_TAGS RESERVED_TAGSr Exceptionr1r;rlr(r(r(r)s2         V