diff --git a/CLASSIFIERS.txt b/CLASSIFIERS.txt new file mode 100644 index 0000000..fc9f4fa --- /dev/null +++ b/CLASSIFIERS.txt @@ -0,0 +1,7 @@ +Development Status :: 5 - Production/Stable +Intended Audience :: Science/Research +Programming Language :: Python :: 3 +Programming Language :: Python :: 3.10 +Programming Language :: Python :: 3.11 +Programming Language :: Python :: 3.12 +Natural Language :: English diff --git a/beampipe/__init__.py b/beampipe/__init__.py new file mode 100644 index 0000000..1ca2fe2 --- /dev/null +++ b/beampipe/__init__.py @@ -0,0 +1,2 @@ +from ._class01_Outline2d import Outline2d as Collection + diff --git a/beampipe/_class00_CSYS.py b/beampipe/_class00_CSYS.py new file mode 100644 index 0000000..161dc2a --- /dev/null +++ b/beampipe/_class00_CSYS.py @@ -0,0 +1,159 @@ +import copy +from typing import Annotated + + +import numpy as np +from datastock import DataStock as Previous + + +# from . import _class00_check as _check +from . import _class00_check as _check +from . import _class00_show as _show +from . import _class00_transform as _transform +from . import _class00_transform_coords as _transform_coords + + +__all__ = ['CSYS'] + + +############################################# +############################################# +# DEFAULT VALUES +############################################# + + +_WHICH_CSYS = 'csys' + + +############################################# +############################################# +# Spectral Lines +############################################# + + +class CSYS(Previous): + + _ddef = copy.deepcopy(Previous._ddef) + _which_csys = _WHICH_CSYS + + # ------------------- + # add csys + # ------------------- + + def add_csys( + self, + key: Annotated[str | None, 'key of the csys to be added'] = None, + # cent + origin: Annotated[np.ndarray | None, "coords of the csys origin"] = None, + # ctype + ctype: Annotated[str | None, "type of csys (cart, ...)"] = None, + # vect + e0: Annotated[np.ndarray | None, "coords of unit vector e0"] = None, + e1: Annotated[np.ndarray | None, "coords of unit vector e1"] = None, + e2: Annotated[np.ndarray | None, "coords of unit vector e2"] = None, + # units + units0: Annotated[str | None, "units of coords along e0"] = None, + units1: Annotated[str | None, "units of coords along e1"] = None, + units2: Annotated[str | None, "units of coords along e2"] = None, + # vector options + ortho: Annotated[bool | None, "Are vectors orthogonal"] = None, + norm: Annotated[bool | None, "Are vectors normalized"] = None, + direct: Annotated[bool | None, "Is the vectors basis direct"] = None, + # ref csys + kcsys0: Annotated[str | None, "csys in which coords are expressed"] = None, + ) -> None: + """ Add a csys + + Can be 1d, 2d or 3d + origin must be an iterable with accordingly 1, 2 or 3 coordinates + ctype is the + 1 2 or 3 base vectors must be provided accordingly + if norm = True => they will be normalized + if direct = True, they should form a direct base (for 2d and 3d only) + + Coordinates are given with respect to a ref csys kcsys0 + + """ + + # ------------ + # check inputs + # ------------ + + dref, ddata, dobj = _check.main(coll=self, **locals()) + + # ------------ + # Populate + # ------------ + + self.update(dref=dref, ddata=ddata, dobj=dobj) + + return + + # ------------------- + # show + # ------------------- + + def _get_show_obj(self, which=None): + if which == self._which_csys: + return _show._show + else: + return super()._get_show_obj(which) + + def _get_show_details(self, which=None): + if which == self._which_csys: + return _show._show_details + else: + return super()._get_show_details(which) + + # ------------------- + # remove csys + # ------------------- + + # ------------------- + # transform + # ------------------- + + def get_csys_transform( + self, + key_in: Annotated[str | None, 'key to csys'] = None, + key_out: Annotated[str | None, 'key to csys'] = None, + ) -> dict: + """ Get the transform needed to get from one to another csys + + Returns a dict + """ + return _transform.main(coll=self, key_in=key_in, key_out=key_out) + + def transform_csys_coords( + self, + key_in=None, + key_out=None, + # coordinates + x0=None, + x1=None, + x2=None, + ): + """ Transform coordinates from a given csys into another + + coordinates have to be broadcastable arrays + can be provided as key to broadcastable ddata + """ + + return _transform_coords.main(coll=self, **locals()) + + # ------------------- + # move within csys + # ------------------- + + def move_translate_by(): + return + + def move_rotate_by(): + return + + # ------------------- + # move to align with something + # ------------------- + + def move_align_to(): + return diff --git a/beampipe/_class00_check.py b/beampipe/_class00_check.py new file mode 100644 index 0000000..ea0deb7 --- /dev/null +++ b/beampipe/_class00_check.py @@ -0,0 +1,429 @@ +import numpy as np +import astropy.units as asunits +import datastock as ds + + +############################################# +############################################# +# DEFAULT VALUES +############################################# + +_LOK_ND = { + 'cart': ['1d', '2d', '3d'], + 'cyl': ['2d', '3d'], + 'sph': ['2d', '3d'], +} + +_ORIGIN = np.r_[0., 0., 0.] + +_DUNITS = { + 'cart': 'm', +} + + +############################################# +############################################# +# main +############################################# + + +def main( + coll=None, + key=None, + # cent + origin=None, + # ctype + ctype=None, + # nd + nd=None, + # vect + e0=None, + e1=None, + e2=None, + # units + units0=None, + units1=None, + units2=None, + # vecto options + ortho=None, + norm=None, + direct=None, + # ref csys + kcsys0=None, + # unused + **kwdargs, +): + + # --------------- + # check inputs + # --------------- + + kwd = _check(**locals()) + + # --------------- + # 1d vs 2d vs 3d + # --------------- + + _unit_vectors(kwd) + + # --------------- + # dobj + # --------------- + + dcsys = { + kk: vv for kk, vv in kwd.items() + if not kk.startswith('e') + and not kk.startswith('units') + and kk != 'key' + } + + le = [kk for kk in ['e0', 'e1', 'e2'] if kwd.get(kk) is not None] + for ie, ke in enumerate(le): + dcsys[ke] = { + 'data': kwd[ke], + 'units': kwd[f'units{ie}'], + } + + wcsys = coll._which_csys + dobj = {wcsys: {kwd['key']: dcsys}} + + return None, None, dobj + + +############################################# +############################################# +# check +############################################# + + +def _check( + **kwdargs, +): + """ check inputs for main() + """ + + # ------------ + # key + # ------------ + + coll = kwdargs['coll'] + wcsys = coll._which_csys + kwdargs['key'] = ds._generic_check._obj_key( + d0=coll.dobj.get(wcsys, {}), + short='csys', + key=kwdargs['key'], + ndigits=None, + ) + + # ------------ + # ctype + # ------------ + + lok = sorted(_LOK_ND.keys()) + kwdargs['ctype'] = ds._generic_check._check_var( + kwdargs['ctype'], 'ctype', + types=str, + allowed=lok, + default=lok[0], + ) + nd_ok = _LOK_ND[kwdargs['ctype']] + + # ------------ + # nd + # ------------ + + dv = { + kk: kwdargs[kk] for kk in ['e0', 'e1', 'e2'] + if kwdargs[kk] is not None + } + nd_min = len(dv) + lok = [f"{ii}d" for ii in range(nd_min, 4)] + lok = [kk for kk in lok if kk in nd_ok] + + nn = list(set([len(vv) for vv in dv.values()])) + if len(nn) > 1: + msg = "Unit vectors do not seem to have consistent size!" + raise Exception(msg) + elif len(nn) == 1: + nd_def = f"{nn[0]}d" + else: + nd_def = lok[-1] + + kwdargs['nd'] = ds._generic_check._check_var( + kwdargs['nd'], 'nd', + types=str, + allowed=lok, + default=nd_def, + ) + size = int(kwdargs['nd'][0]) + + # vs vectors + if len(dv) > 0: + if size - len(dv) > 1: + msg = ( + f"For a csys of nd = '{kwdargs['nd']}' provide either:\n" + "\t- No unit vectors (all default)\n" + "\t- All or all but one (derived) unit vectors\n" + ) + raise Exception(msg) + + # ------------ + # kcsys0 + # ------------ + + lok = [ + kk for kk, vv in coll.dobj.get(wcsys, {}).items() + if vv['ctype'] == kwdargs['ctype'] + and vv['nd'] == kwdargs['nd'] + ] + lref = [kk for kk in lok if coll.dobj[wcsys][kk]['kcsys0'] == kk] + kwdargs['kcsys0'] = ds._generic_check._check_var( + kwdargs['kcsys0'], 'kcsys0', + types=str, + allowed=lok + [kwdargs['key']], + default=(lref + [kwdargs['key']])[0], + ) + + # ------------ + # origin - provided => finite array of proper size + # ------------ + + if kwdargs['origin'] is None: + kwdargs['origin'] = np.copy(_ORIGIN) + + oo = np.atleast_1d(kwdargs['origin']).ravel().astype(float) + + if np.any(~np.isfinite(oo)) or oo.size != size: + msg = ( + "Arg 'origin' must be:\n" + f"\t- a flat np.ndarray of finite values with size = {size}\n" + f"Provided:\n\t{kwdargs['origin']}\n" + ) + raise Exception(msg) + kwdargs['origin'] = oo + + # ------------ + # unit vectors - provided => finite array of proper size + # ------------ + + for kk, vv in dv.items(): + vv = np.atleast_1d(vv).ravel().astype(float) + + if np.any(~np.isfinite(vv)) or vv.size != size: + msg = ( + f"Arg '{kk}' must be:\n" + "\t- a flat np.ndarray of finite values with size = {size}\n" + "Provided:\n\t{kwdargs[kk]}\n" + ) + raise Exception(msg) + kwdargs[kk] = vv + + # ------------ + # bool + # ------------ + + for kk in ['ortho', 'norm', 'direct']: + kwdargs[kk] = ds._generic_check._check_var( + kwdargs[kk], kk, + types=bool, + default=True, + ) + + # ------------ + # clean + # ------------ + + lok = [ + 'key', 'nd', 'ctype', + 'origin', + 'e0', 'e1', 'e2', + 'units0', 'units1', 'units2', + 'ortho', 'norm', 'direct', 'kcsys0', + ] + lout = [kk for kk in kwdargs.keys() if kk not in lok] + for kk in lout: + del kwdargs[kk] + + return kwdargs + + +############################################# +############################################# +# Unit vectors +############################################# + + +def _unit_vectors(kwd): + + # ---------- + # all default + # ---------- + + lv = ['e0', 'e1', 'e2'] + if all([kwd[kk] is None for kk in lv]): + if kwd['nd'] == '1d': + kwd['e0'] = np.r_[0.] + elif kwd['nd'] == '2d': + kwd['e0'] = np.r_[1, 0.] + kwd['e1'] = np.r_[0, 1.] + else: + kwd['e0'] = np.r_[1, 0., 0] + kwd['e1'] = np.r_[0, 1., 0] + kwd['e2'] = np.r_[0, 0., 1] + + # ---------- + # Not all default + # ---------- + + # 2d + if kwd['nd'] == '2d': + if kwd['e0'] is None: + kwd['e0'] = np.r_[kwd['e1'][1], -kwd['e1'][0]] + elif kwd['e1'] is None: + kwd['e1'] = np.r_[-kwd['e0'][1], kwd['e0'][0]] + + # 3d + elif kwd['nd'] == '3d': + if kwd['e0'] is None: + kwd['e0'] = np.cross(kwd['e1'], kwd['e2']) + elif kwd['e1'] is None: + kwd['e1'] = np.cross(kwd['e2'], kwd['e0']) + elif kwd['e2'] is None: + kwd['e2'] = np.cross(kwd['e0'], kwd['e1']) + + # ---------- + # clean extra + # ---------- + + size = int(kwd['nd'][0]) + for ii in range(size, 3): + estr = f"e{ii}" + if kwd[estr] is not None: + msg = ( + f"Arg '{estr}' provided for a '{kwd['nd']}' csys!\n" + f"Provided: {kwd[estr]}\n" + ) + raise Exception(msg) + + # ---------- + # basis - not colinear + # ---------- + + # nd + if kwd['nd'] != '1d': + if kwd['nd'] == '2d': + dcross = {'e0 x e1': np.cross(kwd['e0'], kwd['e1'])} + emax = np.max([kwd['e0'], kwd['e1']]) + elif kwd['nd'] == '3d': + dcross = { + 'e0 x e1': np.linalg.norm(np.cross(kwd['e0'], kwd['e1'])), + 'e1 x e2': np.linalg.norm(np.cross(kwd['e1'], kwd['e2'])), + 'e2 x e0': np.linalg.norm(np.cross(kwd['e2'], kwd['e0'])), + } + emax = np.max([kwd['e0'], kwd['e1'], kwd['e2']]) + + # check colinearity + dfail = { + kk: vv for kk, vv in dcross.items() if np.abs(vv) < 1e-9 * emax + } + if len(dfail) > 0: + lstr = [f"\t- {kk} = vv" for kk, vv in dfail.items()] + msg = ( + "Unit vectors must not be co-linear!\n" + + "\n".join(lstr) + ) + raise Exception(msg) + + # ---------- + # norm + # ---------- + + if kwd['norm'] is True: + for ii in range(size): + estr = f"e{ii}" + kwd[estr] = kwd[estr] / np.sqrt(np.sum(kwd[estr]**2)) + + # ----------- + # ortho + # ----------- + + if kwd['ortho'] is True and kwd['nd'] != '1d': + if kwd['nd'] == '2d': + lsca = [np.sum(kwd['e0'] * kwd['e1'])] + elif kwd['nd'] == '3d': + lsca = [ + np.sum(kwd['e0'] * kwd['e1']), + np.sum(kwd['e1'] * kwd['e2']), + np.sum(kwd['e2'] * kwd['e0']), + ] + + # check perpendicularity + if np.any(np.abs(lsca) > 1e-9 * emax): + msg = ( + "Unit vectors must be perpendicular!" + ) + raise Exception(msg) + + # ---------- + # direct + # ---------- + + if kwd['direct'] is True and kwd['nd'] != '1d': + if kwd['nd'] == '2d': + vect = np.r_[-kwd['e0'][1], kwd['e0'][0]] + sca = np.sum(vect * kwd['e1']) + if sca < 0.: + msg = ( + "The vector basis must be direct!" + ) + raise Exception(msg) + elif kwd['nd'] == '3d': + vect = np.cross(kwd['e0'], kwd['e1']) + sca = np.sum(vect * kwd['e2']) + if sca < 0.: + msg = ( + "The vector basis must be direct!" + ) + raise Exception(msg) + + # ---------- + # units + # ---------- + + lunits = ['units0', 'units1', 'units2'] + for ii in range(size): + kwd[lunits[ii]] = ds._generic_check._check_var( + kwd[lunits[ii]], lunits[ii], + types=str, + default=_DUNITS[kwd['ctype']], + ) + try: + kwd[lunits[ii]] = asunits.Unit(kwd[lunits[ii]]) + except Exception: + pass + + # clean + for ii in range(size, 3): + if kwd[lunits[ii]] is not None: + msg = ( + f"Arg '{lunits[ii]}' provided for a '{kwd['nd']}' csys!\n" + f"Provided: {kwd[lunits[ii]]}\n" + ) + raise Exception(msg) + + # uniformity + lunits = [kwd[f'units{ii}'] for ii in range(size)] + if len(set(lunits)) != 1: + msg = "Non-uniform units!" + raise Exception(msg) + + # --------------- + # clean and check + # --------------- + + lNone = [kk for kk, vv in kwd.items() if vv is None] + assert len(lNone) == (3-size)*2 + for kk in lNone: + del kwd[kk] + + return diff --git a/beampipe/_class00_show.py b/beampipe/_class00_show.py new file mode 100644 index 0000000..bcc9bd5 --- /dev/null +++ b/beampipe/_class00_show.py @@ -0,0 +1,119 @@ +# -*- coding: utf-8 -*- + + +############################################# +############################################# +# DEFAULTS +############################################# + + +_LORDER = [ + 'nd', 'ctype', + 'ortho', 'norm', 'direct', + 'units', + 'kcsys0', +] + + +############################################# +############################################# +# Show +############################################# + + +def _show(coll=None, which=None, lcol=None, lar=None, show=None): + + # --------------------------- + # column names + # --------------------------- + + lcol.append([which] + _LORDER) + + # --------------------------- + # data + # --------------------------- + + lkey = [ + k1 for k1 in coll._dobj.get(which, {}).keys() + if show is None or k1 in show + ] + + lar0 = [] + for k0 in lkey: + + # initialize with key + arr = [k0] + + # loop + for k1 in _LORDER: + + # parameters + if k1 == 'kcsys0' and coll.dobj[which][k0][k1] == k0: + nn = '' + else: + nn = str(coll.dobj[which][k0].get(k1)) + + # units + if k1 == 'units': + size = int(coll.dobj[which][k0]['nd'][0]) + nn = str(tuple([ + str(coll.dobj[which][k0][f'e{ii}']['units']) + for ii in range(size) + ])) + + arr.append(nn) + + lar0.append(arr) + + lar.append(lar0) + + return lcol, lar + + +############################################# +############################################# +# Show single diag +############################################# + + +def _show_details(coll=None, key=None, lcol=None, lar=None, show=None): + + wcsys = coll._which_csys + size = int(coll.dobj[wcsys][key]['nd'][0]) + + # --------------------------- + # column names + # --------------------------- + + lcol.append([ + 'attr', 'x0 (kcsys0_e0)', 'x1 (kcsys0_e1)', 'x2 (kcsys0_e2)', + ]) + + # --------------------------- + # data + # --------------------------- + + lar0 = [] + lk = ['origin', 'e0', 'e1', 'e2'] + for kk in lk: + + # initialize with key, type + arr = [kk] + + # is2d + for ii in range(3): + if ii < size: + if kk == 'origin': + nn = f"{coll.dobj[wcsys][key][kk][ii]:4.3e}" + else: + nn = f"{coll.dobj[wcsys][key][kk]['data'][ii]:4.3e}" + else: + nn = '' + arr.append(nn) + + # aggregate + lar0.append(arr) + + lar.append(lar0) + + return lcol, lar diff --git a/beampipe/_class00_transform.py b/beampipe/_class00_transform.py new file mode 100644 index 0000000..e2e42ef --- /dev/null +++ b/beampipe/_class00_transform.py @@ -0,0 +1,161 @@ +import warnings + + +import numpy as np +import datastock as ds + + +# ############################################# +# ############################################# +# main +# ############################################# + + +def main( + coll=None, + key_in=None, + key_out=None, +): + + # ----------- + # check + # ----------- + + kwd = _check(coll=coll, key_in=key_in, key_out=key_out) + + # ----------- + # transform + # ----------- + + return _get_transform( + coll=coll, + kwd=kwd, + ) + + +# ############################################# +# ############################################# +# check +# ############################################# + + +def _check(**kwd): + + # ----------- + # key_in vs key_out + # ----------- + + coll = kwd['coll'] + wcsys = coll._which_csys + lok = list(coll.dobj.get(wcsys, {}).keys()) + + # key_in + kwd['key_in'] = ds._generic_check._check_var( + kwd['key_in'], 'key_in', + types=str, + allowed=lok, + ) + + # ctype, nd + ctype = coll.dobj[wcsys][kwd['key_in']]['ctype'] + nd = coll.dobj[wcsys][kwd['key_in']]['nd'] + kcsys0 = coll.dobj[wcsys][kwd['key_in']]['kcsys0'] + units = coll.dobj[wcsys][kwd['key_in']]['e0']['units'] + + # key_out + lok = [ + kk for kk in lok + if kk != kwd['key_in'] + if coll.dobj[wcsys][kk]['ctype'] == ctype + and coll.dobj[wcsys][kk]['nd'] == nd + and coll.dobj[wcsys][kk]['kcsys0'] == kcsys0 + and coll.dobj[wcsys][kk]['e0']['units'] == units + ] + kwd['key_out'] = ds._generic_check._check_var( + kwd['key_out'], 'key_out', + types=str, + allowed=lok, + ) + + # ----------- + # implemented ? + # ----------- + + if ctype != 'cart': + msg = f"csys transform not implement for ctype = '{ctype}'\n" + raise Exception(msg) + + return kwd + + +# ############################################# +# ############################################# +# transform +# ############################################# + + +def _get_transform(coll=None, kwd=None): + + # ------------ + # basics + # ------------ + + wcsys = coll._which_csys + ctype = coll.dobj[wcsys][kwd['key_in']]['ctype'] + nd = coll.dobj[wcsys][kwd['key_in']]['nd'] + size = int(nd[0]) + + lx = ['x0', 'x1', 'x2'] + + # ------------ + # data + # ------------ + + origin_in = coll.dobj[wcsys][kwd['key_in']]['origin'] + origin_out = coll.dobj[wcsys][kwd['key_out']]['origin'] + dorigin = origin_in - origin_out + + de = {} + for ii in range(size): + de[f'e{ii}_in'] = coll.dobj[wcsys][kwd['key_in']][f'e{ii}']['data'] + de[f'e{ii}_out'] = coll.dobj[wcsys][kwd['key_out']][f'e{ii}']['data'] + + # ------------ + # units + # ------------ + + lunits = [ + coll.dobj[wcsys][kwd['key_out']][f'e{ii}']['units'] + for ii in range(size) + ] + if len(set(lunits)) > 1: + msg = f"Units are different for each coordinates: {lunits}\n" + warnings.warn(msg) + units = None + elif len(set(lunits)) == 1: + units = lunits[0] + else: + units = None + + # ------------ + # cartesian + # ------------ + + dout = {} + if ctype == 'cart': + for ii in range(size): + + # translation + dout[f'd{lx[ii]}'] = { + 'data': np.sum(dorigin * de[f'e{ii}_out']), + 'units': units, + } + + # rotation + for jj in range(size): + dout[f"cos_e{jj}_e{ii}"] = { + 'data': np.sum(de[f'e{jj}_in'] * de[f"e{ii}_out"]), + 'units': None, + } + + return dout diff --git a/beampipe/_class00_transform_coords.py b/beampipe/_class00_transform_coords.py new file mode 100644 index 0000000..c59aba8 --- /dev/null +++ b/beampipe/_class00_transform_coords.py @@ -0,0 +1,244 @@ +import warnings + + +import numpy as np + + +# ############################################# +# ############################################# +# main +# ############################################# + + +def main( + coll=None, + key_in=None, + key_out=None, + # coordinates + x0=None, + x1=None, + x2=None, + # unused + **kwdargs, +): + + # ----------- + # get transform + # ----------- + + dtrans = coll.get_csys_transform(key_in=key_in, key_out=key_out) + + # ----------- + # check + # ----------- + + kwd = _check(**locals()) + + # ----------- + # transform + # ----------- + + return _transform( + coll=coll, + kwd=kwd, + dtrans=dtrans, + ) + + +# ############################################# +# ############################################# +# check +# ############################################# + + +def _check(**kwd): + + # ----------- + # key_in vs key_out + # ----------- + + coll = kwd['coll'] + wcsys = coll._which_csys + + # ctype, nd + nd = coll.dobj[wcsys][kwd['key_in']]['nd'] + size = int(nd[0]) + + # ------------- + # coordinates + # ------------- + + dfail = {} + lx = ['x0', 'x1', 'x2'] + for ii in range(size): + + # None + if kwd[lx[ii]] is None: + dfail[lx[ii]] = "is None" + + # str + elif isinstance(kwd[lx[ii]], str): + if kwd[lx[ii]] not in coll.ddata.keys(): + dfail[lx[ii]] = f"not found in ddata ({kwd[lx[ii]]})" + + # array + else: + try: + kwd[lx[ii]] = np.atleast_1d(kwd[lx[ii]]) + except Exception: + dfail[lx[ii]] = ( + f"not convertible to np.ndarray ({type(kwd[lx[ii]])})" + ) + + # errors + if len(dfail) > 0: + lstr = [f"\t- {kk}: vv" for kk, vv in dfail.items()] + msg = ( + "Coordinates are not valid:\n" + + "\n".join(lstr) + ) + raise Exception(msg) + + # clean-up + for ii in range(size, 3): + kwd[lx[ii]] = None + + # ------------- + # broadcastable coordinates + # ------------- + + dshapes = { + lx[ii]: kwd[lx[ii]].shape if isinstance(kwd[lx[ii]], np.ndarray) + else coll.ddata[lx[ii]]['data'].shape + for ii in range(size) + } + + try: + _ = np.broadcast_shapes(*list(dshapes.values())) + except Exception: + lstr = [f"\t- {kk}: vv" for kk, vv in dshapes.items()] + msg = ( + "All coordinates must be broadcastable!\n" + + "\n".join(lstr) + ) + raise Exception(msg) + + # ------------- + # clean + # ------------- + + lok = ['key_in', 'key_out', 'x0', 'x1', 'x2'] + lout = [kk for kk in kwd.keys() if kk not in lok] + for kk in lout: + del kwd[kk] + for ii in range(size, 3): + del kwd[f"x{ii}"] + + return kwd + + +# ############################################# +# ############################################# +# transform +# ############################################# + + +def _transform(coll=None, kwd=None, dtrans=None): + + # ------------ + # basics + # ------------ + + wcsys = coll._which_csys + ctype = coll.dobj[wcsys][kwd['key_in']]['ctype'] + nd = coll.dobj[wcsys][kwd['key_in']]['nd'] + size = int(nd[0]) + + lx = ['x0', 'x1', 'x2'] + + # ------------ + # ref + # ------------ + + lref = [ + coll.ddata[kwd[lx[ii]]]['ref'] for ii in range(size) + if isinstance(kwd[lx[ii]], str) + ] + + if len(lref) > 0: + if len(set(lref)) > 1: + msg = "Coordinates do not share the same ref!" + warnings.warn(msg) + ref = None + else: + ref = lref[0] + else: + ref = None + + # ------------ + # units + # ------------ + + units = coll.dobj[wcsys][kwd['key_in']]['e0']['units'] + + lunits = [ + coll.ddata[kwd[lx[ii]]]['units'] for ii in range(size) + if isinstance(kwd[lx[ii]], str) + ] + + if len(lunits) > 0: + + if len(set(lunits)) > 1: + msg = "Coordinates do not share the same units!" + raise Exception(msg) + + elif lunits[0] != units: + msg = "Coordinates do not share the same units as unit vectors!" + raise Exception(msg) + + + # ------------ + # values + # ------------ + + dval = { + lx[ii]: kwd[lx[ii]] if isinstance(kwd[lx[ii]], np.ndarray) + else coll.ddata[kwd[lx[ii]]]['data'] + for ii in range(size) + } + + # ------------ + # dout + # ------------ + + dout = { + lx[ii]: { + 'data': None, + 'units': units, + 'ref': ref, + } + for ii in range(size) + } + + # ------------ + # cartesian + # ------------ + + if ctype == 'cart': + for ii in range(size): + dout[lx[ii]]['data'] = ( + dtrans[f"d{lx[ii]}"]['data'] + + np.sum( + [ + dval[lx[jj]] * dtrans[f'cos_e{jj}_e{ii}']['data'] + for jj in range(size) + ], + axis=0, + ) + ) + + else: + msg = f"tranform for csys of ctype '{ctype}' not implemented yet!" + raise NotImplementedError(msg) + + return dout diff --git a/beampipe/_class01_Outline2d.py b/beampipe/_class01_Outline2d.py new file mode 100644 index 0000000..4ad3d0e --- /dev/null +++ b/beampipe/_class01_Outline2d.py @@ -0,0 +1,68 @@ +import copy +from typing import Annotated + + +from ._class00_CSYS import CSYS as Previous +from . import _class01_check as _check +# from . import _class01_show as _show + + +__all__ = ['Outline2d'] + + +############################################# +############################################# +# DEFAULT VALUES +############################################# + + +_WHICH_OUTLINE2D = 'outline2d' + + +############################################# +############################################# +# Spectral Lines +############################################# + + +class Outline2d(Previous): + + _ddef = copy.deepcopy(Previous._ddef) + _which_outline2d = _WHICH_OUTLINE2D + + # ------------------- + # add csys + # ------------------- + + def add_outline2d( + self, + key: Annotated[str | None, 'key of the outline2d to be added'] = None, + key_csys: Annotated[str | None, 'key of the csys'] = None, + # circle + center=None, + radius=None, + # polygon + outline_x0=None, + outline_x1=None, + # from svg + ) -> None: + """ Add a outline2d + + Coordinates are given with respect to key_csys + + """ + + # ------------ + # check inputs + # ------------ + + dref, ddata, dobj = _check.main(coll=self, **locals()) + + # ------------ + # Populate + # ------------ + + self.update(dref=dref, ddata=ddata, dobj=dobj) + + return + diff --git a/beampipe/_class01_check.py b/beampipe/_class01_check.py new file mode 100644 index 0000000..d5fea92 --- /dev/null +++ b/beampipe/_class01_check.py @@ -0,0 +1,66 @@ +import numpy as np +import datastock as ds + + +############################################# +############################################# +# DEFAULT VALUES +############################################# + + +############################################# +############################################# +# main +############################################# + + +def main( + coll=None, + key=None, + key_csys=None, + # circle + center=None, + radius=None, + # polygon + outline_x0=None, + outline_x1=None, + # unused + **kwdargs, +): + + # --------------- + # check inputs + # --------------- + + kwd = _check(**locals()) + + # --------------- + # dref, ddata + # --------------- + + dref = None + ddata = None + + # --------------- + # dobj + # --------------- + + wout2d = coll._which_outline2d + dobj = {wout2d: {kwd['key']: kwd}} + + return dref, ddata, dobj + + +############################################# +############################################# +# check +############################################# + + +def _check(kwd): + + # ---------- + # key + # ---------- + + return kwd diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..63b7898 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,64 @@ +[build-system] +requires = ["setuptools", "setuptools_scm"] +build-backend = "setuptools.build_meta" + + +#[tool.setuptools.packages.find] +#where = ["datastock"] +#include = ["datastock*"] +#namespaces = false + +[tool.setuptools] +packages = ["beampipe"] # , "beampipe.tests"] + + +[tool.setuptools_scm] +version_file = "beampipe/_version.py" + +[tool.setuptools.package-data] +mypkg = ["*.txt", "*.npz"] + + +[tool.setuptools.dynamic] +classifiers = {file = ["CLASSIFIERS.txt"]} + + +[project] +name = "beampipe" +readme = "README.md" +license = {text = "MIT"} +dynamic = ["version", "classifiers"] +description = "Generic handler for 3d pipes" +authors = [ + {name = "Didier VEZINET", email = "didier.vezinet@gmail.com"}, +] +maintainers = [ + {name = "Didier VEZINET", email = "didier.vezinet@gmail.com"}, +] +keywords = [ + "Collection", "modelling", +] +requires-python = ">=3.10" +dependencies = [ + 'datastock>=0.0.56', +] + + +[project.urls] +Homepage = "https://github.com/ToFuProject/beampipe" +Issues = "https://github.com/ToFuProject/beampipe/issues" + + +[dependency-groups] +dev = [ + "pytest", +] + + +[project.optional-dependencies] +linting = [ + 'ruff' +] +formatting = [ + 'ruff' +]