mrio-toolbox 1.0.0__py3-none-any.whl

mrio_toolbox/_parts/_Part.py

@@ -0,0 +1,1504 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Apr 14 14:13:17 2023
+
+@author: beaufils
+"""
+
+import os
+import itertools
+import numpy as np
+import pandas as pd
+import xarray as xr
+from mrio_toolbox._parts._Axe import Axe
+import logging
+from mrio_toolbox.utils import converters
+from mrio_toolbox.utils.loaders import make_loader
+from mrio_toolbox.utils.savers import save_part_to_folder,save_to_nc
+from mrio_toolbox._parts import part_operations
+
+log = logging.getLogger(__name__)
+
+def load_part(
+        **kwargs
+        ):
+    loader = make_loader(**kwargs)
+    return Part(**loader.load_part(**kwargs))
+
+class Part:
+    def __init__(self,data=None,
+                 labels=None,
+                 axes=None,
+                 **kwargs):
+        """MRIO Parts object
+
+        MRIO Parts are the basic building blocks of the MRIO toolbox.
+        A Part is built from a numpy array and a set of Axes,
+        corresponding to the dimensions of the array.
+        The Axes hold the labels of the Part in the different dimensions
+        and are used to perform advanced indexing and operations on the Part.
+
+        Axes support multi-level indexing and groupings.
+
+        Parameters
+        ----------
+        data : numpy array
+            Numerical data of the part.
+            If left empty, a Part of zeros (or any other fill value) is created
+            with a shape matching the axes.
+        groupings : dict of label level : dict
+            Groupings of the labels of the Part, for each label defined.
+            The groupings are passed to the Axe objects.
+        labels : list of str or dict, optional
+            Labels of the axes.
+            The upper level of the list correspond to each axe (numpy dimension)
+              of the part
+            The lower level correspond to the dict of labels for each axe.
+            Remember that an Axe can have different levels of labels.
+        axes : list of Axe instances, optional
+            Custom Axes for the Part.
+            If left empty, the axes are created from the labels.
+        kwargs : dict
+            Additional metadata of the Part.
+            (e.g. path, name, multiplier, unit...)
+
+        Returns
+        -------
+        None.
+
+        """
+        
+        if data is not None:
+            if isinstance(data,Part):
+                data = data.data
+            if isinstance(data,(xr.DataArray,xr.Dataset)):
+                self.__init__(
+                    **converters.xarray.make_part(data)
+                    )
+                return
+            if isinstance(data,pd.DataFrame):
+                self.__init__(**converters.pandas.make_part(
+                    data)
+                    )
+                return
+            self.data = data
+            self.ndim = data.ndim
+            self.shape = self.data.shape
+
+        self.name = kwargs.pop("name","new_part")
+        log.debug("Create Part instance " + self.name)
+        
+        self.groupings = kwargs.get("groupings",dict())
+        self.metadata = kwargs.get("metadata",dict())
+        self.metadata = {**self.metadata,**kwargs}
+
+        if axes is None:
+            self.axes = []
+            self._create_axis(labels)
+        else:
+            self.axes = axes
+
+        if data is None:
+            log.debug("Create empty Part")
+            data = np.zeros([len(ax) for ax in self.axes])
+            self.data = data
+            self.ndim = data.ndim
+            self.shape = self.data.shape
+
+        self.fix_dims()
+        for dim in range(self.ndim):
+            if len(self.axes[dim]) != self.shape[dim] and self.shape[dim] != 1:
+                log.critical(f"Length of label {dim} does not match data: "+\
+                                f"{len(self.axes[dim])} and {self.shape[dim]}")
+                raise ValueError(f"Length of label {dim} does not match data: "+\
+                                 f"{len(self.axes[dim])} and {self.shape[dim]}")
+
+        if "_original_dimensions" in self.metadata.keys():
+            #This set of instructions is intended to handle
+            #Data loaded from netcdf files
+            log.info("Checking if a reformatting is needed.")
+            original = self.metadata.pop("_original_dimensions")
+            new_dims = [[]]
+            for dim in original:
+                #Decode the original dimensions
+                #Because netcdf files do not support multi-level attributes
+                if dim == "_sep_":
+                    new_dims.append([])
+                else:
+                    new_dims[-1].append(dim)
+
+            if new_dims != self.get_dimensions():
+                log.info("Reformat the Part")
+                new_part = self.reformat(new_dims)
+                self.data = new_part.data
+                self.axes = new_part.axes
+                self.ndim = self.data.ndim
+                self.shape = self.data.shape
+        
+        self._store_labels()
+
+    
+    def alias(self,**kwargs):
+        """Create a new Part in which only prescribed parameters are changed
+        
+        The current Part is taken as reference: all arguments not explicitely
+        set are copied from the current part."""
+        data = kwargs.get("data",self.data).copy()
+        name = kwargs.get("name",self.name+"_alias")
+        groupings = kwargs.get("groupings",self.groupings)
+        axes = kwargs.get("axes",self.axes)
+        labels = kwargs.get("labels",None)
+        metadata = kwargs.get("metadata",self.metadata)
+        return Part(
+            data=data,
+            name=name,
+            groupings = groupings,
+            labels=labels,
+            axes=axes,
+            metadata = metadata)
+        
+    def _create_axis(self,labels):
+        """
+        Create an Axe object based on a tuple of lists of indices
+
+        Parameters
+        ----------
+        *args : tuple of lists of str, list of str
+            Labels of the axe.
+            The first argument is used as the main label.
+            The second argument (if any) is used as secondary label.
+            If left empty, the axe is labelled by indices only
+
+        Raises
+        ------
+        TypeError
+            Raised if the arguments types differs from the number of dimensions
+            or if input labels are incorrect.
+        ValueError
+            Raised if the label length does not match the data.
+
+        Returns
+        -------
+        None.
+
+        """
+        self.axes = []
+        
+        if isinstance(labels,(tuple,list)) and len(labels)>0 and\
+            isinstance(labels[0],(str,int,float)):
+            #If the first item of the labels is not iterable,
+            #we assume the label is an axis label
+            labels = [labels]
+            #We add a dimension to the labels such that enumeration works properly
+
+        if "data" in self.__dict__.keys():
+            enum = self.ndim
+        else:
+            if labels is None:
+                raise ValueError("Cannot create axes without data, axes or labels")
+            enum = len(labels)
+            
+        for dim in range(enum):
+            if labels is None or dim > len(labels):
+                #Fill empty labels with indices
+                self.axes.append(
+                    Axe([i for i in range(self.shape[dim])],
+                        groupings = self.groupings)
+                )
+            elif isinstance(labels,dict):
+                axname = list(labels.keys())[dim]
+                self.axes.append(
+                    Axe(labels[axname],groupings=self.groupings,name=axname)
+                )
+            elif isinstance(labels,(list,tuple)):
+                self.axes.append(
+                    Axe(labels[dim],groupings=self.groupings)
+                )
+            else:
+                log.critical("Unkown label type: "+type(labels))
+                raise TypeError("Unknown label type: "+type(labels))
+            log.debug(f"Create ax {dim} with len {self.axes[-1]}")
+        
+    def fix_dims(self,
+                 skip_labels=False,
+                 skip_data=False):
+        """Align the number of axes with the number of dimensions
+        
+        If one length exceeds the other, axes and/or data are squeezed,
+        i.e. dimensions of length 1 are removed."""
+        if len(self.axes) == self.ndim:
+            return
+        log.warning(
+            f"The number of axes ({len(self.axes)})"\
+                    +f" does not match data dimensions ({self.ndim})"
+            )
+
+        if len(self.axes) > self.ndim and not skip_labels:
+            log.debug(
+                "Try to squeeze axe(s) of len 1"
+            )
+            counter = [len(ax) != 1 for ax in self.axes]
+            self.axes = self.axes[counter]
+            return self.fix_dims(
+                    skip_labels=True,
+                    skip_data=skip_data)
+        if self.ndim > len(self.axes) and not skip_data:
+            self.data = self.data.squeeze()
+            self.shape = self.data.shape
+            self.ndim = self.data.ndim
+            return self.fix_dims(
+                skip_labels=skip_labels,
+                skip_data=True)
+        
+        log.critical("Cannot reconcile data of dims "+ self.ndim+\
+                            " with axes of dim " +len(self.axes))
+        raise IndexError("Cannot reconcile data of dims "+ self.ndim+\
+                        " with axes of dim " +len(self.axes))
+        
+
+    def __getitem__(self,args):
+        if isinstance(args,str) or isinstance(args,int) or isinstance(args,np.integer):
+            args = (args,)
+        return self.get(*args)
+    
+    def __setitem__(self,args,value):
+        if isinstance(value,Part):
+            value = value.data
+        self.setter(value,*args)
+    
+    def setter(self,value,*args):
+        """
+        Change the value of a data selection
+
+        Parameters
+        ----------
+        value : float or numpy like
+            Value to set.
+        *args : list of tuples
+            Indices along the respective axes.
+
+        Returns
+        -------
+        None.
+        Modification is applied to the current Part object
+
+        """
+        sels = []
+        try:
+            #First tries to interpret one arg per ax
+            for i,arg in enumerate(args):
+                sels.append(self.axes[i].get(arg))
+        except (IndexError,ValueError):
+            #Otherwise, tries to interpret all args on the first ax
+            sels = []
+            sels = [self.axes[0].get(args)]
+        if isinstance(value,(np.ndarray,Part)) and len(sels)!=value.ndim:
+            if len(sels) < value.ndim:
+                value = value.squeeze() 
+            if len(sels) > value.ndim:
+                target_shape = [len(sel) for sel in sels]
+                value = np.reshape(value,target_shape)   
+        self.data[np.ix_(*sels)] = value
+    
+    def get(self,*args,aspart=True,squeeze=False):
+        """
+        Extract data from the current Part object
+
+        Parameters
+        ----------
+        *args : list of tuples
+            Selection along the Axes of the Part.
+        aspart : bool, optional
+            Whether to return the selection as a Part object. 
+            If False, the selection is returned as a numpy object.
+            The default is True.
+        squeeze : bool, optional
+            Whether to remove dimensions of length 1.
+            The default is False
+
+        Returns
+        -------
+        New Part object or numpy object
+
+        """
+        sels = []
+        axes = []
+
+        #Extract the indices for the selection
+        try:
+            #First tries to interpret one arg per ax
+            for i,arg in enumerate(args):
+                datasel,labs,groupings = self.axes[i].get(arg,True)
+                if not squeeze or len(datasel) > 1:
+                    axes.append(Axe(labs,groupings))
+                sels.append(datasel)  
+        except (ValueError, IndexError):
+            sels = []
+            axes = []
+            #Try interpreting all args on the first ax
+            datasel,labs,groupings = self.axes[0].get(args,True)
+            if not squeeze or len(datasel) > 1:
+                axes.append(Axe(labs,groupings))
+            sels.append(datasel)
+
+        #If the selection is not complete, fill with all
+        if len(sels)<self.ndim:
+            for i in range(len(sels),self.ndim):
+                datasel,labs,groupings = self.axes[i].get("all",True)
+                if not squeeze or len(datasel) > 1:
+                    axes.append(Axe(labs,groupings))
+                sels.append(datasel)
+
+        #Execute the selection
+        data = self.data[np.ix_(*sels)]
+
+        #Return the selection
+        if squeeze:
+            data = data.squeeze()
+        if aspart:
+            return Part(data=data,name=f"sel_{self.name}",
+                        groupings=self.groupings,axes = axes)
+        return data
+        
+    def develop(self,axis=None,on=None,squeeze=True):
+        """
+        Reshape a Part to avoid double labels
+
+        Parameters
+        ----------
+        axis : int or list of int, optional
+            Axis to develop. 
+            If left empty, all axes are developed.
+            The default is None.
+        on : str or list of str, optional
+            Dimensions to develop.
+            If left empty, all dimensions are developed.
+            Note that the develop method does not support the developping of
+            non-contiguous dimensions.
+            The default is None.
+        squeeze : bool, optional
+            Whether to remove dimensions of length 1.
+            The default is True.
+        
+        Returns
+        -------
+        Part object
+            Developped Part
+
+        """
+        if isinstance(on,str):
+            on = [on]
+        if isinstance(axis,int):
+            axis = [axis]
+        axes = []
+        for i,ax in enumerate(self.axes):
+            if axis is None or i in axis:
+                labels = ax.labels.copy()
+                for dim in ax.dimensions:
+                    if on is None or dim in on:
+                        #Add dimension that needs to be developed
+                        axes.append(Axe({dim:labels[dim]},groupings=ax.groupings))
+                        labels.pop(dim)
+                if len(labels) > 0:
+                    #Keep remaining dimensions together
+                    axes.append(Axe(labels,groupings=ax.groupings))
+            else:
+                axes.append(ax)
+
+        
+        old_dim_order = [dim for ax in self.axes for dim in ax.dimensions]
+        new_dim_order = [dim for ax in axes for dim in ax.dimensions]
+        if new_dim_order != old_dim_order:
+            raise NotImplementedError(
+                "Developping the part misaligns the dimensions. "+\
+                    "This operation is not yet supported."
+                )
+        #If the order of the dimensions is unchanged, we can simply reshape
+        shape = [len(ax) for ax in axes]
+        data = self.data.reshape(shape)
+        if squeeze:
+            return Part(data=data,
+                        name=f"developped_{self.name}",
+                        groupings=self.groupings,axes=axes).squeeze()
+        return Part(data=data,name=f"developped_{self.name}",
+                    groupings=self.groupings,axes=axes)
+
+    def reformat(self,new_dimensions):
+        """
+        Reshape a Part to match a new dimensions combination
+
+        Equivalent to a combination of the develop and combine_axes methods.
+
+        This only works for contiguous dimensions in the current Part,
+        without overlapping dimensions.
+        For example, if the Part has dimensions:
+            [["countries"],["sectors"],["sectors"]]
+        The following is allowed:
+            [["countries","sectors"],["sectors"]]
+        The following is not allowed:
+            [["countries"],["sectors","sectors"]]
+            [["sectors"],["countries","sectors"]]
+            [["sectors","countries"],["sectors"]]
+
+        Parameters
+        ----------
+        dimensions : list of list of str
+            Original dimensions of the Part
+
+        Returns
+        -------
+        data : numpy array
+            Reshaped data
+        axes : list of Axe instances
+            Reshaped axes
+
+        """
+        return part_operations.reformat(self,new_dimensions)
+    
+    def combine_axes(self,start=0,end=None,in_place=False):
+        """
+        Combine axes of a Part into a single one.
+
+        The order of dimensions is preserved in the new axis.
+        Only consecutive axes can be combined.
+        The method can be used to revert the develop method.
+
+        Parameters
+        ----------
+        start : int, optional
+            Index of the first axis to combine, by default 0
+        end : int, optional
+            Index of the final axis to combine, by default None,
+            all axis are combined, i.e. the Part is flattened.
+
+        Returns
+        -------
+        Part instance
+
+        Raises
+        ------
+        IndexError
+            Axes should have no overlapping dimensions.
+        """
+        axes = []
+        covered = []
+        labels,groupings = dict(),dict()
+        if end is None:
+            end = self.ndim - 1
+        for i,ax in enumerate(self.axes):
+            if i in range(start,end+1):
+                for name in ax.dimensions:
+                    if name in covered:
+                        raise IndexError(
+                            "Cannot undevelop axes with overlapping dimensions"
+                        )
+                    covered.append(name)
+                    labels[name] = ax.labels[name]
+                    if name in ax.groupings.keys():
+                        groupings[name] = ax.groupings[name]
+            else:
+                axes.append(ax)
+            if i == end:
+                axes.append(Axe(labels,groupings))
+        
+        new_shape = [len(ax) for ax in axes]
+        data = self.data.reshape(new_shape)
+        if in_place:
+            return data,axes
+        return self.alias(data=data,name=f"combined_{self.name}",
+                    axes=axes)
+    
+    def swap_axes(self,axis1,axis2):
+        """Swap two axes of a Part
+        
+        Parameters
+        ----------
+        axis1 : int
+            First axis to swap.
+        axis2 : int
+            Second axis to swap.
+
+        Returns
+        -------
+        Part instance
+            Part with swapped axes.
+
+        """
+        axes = self.axes.copy()
+        axes[axis1],axes[axis2] = axes[axis2],axes[axis1]
+        data = self.data.swapaxes(axis1,axis2)
+        return Part(data=data,name=f"swapped_{self.name}",axes=axes)
+    
+    def swap_ax_levels(self,axis,dim1,dim2):
+        """Swap two levels of an axis
+        
+        Parameters
+        ----------
+        axis : int
+            Axis to modify.
+        dim1 : str
+            First dimension to swap.
+        dim2 : str
+            Second dimension to swap.
+
+        Returns
+        -------
+        Part instance
+            Part with swapped levels.
+
+        """
+        len1,len2 = len(self.axes[axis].labels[dim1]),len(self.axes[axis].labels[dim2])
+        if len1 == 1 or len2 == 1:
+            axes = self.axes.copy()
+            axes[axis].swap_levels(dim1,dim2)
+            return Part(data=self.data,name=f"swapped_{self.name}",axes=axes)
+        dimensions = self.axes[axis].dimensions
+        id1,id2 = dimensions.index(dim1),dimensions.index(dim2)
+        offset = sum([len(ax.dimensions) for ax in self.axes[:axis]])
+        dev = self.develop(axis)
+        dev = dev.swap_axes(id1+offset,id2+offset)
+        dev = dev.combine_axes(axis,axis+len(dimensions)-1)
+        dev.name = f"swapped_{self.name}"
+        return dev
+    
+    def flatten(self,invert=False):
+        """Flatten a 2D Part into a 1D Part
+        
+        Parameters
+        ----------
+        inverse : bool, optional
+            Whether to in the inverse level order.
+        """
+        if self.ndim != 2:
+            raise ValueError(f"Cannot flatten Part with {self.ndim} dimensions")
+        if invert:
+            labels = {
+                dimension : self.axes[i].labels[dimension] \
+                    for i in range(self.ndim,0,-1) \
+                        for dimension in self.axes[i].dimensions
+            }
+            order = "C"
+        else:
+            labels = {
+                dimension : self.axes[i].labels[dimension] \
+                    for i in range(self.ndim) \
+                        for dimension in self.axes[i].dimensions
+            }
+            order = "F"
+        ax = Axe(labels,self.groupings)
+        return self.alias(data=self.data.flatten(order=order),
+                            name=f"flattened_{self.name}",
+                            axes =[ax])
+        
+        
+    def squeeze(self):
+        axes = []
+        for ax in self.axes:
+            if len(ax) > 1:
+                axes.append(ax)
+        return self.alias(data=np.squeeze(self.data),axes=axes,
+                          name=f"squeezed_{self.name}")
+    
+    def expand_dims(self,axis,copy=None):
+        """Add dimensions to a Part instance
+        
+        Parameters
+        ----------
+        axis : int
+            Position of the new axis.
+        copy : int, optional
+            Axis to copy the labels from.
+            If left empty, the axis is created without labels.
+        """
+        axes = self.axes.copy()
+        if copy is not None:
+            axes.insert(axis,self.axes[copy])
+        else:
+            axes.insert(axis,Axe({"expanded":[0]}))
+        return self.alias(data = np.expand_dims(self.data,axis),
+                          axes=axes,name=f"expanded_{self.name}")
+    
+    def copy(self):
+        """Return a copy of the current Part object"""
+        return self.alias()
+        
+    def extraction(self,
+                   dimensions,
+                   labels=["all"],
+                   on_groupings=True,
+                   domestic_only=False,
+                   axis="all"):
+        """
+        Set labels over dimension(s) to 0.
+
+        Parameters
+        ----------
+        dimensions : str, list of str, dict
+            Name of the dimensions on which the extraction is done.
+            If dict is passed, the keys are interpreted as the dimensions
+            and the values as the labels
+        labels : list of (list of) str, optional
+            Selection on the dimension to put to 0.
+        on_groupings : bool, optional
+            Whether to use the groupings to select the labels.
+            This matters only when the domestic_only argument is set to True.
+        domestic_only : bool, optional
+            If yes, only domestic transactions are set to 0 and trade flows 
+            are left untouched. The default is False.
+        axis : list of ints, optional
+            Axis along which the extraction is done. The default is "all".
+            In any case, the extraction only applies to axis allowing it, that
+            is in axis containing zones or countries labels corresponding to 
+            the zone selection.
+
+        Returns
+        -------
+        Part object
+            New Part with selection set to 0.
+
+        """
+        if isinstance(dimensions,str):
+            dimensions = [dimensions]
+        if isinstance(labels,str) and labels!="all":
+            labels = [labels]
+        if len(labels) != len(dimensions):
+            if len(dimensions)==1:
+                #If only one dimension is passed, we broadcast the labels
+                labels = [labels]
+            else:
+                #Raise an error for ambiguous cases
+                log.critical("Number of dimensions and labels do not match for extraction")
+                raise ValueError("Number of dimensions and labels do not match for extraction")
+        if isinstance(dimensions,dict):
+            to_select = dimensions
+            labels = list(to_select.values())
+            dimensions = list(dimensions.keys())
+        else:
+            to_select = dict()
+            for dim,label in zip(dimensions,labels):
+                to_select[dim] = label
+
+        allowed = []
+        for i,ax in enumerate(self.axes):
+            if all(dimension in ax.dimensions for dimension in dimensions):
+                allowed.append(i)
+        if len(allowed) == 0:
+            log.critical("No axis found for extraction on "+str(dimensions))
+            raise ValueError("No axis found for extraction on "+str(dimensions))
+        if axis == "all":
+            log.info(f"Extract {to_select} on axes "+ str(allowed))
+            axis = allowed
+        if isinstance(axis,int):
+            axis = [axis]
+        if not all(ax in allowed for ax in axis):
+            wrong = [ax for ax in axis if ax not in allowed]
+            log.critical(f"Cannot extract {dimensions} on axis {wrong}")
+            raise ValueError(f"Cannot extract {dimensions} on axis {wrong}")
+
+        if not domestic_only:
+            #If no domestic_only, we can simply set the selection to 0
+            sel = ["all"]*self.ndim
+            for i,ax in enumerate(self.axes):
+                if i in axis:
+                    sel[i] = to_select
+            output = self.copy()
+            output[sel] = 0
+            return output
+        
+        if not on_groupings:
+            #If no groupings, develop the selected groupings
+            for i,dim in enumerate(dimensions):
+                if dim in self.groupings.keys():
+                    for label in labels[i]:
+                        if label in self.groupings[dim].keys():
+                            labels[i].append(self.groupings[dim][label])
+                            labels[i].remove(label)
+        
+        for i,label in enumerate(labels):
+            if label == "all":
+                if on_groupings and dimensions[i] in self.groupings.keys():
+                    labels[i] = list(
+                        self.axes[axis[i]].groupings[dimensions[i]].keys()
+                        )
+                else:
+                    labels[i] = self.axes[axis[i]].labels[dimensions[i]]
+
+        output = self.copy()
+        for label in itertools.product(*labels):
+            #Iteratively set domestic selections to 0
+            seldict = dict(zip(dimensions,label))
+            sel = ["all"]*self.ndim
+            for i in range(self.ndim):
+                if i in axis:
+                    sel[i] = seldict
+            output[sel] = 0
+        return output
+    
+    def leontief_inversion(self):
+        if self.ndim == 2 and self.issquare():
+            data = np.linalg.inv(np.identity(len(self.axes[0])) - self.data)
+            return self.alias(name=f"l_{self.name}",data=data)
+        raise ValueError("Can only compute the Leontief inverse on"+\
+                         " square parts")
+    
+    def zone(self):
+        """
+        Apply a grouping by zone dependent on the shape of the Part.
+        
+        Final demand Parts are summed over zones.
+        Horizontal extensions are expanded by zone.
+
+        Raises
+        ------
+        AttributeError
+            Parts with other shapes are rejected.
+
+        Returns
+        -------
+        Part object
+            Grouped part.
+
+        """
+        log.warning("This function is deprecated as it returns different "+\
+                    "results depending on the shape of the Part. "+\
+                        "Use group or expand instead.")
+        if self.ndim == 2 and not self.issquare():
+            if "countries" in self.axes[1].dimensions:
+                #Expand normal parts
+                return self.group(1,"countries")
+        if self.ndim == 1:
+            #Expand horizontal parts
+            return self.expand("countries")
+        raise AttributeError(f"Part {self.name} has no predefined grouping.")
+    
+    def update_groupings(self,groupings,ax=None):
+        """Update the groupings of the current Part object
+        
+        groupings : dict
+            Description of the groupings
+        ax: int, list of int
+            Axes to update. If left empty, all axes are updated.
+        """
+        self.groupings = groupings
+        if ax is None:
+            ax = range(self.ndim)
+        for axe in list(ax):
+            self.axes[axe].update_groupings(groupings)
+    
+    def aggregate(self,on="countries",axis=None):
+        """Aggregate dimensions along one or several axis.
+
+        If groupings are defined, these are taken into account.
+        If you want to sum over the dimension of an axis, use the sum method.
+
+        If no axis is specified, the operation is applied to all axes.
+
+        Parameters
+        ----------
+        axis : str or list of str, optional
+            List of axis along which countries are grouped.
+            If left emtpy, countries are grouped along all possible axis.
+        on : str or dict, optional
+            Indicate wether the grouping should be done by zones ("zones")
+            or by sector ("sectors"), or both ("both").
+            The default is "zones".
+            If both, the operation is equivalent to summing over an axis
+
+        Raises
+        ------
+        ValueError
+            Raised if a selected Axe cannot be grouped.
+
+        Returns
+        -------
+        Part object
+            Part grouped by zone.
+
+        """
+        log.debug(f"Aggregate Part {self.name} along axis {axis} on {on}")
+        
+        
+        if isinstance(on,list):
+            for item in on:
+                self = self.aggregate(axis,item)
+            return self
+        if on not in self.groupings.keys():
+            raise ValueError(f"No groupings defined for dimensions {on}")
+        
+        if axis is None:
+            axis = self.hasax(
+                on
+            )
+        if isinstance(axis,int):
+            axis = [axis]
+            
+        output = self.alias()
+
+        for ax in axis:
+            output = output.aggregate_on(on,ax)
+
+        output.name = f"{on}_grouped_{self.name}"
+        return output
+    
+    def aggregate_on(self,on,axis):
+        """Aggregate a Part along a given axis
+        
+        Parameters
+        ----------
+        on : str
+            Dimension to aggregate on
+        axis : int
+            Axis to aggregate
+        
+        Returns
+        -------
+        Part instance
+            Aggregated Part
+        """
+        if on not in self.axes[axis].dimensions:
+            raise ValueError(f"Dimension {on} not found in axis {axis}")
+        
+        new_labels = self.axes[axis].labels.copy()
+        new_labels[on] = list(self.axes[axis].groupings[on].keys())
+        new_groupings = self.groupings.copy()
+        new_groupings[on] = {
+            item : [item] for item in self.groupings[on]
+        }
+
+        new_axis = self.axes.copy()
+        new_axis[axis] = Axe(new_labels,new_groupings)
+        new_shape = [len(ax) for ax in new_axis]
+
+        output = Part(axes=new_axis)
+        idsum = new_axis[axis].dimensions.index(on) #Index of the dimension to sum on
+        ref_dev = self.develop(axis)
+        new_dev = output.develop(axis)
+        selector = ["all"]*ref_dev.ndim
+        for label in new_labels[on]:
+            selector[axis+idsum] = label
+            new_dev[selector] = ref_dev[selector].sum(
+                axis=axis+idsum,
+                keepdims=True
+            )
+        output = new_dev.data.reshape(new_shape)
+        return self.alias(data=output,name=f"{on}_grouped_{self.name}",
+                                axes=new_axis)
+
+ 
+    def get_labels(self,axis=None):
+        """
+        Returns the dictionnary of the Part labels
+
+        Parameters
+        ----------
+        axis : int or list of int, optional
+            Axis to investigate, by default None,
+            All axes are investigated.
+        
+        Returns
+        -------
+        list
+            Labels used in the part.
+        """
+        labels = []
+        if axis is None:
+            axis = range(self.ndim)
+        if isinstance(axis,int):
+            axis = [axis]
+            #Make sure the axis is iterable
+        for ax in axis:
+            labels.append(self.axes[ax].labels)
+        return labels
+    
+    def list_labels(self):
+        """List the labels of the Part"""
+        labels = dict()
+        ax_labels = self.get_labels()
+        for ax in ax_labels:
+            for label in ax.keys():
+                if label not in labels.keys():
+                    labels[label] = ax[label]
+        return labels
+    
+    def get_dimensions(self,axis=None):
+        """
+        Returns the list dimensions of the Part
+
+        Parameters
+        ----------
+        axis : int or list of int, optional
+            Axis to investigate, by default None,
+            All axes are investigated.
+        
+        Returns
+        -------
+        list
+            Dimensions of the axes.
+        """
+        dimensions = []
+        if axis is None:
+            axis = range(self.ndim)
+        if isinstance(axis,int):
+            axis = [axis]
+            #Make sure the axis is iterable
+        for ax in axis:
+            dimensions.append(self.axes[ax].dimensions)
+        return dimensions
+    
+    def rename_labels(self,old,new):
+        """
+        Rename some labels of the Part
+
+        Parameters
+        ----------
+        old : str
+            Name of the label to change.
+        new : str
+            New label name.
+        """
+        for ax in self.axes:
+            if old in ax.dimensions:
+                ax.rename_labels(old,new)
+        self._store_labels()
+
+    def replace_labels(self,name,labels,axis=None):
+        """
+        Update a label of the part
+
+        Parameters
+        ----------
+        name : str
+            Name of the label to update, by default None
+        labels : dict or list
+            New labels for the corresponding ax.
+            If a list is passed, the former label name is used.
+        axis : int, list of int, optional
+            List of axis on which the label is changed.
+            By default None, all possible axes are updated.
+        """
+        if axis is None:
+            axis = range(self.ndim)
+        if isinstance(axis,int):
+            axis = [axis]
+        if isinstance(labels,list):
+            labels = {name:labels}
+        for ax in axis:
+            if name in self.axes[ax].dimensions:
+                self.axes[ax].replace_labels(name,labels)
+        self._store_labels()
+    
+    def set_labels(self,labels,axis=None):
+        """
+        Change the labels of the Part
+
+        Parameters
+        ----------
+        labels : dict or nested list
+            New labels of the axes.
+            If a nested list is passed, the first level corresponds to the axes
+        axis : str, optional
+            Axis on which the labels are changes, by default None,
+            all axes are updated.
+        """
+        if axis is None:
+            axis = range(self.ndim)
+        if isinstance(axis,int):
+            axis = [axis]
+        if isinstance(labels,list) and len(labels) == self.ndim:
+            labels = {i:labels[i] for i in range(self.ndim)}
+        for ax in axis:
+            self.axes[ax].set_labels(labels[ax])
+        self._store_labels()
+    
+    def _store_labels(self):
+        """Store the labels of the Part"""
+        self.labels = self.list_labels()
+
+    def add_labels(self,labels,dimension=None,axes=None,
+                    fill_value=0):
+        """
+        Add indices to one or multiple Part axes.
+
+        Parameters
+        ----------
+        new_labels : list of str or dict
+            List of indices to add
+        dimension : str, optional
+            Labels the new indices should be appended to,
+            in case new_labels is not a dict.
+            If new_labels is a dict, dimension is ignored.
+        axes : int or set of ints, optional
+            Axes or list of axes to modify.
+            In case it is not specified, the axes are detected
+            by looking for the dimension (or new_labels keys) in each ax.
+        fill_value : float, optional
+            Value used to initialize the new Part
+
+        Returns
+        ----------
+        Part instance
+            Part instance with the additional ax indices.
+
+        Raise
+        ----------
+        ValueError
+            A Value Error is raised if neither the axes nor the 
+            ref_set arguments are set.
+        """
+        if isinstance(labels,list):
+            labels = {dimension:labels}
+        dimension = list(labels.keys())[0]
+        if axes is None:
+            #Identify the axes with the ref_set in labels
+            axes = self.hasax(dimension)
+        elif isinstance(axes,int):
+            axes = [axes]
+
+        new_axes = self.axes.copy()
+        sel = ["all"]*self.ndim
+        for ax in axes:
+            log.debug("Add labels to axis "+str(ax))
+            sel[ax] = dict()
+            old_labels = self.axes[ax].labels
+            new_labels = self.axes[ax].labels.copy()
+            new_labels[dimension] = old_labels[dimension] + labels[dimension]
+            new_axes[ax] = Axe(new_labels,self.groupings)
+            sel[ax] = old_labels
+
+        new_shape = [len(ax) for ax in new_axes]
+        output = self.alias(data=np.full(new_shape,fill_value,dtype="float64"),
+                            axes=new_axes)
+        
+        #Put original data back in place
+        output[sel] = self.data
+        return output
+    
+    def expand(self,axis=None,over="countries"):
+        """
+        Expand an axis of the Part
+        
+        Create a new Axes with a unique dimension.
+        Note that this operation significantly expands the size of the Part.
+        It is recommended to use this method with Extension parts only.
+
+        Parameters
+        ----------
+        axis : int, optional
+            Axe to extend. 
+            If left empty, the first suitable axe is expanded.
+        over : str, optional
+            Axe dimension to expand the Part by.
+            The default is "countries".
+
+        Returns
+        -------
+        Part object
+            New Part object with an additional dimension.
+
+        """
+        if axis is None:
+            axis = self.hasax(over)
+        
+        for ax in axis:
+            ref_ax = self.axes[ax]
+            new_ax = Axe(ref_ax.labels[over],groupings=self.groupings)
+            axes = self.axes.copy()
+            axes.insert(ax,new_ax)
+            new_shape = list(self.shape)
+            new_shape.insert(axis,len(new_ax))
+            output = np.zeros(new_shape)
+            selector = [slice(None)]*self.ndim
+            for item in ref_ax.labels[over]:
+                newsel,refsel = selector.copy(),selector.copy()
+                newsel.insert(ax,new_ax.sel(item))
+                newsel[ax+1] = ref_ax.sel(item)
+                refsel[ax]= ref_ax.sel(item)
+                output[tuple(newsel)] = self.data[tuple(refsel)]
+        return self.alias(data=output,name=f"expanded_{self.name}",axes=axes)
+    
+    def issquare(self):
+        """Assert wether the Part is square"""
+        return self.ndim == 2 and len(self.axes[0])==len(self.axes[1])
+    
+    def hasneg(self):
+        """Test whether Part has negative elements"""
+        if np.any(self.data<0):
+            return True
+        return False
+
+    def hasax(self,name=None):
+        """Returns the dimensions along which a Part has given labels
+        
+        If no axis can be found, an empty list is returned empty.
+        This method can be used to assert the existence of a given dimension
+        in the part.
+
+        Parameters
+        ----------
+        name : int, optional
+            Name of the label to look for.
+            If no name is given, all axes are returned.
+
+        Returns
+        -------
+        axes : list of ints
+            Dimensions along which the labels are found.
+
+        """
+        if name == "any" or name is None:
+            return [i for i in range(self.ndim)]
+        axes = []
+        for i,ax in enumerate(self.axes):
+            if name in ax.dimensions:
+                axes.append(i)
+        return axes
+    
+    def __str__(self):
+        return f"{self.name} Part object with {self.ndim} dimensions"
+    
+    def sum(self,axis=None,on=None,keepdims=False):
+        """
+        Sum the Part along one or several axis, and/or on a given dimension.
+
+        Parameters
+        ----------
+        axis : int or list of int, optional
+            Axe along which the sum is evaluated.
+            By default None, the sum of all coefficients of the Part is returned
+        on : str, optional
+            name of the dimension to be summed on.
+            If no axis is defined, the Part is summed over all axis having 
+            the corresponding dimension.
+            By default None, the full ax is summed
+        keepdims : bool, optional
+            Whether to keep the number of dimensions of the original.
+            By default False, the dimensions of lenght 1 are removed.
+
+        Returns
+        -------
+        Part instance or float
+            Result of the sum.
+        """
+        if axis is None:
+            if on is None:
+                return self.data.sum()
+            if not keepdims:
+                self = self.squeeze()
+            axis = self.hasax(on)
+        if isinstance(axis,int):
+            if on is not None:
+                return self._sum_on(axis,on,keepdims)
+            ax = self.axes.copy()
+            if not keepdims:
+                del ax[axis]
+            else:
+                ax[axis] = Axe(["all"])
+            return self.alias(
+                data=self.data.sum(axis,keepdims=keepdims),
+                name=f"{self.name}_sum_{axis}",
+                axes = ax
+            )
+        axis = sorted(axis)
+        for ax in axis[::-1]:
+            self = self.sum(ax,on,keepdims)
+        return self
+
+    def _sum_on(self,axis,on,keepdims=False):
+        """
+        Sum a Part along an axis on a given dimension
+        """
+        ax = self.axes[axis]
+        if on not in ax.dimensions:
+            raise ValueError(f"Cannot sum on {on} as it is not a dimension of axis {axis}")
+        if ax.levels == 1:
+            #If the axis has a single level, this is a simple sum
+            return self.alias(
+                data = self.data.sum(axis,keepdims=keepdims),
+                name=f"{self.name}_sum_{axis}",
+                axes = self.axes
+            )
+        #Otherwise, sum on the relevant levels
+        idsum = ax.dimensions.index(on) #Index of the dimension to sum on
+        dev = self.develop(axis,squeeze=False)
+        dev = dev.sum(axis+idsum,keepdims=keepdims)
+        if keepdims:
+            dev = dev.combine_axes(axis,axis+idsum)
+        dev.name = f"{self.name}_sum_on_{on}_{axis}"
+        return dev
+        
+    def save(self,
+             file=None,
+             name=None,
+             extension=".npy",
+             overwrite=False,
+             include_labels=False,
+             write_instructions=False,
+             **kwargs):
+        """
+        Save the Part object to a file
+
+        Parameters
+        ----------
+        name : str, optional
+            Name under which the Part is saved.
+            By default, the current part is used.
+        path : Path-like, optional
+            Directory in which the Path should be saved, 
+            by default None, the dir from which the part was loaded.
+        extension : str, optional
+            Format under which the part is saved. The default ".npy"
+            If ".csv" is chosen, the part is saved as a csv file with labels
+        file : path-like, optional
+            Full path to the file to save the Part to.
+            This overrides the path, name and extension arguments.
+        overwrite : boolm optional
+            Whether to overwrite an existing file. 
+            If set False, the file is saved with a new name.
+            The default is False.
+        write_instructions : bool, optional
+            Whether to write the loading instructions to a yaml file.
+            The default is False.
+        include_labels : bool, optional
+            Whether to include the labels in the saved file.
+            Only applicable to .csv and .xlsx files.
+        **kwargs : dict
+            Additional arguments to pass to the saving function
+
+        Raises
+        ------
+        FileNotFoundError
+            _description_
+        """
+        path = kwargs.get("path",None)
+        if file is not None:
+            path,name = os.path.split(file)
+            name,possible_extension = os.path.splitext(name)
+            if possible_extension != "":
+                extension = possible_extension
+        if name is None:
+            name = self.name
+        if path is None:
+            raise FileNotFoundError("No path specified for saving the Part")
+        if extension == ".nc":
+            path = os.path.join(path,name)
+            save_to_nc(self,path,overwrite,
+                       write_instructions=write_instructions,
+                       **kwargs)
+        else:
+            save_part_to_folder(
+                self,
+                path = path,
+                name = name,
+                extension = extension,
+                overwrite = overwrite,
+                include_labels=include_labels,
+                write_instructions = write_instructions,
+                **kwargs
+            )
+
+    def to_pandas(self):
+        """Return the current Part object as a Pandas DataFrame
+        
+        Only applicable to Parts objects with 1 or 2 dimensions.
+        """
+        return converters.pandas.to_pandas(self)
+    
+    def to_xarray(self):
+        """
+        Save the Part object to an xarray DataArray
+
+        Labels are directly passed to the DataArray as coords.
+        Note that data will be flattened.
+        The dimension order will be saved as an attribute.
+        If you're loading the data back,
+        the Part will be automatically reshaped to its original dimensions.
+
+        Returns
+        -------
+        xr.DataArray
+            Corresponding DataArray
+        """
+        return converters.xarray.to_DataArray(self)
+    
+    def mean(self,axis=None):
+        return self.data.mean(axis)
+    
+    def min(self,axis=None):
+        return self.data.min(axis)
+
+    def max(self,axis=None):
+        return self.data.max(axis)
+        
+    def mul(self,a,propagate_labels=True):
+        """
+        Matrix multiplication between parts with labels propagation
+
+        Parameters
+        ----------
+        a : Part or numpy array
+            Right-hand multiplicator.
+        propagate_labels : bool, optional
+            Whether to try propagating the labels from the right hand multiplicator
+            By default True.
+            If right-hand multiplicator is not a Part object, becomes False.
+
+        Returns
+        -------
+        Part instance
+            result of the multiplication
+        """
+        if not isinstance(a,Part):
+            propagate_labels = False
+            name="array"
+        else:
+            name = a.name
+        data = np.matmul(self.data,a.data)
+        axes = [self.axes[i] for i in range(self.ndim-1)]
+        for ax in range(a.ndim-1):
+            if propagate_labels:
+                axes.append(a.axes[ax+1])
+            else:
+                axes.append(Axe([i for i in range(a.shape[ax+1])]))
+        return self.alias(data=data,name=f"{self.name}.{name}",axes=axes)
+    
+    def filter(self,threshold,fill_value=0):
+        """
+        Set to 0 the values below a given threshold
+
+        Parameters
+        ----------
+        threshold : float
+            Threshold value.
+        fill_value : float, optional
+            Value to replace the filtered values with.
+            The default is 0.
+
+        Returns
+        -------
+        Part instance
+            Filtered Part.
+
+        """
+        data = self.data.copy()
+        data[data<threshold] = fill_value
+        return self.alias(data=data,name=f"filtered_{self.name}_{threshold}")
+    
+    def diag(self):
+        if self.ndim == 1:
+            return self.alias(data=np.diag(self.data),
+                              name=f"diag_{self.name}",
+                              axes = self.axes*2)
+                              
+    def __add__(self,a):
+        if isinstance(a,Part):
+            name = a.name
+            a = a.data
+        else:
+            name=""
+        if isinstance(a,np.ndarray) and self.ndim != a.ndim:
+            a = a.squeeze()
+            self = self.squeeze()
+        return self.alias(data=a+self.data,name=f"{self.name}+{name}")
+    
+    def __radd__(self,a):
+        return self.__add__(a)
+    
+    def __rmul__(self,a):
+        return self.__mul__(a)
+    
+    def __mul__(self,a):
+        if isinstance(a,Part):
+            name = "{a.name}*{self.name}"
+            a = a.data 
+        else:
+            if isinstance(a,int):
+                name = f"{a}*{self.name}"
+            else:
+                name = f"array*{self.name}"
+        data = self.data*a
+        if data.ndim!=self.ndim:
+            data = data.squeeze()
+        #Trust numpy to broadcast the multiplication
+        #Squeeze to get rid of unused dimensions
+        return self.alias(data=data,
+                          name=name)
+    
+    def __neg__(self):
+        return self.alias(data=-self.data,name=f"-{self.name}")
+    
+    def __sub__(self,a):
+        if isinstance(a,Part):
+            name = a.name
+            a = a.data
+        else:
+            name=""
+        return self.alias(data=self.data-a,name=f"{self.name}-{name}")
+    
+    def __rsub__(self,a):
+        if isinstance(a,Part):
+            name = a.name
+            a = a.data
+        else:
+            name=""
+        return self.alias(data=a-self.data,name=f"{name}-{self.name}")
+    
+    def power(self,a):
+        if isinstance(a,Part):
+            a = a.data
+            name = f"{self.name}**{a.name}"
+        elif isinstance(a,int) or isinstance(a,float):
+            name = f"{self.name}**{a}"
+        else:
+            name = f"{self.name}**array"
+        return self.alias(data=np.power(self,a),name=name)
+
+    def __pow__(self,a):
+        return self.power(a)
+
+    def __eq__(self,a):
+        if isinstance(a,Part):
+            return np.all(self.data==a.data)
+        return False
+    
+    def __rtruediv__(self,a):
+        if isinstance(a,Part):
+            name = f"{self.name}/{a.name}"
+            a = a.data
+        else:
+            if isinstance(a,int):
+                name = f"{a}/{self.name}"
+            else:
+                name= f"array/{self.name}"
+        if np.sum(self.data==0)!=0:
+            log.warning("Division by zero in "+name)
+        return self.alias(data=a/self.data,
+                          name=name)
+    
+    def __truediv__(self,a):
+        if isinstance(a,Part):
+            name = f"{a.name}/{self.name}"
+            a = a.data
+        else:
+            if isinstance(a,int):
+                name = f"{self.name}/{a}"
+            else:
+                name= f"{self.name}/array"
+        if np.sum(a==0)!=0:
+            log.warning("Division by zero in "+name)
+        return self.alias(data=self.data/a,
+                          name=name)
+    
+    def __getattr__(self,name):
+        name = name.casefold()
+        try:
+            return self.metadata[name]
+        except:
+            pass
+        raise AttributeError(f"Attribute {name} not found")
+    
+    def transpose(self):
+        return self.alias(data=self.data.transpose(),
+                          name=f"transposed_{self.name}",
+                          axes=self.axes[::-1])
+        
+