Better handling of floats and arrays with units (Improve `UnitsDict`)

[imaliyov]

Currently, we use the previously introduced UnitsDict class that is tailored to dictionaries (it inherits from a dict), and adds a units attribute.

It seems to work for cases like this:

        for phidx in ephmat_dat.keys():
            phdisp[phidx] = ephmat_dat[phidx].pop('phonon energy')
            defpot[phidx] = np.array(ephmat_dat[phidx].pop('deformation potential')).reshape(N, M)
            ephmat[phidx] = np.array(ephmat_dat[phidx].pop('e-ph matrix elements')).reshape(N, M)

        self.phdisp = UnitsDict.from_dict(phdisp, phdisp_units)
        self.defpot = UnitsDict.from_dict(defpot, defpot_units)
        self.ephmat = UnitsDict.from_dict(ephmat, ephmat_units)

However, we do not have (to my knowledge) a way to store floats with units. The only workaround I found, would be

        input_dict = {'pump_energy':  5}
        self.pump_energy = UnitsDict.from_dict(input_dict, units='eV')

Then printing the attribute gives you:

print(self.pump_energy)
{'pump_energy': 1.5}

So, I first use the keyword pump_energy as the attribute name, then it apprears in the dict which is definitely redundant.

Another thing is that the data attribute is no in the UnitsDict, evethough it is specified so.

[hurricane642]

Hey Ivan!
Sounds great, thank you a lot! We'll work on this!

[imaliyov]

Another thought is that handling large data with dicts is not efficient in Python. While in cases like temper, with a couple of temperatures, it is convenient, for, say any quantity in a BZ, with a million of k-points, that would be a problem.

One solution to this:

UnitsDict(dict) inherits from a dict, so we can do, say, UnitsArray(np.ndarray) with similar behavior. But then we need to ensure that all the numpy operations within an array and between arrays are supported, which is tricky.

[imaliyov]

Nevertheless, something like this would work I think:

class UnitsArray(np.ndarray):
    """
    A numpy.ndarray subclass that stores an additional 'units' attribute.
    
    Parameters
    ----------
    input_data : array_like, float, or int
        Data to be converted to a UnitsArray. If it's a float or int, 
        it will be promoted to a 0D array (scalar array).
    units : str, optional
        Physical units associated with the array contents (e.g., "eV", "K", etc.)
    dtype : str or numpy.dtype, optional
        Desired data type of the result. If None, infer from input_data.
    order : {'C', 'F'}, optional
        Memory layout of the result. 'C' means C-order, 'F' means Fortran-order.
    """

    def __new__(cls, input_data, units=None, dtype=None, order=None):
        # Convert input_data to a numpy array if it's just a float, list, etc.
        arr = np.array(input_data, dtype=dtype, order=order)
        
        # View it as a UnitsArray instead of a normal ndarray
        obj = arr.view(cls)
        
        # Attach the units as an attribute
        obj.units = units
        
        return obj

    def __array_finalize__(self, obj):
        """
        This method is called every time a new UnitsArray is created 
        (e.g., via slicing, view casting, etc.).
        """
        if obj is None:
            return
        # Copy over the units attribute from the original
        self.units = getattr(obj, 'units', None)

    def __repr__(self):
        """
        String representation that includes both the array contents and the units.
        """
        # Use the base class's __repr__ for the array part
        base_repr = super().__repr__()
        return f"{base_repr}\nUnits: {self.units}"

[imaliyov]

With some examples:

>>> arr_float = UnitsArray(3.5, units="eV")
>>> arr_float
UnitsArray(3.5)
Units: eV
>>> arr_float.shape
()   # It's a 0D array (scalar array in NumPy)
>>> arr_float + 2.0
UnitsArray(5.5)
Units: eV

from a list:

>>> arr_list = UnitsArray([1.0, 2.0, 3.0], units="K")
>>> arr_list
UnitsArray([1., 2., 3.])
Units: K
>>> arr_list.shape
(3,)
>>> arr_list + 10
UnitsArray([11., 12., 13.])
Units: K

slicing:

>>> sub_arr = arr_list[:2]
>>> sub_arr
UnitsArray([1., 2.])
Units: K
>>> sub_arr.units  # The units carried over
'K'

[imaliyov]

But we are entering dangerous grounds this direction :)

[imaliyov]

We can use an external library that handles units. Last time we looked with Shaelyn, there was no reliable and maintainable one, even though, there were some... but it will add a dependency, version control, bla bla bla

[imaliyov]

Another solution is to forget about fancy types and attributes and just store the units of important attributes in a self.units_dict, separately from the arrays or dicts themselves.

Or, we can even play with the obj.dict itself, for example, have a function set_units(att_name)
that creates a key in dict: dict['units'][attr_name] = units.
(and dict['units'] = {} would be created at the init stage)

[imaliyov]

So, I see the following options with handling the units in PerturboPy:

1. No changes: use UnitsDict whenever we can, use numpy array where we should, but this way it will be inconsistent I think...
2. Use an external units handling libraries (pros: everything is though for us already, cons: another dependency, maintainability of that library, etc.)
3. Do a similar thing as UnitsDict, but UnitsArray (see above). Pros: units are handled for numpy arrays, lists and floats, cons: maybe we will not include easily all the complexity of the numpy arrays and operations between them.
4. 'Manual' handling of units in a units_dict attribute or in dict['attr_name units']

I prefer the 4th option :)

Let's discuss!
@hurricane642 @alex2shiyu @ltan01