nilmtools/nilmtools/sinefit.py

#!/usr/bin/python

# Sine wave fitting.
from nilmdb.utils.printf import *
import nilmtools.filter
import nilmtools.math
import nilmdb.client
from nilmdb.utils.time import (timestamp_to_human,
                               timestamp_to_seconds,
                               seconds_to_timestamp)

from numpy import *
from scipy import *
#import pylab as p
import sys

def main(argv = None):
    f = nilmtools.filter.Filter()
    parser = f.setup_parser("Sine wave fitting")
    group = parser.add_argument_group("Sine fit options")
    group.add_argument('-c', '--column', action='store', type=int,
                       help='Column number (first data column is 1)')
    group.add_argument('-f', '--frequency', action='store', type=float,
                       default=60.0,
                       help='Approximate frequency (default: %(default)s)')
    group.add_argument('-m', '--min-freq', action='store', type=float,
                       help='Minimum valid frequency '
                       '(default: approximate frequency / 2))')
    group.add_argument('-M', '--max-freq', action='store', type=float,
                       help='Maximum valid frequency '
                       '(default: approximate frequency * 2))')
    group.add_argument('-a', '--min-amp', action='store', type=float,
                       default=20.0,
                       help='Minimum signal amplitude (default: %(default)s)')

    # Parse arguments
    try:
        args = f.parse_args(argv)
    except nilmtools.filter.MissingDestination as e:
        rec = "float32_3"
        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 args.column is None or args.column < 1:
        parser.error("need a column number >= 1")
    if args.frequency < 0.1:
        parser.error("frequency must be >= 0.1")
    if args.min_freq is None:
        args.min_freq = args.frequency / 2
    if args.max_freq is None:
        args.max_freq = args.frequency * 2
    if (args.min_freq > args.max_freq or
        args.min_freq > args.frequency or
        args.max_freq < args.frequency):
        parser.error("invalid min or max frequency")
    if args.min_amp < 0:
        parser.error("min amplitude must be >= 0")

    f.check_dest_metadata({ "sinefit_source": f.src.path,
                            "sinefit_column": args.column })
    f.process_numpy(process, args = (args.column, args.frequency, args.min_amp,
                                     args.min_freq, args.max_freq))

class SuppressibleWarning(object):
    def __init__(self, maxcount = 10, maxsuppress = 100):
        self.maxcount = maxcount
        self.maxsuppress = maxsuppress
        self.count = 0
        self.last_msg = ""

    def _write(self, sec, msg):
        if sec:
            now = timestamp_to_human(seconds_to_timestamp(sec)) + ": "
        else:
            now = ""
        sys.stderr.write(now + msg)

    def warn(self, msg, seconds = None):
        self.count += 1
        if self.count <= self.maxcount:
            self._write(seconds, msg)
        if (self.count - self.maxcount) >= self.maxsuppress:
            self.reset(seconds)

    def reset(self, seconds = None):
        if self.count > self.maxcount:
            self._write(seconds, sprintf("(%d warnings suppressed)\n",
                                         self.count - self.maxcount))
        self.count = 0

def process(data, interval, args, insert_function, final):
    (column, f_expected, a_min, f_min, f_max) = args
    rows = data.shape[0]

    # Estimate sampling frequency from timestamps
    fs = (rows-1) / (timestamp_to_seconds(data[-1][0]) -
                     timestamp_to_seconds(data[0][0]))

    # Pull out about 3.5 periods of data at once;
    # we'll expect to match 3 zero crossings in each window
    N = max(int(3.5 * fs / f_expected), 10)

    # If we don't have enough data, don't bother processing it
    if rows < N:
        return 0

    warn = SuppressibleWarning(3, 1000)

    # Process overlapping windows
    start = 0
    num_zc = 0
    last_inserted_timestamp = None
    while start < (rows - N):
        this = data[start:start+N, column]
        t_min = timestamp_to_seconds(data[start, 0])
        t_max = timestamp_to_seconds(data[start+N-1, 0])

        # Do 4-parameter sine wave fit
        (A, f0, phi, C) = nilmtools.math.sfit4(this, fs)

        # Check bounds.  If frequency is too crazy, ignore this window
        if f0 < f_min or f0 > f_max:
            warn.warn(sprintf("frequency %s outside valid range %s - %s\n",
                              str(f0), str(f_min), str(f_max)), t_min)
            start += N
            continue

        # If amplitude is too low, results are probably just noise
        if A < a_min:
            warn.warn(sprintf("amplitude %s below minimum threshold %s\n",
                              str(A), str(a_min)), t_min)
            start += N
            continue

        #p.plot(arange(N), this)
        #p.plot(arange(N), A * sin(f0/fs * 2 * pi * arange(N) + phi) + C, 'g')

        # Period starts when the argument of sine is 0 degrees,
        # so we're looking for sample number:
        #     n = (0 - phi) / (f0/fs * 2 * pi)
        zc_n = (0 - phi) / (f0 / fs * 2 * pi)
        period_n = fs/f0

        # Add periods to make N positive
        while zc_n < 0:
            zc_n += period_n

        last_zc = None
        # Mark the zero crossings until we're a half period away
        # from the end of the window
        while zc_n < (N - period_n/2):
            #p.plot(zc_n, C, 'ro')
            t = t_min + zc_n / fs
            if (last_inserted_timestamp is None or
                t > last_inserted_timestamp):
                insert_function([[seconds_to_timestamp(t), f0, A, C]])
                last_inserted_timestamp = t
                warn.reset(t)
            else:
                warn.warn("timestamp overlap\n", t)
            num_zc += 1
            last_zc = zc_n
            zc_n += period_n

        # Advance the window one quarter period past the last marked
        # zero crossing, or advance the window by half its size if we
        # didn't mark any.
        if last_zc is not None:
            advance = min(last_zc + period_n/4, N)
        else:
            advance = N/2
        #p.plot(advance, C, 'go')
        #p.show()

        start = int(round(start + advance))

    # Return the number of rows we've processed
    warn.reset(last_inserted_timestamp)
    if last_inserted_timestamp:
        now = timestamp_to_human(seconds_to_timestamp(
            last_inserted_timestamp)) + ": "
    else:
        now = ""
    printf("%sMarked %d zero-crossings in %d rows\n", now, num_zc, start)
    return start

if __name__ == "__main__":
    main()