diff --git a/tofu/data/__init__.py b/tofu/data/__init__.py index 0b7d5451f..8ec39e086 100644 --- a/tofu/data/__init__.py +++ b/tofu/data/__init__.py @@ -6,6 +6,8 @@ from ._DataCollection_class1_interactivity import * from ._class01_eqdsk import * from ._spectrallines_class import * +from ._class00_poly2d_check import check as poly2d_check +from ._class00_poly2d_sample import main as poly2d_sample from ._class10_algos import get_available_inversions_algo from ._class10_Inversion import Inversion as Collection from ._spectralunits import * diff --git a/tofu/data/_class00_poly2d_check.py b/tofu/data/_class00_poly2d_check.py new file mode 100644 index 000000000..410ab798e --- /dev/null +++ b/tofu/data/_class00_poly2d_check.py @@ -0,0 +1,189 @@ +# -*- coding: utf-8 -*- +""" Basic tools for formatting 2d polygons + +""" + + +import numpy as np +import datastock as ds + + +# ########################################################### +# ########################################################### +# check +# ########################################################### + + +def check( + x0=None, + x1=None, + key=None, + # options + closed=None, + clockwise=None, +): + """ High-level routine to format a 2d polygon + + + Parameters + ---------- + x0 : sequence + 1st coordinate of the polygon + x1 : sequence + 2nd coordinate of the polygon + key : str / None + To inform error messages + closed : bool / None + whether the polygon should be closed + clockwise : bool / None + whether the polygon should be clockwise + + Returns + ------- + dout : dict + { + 'x0': np.ndarray, + 'x0': np.ndarray, + 'closed': bool, + 'clockwise': bool, + 'key': str, + } + + """ + + # ------------- + # check inputs + # ------------- + + key, close, clockwise = _check( + key=key, + closed=closed, + clockwise=clockwise, + ) + + # ------------- + # check x0, x1 + # ------------- + + x0 = ds._generic_check._check_flat1darray( + x0, 'x0', + dtype=float, + can_be_None=False, + extra_msg=f"x0 of polygon '{key}'", + ) + + x1 = ds._generic_check._check_flat1darray( + x1, 'x1', + dtype=float, + can_be_None=False, + size=x0.size, + extra_msg=f"x1 of polygon '{key}'", + ) + + # ------------- + # closed + # ------------- + + is_closed = np.allclose(np.r_[x0[0], x1[0]], np.r_[x0[-1], x1[-1]]) + if is_closed is True: + x0_closed = x0 + x1_closed = x1 + + else: + ind_closed = np.r_[np.arange(0, x0.size), 0] + x0_closed = x0[ind_closed] + x1_closed = x1[ind_closed] + + # ------------- + # no duplicates + # ------------- + + uni = np.unique([x0_closed[:-1], x1_closed[:-1]], axis=1) + if uni.shape[1] < (x0_closed.size - 1): + ndup = x0.size - 1 - uni.shape[1] + msg = ( + f"Polygon 2d '{key}' seems to have {ndup} duplicate points!\n" + "\t- x0 = {x0_closed[:-1]}\n" + "\t- x1 = {x1_closed[:-1]}\n" + ) + raise Exception(msg) + + # ------------- + # clockwise + # ------------- + + # is already ? + is_cw = is_clockwise(x0_closed, x1_closed) + + # adjust + if is_cw != clockwise: + x0_closed, x1_closed = x0_closed[::-1], x1_closed[::-1] + + # ------------- + # return + # ------------- + + dout = { + 'x0': x0_closed if closed else x0_closed[:-1], + 'x1': x1_closed if closed else x1_closed[:-1], + 'key': key, + 'closed': closed, + 'clockwise': clockwise, + } + + return dout + +# ########################################################### +# ########################################################### +# check +# ########################################################### + + +def _check( + key=None, + closed=None, + clockwise=None, +): + + # --------------- + # key + # --------------- + + key = ds._generic_check._check_var( + key, 'key', + types=str, + default='', + ) + + # --------------- + # closed + # --------------- + + closed = ds._generic_check._check_var( + closed, 'closed', + types=bool, + default=False, + ) + + # --------------- + # clockwise + # --------------- + + clockwise = ds._generic_check._check_var( + clockwise, 'clockwise', + types=bool, + default=True, + ) + + return key, closed, clockwise + + +# ########################################################### +# ########################################################### +# is clockwise +# ########################################################### + + +def is_clockwise(x0, x1): + area_signed = np.sum((x0[1:] - x0[:-1]) * (x1[1:] + x1[:-1])) + return area_signed > 0 \ No newline at end of file diff --git a/tofu/data/_class00_poly2d_sample.py b/tofu/data/_class00_poly2d_sample.py new file mode 100644 index 000000000..5879f9e68 --- /dev/null +++ b/tofu/data/_class00_poly2d_sample.py @@ -0,0 +1,442 @@ +# -*- coding: utf-8 -*- +""" +Polygons are assumed to be: + - finite + - simple (non self-intersecting) + - non-explicitly closed + - have no repeated points + +""" + + +import numpy as np +from matplotlib.path import Path +import datastock as ds + + +from . import _class00_poly2d_check as _poly2d_check + + +__all__ = ['main'] + + +# ########################################################### +# ########################################################### +# main +# ########################################################### + + +def main( + x0=None, + x1=None, + key=None, + # options for edges + dedge=None, + dsurface=None, +): + + # ---------------- + # check inputs + # ---------------- + + # polygon formatting + din = _poly2d_check.check( + x0=x0, + x1=x1, + key=key, + # options + closed=True, + clockwise=True, + ) + + # dedge, dsurface + dedge, dsurface = _check( + dedge=dedge, + dsurface=dsurface, + ) + + # ---------------- + # sample edges + # ---------------- + + if dedge is not None: + dout_edge = _edges( + x0_closed=din['x0'], + x1_closed=din['x1'], + key=din['key'], + # options + **{k0: dedge.get(k0) for k0 in ['res', 'factor']}, + ) + else: + dout_edge = { + 'x0': [], + 'x1': [], + } + + # ---------------- + # sample surface + # ---------------- + + if dsurface is not None: + dout_surf = _surface( + x0_closed=din['x0'], + x1_closed=din['x1'], + key=din['key'], + # options + **{k0: dsurface.get(k0) for k0 in ['res', 'nb']}, + ) + else: + dout_surf = { + 'x0': [], + 'x1': [], + } + + # ---------------- + # combine + # ---------------- + + dout = { + 'x0': np.r_[dout_edge['x0'], dout_surf['x0']], + 'x1': np.r_[dout_edge['x1'], dout_surf['x1']], + 'edge_res': dout_edge.get('res'), + 'edge_factor': dout_edge.get('factor'), + 'surface_nb': dout_surf.get('nb'), + 'key': din['key'], + } + + return dout + + +# ########################################################### +# ########################################################### +# check main +# ########################################################### + + +def _check( + dedge=None, + dsurface=None, +): + + # -------------- + # dedge + # -------------- + + if dedge is not None: + lk = ['res', 'factor'] + c0 = ( + isinstance(dedge, dict) + and all([kk in lk for kk in dedge.keys()]) + ) + + if not c0: + lstr = [f"\t- {kk}" for kk in lk] + msg = ( + "Arg dedge must be None or a dict with keys:\n" + + "\n".join(lstr) + ) + raise Exception(msg) + + # -------------- + # dsurface + # -------------- + + if dsurface is not None: + lk = ['res', 'nb'] + c0 = ( + isinstance(dsurface, dict) + and all([kk in lk for kk in dsurface.keys()]) + ) + + if not c0: + lstr = [f"\t- {kk}" for kk in lk] + msg = ( + "Arg dsurface must be None or a dict with keys:\n" + + "\n".join(lstr) + ) + raise Exception(msg) + + return dedge, dsurface + + +# ########################################################### +# ########################################################### +# Sample edges +# ########################################################### + + +def _edges( + x0_closed=None, + x1_closed=None, + key=None, + # options + res=None, + factor=None, +): + + # ---------------- + # check inputs + # ---------------- + + # options + res, factor = _check_edges( + res=res, + factor=factor, + ) + + # ---------------- + # prepare + # ---------------- + + # get all lengths + dx0 = x0_closed[1:] - x0_closed[:-1] + dx1 = x1_closed[1:] - x1_closed[:-1] + dist = np.hypot(dx0, dx1) + + # ---------------- + # determine samplung res + # ---------------- + + if isinstance(res, str): + + # get res + if res == 'min': + res = np.min(dist) + elif res == 'max': + res = np.max(dist) + else: + raise NotImplementedError() + + # adjust + res = res * factor + + # ---------------- + # sample res + # ---------------- + + # nb of points to be inserted in each segment + npts = np.round(dist / res, decimals=0).astype(int) + + # interpolation for x0 + out0 = np.ravel([ + np.r_[x0_closed[ii]] + if npts[ii] <= 1 else + x0_closed[ii] + dx0[ii] * np.linspace(0, 1, npts[ii], endpoint=False) + for ii in range(x0_closed.size-1) + ]) + + # interpolation for x1 + out1 = np.ravel([ + np.r_[x1_closed[ii]] + if npts[ii] <= 1 else + x1_closed[ii] + dx1[ii] * np.linspace(0, 1, npts[ii], endpoint=False) + for ii in range(x1_closed.size-1) + ]) + + # ---------------- + # return + # ---------------- + + dout = { + 'x0': out0, + 'x1': out1, + 'res': res, + 'factor': factor, + } + + return dout + + +# ########################################################### +# ########################################################### +# check edges +# ########################################################### + + +def _check_edges( + res=None, + factor=None, +): + + # ------------- + # res + # ------------- + + res = ds._generic_check._check_var( + res, 'res', + types=(str, float), + ) + + if isinstance(res, str): + res = ds._generic_check._check_var( + res, 'res', + types=str, + default='min', + allowed=['min', 'max'], + ) + + else: + res = ds._generic_check._check_var( + res, 'res', + types=float, + sign='>0', + ) + + # ------------- + # factor + # ------------- + + if isinstance(res, str): + factor = float(ds._generic_check._check_var( + factor, 'factor', + types=(int, float), + default=1, + sign='>0', + )) + + else: + factor = None + + return res, factor + + +# ########################################################### +# ########################################################### +# Sample surface +# ########################################################### + + +def _surface( + x0_closed=None, + x1_closed=None, + key=None, + # options + res=None, + nb=None, +): + + # ---------------- + # check inputs + # ---------------- + + # Dx0 + x0_min = x0_closed.min() + x0_max = x0_closed.max() + Dx0 = (x0_max - x0_min) + + # Dx1 + x1_min = x1_closed.min() + x1_max = x1_closed.max() + Dx1 = (x1_max - x1_min) + + # options + nb = _check_surfaces( + res=res, + nb=nb, + key=key, + Dx0=Dx0, + Dx1=Dx1, + ) + + # ---------------- + # prepare + # ---------------- + + # out0 + dx0 = Dx0 / nb[0] + out0 = np.linspace(x0_min + dx0/2, x0_max - dx0/2, nb[0]) + + # out1 + dx1 = Dx1 / nb[1] + out1 = np.linspace(x1_min + dx1/2, x1_max - dx1/2, nb[1]) + + # ---------------- + # Check in polygon + # ---------------- + + pts0 = np.repeat(out0[:, None], out1.size, axis=1).ravel() + pts1 = np.repeat(out1[None, :], out1.size, axis=0).ravel() + + pp = Path(np.array([x0_closed, x1_closed]).T) + ind = pp.contains_points(np.array([pts0, pts1]).T) + + # ---------------- + # return + # ---------------- + + dout = { + 'x0': pts0[ind], + 'x1': pts1[ind], + 'nb': nb, + } + + return dout + + +# ########################################################### +# ########################################################### +# check surfaces +# ########################################################### + + +def _check_surfaces( + res=None, + nb=None, + key=None, + Dx0=None, + Dx1=None, +): + + # ------------- + # res vs nb + # ------------- + + lc = [ + nb is not None and res is not None, + nb is None and res is None, + ] + if any(lc): + msg = ( + "Polygon '{key}' surface sampling, please provide res xor nb!\n" + f"\t- res = {res} \n" + f"\t- nb = {nb}\n" + ) + raise Exception(msg) + + # ------------- + # res + # ------------- + + if res is not None: + + if np.isscalar(res): + res = [res, res] + + res = ds._generic_check._check_var_iter( + res, 'res', + types=(list, tuple), + types_iter=(int, float), + size=2, + ) + + res = tuple([float(rr) for rr in res]) + + nb = [np.ceil(Dx0/res[0]), np.ceil(Dx1/res[1])] + + # ------------- + # nb + # ------------- + + if np.isscalar(nb): + nb = [nb, nb] + + nb = ds._generic_check._check_var_iter( + nb, 'nb', + types=(list, tuple), + types_iter=(int, float), + size=2, + ) + + nb = tuple([int(nn) for nn in nb]) + + return nb \ No newline at end of file diff --git a/tofu/data/_class08_Diagnostic.py b/tofu/data/_class08_Diagnostic.py index 50eca61d8..224d75e3e 100644 --- a/tofu/data/_class08_Diagnostic.py +++ b/tofu/data/_class08_Diagnostic.py @@ -67,6 +67,7 @@ def add_diagnostic( etendue=None, # config for los config=None, + strict=None, length=None, # reflections reflections_nb=None, @@ -112,6 +113,7 @@ def add_diagnostic( check=False, # los config=config, + strict=strict, length=length, reflections_nb=reflections_nb, reflections_type=reflections_type, @@ -232,6 +234,7 @@ def compute_diagnostic_etendue_los( convex=None, # for storing los config=None, + strict=None, length=None, reflections_nb=None, reflections_type=None, @@ -290,6 +293,7 @@ def compute_diagnostic_etendue_los( key=key, # los config=config, + strict=strict, length=length, reflections_nb=reflections_nb, reflections_type=reflections_type, @@ -382,21 +386,31 @@ def compute_diagnostic_solidangle_from_plane( def add_rays_from_diagnostic( self, key=None, - strategy=None, - nrays=None, + # sampling + dsampling_pixel=None, + dsampling_optics=None, + # optics (to restrain to certain optics only for faster) + optics=None, + # computing + config=None, # storing store=None, - config=None, + key_rays=None, overwrite=None, ): return _generate_rays.main( coll=self, key=key, - strategy=strategy, - nrays=nrays, + # sampling + dsampling_pixel=dsampling_pixel, + dsampling_optics=dsampling_optics, + # optics (to restrain to certain optics only for faster) + optics=optics, + # computing + config=config, # storing store=store, - config=config, + key_rays=key_rays, overwrite=overwrite, ) diff --git a/tofu/data/_class08_generate_rays.py b/tofu/data/_class08_generate_rays.py index daaf9740b..4b0f24cf7 100644 --- a/tofu/data/_class08_generate_rays.py +++ b/tofu/data/_class08_generate_rays.py @@ -12,6 +12,7 @@ from ..geom import _GG +from ._class00_poly2d_sample import main as poly2d_sample # ############################################################### @@ -23,25 +24,69 @@ def main( coll=None, key=None, - strategy=None, - nrays=None, + # to sample on a single optics + key_optics=None, + # sampling + dsampling_pixel=None, + dsampling_optics=None, + # optics (to restrain to certain optics only for faster) + optics=None, + # computing + config=None, # storing store=None, - config=None, + key_rays=None, overwrite=None, ): + """ + + Parameters + ---------- + coll : TYPE, optional + DESCRIPTION. The default is None. + key : TYPE, optional + DESCRIPTION. The default is None. + # to sample on a single optics key_optics : TYPE, optional + DESCRIPTION. The default is None. + # sampling dsampling_pixel : TYPE, optional + DESCRIPTION. The default is None. + dsampling_optics : TYPE, optional + DESCRIPTION. The default is None. + # computing config : TYPE, optional + DESCRIPTION. The default is None. + # storing store : TYPE, optional + DESCRIPTION. The default is None. + overwrite : TYPE, optional + DESCRIPTION. The default is None. + + Returns + ------- + dout : TYPE + DESCRIPTION. + + """ # ------------- # check # ------------- - key, strategy, nrays, store, overwrite = _check( + ( + key, + dsampling_pixel, + dsampling_optics, + optics, + store, key_rays, overwrite, + ) = _check( coll=coll, key=key, - strategy=strategy, - nrays=nrays, + # sampling + dsampling_pixel=dsampling_pixel, + dsampling_optics=dsampling_optics, + # optics + optics=optics, # storing store=store, + key_rays=key_rays, overwrite=overwrite, ) @@ -53,19 +98,6 @@ def main( key_cam = coll.dobj[wdiag][key]['camera'] doptics = coll.dobj[wdiag][key]['doptics'] - # ------------------------- - # trivial: nrays=1 => LOS - # ------------------------- - - if nrays == 1: - - dout = { - kcam: {} - for kcam in key_cam - } - - store = False - # --------------- # compute # --------------- @@ -76,35 +108,20 @@ def main( # loop on cameras for kcam in key_cam: - # -------------- - # prepare pixel outline - - out0, out1 = coll.dobj['camera'][kcam]['dgeom']['outline'] - out0 = coll.ddata[out0]['data'] - out1 = coll.ddata[out1]['data'] - out_arr = np.array([out0, out1]).T - - # --------------- - # call routine + # -------------------- + # sample generic pixel - if strategy == 'random': - - dout[kcam] = _random( - coll=coll, - kcam=kcam, - doptics=doptics[kcam], - out_arr=out_arr, - nrays=nrays, - strategy=strategy, - ) - - elif strategy == 'outline': - - dout[kcam] = _outline() - - elif strategy == 'mesh': - - dout[kcam] = _mesh() + dout[kcam] = _generic( + coll=coll, + key=key, + kcam=kcam, + doptics=doptics[kcam], + # sampling + dsampling_pixel=dsampling_pixel, + dsampling_optics=dsampling_optics, + # optics + optics=optics.get(kcam), + ) # --------------- # store @@ -114,6 +131,7 @@ def main( _store( coll=coll, dout=dout, + key_rays=key_rays, config=config, overwrite=overwrite, ) @@ -131,10 +149,14 @@ def main( def _check( coll=None, key=None, - strategy=None, - nrays=None, + # sampling + dsampling_pixel=None, + dsampling_optics=None, + # optics + optics=None, # storing store=None, + key_rays=None, overwrite=None, ): @@ -150,38 +172,55 @@ def _check( allowed=lok, ) + key_cam = coll.dobj[wdiag][key]['camera'] + # spectro ? spectro = coll.dobj[wdiag][key]['spectro'] if spectro is True: raise NotImplementedError() # ------------ - # strategy + # dsampling_pixel # ------------ - strategy = ds._generic_check._check_var( - strategy, 'strategy', - types=str, - default='random', - allowed=['random', 'mesh', 'outline'], - ) + ldict = [ + ('dsampling_pixel', dsampling_pixel), + ('dsampling_optics', dsampling_optics), + ] + lk = [('dedge', ['res', 'factor']), ('dsurface', ['res', 'nb'])] + kfunc = 'tf.data.poly2d_sample()' + for ii, (kdict, vdict) in enumerate(ldict): + c0 = ( + isinstance(vdict, dict) + and any([kk in vdict.keys() for (kk, largs) in lk]) + ) + if not c0: + lstr = [f"\t- {kk}: dict of {largs}" for (kk, largs) in lk] + msg = ( + f"Arg '{kdict}' must be a dict with at least one of:\n" + + "\n".join(lstr) + + f"\nFed to {kfunc}\n" + ) + raise Exception(msg) # ------------ - # nrays + # optics # ------------ - if strategy == 'custom': - nrdef = 10 - else: - nrdef = 10 + if optics is not None: + # index => position in optics list + optics0 = { + kc: _check_optics_for_kcam( + doptics=coll.dobj[wdiag][key]['doptics'][kc], + optics=optics.get(kc) if isinstance(optics, dict) else optics, + shape_cam=coll.dobj['camera'][kc]['dgeom']['shape'], + ) + for kc in key_cam + } - nrays = int(ds._generic_check._check_var( - nrays, 'nrays', - types=(int, float), - default=nrdef, - sign='>0', - )) + else: + optics0 = {} # ------------ # store @@ -193,6 +232,41 @@ def _check( default=True, ) + # ------------ + # key_rays + # ------------ + + if store is True: + + if key_rays is None: + key_rays = {kcam: f"{kcam}_rays" for kcam in key_cam} + + else: + if len(key_cam) == 1 and isinstance(key_rays, str): + key_rays[key_cam[0]] = key_rays + + ncam = len(key_cam) + c0 = ( + isinstance(key_rays, dict) + and all([ + isinstance(key_rays.get(kcam), str) + and len(set([key_rays[kcam] for kcam in key_cam])) == ncam + and all([key_rays[kcam] not in coll.dobj['camera'].keys()]) + for kcam in key_cam + ]) + ) + if not c0: + msg = ( + "Arg key_rays must be a dict of key_ray for each camera:\n" + f"\t- key_diag = '{key}'\n" + f"\t- key_cam = {key_cam}\n" + f"\t- key_rays = {key_rays}\n" + ) + raise Exception(msg) + + else: + key_rays = None + # ------------ # overwrite # ------------ @@ -203,7 +277,61 @@ def _check( default=False, ) - return key, strategy, nrays, store, overwrite + return ( + key, + dsampling_pixel, dsampling_optics, + optics0, + store, key_rays, overwrite, + ) + + +def _check_optics_for_kcam( + doptics=None, + optics=None, + shape_cam=None, +): + + if isinstance(optics, int): + + ind0 = np.arange(0, len(doptics['optics'])) + if doptics['pinhole'] is True: + lok = np.r_[ind0, -1] + else: + nop = doptics['paths'].sum(axis=1) + lok = np.r_[np.arange(0, np.min(nop)), -np.arange(1, np.min(nop)+1)] + + optics = ds._generic_check._check_var( + optics, 'optics', + types=int, + allowed=lok, + ) + + if doptics['pinhole'] is True: + optics = [doptics['optics'][optics]] + else: + assert doptics.get('paths') is not None + paths = doptics['paths'] + + optics = [ + doptics['optics'][ind0[paths[ii, :]][optics]] + for ii in range(shape_cam[0]) + ] + + # ------------------- + # key or list of keys + + if isinstance(optics, str): + optics = [optics] + + optics = ds._generic_check._check_var_iter( + optics, 'optics', + types=(list, tuple), + types_iter=str, + allowed=doptics['optics'], + ) + + # eliminate redundants (ex. collimator with 1 common optic) + return list(set(optics)) # ############################################################### @@ -469,12 +597,317 @@ def _seed_optics( # ############################################################### # ############################################################### -# outline +# generic # ############################################################### -def _outline(): +def _generic( + coll=None, + key=None, + kcam=None, + doptics=None, + # sampling options + dsampling_pixel=None, + dsampling_optics=None, + # optics + optics=None, +): + + # ------------------- + # optics + # ------------------- + + if optics is None: + optics = doptics['optics'] + + # --------------------- + # generic pixel outline + # --------------------- + + # get pixel outline + kout0, kout1 = coll.dobj['camera'][kcam]['dgeom']['outline'] + + # sample pixel + dout_pixel = poly2d_sample( + coll.ddata[kout0]['data'], + coll.ddata[kout1]['data'], + dedge=dsampling_pixel.get('dedge'), + dsurface=dsampling_pixel.get('dsurface'), + ) + + nstart = dout_pixel['x0'].size + + # --------------------- + # prepare + # --------------------- + # shared apertures ? + parallel = coll.dobj['camera'][kcam]['dgeom']['parallel'] + + # shape camera + shape_cam = coll.dobj['camera'][kcam]['dgeom']['shape'] + shape_start = shape_cam + (nstart,) + + # camera vector + # get camera vectors + dvect = coll.get_camera_unit_vectors(kcam) + lc = ['x', 'y', 'z'] + if parallel is True: + e0i_x, e0i_y, e0i_z = [dvect[f"e0_{kk}"] for kk in lc] + e1i_x, e1i_y, e1i_z = [dvect[f"e1_{kk}"] for kk in lc] + + cents_x, cents_y, cents_z = coll.get_camera_cents_xyz(kcam) + + # ----------------------- + # prepare buffer for cx, cy, cz + # ----------------------- + + # cx, cy, cz (3d poitns for each pixel) + cx = np.full((nstart,), np.nan) + cy = np.full((nstart,), np.nan) + cz = np.full((nstart,), np.nan) + + # --------------------- + # Only 1 optic (pinhole or collimator with 1 common optics) + # --------------------- + + if len(optics) == 1: + + # ------------------------ + # get end points on optics + + endx, endy, endz = _get_end_optics( + coll=coll, + kop=optics[0], + dsampling_optics=dsampling_optics, + ) + nend = endx.size + + # -------------- + # prepare output + + # rays shape for each pixel + shape_each = (nstart, nend) + vx = np.full(shape_each, np.nan) + vy = np.full(shape_each, np.nan) + vz = np.full(shape_each, np.nan) + + # oreall + shape_out = shape_start + (nend,) + + dout = { + 'start_x': np.full(shape_out, np.nan), + 'start_y': np.full(shape_out, np.nan), + 'start_z': np.full(shape_out, np.nan), + 'vect_x': np.full(shape_out, np.nan), + 'vect_y': np.full(shape_out, np.nan), + 'vect_z': np.full(shape_out, np.nan), + # 'dsang': np.full(shape_out, np.nan), + } + + # ------------------- + # assign output + + for ii, ind in enumerate(np.ndindex(shape_cam)): + + if not parallel: + e0i_x, e0i_y, e0i_z = [dvect[f"e0_{kk}"][ind] for kk in lc] + e1i_x, e1i_y, e1i_z = [dvect[f"e1_{kk}"][ind] for kk in lc] + + # update cx, cy, cz + _update_cxyz( + parallel, + e0i_x, e0i_y, e0i_z, + e1i_x, e1i_y, e1i_z, + cents_x, cents_y, cents_z, + ind, lc, + dout_pixel, + cx, cy, cz, + dvect, + ) + + # vect + vx[...] = endx[None, :] - cx[:, None] + vy[...] = endy[None, :] - cy[:, None] + vz[...] = endz[None, :] - cz[:, None] + + vnorm_inv = 1./np.sqrt(vx**2 + vy**2 + vz**2) + + # slicing + sli = ind + (slice(None), slice(None)) + + # assigning + dout['start_x'][sli] = np.copy(cx)[:, None] + dout['start_y'][sli] = np.copy(cy)[:, None] + dout['start_z'][sli] = np.copy(cz)[:, None] + + dout['vect_x'][sli] = vx * vnorm_inv + dout['vect_y'][sli] = vy * vnorm_inv + dout['vect_z'][sli] = vz * vnorm_inv + + # --------------------- + # loop on pixels + # --------------------- + + else: + + dout = { + 'start_x': {}, + 'start_y': {}, + 'start_z': {}, + 'vect_x': {}, + 'vect_y': {}, + 'vect_z': {}, + # 'dsang': np.full(shape_out, np.nan), + } + + for ii, ind in enumerate(np.ndindex(shape_cam)): + + # check optics per pixel + lop = np.array(doptics['optics'])[doptics['paths'][ii, :]] + lop = [ + kop for kop in lop + if (kop in optics or optics is None) + ] + + # check number of optics + if len(lop) > 1: + msg = ( + "Multiple optics not handled yet!\n" + f"\t- key_diag = '{key}'\n" + f"\t- kcam = '{kcam}'\n" + f"\t- ii, ind = {ii}, {ind}\n" + f"\t- optics = {optics}\n" + f"\t- lop = {lop}\n" + ) + raise NotImplementedError(msg) + + # ------------------------ + # get end points on optics + + endx, endy, endz = _get_end_optics( + coll=coll, + kop=lop[0], + dsampling_optics=dsampling_optics, + ) + nend = endx.size + + # ------------------ + # update pixel cx, cy, cz + + if not parallel: + e0i_x, e0i_y, e0i_z = [dvect[f"e0_{kk}"][ind] for kk in lc] + e1i_x, e1i_y, e1i_z = [dvect[f"e1_{kk}"][ind] for kk in lc] + + _update_cxyz( + parallel, + e0i_x, e0i_y, e0i_z, + e1i_x, e1i_y, e1i_z, + cents_x, cents_y, cents_z, + ind, lc, + dout_pixel, + cx, cy, cz, + dvect, + ) + + # vector + vx = endx[None, :] - cx[:, None] + vy = endy[None, :] - cy[:, None] + vz = endz[None, :] - cz[:, None] + + vnorm_inv = 1./np.sqrt(vx**2 + vy**2 + vz**2) + + # assigning + dout['start_x'][ind] = np.copy(cx)[:, None] + dout['start_y'][ind] = np.copy(cy)[:, None] + dout['start_z'][ind] = np.copy(cz)[:, None] + + dout['vect_x'][ind] = vx * vnorm_inv + dout['vect_y'][ind] = vy * vnorm_inv + dout['vect_z'][ind] = vz * vnorm_inv + + # ------------- + # adjust + # ------------- + + if isinstance(dout['start_x'], dict): + + dnrays = {k0: v0.shape[1] for k0, v0 in dout['vect_x'].items()} + nraysu = np.unique([v0 for v0 in dnrays.values()]) + + lkstart = [f'start_{cc}' for cc in lc] + lkvect = [f'vect_{cc}' for cc in lc] + + # initialize + for kk in lkstart + lkvect: + oo = np.full(shape_cam + (nstart, nraysu.max()), np.nan) + + for ii, ind in enumerate(np.ndindex(shape_cam)): + sli = ind + (slice(None), np.arange(0, dnrays[ind])) + oo[sli] = dout[kk][ind] + + dout[kk] = oo + + # --------------- + # return + # --------------- + + return dout + + +def _get_end_optics( + coll=None, + kop=None, + dsampling_optics=None, +): + + kop, clsop = coll.get_optics_cls(kop) + clsop, kop = clsop[0], kop[0] + kout0, kout1 = coll.dobj[clsop][kop]['dgeom']['outline'] + + dout_optics = poly2d_sample( + coll.ddata[kout0]['data'], + coll.ddata[kout1]['data'], + dedge=dsampling_optics.get('dedge'), + dsurface=dsampling_optics.get('dsurface'), + ) + + # get 3d optics end points coordinates + func = coll.get_optics_x01toxyz(key=kop, asplane=False) + return func(dout_optics['x0'], dout_optics['x1']) + + +def _update_cxyz( + parallel, + e0i_x, e0i_y, e0i_z, + e1i_x, e1i_y, e1i_z, + cents_x, cents_y, cents_z, + ind, lc, + dout_pixel, + cx, cy, cz, + dvect, +): + + if parallel is not True: + e0i_x, e0i_y, e0i_z = [dvect[f"e0_{kk}"][ind] for kk in lc] + e1i_x, e1i_y, e1i_z = [dvect[f"e1_{kk}"][ind] for kk in lc] + + # start + cx[...] = ( + cents_x[ind] + + dout_pixel['x0'] * e0i_x + + dout_pixel['x1'] * e1i_x + ) + cy[...] = ( + cents_y[ind] + + dout_pixel['x0'] * e0i_y + + dout_pixel['x1'] * e1i_y + ) + cz[...] = ( + cents_z[ind] + + dout_pixel['x0'] * e0i_z + + dout_pixel['x1'] * e1i_z + ) return @@ -501,27 +934,38 @@ def _store( coll=None, kdiag=None, dout=None, + key_rays=None, config=None, overwrite=None, ): - # ----------------- - # add ref - # ----------------- - - nrays = list(dout.values())[0]['start_x'].shape[-1] - - kref = f"{kdiag}_nrays" - coll.add_ref(key=kref, size=nrays) # -------------- - # add rays + # store # -------------- for kcam, v0 in dout.items(): - ref = coll.dobj['camera'][kcam]['dgeom']['ref'] + (kref,) + key = key_rays[kcam] + + # ----------------- + # add ref + + nstart, nends = v0['vect_x'].shape[-2:] + + krstart = f"{key}_nstart" + coll.add_ref(key=krstart, size=nstart) + + krend = f"{key}_nend" + coll.add_ref(key=krend, size=nends) + + # ----------------- + # add rays + + ref = coll.dobj['camera'][kcam]['dgeom']['ref'] + (krstart, krend) + coll.add_rays( + key=key, ref=ref, config=config, **v0 diff --git a/tofu/data/_class2_check.py b/tofu/data/_class2_check.py index bbabd2f96..ab0de3b8d 100644 --- a/tofu/data/_class2_check.py +++ b/tofu/data/_class2_check.py @@ -130,20 +130,27 @@ def _check_inputs( # --------------------- # pts vs config vs diag + lkvpts = [('pts_x', pts_x), ('pts_y', pts_y), ('pts_z', pts_z)] + lkvvect = [('vect_x', vect_x), ('vect_y', vect_y), ('vect_z', vect_z)] + lkv_rt = [('config', config), ('length', length), ('diag', diag)] + lc = [ - pts_x is not None, - vect_x is not None - and ( - config is not None - or length is not None - or diag is not None - ), + all([vv is not None for (kk, vv) in lkvpts]), + all([vv is not None for (kk, vv) in lkvvect]) + and any([vv is not None for (kk, vv) in lkv_rt]) ] if np.sum(lc) != 1: + lstr0 = [f"\t\t- {kk} is None: {vv is None}" for (kk, vv) in lkvpts] + lstr1 = [f"\t\t- {kk} is None: {vv is None}" for (kk, vv) in lkvvect] + lstr2 = [f"\t\t- {kk} is None: {vv is None}" for (kk, vv) in lkv_rt] msg = ( "Please provide either:\n" - "\t- pts_x, pts_y, pts_z: directly specify end points\n" - "\t- config and / or diag: ray-tracing" + "\t- pts_x, pts_y, pts_z: to directly specify end points\n" + + "\n".join(lstr0) + + "\n\t- vect_x, vect_y, vect_z" + + " and (config or length or diag): for ray-tracing\n" + + "\n".join(lstr1) + + "\n" + "\n".join(lstr2) ) raise Exception(msg) diff --git a/tofu/data/_class8_los_angles.py b/tofu/data/_class8_los_angles.py index 969ca5740..f343121c3 100644 --- a/tofu/data/_class8_los_angles.py +++ b/tofu/data/_class8_los_angles.py @@ -32,6 +32,7 @@ def compute_los_angles( dcompute=None, # for storing los config=None, + strict=None, length=None, reflections_nb=None, reflections_type=None, @@ -121,6 +122,7 @@ def compute_los_angles( key_cam=key_cam, v0=v0, config=config, + strict=strict, res=res, overwrite=overwrite, ) @@ -190,6 +192,7 @@ def _vos_from_los( key_cam=None, v0=None, config=None, + strict=None, res=None, overwrite=None, ): @@ -286,6 +289,7 @@ def _vos_from_los( coords=coll.get_optics_x01toxyz(key=optics[iref]), lspectro=lspectro, config=config, + strict=strict, # debug key=key, ) @@ -608,6 +612,7 @@ def _get_rays_from_pix( coords=None, lspectro=None, config=None, + strict=None, # debug key=None, ): @@ -677,7 +682,7 @@ def _get_rays_from_pix( Name='', Diag='', Exp='', - strict=True, + strict=strict, ) # pin diff --git a/tofu/data/_class8_plot.py b/tofu/data/_class8_plot.py index 1edae311a..9bb0d1698 100644 --- a/tofu/data/_class8_plot.py +++ b/tofu/data/_class8_plot.py @@ -620,7 +620,7 @@ def _plot_diagnostic( nan_los, nan_los, nan_los, - c=color_dict['x'][ii], + c=color_dict['x'][ii%len(color_dict['x'])], ls='-', lw=1., ) @@ -775,7 +775,7 @@ def _plot_diagnostic( l0, = ax.plot( dataz, ddata[k0][sli], - c=color_dict['x'][ii], + c=color_dict['x'][ii%len(color_dict['x'])], lw=1., ls='-', ) @@ -1185,7 +1185,7 @@ def _add_camera_los_cross( l0, = ax.plot( nan_los, nan_los, - c=color_dict['x'][ii], + c=color_dict['x'][ii%len(color_dict['x'])], ls='-', lw=1., ) @@ -1209,7 +1209,7 @@ def _add_camera_los_cross( l0, = ax.fill( nan_vos, nan_vos, - fc=color_dict['x'][ii], + fc=color_dict['x'][ii%len(color_dict['x'])], alpha=alpha, ls='None', lw=0., @@ -1251,7 +1251,7 @@ def _add_camera_los_hor( l0, = ax.plot( nan_los, nan_los, - c=color_dict['x'][ii], + c=color_dict['x'][ii%len(color_dict['x'])], ls='-', lw=1., ) @@ -1277,7 +1277,7 @@ def _add_camera_los_hor( l0, = ax.fill( nan_vos, nan_vos, - fc=color_dict['x'][ii], + fc=color_dict['x'][ii%len(color_dict['x'])], alpha=alpha, ls='None', lw=0., @@ -1314,6 +1314,8 @@ def _add_camera_vlines_marker( suffix=None, ): + ncolx = len(color_dict['x']) + ncoly = len(color_dict['y']) if suffix is None: suffix = '' @@ -1324,7 +1326,7 @@ def _add_camera_vlines_marker( ddatay[k0][0:1], marker='s', ms=6, - markeredgecolor=color_dict['x'][ii], + markeredgecolor=color_dict['x'][ii%ncolx], markerfacecolor='None', ) @@ -1350,7 +1352,7 @@ def _add_camera_vlines_marker( for ii in range(nlos): lv = ax.axvline( - ddatax[k0][0], c=color_dict['y'][ii], lw=1., ls='-', + ddatax[k0][0], c=color_dict['y'][ii%ncoly], lw=1., ls='-', ) kv = f'{k0}_v{ii:02.0f}{suffix}' coll2.add_mobile( diff --git a/tofu/data/_class8_vos.py b/tofu/data/_class8_vos.py index f517343e2..b2ca6ad54 100644 --- a/tofu/data/_class8_vos.py +++ b/tofu/data/_class8_vos.py @@ -732,6 +732,13 @@ def _get_user_limits( np.full(shape_cam, user_limits['Dphi'][1]), ]) + # ----------- + # clean-up + # ---------- + + if len(user_limits) == 0: + user_limits = None + return user_limits diff --git a/tofu/data/_class8_vos_broadband.py b/tofu/data/_class8_vos_broadband.py index c5491a0bc..c73e701db 100644 --- a/tofu/data/_class8_vos_broadband.py +++ b/tofu/data/_class8_vos_broadband.py @@ -119,6 +119,13 @@ def _vos( phor1 = user_limits['phor1'][key_cam] dphi = user_limits['dphi'][key_cam] + else: + msg = ( + "Something weird with pcross0:\n" + f"user_limits: {user_limits}\n" + ) + raise Exception(msg) + else: # get temporary vos diff --git a/tofu/data/_utils_surface3d.py b/tofu/data/_utils_surface3d.py index 5f12ce70b..41141f24b 100644 --- a/tofu/data/_utils_surface3d.py +++ b/tofu/data/_utils_surface3d.py @@ -96,7 +96,7 @@ def _surface3d( outline_x0, outline_x1, area = _check_polygon_2d( poly_x=outline_x0, poly_y=outline_x1, - poly_name=f'{key}-outline', + poly_name=f'{key}_outline', can_be_None=False, closed=False, counter_clockwise=True, @@ -435,4 +435,4 @@ def _get_outline_from_poly( outline_x0, outline_x1 = None, None area = np.nan - return gtype, cent, outline_x0, outline_x1, area + return gtype, cent, outline_x0, outline_x1, area \ No newline at end of file diff --git a/tofu/tests/tests08_diagnostics/test_01_diagnostics.py b/tofu/tests/tests08_diagnostics/test_01_diagnostics.py index 15a741ba1..5748d0084 100644 --- a/tofu/tests/tests08_diagnostics/test_01_diagnostics.py +++ b/tofu/tests/tests08_diagnostics/test_01_diagnostics.py @@ -819,7 +819,21 @@ def test06_plot_coverage(self): plt.close('all') - def test07_reverse_ray_tracing(self): + def test07_add_rays_from_diagnostic(self): + for ii, (k0, v0) in enumerate(self.coll.dobj['diagnostic'].items()): + noptics = any([len(v1['optics']) == 0 for v1 in v0['doptics'].values()]) + if v0['is2d'] or v0['spectro'] or noptics: + continue + dout = self.coll.add_rays_from_diagnostic( + key=k0, + dsampling_pixel={'dedge': {'res': 'max'}, 'dsurface': {'nb': 3}}, + dsampling_optics={'dedge': {'res': 'max'}, 'dsurface': {'nb': 3}}, + optics=-1, + config=self.conf, + store=ii%2 == 0, + ) + + def test08_reverse_ray_tracing(self): for ii, (k0, v0) in enumerate(self.coll.dobj['diagnostic'].items()): lcam = self.coll.dobj['diagnostic'][k0]['camera'] doptics = self.coll.dobj['diagnostic'][k0]['doptics'] @@ -866,7 +880,7 @@ def test07_reverse_ray_tracing(self): colorbar=None, ) - def test08_save_to_json(self): + def test09_save_to_json(self): for ii, (k0, v0) in enumerate(self.coll.dobj['diagnostic'].items()):