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 884a0d7..1101e45 100644 --- a/bsplines2d/_class01_Mesh2D.py +++ b/bsplines2d/_class01_Mesh2D.py @@ -17,6 +17,7 @@ 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 @@ -544,6 +545,53 @@ def get_sample_mesh_3d_func( 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 # ------------------ diff --git a/bsplines2d/_class01_sample3d.py b/bsplines2d/_class01_sample3d.py index 82c3671..029cb70 100644 --- a/bsplines2d/_class01_sample3d.py +++ b/bsplines2d/_class01_sample3d.py @@ -339,9 +339,8 @@ def func( )[0] else: iphi = np.nonzero( - (phii >= Dphi[1]) | (phii <= Dphi[0]) + (phii >= Dphi[0]) | (phii <= Dphi[1]) )[0] - elif phor0 is not None: phii = np.pi*np.linspace(-1, 1, nphi[iri]) pts = np.array([ @@ -490,5 +489,9 @@ def _check_domain( 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/tests/test_01_Mesh2D.py b/bsplines2d/tests/test_01_Mesh2D.py index f5a7d2f..a618ceb 100644 --- a/bsplines2d/tests/test_01_Mesh2D.py +++ b/bsplines2d/tests/test_01_Mesh2D.py @@ -168,146 +168,153 @@ def test18_sample_mesh_tri(self): 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 test20_plot_mesh_1d(self): + def test21_plot_mesh_1d(self): test_input._plot_mesh(self.bs, nd='1d', kind=None) - def test21_plot_mesh_rect(self): + def test22_plot_mesh_rect(self): test_input._plot_mesh(self.bs, nd='2d', kind='rect') - def test22_plot_mesh_tri(self): + def test23_plot_mesh_tri(self): test_input._plot_mesh(self.bs, nd='2d', kind='tri') # ############## # add bsplines # ############## - def test23_add_bsplines_1d(self): + def test24_add_bsplines_1d(self): test_input._add_bsplines(self.bs, nd='1d') - def test24_add_bsplines_2d_rect(self): + def test25_add_bsplines_2d_rect(self): test_input._add_bsplines(self.bs, kind='rect') - def test25_add_bsplines_2d_tri(self): + def test26_add_bsplines_2d_tri(self): test_input._add_bsplines(self.bs, kind='tri') # ############## # select bsplines # ############## - def test26_select_bsplines_1d(self): + def test27_select_bsplines_1d(self): test_input._select_bsplines(self.bs, nd='1d', kind=None) - def test27_select_bsplines_2d_rect(self): + def test28_select_bsplines_2d_rect(self): test_input._select_bsplines(self.bs, nd='2d', kind='rect') - def test28_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 test29_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 test30_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 test31_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 test32_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 test33_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 test34_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 test35_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 test36_interpolate_bsplines_1d(self): + def test37_interpolate_bsplines_1d(self): test_input._interpolate(self.bs, nd='1d', kind=None, details=False) - def test37_interpolate_bsplines_1d_details(self): + def test38_interpolate_bsplines_1d_details(self): test_input._interpolate(self.bs, nd='1d', kind=None, details=True) - def test38_interpolate_bsplines_2d_rect(self): + def test39_interpolate_bsplines_2d_rect(self): test_input._interpolate(self.bs, nd='2d', kind='rect', details=False) - def test39_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 test40_interpolate_bsplines_2d_tri(self): + def test41_interpolate_bsplines_2d_tri(self): test_input._interpolate(self.bs, nd='2d', kind='tri', details=False) - def test41_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 test42_binning_1d(self): + def test43_binning_1d(self): test_input._bin_bs(self.bs, nd='1d', kind=None) # ############## # plot bsplines # ############## - def test43_plot_bsplines_1d(self): + def test44_plot_bsplines_1d(self): pass - def test44_plot_bsplines_2d_rect(self): + def test45_plot_bsplines_2d_rect(self): pass - def test45_plot_bsplines_2d_tri(self): + def test46_plot_bsplines_2d_tri(self): pass # #################### # add mesh with subkey # #################### - def test46_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 test47_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 test48_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 test49_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 test50_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 test51_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 test52_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 test53_add_data_subkey(self): + def test54_add_data_subkey(self): test_input._add_data_multibs_arrays( self.bs, nd=None, @@ -320,7 +327,7 @@ def test53_add_data_subkey(self): # interpolate data with subkey # ################################ - def test54_interpolate_subkey_1d(self): + def test55_interpolate_subkey_1d(self): test_input._interpolate( self.bs, nd='1d', @@ -337,14 +344,14 @@ def test54_interpolate_subkey_1d(self): # interpolate data with subkey # ################################ - def test55_plot_as_profile2d(self): + def test56_plot_as_profile2d(self): test_input._plot_as_profile2d( self.bs, nd='2d', kind=None, ) - def test56_plot_as_profile2d_compare(self): + def test57_plot_as_profile2d_compare(self): test_input._plot_as_profile2d_compare( self.bs, nd='2d', @@ -355,31 +362,31 @@ def test56_plot_as_profile2d_compare(self): # operators # ################################ - def test57_operators_1d(self): + def test58_operators_1d(self): test_input._get_operators( self.bs, nd='1d', kind=None, ) - def test58_operators_2d_rect(self): + def test59_operators_2d_rect(self): test_input._get_operators( self.bs, nd='2d', kind='rect', ) - def test59_operators_2d_tri(self): + def test60_operators_2d_tri(self): pass - def test60_operators_1d_subkey(self): + def test61_operators_1d_subkey(self): pass # ################################ # saving / loading # ################################ - def test61_saveload_equal(self): + def test62_saveload_equal(self): # save pfe = self.bs.save(return_pfe=True) @@ -393,7 +400,7 @@ def test61_saveload_equal(self): # equal assert self.bs == out - def test62_saveload_coll(self): + def test63_saveload_coll(self): # save pfe = self.bs.save(return_pfe=True) diff --git a/bsplines2d/tests/test_input.py b/bsplines2d/tests/test_input.py index fa3446e..ca9b8f3 100644 --- a/bsplines2d/tests/test_input.py +++ b/bsplines2d/tests/test_input.py @@ -528,6 +528,37 @@ def _sample_mesh_3d_func(coll, nd=None, kind=None): 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