Merge pull request #124 from ToFuProject/Issue123_BetterSample

[#123] get_sample_mesh_3d_slice(adjust_phi=bool) implemented + minor …

Author: Didou09

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,

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
+    return cropbs

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,17 +525,25 @@ 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]))
             x0 = np.linspace(min0, max0, n0)
         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,13 +560,17 @@ def _sample_2d(
                 knots1[-1:],
             ))
 
+    # --------
+    # scalar
+
     if np.isscalar(x0):
         x0 = np.full(x1.shape, x0)
     if np.isscalar(x1):
         x1 = np.full(x0.shape, x1)
 
     # -----------
     # 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
 

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]