geomag.py

#! /usr/bin/env python

import argparse
import sys

# ensure geomag is on the path before importing
try:
    import geomagio
except:
    from os import path
    script_dir = path.dirname(path.abspath(__file__))
    sys.path.append(path.normpath(path.join(script_dir, '..')))

import geomagio.iaga2002 as iaga2002
import geomagio.edge as edge
from geomagio.Algorithm import Algorithm
from geomagio.XYZAlgorithm import XYZAlgorithm
from geomagio.Controller import Controller
from geomagio.iaga2002.IAGA2002Factory import IAGA_FILE_PATTERN
from obspy.core.utcdatetime import UTCDateTime


def main():
    """command line factory for geomag algorithms

    Inputs
    ------
    --input: string
        the type of data for input
        currently either iaga or edge.
    --output: string
        the type of data for ouput
        currently either iaga or edge.
    --starttime: string
        formatted as a obspy.core.UTCDateTime object
        the starttime for data input/output
    --endtime: string
        formatted as a obspy.core.UTCDateTime object
        the endtime for data input/output
    --observatory:string

    --channels: array_like
        list of channels
    --type: string
        data type
    --invterval: string
        data interval.
    --algorithm: string
        name of an algorithm to use.
    --xyz-informat: string
        The input format/coordinate system of the input file.
            geo: geographic coordinate system (xyzf)
            mag: magnetic north coordinate system (hdzf)
            obs: observatory coordinate system (hezf)
            obsd: observatory coordinate system (hdzf)
    --xyz-outformat: string
        The ouput format/coordinate system of the output file.
            geo: geographic coordinate system (xyzf)
            mag: magnetic north coordinate system (hdzf)
            obs: observatory coordinate system (hezf or hdzf)
    --input_iaga_magweb: boolean
        indicates to use http://magweb.cr.usgs.gov/data/magnetometer/ as the
        source of iaga2002 files.
    --input_iaga_url: string
        url of iaga2002 files to use as the data source.
    --input-iaga-urltemplate: string
        template for the subdirectories that files are found in.
        example: %(OBS)s/%(interval)s%(type)s/
    --input-iaga-filetemplate: string
        template for the file name
        example: %(obs)s%(ymd)s%(t)s%(i)s.%(i)s
    --input-iaga-file: string
        the filename of the Iaga2002 file to be read from
    --input-iaga-stdin: boolean
        indicates the file will be coming from stdin
    --output_iaga_file: string
        the filename of a new Iaga2002 file to be read to
    --output-iaga-url: string
        url of directory to write output files in.
    --output-iaga-urltemplate: string
        template for the subdirectories that files are to be written in.
        example: %(OBS)s/%(interval)s%(type)s/
    --output-iaga-filetemplate: string
        template for the file name
        example: %(obs)s%(ymd)s%(t)s%(i)s.%(i)s
    --output-iaga-stdout: boolen
        indicates output will go to stdout
    """

    args = parse_args()

    if args.input == 'iaga':
        if args.input_iaga_magweb:
            inputfactory = iaga2002.MagWebFactory(
                    observatory=args.observatory,
                    type=args.type,
                    interval=args.interval)
        elif args.input_iaga_url is not None:
            inputfactory = iaga2002.IAGA2002Factory(
                    urlTemplate=_get_iaga_input_url(args),
                    observatory=args.observatory,
                    type=args.type,
                    interval=args.interval)
        elif args.input_iaga_file is not None:
            iagaFile = open(args.input_iaga_file, 'r').read()
            inputfactory = iaga2002.StreamIAGA2002Factory(
                stream=iagaFile,
                observatory=args.observatory,
                type=args.type,
                interval=args.interval)
        elif args.input_iaga_stdin:
            print >> sys.stderr, "Iaga Input waiting for data from stdin"
            iagaFile = sys.stdin.read()
            inputfactory = iaga2002.StreamIAGA2002Factory(
                stream=iagaFile,
                observatory=args.observatory,
                type=args.type,
                interval=args.interval)
        else:
            print >> sys.stderr, "Iaga Input was missing needed arguments"

    elif args.input == 'edge':
        inputfactory = edge.EdgeFactory(
                host=args.input_edge_host,
                port=args.input_edge_port,
                observatory=args.observatory,
                type=args.type,
                interval=args.interval)

    if args.output == 'iaga':
        if args.output_iaga_url is not None:
            outputfactory = iaga2002.IAGA2002Factory(
                    urlTemplate=_get_iaga_output_url(args),
                    observatory=args.observatory,
                    type=args.type,
                    interval=args.interval)
        elif args.output_iaga_file is not None:
            iagaFile = open(args.output_iaga_file, 'w')
            outputfactory = iaga2002.StreamIAGA2002Factory(
                stream=iagaFile,
                observatory=args.observatory,
                type=args.type,
                interval=args.interval)
        elif args.output_iaga_stdout:
            iagaFile = sys.stdout
            outputfactory = iaga2002.StreamIAGA2002Factory(
                    stream=iagaFile,
                    observatory=args.observatory,
                    type=args.type,
                    interval=args.interval)
        else:
            print >> sys.stderr, "Iaga Output was missing needed arguments"

    if args.algorithm == 'xyz':
        algorithm = XYZAlgorithm(args.xyz_informat, args.xyz_outformat)
    else:
        algorithm = Algorithm(channels=args.channels)

    # TODO check for unused arguments.

    controller = Controller(inputfactory, outputfactory, algorithm)

    controller.run(UTCDateTime(args.starttime), UTCDateTime(args.endtime))

def parse_args():
    """parse input arguments

    Returns
    -------
    argparse.Namespace
        dictionary like object containing arguments.
    """
    parser = argparse.ArgumentParser(
        description='Use @ to read commands from a file.',
        fromfile_prefix_chars='@',)

    parser.add_argument('--input', choices=['iaga', 'edge'],
            help = 'Input type.', required=True)

    parser.add_argument('--output', choices=['iaga'],
            help = 'Input type.', required=True)

    parser.add_argument('--starttime', default=UTCDateTime(),
            help='UTC date YYYY-MM-DD HH:MM:SS')
    parser.add_argument('--endtime', default=UTCDateTime(),
            help='UTC date YYYY-MM-DD HH:MM:SS')

    parser.add_argument('--observatory',
            help='Observatory code ie BOU, CMO, etc')
    parser.add_argument('--channels', nargs='*',
            help='Channels H, E, Z, etc')
    parser.add_argument('--type', default='variation',
            choices=['variation', 'quasi-definitive', 'definitive'])
    parser.add_argument('--interval', default='minute',
            choices=['minute', 'second'])

    parser.add_argument('--algorithm', choices='xyz')

    # xyz algorithm arguments
    parser.add_argument('--xyz-informat',
            choices=['geo', 'mag', 'obs', 'obsd'])
    parser.add_argument('--xyz-outformat',
            choices=['geo', 'mag', 'obs', 'obsd'])

    # iaga2002 input arguments
    parser.add_argument('--input-iaga-file',
            help='Iaga2002 filename')
    parser.add_argument('--input-iaga-magweb',
            action="store_true", default=False,
            help='Indicates iaga2002 files will be read from \
            http://magweb.cr.usgs.gov/data/magnetometer/')
    parser.add_argument('--input-iaga-stdin',
            action="store_true", default=False,
            help='Indicates file will be redirected from stdin')
    parser.add_argument('--input-iaga-url',
            help='Url or Directory where Iaga2002 files can be read from')
    parser.add_argument('--input-iaga-urltemplate',
            help='Template for directory matching')
    parser.add_argument('--input-iaga-filetemplate',
            help='Template for iaga filenames')

    parser.add_argument('--output-iaga-url',
            help='Url or Directory where IAGA2002 files should be written to')
    parser.add_argument('--output-iaga-stdout',
            action="store_true", default=False,
            help='Indicates file will be directed to stdout')
    parser.add_argument('--output-iaga-urltemplate',
            help='Template for subdirectories')
    parser.add_argument('--output-iaga-filetemplate',
            help='Template for iaga filenames')
    parser.add_argument('--output-iaga-file',
            help='Output file name for single iaga file.')

    # Edge input arguments
    parser.add_argument('--input-edge-host',
            help='ip address of the edge input server')
    parser.add_argument('--input-edge-port', type=int,
            help='port number of the edge input server')

    return parser.parse_args()

def _get_iaga_input_url(args):
    url = args.input_iaga_url or 'file://./'
    urltemplate = args.input_iaga_urltemplate or ''
    filetemplate = args.input_iaga_filetemplate or IAGA_FILE_PATTERN
    return url + urltemplate + filetemplate

def _get_iaga_output_url(args):
    url = args.output_iaga_url or 'file://./'
    urltemplate = args.output_iaga_urltemplate or ''
    filetemplate = args.output_iaga_filetemplate or IAGA_FILE_PATTERN
    return url + urltemplate + filetemplate

if __name__ == '__main__':
    main()