diff --git a/pyproject.toml b/pyproject.toml index 5e6afd4b152ec710373b7969320224b8e853b6d1..ec1e99442bd25b9724ad4fb327dc34248aff9c83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "gvec-to-python" -version = "1.2.0" +version = "1.2.1" readme = "README.md" requires-python = ">=3.7" license = {file = "LICENSE"} diff --git a/src/gvec_to_python/GVEC_functions.py b/src/gvec_to_python/GVEC_functions.py index dab5af44e85171d982e36185a70c71e320341892..f1474b382b9ee485d903a3f4f96c35b060a11c5d 100644 --- a/src/gvec_to_python/GVEC_functions.py +++ b/src/gvec_to_python/GVEC_functions.py @@ -206,268 +206,314 @@ class GVEC: return self._g_inv # direct access to MHD equilibrium - def p0(self, s, a1, a2): + def p0(self, s, a1, a2, flat_eval=False): '''Pressure as 0-form in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. + If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A float (single-point evaluation) or 3d numpy array (meshgrid evaluation). + A float (single-point evaluation) or 3d numpy array (meshgrid evaluation). ''' - return self.profiles.profile(s, a1, a2, name='pressure') + return self.profiles.profile(s, a1, a2, name='pressure', flat_eval=flat_eval) - def p3(self, s, a1, a2): + def p3(self, s, a1, a2, flat_eval=False): '''Pressure as 3-form in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A float (single-point evaluation) or 3d numpy array (meshgrid evaluation). + A float (single-point evaluation) or 3d numpy array (meshgrid evaluation). ''' - return self.det_df(s, a1, a2) * self.p0(s, a1, a2) + return self.det_df(s, a1, a2, flat_eval=flat_eval) * self.p0(s, a1, a2, flat_eval=flat_eval) - def p_cart(self, s, a1, a2): + def p_cart(self, s, a1, a2, flat_eval=False): '''Pressure in Cartesian coordinates, evaluated at one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 2-tuple: first entry is either a float (single-point evaluation) or 3d numpy array (meshgrid evaluation). - Second entry are the Cartesian coordinates of evaluation points. + A 2-tuple: first entry is either a float (single-point evaluation) or 3d numpy array (meshgrid evaluation). + Second entry are the Cartesian coordinates of evaluation points. ''' - return self.p0(s, a1, a2), self.f(s, a1, a2) + return self.p0(s, a1, a2, flat_eval=flat_eval), self.f(s, a1, a2, flat_eval=flat_eval) - def bv(self, s, a1, a2): + def bv(self, s, a1, a2, flat_eval=False): '''Contra-variant components (vector-field) of the magnetic field in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' if self.l is None: - return self._bv_gvec(s, a1, a2) + return self._bv_gvec(s, a1, a2, flat_eval=flat_eval) else: - return self.transform(self._bv_gvec(*self.l(s, a1, a2)), s, a1, a2, kind='pull_vector', full_map=False) + return self.transform(self._bv_gvec(*self.l(s, a1, a2, flat_eval=flat_eval), flat_eval=flat_eval), s, a1, a2, kind='pull_vector', full_map=False, flat_eval=flat_eval) - def b1(self, s, a1, a2): + def b1(self, s, a1, a2, flat_eval=False): '''Co-variant components (1-form) of the magnetic field in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.transform(self.bv(s, a1, a2), s, a1, a2, kind='v_to_1') + return self.transform(self.bv(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='v_to_1', flat_eval=flat_eval) - def b2(self, s, a1, a2): + def b2(self, s, a1, a2, flat_eval=False): '''2-form components of the magnetic field as in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.transform(self.bv(s, a1, a2), s, a1, a2, kind='v_to_2') + return self.transform(self.bv(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='v_to_2', flat_eval=flat_eval) - def b_cart(self, s, a1, a2): + def b_cart(self, s, a1, a2, flat_eval=False): '''Cartesian components of the magnetic field evaluated at one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). - Second entry are the Cartesian coordinates of evaluation points. + A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + Second entry are the Cartesian coordinates of evaluation points. ''' - return self.transform(self.bv(s, a1, a2), s, a1, a2, kind='push_vector'), self.f(s, a1, a2) + return self.transform(self.bv(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='push_vector', flat_eval=flat_eval), self.f(s, a1, a2, flat_eval=flat_eval) - def av(self, s, a1, a2): + def av(self, s, a1, a2, flat_eval=False): '''Contra-variant components (vector-field) of the vector potential in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.transform(self.a1(s, a1, a2), s, a1, a2, kind='1_to_v') + return self.transform(self.a1(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='1_to_v', flat_eval=flat_eval) - def a1(self, s, a1, a2): + def a1(self, s, a1, a2, flat_eval=False): '''Co-variant components (1-form) of the vector potential in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' if self.l is None: - return self._a1_gvec(s, a1, a2) + return self._a1_gvec(s, a1, a2, flat_eval=flat_eval) else: - return self.transform(self._a1_gvec(*self.l(s, a1, a2)), s, a1, a2, kind='pull_1form', full_map=False) + return self.transform(self._a1_gvec(*self.l(s, a1, a2, flat_eval=flat_eval), flat_eval=flat_eval), s, a1, a2, kind='pull_1form', full_map=False, flat_eval=flat_eval) - def a2(self, s, a1, a2): + def a2(self, s, a1, a2, flat_eval=False): '''2-form components of the vector potential as in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.transform(self.a1(s, a1, a2), s, a1, a2, kind='1_to_2') + return self.transform(self.a1(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='1_to_2', flat_eval=flat_eval) - def a_cart(self, s, a1, a2): + def a_cart(self, s, a1, a2, flat_eval=False): '''Cartesian components of vector potential evaluated at one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). - Second entry are the Cartesian coordinates of evaluation points. + A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + Second entry are the Cartesian coordinates of evaluation points. ''' - return self.transform(self.a1(s, a1, a2), s, a1, a2, kind='push_1form'), self.f(s, a1, a2) + return self.transform(self.a1(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='push_1form', flat_eval=flat_eval), self.f(s, a1, a2, flat_eval=flat_eval) - def jv(self, s, a1, a2): + def jv(self, s, a1, a2, flat_eval=False): '''Contra-variant components (vector-field) of the equilibirum current curl B / mu0 in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.transform(self.j2(s, a1, a2), s, a1, a2, kind='2_to_v') + return self.transform(self.j2(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='2_to_v', flat_eval=flat_eval) - def j1(self, s, a1, a2): + def j1(self, s, a1, a2, flat_eval=False): '''Co-variant components (1-form) of the equilibirum current curl B / mu0 in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.transform(self.j2(s, a1, a2), s, a1, a2, kind='2_to_1') + return self.transform(self.j2(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='2_to_1', flat_eval=flat_eval) - def j2(self, s, a1, a2): + def j2(self, s, a1, a2, flat_eval=False): '''2-form components of the equilibirum current curl B / mu0 as in one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' if self.l is None: - return self._j2_gvec(s, a1, a2) + return self._j2_gvec(s, a1, a2, flat_eval=flat_eval) else: - return self.transform(self._j2_gvec(*self.l(s, a1, a2)), s, a1, a2, kind='pull_2form', full_map=False) + return self.transform(self._j2_gvec(*self.l(s, a1, a2, flat_eval=flat_eval), flat_eval=flat_eval), s, a1, a2, kind='pull_2form', full_map=False, flat_eval=flat_eval) - def j_cart(self, s, a1, a2): + def j_cart(self, s, a1, a2, flat_eval=False): '''Cartesian components of equilibirum current curl B / mu0 evaluated at one of the coordinates "gvec", "pest", "unit" or unit_pest. Coordinates can be chosen using the mapping setter. Parameters ---------- - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). - Second entry are the Cartesian coordinates of evaluation points. + A 2-tuple: first entry is a 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + Second entry are the Cartesian coordinates of evaluation points. ''' - j = self.jv(s, a1, a2) - return self.transform(self.jv(s, a1, a2), s, a1, a2, kind='push_vector'), self.f(s, a1, a2) + j = self.jv(s, a1, a2, flat_eval=flat_eval) + return self.transform(self.jv(s, a1, a2, flat_eval=flat_eval), s, a1, a2, kind='push_vector', flat_eval=flat_eval), self.f(s, a1, a2, flat_eval=flat_eval) # test equilibirum condition grad p = J x B / mu0 - def assert_equil(self, s, a1, a2): + def assert_equil(self, s, a1, a2, flat_eval=False): '''Test the equilibirum condition grad p = J x B / mu0.''' - p = self.profiles.profile(s, a1, a2, name='pressure', der=None) - p_s = self.profiles.profile(s, a1, a2, name='pressure', der='s') - j = self.j2(s, a1, a2) - b = self.bv(s, a1, a2) + p = self.profiles.profile(s, a1, a2, name='pressure', der=None, flat_eval=flat_eval) + p_s = self.profiles.profile(s, a1, a2, name='pressure', der='s', flat_eval=flat_eval) + j = self.j2(s, a1, a2, flat_eval=flat_eval) + b = self.bv(s, a1, a2, flat_eval=flat_eval) abs_j = np.sqrt(j[0]**2 + j[1]**2 + j[2]**2) abs_b = np.sqrt(b[0]**2 + b[1]**2 + b[2]**2) @@ -491,76 +537,86 @@ class GVEC: # Internal methods (gvec fields) #--------------------------------- - def _a1_gvec(self, s, th, ze): + def _a1_gvec(self, s, th, ze, flat_eval=False): '''Vector potential as 1-form (co-variant components) in gvec standard coordinates (map f). Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid). + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid). + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' - phi = self.profiles.profile(s, th, ze, name='phi') - dphi = self.profiles.profile(s, th, ze, name='phi', der='s') - chi = self.profiles.profile(s, th, ze, name='chi') + phi = self.profiles.profile(s, th, ze, name='phi', flat_eval=flat_eval) + dphi = self.profiles.profile(s, th, ze, name='phi', der='s', flat_eval=flat_eval) + chi = self.profiles.profile(s, th, ze, name='chi', flat_eval=flat_eval) - return (- self._lambda(s, th, ze) * dphi, phi, -chi) + return (- self._lambda(s, th, ze, flat_eval=flat_eval) * dphi, phi, -chi) - def _bv_gvec(self, s, th, ze): + def _bv_gvec(self, s, th, ze, flat_eval=False): '''Magnetic field as vector-field (contra-variant components) in gvec standard coordinates (map f). Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid). + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid). + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' - det_df = self.domain.det_df_gvec(s, th, ze) + det_df = self.domain.det_df_gvec(s, th, ze, flat_eval=flat_eval) - return (0.*det_df, self._b_small_th(s, th, ze) / det_df, self._b_small_ze(s, th, ze) / det_df) + return (0.*det_df, self._b_small_th(s, th, ze, flat_eval=flat_eval) / det_df, self._b_small_ze(s, th, ze, flat_eval=flat_eval) / det_df) - def _j2_gvec(self, s, th, ze): + def _j2_gvec(self, s, th, ze, flat_eval=False): '''Equilibirum current curl B / mu0, as 2-form in gvec standard coordinates (map f). Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid). + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d or 3d (from meshgrid). + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' mu0 = 1.2566370621219e-6 - b = np.ascontiguousarray(self._bv_gvec(s, th, ze)) + b = np.ascontiguousarray(self._bv_gvec(s, th, ze, flat_eval=flat_eval)) + b = GVEC_domain.prepare_batch_vector(b) - grad_bt = self._grad_b_theta(s, th, ze) - grad_bz = self._grad_b_zeta(s, th, ze) + grad_bt = self._grad_b_theta(s, th, ze, flat_eval=flat_eval) + grad_bz = self._grad_b_zeta(s, th, ze, flat_eval=flat_eval) + db_ds = np.ascontiguousarray((0.*grad_bt[0], grad_bt[0], grad_bz[0])) db_dt = np.ascontiguousarray((0.*grad_bt[1], grad_bt[1], grad_bz[1])) db_dz = np.ascontiguousarray((0.*grad_bt[2], grad_bt[2], grad_bz[2])) - - g = self.domain.g_gvec(s, th, ze) - dg_ds = self.domain.dg_gvec(s, th, ze, der='s') - dg_dt = self.domain.dg_gvec(s, th, ze, der='th') - dg_dz = self.domain.dg_gvec(s, th, ze, der='ze') - - # meshgrid evaluation - if g.ndim == 5: - dgb_ds = GVEC_domain.matvec_5d(dg_ds, b, use_pyccel=self.use_pyccel) + GVEC_domain.matvec_5d(g, db_ds, use_pyccel=self.use_pyccel) - dgb_dt = GVEC_domain.matvec_5d(dg_dt, b, use_pyccel=self.use_pyccel) + GVEC_domain.matvec_5d(g, db_dt, use_pyccel=self.use_pyccel) - dgb_dz = GVEC_domain.matvec_5d(dg_dz, b, use_pyccel=self.use_pyccel) + GVEC_domain.matvec_5d(g, db_dz, use_pyccel=self.use_pyccel) - # single point evaluation - else: - dgb_ds = dg_ds @ b + g @ db_ds - dgb_dt = dg_dt @ b + g @ db_dt - dgb_dz = dg_dz @ b + g @ db_dz + db_ds = GVEC_domain.prepare_batch_vector(db_ds) + db_dt = GVEC_domain.prepare_batch_vector(db_dt) + db_dz = GVEC_domain.prepare_batch_vector(db_dz) + + g = self.domain.g_gvec(s, th, ze, flat_eval=flat_eval) + dg_ds = self.domain.dg_gvec(s, th, ze, der='s', flat_eval=flat_eval) + dg_dt = self.domain.dg_gvec(s, th, ze, der='th', flat_eval=flat_eval) + dg_dz = self.domain.dg_gvec(s, th, ze, der='ze', flat_eval=flat_eval) + + dgb_ds = dg_ds @ b + g @ db_ds + dgb_dt = dg_dt @ b + g @ db_dt + dgb_dz = dg_dz @ b + g @ db_dz + dgb_ds = GVEC_domain.finalize_batch_vector(dgb_ds) + dgb_dt = GVEC_domain.finalize_batch_vector(dgb_dt) + dgb_dz = GVEC_domain.finalize_batch_vector(dgb_dz) j1 = (dgb_dt[2] - dgb_dz[1]) / mu0 j2 = (dgb_dz[0] - dgb_ds[2]) / mu0 @@ -568,112 +624,112 @@ class GVEC: return j1, j2, j3 - def _grad_b_theta(self, s, th, ze): + def _grad_b_theta(self, s, th, ze, flat_eval=False): '''Gradient w.r.t (s, theta, zeta) of the second component of _bv_gvec, returned as 3-tuple. ''' - det_df = self.domain.det_df_gvec(s, th, ze) - dphi = self.profiles.profile(s, th, ze, name='phi', der='s') - dchi = self.profiles.profile(s, th, ze, name='chi', der='s') - ddphi = self.profiles.profile(s, th, ze, name='phi', der='ss') - ddchi = self.profiles.profile(s, th, ze, name='chi', der='ss') - - term1 = ddchi - ddphi * self._lambda_ze(s, th, ze) - dphi * self._lambda_sze(s, th, ze) - term2 = - self._b_small_th(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='s') / det_df + det_df = self.domain.det_df_gvec(s, th, ze, flat_eval=flat_eval) + dphi = self.profiles.profile(s, th, ze, name='phi', der='s', flat_eval=flat_eval) + dchi = self.profiles.profile(s, th, ze, name='chi', der='s', flat_eval=flat_eval) + ddphi = self.profiles.profile(s, th, ze, name='phi', der='ss', flat_eval=flat_eval) + ddchi = self.profiles.profile(s, th, ze, name='chi', der='ss', flat_eval=flat_eval) + + term1 = ddchi - ddphi * self._lambda_ze(s, th, ze, flat_eval=flat_eval) - dphi * self._lambda_sze(s, th, ze, flat_eval=flat_eval) + term2 = - self._b_small_th(s, th, ze, flat_eval=flat_eval) * self.domain.dJ_gvec(s, th, ze, der='s', flat_eval=flat_eval) / det_df db_ds = (term1 + term2) / det_df - term1 = - dphi * self._lambda_thze(s, th, ze) - term2 = - self._b_small_th(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='th') / det_df + term1 = - dphi * self._lambda_thze(s, th, ze, flat_eval=flat_eval) + term2 = - self._b_small_th(s, th, ze, flat_eval=flat_eval) * self.domain.dJ_gvec(s, th, ze, der='th', flat_eval=flat_eval) / det_df db_dt = (term1 + term2) / det_df - term1 = - dphi * self._lambda_zeze(s, th, ze) - term2 = - self._b_small_th(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='ze') / det_df + term1 = - dphi * self._lambda_zeze(s, th, ze, flat_eval=flat_eval) + term2 = - self._b_small_th(s, th, ze, flat_eval=flat_eval) * self.domain.dJ_gvec(s, th, ze, der='ze', flat_eval=flat_eval) / det_df db_dz = (term1 + term2) / det_df return db_ds, db_dt, db_dz - def _grad_b_zeta(self, s, th, ze): + def _grad_b_zeta(self, s, th, ze, flat_eval=False): '''Gradient w.r.t (s, theta, zeta) of the third component of _bv_gvec, returned as 3-tuple. ''' - det_df = self.domain.det_df_gvec(s, th, ze) - dphi = self.profiles.profile(s, th, ze, name='phi', der='s') - ddphi = self.profiles.profile(s, th, ze, name='phi', der='ss') + det_df = self.domain.det_df_gvec(s, th, ze, flat_eval=flat_eval) + dphi = self.profiles.profile(s, th, ze, name='phi', der='s', flat_eval=flat_eval) + ddphi = self.profiles.profile(s, th, ze, name='phi', der='ss', flat_eval=flat_eval) - term1 = ddphi * (1. + self._lambda_th(s, th, ze)) + dphi * self._lambda_sth(s, th, ze) - term2 = - self._b_small_ze(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='s') / det_df + term1 = ddphi * (1. + self._lambda_th(s, th, ze, flat_eval=flat_eval)) + dphi * self._lambda_sth(s, th, ze, flat_eval=flat_eval) + term2 = - self._b_small_ze(s, th, ze, flat_eval=flat_eval) * self.domain.dJ_gvec(s, th, ze, der='s', flat_eval=flat_eval) / det_df db_ds = (term1 + term2) / det_df - term1 = dphi * self._lambda_thth(s, th, ze) - term2 = - self._b_small_ze(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='th') / det_df + term1 = dphi * self._lambda_thth(s, th, ze, flat_eval=flat_eval) + term2 = - self._b_small_ze(s, th, ze, flat_eval=flat_eval) * self.domain.dJ_gvec(s, th, ze, der='th', flat_eval=flat_eval) / det_df db_dt = (term1 + term2) / det_df - term1 = dphi * self._lambda_thze(s, th, ze) - term2 = - self._b_small_ze(s, th, ze) * self.domain.dJ_gvec(s, th, ze, der='ze') / det_df + term1 = dphi * self._lambda_thze(s, th, ze, flat_eval=flat_eval) + term2 = - self._b_small_ze(s, th, ze, flat_eval=flat_eval) * self.domain.dJ_gvec(s, th, ze, der='ze', flat_eval=flat_eval) / det_df db_dz = (term1 + term2) / det_df return db_ds, db_dt, db_dz - def _b_small_th(self, s, th, ze): + def _b_small_th(self, s, th, ze, flat_eval=False): '''B^theta * J''' - dphi = self.profiles.profile(s, th, ze, name='phi', der='s') - dchi = self.profiles.profile(s, th, ze, name='chi', der='s') + dphi = self.profiles.profile(s, th, ze, name='phi', der='s', flat_eval=flat_eval) + dchi = self.profiles.profile(s, th, ze, name='chi', der='s', flat_eval=flat_eval) - return dchi - dphi * self._lambda_ze(s, th, ze) + return dchi - dphi * self._lambda_ze(s, th, ze, flat_eval=flat_eval) - def _b_small_ze(self, s, th, ze): + def _b_small_ze(self, s, th, ze, flat_eval=False): '''B^zeta * J''' - dphi = self.profiles.profile(s, th, ze, name='phi', der='s') + dphi = self.profiles.profile(s, th, ze, name='phi', der='s', flat_eval=flat_eval) - return dphi * (1. + self._lambda_th(s, th, ze)) + return dphi * (1. + self._lambda_th(s, th, ze, flat_eval=flat_eval)) # lambda function - def _lambda(self, s, th, ze): + def _lambda(self, s, th, ze, flat_eval=False): '''The lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef) + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, flat_eval=flat_eval) - def _lambda_s(self, s, th, ze): + def _lambda_s(self, s, th, ze, flat_eval=False): '''s-derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='s') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='s', flat_eval=flat_eval) - def _lambda_th(self, s, th, ze): + def _lambda_th(self, s, th, ze, flat_eval=False): '''theta-derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='th') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='th', flat_eval=flat_eval) - def _lambda_ze(self, s, th, ze): + def _lambda_ze(self, s, th, ze, flat_eval=False): '''zeta-derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='ze') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='ze', flat_eval=flat_eval) - def _lambda_ss(self, s, th, ze): + def _lambda_ss(self, s, th, ze, flat_eval=False): '''Second derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='ss') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='ss', flat_eval=flat_eval) - def _lambda_thth(self, s, th, ze): + def _lambda_thth(self, s, th, ze, flat_eval=False): '''Second derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='thth') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='thth', flat_eval=flat_eval) - def _lambda_zeze(self, s, th, ze): + def _lambda_zeze(self, s, th, ze, flat_eval=False): '''Second derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='zeze') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='zeze', flat_eval=flat_eval) - def _lambda_sth(self, s, th, ze): + def _lambda_sth(self, s, th, ze, flat_eval=False): '''Second derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='sth') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='sth', flat_eval=flat_eval) - def _lambda_sze(self, s, th, ze): + def _lambda_sze(self, s, th, ze, flat_eval=False): '''Second derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='sze') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='sze', flat_eval=flat_eval) - def _lambda_thze(self, s, th, ze): + def _lambda_thze(self, s, th, ze, flat_eval=False): '''Second derivative of the lambda function from gvec.''' - return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='thze') + return self.domain.LA_base.eval_stz(s, th, ze, self.domain.LA_coef, der='thze', flat_eval=flat_eval) #--------------- # Helper methods #--------------- - def transform(self, b, s, a1, a2, kind, full_map=True): + def transform(self, b, s, a1, a2, kind, full_map=True, flat_eval=False): '''Transformations of vector-valued functions under a map m: pull_vector: out = dm_inv * b (dm is the Jacobian of m) @@ -690,23 +746,26 @@ class GVEC: Parameters ---------- - b : tuple[float/array] - 3-tuple of input vector-valued function evaluated as float or array. - - s, a1, a2 : float or array-like - Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). - s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. - - kind : str - Which transformation to use: "pull_vector", "pull_1form", "pull_2form", "push_vector", "push_1form", or "push_2form". - - full_map : bool - True: use the full map "f" to Cartesian coordinates or False: use a map "l" between logical coordinates. - Both are set in the mapping setter. - + b : tuple[float/array] + 3-tuple of input vector-valued function evaluated as float or array. + + s, a1, a2 : float or array-like + Coordinates in logical space. If list or np.array, must be either 1d or 3d (from meshgrid). + s in [0, 1] and a1, a2 are angle-like, either in [0, 2*pi]^2 or in [0, 1]^2. + + kind : str + Which transformation to use: "pull_vector", "pull_1form", "pull_2form", "push_vector", "push_1form", or "push_2form". + + full_map : bool + True: use the full map "f" to Cartesian coordinates or False: use a map "l" between logical coordinates. + Both are set in the mapping setter. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D array arguments are given. + Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation).''' assert isinstance(b, tuple) assert len(b) == 3 @@ -719,7 +778,7 @@ class GVEC: if self.mapping == 'gvec': return b else: - mat = self.dl_inv(s, a1, a2) + mat = self.dl_inv(s, a1, a2, flat_eval=flat_eval) transpose = False elif kind == 'pull_1form': @@ -729,7 +788,7 @@ class GVEC: if self.mapping == 'gvec': return b else: - mat = self.dl(s, a1, a2) + mat = self.dl(s, a1, a2, flat_eval=flat_eval) transpose = True elif kind == 'pull_2form': @@ -739,20 +798,20 @@ class GVEC: if self.mapping == 'gvec': return b else: - mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.dl_inv(s, a1, a2)) * self.det_dl(s, a1, a2)) + mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.dl_inv(s, a1, a2, flat_eval=flat_eval)) * self.det_dl(s, a1, a2, flat_eval=flat_eval)) transpose = False # push to physical elif kind == 'push_vector': if full_map: - mat = self.df(s, a1, a2) + mat = self.df(s, a1, a2, flat_eval=flat_eval) transpose = False else: raise NotImplementedError elif kind == 'push_1form': if full_map: - mat = self.df_inv(s, a1, a2) + mat = self.df_inv(s, a1, a2, flat_eval=flat_eval) transpose = True else: raise NotImplementedError @@ -760,57 +819,55 @@ class GVEC: # transform to other logical elif kind == 'v_to_1': if full_map: - mat = self.g(s, a1, a2) + mat = self.g(s, a1, a2, flat_eval=flat_eval) transpose = False else: raise NotImplementedError elif kind == '1_to_v': if full_map: - mat = self.g_inv(s, a1, a2) + mat = self.g_inv(s, a1, a2, flat_eval=flat_eval) transpose = False else: raise NotImplementedError elif kind == 'v_to_2': if full_map: - out = self.det_df(s, a1, a2) * np.ascontiguousarray(b) + out = self.det_df(s, a1, a2, flat_eval=flat_eval) * np.ascontiguousarray(b) return out[0], out[1], out[2] else: raise NotImplementedError elif kind == '2_to_v': if full_map: - out = np.ascontiguousarray(b) / self.det_df(s, a1, a2) + out = np.ascontiguousarray(b) / self.det_df(s, a1, a2, flat_eval=flat_eval) return out[0], out[1], out[2] else: raise NotImplementedError elif kind == '1_to_2': if full_map: - mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.g_inv(s, a1, a2)) * self.det_df(s, a1, a2)) + mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.g_inv(s, a1, a2, flat_eval=flat_eval)) * self.det_df(s, a1, a2, flat_eval=flat_eval)) transpose = False else: raise NotImplementedError elif kind == '2_to_1': if full_map: - mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.g(s, a1, a2)) / self.det_df(s, a1, a2)) + mat = GVEC_domain.swap_J_axes(GVEC_domain.swap_J_back(self.g(s, a1, a2, flat_eval=flat_eval)) / self.det_df(s, a1, a2, flat_eval=flat_eval)) transpose = False else: raise NotImplementedError else: raise NotImplementedError - - # meshgrid evaluation - if mat.ndim == 5: - out = GVEC_domain.matvec_5d(mat, np.ascontiguousarray(b), transpose_a=transpose, use_pyccel=self.use_pyccel) - # single-point evaluation + + b = GVEC_domain.prepare_batch_vector(np.ascontiguousarray(b)) + if transpose: + matT = mat.swapaxes(-1, -2) + out = matT @ b else: - if transpose: - out = mat.T @ b - else: - out = mat @ b + out = mat @ b + out = GVEC_domain.finalize_batch_vector(out) return out[0], out[1], out[2] diff --git a/src/gvec_to_python/Makefile b/src/gvec_to_python/Makefile index c2c2d60009843c8c727c7daee187fe78041e9d77..e7fdea976f6e12ec04a39909ceda514765501ee2 100644 --- a/src/gvec_to_python/Makefile +++ b/src/gvec_to_python/Makefile @@ -21,7 +21,6 @@ SOURCES := $(BK).py $(BEV1).py OUTPUTS := $(SOURCES:.py=$(SO_EXT)) - #-------------------------------------- # PYCCELIZE #-------------------------------------- @@ -35,7 +34,6 @@ $(BK)$(SO_EXT) : $(BK).py $(BEV1)$(SO_EXT) : $(BEV1).py $(BK)$(SO_EXT) pyccel $(FLAGS) $< - #-------------------------------------- # CLEAN UP #-------------------------------------- diff --git a/src/gvec_to_python/base/base.py b/src/gvec_to_python/base/base.py index 985b1b4eb5ccddd4b851f686fd364c4ca9f37a34..5c82e79ef671ad7c9c2b197eec0f03d894bad720 100644 --- a/src/gvec_to_python/base/base.py +++ b/src/gvec_to_python/base/base.py @@ -61,7 +61,7 @@ class SplineFourierBases: '''Index range of sine modes in `mn` list.''' return self._range_sin - def eval_stz(self, s, th, ze, coef, der=None): + def eval_stz(self, s, th, ze, coef, der=None, flat_eval=False): """Evaluate a B-spline x Fourier basis, or its first or second derivatives, at `s`, `theta`, `zeta`. Parameters @@ -81,13 +81,16 @@ class SplineFourierBases: der : str Which derivative to evaluate: None, "s", "th", "ze", "ss", "thth", "zeze", "sth", "sze" or "thze". + flat_eval : bool + Whether to do flat (marker) evaluation. + Returns ------- float or meshgrid numpy.ndarray Evaluated coordinate. """ - s, th, ze = GVEC_domain.prepare_args(s, th, ze) + s, th, ze = GVEC_domain.prepare_args(s, th, ze, flat_eval=flat_eval) if isinstance(s, np.ndarray): vals = np.zeros((s + th + ze).shape) else: diff --git a/src/gvec_to_python/equilibrium/profiles.py b/src/gvec_to_python/equilibrium/profiles.py index 5ab2656478e6a7b9a2ce3d64d0587b5d848ad60a..24cab3888482e4fea829583b5160884d7c01dc65 100644 --- a/src/gvec_to_python/equilibrium/profiles.py +++ b/src/gvec_to_python/equilibrium/profiles.py @@ -69,20 +69,24 @@ class GVEC_profiles: '''Spline coefficients of pressure profile for basis of spline_space .''' return self._pres_coef - def profile(self, *args, name='phi', der=None): + def profile(self, *args, name='phi', der=None, flat_eval=False): '''Profile (derivatives) as a function of s in [0, 1]. Parameters ---------- *args : float or array-like - Arguments of the profile function. If one array argument is given, it must be 1d. - If three arguments are given (s, a1, a2), a meshgrid evaluation is performed. + Arguments of the profile function. + If one array argument is given, it must be 1d. + If three arguments are given (s, a1, a2), a meshgrid or flat evaluation is performed. name : str Profile identifier; must be one of "phi" (toroidal flux), "chi" (poloidal flux) "iota", or "pressure". der : str Which derivative to evaluate: None, "s" or "ss". + + flat_eval : bool + Whether to do flat (marker) evaluation when three array arguments are given. Returns ------- @@ -110,12 +114,19 @@ class GVEC_profiles: if len(args) == 1: return _fun(args[0], coef) + elif len(args) == 3: - s, a1, a2 = GVEC_domain.prepare_args(args[0], args[1], args[2], sparse=False) - if isinstance(s, np.ndarray): - vals = _fun(s.flatten(), coef).reshape(s.shape) + if flat_eval: + assert np.shape(args[0]) == np.shape(args[1]) == np.shape(args[2]) + assert np.ndim(args[0]) == 1 + vals = _fun(args[0], coef) else: - vals = _fun(s, coef) + s, a1, a2 = GVEC_domain.prepare_args(args[0], args[1], args[2], sparse=False) + if isinstance(s, np.ndarray): + vals = _fun(s.flatten(), coef).reshape(s.shape) + else: + vals = _fun(s, coef) + return vals else: ValueError(f'Number of independent variables must be 1 or 3, is {len(args)}.') diff --git a/src/gvec_to_python/geometry/domain.py b/src/gvec_to_python/geometry/domain.py index 175505ecdd71b91a49cfd7c7763b89be9a7a01d7..d0a1e9ca034856102064cb2fc4f532fc3643565f 100644 --- a/src/gvec_to_python/geometry/domain.py +++ b/src/gvec_to_python/geometry/domain.py @@ -1,8 +1,6 @@ import numpy as np from scipy.optimize import newton -from gvec_to_python.geometry.kernels import matmul_5d_kernel, matvec_5d_kernel, transpose_5d_kernel - class GVEC_domain: """Provides four different mappings from a logical (computational) domain to the physical plasma domain (in Euclidean space). @@ -59,7 +57,7 @@ class GVEC_domain: Evaluation points can be float or array-like (1d or 3d). - There are also three @staticmethods prepare_args, swap_J_axes and matmul_5d. + There are also three @staticmethods prepare_args, swap_J_axes. Parameters ---------- @@ -166,117 +164,138 @@ class GVEC_domain: return self._use_pyccel # standard gvec coordinates - def f_gvec(self, s, th, ze): + def f_gvec(self, s, th, ze, flat_eval=False): '''The map f_gvec:(s, th, ze) -> (x, y, z), where (s, th, ze) in [0, 1] x [0, 2*pi]^2 and (x, y, z) are the Cartesian coordinates. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays ''' - return self.hmap(*self.Xmap(s, th, ze)) + return self.hmap(*self.Xmap(s, th, ze, flat_eval=flat_eval)) - def df_gvec(self, s, th, ze): + def df_gvec(self, s, th, ze, flat_eval=False): '''The Jacobian matrix of f_gvec. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array, - with the last two indices referring to the Jacobian entries at each point.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array, + with the last two indices referring to the Jacobian entries at each point.''' - tmp1 = self.dX(s, th, ze) # Jacobian of the Xmap - tmp2 = self.dh(*self.Xmap(s, th, ze)) # Jacobian of the hmap - - # meshgrid evaluation - if tmp1.ndim == 5: - out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp2 @ tmp1 + tmp1 = self.dX(s, th, ze, flat_eval=flat_eval) # Jacobian of the Xmap + tmp2 = self.dh(*self.Xmap(s, th, ze, flat_eval=flat_eval)) # Jacobian of the hmap - return out + return tmp2 @ tmp1 - def det_df_gvec(self, s, th, ze): + def det_df_gvec(self, s, th, ze, flat_eval=False): '''The Jacobian determinant of f_gvec. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation).''' + A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation).''' - tmp1 = self.det_dX(s, th, ze) - tmp2 = self.det_dh(*self.Xmap(s, th, ze)) + tmp1 = self.det_dX(s, th, ze, flat_eval=flat_eval) + tmp2 = self.det_dh(*self.Xmap(s, th, ze, flat_eval=flat_eval)) return tmp1 * tmp2 - def df_gvec_inv(self, s, th, ze): + def df_gvec_inv(self, s, th, ze, flat_eval=False): '''The inverse Jacobian matrix of f_gvec. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.df_gvec(s, th, ze)) + return np.linalg.inv(self.df_gvec(s, th, ze, flat_eval=flat_eval)) - def g_gvec(self, s, th, ze): + def g_gvec(self, s, th, ze, flat_eval=False): '''The metric tensor of f_gvec. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - tmp = np.ascontiguousarray(self.df_gvec(s, th, ze)) # Jacobian of f_gvec - - # meshgrid evaluation - if tmp.ndim == 5: - out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp.T @ tmp + tmp = np.ascontiguousarray(self.df_gvec(s, th, ze, flat_eval=flat_eval)) # Jacobian of f_gvec + tmpT = tmp.swapaxes(-1, -2) - return out + return tmpT @ tmp - def g_gvec_inv(self, s, th, ze): + def g_gvec_inv(self, s, th, ze, flat_eval=False): '''Inverse metric tensor of f_gvec. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + + return np.linalg.inv(self.g_gvec(s, th, ze, flat_eval=flat_eval)) + + def dg_gvec(self, s, th, ze, der=None, flat_eval=False): + '''Partial derivative of the metric tensor. der in {"s", "th", "ze"} indicates which derivative to evaluate. + + Parameters + ---------- + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + der : str + Which partial derivative to perform, one of "s", "th" or "ze", default is None. - return np.linalg.inv(self.g_gvec(s, th, ze)) + flat_eval : bool + Whether to do flat (marker) evaluation. - def dg_gvec(self, s, th, ze, der=None): - '''Partial derivative of the metric tensor. der in {"s", "th", "ze"} indicates which derivative to evaluate.''' + Returns + ------- + A 2d numpy array (single-point evaluation) or a 5d numpy array. + ''' if der is None: - return self.g_gvec(s, th, ze) + return self.g_gvec(s, th, ze, flat_eval=flat_eval) elif der == 's': comp = 0 elif der == 'th': @@ -286,44 +305,49 @@ class GVEC_domain: else: raise ValueError(f'Derivative {der} not defined.') - dX = self.dX(s, th, ze) - ddX_di = self.ddX(s, th, ze)[comp] + dX = self.dX(s, th, ze, flat_eval=flat_eval) + ddX_di = self.ddX(s, th, ze, flat_eval=flat_eval)[comp] - gh = self.gh(*self.Xmap(s, th, ze)) - dgh_dq1 = self.dgh_dq1(*self.Xmap(s, th, ze)) + gh = self.gh(*self.Xmap(s, th, ze, flat_eval=flat_eval)) + dgh_dq1 = self.dgh_dq1(*self.Xmap(s, th, ze, flat_eval=flat_eval)) - # meshgrid evaluation - if dX.ndim == 5: - dq1_di = dX[:, :, :, 0, comp] - dgh_di = self.swap_J_axes(self.swap_J_back(dgh_dq1) * dq1_di) + tmp = self.swap_J_back(dX) + dq1_di = tmp[0, comp] + dgh_di = self.swap_J_axes(self.swap_J_back(dgh_dq1) * dq1_di) + + tmp1 = gh @ dX + tmp2 = dgh_di @ dX + tmp3 = gh @ ddX_di + + ddX_diT = ddX_di.swapaxes(-1, -2) + dXT = dX.swapaxes(-1, -2) - tmp1 = self.matmul_5d(gh, dX, use_pyccel=self.use_pyccel) - tmp2 = self.matmul_5d(dgh_di, dX, use_pyccel=self.use_pyccel) - tmp3 = self.matmul_5d(gh, ddX_di, use_pyccel=self.use_pyccel) - - term1 = self.matmul_5d(ddX_di, tmp1, transpose_a=True, use_pyccel=self.use_pyccel) - term2 = self.matmul_5d(dX, tmp2, transpose_a=True, use_pyccel=self.use_pyccel) - term3 = self.matmul_5d(dX, tmp3, transpose_a=True, use_pyccel=self.use_pyccel) - # single point evaluation - else: - dq1_di = dX[0, comp] - dgh_di = dgh_dq1 * dq1_di + term1 = ddX_diT @ tmp1 + term2 = dXT @ tmp2 + term3 = dXT @ tmp3 + + return term1 + term2 + term3 + + def dJ_gvec(self, s, th, ze, der=None, flat_eval=False): + '''Partial derivative of Jacobian determinant. der in {"s", "th", "ze"} indicates which derivative to evaluate. - tmp1 = gh @ dX - tmp2 = dgh_di @ dX - tmp3 = gh @ ddX_di + Parameters + ---------- + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. - term1 = ddX_di.T @ tmp1 - term2 = dX.T @ tmp2 - term3 = dX.T @ tmp3 + der : str + Which partial derivative to perform, one of "s", "th" or "ze", default is None. - return term1 + term2 + term3 + flat_eval : bool + Whether to do flat (marker) evaluation. - def dJ_gvec(self, s, th, ze, der=None): - '''Partial derivative of Jacobian determinant. der in {"s", "th", "ze"} indicates which derivative to evaluate.''' + Returns + ------- + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' if der is None: - return self.det_df_gvec(s, th, ze) + return self.det_df_gvec(s, th, ze, flat_eval=flat_eval) elif der == 's': comp = 0 elif der == 'th': @@ -333,432 +357,451 @@ class GVEC_domain: else: raise ValueError(f'Derivative {der} not defined.') - dX = self.dX(s, th, ze) + dX = self.dX(s, th, ze, flat_eval=flat_eval) + tmp = self.swap_J_back(dX) - # meshgrid evaluation - if dX.ndim == 5: - dq1_di = dX[:, :, :, 0, comp] - # single point evaluation - else: - dq1_di = dX[0, comp] + dq1_di = tmp[0, comp] - q1, q2, ze2 = self.Xmap(s, th, ze) + q1, q2, ze2 = self.Xmap(s, th, ze, flat_eval=flat_eval) - return dq1_di * self.det_dX(s, th, ze) + self.det_dh(q1, q2, ze2) * self.dJ_X(s, th, ze)[comp] + return dq1_di * self.det_dX(s, th, ze, flat_eval=flat_eval) + self.det_dh(q1, q2, ze2) * self.dJ_X(s, th, ze, flat_eval=flat_eval)[comp] # gvec_pest coordinates - def f_pest(self, s2, th2, ze2): + def f_pest(self, s2, th2, ze2, flat_eval=False): '''The map f_pest:(s2, th2, ze2) -> (x, y, z), where (s2, th2, ze2) in [0, 1] x [0, 2*pi]^2 are straight-field-line coordinates (PEST) and (x, y, z) are the Cartesian coordinates. Parameters ---------- - s2, th2, ze2 : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s2, th2, ze2 : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.f_gvec(*self.Pmap(s2, th2, ze2)) + return self.f_gvec(*self.Pmap(s2, th2, ze2, flat_eval=flat_eval), flat_eval=flat_eval) - def df_pest(self, s2, th2, ze2): + def df_pest(self, s2, th2, ze2, flat_eval=False): '''The Jacobian matrix of f_pest. Parameters ---------- - s2, th2, ze2 : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s2, th2, ze2 : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array, - with the last two indices referring to the Jacobian entries at each point.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array, + with the last two indices referring to the Jacobian entries at each point.''' + tmp1 = self.dP(s2, th2, ze2, flat_eval=flat_eval) # Jacobian of the Pmap + tmp2 = self.df_gvec(*self.Pmap(s2, th2, ze2, flat_eval=flat_eval), flat_eval=flat_eval) # Jacobian of f_gvec - tmp1 = self.dP(s2, th2, ze2) # Jacobian of the Pmap - tmp2 = self.df_gvec(*self.Pmap(s2, th2, ze2)) # Jacobian of f_gvec + return tmp2 @ tmp1 - # meshgrid evaluation - if tmp1.ndim == 5: - out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp2 @ tmp1 - - return out - - def det_df_pest(self, s2, th2, ze2): + def det_df_pest(self, s2, th2, ze2, flat_eval=False): '''The Jacobian determinant of f_pest. Parameters ---------- - s2, th2, ze2 : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s2, th2, ze2 : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation).''' - - tmp1 = self.det_dP(s2, th2, ze2) - tmp2 = self.det_df_gvec(*self.Pmap(s2, th2, ze2)) + A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation).''' + tmp1 = self.det_dP(s2, th2, ze2, flat_eval=flat_eval) + tmp2 = self.det_df_gvec(*self.Pmap(s2, th2, ze2, flat_eval=flat_eval), flat_eval=flat_eval) return tmp1 * tmp2 - def df_pest_inv(self, s2, th2, ze2): + def df_pest_inv(self, s2, th2, ze2, flat_eval=False): '''The inverse Jacobian matrix of f_pest. Parameters ---------- - s2, th2, ze2 : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s2, th2, ze2 : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - - return np.linalg.inv(self.df_pest(s2, th2, ze2)) + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + return np.linalg.inv(self.df_pest(s2, th2, ze2, flat_eval=flat_eval)) - def g_pest(self, s2, th2, ze2): + def g_pest(self, s2, th2, ze2, flat_eval=False): '''The metric tensor of f_pest. Parameters ---------- - s2, th2, ze2 : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s2, th2, ze2 : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - - tmp = np.ascontiguousarray(self.df_pest(s2, th2, ze2)) # Jacobian of f_gvec - - # meshgrid evaluation - if tmp.ndim == 5: - out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp.T @ tmp + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return out + tmp = np.ascontiguousarray(self.df_pest(s2, th2, ze2, flat_eval=flat_eval)) # Jacobian of f_gvec + tmpT = tmp.swapaxes(-1, -2) + + return tmpT @ tmp - def g_pest_inv(self, s2, th2, ze2): + def g_pest_inv(self, s2, th2, ze2, flat_eval=False): '''Inverse metric tensor of f_pest. Parameters ---------- - s2, th2, ze2 : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s2, th2, ze2 : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - - return np.linalg.inv(self.g_pest(s2, th2, ze2)) + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + return np.linalg.inv(self.g_pest(s2, th2, ze2, flat_eval=flat_eval)) # unit cube coordinates - def f_unit(self, s, u, v): + def f_unit(self, s, u, v, flat_eval=False): '''The map f_unit:(s, u, v) -> (x, y, z), where (s, u, v) in [0, 1]^3 and (x, y, z) are the Cartesian coordinates. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.f_gvec(*self.Lmap(s, u, v)) + return self.f_gvec(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) - def df_unit(self, s, u, v): + def df_unit(self, s, u, v, flat_eval=False): '''The Jacobian matrix of f_unit. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + flat_eval : bool + Whether to do flat (marker) evaluation. + Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array, - with the last two indices referring to the Jacobian entries at each point.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array, + with the last two indices referring to the Jacobian entries at each point.''' - tmp1 = self.dL(s, u, v) # Jacobian of the Lmap - tmp2 = self.df_gvec(*self.Lmap(s, u, v)) # Jacobian of f_gvec + tmp1 = self.dL(s, u, v, flat_eval=flat_eval) # Jacobian of the Lmap + tmp2 = self.df_gvec(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) # Jacobian of f_gvec - # meshgrid evaluation - if tmp1.ndim == 5: - out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp2 @ tmp1 - - return out + return tmp2 @ tmp1 - def det_df_unit(self, s, u, v): + def det_df_unit(self, s, u, v, flat_eval=False): '''The Jacobian determinant of f_unit. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation). + A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation). ''' - return self.det_df_gvec(*self.Lmap(s, u, v)) * (2*np.pi)**2*self.tor_fraction + return self.det_df_gvec(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) * (2*np.pi)**2*self.tor_fraction - def df_unit_inv(self, s, u, v): + def df_unit_inv(self, s, u, v, flat_eval=False): '''The inverse Jacobian matrix of f_unit. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.df_unit(s, u, v)) + return np.linalg.inv(self.df_unit(s, u, v, flat_eval=flat_eval)) - def g_unit(self, s, u, v): + def g_unit(self, s, u, v, flat_eval=False): '''The metric tensor of f_unit. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - - tmp = np.ascontiguousarray(self.df_unit(s, u, v)) # Jacobian of f_unit + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - # meshgrid evaluation - if tmp.ndim == 5: - out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp.T @ tmp - - return out + tmp = np.ascontiguousarray(self.df_unit(s, u, v, flat_eval=flat_eval)) # Jacobian of f_unit + tmpT = tmp.swapaxes(-1, -2) + + return tmpT @ tmp - def g_unit_inv(self, s, u, v): + def g_unit_inv(self, s, u, v, flat_eval=False): '''The inverse metric tensor of f_unit. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.g_unit(s, u, v)) + return np.linalg.inv(self.g_unit(s, u, v, flat_eval=flat_eval)) # unit cube pest coordinates - def f_unit_pest(self, s, u, v): + def f_unit_pest(self, s, u, v, flat_eval=False): '''The map f_unit_pest:(s, u, v) -> (x, y, z), where (s, u, v) in [0, 1]^3 are straight-field-line coordinates (PEST) and (x, y, z) are the Cartesian coordinates. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays (meshgrid evaluation). ''' - return self.f_pest(*self.Lmap(s, u, v)) + return self.f_pest(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) - def df_unit_pest(self, s, u, v): + def df_unit_pest(self, s, u, v, flat_eval=False): '''The Jacobian matrix of f_unit_pest. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array, - with the last two indices referring to the Jacobian entries at each point.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array, + with the last two indices referring to the Jacobian entries at each point.''' - tmp1 = self.dL(s, u, v) # Jacobian of the Lmap - tmp2 = self.df_pest(*self.Lmap(s, u, v)) # Jacobian of f_pest + tmp1 = self.dL(s, u, v, flat_eval=flat_eval) # Jacobian of the Lmap + tmp2 = self.df_pest(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) # Jacobian of f_pest - # meshgrid evaluation - if tmp1.ndim == 5: - out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp2 @ tmp1 - - return out + return tmp2 @ tmp1 - def det_df_unit_pest(self, s, u, v): + def det_df_unit_pest(self, s, u, v, flat_eval=False): '''The Jacobian determinant of f_unit_pest. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation). + A float (single-point evaluation) or a 3d numpy array (meshgrid evaluation). ''' - return self.det_df_pest(*self.Lmap(s, u, v)) * (2*np.pi)**2*self.tor_fraction + return self.det_df_pest(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) * (2*np.pi)**2*self.tor_fraction - def df_unit_pest_inv(self, s, u, v): + def df_unit_pest_inv(self, s, u, v, flat_eval=False): '''The inverse Jacobian matrix of f_unit_pest. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.df_unit_pest(s, u, v)) + return np.linalg.inv(self.df_unit_pest(s, u, v, flat_eval=flat_eval)) - def g_unit_pest(self, s, u, v): + def g_unit_pest(self, s, u, v, flat_eval=False): '''The metric tensor of f_unit_pest. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - - tmp = np.ascontiguousarray(self.df_unit_pest(s, u, v)) # Jacobian of f_unit_pest + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - # meshgrid evaluation - if tmp.ndim == 5: - out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp.T @ tmp + tmp = np.ascontiguousarray(self.df_unit_pest(s, u, v, flat_eval=flat_eval)) # Jacobian of f_unit_pest + tmpT = tmp.swapaxes(-1, -2) - return out + return tmpT @ tmp - def g_unit_pest_inv(self, s, u, v): + def g_unit_pest_inv(self, s, u, v, flat_eval=False): '''The inverse metric tensor of f_unit_pest. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.g_unit_pest(s, u, v)) + return np.linalg.inv(self.g_unit_pest(s, u, v, flat_eval=flat_eval)) # R-Z-phi coordinates - def f_wo_hmap(self, s, th, ze): + def f_wo_hmap(self, s, th, ze, flat_eval=False): '''The map f_wo_hmap:(s, th, ze) -> (q1, q2, ze), where (s, th, ze) in [0, 1] x [0, 2*pi]^2 and (q1, q2, ze) R-Z-phi coordinates. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays. + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays. ''' - return self.Xmap(s, th, ze) + return self.Xmap(s, th, ze, flat_eval=flat_eval) - def df_wo_hmap(self, s, th, ze): + def df_wo_hmap(self, s, th, ze, flat_eval=False): '''The Jacobian matrix of f_wo_hmap. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array, - with the last two indices referring to the Jacobian entries at each point. + A 2d numpy array (single-point evaluation) or a 5d numpy array, + with the last two indices referring to the Jacobian entries at each point. ''' - return self.dX(s, th, ze) + return self.dX(s, th, ze, flat_eval=flat_eval) - def det_df_wo_hmap(self, s, th, ze): + def det_df_wo_hmap(self, s, th, ze, flat_eval=False): '''The Jacobian determinant of f_wo_hmap. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A float (single-point evaluation) or a 3d numpy array. + A float (single-point evaluation) or a 3d numpy array. ''' - return self.det_dX(s, th, ze) + return self.det_dX(s, th, ze, flat_eval=flat_eval) - def df_wo_hmap_inv(self, s, th, ze): + def df_wo_hmap_inv(self, s, th, ze, flat_eval=False): '''The inverse Jacobian matrix of f_wo_hmap. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.df_wo_hmap(s, th, ze)) + return np.linalg.inv(self.df_wo_hmap(s, th, ze, flat_eval=flat_eval)) - def g_wo_hmap(self, s, th, ze): + def g_wo_hmap(self, s, th, ze, flat_eval=False): '''The metric tensor of f_wo_hmap. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - tmp = np.ascontiguousarray(self.df_wo_hmap(s, th, ze)) # Jacobian of f_wo_hmap + tmp = np.ascontiguousarray(self.df_wo_hmap(s, th, ze, flat_eval=flat_eval)) # Jacobian of f_wo_hmap + tmpT = tmp.swapaxes(-1, -2) - # meshgrid evaluation - if tmp.ndim == 5: - out = self.matmul_5d(tmp, tmp, transpose_a=True, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp.T @ tmp + return tmpT @ tmp - return out - - def g_wo_hmap_inv(self, s, th, ze): + def g_wo_hmap_inv(self, s, th, ze, flat_eval=False): '''Inverse metric tensor of f_wo_hmap. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 2d numpy array (single-point evaluation) or a 5d numpy array.''' + A 2d numpy array (single-point evaluation) or a 5d numpy array.''' - return np.linalg.inv(self.g_wo_hmap(s, th, ze)) + return np.linalg.inv(self.g_wo_hmap(s, th, ze, flat_eval=flat_eval)) # --------------- @@ -766,37 +809,40 @@ class GVEC_domain: # --------------- # Xmap - def Xmap(self, s, th, ze): + def Xmap(self, s, th, ze, flat_eval=False): '''The map (q1, q2, ze) = X(s, th, ze) computed by GVEC's minimization algorithm. Parameters ---------- - s, th, ze : float or array-like - Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, th, ze : float or array-like + Coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays.''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays.''' - s, th, ze = self.prepare_args(s, th, ze, sparse=False) + s, th, ze = self.prepare_args(s, th, ze, sparse=False, flat_eval=flat_eval) - q1 = self.X1_base.eval_stz(s, th, ze, self.X1_coef) - q2 = self.X2_base.eval_stz(s, th, ze, self.X2_coef) + q1 = self.X1_base.eval_stz(s, th, ze, self.X1_coef, flat_eval=flat_eval) + q2 = self.X2_base.eval_stz(s, th, ze, self.X2_coef, flat_eval=flat_eval) ze = ze + 0*q1 # broadcast to correct shape return q1, q2, ze - def dX(self, s, th, ze): + def dX(self, s, th, ze, flat_eval=False): '''Jacobian matrix of the Xmap.''' - s, th, ze = self.prepare_args(s, th, ze) + s, th, ze = self.prepare_args(s, th, ze, flat_eval=flat_eval) - dX1ds = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='s') - dX1dt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='th') - dX1dz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='ze') + dX1ds = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='s', flat_eval=flat_eval) + dX1dt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='th', flat_eval=flat_eval) + dX1dz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='ze', flat_eval=flat_eval) - dX2ds = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='s') - dX2dt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='th') - dX2dz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='ze') + dX2ds = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='s', flat_eval=flat_eval) + dX2dt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='th', flat_eval=flat_eval) + dX2dz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='ze', flat_eval=flat_eval) zmat = np.zeros_like(dX1ds) omat = np.ones_like(dX1ds) @@ -807,18 +853,18 @@ class GVEC_domain: return self.swap_J_axes(J) - def ddX(self, s, th, ze): + def ddX(self, s, th, ze, flat_eval=False): '''Returns a 3-tuple for the three partial derivatives of dX w.r.t s, th and ze.''' - s, th, ze = self.prepare_args(s, th, ze) + s, th, ze = self.prepare_args(s, th, ze, flat_eval=flat_eval) # partial derivative dJ/ds - ddX1_dds = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='ss') - ddX1_dsdt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='sth') - ddX1_dsdz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='sze') + ddX1_dds = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='ss', flat_eval=flat_eval) + ddX1_dsdt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='sth', flat_eval=flat_eval) + ddX1_dsdz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='sze', flat_eval=flat_eval) - ddX2_dds = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='ss') - ddX2_dsdt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='sth') - ddX2_dsdz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='sze') + ddX2_dds = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='ss', flat_eval=flat_eval) + ddX2_dsdt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='sth', flat_eval=flat_eval) + ddX2_dsdz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='sze', flat_eval=flat_eval) zmat = np.zeros_like(ddX1_dds) @@ -827,20 +873,20 @@ class GVEC_domain: ( zmat, zmat, zmat))) # partial derivative dJ/dth - ddX1_ddt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='thth') - ddX1_dtdz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='thze') + ddX1_ddt = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='thth', flat_eval=flat_eval) + ddX1_dtdz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='thze', flat_eval=flat_eval) - ddX2_ddt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='thth') - ddX2_dtdz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='thze') + ddX2_ddt = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='thth', flat_eval=flat_eval) + ddX2_dtdz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='thze', flat_eval=flat_eval) J_th = np.array(((ddX1_dsdt, ddX1_ddt, ddX1_dtdz), (ddX2_dsdt, ddX2_ddt, ddX2_dtdz), ( zmat, zmat, zmat))) # partial derivative dJ/dze - ddX1_ddz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='zeze') + ddX1_ddz = self.X1_base.eval_stz(s, th, ze, self.X1_coef, der='zeze', flat_eval=flat_eval) - ddX2_ddz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='zeze') + ddX2_ddz = self.X2_base.eval_stz(s, th, ze, self.X2_coef, der='zeze', flat_eval=flat_eval) J_ze = np.array(((ddX1_dsdz, ddX1_dtdz, ddX1_ddz), (ddX2_dsdz, ddX2_dtdz, ddX2_ddz), @@ -848,48 +894,34 @@ class GVEC_domain: return self.swap_J_axes(J_s), self.swap_J_axes(J_th), self.swap_J_axes(J_ze) - def dX_inv(self, s, th, ze): + def dX_inv(self, s, th, ze, flat_eval=False): '''Inverse Jacobian matrix of the Xmap.''' - return np.linalg.inv(self.dX(s, th, ze)) + return np.linalg.inv(self.dX(s, th, ze, flat_eval=flat_eval)) - def det_dX(self, s, th, ze): + def det_dX(self, s, th, ze, flat_eval=False): '''Jacobian determinant of the Xmap.''' - return np.linalg.det(self.dX(s, th, ze)) + return np.linalg.det(self.dX(s, th, ze, flat_eval=flat_eval)) - def dJ_X(self, s, th, ze): + def dJ_X(self, s, th, ze, flat_eval=False): '''Returns a 3-tuple for the three partial derivatives of det_dX w.r.t s, th and ze.''' - dX = self.dX(s, th, ze) + dX = self.dX(s, th, ze, flat_eval=flat_eval) + tmp = self.swap_J_back(dX) - # meshgrid evaluation - if dX.ndim == 5: - dX1_ds = dX[:, :, :, 0, 0] - dX2_ds = dX[:, :, :, 1, 0] - dX1_dt = dX[:, :, :, 0, 1] - dX2_dt = dX[:, :, :, 1, 1] - # single point evaluation - else: - dX1_ds = dX[0, 0] - dX2_ds = dX[1, 0] - dX1_dt = dX[0, 1] - dX2_dt = dX[1, 1] + dX1_ds = tmp[0, 0] + dX2_ds = tmp[1, 0] + dX1_dt = tmp[0, 1] + dX2_dt = tmp[1, 1] - ddX = self.ddX(s, th, ze) + ddX = self.ddX(s, th, ze, flat_eval=flat_eval) - tmp = [] + out = [] for n in range(3): - # meshgrid evaluation - if dX.ndim == 5: - tmp += [ddX[n][:, :, :, 0, 0] * dX2_dt + ddX[n][:, :, :, 1, 1] * dX1_ds - - ddX[n][:, :, :, 1, 0] * dX1_dt - ddX[n][:, :, :, 0, 1] * dX2_ds] - # single point evaluation - # X1ds/di * dX2_dt + dX1dt/di * dX1_ds - # -X2ds/di * dX1_dt - dX2dt/di * dX2_ds - else: - tmp += [ddX[n][0, 0] * dX2_dt + ddX[n][1, 1] * dX1_ds - - ddX[n][1, 0] * dX1_dt - ddX[n][0, 1] * dX2_ds] - - return tmp[0], tmp[1], tmp[2] + tmp2 = self.swap_J_back(ddX[n]) + out += [tmp2[0, 0] * dX2_dt + tmp2[1, 1] * dX1_ds + - tmp2[1, 0] * dX1_dt - tmp2[0, 1] * dX2_ds] + + return out[0], out[1], out[2] # hmap def hmap(self, q1, q2, ze): @@ -953,26 +985,29 @@ class GVEC_domain: return self.swap_J_axes(J) # Lmap - def Lmap(self, s, u, v): + def Lmap(self, s, u, v, flat_eval=False): '''The map (s, th, ze) = L(s, u, v) defined by th = 2*pi*u and ze = 2*pi*v/nfp with Jacobian determinant 4*pi^2/nfp. Here, the integer nfp>=1 indicates whether to use a toroidal field period in the mapping. Parameters ---------- - s, u, v : float or array-like - Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s, u, v : float or array-like + Coordinates in [0, 1]^3. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays.''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays.''' - s, u, v = self.prepare_args(s, u, v, sparse=False) + s, u, v = self.prepare_args(s, u, v, sparse=False, flat_eval=flat_eval) return s, (2*np.pi)*u, (2*np.pi)*self.tor_fraction*v - def dL(self, s, u, v): + def dL(self, s, u, v, flat_eval=False): '''Jacobian matrix of the Lmap.''' - s, u, v = self.prepare_args(s, u, v, sparse=False) + s, u, v = self.prepare_args(s, u, v, sparse=False, flat_eval=flat_eval) zmat = np.zeros_like(s) omat = np.ones_like(s) @@ -982,10 +1017,10 @@ class GVEC_domain: return self.swap_J_axes(J) - def dL_inv(self, s, u, v): + def dL_inv(self, s, u, v, flat_eval=False): '''Inverse Jacobian matrix of the Lmap.''' - s, u, v = self.prepare_args(s, u, v, sparse=False) + s, u, v = self.prepare_args(s, u, v, sparse=False, flat_eval=flat_eval) zmat = np.zeros_like(s) omat = np.ones_like(s) @@ -995,46 +1030,49 @@ class GVEC_domain: return self.swap_J_axes(J) - def det_dL(self, s, u, v): + def det_dL(self, s, u, v, flat_eval=False): '''Jacobian determinant of the Lmap.''' - s, u, v = self.prepare_args(s, u, v, sparse=False) + s, u, v = self.prepare_args(s, u, v, sparse=False, flat_eval=flat_eval) return np.ones_like(s) * (2*np.pi)**2 * self.tor_fraction # pest map - def Pmap(self, s_in, th2_in, ze_in): + def Pmap(self, s_in, th2_in, ze_in, flat_eval=False): '''The map (s, th, ze) = P(s, th2, ze) giving the gvec coordinates (s, th, ze) in [0, 1] x [0, 2*pi]^2 in terms of straight-field-line (PEST) coordinates (s, th2, ze) in [0, 1] x [0, 2*pi]^2. Parameters ---------- - s_in, th2_in, ze_in : float or array-like - PEST coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + s_in, th2_in, ze_in : float or array-like + PEST coordinates in [0, 1] x [0, 2*pi]^2. If list or np.array, must be either 1d (meshgrid evaluation) or 3d numpy arrays. + + flat_eval : bool + Whether to do flat (marker) evaluation. Returns ------- - A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays.''' + A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays.''' def find_th(s_val,th2_val,ze_val): ''' 1D root finder solving th+lambda(s,th,ze)-th2=0 for th, given s,th2,ze ''' - func = lambda th: th + self.LA_base.eval_stz(s_val, th, ze_val, self.LA_coef) - th2_val - fprime = lambda th: 1 + self.LA_base.eval_stz(s_val, th, ze_val, self.LA_coef, der='th') + func = lambda th: th + self.LA_base.eval_stz(s_val, th, ze_val, self.LA_coef, flat_eval=flat_eval) - th2_val + fprime = lambda th: 1 + self.LA_base.eval_stz(s_val, th, ze_val, self.LA_coef, der='th', flat_eval=flat_eval) return newton(func, th2_val, fprime) #th2_val as startvalue - s, th2, ze = self.prepare_args(s_in, th2_in, ze_in, sparse=True) + s, th2, ze = self.prepare_args(s_in, th2_in, ze_in, sparse=True, flat_eval=flat_eval) vfind_th=np.vectorize(find_th, otypes=[float]) th=vfind_th(s,th2,ze) return s + 0*th, th, ze + 0*th - def dP(self, s2, th2, ze2): + def dP(self, s2, th2, ze2, flat_eval=False): '''Jacobian of the Pmap.''' - s, th, ze = self.Pmap(s2, th2, ze2) + s, th, ze = self.Pmap(s2, th2, ze2, flat_eval=flat_eval) - dlambda_ds = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='s') - dlambda_dth = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='th') - dlambda_dze = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='ze') + dlambda_ds = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='s', flat_eval=flat_eval) + dlambda_dth = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='th', flat_eval=flat_eval) + dlambda_dze = self.LA_base.eval_stz(s, th, ze, self.LA_coef, der='ze', flat_eval=flat_eval) dth_ds2 = - dlambda_ds / (1. + dlambda_dth) dth_dth2 = 1. / (1. + dlambda_dth) @@ -1049,16 +1087,16 @@ class GVEC_domain: return self.swap_J_axes(J) - def dP_inv(self, s2, th2, ze2): + def dP_inv(self, s2, th2, ze2, flat_eval=False): '''Inverse Jacobian of the Pmap.''' - return np.linalg.inv(self.dP(s2, th2, ze2)) + return np.linalg.inv(self.dP(s2, th2, ze2, flat_eval=flat_eval)) - def det_dP(self, s2, th2, ze2): + def det_dP(self, s2, th2, ze2, flat_eval=False): '''Jacobian determinant of the Pmap.''' - return np.linalg.det(self.dP(s2, th2, ze2)) + return np.linalg.det(self.dP(s2, th2, ze2, flat_eval=flat_eval)) # combine Lmap and pest map (only map, Jacobian and inverse) - def PLmap(self, s, u, v): + def PLmap(self, s, u, v, flat_eval=False): '''The composition P ° L. Parameters @@ -1070,34 +1108,35 @@ class GVEC_domain: ------- A 3-tuple of either floats (single-point evaluation) or 3d numpy arrays. ''' - return self.Pmap(*self.Lmap(s, u, v)) + return self.Pmap(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) - def dPL(self, s, u, v): + def dPL(self, s, u, v, flat_eval=False): '''Jacobian matrix of the composition P ° L.''' - tmp1 = self.dL(s, u, v) # Jacobian of the Lmap - tmp2 = self.dP(*self.Lmap(s, u, v)) # Jacobian of the Pmap + tmp1 = self.dL(s, u, v, flat_eval=flat_eval) # Jacobian of the Lmap + tmp2 = self.dP(*self.Lmap(s, u, v, flat_eval=flat_eval), flat_eval=flat_eval) # Jacobian of the Pmap - # meshgrid evaluation - if tmp1.ndim == 5: - out = self.matmul_5d(tmp2, tmp1, use_pyccel=self.use_pyccel) - # single-point evaluation - else: - out = tmp2 @ tmp1 - - return out + return tmp2 @ tmp1 - def dPL_inv(self, s, u, v): + def dPL_inv(self, s, u, v, flat_eval=False): '''Inverse Jacobian matrix of the composition P ° L.''' - return np.linalg.inv(self.dPL(s, u, v)) + return np.linalg.inv(self.dPL(s, u, v, flat_eval=flat_eval)) - def det_dPL(self, s, u, v): + def det_dPL(self, s, u, v, flat_eval=False): '''Jacobian determinant of the composition P ° L.''' - return np.linalg.det(self.dPL(s, u, v)) + return np.linalg.det(self.dPL(s, u, v, flat_eval=flat_eval)) @staticmethod - def prepare_args(s, u, v, sparse=False): - '''Checks consistency, prepares arguments and performs meshgrid if necessary.''' + def prepare_args(s, u, v, sparse=False, flat_eval=False): + '''Checks consistency, prepares arguments and performs meshgrid if necessary. + + Parameters + ---------- + sparse: bool + Whether to use sparse meshgrid. + + flat_eval : bool + Whether to do flat (marker) evaluation when three 1D arguments are given. ''' if isinstance(s, (list, tuple)): s = np.array(s) @@ -1119,7 +1158,8 @@ class GVEC_domain: s.ndim, u.ndim) assert s.ndim == v.ndim, '3rd argument has different dimensions than the 1st. Expected {}, got {} instead.'.format( s.ndim, v.ndim) - if s.ndim == 1: + + if s.ndim == 1 and not flat_eval: s, u, v = np.meshgrid(s, u, v, indexing='ij', sparse=sparse) return s, u, v @@ -1169,133 +1209,44 @@ class GVEC_domain: J = np.moveaxis(J, -1, 0) J = np.moveaxis(J, -1, 0) return J - + @staticmethod - def transpose_5d(a, use_pyccel=False): - '''Performs matrix transposition in the last two indices of the 5-dim array a. - + def prepare_batch_vector(v): + """Swap axes of a batch of vectors and add singelton dimension, + such that it is compatible with numpy's matmul (@). + Parameters ---------- - a : array[float] - 5-dimensional matrix. + v : numpy.ndarray of shape (3) or (3, ...) + A batch of vectors. - use_pyccel : bool - Whether to use an accelerated kernel for the multiplication. - Returns ------- - A 5-dim C-contiguous array with the result.''' - - a_shp = a.shape - - assert isinstance(a, np.ndarray) - assert a.ndim == 5 - - out = np.ascontiguousarray(np.empty(shape=(*a_shp[:3], a_shp[4], a_shp[3]))) - if use_pyccel: - transpose_5d_kernel(np.ascontiguousarray(a), out) - else: - for i in range(a_shp[0]): - for j in range(a_shp[1]): - for k in range(a_shp[2]): - out[i, j, k] = a[i, j, k].T + numpy.ndarray of shape (3) or (..., 3, 1) + A batch of vectors. + """ - return out + if v.ndim > 1 and v.shape[0] == 3: + v = np.moveaxis(v, 0, -1) + v = v.reshape(*v.shape, 1) + return v @staticmethod - def matmul_5d(a, b, transpose_a=False, use_pyccel=False): - '''Performs matrix multiplication a*b in the last two indices of two 5-dim arrays. - If transpose_a=False and a is of shape (s0, s1, s2, l, m), then b must be of shape (s0, s1, s2, m, n). - - Parameters - ---------- - a, b : array[float] - 5-dimensional matrices to multiply. - - transpose_a : bool - Whether to transpose a. - - use_pyccel : bool - Whether to use an accelerated kernel for the multiplication. - - Returns - ------- - A 5-dim C-contiguous array with the result.''' - - a_shp = a.shape - - assert isinstance(a, np.ndarray) - assert isinstance(b, np.ndarray) - assert a.ndim == 5 - assert b.ndim == 5 - assert a.shape[:3] == b.shape[:3] + def finalize_batch_vector(v): + """Swap axes of a batch of vectors and remove singelton dimension. - if transpose_a: - tmp = GVEC_domain.transpose_5d(a, use_pyccel=use_pyccel) - else: - tmp = np.ascontiguousarray(a) - - assert tmp.shape[4] == b.shape[3] - - out = np.ascontiguousarray(np.empty(shape=(*a_shp[:3], tmp.shape[3], b.shape[4]))) - - if use_pyccel: - matmul_5d_kernel(tmp, np.ascontiguousarray(b), out) - else: - for i in range(a_shp[0]): - for j in range(a_shp[1]): - for k in range(a_shp[2]): - out[i, j, k] = tmp[i, j, k] @ b[i, j, k] - - return out - - @staticmethod - def matvec_5d(a, b, transpose_a=False, use_pyccel=False): - '''Performs matrix-vector multiplication a*b with the last two indices of the 5-dim array a. - a must have shape (s0, s1, s2, n, m) and b must have "swapped axes" with shape (m, s0, s1, s2). - Hence the last index of a contracts with the first index of b. - Parameters ---------- - a : array[float] - 5-dimensional matrix to multiply. - - b : array[float] - 4-dimensional matrix to multiply. - - transpose_a : bool - Whether to transpose a. - - use_pyccel : bool - Whether to use an accelerated kernel for the multiplication. - + v : numpy.ndarray of shape (3) or (..., 3, 1) + A batch of vectors. + Returns ------- - A 4-dim C-contiguous array with the result.''' - - a_shp = a.shape - - assert isinstance(a, np.ndarray) - assert isinstance(b, np.ndarray) - assert a.ndim == 5 - assert b.ndim == 4 - assert a.shape[:3] == b.shape[1:] - - if transpose_a: - tmp = GVEC_domain.transpose_5d(a, use_pyccel=use_pyccel) - else: - tmp = np.ascontiguousarray(a) - - assert tmp.shape[4] == b.shape[0] - - out = np.ascontiguousarray(np.empty(shape=(tmp.shape[3], *a_shp[:3]))) - - if use_pyccel: - matvec_5d_kernel(tmp, np.ascontiguousarray(b), out) - else: - for i in range(a_shp[0]): - for j in range(a_shp[1]): - for k in range(a_shp[2]): - out[:, i, j, k] = tmp[i, j, k] @ b[:, i, j, k] + numpy.ndarray of shape (3) or (..., 3, 1) + A batch of vectors. + """ - return out + if v.ndim > 1 and v.shape[-2:] == (3, 1): + v = np.squeeze(v, axis=-1) + v = np.moveaxis(v, -1, 0) + return v \ No newline at end of file diff --git a/src/gvec_to_python/geometry/kernels.py b/src/gvec_to_python/geometry/kernels.py deleted file mode 100644 index 31baca9a13fda8d60fa8ab958d25ed068675cd4e..0000000000000000000000000000000000000000 --- a/src/gvec_to_python/geometry/kernels.py +++ /dev/null @@ -1,38 +0,0 @@ -def matmul_5d_kernel(a: 'float[:,:,:,:,:]', b: 'float[:,:,:,:,:]', out: 'float[:,:,:,:,:]'): - '''Performs matrix multiplication in the last two indices of 5-dim arrays. - If a is shape shape (s0, s1, s2, l, m) then b must be shape shape (s0, s1, s2, m, n).''' - - out[:] = 0. - - for i in range(out.shape[0]): - for j in range(out.shape[1]): - for k in range(out.shape[2]): - for l in range(out.shape[3]): - for n in range(out.shape[4]): - for m in range(a.shape[4]): - out[i, j, k, l, n] += a[i, j, k, l, m] * b[i, j, k, m, n] - -def matvec_5d_kernel(a: 'float[:,:,:,:,:]', b: 'float[:,:,:,:]', out: 'float[:,:,:,:]'): - '''Performs matrix-vector multiplication in the last two indices of the 5-dim array a. - a must have shape (s0, s1, s2, n, m) and b must have "swapped axes" with shape (m, s0, s1, s2). - Hence the last index of a contracts with the first index of b.''' - - out[:] = 0. - - for i in range(out.shape[1]): - for j in range(out.shape[2]): - for k in range(out.shape[3]): - for l in range(out.shape[0]): - for m in range(b.shape[0]): - out[l, i, j, k] += a[i, j, k, l, m] * b[m, i, j, k] - -def transpose_5d_kernel(a: 'float[:,:,:,:,:]', out: 'float[:,:,:,:,:]'): - '''Matrix transpose in the last two indices of a 5-dim array. - If a is shape shape (., ., ., m, n) then out must be shape shape (., ., ., n, m).''' - - for i in range(out.shape[0]): - for j in range(out.shape[1]): - for k in range(out.shape[2]): - for n in range(out.shape[3]): - for m in range(out.shape[4]): - out[i, j, k, n, m] = a[i, j, k, m, n] \ No newline at end of file diff --git a/tests/test_gvec_domain.py b/tests/test_gvec_domain.py index c22582dc44356c10511b4480668f2d379718edcd..993edd4d91296e7f19c248c821fbd6d44684b1cd 100644 --- a/tests/test_gvec_domain.py +++ b/tests/test_gvec_domain.py @@ -79,6 +79,26 @@ def test_output_formats(mapping, use_pyccel,gvec_file): assert gvec.df_inv(s, a1, a2).shape == shp5 assert gvec.g(s, a1, a2).shape == shp5 assert gvec.g_inv(s, a1, a2).shape == shp5 + + # test flat (marker) evaluation + Np = 17 + s = np.linspace(.01, 1, Np) # mappings are singular at s=0 + if 'unit' not in mapping: + a1 = np.linspace(0, 2*np.pi, Np, endpoint=False) + a2 = np.linspace(0, 2*np.pi, Np, endpoint=False) + else: + a1 = np.linspace(0, 1, Np, endpoint=False) + a2 = np.linspace(0, 1, Np, endpoint=False) + + shp3 = s.shape + shp5 = (s.size, 3, 3) + + assert gvec.f(s, a1, a2, flat_eval=True)[0].shape == shp3 + assert gvec.df(s, a1, a2, flat_eval=True).shape == shp5 + assert gvec.det_df(s, a1, a2, flat_eval=True).shape == shp3 + assert gvec.df_inv(s, a1, a2, flat_eval=True).shape == shp5 + assert gvec.g(s, a1, a2, flat_eval=True).shape == shp5 + assert gvec.g_inv(s, a1, a2, flat_eval=True).shape == shp5 @pytest.mark.parametrize('use_pyccel,gvec_file', [(True ,"testcases/ellipstell/newBC_E4D6_M6N6/GVEC_ELLIPSTELL_E4D6_M6N6_State_0001_00200000")]) @@ -156,6 +176,6 @@ def test_values(use_pyccel,gvec_file,unit_tor_domain): if __name__ == '__main__': - test_values(False ,"testcases/ellipstell/newBC_E1D6_M6N6/GVEC_ELLIPSTELL_E1D6_M6N6_State_0000_00200000","full") + #test_values(False ,"testcases/ellipstell/newBC_E1D6_M6N6/GVEC_ELLIPSTELL_E1D6_M6N6_State_0000_00200000","full") test_output_formats('pest', True ,"testcases/ellipstell/newBC_E4D6_M6N6/GVEC_ELLIPSTELL_E4D6_M6N6_State_0001_00200000") test_output_formats('unit', True ,"testcases/ellipstell/newBC_E4D6_M6N6/GVEC_ELLIPSTELL_E4D6_M6N6_State_0001_00200000") diff --git a/tests/test_gvec_equil.py b/tests/test_gvec_equil.py index b112da05cb94af8f3dbc58cecfff75104f88d1f7..5eb3abeb737a465810268447653dae641a9e1f1d 100644 --- a/tests/test_gvec_equil.py +++ b/tests/test_gvec_equil.py @@ -133,6 +133,45 @@ def test_mhd_fields(mapping, use_pyccel,gvec_file): assert np.all([tmp[j].shape == shp3 for j in range(3)]) tmp = gvec.j_cart(s, a1, a2) assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)]) + + # test flat (marker) evaluation + Np = 17 + s = np.linspace(.01, 1, Np) # mappings are singular at s=0 + if 'unit' not in mapping: + a1 = np.linspace(0, 2*np.pi, Np, endpoint=False) + a2 = np.linspace(0, 2*np.pi, Np, endpoint=False) + else: + a1 = np.linspace(0, 1, Np, endpoint=False) + a2 = np.linspace(0, 1, Np, endpoint=False) + + shp3 = s.shape + + assert gvec.p0(s, a1, a2, flat_eval=True).shape == shp3 + assert gvec.p3(s, a1, a2, flat_eval=True).shape == shp3 + tmp = gvec.bv(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.b1(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.b2(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.b_cart(s, a1, a2, flat_eval=True) + assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)]) + tmp = gvec.av(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.a1(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.a2(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.a_cart(s, a1, a2, flat_eval=True) + assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)]) + tmp = gvec.jv(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.j1(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.j2(s, a1, a2, flat_eval=True) + assert np.all([tmp[j].shape == shp3 for j in range(3)]) + tmp = gvec.j_cart(s, a1, a2, flat_eval=True) + assert np.all([tmp[i][j].shape == shp3 for j in range(3) for i in range(2)]) @pytest.mark.parametrize('use_pyccel,gvec_file', [(True ,"testcases/ellipstell/newBC_E4D6_M6N6/GVEC_ELLIPSTELL_E4D6_M6N6_State_0001_00200000")]) @@ -294,5 +333,5 @@ def test_derivs_with_FD(mapping,use_pyccel,gvec_file): if __name__ == '__main__': - test_mhd_fields('gvec', use_pyccel=False,gvec_file='testcases/ellipstell/newBC_E1D6_M6N6/GVEC_ELLIPSTELL_E1D6_M6N6_State_0000_00200000') + #test_mhd_fields('gvec', use_pyccel=False,gvec_file='testcases/ellipstell/newBC_E1D6_M6N6/GVEC_ELLIPSTELL_E1D6_M6N6_State_0000_00200000') test_mhd_fields('pest', use_pyccel=True ,gvec_file='testcases/ellipstell/newBC_E4D6_M6N6/GVEC_ELLIPSTELL_E4D6_M6N6_State_0001_00200000')