Merge pull request #22 from bendudson/jtor-improvements

Jtor improvements

freegs/jtor.py

@@ -191,7 +191,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None):
         psi_norm = (psi - psi_axis)  / (psi_bndry - psi_axis)
 
         # Current profile shape
-        jtorshape = (1. - psi_norm**self.alpha_m)**self.alpha_n
+        jtorshape = (1. - np.clip(psi_norm, 0.0, 1.0)**self.alpha_m)**self.alpha_n
 
         if mask is not None:
             # If there is a masking function (X-points, limiters)
@@ -256,16 +256,18 @@ def pshape(psinorm):
     # Profile functions
     def pprime(self, pn):
         """
-        dp/dpsi as a function of normalised psi
+        dp/dpsi as a function of normalised psi. 0 outside core
+        Calculate pprimeshape inside the core only
         """
-        shape = (1. - pn**self.alpha_m)**self.alpha_n
+        shape = (1. - np.clip(pn, 0.0, 1.0)**self.alpha_m)**self.alpha_n
         return self.L * self.Beta0/self.Raxis * shape
 
     def ffprime(self, pn):
         """
-        f * df/dpsi as a function of normalised psi
+        f * df/dpsi as a function of normalised psi. 0 outside core.
+        Calculate ffprimeshape inside the core only.
         """
-        shape = (1. - pn**self.alpha_m)**self.alpha_n
+        shape = (1. - np.clip(pn, 0.0, 1.0)**self.alpha_m)**self.alpha_n
         return mu0 * self.L * (1-self.Beta0)*self.Raxis * shape
 
         return Jtor, pprime, ffprime
@@ -339,7 +341,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None):
         psi_norm = (psi - psi_axis)  / (psi_bndry - psi_axis)
 
         # Current profile shape
-        jtorshape = (1. - psi_norm**self.alpha_m)**self.alpha_n
+        jtorshape = (1. - np.clip(psi_norm, 0.0, 1.0)**self.alpha_m)**self.alpha_n
 
         if mask is not None:
             # If there is a masking function (X-points, limiters)
@@ -387,16 +389,18 @@ def Jtor(self, R, Z, psi, psi_bndry=None):
     # Profile functions
     def pprime(self, pn):
         """
-        dp/dpsi as a function of normalised psi
+        dp/dpsi as a function of normalised psi. 0 outside core
+        Calculate pprimeshape inside the core only
         """
-        shape = (1. - pn**self.alpha_m)**self.alpha_n
+        shape = (1. - np.clip(pn, 0.0, 1.0)**self.alpha_m)**self.alpha_n
         return self.L * self.Beta0/self.Raxis * shape
 
     def ffprime(self, pn):
         """
-        f * df/dpsi as a function of normalised psi
+        f * df/dpsi as a function of normalised psi. 0 outside core.
+        Calculate ffprimeshape inside the core only.
         """
-        shape = (1. - pn**self.alpha_m)**self.alpha_n
+        shape = (1. - np.clip(pn, 0.0, 1.0)**self.alpha_m)**self.alpha_n
         return mu0 * self.L * (1-self.Beta0)*self.Raxis * shape
 
         return Jtor, pprime, ffprime

freegs/test_jtor.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+from . import jtor
+
+def test_psinorm_range():
+    """Test that the profiles produce finite values outside core"""
+
+    for profiles in [jtor.ConstrainPaxisIp(1e3, 2e5, 2.0),
+                     jtor.ConstrainBetapIp(1.0, 2e5, 2.0)]:
+
+        # Need to give a plasma psi
+        R, Z = np.meshgrid(np.linspace(0.5,1.5,33), np.linspace(-1,1,33), indexing='ij')
+        psi = np.exp((-(R - 1.0)**2 - Z**2)*3) + np.exp((-(R - 1.0)**2 - (Z + 1)**2)*3)
+
+        current_density = profiles.Jtor(R, Z, psi)
+        assert np.all(np.isfinite(current_density))
+
+        assert profiles.pprime(1.0) == 0.0
+        assert profiles.pprime(1.1) == 0.0
+        assert np.isfinite(profiles.pprime(-0.32))
+
+        assert profiles.ffprime(1.0) == 0.0
+        assert profiles.ffprime(1.1) == 0.0
+        assert np.isfinite(profiles.ffprime(-0.32))