text:
import chainer
import chainer.functions as F
import chainer.links as L
import inspect
import ast
import gast
import weakref
from enum import Enum
import re
import numpy as np

from chainer_compiler.elichika.parser import vevaluator
from chainer_compiler.elichika.parser import nodes
from chainer_compiler.elichika.parser import values
from chainer_compiler.elichika.parser import functions
from chainer_compiler.elichika.parser import utils
from chainer_compiler.elichika.parser import core
from chainer_compiler.elichika.parser import config
from chainer_compiler.elichika.parser import canonicalizer

def generate_copied_value(value: 'values.Value'):
    assert(isinstance(value, values.Value))

    if isinstance(value, values.NumberValue):
        copied = values.NumberValue(value.internal_value)
        copied.dtype = value.dtype
        return copied

    if isinstance(value, values.TensorValue):
        copied = values.TensorValue()
        copied.value = value.value
        copied.shape = value.shape
        copied.dtype = value.dtype
        return copied

    if isinstance(value, values.ListValue):
        copied = values.ListValue()
        copied.dtype = value.dtype
        copied.vtype = value.vtype
        if value.internal_value is not None:
            copied.internal_value = value.internal_value.copy()
        return copied

    if isinstance(value, values.NoneValue):
        copied = values.NoneValue()
        return copied

    if isinstance(value, values.BoolValue):
        copied = values.BoolValue(value.internal_value)
        return copied

    if isinstance(value, values.StrValue):
        copied = values.StrValue(value.internal_value)
        return copied

    if isinstance(value, values.RangeValue):
        copied = values.RangeValue()
        return copied

    if isinstance(value, values.TupleValue):
        if value.internal_value is not None:
            copied = values.TupleValue(value.internal_value.copy())
        else:
            copied = values.TupleValue(value.internal_value)
        copied.dtype = value.dtype
        copied.vtype = value.vtype
        return copied

    if config.show_warnings:
        print('Unknown type {} is copied'.format(value))

    return values.Value()

def generate_tensor_value_with_undefined_shape_size(value: 'values.TensorValue'):
    assert(isinstance(value, values.TensorValue))
    ret = values.TensorValue()
    ret.shape = tuple([-1 for v in value.shape])
    return ret

class SuffixType(Enum):
    Unknown = 0,
    Unused = 1,
    Dummy = 2,
    Input = 3,

def generate_value_with_type(value: 'values.Value', type_,  suffix_type = SuffixType.Unknown):
    assert(isinstance(value, values.Value))
    ret = None
    if isinstance(value, values.TensorValue):
        ret = values.TensorValue()
        ret.shape = value.shape
        ret.dtype = type_

    elif isinstance(value, values.NumberValue):
        ret = values.NumberValue(None)
        ret.dtype = type_

    else:
        assert(False)

    if suffix_type == SuffixType.Unknown:
        ret.name = value.name + '_st'
    elif suffix_type == SuffixType.Unused:
        ret.name = value.name + '_unused'
    elif suffix_type == SuffixType.Dummy:
        ret.name = value.name + '_dummy'
    elif suffix_type == SuffixType.Input:
        ret.name = value.name + '_in'
    else:
        assert(False)

    return ret

def generate_value_with_same_type(value: 'values.Value', is_dummy_value = False, suffix_type = SuffixType.Unknown):
    assert(isinstance(value, values.Value))
    ret = None
    if isinstance(value, values.TensorValue):
        ret = values.TensorValue()
        ret.shape = value.shape
        ret.dtype = value.dtype

    elif isinstance(value, values.NumberValue):
        dtype = None
        if value.internal_value is None:
            dtype = value.dtype
        elif isinstance(value.internal_value, int):
            dtype = np.array(value.internal_value).dtype
        elif isinstance(value.internal_value, float):
            dtype = np.array(value.internal_value).dtype

        ret = values.NumberValue(None)
        ret.dtype = dtype

    elif isinstance(value, values.StrValue):
        ret = values.StrValue(None)

    elif isinstance(value, values.BoolValue):
        ret = values.BoolValue(None)

    elif isinstance(value, values.ListValue):
        ret = values.ListValue(None)
        ret.dtype = value.dtype
        ret.vtype = value.vtype

    elif isinstance(value, values.NoneValue):
        ret = values.NoneValue()

    elif isinstance(value, values.TupleValue):
        ret = values.TupleValue()
        ret.dtype = value.dtype
        ret.vtype = value.vtype

    elif isinstance(value, values.RangeValue):
        ret = values.RangeValue()

    elif isinstance(value, values.UnknownValue):
        ret = values.UnknownValue()

    elif ret is None and isinstance(value, values.Value):
        ret = values.Value()

    else:
        assert(False)

    assert(ret is not None)

    ret.is_dummy_value = is_dummy_value
    if suffix_type == SuffixType.Unknown:
        ret.name = value.name + '_st'
    elif suffix_type == SuffixType.Unused:
        ret.name = value.name + '_unused'
    elif suffix_type == SuffixType.Dummy:
        ret.name = value.name + '_dummy'
    elif suffix_type == SuffixType.Input:
        ret.name = value.name + '_in'
    else:
        assert(False)

    return ret

class FunctionArgInput():
    def __init__(self):
        self.inputs = []
        self.keywords = {}

    def get_value(self) -> "FunctionArgValueInput":
        ret = functions.FunctionArgValueInput()
        ret.inputs = [v.get_value() for v in self.inputs]

        keywords_ = {}
        for k, v in self.keywords.items():
            keywords_[k] = v.get_value()
        ret.keywords = keywords_
        return ret

class FunctionArgValueInput():
    def __init__(self):
        self.inputs = [] # List[values.Value]
        self.keywords = {}  # Dict[str,values.Value]

    def get_value(self, key) -> 'values.Value':
        if isinstance(key, int):
            return self.inputs[key]
        if isinstance(key, str) and key in self.keywords.keys():
            return self.keywords[key]
        return None

class FunctionArg():
    def __init__(self, name: 'str' = '', obj: 'values.Object' = None):
        self.name = name
        self.obj = obj

class FunctionArgCollection():
    def __init__(self):
        self.args = {}  # Dict[str,FunctionArg]
        self.args_list = []

    def add_arg(self, name, value):

        if isinstance(value, values.Value):
            value = values.Object(value)

        assert not(name in self.args.keys())

        fa = FunctionArg(name, value)
        self.args_list.append(fa)
        self.args[fa.name] = fa

    def analyze_args(self, func):
        sig = inspect.signature(func)
        # TODO: replace with https://docs.python.org/3.4/library/inspect.html#inspect.getfullargspec
        argspec = inspect.getargspec(func)  # TODO: Doesn't support keyword only args.

        parameter_count = 0
        for k, v in sig.parameters.items():
            # TODO improve it
            if k == 'kwargs':
                continue
            parameter_count += 1

        isSelfRemoved = parameter_count != len(argspec[0])

        if isSelfRemoved:
            self.add_arg(argspec[0][0], None)

        for k, v in sig.parameters.items():
            # TODO improve it
            if k == 'kwargs':
                continue

            self.add_arg(v.name, values.parse_instance(None, v.name, v.default))

    def merge_inputs(self, self_Object, inputs: 'FunctionArgInput') -> 'FunctionArgInput':
        ret = FunctionArgInput()

        for fa in self.get_args():
            ret.inputs.append(fa.obj)
            ret.keywords[fa.name] = fa.obj

        inputs_ = inputs.inputs.copy()
        keywords_ = inputs.keywords.copy()

        if self_Object is not None:
            inputs_ = [self_Object] + inputs_
            keywords_[self.args_list[0].name] = self_Object

        for i in range(len(inputs_)):
            ret.inputs[i] = inputs_[i]
            ret.keywords[self.args_list[i].name] = ret.inputs[i]

        for k, v in keywords_.items():
            if k in ret.keywords.keys():
                ret.keywords[k] = v

            for i in range(len(self.args_list)):
                if self.args_list[i].name == k:
                    ret.inputs[i] = v

        return ret

    def get_value(self, key) -> 'values.Value':
        if isinstance(key, int):
            return self.args_list[key].obj.get_value()
        if isinstance(key, str) and key in self.args.keys():
            return self.args[key].obj.get_value()
        return None

    def get_values(self) -> 'List[values.Value]':
        return [a.obj.get_value() for a in self.args_list]

    def get_args(self) -> 'List[FunctionArg]':
        ret = []

        for fa in self.args_list:
            ret.append(FunctionArg(fa.name, fa.obj))
        return ret

def auto_set_unset(func, flag):
    def decorated(self, *args, **kwargs):
        self.history.append((flag, self.flags[flag]))
        ret = func(self, *args, **kwargs)
        return ret
    return decorated

class StackTrace:
    def __init__(self):
        self.lineprops = []

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.lineprops.pop()
        return False

    def append(self, line):
        self.lineprops.append(line)

class VEvalContext:
    def __init__(self):
        self.stacktrace = StackTrace()
        self.history = []
        self.flags = {
            "eval_as_written_target": False,
            "ignore_branch": False,
            "for_unroll": False
        }
        self.flags_cache = []

        list(map(lambda flag: setattr(self.__class__, '_' + flag, property(fset=lambda obj, value: VEvalContext.generic_setter(obj, flag, value),
                                                                           fget=lambda obj: VEvalContext.generic_getter(obj, flag))),
                 self.flags.keys()))
        list(map(lambda flag: setattr(self.__class__, flag, auto_set_unset(lambda obj, default=True: VEvalContext.generic_setter(obj, flag, default), flag)),
                 self.flags.keys()))

    def generic_setter(self, name, value = True):
        self.flags[name] = value
        return self

    def generic_getter(self, name):
        return self.flags[name]

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        flag, saved_value = self.history.pop()
        self.flags[flag] = saved_value
        return False

class FunctionBase():
    def __init__(self):
        self.name = ''
        self.is_property = False
        self.args = FunctionArgCollection()

        self.base_func = None

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
              context: 'VEvalContext' = None, line=-1):
        return None

class UserDefinedClassConstructorFunction(FunctionBase):
    def __init__(self, classinfo):
        super().__init__()

        members = inspect.getmembers(classinfo)
        init_func = [m[1] for m in members if m[0] == '__init__']
        assert(len(init_func) == 1)

        func = init_func[0]
        self.inst = func
        self.name = func.__name__
        self.filename = inspect.getfile(func)

        sourcelines = inspect.getsourcelines(func)
        if sourcelines is None or len(sourcelines) < 1:
            utils.print_warning('Failed to parase {}'.format(classinfo), utils.LineProperty())
            return

        self.lineno = sourcelines[1]
        self.classinfo = classinfo

        original_code = inspect.getsource(func)
        code = utils.clip_head(original_code)

        self.args.analyze_args(func)

        ast_ = gast.ast_to_gast(ast.parse(code)).body[0]
        self.ast = canonicalizer.Canonicalizer().visit(ast_)

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
              context: 'VEvalContext' = None, line=-1):
        ret = values.Object(values.UserDefinedInstance(
            module, None, self.classinfo))
        inst = ret

        func_field = values.Field()
        func_field.set_module(module)

        # add args
        funcArgs = self.args.merge_inputs(inst, args)

        for k, v in funcArgs.keywords.items():
            func_field.get_field().get_attribute(k, from_module=False).revise(v)

        astc = vevaluator.AstContext(self.ast.body, self.lineno - 1, filename=self.filename)
        vevaluator.veval_ast(astc, func_field, graph, context)

        # dispose because of exit from function
        func_field.dispose()

        return ret

class UserDefinedFunction(FunctionBase):
    def __init__(self, func):
        super().__init__()

        self.inst = func
        self.name = func.__name__
        self.filename = inspect.getfile(func)
        sourcelines = inspect.getsourcelines(func)
        self.lineno = sourcelines[1]
        self.args.analyze_args(func)

        if (func.__name__ == (lambda: None).__name__):
            original_code = utils.lambda_source(func)
            code = 'return ' + original_code[re.search('lambda.*?:', original_code).end():]
            self.ast = gast.ast_to_gast(ast.parse(code))
        else:
            original_code = inspect.getsource(func)
            code = utils.clip_head(original_code)
            ast_ = gast.ast_to_gast(ast.parse(code)).body[0]
            self.ast = canonicalizer.Canonicalizer().visit(ast_)

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
              context: 'VEvalContext' = None, line=-1):

        if context is None:
            context = VEvalContext()

        context.stacktrace.append(line)
        with context.stacktrace:

            func_field = values.Field()
            func_field.set_module(module)

            # add args
            funcArgs = self.args.merge_inputs(inst, args)

            for k, v in funcArgs.keywords.items():
                func_field.get_field().get_attribute(k, from_module=False).revise(utils.try_get_obj(v, self.name, utils.LineProperty()))

            astc = vevaluator.AstContext(self.ast.body, self.lineno - 1, filename=self.filename)
            ret = vevaluator.veval_ast(astc, func_field, graph, context)

            # dispose because of exit from function
            func_field.dispose()

            return ret

class UnimplementedFunction(FunctionBase):
    def __init__(self, func):
        super().__init__()

        if isinstance(func, str):
            self.name = func
        else:
            self.name = func.__name__

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
              context: 'VEvalContext' = None, line=-1):
        raise utils.UnimplementedError('{} is unimplemented.'.format(self.name), utils.LineProperty(line))

class UserDefinedFunctionFromAst(FunctionBase):
    def __init__(self, astc, args, func_field):
        super().__init__()
        assert isinstance(astc.nast, (gast.FunctionDef, gast.Lambda))

        self.name = astc.gast.name if isinstance(astc.nast, gast.FunctionDef) else (lambda: None).__name__
        self.args = args
        self.func_field = func_field
        if isinstance(astc.nast, gast.Lambda):
            astc.nast.body = gast.Return(value=astc.nast.body) # Add return to the body
        self.ast = astc.nast
        self.filename = astc.filename
        self.lineno = astc.lineno

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
              context: 'VEvalContext' = None, line=-1):
        self.func_field.set_module(module)

        # add args
        funcArgs = self.args.merge_inputs(inst, args)

        for k, v in funcArgs.keywords.items():
            self.func_field.get_field().get_attribute(k, from_module=False).revise(utils.try_get_obj(v, self.name, utils.LineProperty()))

        astc = vevaluator.AstContext(self.ast.body, self.lineno - 1, filename=self.filename)
        ret = vevaluator.veval_ast(astc, self.func_field, graph, context)

        # dispose because of exit from function
        self.func_field.dispose()

        return ret

class CheckAttributeValueFunction(FunctionBase):
    def __init__(self):
        super().__init__()
        self.name = 'check_attribute_value'
        self.args.add_arg('actual_value', values.NoneValue())
        self.args.add_arg('expected_value', values.NoneValue())
        self.args.add_arg('func_name', values.NoneValue())
        self.args.add_arg('arg_name', values.NoneValue())

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
              context: 'VEvalContext' = None, line=-1):

        funcArgs = self.args.merge_inputs(inst ,args)
        vargs = funcArgs.get_value().inputs

        if type(vargs[0]) != type(vargs[1]) or vargs[0].get_constant_value() != vargs[1].get_constant_value():
            raise Exception("Value must be {} : {} from {} in {}".format(
                vargs[1].get_constant_value(),
                vargs[3].get_constant_value(),
                vargs[2].get_constant_value(),
                context.stacktrace.lineprops[-1]))

class CheckAttributeScalarFunction(FunctionBase):
    def __init__(self):
        super().__init__()
        self.name = 'check_attribute_scalar'
        self.args.add_arg('value', values.NoneValue())
        self.args.add_arg('func_name', values.NoneValue())
        self.args.add_arg('arg_name', values.NoneValue())

    def vcall(self, module: 'values.Field', graph: 'graphs.Graph', inst: 'values.Object', args: 'FunctionArgInput',
                      context: 'VEvalContext' = None, line=-1):

        funcArgs = self.args.merge_inputs(inst ,args)
        vargs = funcArgs.get_value().inputs

        if not isinstance(vargs[0], values.NumberValue):
            raise Exception("A number is only supported {} : from {} in {}".format(
                vargs[2].get_constant_value(),
                vargs[1].get_constant_value(),
                context.stacktrace.lineprops[-1]))


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import os
import datetime
import random
import numpy as np
import torch
from torch import nn
import torch.multiprocessing as mp
from torch.utils.tensorboard import SummaryWriter

from model.agent import Agent
from model.trainer import Trainer, train_local_mp
from model.cifar_models import CifarModel, CifarwithAttn
from data.cifar_data import CifarData

class CIFARAgent(Agent):
    '''
    CIFARAgent for CIFAR10 and CIFAR100.
    '''
    def __init__(self, global_args, subset=tuple(range(10)),
                 fine='CIFAR10', train_indices=None, test_indices=None):
        super().__init__(global_args, subset, fine, train_indices, test_indices)

    def load_data(self):
        print("=> loading data")
        self.data = CifarData(self.subset, self.fine, self.train_indices, self.test_indices)
        self.train_loader = self.data.get_train_loader(self.batch_size, self.num_workers)
        self.test_loader = self.data.get_test_loader(self.batch_size, self.num_workers)

    def build_model(self):
        print("=> building model")
        if self.fine == 'CIFAR10':
            num_class = 10
        elif self.fine == 'CIFAR100':
            num_class = 100
        else:
            raise ValueError('Invalid dataset choice.')
        if self.fusion == 'none':
            self.model = CifarModel(num_class).to(self.device)
        else:
            self.model = CifarwithAttn(self.fusion, num_class).to(self.device)
        if self.fusion in ['multi', 'single']:
            self.shadow = torch.zeros(self.model.attn.gamma.size(), device=self.device)
        self.criterion = nn.CrossEntropyLoss().to(self.device)
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.lr,
                                         momentum=0.9, weight_decay=5e-4)

class CIFARTrainer(Trainer):
    '''
    CIFAR Trainer.
    '''
    def __init__(self, global_args):
        super().__init__(global_args)

        # init the global model
        self.global_agent = CIFARAgent(global_args, fine=self.fine)
        self.global_agent.load_data()
        self.global_agent.build_model()
        self.global_agent.resume_model(self.resume)

    def build_local_models(self, global_args):
        self.nets_pool = list()
        train_indices = np.random.permutation(50000)
        test_indices = np.random.permutation(10000)
        train_per_local = len(train_indices) // self.num_locals
        test_per_local = len(test_indices) // self.num_locals
        for i in range(self.num_locals):
            t_train_indices = train_indices[i * train_per_local: (i + 1) * train_per_local]
            t_test_indices = test_indices[i * test_per_local: (i + 1) * test_per_local]
            self.nets_pool.append(CIFARAgent(global_args, fine=self.fine, subset=tuple(range(10)),
                train_indices=t_train_indices, test_indices=t_test_indices))
        self.init_local_models()

    def train(self):
        ctx = mp.get_context('forkserver')
        self.num_per_rnd = 5
        for rnd in range(self.rounds):
            random.shuffle(self.nets_pool)
            pool = ctx.Pool(self.num_per_rnd)
            self.q = ctx.Manager().Queue()
            pool.starmap(train_local_mp, [(self.local_epochs, net, rnd, self.q) for net in self.nets_pool[:self.num_per_rnd]])
            pool.close()
            pool.join()
            self.update_global(rnd)

def main():
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    cifar_trainer = CIFARTrainer(args)

    # test
    if args.mode == 'test':
        cifar_trainer.test()
        return

    writer_dir = os.path.join(f'runs/{args.fine}_adam_{args.num_locals}',
                              datetime.datetime.now().strftime('%b%d_%H-%M'))
    cifar_trainer.writer = SummaryWriter(writer_dir)

    cifar_trainer.build_local_models(args)
    cifar_trainer.train()

    cifar_trainer.writer.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--model_file',
        type=str,
        default='model.pth.tar',
        help='File to save model.'
    )
    parser.add_argument(
        '--mode',
        type=str,
        default='train',
        help='train or test.'
    )
    parser.add_argument(
        '--fine',
        type=str,
        default='CIFAR10',
        choices=('CIFAR10', 'CIFAR100'),
        help='Fine choice of dataset.'
    )
    parser.add_argument(
        '--fusion',
        type=str,
        default='none',
        choices=('none', 'multi', 'single', 'conv'),
        help='Method for feature fusion.'
    )
    parser.add_argument(
        '--num_workers',
        type=int,
        default=0,
        help='number of workers to preprocess data, must be 0 for mp agents.'
    )
    parser.add_argument(
        '--num_locals',
        type=int,
        default=2,
        help='number of local agents.'
    )
    parser.add_argument(
        '--rounds',
        type=int,
        default=500,
        help='number of communication rounds.'
    )
    parser.add_argument(
        '--lr',
        type=float,
        default=5e-3,
        help='learning rate.'
    )
    parser.add_argument(
        '--min_lr',
        type=float,
        default=1e-4,
        help='minimum learning rate.'
    )
    parser.add_argument(
        '--decay_rate',
        type=float,
        default=0.99,
        help='lr decay rate.'
    )
    parser.add_argument(
        '--batch_size',
        type=int,
        default=64,
        help='Batch size. (B)'
    )
    parser.add_argument(
        '--local_epochs',
        type=int,
        default=5,
        help='Number of epoch in local. (E)'
    )
    parser.add_argument(
        '--meta_lr',
        type=float,
        default=1e-3,
        help='meta learning rate for model aggregation.'
    )
    parser.add_argument(
        '--resume',
        default='',
        type=str,
        metavar='PATH',
        help='path to resume checkpoint (default: none)'
    )
    parser.add_argument(
        '--model_dir',
        type=str,
        default='models',
        help='Directory for storing checkpoint file.'
    )
    parser.add_argument(
        '--gpu',
        type=str,
        default='2',
        help='Number of gpu to use'
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=1234,
        help='Random seed'
    )
    args = parser.parse_args()
    main()

ext: py

max_stars_repo_name: thu-media/FedFusion

max_issues_repo_name: thu-media/FedFusion

max_forks_repo_name: thu-media/FedFusion

reason: import numpy


text:
"""
Load a recorded .bag streams, and save
the depth stream as a video file.
"""

import pyrealsense2 as rs
import numpy as np
import cv2
import os
import json
import argparse

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--path', '-p', default='../../../Documents/realsense_bag_files', type=str, help="Path to the bag files folder")
    parser.add_argument(
        '--format', '-f', default='mp4', type=str, choices=['mp4', 'avi'])
    parser.add_argument(
        '--json', '-j', default='../config.json', type=str, help="Path to the json config file")
    # parser.add_argument(
    #     '--visual_preset', default="High Accuracy", type=str,
    #     choices=["Custom", "Default", "Hand", "High Accuracy", "High Density"])

    return parser

def get_args(parser):
    # Parse the command line arguments to an object
    args = parser.parse_args()

    # Safety if no parameter have been given
    if not args.path:
        print("No path paramater have been given.")
        print("For help type --help")
        exit()
    return args

def main(args):

    # json config
    json_path = args.json
    jason_obj = json.load(open(json_path))
    json_string= str(jason_obj).replace("'", '\"')
    width = int(jason_obj['viewer']['stream-width'])
    height = int(jason_obj['viewer']['stream-height'])
    FPS = int(jason_obj['viewer']['stream-fps'])

    # set output video encoding
    if args.format == 'mp4':
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    elif args.format == 'avi':
        fourcc = cv2.VideoWriter_fourcc(*'XVID')

    for filename in os.listdir(args.path):

        output_name = os.path.splitext(filename)[0]
        # Check if the file has bag extension
        if os.path.splitext(filename)[1] != ".bag":
            print("The given file is not of correct file format.")
            print("Only .bag files are accepted")
            exit()

        # set output video names
        depth_path = output_name + '_depth.' + args.format
        # set output video writers
        depthwriter = cv2.VideoWriter(depth_path, fourcc, FPS, (width, height), 1)

        # Create pipeline
        pipeline = rs.pipeline()
        # Create a config object
        config = rs.config()
        # Tell config that we will use a recorded device from file to be used by the pipeline through playback.
        config.enable_device_from_file(os.path.join(args.path, filename), repeat_playback=False)
        # Configure the pipeline to stream the depth stream
        # REMEMBER that width, height and FPS should be the same of the recorded stream
        config.enable_stream(rs.stream.depth, width, height, rs.format.z16, FPS)

        # Start streaming from file
        profile = pipeline.start(config)
        device = profile.get_device()

        # Playback is used to find duration of recorded video
        playback = device.as_playback()
        playback.set_real_time(False)
        # duration in nano seconds
        duration = playback.get_duration().total_seconds() * 1e9

        # Load advanced controls settings
        # advnc_mode = rs.rs400_advanced_mode(device)
        # advnc_mode.load_json(json_string)

        # DEPTH SENSOR CONFIGURATION
        # these are the stereo module main parameters
        depth_sensor = device.first_depth_sensor()

        # # set visual preset -- 0=Custom, 1=Default, 2=Hand, 3=High Accuracy, 4=High Density
        # preset_range = depth_sensor.get_option_range(rs.option.visual_preset)
        # for i in range(int(preset_range.max)):
        #     visual_preset = depth_sensor.get_option_value_description(rs.option.visual_preset, i)
        #     if visual_preset == args.visual_preset:
        #         depth_sensor.set_option(rs.option.visual_preset, i)

        # current_preset = depth_sensor.get_option(rs.option.visual_preset)
        # current_visual_preset = depth_sensor.get_option_value_description(rs.option.visual_preset, current_preset)
        # print("Depth visual preset:", current_visual_preset)

        # depth_sensor.set_option(rs.option.enable_auto_exposure, True)
        # depth_sensor.set_option(rs.option.emitter_enabled, 1)  # 1=Laser is the default
        # depth_sensor.set_option(rs.option.hdr_enabled, True)  # DO NOT USE, it makes the image flash

        try:
            # COLORMAPS
            # these are the depth visualization parameters
            colorizer = rs.colorizer()
            colorizer.set_option(rs.option.color_scheme, 0)  # 0 is Jet
            # colorizer.set_option(rs.option.visual_preset, 1)  # 0=Dynamic, 1=Fixed, 2=Near, 3=Far
            value_min = 0
            value_max = 6
            colorizer.set_option(rs.option.min_distance, value_min)
            colorizer.set_option(rs.option.max_distance, value_max)
            colorizer.set_option(rs.option.histogram_equalization_enabled, True)

            # POST PROCESSING FILTERS
            decimation_filter = rs.decimation_filter(magnitude=1)  # Performs downsampling by using the median with specific kernel size
            threshold_filter = rs.threshold_filter(min_dist=0, max_dist=5.2)  # filter out depth values that are either too large or too small, as a software post-processing step
            depth_to_disparity_filter = rs.disparity_transform(transform_to_disparity=True)  # Converts from depth representation to disparity representation and vice
            spatial_filter = rs.spatial_filter(smooth_alpha=0.5, smooth_delta=20, magnitude=2, hole_fill=2)
            # Spatial filter smooths the image by calculating frame with
            # alpha and delta settings. Alpha defines the weight of the current
            # pixel for smoothing, and is bounded within [25..100]%. Delta
            # defines the depth gradient below which the smoothing will occur
            # as number of depth levels.
            temporal_filter = rs.temporal_filter(smooth_alpha=0.4, smooth_delta=20, persistence_control=3)  # persistence_control=3 - Valid in 2 / last 4 - Activated if the pixel was valid in two out of the last 4 frames
            # Temporal filter smooths the image by calculating multiple frames
            # with alpha and delta settings. Alpha defines the weight of
            # current frame, and delta defines thethreshold for edge
            # classification and preserving.
            disparity_to_depth_filter = rs.disparity_transform(transform_to_disparity=False)  # Converts from depth representation to disparity representation and vice

            # Streaming loop
            while True:
                # Get frameset
                frames = pipeline.wait_for_frames()
                curr_pos = playback.get_position()

                # DEPTH
                depth_frame = frames.get_depth_frame()
                if not depth_frame:
                    print("No depth_frame")
                    continue

                # Apply filters to the depth channel
                filtered_depth = depth_frame
                filtered_depth = decimation_filter.process(filtered_depth)
                filtered_depth = threshold_filter.process(filtered_depth)
                filtered_depth = depth_to_disparity_filter.process(filtered_depth)
                filtered_depth = spatial_filter.process(filtered_depth)
                filtered_depth = temporal_filter.process(filtered_depth)
                filtered_depth = disparity_to_depth_filter.process(filtered_depth)

                # Colorize depth frame to jet colormap
                depth_color_frame = colorizer.colorize(filtered_depth)
                # Convert depth_frame to numpy array to render image in opencv
                depth_color_image = np.asanyarray(depth_color_frame.get_data())
                # Save to disk
                depthwriter.write(depth_color_image)
                # Render image in opencv window
                cv2.imshow('Depth', depth_color_image)

                print("Progress:", f"{curr_pos}/{duration}")
                if curr_pos >= duration:
                    print("End of recording reached")
                    break

                # if pressed escape exit program
                if cv2.waitKey(1) in [27, ord("q")]:
                    break
        finally:
            cv2.destroyAllWindows()
            depthwriter.release()
            pipeline.stop()
            print(filename + " done!")

if __name__ == '__main__':
    parser = get_parser()
    args = get_args(parser)
    print(args)
    main(args)


text:
"""
forecast.py
Ross Spicer
created: 2015/09/18

    forecast
"""
from community_data import CommunityData
from diagnostics import Diagnostics
import constants
#~ import plot
#~ import colors

import numpy as np
from pandas import DataFrame, read_csv, concat
import os.path
#~ import copy

#~ from datetime import datetime

class Forecast (object):
    """ Class doc """

    def __init__ (self, community_data, diag = None,
                        scalers = {'kWh consumption':1.0}):
        """
        pre:
            self.cd is a community_data instance.
        post:
            self.start_year < self.end_year are years(ints)
            self.cd is a community_data instance.
        """
        self.diagnostics = diag
        if self.diagnostics == None:
            self.diagnostics = Diagnostics()
        self.cd = community_data

        self.forecast_population()

        #~ ## test block
        #~ self.forecast_consumption(scalers['kWh consumption'])
        #~ self.forecast_generation()
        #~ return

        if self.cd.get_item("community","model electricity") is False:
            pass
        else:
            try:
                self.forecast_consumption(scalers['kWh consumption'])
                self.forecast_generation()

                self.forecast_average_kW()
            except RuntimeError:
                self.cd.set_item("community", "model financial", False)
                self.cd.set_item("community", "model electricity", False)
                self.diagnostics.add_warning('Forecast',
                    (
                        'Electricity data was insufficient to model'
                        ' electricity or financial values. In community data'
                        ' model financial, and model electricity have been'
                        ' set to false.'
                    )
                )

        self.calc_average_diesel_load()

        self.cpi = np.array(self.cd.get_item('community','cpi multipliers'))
            #~ DataFrame([
            #~ 1, 0.9765625, 0.953674316, 0.931322575, 0.909494702, 0.88817842,
            #~ 0.867361738, 0.847032947, 0.827180613, 0.807793567, 0.788860905,
            #~ 0.770371978, 0.752316385, 0.734683969, 0.717464814, 0.700649232,
            #~ 0.684227766, 0.668191178, 0.652530447, 0.637236765, 0.622301528,
            #~ 0.607716336, 0.593472984, 0.579563461, 0.565979942, 0.552714788,
            #~ 0.539760535, 0.527109897, 0.514755759, 0.502691171, 0.490909347,
            #~ 0.479403659, 0.468167636, 0.457194957, 0.44647945, 0.436015088,
            #~ 0.425795984, 0.415816391, 0.406070694, 0.396553412, 0.387259192,
            #~ 0.378182804, 0.369319145, 0.360663227, 0.352210183, 0.343955257,
            #~ 0.335893805, 0.328021294, 0.320333295, 0.312825484, 0.305493636,
            #~ 0.298333629, 0.291341435, 0.28451312, 0.277844844, 0.271332855,
            #~ 0.264973491, 0.258763175, 0.252698413, 0.246775794, 0.240991987],
        #~ )
        #~ last_cpi_year = self.cpi.index.tolist()[-1]
        #~ if self.end_year < last_cpi_year:
            #~ self.cpi = self.cpi[:self.end_year]
        #~ elif self.end_year > last_cpi_year:
            #~ self.diagnostics.add_note("Forecast",
                    #~ "extending cpi past avaiable data")
            #~ last_cpi = float(self.cpi.ix[last_cpi_year])
            #~ for i in range(last_cpi_year,self.end_year+1):
                #~ self.cpi.ix[i] = last_cpi

        #~ self.forecast_households()
        self.heat_demand_cols = []
        self.heating_fuel_cols = []
        self.electric_columns = []

    #~ def calc_electricity_values (self):
        #~ """
        #~ pre:
            #~ 'fc_electricity_used' should contain the kWh used for each key type
        #~ post:
            #~ self.electricity_totals is a array of yearly values of total kWh used
        #~ """
        #~ kWh = self.fc_specs["electricity"]
        #~ print kWh
        #~ years = kWh.T.keys().values
        #~ self.yearly_res_kWh = DataFrame({"year":years,
                          #~ "total":kWh['consumption residential'].values}).set_index("year")
        #~ self.yearly_total_kWh = DataFrame({"year":years,
                          #~ "total":kWh['consumption'].values}).set_index("year")
        #~ self.average_nr_kWh = kWh['consumption non-residential'].values[-3:].mean()
        #~ if np.isnan(self.average_nr_kWh):
            #~ temp = kWh['consumption non-residential']
            #~ self.average_nr_kWh = temp[ np.logical_not(np.isnan(temp))].mean()
        #~ self.yearly_nr_kWh = DataFrame({"year":years,
                          #~ "total":kWh['consumption non-residential'].values}).set_index("year")
        #~ print self.average_nr_kWh

    def forecast_population (self):
        """
        pre:
            tbd.
        post:
            self.population is a array of estimated populations for each
        year between start and end
        """
        # pop forecast is preprocessed now
        population = self.cd.get_item('community',"population")

        self.p_map = DataFrame(population["population_qualifier"])
        self.population = DataFrame(population["population"])

        #~ last_pop_year = self.population.index.tolist()[-1]
        #~ if self.end_year < last_pop_year:
            #~ self.population = self.population[:self.end_year]
        #~ elif self.end_year > last_pop_year:
            #~ self.diagnostics.add_note("Forecast",
                    #~ "extending popultation past avaiable data")
            #~ last_pop = int(round(self.population.ix[last_pop_year]))
            #~ for i in range(last_pop_year,self.end_year+1):
                #~ self.population.ix[i] = last_pop

        if len(self.p_map[self.p_map == "M"] ) < 10:
            msg = "the data range is < 10 for input population "\
                  "check population.csv in the models data directory"
            self.diagnostics.add_warning("forecast", msg)

    def forecast_consumption (self, consumption_scaler = 1.0):
        """
        pre:
            tbd.
        post:
            self.consumption is a array of estimated kWh consumption for each
        year between start and end
        """
        total_consumption = \
            self.cd.get_item('community','utility info')['consumption'] \
            * consumption_scaler
        residential_consumption = \
            self.cd.get_item('community','utility info')\
            ['consumption residential'] * consumption_scaler
        non_residential_consumption = \
            self.cd.get_item('community','utility info')\
            ['consumption non-residential'] * consumption_scaler

        #~ print self.cd.get_item('community','utility info')

        index = list(residential_consumption.index)

        if len(residential_consumption) < 10:
            msg = "the data range is < 10 for input consumption "\
                  "check electricity.csv in the models data directory"
            self.diagnostics.add_warning("Forecast: consumption", msg)

        population = self.population.ix[index]
        if any(population.isnull()):
            v= population.isnull().values.T.tolist()[0]
            for year in population[v].index.values.tolist():
                index.remove(year)

        if any(residential_consumption.isnull()):
            to_remove = \
                residential_consumption[residential_consumption.isnull()].index
            for year in to_remove:
                #~ print year
                index.remove(year)

        self.diagnostics.add_note("forecast",
           "years with measured consumption and population " + str(index) +\
        ". List used to generate fit function")
        population = self.population.ix[index]['population'].values

        #~ self.measured_consumption = self.yearly_total_kWh.ix[idx]
        consumption = residential_consumption.ix[index].values

        if len(population) < 10:
            self.diagnostics.add_warning("forecast",
                  "the data range is < 10 matching years for "\
                  "population and consumption "\
                  "check population.csv and electricity.csv "\
                  "in the models data directory")

        try:
            #~ print population,consumption
            m, b = np.polyfit(population,consumption,1)
        except TypeError:
            raise RuntimeError, "Known population & consumption do not overlap"

        forcasted_consumption = (m * self.population + b) #+ self.average_nr_kWh
        forcasted_consumption.columns = ['consumption residential']

        forcasted_consumption['consumption_qualifier'] = "P"

        forcasted_consumption['consumption residential']\
            [residential_consumption.ix[index].index] = \
            residential_consumption.ix[index]

        if consumption_scaler == 1.0:
            ## if the scaler is not 1 then none of the values are really
            ## measured
            forcasted_consumption['consumption_qualifier'].ix[index] = \
                (residential_consumption * np.nan).fillna("M")

        mean_non_res_con  = non_residential_consumption.values[-3:].mean()
        if np.isnan(mean_non_res_con):
            self.diagnostics.add_note('Forecast',
                ('Average of last 3 known years was null,'
                ' averaging all non null years')
            )
            mean_non_res_con = \
                non_residential_consumption\
                [ ~non_residential_consumption.isnull()].mean()

        #~ print non_residential_consumption.values
        #~ print mean_non_res_con

        forcasted_consumption['consumption non-residential'] = mean_non_res_con
        #~ print forcasted_consumption
        forcasted_consumption['consumption non-residential']\
            [non_residential_consumption.ix[index].index] = \
            non_residential_consumption.ix[index]

        forcasted_consumption['consumption'] = \
            forcasted_consumption['consumption non-residential'] + \
            forcasted_consumption['consumption residential']

        #~ forcasted_consumption['consumption non-residential']\
            #~ [total_consumption.ix[index].index] = total_consumption.ix[index]

        ## don't forecast backwards
        forcasted_consumption = \
            forcasted_consumption.ix[residential_consumption.index[0]:]
        self.consumption = forcasted_consumption
        #~ print self.consumption

    def forecast_generation (self):
        """
        forecast generation
        pre:
            tbd.
            self.consumption should be a float array of kWh/yr values
        post:
            self.generation is a array of estimated kWh generation for each
        year between start and end
        """
        measured_generation = \
            self.cd.get_item('community','utility info')['net generation']
        generation = self.consumption['consumption']/\
                    (1.0-self.cd.get_item('community','line losses')/100.0)
        generation = DataFrame(generation)
        generation.columns = ['generation']
        generation = generation.round(-3)

        index = set(measured_generation.index) & set(generation.index)

        generation['generation'][index] = measured_generation.ix[index]

        types = [u'generation diesel', u'generation hydro',
                u'generation natural gas', u'generation wind',
                u'generation solar', u'generation biomass'
            ]
        generation_types = self.cd.get_item('community','utility info')[types]

        backup_type = 'generation diesel'
        if (generation_types[['generation natural gas']].fillna(0)\
            != 0).any().bool():
            backup_type = 'generation natural gas'

        ## forecast each type
        for fuel in types:
            current_fuel = generation_types[fuel]
            #~ print current_rfuel
            name = fuel.replace('generation ','')
            ## get the hypothetical limit
            try:
                current_generation = self.cd.get_item('community',
                    name + ' generation limit')
            except KeyError:
                current_generation = 0
            ## is there a non-zero limit defined
            if current_generation != 0:

                if fuel == 'generation wind':
                    measured = \
                        current_fuel[current_fuel.notnull()].values[-3:].mean()
                    max_wind_percent = .2
                    if measured > (current_generation * max_wind_percent):
                        pass
                    else:
                        generation[fuel] = current_generation * max_wind_percent
                        generation[fuel].ix[index] = current_fuel.ix[index]
                        continue
                else:
                    generation[fuel] = current_generation
                    generation[fuel].ix[index] = current_fuel.ix[index]
                    continue
            else:
                current_generation = \
                    current_fuel[current_fuel.notnull()].values[-3:].mean()

            ## no generation found
            if current_generation == 0:
                generation[fuel] = current_generation
                continue

            if fuel in [ u'generation wind', u'generation hydro' ]:
                generation[fuel] = current_generation
                generation[fuel].ix[index] = current_fuel.ix[index]
                for i in generation[fuel].ix[current_fuel.index[-1]+1:].index:
                     generation[fuel].ix[i] = \
                        generation[fuel].ix[i-3:i-1].mean()
            else:
                generation[fuel] = current_generation
                generation[fuel].ix[index] = current_fuel.ix[index]

            #~ print fuel, current_generation
        #~ print generation['generation hydro']
        ## scale back any excess
        for fuel in ['generation wind','generation hydro','generation biomass']:
            current_fuel = generation_types[fuel]
            fuel_sums = generation[types].sum(1)
            #~ print fuel_sums
            if any(generation['generation'] < fuel_sums):
                #~ print 'a'
                if any(generation['generation'] < generation[fuel]):
                    #~ print 'b'
                    msg = "scaling generation(" + fuel + ") where geneation(" + \
                            fuel + ") > total generation"
                    self.diagnostics.add_note('forecast', msg)

                    generation[fuel] = generation[fuel] -\
                            (fuel_sums - generation['generation'])

                    generation[fuel][generation[fuel] < 0 ] = 0

                    generation[fuel].ix[index] = current_fuel.ix[index]

        corrected_backup_type = generation['generation'] - \
            (generation[types].sum(1) - generation[backup_type])
        #~ print corrected_backup_type
        corrected_backup_type[corrected_backup_type < 0] = 0

        generation[backup_type ] = corrected_backup_type
        generation[backup_type].ix[index] = generation_types[backup_type].ix[index]

        self.generation = generation

    def forecast_average_kW (self):
        """
        forecast the average kW used per community per year
        pre:
        post:
        """
        self.average_kW = (self.consumption['consumption']/ 8760.0)\
                                /(1-self.cd.get_item('community','line losses'))

        self.average_kW.columns = ["average load"]

    def forecast_households (self):
        """
        forcast # of houselholds
        pre:
        post:
        """
        peeps_per_house = float(self.base_pop) / \
            self.cd.get_item('Residential Buildings',
                'data').ix['total_occupied']
        self.households = \
            np.round(self.population / np.float64(peeps_per_house))
        self.households.columns = ["households"]

    def get_population (self, start, end = None):
        """
        get population values from the population forecast.

        pre:
            start is a year where start >= self.start_year
            end(if provided) is a year where start <= end < self.end_year
        post:
            returns a float or list of floats
        """
        if end is None:
            return self.population.ix[start].T.values[0]

        ## dynamic extension
        existing_len = len(self.population.ix[start:end])
        extend_by = (end + 1) - start - existing_len
        if extend_by > 0:
            extend = DataFrame(
                index=range(
                    self.population.index[-1] + 1,
                    self.population.index[-1] + 1+extend_by
                ),
                columns=['population'])
            extend['population'] = self.population.iloc[-1]['population']
            population = self.population.ix[start:end].append(extend)

        else:
            # -1 to ensure same behavour
            population = self.population.ix[start:end]
        return population['population'].values

    def get_consumption (self, start, end = None):
        """
        get consumption values from the consumption forecast.

        pre:
            start is a year where start >= self.start_year
            end(if provided) is a year where start <= end < self.end_year
        post:
            returns a float or list of floats
        """
        if end is None:
            return self.consumption.ix[start]['consumption']

        ## dynamic extension
        existing_len = len(self.consumption.ix[start:end])
        extend_by = (end + 1) - start - existing_len
        if extend_by > 0:
            extend = DataFrame(
                index=range(
                    self.consumption.index[-1] +1,
                    self.consumption.index[-1]+extend_by+1
                ),
                columns=['consumption'])
            extend['consumption'] = self.consumption.iloc[-1]['consumption']
            consumption = \
                DataFrame(self.consumption.ix[start:end]['consumption']).\
                append(extend)

        else:
            #  -1 to ensure same behavour
            consumption = \
                DataFrame(self.consumption['consumption'].ix[start:end])
        return consumption['consumption'].values

    def get_generation (self, start, end = None):
        """
        get population values from the population forecast.

        pre:
            start is a year where start >= self.start_year
            end(if provided) is a year where start <= end < self.end_year
        post:
            returns a float or list of floats
        """
        if end is None:
            return self.generation.ix[start]['generation']

        ## dynamic extension
        existing_len = len(self.generation.ix[start:end])
        extend_by = (end + 1) - start - existing_len
        if extend_by > 0:
            extend = DataFrame(
                index=range(
                    self.generation.index[-1] + 1 ,
                    self.generation.index[-1]+extend_by +1
                ),
                columns=['generation'])
            extend['generation'] = self.generation.iloc[-1]['generation']
            generation = \
                DataFrame(self.generation.ix[start:end]['generation']).\
                append(extend)

        else:
            generation = \
                DataFrame(self.generation['generation'].ix[start:end])
        return generation['generation'].values

    def get_households (self, start, end = None):
        """
        get households values from the households forecast.

        pre:
            start is a year where start >= self.start_year
            end(if provided) is a year where start <= end < self.end_year
        post:
            returns a float or list of floats
        """
        if end is None:
            return self.households.ix[start].T.values[0]

        ## dynamic extension
        existing_len = len(self.households.ix[start:end])
        extend_by = (end + 1) - start - existing_len
        if extend_by > 0:
            extend = DataFrame(
                index=range(
                    self.households.index[-1] + 1,
                    self.households.index[-1]+extend_by +1
                ),
                columns=['households'])
            extend['households'] = self.households.iloc[-1]['households']
            households = self.households.ix[start:end].append(extend)

        else:
            # -1 to ensure same behavour
            households = self.households.ix[start:end]
        return households['households'].values

    #~ def save_forecast (self, path, png_path = None, do_plots = False):
        #~ """
        #~ save the forecast to a csv
        #~ pre:
            #~ everything needs to be foretasted
            #~ path: path to each communities output dir
            #~ png_path: altenat location for pngs
        #~ post:
            #~ saves 3 files
        #~ """
        #~ tag = self.cd.get_item("community", "name").replace(" ", "_") + "_"
        #~ pathrt = os.path.join(path, tag)
        #~ if png_path is None:
            #~ if do_plots:
                #~ os.makedirs(os.path.join(path,"images"))
            #~ png_path = os.path.join(path, 'images',tag)
            #~ epng_path = png_path
            #~ hdpng_path = png_path
            #~ hfpng_path = png_path;
            #~ gpng_path = png_path
        #~ else:
            #~ epng_path = os.path.join(png_path,'electric_forecast',tag)
            #~ try:
                #~ if do_plots:
                    #~ os.makedirs(os.path.join(png_path,'electric_forecast'))
            #~ except OSError:
                #~ pass
            #~ gpng_path = os.path.join(png_path,'generation_forecast',tag)
            #~ try:
                #~ if do_plots:
                    #~ os.makedirs(os.path.join(png_path,'generation_forecast'))
            #~ except OSError:
                #~ pass
            #~ hdpng_path = os.path.join(png_path,'heat_demand_forecast',tag)
            #~ try:
                #~ if do_plots:
                    #~ os.makedirs(os.path.join(png_path,'heat_demand_forecast'))
            #~ except OSError:
                #~ pass
            #~ hfpng_path = os.path.join(png_path,'heating_fuel_forecast',tag)
            #~ try:
                #~ if do_plots:
                    #~ os.makedirs(os.path.join(png_path,'heating_fuel_forecast'))
            #~ except OSError:
                #~ pass
        #~ if self.cd.intertie != 'child' and \
            #~ self.cd.get_item("community","model electricity"):

            ##start = datetime.now()
            #~ self.save_electric(pathrt, epng_path, do_plots)
            #~ self.save_generation_forecast(pathrt, gpng_path, do_plots)
            ##print "saving electric:" + str(datetime.now() - start)
        #~ if self.cd.intertie != 'parent' and \
            #~ self.cd.get_item("community","model heating fuel"):

            ##start = datetime.now()
            #~ self.save_heat_demand(pathrt, hdpng_path, do_plots)
            ##print "saving heat demand:" + str(datetime.now() - start)

            ##start = datetime.now()
            #~ self.save_heating_fuel(pathrt, hfpng_path, do_plots)
            ##print "saving heating fuel:" + str( datetime.now() - start)

    #~ def add_heat_demand_column (self, key, year_col, data_col):
        #~ """
        #~ add a column to be saved with the heat demand forecast
        #~ pre:
            #~ Key is a string. year_col is a numpy array of values, date_col
        #~ is a corresponding array of years
        #~ post:
            #~ a dataframe is added to self.heat_demand_cols
        #~ """
        #~ self.heat_demand_cols.append(DataFrame({"year":year_col,
                                           #~ key:data_col}).set_index("year"))

    #~ def add_heating_fuel_column (self, key, year_col, data_col):
        #~ """
        #~ add a column to be saved with the heating fuel forecast
        #~ pre:
            #~ Key is a string. year_col is a numpy array of values, date_col
        #~ is a corresponding array of years
        #~ post:
            #~ a dataframe is added to self.heating_fuel_cols
        #~ """
        #~ self.heating_fuel_cols.append(DataFrame({"year":year_col,
                                           #~ key:data_col}).set_index("year"))

    #~ def generate_electric_output_dataframe (self):
        #~ """
        #~ """
        #~ from copy import deepcopy
        #~ kWh_con = deepcopy(self.consumption)
        #~ kWh_con.columns = ["total_electricity_consumed [kWh/year]",
                           #~ 'residential_electricity_consumed [kWh/year]',
                           #~ 'non-residential_electricity_consumed [kWh/year]']
        #~ c_map = copy.deepcopy(self.c_map)
        #~ c_map.columns = ["total_electricity_consumed_qualifier"]

        #~ kWh_gen = self.generation
        #~ kWh_gen.columns = ["total_electricity_generation [kWh/year]"]
        #~ g_map = copy.deepcopy(c_map)
        #~ g_map.columns = ["total_electricity_generation_qualifier"]

        #~ community = copy.deepcopy(self.c_map)
        #~ community.columns = ["community"]
        #~ community.ix[:] = self.cd.get_item("community","name")
        #~ try:
            #~ self.electric_dataframe = \
                    #~ concat([community,self.population.round().astype(int),
                           #~ self.p_map, kWh_con.round().astype(int), c_map,
                           #~ kWh_gen.round().astype(int), g_map] ,axis=1)
        #~ except ValueError:
            #~ self.electric_dataframe = None

    #~ def save_electric (self, csv_path, png_path, do_plots):
        #~ """
        #~ save the electric forecast

        #~ pre:
            #~ self.population, consumption, and generation. should be dataframes
        #~ as calculated in this class. path should be the path to the output
        #~ directory.
        #~ post:
            #~ a file, electricity_forecast.csv, is save in the output directory.
        #~ """
        #~ self.generate_electric_output_dataframe()
        #~ self.save_electric_csv(csv_path)
        #~ if do_plots:
            #~ self.save_electric_png(png_path)

    #~ def save_electric_csv (self, path):
        #~ """
        #~ save the electric forecast

        #~ pre:
            #~ self.population, consumption, and generation. should be dataframes
        #~ as calculated in this class. path should be the path to the output
        #~ directory.
        #~ post:
            #~ a file, electricity_forecast.csv, is save in the output directory.
        #~ """

        #~ if self.electric_dataframe is None:
            #~ self.diagnostics.add_warning("Forecast",
                                #~ ("when saving null values were found in "
                                 #~ "generation/consumption data. Electricity "
                                 #~ "forecast is suspect not saving summary(csv)"))
            #~ return
        #~ data = self.electric_dataframe
        #~ #not joining path adding file name
        #~ f_name = path + "electricity_forecast.csv"

        #~ fd = open(f_name ,"w")
        #~ fd.write("# Electricity Forecast for " + \
                                    #~ self.cd.get_item("community","name") + "\n")
        #~ fd.write("# Qualifier info: \n")
        #~ fd.write("#   M indicates a measured value \n")
        #~ fd.write("#   P indicates a projected value \n")
        #~ fd.write("#   I indicates a value carried over from the input data. May be projected or measured see input data metadata.\n")

        #~ fd.close()

        #~ data.index = data.index.values.astype(int)
        #~ data.to_csv(f_name, index_label="year", mode = 'a')

    #~ def save_electric_png (self, path):
        #~ """ Function doc """
        #~ if self.electric_dataframe is None:
            #~ self.diagnostics.add_warning("Forecast",
                                #~ ("when saving null values were found in "
                                 #~ "generation/consumption data. Electricity "
                                 #~ "forecast is suspect not saving summary(png)"))
            #~ return

        #~ path = path+ "electricity_forecast.png"
        #~ start = self.c_map[self.c_map['consumption_qualifier'] == 'P'].index[0]
        #~ df2 = self.electric_dataframe[['population',
                            #~ 'total_electricity_consumed [kWh/year]',
                            #~ 'total_electricity_generation [kWh/year]',
                            #~ 'residential_electricity_consumed [kWh/year]',
                            #~ 'non-residential_electricity_consumed [kWh/year]']]
        #~ name = self.cd.get_item("community","name") + ' Electricity Forecast'

        #~ fig, ax = plot.setup_fig(name ,'years','population')
        #~ ax1 = plot.add_yaxis(fig,'kWh')

        #~ c_dict = {'population': [min(1,r*4) for r in colors.red],
                   #~ 'total_electricity_consumed [kWh/year]':colors.orange,
                    #~ 'total_electricity_generation [kWh/year]':colors.green,
                     #~ 'residential_electricity_consumed [kWh/year]':colors.yellow,
                    #~ 'non-residential_electricity_consumed [kWh/year]':colors.violet
        #~ }

        #~ plot.plot_dataframe(ax1,df2,ax,['population'],
                    #~ {'population':'population',
                    #~ 'total_electricity_consumed [kWh/year]':'consumption',
                    #~ 'total_electricity_generation [kWh/year]':'generation',
                    #~ 'residential_electricity_consumed [kWh/year]':'residential consumption',
                    #~ 'non-residential_electricity_consumed [kWh/year]':'non residential consumption'},
                    #~ c_dict)
        #~ fig.subplots_adjust(right=.85)
        #~ fig.subplots_adjust(left=.12)
        #~ plot.add_vertical_line(ax,start, 'forecasting starts' )
        #~ plot.add_horizontal_line(ax1,0)

        #~ plot.create_legend(fig,.20)
        #~ plot.save(fig,path)
        #~ plot.clear(fig)

    #~ def generate_generation_forecast_dataframe(self):
        #~ """
        #~ """
        #~ data = copy.deepcopy(self.generation)
        #~ g_map = copy.deepcopy(self.c_map)
        #~ g_map.columns = ["generation_qualifier"]

        #~ order = []
        #~ try:
            #~ for col in ['generation total', 'generation diesel',
                        #~ 'generation hydro',
                        #~ 'generation natural gas', 'generation wind',
                        #~ 'generation solar', 'generation biomass']:
                #~ if col == 'generation total' and \
                    #~ any(np.isnan(data[col].ix[self.start_year:])):
                    #~ raise ValueError, "nans found"

                #~ kWh = col.replace(" ", "_") + " [kWh/year]"
                #~ mmbtu = col.replace(" ", "_") + " [mmbtu/year]"
                #~ data[kWh] = data[col].fillna(0).round().astype(int)
                #~ data[mmbtu] = (data[col] / constants.mmbtu_to_kWh)\
                                                #~ .fillna(0).round().astype(int)
                #~ del data[col]
                #~ order += [kWh, mmbtu]
            #~ data = concat([self.population.round().astype(int),
                           #~ self.p_map,g_map,data],axis=1)
            #~ data["community"] = self.cd.get_item("community","name")

            #~ self.generation_forecast_dataframe = \
                            #~ data[[data.columns[-1]] + data.columns[:-1].tolist()]
            #~ del data
        #~ except ValueError:
            #~ self.generation_forecast_dataframe = None

    #~ def save_generation_forecast (self, csv_path, png_path, do_plots):

        #~ self.generate_generation_forecast_dataframe()

        #~ if self.generation_forecast_dataframe is None:
            #~ self.diagnostics.add_warning("Forecast",
                        #~ ("when saving null values were found in "
                         #~ "generation data. generation generation "
                         #~ "forecast is suspect not saving summary(csv and png)"))
        #~ else:
            #~ self.save_generation_forecast_csv(csv_path)
            #~ if do_plots:
                #~ self.save_generation_forecast_png(png_path)

    #~ def save_generation_forecast_csv (self, path):
        #~ """
        #~ """
        #~ out_file = path + "generation_forecast.csv"
        #~ fd = open(out_file, 'w')
        #~ fd.write("# Generation forecast\n")
        #~ fd.write(("# The column 'generation_qualifier' applies to all "
                                #~ "of the following generation columns"))
        #~ fd.write("# Qualifier info: \n")
        #~ fd.write("#   M indicates a measured value \n")
        #~ fd.write("#   P indicates a projected value \n")
        #~ fd.write("#   I indicates a value carried over from the input data. May be projected or measured see input data metadata.\n")
        #~ fd.close()
        #~ self.generation_forecast_dataframe.to_csv(out_file,
                                        #~ index_label="year", mode = 'a')

    #~ def save_generation_forecast_png (self, path):
        #~ """ Function doc """

        #~ path = path + "generation_forecast.png"

        #~ png_list = ['population',
        #~ 'generation_total [kWh/year]',
        #~ 'generation_diesel [kWh/year]',
        #~ 'generation_hydro [kWh/year]',
        #~ 'generation_natural_gas [kWh/year]',
        #~ 'generation_wind [kWh/year]',
        #~ 'generation_solar [kWh/year]',
        #~ 'generation_biomass [kWh/year]']

        #~ png_dict = {'population':'population',
        #~ 'generation_total [kWh/year]':'total',
        #~ 'generation_diesel [kWh/year]':'diesel',
        #~ 'generation_hydro [kWh/year]':'hydro',
        #~ 'generation_natural_gas [kWh/year]':'natural gas',
        #~ 'generation_wind [kWh/year]':'wind',
        #~ 'generation_solar [kWh/year]':'solar',
        #~ 'generation_biomass [kWh/year]':'biomass'}

        #~ c_dict = {'population': [min(1,r*4) for r in colors.red],
        #~ 'generation_total [kWh/year]':colors.jet,
        #~ 'generation_diesel [kWh/year]':colors.red,
        #~ 'generation_hydro [kWh/year]':colors.blue,
        #~ 'generation_natural_gas [kWh/year]':colors.violet,
        #~ 'generation_wind [kWh/year]':[ min(1,r*2) for r in colors.blue],
        #~ 'generation_solar [kWh/year]':colors.orange,
        #~ 'generation_biomass [kWh/year]':colors.avacado}

        #~ temp = []
        #~ for i in png_list:
            #~ if  all(self.generation_forecast_dataframe[i].fillna(0) == 0):
                #~ continue
            #~ temp.append(i)
        #~ png_list = temp
        #~ if len(png_list) == 3:
            #~ png_list = list(set(png_list).\
                    #~ difference(['generation_total [kWh/year]']))

        #~ df2 = self.generation_forecast_dataframe[png_list]
        #~ plot_name = self.cd.get_item("community","name") +' Generation Forecast'

        #~ fig, ax = plot.setup_fig(plot_name ,'years','population')
        #~ ax1 = plot.add_yaxis(fig,'Generation kWh')

        #~ plot.plot_dataframe(ax1,df2,ax,['population'],png_dict,c_dict)
        #~ fig.subplots_adjust(right=.85)
        #~ fig.subplots_adjust(left=.12)
        #~ start = self.c_map[self.c_map['consumption_qualifier'] == 'P'].index[0]
        #~ plot.add_vertical_line(ax,start, 'forecasting starts' )
        #~ plot.add_horizontal_line(ax1,0)
        #~ plot.create_legend(fig,.20)

        #~ plot.save(fig,path)
        #~ plot.clear(fig)

    #~ def generate_heat_demand_dataframe (self):
        #~ """
        #~ """
        #~ data = concat([self.population.round().astype(int), self.p_map] + \
                                                #~ self.heat_demand_cols,axis=1)

        #~ for key in data.keys():
            #~ try:
                #~ col = data[key]
                #~ idx = np.logical_not(col.apply(np.isnan))
                #~ lv = col[idx].tail(1).values[0]
                #~ yr = col[idx].tail(1).index[0]
                #~ col[col.index > yr] = lv
                #~ data[key] = col.round()
            #~ except (TypeError, IndexError, AttributeError):
                #~ pass

        #~ idx = data.keys()[['mmbtu' in s for s in data.keys()]]

        #~ data["heat_energy_demand_total [mmbtu/year]"] = data[idx].sum(1).round()
        #~ data["community"] = self.cd.get_item("community","name")

        #~ self.heat_demand_dataframe = data[[data.columns[-1]] + data.columns[:-1].tolist()]
        #~ del data

    #~ def save_heat_demand (self,csv_path, png_path, do_plots):
        #~ """ Function doc """
        #~ self.generate_heat_demand_dataframe()
        #~ self.save_heat_demand_csv(csv_path)
        #~ if do_plots:
            #~ self.save_heat_demand_png(png_path)

    #~ def save_heat_demand_csv (self, path):
        #~ """
        #~ save the heat demand forecast

        #~ pre:
            #~ self.population should be a dataframe. self.heat_demad_cols should
        #~ be populated using add_heat_demad_column. path should be the path to the
        #~ output directory.
        #~ post:
            #~ a file, heat_demand_forecast.csv, is save in the output directory.
        #~ """
        #~ f_name = path + "heat_demand_forecast.csv"
        #~ data =  self.heat_demand_dataframe
        #~ fd = open(f_name ,"w")
        #~ fd.write("# Heat Demand Forecast for " + \
                                    #~ self.cd.get_item("community","name") + "\n")
        #~ fd.write("# Qualifier info: \n")
        #~ fd.write("#   M indicates a measured value \n")
        #~ fd.write("#   P indicates a projected value \n")
        #~ fd.write("#   I indicates a value carried over from the input data. May be projected or measured see input data metadata.\n")
        #~ fd.close()

        #~ data.index = data.index.values.astype(int)
        #~ data.to_csv(f_name, index_label="year", mode = 'a')

    #~ def save_heat_demand_png(self,path):
        #~ """
        #~ """
        #~ path = path+ "heat_demand_forecast.png"
        ##start = self.c_map[self.c_map['consumption_qualifier'] == 'P'].index[0]

        #~ cols = ['population']
        #~ for col in self.heat_demand_dataframe.columns:
            #~ if "[mmbtu/year]" in col:
                #~ cols += [col]

        #~ df2 = self.heat_demand_dataframe[cols]
        #~ name = self.cd.get_item("community","name") + ' Heat Demand Forecast'

        #~ fig, ax = plot.setup_fig(name ,'years','population')
        #~ ax1 = plot.add_yaxis(fig,'Heat Demand MMBtu')

        #~ c_dict = {'population': [min(1,r*4) for r in colors.red],
        #~ 'heat_energy_demand_residential [mmbtu/year]':colors.yellow,
        #~ 'heat_energy_demand_water-wastewater [mmbtu/year]':colors.cobalt,
        #~ 'heat_energy_demand_non-residential [mmbtu/year]':colors.green,
                        #~ 'heat_energy_demand_total [mmbtu/year]':colors.jet}

        #~ plot.plot_dataframe(ax1,df2,ax,['population'],
            #~ {'population':'population',
        #~ 'heat_energy_demand_residential [mmbtu/year]':'residential',
        #~ 'heat_energy_demand_water-wastewater [mmbtu/year]':'water & wastewater',
        #~ 'heat_energy_demand_non-residential [mmbtu/year]':'non-residential',
                        #~ 'heat_energy_demand_total [mmbtu/year]':'total'},
                        #~ c_dict)
        #~ fig.subplots_adjust(right=.85)
        #~ fig.subplots_adjust(left=.12)
        ##plot.add_vertical_line(ax,start, 'forecasting starts' )
        #~ plot.add_horizontal_line(ax1,0)

        #~ plot.create_legend(fig,.2)
        #~ plot.save(fig,path)
        #~ plot.clear(fig)

    #~ def generate_heating_fuel_dataframe(self):
        #~ """ """
        #~ data = concat([self.population.round().astype(int), self.p_map] + \
                                                #~ self.heating_fuel_cols, axis=1)

        #~ for key in data.keys():
            #~ try:
                #~ col = data[key]
                #~ idx = np.logical_not(col.apply(np.isnan))
                #~ lv = col[idx].tail(1).values[0]
                #~ yr = col[idx].tail(1).index[0]
                #~ col[col.index > yr] = lv
                #~ data[key] = col.round()
            #~ except (TypeError,IndexError, AttributeError):
                #~ pass

        #~ hf_gal_idx = data.keys()[['gallons' in s for s in data.keys()]]
        #~ hf_gal_idx =  hf_gal_idx[['heating_fuel' in s for s in hf_gal_idx]]

        #~ hf_btu_idx = data.keys()[['mmbtu' in s for s in data.keys()]]
        #~ hf_btu_idx =  hf_btu_idx[['heating_fuel' in s for s in hf_btu_idx]]

        #~ data["heating_fuel_total_consumed [gallons/year]"] = \
                                        #~ data[hf_gal_idx].sum(1).round()
        #~ data["heating_fuel_total_consumed [mmbtu/year]"]= \
                                        #~ data[hf_btu_idx].sum(1).round()
        #~ data["community"] = self.cd.get_item("community","name")

        #~ self.heating_fuel_dataframe = \
                        #~ data[[data.columns[-1]] + data.columns[:-1].tolist()]
        #~ del data

    #~ def save_heating_fuel(self, csv_path, png_path, do_plots):
        #~ """ """
        ##start = datetime.now()
        #~ self.generate_heating_fuel_dataframe()
        ##print "saving heating fuel *1:" + str( datetime.now() - start)
        ##start = datetime.now()
        #~ self.save_heating_fuel_csv(csv_path)
        ##print "saving heating fuel *2:" + str( datetime.now() - start)
        ##start = datetime.now()
        #~ if do_plots:
            #~ self.save_heating_fuel_png(png_path)
        ##print "saving heating fuel *3:" + str( datetime.now() - start)

    #~ def save_heating_fuel_csv (self, path):
        #~ """
        #~ save the heating fuel

        #~ pre:
            #~ self.population should be a dataframe. self.heat_demad_cols should
        #~ be populated using add_heating_fuel_column. path should be the path to
        #~ the output directory.
        #~ post:
            #~ a file, heating_fuel_forecast.csv, is save in the output directory.
        #~ """
        #~ f_name = path + "heating_fuel_forecast.csv"
        #~ data = self.heating_fuel_dataframe

        #~ fd = open(f_name ,"w")
        #~ fd.write("# Heating Fuel Forecast for " + \
                                    #~ self.cd.get_item("community","name") + "\n")
        #~ fd.write("# Qualifier info: \n")
        #~ fd.write("#   M indicates a measured value \n")
        #~ fd.write("#   P indicates a projected value \n")
        #~ fd.write("#   I indicates a value carried over from the input data. May be projected or measured see input data metadata.\n")
        #~ fd.close()

        #~ data.index = data.index.values.astype(int)
        #~ data.to_csv(f_name, index_label="year", mode = 'a')

    #~ def save_heating_fuel_png (self, path):
        #~ """"""
        #~ path = path+ "heating_fuel_forecast.png"
        ##start = self.c_map[self.c_map['consumption_qualifier'] == 'P'].index[0]

        #~ cols = ['population']
        #~ for col in self.heating_fuel_dataframe.columns:
            #~ if "[mmbtu/year]" in col:
                #~ cols += [col]

        #~ df2 = self.heating_fuel_dataframe[cols]

        #~ name = self.cd.get_item("community","name") + ' Heating fuel Forecast'

        #~ fig, ax = plot.setup_fig(name ,'years','population')
        #~ ax1 = plot.add_yaxis(fig,'Heating fuel MMBtu')

        #~ c_dict = {'population': [min(1,r*4) for r in colors.red],
            #~ "heating_fuel_residential_consumed [mmbtu/year]":
                    #~ colors.red,
            #~ "cords_wood_residential_consumed [mmbtu/year]":
                    #~ colors.avacado,
            #~ "gas_residential_consumed [mmbtu/year]":
                    #~ colors.violet,
            #~ "electric_residential_consumed [mmbtu/year]":
                    #~ colors.goldenrod,
            #~ "propane_residential_consumed [mmbtu/year]":
                    #~ colors.orange,
            #~ "heating_fuel_water-wastewater_consumed [mmbtu/year]":
                    #~ colors.cobalt,
            #~ "heating_fuel_non-residential_consumed [mmbtu/year]":
                    #~ colors.green,
            #~ "heating_fuel_total_consumed [mmbtu/year]":colors.jet,}

        #~ plot.plot_dataframe(ax1,df2,ax,['population'],
            #~ {"population":'population',
            #~ "heating_fuel_residential_consumed [mmbtu/year]":
                    #~ 'heating oil - residential',
            #~ "cords_wood_residential_consumed [mmbtu/year]":
                    #~ 'wood - residential',
            #~ "gas_residential_consumed [mmbtu/year]":
                    #~ 'natural gas - residential',
            #~ "electric_residential_consumed [mmbtu/year]":
                    #~ 'electric - residential',
            #~ "propane_residential_consumed [mmbtu/year]":
                    #~ 'propane - residential',
            #~ "heating_fuel_water-wastewater_consumed [mmbtu/year]":
                    #~ 'heating oil - water wastewater',
            #~ "heating_fuel_non-residential_consumed [mmbtu/year]":
                    #~ 'heating fuel(all) - non - residential',
            #~ "heating_fuel_total_consumed [mmbtu/year]":'total'},
            #~ c_dict)
        #~ fig.subplots_adjust(right=.85)
        #~ fig.subplots_adjust(left=.12)
        ##plot.add_vertical_line(ax,start, 'forecasting starts' )
        #~ plot.add_horizontal_line(ax1,0)

        #~ plot.create_legend(fig,.30)
        #~ plot.save(fig,path)
        #~ plot.clear(fig)

    def calc_average_diesel_load (self):
        """
        """
        #~ self.average_diesel_load = 0
        try:
            generation = self.generation['generation diesel']
        except AttributeError:
            self.diagnostics.add_error('Forecast',
                                        ('no generation diesel to caclulate'
                                         ' yearly average load. setting to 0'))
            generation = 0.0

        self.yearly_average_diesel_load = generation / constants.hours_per_year

#~ def test ():
    #~ """
    #~ test the forecast
    #~ """
    #~ manley_data = CommunityData("../test_case/input_data/",
                            #~ "../test_case/baseline_results/config_used.yaml")

    #~ fc = Forecast(manley_data)

    #~ return fc


text:
import numpy as np
from sklearn.base import clone
from sklearn.linear_model import LinearRegression, LassoCV
from sklearn.model_selection import GroupKFold
import scipy
from econml.utilities import cross_product
from statsmodels.tools.tools import add_constant

class HeteroDynamicPanelDML:

    def __init__(self, model_t=LassoCV(cv=3),
                 model_y=LassoCV(cv=3),
                 model_final=LinearRegression(fit_intercept=False),
                 n_cfit_splits=3):
        model_t_copy = clone(model_t, safe=False)
        model_y_copy = clone(model_y, safe=False)
        model_final_copy = clone(model_final, safe=False)
        self._model_t_gen = lambda: clone(model_t_copy, safe=False)
        self._model_y_gen = lambda: clone(model_y_copy, safe=False)
        self._model_final_gen = lambda: clone(model_final_copy, safe=False)
        self._n_cfit_splits = n_cfit_splits
        return

    def fit_nuisances(self, Y, T, X, groups, n_periods):
        ''' Fits all the nuisance models and calculates all residuals for each period and information set
        '''
        resT = {}
        resY = {}
        for kappa in np.arange(n_periods):
            resT[kappa] = {}
            resY[kappa] = np.zeros(self._n_train_units)
            for tau in np.arange(kappa, n_periods):
                resT[kappa][tau] = np.zeros(
                    (self._n_train_units,) + T.shape[1:])

        for train, test in GroupKFold(self._n_cfit_splits).split(X, Y, groups):
            inds_train = train[np.arange(train.shape[0]) % n_periods == 0]
            inds_test = test[np.arange(test.shape[0]) % n_periods == 0]
            for kappa in np.arange(n_periods):
                for tau in np.arange(kappa, n_periods):
                    resT[kappa][tau][inds_test // n_periods] = T[inds_test + tau]\
                        - self._model_t_gen().fit(X[inds_train + kappa],
                                                  T[inds_train + tau]).predict(X[inds_test + kappa])
                    resY[kappa][inds_test // n_periods] = Y[inds_test + n_periods - 1]\
                        - self._model_y_gen().fit(X[inds_train + kappa],
                                                  Y[inds_train + n_periods - 1]).predict(X[inds_test + kappa])
        return resT, resY

    def _fit_cov_matrix(self, resT, resY, models):
        ''' Calculates the covariance (n_periods*n_treatments) x (n_periods*n_treatments) matrix for all the parameters
        '''
        n_periods = len(models)
        n_params = 1 if len(resT[0][0].shape) == 1 else resT[0][0].shape[1]
        M = np.zeros((n_periods * n_params,
                      n_periods * n_params))
        Sigma = np.zeros((n_periods * n_params,
                          n_periods * n_params))
        for kappa in np.arange(n_periods):
            # Calculating the (kappa, kappa) block entry (of size n_treatments x n_treatments) of matrix Sigma
            period = n_periods - 1 - kappa
            res_epsilon = (resY[period]
                           - np.sum([models[tau].predict(resT[period][n_periods - 1 - tau])
                                     for tau in np.arange(kappa + 1)], axis=0)).reshape(-1, 1, 1)
            cur_resT = resT[period][period]
            cov_cur_resT = np.matmul(cur_resT.reshape(-1, n_params, 1),
                                     cur_resT.reshape(-1, 1, n_params))
            sigma_kappa = np.mean((res_epsilon**2) * cov_cur_resT, axis=0)
            Sigma[kappa * n_params:(kappa + 1) * n_params,
                  kappa * n_params:(kappa + 1) * n_params] = sigma_kappa

            for tau in np.arange(kappa + 1):
                # Calculating the (kappa, tau) block entry (of size n_treatments x n_treatments) of matrix M
                m_kappa_tau = np.mean(
                    np.matmul(resT[period][n_periods - 1 - tau].reshape(-1, n_params, 1),
                              cur_resT.reshape(-1, 1, n_params)),
                    axis=0)
                M[kappa * n_params:(kappa + 1) * n_params,
                  tau * n_params:(tau + 1) * n_params] = m_kappa_tau
        self._cov = np.linalg.inv(M) @ Sigma @ np.linalg.inv(M).T
        self._M = M
        self._Sigma = Sigma
        return self

    def fit_final(self, Y, T, X, groups, resT, resY, n_periods, hetero_inds):
        ''' Fits the final lag effect models
        '''
        models = {}
        panelX = X.reshape((X.shape[0] // n_periods, n_periods, -1))
        resTX = {}
        for kappa in np.arange(n_periods):
            resTX[kappa] = {}
            for tau in np.arange(kappa, n_periods):
                resTX[kappa][tau] = cross_product(add_constant(panelX[:, tau, hetero_inds], has_constant='add'),
                                                  resT[kappa][tau].reshape(-1, self._n_treatments))
        for kappa in np.arange(n_periods):
            period = n_periods - 1 - kappa
            Y_cal = resY[period].copy()
            if kappa > 0:
                Y_cal -= np.sum([models[tau].predict(resTX[period][n_periods - 1 - tau])
                                 for tau in np.arange(kappa)],
                                axis=0)

            models[kappa] = self._model_final_gen().fit(
                resTX[period][period], Y_cal)

        self._fit_cov_matrix(resTX, resY, models)
        self.final_models = models
        return self

    def fit(self, Y, T, X, groups, hetero_inds=np.empty(shape=(0,))):
        self.hetero_inds = hetero_inds.astype(int)
        u_periods = np.unique(np.bincount(groups.astype(int)))
        self._n_train_units = len(np.unique(groups))
        self._n_treatments = 1 if len(T.shape[1:]) == 0 else T.shape[1]
        if len(u_periods) > 1:
            raise AttributeError(
                "Imbalanced panel. Method currently expects only panels with equal number of periods. Pad your data")
        self._n_train_periods = u_periods[0]
        resT, resY = self.fit_nuisances(Y, T, X, groups, self._n_train_periods)
        self.fit_final(Y, T, X, groups, resT, resY,
                       self._n_train_periods, self.hetero_inds)
        return self

    @property
    def param(self):
        return np.array([model.coef_ for key, model in self.final_models.items()]).flatten()

    @property
    def param_cov(self):
        return self._cov

    @property
    def param_stderr(self):
        return np.sqrt(np.diag(self._cov) / self._n_train_units)

    def param_interval(self, alpha=.05):
        return np.array([(scipy.stats.norm.ppf(alpha / 2, loc=param, scale=std),
                          scipy.stats.norm.ppf(1 - alpha / 2, loc=param, scale=std)) if std > 0 else (param, param)
                         for param, std in zip(self.param, self.param_stderr)])

    def _policy_effect_var(self, tau):
        return (tau.reshape(1, -1) @ self.param_cov @ tau.reshape(-1, 1))[0, 0]

    def _policy_effect_stderr(self, tau):
        return np.sqrt(self._policy_effect_var(tau) / self._n_train_units)

    def policy_effect(self, tau, subX, groups, alpha=0.05):
        u_periods = np.unique(np.bincount(groups.astype(int)))
        if len(u_periods) > 1 or u_periods[0] != self._n_train_periods:
            raise AttributeError("Invalid period lengths.")
        n_periods = u_periods[0]
        panelX = subX.reshape((subX.shape[0] // n_periods, n_periods, -1))
        resTX = np.zeros((n_periods, (subX.shape[1] + 1) * self._n_treatments))
        for kappa in np.arange(n_periods):
            resTX[kappa] = np.mean(cross_product(add_constant(panelX[:, n_periods - 1 - kappa, :], has_constant='add'),
                                                 np.tile(tau[n_periods - 1 - kappa].reshape(1,
                                                                                self._n_treatments), (panelX.shape[0], 1))), axis=0)

        point = np.dot(self.param, resTX.flatten())
        std = self._policy_effect_stderr(resTX.flatten())
        if std == 0:
            return point, (point, point), 0
        return point, (scipy.stats.norm.ppf(alpha / 2, loc=point, scale=std),
                       scipy.stats.norm.ppf(1 - alpha / 2, loc=point, scale=std)), std

    def adaptive_policy_effect(self, X, groups, policy_gen, alpha=.05):
        """ Assumes that the policy is adaptive only on exogenous states that
        are not affected by the treatmnet.
        """
        u_periods = np.unique(np.bincount(groups.astype(int)))
        if len(u_periods) > 1 or u_periods[0] != self._n_train_periods:
            raise AttributeError("Invalid period lengths.")
        n_periods = u_periods[0]

        panelX = X.reshape((X.shape[0] // n_periods, n_periods, -1))
        tau = np.zeros(
            (panelX.shape[0], n_periods, self._n_treatments))
        for period in range(n_periods):
            if period == 0:
                tau_pre = np.zeros((panelX.shape[0], self._n_treatments))
            else:
                tau_pre = tau[:, period - 1, :]
            tau[:, period, :] = np.array([policy_gen(t_pre, x, period)
                                          for t_pre, x in zip(tau_pre, panelX[:, period, :])])

        resTX = np.zeros(
            (n_periods, (len(self.hetero_inds) + 1) * self._n_treatments))
        for kappa in np.arange(n_periods):
            resTX[kappa] = np.mean(cross_product(add_constant(panelX[:, n_periods - 1 - kappa, self.hetero_inds], has_constant='add'),
                                                 tau[:, n_periods - 1 - kappa, :]), axis=0)

        point = np.dot(self.param, resTX.flatten())
        std = self._policy_effect_stderr(resTX.flatten())
        if std == 0:
            return point, (point, point), 0
        return point, (scipy.stats.norm.ppf(alpha / 2, loc=point, scale=std),
                       scipy.stats.norm.ppf(1 - alpha / 2, loc=point, scale=std)), std

max_stars_repo_name: imatiach-msft/EconML

max_issues_repo_name: imatiach-msft/EconML

max_forks_repo_name: imatiach-msft/EconML

reason: import numpy,import scipy,from statsmodels


text:
#!/usr/bin/env python
# coding: utf-8

# In[79]:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import Lasso
from sklearn.linear_model import Ridge

# In[80]:

### load the dataset

# In[81]:

dataset=pd.read_csv('train.csv',nrows=14000)
print(dataset.head())

# In[82]:

plt.scatter(dataset['Angle of Rays'],dataset['Eclipse Duration (m)'])

# In[ ]:

# In[83]:

### create a plot with gamma and eclipse duration

# In[84]:

plt.scatter(dataset['Gamma'],dataset['Eclipse Duration (m)'])

# In[85]:

### create a plot with magnitude and eclipse duration

# In[86]:

plt.scatter(dataset['Magnitude 1'],dataset['Eclipse Duration (m)'])

# In[87]:

dataset=dataset.drop('Lunation Number',axis=1)

# In[88]:

dataset.head()

# In[89]:

#dataset=dataset.drop('Date of occurrence',axis=1)
dataset=dataset.drop('Day of Month',axis=1)
dataset=dataset.drop('Time of occurrence',axis=1)
dataset=dataset.drop('Time Difference',axis=1)
dataset=dataset.drop('Latitude',axis=1)
dataset=dataset.drop('Longitude',axis=1)

# In[90]:

### remove all unnecessary coluums of data

# In[91]:

dataset.head(10)

# In[92]:

dataset=dataset.drop('Eclipse Type 1',axis=1)

# In[93]:

dataset.head()

# In[94]:

from sklearn.model_selection import train_test_split
x=dataset[['Angle of Rays','Gamma','Magnitude 1']]
y=dataset['Eclipse Duration (m)']
X_train,x_test,Y_train,y_test=train_test_split(x,y,test_size=0.3)

# In[95]:

X_train

# In[96]:

x_test

# In[97]:

clf=LinearRegression()
clf.fit(X_train,Y_train)

# In[98]:

### applied linear regression

# In[99]:

clf.predict(x_test)

# In[100]:

y_test

# In[101]:

clf.score(x_test,y_test)

# In[102]:

clf1=Lasso(alpha=0.001)
clf1.fit(X_train,Y_train)

# In[103]:

clf.predict(x_test)

# In[104]:

clf1.score(x_test,y_test)

# In[105]:

## got the accuracy 81.7 percent in both linear regression lasso regression

# In[106]:

import seaborn as sns; sns.set(style="darkgrid")
sns.relplot(x='Gamma', y='Eclipse Duration (m)', data=dataset)
sns.lmplot(x='Gamma', y='Eclipse Duration (m)', data=dataset)

# In[107]:

sns.relplot(x='Magnitude 1', y='Eclipse Duration (m)', data=dataset)
sns.lmplot(x='Magnitude 1', y='Eclipse Duration (m)', data=dataset)

# In[108]:

sns.relplot(x='Angle of Rays', y='Eclipse Duration (m)', data=dataset)

# In[109]:

from mpl_toolkits import mplot3d
get_ipython().run_line_magic('matplotlib', 'inline')
fig = plt.figure()
ax = plt.axes(projection='3d')
xline=dataset['Gamma']
yline=dataset['Eclipse Duration (m)']
zline=dataset['Magnitude 1']
ax.plot3D(xline, yline, zline, 'blue')

# In[110]:

ax.scatter3D(xline, yline, zline, c=zline, cmap='hsv')

# In[128]:

X=dataset[['Magnitude 1']]
Y=dataset['Eclipse Duration (m)']
X_train,x_test,Y_train,y_test=train_test_split(X,Y,test_size=0.3,shuffle=True)
clf.fit(X_train,Y_train)

# In[112]:

y_test

# In[129]:

clf.score(x_test,y_test)

# In[130]:

clf1.fit(X_train,Y_train)

# In[131]:

clf1.predict(x_test)
clf1.score(x_test,y_test)

# In[132]:

x=dataset[['Angle of Rays','Gamma','Magnitude 1']]
y=dataset['Eclipse Duration (m)']
X_train,x_test,Y_train,y_test=train_test_split(x,y,test_size=0.2)
reg = Ridge(alpha=1)
reg.fit(X_train,Y_train)

# In[133]:

reg.predict(x_test)

# In[134]:

reg.score(x_test,y_test)

# In[ ]:

# In[119]:

dataset.head()

# In[120]:

df=pd.read_csv('test.csv')
df.head()

# In[123]:

x_test121=df[['Angle of Rays','Gamma','Magnitude 1']]
x_test121

# In[135]:

data=x_test121 #an array of dim (188,3)

X=data[:,0:2]
y=data[:,2]
m,n=np.shape(X)
y=y.reshape(m,1)
x=np.c_[np.ones((m,1)),X]
theta=np.zeros((n+1,1))

# In[ ]:

max_stars_repo_path: train_the _model (1).py

max_issues_repo_path: train_the _model (1).py

max_forks_repo_path: train_the _model (1).py


text:
"""
Preprocess module with utilities that aim to uniformize the
naming convention and some other dimension/coordinate operations.

Does not modify the underlying data
"""
import numpy as np
import xarray as xr

def rename_dims(ds):
    """
    Rename dataset dimensions with a more explicit naming scheme
    """
    dims_mapping = {"L": "lead", "Y": "lat", "X": "lon", "M": "member", "S": "time"}
    return ds.rename({k: v for k, v in dims_mapping.items() if k in ds.variables})

def fix_time(ds):
    """
    Fix non-standard calendar to work with cftime
    """
    ds["time"].attrs["calendar"] = "360_day"
    return xr.decode_cf(ds)

def coords_to_int(ds, coords):
    """
    Cast the coordinates to integer values
    """
    coords = [coords] if not isinstance(coords, list) else coords
    for coord in coords:
        _coord_attrs = ds[coord].attrs
        ds[coord] = np.int32(ds[coord])
        ds[coord].attrs = _coord_attrs
    return ds

def capitalize_names(ds):
    """
    Capitalize the 'long_name' attribute for plotting purposes
    """
    for var in ds.variables.values():
        if "long_name" in var.attrs:
            var.attrs["long_name"] = var.attrs["long_name"].capitalize()
    return ds

def preprocess_nmme(ds):
    """
    Core preprocessing function. It combines most of the standalone functions
    defined in nmme_tools.preprocess to provide a unified tool.
    """
    ds = rename_dims(ds)
    ds = fix_time(ds)
    coords = [coord for coord in ["lead", "member"] if coord in ds.variables]
    if coords:
        ds = coords_to_int(ds, coords)
    ds = capitalize_names(ds)
    return ds.sortby("time")