mrio-toolbox 1.0.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

{mrio_toolbox/_parts → _parts}/_Part.py

@@ -10,6 +10,7 @@ import itertools
 import numpy as np
 import pandas as pd
 import xarray as xr
+import copy
 from mrio_toolbox._parts._Axe import Axe
 import logging
 from mrio_toolbox.utils import converters
@@ -26,46 +27,151 @@ def load_part(
     return Part(**loader.load_part(**kwargs))
 
 class Part:
+    """
+    Representation of an MRIO Part 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.
+
+    Instance variables
+    ------------------
+    data : numpy.ndarray
+        Numerical data of the Part.
+    axes : list of Axe instances
+        Axes corresponding to the dimensions of the Part.
+    groupings : dict
+        Groupings of the labels of the Part, for each label defined.
+    metadata : dict
+        Additional metadata of the Part (e.g., path, name, multiplier, unit).
+    name : str
+        Name of the Part.
+    ndim : int
+        Number of dimensions of the Part.
+    shape : tuple
+        Shape of the Part.
+
+    Methods
+    -------
+    __init__(data=None, labels=None, axes=None, **kwargs):
+        Initialize a Part object.
+    alias(**kwargs):
+        Create a new Part with modified parameters.
+    fix_dims(skip_labels=False, skip_data=False):
+        Align the number of axes with the number of dimensions.
+    get(*args, aspart=True, squeeze=False):
+        Extract data from the Part object.
+    setter(value, *args):
+        Change the value of a data selection.
+    develop(axis=None, on=None, squeeze=True):
+        Reshape a Part to avoid double labels.
+    reformat(new_dimensions):
+        Reshape a Part to match a new dimensions combination.
+    combine_axes(start=0, end=None, in_place=False):
+        Combine axes of a Part into a single one.
+    swap_axes(axis1, axis2):
+        Swap two axes of a Part.
+    swap_ax_levels(axis, dim1, dim2):
+        Swap two levels of an axis.
+    flatten(invert=False):
+        Flatten a 2D Part into a 1D Part.
+    squeeze():
+        Remove dimensions of length 1 from the Part.
+    expand_dims(axis, copy=None):
+        Add dimensions to a Part instance.
+    copy():
+        Return a copy of the current Part object.
+    extraction(dimensions, labels=["all"], on_groupings=True, domestic_only=False, axis="all"):
+        Set labels over dimension(s) to 0.
+    leontief_inversion():
+        Compute the Leontief inverse of a square Part.
+    update_groupings(groupings, ax=None):
+        Update the groupings of the current Part object.
+    aggregate(on="countries", axis=None):
+        Aggregate dimensions along one or several axes.
+    aggregate_on(on, axis):
+        Aggregate a Part along a given axis.
+    get_labels(axis=None):
+        Returns a list with the labels of the axes as dictionaries.
+    list_labels():
+        List the labels of the Part.
+    get_dimensions(axis=None):
+        Return the list of dimensions of the Part.
+    rename_labels(old, new):
+        Rename some labels of the Part.
+    replace_labels(name, labels, axis=None):
+        Update a label of the Part.
+    set_labels(labels, axis=None):
+        Change the labels of the Part.
+    add_labels(labels, dimension=None, axes=None, fill_value=0):
+        Add indices to one or multiple Part axes.
+    expand(axis=None, over="countries"):
+        Expand an axis of the Part.
+    issquare():
+        Check whether the Part is square.
+    hasneg():
+        Check whether the Part has negative elements.
+    hasax(name=None):
+        Return the dimensions along which a Part has given labels.
+    sum(axis=None, on=None, keepdims=False):
+        Sum the Part along one or several axes or on a given dimension.
+    save(file=None, name=None, extension=".npy", overwrite=False, include_labels=False, write_instructions=False, **kwargs):
+        Save the Part object to a file.
+    to_pandas():
+        Return the current Part object as a Pandas DataFrame.
+    to_xarray():
+        Save the Part object to an xarray DataArray.
+    mean(axis=None):
+        Compute the mean of the Part along a given axis.
+    min(axis=None):
+        Compute the minimum value of the Part along a given axis.
+    max(axis=None):
+        Compute the maximum value of the Part along a given axis.
+    mul(a, propagate_labels=True):
+        Perform matrix multiplication between Parts with label propagation.
+    filter(threshold, fill_value=0):
+        Set to 0 the values below a given threshold.
+    diag():
+        Create a diagonal Part from a 1D Part or extract the diagonal of a 2D Part.
+    transpose():
+        Transpose the Part object.
+    """
+    
     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.
+        """
+        Initialize a Part object.
 
         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.
+        data : numpy.ndarray, optional
+            Numerical data of the Part. If not provided, a Part filled with 
+            zeros (or another fill value) is created based on the shape of the axes.
         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.
+            Labels for the axes. If provided, the labels define the structure 
+            of the axes. If not provided, axes are created based on the data.
         axes : list of Axe instances, optional
-            Custom Axes for the Part.
-            If left empty, the axes are created from the labels.
+            Custom Axes for the Part. If not provided, axes are created from 
+            the labels or inferred from the data.
         kwargs : dict
-            Additional metadata of the Part.
-            (e.g. path, name, multiplier, unit...)
+            Additional metadata for the Part (e.g., path, name, multiplier, unit).
 
-        Returns
-        -------
-        None.
+        Raises
+        ------
+        ValueError
+            If the length of the labels does not match the data dimensions.
+        TypeError
+            If the provided labels are of an unsupported type.
 
+        Notes
+        -----
+        If both `data` and `axes` are not provided, the method creates an 
+        empty Part with default axes and zero-filled data.
         """
         
         if data is not None:
@@ -128,6 +234,7 @@ class Part:
                     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
@@ -257,13 +364,15 @@ class Part:
         
 
     def __getitem__(self,args):
-        if isinstance(args,str) or isinstance(args,int) or isinstance(args,np.integer):
+        if isinstance(args,str) or isinstance(args,int) or isinstance(args,np.integer) or isinstance(args,dict):
             args = (args,)
         return self.get(*args)
     
     def __setitem__(self,args,value):
         if isinstance(value,Part):
             value = value.data
+        if isinstance(args,str) or isinstance(args,int) or isinstance(args,np.integer) or isinstance(args,dict):
+            args = (args,)
         self.setter(value,*args)
     
     def setter(self,value,*args):
@@ -319,7 +428,6 @@ class Part:
         Returns
         -------
         New Part object or numpy object
-
         """
         sels = []
         axes = []
@@ -379,12 +487,11 @@ class Part:
         squeeze : bool, optional
             Whether to remove dimensions of length 1.
             The default is True.
-        
+            
         Returns
         -------
-        Part object
-            Developped Part
-
+        Developped Part :  Part object
+            The developed part
         """
         if isinstance(on,str):
             on = [on]
@@ -414,6 +521,7 @@ class Part:
                     "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:
@@ -423,35 +531,42 @@ class Part:
         return Part(data=data,name=f"developped_{self.name}",
                     groupings=self.groupings,axes=axes)
 
-    def reformat(self,new_dimensions):
+    def reformat(self, new_dimensions):
         """
-        Reshape a Part to match a new dimensions combination
+        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
+        new_dimensions : list of list of str
+            Target dimensions to reshape into.
 
         Returns
         -------
-        data : numpy array
-            Reshaped data
-        axes : list of Axe instances
-            Reshaped axes
+        data : numpy.ndarray
+            Reshaped data.
+        axes : list of Axe
+            Reshaped axes.
+
+        Examples
+        --------
+        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"]]
         """
         return part_operations.reformat(self,new_dimensions)
     
@@ -561,40 +676,44 @@ class Part:
         dev.name = f"swapped_{self.name}"
         return dev
     
-    def flatten(self,invert=False):
-        """Flatten a 2D Part into a 1D Part
+    def flatten(self):
+        """Flatten a multidimensional Part into a 1D Part
         
-        Parameters
-        ----------
-        inverse : bool, optional
-            Whether to in the inverse level order.
+        Because Parts do not support repeated dimensions over the same axis,
+        Axis dimensions are disambiguated if needed by appending an index to the dimension name.
+
         """
-        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)
+        def disambiguous_dimension(dim,existing,i=1):
+            while f"{dim}_{i}" in existing:
+                return disambiguous_dimension(dim,existing,i+1)
+            
+            return f"{dim}_{i}"
+        
+        enumerator = range(self.ndim)
+        order = "C"
+        labels = dict()
+        groupings = dict()
+        for ax in enumerator:
+            for dim in self.axes[ax].dimensions:
+                dim_name = dim
+                if dim in labels.keys():
+                    dim_name = disambiguous_dimension(dim,labels)
+                    log.info(f"Disambiguate dimension {dim} on axis {ax} into {dim_name}")
+                labels[dim_name] = self.axes[ax].labels[dim]
+                if dim in self.groupings.keys():
+                    groupings[dim_name] = self.axes[ax].groupings[dim]
+        ax = Axe(labels,groupings)
         return self.alias(data=self.data.flatten(order=order),
                             name=f"flattened_{self.name}",
                             axes =[ax])
         
         
-    def squeeze(self):
+    def squeeze(self,drop_ax=True,drop_dims=True):
         axes = []
         for ax in self.axes:
-            if len(ax) > 1:
+            if drop_dims:
+                ax.squeeze()
+            if len(ax) > 1 or not drop_ax:
                 axes.append(ax)
         return self.alias(data=np.squeeze(self.data),axes=axes,
                           name=f"squeezed_{self.name}")
@@ -661,14 +780,6 @@ class Part:
             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())
@@ -677,14 +788,30 @@ class Part:
             to_select = dict()
             for dim,label in zip(dimensions,labels):
                 to_select[dim] = label
+        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")
 
         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 len(dimensions) == 1:
+                log.critical("No axis found for extraction on "+str(dimensions))
+                raise ValueError("No axis found for extraction on "+str(dimensions))
+            log.info(f"No axis found for simultaneous extractions on {dimensions}")
+            log.info(f"Try successive extractions on {dimensions}")
+            for dim,label in zip(dimensions,labels):
+                self.extraction(dim,label,
+                                on_groupings=on_groupings,
+                                domestic_only=domestic_only,
+                                axis=axis)
         if axis == "all":
             log.info(f"Extract {to_select} on axes "+ str(allowed))
             axis = allowed
@@ -785,13 +912,14 @@ class Part:
         for axe in list(ax):
             self.axes[axe].update_groupings(groupings)
     
-    def aggregate(self,on="countries",axis=None):
+    def aggregate(self,on=None,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.
+        If no dimension is specified, the operation is applied to all possible dimensions.
 
         Parameters
         ----------
@@ -816,11 +944,14 @@ class Part:
 
         """
         log.debug(f"Aggregate Part {self.name} along axis {axis} on {on}")
+
+        if on is None:
+            on = list(self.groupings.keys())
         
         
         if isinstance(on,list):
             for item in on:
-                self = self.aggregate(axis,item)
+                self = self.aggregate(on = item, axis=axis)
             return self
         if on not in self.groupings.keys():
             raise ValueError(f"No groupings defined for dimensions {on}")
@@ -871,8 +1002,8 @@ class Part:
 
         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)
+        ref_dev = self.develop(axis, squeeze=False)
+        new_dev = output.develop(axis,squeeze=False)
         selector = ["all"]*ref_dev.ndim
         for label in new_labels[on]:
             selector[axis+idsum] = label
@@ -887,8 +1018,9 @@ class Part:
  
     def get_labels(self,axis=None):
         """
-        Returns the dictionnary of the Part labels
-
+        Returns a list with the labels of each axis 
+        of the part in a the dictionary.
+        
         Parameters
         ----------
         axis : int or list of int, optional
@@ -1035,12 +1167,12 @@ class Part:
             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.
@@ -1073,6 +1205,67 @@ class Part:
         output[sel] = self.data
         return output
     
+    def reorder_data(self,new_labels):
+        """
+        Reorder the data of the Part according to new labels.
+
+        Parameters
+        ----------
+        new_labels : dict
+            New labels for the axes.
+            The keys are the dimensions, the values are the labels.
+
+        Raises
+        ------
+        ValueError
+            If the new labels do not match the current axes.
+        """
+        
+        if not isinstance(new_labels,dict):
+            raise ValueError("New labels should be a dictionary")
+        
+        
+        
+        sels = []
+        for axis in self.axes:
+            old_labels = axis.labels
+            if not set(new_labels.keys()).issubset(set(old_labels.keys())):
+                sels.append(axis.get("all"))
+                continue
+            for key in new_labels.keys():
+                set_old = set(old_labels[key])
+                set_new = set(new_labels[key])
+                if not set_old.issubset(set_new):
+                    raise ValueError(f"The new labels provided for dimension '{key}' is not a superset of the old labels. " +
+                                    f"Old labels: {old_labels[key]}, new labels: {new_labels[key]}. "
+                                    "If you want to rename the labels of this dimensions, use the method 'replace_labels() before reordering the data")
+
+
+            ax_label_dict = {}
+            for key in old_labels.keys():
+                if key in new_labels.keys():
+                    ax_label_dict[key] = new_labels[key]
+                    for lab in new_labels[key]: 
+                        if lab not in old_labels[key]:
+                            # If the label is not in the list, remove it
+                            ax_label_dict[key].remove(lab)
+                else:
+                    ax_label_dict[key] = old_labels[key]
+ 
+            sels.append(axis.get(ax_label_dict))
+        
+        if len(sels) == 0:
+            raise ValueError(f"None of the dimensions provided in the new labels dict {new_labels.keys()} are present "+
+                             f"in the labels of part '{self.name}', which only contains the dimensions {self.get_dimensions()}")
+        
+        #Execute the selection
+        self.data = self.data[np.ix_(*sels)]
+        
+        # Update the axes with the new labels
+        for dim in new_labels.keys():
+            self.replace_labels(name = dim, labels = new_labels[dim])
+     
+    
     def expand(self,axis=None,over="countries"):
         """
         Expand an axis of the Part
@@ -1098,23 +1291,36 @@ class Part:
         """
         if axis is None:
             axis = self.hasax(over)
+        if isinstance(axis,int):
+            axis = [axis]
+
+        output = self.copy()
         
-        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)]
+        for ax in axis[::-1]:
+            output = output._expand_on(ax,over)
+        return output
+
+    def _expand_on(self,ax,over):
+        """Expand over a single axis"""
+        ref_ax = self.axes[ax]
+        new_ax = Axe({over: ref_ax.labels[over]}, groupings=self.groupings)
+        axes = self.axes.copy()
+        axes.insert(ax,new_ax)
+        new_shape = list(self.shape)
+        new_shape.insert(ax,len(new_ax))
+        output = np.zeros(new_shape)
+        selector = [slice(None)]*self.ndim
+        ordering = ref_ax.dimensions.index(over)
+        for item in ref_ax.labels[over]:
+            newsel,refsel = selector.copy(),selector.copy()
+            newsel.insert(ax,new_ax.get(item))
+            ax_selector = ["all" for i in range(ordering)]
+            ax_selector.append(item)
+            newsel[ax+1] = ref_ax.get(ax_selector)
+            refsel[ax]= ref_ax.get(ax_selector)
+            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"""
@@ -1212,10 +1418,13 @@ class Part:
             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
+            axes = self.axes.copy()
+            if not keepdims:
+                del axes[axis]
             return self.alias(
                 data = self.data.sum(axis,keepdims=keepdims),
                 name=f"{self.name}_sum_{axis}",
-                axes = self.axes
+                axes = axes
             )
         #Otherwise, sum on the relevant levels
         idsum = ax.dimensions.index(on) #Index of the dimension to sum on
@@ -1280,7 +1489,7 @@ class Part:
         if path is None:
             raise FileNotFoundError("No path specified for saving the Part")
         if extension == ".nc":
-            path = os.path.join(path,name)
+            path = os.path.join(path,name+extension)
             save_to_nc(self,path,overwrite,
                        write_instructions=write_instructions,
                        **kwargs)
@@ -1385,10 +1594,16 @@ class Part:
     
     def diag(self):
         if self.ndim == 1:
+            log.info("Diagonalize a 1D part")
             return self.alias(data=np.diag(self.data),
                               name=f"diag_{self.name}",
                               axes = self.axes*2)
-                              
+        try:
+            log.info("The part has too many dimensions: try to diagonalize the squeezed part")
+            return self.squeeze().diag()
+        except:
+            raise ValueError("Cannot diagonalize a part with more than 2 dimensions")
+
     def __add__(self,a):
         if isinstance(a,Part):
             name = a.name
@@ -1426,6 +1641,26 @@ class Part:
     def __neg__(self):
         return self.alias(data=-self.data,name=f"-{self.name}")
     
+    def __lt__(self,other):
+        if isinstance(other,Part):
+            return self.data < other.data
+        return self.data < other 
+    
+    def __le__(self,other):
+        if isinstance(other,Part):
+            return self.data <= other.data
+        return self.data <= other
+    
+    def __gt__(self,other):
+        if isinstance(other,Part):
+            return self.data > other.data
+        return self.data > other
+    
+    def __ge__(self,other):
+        if isinstance(other,Part):
+            return self.data >= other.data
+        return self.data >= other
+    
     def __sub__(self,a):
         if isinstance(a,Part):
             name = a.name

_parts/__init__.py

@@ -0,0 +1,7 @@
+"""
+This module provides the Part and Axe classes.
+"""
+
+from ._Part import Part,load_part
+
+__all__ = ['Part','load_part']