AirTerminalDevice/model/_base/_base_device.py

import numpy as np
import pandas as pd
from sklearn.metrics import (
    r2_score,
    mean_absolute_error,
    mean_absolute_percentage_error
)
try:
    import plotnine as gg
except:
    pass

from ._base import BaseModel




class BaseDevice(BaseModel):
    
    def __init__(self) -> None:
        super().__init__()
    
    def predict(self,input_data:pd.DataFrame) -> dict:
        param_posterior = self.model_info['model_ATD']
        res = self.model(
            **{k:input_data.loc[:,v].values for k,v in self.model_input_data_columns.items()},
            engine       = 'numpy',
            components   = self.components,
            param        = param_posterior
        )
        return res
    
    def predict_system(self,input_data : pd.DataFrame) -> pd.DataFrame:
        pred_res = self.predict(input_data)
        system_output = {}
        for equp_name,output_info in pred_res.items():
            for output_name,output_value in output_info.items():
                system_output[f'{equp_name}_{output_name}'] = output_value
        system_output = pd.DataFrame(system_output)
        return system_output
    
    def get_TVP(self,posterior:dict,observed_data:pd.DataFrame):
        TVP_data = []
        for param_name in posterior.keys():
            if param_name.replace('_mu','') not in observed_data.columns:
                continue
            TVP_data.append(
                pd.DataFrame(
                    {
                        'idx'       : observed_data.index,
                        'param_name': param_name.replace('_mu',''),
                        'real'      : observed_data.loc[:,param_name.replace('_mu','')].values,
                        'pred'      : posterior[param_name]
                    }
                )
            )
        TVP_data = pd.concat(TVP_data,axis=0)
        return TVP_data
    
    def get_metric(self,TVP:pd.DataFrame):
        group_by_data = TVP.groupby(['param_name'])[['pred','real']]
        TVP_metric = (
            pd.concat(
                [
                    group_by_data.apply(lambda dt:r2_score(dt.real,dt.pred)),
                    group_by_data.apply(lambda dt:mean_absolute_error(dt.real,dt.pred)),
                    group_by_data.apply(lambda dt:mean_absolute_percentage_error(dt.real,dt.pred)),
                ],
                axis=1
            )
            .set_axis(['R2','MAE','MAPE'],axis=1)
            .sort_values(by='R2',ascending=True)
        )
        return TVP_metric
    
    def plot_TVP(self,TVP):
        plot = (
            TVP
            .pipe(gg.ggplot)
            + gg.aes(x='real',y='pred')
            + gg.geom_point()
            + gg.facet_wrap(facets='param_name',scales='free')
            + gg.geom_abline(intercept=0,slope=1,color='red')
            + gg.theme(figure_size=[10,10])
        )
        return plot
    
    def curve(
        self,
        input_data: pd.DataFrame,
        x         : str,
        y         : str,
        color     : str = None,
        grid_x    : str = None,
        grid_y    : str = None,
    ):
        
        if x not in input_data.columns: 
            raise Exception(f'{x} is not in input_data')
        
        product = [np.linspace(input_data.loc[:,x].min(),input_data.loc[:,x].max(),100)]
        names   = [x]
        
        if color is not None:
            if color not in input_data.columns: 
                raise Exception(f'{color} is not in input_data')
            product.append(np.quantile(input_data.loc[:,color],q=[0.25,0.5,0.75]))
            names.append(color)
        if grid_x is not None:
            if grid_x not in input_data.columns: 
                raise Exception(f'{grid_x} is not in input_data')
            product.append(np.quantile(input_data.loc[:,grid_x],q=[0.25,0.5,0.75]))
            names.append(grid_x)
        if grid_y is not None:
            if grid_y not in input_data.columns: 
                raise Exception(f'{grid_y} is not in input_data')
            product.append(np.quantile(input_data.loc[:,grid_y],q=[0.25,0.5,0.75]))
            names.append(grid_y)
        
        curve_input = (
            pd.MultiIndex.from_product(
                product,
                names=names
            )
            .to_frame(index=False)
        )
        curve_input_all = curve_input.copy(deep=True)
        
        for col in input_data.columns:
            if col not in curve_input_all:
                curve_input_all.loc[:,col] = input_data.loc[:,col].median()
        pred_data  = self.predict_system(curve_input_all)
        if y not in pred_data.columns:
            raise Exception(f'{y} is not in Prediction')
        
        curve_data = pd.concat([curve_input,pred_data],axis=1)
        plot = (
            curve_data
            .round(2)
            .pipe(gg.ggplot)
            + gg.aes(x=x,y=y)
            + gg.geom_line()
        )
        if color is not None:
            plot += gg.aes(color=f'factor({color})',group=color)
            plot += gg.labs(color=color)
        
        if grid_x is not None or grid_y is not None:
            plot += gg.facet_grid(rows=grid_x,cols=grid_y,labeller='label_both')
        
        return plot