power_spectra.py

#!/usr/bin/env python3

import numpy as n
import matplotlib.pyplot as plt
import pyfftw
import sys
import chirp_config as cc
import h5py
import digital_rf as drf
import time
import scipy.signal as ss
import chirp_det as cd
import traceback
import os


def fft(z, l=None):
    """
    wrap fft, so that it can be configured
    """
    if l == None:
        l = len(z)
    return (pyfftw.interfaces.numpy_fft.fft(z, l, planner_effort='FFTW_ESTIMATE'))


def ifft(z, l=None):
    """
    wrap fft, so that it can be configured
    """
    if l == None:
        l = len(z)
    return (pyfftw.interfaces.numpy_fft.ifft(z, l, planner_effort='FFTW_ESTIMATE'))
#    return(sf.ifft(z,l))


if __name__ == "__main__":
    if len(sys.argv) == 2:
        conf = cc.chirp_config(sys.argv[1])
    else:
        print("give configuration file as command line argument")

    sr = 25e6
    fftlen = 1024 * 2 * 2 * 2 * 2 * 2 * 2 * 2
    print("fftlen %d freqeuency resolution %f (kHz)" %
          (fftlen, sr / fftlen / 1e3))
    wfun = n.array(ss.hann(fftlen), dtype=n.float32)
    S = n.zeros(fftlen, dtype=n.float32)

    dt = int(sr) * 60
    nfft = int(n.floor(dt / fftlen))

    d = drf.DigitalRFReader(conf.data_dir)
    b = d.get_bounds(conf.channel)
    i0 = b[1] - dt - int(sr)

    while True:
        d = drf.DigitalRFReader(conf.data_dir)
        b = d.get_bounds(conf.channel)
        n_windows = int(n.floor((b[1] - i0) / dt))
        for wi in range(n_windows):
            t0 = time.time()
            S[:] = 0.0
            for fi in range(nfft):
                F = fft(wfun * d.read_vector_c81d(wi * dt +
                        fi * fftlen + i0, fftlen, conf.channel))
                S += (F * n.conj(F)).real
            t1 = time.time()
            print(t1 - t0)
            print("done")

            try:
                t0 = i0 / sr
                dname = "%s/%s" % (conf.output_dir, cd.unix2dirname(t0))
                if not os.path.exists(dname):
                    os.mkdir(dname)
                ofname = "%s/spec-%s-%1.2f.h5" % (dname, conf.station_name, t0)
                print("Writing to %s" % ofname)
                ho = h5py.File(ofname, "w")
                ho["spec"] = S
                ho["nfft"] = nfft
                ho["f0"] = conf.center_freq
                ho["sr"] = conf.sample_rate
                ho["dt"] = dt
                ho["station_name"] = conf.station_name
                ho["ch"] = conf.channel
                ho.close()
            except:
                traceback.print_exc(file=sys.stdout)
                print("error writing file")

        i0 = i0 + n_windows * dt

        b = d.get_bounds(conf.channel)
        if (b[1] - i0) < dt:
            print("sleeping for more data %d" % (b[1] - i0))
            time.sleep(1)