Source code for schrodinger.application.matsci.mlearn.sklearn_json.decomposition

"""
Module for serializing and de-serializing decomposition models

Copyright Schrodinger, LLC. All rights reserved.
"""

import numpy as np
from sklearn.decomposition import PCA

from . import common


def serialize_pca(model):
    serialized_model = {
        'meta': 'pca',
        'components_': model.components_.tolist(),
        'mean_': model.mean_.tolist(),
        'explained_variance_': model.explained_variance_.tolist(),
        'params': model.get_params()
    }
    common.serialize_params_random_state(serialized_model['params'])

    return serialized_model


def deserialize_pca(model_dict):
    common.deserialize_params_random_state(model_dict['params'])

    model = PCA(**model_dict['params'])
    model.components_ = np.array(model_dict['components_'])
    model.mean_ = np.array(model_dict['mean_'])
    model.explained_variance_ = np.array(model_dict['explained_variance_'])

    return model