Source code for wholecell.utils.unit_struct_array

"""
UnitStructArray

Wraps Numpy struct arrays using Unum units. Will assure that correct units are
being used while loading constants.

TODO: Unum is a defunct project. Its source repo is no longer online. Either
switch to a newer package like Pint or copy and improve the Unum source code
from its Python package.
"""

import numpy as np
import numpy.typing as npt
import unum
from typing import Optional

from wholecell.utils import units as units_pkg


[docs] class UnitStructArray(object): """Wraps Numpy structured arrays using Unum units. Will assure that correct units are being used while loading constants. """ def __init__(self, struct_array: npt.ArrayLike, units: dict[str, Optional[str]]): self._validate(struct_array, units) self.struct_array = struct_array self.units = units
[docs] def _validate(self, struct_array, units): s = "" if not isinstance(struct_array, np.ndarray): s += "UnitStructArray must be initialized with a numpy array!\n" elif not isinstance(units, dict): s += "UnitStructArray must be initialized with a dict storing units!\n" elif set([x[0] for x in struct_array.dtype.descr]) != set(units.keys()): s += "Struct array fields do not match unit fields!\n" if len(s): raise Exception(s)
[docs] def _field(self, fieldname): if not units_pkg.hasUnit(self.units[fieldname]): if self.units[fieldname] is None: return self.struct_array[fieldname] else: raise Exception("Field has incorrect units or unitless designation!\n") else: return self.units[fieldname] * self.struct_array[fieldname]
[docs] def fullArray(self): return self.struct_array
[docs] def fullUnits(self): return self.units
def __getitem__(self, key): if isinstance(key, slice): return UnitStructArray(self.struct_array[key], self.units) elif isinstance(key, np.ndarray) or isinstance(key, list): return UnitStructArray(self.struct_array[key], self.units) elif isinstance(key, int): return self.struct_array[key] else: return self._field(key) def __setitem__(self, key, value): if units_pkg.hasUnit(value): try: self.units[key].matchUnits(value) except unum.IncompatibleUnitsError: raise Exception("Units do not match!\n") self.struct_array[key] = value.asNumber() self.units[key] = units_pkg.getUnit(value) elif isinstance(value, list) or isinstance(value, np.ndarray): if units_pkg.hasUnit(self.units[key]): raise Exception( "Units do not match! Quantity has units your input does not!\n" ) self.struct_array[key] = value self.units[key] = None else: raise Exception( "Can't assign data-type other than unum datatype or list/numpy array!\n" ) def __len__(self): return len(self.struct_array) def __repr__(self): return "STRUCTURED ARRAY:\n{}\nUNITS:\n{}".format( self.struct_array.__repr__(), self.units ) def __eq__(self, other): if type(other) != type(self): return False elif self.struct_array.dtype != other.struct_array.dtype: return False elif not all(self.struct_array == other.struct_array): return False elif self.units != other.units: return False else: return True def __ne__(self, other): return not self.__eq__(other)