org.jpmml.manager.PMMLManager Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2009 University of Tartu
 */
package org.jpmml.manager;

import java.io.*;
import java.util.*;

import org.dmg.pmml.*;

import com.google.common.collect.*;

/**
 * Naming conventions for getter methods:
 * 
 * getXXX() - Required schema elements. For example {@link #getDataDictionary()}
 * 
getOrCreateXXX() - Optional schema elements. When null then a new element instance is created. For example {@link #getOrCreateTransformationDictionary()}
 * 
 */
public class PMMLManager implements Serializable {

	private PMML pmml = null;


	public PMMLManager(){
		this(new PMML(new Header(), new DataDictionary(), "4.1"));
	}

	public PMMLManager(PMML pmml){
		setPmml(pmml);
	}

	public DataField getDataField(FieldName name){
		List dataFields = getDataDictionary().getDataFields();

		return find(dataFields, name);
	}

	public DataField addDataField(FieldName name, String displayName, OpType opType, DataType dataType){
		DataField dataField = new DataField(name, opType, dataType);
		dataField.setDisplayName(displayName);

		List dataFields = getDataDictionary().getDataFields();
		dataFields.add(dataField);

		return dataField;
	}

	public DerivedField resolveField(FieldName name){
		TransformationDictionary transformationDictionary = getOrCreateTransformationDictionary();

		List derivedFields = transformationDictionary.getDerivedFields();

		return find(derivedFields, name);
	}

	public DefineFunction resolveFunction(String name){
		TransformationDictionary transformationDictionary = getOrCreateTransformationDictionary();

		List defineFunctions = transformationDictionary.getDefineFunctions();
		for(DefineFunction defineFunction : defineFunctions){

			if((defineFunction.getName()).equals(name)){
				return defineFunction;
			}
		}

		return null;
	}

	public PMML getPmml(){
		return this.pmml;
	}

	private void setPmml(PMML pmml){
		this.pmml = pmml;
	}

	public Header getHeader(){
		return getPmml().getHeader();
	}

	public DataDictionary getDataDictionary(){
		return getPmml().getDataDictionary();
	}

	public TransformationDictionary getOrCreateTransformationDictionary(){
		PMML pmml = getPmml();

		TransformationDictionary transformationDictionary = pmml.getTransformationDictionary();
		if(transformationDictionary == null){
			transformationDictionary = new TransformationDictionary();

			pmml.setTransformationDictionary(transformationDictionary);
		}

		return transformationDictionary;
	}

	public List getModels(){
		return getPmml().getContent();
	}

	/**
	 * @param modelName The name of the Model to be selected. If null, the first model is selected.
	 *
	 * @see Model#getModelName()
	 */
	public Model getModel(String modelName){
		List models = getModels();

		if(modelName != null){

			for(Model model : models){

				if(modelName.equals(model.getModelName())){
					return model;
				}
			}

			return null;
		} // End if

		if(models.size() > 0){
			return models.get(0);
		}

		return null;
	}

	public ModelManager getModelManager(String modelName){
		return getModelManager(modelName, ModelManagerFactory.getInstance());
	}

	public ModelManager getModelManager(String modelName, ModelManagerFactory modelManagerFactory){
		Model model = getModel(modelName);

		return modelManagerFactory.getModelManager(getPmml(), model);
	}

	@SuppressWarnings (
		value = {"unchecked"}
	)
	static
	public  E find(List objects, Class clazz){

		for(PMMLObject object : objects){

			if(object.getClass().equals(clazz)){
				return (E)object;
			}
		}

		return null;
	}

	@SuppressWarnings (
		value = {"unchecked"}
	)
	static
	public  List findAll(List objects, Class clazz){
		List result = Lists.newArrayList();

		for(PMMLObject object : objects){

			if(object.getClass().equals(clazz)){
				result.add((E)object);
			}
		}

		return result;
	}

	static
	public  E find(Collection objects, FieldName name){

		for(E object : objects){

			if((object.getName()).equals(name)){
				return object;
			}
		}

		return null;
	}
}