Examples with IComponentModel com.amazon.randomcutforest.IComponentModel used on opensource projects

Example 1 with IComponentModel

use of com.amazon.randomcutforest.IComponentModel in project random-cut-forest-by-aws by aws.

the class RandomCutForestMapper method toState.

/**
 * Create a {@link RandomCutForestState} object representing the state of the
 * given forest. If the forest is compact and the {@code saveTreeState} flag is
 * set to true, then structure of the trees in the forest will be included in
 * the state object. If the flag is set to false, then the state object will
 * only contain the sampler data for each tree. If the
 * {@code saveExecutorContext} is true, then the executor context will be
 * included in the state object.
 *
 * @param forest A Random Cut Forest whose state we want to capture.
 * @return a {@link RandomCutForestState} object representing the state of the
 *         given forest.
 * @throws IllegalArgumentException if the {@code saveTreeState} flag is true
 *                                  and the forest is not compact.
 */
@Override
public RandomCutForestState toState(RandomCutForest forest) {
    if (saveTreeStateEnabled) {
        checkArgument(forest.isCompact(), "tree state cannot be saved for noncompact forests");
    }
    RandomCutForestState state = new RandomCutForestState();
    state.setNumberOfTrees(forest.getNumberOfTrees());
    state.setDimensions(forest.getDimensions());
    state.setTimeDecay(forest.getTimeDecay());
    state.setSampleSize(forest.getSampleSize());
    state.setShingleSize(forest.getShingleSize());
    state.setCenterOfMassEnabled(forest.isCenterOfMassEnabled());
    state.setOutputAfter(forest.getOutputAfter());
    state.setStoreSequenceIndexesEnabled(forest.isStoreSequenceIndexesEnabled());
    state.setTotalUpdates(forest.getTotalUpdates());
    state.setCompact(forest.isCompact());
    state.setInternalShinglingEnabled(forest.isInternalShinglingEnabled());
    state.setBoundingBoxCacheFraction(forest.getBoundingBoxCacheFraction());
    state.setSaveSamplerStateEnabled(saveSamplerStateEnabled);
    state.setSaveTreeStateEnabled(saveTreeStateEnabled);
    state.setSaveCoordinatorStateEnabled(saveCoordinatorStateEnabled);
    state.setPrecision(forest.getPrecision().name());
    state.setCompressed(compressionEnabled);
    state.setPartialTreeState(partialTreeStateEnabled);
    if (saveExecutorContextEnabled) {
        ExecutionContext executionContext = new ExecutionContext();
        executionContext.setParallelExecutionEnabled(forest.isParallelExecutionEnabled());
        executionContext.setThreadPoolSize(forest.getThreadPoolSize());
        state.setExecutionContext(executionContext);
    }
    if (saveCoordinatorStateEnabled) {
        PointStoreCoordinator<?> pointStoreCoordinator = (PointStoreCoordinator<?>) forest.getUpdateCoordinator();
        PointStoreMapper mapper = new PointStoreMapper();
        mapper.setCompressionEnabled(compressionEnabled);
        mapper.setNumberOfTrees(forest.getNumberOfTrees());
        PointStoreState pointStoreState = mapper.toState((PointStore) pointStoreCoordinator.getStore());
        state.setPointStoreState(pointStoreState);
    }
    List<CompactSamplerState> samplerStates = null;
    if (saveSamplerStateEnabled) {
        samplerStates = new ArrayList<>();
    }
    List<ITree<Integer, ?>> trees = null;
    if (saveTreeStateEnabled) {
        trees = new ArrayList<>();
    }
    CompactSamplerMapper samplerMapper = new CompactSamplerMapper();
    samplerMapper.setCompressionEnabled(compressionEnabled);
    for (IComponentModel<?, ?> component : forest.getComponents()) {
        SamplerPlusTree<Integer, ?> samplerPlusTree = (SamplerPlusTree<Integer, ?>) component;
        CompactSampler sampler = (CompactSampler) samplerPlusTree.getSampler();
        if (samplerStates != null) {
            samplerStates.add(samplerMapper.toState(sampler));
        }
        if (trees != null) {
            trees.add(samplerPlusTree.getTree());
        }
    }
    state.setCompactSamplerStates(samplerStates);
    if (trees != null) {
        RandomCutTreeMapper treeMapper = new RandomCutTreeMapper();
        List<CompactRandomCutTreeState> treeStates = trees.stream().map(t -> treeMapper.toState((RandomCutTree) t)).collect(Collectors.toList());
        state.setCompactRandomCutTreeStates(treeStates);
    }
    return state;
}