Esempio n. 1
0
def process_grid( rz_grid, mesh, output=None, poorquality=None, 
                  gui=None, parent=None, reverse_bt=None,  
                  curv=None, smoothpressure=None,  
                  smoothhthe=None, smoothcurv=None,  
                  settings=None):
  
    if settings==None :
        # Create an empty structure
        settings = Bunch(dummy=0)
     
    # Check settings
        settings.calcp= -1
        settings.calcbt= -1
        settings.calchthe= -1
        settings.calcjpar= -1
  
   # ;CATCH, err
   # ;IF err NE 0 THEN BEGIN
   # ;  PRINT, "PROCESS_GRID failed"
  	#;  PRINT, "   Error message: "+!ERROR_STATE.MSG
   # ;  CATCH, /cancel
   # ;  RETURN
   # ;ENDIF

    MU = 4.e-7*numpy.pi

    poorquality = 0

    if output==None : output="bout.grd.nc"
  
    # Size of the mesh
    nx = numpy.int(numpy.sum(mesh.nrad))
    ny = numpy.int(numpy.sum(mesh.npol))

    # Find the midplane
    ymid = 0
    status = gen_surface(mesh=mesh) # Start generator
     
    while True:
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        if period :
            rm = numpy.max(mesh.Rxy[xi,yi])
            ymidindx = numpy.argmax(mesh.Rxy[xi,yi])
            ymid = yi[ymidindx]
            break
         
        if last==1: break
  

    Rxy = numpy.asarray(mesh.Rxy)
    Zxy = numpy.asarray(mesh.Zxy)
    psixy = mesh.psixy*mesh.fnorm + mesh.faxis # Non-normalised psi

    pressure = numpy.zeros((nx, ny))
    
    
 
    # Use splines to interpolate pressure profile
    status = gen_surface(mesh=mesh) # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
        if period :
            # Pressure only given on core surfaces
           # pressure[xi,yi] = SPLINE(rz_grid.npsigrid, rz_grid.pres, mesh.psixy[xi,yi[0]], /double)
            sol=interpolate.UnivariateSpline(rz_grid.npsigrid, rz_grid.pres,s=1)
            pressure[xi,yi] =sol(mesh.psixy[xi,yi[0]])

        else:

            pressure[xi,yi] = rz_grid.pres[numpy.size(rz_grid.pres)-1]

        if last==1 : break
        
  
    # Add a minimum amount
    if numpy.min(pressure) < 1.0e-2*numpy.max(pressure) :
        print("****Minimum pressure is very small:", numpy.min(pressure))
        print("****Setting minimum pressure to 1% of maximum")
        pressure = pressure + 1e-2*numpy.max(pressure)
         
  
    if smoothpressure != None :
        p0 = pressure[:,ymid] # Keep initial pressure for comparison
        while True :
            #!P.multi=[0,0,2,0,0]
            fig=figure()
            plot( p0, xtitle="X index", ytitle="pressure at y="+numpy.strip(numpy.str(ymid),2)+" dashed=original", color=1, lines=1)
            plot( pressure[:,ymid], color=1)
            plot( deriv(p0), xtitle="X index", ytitle="DERIV(pressure)", color=1, lines=1)
            plot( deriv(pressure[:,ymid]), color=1 )
            sm = query_yes_no("Smooth pressure profile?")#, gui=gui, dialog_parent=parent)
            if sm :
                # Smooth the pressure profile
        
                p2 = pressure
                for i in range (6) :
                    status = gen_surface(mesh=mesh) # Start generator
                    while True :
                        # Get the next domain
                        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
            
                        if (xi > 0) and (xi < (nx-1)) :
                            for j in range (numpy.size(yi)) :
                                p2[xi,yi[j]] = ( 0.5*pressure[xi,yi[j]] +  
                                                0.25*(pressure[xi-1,yi[j]] + pressure[xi+1,yi[j]]) 
                                                )
                             
                         
            
                        # Make sure it's still constant on flux surfaces
                        p2[xi,yi] = numpy.mean(p2[xi,yi])
                        if last != None : break
                    pressure = p2
                 
             
            if sm == 0 : break
     

    if numpy.min(pressure) < 0.0 :
        print("")
        print("============= WARNING ==============")
        print("Poor quality equilibrium: Pressure is negative")
        print("")
        poorquality = 1
     
  
    dpdpsi = DDX(psixy, pressure)
    
    
    #;IF MAX(dpdpsi)*mesh.fnorm GT 0.0 THEN BEGIN
    #;  PRINT, ""
    #;  PRINT, "============= WARNING =============="
    #;  PRINT, "Poor quality equilibrium: Pressure is increasing radially"
    #;  PRINT, ""
    #;  poorquality = 1
    #;ENDIF

    # Grid spacing
    dx = numpy.zeros((nx, ny))
    for y in range (ny) : 
        dx[0:(nx-1),y] = psixy[1::,y] - psixy[0:(nx-1),y]
        dx[nx-1,y] = dx[nx-2,y]
     
  
    # Sign
    bpsign = 1.
    xcoord = psixy
    if numpy.min(dx) < 0. :
        bpsign = -1.
        dx = -dx # dx always positive
        xcoord = -xcoord
     

    dtheta = 2.*numpy.pi / numpy.float(ny)
    dy = numpy.zeros((nx, ny)) + dtheta
     
    
    # B field components
    # Following signs mean that psi increasing outwards from
    # core to edge results in Bp clockwise in the poloidal plane
    # i.e. in the positive Grad Theta direction.
  
    Brxy = old_div(mesh.dpsidZ, Rxy)
    Bzxy = old_div(-mesh.dpsidR, Rxy)
    Bpxy = numpy.sqrt(Brxy**2 + Bzxy**2)
    

    # Determine direction (dot B with grad y vector)
  
    dot = ( Brxy[0,ymid]*(Rxy[0,ymid+1] - Rxy[0,ymid-1]) + 
            Bzxy[0,ymid]*(Zxy[0,ymid+1] - Zxy[0,ymid-1])
            ) 
  
    if dot < 0. :
        print("**** Poloidal field is in opposite direction to Grad Theta -> Bp negative")
        Bpxy = -Bpxy
        if bpsign > 0 : sys.exit() # Should be negative
        bpsign = -1.0
    else:
        if bpsign < 0 : sys.exit() # Should be positive
        bpsign = 1.
     

  # Get toroidal field from poloidal current function fpol
    Btxy = numpy.zeros((nx, ny))
    fprime = numpy.zeros((nx, ny))
    fp = deriv(rz_grid.npsigrid*(rz_grid.sibdry - rz_grid.simagx), rz_grid.fpol)
    
    
    status = gen_surface(mesh=mesh) # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=period, xi=xi)

        if period :
            # In the core
            #fpol = numpy.interp(rz_grid.fpol, rz_grid.npsigrid, mesh.psixy[xi,yi], /spline)
                        
            sol=interpolate.UnivariateSpline(rz_grid.npsigrid, rz_grid.fpol,s=1)
         #  fpol = SPLINE(rz_grid.npsigrid, rz_grid.fpol, mesh.psixy[xi,yi[0]], 'double')
            fpol = sol(mesh.psixy[xi,yi[0]])
            
            sol=interpolate.UnivariateSpline(rz_grid.npsigrid, fp ,s=1)
           # fprime[xi,yi] = SPLINE(rz_grid.npsigrid, fp, mesh.psixy[xi,yi[0]], 'double')
            fprime[xi,yi] = sol(mesh.psixy[xi,yi[0]])
            
        else:
            # Outside core. Could be PF or SOL
            fpol = rz_grid.fpol[numpy.size(rz_grid.fpol)-1]
            fprime[xi,yi] = 0.
         
        Btxy[xi,yi] = old_div(fpol, Rxy[xi,yi])
        
        if last ==1 : break
  
    # Total B field
    Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)
    

  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Go through the domains to get a starting estimate
  # of hthe
    hthe = numpy.zeros((nx, ny))

  #   Pick a midplane index
    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
    
        if period :
      # In the core
            rmax = numpy.argmax(Rxy[xi,yi])
            ymidplane = yi[rmax]
            break
     
        if last == 1: break

    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
    
        n = numpy.size(yi)
    
        # Get distance along this line
    
        if period :
             
            # Periodic, so can use FFT
            #drdi = REAL_PART(fft_deriv(Rxy[xi, yi]))
            #dzdi = REAL_PART(fft_deriv(Zxy[xi, yi]))
            line=numpy.append(Rxy[xi,yi[n-1::]], Rxy[xi,yi])  
            line=numpy.append(line,Rxy[xi,yi[0:1]])
           
            drdi = deriv(line)[1:n+1]
            
            line=numpy.append(Zxy[xi,yi[n-1::]], Zxy[xi,yi])  
            line=numpy.append(line,Zxy[xi,yi[0:1]])
                              
            dzdi = deriv(line)[1:n+1]
        else:
        # Non-periodic
            drdi = numpy.gradient(Rxy[xi, yi])
            dzdi = numpy.gradient(Zxy[xi, yi])
     
    
        dldi = numpy.sqrt(drdi**2 + dzdi**2)
    
            
        if 0 :

        # Need to smooth to get sensible results
            if period :
                n = numpy.size(dldi)
                line=numpy.append(dldi[(n-2)::], dldi) # once
                line=numpy.append(line,dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n+4)]
                
                line=numpy.append(dldi[(n-2)::], dldi) #twice
                line=numpy.append(line,dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n+4)]
                
                line=numpy.append(dldi[(n-2)::], dldi) # three
                line=numpy.append(line,dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n+4)]
                
            else:
                line = dldi
                dldi = SMOOTH(line, 5)[2:n+2]
                line = dldi
                dldi = SMOOTH(line, 5)[2:n+2]
                line = dldi
                dldi = SMOOTH(dldi, 5)[2:n+2]
        
    
        hthe[xi, yi] = old_div(dldi, dtheta) # First estimate of hthe
    
        # Get outboard midplane
        if period and xi == 0 :
            m = numpy.argmax(Rxy[0,yi])
            ymidplane = yi[m]
         
        if last == 1 : break

    print("Midplane index ", ymidplane)

    fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pressure)
    print("Force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))

    

    
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Correct pressure using hthe
  
    print("Calculating pressure profile from force balance")

    try:

    # Calculate force balance
        dpdx = old_div(( -Bpxy*DDX(xcoord, Bpxy * hthe) - Btxy*hthe*DDX(xcoord, Btxy) - (Btxy*Btxy*hthe/Rxy)*DDX(xcoord, Rxy) ), (MU*hthe))
    
        # Surface average
        dpdx2 = surface_average(dpdx, mesh)
        
        pres = numpy.zeros((nx, ny))
        # Integrate to get pressure
        for i in range (ny) :
            pres[:,i] = int_func(psixy[:,i], dpdx2[:,i])
            pres[:,i] = pres[:,i] - pres[nx-1,i]
         
        
       
        status = gen_surface(mesh=mesh) # Start generator
        while True:
      # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
      
            ma = numpy.max(pres[xi,yi])
           
            for i in range (numpy.size(yi)) :
                pres[:,yi[i]] = pres[:,yi[i]] - pres[xi,yi[i]] + ma
             
            if last == 1 : break
        
    
        pres = pres - numpy.min(pres)
  
    # Some sort of smoothing here?
  
  
        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pres)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))
  
  
       #!P.MULTI=[0,0,2,0,0]
        fig=figure(figsize=(7, 11))
        subplots_adjust(left=.07, bottom=.07, right=0.95, top=0.95,
                wspace=.3, hspace=.25)
        
        SURFACE( pressure, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2,1,1])
        title("Input pressure")
        SURFACE( pres, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2,1,2])
        title("New pressure")
  #  arrange the plot on the screen      
  #      mngr = get_current_fig_manager()
  #      geom = mngr.window.geometry()
  #      x,y,dx,dy = geom.getRect()
  #      mngr.window.setGeometry(0, 0, dx, dy)
  #
        show(block=False)
        
  
        calcp = settings.calcp
    
        if calcp == -1 :
            calcp = query_yes_no("Keep new pressure?")#, gui=gui, dialog_parent=parent)
        else: time.sleep( 2 )
        if calcp == 1 :
            pressure = pres
            dpdpsi = dpdx2
            
         
    except Exception:
        print("WARNING: Pressure profile calculation failed: ")#, !ERROR_STATE.MSG 
        pass

    #CATCH, /cancel
  
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Correct f = RBt using force balance

    calcbt = settings.calcbt
    if calcbt == -1 : calcbt = query_yes_no("Correct f=RBt using force balance?")#, gui=gui, dialog_parent=parent)
    if calcbt == 1 :

        new_Btxy = newton_Bt(psixy, Rxy, Btxy, Bpxy, pres, hthe, mesh)
    
        fb0 = force_balance(psixy, Rxy, Bpxy, new_Btxy, hthe, pressure)
        print("force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))
    
    
        fig=figure(figsize=(7, 11))
        subplots_adjust(left=.07, bottom=.07, right=0.95, top=0.95,
                wspace=.3, hspace=.25)
        
        subplot(211)
        SURFACE( Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2,1,1])
        title("Input Bt")
        subplot(212)
        SURFACE( new_Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2,1,2])
        title("New Bt")
          #  arrange the plot on the screen      
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(600, 0, dx, dy)


        show(block=False)

        calcbt = settings.calcbt
        if calcbt == -1 : calcbt = query_yes_no("Keep new Bt?")#, gui=gui, dialog_parent=parent)
        if calcbt == 1 :
            Btxy = new_Btxy
            Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)
    
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # CALCULATE HTHE
  # Modify hthe to fit force balance using initial guess
  # Does not depend on signs
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  
    calchthe = settings.calchthe
    if calchthe == -1 : calchthe = query_yes_no("Adjust hthe using force balance?")#, gui=gui, dialog_parent=parent) 
    if calchthe == 1 :
        # This doesn't behave well close to the x-points
        fixhthe = numpy.int(old_div(nx, 2))
        nh = correct_hthe(Rxy, psixy, Btxy, Bpxy, hthe, pressure, fixhthe=fixhthe)
    
        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, nh, pressure)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)), numpy.max(numpy.abs(fb0)))
    
        print("numpy.maximum difference in hthe: ", numpy.max(numpy.abs(hthe - nh)))
        print("numpy.maximum percentage difference: ", 100.*numpy.max(numpy.abs(old_div((hthe - nh),hthe))))

       #!P.multi=[0,0,1,0,0]
        fig=figure(figsize=(7, 4))
        title("Poloidal arc length at midplane. line is initial estimate")
        plot( hthe[:,0], '-' )
        plot( nh[:,0], 'r-+' )
                  #  arrange the plot on the screen      
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(0, 1150, dx, dy)

        show(block=False)

        if query_yes_no("Keep new hthe?") == 1:#, gui=gui, dialog_parent=parent) :
            hthe = nh
             
   
  
    if smoothhthe != None :
    # Smooth hthe to prevent large jumps in X or Y. This
    # should be done by creating a better mesh in the first place
    
    # Need to smooth in Y and X otherwise smoothing in X
    # produces discontinuities in Y
        hold = hthe
    
        if 1 :
      # Nonlinear smoothing. Tries to smooth only regions with large
      # changes in gradient
      
            hthe =0.# smooth_nl(hthe, mesh)
      
        else:
      # Just use smooth in both directions
      
            for i in range (ny) :
                hthe[:,i] = SMOOTH(SMOOTH(hthe[:,i],10),10)
         
      
        status = gen_surface(mesh=mesh) # Start generator
        while True:
        # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        
            n = numpy.size(yi)
        
            if period :
                hthe[xi,yi] = (SMOOTH([hthe[xi,yi[(n-4):(n-1)]], hthe[xi,yi], hthe[xi,yi[0:3]]], 4))[4:(n+3)]
            else:
                hthe[xi,yi] = SMOOTH(hthe[xi,yi], 4)
           
            if last == 1: break
            
    

    # Calculate field-line pitch
    pitch = hthe * Btxy / (Bpxy * Rxy)
  
   
    # derivative with psi
    dqdpsi = DDX(psixy, pitch)
      

    qinty, qloop = int_y(pitch, mesh, loop=0, nosmooth='nosmooth', simple='simple') 
    qinty = qinty * dtheta
    qloop = qloop * dtheta
    
    
    sinty = int_y(dqdpsi, mesh, nosmooth='nosmooth', simple='simple') * dtheta
  

  
    # NOTE: This is only valid in the core
    pol_angle = numpy.zeros((nx,ny))
    for i in range (nx) :  pol_angle[i, :] = 2.0*numpy.pi * qinty[i,:] / qloop[i]
    
    
  #;;;;;;;;;;;;;;;;;;; THETA_ZERO ;;;;;;;;;;;;;;;;;;;;;;
  # re-set zshift to be zero at the outboard midplane
  
    print("MIDPLANE INDEX = ", ymidplane)
  
    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
    
        w = numpy.size(numpy.where(yi == ymidplane))
        if w > 0 :
      # Crosses the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi, ymidplane]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi, ymidplane]
        else:
      # Doesn't include a point at the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi,yi[0]]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi,yi[0]]
     
        if last ==1 : break
  
    print("")
    print("==== Calculating curvature ====")
  
  #;;;;;;;;;;;;;;;;;;; CURVATURE ;;;;;;;;;;;;;;;;;;;;;;;
  # Calculating b x kappa
  
    if curv == None :
    
        print("*** Calculating curvature in toroidal coordinates")
    
        thetaxy = numpy.zeros((nx, ny))
        status = gen_surface(mesh=mesh) # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
            thetaxy[xi,yi] = numpy.arange(numpy.size(yi)).astype(float)*dtheta
            if last ==1 : break
    
        
        bxcv = curvature( nx, ny, Rxy,Zxy, Brxy, Bzxy, Btxy,  
                    psixy, thetaxy, hthe,  
                     mesh=mesh)
                             
        bxcvx = bpsign*bxcv.psi 
        bxcvy= bxcv.theta
        bxcvz = bpsign*(bxcv.phi - sinty*bxcv.psi - pitch*bxcv.theta)
        
       
        # x borders
        bxcvx[0,:] = bxcvx[1,:]
        bxcvx[nx-1,:] = bxcvx[nx-2,:]
    
        bxcvy[0,:] = bxcvy[1,:]
        bxcvy[nx-1,:] = bxcvy[nx-2,:]
    
        bxcvz[0,:] = bxcvz[1,:]
        bxcvz[nx-1,:] = bxcvz[nx-2,:]

    elif curv == 1 :
        # Calculate on R-Z mesh and then interpolate onto grid
        # ( cylindrical coordinates)

        print("*** Calculating curvature in cylindrical coordinates")
    
        bxcv = rz_curvature(rz_grid)
    
        # DCT methods cause spurious oscillations
        # Linear interpolation seems to be more robust
        bxcv_psi = numpy.interp(bxcv.psi, mesh.Rixy, mesh.Zixy)
        bxcv_theta = old_div(numpy.interp(bxcv.theta, mesh.Rixy, mesh.Zixy), hthe)
        bxcv_phi = numpy.interp(bxcv.phi, mesh.Rixy, mesh.Zixy)
    
        # If Bp is reversed, then Grad x = - Grad psi
        bxcvx = bpsign*bxcv_psi
        bxcvy = bxcv_theta
        bxcvz = bpsign*(bxcv_phi - sinty*bxcv_psi - pitch*bxcv_theta)
    elif curv == 2 :
        # Curvature from Curl(b/B)
    
        bxcvx = bpsign*(Bpxy * Btxy*Rxy * DDY(old_div(1., Bxy), mesh) / hthe)
        bxcvy = -bpsign*Bxy*Bpxy * DDX(xcoord, Btxy*Rxy/Bxy^2) / (2.*hthe)
        bxcvz = Bpxy^3 * DDX(xcoord, old_div(hthe,Bpxy)) / (2.*hthe*Bxy) - Btxy*Rxy*DDX(xcoord, old_div(Btxy,Rxy)) / (2.*Bxy) - sinty*bxcvx
    
    else:
        # calculate in flux coordinates.
    
        print("*** Calculating curvature in flux coordinates")
    
        dpb = numpy.zeros((nx, ny))      # quantity used for y and z components
    
        for i in range (ny) :
            dpb[:,i] = MU*dpdpsi/Bxy[:,i]
         
        dpb = dpb + DDX(xcoord, Bxy)

        bxcvx = bpsign*(Bpxy * Btxy*Rxy * DDY(old_div(1., Bxy), mesh) / hthe)
        bxcvy = bpsign*(Bpxy*Btxy*Rxy*dpb / (hthe*Bxy^2))
        bxcvz = -dpb - sinty*bxcvx
     
  

    if smoothcurv:
        # Smooth curvature to prevent large jumps
    
        # Nonlinear smoothing. Tries to smooth only regions with large
        # changes in gradient
    
        bz = bxcvz + sinty * bxcvx
    
        print("Smoothing bxcvx...")
        bxcvx = 0.#smooth_nl(bxcvx, mesh)
        print("Smoothing bxcvy...")
        bxcvy = 0.#smooth_nl(bxcvy, mesh)
        print("Smoothing bxcvz...")
        bz = 0.#smooth_nl(bz, mesh)
    
        bxcvz = bz - sinty * bxcvx
   

  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # CALCULATE PARALLEL CURRENT
  # 
  # Three ways to calculate Jpar0:
  # 1. From fprime and pprime
  # 2. From Curl(B) in field-aligned coords
  # 3. From the curvature
  # 
  # Provides a way to check if Btor should be reversed
  #
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  
    print("")
    print("==== Calculating parallel current ====")
    
    jpar0 = - Bxy * fprime / MU - Rxy*Btxy * dpdpsi / Bxy
     
  
    # Set to zero in PF and SOL
    status = gen_surface(mesh=mesh) # Start generator
    while True:
    # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
    
        if period == None : jpar0[xi,yi] = 0.0
        if last == 1 : break
  
  # Curl(B) expression for Jpar0 (very noisy usually)
    j0 = ( bpsign*((Bpxy*Btxy*Rxy/(Bxy*hthe))*( DDX(xcoord, Bxy**2*hthe/Bpxy) - bpsign*Btxy*Rxy*DDX(xcoord,Btxy*hthe/(Rxy*Bpxy)) ) 
        - Bxy*DDX(xcoord, Btxy*Rxy)) / MU )
  

  
  # Create a temporary mesh structure to send to adjust_jpar
    tmp_mesh = Bunch(mesh,  
                           bxcvx=bxcvx, bxcvy=bxcvy,  bxcvz=bxcvz,  
                            Bpxy=Bpxy,  Btxy=Btxy,  Bxy=Bxy,  
                            dx=dx,  dy=dy,  
                            hthe=hthe,  jpar0=jpar0,  pressure=pressure)
    tmp_mesh.psixy = psixy
  
    jpar = adjust_jpar( tmp_mesh, noplot='noplot')
    

   #!P.multi=[0,2,2,0,0]
   
    fig=figure(figsize=(15, 11))
    subplots_adjust(left=.07, bottom=.07, right=0.95, top=0.95,
                wspace=.3, hspace=.25)

    subplot(221)
    SURFACE( jpar0, fig, xtitle="X", ytitle="Y", var='A', sub=[2,2,1])
    title("Jpar from F' and P'")
    
    subplot(222)
    SURFACE( jpar, fig, xtitle="X", ytitle="Y", var='A', sub=[2,2,2])
    title("Jpar from curvature")
 
    subplot(223)
    plot( jpar0[0,:],'-', jpar[0,:] ,'+' )
    ylim([numpy.min([jpar0[0,:],jpar[0,:]]), numpy.max([jpar0[0,:],jpar[0,:]])])
    title("jpar at x=0. Solid from f' and p'")
    
    subplot(224)
    plot(jpar0[:,ymidplane],'-' , jpar[:,ymidplane] , '+' )
    ylim([numpy.min([jpar0[:,ymidplane],jpar[:,ymidplane]]),numpy.max([jpar0[:,ymidplane],jpar[:,ymidplane]])])
        
    title("Jpar at y="+numpy.str(ymidplane)+" Solid from f' and p'")
    
        #  arrange the plot on the screen      
    #mngr = get_current_fig_manager()
    #geom = mngr.window.geometry()
    #x,y,dx,dy = geom.getRect()
    #mngr.window.setGeometry(1350, 0, dx, dy)

    
    show(block=False)
  
 # !P.multi=0
  
    calcjpar = settings.calcjpar
    if calcjpar == -1 : calcjpar = query_yes_no("Use Jpar from curvature?")#, gui=gui, dialog_parent=parent)
    if calcjpar == True :
        jpar0 = jpar
   
  
    if 0 :
    
    # Try smoothing jpar0 in psi, preserving zero points and maxima
        jps = jpar0
        for y in range ( ny ):
            j = jpar0[:,y]
            js = j
            ma = numpy.max(numpy.abs(j))
            ip = numpy.argmax(numpy.abs(j))
            if (ma < 1.e-4) or (ip == 0) :
                jps[:,y] = j
         
            level = 1.
            #i0 = MAX(WHERE(ABS(j[0:ip]) LT level))
            i1 = numpy.min(numpy.where(numpy.abs(j[ip::]) < level))
      
            #IF i0 LE 0 THEN i0 = 1
            i0 = 1
      
            if i1 == -1 :
                i1 = nx-2 
            else: 
                i1 = i1 + ip
      
            if (ip <= i0) or (ip >= i1) :
      
      # Now preserve starting and end points, and peak value
                div = numpy.int(old_div((i1-i0),10))+1 # reduce number of points by this factor
      
                inds = [i0] # first point
                for i in [i0+div, ip-div, div] : inds = [inds, i]
                inds = [inds, ip] # Put in the peak point
      
        # Calculate spline interpolation of inner part
        
                js[0:ip] = spline_mono(inds, j[inds], numpy.arange(ip+1),
                             yp0=(j[i0] - j[i0-1]), ypn_1=0.0)
      
                inds = [ip] # peak point
                for i in [ip+div, i1-div, div] :
                    inds = [inds, i]
                 
      
                inds = [inds, i1]  # Last point
                js[ip:i1] = spline_mono(inds, j[inds], ip+numpy.arange(i1-ip+1),  
                              yp0=0.0, ypn_1=(j[i1+1]-j[i1]))
      
                jps[:,y] = js
     
   
  
  #;;;;;;;;;;;;;;;;;;; TOPOLOGY ;;;;;;;;;;;;;;;;;;;;;;;
  # Calculate indices for backwards-compatibility
  
    nr = numpy.size(mesh.nrad)
    np = numpy.size(mesh.npol)
    if (nr == 2) and (np == 3) :
        print("Single null equilibrium")
    
        ixseps1 = mesh.nrad[0]
        ixseps2 = nx
    
        jyseps1_1 = mesh.npol[0]-1
        jyseps1_2 = mesh.npol[0] + numpy.int(old_div(mesh.npol[1],2))
        ny_inner = jyseps1_2
        jyseps2_1 = jyseps1_2
        jyseps2_2 = ny - mesh.npol[2]-1

    elif (nr == 3) and (np == 6) :
        print("Double null equilibrium")
    
        ixseps1 = mesh.nrad[0]
        ixseps2 = ixseps1 + mesh.nrad[1]
    
        jyseps1_1 = mesh.npol[0]-1
        jyseps2_1 = jyseps1_1 + mesh.npol[1]
    
        ny_inner = jyseps2_1 + mesh.npol[2] + 1
    
        jyseps1_2 = ny_inner + mesh.npol[3] - 1
        jyseps2_2 = jyseps1_2 + mesh.npol[4]
    
    elif (nr == 1) and (np == 1) :
    
        print("Single domain")
    
        ixseps1 = nx
        ixseps2 = nx
    
        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny,2))
        jyseps2_1 = numpy.int(old_div(ny,2))
        ny_inner = numpy.int(old_div(ny,2))
        jyseps2_2 = ny - 1
    
    else:
        print("***************************************") 
        print("* WARNING: Equilibrium not recognised *")
        print("*                                     *")
        print("*  Check mesh carefully!              *")
        print("*                                     *")
        print("*  Contact Ben Dudson                 *")
        print("*      [email protected]     *")
        print("***************************************") 
        ixseps1 = -1
        ixseps2 = -1
    
        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny,2))
        jyseps2_1 = numpy.int(old_div(ny,2))
        ny_inner = numpy.int(old_div(ny,2))
        jyseps2_2 = ny - 1
   

    print("Generating plasma profiles:")
          
    print("  1. Flat temperature profile")
    print("  2. Flat density profile")
    print("  3. Te proportional to density")
    while True:
        opt = input("Profile option:")
        if eval(opt) >= 1 and eval(opt) <= 3 : break

  
    if eval(opt) == 1 :
        # flat temperature profile
    
        print("Setting flat temperature profile")
        while True:
            Te_x = eval(input("Temperature (eV):"))
                
      
        # get density
            Ni = old_div(pressure, (2.* Te_x* 1.602e-19*1.0e20))
      
            print("numpy.maximum density (10^20 m^-3):", numpy.max(Ni))
      
            done = query_yes_no("Is this ok?")
            if done == 1 : break
    
        Te = numpy.zeros((nx, ny))+Te_x
        Ti = Te
        Ni_x = numpy.max(Ni)
        Ti_x = Te_x
    elif eval(opt) == 2 :
        print("Setting flat density profile")
    
        while True:
            Ni_x = eval(input("Density [10^20 m^-3]:"))
      
            # get temperature
            Te = old_div(pressure, (2.* Ni_x * 1.602e-19*1.0e20))
      
            print("numpy.maximum temperature (eV):", numpy.max(Te))
            if query_yes_no("Is this ok?") == 1 : break
    
        Ti = Te
        Ni = numpy.zeros((nx, ny)) + Ni_x
        Te_x = numpy.max(Te)
        Ti_x = Te_x
    else:
        print("Setting te proportional to density")
    
        while True:
            Te_x = eval(input("Maximum temperature [eV]:"))
            
            
            Ni_x = old_div(numpy.max(pressure), (2.*Te_x * 1.602e-19*1.0e20))
      
            print("Maximum density [10^20 m^-3]:", Ni_x)
      
            Te = Te_x * pressure / numpy.max(pressure)
            Ni = Ni_x * pressure / numpy.max(pressure)
            if query_yes_no("Is this ok?") == 1 : break
        Ti = Te
        Ti_x =  Te_x
   
  
    rmag = numpy.max(numpy.abs(Rxy))
    print("Setting rmag = ", rmag)
  
    bmag = numpy.max(numpy.abs(Bxy))
    print("Setting bmag = ", bmag)

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # save to file
    # open a new netCDF file for writing.
    handle = file_open(output) 

    print("Writing grid to file "+output)

    # Size of the grid

    s = file_write(handle, "nx", nx)
    s = file_write(handle, "ny", ny)


    # Topology for original scheme
    s = file_write(handle, "ixseps1", ixseps1)
    s = file_write(handle, "ixseps2", ixseps2)
    s = file_write(handle, "jyseps1_1", jyseps1_1)
    s = file_write(handle, "jyseps1_2", jyseps1_2)
    s = file_write(handle, "jyseps2_1", jyseps2_1)
    s = file_write(handle, "jyseps2_2", jyseps2_2)
    s = file_write(handle, "ny_inner", ny_inner);
  
    # Grid spacing
    
    s = file_write(handle, "dx", dx)
    s = file_write(handle, "dy", dy)
    
    s = file_write(handle, "ShiftAngle", qloop)
    s = file_write(handle, "zShift", qinty)
    s = file_write(handle, "pol_angle", pol_angle)
    s = file_write(handle, "ShiftTorsion", dqdpsi)

    s = file_write(handle, "Rxy",  Rxy)
    s = file_write(handle, "Zxy",  Zxy)
    s = file_write(handle, "Bpxy", Bpxy)
    s = file_write(handle, "Btxy", Btxy)
    s = file_write(handle, "Bxy",  Bxy)
    s = file_write(handle, "hthe", hthe)
    s = file_write(handle, "sinty", sinty)
    s = file_write(handle, "psixy", psixy)
    
    # Topology for general configurations
    s = file_write(handle, "yup_xsplit", mesh.yup_xsplit)
    s = file_write(handle, "ydown_xsplit", mesh.ydown_xsplit)
    s = file_write(handle, "yup_xin", mesh.yup_xin)
    s = file_write(handle, "yup_xout", mesh.yup_xout)
    s = file_write(handle, "ydown_xin", mesh.ydown_xin)
    s = file_write(handle, "ydown_xout", mesh.ydown_xout)
    s = file_write(handle, "nrad", mesh.nrad)
    s = file_write(handle, "npol", mesh.npol)

    # plasma profiles

    s = file_write(handle, "pressure", pressure)
    s = file_write(handle, "Jpar0", jpar0)
    s = file_write(handle, "Ni0", Ni)
    s = file_write(handle, "Te0", Te)
    s = file_write(handle, "Ti0", Ti)
    

    s = file_write(handle, "Ni_x", Ni_x)
    s = file_write(handle, "Te_x", Te_x)
    s = file_write(handle, "Ti_x", Ti_x)
    s = file_write(handle, "bmag", bmag)
    s = file_write(handle, "rmag", rmag)

    # Curvature
    s = file_write(handle, "bxcvx", bxcvx)
    s = file_write(handle, "bxcvy", bxcvy)
    s = file_write(handle, "bxcvz", bxcvz)

    # Psi range
    s = file_write(handle, "psi_axis", mesh.faxis)
    psi_bndry = mesh.faxis + mesh.fnorm
    s = file_write(handle, "psi_bndry", psi_bndry)

    file_close, handle
    print("DONE")
Esempio n. 2
0
def process_grid(rz_grid,
                 mesh,
                 output=None,
                 poorquality=None,
                 gui=None,
                 parent=None,
                 reverse_bt=None,
                 curv=None,
                 smoothpressure=None,
                 smoothhthe=None,
                 smoothcurv=None,
                 settings=None):

    if settings == None:
        # Create an empty structure
        settings = Bunch(dummy=0)

        # Check settings
        settings.calcp = -1
        settings.calcbt = -1
        settings.calchthe = -1
        settings.calcjpar = -1

# ;CATCH, err
# ;IF err NE 0 THEN BEGIN
# ;  PRINT, "PROCESS_GRID failed"
#;  PRINT, "   Error message: "+!ERROR_STATE.MSG
# ;  CATCH, /cancel
# ;  RETURN
# ;ENDIF

    MU = 4.e-7 * numpy.pi

    poorquality = 0

    if output == None: output = "bout.grd.nc"

    # Size of the mesh
    nx = numpy.int(numpy.sum(mesh.nrad))
    ny = numpy.int(numpy.sum(mesh.npol))

    # Find the midplane
    ymid = 0
    status = gen_surface(mesh=mesh)  # Start generator

    while True:
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
        if period:
            rm = numpy.max(mesh.Rxy[xi, yi])
            ymidindx = numpy.argmax(mesh.Rxy[xi, yi])
            ymid = yi[ymidindx]
            break

        if last == 1: break

    Rxy = numpy.asarray(mesh.Rxy)
    Zxy = numpy.asarray(mesh.Zxy)
    psixy = mesh.psixy * mesh.fnorm + mesh.faxis  # Non-normalised psi

    pressure = numpy.zeros((nx, ny))

    # Use splines to interpolate pressure profile
    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)
        if period:
            # Pressure only given on core surfaces
            # pressure[xi,yi] = SPLINE(rz_grid.npsigrid, rz_grid.pres, mesh.psixy[xi,yi[0]], /double)
            sol = interpolate.UnivariateSpline(rz_grid.npsigrid,
                                               rz_grid.pres,
                                               s=1)
            pressure[xi, yi] = sol(mesh.psixy[xi, yi[0]])

        else:

            pressure[xi, yi] = rz_grid.pres[numpy.size(rz_grid.pres) - 1]

        if last == 1: break

    # Add a minimum amount
    if numpy.min(pressure) < 1.0e-2 * numpy.max(pressure):
        print("****Minimum pressure is very small:", numpy.min(pressure))
        print("****Setting minimum pressure to 1% of maximum")
        pressure = pressure + 1e-2 * numpy.max(pressure)

    if smoothpressure != None:
        p0 = pressure[:, ymid]  # Keep initial pressure for comparison
        while True:
            #!P.multi=[0,0,2,0,0]
            fig = figure()
            plot(p0,
                 xtitle="X index",
                 ytitle="pressure at y=" + numpy.strip(numpy.str(ymid), 2) +
                 " dashed=original",
                 color=1,
                 lines=1)
            plot(pressure[:, ymid], color=1)
            plot(deriv(p0),
                 xtitle="X index",
                 ytitle="DERIV(pressure)",
                 color=1,
                 lines=1)
            plot(deriv(pressure[:, ymid]), color=1)
            sm = query_yes_no(
                "Smooth pressure profile?")  #, gui=gui, dialog_parent=parent)
            if sm:
                # Smooth the pressure profile

                p2 = pressure
                for i in range(6):
                    status = gen_surface(mesh=mesh)  # Start generator
                    while True:
                        # Get the next domain
                        period, yi, xi, last = gen_surface(period=period,
                                                           last=last,
                                                           xi=xi)

                        if (xi > 0) and (xi < (nx - 1)):
                            for j in range(numpy.size(yi)):
                                p2[xi,
                                   yi[j]] = (0.5 * pressure[xi, yi[j]] + 0.25 *
                                             (pressure[xi - 1, yi[j]] +
                                              pressure[xi + 1, yi[j]]))

                        # Make sure it's still constant on flux surfaces
                        p2[xi, yi] = numpy.mean(p2[xi, yi])
                        if last != None: break
                    pressure = p2

            if sm == 0: break

    if numpy.min(pressure) < 0.0:
        print("")
        print("============= WARNING ==============")
        print("Poor quality equilibrium: Pressure is negative")
        print("")
        poorquality = 1

    dpdpsi = DDX(psixy, pressure)

    #;IF MAX(dpdpsi)*mesh.fnorm GT 0.0 THEN BEGIN
    #;  PRINT, ""
    #;  PRINT, "============= WARNING =============="
    #;  PRINT, "Poor quality equilibrium: Pressure is increasing radially"
    #;  PRINT, ""
    #;  poorquality = 1
    #;ENDIF

    # Grid spacing
    dx = numpy.zeros((nx, ny))
    for y in range(ny):
        dx[0:(nx - 1), y] = psixy[1::, y] - psixy[0:(nx - 1), y]
        dx[nx - 1, y] = dx[nx - 2, y]

    # Sign
    bpsign = 1.
    xcoord = psixy
    if numpy.min(dx) < 0.:
        bpsign = -1.
        dx = -dx  # dx always positive
        xcoord = -xcoord

    dtheta = 2. * numpy.pi / numpy.float(ny)
    dy = numpy.zeros((nx, ny)) + dtheta

    # B field components
    # Following signs mean that psi increasing outwards from
    # core to edge results in Bp clockwise in the poloidal plane
    # i.e. in the positive Grad Theta direction.

    Brxy = old_div(mesh.dpsidZ, Rxy)
    Bzxy = old_div(-mesh.dpsidR, Rxy)
    Bpxy = numpy.sqrt(Brxy**2 + Bzxy**2)

    # Determine direction (dot B with grad y vector)

    dot = (Brxy[0, ymid] * (Rxy[0, ymid + 1] - Rxy[0, ymid - 1]) +
           Bzxy[0, ymid] * (Zxy[0, ymid + 1] - Zxy[0, ymid - 1]))

    if dot < 0.:
        print(
            "**** Poloidal field is in opposite direction to Grad Theta -> Bp negative"
        )
        Bpxy = -Bpxy
        if bpsign > 0: sys.exit()  # Should be negative
        bpsign = -1.0
    else:
        if bpsign < 0: sys.exit()  # Should be positive
        bpsign = 1.

# Get toroidal field from poloidal current function fpol
    Btxy = numpy.zeros((nx, ny))
    fprime = numpy.zeros((nx, ny))
    fp = deriv(rz_grid.npsigrid * (rz_grid.sibdry - rz_grid.simagx),
               rz_grid.fpol)

    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=period, xi=xi)

        if period:
            # In the core
            #fpol = numpy.interp(rz_grid.fpol, rz_grid.npsigrid, mesh.psixy[xi,yi], /spline)

            sol = interpolate.UnivariateSpline(rz_grid.npsigrid,
                                               rz_grid.fpol,
                                               s=1)
            #  fpol = SPLINE(rz_grid.npsigrid, rz_grid.fpol, mesh.psixy[xi,yi[0]], 'double')
            fpol = sol(mesh.psixy[xi, yi[0]])

            sol = interpolate.UnivariateSpline(rz_grid.npsigrid, fp, s=1)
            # fprime[xi,yi] = SPLINE(rz_grid.npsigrid, fp, mesh.psixy[xi,yi[0]], 'double')
            fprime[xi, yi] = sol(mesh.psixy[xi, yi[0]])

        else:
            # Outside core. Could be PF or SOL
            fpol = rz_grid.fpol[numpy.size(rz_grid.fpol) - 1]
            fprime[xi, yi] = 0.

        Btxy[xi, yi] = old_div(fpol, Rxy[xi, yi])

        if last == 1: break

    # Total B field
    Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # Go through the domains to get a starting estimate
    # of hthe
    hthe = numpy.zeros((nx, ny))

    #   Pick a midplane index
    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

        if period:
            # In the core
            rmax = numpy.argmax(Rxy[xi, yi])
            ymidplane = yi[rmax]
            break

        if last == 1: break

    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=period, last=last, xi=xi)

        n = numpy.size(yi)

        # Get distance along this line

        if period:

            # Periodic, so can use FFT
            #drdi = REAL_PART(fft_deriv(Rxy[xi, yi]))
            #dzdi = REAL_PART(fft_deriv(Zxy[xi, yi]))
            line = numpy.append(Rxy[xi, yi[n - 1::]], Rxy[xi, yi])
            line = numpy.append(line, Rxy[xi, yi[0:1]])

            drdi = deriv(line)[1:n + 1]

            line = numpy.append(Zxy[xi, yi[n - 1::]], Zxy[xi, yi])
            line = numpy.append(line, Zxy[xi, yi[0:1]])

            dzdi = deriv(line)[1:n + 1]
        else:
            # Non-periodic
            drdi = numpy.gradient(Rxy[xi, yi])
            dzdi = numpy.gradient(Zxy[xi, yi])

        dldi = numpy.sqrt(drdi**2 + dzdi**2)

        if 0:

            # Need to smooth to get sensible results
            if period:
                n = numpy.size(dldi)
                line = numpy.append(dldi[(n - 2)::], dldi)  # once
                line = numpy.append(line, dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n + 4)]

                line = numpy.append(dldi[(n - 2)::], dldi)  #twice
                line = numpy.append(line, dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n + 4)]

                line = numpy.append(dldi[(n - 2)::], dldi)  # three
                line = numpy.append(line, dldi[0:2])
                dldi = SMOOTH(line, 5)[4:(n + 4)]

            else:
                line = dldi
                dldi = SMOOTH(line, 5)[2:n + 2]
                line = dldi
                dldi = SMOOTH(line, 5)[2:n + 2]
                line = dldi
                dldi = SMOOTH(dldi, 5)[2:n + 2]

        hthe[xi, yi] = old_div(dldi, dtheta)  # First estimate of hthe

        # Get outboard midplane
        if period and xi == 0:
            m = numpy.argmax(Rxy[0, yi])
            ymidplane = yi[m]

        if last == 1: break

    print("Midplane index ", ymidplane)

    fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pressure)
    print("Force imbalance: ", numpy.mean(numpy.abs(fb0)),
          numpy.max(numpy.abs(fb0)))

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # Correct pressure using hthe

    print("Calculating pressure profile from force balance")

    try:

        # Calculate force balance
        dpdx = old_div((-Bpxy * DDX(xcoord, Bpxy * hthe) -
                        Btxy * hthe * DDX(xcoord, Btxy) -
                        (Btxy * Btxy * hthe / Rxy) * DDX(xcoord, Rxy)),
                       (MU * hthe))

        # Surface average
        dpdx2 = surface_average(dpdx, mesh)

        pres = numpy.zeros((nx, ny))
        # Integrate to get pressure
        for i in range(ny):
            pres[:, i] = int_func(psixy[:, i], dpdx2[:, i])
            pres[:, i] = pres[:, i] - pres[nx - 1, i]

        status = gen_surface(mesh=mesh)  # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

            ma = numpy.max(pres[xi, yi])

            for i in range(numpy.size(yi)):
                pres[:, yi[i]] = pres[:, yi[i]] - pres[xi, yi[i]] + ma

            if last == 1: break

        pres = pres - numpy.min(pres)

        # Some sort of smoothing here?

        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, hthe, pres)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)),
              numpy.max(numpy.abs(fb0)))

        #!P.MULTI=[0,0,2,0,0]
        fig = figure(figsize=(7, 11))
        subplots_adjust(left=.07,
                        bottom=.07,
                        right=0.95,
                        top=0.95,
                        wspace=.3,
                        hspace=.25)

        SURFACE(pressure, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2, 1, 1])
        title("Input pressure")
        SURFACE(pres, fig, xtitle="X", ytitle="Y", var='Pa', sub=[2, 1, 2])
        title("New pressure")
        #  arrange the plot on the screen
        #      mngr = get_current_fig_manager()
        #      geom = mngr.window.geometry()
        #      x,y,dx,dy = geom.getRect()
        #      mngr.window.setGeometry(0, 0, dx, dy)
        #
        show(block=False)

        calcp = settings.calcp

        if calcp == -1:
            calcp = query_yes_no(
                "Keep new pressure?")  #, gui=gui, dialog_parent=parent)
        else:
            time.sleep(2)
        if calcp == 1:
            pressure = pres
            dpdpsi = dpdx2

    except Exception:
        print("WARNING: Pressure profile calculation failed: "
              )  #, !ERROR_STATE.MSG
        pass

    #CATCH, /cancel

#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
# Correct f = RBt using force balance

    calcbt = settings.calcbt
    if calcbt == -1:
        calcbt = query_yes_no("Correct f=RBt using force balance?"
                              )  #, gui=gui, dialog_parent=parent)
    if calcbt == 1:

        new_Btxy = newton_Bt(psixy, Rxy, Btxy, Bpxy, pres, hthe, mesh)

        fb0 = force_balance(psixy, Rxy, Bpxy, new_Btxy, hthe, pressure)
        print("force imbalance: ", numpy.mean(numpy.abs(fb0)),
              numpy.max(numpy.abs(fb0)))

        fig = figure(figsize=(7, 11))
        subplots_adjust(left=.07,
                        bottom=.07,
                        right=0.95,
                        top=0.95,
                        wspace=.3,
                        hspace=.25)

        subplot(211)
        SURFACE(Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2, 1, 1])
        title("Input Bt")
        subplot(212)
        SURFACE(new_Btxy, fig, xtitle="X", ytitle="Y", var='T', sub=[2, 1, 2])
        title("New Bt")
        #  arrange the plot on the screen
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(600, 0, dx, dy)

        show(block=False)

        calcbt = settings.calcbt
        if calcbt == -1:
            calcbt = query_yes_no(
                "Keep new Bt?")  #, gui=gui, dialog_parent=parent)
        if calcbt == 1:
            Btxy = new_Btxy
            Bxy = numpy.sqrt(Btxy**2 + Bpxy**2)

#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
# CALCULATE HTHE
# Modify hthe to fit force balance using initial guess
# Does not depend on signs
#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    calchthe = settings.calchthe
    if calchthe == -1:
        calchthe = query_yes_no("Adjust hthe using force balance?"
                                )  #, gui=gui, dialog_parent=parent)
    if calchthe == 1:
        # This doesn't behave well close to the x-points
        fixhthe = numpy.int(old_div(nx, 2))
        nh = correct_hthe(Rxy,
                          psixy,
                          Btxy,
                          Bpxy,
                          hthe,
                          pressure,
                          fixhthe=fixhthe)

        fb0 = force_balance(psixy, Rxy, Bpxy, Btxy, nh, pressure)
        print("Force imbalance: ", numpy.mean(numpy.abs(fb0)),
              numpy.max(numpy.abs(fb0)))

        print("numpy.maximum difference in hthe: ",
              numpy.max(numpy.abs(hthe - nh)))
        print("numpy.maximum percentage difference: ",
              100. * numpy.max(numpy.abs(old_div((hthe - nh), hthe))))

        #!P.multi=[0,0,1,0,0]
        fig = figure(figsize=(7, 4))
        title("Poloidal arc length at midplane. line is initial estimate")
        plot(hthe[:, 0], '-')
        plot(nh[:, 0], 'r-+')
        #  arrange the plot on the screen
        #mngr = get_current_fig_manager()
        #geom = mngr.window.geometry()
        #x,y,dx,dy = geom.getRect()
        #mngr.window.setGeometry(0, 1150, dx, dy)

        show(block=False)

        if query_yes_no(
                "Keep new hthe?") == 1:  #, gui=gui, dialog_parent=parent) :
            hthe = nh

    if smoothhthe != None:
        # Smooth hthe to prevent large jumps in X or Y. This
        # should be done by creating a better mesh in the first place

        # Need to smooth in Y and X otherwise smoothing in X
        # produces discontinuities in Y
        hold = hthe

        if 1:
            # Nonlinear smoothing. Tries to smooth only regions with large
            # changes in gradient

            hthe = 0.  # smooth_nl(hthe, mesh)

        else:
            # Just use smooth in both directions

            for i in range(ny):
                hthe[:, i] = SMOOTH(SMOOTH(hthe[:, i], 10), 10)

        status = gen_surface(mesh=mesh)  # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

            n = numpy.size(yi)

            if period:
                hthe[xi, yi] = (SMOOTH([
                    hthe[xi, yi[(n - 4):(n - 1)]], hthe[xi, yi], hthe[xi,
                                                                      yi[0:3]]
                ], 4))[4:(n + 3)]
            else:
                hthe[xi, yi] = SMOOTH(hthe[xi, yi], 4)

            if last == 1: break

    # Calculate field-line pitch
    pitch = hthe * Btxy / (Bpxy * Rxy)

    # derivative with psi
    dqdpsi = DDX(psixy, pitch)

    qinty, qloop = int_y(pitch,
                         mesh,
                         loop=0,
                         nosmooth='nosmooth',
                         simple='simple')
    qinty = qinty * dtheta
    qloop = qloop * dtheta

    sinty = int_y(dqdpsi, mesh, nosmooth='nosmooth', simple='simple') * dtheta

    # NOTE: This is only valid in the core
    pol_angle = numpy.zeros((nx, ny))
    for i in range(nx):
        pol_angle[i, :] = 2.0 * numpy.pi * qinty[i, :] / qloop[i]

    #;;;;;;;;;;;;;;;;;;; THETA_ZERO ;;;;;;;;;;;;;;;;;;;;;;
    # re-set zshift to be zero at the outboard midplane

    print("MIDPLANE INDEX = ", ymidplane)

    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

        w = numpy.size(numpy.where(yi == ymidplane))
        if w > 0:
            # Crosses the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi, ymidplane]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi, ymidplane]
        else:
            # Doesn't include a point at the midplane
            qinty[xi, yi] = qinty[xi, yi] - qinty[xi, yi[0]]
            sinty[xi, yi] = sinty[xi, yi] - sinty[xi, yi[0]]

        if last == 1: break

    print("")
    print("==== Calculating curvature ====")

    #;;;;;;;;;;;;;;;;;;; CURVATURE ;;;;;;;;;;;;;;;;;;;;;;;
    # Calculating b x kappa

    if curv == None:

        print("*** Calculating curvature in toroidal coordinates")

        thetaxy = numpy.zeros((nx, ny))
        status = gen_surface(mesh=mesh)  # Start generator
        while True:
            # Get the next domain
            period, yi, xi, last = gen_surface(period=None, last=None, xi=None)
            thetaxy[xi,
                    yi] = numpy.arange(numpy.size(yi)).astype(float) * dtheta
            if last == 1: break

        bxcv = curvature(nx,
                         ny,
                         Rxy,
                         Zxy,
                         Brxy,
                         Bzxy,
                         Btxy,
                         psixy,
                         thetaxy,
                         hthe,
                         mesh=mesh)

        bxcvx = bpsign * bxcv.psi
        bxcvy = bxcv.theta
        bxcvz = bpsign * (bxcv.phi - sinty * bxcv.psi - pitch * bxcv.theta)

        # x borders
        bxcvx[0, :] = bxcvx[1, :]
        bxcvx[nx - 1, :] = bxcvx[nx - 2, :]

        bxcvy[0, :] = bxcvy[1, :]
        bxcvy[nx - 1, :] = bxcvy[nx - 2, :]

        bxcvz[0, :] = bxcvz[1, :]
        bxcvz[nx - 1, :] = bxcvz[nx - 2, :]

    elif curv == 1:
        # Calculate on R-Z mesh and then interpolate onto grid
        # ( cylindrical coordinates)

        print("*** Calculating curvature in cylindrical coordinates")

        bxcv = rz_curvature(rz_grid)

        # DCT methods cause spurious oscillations
        # Linear interpolation seems to be more robust
        bxcv_psi = numpy.interp(bxcv.psi, mesh.Rixy, mesh.Zixy)
        bxcv_theta = old_div(numpy.interp(bxcv.theta, mesh.Rixy, mesh.Zixy),
                             hthe)
        bxcv_phi = numpy.interp(bxcv.phi, mesh.Rixy, mesh.Zixy)

        # If Bp is reversed, then Grad x = - Grad psi
        bxcvx = bpsign * bxcv_psi
        bxcvy = bxcv_theta
        bxcvz = bpsign * (bxcv_phi - sinty * bxcv_psi - pitch * bxcv_theta)
    elif curv == 2:
        # Curvature from Curl(b/B)

        bxcvx = bpsign * (Bpxy * Btxy * Rxy * DDY(old_div(1., Bxy), mesh) /
                          hthe)
        bxcvy = -bpsign * Bxy * Bpxy * DDX(xcoord,
                                           Btxy * Rxy / Bxy ^ 2) / (2. * hthe)
        bxcvz = Bpxy ^ 3 * DDX(xcoord, old_div(
            hthe, Bpxy)) / (2. * hthe * Bxy) - Btxy * Rxy * DDX(
                xcoord, old_div(Btxy, Rxy)) / (2. * Bxy) - sinty * bxcvx

    else:
        # calculate in flux coordinates.

        print("*** Calculating curvature in flux coordinates")

        dpb = numpy.zeros((nx, ny))  # quantity used for y and z components

        for i in range(ny):
            dpb[:, i] = MU * dpdpsi / Bxy[:, i]

        dpb = dpb + DDX(xcoord, Bxy)

        bxcvx = bpsign * (Bpxy * Btxy * Rxy * DDY(old_div(1., Bxy), mesh) /
                          hthe)
        bxcvy = bpsign * (Bpxy * Btxy * Rxy * dpb / (hthe * Bxy ^ 2))
        bxcvz = -dpb - sinty * bxcvx

    if smoothcurv:
        # Smooth curvature to prevent large jumps

        # Nonlinear smoothing. Tries to smooth only regions with large
        # changes in gradient

        bz = bxcvz + sinty * bxcvx

        print("Smoothing bxcvx...")
        bxcvx = 0.  #smooth_nl(bxcvx, mesh)
        print("Smoothing bxcvy...")
        bxcvy = 0.  #smooth_nl(bxcvy, mesh)
        print("Smoothing bxcvz...")
        bz = 0.  #smooth_nl(bz, mesh)

        bxcvz = bz - sinty * bxcvx

#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
# CALCULATE PARALLEL CURRENT
#
# Three ways to calculate Jpar0:
# 1. From fprime and pprime
# 2. From Curl(B) in field-aligned coords
# 3. From the curvature
#
# Provides a way to check if Btor should be reversed
#
#;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    print("")
    print("==== Calculating parallel current ====")

    jpar0 = -Bxy * fprime / MU - Rxy * Btxy * dpdpsi / Bxy

    # Set to zero in PF and SOL
    status = gen_surface(mesh=mesh)  # Start generator
    while True:
        # Get the next domain
        period, yi, xi, last = gen_surface(period=None, last=None, xi=None)

        if period == None: jpar0[xi, yi] = 0.0
        if last == 1: break

# Curl(B) expression for Jpar0 (very noisy usually)
    j0 = (bpsign * ((Bpxy * Btxy * Rxy / (Bxy * hthe)) *
                    (DDX(xcoord, Bxy**2 * hthe / Bpxy) - bpsign * Btxy * Rxy *
                     DDX(xcoord, Btxy * hthe /
                         (Rxy * Bpxy))) - Bxy * DDX(xcoord, Btxy * Rxy)) / MU)

    # Create a temporary mesh structure to send to adjust_jpar
    tmp_mesh = Bunch(mesh,
                     bxcvx=bxcvx,
                     bxcvy=bxcvy,
                     bxcvz=bxcvz,
                     Bpxy=Bpxy,
                     Btxy=Btxy,
                     Bxy=Bxy,
                     dx=dx,
                     dy=dy,
                     hthe=hthe,
                     jpar0=jpar0,
                     pressure=pressure)
    tmp_mesh.psixy = psixy

    jpar = adjust_jpar(tmp_mesh, noplot='noplot')

    #!P.multi=[0,2,2,0,0]

    fig = figure(figsize=(15, 11))
    subplots_adjust(left=.07,
                    bottom=.07,
                    right=0.95,
                    top=0.95,
                    wspace=.3,
                    hspace=.25)

    subplot(221)
    SURFACE(jpar0, fig, xtitle="X", ytitle="Y", var='A', sub=[2, 2, 1])
    title("Jpar from F' and P'")

    subplot(222)
    SURFACE(jpar, fig, xtitle="X", ytitle="Y", var='A', sub=[2, 2, 2])
    title("Jpar from curvature")

    subplot(223)
    plot(jpar0[0, :], '-', jpar[0, :], '+')
    ylim([
        numpy.min([jpar0[0, :], jpar[0, :]]),
        numpy.max([jpar0[0, :], jpar[0, :]])
    ])
    title("jpar at x=0. Solid from f' and p'")

    subplot(224)
    plot(jpar0[:, ymidplane], '-', jpar[:, ymidplane], '+')
    ylim([
        numpy.min([jpar0[:, ymidplane], jpar[:, ymidplane]]),
        numpy.max([jpar0[:, ymidplane], jpar[:, ymidplane]])
    ])

    title("Jpar at y=" + numpy.str(ymidplane) + " Solid from f' and p'")

    #  arrange the plot on the screen
    #mngr = get_current_fig_manager()
    #geom = mngr.window.geometry()
    #x,y,dx,dy = geom.getRect()
    #mngr.window.setGeometry(1350, 0, dx, dy)

    show(block=False)

    # !P.multi=0

    calcjpar = settings.calcjpar
    if calcjpar == -1:
        calcjpar = query_yes_no(
            "Use Jpar from curvature?")  #, gui=gui, dialog_parent=parent)
    if calcjpar == True:
        jpar0 = jpar

    if 0:

        # Try smoothing jpar0 in psi, preserving zero points and maxima
        jps = jpar0
        for y in range(ny):
            j = jpar0[:, y]
            js = j
            ma = numpy.max(numpy.abs(j))
            ip = numpy.argmax(numpy.abs(j))
            if (ma < 1.e-4) or (ip == 0):
                jps[:, y] = j

            level = 1.
            #i0 = MAX(WHERE(ABS(j[0:ip]) LT level))
            i1 = numpy.min(numpy.where(numpy.abs(j[ip::]) < level))

            #IF i0 LE 0 THEN i0 = 1
            i0 = 1

            if i1 == -1:
                i1 = nx - 2
            else:
                i1 = i1 + ip

            if (ip <= i0) or (ip >= i1):

                # Now preserve starting and end points, and peak value
                div = numpy.int(old_div(
                    (i1 - i0),
                    10)) + 1  # reduce number of points by this factor

                inds = [i0]  # first point
                for i in [i0 + div, ip - div, div]:
                    inds = [inds, i]
                inds = [inds, ip]  # Put in the peak point

                # Calculate spline interpolation of inner part

                js[0:ip] = spline_mono(inds,
                                       j[inds],
                                       numpy.arange(ip + 1),
                                       yp0=(j[i0] - j[i0 - 1]),
                                       ypn_1=0.0)

                inds = [ip]  # peak point
                for i in [ip + div, i1 - div, div]:
                    inds = [inds, i]

                inds = [inds, i1]  # Last point
                js[ip:i1] = spline_mono(inds,
                                        j[inds],
                                        ip + numpy.arange(i1 - ip + 1),
                                        yp0=0.0,
                                        ypn_1=(j[i1 + 1] - j[i1]))

                jps[:, y] = js

#;;;;;;;;;;;;;;;;;;; TOPOLOGY ;;;;;;;;;;;;;;;;;;;;;;;
# Calculate indices for backwards-compatibility

    nr = numpy.size(mesh.nrad)
    np = numpy.size(mesh.npol)
    if (nr == 2) and (np == 3):
        print("Single null equilibrium")

        ixseps1 = mesh.nrad[0]
        ixseps2 = nx

        jyseps1_1 = mesh.npol[0] - 1
        jyseps1_2 = mesh.npol[0] + numpy.int(old_div(mesh.npol[1], 2))
        ny_inner = jyseps1_2
        jyseps2_1 = jyseps1_2
        jyseps2_2 = ny - mesh.npol[2] - 1

    elif (nr == 3) and (np == 6):
        print("Double null equilibrium")

        ixseps1 = mesh.nrad[0]
        ixseps2 = ixseps1 + mesh.nrad[1]

        jyseps1_1 = mesh.npol[0] - 1
        jyseps2_1 = jyseps1_1 + mesh.npol[1]

        ny_inner = jyseps2_1 + mesh.npol[2] + 1

        jyseps1_2 = ny_inner + mesh.npol[3] - 1
        jyseps2_2 = jyseps1_2 + mesh.npol[4]

    elif (nr == 1) and (np == 1):

        print("Single domain")

        ixseps1 = nx
        ixseps2 = nx

        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny, 2))
        jyseps2_1 = numpy.int(old_div(ny, 2))
        ny_inner = numpy.int(old_div(ny, 2))
        jyseps2_2 = ny - 1

    else:
        print("***************************************")
        print("* WARNING: Equilibrium not recognised *")
        print("*                                     *")
        print("*  Check mesh carefully!              *")
        print("*                                     *")
        print("*  Contact Ben Dudson                 *")
        print("*      [email protected]     *")
        print("***************************************")
        ixseps1 = -1
        ixseps2 = -1

        jyseps1_1 = -1
        jyseps1_2 = numpy.int(old_div(ny, 2))
        jyseps2_1 = numpy.int(old_div(ny, 2))
        ny_inner = numpy.int(old_div(ny, 2))
        jyseps2_2 = ny - 1

    print("Generating plasma profiles:")

    print("  1. Flat temperature profile")
    print("  2. Flat density profile")
    print("  3. Te proportional to density")
    while True:
        opt = input("Profile option:")
        if eval(opt) >= 1 and eval(opt) <= 3: break

    if eval(opt) == 1:
        # flat temperature profile

        print("Setting flat temperature profile")
        while True:
            Te_x = eval(input("Temperature (eV):"))

            # get density
            Ni = old_div(pressure, (2. * Te_x * 1.602e-19 * 1.0e20))

            print("numpy.maximum density (10^20 m^-3):", numpy.max(Ni))

            done = query_yes_no("Is this ok?")
            if done == 1: break

        Te = numpy.zeros((nx, ny)) + Te_x
        Ti = Te
        Ni_x = numpy.max(Ni)
        Ti_x = Te_x
    elif eval(opt) == 2:
        print("Setting flat density profile")

        while True:
            Ni_x = eval(input("Density [10^20 m^-3]:"))

            # get temperature
            Te = old_div(pressure, (2. * Ni_x * 1.602e-19 * 1.0e20))

            print("numpy.maximum temperature (eV):", numpy.max(Te))
            if query_yes_no("Is this ok?") == 1: break

        Ti = Te
        Ni = numpy.zeros((nx, ny)) + Ni_x
        Te_x = numpy.max(Te)
        Ti_x = Te_x
    else:
        print("Setting te proportional to density")

        while True:
            Te_x = eval(input("Maximum temperature [eV]:"))

            Ni_x = old_div(numpy.max(pressure),
                           (2. * Te_x * 1.602e-19 * 1.0e20))

            print("Maximum density [10^20 m^-3]:", Ni_x)

            Te = Te_x * pressure / numpy.max(pressure)
            Ni = Ni_x * pressure / numpy.max(pressure)
            if query_yes_no("Is this ok?") == 1: break
        Ti = Te
        Ti_x = Te_x

    rmag = numpy.max(numpy.abs(Rxy))
    print("Setting rmag = ", rmag)

    bmag = numpy.max(numpy.abs(Bxy))
    print("Setting bmag = ", bmag)

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # save to file
    # open a new netCDF file for writing.
    handle = file_open(output)

    print("Writing grid to file " + output)

    # Size of the grid

    s = file_write(handle, "nx", nx)
    s = file_write(handle, "ny", ny)

    # Topology for original scheme
    s = file_write(handle, "ixseps1", ixseps1)
    s = file_write(handle, "ixseps2", ixseps2)
    s = file_write(handle, "jyseps1_1", jyseps1_1)
    s = file_write(handle, "jyseps1_2", jyseps1_2)
    s = file_write(handle, "jyseps2_1", jyseps2_1)
    s = file_write(handle, "jyseps2_2", jyseps2_2)
    s = file_write(handle, "ny_inner", ny_inner)

    # Grid spacing

    s = file_write(handle, "dx", dx)
    s = file_write(handle, "dy", dy)

    s = file_write(handle, "ShiftAngle", qloop)
    s = file_write(handle, "zShift", qinty)
    s = file_write(handle, "pol_angle", pol_angle)
    s = file_write(handle, "ShiftTorsion", dqdpsi)

    s = file_write(handle, "Rxy", Rxy)
    s = file_write(handle, "Zxy", Zxy)
    s = file_write(handle, "Bpxy", Bpxy)
    s = file_write(handle, "Btxy", Btxy)
    s = file_write(handle, "Bxy", Bxy)
    s = file_write(handle, "hthe", hthe)
    s = file_write(handle, "sinty", sinty)
    s = file_write(handle, "psixy", psixy)

    # Topology for general configurations
    s = file_write(handle, "yup_xsplit", mesh.yup_xsplit)
    s = file_write(handle, "ydown_xsplit", mesh.ydown_xsplit)
    s = file_write(handle, "yup_xin", mesh.yup_xin)
    s = file_write(handle, "yup_xout", mesh.yup_xout)
    s = file_write(handle, "ydown_xin", mesh.ydown_xin)
    s = file_write(handle, "ydown_xout", mesh.ydown_xout)
    s = file_write(handle, "nrad", mesh.nrad)
    s = file_write(handle, "npol", mesh.npol)

    # plasma profiles

    s = file_write(handle, "pressure", pressure)
    s = file_write(handle, "Jpar0", jpar0)
    s = file_write(handle, "Ni0", Ni)
    s = file_write(handle, "Te0", Te)
    s = file_write(handle, "Ti0", Ti)

    s = file_write(handle, "Ni_x", Ni_x)
    s = file_write(handle, "Te_x", Te_x)
    s = file_write(handle, "Ti_x", Ti_x)
    s = file_write(handle, "bmag", bmag)
    s = file_write(handle, "rmag", rmag)

    # Curvature
    s = file_write(handle, "bxcvx", bxcvx)
    s = file_write(handle, "bxcvy", bxcvy)
    s = file_write(handle, "bxcvz", bxcvz)

    # Psi range
    s = file_write(handle, "psi_axis", mesh.faxis)
    psi_bndry = mesh.faxis + mesh.fnorm
    s = file_write(handle, "psi_bndry", psi_bndry)

    file_close, handle
    print("DONE")
Esempio n. 3
0
def create_grid( F, R, Z, in_settings, critical,
                      boundary=None,  
                      iter=None, 
                      fpsi=None, fast=None):#, # f(psi) = R*Bt current function
                      #nrad_flexible,
                      #single_rad_grid, 
                      #xpt_mindist, xpt_mul, strictbndry,debug):

   # if size(nrad_flexible) == 0 :
    #    nrad_flexible = 0

    if iter==None:
        iter = 0
    if iter > 3:
        print("ERROR: Too many iterations")
        return #, {error:1}


   
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  # Check the settings
  # If a setting is missing, set a default value
  
   # inspect.getargspec(create_grid)
    
 # if N_PARAMS() LT 3 THEN BEGIN
  #  PRINT, "ERROR: Need at least a 2D array of psi values, R and Z arrays"
  #  RETURN, {error:1}
  #ENDIF ELSE IF N_PARAMS() LT 4 THEN BEGIN
  #  ; Settings omitted. Set defaults
  #  print "Settings not given -> using default values"
   # settings = {psi_inner:0.9, 
    #            psi_outer:1.1, 
     #           nrad:36, 
      #          npol:64, 
       #         rad_peaking:0.0, 
        #        pol_peaking:0.0, 
         #       parweight:0.0}
 # ENDIF ELSE BEGIN
  # print "Checking settings"
    settings = in_settings # So the input isn't changed
  #  str_check_present, settings, 'psi_inner', 0.9
  #  str_check_present, settings, 'psi_outer', 1.1
  #  str_check_present, settings, 'nrad', 36
  #  str_check_present, settings, 'npol', 64
  #  str_check_present, settings, 'rad_peaking', 0.0
  #  str_check_present, settings, 'pol_peaking', 0.0
  #  str_check_present, settings, 'parweight', 0.0
  #ENDELSE

    s = numpy.ndim(F)
    s1 = numpy.shape(F)
    if s != 2:
        print("ERROR: First argument must be 2D array of psi values")
        return #, {error:1}
    nx = s1[0]
    ny = s1[1]
  
    s = numpy.ndim(R)
    s1 = numpy.size(R)
    if s != 1  or s1 != nx :
        print("ERROR: Second argument must be 1D array of major radii")
        return # {error:1}
  
    s = numpy.ndim(Z)
    s1 = numpy.size(Z)
    if s != 1  or s1 != ny:
        print("ERROR: Second argument must be 1D array of heights")
        return # {error:1}


  # Get an even number of points for efficient FFTs
    if nx % 2 == 1:
    # odd number of points in R. Cut out last point
        R = R[0:(nx-1)]
        F = F[0:(nx-1), :]
        nx = nx - 1
  
    if ny % 2 == 1:
    # Same for Z
        Z = Z[0:(ny-1)]
        F = F[:,0:(ny-1)]
        ny = ny - 1
  


    if boundary != None:
        s = numpy.ndim(boundary)
        s1= numpy.shape(boundary)
        if s != 2  or s1[0] != 2:
            print("WARNING: boundary must be a 2D array: [2, n]. Ignoring")
            boundary = 0
        else:       
        # Calculate indices
            bndryi = numpy.zeros((2,1188))
            bndryi[0,:] = numpy.interp(bndryi[0,:], R, numpy.arange(0.,nx))
            bndryi[1,:] = numpy.interp(bndryi[1,:], Z, numpy.arange(0.,ny))

  
        if bndryi == None :
            bndryi = numpy.zeros((2,4))
            bndryi[0,:] = [1, nx-1, nx-1, 1]
            bndryi[1,:] = [1, 1, ny-1, ny-1]
 
  
  #;;;;;;;;;;;;;; Psi interpolation data ;;;;;;;;;;;;;;
  
    interp_data = Bunch(nx=nx, ny=ny, 
                 method=0, 
                 f= F)       # Always include function
               
    if fast == 'fast':
        print("Using Fast settings")
        interp_data.method = 2
  

  #;;;;;;;;;;;;;;; First plot ;;;;;;;;;;;;;;;;

    nlev = 100
    minf = numpy.min(F)
    maxf = numpy.max(F)
    levels = numpy.arange(numpy.float(nlev))*(maxf-minf)/numpy.float(nlev-1) + minf

    Rr=numpy.tile(R,ny).reshape(ny,nx).T
    Zz=numpy.tile(Z,nx).reshape(nx,ny)


    contour( Rr, Zz, F, levels=levels)
    
          #  arrange the plot on the screen      
    #mngr = get_current_fig_manager()
    #geom = mngr.window.geometry()
    #x,y,dx,dy = geom.getRect()
    #mngr.window.setGeometry(100, 100, dx, dy)


  
  #  if boundary != None :
  #      plot(boundary[0,:],boundary[1,:],'r-')
  
    show(block=False)  
    
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  

    n_opoint = critical.n_opoint
    n_xpoint = critical.n_xpoint
    primary_opt = critical.primary_opt
    inner_sep   = critical.inner_sep
    opt_ri = critical.opt_ri
    opt_zi = critical.opt_zi
    opt_f  = critical.opt_f
    xpt_ri = numpy.array(critical.xpt_ri).flatten()
    xpt_zi = numpy.array(critical.xpt_zi).flatten()
    xpt_f  = numpy.array(critical.xpt_f).flatten()
    
    


  # Overplot the separatrices, O-points
  #oplot_critical, F, R, Z, critical

  # Psi normalisation factors

    faxis = opt_f[primary_opt]
    
    fnorm = xpt_f[inner_sep] - opt_f[primary_opt]
    


  # From normalised psi, get range of f
    f_inner = faxis + numpy.min(settings.psi_inner)*fnorm
    f_outer = faxis + numpy.max(settings.psi_outer)*fnorm
    
  
  # Check the number of x-points
    if critical.n_xpoint == 0 :
    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # Grid entirely in the core
    	print("Generating grid entirely in the core")

    nrad = numpy.sum(settings.nrad) # Add up all points
    npol = numpy.sum(settings.npol)
    rad_peaking = settings.rad_peaking[0] # Just the first region
    pol_peaking = settings.pol_peaking[0]

    # work out where to put the surfaces in the core
    fvals = radial_grid(nrad, f_inner, f_outer, 1, 1, [xpt_f[inner_sep]], rad_peaking)

    fvals = fvals.flatten()
    # Create a starting surface
    sind = numpy.int(old_div(nrad, 2))
    start_f = fvals[sind]
    

  #  contour_lines( F, numpy.arange(nx).astype(float), numpy.arange(ny).astype(float), levels=[start_f])
    cs=contour( Rr, Zz, F,  levels=[start_f])

    p = cs.collections[0].get_paths()
 #
 #  You might get more than one contours for the same start_f. We need to keep the closed one  
    vn=numpy.zeros(numpy.size(p))
    
      
    v = p[0].vertices
    vn[0]=numpy.shape(v)[0]
    xx=[v[:,0]]
    yy=[v[:,1]]
    
    if numpy.shape(vn)[0] > 1:
        for i in range(1,numpy.shape(vn)[0]):
            v = p[i].vertices
            vn[i]=numpy.shape(v)[0]
            xx.append(v[:,0])
            yy.append(v[:,1])
            #xx = [xx,v[:,0]]
            #yy = [yy,v[:,1]]

    print("PRIMARY: ", primary_opt, opt_ri[primary_opt], opt_zi[primary_opt])

    if numpy.shape(vn)[0] > 1 :
        # Find the surface closest to the o-point
        opt_r = numpy.interp(opt_ri[primary_opt], numpy.arange(len(R)), R)
        opt_z = numpy.interp(opt_zi[primary_opt], numpy.arange(len(Z)), Z)
        
        ind = closest_line(xx, yy, opt_r, opt_z)
        
        x=xx[ind]
        y=yy[ind]
        print("Contour: ", ind)
    else:
        ind = 0
        x=xx[0]
        y=yy[0]
          

# plot the start_f line     
    zc = cs.collections[0]
    setp(zc, linewidth=4)

    clabel(cs, [start_f],  # label the level
           inline=1,
           fmt='%9.6f',
           fontsize=14)
     
    draw()             
   
    show(block=False)  
   
#

    ans=query_yes_no('Press enter to create grid')  
    
    if ans != 1 : 
        show()
	sys.exit()
      
    start_ri, start_zi=transform_xy(x,y,R,Z)
    
    ## Make sure that the line goes clockwise
    #
    m = numpy.argmax(numpy.interp(start_zi,numpy.arange(Z.size).astype(float), Z))
    if (numpy.gradient(numpy.interp(start_ri, numpy.arange(R.size).astype(float), R)))[m] < 0.0:
      # R should be increasing at the top. Need to reverse
        start_ri = start_ri[::-1]
        start_zi = start_zi[::-1]
        print('points reversed')
        

    ## Last point should be the same as the first
    #
    # Smooth and refine the starting location
    np = numpy.size(start_ri)
    s = 3
    ar=numpy.append(numpy.append(start_ri[(np-s-1):(np-1)], start_ri), start_ri[1:s+1])
    start_ri = SMOOTH(ar, window_len=s)[s+1:(np+s+1)]
    az=numpy.append(numpy.append(start_zi[(np-s-1):(np-1)], start_zi), start_zi[1:s+1])
    start_zi = SMOOTH(az, window_len=s)[s+1:(np+s+1)]
    
    
    for i in range (np) :
        ri1=0.
        zi1=0.
        out=follow_gradient( interp_data, R, Z, start_ri[i], start_zi[i], start_f, ri1, zi1 )
        status=out.status
        ri1=out.rinext
        zi1=out.zinext

        start_ri[i] = ri1
        start_zi[i] = zi1
    
        
    a = grid_region(interp_data, R, Z,
                    start_ri, start_zi, 
                    fvals, 
                    sind, 
                    npol, 
                    boundary=boundary,
                    fpsi=fpsi, 
                    parweight=settings.parweight, 
                    oplot='oplot')
    
    
    plot( numpy.append(a.rxy[0,:], a.rxy[0,0]), numpy.append(a.zxy[0,:], a.zxy[0,0]), 'r')
    
          
    for i in range (1, nrad) :
        plot( numpy.append(a.rxy[i,:], a.rxy[i,0]), numpy.append(a.zxy[i,:], a.zxy[i,0]), 'r')


    for i in range (0, npol-1) :
        plot( a.rxy[:,i], a.zxy[:,i], 'r')
    
    draw()

    # Get other useful variables
    psixy = numpy.zeros((nrad, npol))
    for i in range (0, npol) :
        psixy[:,i] = old_div((fvals - faxis),fnorm) # to get normalised psi
    
    # Calculate magnetic field components
    dpsidR = numpy.zeros((nrad, npol))
    dpsidZ = numpy.zeros((nrad, npol))
    
    interp_data.method = 2
    
    

    for i in range (nrad) :
        for j in range (npol) :
            out = local_gradient( interp_data, a.rixy[i,j], a.zixy[i,j], status=0, dfdr=0., dfdz=0.)
            status=out.status
            dfdr=out.dfdr[0][0]
            dfdz=out.dfdz[0][0]
    # dfd* are derivatives wrt the indices. Need to multiply by dr/di etc
            dpsidR[i,j] = old_div(dfdr,numpy.interp(a.rixy[i,j], numpy.arange(R.size).astype(float),numpy.gradient(R))) 
            dpsidZ[i,j] = old_div(dfdz,numpy.interp(a.zixy[i,j], numpy.arange(Z.size).astype(float),numpy.gradient(Z))) 
                 

    # Set topology to connect in the core
    yup_xsplit = [nrad]
    ydown_xsplit = [nrad]
    yup_xin = [0]
    yup_xout = [-1]
    ydown_xin = [0]
    ydown_xout = [-1]

    #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    # Create result structure

    result = Bunch(error=0, # Signal success
              psi_inner=settings.psi_inner, psi_outer=settings.psi_outer, # Unchanged psi range
              nrad=nrad, npol=npol, #Number of points in radial and poloidal direction
              Rixy=a.rixy, Zixy=a.zixy, # Indices into R and Z of each point
              Rxy=a.rxy, Zxy=a.zxy, # Location of each grid point
              psixy=psixy, # Normalised psi for each point
              dpsidR=dpsidR, dpsidZ=dpsidZ, # Psi derivatives (for Bpol)
              faxis=faxis, fnorm=fnorm, # Psi normalisation factors
              settings=settings, # Settings used to create grid
              critical=critical, # Critical points
              yup_xsplit=yup_xsplit, # X index where domain splits (number of points in xin)
              ydown_xsplit=ydown_xsplit, 
              yup_xin=yup_xin, yup_xout=yup_xout, # Domain index to connect to
              ydown_xin=ydown_xin, ydown_xout=ydown_xout)

    return result