170 lines
6.6 KiB
Python
Executable File
170 lines
6.6 KiB
Python
Executable File
#!/usr/bin/python
|
|
|
|
# Spectral envelope preprocessor.
|
|
# Requires two streams as input: the original raw data, and sinefit data.
|
|
|
|
import nilmtools.filter
|
|
import nilmdb.client
|
|
from numpy import *
|
|
import scipy.fftpack
|
|
import scipy.signal
|
|
#from matplotlib import pyplot as p
|
|
import bisect
|
|
|
|
def main(argv = None):
|
|
# Set up argument parser
|
|
f = nilmtools.filter.Filter()
|
|
parser = f.setup_parser("Spectral Envelope Preprocessor", skip_paths = True)
|
|
group = parser.add_argument_group("Prep options")
|
|
group.add_argument("-c", "--column", action="store", type=int,
|
|
help="Column number (first data column is 1)")
|
|
group.add_argument("-n", "--nharm", action="store", type=int, default=4,
|
|
help="number of odd harmonics to compute (default 4)")
|
|
group.add_argument("-N", "--nshift", action="store", type=int, default=1,
|
|
help="number of shifted FFTs per period (default 1)")
|
|
exc = group.add_mutually_exclusive_group()
|
|
exc.add_argument("-r", "--rotate", action="store", type=float,
|
|
help="rotate FFT output by this many degrees (default 0)")
|
|
exc.add_argument("-R", "--rotate-rad", action="store", type=float,
|
|
help="rotate FFT output by this many radians (default 0)")
|
|
|
|
group.add_argument("srcpath", action="store",
|
|
help="Path of raw input, e.g. /foo/raw")
|
|
group.add_argument("sinepath", action="store",
|
|
help="Path of sinefit input, e.g. /foo/sinefit")
|
|
group.add_argument("destpath", action="store",
|
|
help="Path of prep output, e.g. /foo/prep")
|
|
|
|
# Parse arguments
|
|
try:
|
|
args = f.parse_args(argv)
|
|
except nilmtools.filter.MissingDestination as e:
|
|
rec = "float32_%d" % (e.parsed_args.nharm * 2)
|
|
print "Source is %s (%s)" % (e.src.path, e.src.layout)
|
|
print "Destination %s doesn't exist" % (e.dest.path)
|
|
print "You could make it with a command like:"
|
|
print " nilmtool -u %s create %s %s" % (e.dest.url, e.dest.path, rec)
|
|
raise SystemExit(1)
|
|
|
|
if f.dest.layout_count != args.nharm * 2:
|
|
print "error: need", args.nharm*2, "columns in destination stream"
|
|
raise SystemExit(1)
|
|
|
|
# Check arguments
|
|
if args.column is None or args.column < 1:
|
|
parser.error("need a column number >= 1")
|
|
|
|
if args.nharm < 1 or args.nharm > 32:
|
|
parser.error("number of odd harmonics must be 1-32")
|
|
|
|
if args.nshift < 1:
|
|
parser.error("number of shifted FFTs must be >= 1")
|
|
|
|
if args.rotate is not None:
|
|
rotation = args.rotate * 2.0 * pi / 360.0
|
|
else:
|
|
rotation = args.rotate_rad or 0.0
|
|
|
|
# Check the sine fit stream
|
|
client_sinefit = nilmdb.client.Client(args.url)
|
|
sinefit = nilmtools.filter.get_stream_info(client_sinefit, args.sinepath)
|
|
if not sinefit:
|
|
raise Exception("sinefit data not found")
|
|
if sinefit.layout != "float32_3":
|
|
raise Exception("sinefit data type is " + sinefit.layout
|
|
+ "; expected float32_3")
|
|
|
|
# Check and set metadata in prep stream
|
|
f.check_dest_metadata({ "prep_raw_source": f.src.path,
|
|
"prep_sinefit_source": sinefit.path,
|
|
"prep_column": args.column })
|
|
|
|
# Run the processing function on all data
|
|
f.process_numpy(process, args = (client_sinefit, sinefit.path, args.column,
|
|
args.nharm, rotation, args.nshift))
|
|
|
|
def process(data, interval, args, insert_function, final):
|
|
(client, sinefit_path, column, nharm, rotation, nshift) = args
|
|
rows = data.shape[0]
|
|
data_timestamps = data[:,0]
|
|
|
|
if rows < 2:
|
|
return 0
|
|
|
|
last_inserted = [nilmdb.utils.time.min_timestamp]
|
|
def insert_if_nonoverlapping(data):
|
|
"""Call insert_function to insert data, but only if this
|
|
data doesn't overlap with other data that we inserted."""
|
|
if data[0][0] <= last_inserted[0]:
|
|
return
|
|
last_inserted[0] = data[-1][0]
|
|
insert_function(data)
|
|
|
|
processed = 0
|
|
out = zeros((1, nharm * 2 + 1))
|
|
# Pull out sinefit data for the entire time range of this block
|
|
for sinefit_line in client.stream_extract(sinefit_path,
|
|
data[0, 0], data[rows-1, 0]):
|
|
|
|
def prep_period(t_min, t_max, rot):
|
|
"""
|
|
Compute prep coefficients from time t_min to t_max, which
|
|
are the timestamps of the start and end of one period.
|
|
Results are rotated by an additional extra_rot before
|
|
being inserted into the database. Returns the maximum
|
|
index processed, or None if the period couldn't be
|
|
processed.
|
|
"""
|
|
# Find the indices of data that correspond to (t_min, t_max)
|
|
idx_min = bisect.bisect_left(data_timestamps, t_min)
|
|
idx_max = bisect.bisect_left(data_timestamps, t_max)
|
|
if idx_min >= idx_max or idx_max >= len(data_timestamps):
|
|
return None
|
|
|
|
# Perform FFT over those indices
|
|
N = idx_max - idx_min
|
|
d = data[idx_min:idx_max, column]
|
|
F = scipy.fftpack.fft(d) * 2.0 / N
|
|
|
|
# If we wanted more harmonics than the FFT gave us, pad with zeros
|
|
if N < (nharm * 2):
|
|
F = r_[F, zeros(nharm * 2 - N)]
|
|
|
|
# Fill output data.
|
|
out[0, 0] = round(t_min)
|
|
for k in range(nharm):
|
|
Fk = F[2 * k + 1] * e**(rot * 1j * (k+1))
|
|
out[0, 2 * k + 1] = -imag(Fk) # Pk
|
|
out[0, 2 * k + 2] = real(Fk) # Qk
|
|
|
|
insert_if_nonoverlapping(out)
|
|
return idx_max
|
|
|
|
# Extract sinefit data to get zero crossing timestamps.
|
|
# t_min = beginning of period
|
|
# t_max = end of period
|
|
(t_min, f0, A, C) = [ float(x) for x in sinefit_line.split() ]
|
|
t_max = t_min + 1e6 / f0
|
|
|
|
# Compute prep over shifted windows of the period
|
|
# (nshift is typically 1)
|
|
for n in range(nshift):
|
|
# Compute timestamps and rotations for shifted window
|
|
time_shift = n * (t_max - t_min) / nshift
|
|
shifted_min = t_min + time_shift
|
|
shifted_max = t_max + time_shift
|
|
angle_shift = n * 2 * pi / nshift
|
|
shifted_rot = rotation - angle_shift
|
|
|
|
# Run prep computation
|
|
idx_max = prep_period(shifted_min, shifted_max, shifted_rot)
|
|
if not idx_max:
|
|
break
|
|
processed = idx_max
|
|
|
|
print "Processed", processed, "of", rows, "rows"
|
|
return processed
|
|
|
|
if __name__ == "__main__":
|
|
main()
|