def count_amino_acids(all_aa_chains):
    """
    Count the number of aa *per chain*
    Take a dict from entire aa sequence and return a dictionary of {chain: int}
    """

    # Count the spaces and add 1 to get number of amino acids
    return dict(
        (chain, all_aa_chains[chain].count(' ')) for chain in all_aa_chains)