Workaround for astropy quantities with mpi comm.gather

examples/mhs_atmosphere/hmi_atmosphere.py

@@ -31,6 +31,10 @@
 import pysac.mhs_atmosphere as atm
 import astropy.units as u
 from pysac.mhs_atmosphere.parameters.model_pars import hmi_model as model_pars
+import astropy
+l_astropy_0=True
+if astropy.__version__[0]=='1':
+    l_astropy_0=False
 #==============================================================================
 #check whether mpi is required and the number of procs = size
 #==============================================================================
@@ -41,9 +45,11 @@
     rank = comm.Get_rank()
     size = comm.Get_size()
 
+    collective=True
     l_mpi = True
     l_mpi = l_mpi and (size != 1)
 except ImportError:
+    collective=False
     l_mpi = False
     rank = 0
     size = 1
@@ -57,11 +63,29 @@
 scales, physical_constants = \
     atm.units_const.get_parameters()
 
+if not l_astropy_0:
+    dataset = 'hmi_m_45s_2014_07_06_00_00_45_tai_magnetogram_fits'
+else:
+    dataset = 'hmi_m_45s_2014_07_06_00_00_45_tai_magnetogram.fits'
+l_newdata = True # change this to False if a local copy already exists at ~/sunpy/data/
 #obtain code coordinates and model parameters in astropy units
-coords = atm.model_pars.get_hmi_map(model_pars, option_pars, indx = [1787,1798,1818,1822], 
-                dataset = 'hmi_m_45s_2014_07_06_00_00_45_tai_magnetogram_fits',
-                l_newdata = False
+model_pars['Nxyz'] = [64,64,64] # 3D grid
+model_pars['xyz']  = [-0.63*u.Mm,0.63*u.Mm,-0.63*u.Mm,0.63*u.Mm,0.0*u.Mm,3.8*u.Mm]
+indx = [1785,1800,1814,1829]
+interpfactor=5
+frac = [0.25,0.25]
+coords = atm.model_pars.get_hmi_coords(
+                model_pars['Nxyz'],
+                model_pars['xyz'],
+                indx = indx,
+                dataset = dataset,
+                l_newdata = False,
+                interpfactor=interpfactor,
+                frac=frac,
+                rank=rank,
+                lmpi=l_mpi
                )
+model_pars['xyz'][0:4] = coords['xmin'],coords['xmax'],coords['ymin'],coords['ymax']
 
 #interpolate the hs 1D profiles from empirical data source[s]
 empirical_data = atm.hs_atmosphere.read_VAL3c_MTW(mu=physical_constants['mu'])
@@ -92,30 +116,41 @@
 # load flux tube footpoint parameters
 #==============================================================================
 # axial location and value of Bz at each footpoint
-model_pars['B_corona']/=model_pars['nftubes']
-xi, yi, Si = atm.flux_tubes.get_flux_tubes(
-                                model_pars,
-                                coords,
-                                option_pars
-                               )
+Stmp,xtmp,ytmp,FWHM,sdummy,xdummy,ydummy,cmax,cmin \
+               =atm.parameters.model_pars.get_hmi_map(
+                indx,
+                dataset = dataset,
+                l_newdata = False,
+                interpfactor=interpfactor,
+                frac=frac,
+                rank=rank,
+                lmpi=l_mpi
+               )
+xi, yi, Si = xtmp.to(u.Mm).reshape(xtmp.size),\
+             ytmp.to(u.Mm).reshape(ytmp.size), Stmp.reshape(Stmp.size)
+model_pars['radial_scale'] = 0.5*FWHM.to(u.Mm)/np.sqrt(2*np.log(2))
 #==============================================================================
 # split domain into processes if mpi
 #==============================================================================
 ax, ay, az = np.mgrid[coords['xmin']:coords['xmax']:1j*model_pars['Nxyz'][0],
                       coords['ymin']:coords['ymax']:1j*model_pars['Nxyz'][1],
                       coords['zmin']:coords['zmax']:1j*model_pars['Nxyz'][2]]
 
+axindex=np.arange(model_pars['Nxyz'][0])
 # split the grid between processes for mpi
 if l_mpi:
     x_chunks = np.array_split(ax, size, axis=0)
     y_chunks = np.array_split(ay, size, axis=0)
     z_chunks = np.array_split(az, size, axis=0)
+    i_chunks = np.array_split(axindex, size, axis=0)
 
     x = comm.scatter(x_chunks, root=0)
     y = comm.scatter(y_chunks, root=0)
     z = comm.scatter(z_chunks, root=0)
+    xindex=i_chunks[rank]
 else:
     x, y, z = ax, ay, az
+    xindex=axindex
 
 x = u.Quantity(x, unit=coords['xmin'].unit)
 y = u.Quantity(y, unit=coords['ymin'].unit)
@@ -137,8 +172,8 @@
 #==============================================================================
 #calculate the magnetic field and pressure/density balancing expressions
 #==============================================================================
-for i in range(0,model_pars['nftubes']):
-    for j in range(i,model_pars['nftubes']):
+for i in range(0,Si.size):
+    for j in range(i,Si.size):
         if rank == 0:
             print'calculating ij-pair:',i,j
         if i == j:
@@ -175,6 +210,19 @@
 
 # clear some memory
 del pressure_mi, rho_mi, Bxi, Byi ,Bzi, B2x, B2y
+
+import matplotlib.pyplot as plt
+plt.figure()
+plt.pcolormesh(x[:,:,0].T.value,y[:,:,0].T.value,Bz[:,:,0].T.value,vmin=cmin,vmax=cmax)
+plt.xlabel('lon [Mm]')
+plt.ylabel('lat [Mm]')
+plt.axis([x.min().value,x.max().value,y.min().value,y.max().value])
+cbar = plt.colorbar()
+cbar.ax.set_ylabel(r'$B_z$ [T]')
+cbar.solids.set_edgecolor("face")
+plt.savefig('model_derived_hmi.png')
+plt.close()
+print 'Bz[:,:,0].min(), Bz[:,:,0].max()=', Bz[:,:,0].min(), Bz[:,:,0].max()
 #==============================================================================
 # Construct 3D hs arrays and then add the mhs adjustments to obtain atmosphere
 #==============================================================================
@@ -192,11 +240,27 @@
                                   )
 magp = (Bx**2 + By**2 + Bz**2)/(2.*physical_constants['mu0'])
 if rank ==0:
-    print'max B corona = ',magp[:,:,-1].max().decompose()
-    print'min B corona = ',magp[:,:,-1].min().decompose()
+    print'max pB corona = ',magp[:,:,-1].max().decompose()
+    print'min pB corona = ',magp[:,:,-1].min().decompose()
 energy = atm.mhs_3D.get_internal_energy(pressure,
                                                   magp,
                                                   physical_constants)
+Rgas = u.Quantity(np.zeros(x.shape), unit=Rgas_z.unit)
+Rgas[:] = Rgas_z
+temperature = pressure/rho/Rgas
+if not option_pars['l_hdonly']:
+    inan = np.where(magp <=1e-7*pressure.min())
+    magpbeta = magp
+    magpbeta[inan] = 1e-7*pressure.min()  # low pressure floor to avoid NaN
+    pbeta  = pressure/magpbeta
+else:
+    pbeta  = magp+1.0    #dummy to avoid NaN
+alfven = np.sqrt(2.*magp/rho)
+#if rank == 0:
+#    print'Alfven speed Z.min to Z.max =',\
+#    alfven[model_pars['Nxyz'][0]/2,model_pars['Nxyz'][1]/2, 0].decompose(),\
+#    alfven[model_pars['Nxyz'][0]/2,model_pars['Nxyz'][1]/2,-1].decompose()
+cspeed = np.sqrt(physical_constants['gamma']*pressure/rho)
 #============================================================================
 # Save data for SAC and plotting
 #============================================================================
@@ -212,66 +276,67 @@
 
 # save the variables for the initialisation of a SAC simulation
 atm.mhs_snapshot.save_SACvariables(
-              filename,
-              rho,
-              Bx,
-              By,
-              Bz,
-              energy,
-              option_pars,
-              physical_constants,
-              coords,
-              model_pars['Nxyz']
-             )
+          filename,
+          rho,
+          Bx,
+          By,
+          Bz,
+          energy,
+          option_pars,
+          physical_constants,
+          coords,
+          model_pars['Nxyz'],
+          xindex,
+          rank=rank,
+          collective=collective
+         )
 # save the balancing forces as the background source terms for SAC simulation
 atm.mhs_snapshot.save_SACsources(
-              sourcefile,
-              Fx,
-              Fy,
-              option_pars,
-              physical_constants,
-              coords,
-              model_pars['Nxyz']
-             )
+          sourcefile,
+          Fx,
+          Fy,
+          option_pars,
+          physical_constants,
+          coords,
+          model_pars['Nxyz'],
+          xindex,
+          rank=rank,
+          collective=collective
+         )
 # save auxilliary variable and 1D profiles for plotting and analysis
-Rgas = u.Quantity(np.zeros(x.shape), unit=Rgas_z.unit)
-Rgas[:] = Rgas_z
-temperature = pressure/rho/Rgas
-if not option_pars['l_hdonly']:
-    inan = np.where(magp <=1e-7*pressure.min())
-    magpbeta = magp
-    magpbeta[inan] = 1e-7*pressure.min()  # low pressure floor to avoid NaN
-    pbeta  = pressure/magpbeta
-else:
-    pbeta  = magp+1.0    #dummy to avoid NaN
-alfven = np.sqrt(2.*magp/rho)
-#if rank == 0:
-#    print'Alfven speed Z.min to Z.max =',\
-#    alfven[model_pars['Nxyz'][0]/2,model_pars['Nxyz'][1]/2, 0].decompose(),\
-#    alfven[model_pars['Nxyz'][0]/2,model_pars['Nxyz'][1]/2,-1].decompose()
-cspeed = np.sqrt(physical_constants['gamma']*pressure/rho)
 atm.mhs_snapshot.save_auxilliary3D(
-              aux3D,
-              pressure_m,
-              rho_m,
-              temperature,
-              pbeta,
-              alfven,
-              cspeed,
-              Btensx,
-              Btensy,
-              option_pars,
-              physical_constants,
-              coords,
-              model_pars['Nxyz']
-             )
-atm.mhs_snapshot.save_auxilliary1D(
-              aux1D,
-              pressure_Z,
-              rho_Z,
-              Rgas_Z,
-              option_pars,
-              physical_constants,
-              coords,
-              model_pars['Nxyz']
-             )
+          aux3D,
+          pressure_m,
+          rho_m,
+          temperature,
+          pbeta,
+          alfven,
+          cspeed,
+          Btensx,
+          Btensy,
+          option_pars,
+          physical_constants,
+          coords,
+          model_pars['Nxyz'],
+          xindex,
+          rank=rank,
+          collective=collective
+         )
+if rank==0:
+    atm.mhs_snapshot.save_auxilliary1D(
+          aux1D,
+          pressure_Z,
+          rho_Z,
+          Rgas_Z,
+          option_pars,
+          physical_constants,
+          coords,
+          model_pars['Nxyz'],
+          rank=rank,
+          collective=False
+         )
+if rho.min()<0 or pressure.min()<0:
+    print"FAIL rank {}: negative rho.min() {} and/or pressure.min() {}.".format(
+                            rank,rho.min(),pressure.min())
+FWHM = 2*np.sqrt(np.log(2))*model_pars['radial_scale']
+print'FWHM(0) =',FWHM