ebm 0.99.3__py3-none-any.whl

ebm/model/file_handler.py

@@ -0,0 +1,388 @@
+import os
+import pathlib
+import shutil
+import typing
+
+import pandas as pd
+from loguru import logger
+from pandera.errors import SchemaErrors, SchemaError
+
+import ebm.validators as validators
+from ebm.model.defaults import default_calibrate_heating_rv, default_calibrate_energy_consumption
+
+
+class FileHandler:
+    """
+    Handles file operations.
+    """
+
+    # Filenames
+    BUILDING_CONDITIONS = 'building_conditions.csv'
+    BUILDING_CODE_PARAMS = 'building_code_parameters.csv'
+    S_CURVE = 's_curve.csv'
+    POPULATION_FORECAST = 'population_forecast.csv'
+    NEW_BUILDINGS_RESIDENTIAL = 'new_buildings_residential.csv'
+    AREA_NEW_RESIDENTIAL_BUILDINGS = 'area_new_residential_buildings.csv'
+    AREA = 'area.csv'
+    BEHAVIOUR_FACTOR = 'energy_need_behaviour_factor.csv'
+    ENERGY_NEED_ORIGINAL_CONDITION = 'energy_need_original_condition.csv'
+    IMPROVEMENT_BUILDING_UPGRADE = 'improvement_building_upgrade.csv'
+    ENERGY_NEED_YEARLY_IMPROVEMENTS = 'energy_need_improvements.csv'
+    HOLIDAY_HOME_STOCK = 'holiday_home_stock.csv'
+    HOLIDAY_HOME_ENERGY_CONSUMPTION = 'holiday_home_energy_consumption.csv'
+    AREA_PER_PERSON = 'area_per_person.csv'
+    HEATING_SYSTEM_INITIAL_SHARES = 'heating_system_initial_shares.csv'
+    HEATING_SYSTEM_EFFICIENCIES = 'heating_system_efficiencies.csv'
+    HEATING_SYSTEM_FORECAST = 'heating_system_forecast.csv'
+    CALIBRATE_ENERGY_REQUIREMENT = 'calibrate_heating_rv.xlsx'
+    CALIBRATE_ENERGY_CONSUMPTION = 'calibrate_energy_consumption.xlsx'
+
+    input_directory: pathlib.Path
+
+    def __init__(self, directory: typing.Union[str, pathlib.Path, None] = None):
+        """
+        Constructor for FileHandler Object. Sets FileHandler.input_directory.
+
+        Parameters
+        ----------
+        directory : pathlib.Path | None | (str)
+            When directory is None the constructor will attempt to read directory location from
+                environment variable EBM_INPUT_DIRECTORY
+        """
+        if directory is None:
+            # Use 'input' as fall back when EBM_INPUT_DIRECTORY is not set in environment.
+            directory = os.environ.get('EBM_INPUT_DIRECTORY', 'input')
+
+        self.input_directory = directory if isinstance(directory, pathlib.Path) else pathlib.Path(directory)
+        self.files_to_check = [self.BUILDING_CODE_PARAMS, self.S_CURVE, self.POPULATION_FORECAST,
+                               self.NEW_BUILDINGS_RESIDENTIAL, self.AREA_NEW_RESIDENTIAL_BUILDINGS,
+                               self.AREA, self.BEHAVIOUR_FACTOR, self.ENERGY_NEED_ORIGINAL_CONDITION,
+                               self.IMPROVEMENT_BUILDING_UPGRADE, self.ENERGY_NEED_YEARLY_IMPROVEMENTS,
+                               self.HOLIDAY_HOME_ENERGY_CONSUMPTION, self.HOLIDAY_HOME_STOCK,
+                               self.AREA_PER_PERSON, self.HEATING_SYSTEM_INITIAL_SHARES, self.HEATING_SYSTEM_EFFICIENCIES, self.HEATING_SYSTEM_FORECAST]
+
+    def __repr__(self):
+        return f'FileHandler(input_directory="{self.input_directory}")'
+
+    def __str__(self):
+        return repr(self)
+
+    @staticmethod
+    def default_data_directory() -> pathlib.Path:
+        """
+        Returns the path for ebm default data. The function is used when content is needed for a new input directory.
+        Not to be confused with FileHandler.input_directory.
+
+        Returns
+        -------
+        pathlib.Path
+
+        See Also
+        --------
+        create_missing_input_files
+        """
+        return pathlib.Path(__file__).parent.parent / 'data'
+
+    def get_file(self, file_name: str) -> pd.DataFrame:
+        """
+        Finds and returns a file by searching in the folder defined by self.input_folder.
+
+        Parameters:
+        - file_name (str): Name of the file to retrieve.
+
+        Returns:
+        - file_df (pd.DataFrame): DataFrame containing file data.
+        """
+        logger.debug(f'get_file {file_name}')
+        file_path: pathlib.Path = pathlib.Path(self.input_directory) / file_name
+        logger.debug(f'{file_path=}')
+
+        try:
+            if file_path.suffix == '.xlsx':
+                file_df = pd.read_excel(file_path)
+            elif file_path.suffix == '.csv':
+                file_df = pd.read_csv(file_path)
+            else:
+                msg = f'{file_name} is not of type xlsx or csv'
+                logger.error(msg)
+                raise ValueError(msg)
+            return file_df
+        except FileNotFoundError as ex:
+            logger.exception(ex)
+            logger.debug(f'Current directory is {os.getcwd()}')
+            logger.error(f'Unable to open {file_path}. File not found.')
+            raise
+        except PermissionError as ex:
+            logger.exception(ex)
+            logger.error(f'Unable to open {file_path}. Permission denied.')
+            raise
+        except IOError as ex:
+            logger.exception(ex)
+            logger.error(f'Unable to open {file_path}. Unable to read file.')
+            raise
+
+    def get_building_code(self) -> pd.DataFrame:
+        """
+        Get TEK parameters DataFrame.
+
+        Returns:
+        - building_code_params (pd.DataFrame): DataFrame containing TEK parameters.
+        """
+        building_code_params = self.get_file(self.BUILDING_CODE_PARAMS)
+        return building_code_params
+    
+    def get_s_curve(self) -> pd.DataFrame:
+        """
+        Get S-curve parameters DataFrame.
+
+        Returns:
+        - scurve_params (pd.DataFrame): DataFrame containing S-curve parameters.
+        """
+        scurve_params = self.get_file(self.S_CURVE)
+        return scurve_params
+
+    def get_construction_population(self) -> pd.DataFrame:
+        """
+        Get population and household size DataFrame from a file.
+
+        Returns:
+        - construction_population (pd.DataFrame): Dataframe containing population numbers
+          year population household_size
+        """
+        return self.get_file(self.POPULATION_FORECAST)
+
+    def get_population(self) -> pd.DataFrame:
+        """
+        Loads population data from population.csv as float64
+
+        Should probably be merged with get_construction_population
+
+        Returns population : pd.DataFrame
+            dataframe with population
+        -------
+
+        """
+        file_path = self.input_directory / self.POPULATION_FORECAST
+        logger.debug(f'{file_path=}')
+        return pd.read_csv(file_path, dtype={"household_size": "float64"})
+
+    def get_construction_building_category_share(self) -> pd.DataFrame:
+        """
+        Get building category share by year DataFrame from a file.
+
+        The number can be used in conjunction with number of households to calculate total number
+        of buildings of category house and apartment block
+
+        Returns:
+        - construction_population (pd.DataFrame): Dataframe containing population numbers
+          "year", "Andel nye småhus", "Andel nye leiligheter", "Areal nye småhus", "Areal nye leiligheter"
+        """
+        return self.get_file(self.NEW_BUILDINGS_RESIDENTIAL)
+
+    def get_building_category_area(self) -> pd.DataFrame:
+        """
+        Get population and household size DataFrame from a file.
+
+        Returns:
+        - construction_population (pd.DataFrame): Dataframe containing population numbers
+          "area","type of building","2010","2011"
+        """
+        file_path = self.input_directory / self.AREA_NEW_RESIDENTIAL_BUILDINGS
+        logger.debug(f'{file_path=}')
+        return pd.read_csv(file_path,
+                           index_col=0, header=0)
+
+    def get_area_parameters(self) -> pd.DataFrame:
+        """
+        Get dataframe with area parameters.
+
+        Returns:
+        - area_parameters (pd.DataFrame): Dataframe containing total area (m^2) per
+                                          building category and TEK.
+        """
+        return self.get_file(self.AREA)
+    
+    def get_energy_req_original_condition(self) -> pd.DataFrame:
+        """
+        Get dataframe with energy requirement (kWh/m^2) for floor area in original condition.
+
+        Returns
+        -------
+        pd.DataFrame
+            Dataframe containing energy requirement (kWh/m^2) for floor area in original condition,
+            per building category and purpose.
+        """
+        return self.get_file(self.ENERGY_NEED_ORIGINAL_CONDITION)
+    
+    def get_energy_req_reduction_per_condition(self) -> pd.DataFrame:
+        """
+        Get dataframe with shares for reducing the energy requirement of the different building conditions.
+
+        Returns
+        -------
+        pd.DataFrame
+            Dataframe containing energy requirement reduction shares for the different building conditions, 
+            per building category, TEK and purpose.
+        """
+        return self.get_file(self.IMPROVEMENT_BUILDING_UPGRADE)
+    
+    def get_energy_need_yearly_improvements(self) -> pd.DataFrame:
+        """
+        Get dataframe with yearly efficiency rates for energy requirement improvements.
+
+        Returns
+        -------
+        pd.DataFrame
+            Dataframe containing yearly efficiency rates (%) for energy requirement improvements,
+            per building category, tek and purpose.
+        """
+        return self.get_file(self.ENERGY_NEED_YEARLY_IMPROVEMENTS)
+
+    def get_holiday_home_energy_consumption(self) -> pd.DataFrame:
+        return self.get_file(self.HOLIDAY_HOME_ENERGY_CONSUMPTION)
+
+    def get_holiday_home_by_year(self) -> pd.DataFrame:
+        return self.get_file(self.HOLIDAY_HOME_STOCK)
+
+    def get_area_per_person(self):
+        return self.get_file(self.AREA_PER_PERSON)
+
+    def get_calibrate_heating_rv(self) -> pd.DataFrame:
+        calibrate_heating_rv = self.input_directory / self.CALIBRATE_ENERGY_REQUIREMENT
+        if calibrate_heating_rv.is_file():
+            return self.get_file(calibrate_heating_rv.name)
+        return default_calibrate_heating_rv()
+
+    def get_calibrate_heating_systems(self) -> pd.DataFrame:
+        calibrate_energy_consumption = self.input_directory / self.CALIBRATE_ENERGY_CONSUMPTION
+        if calibrate_energy_consumption.is_file():
+            return self.get_file(calibrate_energy_consumption.name)
+        return default_calibrate_energy_consumption()
+
+    def get_heating_systems_shares_start_year(self) -> pd.DataFrame:
+        """
+        """
+        return self.get_file(self.HEATING_SYSTEM_INITIAL_SHARES)
+    
+    def get_heating_system_efficiencies(self) -> pd.DataFrame:
+        """Load heating_system_efficiencies.csv from file into a dataframe
+        
+        Returns
+        -------
+        heating_system_efficiencies : pd.DataFrame
+            pandas DataFrame with heating system efficiencies
+        """
+
+        return self.get_file(self.HEATING_SYSTEM_EFFICIENCIES)
+
+    def get_heating_system_forecast(self) -> pd.DataFrame:
+        """
+        """
+        return self.get_file(self.HEATING_SYSTEM_FORECAST)
+
+    def _check_is_file(self, filename: str) -> bool:
+        """
+        Check if the filename is a file in self.input_folder
+
+        Parameters
+        ----------
+        filename : str
+
+        Returns
+        -------
+        file_exists : bool
+        """
+        return (pathlib.Path(self.input_directory) / filename).is_file()
+
+    def check_for_missing_files(self) -> typing.List[str]:
+        """
+        Returns a list of required files that are not present in self.input_folder
+
+        Returns
+        -------
+        missing_files : List[str]
+
+        Raises
+        ------
+        FileNotFoundError
+            If FileHandler::input_directory not found
+        NotADirectoryError
+            If FileHandler::input_directory is not a directory
+        """
+        if not self.input_directory.exists():
+            msg=f'{self.input_directory.absolute()} not found'
+            logger.error(msg)
+            raise FileNotFoundError(f'Input Directory Not Found')
+        if not self.input_directory.is_dir():
+            raise NotADirectoryError(f'{self.input_directory} is not a directory')
+
+        missing_files = [file for file in self.files_to_check if not self._check_is_file(file)]
+        if missing_files:
+            plural = 's' if len(missing_files) != 1 else ''
+            msg = f'{len(missing_files)} required file{plural} missing from {self.input_directory}'
+            logger.error(msg)
+            for f in missing_files:
+                logger.error(f'Could not find {f}')
+        return missing_files
+
+
+    def create_missing_input_files(self, source_directory: (pathlib.Path | None)=None) -> None:
+        """
+        Creates any input files missing in self.input_directory. When source is omitted FileHandler
+
+        Parameters
+        ----------
+        source_directory : pathlib.Path, optional
+                           Optional directory for sourcing files to copy.
+
+        Returns
+        -------
+        None
+
+        See Also
+        --------
+        default_data_directory : default source for data files
+        """
+        source = FileHandler.default_data_directory() if not source_directory else source_directory
+
+        if not source.is_dir():
+            raise NotADirectoryError(f'{self.input_directory} is not a directory')
+        if not self.input_directory.is_dir():
+            logger.info(f'Creating directory {self.input_directory}')
+            self.input_directory.mkdir()
+        for file in self.files_to_check:
+            logger.debug(f'Create input file {file}')
+            self.create_input_file(file, source_directory=source)
+
+    def create_input_file(self, file, source_directory=None):
+        source_directory = FileHandler.default_data_directory() if not source_directory else source_directory
+
+        source_file = source_directory / file
+        target_file = self.input_directory / file
+        if target_file.is_file():
+            logger.debug(f'Skipping existing file {target_file}')
+        elif not source_file.is_file():
+            logger.error(f'Source file {source_file} does not exist!')
+        else:
+            shutil.copy(source_file, target_file)
+            logger.info( f'Creating missing file  {target_file}')
+
+    def validate_input_files(self):
+        """
+        Validates the input files for correct formatting and content using the validators module
+
+        Raises
+        ------
+        pa.errors.SchemaErrors
+            If any invalid data for formatting is found when validating files. The validation is lazy, meaning
+            multiple errors may be listed in the exception.
+        """
+        for file_to_validate in self.files_to_check:
+            df = self.get_file(file_to_validate)
+            validator = getattr(validators, file_to_validate[:-4].lower())
+
+            try:
+                validator.validate(df, lazy=True)
+            except (SchemaErrors, SchemaError):
+                logger.error(f'Got error while validating {file_to_validate}')
+                raise

ebm/model/filter_scurve_params.py

@@ -0,0 +1,83 @@
+import typing
+
+import pandas as pd
+
+from .building_category import BuildingCategory
+from .building_condition import BuildingCondition
+from .data_classes import ScurveParameters
+
+
+class FilterScurveParams():
+    """
+    A utility class for filtering S-curve parameters based on building category and condition.
+
+    This class provides a static method to filter a DataFrame containing S-curve parameters, 
+    extracting the relevant data for a specific building category and set of conditions.
+    """
+
+    COL_BUILDING_CATEGORY = 'building_category'
+    COL_BUILDING_CONDITION = 'condition'
+    COL_EARLIEST_AGE = 'earliest_age_for_measure'
+    COL_AVERAGE_AGE = 'average_age_for_measure'
+    COL_LAST_AGE = 'last_age_for_measure'
+    COL_RUSH_YEARS = 'rush_period_years'
+    COL_RUSH_SHARE = 'rush_share'
+    COL_NEVER_SHARE = 'never_share'
+
+    @staticmethod
+    def filter(building_category: BuildingCategory,
+               scurve_condition_list: typing.List[str],
+               scurve_params: pd.DataFrame) -> typing.Dict[str, ScurveParameters]:
+        """
+        Filters S-curve parameters by building category and condition.
+
+        This method filters a DataFrame containing S-curve parameters to extract data specific to 
+        the provided building category and conditions listed in `scurve_condition_list`. The filtered 
+        data is then converted into a dictionary of `ScurveParameters` dataclass instances, each 
+        representing the S-curve parameters for a particular condition.
+
+        Parameters:
+        - building_category (BuildingCategory): The building category for which the S-curve parameters are being filtered.
+        - scurve_condition_list (List[str]): A list of conditions (as strings) for which the S-curve parameters are needed.
+        - scurve_params (pd.DataFrame): DataFrame containing the S-curve parameters, with columns for building category, condition, and various age-related metrics.
+
+        Returns:
+        - filtered_scurve_params (Dict[str, ScurveParameters]): A dictionary where the keys are conditions (str) and the values 
+                                                                are `ScurveParameters` dataclass instances containing the 
+                                                                corresponding S-curve parameters for each condition.
+
+        Raises:
+        - KeyError: If the provided building category is not found in the S-curve parameters DataFrame.
+        """
+        filtered_scurve_params = {}
+
+        for condition in scurve_condition_list:
+            if not scurve_params.building_category.str.contains(building_category).any():
+                msg = 'Unknown building_category "{}" encountered when setting up scurve parameters'.format(building_category)
+                raise KeyError(msg)
+
+            # Filter dataframe on building category and condition
+            scurve_params_filtered = scurve_params[(scurve_params[FilterScurveParams.COL_BUILDING_CATEGORY] == building_category) &
+                                                   (scurve_params[FilterScurveParams.COL_BUILDING_CONDITION] == condition)]
+
+            # Assuming there is only one row in the filtered DataFrame
+            scurve_params_row = scurve_params_filtered.iloc[0]
+
+            # Convert the single row to a dictionary
+            scurve_params_dict = scurve_params_row.to_dict()
+            
+            # Map the dictionary values to the dataclass attributes
+            scurve_parameters = ScurveParameters(
+                building_category=scurve_params_dict[FilterScurveParams.COL_BUILDING_CATEGORY],
+                condition=scurve_params_dict[FilterScurveParams.COL_BUILDING_CONDITION],
+                earliest_age=scurve_params_dict[FilterScurveParams.COL_EARLIEST_AGE],
+                average_age=scurve_params_dict[FilterScurveParams.COL_AVERAGE_AGE], 
+                rush_years=scurve_params_dict[FilterScurveParams.COL_RUSH_YEARS], 
+                last_age=scurve_params_dict[FilterScurveParams.COL_LAST_AGE],
+                rush_share=scurve_params_dict[FilterScurveParams.COL_RUSH_SHARE],
+                never_share=scurve_params_dict[FilterScurveParams.COL_NEVER_SHARE],
+            )
+
+            filtered_scurve_params[condition] = scurve_parameters
+
+        return filtered_scurve_params 

ebm/model/filter_tek.py

@@ -0,0 +1,152 @@
+import typing
+
+import pandas as pd
+from loguru import logger
+
+from ebm.model.building_category import BuildingCategory
+from ebm.model.data_classes import TEKParameters
+
+
+class FilterTek:
+    """
+    Utility class for filtering TEK lists and parameters. 
+    """
+
+    CATEGORY_APARTMENT = 'apartment_block'
+    CATEGORY_HOUSE = 'house'
+    COMMERCIAL_BUILDING = 'COM'
+    RESIDENTIAL_BUILDING = 'RES'
+    PRE_TEK49_APARTMENT = 'PRE_TEK49_RES_1950'
+    PRE_TEK49_HOUSE = 'PRE_TEK49_RES_1940'
+
+    @staticmethod
+    def get_filtered_list(building_category: BuildingCategory, building_code_list: typing.List[str]) -> typing.List[str]:
+        """
+        Filters the provided TEK list based on the building category.
+
+        Parameters:
+        - building_category (BuildingCategory): The category of the building.
+        - building_code_list (List[str]): List of TEK strings to be filtered.
+
+        Returns:
+        - filtered_building_code_list (List[str]): Filtered list of TEK strings.
+        """
+        residential_building_list = [FilterTek.CATEGORY_APARTMENT, FilterTek.CATEGORY_HOUSE]
+        
+        if building_category in residential_building_list:
+            # Filter out all TEKs associated with commercial buildings
+            filtered_building_code_list = [tek for tek in building_code_list if FilterTek.COMMERCIAL_BUILDING not in tek]
+
+            # Further filtering based on the specific residential building category
+            if building_category == FilterTek.CATEGORY_APARTMENT:
+                filtered_building_code_list = [tek for tek in filtered_building_code_list if tek != FilterTek.PRE_TEK49_HOUSE]
+            elif building_category == FilterTek.CATEGORY_HOUSE:
+                filtered_building_code_list = [tek for tek in filtered_building_code_list if tek != FilterTek.PRE_TEK49_APARTMENT]
+
+        else:
+            # Filter out all TEKs associated with residential buildings
+            filtered_building_code_list = [tek for tek in building_code_list if FilterTek.RESIDENTIAL_BUILDING not in tek]
+
+        return filtered_building_code_list
+    
+    # This method is only needed if building_code_params are initialized in the Buildings class
+    @staticmethod
+    def get_filtered_params(building_code_list: typing.List[str],
+                            building_code_params: typing.Dict[str, TEKParameters]) -> typing.Dict[str, TEKParameters]:
+        """
+        Filters the TEK parameters to include only those relevant to the provided TEK list.
+
+        This method takes a dictionary of TEK parameters and filters it to include only 
+        the parameters for TEKs that are present in the `building_code_list`. This ensures that 
+        only the relevant TEK parameters are retained for use in subsequent calculations.
+
+        Parameters:
+        - building_code_list (List[str]): A list of TEK identifiers to filter by.
+        - building_code_params (Dict[str, TEKParameters]): A dictionary where the keys are TEK identifiers 
+                                                    and the values are TEKParameters objects containing
+                                                    the parameters for each TEK.
+
+        Returns:
+        - filtered_building_code_params (Dict[str, TEKParameters]): A dictionary containing only the TEK parameters
+                                                        for the TEKs present in the `building_code_list`.
+        """
+        filtered_building_code_params = {}
+        for tek in building_code_list:
+            filtered_building_code_params[tek] = building_code_params[tek]    
+
+        return filtered_building_code_params
+
+    @staticmethod
+    def merge_building_code(df: pd.DataFrame,
+                  new_building_code_name: str,
+                  old_building_code_names: typing.List[str],
+                  aggregates: typing.Dict[str, str] = None) -> pd.DataFrame:
+        """
+        Merge rows in a DataFrame based on specified 'tek' names and aggregate their values.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            The input DataFrame with a MultiIndex.
+        new_building_code_name : str
+            The new 'tek' name to assign to the merged rows.
+        old_building_code_names : typing.List[str]
+            A list of 'tek' names to be merged.
+        aggregates : typing.Dict[str, str], optional
+            A dictionary specifying the aggregation functions for each column.
+            If not provided, default aggregations will be used:
+            {'tek': 'max', 'm2': 'first', 'kwh_m2': 'mean', 'energy_requirement': 'sum'}.
+
+        Returns
+        -------
+        pd.DataFrame
+            The DataFrame with the specified 'tek' rows merged and aggregated.
+        """
+        if not isinstance(df, pd.DataFrame):
+            raise ValueError("`df` should be a pandas DataFrame.")
+        if not isinstance(old_building_code_names, list):
+            raise ValueError("`old_building_code_names` should be a list of strings.")
+
+        # Apply default aggregates if the parameter is empty
+        aggregates = aggregates or {'building_code': 'max', 'm2': 'first', 'kwh_m2': 'mean', 'energy_requirement': 'sum'}
+        building_code_values = [tek for tek in old_building_code_names if tek in df.index.get_level_values('building_code')]
+
+        if not building_code_values:
+            return df
+
+        level_values = df.index.get_level_values('building_category')
+        building_categories = [bc for bc in BuildingCategory if bc.is_residential() and bc in level_values]
+        if not building_categories:
+            return df
+
+        residential = df.loc[
+                (building_categories, slice(None), slice(None), slice(None), slice(None))].reset_index()
+
+        building_code_to_merge = residential[residential.building_code.isin(building_code_values)]
+        agg_building_code = building_code_to_merge.groupby(by=['building_category',
+                                           'building_condition',
+                                           'purpose',
+                                           'year']).agg(aggregates)
+        agg_building_code = agg_building_code.reset_index()
+
+        agg_building_code['building_code'] = new_building_code_name
+        rows_to_remove = df.loc[(slice(None), building_code_values, slice(None), slice(None), slice(None))].index
+        df = df.drop(rows_to_remove)
+        df = pd.concat([df, agg_building_code.set_index(['building_category', 'building_code', 'building_condition',  'year', 'purpose'])])
+        df = df.sort_index()
+
+        return df
+
+    @staticmethod
+    def remove_building_code_suffix(df: pd.DataFrame, suffix) -> pd.DataFrame:
+        # Convert MultiIndex to DataFrame
+        index_df = df.index.to_frame(index=False)
+
+        key_name = 'tek' if 'tek' in index_df.keys() else 'building_code'
+        # Remove '_RES' from 'tek' values
+        index_df[key_name] = index_df[key_name].str.replace(suffix, '')
+
+        # Set the modified DataFrame back to MultiIndex
+        df.index = pd.MultiIndex.from_frame(index_df)
+
+        return df

ebm/model/heat_pump.py

@@ -0,0 +1,53 @@
+import pandas as pd
+
+from ebm.energy_consumption import HP_ENERGY_SOURCE, HEAT_PUMP
+
+
+def air_source_heat_pump(heating_systems_parameters: pd.DataFrame):
+    df = heating_systems_parameters.copy()
+    el_slice = df[df['heating_system'] == 'HP'].index
+    df.loc[el_slice, 'pump_factor'] = df.loc[el_slice, 'load_share'] * df.loc[el_slice, 'heating_system_share']
+    df.loc[el_slice, HP_ENERGY_SOURCE] = 'Luft/luft'
+    df.loc[el_slice, 'purpose'] = 'heating_rv'
+
+    return  df.query('heating_system=="HP"')
+
+
+def district_heating_heat_pump(heating_systems_parameters: pd.DataFrame):
+    df = heating_systems_parameters.copy()
+    vann_slice = df[df['heating_system'] == 'HP Central heating'].index
+    df.loc[vann_slice, 'pump_factor'] = df.loc[vann_slice, 'load_share'] * df.loc[vann_slice, 'heating_system_share']
+    df.loc[vann_slice, HP_ENERGY_SOURCE] = 'Vannbåren varme'
+    df.loc[vann_slice, 'purpose'] = 'heating_rv,heating_dhw'
+    df = df.assign(**{'purpose': df['purpose'].str.split(',')}).explode('purpose')
+
+    return df.query('heating_system=="HP Central heating"')
+
+
+def heat_pump_production(energy_need, air_air, district_heating):
+    df_en = energy_need.copy()
+    df_hp = pd.concat([air_air, district_heating])
+
+    df = pd.merge(left=df_en,
+                  left_on=['building_category', 'building_code', 'purpose', 'year'],
+                  right=df_hp,
+                  right_on=['building_category', 'building_code', 'purpose', 'year'])
+
+
+    df[HEAT_PUMP] = df.energy_requirement * df.pump_factor
+
+    return df
+
+
+def heat_prod_hp(production: pd.DataFrame, group_by:list|None=None) -> pd.DataFrame:
+    grouping = ['building_group', 'year'] if not group_by else group_by
+    production.loc[production['building_category'].isin(['house', 'apartment_block']), 'building_group'] = 'Bolig'
+    production.loc[production['building_group'] != 'Bolig', 'building_group'] = 'Yrkesbygg'
+    return production.groupby(by=grouping+['hp_source']).agg({'RV_HP': 'sum'}) / 1_000_000
+
+
+def heat_prod_hp_wide(production: pd.DataFrame) -> pd.DataFrame:
+    df = heat_prod_hp(production)
+    wide = df.reset_index().pivot(columns=['year'], index=['building_group', 'hp_source'], values=['RV_HP']).reset_index()
+    wide.columns = ['building_group', 'hp_source'] + [c for c in wide.columns.get_level_values(1)[2:]]
+    return wide