turbx 1.0.2__py3-none-any.whl

turbx/grid_metric.py

@@ -0,0 +1,272 @@
+import timeit
+
+import numpy as np
+from tqdm import tqdm
+
+from .gradient import gradient
+from .utils import even_print
+
+# ======================================================================
+
+def get_metric_tensor_2d(x2d, y2d, acc=2, edge_stencil='full', **kwargs):
+    '''
+    Compute the grid metric tensor (inverse of grid Jacobian) for a 2D grid
+    -----
+    Computational Fluid Mechanics and Heat Transfer (2012) Pletcher, Tannehill, Anderson
+    p.266-270, 335-337, 652
+    '''
+    
+    verbose = kwargs.get('verbose',False)
+    no_warn = kwargs.get('no_warn',False)
+    
+    if not isinstance(x2d, np.ndarray):
+        raise ValueError('x2d should be of type np.ndarray')
+    if not isinstance(y2d, np.ndarray):
+        raise ValueError('y2d should be of type np.ndarray')
+    
+    if (x2d.ndim!=2):
+        raise ValueError('x2d should have ndim=2 (xy)')
+    if (y2d.ndim!=2):
+        raise ValueError('y2d should have ndim=2 (xy)')
+    
+    if not (x2d.shape==y2d.shape):
+        raise ValueError('x2d.shape!=y2d.shape')
+    
+    nx,ny = x2d.shape
+    
+    ## the 'computational' grid (unit Cartesian)
+    ## --> [x_comp,y_comp]= [ξ,η] = [q1,q2]
+    #x_comp = np.arange(nx, dtype=np.float64)
+    #y_comp = np.arange(ny, dtype=np.float64)
+    x_comp = 1.
+    y_comp = 1.
+    
+    # === get Jacobian :: ∂(x,y)/∂(q1,q2)
+    
+    t_start = timeit.default_timer()
+    
+    dxdx = gradient(x2d, x_comp, axis=0, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dydx = gradient(y2d, x_comp, axis=0, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dxdy = gradient(x2d, y_comp, axis=1, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dydy = gradient(y2d, y_comp, axis=1, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    
+    J = np.stack((np.stack((dxdx, dydx), axis=2),
+                  np.stack((dxdy, dydy), axis=2)), axis=3)
+    
+    t_delta = timeit.default_timer() - t_start
+    if verbose: tqdm.write( even_print('get J','%0.3f [s]'%(t_delta,), s=True) )
+    
+    # === get metric tensor M = J^-1 = ∂(q1,q2)/∂(x,y) = ∂(ξ,η)/∂(x,y)
+    
+    if False: ## method 1
+        
+        t_start = timeit.default_timer()
+        
+        M = np.linalg.inv(J)
+        
+        # M_bak = np.copy(M)
+        # M = np.zeros((nx,ny,2,2),dtype=np.float64)
+        # for i in range(nx):
+        #     for j in range(ny):
+        #         M[i,j,:,:] = sp.linalg.inv( J[i,j,:,:] )
+        # np.testing.assert_allclose(M_bak, M, atol=1e-12, rtol=1e-12)
+        # print('check passed')
+        
+        t_delta = timeit.default_timer() - t_start
+        if verbose: tqdm.write( even_print('get M','%0.3f [s]'%(t_delta,), s=True) )
+    
+    if True: ## method 2
+        
+        if ('M' in locals()):
+            M_bak = np.copy(M)
+            M = None; del M
+        
+        t_start = timeit.default_timer()
+        
+        ## Jacobian determinant
+        Jac_det = dxdx*dydy - dydx*dxdy
+        
+        # Jac_det_bak = np.copy(Jac_det)
+        # Jac_det = None; del Jac_det
+        # Jac_det = np.linalg.det(J)
+        # np.testing.assert_allclose(Jac_det, Jac_det_bak, atol=1e-14, rtol=1e-14)
+        # print('check passed')
+        
+        M = np.zeros((nx,ny,2,2), dtype=np.float64)
+        M[:,:,0,0] = +dydy / Jac_det ## ξ_x
+        M[:,:,0,1] = -dxdy / Jac_det ## ξ_y
+        M[:,:,1,0] = -dydx / Jac_det ## η_x
+        M[:,:,1,1] = +dxdx / Jac_det ## η_y
+        
+        t_delta = timeit.default_timer() - t_start
+        if verbose: tqdm.write( even_print('get M','%0.3f [s]'%(t_delta,), s=True) )
+        
+        if ('M_bak' in locals()):
+            np.testing.assert_allclose(M[:,:,0,0], M_bak[:,:,0,0], atol=1e-14, rtol=1e-14)
+            print('check passed: ξ_x')
+            np.testing.assert_allclose(M[:,:,0,1], M_bak[:,:,0,1], atol=1e-14, rtol=1e-14)
+            print('check passed: ξ_y')
+            np.testing.assert_allclose(M[:,:,1,0], M_bak[:,:,1,0], atol=1e-14, rtol=1e-14)
+            print('check passed: η_x')
+            np.testing.assert_allclose(M[:,:,1,1], M_bak[:,:,1,1], atol=1e-14, rtol=1e-14)
+            print('check passed: η_y')
+            np.testing.assert_allclose(M, M_bak, atol=1e-14, rtol=1e-14)
+            print('check passed: M')
+    
+    return M
+
+def get_metric_tensor_3d(x3d, y3d, z3d, acc=2, edge_stencil='full', **kwargs):
+    '''
+    Compute the grid metric tensor (inverse of grid Jacobian) for a 3D grid
+    -----
+    Computational Fluid Mechanics and Heat Transfer (2012) Pletcher, Tannehill, Anderson
+    p.266-270, 335-337, 652
+    '''
+    
+    verbose = kwargs.get('verbose',False)
+    no_warn = kwargs.get('no_warn',False)
+    
+    if not isinstance(x3d, np.ndarray):
+        raise ValueError('x3d should be of type np.ndarray')
+    if not isinstance(y3d, np.ndarray):
+        raise ValueError('y3d should be of type np.ndarray')
+    if not isinstance(z3d, np.ndarray):
+        raise ValueError('z3d should be of type np.ndarray')
+    
+    if (x3d.ndim!=3):
+        raise ValueError('x3d should have ndim=3 (xyz)')
+    if (y3d.ndim!=3):
+        raise ValueError('y3d should have ndim=3 (xyz)')
+    if (z3d.ndim!=3):
+        raise ValueError('z3d should have ndim=3 (xyz)')
+    
+    if not (x3d.shape==y3d.shape):
+        raise ValueError('x3d.shape!=y3d.shape')
+    if not (y3d.shape==z3d.shape):
+        raise ValueError('y3d.shape!=z3d.shape')
+    
+    nx,ny,nz = x3d.shape
+    
+    ## the 'computational' grid (unit Cartesian)
+    ## --> [x_comp,y_comp,z_comp ]= [ξ,η,ζ] = [q1,q2,q3]
+    #x_comp = np.arange(nx, dtype=np.float64)
+    #y_comp = np.arange(ny, dtype=np.float64)
+    #z_comp = np.arange(nz, dtype=np.float64)
+    x_comp = 1.
+    y_comp = 1.
+    z_comp = 1.
+    
+    # === get Jacobian :: ∂(x,y,z)/∂(q1,q2,q3)
+    
+    t_start = timeit.default_timer()
+    
+    dxdx = gradient(x3d, x_comp, axis=0, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dydx = gradient(y3d, x_comp, axis=0, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dzdx = gradient(z3d, x_comp, axis=0, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    
+    dxdy = gradient(x3d, y_comp, axis=1, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dydy = gradient(y3d, y_comp, axis=1, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dzdy = gradient(z3d, y_comp, axis=1, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    
+    dxdz = gradient(x3d, z_comp, axis=2, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dydz = gradient(y3d, z_comp, axis=2, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    dzdz = gradient(z3d, z_comp, axis=2, d=1, acc=acc, edge_stencil=edge_stencil, no_warn=no_warn)
+    
+    J = np.stack((np.stack((dxdx, dydx, dzdx), axis=3),
+                  np.stack((dxdy, dydy, dzdy), axis=3),
+                  np.stack((dxdz, dydz, dzdz), axis=3)), axis=4)
+    
+    t_delta = timeit.default_timer() - t_start
+    if verbose: tqdm.write( even_print('get J','%0.3f [s]'%(t_delta,), s=True) )
+    
+    # === get metric tensor M = J^-1 = ∂(q1,q2,q3)/∂(x,y,z) = ∂(ξ,η,ζ)/∂(x,y,z)
+    
+    if False: ## method 1
+        
+        t_start = timeit.default_timer()
+        
+        M = np.linalg.inv(J)
+        
+        # M_bak = np.copy(M)
+        # for i in range(nx):
+        #     for j in range(ny):
+        #         for k in range(nz):
+        #             M[i,j,k,:,:] = sp.linalg.inv( J[i,j,k,:,:] )
+        # np.testing.assert_allclose(M_bak, M, atol=1e-12, rtol=1e-12)
+        # print('check passed')
+        
+        t_delta = timeit.default_timer() - t_start
+        if verbose: tqdm.write( even_print('get M','%0.3f [s]'%(t_delta,), s=True) )
+    
+    if True: ## method 2
+        
+        if ('M' in locals()):
+            M_bak = np.copy(M)
+            M = None; del M
+        
+        t_start = timeit.default_timer()
+        
+        a = J[:,:,:,0,0]
+        b = J[:,:,:,0,1]
+        c = J[:,:,:,0,2]
+        d = J[:,:,:,1,0]
+        e = J[:,:,:,1,1]
+        f = J[:,:,:,1,2]
+        g = J[:,:,:,2,0]
+        h = J[:,:,:,2,1]
+        i = J[:,:,:,2,2]
+        
+        # a = J[:,:,:,0,0]
+        # b = J[:,:,:,1,0]
+        # c = J[:,:,:,2,0]
+        # d = J[:,:,:,0,1]
+        # e = J[:,:,:,1,1]
+        # f = J[:,:,:,2,1]
+        # g = J[:,:,:,0,2]
+        # h = J[:,:,:,1,2]
+        # i = J[:,:,:,2,2]
+        
+        Jac_det = ( + a*e*i
+                    + b*f*g
+                    + c*d*h
+                    - c*e*g
+                    - b*d*i
+                    - a*f*h )
+        
+        M = np.zeros((nx,ny,nz,3,3), dtype=np.float64)
+        M[:,:,:,0,0] = +( dydy * dzdz - dydz * dzdy ) / Jac_det ## ξ_x
+        M[:,:,:,0,1] = -( dxdy * dzdz - dxdz * dzdy ) / Jac_det ## ξ_y
+        M[:,:,:,0,2] = +( dxdy * dydz - dxdz * dydy ) / Jac_det ## ξ_z
+        M[:,:,:,1,0] = -( dydx * dzdz - dydz * dzdx ) / Jac_det ## η_x
+        M[:,:,:,1,1] = +( dxdx * dzdz - dxdz * dzdx ) / Jac_det ## η_y
+        M[:,:,:,1,2] = -( dxdx * dydz - dxdz * dydx ) / Jac_det ## η_z
+        M[:,:,:,2,0] = +( dydx * dzdy - dydy * dzdx ) / Jac_det ## ζ_x
+        M[:,:,:,2,1] = -( dxdx * dzdy - dxdy * dzdx ) / Jac_det ## ζ_y
+        M[:,:,:,2,2] = +( dxdx * dydy - dxdy * dydx ) / Jac_det ## ζ_z
+        
+        t_delta = timeit.default_timer() - t_start
+        if verbose: tqdm.write( even_print('get M','%0.3f [s]'%(t_delta,), s=True) )
+        
+        if ('M_bak' in locals()):
+            np.testing.assert_allclose(M[:,:,:,0,0], M_bak[:,:,:,0,0], atol=1e-14, rtol=1e-14)
+            print('check passed: ξ_x')
+            np.testing.assert_allclose(M[:,:,:,0,1], M_bak[:,:,:,0,1], atol=1e-14, rtol=1e-14)
+            print('check passed: ξ_y')
+            np.testing.assert_allclose(M[:,:,:,0,2], M_bak[:,:,:,0,2], atol=1e-14, rtol=1e-14)
+            print('check passed: ξ_z')
+            np.testing.assert_allclose(M[:,:,:,1,0], M_bak[:,:,:,1,0], atol=1e-14, rtol=1e-14)
+            print('check passed: η_x')
+            np.testing.assert_allclose(M[:,:,:,1,1], M_bak[:,:,:,1,1], atol=1e-14, rtol=1e-14)
+            print('check passed: η_y')
+            np.testing.assert_allclose(M[:,:,:,1,2], M_bak[:,:,:,1,2], atol=1e-14, rtol=1e-14)
+            print('check passed: η_z')
+            np.testing.assert_allclose(M[:,:,:,2,0], M_bak[:,:,:,2,0], atol=1e-14, rtol=1e-14)
+            print('check passed: ζ_x')
+            np.testing.assert_allclose(M[:,:,:,2,1], M_bak[:,:,:,2,1], atol=1e-14, rtol=1e-14)
+            print('check passed: ζ_y')
+            np.testing.assert_allclose(M[:,:,:,2,2], M_bak[:,:,:,2,2], atol=1e-14, rtol=1e-14)
+            print('check passed: ζ_z')
+            np.testing.assert_allclose(M, M_bak, atol=1e-14, rtol=1e-14)
+            print('check passed: M')
+    
+    return M

turbx/h5.py

@@ -0,0 +1,236 @@
+import sys
+
+import h5py
+import numpy as np
+from tqdm import tqdm
+
+# ======================================================================
+
+def h5_chunk_sizer(nxi, **kwargs):
+    '''
+    Solve for HDF5 dataset chunk size.
+
+    Parameters:
+    ----------
+    nxi : iterable
+        The shape of the full HDF5 dataset.
+    constraint : iterable
+        Per-axis constraint. Each element can be:
+            - None         → flexible
+            - int (>0)     → fixed chunk size
+            - 'full' or -1 → chunk size equals axis size
+            - ('max', int) → chunk size must be ≤ int
+    '''
+    
+    size_kb    = kwargs.get('size_kb'    , 2*1024 ) ## target chunk size in [KB] --> default = 2 [MB]
+    itemsize   = kwargs.get('itemsize'   , 4      ) ## dtype.itemsize --> default single precision i.e. 4 [B]
+    constraint = kwargs.get('constraint' , None   ) ## iterable of nxi constraints --> int,None,'full'/-1
+    base       = kwargs.get('base'       , 2      ) ## axis chunk size = ceil[size/(<int>*base)] where <int> is incremented
+    
+    ## if no constraint given, all axes are fully flexible
+    if constraint is None:
+        constraint = [ None for i in range(len(nxi)) ]
+    
+    ## check inputs
+    if not hasattr(constraint, '__iter__') or len(nxi) != len(constraint):
+        raise ValueError('nxi and constraint must be iterable and the same length')
+    if not isinstance(base,int):
+        raise TypeError('base must be an integer')
+    if (base<1):
+        raise TypeError('base must be an integer')
+    
+    # === increment divisor on largest axis, with divisor=<int>*base
+    
+    nxi = list(nxi)
+    div = [ 1 for i in range(len(nxi)) ] ## divisor vector, initialize with int ones
+    
+    ## list of axes indices which are 'flexible' ... this is updated dynamically in loop
+    i_flexible = [ i for i,c in enumerate(constraint) if c is None or isinstance(c,tuple) ]
+    
+    while True:
+        
+        div_last = list(div) ## make a copy
+        #print(f'div_last = {str(tuple(div_last))}')
+        
+        chunks = []
+        for i in range(len(nxi)):
+            
+            dim = nxi[i]
+            
+            if (constraint[i] is None):
+                C = max( int(np.floor(dim/div[i])) , 1 ) ## divide by divisor
+            elif (constraint[i] == 'full') or (constraint[i] == -1):
+                C = dim ## chunk axis shape is == dset axis shape
+            elif isinstance(constraint[i], int) and (constraint[i]>0):
+                C = constraint[i] ## chunk axis shape is just the constraint
+            elif isinstance(constraint[i], tuple) and (constraint[i][0]=='max') and isinstance(constraint[i][1],int):
+                max_val = constraint[i][1]
+                C = min( max( int(np.floor(dim/div[i])) , 1 ) , max_val )
+            else:
+                raise ValueError(f'problem with constraint[{i:d}] = {str(constraint[i])}')
+            chunks.append(C)
+        
+        #print(f'chunks = {str(tuple(chunks))}')
+        
+        ## recalculate i_flexible
+        i_flexible = []
+        for i,c in enumerate(constraint):
+            if chunks[i] == 1: ## already at min, is not flexible
+                continue
+            if c is None:
+                i_flexible.append(i)
+            elif isinstance(c,tuple) and c[0]=='max':
+                if chunks[i] > 1:
+                    i_flexible.append(i)
+        
+        #print(f'i_flexible = {str(i_flexible)}')
+        
+        ## there are no flexible axes --> exit loop
+        if len(i_flexible)==0:
+            break
+        
+        ## the current size of a chunk
+        chunk_size_kb = np.prod(chunks)*itemsize / 1024.
+        #print(f'chunk size {chunk_size_kb:0.1f} [KB] / {np.prod(chunks)*itemsize:d} [B]')
+        
+        if ( chunk_size_kb <= size_kb ): ## if chunk size is < target, then break
+            break
+        else: ## otherwise, increase the divisor of the greatest 'flexible' axis
+            
+            ## get index of (flexible) axis with greatest size
+            aa = [ i for i,c in enumerate(chunks) if (i in i_flexible) ]
+            bb = [ c for i,c in enumerate(chunks) if (i in i_flexible) ]
+            i_gt = aa[np.argmax(bb)]
+            
+            ## update divisor
+            div[i_gt] *= base
+        
+        #print(f'div = {str(tuple(div))}')
+        #print('---')
+        
+        ## check if in infinite loop (divisor not being updated)
+        if (div_last is not None) and (div == div_last):
+            raise ValueError(f'invalid parameters for h5_chunk_sizer() : constraint={str(constraint)}, size_kb={size_kb:d}, base={base:d}')
+    
+    return tuple(chunks)
+
+def h5_visititems_print_attrs(name, obj):
+    '''
+    callable for input to h5py.Group.visititems() to print names & attributes
+    '''
+    n_slashes = name.count('/')
+    shift = n_slashes*2*' '
+    item_name = name.split('/')[-1]
+    
+    if isinstance(obj,h5py._hl.dataset.Dataset):
+        print(shift + item_name + ' --> shape=%s, dtype=%s'%( str(obj.shape), str(obj.dtype) ) )
+    else:
+        print(shift + item_name)
+    
+    ## print attributes
+    for key, val in obj.attrs.items():
+        try:
+            print(shift + 2*' ' + f'{key} = {str(val)} --> dtype={str(val.dtype)}')
+        except AttributeError:
+            print(shift + 2*' ' + f'{key} = {str(val)} --> type={str(type(val).__name__)}')
+
+def h5_print_contents(h5filehandle):
+    '''
+    Print file-level attributes and recursively print names & attributes of all groups and datasets.
+    '''
+    
+    ## file-level attributes
+    for key, val in h5filehandle.attrs.items():
+        try:
+            print(f'{key} = {str(val)} --> dtype={str(val.dtype)}')
+        except AttributeError:
+            print(f'{key} = {str(val)} --> type={str(type(val).__name__)}')
+    
+    def visitor(name, obj):
+        n_slashes = name.count('/')
+        shift     = n_slashes * 2 * ' '
+        item_name = name.split('/')[-1]
+        
+        if isinstance(obj, h5py._hl.dataset.Dataset):
+            print(shift + item_name + ' --> shape=%s, dtype=%s' % (str(obj.shape), str(obj.dtype)))
+        else:
+            print(shift + item_name)
+        
+        for key, val in obj.attrs.items():
+            try:
+                print(shift + 2 * ' ' + f'{key} = {str(val)} --> dtype={str(val.dtype)}')
+            except AttributeError:
+                print(shift + 2 * ' ' + f'{key} = {str(val)} --> type={str(type(val).__name__)}')
+    
+    # Use the visitor function with visititems
+    h5filehandle.visititems(visitor)
+    
+    return
+
+class h5_visit_container:
+    '''
+    callable for input to h5py.Group.visit() which stores dataset/group names
+    '''
+    def __init__(self):
+        self.names = []
+    def __call__(self, name):
+        if (name not in self.names):
+            self.names.append(name)
+
+def h5_ds_force_allocate_chunks(ds, verbose=False, quick=False):
+    '''
+    Force allocation of all chunks in an ND dataset by writing real data
+    '''
+    if not isinstance(ds, h5py.Dataset):
+        raise TypeError('ds must be a h5py.Dataset object')
+    
+    shape = ds.shape
+    dtype = ds.dtype
+    chunk_shape = ds.chunks
+    rng = np.random.default_rng(seed=1)
+    
+    ## for contiguous datasets, fill the entire array
+    if chunk_shape is None:
+        #ds[...] = np.zeros(shape, dtype=dtype) ## might lead to optimizations under the hood
+        #ds[...] = rng.uniform(-1,+1,size=shape).astype(dtype)
+        ds[...] = rng.random(size=shape, dtype=dtype)
+        return
+    
+    ## info needed for iterating through chunks
+    chunk_starts = [range(0, dim, cdim) for dim, cdim in zip(shape, chunk_shape)]
+    chunk_grid_shape = [len(r) for r in chunk_starts]
+    total_chunks = np.prod(chunk_grid_shape)
+    
+    if verbose:
+        progress_bar = tqdm(
+            total=total_chunks,
+            ncols=100,
+            desc='allocate chunks',
+            leave=True,
+            file=sys.stdout,
+            mininterval=0.1,
+            smoothing=0.,
+            #bar_format="\033[B{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}\033[A\n\b",
+            bar_format="{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}",
+            ascii="░█",
+            colour='#FF6600',
+            )
+    
+    for chunk_idx in np.ndindex(*chunk_grid_shape):
+        starts = [r[i] for r, i in zip(chunk_starts, chunk_idx)]
+        
+        if quick: ## just write a single element to allocate the chunk
+            ds[tuple(starts)] = 0
+        else:
+            slices = tuple(
+                slice(start, min(start + size, dim))
+                for start, size, dim in zip(starts, chunk_shape, shape)
+                )
+            actual_shape = tuple(slc.stop - slc.start for slc in slices)
+            #ds[slices] = np.zeros(actual_shape, dtype=dtype) ## might lead to optimizations under the hood
+            #ds[slices] = rng.uniform(-1,+1,size=actual_shape).astype(dtype)
+            ds[slices] = rng.random(size=actual_shape, dtype=dtype)
+        
+        if verbose: progress_bar.update()
+    if verbose: progress_bar.close()
+    return