Takes a NEXUS file and writes SETS block with CharSet statements that partition the characters by the number of states

normalize-and-partition-by-num-states-obs.py

#!/usr/bin/env python
from collections import defaultdict
import dendropy
import sys
import os
_SCRIPT_NAME = os.path.split(sys.argv[0])[-1]
def error(msg):
    sys.stderr.write('{}: Error: {}\n'.format(_SCRIPT_NAME, msg))

OUT = sys.stdout
def partition_by_num_obs_states(dataset):
    for char_matrix in dataset.char_matrices:
        col2char_state_set = defaultdict(set)
        num_cols = None
        assert len(char_matrix.state_alphabets) == 1
        for taxon, data in char_matrix.items():
            if num_cols is not None:
                assert len(data) == num_cols
            else:
                num_cols = len(data)
            for col_index, datum in enumerate(data):
                if datum.is_fundamental_state:
                    s = datum.fundamental_states[0].symbol
                    col2char_state_set[col_index].add(s)
        alphabet = char_matrix.state_alphabets[0]
        pure_states = [i for i in alphabet if i.is_fundamental_state and not i.is_gap_state]
        pure_state_symbols = [i.symbol for i in pure_states]
        num_states_to_exp_sym_set = defaultdict(set)
        by_num_obs_states = defaultdict(list)
        for k in xrange(num_cols):
            ss = col2char_state_set[k]
            num_obs_states = len(ss)
            ess = num_states_to_exp_sym_set[num_obs_states]
            if not ess:
                ess.update(set(pure_state_symbols[:num_obs_states]))
            if ess != ss:
                print k, ess, ss
                assert ess == ss # if this trips, we have don't have the lowest set of symbols
            by_num_obs_states[num_obs_states].append(k)
        OUT.write('BEGIN SETS;\n')
        for ns in xrange(1 + len(pure_states)):
            cs = by_num_obs_states.get(ns)
            if cs is None:
                continue
            OUT.write('  CHARSET {}_STATE_CHARS ='.format(ns))
            for ind in cs:
                OUT.write(' {}'.format(ind + 1))
            OUT.write(' ;\n')
        OUT.write('END;\n')
        

if __name__ == '__main__':
    try:
        assert len(sys.argv) == 2
        fp = sys.argv[1]
    except:
        sys.exit('{}: Error: Expecting 1 argument a NEXUS filepath\n'.format(_SCRIPT_NAME))
    d = dendropy.DataSet()
    try:
        d.read(path=fp, schema='nexus')
    except:
        error('Error reading NEXUS file "{}"'.format(fp))
        raise
    partition_by_num_obs_states(d)