com.expleague.ml.models.multiclass.HierarchicalModel Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of jmll

Various ML methods implemented by myself and my students

There is a newer version: 1.4.9

package com.expleague.ml.models.multiclass;

import com.expleague.commons.math.vectors.Mx;
import com.expleague.commons.math.vectors.Vec;
import gnu.trove.list.TIntList;
import gnu.trove.map.TIntObjectMap;
import gnu.trove.map.hash.TIntObjectHashMap;
import org.jetbrains.annotations.Nullable;

import java.util.Arrays;

/**
 * User: qdeee
 * Date: 06.02.14
 */
public class HierarchicalModel extends MCModel.Stub {
  protected final MCModel basedOn;
  protected final TIntList classLabels;
  protected final TIntObjectMap label2childModel;

  public HierarchicalModel(final MCModel basedOn, final TIntList classLabels) {
    this.basedOn = basedOn;
    this.classLabels = classLabels;
    this.label2childModel = new TIntObjectHashMap<>(classLabels.size());
  }

  public void addChild(final HierarchicalModel child, final int label) {
    label2childModel.put(label, child);
  }

  @Nullable
  public HierarchicalModel getChild(final int label) {
    return label2childModel.get(label);
  }

  @Override
  public int countClasses() {
    return basedOn.countClasses();
  }

  @Override
  public Vec probs(final Vec x) {
    return basedOn.probs(x);
  }

  @Override
  public int bestClass(final Vec x) {
    final int c = basedOn.bestClass(x);
    final int label = classLabels.get(c);
    return label2childModel.containsKey(label)? label2childModel.get(label).bestClass(x) : label;
  }

  @Override
  public Vec bestClassAll(final Mx x) {
    return transAll(x);
  }

  @Override
  public int dim() {
    return basedOn.dim();
  }

  @Override
  public String toString() {
    return "splits " + classLabels.toString() + " classes, has child models for " +
        Arrays.toString(label2childModel.keys());
  }
}