diff --git a/bsplines2d/_class01_Mesh2D.py b/bsplines2d/_class01_Mesh2D.py index 1101e45..6828ad4 100644 --- a/bsplines2d/_class01_Mesh2D.py +++ b/bsplines2d/_class01_Mesh2D.py @@ -25,10 +25,10 @@ __all__ = ['Mesh2D'] -# ############################################################################# -# ############################################################################# -# -# ############################################################################# +# ############################################################# +# ############################################################# +# Class +# ############################################################# class Mesh2D(ds.DataStock): @@ -487,6 +487,7 @@ def get_sample_mesh( mode=None, x0=None, x1=None, + # options Dx0=None, Dx1=None, imshow=None, @@ -505,6 +506,7 @@ def get_sample_mesh( mode=mode, x0=x0, x1=x1, + # options Dx0=Dx0, Dx1=Dx1, imshow=imshow, @@ -558,6 +560,7 @@ def get_sample_mesh_3d_slice( Dphi=None, # option reshape_2d=None, + adjust_phi=None, # plot plot=None, dax=None, @@ -586,6 +589,7 @@ def get_sample_mesh_3d_slice( Dphi=Dphi, # option reshape_2d=reshape_2d, + adjust_phi=adjust_phi, # plot plot=plot, dax=dax, diff --git a/bsplines2d/_class01_cropping.py b/bsplines2d/_class01_cropping.py index efc6728..f3bb5ec 100644 --- a/bsplines2d/_class01_cropping.py +++ b/bsplines2d/_class01_cropping.py @@ -14,10 +14,10 @@ from . import _class01_select as _select -# ############################################################################# -# ############################################################################# +# ################################################################## +# ################################################################## # crop rect mesh -# ############################################################################# +# ################################################################## def crop( @@ -58,7 +58,10 @@ def crop( # ------------ key, mtype, cropbool, thresh_in, remove_isolated = _crop_check( - coll=coll, key=key, crop=crop, thresh_in=thresh_in, + coll=coll, + key=key, + crop=crop, + thresh_in=thresh_in, remove_isolated=remove_isolated, ) @@ -206,21 +209,18 @@ def _crop_check( cropbool = crop.dtype == np.bool_ + # -------------------- # thresh_in and maxth - if thresh_in is None: - thresh_in = 3 + # -------------------- + maxth = 5 if coll.dobj[wm][key]['type'] == 'rect' else 4 - c0 = ( - isinstance(thresh_in, (int, np.integer)) - and (1 <= thresh_in <= maxth) - ) - if not c0: - msg = ( - f"Arg thresh_in must be a int in 1 <= thresh_in <= {maxth}\n" - f"Provided: {thresh_in}" - ) - raise Exception(msg) + thresh_in = int(ds._generic_check._check_var( + thresh_in, 'thresh_in', + default=3, + types=(int, float), + sign=['>=1', f"<={maxth}"], + )) # ---------------- # remove_isolated @@ -328,4 +328,4 @@ def _get_cropbs_from_crop(coll=None, crop=None, keybs=None): for ii in range(shapebs[0]) ], dtype=bool) - return cropbs \ No newline at end of file + return cropbs diff --git a/bsplines2d/_class01_sample.py b/bsplines2d/_class01_sample.py index 21fc852..f2c16f6 100644 --- a/bsplines2d/_class01_sample.py +++ b/bsplines2d/_class01_sample.py @@ -457,8 +457,8 @@ def _get_res( ) if not c0: msg = ( - f"Arg res must be a list of 2 positive floats!\n" - "Provided: {res}" + "Arg res must be a list of 2 positive floats!\n" + f"Provided: {res}" ) raise Exception(msg) @@ -525,10 +525,14 @@ def _sample_2d( # -------- # compute + # -------- min0, max0 = knots0.min(), knots0.max() min1, max1 = knots1.min(), knots1.max() + # -------- + # absolute + if mode == 'abs': if x0 is None: n0 = int(np.ceil((max0 - min0) / res[0])) @@ -536,6 +540,10 @@ def _sample_2d( if x1 is None: n1 = int(np.ceil((max1 - min1) / res[1])) x1 = np.linspace(min1, max1, n1) + + # -------- + # relative + else: if x0 is None: n0 = int(np.ceil(1./res[0])) @@ -552,6 +560,9 @@ def _sample_2d( knots1[-1:], )) + # -------- + # scalar + if np.isscalar(x0): x0 = np.full(x1.shape, x0) if np.isscalar(x1): @@ -559,6 +570,7 @@ def _sample_2d( # ----------- # prepare ind + # ----------- # x0, x1 if grid or in_mesh or Dx0 is not None: @@ -688,7 +700,6 @@ def _store_2d( kx1=None, ): - # ----------- # check key diff --git a/bsplines2d/_class01_slice3d.py b/bsplines2d/_class01_slice3d.py index d196366..d0559b8 100644 --- a/bsplines2d/_class01_slice3d.py +++ b/bsplines2d/_class01_slice3d.py @@ -27,6 +27,7 @@ def main( Dphi=None, # option reshape_2d=None, + adjust_phi=None, # plot plot=None, dax=None, @@ -38,7 +39,7 @@ def main( # -------------- ( - res, Z, phi, domain, reshape_2d, plot, + res, Z, phi, domain, reshape_2d, adjust_phi, plot, ) = _check( res=res, # slice @@ -50,6 +51,7 @@ def main( Dphi=Dphi, # option reshape_2d=reshape_2d, + adjust_phi=adjust_phi, plot=plot, ) @@ -107,6 +109,13 @@ def main( indphi=indphi, ) + # ------------ + # adjust phi + # ------------ + + if adjust_phi is True: + pts_phi = phi + # -------------- # output # -------------- @@ -171,6 +180,7 @@ def _check( Dphi=None, # option reshape_2d=None, + adjust_phi=None, # plot plot=None, color=None, @@ -258,6 +268,18 @@ def _check( default=True, ) + # --------- + # adjust_phi + # --------- + + adjust_phi = ds._generic_check._check_var( + adjust_phi, 'adjust_phi', + types=bool, + default=False, + ) + if phi is None: + adjust_phi = False + # --------- # plot # --------- @@ -268,7 +290,7 @@ def _check( default=False, ) - return res, Z, phi, domain, reshape_2d, plot + return res, Z, phi, domain, reshape_2d, adjust_phi, plot # ############################################## @@ -337,7 +359,7 @@ def _poloidal_slice( # domain Z # ----------- - dphi = np.pi/8 + dphi = np.pi/12 Dphi = (phi - dphi, phi + dphi) indr, indz, indphi = func_ind_from_domain( @@ -370,9 +392,19 @@ def _poloidal_slice( # --------- # reshape # --------- + ir = np.concatenate(indr_new) iz = np.concatenate(indz_new) iphi = np.concatenate(indphi_new) + + # ------------ + # safety check + + assert np.unique([ir, iz], axis=1).shape[1] == ir.size + + # ---------------- + # optional reshape + if reshape_2d is True: i0u = np.unique(ir) i1u = np.unique(iz) @@ -386,7 +418,7 @@ def _poloidal_slice( indsz = np.argsort(iz[ind]) iiz = np.searchsorted(i1u, np.sort(iz[ind])) - sli = (iri, iiz) + sli = (ii, iiz) indr[sli] = iri indz[sli] = iz[ind][indsz]