diff --git a/.github/workflows/python-testing-matrix.yml b/.github/workflows/python-testing-matrix.yml index 0c3ddd8..047e69e 100644 --- a/.github/workflows/python-testing-matrix.yml +++ b/.github/workflows/python-testing-matrix.yml @@ -38,4 +38,4 @@ jobs: # Run tests - name: Run tests # For example, using `pytest` - run: uv run pytest bsplines2d/tests + run: uv run pytest bsplines2d/tests -v -x diff --git a/bsplines2d/_class01_Mesh2D.py b/bsplines2d/_class01_Mesh2D.py index 2726c24..1101e45 100644 --- a/bsplines2d/_class01_Mesh2D.py +++ b/bsplines2d/_class01_Mesh2D.py @@ -16,6 +16,8 @@ from . import _class01_cropping as _cropping from . import _class01_select as _select from . import _class01_sample as _sample +from . import _class01_sample3d as _sample3d +from . import _class01_slice3d as _slice3d from . import _class01_outline as _outline from . import _class01_plot as _plot @@ -513,6 +515,83 @@ def get_sample_mesh( kx1=kx1, ) + def get_sample_mesh_3d_func( + self, + key=None, + res_RZ=None, + mode=None, + res_phi=None, + ): + """ Return 2 functions + + Used for vos computation + + func_RZphi_from_ind(ind) + where in = (3, npts) indices of a 3d sampled mesh + return R, Z, phi + + func_ind_from_domain(DR, DZ, Dphi, pcross, phor) + where DR, DZ, Dphi define a rectangular domain (optional) + where (pcross0, pcross1) and (phor0, phor1) define ploygons + return ind (to be used by func_RZphi_from_ind) + + """ + + return _sample3d.main( + coll=self, + key=key, + res_RZ=res_RZ, + mode=mode, + res_phi=res_phi, + ) + + def get_sample_mesh_3d_slice( + self, + key=None, + res=None, + # slice + phi=None, + Z=None, + # domain + DR=None, + DZ=None, + Dphi=None, + # option + reshape_2d=None, + # plot + plot=None, + dax=None, + color=None, + ): + """ Return a dict continaing pts coordinates on a plane (slice) + + Slice can be either horizontal (Z) or poloidal (phi) + + A subset of the mesh can be defined using: + - horizontal: DR and Dphi + - poloidal: DR and DZ + + """ + + return _slice3d.main( + coll=self, + key=key, + res=res, + # slice + phi=phi, + Z=Z, + # domain + DR=DR, + DZ=DZ, + Dphi=Dphi, + # option + reshape_2d=reshape_2d, + # plot + plot=plot, + dax=dax, + color=color, + ) + # ----------------- # outline # ------------------ @@ -561,4 +640,4 @@ def plot_mesh( fs=fs, dleg=dleg, connect=connect, - ) \ No newline at end of file + ) diff --git a/bsplines2d/_class01_sample.py b/bsplines2d/_class01_sample.py index 23b01c1..21fc852 100644 --- a/bsplines2d/_class01_sample.py +++ b/bsplines2d/_class01_sample.py @@ -11,10 +11,10 @@ from . import _generic_mesh -# ############################################################## -# ############################################################## -# Main -# ############################################################## +# ###################################################### +# ###################################################### +# Main +# ###################################################### def sample_mesh( @@ -151,10 +151,10 @@ def sample_mesh( return ddata -# ################################################################## -# ################################################################## -# checks -# ################################################################## +# ########################################################## +# ########################################################## +# checks +# ########################################################## def _check( @@ -500,10 +500,10 @@ def _sample_1d( return xx -# ################################################################## -# ################################################################## -# sample 2d -# ################################################################## +# ######################################################### +# ######################################################### +# sample 2d +# ######################################################### def _sample_2d( @@ -767,4 +767,4 @@ def _store(coll=None, dref=None, ddata=None): # ddata for k0, v0 in ddata.items(): - coll.add_data(**v0) \ No newline at end of file + coll.add_data(**v0) diff --git a/bsplines2d/_class01_sample3d.py b/bsplines2d/_class01_sample3d.py new file mode 100644 index 0000000..029cb70 --- /dev/null +++ b/bsplines2d/_class01_sample3d.py @@ -0,0 +1,497 @@ +# -*- coding: utf-8 -*- + + +# Common +import numpy as np +import datastock as ds +import matplotlib.pyplot as plt +from matplotlib.path import Path + + +# tofu +from . import _generic_mesh + + +# ############################################################## +# ############################################################## +# Main +# ############################################################## + + +def main( + coll=None, + key=None, + res_RZ=None, + mode=None, + res_phi=None, +): + + # ------------- + # check inputs + + ( + key, res_phi, + ) = _check( + coll=coll, + key=key, + res_phi=res_phi, + ) + + # --------------------- + # prepare dsamp + # --------------------- + + dsamp = coll.get_sample_mesh( + key=key, + res=res_RZ, + grid=False, + mode=mode, + x0=None, + x1=None, + Dx0=None, + Dx1=None, + imshow=False, + in_mesh=True, + store=False, + ) + + x0u = dsamp['x0']['data'] + x1u = dsamp['x1']['data'] + index = dsamp['ind']['data'] + + # ------------------ + # phi from res_phi + # ------------------ + + nphi = np.ceil(3 + x0u * (2*np.pi) / res_phi).astype(int) + + # ----------------- + # func_pts_from_ind + # ----------------- + + func_RZphi_from_ind = _get_func_RZphi_from_ind( + x0u=x0u, + x1u=x1u, + nphi=nphi, + ) + + # ----------------- + # func_pts_from_domain + # ----------------- + + func_ind_from_domain = _get_func_ind_from_domain( + x0u=x0u, + x1u=x1u, + index=index, + nphi=nphi, + # debug + func_RZphi_from_ind=func_RZphi_from_ind, + ) + + return func_RZphi_from_ind, func_ind_from_domain + + +# ########################################################## +# ########################################################## +# checks +# ########################################################## + + +def _check( + coll=None, + key=None, + res_phi=None, +): + + # --------- + # mesh + # --------- + + # key + wm = coll._which_mesh + key, _, cat = _generic_mesh._get_key_mesh_vs_bplines( + coll=coll, + key=key, + which=wm, + ) + + lok = ['rect', 'tri'] + if coll.dobj[wm][key]['type'] not in lok: + msg = ( + f"Only accepts mtype in {lok}\n" + f"Provided: {key}\n" + ) + raise Exception(msg) + + # --------- + # res_phi + # --------- + + res_phi = float(ds._generic_check._check_var( + res_phi, 'res_phi', + types=(int, float), + sign='>0', + )) + + return key, res_phi + + +# ########################################################## +# ########################################################## +# pts_from_ind +# ########################################################## + + +def _get_func_RZphi_from_ind( + x0u=None, + x1u=None, + nphi=None, +): + + # -------------- + # func + # ------------- + + def func( + indr=None, + indz=None, + indphi=None, + # resources + x0u=x0u, + x1u=x1u, + nphi=nphi, + ): + + # ---------- + # check + # ---------- + + if indphi is not None: + c0 = (indr is not None) and (indr.shape == indphi.shape) + if not c0: + indr_str = 'None' if indr is None else indr.shape + msg = ( + "Arg indr and indphi must be the same shape!\n" + f"Provided:\n" + f"\t- indr.shape = {indr_str}\n" + f"\t- indphi.shape = {indphi.shape}\n" + ) + raise Exception(msg) + + # ---------- + # R, Z + # ---------- + + # R + if indr is not None: + R = x0u[indr] + else: + R = None + + # Z + if indz is not None: + Z = x1u[indz] + else: + Z = None + + # dS + dS = (x0u[1] - x0u[0]) * (x1u[1] - x1u[0]) + + # ------------- + # phi + # ------------- + + dV = None + phi = None + if indr is not None: + dV = np.full(indr.shape, np.nan) + iru = np.unique(indr) + + if indphi is None: + for iri in iru: + phii = _get_phi_from_nphi(nphi[iri]) + + iok = indr == iri + dV[iok] = dS * x0u[iri] * (phii[1] - phii[0]) + + else: + phi = np.full(indr.shape, np.nan) + for iri in iru: + phii = _get_phi_from_nphi(nphi[iri]) + + iok = indr == iri + dV[iok] = dS * x0u[iri] * (phii[1] - phii[0]) + phi[iok] = phii[indphi[iok]] + + return R, Z, phi, dV + + return func + + +# ########################################################## +# ########################################################## +# pts_ind_from_domain +# ########################################################## + + +def _get_func_ind_from_domain( + x0u=None, + x1u=None, + nphi=None, + index=None, + # debug + func_RZphi_from_ind=None, +): + + x0f = np.repeat(x0u[:, None], x1u.size, axis=1) + x1f = np.repeat(x1u[None, :], x0u.size, axis=0) + ptsf = np.array([x0f.ravel(), x1f.ravel()]).T + shapef = x0f.shape + + # --------------------- + # get polygons - cross + # --------------------- + + # --------------------- + # get cross-section polygon with margin + + def func( + # from domain + DR=None, + DZ=None, + Dphi=None, + # from poly + pcross0=None, + pcross1=None, + phor0=None, + phor1=None, + # resources + index=index, + x0u=x0u, + x1u=x1u, + ptsf=ptsf, + shapef=shapef, + nphi=nphi, + # debug + debug=None, + debug_msg=None, + func_RZphi_from_ind=func_RZphi_from_ind, + ): + # ------------------ + # check + # ------------------ + + index = np.copy(index) + + DR, DZ, Dphi = _check_domain( + DR=DR, + DZ=DZ, + Dphi=Dphi, + ) + + # ------------------ + # limits from domain + # ------------------ + + if DR is not None: + index &= (x0u[:, None] >= DR[0]) & (x0u[:, None] <= DR[1]) + if DZ is not None: + index &= (x1u[None, :] >= DZ[0]) & (x1u[None, :] <= DZ[1]) + + # ------------------ + # limits from pcross + # ------------------ + + if pcross0 is not None: + pcross = Path(np.array([pcross0, pcross1]).T) + index &= pcross.contains_points(ptsf).reshape(shapef) + + # R and Z indices + ir, iz = index.nonzero() + iru = np.unique(ir) + + # ------------ + # indices + # ------------ + + ind = np.empty((3, 0), dtype=int) + for ii, iri in enumerate(iru): + + # --- + # iz + + izi = np.unique(iz[ir == iri]) + if izi.size == 0: + print(f"SKIP: {ii} {iri}") + continue + + # ----- + # iphi + + if phor0 is None and Dphi is None: + iphi = np.arange(0, nphi[iri]) + + elif Dphi is not None: + phii = _get_phi_from_nphi(nphi[iri]) + if Dphi[0] < Dphi[1]: + iphi = np.nonzero( + (phii >= Dphi[0]) & (phii <= Dphi[1]) + )[0] + else: + iphi = np.nonzero( + (phii >= Dphi[0]) | (phii <= Dphi[1]) + )[0] + elif phor0 is not None: + phii = np.pi*np.linspace(-1, 1, nphi[iri]) + pts = np.array([ + x0u[iri]*np.cos(phii), + x0u[iri]*np.sin(phii), + ]).T + path = Path(np.array([phor0, phor1]).T) + iphi = path.contains_points(pts).nonzero()[0] + + # --------- + # group ind + + indi = np.array( + [ + np.full((iphi.size*izi.size,), iri), + np.repeat(izi, iphi.size), + np.tile(iphi, izi.size), + ], + dtype=int, + ) + + ind = np.concatenate((ind, indi), axis=1) + + # ---------- + # debug + + if debug is True: + + # coordinates + rr, zz, pp, dV = func_RZphi_from_ind(ind) + + # title + tit = ( + "func_ind_from_domain\n" + f"pcross0 {pcross0 is not None} - phor0 {phor0 is not None}\n" + f"DR {DR is not None} - DZ {DZ is not None}" + f" - Dphi{Dphi is not None}" + ) + if debug_msg is not None: + tit += f"\n{debug_msg}" + + # figure and axes + fig = plt.figure(figsize=(14, 8)) + fig.suptitle(tit, size=14, fontweight='bold') + + ax0 = fig.add_subplot(1, 2, 1, aspect='equal') + ax0.set_xlabel("R (m)", size=12, fontweight='bold') + ax0.set_ylabel("Z (m)", size=12, fontweight='bold') + + ax1 = fig.add_subplot(1, 2, 2, aspect='equal') + ax1.set_xlabel("X (m)", size=12, fontweight='bold') + ax1.set_ylabel("Y (m)", size=12, fontweight='bold') + + # domain + if DR is not None: + ax0.axvline(DR[0], c='k', ls='--') + ax0.axvline(DR[1], c='k', ls='--') + + if DZ is not None: + ax0.axhline(DZ[0], c='k', ls='--') + ax0.axhline(DZ[1], c='k', ls='--') + + if Dphi is not None: + maxr = np.max(rr) + ax1.plot( + np.r_[0, maxr]*np.cos(Dphi[0]), + np.r_[0, maxr]*np.sin(Dphi[0]), + c='k', + ls='--', + ) + ax1.plot( + np.r_[0, maxr]*np.cos(Dphi[1]), + np.r_[0, maxr]*np.sin(Dphi[1]), + c='k', + ls='--', + ) + + # polygons + if pcross0 is not None: + ax0.fill(pcross0, pcross1, fc=(0.5, 0.5, 0.5, 0.5)) + if phor0 is not None: + ax1.fill(phor0, phor1, fc=(0.5, 0.5, 0.5, 0.5)) + + # points + ax0.plot(rr, zz, '.') + ax1.plot(rr*np.cos(pp), rr*np.sin(pp), '.') + + return ind[0, :], ind[1, :], ind[2, :] + + return func + + +# ########################################################## +# ########################################################## +# unique get phi function +# ########################################################## + + +def _get_phi_from_nphi(nphi): + return np.pi * np.linspace(-1, 1, nphi, endpoint=False) + + +# ########################################################## +# ########################################################## +# _check_domain +# ########################################################## + + +def _check_domain( + DR=None, + DZ=None, + Dphi=None, +): + # ----------- + # DR + # ----------- + + if DR is not None: + DR = ds._generic_check._check_flat1darray( + DR, 'DR', + size=2, + dtype=float, + unique=True, + sign='>=0', + ) + + # ----------- + # DZ + # ----------- + + if DZ is not None: + DZ = ds._generic_check._check_flat1darray( + DZ, 'DZ', + size=2, + dtype=float, + unique=True, + ) + + # ----------- + # Dphi + # ----------- + + if Dphi is not None: + Dphi = ds._generic_check._check_flat1darray( + Dphi, 'Dphi', + size=2, + dtype=float, + ) + Dphi = [ + np.arctan2(np.sin(Dphi[0]), np.cos(Dphi[0])), + np.arctan2(np.sin(Dphi[1]), np.cos(Dphi[1])), + ] + + return DR, DZ, Dphi diff --git a/bsplines2d/_class01_slice3d.py b/bsplines2d/_class01_slice3d.py new file mode 100644 index 0000000..d196366 --- /dev/null +++ b/bsplines2d/_class01_slice3d.py @@ -0,0 +1,519 @@ +# -*- coding: utf-8 -*- + + +# Common +import numpy as np +import datastock as ds +import matplotlib.pyplot as plt +import matplotlib.gridspec as gridspec + + +# ################################################## +# ################################################## +# Main +# ################################################## + + +def main( + coll=None, + key=None, + res=None, + # slice + Z=None, + phi=None, + # domain + DR=None, + DZ=None, + Dphi=None, + # option + reshape_2d=None, + # plot + plot=None, + dax=None, + color=None, +): + + # -------------- + # check inputs + # -------------- + + ( + res, Z, phi, domain, reshape_2d, plot, + ) = _check( + res=res, + # slice + Z=Z, + phi=phi, + # domain + DR=DR, + DZ=DZ, + Dphi=Dphi, + # option + reshape_2d=reshape_2d, + plot=plot, + ) + + # -------------- + # prepare + # -------------- + + ( + func_RZphi_from_ind, + func_ind_from_domain, + ) = coll.get_sample_mesh_3d_func( + key=key, + res_RZ=res, + res_phi=res, + mode='abs', + ) + + # -------------- + # compute + # -------------- + + if phi is None: + indr, indz, indphi = _horizontal_slice( + res=res, + # sampling + func_ind_from_domain=func_ind_from_domain, + func_RZphi_from_ind=func_RZphi_from_ind, + # slice + Z=Z, + # domain + **domain, + ) + + else: + indr, indz, indphi = _poloidal_slice( + res=res, + # sampling + func_ind_from_domain=func_ind_from_domain, + func_RZphi_from_ind=func_RZphi_from_ind, + # slice + phi=phi, + # option + reshape_2d=reshape_2d, + # domain + **domain, + ) + + # ------------ + # get points + # ------------ + + pts_r, pts_z, pts_phi, _ = func_RZphi_from_ind( + indr=indr, + indz=indz, + indphi=indphi, + ) + + # -------------- + # output + # -------------- + + dout = { + # indices + 'indr': { + 'data': indr, + 'units': 'index', + }, + 'indz': { + 'data': indz, + 'units': 'index', + }, + 'indphi': { + 'data': indphi, + 'units': 'index', + }, + # points + 'pts_r': { + 'data': pts_r, + 'units': 'm', + }, + 'pts_z': { + 'data': pts_z, + 'units': 'm', + }, + 'pts_phi': { + 'data': pts_phi, + 'units': 'rad', + }, + } + + # ----------------- + # optional plotting + # ----------------- + + if plot is True: + _plot( + dout, + dax=dax, + color=color, + ) + + return dout + + +# ############################################## +# ############################################## +# Check +# ############################################## + + +def _check( + res=None, + # slice + Z=None, + phi=None, + # domain + DR=None, + DZ=None, + Dphi=None, + # option + reshape_2d=None, + # plot + plot=None, + color=None, +): + + # --------- + # res + # --------- + + res = float(ds._generic_check._check_var( + res, 'res', + types=(int, float), + sign=">0", + )) + + # --------- + # Z vs phi + # --------- + + lc = [ + Z is not None, + phi is not None, + ] + if np.sum(lc) != 1: + msg = ( + "Please provide Z xor phi!\n" + "\t- Z: height of horizontal slice\n" + "\t- phi: toroidal angle of poloidal slice\n" + "Provided:\n" + f"\t- Z = {Z}\n" + f"\t- phi = {phi}\n" + ) + raise Exception(msg) + + if phi is None: + Z = float(ds._generic_check._check_var( + Z, 'Z', + types=(int, float), + )) + else: + phi = float(ds._generic_check._check_var( + phi, 'phi', + types=(int, float), + )) + phi = np.arctan2(np.sin(phi), np.cos(phi)) + + # --------- + # domain + # --------- + + domain = { + 'DR': DR, + 'DZ': DZ, + 'Dphi': Dphi, + } + + for k0, v0 in domain.items(): + if v0 is not None: + v0 = ds._generic_check._check_flat1darray( + v0, k0, + dtype=float, + unique=True, + size=2, + ) + + if v0[0] >= v0[1]: + msg = ( + f"If provided, arg '{k0}' must be strictly increasing!\n" + + "For Dphi, use +2pi if needed\n" if k0 == 'Dphi' else '' + ) + raise Exception(msg) + + if phi is None: + del domain['DZ'] + else: + del domain['Dphi'] + + # --------- + # reshape_2d + # --------- + + reshape_2d = ds._generic_check._check_var( + reshape_2d, 'reshape_2d', + types=bool, + default=True, + ) + + # --------- + # plot + # --------- + + plot = ds._generic_check._check_var( + plot, 'plot', + types=bool, + default=False, + ) + + return res, Z, phi, domain, reshape_2d, plot + + +# ############################################## +# ############################################## +# Horizontal slice +# ############################################## + + +def _horizontal_slice( + res=None, + # pre-sampling + func_ind_from_domain=None, + func_RZphi_from_ind=None, + # slice + Z=None, + # domain + DR=None, + Dphi=None, +): + # ----------- + # domain Z + # ----------- + + dZ = 1.5*res + DZ = (Z - dZ, Z + dZ) + + indr, indz, indphi = func_ind_from_domain( + DR=DR, + DZ=DZ, + Dphi=Dphi, + ) + + # ------------ + # select plane + # ------------ + + izu = np.unique(indz) + _, zz, _, _ = func_RZphi_from_ind(indz=izu) + + iz = np.argmin(np.abs(zz - Z)) + iok = indz == izu[iz] + + return indr[iok], indz[iok], indphi[iok] + + +# ############################################## +# ############################################## +# Poloidal slice +# ############################################## + + +def _poloidal_slice( + res=None, + # pre-sampling + func_ind_from_domain=None, + func_RZphi_from_ind=None, + # slice + phi=None, + # domain + DR=None, + DZ=None, + # option + reshape_2d=None, +): + # ----------- + # domain Z + # ----------- + + dphi = np.pi/8 + Dphi = (phi - dphi, phi + dphi) + + indr, indz, indphi = func_ind_from_domain( + DR=DR, + DZ=DZ, + Dphi=Dphi, + ) + + # ------------ + # select plane + # ------------ + + iru = np.unique(indr) + iphi = np.zeros(iru.shape) + indr_new, indz_new, indphi_new = [], [], [] + for ii, iri in enumerate(iru): + ind = indr == iri + rr, _, phii, dV = func_RZphi_from_ind( + indr=indr[ind], + indphi=indphi[ind] + ) + + iphi[ii] = indphi[ind][np.argmin(np.abs(phii - phi))] + + ind[ind] = indphi[ind] == iphi[ii] + indr_new.append(indr[ind]) + indz_new.append(indz[ind]) + indphi_new.append(indphi[ind]) + + # --------- + # reshape + # --------- + ir = np.concatenate(indr_new) + iz = np.concatenate(indz_new) + iphi = np.concatenate(indphi_new) + if reshape_2d is True: + i0u = np.unique(ir) + i1u = np.unique(iz) + + indr = -np.ones((i0u.size, i1u.size), dtype=int) + indz = -np.ones((i0u.size, i1u.size), dtype=int) + indphi = -np.ones((i0u.size, i1u.size), dtype=int) + + for ii, iri in enumerate(i0u): + ind = ir == iri + + indsz = np.argsort(iz[ind]) + iiz = np.searchsorted(i1u, np.sort(iz[ind])) + sli = (iri, iiz) + + indr[sli] = iri + indz[sli] = iz[ind][indsz] + indphi[sli] = iphi[ind][indsz] + + else: + indr = ir + indz = iz + indphi = iphi + + return indr, indz, indphi + + +# ############################################## +# ############################################## +# Plot +# ############################################## + + +def _plot( + dout=None, + dax=None, + color=None, +): + # -------------- + # prepare data + # -------------- + + iok = dout['indr']['data'] >= 0 + xx = dout['pts_r']['data'][iok] * np.cos(dout['pts_phi']['data'][iok]) + yy = dout['pts_r']['data'][iok] * np.sin(dout['pts_phi']['data'][iok]) + + # -------------- + # prepare figure + # -------------- + + if dax is None: + + dmargin = { + 'left': 0.05, 'right': 0.95, + 'bottom': 0.06, 'top': 0.90, + 'hspace': 0.20, 'wspace': 0.20, + } + + fig = plt.figure(figsize=(14, 8)) + gs = gridspec.GridSpec(ncols=3, nrows=1, **dmargin) + + # -------------- + # prepare axes + + ax0 = fig.add_subplot(gs[:, 0], aspect='equal', adjustable='box') + ax0.set_ylabel('Z (m)', size=12, fontweight='bold') + ax0.set_xlabel('R (m)', size=12, fontweight='bold') + ax0.set_title("cross-section", size=14, fontweight='bold') + + ax1 = fig.add_subplot(gs[:, 1], aspect='equal', adjustable='box') + ax1.set_xlabel('X (m)', size=12, fontweight='bold') + ax1.set_ylabel('Y (m)', size=12, fontweight='bold') + ax1.set_title("horizontal", size=14, fontweight='bold') + + ax2 = fig.add_subplot( + gs[:, 2], + aspect='equal', + adjustable='box', + projection='3d', + ) + ax2.set_xlabel('X (m)', size=12, fontweight='bold') + ax2.set_ylabel('Y (m)', size=12, fontweight='bold') + ax2.set_zlabel('Z (m)', size=12, fontweight='bold') + ax2.set_title("3d", size=14, fontweight='bold') + + dax = { + 'cross': {'handle': ax0}, + 'hor': {'handle': ax1}, + '3d': {'handle': ax2}, + } + + # -------------- + # plot cross + # -------------- + + kax = 'cross' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + ax.plot( + dout['pts_r']['data'][iok], + dout['pts_z']['data'][iok], + marker='.', + linestyle='None', + ms=6, + color=color, + ) + + # -------------- + # plot hor + # -------------- + + kax = 'hor' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + ax.plot( + xx, + yy, + marker='.', + linestyle='None', + ms=6, + color=color, + ) + + # -------------- + # plot 3d + # -------------- + + kax = '3d' + if dax.get(kax) is not None: + ax = dax[kax]['handle'] + + ax.plot( + xx, + yy, + dout['pts_z']['data'][iok], + marker='.', + linestyle='None', + ms=6, + color=color, + ) + + return diff --git a/bsplines2d/_generic_mesh.py b/bsplines2d/_generic_mesh.py index e095fe0..5fca062 100644 --- a/bsplines2d/_generic_mesh.py +++ b/bsplines2d/_generic_mesh.py @@ -14,10 +14,10 @@ _ELEMENTS = 'knots' -# ############################################################################# -# ############################################################################# -# names of knots and cents -# ############################################################################# +# ############################################################# +# ############################################################# +# names of knots and cents +# ############################################################# def names_knots_cents(key=None, knots_name=''): @@ -51,10 +51,10 @@ def _get_kwdargs_2d(kwdargs, latt=None): return dim, quant, name, units -# ############################################################### -# ############################################################### -# mesh vs bsplines -# ############################################################### +# ################################################### +# ################################################### +# mesh vs bsplines +# ################################################### def _get_key_mesh_vs_bplines( @@ -100,10 +100,10 @@ def _get_key_mesh_vs_bplines( return keym, keybs, cat -# ############################################################################# -# ############################################################################# -# Mesh2DRect - select -# ############################################################################# +# ################################################### +# ################################################### +# Mesh2DRect - select +# ################################################### def _select_ind_check( diff --git a/bsplines2d/tests/prepublish.py b/bsplines2d/tests/prepublish.py index d8d3ab4..6433fc6 100644 --- a/bsplines2d/tests/prepublish.py +++ b/bsplines2d/tests/prepublish.py @@ -1,4 +1,4 @@ print('test import') -import bsplines2d as dbs2 +import bsplines2d as bs2 print('test version') print(bs2.__version__) diff --git a/bsplines2d/tests/test_01_Mesh2D.py b/bsplines2d/tests/test_01_Mesh2D.py index e6b506f..a618ceb 100644 --- a/bsplines2d/tests/test_01_Mesh2D.py +++ b/bsplines2d/tests/test_01_Mesh2D.py @@ -161,146 +161,160 @@ def test17_sample_mesh_rect(self): def test18_sample_mesh_tri(self): test_input._sample_mesh(self.bs, nd='2d', kind='tri') + # ############## + # sample 3d + # ############## + + def test19_sample_3d_func(self): + test_input._sample_mesh_3d_func(self.bs, nd='2d', kind='rect') + + # ############## + # slice 3d + # ############## + + def test20_slice_3d(self): + test_input._slice_mesh_3d(self.bs, nd='2d', kind='rect') + # ############## # plot # ############## - def test19_plot_mesh_1d(self): + def test21_plot_mesh_1d(self): test_input._plot_mesh(self.bs, nd='1d', kind=None) - def test20_plot_mesh_rect(self): + def test22_plot_mesh_rect(self): test_input._plot_mesh(self.bs, nd='2d', kind='rect') - def test21_plot_mesh_tri(self): + def test23_plot_mesh_tri(self): test_input._plot_mesh(self.bs, nd='2d', kind='tri') # ############## # add bsplines # ############## - def test22_add_bsplines_1d(self): + def test24_add_bsplines_1d(self): test_input._add_bsplines(self.bs, nd='1d') - def test23_add_bsplines_2d_rect(self): + def test25_add_bsplines_2d_rect(self): test_input._add_bsplines(self.bs, kind='rect') - def test24_add_bsplines_2d_tri(self): + def test26_add_bsplines_2d_tri(self): test_input._add_bsplines(self.bs, kind='tri') # ############## # select bsplines # ############## - def test25_select_bsplines_1d(self): + def test27_select_bsplines_1d(self): test_input._select_bsplines(self.bs, nd='1d', kind=None) - def test26_select_bsplines_2d_rect(self): + def test28_select_bsplines_2d_rect(self): test_input._select_bsplines(self.bs, nd='2d', kind='rect') - def test27_select_bsplines_2d_tri(self): + def test29_select_bsplines_2d_tri(self): test_input._select_bsplines(self.bs, nd='2d', kind='tri') # ############### # add data vs bs # ############### - def test28_add_data_1bs_fix_1d(self, remove=True): + def test30_add_data_1bs_fix_1d(self, remove=True): test_input._add_data_1bs_fix(self.bs, nd='1d', kind=None, remove=remove) - def test29_add_data_1bs_fix_2d_rect(self, remove=True): + def test31_add_data_1bs_fix_2d_rect(self, remove=True): test_input._add_data_1bs_fix(self.bs, nd='2d', kind='rect', remove=remove) - def test30_add_data_1bs_fix_2d_tri(self, remove=True): + def test32_add_data_1bs_fix_2d_tri(self, remove=True): test_input._add_data_1bs_fix(self.bs, nd='2d', kind='tri', remove=remove) - def test31_add_data_1bs_arrays_1d(self, remove=False): + def test33_add_data_1bs_arrays_1d(self, remove=False): test_input._add_data_1bs_arrays(self.bs, nd='1d', kind=None, remove=remove) - def test32_add_data_1bs_arrays_2d_rect(self, remove=False): + def test34_add_data_1bs_arrays_2d_rect(self, remove=False): test_input._add_data_1bs_arrays(self.bs, nd='2d', kind='rect', remove=remove) - def test33_add_data_1bs_arrays_2d_tri(self, remove=False): + def test35_add_data_1bs_arrays_2d_tri(self, remove=False): test_input._add_data_1bs_arrays(self.bs, nd='2d', kind='tri', remove=remove) - def test34_add_data_multibs_arrays(self, remove=False): + def test36_add_data_multibs_arrays(self, remove=False): test_input._add_data_multibs_arrays(self.bs, remove=remove) # ############## # interp bs # ############## - def test35_interpolate_bsplines_1d(self): + def test37_interpolate_bsplines_1d(self): test_input._interpolate(self.bs, nd='1d', kind=None, details=False) - def test36_interpolate_bsplines_1d_details(self): + def test38_interpolate_bsplines_1d_details(self): test_input._interpolate(self.bs, nd='1d', kind=None, details=True) - def test37_interpolate_bsplines_2d_rect(self): + def test39_interpolate_bsplines_2d_rect(self): test_input._interpolate(self.bs, nd='2d', kind='rect', details=False) - def test38_interpolate_bsplines_2d_rect_details(self): + def test40_interpolate_bsplines_2d_rect_details(self): test_input._interpolate(self.bs, nd='2d', kind='rect', details=True) - def test39_interpolate_bsplines_2d_tri(self): + def test41_interpolate_bsplines_2d_tri(self): test_input._interpolate(self.bs, nd='2d', kind='tri', details=False) - def test40_interpolate_bsplines_2d_tri_details(self): + def test42_interpolate_bsplines_2d_tri_details(self): test_input._interpolate(self.bs, nd='2d', kind='tri', details=True) # ############## # binning 1d # ############## - def test41_binning_1d(self): + def test43_binning_1d(self): test_input._bin_bs(self.bs, nd='1d', kind=None) # ############## # plot bsplines # ############## - def test42_plot_bsplines_1d(self): + def test44_plot_bsplines_1d(self): pass - def test43_plot_bsplines_2d_rect(self): + def test45_plot_bsplines_2d_rect(self): pass - def test44_plot_bsplines_2d_tri(self): + def test46_plot_bsplines_2d_tri(self): pass # #################### # add mesh with subkey # #################### - def test45_add_mesh_1d_subkey_1d(self): + def test47_add_mesh_1d_subkey_1d(self): test_input._add_mesh_1d_subkey(self.bs, nd='1d', kind=None) - def test46_add_mesh_1d_subkey_rect(self): + def test48_add_mesh_1d_subkey_rect(self): test_input._add_mesh_1d_subkey(self.bs, nd='2d', kind='rect') - def test47_add_mesh_1d_subkey_tri(self): + def test49_add_mesh_1d_subkey_tri(self): test_input._add_mesh_1d_subkey(self.bs, nd='2d', kind='tri') - def test48_add_mesh_2d_rect_subkey_rect(self): + def test50_add_mesh_2d_rect_subkey_rect(self): test_input._add_mesh_2d_rect_subkey(self.bs, nd='2d', kind='rect') - def test49_add_mesh_2d_rect_subkey_tri(self): + def test51_add_mesh_2d_rect_subkey_tri(self): test_input._add_mesh_2d_rect_subkey(self.bs, nd='2d', kind='tri') - def test50_add_mesh_2d_rect_var_subkey_rect(self): + def test52_add_mesh_2d_rect_var_subkey_rect(self): test_input._add_mesh_2d_rect_subkey(self.bs, nd='2d', kind='rect') # ################################ # add bsplines on mesh with subkey # ################################ - def test51_add_bsplines_subkey(self): + def test53_add_bsplines_subkey(self): test_input._add_bsplines(self.bs, subkey=True) # ################################ # add data on bsplines with subkey # ################################ - def test52_add_data_subkey(self): + def test54_add_data_subkey(self): test_input._add_data_multibs_arrays( self.bs, nd=None, @@ -313,7 +327,7 @@ def test52_add_data_subkey(self): # interpolate data with subkey # ################################ - def test53_interpolate_subkey_1d(self): + def test55_interpolate_subkey_1d(self): test_input._interpolate( self.bs, nd='1d', @@ -330,14 +344,14 @@ def test53_interpolate_subkey_1d(self): # interpolate data with subkey # ################################ - def test54_plot_as_profile2d(self): + def test56_plot_as_profile2d(self): test_input._plot_as_profile2d( self.bs, nd='2d', kind=None, ) - def test55_plot_as_profile2d_compare(self): + def test57_plot_as_profile2d_compare(self): test_input._plot_as_profile2d_compare( self.bs, nd='2d', @@ -348,31 +362,31 @@ def test55_plot_as_profile2d_compare(self): # operators # ################################ - def test56_operators_1d(self): + def test58_operators_1d(self): test_input._get_operators( self.bs, nd='1d', kind=None, ) - def test57_operators_2d_rect(self): + def test59_operators_2d_rect(self): test_input._get_operators( self.bs, nd='2d', kind='rect', ) - def test58_operators_2d_tri(self): + def test60_operators_2d_tri(self): pass - def test59_operators_1d_subkey(self): + def test61_operators_1d_subkey(self): pass # ################################ # saving / loading # ################################ - def test60_saveload_equal(self): + def test62_saveload_equal(self): # save pfe = self.bs.save(return_pfe=True) @@ -386,7 +400,7 @@ def test60_saveload_equal(self): # equal assert self.bs == out - def test61_saveload_coll(self): + def test63_saveload_coll(self): # save pfe = self.bs.save(return_pfe=True) @@ -399,4 +413,4 @@ def test61_saveload_coll(self): os.remove(pfe) # equal - assert self.bs == coll \ No newline at end of file + assert self.bs == coll diff --git a/bsplines2d/tests/test_input.py b/bsplines2d/tests/test_input.py index 9b059a5..ca9b8f3 100644 --- a/bsplines2d/tests/test_input.py +++ b/bsplines2d/tests/test_input.py @@ -458,6 +458,13 @@ def _select_mesh_elements(bsplines, nd=None, kind=None): ) +####################################################### +# +# Samplings +# +####################################################### + + def _sample_mesh(bsplines, nd=None, kind=None): lkm = _get_mesh(bsplines, nd=nd, kind=kind) @@ -493,6 +500,72 @@ def _sample_mesh(bsplines, nd=None, kind=None): ) +def _sample_mesh_3d_func(coll, nd=None, kind=None): + lkm = _get_mesh(coll, nd=nd, kind=kind) + + lres = [0.1] + lmode = ['abs', 'rel'] + lres_phi = [0.05] + + for km in lkm: + for ii, comb in enumerate(itt.product(lres, lmode, lres_phi)): + ( + func_RZphi_from_ind, + func_ind_from_domain, + ) = coll.get_sample_mesh_3d_func( + key=km, + res_RZ=comb[0], + mode=comb[1], + res_phi=comb[2], + ) + + # Call + indr, indz, indphi = func_ind_from_domain( + DR=[1, 2], + DZ=[-0.2, 0.2], + Dphi=[-0.1, 0.1], + ) + + rr, zz, pp, dV = func_RZphi_from_ind(indr, indz, indphi) + + +def _slice_mesh_3d(coll, nd=None, kind=None): + lkm = _get_mesh(coll, nd=nd, kind=kind) + + lres = [0.05, 0.05, 0.05, 0.05] + Z = [0, 0.5, None, None] + phi = [None, None, np.pi/4, 3*np.pi/4] + Dphi = [None, np.pi*np.r_[3/4, 5/4], None, None] + DZ = [None, None, [-0.5, 0.5], None] + lreshape_2d = [None, None, True, False] + + lparam = [lres, Z, phi, Dphi, DZ, lreshape_2d] + + for km in lkm: + for ii, comb in enumerate(zip(*lparam)): + + dout = coll.get_sample_mesh_3d_slice( + key=km, + res=comb[0], + Z=comb[1], + phi=comb[2], + Dphi=comb[3], + DZ=comb[4], + reshape_2d=comb[5], + plot=True, + ) + assert isinstance(dout, dict) + + plt.close('all') + + +####################################################### +# +# Plotting +# +####################################################### + + def _plot_mesh(bsplines, nd=None, kind=None): lkm = _get_mesh(bsplines, nd=nd, kind=kind) @@ -1101,4 +1174,4 @@ def _get_data(bs, nd=None, kind=None, submesh=None, maxref=None): else: dkd[k0]['bs'].append(kb) - return dkd \ No newline at end of file + return dkd diff --git a/pyproject.toml b/pyproject.toml index b2eaaa6..0dcc680 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ keywords = [ ] requires-python = ">=3.8" dependencies = [ - "contourpy", + "contourpy", 'datastock>=0.0.54', ]