name = './test_hdf5'

# skip over if output already exists, or delete and rerun

# write input radiative transfer parameters

mu = np.cos(ang*np.pi/180.)

freq = freq_ghz*1.e9

uout = 1./RAYTRACESIZE

x=gr.grtrans()

npol=1

x.write_grtrans_inputs(name + '.in', oname=name+'.out',

fscalefac=RESCALE, sigcut=sigma_cut,

fname='KORALH5',phi0=0.,

nfreq=1,fmin=freq,fmax=freq,

ename='SYNCHTHAV',

nvals=npol,

gmin=-3, # confusingly, this is rhigh. #-2-->vladimir model, -3--> michael model

spin=SPIN,standard=1,

uout=uout,

mbh=MBH,

mdotmin=1.,mdotmax=1.,nmdot=1,

#mdotmin=1.57e15,mdotmax=1.57e15,nmdot=1,

nmu=1,mumin=mu,mumax=mu,

gridvals=[-size,size,-size,size],

nn=[NPIX_IM,NPIX_IM,NGEO],

hhfile=hfile, hdfile=dfile,

hindf=1,hnt=1,

muval=1.)

run=True

if os.path.exists(name+'.out'):

run = False

if RERUN:

run=True

os.remove(name+'.out')

# run grtrans

if run:

x.run_grtrans()

# Read grtrans output

try:

x.read_grtrans_output()

except:# IOError:

return None

# pixel sizes

da = x.ab[x.nx,0]-x.ab[0,0]

db = x.ab[1,1]-x.ab[0,1]

if (da!=db): raise Exception("pixel da!=db")

psize = da*(lbh/DTOBH)

#image values

if npol==4:

ivals = x.ivals[:,0,0]*fac*da*db*LumtoJy

qvals = x.ivals[:,1,0]*fac*da*db*LumtoJy

uvals = x.ivals[:,2,0]*fac*da*db*LumtoJy

vvals = x.ivals[:,3,0]*fac*da*db*LumtoJy

# mask nan failure points with zeros

ivals = np.array(ivals)

qvals = np.array(qvals)

uvals = np.array(uvals)

vvals = np.array(vvals)

imask = np.isnan(ivals)

qumask = ~(~imask * ~np.isnan(qvals) * ~np.isnan(uvals))

vmask = ~(~imask * ~np.isnan(vvals))

ivals[imask] = 0.

qvals[qumask] = 0.

uvals[qumask] = 0.

vvals[vmask] = 0.

ivals = (np.flipud(np.transpose(ivals.reshape((NPIX_IM,NPIX_IM))))).flatten()

qvals = -(np.flipud(np.transpose(qvals.reshape((NPIX_IM,NPIX_IM))))).flatten()

uvals = -(np.flipud(np.transpose(uvals.reshape((NPIX_IM,NPIX_IM))))).flatten()

vvals = (np.flipud(np.transpose(vvals.reshape((NPIX_IM,NPIX_IM))))).flatten()

else:

ivals = x.ivals[:,0,0]*fac*da*db*LumtoJy

ivals = np.array(ivals)

imask = np.isnan(ivals)

ivals[imask] = 0.

ivals = (np.flipud(np.transpose(ivals.reshape((NPIX_IM,NPIX_IM))))).flatten()

qvals = uvals = vvals = 0*ivals

print('total flux', np.sum(ivals))

save_im_fits((ivals,qvals,uvals,vvals, psize))

fname = 'hdf5test.fits'

# unpack the image data

(imvec, qvec, uvec, vvec, psize) = imdata

# Create header and fill in some values

header = fits.Header()

header['OBJECT'] = source

header['CTYPE1'] = 'RA---SIN'

header['CTYPE2'] = 'DEC--SIN'

header['CDELT1'] = -psize/DEGREE

header['CDELT2'] = psize/DEGREE

header['OBSRA'] = ra * 180/12.

header['OBSDEC'] = dec

header['FREQ'] = freq_ghz * 1.e9

#TODO these are the default values for centered images

#TODO support for arbitrary CRPIX?

header['CRPIX1'] = NPIX_IM/2. + .5

header['CRPIX2'] = NPIX_IM/2. + .5

mjd += (time/24.)

header['MJD'] = float(mjd)

header['TELESCOP'] = 'VLBI'

header['BUNIT'] = 'JY/PIXEL'

header['STOKES'] = 'I'

# Create the fits image

image = np.reshape(imvec,(NPIX_IM,NPIX_IM))[::-1,:] #flip y axis!

hdu = fits.PrimaryHDU(image, header=header)

hdulist = [hdu]

if len(qvec):

qimage = np.reshape(qvec,(NPIX_IM,NPIX_IM))[::-1,:]

uimage = np.reshape(uvec,(NPIX_IM,NPIX_IM))[::-1,:]

header['STOKES'] = 'Q'

hduq = fits.ImageHDU(qimage, name='Q', header=header)

header['STOKES'] = 'U'

hduu = fits.ImageHDU(uimage, name='U', header=header)

hdulist = [hdu, hduq, hduu]

if len(vvec):

vimage = np.reshape(vvec,(NPIX_IM,NPIX_IM))[::-1,:]

header['STOKES'] = 'V'

hduv = fits.ImageHDU(vimage, name='V', header=header)

hdulist.append(hduv)

hdulist = fits.HDUList(hdulist)

# Save fits

try:

hdulist.writeto(fname, overwrite=True)

except:

hdulist.writeto(fname, clobber=True)