Return to soshdata.py CVS log

Up to [Development] / JSOC / proj / globalhs / sosh

new versions

import numpy as np
import os
import re
import drms
from astropy.io import fits
from urllib.request import urlretrieve
import soundfile as sf
from scipy import fftpack

import warnings 
warnings.simplefilter(action='ignore', category=FutureWarning)

#datadir = '/home/tplarson/solar/data'
datadir = '../data'

def getmodeparms(instrument=None, day=None, lmin=0, lmax=300, makemsplit=True):

  if not os.path.exists(datadir):
    os.makedirs(datadir)

  if (instrument == None):
    inst=input("Enter name of instrument: ")
  else:
    inst=instrument

  if (day == None):
    daystr=input("Enter day number: ")
  else:
    daystr=str(day)

  mfile=getmfile(inst.lower(), daystr)
  if (mfile == None):
    return None

  if (makemsplit):
    msfile=mfile+'.msplit'

    if os.path.exists(msfile):
      print("Msplit file already exists.")
    else:
      msfile=mkmsplit(mfile,lmin=lmin,lmax=lmax)

    return mfile,msfile
  else:
    return mfile


def getmfile(inst, daystr):

  if (inst == 'hmi'):
    qblank = 'hmi.V_sht_modes[{:d}d][0][300][138240]'
    avgurl = 'http://sun.stanford.edu/~tplarson/audio/HMI/hmi.average.modes'
  elif (inst == 'mdi'):
    qblank = 'mdi.vw_V_sht_modes[{:d}d][0][300][103680]'
    avgurl = 'http://sun.stanford.edu/~tplarson/audio/MDI/mdi.average.modes'
  else:
    print("Invalid instrument, try again.")
    return None

  if (daystr == 'average'):
    mblank = '%s.average.modes'
    mfile=os.path.join(datadir,mblank % inst)
    if os.path.exists(mfile):
      print("Mode fits file already exists.")
    else:
      try:
        urlretrieve(avgurl, mfile)
      except:
        print("Failure of urlretrieve() for %s." % avgurl)
        return None
  else:
    mblank = '%s.%sd.modes'
    mfile=os.path.join(datadir,mblank % (inst, daystr))
    if os.path.exists(mfile):
      print("Mode fits file already exists.")
    else:
      try:
        day=int(daystr)
      except:
        print("Invalid number, try again.")
        return None
      qstr=qblank.format(day)
      c = drms.Client()
      m = c.query(qstr, seg='m6')
      if (len(m) == 0):
        print("No data found for that day number, try again.")
        return None
      url = 'http://jsoc.stanford.edu' + m.m6[0]
      try:
        urlretrieve(url, mfile)
      except:
        print("Failure of urlretrieve() for %s." % url)
        return None

  return mfile


def mkmsplit(mfile, lmin=0, lmax=300):

  modeparms=np.loadtxt(mfile)
  lmod=modeparms[:,0]
  nmod=modeparms[:,1]

  outfmt='{:d} {:d} {:d} {:f} {:f}'
  outfile = mfile + '.msplit'
  file=open(outfile,'w')

  l=lmin
  while (l <= lmax):
    pols=apols(l,6)
    ind = (lmod == l)
    nlist = nmod[ind]
    for n in nlist:
      ind2 = (lmod == l) & (nmod == n)
      ai=np.append([0.0],modeparms[ind2,12:18]/1000)
      ei=np.append([0.0],modeparms[ind2,18:24]/1000)
      fx=np.matmul(pols,ai)
      freq=modeparms[ind2,2]+fx
      var=np.matmul((pols*pols),(ei*ei))
      sig=modeparms[ind2,7]+np.sqrt(var)
      for m in range(-l,l+1):
        outstr=outfmt.format(int(l),int(n),int(m),freq[m+l],sig[m+l])
        file.write('%s\n' % (outstr))
    l+=1
  file.close()
  return outfile


def getwavfiles(instrument=None, day=None, l=None, m=None, delfits=False):

  if not os.path.exists(datadir):
    os.makedirs(datadir)

  if (instrument == None):
    inst=input("Enter name of instrument: ")
  else:
    inst=instrument

  if (inst.lower() == 'hmi'):
    hind=1
  elif (inst.lower() == 'mdi'):
    hind=0
  else:
    print("Invalid instrument, try again.")
    return None

  if (day == None):
    daystr=input("Enter day number: ")
    try:
      day=int(daystr)
    except:
      print("Invalid number, try again.")
      return None
  if (l == None):
    lstr=input("Enter spherical harmonic degree: ")
    try:
      l=int(lstr)
    except:
      print("Invalid number, try again.")
      return None
  if (m == None):
    mstr=input("Enter azimuthal order: ")
    try:
      m=int(mstr)
      m=np.abs(m)
      if (m > l):
        print("Abs(m) must be less than l, try again.")
        return None
    except:
      print("Invalid number, try again.")
      return None

  wblank = '%s.%id_72d.l=%i_m=%i_data%s.wav'
  wfile=os.path.join(datadir,wblank % (inst.lower(),day,l,m,'r'))
  wfile2=os.path.join(datadir,wblank % (inst.lower(),day,l,m,'i'))
  if os.path.exists(wfile) and (m==0 or os.path.exists(wfile2)):
    if (m == 0):
      print("Wav file already exists.")
      return wfile
    else:
      print("Wav files already exists.")
      return wfile,wfile2

  ffile=getfitsfile(inst.lower(),day,l)
  if (ffile == None):
    return None

  wf=convertfitstowav(ffile, hind, m, wfile)
  if (wf == None):
    return None

  if delfits:
    os.remove(ffile)
  return wf


def getfitsfile(inst, day, l):

  if (inst == 'hmi'):
    qblank = 'hmi.V_sht_gf_72d[{:d}d][{:d}][{:d}][138240]'
  elif (inst == 'mdi'):
    qblank = 'mdi.vw_V_sht_gf_72d[{:d}d][{:d}][{:d}][103680]'
  else:
    print("Invalid instrument, try again.")
    return None

  fblank='%s.%id.%i.fits'
  ffile=os.path.join(datadir,fblank % (inst,day,l))
  if os.path.exists(ffile):
    print("Fits file already exists.")
  else:
    qstr=qblank.format(day,l,l)
    c = drms.Client()
    f = c.query(qstr,seg='data')
    if (len(f) == 0):
      print("No data found for that day number and degree, try again.")
      return None
    url = 'http://jsoc.stanford.edu' + f.data[0]
    try:
      urlretrieve(url, ffile)
    except:
      print("Failure of urlretrieve() for %s." % url)
      return None

  return ffile


def convertfitstowav(ffile, hind, m, wfile):

  hdul = fits.open(ffile)
  h=hdul[hind]
  d=h.data
  x=d[m,0::2]/np.abs(d[m,:]).max()
  sf.write(wfile, x, 8000, subtype='FLOAT')

  if (m > 0):
    wfile2=wfile.replace('datar','datai')
    x=d[m,1::2]/np.abs(d[m,:]).max()
    sf.write(wfile2, x, 8000, subtype='FLOAT')

  hdul.close()
  if (m==0):
    return wfile
  else:
    return wfile,wfile2


def apols(lin, amaxin):

# adapted from IDL procedure apols.pro written by Jesper Schou

  l=np.long(lin)
  amax=np.minimum(amaxin,2*l)

  pols=np.zeros((2*l+1,amaxin+1))
# It is ESSENTIAL that x is set up such that x(-m)=-x(m) to full machine
# accuracy or that the re-orthogonalization is done with respect to all
# previous polynomials (second option below).
# x=(dindgen(2*l+1)-l)/l
  x=np.linspace(-1,1,2*l+1)

  pols[:,0]=1/np.sqrt(2*l+1.0)
  if (amax >= 1):
    pols[:,1]=x/np.sqrt((l+1.0)*(2*l+1.0)/(3.0*l))
# for n=2l,amax do begin
# Set up polynomials using exact recurrence relation.
  for n in range(2,amax+1):
    a0=2.0*l*(2*n-1.0)/(n*(2*l-n+1))
    b0=-(n-1.0)/n*(2*l+n)/(2*l-n+1)
    a=a0*np.sqrt((2*n+1.0)/(2*n-1.0)*(2*l-n+1.0)/(2*l+n+1.0))
    b=b0*np.sqrt((2*n+1.0)/(2*n-3.0)*(2*l-n+1.0)*(2*l-n+2.0)/(2*l+n+1.0)/(2*l+n))
    help=a*x*pols[:,n-1]+b*pols[:,n-2]
# Turns out that roundoff kills the algorithm, so we have to re-orthogonalize.
# Two choices here. First one is twice as fast and more accurate, but
# depends on the rounding being done in the standard IEEE way.
#  for j=n-2,0,-2 do begin 
    for j in range(n-2,-1,-2):    
# This choice is robust to the roundoff, but twice as slow and generally
# somewhat less accurate
# for j=n-1,0,-1 do begin
      help=help-pols[:,j]*np.sum(help*pols[:,j])
#  end
# Reset norm to 1.
    pols[:,n]=help/np.sqrt(np.sum(np.square(help)))

# Reset polynomials to have P_l(l)=l by setting overall norm.
# Note that this results in more accurate overall normalization, but
# that P_l(l) may be very far from l. This is the key to the algorithm.
  c=l**2*(2*l+1.0)
# for n=0l,amax do begin
  for n in range(0,amax+1):
    c=c*(2*l+n+1.0)/(2*l-n+1.0)
    pols[:,n]=pols[:,n]*np.sqrt(c/(2*n+1))

  return pols


def splitwavfile(wavfile=None, splitfactor=2, splitboth=True):

  wblank = '%s.%id_%id.l=%i_m=%i_data%s.wav'
  outlist=[]

  if (wavfile == None):
    inst=input("Enter name of instrument: ")
    inst=inst.lower()
    if (inst != 'mdi' and inst != 'hmi'):
      print("Invalid instrument, try again.")
      return None
    daystr=input("Enter day number: ")
    try:
      firstday=int(daystr)
    except:
      print("Invalid number, try again.")
      return None
    ndstr=input("Enter number of days: ")
    try:
      ndaysin=int(ndstr)
    except:
      print("Invalid number, try again.")
      return None
    lstr=input("Enter spherical harmonic degree: ")
    try:
      l=int(lstr)
      if (l < 0):
        print("Degree l must be nonnegative, try again.")
        return None
    except:
      print("Invalid number, try again.")
      return None
    mstr=input("Enter azimuthal order: ")
    try:
      m=int(mstr)
      m=np.abs(m)
      if (m > l):
        print("Abs(m) must be less than l, try again.")
        return None
    except:
      print("Invalid number, try again.")
      return None
    ri='r'
    wavfile=os.path.join(datadir,wblank % (inst,firstday,ndaysin,l,m,ri))
  else:
    regmatch = re.search(r"(\S+)\.(\d+)d_(\d+)d\.l=(\d+)_m=(\d+)_data(\S+).wav", wavfile)
    if (regmatch == None):
      print("Invalid file name, try again.")
      return None
    else:
      inst=regmatch.group(1)
      firstday=int(regmatch.group(2))
      ndaysin=int(regmatch.group(3))
      l=int(regmatch.group(4))
      m=int(regmatch.group(5))
      ri=regmatch.group(6)

  test= ndaysin % splitfactor
  if (test != 0):
    print("Timeseries not divisible by that factor.")
    return None

#  wfile = os.path.join(datadir,wavfile)
  try:
    x, sr = sf.read(wavfile)
  except:
    print("Invalid file, try again.")
    return None

  if (inst == 'mdi'):
    nperday = 1440
  elif (inst == 'hmi'):
    nperday = 1920

  ndaysout = ndaysin / splitfactor
  ndt=int(ndaysout*nperday)
  day = firstday
  index=0
  while (day < firstday+ndaysin):
    xout=x[index:index+ndt]
    wfileout = os.path.join(datadir,wblank % (inst,day,ndaysout,l,m,ri))
    sf.write(wfileout, xout, sr, subtype='FLOAT')
    day += ndaysout
    index += ndt
    outlist.append(wfileout)

  if (splitboth and m > 0):
    if (ri == 'i'):
      ir = 'r'
    elif (ri == 'r'):
      ir = 'i'
    wfile=wblank % (inst,firstday,ndaysin,l,m,ir)
    wavfile = os.path.join(datadir,wfile)
    try:
      x, sr = sf.read(wavfile)
    except:
      print("Invalid file, try again.")
      return None
    day=firstday
    index=0
    while (day < firstday+ndaysin):
      xout=x[index:index+ndt]
      wfileout = os.path.join(datadir,wblank % (inst,day,ndaysout,l,m,ir))
      sf.write(wfileout, xout, sr, subtype='FLOAT')
      day += ndaysout
      index += ndt
      outlist.append(wfileout)

  return outlist


def scanspectrum(wavfile=None,downshift=4,firstbin=None,lastbin=None,nbins=200,ndt=10000,binstep=None,ramp=50):

  if (binstep == None):
    binstep = int(nbins/2)

  wblank = '%s.%id_%id.l=%i_m=%i_data%s.wav'
  if (wavfile == None):
    inst=input("Enter name of instrument: ")
    inst=inst.lower()
    if (inst != 'mdi' and inst != 'hmi'):
      print("Invalid instrument, try again.")
      return None
    daystr=input("Enter day number: ")
    try:
      firstday=int(daystr)
    except:
      print("Invalid number, try again.")
      return None
    ndstr=input("Enter number of days: ")
    try:
      ndaysin=int(ndstr)
    except:
      print("Invalid number, try again.")
      return None
    lstr=input("Enter spherical harmonic degree: ")
    try:
      l=int(lstr)
      if (l < 0):
        print("Degree l must be nonnegative, try again.")
        return None
    except:
      print("Invalid number, try again.")
      return None
    mstr=input("Enter azimuthal order: ")
    try:
      m=int(mstr)
      m=np.abs(m)
      if (m > l):
        print("Abs(m) must be less than l, try again.")
        return None
    except:
      print("Invalid number, try again.")
      return None
    ri='r'
    wavfile=os.path.join(datadir,wblank % (inst,firstday,ndaysin,l,m,ri))
    print("Wav file is %s" % wavfile)
  else:
    regmatch = re.search(r"(\S+)\.(\d+)d_(\d+)d\.l=(\d+)_m=(\d+)_data(\S+).wav", wavfile)
    if (regmatch == None):
      print("Invalid file name, try again.")
      return None
    else:
      inst=regmatch.group(1)
      firstday=int(regmatch.group(2))
      ndaysin=int(regmatch.group(3))
      l=int(regmatch.group(4))
      m=int(regmatch.group(5))
      ri=regmatch.group(6)

  try:
    x, sr = sf.read(wavfile)
  except:
    print("Invalid file, try again.")
    return None

  z=np.zeros((len(x)),dtype=np.complex_)
  if (ri == 'r'):
    z.real=x
  else:
    z.imag=x

  if (m > 0):
    if (ri == 'i'):
      ir = 'r'
    elif (ri == 'r'):
      ir = 'i'
    wfile=wblank % (inst,firstday,ndaysin,l,m,ir)
    wavfile = os.path.join(datadir,wfile)
    try:
      x, sr = sf.read(wavfile)
    except:
      print("Invalid file, try again.")
      return None
  else:
    x=0.0*x

  if (ri == 'r'):
    z.imag=x
  else:
    z.real=x

  window=np.ones(ndt)
  if (ramp != 0):
    rampsamples=int(sr*ramp/1000)
    if (rampsamples > ndt/2):
      print("Ramp longer than ndt/2, try again.")
      return None
    window[0:rampsamples]=np.linspace(0.0,1.0,rampsamples,endpoint=False)
    window[ndt-1:ndt-rampsamples-1:-1]=np.linspace(0.0,1.0,rampsamples,endpoint=False)

  nx=len(x)
  nx2=int(nx/2)
  if (firstbin == None):
    firstbin = int(nx2/6)
  if (lastbin == None):
    lastbin = int(0.9*nx2)
  ftz = fftpack.fft(z)
  freq = fftpack.fftfreq(nx) #*sr
  if (m >= 0):
    mag = np.abs(ftz[0:nx2])
    phase = np.arctan2(ftz[0:nx2].imag,ftz[0:nx2].real)
    f = freq[0:nx2]/downshift
  else:
    mag = np.abs(ftz[nx:nx2:-1])
    phase = np.arctan2(-1*ftz[nx:nx2:-1].imag,ftz[nx:nx2:-1].real)
    f = -1*freq[nx:nx2:-1]/downshift

  nstep=int((lastbin-firstbin)/binstep)
  tlen=nstep*ndt
  t=np.linspace(0,ndt-1,ndt)
  xout=np.zeros((tlen))
  bindex=firstbin
  tindex=0
  t0in=0
  while (bindex + nbins <= lastbin):
    for i in range(nbins):
      xout[tindex:tindex+ndt]+=mag[bindex+i]*np.sin(2*np.pi*f[bindex+i]*(t+tindex)+phase[bindex+i])
    xout[tindex:tindex+ndt]*=window
    bindex+=binstep
    tindex+=ndt

  mx=np.abs(xout).max()
  xout /= mx
  sf.write('output.wav', xout, sr, subtype='FLOAT')

  return xout


#def plotspectrum(wavfile):
#  x, sr = sf.read(wavfile)
#  ftx = fftpack.fft(x)
#  freqs = fftpack.fftfreq(len(x)) * sr