<commit_before># Copyright (c) 2016. Mount Sinai School of Medicine
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import, print_function, division
import logging
from collections import OrderedDict

import pandas as pd
from isovar.protein_sequences import (
    reads_generator_to_protein_sequences_generator,
)

from .mutant_protein_fragment import MutantProteinFragment
from .epitope_prediction import predict_epitopes, slice_epitope_predictions
from .vaccine_peptide import VaccinePeptide

logger = logging.getLogger(__name__)

def vaccine_peptides_for_variant(
        variant,
        isovar_protein_sequences,
        mhc_predictor,
        vaccine_peptide_length,
        padding_around_mutation,
        max_vaccine_peptides_per_variant,
        min_epitope_score):
    """
    Returns list containing (MutantProteinFragment, VaccinePeptideMetrics) pairs
    """
    isovar_protein_sequences = list(isovar_protein_sequences)
    if len(isovar_protein_sequences) == 0:
        logger.info("No protein sequences for %s", variant)
        return []

    protein_fragment = MutantProteinFragment.from_isovar_protein_sequence(
        variant=variant,
        protein_sequence=isovar_protein_sequences[0])

    logger.info(
        "Mutant protein fragment for %s: %s",
        variant,
        protein_fragment)

    epitope_predictions = predict_epitopes(
        mhc_predictor=mhc_predictor,
        protein_fragment=protein_fragment,
        min_epitope_score=min_epitope_score)

    candidate_vaccine_peptides = []

    for offset, candidate_fragment in protein_fragment.top_k_subsequences(
            subsequence_length=vaccine_peptide_length,
            k=2 * padding_around_mutation + 1):
        subsequence_epitope_predictions = slice_epitope_predictions(
            epitope_predictions,
            start_offset=offset,
            end_offset=offset + len(candidate_fragment))

        assert all(
            p.source_sequence == candidate_fragment.amino_acids
            for p in subsequence_epitope_predictions)

        candidate_vaccine_peptide = VaccinePeptide(
            mutant_protein_fragment=candidate_fragment,
            epitope_predictions=subsequence_epitope_predictions)

        logger.debug(
            "%s, combined score: %0.4f",
            candidate_vaccine_peptide,
            candidate_vaccine_peptide.combined_score)
        candidate_vaccine_peptides.append(candidate_vaccine_peptide)

    max_score = max(vp.combined_score for vp in candidate_vaccine_peptides)
    n_total_candidates = len(candidate_vaccine_peptides)

    # only keep candidate vaccines that are within 1% of the maximum
    # combined score
    filtered_candidate_vaccine_peptides = [
        vp
        for vp in candidate_vaccine_peptides
        if vp.combined_score / max_score > 0.99
    ]

    logger.info(
        "Keeping %d/%d vaccine peptides for %s",
        len(filtered_candidate_vaccine_peptides),
        n_total_candidates,
        variant)
    filtered_candidate_vaccine_peptides.sort(key=VaccinePeptide.lexicographic_sort_key)

    if len(filtered_candidate_vaccine_peptides) > 0:
        logger.debug("Top vaccine peptides for %s:", variant)
        for i, vaccine_peptide in enumerate(filtered_candidate_vaccine_peptides):
            logger.debug(
                "%d) %s (combined score = %0.4f)",
                i + 1,
                vaccine_peptide,
                vaccine_peptide.combined_score)
    return filtered_candidate_vaccine_peptides[:max_vaccine_peptides_per_variant]

def generate_vaccine_peptides(
        reads_generator,
        mhc_predictor,
        vaccine_peptide_length,
        padding_around_mutation,
        max_vaccine_peptides_per_variant,
        min_reads_supporting_cdna_sequence,
        min_epitope_score):
    """
    Returns dictionary mapping each variant to list of VaccinePeptide objects.
    """

    # total number of amino acids is the vaccine peptide length plus the
    # number of off-center windows around the mutation
    protein_fragment_sequence_length = (
        vaccine_peptide_length + 2 * padding_around_mutation)

    protein_sequences_generator = reads_generator_to_protein_sequences_generator(
        reads_generator,
        transcript_id_whitelist=None,
        protein_sequence_length=protein_fragment_sequence_length,
        min_reads_supporting_cdna_sequence=min_reads_supporting_cdna_sequence,
        max_protein_sequences_per_variant=1)

    result_dict = {}
    for variant, isovar_protein_sequences in protein_sequences_generator:
        vaccine_peptides = vaccine_peptides_for_variant(
            variant=variant,
            isovar_protein_sequences=isovar_protein_sequences,
            mhc_predictor=mhc_predictor,
            vaccine_peptide_length=vaccine_peptide_length,
            padding_around_mutation=padding_around_mutation,
            max_vaccine_peptides_per_variant=max_vaccine_peptides_per_variant,
            min_epitope_score=min_epitope_score)
        result_dict[variant] = vaccine_peptides
    return result_dict

def ranked_vaccine_peptides(
        reads_generator,
        mhc_predictor,
        vaccine_peptide_length,
        padding_around_mutation,
        max_vaccine_peptides_per_variant,
        min_reads_supporting_cdna_sequence,
        min_epitope_score):
    """
    Returns sorted list whose first element is a Variant and whose second
    element is a list of VaccinePeptide objects.
    """
    variants_to_vaccine_peptides_dict = generate_vaccine_peptides(
        reads_generator=reads_generator,
        mhc_predictor=mhc_predictor,
        vaccine_peptide_length=vaccine_peptide_length,
        padding_around_mutation=padding_around_mutation,
        max_vaccine_peptides_per_variant=max_vaccine_peptides_per_variant,
        min_reads_supporting_cdna_sequence=min_reads_supporting_cdna_sequence,
        min_epitope_score=min_epitope_score)
    result_list = list(variants_to_vaccine_peptides_dict.items())
    result_list.sort(
        key=lambda x: x[1][0].combined_score if len(x[1]) > 0 else 0.0,
        reverse=True)
    return result_list

def dataframe_from_ranked_list(ranked_list_of_variants_with_vaccine_peptides):
    """
    Converts list of (Variant, [VaccinePeptide]) into a DataFrame, dropping
    any variants without vaccine peptide predictions.
    """
    columns = OrderedDict([
        ("genome", []),
        ("chr", []),
        ("pos", []),
        ("ref", []),
        ("alt", []),

    ])
    for field in MutantProteinFragment._fields:
        columns[field] = []

    columns["variant_rank"] = []
    columns["peptide_secondary_rank"] = []
    columns["mutant_epitope_score"] = []
    columns["wildtype_epitope_score"] = []
    columns["expression_score"] = []
    columns["combined_score"] = []

    for i, (variant, vaccine_peptides) in enumerate(
            ranked_list_of_variants_with_vaccine_peptides):
        for j, vaccine_peptide in enumerate(vaccine_peptides):
            columns["genome"].append(variant.reference_name)
            columns["chr"].append(variant.contig)
            columns["pos"].append(variant.original_start)
            columns["ref"].append(variant.original_ref)
            columns["alt"].append(variant.original_alt)

            for field in MutantProteinFragment._fields:
                columns[field].append(getattr(
                    vaccine_peptide.mutant_protein_fragment, field))
            columns["variant_rank"].append(i + 1)
            columns["peptide_secondary_rank"].append(j + 1)
            columns["mutant_epitope_score"].append(vaccine_peptide.mutant_epitope_score)
            columns["wildtype_epitope_score"].append(vaccine_peptide.wildtype_epitope_score)
            columns["expression_score"].append(vaccine_peptide.expression_score)
            columns["combined_score"].append(vaccine_peptide.combined_score)
    return pd.DataFrame(columns, columns=columns.keys())
<commit_msg>Gave min_epitope_score default value of 0<commit_after># Copyright (c) 2016. Mount Sinai School of Medicine
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import, print_function, division
import logging
from collections import OrderedDict

import pandas as pd
from isovar.protein_sequences import (
    reads_generator_to_protein_sequences_generator,
)

from .mutant_protein_fragment import MutantProteinFragment
from .epitope_prediction import predict_epitopes, slice_epitope_predictions
from .vaccine_peptide import VaccinePeptide

logger = logging.getLogger(__name__)

def vaccine_peptides_for_variant(
        variant,
        isovar_protein_sequences,
        mhc_predictor,
        vaccine_peptide_length,
        padding_around_mutation,
        max_vaccine_peptides_per_variant,
        min_epitope_score):
    """
    Returns list containing (MutantProteinFragment, VaccinePeptideMetrics) pairs
    """
    isovar_protein_sequences = list(isovar_protein_sequences)
    if len(isovar_protein_sequences) == 0:
        logger.info("No protein sequences for %s", variant)
        return []

    protein_fragment = MutantProteinFragment.from_isovar_protein_sequence(
        variant=variant,
        protein_sequence=isovar_protein_sequences[0])

    logger.info(
        "Mutant protein fragment for %s: %s",
        variant,
        protein_fragment)

    epitope_predictions = predict_epitopes(
        mhc_predictor=mhc_predictor,
        protein_fragment=protein_fragment,
        min_epitope_score=min_epitope_score)

    candidate_vaccine_peptides = []

    for offset, candidate_fragment in protein_fragment.top_k_subsequences(
            subsequence_length=vaccine_peptide_length,
            k=2 * padding_around_mutation + 1):
        subsequence_epitope_predictions = slice_epitope_predictions(
            epitope_predictions,
            start_offset=offset,
            end_offset=offset + len(candidate_fragment))

        assert all(
            p.source_sequence == candidate_fragment.amino_acids
            for p in subsequence_epitope_predictions)

        candidate_vaccine_peptide = VaccinePeptide(
            mutant_protein_fragment=candidate_fragment,
            epitope_predictions=subsequence_epitope_predictions)

        logger.debug(
            "%s, combined score: %0.4f",
            candidate_vaccine_peptide,
            candidate_vaccine_peptide.combined_score)
        candidate_vaccine_peptides.append(candidate_vaccine_peptide)

    max_score = max(vp.combined_score for vp in candidate_vaccine_peptides)
    n_total_candidates = len(candidate_vaccine_peptides)

    # only keep candidate vaccines that are within 1% of the maximum
    # combined score
    filtered_candidate_vaccine_peptides = [
        vp
        for vp in candidate_vaccine_peptides
        if vp.combined_score / max_score > 0.99
    ]

    logger.info(
        "Keeping %d/%d vaccine peptides for %s",
        len(filtered_candidate_vaccine_peptides),
        n_total_candidates,
        variant)
    filtered_candidate_vaccine_peptides.sort(key=VaccinePeptide.lexicographic_sort_key)

    if len(filtered_candidate_vaccine_peptides) > 0:
        logger.debug("Top vaccine peptides for %s:", variant)
        for i, vaccine_peptide in enumerate(filtered_candidate_vaccine_peptides):
            logger.debug(
                "%d) %s (combined score = %0.4f)",
                i + 1,
                vaccine_peptide,
                vaccine_peptide.combined_score)
    return filtered_candidate_vaccine_peptides[:max_vaccine_peptides_per_variant]

def generate_vaccine_peptides(
        reads_generator,
        mhc_predictor,
        vaccine_peptide_length,
        padding_around_mutation,
        max_vaccine_peptides_per_variant,
        min_reads_supporting_cdna_sequence,
        min_epitope_score):
    """
    Returns dictionary mapping each variant to list of VaccinePeptide objects.
    """

    # total number of amino acids is the vaccine peptide length plus the
    # number of off-center windows around the mutation
    protein_fragment_sequence_length = (
        vaccine_peptide_length + 2 * padding_around_mutation)

    protein_sequences_generator = reads_generator_to_protein_sequences_generator(
        reads_generator,
        transcript_id_whitelist=None,
        protein_sequence_length=protein_fragment_sequence_length,
        min_reads_supporting_cdna_sequence=min_reads_supporting_cdna_sequence,
        max_protein_sequences_per_variant=1)

    result_dict = {}
    for variant, isovar_protein_sequences in protein_sequences_generator:
        vaccine_peptides = vaccine_peptides_for_variant(
            variant=variant,
            isovar_protein_sequences=isovar_protein_sequences,
            mhc_predictor=mhc_predictor,
            vaccine_peptide_length=vaccine_peptide_length,
            padding_around_mutation=padding_around_mutation,
            max_vaccine_peptides_per_variant=max_vaccine_peptides_per_variant,
            min_epitope_score=min_epitope_score)
        result_dict[variant] = vaccine_peptides
    return result_dict

def ranked_vaccine_peptides(
        reads_generator,
        mhc_predictor,
        vaccine_peptide_length,
        padding_around_mutation,
        max_vaccine_peptides_per_variant,
        min_reads_supporting_cdna_sequence,
        min_epitope_score=0):
    """
    Returns sorted list whose first element is a Variant and whose second
    element is a list of VaccinePeptide objects.
    """
    variants_to_vaccine_peptides_dict = generate_vaccine_peptides(
        reads_generator=reads_generator,
        mhc_predictor=mhc_predictor,
        vaccine_peptide_length=vaccine_peptide_length,
        padding_around_mutation=padding_around_mutation,
        max_vaccine_peptides_per_variant=max_vaccine_peptides_per_variant,
        min_reads_supporting_cdna_sequence=min_reads_supporting_cdna_sequence,
        min_epitope_score=min_epitope_score)
    result_list = list(variants_to_vaccine_peptides_dict.items())
    result_list.sort(
        key=lambda x: x[1][0].combined_score if len(x[1]) > 0 else 0.0,
        reverse=True)
    return result_list

def dataframe_from_ranked_list(ranked_list_of_variants_with_vaccine_peptides):
    """
    Converts list of (Variant, [VaccinePeptide]) into a DataFrame, dropping
    any variants without vaccine peptide predictions.
    """
    columns = OrderedDict([
        ("genome", []),
        ("chr", []),
        ("pos", []),
        ("ref", []),
        ("alt", []),

    ])
    for field in MutantProteinFragment._fields:
        columns[field] = []

    columns["variant_rank"] = []
    columns["peptide_secondary_rank"] = []
    columns["mutant_epitope_score"] = []
    columns["wildtype_epitope_score"] = []
    columns["expression_score"] = []
    columns["combined_score"] = []

    for i, (variant, vaccine_peptides) in enumerate(
            ranked_list_of_variants_with_vaccine_peptides):
        for j, vaccine_peptide in enumerate(vaccine_peptides):
            columns["genome"].append(variant.reference_name)
            columns["chr"].append(variant.contig)
            columns["pos"].append(variant.original_start)
            columns["ref"].append(variant.original_ref)
            columns["alt"].append(variant.original_alt)

            for field in MutantProteinFragment._fields:
                columns[field].append(getattr(
                    vaccine_peptide.mutant_protein_fragment, field))
            columns["variant_rank"].append(i + 1)
            columns["peptide_secondary_rank"].append(j + 1)
            columns["mutant_epitope_score"].append(vaccine_peptide.mutant_epitope_score)
            columns["wildtype_epitope_score"].append(vaccine_peptide.wildtype_epitope_score)
            columns["expression_score"].append(vaccine_peptide.expression_score)
            columns["combined_score"].append(vaccine_peptide.combined_score)
    return pd.DataFrame(columns, columns=columns.keys())
<|endoftext|>


<commit_before>from typing import Dict, Tuple

from hwt.hdl.assignment import Assignment
from hwt.hdl.operator import Operator
from hwt.hdl.operatorDefs import AllOps
from hwt.pyUtils.arrayQuery import arr_any
from hwt.serializer.resourceUsageResolver.resourceTypes import Unconnected, \
    ResourceFF, ResourceMUX, ResourceFFwithMux, ResourceLatch, ResourceRAM, \
    ResourceROM, ResourceLatchWithMux, ResourceAsyncRAM, ResourceAsyncROM
from hwt.synthesizer.rtlLevel.rtlSignal import RtlSignal

# tables for resolving of resource type updates
# sequential statements merging
resourceTransitions_override = {
    (Assignment, Assignment): Assignment,  # overwrite
    (Unconnected, Unconnected): Unconnected,

    (Assignment, Unconnected): Assignment,  # written in another command
    (Unconnected, Assignment): Assignment,

    (Assignment, ResourceMUX): ResourceMUX,  # written in another command
    (ResourceMUX, Assignment): ResourceMUX,

    (ResourceFF, Unconnected): ResourceFF,  # written in another command
    (Unconnected, ResourceFF): ResourceFF,

    (ResourceMUX, Unconnected): ResourceMUX,  # written in another command
    (Unconnected, ResourceMUX): ResourceMUX,

    (ResourceFFwithMux, Assignment): ResourceFFwithMux,  # written in another command
    (Assignment, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceFFwithMux, Unconnected): ResourceFFwithMux,  # written in another command
    (Unconnected, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceMUX, ResourceLatch): ResourceMUX,  # written in another command
    (ResourceLatch, ResourceMUX): ResourceMUX,

    (ResourceRAM, Unconnected): ResourceRAM,
    (Unconnected, ResourceRAM): ResourceRAM,

    (ResourceROM, Unconnected): ResourceROM,
    (Unconnected, ResourceROM): ResourceROM,

    (Unconnected, ResourceLatchWithMux): ResourceLatchWithMux,
    (Unconnected, ResourceLatch): ResourceLatchWithMux,

    (Assignment, ResourceLatchWithMux): ResourceMUX,

    }

resourceTransitions_sameBranchLevel = {
    # current : now discovered
    (Assignment, Assignment): ResourceMUX,

    (Unconnected, Unconnected): Unconnected,
    (Assignment, Unconnected): ResourceLatch,
    (Unconnected, Assignment): ResourceLatch,

    (Assignment, ResourceMUX): ResourceMUX,
    (ResourceMUX, Assignment): ResourceMUX,

    (Assignment, Unconnected): ResourceLatchWithMux,
    (Unconnected, Assignment): ResourceLatchWithMux,

    (Assignment, ResourceLatch): ResourceLatchWithMux,
    (ResourceLatch, Assignment): ResourceLatchWithMux,

    (Assignment, ResourceLatchWithMux): ResourceLatchWithMux,
    (ResourceLatchWithMux, Assignment): ResourceLatchWithMux,

    (Unconnected, ResourceLatchWithMux): ResourceLatchWithMux,
    (ResourceLatchWithMux, Unconnected): ResourceLatchWithMux,

    (ResourceFF, ResourceFF): ResourceFFwithMux,
    (ResourceFF, Unconnected): ResourceFF,

    (Assignment, ResourceFF): ResourceFFwithMux,
    (ResourceFF, Assignment): ResourceFFwithMux,

    (ResourceMUX, ResourceFF): ResourceFFwithMux,
    (ResourceFF, ResourceMUX): ResourceFFwithMux,

    (ResourceFFwithMux, Unconnected): ResourceFFwithMux,
    (Unconnected, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceRAM, Unconnected): ResourceRAM,
    (ResourceROM, Unconnected): ResourceROM,
    (ResourceAsyncRAM, Unconnected): ResourceAsyncRAM,
    (ResourceAsyncROM, Unconnected): ResourceAsyncROM,

    }

# resourceTransitions_clockBind = {
#    Assignment: ResourceFF,
#    ResourceFF: ResourceError("Register sensitive to more clock at one time"),
#    ResourceFFwithMux: ResourceError("Register sensitive to more clock at one time"),
#    ResourceLatch: ResourceFF,
#    ResourceLatchWithMux: ResourceFFwithMux,
#
#    ResourceAsyncRAM: ResourceRAM,
#    ResourceAsyncROM: ResourceROM,
#    ResourceRAM: ResourceError("RAM port sensitive to more clock at one time"),
#    ResourceROM: ResourceError("RAM port sensitive to more clock at one time"),
#    }
#
resourceTransitions_memoryAnotherInput = {
    ResourceAsyncROM: ResourceAsyncRAM,
    ResourceAsyncRAM: ResourceAsyncRAM,
    ResourceROM: ResourceRAM,
    ResourceRAM: ResourceRAM,
    }

def updateGuesFromAssignment(gues:Dict, assignment: Assignment) -> None:
    """
    Apply next assignment to this signal on current resoruce gues
    for this signal
    """
    sig = assignment.dst
    if assignment.indexes:
        isRam = arr_any(assignment.indexes,
                        lambda x:isinstance(x, RtlSignal) and not x._const)
    else:
        isRam = False

    try:
        current = gues[sig]
    except KeyError:
        if isRam:
            if assignment.isEventDependent:
                gues[sig] = ResourceRAM
            else:
                gues[sig] = ResourceAsyncRAM
        else:
            if assignment.isEventDependent:
                gues[sig] = ResourceFF
            else:
                gues[sig] = Assignment
        return

    if isRam:
        # [TODO] check for event dependency
        nextType = resourceTransitions_memoryAnotherInput[current]
    else:
        if assignment.isEventDependent:
            g = ResourceFF
        else:
            g = Assignment

        nextType = resourceTransitions_sameBranchLevel[(current, g)]
    gues[sig] = nextType

def mergeGues(gues: Dict, otherGues: Dict) -> None:
    """
    Merge resource gues dicts, "otherGues" into "gues"
    """
    update = {}
    for sig, g0 in gues.items():
        g1 = otherGues[sig]
        nextG = resourceTransitions_sameBranchLevel[(g0, g1)]
        if g0 is not nextG:
            update[sig] = nextG

    gues.update(update)

conversions = {AllOps.BitsAsSigned,
               AllOps.BitsAsUnsigned,
               AllOps.BitsAsVec,
               AllOps.BitsToInt,
               AllOps.IntToBits }

operatorsWithoutResource = conversions.union(
            {AllOps.CONCAT,
             AllOps.CALL,
             AllOps.FALLIGN_EDGE,
             AllOps.RISING_EDGE})

RamOrRomResources = {ResourceAsyncRAM,
                     ResourceAsyncROM,
                     ResourceRAM,
                     ResourceROM}

def findAssingmentOf(sig: RtlSignal) -> Assignment:
    """
    Walk endpoints to neares assignment
    (undirect/direct driver of this signal)
    """
    for ep in sig.endpoints:
        if isinstance(ep, Assignment):
            return ep
        elif isinstance(ep, Operator):
            return findAssingmentOf(ep.result)
        else:
            raise NotImplementedError()

class ResourceContext():
    """
    Container of informations about resources used in architecture

    :ivar unit: optional unit for which is this context build
    :ivar seen: set of seen objects
    :ivar resources: dictionary {type of resource: cnt}
    :ivar discoveredRamSignals: set of signals which seems to be some kind of RAM/ROM memory
    """

    def __init__(self, unit):
        self.unit = unit
        self.seen = set()
        self.resources = {}
        self.discoveredRamSignals = set()

    def registerOperator(self, op: Operator):
        w = op.operands[0]._dtype.bit_length()
        res = self.resources
        k = (op.operator, w)
        r = res.get(k, 0)
        res[k] = r + 1

    def registerMUX_known(self, width, inputsCnt):
        """
        Register MUX of known properties
        """
        k = (ResourceMUX, width, inputsCnt)
        muxs = self.resources.get(k, 0)
        self.resources[k] = muxs + 1

    def registerMUX(self, mux):
        """
        mux record is in format (self.MUX, n, m)
        where n is number of bits of this mux
        and m is number of possible inputs
        """
        inputs = len(mux.drivers)
        if inputs > 1:
            w = mux._dtype.bit_length()
            k = (ResourceMUX, w, inputs)
            muxs = self.resources.get(k, 0)
            self.resources[k] = muxs + 1

    def registerFF(self, ff):
        res = self.resources
        w = ff._dtype.bit_length()
        ffs = res.get(ResourceFF, 0)
        res[ResourceFF] = ffs + w

    def registerLatch(self, latch):
        res = self.resources
        w = latch._dtype.bit_length()
        latches = res.get(ResourceLatch, 0)
        res[ResourceLatch] = latches + w

    def _extractRamPorts(self, mem: RtlSignal) -> Dict[RtlSignal, Tuple[int, int, int, int]]:
        """
        Resolve address signals and read/write ports of memory
        :return: dict  address signal : [syncReads, syncWrites, asyncReads, asyncWrites]
        """

        addressSignals = {}
        # collect write ports
        for d in mem.drivers:
            assert isinstance(d, Assignment)
            assert len(d.indexes) == 1
            index = d.indexes[0]
            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            if d.isEventDependent:
                ports[1] += 1
            else:
                ports[3] += 1
        # collect read ports
        for e in mem.endpoints:
            if isinstance(e, Assignment):
                assert len(e.indexes) == 1
                index = e.indexes
                isEventDependent = e.isEventDependent
            elif isinstance(e, Operator) and e.operator == AllOps.INDEX:
                index = e.operands[1]
                a = findAssingmentOf(e.result)
                isEventDependent = a.isEventDependent
            else:
                raise NotImplementedError(e)

            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            # readportscnt++
            if isEventDependent:
                ports[0] += 1
            else:
                ports[2] += 1

        return addressSignals

    def registerMem(self, mem):
        # [TODO] find clks
        # walk all endpoints and drivers
        # for each address signal collect number of read/write pors
        # merge maximum or r/w ports to rw port for each address signal
        width = mem._dtype.elmType.bit_length()
        items = int(mem._dtype.size)

        # resolve conts of r/w ports
        rwSyncPorts = 0
        rSyncPorts = 0
        wSyncPorts = 0

        rwAsyncPorts = 0
        rAsyncPorts = 0
        wAsyncPorts = 0

        addressSignals = self._extractRamPorts(mem)
        for _, (r, w, asyncR, asyncW) in addressSignals.items():
            rw = min(r, w)
            rwSyncPorts += rw
            rSyncPorts += r - rw
            wSyncPorts += w - rw

            rw = min(asyncR, asyncW)
            rwAsyncPorts += rw
            rAsyncPorts += asyncR - rw
            wAsyncPorts += asyncW - rw

        # register ram resource for this mem
        m = ResourceRAM(width, items,
                        rwSyncPorts, rSyncPorts, wSyncPorts,
                        rwAsyncPorts, rAsyncPorts, wAsyncPorts)
        cnt = self.resources.get(m, 0)
        self.resources[m] = cnt + 1

        # if has sync reads merge FFs into this resource
        FFsInRam = (rwSyncPorts + rSyncPorts) * width
        if FFsInRam:
            ffs = self.resources[ResourceFF]
            if ffs == FFsInRam:
                del self.resources[ResourceFF]
            elif ffs > FFsInRam:
                self.resources[ResourceFF] = ffs - FFsInRam
            else:
                raise Exception("Incompatible ram description (read port did not found FFs as expected)")

    def register(self, sig: RtlSignal, resourceGues):
        """
        Register resource(s) for the signal
        """
        if resourceGues is Assignment:
            pass
        elif resourceGues is ResourceMUX:
            self.registerMUX(sig)
        elif resourceGues is ResourceFF:
            self.registerFF(sig)
        elif resourceGues is ResourceFFwithMux:
            self.registerFF(sig)
            self.registerMUX(sig)
        elif resourceGues is ResourceLatch:
            self.registerLatch(sig)
        elif resourceGues is ResourceLatchWithMux:
            self.registerLatch(sig)
            self.registerMUX(sig)
        elif resourceGues in RamOrRomResources:
            self.discoveredRamSignals.add(sig)
        else:
            raise NotImplementedError(resourceGues)

    def finalize(self):
        """
        Resolve ports of discovered memories
        """
        for sig in self.discoveredRamSignals:
            self.registerMem(sig)
        self.gues = set()
<commit_msg>refactor updateGuesFromAssignment<commit_after>from typing import Dict, Tuple

from hwt.hdl.assignment import Assignment
from hwt.hdl.operator import Operator
from hwt.hdl.operatorDefs import AllOps
from hwt.pyUtils.arrayQuery import arr_any
from hwt.serializer.resourceUsageResolver.resourceTypes import Unconnected, \
    ResourceFF, ResourceMUX, ResourceFFwithMux, ResourceLatch, ResourceRAM, \
    ResourceROM, ResourceLatchWithMux, ResourceAsyncRAM, ResourceAsyncROM
from hwt.synthesizer.rtlLevel.rtlSignal import RtlSignal

# tables for resolving of resource type updates
# sequential statements merging
resourceTransitions_override = {
    (Assignment, Assignment): Assignment,  # overwrite
    (Unconnected, Unconnected): Unconnected,

    (Assignment, Unconnected): Assignment,  # written in another command
    (Unconnected, Assignment): Assignment,

    (Assignment, ResourceMUX): ResourceMUX,  # written in another command
    (ResourceMUX, Assignment): ResourceMUX,

    (ResourceFF, Unconnected): ResourceFF,  # written in another command
    (Unconnected, ResourceFF): ResourceFF,

    (ResourceMUX, Unconnected): ResourceMUX,  # written in another command
    (Unconnected, ResourceMUX): ResourceMUX,

    (ResourceFFwithMux, Assignment): ResourceFFwithMux,  # written in another command
    (Assignment, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceFFwithMux, Unconnected): ResourceFFwithMux,  # written in another command
    (Unconnected, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceMUX, ResourceLatch): ResourceMUX,  # written in another command
    (ResourceLatch, ResourceMUX): ResourceMUX,

    (ResourceRAM, Unconnected): ResourceRAM,
    (Unconnected, ResourceRAM): ResourceRAM,

    (ResourceROM, Unconnected): ResourceROM,
    (Unconnected, ResourceROM): ResourceROM,

    (Unconnected, ResourceLatchWithMux): ResourceLatchWithMux,
    (Unconnected, ResourceLatch): ResourceLatchWithMux,

    (Assignment, ResourceLatchWithMux): ResourceMUX,

    }

resourceTransitions_sameBranchLevel = {
    # current : now discovered
    (Assignment, Assignment): ResourceMUX,

    (Unconnected, Unconnected): Unconnected,
    (Assignment, Unconnected): ResourceLatch,
    (Unconnected, Assignment): ResourceLatch,

    (Assignment, ResourceMUX): ResourceMUX,
    (ResourceMUX, Assignment): ResourceMUX,

    (Assignment, Unconnected): ResourceLatchWithMux,
    (Unconnected, Assignment): ResourceLatchWithMux,

    (Assignment, ResourceLatch): ResourceLatchWithMux,
    (ResourceLatch, Assignment): ResourceLatchWithMux,

    (Assignment, ResourceLatchWithMux): ResourceLatchWithMux,
    (ResourceLatchWithMux, Assignment): ResourceLatchWithMux,

    (Unconnected, ResourceLatchWithMux): ResourceLatchWithMux,
    (ResourceLatchWithMux, Unconnected): ResourceLatchWithMux,

    (ResourceFF, ResourceFF): ResourceFFwithMux,
    (ResourceFF, Unconnected): ResourceFF,

    (Assignment, ResourceFF): ResourceFFwithMux,
    (ResourceFF, Assignment): ResourceFFwithMux,

    (ResourceMUX, ResourceFF): ResourceFFwithMux,
    (ResourceFF, ResourceMUX): ResourceFFwithMux,

    (ResourceFFwithMux, Unconnected): ResourceFFwithMux,
    (Unconnected, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceRAM, Unconnected): ResourceRAM,
    (ResourceROM, Unconnected): ResourceROM,
    (ResourceAsyncRAM, Unconnected): ResourceAsyncRAM,
    (ResourceAsyncROM, Unconnected): ResourceAsyncROM,

    }

resourceTransitions_memoryAnotherInput = {
    ResourceAsyncROM: ResourceAsyncRAM,
    ResourceAsyncRAM: ResourceAsyncRAM,
    ResourceROM: ResourceRAM,
    ResourceRAM: ResourceRAM,
    }

def updateGuesFromAssignment(gues:Dict, assignment: Assignment) -> None:
    """
    Apply next assignment to this signal on current resoruce gues
    for this signal
    """
    sig = assignment.dst
    if assignment.indexes:
        isRam = arr_any(assignment.indexes,
                        lambda x:isinstance(x, RtlSignal) and not x._const)
    else:
        isRam = False

    try:
        current = gues[sig]
    except KeyError:
        if isRam:
            if assignment.isEventDependent:
                g = ResourceRAM
            else:
                g = ResourceAsyncRAM
        else:
            if assignment.isEventDependent:
                g = ResourceFF
            else:
                g = Assignment
        gues[sig] = g
        return

    if isRam:
        # [TODO] check for event dependency
        g = resourceTransitions_memoryAnotherInput[current]
    else:
        if assignment.isEventDependent:
            g = ResourceFF
        else:
            g = Assignment

        g = resourceTransitions_sameBranchLevel[(current, g)]
    gues[sig] = g

def mergeGues(gues: Dict, otherGues: Dict) -> None:
    """
    Merge resource gues dicts, "otherGues" into "gues"
    """
    update = {}
    for sig, g0 in gues.items():
        g1 = otherGues[sig]
        nextG = resourceTransitions_sameBranchLevel[(g0, g1)]
        if g0 is not nextG:
            update[sig] = nextG

    gues.update(update)

conversions = {AllOps.BitsAsSigned,
               AllOps.BitsAsUnsigned,
               AllOps.BitsAsVec,
               AllOps.BitsToInt,
               AllOps.IntToBits }

operatorsWithoutResource = conversions.union(
            {AllOps.CONCAT,
             AllOps.CALL,
             AllOps.FALLIGN_EDGE,
             AllOps.RISING_EDGE})

RamOrRomResources = {ResourceAsyncRAM,
                     ResourceAsyncROM,
                     ResourceRAM,
                     ResourceROM}

def findAssingmentOf(sig: RtlSignal) -> Assignment:
    """
    Walk endpoints to neares assignment
    (undirect/direct driver of this signal)
    """
    for ep in sig.endpoints:
        if isinstance(ep, Assignment):
            return ep
        elif isinstance(ep, Operator):
            return findAssingmentOf(ep.result)
        else:
            raise NotImplementedError()

class ResourceContext():
    """
    Container of informations about resources used in architecture

    :ivar unit: optional unit for which is this context build
    :ivar seen: set of seen objects
    :ivar resources: dictionary {type of resource: cnt}
    :ivar discoveredRamSignals: set of signals which seems to be some kind of RAM/ROM memory
    """

    def __init__(self, unit):
        self.unit = unit
        self.seen = set()
        self.resources = {}
        self.discoveredRamSignals = set()

    def registerOperator(self, op: Operator):
        w = op.operands[0]._dtype.bit_length()
        res = self.resources
        k = (op.operator, w)
        r = res.get(k, 0)
        res[k] = r + 1

    def registerMUX_known(self, width, inputsCnt):
        """
        Register MUX of known properties
        """
        k = (ResourceMUX, width, inputsCnt)
        muxs = self.resources.get(k, 0)
        self.resources[k] = muxs + 1

    def registerMUX(self, mux):
        """
        mux record is in format (self.MUX, n, m)
        where n is number of bits of this mux
        and m is number of possible inputs
        """
        inputs = len(mux.drivers)
        if inputs > 1:
            w = mux._dtype.bit_length()
            k = (ResourceMUX, w, inputs)
            muxs = self.resources.get(k, 0)
            self.resources[k] = muxs + 1

    def registerFF(self, ff):
        res = self.resources
        w = ff._dtype.bit_length()
        ffs = res.get(ResourceFF, 0)
        res[ResourceFF] = ffs + w

    def registerLatch(self, latch):
        res = self.resources
        w = latch._dtype.bit_length()
        latches = res.get(ResourceLatch, 0)
        res[ResourceLatch] = latches + w

    def _extractRamPorts(self, mem: RtlSignal) -> Dict[RtlSignal, Tuple[int, int, int, int]]:
        """
        Resolve address signals and read/write ports of memory
        :return: dict  address signal : [syncReads, syncWrites, asyncReads, asyncWrites]
        """

        addressSignals = {}
        # collect write ports
        for d in mem.drivers:
            assert isinstance(d, Assignment)
            assert len(d.indexes) == 1
            index = d.indexes[0]
            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            if d.isEventDependent:
                ports[1] += 1
            else:
                ports[3] += 1
        # collect read ports
        for e in mem.endpoints:
            if isinstance(e, Assignment):
                assert len(e.indexes) == 1
                index = e.indexes
                isEventDependent = e.isEventDependent
            elif isinstance(e, Operator) and e.operator == AllOps.INDEX:
                index = e.operands[1]
                a = findAssingmentOf(e.result)
                isEventDependent = a.isEventDependent
            else:
                raise NotImplementedError(e)

            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            # readportscnt++
            if isEventDependent:
                ports[0] += 1
            else:
                ports[2] += 1

        return addressSignals

    def registerMem(self, mem):
        # [TODO] find clks
        # walk all endpoints and drivers
        # for each address signal collect number of read/write pors
        # merge maximum or r/w ports to rw port for each address signal
        width = mem._dtype.elmType.bit_length()
        items = int(mem._dtype.size)

        # resolve conts of r/w ports
        rwSyncPorts = 0
        rSyncPorts = 0
        wSyncPorts = 0

        rwAsyncPorts = 0
        rAsyncPorts = 0
        wAsyncPorts = 0

        addressSignals = self._extractRamPorts(mem)
        for _, (r, w, asyncR, asyncW) in addressSignals.items():
            rw = min(r, w)
            rwSyncPorts += rw
            rSyncPorts += r - rw
            wSyncPorts += w - rw

            rw = min(asyncR, asyncW)
            rwAsyncPorts += rw
            rAsyncPorts += asyncR - rw
            wAsyncPorts += asyncW - rw

        # register ram resource for this mem
        m = ResourceRAM(width, items,
                        rwSyncPorts, rSyncPorts, wSyncPorts,
                        rwAsyncPorts, rAsyncPorts, wAsyncPorts)
        cnt = self.resources.get(m, 0)
        self.resources[m] = cnt + 1

        # if has sync reads merge FFs into this resource
        FFsInRam = (rwSyncPorts + rSyncPorts) * width
        if FFsInRam:
            ffs = self.resources[ResourceFF]
            if ffs == FFsInRam:
                del self.resources[ResourceFF]
            elif ffs > FFsInRam:
                self.resources[ResourceFF] = ffs - FFsInRam
            else:
                raise Exception("Incompatible ram description (read port did not found FFs as expected)")

    def register(self, sig: RtlSignal, resourceGues):
        """
        Register resource(s) for the signal
        """
        if resourceGues is Assignment:
            pass
        elif resourceGues is ResourceMUX:
            self.registerMUX(sig)
        elif resourceGues is ResourceFF:
            self.registerFF(sig)
        elif resourceGues is ResourceFFwithMux:
            self.registerFF(sig)
            self.registerMUX(sig)
        elif resourceGues is ResourceLatch:
            self.registerLatch(sig)
        elif resourceGues is ResourceLatchWithMux:
            self.registerLatch(sig)
            self.registerMUX(sig)
        elif resourceGues in RamOrRomResources:
            self.discoveredRamSignals.add(sig)
        else:
            raise NotImplementedError(resourceGues)

    def finalize(self):
        """
        Resolve ports of discovered memories
        """
        for sig in self.discoveredRamSignals:
            self.registerMem(sig)
        self.gues = set()
<|endoftext|>


<commit_before>from typing import Dict

from hwt.hdl.assignment import Assignment
from hwt.hdl.operator import Operator
from hwt.hdl.operatorDefs import AllOps
from hwt.pyUtils.arrayQuery import arr_any
from hwt.serializer.resourceUsageResolver.resourceTypes import Unconnected, \
    ResourceFF, ResourceMUX, ResourceFFwithMux, ResourceLatch, ResourceRAM, \
    ResourceROM, ResourceLatchWithMux, ResourceAsyncRAM, ResourceAsyncROM
from hwt.synthesizer.rtlLevel.rtlSignal import RtlSignal

# tables for resolving of resource type updates
# sequential statements merging
resourceTransitions_override = {
    (Assignment, Assignment): Assignment,  # overwrite
    (Unconnected, Unconnected): Unconnected,

    (Assignment, Unconnected): Assignment,  # written in another command
    (Unconnected, Assignment): Assignment,

    (Assignment, ResourceMUX): ResourceMUX,  # written in another command
    (ResourceMUX, Assignment): ResourceMUX,

    (ResourceFF, Unconnected): ResourceFF,  # written in another command
    (Unconnected, ResourceFF): ResourceFF,

    (ResourceMUX, Unconnected): ResourceMUX,  # written in another command
    (Unconnected, ResourceMUX): ResourceMUX,

    (ResourceFFwithMux, Assignment): ResourceFFwithMux,  # written in another command
    (Assignment, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceFFwithMux, Unconnected): ResourceFFwithMux,  # written in another command
    (Unconnected, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceMUX, ResourceLatch): ResourceMUX,  # written in another command
    (ResourceLatch, ResourceMUX): ResourceMUX,

    (ResourceRAM, Unconnected): ResourceRAM,
    (Unconnected, ResourceRAM): ResourceRAM,

    (ResourceROM, Unconnected): ResourceROM,
    (Unconnected, ResourceROM): ResourceROM,

    }

resourceTransitions_sameBranchLevel = {
    # current : now discovered
    (Assignment, Assignment): ResourceMUX,

    (Unconnected, Unconnected): Unconnected,
    (Assignment, Unconnected): ResourceLatch,
    (Unconnected, Assignment): ResourceLatch,

    (Assignment, ResourceMUX): ResourceMUX,
    (ResourceMUX, Assignment): ResourceMUX,

    (Assignment, Unconnected): ResourceLatchWithMux,
    (Unconnected, Assignment): ResourceLatchWithMux,

    (Assignment, ResourceLatch): ResourceLatchWithMux,
    (ResourceLatch, Assignment): ResourceLatchWithMux,

    (ResourceFF, ResourceFF): ResourceFFwithMux,

    (Assignment, ResourceFF): ResourceFFwithMux,
    (ResourceFF, Assignment): ResourceFFwithMux,

    (ResourceMUX, ResourceFF): ResourceFFwithMux,
    (ResourceFF, ResourceMUX): ResourceFFwithMux,
    }

# resourceTransitions_clockBind = {
#    Assignment: ResourceFF,
#    ResourceFF: ResourceError("Register sensitive to more clock at one time"),
#    ResourceFFwithMux: ResourceError("Register sensitive to more clock at one time"),
#    ResourceLatch: ResourceFF,
#    ResourceLatchWithMux: ResourceFFwithMux,
#
#    ResourceAsyncRAM: ResourceRAM,
#    ResourceAsyncROM: ResourceROM,
#    ResourceRAM: ResourceError("RAM port sensitive to more clock at one time"),
#    ResourceROM: ResourceError("RAM port sensitive to more clock at one time"),
#    }
#
resourceTransitions_memoryAnotherInput = {
    ResourceAsyncROM: ResourceAsyncRAM,
    ResourceAsyncRAM: ResourceAsyncRAM,
    ResourceROM: ResourceRAM,
    ResourceRAM: ResourceRAM,
    }

def updateGuesFromAssignment(gues:Dict, assignment: Assignment) -> None:
    """
    Apply next assignment to this signal on current resoruce gues
    for this signal
    """
    sig = assignment.dst
    if assignment.indexes:
        isRam = arr_any(assignment.indexes,
                        lambda x:isinstance(x, RtlSignal) and not x._const)
    else:
        isRam = False

    try:
        current = gues[sig]
    except KeyError:
        if isRam:
            if assignment.isEventDependent:
                gues[sig] = ResourceRAM
            else:
                gues[sig] = ResourceAsyncRAM
        else:
            if assignment.isEventDependent:
                gues[sig] = ResourceFF
            else:
                gues[sig] = Assignment
        return

    if isRam:
        # [TODO] check for event dependency
        nextType = resourceTransitions_memoryAnotherInput[current]
    else:
        if assignment.isEventDependent:
            g = ResourceFF
        else:
            g = Assignment

        nextType = resourceTransitions_sameBranchLevel[(current, g)]
    gues[sig] = nextType

def mergeGues(gues: Dict, otherGues: Dict) -> None:
    """
    Merge resource gues dicts, "otherGues" into "gues"
    """
    update = {}
    for sig, g0 in gues.items():
        g1 = otherGues[sig]
        nextG = resourceTransitions_sameBranchLevel[(g0, g1)]
        if g0 is not nextG:
            update[sig] = nextG

    gues.update(update)

conversions = {AllOps.BitsAsSigned,
               AllOps.BitsAsUnsigned,
               AllOps.BitsAsVec,
               AllOps.BitsToInt,
               AllOps.IntToBits }

operatorsWithoutResource = conversions.union(
            {AllOps.CONCAT,
             AllOps.CALL,
             AllOps.FALLIGN_EDGE,
             AllOps.RISING_EDGE})

def findAssingmentOf(sig: RtlSignal) -> Assignment:
    """
    Walk endpoints to neares assignment
    (undirect/direct driver of this signal)
    """
    for ep in sig.endpoints:
        if isinstance(ep, Assignment):
            return ep
        elif isinstance(ep, Operator):
            return findAssingmentOf(ep.result)
        else:
            raise NotImplementedError()

class ResourceContext():
    """
    Container of informations about resources used in architecture

    :ivar unit: optional unit for which is this context build
    :ivar seen: set of seen objects
    :ivar resources: dictionary {type of resource: cnt}
    :ivar discoveredRamSignals: set of signals which seems to be some kind of RAM/ROM memory
    """

    def __init__(self, unit):
        self.unit = unit
        self.seen = set()
        self.resources = {}
        self.discoveredRamSignals = set()

    def registerOperator(self, op: Operator):
        w = op.operands[0]._dtype.bit_length()
        res = self.resources
        k = (op.operator, w)
        r = res.get(k, 0)
        res[k] = r + 1

    def registerMUX_known(self, width, inputsCnt):
        """
        Register MUX of known properties
        """
        k = (ResourceMUX, width, inputsCnt)
        muxs = self.resources.get(k, 0)
        self.resources[k] = muxs + 1

    def registerMUX(self, mux):
        """
        mux record is in format (self.MUX, n, m)
        where n is number of bits of this mux
        and m is number of possible inputs
        """
        inputs = len(mux.drivers)
        w = mux._dtype.bit_length()
        k = (ResourceMUX, w, inputs)
        muxs = self.resources.get(k, 0)
        self.resources[k] = muxs + 1

    def registerFF(self, ff):
        res = self.resources
        w = ff._dtype.bit_length()
        ffs = res.get(ResourceFF, 0)
        res[ResourceFF] = ffs + w

    def registerLatch(self, latch):
        res = self.resources
        w = latch._dtype.bit_length()
        latches = res.get(ResourceLatch, 0)
        res[ResourceLatch] = latches + w

    def registerMem(self, mem):
        # [TODO] find clks
        # walk all endpoints and drivers
        # for each address signal collect number of read/write pors
        # merge maximum or r/w ports to rw port for each address signal
        width = mem._dtype.elmType.bit_length()
        items = int(mem._dtype.size)
        # address signal : [reads, writes]
        addressSignals = {}
        for d in mem.drivers:
            assert isinstance(d, Assignment)
            assert len(d.indexes) == 1
            index = d.indexes[0]
            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            if d.isEventDependent:
                ports[1] += 1
            else:
                ports[3] += 1

        for e in mem.endpoints:
            if isinstance(e, Assignment):
                assert len(e.indexes) == 1
                index = e.indexes
                isEventDependent = e.isEventDependent
            elif isinstance(e, Operator) and e.operator == AllOps.INDEX:
                index = e.operands[1]
                a = findAssingmentOf(e.result)
                isEventDependent = a.isEventDependent
            else:
                raise NotImplementedError(e)

            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            if isEventDependent:
                ports[0] += 1
            else:
                ports[2] += 1

        rwSyncPorts = 0
        rSyncPorts = 0
        wSyncPorts = 0

        rwAsyncPorts = 0
        rAsyncPorts = 0
        wAsyncPorts = 0

        for _, (r, w, asyncR, asyncW) in addressSignals.items():
            rw = min(r, w)
            rwSyncPorts += rw
            rSyncPorts += r - rw
            wSyncPorts += w - rw

            rw = min(asyncR, asyncW)
            rwSyncPorts += rw
            rSyncPorts += asyncR - rw
            wSyncPorts += asyncW - rw

        m = ResourceRAM(width, items,
                        rwSyncPorts, rSyncPorts, wSyncPorts,
                        rwAsyncPorts, rAsyncPorts, wAsyncPorts)
        cnt = self.resources.get(m, 0)
        self.resources[m] = cnt + 1

    def register(self, sig: RtlSignal, resourceGues):
        """
        Register resource(s) for the signal
        """
        if resourceGues is Assignment:
            pass
        elif resourceGues is ResourceMUX:
            self.registerMUX(sig)
        elif resourceGues is ResourceFF:
            self.registerFF(sig)
        elif resourceGues is ResourceFFwithMux:
            self.registerFF(sig)
            self.registerMUX(sig)
        elif resourceGues is ResourceLatch:
            self.registerLatch(sig)
        elif resourceGues is ResourceFFwithMux:
            self.registerLatch(sig)
            self.registerMUX(sig)
        elif resourceGues in {ResourceAsyncRAM,
                              ResourceAsyncROM,
                              ResourceRAM,
                              ResourceROM}:
            self.discoveredRamSignals.add(sig)
        else:
            raise NotImplementedError(resourceGues)

    def finalize(self):
        """
        Resolve ports of discovered memories
        """
        for sig in self.discoveredRamSignals:
            self.registerMem(sig)
        self.gues = set()
<commit_msg>ResourceContext: fix ram ports discovery<commit_after>from typing import Dict

from hwt.hdl.assignment import Assignment
from hwt.hdl.operator import Operator
from hwt.hdl.operatorDefs import AllOps
from hwt.pyUtils.arrayQuery import arr_any
from hwt.serializer.resourceUsageResolver.resourceTypes import Unconnected, \
    ResourceFF, ResourceMUX, ResourceFFwithMux, ResourceLatch, ResourceRAM, \
    ResourceROM, ResourceLatchWithMux, ResourceAsyncRAM, ResourceAsyncROM
from hwt.synthesizer.rtlLevel.rtlSignal import RtlSignal

# tables for resolving of resource type updates
# sequential statements merging
resourceTransitions_override = {
    (Assignment, Assignment): Assignment,  # overwrite
    (Unconnected, Unconnected): Unconnected,

    (Assignment, Unconnected): Assignment,  # written in another command
    (Unconnected, Assignment): Assignment,

    (Assignment, ResourceMUX): ResourceMUX,  # written in another command
    (ResourceMUX, Assignment): ResourceMUX,

    (ResourceFF, Unconnected): ResourceFF,  # written in another command
    (Unconnected, ResourceFF): ResourceFF,

    (ResourceMUX, Unconnected): ResourceMUX,  # written in another command
    (Unconnected, ResourceMUX): ResourceMUX,

    (ResourceFFwithMux, Assignment): ResourceFFwithMux,  # written in another command
    (Assignment, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceFFwithMux, Unconnected): ResourceFFwithMux,  # written in another command
    (Unconnected, ResourceFFwithMux): ResourceFFwithMux,

    (ResourceMUX, ResourceLatch): ResourceMUX,  # written in another command
    (ResourceLatch, ResourceMUX): ResourceMUX,

    (ResourceRAM, Unconnected): ResourceRAM,
    (Unconnected, ResourceRAM): ResourceRAM,

    (ResourceROM, Unconnected): ResourceROM,
    (Unconnected, ResourceROM): ResourceROM,

    }

resourceTransitions_sameBranchLevel = {
    # current : now discovered
    (Assignment, Assignment): ResourceMUX,

    (Unconnected, Unconnected): Unconnected,
    (Assignment, Unconnected): ResourceLatch,
    (Unconnected, Assignment): ResourceLatch,

    (Assignment, ResourceMUX): ResourceMUX,
    (ResourceMUX, Assignment): ResourceMUX,

    (Assignment, Unconnected): ResourceLatchWithMux,
    (Unconnected, Assignment): ResourceLatchWithMux,

    (Assignment, ResourceLatch): ResourceLatchWithMux,
    (ResourceLatch, Assignment): ResourceLatchWithMux,

    (ResourceFF, ResourceFF): ResourceFFwithMux,

    (Assignment, ResourceFF): ResourceFFwithMux,
    (ResourceFF, Assignment): ResourceFFwithMux,

    (ResourceMUX, ResourceFF): ResourceFFwithMux,
    (ResourceFF, ResourceMUX): ResourceFFwithMux,
    }

# resourceTransitions_clockBind = {
#    Assignment: ResourceFF,
#    ResourceFF: ResourceError("Register sensitive to more clock at one time"),
#    ResourceFFwithMux: ResourceError("Register sensitive to more clock at one time"),
#    ResourceLatch: ResourceFF,
#    ResourceLatchWithMux: ResourceFFwithMux,
#
#    ResourceAsyncRAM: ResourceRAM,
#    ResourceAsyncROM: ResourceROM,
#    ResourceRAM: ResourceError("RAM port sensitive to more clock at one time"),
#    ResourceROM: ResourceError("RAM port sensitive to more clock at one time"),
#    }
#
resourceTransitions_memoryAnotherInput = {
    ResourceAsyncROM: ResourceAsyncRAM,
    ResourceAsyncRAM: ResourceAsyncRAM,
    ResourceROM: ResourceRAM,
    ResourceRAM: ResourceRAM,
    }

def updateGuesFromAssignment(gues:Dict, assignment: Assignment) -> None:
    """
    Apply next assignment to this signal on current resoruce gues
    for this signal
    """
    sig = assignment.dst
    if assignment.indexes:
        isRam = arr_any(assignment.indexes,
                        lambda x:isinstance(x, RtlSignal) and not x._const)
    else:
        isRam = False

    try:
        current = gues[sig]
    except KeyError:
        if isRam:
            if assignment.isEventDependent:
                gues[sig] = ResourceRAM
            else:
                gues[sig] = ResourceAsyncRAM
        else:
            if assignment.isEventDependent:
                gues[sig] = ResourceFF
            else:
                gues[sig] = Assignment
        return

    if isRam:
        # [TODO] check for event dependency
        nextType = resourceTransitions_memoryAnotherInput[current]
    else:
        if assignment.isEventDependent:
            g = ResourceFF
        else:
            g = Assignment

        nextType = resourceTransitions_sameBranchLevel[(current, g)]
    gues[sig] = nextType

def mergeGues(gues: Dict, otherGues: Dict) -> None:
    """
    Merge resource gues dicts, "otherGues" into "gues"
    """
    update = {}
    for sig, g0 in gues.items():
        g1 = otherGues[sig]
        nextG = resourceTransitions_sameBranchLevel[(g0, g1)]
        if g0 is not nextG:
            update[sig] = nextG

    gues.update(update)

conversions = {AllOps.BitsAsSigned,
               AllOps.BitsAsUnsigned,
               AllOps.BitsAsVec,
               AllOps.BitsToInt,
               AllOps.IntToBits }

operatorsWithoutResource = conversions.union(
            {AllOps.CONCAT,
             AllOps.CALL,
             AllOps.FALLIGN_EDGE,
             AllOps.RISING_EDGE})

def findAssingmentOf(sig: RtlSignal) -> Assignment:
    """
    Walk endpoints to neares assignment
    (undirect/direct driver of this signal)
    """
    for ep in sig.endpoints:
        if isinstance(ep, Assignment):
            return ep
        elif isinstance(ep, Operator):
            return findAssingmentOf(ep.result)
        else:
            raise NotImplementedError()

class ResourceContext():
    """
    Container of informations about resources used in architecture

    :ivar unit: optional unit for which is this context build
    :ivar seen: set of seen objects
    :ivar resources: dictionary {type of resource: cnt}
    :ivar discoveredRamSignals: set of signals which seems to be some kind of RAM/ROM memory
    """

    def __init__(self, unit):
        self.unit = unit
        self.seen = set()
        self.resources = {}
        self.discoveredRamSignals = set()

    def registerOperator(self, op: Operator):
        w = op.operands[0]._dtype.bit_length()
        res = self.resources
        k = (op.operator, w)
        r = res.get(k, 0)
        res[k] = r + 1

    def registerMUX_known(self, width, inputsCnt):
        """
        Register MUX of known properties
        """
        k = (ResourceMUX, width, inputsCnt)
        muxs = self.resources.get(k, 0)
        self.resources[k] = muxs + 1

    def registerMUX(self, mux):
        """
        mux record is in format (self.MUX, n, m)
        where n is number of bits of this mux
        and m is number of possible inputs
        """
        inputs = len(mux.drivers)
        w = mux._dtype.bit_length()
        k = (ResourceMUX, w, inputs)
        muxs = self.resources.get(k, 0)
        self.resources[k] = muxs + 1

    def registerFF(self, ff):
        res = self.resources
        w = ff._dtype.bit_length()
        ffs = res.get(ResourceFF, 0)
        res[ResourceFF] = ffs + w

    def registerLatch(self, latch):
        res = self.resources
        w = latch._dtype.bit_length()
        latches = res.get(ResourceLatch, 0)
        res[ResourceLatch] = latches + w

    def registerMem(self, mem):
        # [TODO] find clks
        # walk all endpoints and drivers
        # for each address signal collect number of read/write pors
        # merge maximum or r/w ports to rw port for each address signal
        width = mem._dtype.elmType.bit_length()
        items = int(mem._dtype.size)
        # address signal : [reads, writes]
        addressSignals = {}
        for d in mem.drivers:
            assert isinstance(d, Assignment)
            assert len(d.indexes) == 1
            index = d.indexes[0]
            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            if d.isEventDependent:
                ports[1] += 1
            else:
                ports[3] += 1

        for e in mem.endpoints:
            if isinstance(e, Assignment):
                assert len(e.indexes) == 1
                index = e.indexes
                isEventDependent = e.isEventDependent
            elif isinstance(e, Operator) and e.operator == AllOps.INDEX:
                index = e.operands[1]
                a = findAssingmentOf(e.result)
                isEventDependent = a.isEventDependent
            else:
                raise NotImplementedError(e)

            try:
                ports = addressSignals[index]
            except KeyError:
                ports = [0, 0, 0, 0]
                addressSignals[index] = ports

            # readportscnt++
            if isEventDependent:
                ports[0] += 1
            else:
                ports[2] += 1

        rwSyncPorts = 0
        rSyncPorts = 0
        wSyncPorts = 0

        rwAsyncPorts = 0
        rAsyncPorts = 0
        wAsyncPorts = 0

        for _, (r, w, asyncR, asyncW) in addressSignals.items():
            rw = min(r, w)
            rwSyncPorts += rw
            rSyncPorts += r - rw
            wSyncPorts += w - rw

            rw = min(asyncR, asyncW)
            rwAsyncPorts += rw
            rAsyncPorts += asyncR - rw
            wAsyncPorts += asyncW - rw

        m = ResourceRAM(width, items,
                        rwSyncPorts, rSyncPorts, wSyncPorts,
                        rwAsyncPorts, rAsyncPorts, wAsyncPorts)
        cnt = self.resources.get(m, 0)
        self.resources[m] = cnt + 1
        FFsInRam = (rwSyncPorts + rSyncPorts) * width
        if FFsInRam:
            ffs = self.resources[ResourceFF]
            if ffs == FFsInRam:
                del self.resources[ResourceFF]
            elif ffs > FFsInRam:
                self.resources[ResourceFF] = ffs - FFsInRam
            else:
                raise Exception("Incompatible ram description (read port did not found FFs as expected)")

    def register(self, sig: RtlSignal, resourceGues):
        """
        Register resource(s) for the signal
        """
        if resourceGues is Assignment:
            pass
        elif resourceGues is ResourceMUX:
            self.registerMUX(sig)
        elif resourceGues is ResourceFF:
            self.registerFF(sig)
        elif resourceGues is ResourceFFwithMux:
            self.registerFF(sig)
            self.registerMUX(sig)
        elif resourceGues is ResourceLatch:
            self.registerLatch(sig)
        elif resourceGues is ResourceFFwithMux:
            self.registerLatch(sig)
            self.registerMUX(sig)
        elif resourceGues in {ResourceAsyncRAM,
                              ResourceAsyncROM,
                              ResourceRAM,
                              ResourceROM}:
            self.discoveredRamSignals.add(sig)
        else:
            raise NotImplementedError(resourceGues)

    def finalize(self):
        """
        Resolve ports of discovered memories
        """
        for sig in self.discoveredRamSignals:
            self.registerMem(sig)
        self.gues = set()
<|endoftext|>


<commit_before>#!/usr/bin/python
import sys, argparse, gzip, re
from Bio.Format.GPD import GPD
from multiprocessing import cpu_count, Pool, Lock

glock = Lock()
of = sys.stdout
sys.setrecursionlimit(10000)

def main():
  parser = argparse.ArgumentParser(description="",formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument('input',help="Use - for STDIN")
  parser.add_argument('--output','-o',help="output file otherwise STDOUT")
  parser.add_argument('--threads',type=int,default=cpu_count(),help="Number of threads to convert names")
  parser.add_argument('-r','--reference',required=True,help="reference gpd")
  args = parser.parse_args()

  global of
  if args.output:
    if re.search('\.gz$',args.output):
      of = gzip.open(args.output,'w')
    else:
      of = open(args.output,'w')

  #read the reference gpd
  rinf = None
  txome = {}
  if re.search('\.gz$',args.reference):
    rinf = gzip.open(args.reference)
  else:
    rinf = open(args.reference)
  sys.stderr.write("Reading in reference\n")
  z = 0
  for line in rinf:
    z += 1
    gpd = GPD(line)
    if z%100 == 0:  sys.stderr.write(str(z)+"          \r")
    if gpd.value('chrom') not in txome: txome[gpd.value('chrom')] = []
    r = gpd.get_range()
    r.set_payload(gpd)
    txome[gpd.value('chrom')].append(r)
  rinf.close()
  sys.stderr.write(str(z)+"          \r")
  sys.stderr.write("\n")
  inf = sys.stdin
  if args.input != '-':
    if re.search('\.gz$',args.input):
      inf = gzip.open(args.input)
    else:
      inf = open(args.input)
  z = 0
  buffer = []
  bsize = 100000
  if args.threads > 1:
    p = Pool(processes=args.threads)
  for line in inf:
    z+=1
    if len(buffer) >= bsize:
      if args.threads > 1:
        p.apply_async(do_buffer,args=(buffer,txome,args),callback=do_out)
      else:
        v = do_buffer(buffer,txome,args)
        do_out(v)
      buffer = []
    buffer.append([line,z])
  if len(buffer) > 0:
      if args.threads > 1:
        p.apply_async(do_buffer,args=(buffer,txome,args),callback=do_out)
      else:
        v = do_buffer(buffer,txome,args)
        do_out(v)
  if args.threads > 1:
    p.close()
    p.join()
  inf.close()
  of.close()

def do_out(results):
  global glock
  global of
  glock.acquire()
  for res in results:
    of.write(res)
  glock.release()

def do_buffer(buffer,txome,args):
  results = []
  for line_z in buffer:
    z = line_z[1]
    line = line_z[0]
    gpd = GPD(line)
    v = annotate_line(gpd,txome,args)
    if not v: continue
    type = 'partial'
    if v[0]: type = 'full'
    exon_count = v[2]
    most_consecutive_exons = v[3]
    read_exon_count = v[4]
    tx_exon_count = v[5]
    overlap_size = v[6]
    read_length = v[7]
    tx_length = v[8]
    results.append(str(z)+"\t"+gpd.get_gene_name()+"\t"+v[9].get_gene_name()+"\t"+v[9].get_transcript_name()+"\t"+type+"\t"+\
          str(exon_count)+"\t"+str(most_consecutive_exons)+"\t"+str(read_exon_count)+"\t"+str(tx_exon_count)+"\t"+\
          str(overlap_size)+"\t"+str(read_length)+"\t"+str(tx_length)+"\n")
  return results

def annotate_line(gpd,txome,args):
  v = gpd.get_range()
  if v.chr not in txome: return None
  possible = [x.get_payload() for x in txome[v.chr] if x.overlaps(v)]
  candidates = []
  if len(possible) == 0: return None
  for tx in possible:
    eo = None
    full = False
    subset = False
    econsec = 1
    if tx.get_exon_count() == 1 or gpd.get_exon_count() == 1:
      eo = gpd.exon_overlap(tx,single_minover=100,single_frac=0.5)
    else:
      eo = gpd.exon_overlap(tx,multi_minover=10,multi_endfrac=0,multi_midfrac=0.8,multi_consec=False)
      if eo.is_full_overlap():
        full = True
      if eo.is_subset():
        subset = True
      if eo:
        econsec = eo.consecutive_exon_count()
    if not eo: continue
    ecnt = eo.match_exon_count()
    osize = gpd.overlap_size(tx)
    candidates.append([full,subset,ecnt,econsec,gpd.get_exon_count(),tx.get_exon_count(),osize,gpd.get_length(),tx.get_length(),tx])
  if len(candidates)==0: return None
  bests = sorted(candidates,key=lambda x: (-x[0],-x[1],-x[3],-x[2],-min(float(x[6])/float(x[7]),float(x[6])/float(x[8]))))
  return bests[0]

if __name__=="__main__":
  main()
<commit_msg>optimizations for annotating reads in gpd format<commit_after>#!/usr/bin/python
import sys, argparse, gzip, re
from Bio.Format.GPD import GPD
from multiprocessing import cpu_count, Pool, Lock

glock = Lock()
of = sys.stdout
sys.setrecursionlimit(10000)
chrdone  = 0
chrtotal = 0
def main():
  parser = argparse.ArgumentParser(description="",formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument('input',help="Use - for STDIN")
  parser.add_argument('--output','-o',help="output file otherwise STDOUT")
  parser.add_argument('--threads',type=int,default=cpu_count(),help="Number of threads to convert names")
  parser.add_argument('-r','--reference',required=True,help="reference gpd")
  args = parser.parse_args()

  global of
  if args.output:
    if re.search('\.gz$',args.output):
      of = gzip.open(args.output,'w')
    else:
      of = open(args.output,'w')

  #read the reference gpd
  rinf = None
  txome = {}
  if re.search('\.gz$',args.reference):
    rinf = gzip.open(args.reference)
  else:
    rinf = open(args.reference)
  sys.stderr.write("Reading in reference\n")
  z = 0
  for line in rinf:
    z += 1
    gpd = GPD(line)
    if z%100 == 0:  sys.stderr.write(str(z)+"          \r")
    if gpd.value('chrom') not in txome: txome[gpd.value('chrom')] = []
    r = gpd.get_range()
    r.set_payload(gpd)
    txome[gpd.value('chrom')].append(r)
  rinf.close()
  sys.stderr.write(str(z)+"          \r")
  sys.stderr.write("\n")
  inf = sys.stdin
  if args.input != '-':
    if re.search('\.gz$',args.input):
      inf = gzip.open(args.input)
    else:
      inf = open(args.input)
  z = 0
  chroms = {}
  sys.stderr.write("Buffering reads\n")
  for line in inf:
    z += 1
    m = re.match('[^\t]*\t[^\t]*\t([^\t]+)',line)
    chrom = m.group(1)
    if z%100==0:  sys.stderr.write(str(z)+"      \r")
    if chrom not in chroms:
      chroms[chrom] = []
    chroms[chrom].append([line,z])
  sys.stderr.write("\n")
  sys.stderr.write("Finished buffering reads\n")
  if args.threads > 1:
    p = Pool(processes=args.threads)
  results = []
  global chrtotal
  chrtotal = len(chroms)
  for chrom in chroms:
    if chrom not in txome: continue
    if args.threads > 1:
      v = p.apply_async(do_buffer,args=(chroms[chrom],{chrom:txome[chrom]},args),callback=do_out)
      results.append(v)
    else:
      v = do_buffer(chroms[chrom],{chrom:txome[chrom]},args)
      results.append(Queue(v))
      do_out(v)
  if args.threads > 1:
    p.close()
    p.join()
  sys.stderr.write("\n")
  for res in [x.get() for x in results]:
    for oline in res:
      of.write(oline)
  inf.close()
  of.close()

class Queue:
  def __init__(self,val):
    self.val = val
  def get(self):
    return self.val

def do_out(results):
  global glock
  global chrtotal
  global chrdone
  glock.acquire()
  chrdone+=1
  sys.stderr.write(str(chrdone)+'/'+str(chrtotal)+" groups finished    \r")
  glock.release()

def do_buffer(buffer,txome,args):
  results = []
  for line_z in buffer:
    z = line_z[1]
    line = line_z[0]
    gpd = GPD(line)
    v = annotate_line(gpd,txome,args)
    if not v: continue
    type = 'partial'
    if v[0]: type = 'full'
    exon_count = v[2]
    most_consecutive_exons = v[3]
    read_exon_count = v[4]
    tx_exon_count = v[5]
    overlap_size = v[6]
    read_length = v[7]
    tx_length = v[8]
    results.append(str(z)+"\t"+gpd.get_gene_name()+"\t"+v[9].get_gene_name()+"\t"+v[9].get_transcript_name()+"\t"+type+"\t"+\
          str(exon_count)+"\t"+str(most_consecutive_exons)+"\t"+str(read_exon_count)+"\t"+str(tx_exon_count)+"\t"+\
          str(overlap_size)+"\t"+str(read_length)+"\t"+str(tx_length)+"\t"+gpd.get_range().get_range_string()+"\t"+v[9].get_range().get_range_string()+"\n")
  return results

def annotate_line(gpd,txome,args):
  v = gpd.get_range()
  if v.chr not in txome: return None
  possible = [x.get_payload() for x in txome[v.chr] if x.overlaps(v)]
  candidates = []
  if len(possible) == 0: return None
  for tx in possible:
    eo = None
    full = False
    subset = False
    econsec = 1
    if tx.get_exon_count() == 1 or gpd.get_exon_count() == 1:
      eo = gpd.exon_overlap(tx,single_minover=100,single_frac=0.5)
    else:
      eo = gpd.exon_overlap(tx,multi_minover=10,multi_endfrac=0,multi_midfrac=0.8,multi_consec=False)
      if eo.is_full_overlap():
        full = True
      if eo.is_subset():
        subset = True
      if eo:
        econsec = eo.consecutive_exon_count()
    if not eo: continue
    ecnt = eo.match_exon_count()
    osize = gpd.overlap_size(tx)
    candidates.append([full,subset,ecnt,econsec,gpd.get_exon_count(),tx.get_exon_count(),osize,gpd.get_length(),tx.get_length(),tx])
  if len(candidates)==0: return None
  bests = sorted(candidates,key=lambda x: (-x[0],-x[1],-x[3],-x[2],-min(float(x[6])/float(x[7]),float(x[6])/float(x[8]))))
  return bests[0]

if __name__=="__main__":
  main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*
#
# Copyright (c) 2013-2014 Christian Brueffer (ORCID: 0000-0002-3826-0989)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
# This script fixes the reads in a TopHat unmapped.bam to make them compatible
# with downstream tools (i.e., Picard and samtools).

import os
import pysam
import sys

def get_index_pos(index, unmapped_reads, read):
    """Returns the position of a read in the index or None."""

    if read.qname in index:
        return index[read.qname]
    else:
        return None

def main(path, outdir, mapped_file="accepted_hits.bam", unmapped_file="unmapped.bam"):
    # Fix things that relate to all unmapped reads.
    unmapped_dict = {}
    unmapped_index = {}
    with pysam.Samfile(os.path.join(path, unmapped_file)) as bam_unmapped:
        unmapped_reads = list(bam_unmapped.fetch(until_eof=True))
        for i in range(len(unmapped_reads)):
            read = unmapped_reads[i]

            # remove /1 and /2 suffixes
            if read.qname.find("/") != -1:
                read.qname = read.qname[:-2]

            unmapped_index[read.qname] = i

            # work around "mate is unmapped" bug in TopHat
            if read.qname in unmapped_dict:
                unmapped_reads[unmapped_dict[read.qname]].mate_is_unmapped = True
                read.mate_is_unmapped = True
            else:
                unmapped_dict[read.qname] = i

            read.mapq = 0

            unmapped_reads[i] = read

        # Fix things that relate only to unmapped reads with a mapped mate.
        with pysam.Samfile(os.path.join(path, mapped_file)) as bam_mapped:
            for mapped in bam_mapped:
                if mapped.mate_is_unmapped:
                    i = get_index_pos(unmapped_index, unmapped_reads, mapped)
                    if i is not None:
                        unmapped = unmapped_reads[i]

                        # map chromosome TIDs from mapped to unmapped file
                        mapped_rname = bam_mapped.getrname(mapped.tid)
                        unmapped_new_tid = bam_unmapped.gettid(mapped_rname)

                        unmapped.tid = unmapped_new_tid
                        unmapped.rnext = unmapped_new_tid
                        unmapped.pos = mapped.pos
                        unmapped.pnext = 0

                        unmapped_reads[i] = unmapped

        # for the output file, take the headers from the unmapped file
        base, ext = os.path.splitext(unmapped_file)
        out_filename = "".join([base, "_fixup", ext])

        with pysam.Samfile(os.path.join(outdir, out_filename), "wb",
                           template=bam_unmapped) as bam_out:
            for read in unmapped_reads:
                bam_out.write(read)

def usage(scriptname):
    print "Usage:"
    print scriptname, "tophat_output_dir"
    sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) == 2:
        path = sys.argv[1]
        if os.path.exists(path) and os.path.isdir(path):
            # no tmpdir specified, use the bam dir
            main(path, path)
        else:
            usage(sys.argv[0])
    elif len(sys.argv) == 3:
        path = sys.argv[1]
        outdir = sys.argv[2]
        if os.path.exists(path) and os.path.isdir(path) and os.path.exists(outdir) and os.path.isdir(outdir):
            main(path, outdir)
        else:
            usage(sys.argv[0])
    else:
        usage(sys.argv[0])
<commit_msg>Use a getopt based commandline frontend.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*
#
# Copyright (c) 2013-2014 Christian Brueffer (ORCID: 0000-0002-3826-0989)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
# This script fixes the reads in a TopHat unmapped.bam to make them compatible
# with downstream tools (i.e., Picard and samtools).

import errno
import os
import pysam
import sys

def get_index_pos(index, unmapped_reads, read):
    """Returns the position of a read in the index or None."""

    if read.qname in index:
        return index[read.qname]
    else:
        return None

def fix_unmapped_reads(path, outdir, mapped_file="accepted_hits.bam", unmapped_file="unmapped.bam"):
    # Fix things that relate to all unmapped reads.
    unmapped_dict = {}
    unmapped_index = {}
    with pysam.Samfile(os.path.join(path, unmapped_file)) as bam_unmapped:
        unmapped_reads = list(bam_unmapped.fetch(until_eof=True))
        for i in range(len(unmapped_reads)):
            read = unmapped_reads[i]

            # remove /1 and /2 suffixes
            if read.qname.find("/") != -1:
                read.qname = read.qname[:-2]

            unmapped_index[read.qname] = i

            # work around "mate is unmapped" bug in TopHat
            if read.qname in unmapped_dict:
                unmapped_reads[unmapped_dict[read.qname]].mate_is_unmapped = True
                read.mate_is_unmapped = True
            else:
                unmapped_dict[read.qname] = i

            read.mapq = 0

            unmapped_reads[i] = read

        # Fix things that relate only to unmapped reads with a mapped mate.
        with pysam.Samfile(os.path.join(path, mapped_file)) as bam_mapped:
            for mapped in bam_mapped:
                if mapped.mate_is_unmapped:
                    i = get_index_pos(unmapped_index, unmapped_reads, mapped)
                    if i is not None:
                        unmapped = unmapped_reads[i]

                        # map chromosome TIDs from mapped to unmapped file
                        mapped_rname = bam_mapped.getrname(mapped.tid)
                        unmapped_new_tid = bam_unmapped.gettid(mapped_rname)

                        unmapped.tid = unmapped_new_tid
                        unmapped.rnext = unmapped_new_tid
                        unmapped.pos = mapped.pos
                        unmapped.pnext = 0

                        unmapped_reads[i] = unmapped

        # for the output file, take the headers from the unmapped file
        base, ext = os.path.splitext(unmapped_file)
        out_filename = "".join([base, "_fixup", ext])

        with pysam.Samfile(os.path.join(outdir, out_filename), "wb",
                           template=bam_unmapped) as bam_out:
            for read in unmapped_reads:
                bam_out.write(read)

def usage(scriptname, errcode=errno.EINVAL):
    print "Usage:\n"
    print scriptname, "[-hv] [tophat_output_dir [result_dir]]\n"
    print "-h                 print this usage text and exit"
    print "-v                 print the script name and version, and exit"
    print "tophat_output_dir: directory containing accepted_hits.bam and unmapped.bam"
    print "result_dir:        directory to write unmapped_fixup.bam to (default: tophat_output_dir)"
    sys.exit(errcode)

if __name__ == "__main__":
    import getopt

    scriptname = os.path.basename(sys.argv[0])
    version = "0.1"

    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], "dhv")
    except getopt.GetoptError, err:
        # print help information and exit
        print >>sys.stderr, str(err)
        usage(scriptname, errcode=errno.EINVAL)

    debug = None
    for o, a in opts:
        if o in ("-d"):
            usage(scriptname, errcode=0)
        elif o in ("-h"):
            usage(scriptname, errcode=0)
        elif o in ("-v"):
            print scriptname, version
            sys.exit(0)
        else:
            assert False, "unhandled option"
            sys.exit(errno.EINVAL)

    if len(args) == 0 or len(args) > 2:
        usage(scriptname, errcode=errno.EINVAL)

    bamdir = args.pop(0)
    if not os.path.isdir(bamdir):
        print >>sys.stderr, "Specified tophat_output_dir does not exist or is not a directory: %s" % bamdir
        sys.exit(errno.EINVAL)

    if len(args) > 0:
        resultdir = args.pop(0)
        if not os.path.isdir(resultdir):
            print >>sys.stderr, "Specified result_dir does not exist or is not a directory: %s" % resultdir
            sys.exit(errno.EINVAL)
    else:
        resultdir = bamdir

    try:
        fix_unmapped_reads(bamdir, resultdir)
    except KeyboardInterrupt:
        print "Program interrupted by user, exiting."
        sys.exit(errno.EINTR)
    except Exception as e:
        if debug:
            import traceback
            print >>sys.stderr, traceback.format_exc()

        print >>sys.stderr, "Error: %s" % str(e)
        if hasattr(e, "errno"):
            sys.exit(e.errno)
        else:
            sys.exit(1)
<|endoftext|>


<commit_before>import praw
from time import ctime
from re import findall
import booknames
from collections import OrderedDict

currentComment = ''
bookNumber = 0

nivbible = OrderedDict()
esvbible = OrderedDict()
kjvbible = OrderedDict()
nrsvbible = OrderedDict()

def constructComment(commands, comment, niv, esv, kjv, nrsv):
    global currentComment, bookNumber, nivbible, esvbible, kjvbible, nrsvbible

    nivbible = niv
    esvbible = esv
    kjvbible = kjv
    nrsvbible = nrsv

    currentComment = ''
    commentFooter = '\n[[Source Code](https://github.com/matthieugrieger/versebot)] [[Feedback](https://github.com/matthieugrieger/versebot/issues)] [[Contact Dev](http://www.reddit.com/message/compose/?to=mgrieger)]'
    for command in commands:
        bookNumber = booknames.getBookNumber(command.lower())
        if bookNumber != False:
            nextCommand = parseCommand(command.lower())
            if nextCommand[0] != False:
                if nextCommand[2] != '0':
                    currentComment += ('**' + booknames.getBookTitle(bookNumber) + ' ' + str(nextCommand[1]) + ':' + str(nextCommand[2]) + ' (*' +
                                       getBibleTranslation(command.lower())[1] + '*)**\n>' + nextCommand[0]) + '\n\n'
                else:
                    currentComment += ('**' + booknames.getBookTitle(bookNumber) + ' ' + str(nextCommand[1]) + ' (*' + getBibleTranslation(command.lower())[1] +
                                       '*)**\n>' + nextCommand[0]) + '\n\n'
    currentComment += commentFooter
    if currentComment != commentFooter:
        if len(currentComment) <= 2000: # Only posts generated response if it is less than or equal to 2000 characters in length
            newComment = comment.reply(currentComment).id
        else:
            errorMessage = 'Oops! It seems that the verses you tried to quote were too long. Instead, here are links to the verses on BibleGateway!\n\n'
            for command in commands:
                linkVerse = '0'
                linkBook = booknames.getBookTitle(booknames.getBookNumber(command.lower()))
                linkCommand = parseCommand(command.lower(), True)
                linkChapter = str(linkCommand[1])
                linkVerse = str(linkCommand[2])
                linkTranslation = getBibleTranslation(command.lower())[1]
                if linkVerse != '0':
                    errorLink = ('http://www.biblegateway.com/passage/?search=' + linkBook + '%20' + linkChapter + ':' + linkVerse + '&version=' + linkTranslation).replace(' ', '%20')
                    errorMessage += ('- [' + linkBook + ' ' + linkChapter + ':' + linkVerse + ' (' + linkTranslation + ')](' + errorLink + ')\n')
                else:
                    errorLink = ('http://www.biblegateway.com/passage/?search=' + linkBook + '%20' + linkChapter + '&version=' + linkTranslation).replace(' ', '%20')
                    errorMessage += ('- [' + linkBook + ' ' + linkChapter + ' (' + linkTranslation + ')](' + errorLink + ')\n')
            errorMessage += commentFooter
            newComment = comment.reply(errorMessage).id
        print('Comment posted on ' + ctime() + '.')
        return newComment # Returns comment id of reply to keep bot from replying to itself
    else:
        return False

def parseCommand(command, error = False):
    global currentComment
    currentChapter = '0'
    currentVerse = '0'
    validComment = True

    if ':' in command:
        chapterAndVerse = str(findall(r'\d+:\d*(?:-\d+)?', command))
        currentChapter = (chapterAndVerse.partition(':')[0])[2:]
        currentVerse = (chapterAndVerse.partition(':')[2])[:-2]
    else:
        currentChapter = str(findall(r'\s\d+', command)).lstrip(' ')
        currentChapter = currentChapter[3:-2]
    if not error:
        if currentVerse != '0' and currentVerse != None:
            try:
                validComment = lookupPassage(booknames.getBookNumber(command.lower()), currentChapter, currentVerse, getBibleTranslation(command.lower())[0])
            except KeyError:
                validComment = False
        else:
            try:
                validComment = lookupPassage(booknames.getBookNumber(command.lower()), currentChapter, False, getBibleTranslation(command.lower())[0])
            except KeyError:
                validComment = False

    return validComment, currentChapter, currentVerse

def lookupPassage(book = False, chapter = False, verse = False, bible = False):
    verseText = ''
    currentSelection = ''

    if book and chapter:
        if verse:
            if '-' in verse:
                startingVerse = verse.partition('-')[0]
                endingVerse = verse.partition('-')[2]
                for ver in range(int(startingVerse), int(endingVerse) + 1):
                    verseText += '[**' + str(ver) + '**] ' + (bible[str(book)][int(chapter)][ver] + ' ')
            else:
                verseText = '[**' + verse + '**] ' + bible[str(book)][int(chapter)][int(verse)]
        else:
            for ver in bible[str(book)][int(chapter)]:
                verseText += '[**' + str(ver) + '**] ' + (bible[str(book)][int(chapter)][ver] + ' ')
        currentSelection += verseText
        return currentSelection

def getBibleTranslation(commentText):
    global nivbible, esvbible, kjvbible, nrsvbible
    if 'niv' in commentText:
        return nivbible, 'NIV'
    elif 'esv' in commentText:
        return esvbible, 'ESV'
    elif 'kjv' in commentText:
        return kjvbible, 'KJV'
    elif 'nrsv' in commentText:
        return nrsvbible, 'NRSV'
    else: # Defaults to ESV if no translation is specified
        return esvbible, 'ESV'
<commit_msg>Added FAQ to comment footer<commit_after>import praw
from time import ctime
from re import findall
import booknames
from collections import OrderedDict

currentComment = ''
bookNumber = 0

nivbible = OrderedDict()
esvbible = OrderedDict()
kjvbible = OrderedDict()
nrsvbible = OrderedDict()

def constructComment(commands, comment, niv, esv, kjv, nrsv):
    global currentComment, bookNumber, nivbible, esvbible, kjvbible, nrsvbible

    nivbible = niv
    esvbible = esv
    kjvbible = kjv
    nrsvbible = nrsv

    currentComment = ''
    commentFooter = '\n[[Source Code](https://github.com/matthieugrieger/versebot)] [[Feedback](https://github.com/matthieugrieger/versebot/issues)] [[Contact Dev](http://www.reddit.com/message/compose/?to=mgrieger)] [[What is this/how do I use it?](https://github.com/matthieugrieger/versebot/blob/master/docs/VerseBot%20Info.md#faq)]'
    for command in commands:
        bookNumber = booknames.getBookNumber(command.lower())
        if bookNumber != False:
            nextCommand = parseCommand(command.lower())
            if nextCommand[0] != False:
                if nextCommand[2] != '0':
                    currentComment += ('**' + booknames.getBookTitle(bookNumber) + ' ' + str(nextCommand[1]) + ':' + str(nextCommand[2]) + ' (*' +
                                       getBibleTranslation(command.lower())[1] + '*)**\n>' + nextCommand[0]) + '\n\n'
                else:
                    currentComment += ('**' + booknames.getBookTitle(bookNumber) + ' ' + str(nextCommand[1]) + ' (*' + getBibleTranslation(command.lower())[1] +
                                       '*)**\n>' + nextCommand[0]) + '\n\n'
    currentComment += commentFooter
    if currentComment != commentFooter:
        if len(currentComment) <= 2000: # Only posts generated response if it is less than or equal to 2000 characters in length
            newComment = comment.reply(currentComment).id
        else:
            errorMessage = 'Oops! It seems that the verses you tried to quote were too long. Instead, here are links to the verses on BibleGateway!\n\n'
            for command in commands:
                linkVerse = '0'
                linkBook = booknames.getBookTitle(booknames.getBookNumber(command.lower()))
                linkCommand = parseCommand(command.lower(), True)
                linkChapter = str(linkCommand[1])
                linkVerse = str(linkCommand[2])
                linkTranslation = getBibleTranslation(command.lower())[1]
                if linkVerse != '0':
                    errorLink = ('http://www.biblegateway.com/passage/?search=' + linkBook + '%20' + linkChapter + ':' + linkVerse + '&version=' + linkTranslation).replace(' ', '%20')
                    errorMessage += ('- [' + linkBook + ' ' + linkChapter + ':' + linkVerse + ' (' + linkTranslation + ')](' + errorLink + ')\n')
                else:
                    errorLink = ('http://www.biblegateway.com/passage/?search=' + linkBook + '%20' + linkChapter + '&version=' + linkTranslation).replace(' ', '%20')
                    errorMessage += ('- [' + linkBook + ' ' + linkChapter + ' (' + linkTranslation + ')](' + errorLink + ')\n')
            errorMessage += commentFooter
            newComment = comment.reply(errorMessage).id
        print('Comment posted on ' + ctime() + '.')
        return newComment # Returns comment id of reply to keep bot from replying to itself
    else:
        return False

def parseCommand(command, error = False):
    global currentComment
    currentChapter = '0'
    currentVerse = '0'
    validComment = True

    if ':' in command:
        chapterAndVerse = str(findall(r'\d+:\d*(?:-\d+)?', command))
        currentChapter = (chapterAndVerse.partition(':')[0])[2:]
        currentVerse = (chapterAndVerse.partition(':')[2])[:-2]
    else:
        currentChapter = str(findall(r'\s\d+', command)).lstrip(' ')
        currentChapter = currentChapter[3:-2]
    if not error:
        if currentVerse != '0' and currentVerse != None:
            try:
                validComment = lookupPassage(booknames.getBookNumber(command.lower()), currentChapter, currentVerse, getBibleTranslation(command.lower())[0])
            except KeyError:
                validComment = False
        else:
            try:
                validComment = lookupPassage(booknames.getBookNumber(command.lower()), currentChapter, False, getBibleTranslation(command.lower())[0])
            except KeyError:
                validComment = False

    return validComment, currentChapter, currentVerse

def lookupPassage(book = False, chapter = False, verse = False, bible = False):
    verseText = ''
    currentSelection = ''

    if book and chapter:
        if verse:
            if '-' in verse:
                startingVerse = verse.partition('-')[0]
                endingVerse = verse.partition('-')[2]
                for ver in range(int(startingVerse), int(endingVerse) + 1):
                    verseText += '[**' + str(ver) + '**] ' + (bible[str(book)][int(chapter)][ver] + ' ')
            else:
                verseText = '[**' + verse + '**] ' + bible[str(book)][int(chapter)][int(verse)]
        else:
            for ver in bible[str(book)][int(chapter)]:
                verseText += '[**' + str(ver) + '**] ' + (bible[str(book)][int(chapter)][ver] + ' ')
        currentSelection += verseText
        return currentSelection

def getBibleTranslation(commentText):
    global nivbible, esvbible, kjvbible, nrsvbible
    if 'niv' in commentText:
        return nivbible, 'NIV'
    elif 'esv' in commentText:
        return esvbible, 'ESV'
    elif 'kjv' in commentText:
        return kjvbible, 'KJV'
    elif 'nrsv' in commentText:
        return nrsvbible, 'NRSV'
    else: # Defaults to ESV if no translation is specified
        return esvbible, 'ESV'
<|endoftext|>