diff --git a/tofu/data/_class08_Diagnostic.py b/tofu/data/_class08_Diagnostic.py index 32a618716..289899c07 100644 --- a/tofu/data/_class08_Diagnostic.py +++ b/tofu/data/_class08_Diagnostic.py @@ -565,6 +565,7 @@ def compute_diagnostic_resolution( def add_rays_from_diagnostic( self, key=None, + key_cam=None, # sampling dsampling_pixel=None, dsampling_optics=None, @@ -581,6 +582,7 @@ def add_rays_from_diagnostic( return _generate_rays.main( coll=self, key=key, + key_cam=key_cam, # sampling dsampling_pixel=dsampling_pixel, dsampling_optics=dsampling_optics, diff --git a/tofu/data/_class08_generate_rays.py b/tofu/data/_class08_generate_rays.py index 1b3d90edb..91c048e0f 100644 --- a/tofu/data/_class08_generate_rays.py +++ b/tofu/data/_class08_generate_rays.py @@ -24,6 +24,7 @@ def main( coll=None, key=None, + key_cam=None, # to sample on a single optics key_optics=None, # sampling @@ -73,6 +74,7 @@ def main( ( key, + key_cam, dsampling_pixel, dsampling_optics, optics, @@ -80,6 +82,7 @@ def main( ) = _check( coll=coll, key=key, + key_cam=key_cam, # sampling dsampling_pixel=dsampling_pixel, dsampling_optics=dsampling_optics, @@ -97,7 +100,6 @@ def main( # --------------- wdiag = 'diagnostic' - key_cam = coll.dobj[wdiag][key]['camera'] doptics = coll.dobj[wdiag][key]['doptics'] # --------------- @@ -152,6 +154,7 @@ def main( def _check( coll=None, key=None, + key_cam=None, # sampling dsampling_pixel=None, dsampling_optics=None, @@ -168,6 +171,7 @@ def _check( # key # ------------ + # key wdiag = 'diagnostic' lok = list(coll.dobj.get(wdiag, {}).keys()) key = ds._generic_check._check_var( @@ -176,7 +180,18 @@ def _check( allowed=lok, ) - key_cam = coll.dobj[wdiag][key]['camera'] + # key_cam + lok = coll.dobj[wdiag][key]['camera'] + if key_cam is None: + key_cam = lok + if isinstance(key_cam, str): + key_cam = [key_cam] + key_cam = ds._generic_check._check_var_iter( + key_cam, 'key_cam', + types=(list, tuple), + types_iter=str, + allowed=lok, + ) # spectro ? spectro = coll.dobj[wdiag][key]['spectro'] @@ -295,6 +310,7 @@ def _check( return ( key, + key_cam, dsampling_pixel, dsampling_optics, optics0, store, strict, key_rays, overwrite, diff --git a/tofu/data/_class8_check.py b/tofu/data/_class8_check.py index 37ec7bb0a..313dfc650 100644 --- a/tofu/data/_class8_check.py +++ b/tofu/data/_class8_check.py @@ -339,7 +339,7 @@ def _check_doptics_basics( key=key, doptics=doptics, returnas=True, - shape_cam=emsg, + extra_msg=emsg, ) raise err0 from err diff --git a/tofu/physics_tools/__init__.py b/tofu/physics_tools/__init__.py index 406c53a9b..e32375388 100644 --- a/tofu/physics_tools/__init__.py +++ b/tofu/physics_tools/__init__.py @@ -3,3 +3,4 @@ from . import electrons from .transmission import get_xray_transmission from .heat_transport_1d import main as heat_transport_1d +from . import magnetic diff --git a/tofu/physics_tools/magnetic/__init__.py b/tofu/physics_tools/magnetic/__init__.py new file mode 100644 index 000000000..53e2a2731 --- /dev/null +++ b/tofu/physics_tools/magnetic/__init__.py @@ -0,0 +1 @@ +from . import shielding diff --git a/tofu/physics_tools/magnetic/shielding/__init__.py b/tofu/physics_tools/magnetic/shielding/__init__.py new file mode 100644 index 000000000..d5ec7790e --- /dev/null +++ b/tofu/physics_tools/magnetic/shielding/__init__.py @@ -0,0 +1,2 @@ +from ._single import main as single +from ._double import main as double diff --git a/tofu/physics_tools/magnetic/shielding/_double.py b/tofu/physics_tools/magnetic/shielding/_double.py new file mode 100644 index 000000000..5a9db5d4b --- /dev/null +++ b/tofu/physics_tools/magnetic/shielding/_double.py @@ -0,0 +1,402 @@ +""" Analytical formulation of a 2d shielding cylinder + + +""" + + +import os + + +import numpy as np +import matplotlib.pyplot as plt + + +# ################################################################ +# ################################################################ +# DEFAULTS +# ################################################################ + + +_PATH_HERE = os.path.dirname(__file__) + +_DMUR = { + 'SS1008': 2000, + 'mumetal': 1, # 470000, +} +# _B0T = 0.015 +# _B0A = 0.003 +_B0T = 0.011 +_B0A = 0.002 + + +_DGEOM = { + # 'a': 66e-3, + 'a': 69e-3, + 'a_tol': 0.5e-3, + # 'b': 139.725e-3, + 'b': 140e-3, + 'b_tol': 0.1e-3, + 'thick': 1.5e-3, + # length + # 'len0': 250e-3, # S1008 + 'len0': 230e-3, + 'len1': (240-24.5)*1e-3, # mumetal +} + + +# ################################################################ +# ################################################################ +# Compute double shield +# ################################################################ + + +def main( + dmur=_DMUR, + dgeom=_DGEOM, + B0T=_B0T, + B0A=_B0A, +): + + # -------------- + # compute + # -------------- + + dshield = { + 'SS1008': shielding_trans( + a=dgeom['a'] + dgeom['a_tol'], + b=dgeom['b'] - _DGEOM['b_tol'], + mur=dmur['SS1008'], + ), + 'mumetal': shielding_trans( + a=dgeom['a'] - dgeom['a_tol'] - dgeom['thick'], + b=dgeom['a'] - dgeom['a_tol'], + mur=dmur['mumetal'], + ), + } + + shield = dshield['mumetal'] * dshield['SS1008'] + ST = int(1/shield) + + # -------------- + # axial component + # -------------- + + p_S1008 = dgeom['len0'] / dgeom['b'] + p_mumetal = dgeom['len1'] / dgeom['a'] + + N_S1008 = shielding_axial(p=p_S1008) + N_mumetal = shielding_axial(p=p_mumetal) + + SA_S1008 = 4*N_S1008/dshield['SS1008'] + 1 + SA_mumetal = 4*N_mumetal/dshield['mumetal'] + 1 + SA = SA_S1008 * SA_mumetal + + # --------------- + # prepare figure + # --------------- + + fig = plt.figure(figsize=(15, 8)) + + # transverse + axT = fig.add_axes( + [.06, 0.50, 0.4, 0.40], + aspect='equal', + adjustable='datalim', + ) + + # axial + axA = fig.add_axes( + [.56, 0.50, 0.4, 0.40], + adjustable='datalim', + ) + + # axial - shapes + axS = fig.add_axes( + [.70, 0.60, 0.25, 0.25], + aspect='equal', + adjustable='datalim', + ) + + # transverse title + ref0 = "[1] A. Mager, IEEE Transactions on Magnetics, vol. 6, 1970" + axT.set_title( + "B field in a 2d infinite shielding cylinder\n" + f"External transverse field B0 = {B0T} T ({int(B0T*1e4)} gauss)\n" + + ref0, + size=14, + fontweight='bold', + ) + + ref1 = "[2] A. Mager, Journal of Applied Physics, vol. 39, 1968" + axA.set_title( + "B field in a 2d finite shielding cylinder\n" + f"External axial field B0 = {B0A} T ({int(B0A*1e4)} gauss)\n" + + ref1, + size=14, + fontweight='bold', + ) + + # xlabel + axT.set_xlabel( + 'x (m)', + size=12, + fontweight='bold', + ) + + # ylabel + axT.set_ylabel( + "y (m)", + size=12, + fontweight='bold', + ) + + # xlabel + axA.set_xlabel( + 'L/D', + size=12, + fontweight='bold', + ) + + # ylabel + axA.set_ylabel( + "N", + size=12, + fontweight='bold', + ) + + axA.set_ylim(0, 1) + + # xlabel + axS.set_xlabel( + 'l (m)', + size=12, + fontweight='bold', + ) + + # ylabel + axS.set_ylabel( + "y (m)", + size=12, + fontweight='bold', + ) + + # --------------- + # add circles + # --------------- + + theta = np.pi * np.linspace(-1, 1, 101) + + # SS1008 + ss_x_out = 0.5 * dgeom['b'] * np.cos(theta) + ss_y_out = 0.5 * dgeom['b'] * np.sin(theta) + ss_x_in = 0.5 * (dgeom['a'] + dgeom['a_tol']) * np.cos(theta) + ss_y_in = 0.5 * (dgeom['a'] + dgeom['a_tol']) * np.sin(theta) + ss_x = np.r_[ss_x_out, ss_x_in[::-1], ss_x_out[0]] + ss_y = np.r_[ss_y_out, ss_y_in[::-1], ss_y_out[0]] + + axT.fill( + ss_x, + ss_y, + fc=(0.8, 0.8, 0.8, 0.5), + ) + + # mumetal + mm_x_out = 0.5 * (dgeom['a'] - dgeom['a_tol']) * np.cos(theta) + mm_y_out = 0.5 * (dgeom['a'] - dgeom['a_tol']) * np.sin(theta) + mm_x_in = 0.5 * (dgeom['a'] - dgeom['a_tol'] - dgeom['thick']) * np.cos(theta) + mm_y_in = 0.5 * (dgeom['a'] - dgeom['a_tol'] - dgeom['thick']) * np.sin(theta) + mm_x = np.r_[mm_x_out, mm_x_in[::-1], mm_x_out[0]] + mm_y = np.r_[mm_y_out, mm_y_in[::-1], mm_y_out[0]] + + axT.fill( + mm_x, + mm_y, + fc=(0.8, 0., 0.8, 0.5), + ) + + # --------------- + # add shapes - axial + # --------------- + + axS.fill( + dgeom['len0'] * np.r_[0, 1, 1, 0], + dgeom['b'] * np.r_[0, 0, 1, 1], + fc=(0.8, 0.8, 0.8), + ) + axS.fill( + dgeom['len1'] * np.r_[0, 1, 1, 0] + 0.5*(dgeom['len0'] - dgeom['len1']), + dgeom['a'] * np.r_[0, 0, 1, 1] + 0.5*(dgeom['b'] - dgeom['a']), + fc=(0.8, 0., 0.8), + ) + + # --------------- + # add expression - transverse + # --------------- + + exp0 = ( + r"$S_T = \frac{B_0}{B_{in}} = $" + + r"$\frac{(1+\mu_r)^2 - \left(\frac{a}{b}\right)^2(1-\mu_r)^2}{4\mu_r}$" + + "\n " + ) + + exp1 = ( + r"$ = \frac{B_0}{B_{inter}}\frac{B_{inter}}{B_{in}}$" + + "\n " + + f" = {int(1/dshield['mumetal'])} " + r"$\times$" + f" {int(1/dshield['SS1008'])}" + + "\n " + + f" = {ST}" + ) + + axT.text( + -0.1, -0.2, exp0 + exp1, + horizontalalignment='left', + verticalalignment='top', + transform=axT.transAxes, + size=18, + ) + + exp2 = ( + r"$B_{in} = \frac{B_0}{" + f"{ST}" + r"}$" + + r"$ = \frac{" + + f"{B0T*dshield['SS1008']*1e4:3.2} gauss" + + r"}{" + f"{int(1/dshield['mumetal'])}" + r"}$" + + "\n" + r"$B_{in} = " + f"{B0T/ST * 1e4:3.2e}" + r"\text{ gauss}$" + ) + + axT.text( + -0.1, -0.8, exp2, + horizontalalignment='left', + verticalalignment='top', + transform=axT.transAxes, + size=18, + ) + + # --------------- + # add constants + # --------------- + + const = ( + "S1008\n" + f" length = {dgeom['len0'] * 1e3} mm\n" + f" outer diam = {dgeom['b']*1e3} +/- {dgeom['b_tol']*1e3} mm\n" + f" inner diam = {dgeom['a']*1e3} +/- {dgeom['a_tol']*1e3} mm\n" + f" mur = {dmur['SS1008']}\n\n" + "mumetal\n" + f" length = {dgeom['len1'] * 1e3} mm\n" + f" outer diam = {dgeom['a']*1e3} +/- {dgeom['a_tol']*1e3} mm\n" + f" thickness = {dgeom['thick']*1e3} mm\n" + f" mur = {dmur['mumetal']}\n" + ) + + axT.text( + 0.6, -0.2, const, + horizontalalignment='left', + verticalalignment='top', + transform=axT.transAxes, + size=14, + ) + + # --------------- + # add plot - axial + # --------------- + + pp = np.linspace(1.000001, 15, 101) + NN = shielding_axial(pp) + + axA.plot( + pp, + NN, + ls='-', + c='k', + ) + + axA.axvline(p_S1008, c='k', ls='--', lw=1) + axA.axvline(p_mumetal, c='k', ls='--', lw=1) + + axA.text( + p_S1008, 1, 'S1008', + horizontalalignment='left', + verticalalignment='top', + transform=axA.transData, + rotation=90, + size=14, + ) + + axA.text( + p_mumetal, 1, 'mumetal', + horizontalalignment='left', + verticalalignment='top', + transform=axA.transData, + rotation=90, + size=14, + ) + + # --------------- + # add expression - axial + # --------------- + + exp0 = ( + r"$S_A = 4NS_T + 1$" + + "\n" + r"$p = \frac{L}{D}$" + + "\n" + + r"$N = \frac{1}{p^2-1}\left(\frac{p}{\sqrt{p^2-1}}\log\left(p + \sqrt{p^2-1}\right)-1\right)$" + + "\n " + ) + + axA.text( + 0.1, -0.2, exp0, + horizontalalignment='left', + verticalalignment='top', + transform=axA.transAxes, + size=18, + ) + + exp1 = ( + r"$S_A = $" + f"{int(SA_S1008)} x {int(SA_mumetal)} = {int(SA_S1008*SA_mumetal)}" + + "\n\n" + + r"$B_{in} = $" + f"{B0A / SA * 1e4: 3.2e} gauss" + ) + + axA.text( + 0.1, -0.8, exp1, + horizontalalignment='left', + verticalalignment='top', + transform=axA.transAxes, + size=18, + ) + + return dshield + + +# ################################################################ +# ################################################################ +# shielding formula +# ################################################################ + + +def shielding_trans(a=None, b=None, mur=None): + + num = 4 * mur + denom0 = (1 + mur)**2 + denom1 = (a/b)**2 * (1 - mur)**2 + + return num / (denom0 - denom1) + + +def shielding_axial(p=None): + + p2m1 = p**2 - 1 + sqp2m1 = np.sqrt(p2m1) + + return (1./p2m1) * ((p / sqp2m1) * np.log(p + sqp2m1) - 1) + + +# ################################################################ +# ################################################################ +# __main__ +# ################################################################ + + +if __name__ == '__main__': + dshield = main() diff --git a/tofu/physics_tools/magnetic/shielding/_single.py b/tofu/physics_tools/magnetic/shielding/_single.py new file mode 100644 index 000000000..906dba0a7 --- /dev/null +++ b/tofu/physics_tools/magnetic/shielding/_single.py @@ -0,0 +1,569 @@ +""" Analytical formulation of a 2d shielding cylinder + + +""" + + +import os +import copy + + +import numpy as np +import matplotlib.pyplot as plt +import datastock as ds + + +# ################################################################ +# ################################################################ +# DEFAULTS +# ################################################################ + + +_PATH_HERE = os.path.dirname(__file__) + + +# MAG FIELD +_B0_T = 0.010 +_BIN_LIM_T = 0.001 + + +# MAG PERMEABILITY +_MUR0 = 'SS1008' +_DMUR = { + 'SS1008': 2000, + 'mumetal': 1, # 470000, +} +_MUR = np.linspace(200, 5000, 51) + + +# DIAMETERS +_DGEOM = { + # 'a': 66e-3, + 'a': 69e-3, + 'a_tol': 0.5e-3, + # 'b': 139.725e-3, + 'b': 140e-3, + 'b_tol': 0.1e-3, + 'thick': 1.5e-3, + # length + # 'len0': 250e-3, # S1008 + 'len0': 230e-3, + 'len1': (240-24.5)*1e-3, # mumetal +} + + +# BIBLIO REF +_REF = "[1] A. Mager, IEEE Transactions on Magnetics, vol. 6, 1970" + + +# ################################################################ +# ################################################################ +# Scan inner diameter of single shield +# ################################################################ + + +def main( + # inner diameter + diam_inner=None, + diam_inner_ref=None, + diam_inner_ref_tol=None, + # outer diameter + diam_outer=None, + diam_outer_ref=None, + diam_outer_ref_tol=None, + # mag fields + B0_T=None, + Bin_lim_T=None, + Bin_levels=None, + # mag permeability + dmur0='SS1008', + # bool + plot=None, + save=None, + pfe_save=None, +): + """ Compute the magnetic shielding of an infitine hollow cylinder + + Scans the magnetic permeability and either the inner or outer diameter + The other diameter is set + + Additionally, any number of reference magnetic permeabiliies + and diameter can be indicated + + Optionally plots and saves a figure + + All diameters provided in m + All B fields provided in T + + """ + + # --------------- + # check + # --------------- + + kwd = _check(**locals()) + + # --------------- + # Compute + # --------------- + + # prepare + if kwd['kd_vary'] == 'inner': + diam_inner = kwd['diam_inner'][:, None] + diam_outer = kwd['diam_outer_ref'] + else: + diam_inner = kwd['diam_inner_ref'] + diam_outer = kwd['diam_outer'][:, None] + + # compute + shield = shielding_trans( + diam_inner=diam_inner, + diam_outer=diam_outer, + mur=kwd['mur'][None, :], + ) + + # derive Bin + Bin = kwd['B0_T'] * shield + + # prepare output + dout = copy.deepcopy(kwd) + dout.update({'Bin_T': Bin, 'shield': shield}) + + # --------------- + # plot & save + # --------------- + + dax = None + if kwd['plot'] is True: + dax = _plot( + dout=dout, + ) + + if kwd['save'] is True: + fig = dax['main'].figure + fig.savefig(kwd['pfe_save'], dpi=300) + msg = "Saved figure in:\n{kwd['pfe_save']}" + print(msg) + + return dout, dax + + +# ################################################################ +# ################################################################ +# Check +# ################################################################ + + +def _check(**kwd): + + # ---------------- + # mag field + # ---------------- + + # Background + kwd['B0_T'] = ds._generic_check._check_var( + kwd.get('B0_T'), 'B0_T', + types=(float, int), + default=_B0_T, + sign='>0', + ) + + # target + kwd['Bin_lim_T'] = ds._generic_check._check_var( + kwd.get('Bin_lim_T'), 'Bin_lim_T', + types=(float, int), + default=_BIN_LIM_T, + sign='>0', + ) + + # levels + if kwd.get('Bin_levels') is None: + pow10_base = np.floor(np.log10(kwd['Bin_lim_T'])) + pow10_min = pow10_base - 1 + pow10_max = pow10_base + 1 + kwd['Bin_levels'] = np.r_[ + 10**(pow10_min-1)*np.r_[1, 2, 5], + 10**pow10_min*np.r_[1, 2, 5], + 10**pow10_base*np.r_[1, 2, 5], + 10**pow10_max, + ] + kwd['Bin_levels'] = ds._generic_check._check_flat1darray( + kwd.get('Bin_levels'), 'Bin_levels', + dtype=float, + unique=True, + sign='>0', + ) + + # ---------------- + # mag permeability + # ---------------- + + # ----------- + # dmur0 + + dmur0 = kwd.get('dmur0') + + if dmur0 is None: + dmur0 = _MUR0 + + if isinstance(dmur0, str): + assert dmur0 in _DMUR.keys(), dmur0 + dmur0 = {dmur0: _DMUR[dmur0]} + elif isinstance(dmur0, (int, float)): + dmur0 = {'custom': dmur0} + + c0 = ( + isinstance(dmur0, dict) + and all([ + isinstance(k0, str) and isinstance(v0, (float, int)) + for k0, v0 in dmur0.items() + ]) + ) + if not c0: + msg = ( + "Arg 'dmur0' must be a dict of the form:\n" + "{\n" + "\t- 'key0': float,\n" + "\t- ... : float,\n" + "\t- 'keyn': float,\n" + "}\n" + ) + raise Exception(msg) + kwd['dmur0'] = dmur0 + + # ----------- + # mur + + mur = kwd.get('mur') + if mur is None: + mur = _MUR + + kwd['mur'] = ds._generic_check._check_flat1darray( + mur, 'mur', + dtype=float, + unique=True, + sign='>0', + ) + + # ---------------- + # inner vs outer diameter scan + # ---------------- + + lc = [ + kwd.get('diam_inner') is not None + and kwd.get('diam_outer_ref') is not None, + kwd.get('diam_outer') is not None + and kwd.get('diam_inner_ref') is not None, + ] + if np.sum(lc) != 1: + lk = ['diam_inner', 'diam_outer_ref', 'diam_outer', 'diam_inner_ref'] + lstr = [f"\t- {k0}: {kwd[k0]}" for k0 in lk] + msg = ( + "Provide either (xor):\n" + "\t- diam_inner (array) and diam_outer_ref (float)\n" + "\t- diam_outer (array) and diam_inner_ref (float)\n" + "Provided:\n" + + "\n".join(lstr) + ) + raise Exception(msg) + + # ---------------- + # Check each diam + # ---------------- + + ldiam = ['inner', 'outer'] + idiam_vary = lc.index(True) + kd_vary = ldiam[idiam_vary] + kd_fix = ldiam[1 - idiam_vary] + + # set to None unused + kwd[f'diam_{kd_fix}'] = None + kwd[f'diam_{kd_fix}_ref_tol'] = None + + # -------------- + # diam_vary + + kd = f"diam_{kd_vary}" + kwd[kd] = ds._generic_check._check_flat1darray( + kwd[kd], kd, + dtype=float, + unique=True, + sign='>0', + ) + + kd = f"diam_{kd_vary}_ref" + kwd[kd] = float(ds._generic_check._check_var( + kwd[kd], kd, + types=(int, float), + default=np.mean(kwd[f"diam_{kd_vary}"]), + sign='>0', + )) + + kd = f"diam_{kd_vary}_ref_tol" + kwd[kd] = float(ds._generic_check._check_var( + kwd[kd], kd, + types=(int, float), + default=0.001, + sign='>0', + )) + + # -------------- + # diam_fix + + kd = f"diam_{kd_fix}_ref" + kwd[kd] = float(ds._generic_check._check_var( + kwd[kd], kd, + types=(int, float), + sign='>0', + )) + + # -------------- + # store keys + + kwd['kd_vary'] = kd_vary + kwd['kd_fix'] = kd_fix + + # ---------------- + # bool + # ---------------- + + # plot + kwd['plot'] = ds._generic_check._check_var( + kwd.get('plot'), 'plot', + types=bool, + default=True, + ) + + # save + kwd['save'] = ds._generic_check._check_var( + kwd.get('save'), 'save', + types=bool, + default=False, + ) + + # pfe_save + if kwd['save'] is True: + path = os.path.abs('.') + name = 'MagneticShield_Transverse.png' + kwd['pfe_save'] = ds._generic_check_check_var( + kwd.get('pfe_save'), 'pfe_save', + types=str, + default=os.path.join(path, name), + ) + + if not os.path.isdir(kwd['pfe_save'].split()[0]): + msg = ( + "Arg 'pfe_save' seems to point to a non-existing dir:\n" + f"\t{kwd['pfe_save']}\n" + ) + raise Exception(msg) + + else: + kwd['pfe_save'] = None + + return kwd + + +# ################################################################ +# ################################################################ +# shielding formula +# ################################################################ + + +def shielding_trans(diam_inner=None, diam_outer=None, mur=None): + + num = 4 * mur + denom0 = (1 + mur)**2 + denom1 = (diam_inner / diam_outer)**2 * (1 - mur)**2 + + return num / (denom0 - denom1) + + +def shielding_axial(p=None): + + p2m1 = p**2 - 1 + sqp2m1 = np.sqrt(p2m1) + + return (1./p2m1) * ((p / sqp2m1) * np.log(p + sqp2m1) - 1) + + +# ################################################################ +# ################################################################ +# Plot +# ################################################################ + + +def _plot( + dout=None, + # figure + figsize=(12, 6), + # unused + **kwdargs, +): + + # --------------- + # prepare data + # --------------- + + kvary = dout['kd_vary'] + kfix = dout['kd_fix'] + diam_vary = dout[f"diam_{kvary}"] + diam_vary_ref = dout[f"diam_{kvary}_ref"] + diam_vary_ref_tol = dout[f"diam_{kvary}_ref_tol"] + diam_fix_ref = dout[f"diam_{kfix}_ref"] + + # --------------- + # prepare figure + # --------------- + + fig = plt.figure(figsize=figsize) + ax = fig.add_axes([.06, 0.08, 0.6, 0.80]) + + # title + ax.set_title( + "Transverse B field in a 2d infinite shielding cylinder\n" + f"{_REF}\n" + f"{kfix.capitalize()} diameter = {diam_fix_ref} m\n" + f"External field B0 = {dout['B0_T']} T ({int(dout['B0_T']*1e3)} mT)", + size=14, + fontweight='bold', + ) + + # xlabel + ax.set_xlabel( + f"{kvary.capitalize()} diameter (m)", + size=12, + fontweight='bold', + ) + + # ylabel + ax.set_ylabel( + "Relative permeability " + r"$\mu_r$", + size=12, + fontweight='bold', + ) + + # --------------- + # contours + # --------------- + + # contourf plot + ax.contourf( + diam_vary, + dout['mur'], + dout['Bin_T'].T*1e3, + [0, dout['Bin_lim_T']*1e3, np.max(dout['Bin_T']*1e3)], + cmap='RdYlGn_r', + ) + + # contour plot + cs = ax.contour( + diam_vary, + dout['mur'], + dout['Bin_T'].T*1e3, + dout['Bin_levels']*1e3, + colors='k', + ) + + # labelled contours + ax.clabel( + cs, + cs.levels, + inline=True, + fontsize=14, + fmt=lambda x: f'{x} mT', + ) + + # --------------- + # add mur0 + # --------------- + + for k0, v0 in dout['dmur0'].items(): + ax.axhline( + v0, + ls='--', + c='k', + lw=1, + ) + + ax.text( + diam_vary[-1], v0, k0, + horizontalalignment='right', + verticalalignment='bottom', + transform=ax.transData, + size=12, + fontweight='bold', + ) + + # ------------------ + # add inner diameter + # ------------------ + + ax.axvspan( + diam_vary_ref - diam_vary_ref_tol, + diam_vary_ref + diam_vary_ref_tol, + fc=(0.8, 0.8, 0.8, 0.5), + ) + + # --------------- + # add expression + # --------------- + + exp = ( + r"$\frac{B_{in}}{B_0} = $" + + r"$\frac{4\mu_r}{(1+\mu_r)^2 - \left(\frac{a}{b}\right)^2(1-\mu_r)^2}$" + ) + + ax.text( + 1.28, 0.3, exp, + horizontalalignment='center', + transform=ax.transAxes, + size=24, + ) + + # shield value + kmur0 = list(dout['dmur0'].keys())[0] + imur = np.argmin(np.abs(dout['mur'] - dout['dmur0'][kmur0])) + ia = np.argmin(np.abs(diam_vary - diam_vary_ref)) + + ax.plot( + [diam_vary[ia]], + [dout['mur'][imur]], + marker='s', + ms=12, + c='r' + ) + + exp = ( + r"$\frac{B_{in}}{B_0} = $" + + f"{dout['shield'][ia, imur]:3.2e}" + ) + + ax.text( + 1.28, 0.05, exp, + horizontalalignment='center', + transform=ax.transAxes, + size=24, + ) + + # ----------------- + # add image + # ----------------- + + pfe = os.path.join(_PATH_HERE, '2dproblem.png') + image = plt.imread(pfe) + axim = fig.add_axes([0.72, 0.4, 0.25, 0.55]) + axim.imshow(image) + axim.axis('off') + + return {'main': ax} + + +# ################################################################ +# ################################################################ +# __main__ +# ################################################################ + + +if __name__ == '__main__': + main() diff --git a/tofu/spectro/__init__.py b/tofu/spectro/__init__.py index 7eeedda33..e98310657 100644 --- a/tofu/spectro/__init__.py +++ b/tofu/spectro/__init__.py @@ -7,4 +7,4 @@ from ._plot import * from ._rockingcurve import * from ._rockingcurve_tools import * -from ._spectralrange2d import * +from ._spectralrange2d import main as spectral_range_2d diff --git a/tofu/spectro/_spectralrange2d.py b/tofu/spectro/_spectralrange2d.py index 95d90eb80..ecf100a52 100644 --- a/tofu/spectro/_spectralrange2d.py +++ b/tofu/spectro/_spectralrange2d.py @@ -2,12 +2,11 @@ import numpy as np -import matplotlib.pyplot as plt -import matplotlib.gridspec as gridspec -import datastock as ds -__all__ = ['spectral_range_2d'] +from ._spectralrange2d_check import main as _check +from ._spectralrange2d_compute import main as _compute +from . import _spectralrange2d_plot as _plot # ################################################################# @@ -16,23 +15,17 @@ # ################################################################# -def spectral_range_2d( - # crystal - lamb0=None, - bragg0=None, +def main( + # optics + dap=None, + dcrystals=None, + dcam=None, + # matching + dmatch=None, + # large scans + dscans=None, # geometry basis - ap=None, - ex=None, - ey=None, beta_max=None, - # geometry - xx=None, - length=None, - rcurve=None, - varrad_b=None, - dist=None, - # camera - dcam=None, # options npts=None, # plotting @@ -56,747 +49,103 @@ def spectral_range_2d( beta_max: maximum angular opening from point source (optionnal) npts: nb of rays from point source to crystals length: crystal length + varrad: for variable-radii spiral """ - # -------- - # check + # ------------- + # check inputs + # ------------- - din, npts, plot, save, pfe_fig, pfe_npz = _check(**locals()) + ( + dap, dcrystals, dcam, dmatch, + dscans, npts, + plot, save, pfe_fig, pfe_npz, + ) = _check(**locals()) # -------------- # compute + # -------------- - crystx, crysty, endx, endy, lamb = _compute( - npts=npts, - beta_max=beta_max, - dcam=dcam, - **din, - ) + dout = _compute(npts=npts, **dscans) - # ------------- - # format output - - - ilamb_min = np.full((lamb.shape[1],), -1) - ilamb_max = np.full((lamb.shape[1],), -1) - iok = np.any(np.isfinite(lamb), axis=0) - ilamb_min[iok] = np.nanargmin(lamb[:, iok], axis=0) - ilamb_max[iok] = np.nanargmax(lamb[:, iok], axis=0) - - lamb_min = np.array([ - lamb[imin, ii] if imin >= 0 else np.nan - for ii, imin in enumerate(ilamb_min) - ]) - lamb_max = np.array([ - lamb[imax, ii] if imax >= 0 else np.nan - for ii, imax in enumerate(ilamb_max) - ]) - - dout = dict(din) - dout.update({ - 'beta_max': beta_max, - 'crystx': crystx, - 'crysty': crysty, - 'endx': endx, - 'endy': endy, - 'lamb': lamb, - 'ilamb_min': ilamb_min, - 'ilamb_max': ilamb_max, - 'lamb_min': lamb_min, - 'lamb_max': lamb_max, - 'Dlamb': lamb_max - lamb_min, - }) - - if dcam is not None: - dout['dcam'] = dcam + # -------------- + # extract + # -------------- + + dout_match = None + if dmatch is not None: + dout_match = _dout_match(dout=dout, dmatch=dmatch) # --------- # plot + # --------- if plot is True: - dax = _plot( - dax=dax, - pfe_fig=pfe_fig, - **dout, - ) + if dmatch is None: + dax = _plot.scans( + dax=dax, + pfe_fig=pfe_fig, + **dout, + ) + else: + dax = _plot.match( + dax=dax, + pfe_fig=pfe_fig, + dout=dout, + dmatch=dmatch, + dap=dap, + dcam=dcam, + dscans=dscans, + ) # ---------- # save + # ---------- + + dout0 = dout if dmatch is None else dout_match if save is True: - np.savez(pfe_npz, **dout) + np.savez(pfe_npz, **dout0) + msg = f"Saved in:\n\t{pfe_npz}" + print(msg) # --------- # return if plot is True: - return dout, dax + return dout0, dax else: - return dout - - -# ################################################################# -# ################################################################# -# Check -# ################################################################# - - -def _check( - # crystal - lamb0=None, - bragg0=None, - # geometry basis - ap=None, - ex=None, - ey=None, - # geometry - xx=None, - length=None, - rcurve=None, - varrad_b=None, - dist=None, - # options - npts=None, - # plotting - plot=None, - ax=None, - # saving - save=None, - pfe_fig=None, - pfe_npz=None, - # unused - **kwdargs, -): - - # -------------- - # geometry basis - - basis_def = {'ap': np.r_[0, 0], 'ex': np.r_[1, 0], 'ey': np.r_[0, 1]} - din_basis = {'ap': ap, 'ex': ex, 'ey': ey} - for k0, v0 in din_basis.items(): - - if v0 is None: - din_basis[k0] = basis_def[k0] - - din_basis[k0] = np.atleast_1d(din_basis[k0]).ravel().astype(float)[:2] - - # normalize ex - din_basis['ex'] = din_basis['ex'] / np.linalg.norm(din_basis['ex']) - - # perpendicular + normalize ey - sca = np.sum(din_basis['ex']*din_basis['ey']) - din_basis['ey'] = din_basis['ey'] - sca * din_basis['ex'] - din_basis['ey'] = din_basis['ey'] / np.linalg.norm(din_basis['ey']) - - # ----------------- - # initialize dict - - din = { - 'lamb0': lamb0, - 'bragg0': bragg0, - # geometry - 'xx': xx, - 'length': length, - 'rcurve': rcurve, - 'varrad_b': varrad_b, - 'dist': dist, - } - - # ------------- - # get size - - # make all arrays - for k0, v0 in din.items(): - if v0 is None: - din[k0] = np.nan - din[k0] = np.atleast_1d(din[k0]).ravel().astype(float) - - # sizes - lsizes = list(set([v0.size for v0 in din.values()])) - if len(lsizes) == 1: - pass - elif 1 in lsizes and len(lsizes) == 2: - size = [ss for ss in lsizes if ss != 1][0] - for k0, v0 in din.items(): - if v0.size == 1: - din[k0] = np.full((size,), v0[0]) - else: - lstr = [f"\t- '{k0}': {v0.size}" for k0, v0 in din.items()] - msg = ( - "All args must be either scalar or 1d arrays of the same size:\n" - + "\n".join(lstr) - ) - raise Exception(msg) - - # ------- - # values - - for k0, v0 in din.items(): - if k0 not in ['rcurve', 'varrad_b']: - c0 = np.all(np.isfinite(v0)) and np.all(v0 >= 0.) - - if not c0: - msg = ( - f"Arg '{k0}' must be finite and positive\n" - f"Provided: {v0}" - ) - raise Exception(msg) - - # ------------ - # add basis - - din.update(din_basis) - - # --------- - # npts - - if npts is None: - npts = 101 - npts = int(npts) - if npts % 2 == 0: - npts += 1 - - # --------- - # plot - - # plot - plot = ds._generic_check._check_var( - plot, 'plot', - types=bool, - default=True, - ) - - # --------- - # save - - # save - save = ds._generic_check._check_var( - save, 'save', - types=bool, - default=False, - ) - - return din, npts, plot, save, pfe_fig, pfe_npz - - -# ################################################################# -# ################################################################# -# Compute -# ################################################################# - - -def _compute( - # crystal - lamb0=None, - bragg0=None, - # geometry basis - ap=None, - ex=None, - ey=None, - beta_max=None, - # geometry - xx=None, - length=None, - rcurve=None, - varrad_b=None, - dist=None, - # options - npts=None, - # camera - dcam=None, -): - - # ------------ - # initialize - - size = lamb0.size - - crystx = np.full((npts, size), np.nan) - crysty = np.full((npts, size), np.nan) - vnx = np.full((npts, size), np.nan) - vny = np.full((npts, size), np.nan) - - # ---------------- - # compute geometry - # ---------------- - - # 2d - d2 = lamb0 / np.sin(bragg0) - - # summit of crystal - sx = ap[0] + xx * ex[0] - sy = ap[1] + xx * ex[1] - - # ------------------------ - # indices of crystal types + return dout0 - # variable-radii sinusoidal spiral - indb = np.isfinite(varrad_b) - - # indices of curved crystals - indc = np.isfinite(rcurve) & (~indb) - - # flat - indf = (~indc) & (~indb) - - # safety check - if not np.all(np.sum([indc, indb, indf], axis=0) == 1): - msg = ( - "Some undetermined 2d crystal shapes:\n" - f"\t- indc = {indc}\n" - f"\t- indb = {indb}\n" - f"\t- indf = {indf}\n" - ) - raise Exception(msg) - - # --------------------- - # curved crystals - - # center of curvature - ecx = np.sin(bragg0[indc]) * ex[0] - np.cos(bragg0[indc]) * ey[0] - ecy = np.sin(bragg0[indc]) * ex[1] - np.cos(bragg0[indc]) * ey[1] - ecx_p = -ecy - ecy_p = ecx - - cx = sx[indc] - rcurve[indc] * ecx - cy = sy[indc] - rcurve[indc] * ecy - - # half angular opening of crystal - dalpha = 0.5*length[indc] / rcurve[indc] - theta = dalpha * np.linspace(-1, 1, npts)[:, None] - - # crystal plotting - curved - ethetax = np.cos(theta) * ecx[None, :] + np.sin(theta) * ecx_p[None, :] - ethetay = np.cos(theta) * ecy[None, :] + np.sin(theta) * ecy_p[None, :] - - crystx[:, indc] = cx[None, :] + rcurve[indc][None, :] * ethetax - crysty[:, indc] = cy[None, :] + rcurve[indc][None, :] * ethetay - - # local normal vectors - vnx[:, indc] = -ethetax - vny[:, indc] = -ethetay - - # ----------------------- - # flat crystals - - # crystal plotting - straight - estraightx = np.cos(bragg0)[indf] * ex[0] + np.sin(bragg0)[indf] * ey[0] - estraighty = np.cos(bragg0)[indf] * ex[1] + np.sin(bragg0)[indf] * ey[1] - - ll = 0.5 * length[None, indf] * np.linspace(-1, 1, npts)[:, None] - crystx[:, indf] = sx[None, indf] + ll*estraightx[None, :] - crysty[:, indf] = sy[None, indf] + ll*estraighty[None, :] - - # local normal vectors - vnx[:, indf] = -estraighty - vny[:, indf] = estraightx - - # ----------------------- - # variable radii crystals - - # main parameters - gam0 = bragg0[indb] - r0 = rcurve[indb] - ix = ~np.isfinite(r0) - r0[ix] = xx[indb][ix] - b = varrad_b[indb] - - # local radius of curvature at center - # rc0 = r0 / (b * np.sin(gam0)) - - # dOMx = r / (b-1) * (cos(phi) / tan(gam) - sin(phi)) - # dOMy = r / (b-1) * (sin(phi) / tan(gam) + cos(phi)) - # dL = r/(b-1) * 1 / sin(gam) - # dL ~ r0/(b-1) * 1/sin(gam0) * Dgam - - # half angular opening of crystal (approximative) - # dgam = 0.5*length / rc0 - dgam = 1.1 * length[indb] * np.sin(gam0) * (b-1) / r0 / 2 - - # gam - gam = gam0[None, :] + dgam[None, :] * np.linspace(-1, 1, npts)[:, None] - - # r - r = r0[None, :] * (np.sin(gam) / np.sin(gam0)[None, :])**(1/(b[None, :]-1)) - - # phi - phi = (gam - gam0[None, :]) / (b[None, :]-1) - - # cryst - - crystx[:, indb] = ( - ap[0] - + (xx[indb] - r0) * ex[0] - + r * (np.cos(phi) * ex[0] + np.sin(phi) * ey[0]) - ) - crysty[:, indb] = ( - ap[1] - + (xx[indb] - r0) * ex[1] - + r * (np.cos(phi) * ex[1] + np.sin(phi) * ey[1]) - ) - - # derivative - c0 = r / (b[None, :] - 1) - c1 = np.cos(gam) / np.sin(gam) - dOMxx = c0 * (c1 * np.cos(phi) - np.sin(phi)) - dOMyy = c0 * (c1 * np.sin(phi) + np.cos(phi)) - dOMx = dOMxx * ex[0] + dOMyy * ey[0] - dOMy = dOMxx * ex[1] + dOMyy * ey[1] - - # local normal vectors - vnx[:, indb] = dOMy / np.sqrt(dOMx**2 + dOMy**2) - vny[:, indb] = -dOMx / np.sqrt(dOMx**2 + dOMy**2) - - # ---------------- - # compute rays - # ---------------- - - # vectors of incident rays - vix = crystx - ap[0] - viy = crysty - ap[1] - vin = np.sqrt(vix**2 + viy**2) - vix = vix / vin - viy = viy / vin - - - # reflected vectors - sca = vix*vnx + viy*vny - vrx = vix - 2.*sca*vnx - vry = viy - 2.*sca*vny - - # end of rays at dist - endx = crystx + dist * vrx - endy = crysty + dist * vry - - # ---------------------- - # compute spectral range - - # get local bragg angle - top and bottom - bragg = np.arccos(sca) - np.pi/2. - - # lamb - lamb = d2 * np.sin(bragg) - - # beta_max - if beta_max is not None: - dvx, dvy = crystx - ap[0], crysty - ap[1] - beta = np.arctan2(dvx*ey[0] + dvy*ey[1], dvx*ex[0] + dvy*ex[1]) - ind = np.abs(beta) > beta_max - endx[ind] = np.nan - endy[ind] = np.nan - lamb[ind] = np.nan - - # ----------------- - # impacts on camera - # ----------------- - - if dcam is not None: - ninx, niny = dcam['nin'][:2] - ninn = np.sqrt(ninx**2 + niny**2) - ninx, niny = ninx/ninn, niny/ninn - - ninx_r = ninx * ex[0] + niny * ey[0] - niny_r = ninx * ex[1] + niny * ey[1] - - camx = ap[0] + dcam['cent'][0] * ex[0] + dcam['cent'][1] * ey[0] - camy = ap[1] + dcam['cent'][0] * ex[1] + dcam['cent'][1] * ey[1] - - sca_up = (camx - crystx) * ninx_r + (camy - crysty) * niny_r - sca_bot = vrx*ninx_r + vry*niny_r - - kk = sca_up / sca_bot - ptsx = crystx + kk * vrx - ptsy = crysty + kk * vry - - e0x = -niny_r - e0y = ninx_r - x0 = (ptsx - camx) * e0x + (ptsy - camy) * e0y - - if beta_max is not None: - x0[ind] = np.nan - - dcam['x0'] = x0 - dcam['cent_r'] = np.r_[camx, camy] - dcam['nin_r'] = np.r_[ninx_r, niny_r] - - return crystx, crysty, endx, endy, lamb - - -# ################################################################# -# ################################################################# -# Plot -# ################################################################# - - -def _plot( - # crystal - lamb0=None, - bragg0=None, - # geometry - xx=None, - length=None, - rcurve=None, - dist=None, - beta_max=None, - # computed - ap=None, - crystx=None, - crysty=None, - endx=None, - endy=None, - ilamb_min=None, - ilamb_max=None, - lamb_min=None, - lamb_max=None, - Dlamb=None, - x0=None, - # camera - dcam=None, - # plotting - dax=None, - # saving - pfe_fig=None, - # unused - **kwdargs, -): - - # ---------- - # prepare - - npts, size = crystx.shape - - # envelop - iok = np.isfinite(endx) - i0 = tuple([iok[:, ii].nonzero()[0][0] for ii in range(size)]) - i1 = tuple([iok[:, ii].nonzero()[0][-1] for ii in range(size)]) - nind = tuple(range(size)) - - # envelop - envx = np.array([ - endx[i1, nind], crystx[i1, nind], - np.full((size,), ap[0]), - crystx[i0, nind], endx[i0, nind], - ]) - envy = np.array([ - endy[i1, nind], crysty[i1, nind], - np.full((size,), ap[1]), - crysty[i0, nind], endy[i0, nind], - ]) - - # central rays - ind = int((npts-1)/2) - raycx = np.array([np.full((size,), ap[0]), crystx[ind, :], endx[ind, :]]) - raycy = np.array([np.full((size,), ap[1]), crysty[ind, :], endy[ind, :]]) - - # dcam - if dcam is not None: - ninx, niny = dcam['nin_r'][:2] - ninn = np.sqrt(ninx**2 + niny**2) - ninx, niny = ninx/ninn, niny/ninn - e0x, e0y = -niny, ninx - e0n = np.sqrt(e0x**2 + e0y**2) - e0x, e0y = e0x/e0n, e0y/e0n - clen = dcam['length'] - camx = dcam['cent_r'][0] + 0.5*clen*np.r_[-1, 1] * e0x - camy = dcam['cent_r'][1] + 0.5*clen*np.r_[-1, 1] * e0y - - # -------------- - # prepare figure - - if dax is None: - dax = _get_axes() - - # ----------- - # plot - - color = None - for ii in range(size): - - # --- - # hor - - kax = 'hor' - if dax.get(kax) is not None: - ax = dax[kax]['handle'] - - # crystals - ll, = ax.plot( - crystx[:, ii], - crysty[:, ii], - ls='-', - lw=2, - marker='None', - ) - color = ll.get_color() - - # central rays - ax.plot( - raycx[:, ii], - raycy[:, ii], - ls='--', - lw=1, - marker='None', - c=color, - ) - - # edge rays - ax.plot( - envx[:, ii], - envy[:, ii], - ls='-', - lw=1, - marker='None', - c=color, - label=( - f"r = {rcurve[ii]} m\t" - + r"$\lambda_0$" + f" = {lamb0[ii]*1e10:5.3f} AA\t" - + r"$\beta_0$" + f" = {bragg0[ii]*180/np.pi:5.2f} deg" - ), - ) - - kax = 'cam' - if dcam is not None and dax.get(kax) is not None: - ax = dax[kax]['handle'] - - # images - ax.plot( - dcam['x0'][:, ii], - np.full((npts,), ii+1), - ls='None', - marker='.', - color=color, - ms=6, - ) - - # lamb min, max - ax.text( - dcam['x0'][ilamb_min[ii], ii], - ii + 1 - 0.1, - f'{lamb_min[ii]*1e10:2.3} AA', - color=color, - size=8, - horizontalalignment='center', - verticalalignment='top', - ) - - ax.text( - dcam['x0'][ilamb_max[ii], ii], - ii + 1 - 0.1, - f'{lamb_max[ii]*1e10:2.3} AA', - color=color, - size=8, - horizontalalignment='center', - verticalalignment='top', - ) - - # --------------- - # plot input data - - kax = 'hor' - if dax.get(kax) is not None: - ax = dax[kax]['handle'] - ax.legend(fontsize=12) - - if beta_max is None: - beta_str = 'None' - else: - beta_str = f'{beta_max*180/np.pi:5.3} deg' - - msg = ( - f"beta_max = {beta_str}\n" - ) - - ax.text( - 0.8, - 0.4, - msg, - color='k', - size=10, - horizontalalignment='center', - verticalalignment='top', - transform=ax.figure.transFigure, - ) - - - # ------------ - # camera - - if dcam is not None: - - kax = 'hor' - if dax.get(kax) is not None: - ax = dax[kax]['handle'] - - ax.plot( - camx, - camy, - ls='-', - lw=2., - marker='None', - c='k', - ) - - kax = 'cam' - if dax.get(kax) is not None: - ax = dax[kax]['handle'] - ax.axvline(-0.5*dcam['length'], c='k', ls='-', lw=1.) - ax.axvline(0.5*dcam['length'], c='k', ls='-', lw=1.) - ax.set_ylim(0, size + 1) - - # ---------- - # saving - - if pfe_fig is not None: - dax['hor']['handle'].figure.savefig(pfe_fig, format='png', dpi=200) - - return dax - - -def _get_axes(): - - # -------------- - # prepare figure - - dmargin = { - 'left': 0.08, 'right': 0.98, - 'bottom': 0.08, 'top': 0.90, - 'hspace': 0.20, 'wspace': 0.25, - } - - fig = plt.figure(figsize=(13, 8)) - fig.suptitle('2d ray-tracing model') - gs = gridspec.GridSpec(ncols=3, nrows=2, **dmargin) - - # ---------- - # make axes - # ax0 - hor - ax0 = fig.add_subplot( - gs[:, :-1], - aspect='equal', - adjustable='datalim', - ) +# ############################################ +# ############################################ +# Extract dmatch +# ############################################ - ax0.set_xlabel("x (m)", size=12) - ax0.set_ylabel("y (m)", size=12) - ax0.set_title("2d ray tracing", size=12, fontweight='bold') - # ax1 - cam - ax1 = fig.add_subplot( - gs[0, -1], - aspect='auto', - ) +def _dout_match(dout=None, dmatch=None): - ax1.set_xlabel("x0 (m)", size=12) - ax1.set_title("Image on camera", size=12, fontweight='bold') + dout_match = {k0: {k1: {} for k1 in dout.keys()} for k0 in dmatch.keys()} + for k0, v0 in dmatch.items(): + for k1, v1 in dout.items(): - # ------------ - # dict + # array + if isinstance(v1, np.ndarray): + if v1.ndim == dout['cryst0'].ndim: + sli = (slice(None),) + v0['ind'] + else: + sli = v0['ind'] + dout_match[k0][k1] = v1[sli] - dax = { - 'hor': {'handle': ax0}, - 'cam': {'handle': ax1}, - } + # dict + else: + for k2, v2 in v1.items(): + if v2.ndim == dout['cryst0'].ndim: + sli = (slice(None),) + v0['ind'] + else: + sli = v0['ind'] + dout_match[k0][k1][k2] = v2[sli] - return dax \ No newline at end of file + return dout_match diff --git a/tofu/spectro/_spectralrange2d_check.py b/tofu/spectro/_spectralrange2d_check.py new file mode 100644 index 000000000..4d06257a0 --- /dev/null +++ b/tofu/spectro/_spectralrange2d_check.py @@ -0,0 +1,1083 @@ +import os + + +import numpy as np +import matplotlib.colors as mcolors +import datastock as ds + + +# ###################################### +# ###################################### +# DEFAULTS +# ###################################### + + +_NPTS = 21 + + +_DSCANS = { + # aperture + 'ap0': (float, 0), + 'ap1': (float, 0), + 'ex0': (float, 1), + 'ex1': (float, 0), + 'ey0': (float, 0), + 'ey1': (float, 1), + 'semi_angle_max': (float, np.nan), + # crystal + 'dist_from_ap': (float, '>0'), + 'lamb0': (float, '>0'), + 'bragg0': (float, '>0'), + 'rcurve': (float, np.inf), + 'length': (float, '>0'), + 'varrad_b': (float, np.nan), + 'lamb0_min': (float, np.nan), + 'lamb0_max': (float, np.nan), + # camera + 'cam_c0': (float,), + 'cam_c1': (float,), + 'cam_nin0': (float,), + 'cam_nin1': (float,), + 'cam_length': (float,), + # options + # 'npts': (int, 31), +} + + +# ###################################### +# ###################################### +# Main check function +# ###################################### + + +def main( + # apertures, crystals, cameras + dap=None, + dcrystals=None, + dcam=None, + # matches + dmatch=None, + # large scans + dscans=None, + # options + npts=None, + # plotting + plot=None, + dax=None, + # saving + save=None, + pfe_fig=None, + pfe_npz=None, + # unused + **kwdargs, +): + + # --------- + # npts + # --------- + + npts = int(ds._generic_check._check_var( + npts, 'npts', + types=(float, int), + sign='>0', + default=_NPTS, + )) + if npts % 2 == 0: + npts += 1 + + # -------------- + # dscans vs the rest + # -------------- + + lc = [ + dscans is not None, + all([dd is not None for dd in [dap, dcrystals, dcam]]), + ] + if np.sum(lc) != 1: + msg = ( + "Provide either (xor):\n" + "\t- dscans: dict of numpy arrays for large sets\n" + "\t- {dap dcrystals, dcam, (dmatch)}: for details\n" + ) + raise Exception(msg) + + # -------------- + # details => derive dscans + # -------------- + + if lc[1]: + + # -------------- + # dap + + _dap(dap) + + # ----------------- + # dcrystals + + _dcrystals(dcrystals) + + # -------------- + # dcam + + # (cent, nin) or from_cryst[dist, angle] + _dcam(dcam, dap=dap, dcrystals=dcrystals) + + # -------------- + # dmatch + + dmatch = _dmatch( + dmatch=dmatch, + dap=dap, + dcrystals=dcrystals, + dcam=dcam, + npts=npts, + ) + + # --------- + # derive dscans + + dscans = _derive_dscans( + dap=dap, + dcrystals=dcrystals, + dcam=dcam, + dmatch=dmatch, + ) + + # -------------- + # check dscans + # -------------- + + _dscans(dscans) + + # --------- + # plot + # --------- + + # plot + plot = ds._generic_check._check_var( + plot, 'plot', + types=bool, + default=True, + ) + + # --------- + # save + # --------- + + # save + save = ds._generic_check._check_var( + save, 'save', + types=bool, + default=False, + ) + + # --------- + # pfe + # --------- + + if save is True: + pfe_fig, pfe_npz = _pfe( + pfe_fig=pfe_fig, + pfe_npz=pfe_npz, + ) + else: + pfe_fig = None + pfe_npz = None + + return ( + dap, dcrystals, dcam, dmatch, + dscans, npts, + plot, save, pfe_fig, pfe_npz, + ) + + +# ###################################### +# ###################################### +# Apertures check function +# ###################################### + + +def _dap(dap): + + # ---------------- + # basics + # ---------------- + + c0 = ( + isinstance(dap, dict) + and all([isinstance(v0, dict) for v0 in dap.values()]) + ) + if not c0: + _err_dap(dap) + + # ------------------- + # loop on key, values + # ------------------- + + dfail = {} + for i0, (k0, v0) in enumerate(dap.items()): + + try: + + # --------------- + # cent + + if dap[k0].get('cent') is None: + dap[k0]['cent'] = np.r_[0, 0] + + dap[k0]['cent'] = ds._generic_check._check_flat1darray( + dap[k0]['cent'], + f"dap['{k0}']['cent']", + dtype=float, + size=2, + ) + + # --------------- + # ex + + if dap[k0].get('ex') is None: + dap[k0]['ex'] = np.r_[1, 0] + + dap[k0]['ex'] = ds._generic_check._check_flat1darray( + dap[k0]['ex'], + f"dap['{k0}']['ex']", + dtype=float, + size=2, + norm=True, + ) + + # --------------- + # ey + + if dap[k0].get('ey') is None: + dap[k0]['ey'] = np.r_[-dap[k0]['ex'][1], dap[k0]['ex'][0]] + + dap[k0]['ey'] = ds._generic_check._check_flat1darray( + dap[k0]['ey'], + f"dap['{k0}']['ey']", + dtype=float, + size=2, + norm=True, + ) + + dap[k0]['ey'] -= np.sum(dap[k0]['ey']*dap[k0]['ex'])*dap[k0]['ex'] + dap[k0]['ey'] = dap[k0]['ey'] / np.linalg.norm(dap[k0]['ey']) + + # --------------- + # semi_angle_max + + if dap[k0].get('semi_angle_max') is not None: + dap[k0]['semi_angle_max'] = float( + ds._generic_check._check_var( + dap[k0]['semi_angle_max'], + f"dap['{k0}']['semi_angle_max']", + types=(float, int), + sign=['>0', '<1.57'], + ) + ) + + # --------------- + # label + + dap[k0]['label'] = ds._generic_check._check_var( + dap[k0].get('label'), + f"dap['{k0}']['label']", + types=str, + default=str(k0), + ) + + # --------------- + # color + + if dap[k0].get('color') is None: + dap[k0]['color'] = 'k' + if not mcolors.is_color_like(dap[k0]['color']): + msg = f"dap['{k0}']['color'] not color-like!" + raise Exception(msg) + dap[k0]['color'] = mcolors.to_rgba(dap[k0]['color']) + + except Exception as err: + dfail[k0] = str(err) + + # ------------------- + # raise errors if any + # ------------------- + + if len(dfail) > 0: + lstr = [f"\t- {k0}: {v0}" for k0, v0 in dfail.items()] + msg = "\n".join(lstr) + _err_dap(dap, errstr=msg) + + return + + +def _err_dap(dap, errstr=''): + msg = ( + "Arg dap must be a dict of sub-dicts of the form:\n" + "\t- 'key0': {\n" + "\t\t'cent': array of 2 floats, (default to [0, 0])\n" + "\t\t'ex': array of 2 floats, normalized (default to [1, 0])\n" + "\t\t'ey': array of 2 floats, normalized (default to [0, 1])\n" + "\t\t'semi_angle_max': None / float, max opening of ap\n" + "\t\t'color': color-like, (optional)\n" + "\t\t'label': str, (optional)\n" + "\t}\n\n" + + errstr + + f"\n\nProvided:\n{dap}\n" + ) + raise Exception(msg) + + +# ###################################### +# ###################################### +# Crystals check function +# ###################################### + + +def _dcrystals(dcrystals): + + # ---------------- + # basics + # ---------------- + + c0 = ( + isinstance(dcrystals, dict) + and all([isinstance(v0, dict) for v0 in dcrystals.values()]) + ) + if not c0: + _err_dcrystals(dcrystals) + + # ------------------- + # loop on key, values + # ------------------- + + dfail = {} + for i0, (k0, v0) in enumerate(dcrystals.items()): + + try: + # --------------- + # bragg0 + + dcrystals[k0]['bragg0'] = float(ds._generic_check._check_var( + dcrystals[k0].get('bragg0'), + f"dcrystals['{k0}']['bragg0']", + types=(int, float), + sign=[">0", "<1.5708"], + )) + + # --------------- + # lamb0 + + dcrystals[k0]['lamb0'] = float(ds._generic_check._check_var( + dcrystals[k0].get('lamb0'), + f"dcrystals['{k0}']['lamb0']", + types=(int, float), + sign=[">0"], + )) + + # --------------- + # lamb0_min + + if dcrystals[k0].get('lamb0_min') is not None: + dcrystals[k0]['lamb0_min'] = float( + ds._generic_check._check_var( + dcrystals[k0].get('lamb0_min'), + f"dcrystals['{k0}']['lamb0_min']", + types=(int, float), + sign=[">0", f"<{dcrystals[k0]['lamb0']}"], + ) + ) + else: + dcrystals[k0]['lamb0_min'] = np.nan + + # --------------- + # lamb0_max + + if dcrystals[k0].get('lamb0_max') is not None: + dcrystals[k0]['lamb0_max'] = float( + ds._generic_check._check_var( + dcrystals[k0].get('lamb0_max'), + f"dcrystals['{k0}']['lamb0_max']", + types=(int, float), + sign=[">0", f">{dcrystals[k0]['lamb0']}"], + ) + ) + else: + dcrystals[k0]['lamb0_max'] = np.nan + + # --------------- + # rcurve + + dcrystals[k0]['rcurve'] = float(ds._generic_check._check_var( + dcrystals[k0].get('rcurve'), + f"dcrystals['{k0}']['rcurve']", + types=(int, float), + default=np.inf, + )) + + # --------------- + # xx + + dcrystals[k0]['dist_from_ap'] = float( + ds._generic_check._check_var( + dcrystals[k0].get('dist_from_ap'), + f"dcrystals['{k0}']['dist_from_ap']", + types=(int, float), + sign='>0.', + ) + ) + + # --------------- + # length + + dcrystals[k0]['length'] = float(ds._generic_check._check_var( + dcrystals[k0].get('length'), + f"dcrystals['{k0}']['length']", + types=(int, float), + sign='>0.', + )) + + # --------------- + # varrad_b + + dcrystals[k0]['varrad_b'] = float(ds._generic_check._check_var( + dcrystals[k0].get('varrad_b'), + f"dcrystals['{k0}']['varrad_b']", + types=(int, float), + default=np.nan, + )) + + # --------------- + # label + + dcrystals[k0]['label'] = ds._generic_check._check_var( + dcrystals[k0].get('label'), + f"dcrystals['{k0}']['label']", + types=str, + default=str(k0), + ) + + # --------------- + # color + + if dcrystals[k0].get('color') is None: + dcrystals[k0]['color'] = 'k' + if not mcolors.is_color_like(dcrystals[k0]['color']): + msg = f"dcrystals['{k0}']['color'] not color-like!" + raise Exception(msg) + dcrystals[k0]['color'] = mcolors.to_rgba(dcrystals[k0]['color']) + + except Exception as err: + dfail[k0] = str(err) + + # ------------------- + # raise errors if any + # ------------------- + + if len(dfail) > 0: + lstr = [f"\t- {k0}: {v0}" for k0, v0 in dfail.items()] + msg = "\n".join(lstr) + _err_dcrystals(dcrystals, errstr=msg) + + return + + +def _err_dcrystals(dcrystals, errstr=''): + msg = ( + "Arg dcrystals must be a dict of sub-dicts of the form:\n" + "\t- 'key0': {\n" + "\t\t'lamb0': float, (m)\n" + "\t\t'lamb0_min': float, (optional)\n" + "\t\t'lamb0_max': float, (optional)\n" + "\t\t'bragg0': float, (rad)\n" + "\t\t'rcurve': float, (inf if flat, +/-float if concave/convex)\n" + "\t\t'xx': float, (m, distance from aperture)\n" + "\t\t'length': float, (m, crystal length)\n" + "\t\t'dist': float, (m, crystal-to-camera distance)\n" + "\t\t'varrad_b': float, (m, ??)\n" + "\t\t'color': color-like, (optional)\n" + "\t\t'label': str, (optional)\n" + "\t\t'yy': str, (optional, height on camera image, ii)\n" + "\t}\n\n" + + errstr + + f"\n\nProvided:\n{dcrystals}\n" + ) + raise Exception(msg) + + +# ###################################### +# ###################################### +# Cameras check function +# ###################################### + + +def _dcam(dcam, dap=None, dcrystals=None): + + # ---------------- + # basics + # ---------------- + + c0 = ( + isinstance(dcam, dict) + and all([isinstance(v0, dict) for v0 in dcam.values()]) + ) + if not c0: + _err_dcam(dcam) + + # ---------------- + # prepare + # ---------------- + + lok_ap = list(dap.keys()) + lok_cryst = list(dcrystals.keys()) + + # ---------------- + # loop on keys + # ---------------- + + dfail = {} + for i0, (k0, v0) in enumerate(dcam.items()): + + try: + # --------------- + # from_dist vs (cent, nin) + + lc = [ + v0.get('from_crystal') is not None + and isinstance(v0['from_crystal'], dict), + all([v0.get(kk) is not None for kk in ['cent', 'nin']]) + ] + if np.sum(lc) != 1: + msg = "Provide either 'from_crystal' or {'cent', 'nin'}" + dfail[k0] = msg + continue + + # --------------- + # from_crystal + + if lc[0]: + + # key + dcam[k0]['from_crystal']['key'] = ds._generic_check._check_var( + dcam[k0]['from_crystal'].get('key'), + f"dcam['{k0}']['from_crystal']['key']", + types=str, + allowed=lok_cryst, + ) + + # dist + dcam[k0]['from_crystal']['dist'] = float( + ds._generic_check._check_var( + dcam[k0]['from_crystal'].get('dist'), + f"dcam['{k0}']['from_crystal']['dist']", + types=(float, int), + sign='>0.', + ) + ) + + # angle + dcam[k0]['from_crystal']['angle'] = float( + ds._generic_check._check_var( + dcam[k0]['from_crystal'].get('angle'), + f"dcam['{k0}']['from_crystal']['angle']", + types=(float, int), + ) + ) + + dcam[k0]['ref_frame'] = None + + # --------------- + # cent, nin + + else: + + # cent + dcam[k0]['cent'] = ds._generic_check._check_flat1darray( + dcam[k0]['cent'], + f"dcam['{k0}']['cent']", + dtype=float, + size=2, + ) + + # nin + dcam[k0]['nin'] = ds._generic_check._check_flat1darray( + dcam[k0]['nin'], + f"dcam['{k0}']['nin']", + dtype=float, + size=2, + norm=True, + ) + + # ref_frame + dcam[k0]['ref_frame'] = ds._generic_check._check_var( + dcam[k0]['ref_frame'], + f"dcam['{k0}']['ref_frame']", + types=str, + default='abs', + allowed=lok_ap + ['abs'], + ) + + if dcam[k0]['ref_frame'] != 'abs': + dapi = dap[dcam[k0]['ref_frame']] + + cent = ( + dapi['cent'] + + dcam[k0]['cent'][0] * dapi['cent']['ex'] + + dcam[k0]['cent'][1] * dapi['cent']['ey'] + ) + nin = ( + dcam[k0]['nin'][0] * dapi['cent']['ex'] + + dcam[k0]['nin'][1] * dapi['cent']['ey'] + ) + dcam[k0]['cent'] = cent + dcam[k0]['nin'] = nin + dcam[k0]['ref_frame'] = 'abs' + + # --------------- + # length + + dcam[k0]['length'] = float( + ds._generic_check._check_var( + dcam[k0].get('length'), + f"dcam['{k0}']['length']", + types=(float, int), + sign='>0.', + ) + ) + + # --------------- + # label + + dcam[k0]['label'] = ds._generic_check._check_var( + dcam[k0].get('label'), + f"dcam['{k0}']['label']", + types=str, + default=str(k0), + ) + + # --------------- + # color + + if dcam[k0].get('color') is None: + dcam[k0]['color'] = 'k' + if not mcolors.is_color_like(dcam[k0]['color']): + msg = f"dcam['{k0}']['color'] not color-like!" + raise Exception(msg) + dcam[k0]['color'] = mcolors.to_rgba(dcam[k0]['color']) + + except Exception as err: + dfail[k0] = str(err) + + # ------------------- + # raise errors if any + # ------------------- + + if len(dfail) > 0: + lstr = [f"\t- {k0}: {v0}" for k0, v0 in dfail.items()] + msg = "\n".join(lstr) + _err_dcam(dcam, errstr=msg) + + return + + +def _err_dcam(dcam, errstr=''): + msg = ( + "Arg dcam must be a dict of sub-dicts of the form:\n" + "\t- 'key0': {\n" + "\t\t'from_crystal': {'key': str, 'dist': float, 'angle': float}\n" + "\t\t'cent': array of 2 floats, in ref_frame\n" + "\t\t'nin': array of 2 floats, in ref_frame\n" + "\t\t'ref_frame': None / str, (absolute or kap)\n" + "\t\t'length': float,\n" + "\t\t'color': color-like, (optional)\n" + "\t\t'label': str, (optional)\n" + "\t}\n\n" + "Provide either 'from_crystal' xor ('cent', 'nin', 'ref_frame')\n" + + errstr + + f"\n\nProvided:\n{dcam}\n" + ) + raise Exception(msg) + + +# ###################################### +# ###################################### +# dmatch check functions +# ###################################### + + +def _dmatch(dmatch, dcam=None, dap=None, dcrystals=None, npts=None): + + # ---------------- + # prepare + # ---------------- + + lok = { + 'aperture': list(dap.keys()), + 'crystal': list(dcrystals.keys()), + 'cam': list(dcam.keys()), + } + + # ---------------- + # if None => all + # ---------------- + + if dmatch is None: + dmatch = {} + for kap in lok['aperture']: + for kcryst in lok['crystal']: + for kcam in lok['cam']: + + key = f"{kap}_{kcryst}_{kcam}" + dmatch[key] = { + 'keys': { + 'aperture': kap, + 'crystal': kcryst, + 'cam': kcam, + }, + 'npts': None, + 'color': None, + 'label': None, + } + + # ---------------- + # basics + # ---------------- + + c0 = ( + isinstance(dmatch, dict) + and all([isinstance(v0, dict) for v0 in dmatch.values()]) + ) + if not c0: + _err_dmatch(dmatch) + + # ---------------- + # loop on keys + # ---------------- + + dfail = {} + lcolor = ['blue', 'orange', 'green', 'red', 'purple', 'brown', 'pink'] + lcolor = [f"tab:{cc}" for cc in lcolor] + for i0, (k0, v0) in enumerate(dmatch.items()): + + try: + # --------------- + # keys + + if not isinstance(v0.get('keys'), dict): + dfail[k0] = "keys must be a dict" + continue + + for kk in ['aperture', 'crystal', 'cam']: + dmatch[k0]['keys'][kk] = ds._generic_check._check_var( + dmatch[k0]['keys'].get(kk), + f"dmatch['{k0}']['keys']['{kk}']", + types=str, + allowed=lok[kk], + ) + + # --------------- + # ycam + + dmatch[k0]['ycam'] = float(ds._generic_check._check_var( + dmatch[k0].get('ycam'), + f"dmatch['{k0}']['ycam']", + types=(int, float), + default=i0, + )) + + # --------------- + # npts + + dmatch[k0]['npts'] = ds._generic_check._check_var( + dmatch[k0].get('npts'), + f"dmatch['{k0}']['npts']", + types=(int, float), + sign='>0.', + default=npts, + ) + + # --------------- + # label + + dmatch[k0]['label'] = ds._generic_check._check_var( + dmatch[k0].get('label'), + f"dmatch['{k0}']['label']", + types=str, + default=str(k0), + ) + + # --------------- + # color + + if dmatch[k0].get('color') is None: + dmatch[k0]['color'] = lcolor[i0 % len(lcolor)] + if not mcolors.is_color_like(dmatch[k0]['color']): + msg = f"dcam['{k0}']['color'] not color-like!" + raise Exception(msg) + dmatch[k0]['color'] = mcolors.to_rgba(dmatch[k0]['color']) + + except Exception as err: + dfail[k0] = str(err) + + # ------------------- + # raise errors if any + # ------------------- + + if len(dfail) > 0: + lstr = [f"\t- {k0}: {v0}" for k0, v0 in dfail.items()] + msg = "\n".join(lstr) + _err_dmatch(dmatch, errstr=msg) + + return dmatch + + +def _err_dmatch(dmatch, errstr=''): + msg = ( + "Arg dcam must be a dict of sub-dicts of the form:\n" + "\t- 'key0': {\n" + "\t\t'cent': array of 2 floats, (default to [0, 0])\n" + "\t\t'color': color-like, (optional)\n" + "\t\t'label': str, (optional)\n" + "\t}\n\n" + + errstr + + f"\n\nProvided:\n{dmatch}\n" + ) + raise Exception(msg) + + +# ###################################### +# ###################################### +# dscans check function +# ###################################### + + +def _dscans( + dscans=None, +): + + # ---------------- + # basics + # ---------------- + + if not isinstance(dscans, dict): + _err_dscans(dscans) + + # ---------------- + # loop on keys + # ---------------- + + dfail = {} + for k0, v0 in _DSCANS.items(): + + # ------------------- + # set values as array + + try: + if dscans.get(k0) is None: + if len(v0) >= 2: + if not isinstance(v0[1], str): + dscans[k0] = np.atleast_1d(v0[1]) + else: + msg = ( + f"Arg dscans['{k0}'] must be provided!\n" + ) + raise Exception(msg) + + # set + dscans[k0] = np.atleast_1d(dscans[k0]).astype(v0[0]) + + except Exception as err: + dfail[k0] = str(err) + + # ------------------- + # check broadcastable + # ------------------- + + try: + shape = np.broadcast_shapes(*[vv.shape for vv in dscans.values()]) + for k0, v0 in dscans.items(): + dscans[k0] = np.broadcast_to(v0, shape) + except Exception: + lstr = [f"\t- {k0}: {v0.shape}" for k0, v0 in dscans.items()] + dfail["broadcastable"] = "\n".join(lstr) + + # ------------------- + # raise errors if any + # ------------------- + + if len(dfail) > 0: + lstr = [f"\t- {k0}: {v0}" for k0, v0 in dfail.items()] + msg = "\n".join(lstr) + _err_dmatch(dscans, errstr=msg) + + return + + +def _err_dscans(dscans, errstr=''): + lstr = [f"\t- {k0}: {v0}" for k0, v0 in _DSCANS.items()] + msg = ( + "Arg dscans must be a dict of sub-dicts of the form:\n" + "{" + + "\n".join(lstr) + + "}\n\n" + + "Where all arrays must be broadcastable with each other\n" + + errstr + + f"\n\nProvided:\n{dscans}\n" + ) + raise Exception(msg) + + +# ###################################### +# ###################################### +# derive dscans +# ###################################### + + +def _derive_dscans( + dap=None, + dcrystals=None, + dcam=None, + dmatch=None, +): + + # ---------------- + # prepare + # ---------------- + + shape = (len(dmatch),) + dscans = {k0: np.full(shape, np.nan) for k0 in _DSCANS.keys()} + + # ---------------- + # loop on matches + # ---------------- + + for i0, (k0, v0) in enumerate(dmatch.items()): + + # -------- + # prepare + + dmatch[k0]['ind'] = (i0,) + + dapi = dap[v0['keys']['aperture']] + dcrysti = dcrystals[v0['keys']['crystal']] + dcami = dcam[v0['keys']['cam']] + + # ---------- + # aperture + + dscans['ap0'][i0] = dapi['cent'][0] + dscans['ap1'][i0] = dapi['cent'][1] + dscans['ex0'][i0] = dapi['ex'][0] + dscans['ex1'][i0] = dapi['ex'][1] + dscans['ey0'][i0] = dapi['ey'][0] + dscans['ey1'][i0] = dapi['ey'][1] + dscans['semi_angle_max'][i0] = dapi['semi_angle_max'] + + # ---------- + # crystal + + dscans['dist_from_ap'][i0] = dcrysti['dist_from_ap'] + dscans['lamb0'][i0] = dcrysti['lamb0'] + dscans['lamb0_min'][i0] = dcrysti['lamb0_min'] + dscans['lamb0_max'][i0] = dcrysti['lamb0_max'] + dscans['bragg0'][i0] = dcrysti['bragg0'] + dscans['rcurve'][i0] = dcrysti['rcurve'] + dscans['length'][i0] = dcrysti['length'] + dscans['varrad_b'][i0] = dcrysti['varrad_b'] + + # ---------- + # cam + + # from crystal + if dcami.get('frame_ref') is None: + kcryst = dcami['from_crystal']['key'] + cc = dapi['cent'] + dapi['ex'] * dcrystals[kcryst]['dist_from_ap'] + vc = ( + dapi['ex'] * np.cos(2. * dcrystals[kcryst]['bragg0']) + + dapi['ey'] * np.sin(2. * dcrystals[kcryst]['bragg0']) + ) + cent = cc + vc * dcami['from_crystal']['dist'] + nin = -( + np.cos(dcami['from_crystal']['angle']) * vc + + np.sin(dcami['from_crystal']['angle']) * np.r_[-vc[1], vc[0]] + ) + + # update dcam + dcam[v0['keys']['cam']]['cent'] = cent + dcam[v0['keys']['cam']]['nin'] = nin + else: + assert dcami['frame_ref'] == 'abs' + cent = dcami['cent'] + + # update dscans + nin = dcami['nin'] + dscans['cam_c0'][i0] = cent[0] + dscans['cam_c1'][i0] = cent[1] + dscans['cam_nin0'][i0] = nin[0] + dscans['cam_nin1'][i0] = nin[1] + dscans['cam_length'][i0] = dcami['length'] + + return dscans + + +# ###################################### +# ###################################### +# save, pfe check functions +# ###################################### + + +def _pfe( + dmatch=None, + pfe_fig=None, + pfe_npz=None, +): + + # ----------- + # pfe_fig + # ----------- + + # defaults + path = os.path.abspath('.') + name = f"spectral_range_2d_{len(dmatch)}cases.png" + pfe_fig_def = os.path.join(path, name) + + # check 1 + pfe_fig = ds._generic_check._check_var( + pfe_fig, 'pfe_fig', + types=str, + default=pfe_fig_def, + ) + + # check 2 + if not os.path.isdir(os.path.split(pfe_fig)[0]): + msg = ( + "Arg 'pfe_fig' points to a non-existing dir!\n" + f"Provided:\n{pfe_fig}\n" + ) + raise Exception(msg) + + # ----------- + # pfe_npz + # ----------- + + # defaults + name = f"spectral_range_2d_{len(dmatch)}cases.npz" + pfe_npz_def = os.path.join(path, name) + + # check 1 + pfe_npz = ds._generic_check._check_var( + pfe_npz, 'pfe_npz', + types=str, + default=pfe_npz_def, + ) + + # check 1 + if not os.path.isdir(os.path.split(pfe_npz)[0]): + msg = ( + "Arg 'pfe_fig' points to a non-existing dir!\n" + f"Provided:\n{pfe_npz}\n" + ) + raise Exception(msg) + + return pfe_fig, pfe_npz diff --git a/tofu/spectro/_spectralrange2d_compute.py b/tofu/spectro/_spectralrange2d_compute.py new file mode 100644 index 000000000..fc4353911 --- /dev/null +++ b/tofu/spectro/_spectralrange2d_compute.py @@ -0,0 +1,474 @@ +import numpy as np + + +# ################################################################# +# ################################################################# +# Compute +# ################################################################# + + +def main( + # aperture + ap0=None, + ap1=None, + ex0=None, + ex1=None, + ey0=None, + ey1=None, + semi_angle_max=None, + # crystal + dist_from_ap=None, + lamb0=None, + bragg0=None, + rcurve=None, + length=None, + varrad_b=None, + lamb0_min=None, + lamb0_max=None, + # camera + cam_c0=None, + cam_c1=None, + cam_nin0=None, + cam_nin1=None, + cam_length=None, + # options + npts=None, +): + + # ------------ + # crystal's summit + # ------------ + + csummit0 = ap0 + dist_from_ap * ex0 + csummit1 = ap1 + dist_from_ap * ex1 + + # d2 + d2 = lamb0 / np.sin(bragg0) + + # ----------------- + # sort by crystal type + # ----------------- + + # variable-radii sinusoidal spiral + ispiral = np.isfinite(varrad_b) + + # indices of curved crystals + icurve = np.isfinite(rcurve) & (~ispiral) + + # flat + iflat = (~icurve) & (~ispiral) + + # safety check + if not np.all(np.sum([icurve, iflat, ispiral], axis=0) == 1): + msg = ( + "Some undetermined 2d crystal shapes:\n" + f"\t- iflat = {iflat}\n" + f"\t- icurve = {icurve}\n" + f"\t- ispiral = {ispiral}\n" + ) + raise Exception(msg) + + # ---------------- + # initialize + # ---------------- + + shape = (npts,) + ex0.shape + dout = { + 'cryst0': np.full(shape, np.nan), + 'cryst1': np.full(shape, np.nan), + 'vn0': np.full(shape, np.nan), + 'vn1': np.full(shape, np.nan), + 'end0': np.full(shape, np.nan), + 'end1': np.full(shape, np.nan), + 'lamb': np.full(shape, np.nan), + 'cam_coord': np.full(shape, np.nan), + 'dmask': { + 'semi_angle_max': np.ones(shape, dtype=bool), + 'crystal': np.ones(shape, dtype=bool), + 'camera': np.ones(shape, dtype=bool), + }, + 'dind_lamb': { + 'lamb0': np.zeros(shape[1:], dtype=int), + 'lamb0_min': np.zeros(shape[1:], dtype=int), + 'lamb0_max': np.zeros(shape[1:], dtype=int), + }, + } + + # ---------------- + # sample rays on crytals + # ---------------- + + # sample + 2 extra points + kpts = np.linspace(-1., 1., npts-2) + dk = kpts[1] - kpts[0] + kpts = np.r_[kpts[0] - dk, kpts, kpts[-1] + dk] + + # crystal types + lif = [ + (iflat, _compute_flat), + (icurve, _compute_curve), + (ispiral, _compute_spiral), + ] + + for ind, func in lif: + if np.any(ind): + ( + dout['cryst0'][:, ind], + dout['cryst1'][:, ind], + dout['vn0'][:, ind], + dout['vn1'][:, ind], + dout['dmask']['crystal'][:, ind], + ) = func( + csummit0=csummit0[ind], + csummit1=csummit1[ind], + bragg0=bragg0[ind], + ex0=ex0[ind], + ex1=ex1[ind], + ey0=ey0[ind], + ey1=ey1[ind], + length=length[ind], + kpts=kpts, + # curved + rcurve=rcurve[ind], + # spiral + ap0=ap0[ind], + ap1=ap1[ind], + dist_from_ap=dist_from_ap[ind], + varrad_b=varrad_b[ind], + # lamb_min, lamb_max + d2=d2[ind], + lamb0_min=lamb0_min[ind], + lamb0_max=lamb0_max[ind], + npts=npts, + ) + + # ----------- + # semi_angle_max + # ----------- + + iout = np.isfinite(semi_angle_max) + if np.any(iout): + dv0 = (dout['cryst0'] - ap0[None, ...]) + dv1 = (dout['cryst1'] - ap1[None, ...]) + semi_angle = np.arctan2(dv0*ey0 + dv1*ey1, dv0*ex0 + dv1*ex1) + ind = np.abs(semi_angle[:, iout]) <= semi_angle_max[None, iout] + dout['dmask']['semi_angle_max'][:, iout] = ind + + # ---------------- + # compute rays + # ---------------- + + # vectors of incident rays + vi0 = dout['cryst0'] - ap0 + vi1 = dout['cryst1'] - ap1 + vin = np.sqrt(vi0**2 + vi1**2) + vi0 = vi0 / vin + vi1 = vi1 / vin + + # reflected vectors + sca = vi0*dout['vn0'] + vi1*dout['vn1'] + vr0 = vi0 - 2.*sca*dout['vn0'] + vr1 = vi1 - 2.*sca*dout['vn1'] + + # ---------------------- + # compute spectral range + # ---------------------- + + # get local bragg angle - top and bottom + bragg = np.arccos(sca) - np.pi/2. + + # lamb + dout['lamb'] = d2 * np.sin(bragg) + + # --------------------- + # dind_lamb + # --------------------- + + dlamb = np.abs(np.mean(np.diff(dout['lamb'], axis=0), axis=0)) + + ll = [('lamb0', lamb0), ('lamb0_min', lamb0_min), ('lamb0_max', lamb0_max)] + for (k0, v0) in ll: + iok = np.isfinite(v0) + slil = (slice(None), iok) + sliv = (None, iok) + ind = np.argmin(np.abs(dout['lamb'][slil] - v0[sliv]), axis=0) + sli = (ind, iok) + c0 = np.abs(dout['lamb'][sli] - v0[sliv]) <= 0.5 * dlamb[iok] + assert np.all(c0) + dout['dind_lamb'][k0][iok] = ind + + # ---------------- + # intersection with camera plane + # ---------------- + + ( + dout['end0'], + dout['end1'], + dout['cam_coord'], + dout['dmask']['camera'], + ) = _camera_plane( + ex0=ex0, + ex1=ex1, + cryst0=dout['cryst0'], + cryst1=dout['cryst1'], + vr0=vr0, + vr1=vr1, + cam_c0=cam_c0, + cam_c1=cam_c1, + cam_nin0=cam_nin0, + cam_nin1=cam_nin1, + cam_length=cam_length, + ) + + return dout + + +# ################################################################# +# ################################################################# +# Compute by crystal type +# ################################################################# + + +def _compute_flat( + csummit0=None, + csummit1=None, + bragg0=None, + ex0=None, + ex1=None, + ey0=None, + ey1=None, + length=None, + kpts=None, + # lamb_min, max + d2=None, + dist_from_ap=None, + lamb0_min=None, + lamb0_max=None, + npts=None, + # unused + **kwdargs, +): + + # -------------------- + # lamb_min, max + # -------------------- + + imin = np.isfinite(lamb0_min) + k_lambmin = np.full(lamb0_min.shape, np.nan) + bragg0_min = np.arcsin(lamb0_min[imin] / d2[imin]) + k_lambmin[imin] = ( + dist_from_ap[imin] + * np.sin(bragg0[imin] - bragg0_min) / np.sin(bragg0_min) + ) + + imax = np.isfinite(lamb0_max) + k_lambmax = np.full(lamb0_max.shape, np.nan) + bragg0_max = np.arcsin(lamb0_max[imax] / d2[imax]) + k_lambmax[imax] = ( + dist_from_ap[imax] + * np.sin(bragg0[imax] - bragg0_max) / np.sin(bragg0_max) + ) + + kmin = np.nanmin([-0.5 * length, k_lambmin, k_lambmax], axis=0) + kmax = np.nanmax([0.5 * length, k_lambmin, k_lambmax], axis=0) + kk = np.linspace(kmin, kmax, npts, axis=0) + + mask_cryst = np.abs(kk) <= 0.5 * length + + # -------------------- + # prepare + # -------------------- + + # crystal plotting - straight + estraight0 = np.cos(bragg0) * ex0 + np.sin(bragg0) * ey0 + estraight1 = np.cos(bragg0) * ex1 + np.sin(bragg0) * ey1 + + # pts on crystal surface + cryst0 = csummit0[None, ...] + kk * estraight0[None, ...] + cryst1 = csummit1[None, ...] + kk * estraight1[None, ...] + + # local normal vectors + vn0 = -estraight1 + vn1 = estraight0 + + return cryst0, cryst1, vn0, vn1, mask_cryst + + +def _compute_curve( + csummit0=None, + csummit1=None, + bragg0=None, + ex0=None, + ex1=None, + ey0=None, + ey1=None, + rcurve=None, + length=None, + kpts=None, + # unused + **kwdargs, +): + + # center of curvature + ec0 = np.sin(bragg0) * ex0 - np.cos(bragg0) * ey0 + ec1 = np.sin(bragg0) * ex1 - np.cos(bragg0) * ey1 + ec0_p = -ec1 + ec1_p = ec0 + + # crystal center of curvature + cc0 = csummit0 - rcurve * ec0 + cc1 = csummit1 - rcurve * ec1 + + # half angular opening of crystal + dalpha = 0.5 * length / rcurve + sli = (slice(None),) + (None,) * ex0.ndim + theta = dalpha * kpts[sli] + + # mask_cryst + mask_cryst = np.abs(theta) <= dalpha + + # crystal plotting - curved + etheta0 = np.cos(theta) * ec0[None, :] + np.sin(theta) * ec0_p[None, :] + etheta1 = np.cos(theta) * ec1[None, :] + np.sin(theta) * ec1_p[None, :] + + cryst0 = cc0[None, :] + rcurve[None, :] * etheta0 + cryst1 = cc1[None, :] + rcurve[None, :] * etheta1 + + # local normal vectors + vn0 = -etheta0 + vn1 = -etheta1 + + return cryst0, cryst1, vn0, vn1, mask_cryst + + +def _compute_spiral( + csummit0=None, + csummit1=None, + bragg0=None, + ex0=None, + ex1=None, + ey0=None, + ey1=None, + rcurve=None, + length=None, + kpts=None, + dist_from_ap=None, + varrad_b=None, + ap0=None, + ap1=None, + # unused + **kwdargs, +): + + # main parameters + r0 = rcurve + ix = ~np.isfinite(r0) + r0[ix] = dist_from_ap[ix] + + # local radius of curvature at center + # rc0 = r0 / (b * np.sin(bragg0)) + + # dOMx = r / (b-1) * (cos(phi) / tan(gam) - sin(phi)) + # dOMy = r / (b-1) * (sin(phi) / tan(gam) + cos(phi)) + # dL = r/(b-1) * 1 / sin(gam) + # dL ~ r0/(b-1) * 1/sin(bragg0) * Dgam + + # half angular opening of crystal (approximative) + # dgam = 0.5*length / rc0 + dgam = 1.1 * length * np.sin(bragg0) * (b-1) / r0 / 2 + + # gam + sli = (slice(None),) + (None,) * ex0.ndim + gam = bragg0[None, ...] + dgam[None, ...] * kpts[sli] + + # rr + rr = ( + r0[None, ...] + * (np.sin(gam) / np.sin(bragg0)[None, ...])**( + 1 / (varrad_b[None, ...] - 1) + ) + ) + + # phi + phi = (gam - bragg0[None, ...]) / (varrad_b[None, ...] - 1) + + # pts on cryst + cryst0 = ( + ap0 + + (dist_from_ap - r0) * ex0 + + rr * (np.cos(phi) * ex0 + np.sin(phi) * ey0) + ) + cryst1 = ( + ap1 + + (dist_from_ap - r0) * ex1 + + rr * (np.cos(phi) * ex1 + np.sin(phi) * ey1) + ) + + # derivative + c0 = rr / (varrad_b[None, ...] - 1.) + c1 = np.cos(gam) / np.sin(gam) + dOMxx = c0 * (c1 * np.cos(phi) - np.sin(phi)) + dOMyy = c0 * (c1 * np.sin(phi) + np.cos(phi)) + dOMx = dOMxx * ex0 + dOMyy * ey0 + dOMy = dOMxx * ex1 + dOMyy * ey1 + + # local normal vectors + vn0 = dOMy / np.sqrt(dOMx**2 + dOMy**2) + vn1 = -dOMx / np.sqrt(dOMx**2 + dOMy**2) + + return cryst0, cryst1, vn0, vn1 + + +# ################################################################# +# ################################################################# +# camera plane +# ################################################################# + + +def _camera_plane( + ex0=None, + ex1=None, + cryst0=None, + cryst1=None, + vr0=None, + vr1=None, + cam_c0=None, + cam_c1=None, + cam_nin0=None, + cam_nin1=None, + cam_length=None, +): + + # ----------- + # end points + # ----------- + + kk = ( + (cam_c0 - cryst0) * cam_nin0 + + (cam_c1 - cryst1) * cam_nin1 + ) / (vr0 * cam_nin0 + vr1 * cam_nin1) + + # end of rays at camera + end0 = cryst0 + kk * vr0 + end1 = cryst1 + kk * vr1 + + # ----------- + # coordinates on cameras + # ----------- + + # get lateral cam unit vector + cam_e00 = -cam_nin1 + cam_e01 = np.copy(cam_nin0) + ineg = (cam_e00 * ex0 + cam_e01 * ex1) < 0. + cam_e00[ineg] = -cam_e00[ineg] + cam_e01[ineg] = -cam_e01[ineg] + + cam_coord = (end0 - cam_c0) * cam_e00 + (end1 - cam_c1) * cam_e01 + + # ---------- + # mask + # ---------- + + mask_cam = np.abs(cam_coord) < cam_length * 0.5 + + return end0, end1, cam_coord, mask_cam diff --git a/tofu/spectro/_spectralrange2d_plot.py b/tofu/spectro/_spectralrange2d_plot.py new file mode 100644 index 000000000..7b545edee --- /dev/null +++ b/tofu/spectro/_spectralrange2d_plot.py @@ -0,0 +1,507 @@ + + +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.colors as mcolors +import matplotlib.lines as mlines +import matplotlib.gridspec as gridspec +import datastock as ds + + +# #################################### +# #################################### +# DEFAULTS +# #################################### + + +_DMARKER = { + 'semi_angle_max': '^', + 'crystal': 'D', + 'camera': 's', +} + + +# #################################### +# #################################### +# Plot match function +# #################################### + + +def match( + dap=None, + dcam=None, + dmatch=None, + dscans=None, + dout=None, + # plotting + dax=None, + # saving + pfe_fig=None, + # unused + **kwdargs, +): + + # ------------- + # prepare rays + # ------------ + + drays = _prepare_rays( + dap=dap, + dmatch=dmatch, + dscans=dscans, + dout=dout, + ) + + # ------------- + # prepare img + # ------------ + + dimg = _prepare_img( + dap=dap, + dmatch=dmatch, + dscans=dscans, + dout=dout, + ) + + # -------------- + # prepare figure + # -------------- + + if dax is None: + dax = _match_dax() + + dax = ds._generic_check._check_dax(dax) + + # --------------------- + # plot - loop on match + # --------------------- + + for i0, (k0, v0) in enumerate(dmatch.items()): + + # --- + # hor + + kax = 'hor' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + # rays - lamb + semi_angle_max + for kr, vr in drays.items(): + + for k1, v1 in vr.items(): + ax.plot( + v1['x0'], + v1['x1'], + **v1['prop'], + ) + + # --- + # img + + kax = 'cam' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + # plot cam_coords + for k1, v1 in dimg[k0].items(): + ax.plot( + v1['x'], + v1['y'], + **v1['prop'], + ) + + # add text + # ax.text( + # dcam['x0'][ilamb_max[ii], ii], + # v1['y'][0] - 0.1, + # f'{lamb_max[ii]*1e10:2.3} AA', + # color=color, + # size=8, + # horizontalalignment='center', + # verticalalignment='top', + # ) + + # --------------- + # plot input data + # --------------- + + ax.text( + 0.8, + 0.4, + "beta_max = {beta_str}\n", + color='k', + size=10, + horizontalalignment='center', + verticalalignment='top', + transform=ax.figure.transFigure, + ) + + # ------------ + # camera + # ------------ + + lc = [] + for k0, v0 in dmatch.items(): + kc = v0['keys']['cam'] + + if kc not in lc: + + # --------- + # hor + + kax = 'hor' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + dx0 = 0.5 * dcam[kc]['length'] * np.r_[-1, 1] + dx1 = 0.5 * dcam[kc]['length'] * np.r_[-1, 1] + ax.plot( + dcam[kc]['cent'][0] + dx0 * (-dcam[kc]['nin'][1]), + dcam[kc]['cent'][1] + dx1 * dcam[kc]['nin'][0], + ls='-', + lw=2., + marker='None', + c='k', + label=kc, + ) + + # --------- + # cam + + kax = 'cam' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + ax.axvline(-0.5*dcam[kc]['length'], c='k', ls='-', lw=1.) + ax.axvline(0.5*dcam[kc]['length'], c='k', ls='-', lw=1.) + # ax.set_ylim(0, size + 1) + + # --------------- + # decorate + # --------------- + + kax = 'hor' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + lh = [ + mlines.Line2D([], [], c=v0['color'], ls='-', label=k0) + for k0, v0 in dmatch.items() + ] + ax.legend(handles=lh, loc='upper right', fontsize=12) + + # ---------- + # saving + + if pfe_fig is not None: + dax['hor']['handle'].figure.savefig(pfe_fig, format='png', dpi=300) + + return dax + + +# #################################### +# #################################### +# Prepare rays +# #################################### + + +def _prepare_rays( + dap=None, + dmatch=None, + dscans=None, + dout=None, +): + + drays = { + 'semi_angle_max': {}, + 'envelop': {}, + 'lamb': {}, + } + + # --------------- + # semi_angle_max + # --------------- + + lmax = np.max(dscans['dist_from_ap']) * 0.2 + for k0, v0 in dmatch.items(): + + kap = v0['keys']['aperture'] + if drays['semi_angle_max'].get(kap) is not None: + continue + if not np.isfinite(dap[kap]['semi_angle_max']): + continue + + cent = dap[kap]['cent'] + vup = ( + np.cos(dap[kap]['semi_angle_max']) * dap[kap]['ex'] + + np.sin(dap[kap]['semi_angle_max']) * dap[kap]['ey'] + ) + vdown = ( + np.cos(dap[kap]['semi_angle_max']) * dap[kap]['ex'] + - np.sin(dap[kap]['semi_angle_max']) * dap[kap]['ey'] + ) + drays['semi_angle_max'][kap] = { + 'x0': cent[0] + lmax * np.r_[vup[0], 0, vdown[0]], + 'x1': cent[1] + lmax * np.r_[vup[1], 0, vdown[1]], + 'prop': { + 'color': dap[kap]['color'], + 'ls': '-', + 'lw': 2, + 'label': f'semi_angle_max - {kap}', + } + } + + # --------------- + # lamb - rays + # --------------- + + for k0, v0 in dmatch.items(): + kap = v0['keys']['aperture'] + cent = dap[kap]['cent'] + + x0 = [] + x1 = [] + for klamb, ilamb in dout['dind_lamb'].items(): + + if np.isfinite(dscans[klamb][v0['ind']]): + sli = (ilamb[v0['ind']],) + v0['ind'] + + # cryst + cryst0 = dout['cryst0'][sli] + cryst1 = dout['cryst1'][sli] + + # end + end0 = dout['end0'][sli] + end1 = dout['end1'][sli] + + x0 += [cent[0], cryst0, end0, np.nan] + x1 += [cent[1], cryst1, end1, np.nan] + + drays['lamb'][k0] = { + 'x0': x0, + 'x1': x1, + 'prop': { + 'color': v0['color'], + 'ls': '--', + 'lw': 1, + 'label': f'rays - lamb - {k0}', + } + } + + # --------------- + # envelop + # --------------- + + for k0, v0 in dmatch.items(): + kap = v0['keys']['aperture'] + cent = dap[kap]['cent'] + + x0 = [] + x1 = [] + for klamb, ilamb in dout['dind_lamb'].items(): + pass + + # npts, size = crystx.shape + + # # envelop + # iok = np.isfinite(endx) + # i0 = tuple([iok[:, ii].nonzero()[0][0] for ii in range(size)]) + # i1 = tuple([iok[:, ii].nonzero()[0][-1] for ii in range(size)]) + # nind = tuple(range(size)) + + # # envelop + # envx = np.array([ + # endx[i1, nind], crystx[i1, nind], + # np.full((size,), ap[0]), + # crystx[i0, nind], endx[i0, nind], + # ]) + # envy = np.array([ + # endy[i1, nind], crysty[i1, nind], + # np.full((size,), ap[1]), + # crysty[i0, nind], endy[i0, nind], + # ]) + + # # central rays + # ind = int((npts-1)/2) + # raycx = np.array([np.full((size,), ap[0]), crystx[ind, :], endx[ind, :]]) + # raycy = np.array([np.full((size,), ap[1]), crysty[ind, :], endy[ind, :]]) + + # # --------------- + # # envelop + # # --------------- + + # # dcam + # if dcam is not None: + # ninx, niny = dcam['nin_r'][:2] + # ninn = np.sqrt(ninx**2 + niny**2) + # ninx, niny = ninx/ninn, niny/ninn + # e0x, e0y = -niny, ninx + # e0n = np.sqrt(e0x**2 + e0y**2) + # e0x, e0y = e0x/e0n, e0y/e0n + # clen = dcam['length'] + # camx = dcam['cent_r'][0] + 0.5*clen*np.r_[-1, 1] * e0x + # camy = dcam['cent_r'][1] + 0.5*clen*np.r_[-1, 1] * e0y + + return drays + + +# #################################### +# #################################### +# match -Get dax +# #################################### + + +def _match_dax(): + + # -------------- + # prepare figure + + dmargin = { + 'left': 0.08, 'right': 0.98, + 'bottom': 0.08, 'top': 0.90, + 'hspace': 0.20, 'wspace': 0.25, + } + + fig = plt.figure(figsize=(13, 8)) + fig.suptitle('2d ray-tracing model') + gs = gridspec.GridSpec(ncols=3, nrows=2, **dmargin) + + # ---------- + # make axes + + # ax0 - hor + ax0 = fig.add_subplot( + gs[:, :-1], + aspect='equal', + adjustable='datalim', + ) + + ax0.set_xlabel("x (m)", size=12) + ax0.set_ylabel("y (m)", size=12) + ax0.set_title("2d ray tracing", size=12, fontweight='bold') + + # ax1 - cam + ax1 = fig.add_subplot( + gs[0, -1], + aspect='auto', + ) + + ax1.set_xlabel("x0 (m)", size=12) + ax1.set_title("Image on camera", size=12, fontweight='bold') + + # ------------ + # dict + + dax = { + 'hor': {'handle': ax0}, + 'cam': {'handle': ax1}, + } + + return dax + + +# #################################### +# #################################### +# Prepare img +# #################################### + + +def _prepare_img( + dap=None, + dmatch=None, + dscans=None, + dout=None, +): + + # --------------- + # prepare + # --------------- + + dimg = {} + + # --------------- + # semi_angle_max + # --------------- + + for i0, (k0, v0) in enumerate(dmatch.items()): + + # -------- + # prepare + + sli = (slice(None),) + v0['ind'] + dmask = {kk: vv[sli] for kk, vv in dout['dmask'].items()} + + # ------------ + # bool indices + + iin = np.all([vv for vv in dmask.values()], axis=0) + iout_ap = (~dmask['semi_angle_max']) + iout_cryst = (~dmask['crystal']) & dmask['semi_angle_max'] + iout_cam = ( + (~dmask['camera']) & dmask['semi_angle_max'] & dmask['crystal'] + ) + diout = { + 'semi_angle_max': iout_ap, + 'crystal': iout_cryst, + 'camera': iout_cam, + } + + # ------------ + # x1 + + x1 = v0['ycam'] + + # ----------- + # store + + # iin + dimg[k0] = { + 'in': { + 'x': dout['cam_coord'][sli][iin], + 'y': np.full((iin.sum(),), x1), + 'prop': { + 'color': v0['color'], + 'ls': '-', + 'lw': 2, + 'marker': '.', + 'ms': 8, + 'label': f"{k0} - in", + }, + }, + } + + # iout + for ko, vo in diout.items(): + dimg[k0][ko] = { + 'x': dout['cam_coord'][sli][vo], + 'y': np.full((vo.sum(),), x1), + 'prop': { + 'color': mcolors.to_rgb(v0['color']) + (0.5,), + 'ls': 'None', + 'lw': 1., + 'marker': _DMARKER[ko], + 'ms': 8, + 'label': f"{k0} - out {ko}", + }, + } + + # lamb + ind = [] + for kl, vl in dout['dind_lamb'].items(): + if np.isfinite(dscans[kl][v0['ind']]): + ind.append(vl[v0['ind']]) + + dimg[k0]['lamb'] = { + 'x': dout['cam_coord'][sli][np.array(ind).astype(int)], + 'y': np.full((len(ind),), x1), + 'prop': { + 'color': v0['color'], + 'ls': 'None', + 'lw': 2., + 'marker': 'o', + 'ms': 10, + 'label': f"{k0} - lamb", + } + } + + return dimg