io.repseq.gen.dist.DJDependentVDJCGenesModel Maven / Gradle / Ivy

package io.repseq.gen.dist;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import io.repseq.core.GeneType;
import io.repseq.core.VDJCGene;
import io.repseq.core.VDJCLibrary;
import io.repseq.gen.VDJCGenes;
import org.apache.commons.math3.distribution.EnumeratedDistribution;
import org.apache.commons.math3.random.RandomGenerator;
import org.apache.commons.math3.util.Pair;

import java.util.*;

import static io.repseq.gen.dist.IndependentVDJCGenesModel.geneOrNull;
import static io.repseq.gen.dist.IndependentVDJCGenesModel.toDistribution;

public final class DJDependentVDJCGenesModel implements VDJCGenesModel {
    final Map v, dj, c;

    @JsonCreator
    public DJDependentVDJCGenesModel(@JsonProperty("v") Map v,
                                     @JsonProperty("dj") Map dj,
                                     @JsonProperty("c") Map c) {
        this.v = v;
        this.dj = dj;
        this.c = c;
    }

    public static EnumeratedDistribution toDistribution2(RandomGenerator random,
                                                                    VDJCLibrary library,
                                                                    Map distMap,
                                                                    GeneType geneType1, GeneType geneType2) {
        List> ps = new ArrayList<>();
        for (Map.Entry e : distMap.entrySet()) {
            String[] split = e.getKey().split("\\|");
            if (split.length != 2)
                throw new IllegalArgumentException("Wrong key format (expected single '|' symbol): " + e.getKey());

            ps.add(new Pair<>(new VDJCGene2(geneOrNull(library, split[0], geneType1),
                    geneOrNull(library, split[1], geneType2)), e.getValue()));
        }
        return new EnumeratedDistribution<>(random, ps);
    }

    public static final class VDJCGene2 {
        public final VDJCGene gene1, gene2;

        public VDJCGene2(VDJCGene gene1, VDJCGene gene2) {
            this.gene1 = gene1;
            this.gene2 = gene2;
        }

        public VDJCGene get(int index) {
            switch (index) {
                case 0:
                    return gene1;
                case 1:
                    return gene2;
                default:
                    throw new IndexOutOfBoundsException();
            }
        }
    }

    @Override
    public VDJCGenesGenerator create(RandomGenerator random, VDJCLibrary library) {
        final EnumeratedDistribution vDist = toDistribution(random, library, v, GeneType.Variable);
        final EnumeratedDistribution djDist = toDistribution2(random, library, dj,
                GeneType.Diversity, GeneType.Joining);
        final EnumeratedDistribution cDist = toDistribution(random, library, c, GeneType.Constant);
        return new VDJCGenesGenerator() {
            @Override
            public List genes(GeneType gt) {
                switch (gt) {
                    case Variable:
                        return IndependentVDJCGenesModel.genes(vDist);
                    case Diversity:
                        return DJDependentVDJCGenesModel.genes(djDist, 0);
                    case Joining:
                        return DJDependentVDJCGenesModel.genes(djDist, 1);
                    case Constant:
                        return IndependentVDJCGenesModel.genes(cDist);
                }
                throw new IllegalArgumentException();
            }

            @Override
            public VDJCGenes sample() {
                VDJCGene2 dj = djDist.sample();
                return new VDJCGenes(vDist.sample(), dj.gene1, dj.gene2, cDist.sample());
            }
        };
    }

    public static List genes(EnumeratedDistribution d, int index) {
        Set result = new HashSet<>();
        for (Pair p : d.getPmf())
            if (p.getFirst() != null)
                result.add(p.getFirst().get(index));
        return new ArrayList<>(result);
    }
}