cdxcore 0.1.5__py3-none-any.whl

cdxcore/npio.py

@@ -0,0 +1,270 @@
+"""
+Numpy fast IO
+Hans Buehler 2023
+"""
+
+from .util import fmt_digits, fmt as txtfmt
+import numpy as np
+import numba as numba#NOQA
+import warnings as warnings
+
+def error( text, *args, exception = RuntimeError, **kwargs ):
+    raise exception( txtfmt(text, *args, **kwargs) )
+def verify( cond, text, *args, exception = RuntimeError, **kwargs ):
+    if not cond:
+        error( text, *args, **kwargs, exception=exception )
+def warn( text, *args, warning=warnings.RuntimeWarning, stack_level=1, **kwargs ):    
+    warnings.warn( txtfmt(text, *args, **kwargs), warning, stack_level=stack_level ) 
+def warn_if( cond, text, *args, warning=warnings.RuntimeWarning, stack_level=1, **kwargs ):    
+    if cond:
+        warn( text, *args, warning=warning, stack_level=stack_level, **kwargs )
+
+dtype_map = {
+        "bool"   :  0,
+        "int8"    :  1,
+        "int16"   :  2,
+        "int32"   :  3,
+        "int64"   :  4,
+        "uint16"  :  5,
+        "uint32"  :  6,
+        "uint64"  :  7,
+        "float16" :  8,
+        "float32" :  9,
+        "float64" :  10,
+        "complex64"  : 11,
+        "complex128" : 12,
+        "datetime64" : 13,
+        "timedelta64": 14
+    }
+dtype_rev = { v:k for k,v in dtype_map.items() }
+
+def _write_int(f,x,lbytes):
+    x = int(x).to_bytes(lbytes,"big")
+    w = f.write( x )
+    if w != len(x):
+        raise IOError(f"could only write {w} bytes, not {len(x)}.")
+
+def _tofile(f, array : np.ndarray, dtype_map : dict ):
+    # split into chunks
+    array    = np.asarray( array )
+    dtypec   = np.int8(dtype_map[ str(array.dtype) ] )
+    length   = np.int64( np.product( array.shape, dtype=np.uint64 ) )
+    shape32  = tuple( [np.int32(i) for i in array.shape])
+    array    = np.reshape( array, (length,) )  # this operation should not reallocate any memory
+    dsize    = int(array.itemsize)   
+    max_size = int(1024*1024*1024//dsize)
+    num      = int(length-1)//max_size+1        
+    saved    = 0
+
+    # write shape
+    _write_int( f, len(shape32), 2 )  # max 32k dimension
+    for i in shape32:
+        _write_int(f, i, 4) # max 32 bit resolution
+    # write dtype
+    _write_int(f, dtypec, 1)
+    # write object      
+    for j in range(num):
+        s   = j*max_size
+        e   = min(s+max_size, length)
+        bts = array.data[s:e]
+        nw  = f.write( bts )
+        if nw != (e-s)*dsize:
+            raise IOError(asked= (e-s)*dsize, recevied=nw)
+        saved += nw
+    if saved != length*dsize:
+        return IOError(asked=length*dsize, recevied=saved)
+
+def tofile( file,
+            array        : np.ndarray, *,
+            buffering    : int = -1
+            ):
+    """
+    Write 'array' into file using a binary format.
+    This function will work for unbuffered files exceeding 2GB which is the usual unbuffered write() limitation on Linux.
+    This function will only work with the types contained in 'dtype_map'
+    
+    Parameters
+    ----------
+        file  : file name passed to open() or an open file handle
+        array : numpy or sharedarray
+        buffering : see open(). Use 0 to turn off buffering.
+    """
+    if isinstance(file, str):
+        with open( file, "wb", buffering=buffering ) as f:
+            return tofile(f, array, buffering=buffering)
+    f = file
+    del file
+    
+    if not array.data.contiguous:
+        warn("Array is not 'contiguous'. Is that an issue??")
+        array    = np.ascontiguousarray( array, dtype=array.dtype ) if not array.data.contiguous else array
+        
+    try:
+        _tofile(f, array=array, dtype_map=dtype_map )
+    except IOError as e:
+        raise IOError(f"Could not write all {fmt_digits(array.nbytes)} bytes to {f.name}: {str(e)}", e)
+
+def _read_int(f, lbytes) -> int:
+    x = f.read(lbytes)
+    if len(x) != lbytes:
+        raise IOError(f"could only read {len(x)} bytes not {lbytes}.")
+    x = int.from_bytes(x,"big")
+    return int(x)
+
+def _readfromfile( f, array ):
+    # split into chunks
+    shape    = array.shape
+    length   = int( np.product( array.shape, dtype=np.uint64 ) )
+    array    = np.reshape( array, (length,) )
+    dsize    = int(array.itemsize)
+    max_size = int(1024*1024*1024//dsize)
+    num      = int(length-1)//max_size+1
+    read     = 0
+    # read        
+    for j in range(num):
+        s   = j*max_size
+        e   = min(s+max_size, length)
+        nr  = f.readinto( array.data[s:e] )
+        if nr != (e-s)*dsize:
+            raise IOError(f"could only read {fmt_digits(nr)} of {fmt_digits((e-s)*dsize)} bytes.")
+        read += nr
+    if read != length*dsize:
+        raise IOError(f"could only read {fmt_digits(read)} of {fmt_digits(length*dsize)} bytes.")
+    return np.reshape( array, shape )  # no copy
+
+def _readheader(f):
+    """
+    Read shape, dtype
+    """
+    shape_len  = _read_int(f,2)
+    shape      = tuple( [ int(_read_int(f,4)) for _ in range(shape_len) ] )
+    dtype      = dtype_rev[_read_int(f,1)]
+    return shape, dtype
+
+def readfromfile( file, 
+                  target         : np.ndarray, *, 
+                  read_only      : bool = False,
+                  buffering      : int  = -1,
+                  validate_dtype : type = None,
+                  validate_shape : tuple = None
+                  ) -> np.ndarray:
+    """
+    Read array from disk into an existing array or into a new array.
+    See readinto and fromfile for a simpler interface.
+    
+    Parameters
+    ----------
+        file      : file name passed to open(), or a file handle from open()
+        target    : either an array, or a function which returns an array for a given shape and dtype
+                    def create( shape ):
+                        return np.empty( shape, dtype )
+        read_only : whether to clear the 'writable' flag of the array 
+        buffering : see open(); -1 is the default, 0 for no buffering.
+        validate_dtype: if specified, check that the array has the specified dtype
+        validate_shape: if specified, check that the array has the specified shape
+        
+    Returns
+    -------
+        The array
+    """
+    if isinstance(file, str):
+        with open( file, "rb", buffering=buffering ) as f:
+            return readfromfile( f, target, 
+                                 read_only=read_only,
+                                 buffering=buffering,
+                                 validate_dtype=validate_dtype,
+                                 validate_shape=validate_shape )
+    f = file
+    del file
+        
+    # read shape
+    shape, dtype = _readheader(f)
+
+    if not validate_dtype is None and validate_dtype != dtype:
+        raise IOError(f"Failed to read {f.name}: found type {dtype} expected {validate_dtype}.")
+    if not validate_shape is None and validate_shape != shape:
+        raise IOError(f"Failed to read {f.name}: found type {shape} expected {validate_shape}.")
+
+    # handle array
+    if isinstance(target, np.ndarray):
+        if target.shape != shape or target.dtype.base != dtype:
+            raise IOError(f"File {f.name} read error: expected shape {target.shape}/{str(target.dtype)} but found {shape}/{str(dtype)}.")
+        array = target
+        
+    else:
+        array = target( shape=shape, dtype=dtype ) 
+        assert not array is None, ("'target' function returned None")
+        assert array.shape == shape and array.dtype == dtype, ("'target' function returned wrong array; shape:", array.shape, shape, "; dtype:", array.dtype, dtype)
+    del target
+
+    try:
+        _readfromfile(f, array)
+    except IOError as e:
+        raise IOError(f"Cannot read from {f.name}: {str(e)}", e)
+    if read_only:
+        array.flags.writeable  = False
+
+    assert array.flags.writeable == (not read_only), ("Internal flag error", array.flags.writeable, read_only, not read_only )
+    return array
+
+def read_shape_dtype( file, buffering : int = -1 ) -> tuple:
+    """
+    Read shape and dtype from a numpy binary file.
+    
+    Parameters
+    ----------
+        file      : file name passed to open(), or a file handle from open()
+        
+    Returns
+    -------
+        shape, dtype
+    """
+    if isinstance(file, str):
+        with open( file, "rb", buffering=buffering ) as f:
+            return read_shape_dtype( f, buffering=buffering )
+    return _readheader(file)
+
+def readinto( file, array : np.ndarray, *, read_only : bool = False, buffering : int = -1 ):
+    """
+    Read an array from disk into an existing array.    
+    The receiving array must have the same shape and dtype as the array on disk. 
+
+    Parameters
+    ----------
+        file  : file name passed to open(), or an open file
+        array : target array to write into. This array must have the same shape and dtype as the source data.
+        read_only : whether to clear the 'writable' flag of the array after the file was read
+        buffering : see open(); -1 is the default, 0 for no buffering.
+        
+    Returns
+    -------
+        The array.
+    """
+    return readfromfile( file, target = array, read_only=read_only, buffering=buffering )
+
+def fromfile( file, *, validate_dtype = None, validate_shape = None, read_only : bool = False, buffering : int = -1  ) -> np.ndarray:
+    """
+    Read array from disk into a new numpy array.
+    Use sharedarray.shared_fromfile() to create a shared array
+
+    Parameters
+    ----------
+        file     : file name passed to open(), or an open file
+        read_only: if True, clears the 'writable' flag for the returned array
+        validate_dtype: if specified, check that the array has the specified dtype
+        validate_shape: if specified, check that the array has the specified shape
+        buffering : see open(); -1 is the default, 0 for no buffering.
+
+    Returns
+    -------
+        Newly created numpy array
+    """
+    return readfromfile( file,
+                         target=lambda shape, dtype : np.empty( shape=shape, dtype=dtype ),
+                         read_only = read_only, 
+                         validate_dtype=validate_dtype, 
+                         validate_shape=validate_shape,
+                         buffering=buffering )
+        
+
+

cdxcore/prettydict.py

@@ -0,0 +1,388 @@
+"""
+prettyict
+Dictionaries with member synthax access
+Hans Buehler 2022
+"""
+
+from collections import OrderedDict
+from sortedcontainers import SortedDict
+import dataclasses as dataclasses
+from dataclasses import Field
+import types as types
+from collections.abc import Mapping
+
+class PrettyDict(dict):
+    """
+    Dictionary which allows accessing its members with member notation, e.g.        
+        pdct = PrettyDict()
+        pdct.x = 1
+        x = pdct.x
+        
+    Functions will be made members, i.e the following works as expected
+        def mult_x(self, a):
+            return self.x * a
+        pdct.mult_x = mult_x
+        pdct.mult_x(2) --> 2
+        
+    To assign a static member use []:
+        def mult(a,b):
+            return a*b
+        pdct['mult'] = mult
+        pdct.mult(1,3) --> 3
+        
+    IMPORTANT
+    Attributes starting with '__' are handled as standard attributes.
+    In other words, 
+        pdct = PrettyDict()
+        pdct.__x = 1
+        _ = pdct['__x']   <- throws an exception
+    This allows re-use of general operator handling.        
+    """
+    
+    def __getattr__(self, key : str): 
+        """ Equyivalent to self[key] """
+        if key[:2] == "__": raise AttributeError(key) # you cannot treat private members as dictionary members
+        return self[key]
+    def __delattr__(self, key : str):
+        """ Equyivalent to del self[key] """
+        if key[:2] == "__": raise AttributeError(key) # you cannot treat private members as dictionary members
+        del self[key]
+    def __setattr__(self, key : str, value):
+        """ Equivalent to self[key] = value """
+        if key[:2] == "__":
+            return dict.__setattr__(self, key, value)
+        if isinstance(value,types.FunctionType):
+            # bind function to this object
+            value = types.MethodType(value,self)
+        elif isinstance(value,types.MethodType):
+            # re-point the method to the current instance
+            value = types.MethodType(value.__func__,self)
+        self[key] = value
+    def __call__(self, key : str, *default):
+        """ Equivalent of self.get(key,default) """
+        if len(default) > 1:
+            raise NotImplementedError("Cannot pass more than one default parameter.")
+        return self.get(key,default[0]) if len(default) == 1 else  self.get(key)
+
+    def as_field(self) -> Field:
+        """
+        Returns a PrettyDictField wrapper around self for use in dataclasses
+        See PrettyDictField documentation for an example
+        """
+        return PrettyDictField(self)
+        
+class PrettyOrderedDict(OrderedDict):
+    """
+    Ordered dictionary which allows accessing its members with member notation, e.g.        
+        pdct = PrettyDict()
+        pdct.x = 1
+        x = pdct.x
+        
+    Functions will be made members, i.e the following works as expected
+        def mult_x(self, a):
+            return self.x * a
+        pdct.mult_x = mult_x
+        pdct.mult_x(2) --> 2
+        
+    To assign a static member use []:
+        def mult(a,b):
+            return a*b
+        pdct['mult'] = mult
+        pdct.mult(1,3) --> 3        
+
+    IMPORTANT
+    Attributes starting with '__' are handled as standard attributes.
+    In other words, 
+        pdct = PrettyOrderedDict()
+        pdct.__x = 1
+        _ = pdct['__x']   <- throws an exception
+    This allows re-use of general operator handling.        
+    """
+
+    def __getattr__(self, key : str):
+        """ Equyivalent to self[key] """
+        if key[:2] == "__": raise AttributeError(key) # you cannot treat private members as dictionary members
+        return self[key]
+    def __delattr__(self, key : str):
+        """ Equyivalent to del self[key] """
+        if key[:2] == "__": raise AttributeError(key) # you cannot treat private members as dictionary members
+        del self[key]
+    def __setattr__(self, key : str, value):
+        """ Equivalent to self[key] = value """
+        if key[:2] == "__":
+            return OrderedDict.__setattr__(self, key, value)
+        if isinstance(value,types.FunctionType):
+            # bind function to this object
+            value = types.MethodType(value,self)
+        elif isinstance(value,types.MethodType):
+            # re-point the method to the current instance
+            value = types.MethodType(value.__func__,self)
+        self[key] = value
+    def __call__(self, key : str, *default):
+        """ Equivalent of self.get(key,default) """
+        if len(default) > 1:
+            raise NotImplementedError("Cannot pass more than one default parameter.")
+        return self.get(key,default[0]) if len(default) == 1 else  self.get(key)
+
+    # pickling    
+    def __getstate__(self):
+        """ Return state to pickle """
+        return self.__dict__
+    def __setstate__(self, state):
+        """ Restore pickle """
+        self.__dict__.update(state)
+    
+    def as_field(self) -> Field:
+        """
+        Returns a PrettyDictField wrapper around 'self' for use in dataclasses
+        See PrettyDictField documentation for an example
+        """
+        return PrettyDictField(self)
+
+    @property
+    def at_pos(self):
+        """
+        Element access
+        
+        at_pos[position] returns an element or elements at an ordinal position.
+            It returns a single element if 'position' refers to only one field.
+            If 'position' is a slice then the respecitve list of fields is returned
+            
+        at_pos[position] = item assigns an item or an ordinal position
+            If 'position' refers to a single element, 'item' must be that item
+            If 'position' is a slice then 'item' must resolve to a list of the required size.
+            
+        Key access
+        
+        at_pos.keys[position] returns the key or keys at 'position'
+        
+        at_pos.items[position] returns the tuple (key, element) or a list thereof for `position`            
+        """
+        class Access:
+            """ 
+            Wrapper object to allow index access for at_pos
+            """
+            def __init__(self):
+                self.__keys = None
+            
+            def __getitem__(_, position):
+                key = _.keys[position]
+                return self[key] if not isinstance(key,list) else [ self[k] for k in key ]
+            def __setitem__(_, position, item ):
+                key = _.keys[position]
+                if not isinstance(key,list):
+                    self[key] = item
+                else:
+                    for k, i in zip(key, item):
+                        self[k] = i
+            @property
+            def keys(_) -> list:
+                """ Returns the list of keys of the original dictionary """
+                if _.__keys is None:
+                    _.__keys = list(self.keys())
+                return _.__keys
+            @property
+            def items(_) -> list:
+                """ Returns the list of keys of the original dictionary """
+                class ItemAccess(object):
+                    def __getitem__(_x, position):
+                        key = _.keys[position]
+                        return (key, self[key]) if not isinstance(key,list) else [ (k,self[k]) for k in key ]                
+                return ItemAccess()
+                
+        return Access()
+
+    @property
+    def key_at_pos(self) -> list:
+        """
+        key_at_pos returns the list of keys, hence key_at_pos[i] returns the ith key
+        """
+        return list(self)
+
+    @property
+    def item_at_pos(self) -> list:
+        """
+        key_at_pos returns the list of keys, hence key_at_pos[i] returns the ith key
+        """
+        return list(self)
+    
+pdct = PrettyOrderedDict
+    
+class BETA_PrettySortedDict(SortedDict):
+    """
+    *NOT WORKING WELL*
+    Sorted dictionary which allows accessing its members with member notation, e.g.        
+        pdct = PrettyDict()
+        pdct.x = 1
+        x = pdct.x
+        
+    Functions will be made members, i.e the following works as expected
+        def mult_x(self, a):
+            return self.x * a
+        pdct.mult_x = mult_x
+        pdct.mult_x(2) --> 2
+        
+    To assign a static member use []:
+        def mult(a,b):
+            return a*b
+        pdct['mult'] = mult
+        pdct.mult(1,3) --> 3       
+        
+    IMPORTANT
+    Attributes starting with '_' (one underscore) are handled as standard attributes.
+    In other words, 
+        pdct = PrettyOrderedDict()
+        pdct._x = 1
+        _ = pdct['_x']   <- throws an exception
+    This allows re-use of general operator handling.   
+    The reason the sorted class disallow '_' (as opposed to the other two classes who merely disallow '__')     
+    is that SortedDict() uses protected members.
+    """
+
+    def __getattr__(self, key : str):
+        """ Equyivalent to self[key] """
+        if key[:1] == "_": raise AttributeError(key) # you cannot treat protected or private members as dictionary members
+        return self[key]
+    def __delattr__(self, key : str):
+        """ Equyivalent to del self[key] """
+        if key[:2] == "__": raise AttributeError(key) # you cannot treat private members as dictionary members
+        del self[key]
+    def __setattr__(self, key : str, value):
+        """ Equivalent to self[key] = value """
+        if key[:1] == "_":
+            return SortedDict.__setattr__(self, key, value)
+        if isinstance(value,types.FunctionType):
+            # bind function to this object
+            value = types.MethodType(value,self)
+        elif isinstance(value,types.MethodType):
+            # re-point the method to the current instance
+            value = types.MethodType(value.__func__,self)
+        self[key] = value
+    def __call__(self, key : str, *default):
+        """ Equivalent of self.get(key,default) """
+        if len(default) > 1:
+            raise NotImplementedError("Cannot pass more than one default parameter.")
+        return self.get(key,default[0]) if len(default) == 1 else  self.get(key)
+    
+class PrettyDictField(object):
+    """
+    Simplististc 'read only' wrapper for PrettyOrderedDict objects.
+    Useful for Flax
+
+        import dataclasses as dataclasses
+        import jax.numpy as jnp
+        import jax as jax
+        from options.cdxbasics.config import Config, ConfigField
+        import types as types
+        
+        class A( nn.Module ):
+            pdct : PrettyOrderedDictField = PrettyOrderedDictField.Field()
+        
+            def setup(self):
+                self.dense = nn.Dense(1)
+        
+            def __call__(self, x):
+                a = self.pdct.a
+                return self.dense(x)*a
+        
+        r = PrettyOrderedDict(a=1.)
+        a = A( r.as_field() )
+        
+        key1, key2 = jax.random.split(jax.random.key(0))
+        x = jnp.zeros((10,10))
+        param = a.init( key1, x )
+        y = a.apply( param, x )
+        
+    The class will traverse pretty dictionaries of pretty dictionaries correctly.
+    However, it has some limitations as it does not handle custom lists of pretty dicts.
+    """
+    def __init__(self, pdct : PrettyOrderedDict = None, **kwargs):
+        """ Initialize with an input dictionary and potential overwrites """
+        if not pdct is None:
+            if type(pdct).__name__ == type(self).__name__ and len(kwargs) == 0:
+                # copy
+                self.__pdct = PrettyOrderedDict( pdct.__pdct )
+                return
+            if not isinstance(pdct, Mapping): raise ValueError("'pdct' must be a Mapping")
+            self.__pdct = PrettyOrderedDict(pdct)
+            self.__pdct.update(kwargs)
+        else:
+            self.__pdct = PrettyOrderedDict(**kwargs) 
+        def rec(x):
+            for k, v in x.items():
+                if isinstance(v, (PrettyDict, PrettyOrderedDict)):
+                    x[k] = PrettyDictField(v)
+                elif isinstance(v, Mapping):
+                    rec(v)
+        rec(self.__pdct)
+            
+    def as_dict(self) -> PrettyOrderedDict:
+        """ Return copy of underlying dictionary """
+        return PrettyOrderedDict( self.__pdct )
+
+    def as_field(self) -> Field:
+        """
+        Returns a PrettyDictField wrapper around self for use in dataclasse
+        This function makes a (shallow enough) copy of the current field.
+        It is present so iterative applications of as_field() are convenient.        
+        """
+        return PrettyDictField(self)
+            
+    @staticmethod
+    def default():
+        return PrettyDictField()
+    
+    @staticmethod
+    def Field( default : PrettyOrderedDict = None, **kwargs):
+        """
+        Returns a dataclasses.field for PrettyDictField
+        """
+        if default is None and len(kwargs) == 0:
+            return dataclasses.field( default_factory=PrettyDictField )
+        
+        if not default is None:
+            default.upate(kwargs)
+        else:
+            default = kwargs
+        def factory():
+            return PrettyDictField(default)
+        return dataclasses.field( default_factory=factory )
+
+    # mimic the underlying dictionary
+    # -------------------------------
+    
+    def __getattr__(self, key):
+        if key[:2] == "__":
+            return object.__getattr__(self,key)
+        return self.__pdct.__getattr__(key)
+    def __getitem__(self, key):
+        return self.__pdct[key]
+    def __call__(self, *kargs, **kwargs):
+        return self.__pdct(*kargs, **kwargs)
+    def __eq__(self, other):
+        if type(other).__name__ == "PrettyOrderedDict":
+            return self.__pdct == other
+        else:
+            return self.__pdct == other.pdct
+    def keys(self):
+        return self.__pdct.keys()
+    def items(self):
+        return self.__pdct.items()
+    def values(self):
+        return self.__pdct.values()
+    def __hash__(self):
+        h = 0
+        for k, v in self.items():
+            h ^= hash(k) ^ hash(v)
+        return h
+    def __iter__(self):
+        return self.__pdct.__iter__()
+    def __contains__(self, key):
+        return self.__pdct.__contains__(key)
+    def __len__(self):
+        return self.__pdct.__len__()
+    def __str__(self):
+        return self.__pdct.__str__()
+    def __repr__(self):
+        return self.__pdct.__repr__()
+