# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/CLI/plot.ipynb.

# %% auto 0
__all__ = ['zarr_stack_info', 'ras_pyramid', 'hv_ras_callback', 'ras_plot', 'hv_ras_stack_callback', 'ras_stack_plot',
           'pc_pyramid', 'is_nan_range', 'hv_pc_Image_callback', 'hv_pc_Points_callback', 'pc_plot',
           'hv_pc_Image_stack_callback', 'hv_pc_Points_stack_callback', 'pc_stack_plot']

# %% ../../nbs/CLI/plot.ipynb 4
import logging
import zarr
import numpy as np
import math
from pathlib import Path
import pandas as pd
from tqdm import tqdm
import sys
from functools import partial
from typing import Callable
import numpy as np
from numba import prange
import holoviews as hv
from holoviews import streams

import toml
from ..utils_ import ngpjit
from ..rtree import HilbertRtree
from .logging import mc_logger
from ..coord_ import Coord

# %% ../../nbs/CLI/plot.ipynb 5
def zarr_stack_info(zarr_path_list, #list of zarr path
                   ):
    shape_list = []; chunks_list = []; dtype_list = []
    for zarr_path in zarr_path_list:
        zarr_data = zarr.open(zarr_path,'r')
        shape_list.append(zarr_data.shape)
        chunks_list.append(zarr_data.chunks)
        dtype_list.append(zarr_data.dtype)
    df = pd.DataFrame({'path':zarr_path_list,'shape':shape_list,'chunks':chunks_list,'dtype':dtype_list})
    return df

# %% ../../nbs/CLI/plot.ipynb 7
@mc_logger
def ras_pyramid(ras:str, # path to input data, 2D zarr array (one single raster) or 3D zarr array (a stack of rasters)
                out_dir:str, # output directory to store rendered data
                ):
    '''render raster data to pyramid of difference zoom levels.'''
    logger = logging.getLogger(__name__)
    out_dir = Path(out_dir); out_dir.mkdir(exist_ok=True)
    ras_zarr = zarr.open(ras,'r')
    logger.zarr_info(ras, ras_zarr)
    
    ny, nx = ras_zarr.shape[0:2]
    maxlevel = math.floor(math.log2(min(nx,ny))) # so at least 2 pixels
    logger.info(f'rendered raster pyramid with zoom level ranging from 0 (finest resolution) to {maxlevel} (coarsest resolution).')

    ndim = ras_zarr.ndim
    o_zarr_path_list = []; o_zarr_list = []
    for level in range(maxlevel+1):
        if ndim == 3:
            shape = (math.ceil(ny/(2**level)), math.ceil(nx/(2**level)), ras_zarr.shape[2])
            chunks = (256, 256, 1)
        elif ndim == 2:
            shape = (math.ceil(ny/(2**level)), math.ceil(nx/(2**level)))
            chunks = (256, 256)
        else:
            raise ValueError('ndim of input should be 2 or 3.')
        o_zarr_path = out_dir/f'{level}.zarr'
        o_zarr = zarr.open(zarr.NestedDirectoryStore(o_zarr_path),
                           'w',shape=shape,chunks=chunks,dtype=ras_zarr.dtype)
        o_zarr_path_list.append(o_zarr_path); o_zarr_list.append(o_zarr)

    df_disp = zarr_stack_info(o_zarr_path_list)
    logger.info(f'pyramid to be rendered: \n {df_disp}')
    
    logger.info(f'pyramid rendering starts.')
    if ndim == 2:
        ras_data = ras_zarr[:]
        for level in tqdm(range(maxlevel+1),file=sys.stdout,desc='Tiles'):
            o_zarr_list[level][:,:] = ras_data[::2**level,::2**level]
    else:
        for i in tqdm(range(ras_zarr.shape[2]),file=sys.stdout,desc='Images'):
            ras_data = ras_zarr[:,:,i]
            for level in range(maxlevel+1):
            # for level in tqdm(range(maxlevel+1),file=sys.stdout,desc='Tiles',position=0):
                o_zarr_list[level][:,:,i] = ras_data[::2**level,::2**level]
    logger.info(f'tiles rendering finished.')

# %% ../../nbs/CLI/plot.ipynb 13
def hv_ras_callback(x_range,y_range,width,height,scale,data_dir,post_proc,coord):
    # start = time.time()
    if x_range is None:
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if y_range is None:
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res,y_res)))
    level = sorted((0, level, coord.maxlevel))[1]
    data_zarr = zarr.open(data_dir/f'{level}.zarr','r')
    xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
    coord_bbox = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
    # decide_slice = time.time()
    # print(f"It takes {decide_slice-start} to decide the data range", file = sourceFile) 
    data = data_zarr[yi0:yim+1,xi0:xim+1]
    # load_data = time.time()
    # print(f"It takes {load_data-decide_slice} to load the data", file = sourceFile) 
    data = post_proc(data)
    # post_proc_data = time.time()
    # print(f"It takes {post_proc_data-load_data} to post_proc the data", file = sourceFile)
    ### test shows data read takes only 0.006 s, post_proc and data_range takes only 0.001s
    ### the majority of time is used by holoviews that I can not optimize.
    return hv.Image(data[::-1,:],bounds=coord_bbox)

# %% ../../nbs/CLI/plot.ipynb 14
def ras_plot(rendered_tiles_dir:str, # directory to the rendered images
             post_proc:Callable=None, # function for the post processing
             bounds:tuple=None, # bounding box (x0, y0, x_max, y_max)
            ):
    '''plot rendered ras tiles.'''
    rendered_tiles_dir = Path(rendered_tiles_dir)
    data_zarr = zarr.open(rendered_tiles_dir/'0.zarr','r')
    ny, nx = data_zarr.shape
    
    if post_proc is None: post_proc = lambda x: x
    if bounds is None:
        x0 = 0; dx = 1; y0 = 0; dy = 1
    else:
        x0, y0, xm, ym = bounds
        dx = (xm-x0)/(nx-1); dy = (ym-y0)/(ny-1)
    coord = Coord(x0,dx,nx,y0,dy,ny)
    
    rangexy = streams.RangeXY()
    plotsize = streams.PlotSize()
    images = hv.DynamicMap(partial(hv_ras_callback,data_dir=rendered_tiles_dir,post_proc=post_proc,coord=coord),streams=[rangexy,plotsize])
    return images

# %% ../../nbs/CLI/plot.ipynb 24
def hv_ras_stack_callback(x_range,y_range,width,height,scale,data_dir,post_proc,coord,i=0):
    if x_range is None:
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if y_range is None:
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res,y_res)))
    level = sorted((0, level, coord.maxlevel))[1]
    data_zarr = zarr.open(data_dir/f'{level}.zarr','r')
    xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
    coord_bbox = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
    data = post_proc(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1),i)
    return hv.Image(data[::-1,:],bounds=coord_bbox)

# %% ../../nbs/CLI/plot.ipynb 25
def ras_stack_plot(rendered_tiles_dir:str, # directory to the rendered images
                   post_proc:Callable=None, # function for the post processing
                   bounds:tuple=None, # bounding box (x0, y0, x_max, y_max)
                  ):
    '''plot rendered stack of ras tiles.'''
    rendered_tiles_dir = Path(rendered_tiles_dir)
    data_zarr = zarr.open(rendered_tiles_dir/'0.zarr','r')
    ny, nx = data_zarr.shape[:2]
    
    if post_proc is None: post_proc = lambda data_zarr,x_slice,y_slice,i: data_zarr[y_slice,x_slice,i]
    if bounds is None:
        x0 = 0; dx = 1; y0 = 0; dy = 1
    else:
        x0, y0, xm, ym = bounds
        dx = (xm-x0)/(nx-1); dy = (ym-y0)/(ny-1)
    coord = Coord(x0,dx,nx,y0,dy,ny)
    
    rangexy = streams.RangeXY()
    plotsize = streams.PlotSize()
    images = hv.DynamicMap(partial(hv_ras_stack_callback,data_dir=rendered_tiles_dir,post_proc=post_proc,coord=coord),streams=[rangexy,plotsize],kdims='i')
    return images

# %% ../../nbs/CLI/plot.ipynb 39
@ngpjit
def _next_level_idx_from_raster_of_integer(pc_idx, nan_value):
    '''return the raster indices to the next level of raster'''
    assert pc_idx.ndim == 2
    ny, nx = pc_idx.shape
    next_ny, next_nx = math.ceil(ny/2), math.ceil(nx/2)
    xi = np.empty((next_ny,next_nx), dtype=np.int32)
    yi = np.empty((next_ny,next_nx), dtype=np.int32)

    for i in range(next_ny):
        for j in prange(next_nx):
            # Select a 2x2 box from the original array
            box = pc_idx[i*2:min(i*2+2, ny), j*2:min(j*2+2, nx)]
            idx_ = np.argwhere(box != nan_value)
            if len(idx_) == 0:
                yi[i,j]= i*2
                xi[i,j] = j*2
            else:
                yi[i,j] = idx_[0,0] + i*2
                xi[i,j] = idx_[0,1] + j*2
    return yi, xi

# %% ../../nbs/CLI/plot.ipynb 40
@ngpjit
def _next_level_idx_from_raster_of_noninteger(pc_data):
    '''return the raster indices to the next level of raster'''
    assert pc_data.ndim == 2
    ny, nx = pc_data.shape
    next_ny, next_nx = math.ceil(ny/2), math.ceil(nx/2)
    xi = np.empty((next_ny,next_nx), dtype=np.int32)
    yi = np.empty((next_ny,next_nx), dtype=np.int32)

    for i in range(next_ny):
        for j in prange(next_nx):
            # Select a 2x2 box from the original array
            box = pc_data[i*2:min(i*2+2, ny), j*2:min(j*2+2, nx)]
            idx_ = np.argwhere(~np.isnan(box))
            if len(idx_) == 0:
                yi[i,j]= i*2
                xi[i,j] = j*2
            else:
                yi[i,j] = idx_[0,0] + i*2
                xi[i,j] = idx_[0,1] + j*2
    return yi, xi

# %% ../../nbs/CLI/plot.ipynb 42
@mc_logger
def pc_pyramid(pc:str, # path to point cloud data, 1D array (one single pc image) or 2D zarr array (a stack of pc images)
               x:str, # path to x coordinate, e.g., longitude or web mercator x
               y:str, # path to y coordinate, e.g., latitude or web mercator y
               out_dir:str, # output directory to store rendered data
               ras_resolution:float=5, # minimum resolution of rendered raster data,
               ):
    '''render point cloud data to pyramid of difference zoom levels.'''
    logger = logging.getLogger(__name__)
    out_dir = Path(out_dir); out_dir.mkdir(exist_ok=True)
    pc_zarr = zarr.open(pc,'r')
    logger.zarr_info(pc, pc_zarr)
    
    n_pc = pc_zarr.shape[0]
    logger.info(f'rendering point cloud data coordinates:')
    x, y = zarr.open(x,'r')[:], zarr.open(y,'r')[:]
    o_x_zarr = zarr.open(out_dir/f'x.zarr','w',shape=x.shape, dtype=x.dtype, chunks=(256*256,))
    o_y_zarr = zarr.open(out_dir/f'y.zarr','w',shape=y.shape, dtype=y.dtype, chunks=(256*256,))
    o_x_zarr[:] = x; o_y_zarr[:]=y
    logger.info('pc data coordinates rendering ends.')

    x0, xm, y0, ym = x.min(), x.max(), y.min(), y.max()
    nx, ny = math.ceil((xm-x0)/ras_resolution), math.ceil((ym-y0)/ras_resolution)
    coord = Coord(x0, ras_resolution, nx, y0, ras_resolution, ny)
    bounds = {'bounds':[x0, y0, coord.xm, coord.ym]}
    logger.info(f"rasterizing point cloud data to grid with bounds: {bounds['bounds']}.")
    with open(out_dir/'bounds.toml','w') as f:
        toml.dump(bounds, f, encoder=toml.TomlNumpyEncoder())

    ndim = pc_zarr.ndim
    o_pc_zarr_path = out_dir/f'pc.zarr'
    o_pc_chunks = (256*256,1) if ndim == 2 else (256*256,)
    o_pc_zarr = zarr.open(o_pc_zarr_path,'w',shape=pc_zarr.shape,chunks=o_pc_chunks,dtype=pc_zarr.dtype)
    logger.info(f'pc pyramid pc part (level -1):')
    logger.zarr_info(o_pc_zarr_path, o_pc_zarr)

    gix = coord.coords2gixs((y,x)); del y; del x
    maxlevel = coord.maxlevel # so at least 2 pixels
    logger.info(f'rendered raster pyramid with zoom level ranging from 0 (finest resolution) to {maxlevel} (coarsest resolution).')

    ras_zarr_path_list = []; ras_zarr_list = []; ras_idx_zarr_path_list = []; ras_idx_zarr_list = []
    for level in range(maxlevel+1):
        shape = (math.ceil(ny/(2**level)), math.ceil(nx/(2**level)))
        chunks = (256, 256)
        ras_idx_zarr_path = out_dir/f'idx_{level}.zarr'
        ras_idx_zarr = zarr.open(zarr.NestedDirectoryStore(ras_idx_zarr_path),
                            'w',shape=shape,chunks=chunks,dtype=np.int64)
        ras_idx_zarr_path_list.append(ras_idx_zarr_path); ras_idx_zarr_list.append(ras_idx_zarr)
        if ndim == 2:
            shape = (math.ceil(ny/(2**level)), math.ceil(nx/(2**level)), pc_zarr.shape[1])
            chunks = (256, 256, 1)
        ras_zarr_path = out_dir/f'{level}.zarr'
        ras_zarr = zarr.open(zarr.NestedDirectoryStore(ras_zarr_path),
                            'w',shape=shape,chunks=chunks,dtype=pc_zarr.dtype)
        ras_zarr_path_list.append(ras_zarr_path); ras_zarr_list.append(ras_zarr)

    df_ras = zarr_stack_info(ras_zarr_path_list)
    logger.info(f'pc pyramid raster part to be rendered: \n {df_ras}')
    df_idx = zarr_stack_info(ras_idx_zarr_path_list)
    logger.info(f'pc pyramid index to be randered: \n {df_idx}')

    logger.info(f'pyramid raster indexing rendering starts.')
    current_idx = coord.rasterize_iidx(gix)
    next_yi_list = []; next_xi_list = []
    for level in tqdm(range(maxlevel+1),file=sys.stdout,desc='Tiles'):
        ras_idx_zarr_list[level][:] = current_idx
        next_yi, next_xi = _next_level_idx_from_raster_of_integer(current_idx,-1)
        next_yi_list.append(next_yi); next_xi_list.append(next_xi)
        current_idx = current_idx[next_yi, next_xi]

    logger.info(f'pyramid data rendering starts.')
    if ndim == 1:
        logger.info(f'pyramid pc part rendering starts.')
        pc_data = pc_zarr[:]
        o_pc_zarr[:] = pc_data; logger.info(f'pyramid pc part rendering ends.')
        logger.info(f'pyramid raster part rendering starts.')
        current_ras = coord.rasterize(pc_data, gix)
        for level in tqdm(range(maxlevel+1),file=sys.stdout,desc='Tiles'):
            ras_zarr_list[level][:,:] = current_ras
            current_ras = current_ras[next_yi_list[level],next_xi_list[level]]
    else:
        for i in tqdm(range(ras_zarr.shape[2]),file=sys.stdout,desc='Images'):
            pc_data = pc_zarr[:,i]
            o_pc_zarr[:,i] = pc_data
            current_ras = coord.rasterize(pc_data, gix)
            for level in range(maxlevel+1):
                ras_zarr_list[level][:,:,i] = current_ras
                current_ras = current_ras[next_yi_list[level],next_xi_list[level]]
    logger.info(f'pyramid rendering finished.')

# %% ../../nbs/CLI/plot.ipynb 47
def is_nan_range(x_range):
    if x_range is None:
        return True
    if np.isnan(x_range[0]):
        return True
    return False

# %% ../../nbs/CLI/plot.ipynb 50
def hv_pc_Image_callback(x_range,y_range,width,height,scale,data_dir,post_proc_ras,coord):
    if is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        data_zarr = zarr.open(data_dir/f'{level}.zarr','r')
        idx_zarr = zarr.open(data_dir/f'idx_{level}.zarr','r')
        xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
        x0, y0, xm, ym = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
        data = data_zarr[yi0:yim+1,xi0:xim+1]
        data = post_proc_ras(data)
        idx = idx_zarr[yi0:yim+1,xi0:xim+1]
        return hv.Image((np.linspace(x0,xm,data.shape[1]), np.linspace(y0,ym,data.shape[0]),data,idx),vdims=['z','idx'])
    else:
        return hv.Image([],vdims=['z','idx'])

# %% ../../nbs/CLI/plot.ipynb 51
def hv_pc_Points_callback(x_range,y_range,width,height,scale,data_dir,post_proc_pc,coord,rtree):
    if is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        return hv.Points([],vdims=['z','idx'])
    else:
        data_zarr, x_zarr, y_zarr = (zarr.open(data_dir/file,'r') for file in ('pc.zarr', 'x.zarr', 'y.zarr'))
        idx = rtree.bbox_query((x0, y0, xm, ym), x_zarr, y_zarr)
        data, x, y = (zarr_[idx] for zarr_ in (data_zarr, x_zarr, y_zarr))
        data = post_proc_pc(data)
        return hv.Points((x,y,data,idx),vdims=['z','idx'])

# %% ../../nbs/CLI/plot.ipynb 52
def pc_plot(pyramid_dir:str, # directory to the rendered point cloud pyramid
            post_proc_ras:Callable=None, # function for the post processing
            post_proc_pc:Callable=None, # function for the post processing
            rtree=None, # rtree
           ):
    '''plot rendered point cloud pyramid.'''    
    if post_proc_ras is None: post_proc_ras = lambda x: x
    if post_proc_pc is None: post_proc_pc = lambda x: x

    pyramid_dir = Path(pyramid_dir)
    data_zarr = zarr.open(pyramid_dir/'0.zarr','r')
    ny, nx = data_zarr.shape
    with open(pyramid_dir/'bounds.toml','r') as f:
        x0, y0, xm, ym = toml.load(f)['bounds']

    dx = (xm-x0)/(nx-1); dy = (ym-y0)/(ny-1)
    coord = Coord(x0,dx,nx,y0,dy,ny)
    
    if rtree is None:
        x = zarr.open(pyramid_dir/'x.zarr','r')[:]
        y = zarr.open(pyramid_dir/'y.zarr','r')[:]
        rtree = HilbertRtree.build(x,y,page_size=512)
    rangexy = streams.RangeXY()
    plotsize = streams.PlotSize()
    images = hv.DynamicMap(partial(hv_pc_Image_callback,data_dir=pyramid_dir,
                                   post_proc_ras=post_proc_ras,coord=coord),
                           streams=[rangexy,plotsize])
    points = hv.DynamicMap(partial(hv_pc_Points_callback,data_dir=pyramid_dir,
                                   post_proc_pc=post_proc_pc,coord=coord,rtree=rtree),
                           streams=[rangexy,plotsize])
    return images*points

# %% ../../nbs/CLI/plot.ipynb 64
def hv_pc_Image_stack_callback(x_range,y_range,width,height,scale,data_dir,post_proc_ras,coord,i):
    if is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        data_zarr = zarr.open(data_dir/f'{level}.zarr','r')
        idx_zarr = zarr.open(data_dir/f'idx_{level}.zarr','r')
        xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
        x0, y0, xm, ym = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
        data = data_zarr[yi0:yim+1,xi0:xim+1]
        data = post_proc_ras(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1),i)
        idx = idx_zarr[yi0:yim+1,xi0:xim+1]
        return hv.Image((np.linspace(x0,xm,data.shape[1]), np.linspace(y0,ym,data.shape[0]),data,idx),vdims=['z','idx'])
    else:
        return hv.Image([],vdims=['z','idx'])

# %% ../../nbs/CLI/plot.ipynb 65
def hv_pc_Points_stack_callback(x_range,y_range,width,height,scale,data_dir,post_proc_pc,coord,rtree,i):
    if is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        return hv.Points([],vdims=['z','idx'])
    else:
        data_zarr, x_zarr, y_zarr = (zarr.open(data_dir/file,'r') for file in ('pc.zarr', 'x.zarr', 'y.zarr'))
        idx = rtree.bbox_query((x0, y0, xm, ym), x_zarr, y_zarr)
        x, y = (zarr_[idx] for zarr_ in (x_zarr, y_zarr))
        data = post_proc_pc(data_zarr,idx,i)
        return hv.Points((x,y,data,idx),vdims=['z','idx'])

# %% ../../nbs/CLI/plot.ipynb 66
def pc_stack_plot(pyramid_dir:str, # directory to the rendered point cloud pyramid
                  post_proc_ras:Callable=None, # function for the post processing
                  post_proc_pc:Callable=None, # function for the post processing
                  rtree=None, # rtree
                 ):
    '''plot rendered point cloud pyramid.'''    
    if post_proc_ras is None: post_proc_ras = lambda data_zarr,x_slice,y_slice,i: data_zarr[y_slice, x_slice, i]
    if post_proc_pc is None: post_proc_pc = lambda data_zarr,idx_array,i: data_zarr[idx_array,i]

    pyramid_dir = Path(pyramid_dir)
    data_zarr = zarr.open(pyramid_dir/'idx_0.zarr','r')
    ny, nx = data_zarr.shape
    with open(pyramid_dir/'bounds.toml','r') as f:
        x0, y0, xm, ym = toml.load(f)['bounds']

    dx = (xm-x0)/(nx-1); dy = (ym-y0)/(ny-1)
    coord = Coord(x0,dx,nx,y0,dy,ny)
    
    if rtree is None:
        x = zarr.open(pyramid_dir/'x.zarr','r')[:]
        y = zarr.open(pyramid_dir/'y.zarr','r')[:]
        rtree = HilbertRtree.build(x,y,page_size=512)
    rangexy = streams.RangeXY()
    plotsize = streams.PlotSize()
    images = hv.DynamicMap(partial(hv_pc_Image_stack_callback,data_dir=pyramid_dir,
                                   post_proc_ras=post_proc_ras,coord=coord),
                           streams=[rangexy,plotsize],kdims='i')
    points = hv.DynamicMap(partial(hv_pc_Points_stack_callback,data_dir=pyramid_dir,
                                   post_proc_pc=post_proc_pc,coord=coord,rtree=rtree),
                           streams=[rangexy,plotsize],kdims='i')
    return images*points
