text:
import itertools
import logging
import pathlib
from typing import Tuple

import atlas_mpl_style as ampl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yaml
from seaborn.matrix import heatmap
from sklearn.metrics import (classification_report, confusion_matrix,
                             precision_recall_curve, roc_curve)

import hepynet.common.hepy_type as ht

logger = logging.getLogger("hepynet")

def make_metrics_plot(
    df_raw: pd.DataFrame,
    df: pd.DataFrame,
    job_config: ht.config,
    save_dir: ht.pathlike,
):
    """Plots PR curve."""
    ac = job_config.apply.clone()
    if ac.sample_large_inputs and df.shape[0] > ac.metric_max_events:
        logger.warn(
            f"Too large input detected ({df.shape[0]} rows), randomly sampling {ac.metric_max_events} rows for metrics calculation"
        )
        df = df.sample(n=ac.metric_max_events)
        df_raw =df_raw.loc[df.index]
    train_index = df["is_train"] == True
    test_index = df["is_train"] == False
    y_train = df.loc[train_index, ["y"]].values
    y_train_pred = df.loc[train_index, ["y_pred"]].values
    y_test = df.loc[test_index, ["y"]].values
    y_test_pred = df.loc[test_index, ["y_pred"]].values
    wt_raw_train = df_raw.loc[train_index, "weight"].values  # need raw weights
    wt_raw_test = df_raw.loc[test_index, "weight"].values
    train_inputs = (y_train, y_train_pred, wt_raw_train)
    test_inputs = (y_test, y_test_pred, wt_raw_test)

    save_dir = pathlib.Path(save_dir) / "metrics"
    save_dir.mkdir(parents=True, exist_ok=True)
    if ac.book_confusion_matrix:
        make_confusion_matrix_plot(
            train_inputs, test_inputs, job_config, save_dir
        )
    if ac.book_pr:
        make_pr_curve_plot(train_inputs, test_inputs, job_config, save_dir)
    if ac.book_roc:
        roc_auc = make_roc_curve_plot(train_inputs, test_inputs, job_config, save_dir)
        #roc_auc = {ky: int(val) for ky, val in roc_auc.items()}
        with open(save_dir / "roc_auc.yaml", "w") as f:
            yaml.dump({"roc_auc": roc_auc}, f, default_flow_style=False)

def make_confusion_matrix_plot(
    train_inputs: Tuple[np.ndarray],
    test_inputs: Tuple[np.ndarray],
    job_config: ht.config,
    save_dir: ht.pathlike,
):
    tc = job_config.train.clone()
    ac = job_config.apply.clone()
    score_cut = ac.cfg_confusion_matrix.dnn_cut
    save_dir = pathlib.Path(save_dir) / f"confusion_matrix"
    save_dir.mkdir(exist_ok=True, parents=True)
    # prepare
    output_bkg_node_names = tc.output_bkg_node_names
    output_bkg_node_names = tc.output_bkg_node_names
    all_nodes = ["sig"] + output_bkg_node_names
    y_train, y_train_pred, wt_train = train_inputs
    y_test, y_test_pred, wt_test = test_inputs
    # plot node by node
    row_label = ["Negative Class", "Positive Class"]
    col_label = ["Negative Prediction", "Positive Prediction"]
    for node_num, node in enumerate(all_nodes):
        # train dataset
        con_matrix_train = confusion_matrix(
            y_train[:, node_num],
            y_train_pred[:, node_num] > score_cut,
            sample_weight=wt_train,
        )
        matrix_df = pd.DataFrame(con_matrix_train)
        matrix_df.columns = col_label
        matrix_df.index = row_label
        matrix_df["Total"] = matrix_df.sum(axis=1, numeric_only=True)
        matrix_df.loc["Total"] = matrix_df.sum(numeric_only=True)
        matrix_df.to_csv(save_dir / f"node_{node}_cut_{score_cut}_train.csv")
        fig, ax = plt.subplots()
        heatmap(con_matrix_train, cmap="coolwarm_r", annot=True, ax=ax)
        ax.set_title(f"node {node} - cut {score_cut} - train")
        ax.set_xlabel("Predicted Classes")
        ax.set_ylabel("Real Classes")
        fig.savefig(save_dir / f"node_{node}_cut_{score_cut}_train.png")
        with open(
            save_dir / f"node_{node}_cut_{score_cut}_train_report.txt"
        ) as report_file:
            report = classification_report(
                y_train, y_train_pred[:, node_num] > score_cut
            )
            print(report, file=report_file)
        # test dataset
        con_matrix_test = confusion_matrix(
            y_test[:, node_num],
            y_test_pred[:, node_num] > score_cut,
            sample_weight=wt_test,
        )
        matrix_df = pd.DataFrame(con_matrix_test)
        matrix_df.columns = col_label
        matrix_df.index = row_label
        matrix_df["Total"] = matrix_df.sum(axis=1, numeric_only=True)
        matrix_df.loc["Total"] = matrix_df.sum(numeric_only=True)
        matrix_df.to_csv(save_dir / f"node_{node}_cut_{score_cut}_test.csv")
        fig, ax = plt.subplots()
        heatmap(con_matrix_test, cmap="coolwarm_r", annot=True, ax=ax)
        ax.set_title(f"node {node} - cut {score_cut} - test")
        ax.set_xlabel("Predicted Classes")
        ax.set_ylabel("Real Classes")
        fig.savefig(save_dir / f"node_{node}_cut_{score_cut}_test.png")
        with open(
            save_dir / f"node_{node}_cut_{score_cut}_test_report.txt"
        ) as report_file:
            report = classification_report(
                y_test, y_test_pred[:, node_num] > score_cut
            )
            print(report, file=report_file)

def make_pr_curve_plot(
    train_inputs: Tuple[np.ndarray],
    test_inputs: Tuple[np.ndarray],
    job_config: ht.config,
    save_dir: ht.pathlike,
):
    """Plots PR curve."""
    logger.info("Plotting train/test PR curve.")
    tc = job_config.train.clone()
    ac = job_config.apply.clone()
    # prepare
    output_bkg_node_names = tc.output_bkg_node_names
    output_bkg_node_names = tc.output_bkg_node_names
    all_nodes = ["sig"] + output_bkg_node_names
    y_train, y_train_pred, wt_train = train_inputs
    y_test, y_test_pred, wt_test = test_inputs
    # plot node by node
    num_nodes = len(all_nodes)
    colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    color_cycle = itertools.cycle(colors)
    auc_labels = []
    auc_contents = []
    fig, ax = plt.subplots()
    for node_num in range(num_nodes):
        color = next(color_cycle)
        # plot roc for train dataset without reseting mass
        auc_train, _, _ = plot_single_pr(
            ax,
            y_train[:, node_num],
            y_train_pred[:, node_num],
            wt_train,
            node_num=node_num,
            color=color,
            linestyle="dashed",
        )
        # plot roc for test dataset without reseting mass
        auc_test, _, _ = plot_single_pr(
            ax,
            y_test[:, node_num],
            y_test_pred[:, node_num],
            wt_test,
            node_num=node_num,
            color=color,
            linestyle="solid",
        )
        auc_labels += [
            f"tr-{all_nodes[node_num]} (AUC: {round(auc_train, 5)})",
            f"te-{all_nodes[node_num]} (AUC: {round(auc_test, 5)})",
        ]
        auc_contents += [round(auc_train, 5), round(auc_test, 5)]
    # extra plot config
    ax.legend(auc_labels, loc="lower right")
    ax.grid()
    # collect meta data
    auc_dict = {}
    auc_dict["auc_train_original"] = float(auc_train)
    auc_dict["auc_test_original"] = float(auc_test)
    # make plots
    ax.set_ylim(0, 1.4)
    ax.set_yscale("linear")
    if ac.plot_atlas_label:
        ampl.plot.draw_atlas_label(
            0.05, 0.95, ax=ax, **(ac.atlas_label.get_config_dict())
        )
    ## save linear scale plot
    fig.savefig(save_dir / "pr_linear.png")
    ## log scale x
    ax.set_xlim(1e-5, 1)
    ax.set_xscale("log")
    fig.savefig(save_dir / "pr_logx.png")
    return auc_dict

def make_roc_curve_plot(
    train_inputs: Tuple[np.ndarray],
    test_inputs: Tuple[np.ndarray],
    job_config: ht.config,
    save_dir: ht.pathlike,
):
    """Plots ROC curve."""
    logger.info("Plotting train/test ROC curve.")
    tc = job_config.train.clone()
    ac = job_config.apply.clone()
    # prepare
    output_bkg_node_names = tc.output_bkg_node_names
    output_bkg_node_names = tc.output_bkg_node_names
    all_nodes = ["sig"] + output_bkg_node_names
    y_train, y_train_pred, wt_train = train_inputs
    y_test, y_test_pred, wt_test = test_inputs
    # plot node by node
    num_nodes = len(all_nodes)
    colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    color_cycle = itertools.cycle(colors)
    auc_labels = []
    auc_contents = []
    fig, ax = plt.subplots()
    for node_num in range(num_nodes):
        color = next(color_cycle)
        # plot roc for train dataset without reseting mass
        auc_train, _, _ = plot_single_roc(
            ax,
            y_train[:, node_num],
            y_train_pred[:, node_num],
            wt_train,
            node_num=node_num,
            color=color,
            linestyle="dashed",
        )
        # plot roc for test dataset without reseting mass
        auc_test, _, _ = plot_single_roc(
            ax,
            y_test[:, node_num],
            y_test_pred[:, node_num],
            wt_test,
            node_num=node_num,
            color=color,
            linestyle="solid",
        )
        auc_labels += [
            f"tr-{all_nodes[node_num]} (AUC: {round(auc_train, 5)})",
            f"te-{all_nodes[node_num]} (AUC: {round(auc_test, 5)})",
        ]
        auc_contents += [round(auc_train, 5), round(auc_test, 5)]
    # extra plot config
    ax.legend(auc_labels, loc="lower right")
    ax.grid()
    # collect meta data
    auc_dict = {}
    auc_dict["auc_train_original"] = float(auc_train)
    auc_dict["auc_test_original"] = float(auc_test)
    # make plots
    ax.set_ylim(0, 1.4)
    ax.set_yscale("linear")
    if ac.plot_atlas_label:
        ampl.plot.draw_atlas_label(
            0.05, 0.95, ax=ax, **(ac.atlas_label.get_config_dict())
        )
    ## save linear scale plot
    fig.savefig(save_dir / "roc_linear.png")
    ## log scale x
    ax.set_xlim(1e-5, 1)
    ax.set_xscale("log")
    fig.savefig(save_dir / "roc_logx.png")
    return auc_dict

def plot_single_pr(
    ax: ht.ax,
    y_true: np.ndarray,
    y_pred: np.ndarray,
    weights: np.ndarray,
    node_num: int = 0,
    color: str = "blue",
    linestyle: str = "solid",
    yscal: str = "linear",
    ylim: ht.bound = (0, 1),
):
    """Plots PR (Precision-Recall) curve on given axes."""
    if y_true.ndim == 1:
        y_true = y_true.reshape((-1, 1))
    if y_pred.ndim == 1:
        y_pred = y_pred.reshape((-1, 1))
    # Make plots
    precision, recall, _ = precision_recall_curve(
        y_true[:, node_num], y_pred[:, node_num], sample_weight=weights
    )
    # precision = np.concatenate(([0], precision, [1]))
    # recall = np.concatenate(([1], recall, [0]))
    ax.plot
    ax.plot(recall, precision, color=color, linestyle=linestyle)
    ax.set_title("PR Curve")
    ax.set_xlabel("recall")
    ax.set_ylabel("precision")
    ax.set_ylim(ylim[0], ylim[-1])
    ax.set_yscale(yscal)
    # auc_value = my_roc_auc(y_true[:, node_num], y_pred[:, node_num], weights)
    auc_value = np.trapz(
        y=precision, x=recall
    )  # use Trapezium rule for simplicity
    return auc_value, precision, recall

def plot_single_roc(
    ax: ht.ax,
    y_true: np.ndarray,
    y_pred: np.ndarray,
    weights: np.ndarray,
    node_num: int = 0,
    color: str = "blue",
    linestyle: str = "solid",
    yscal: str = "linear",
    ylim: ht.bound = (0, 1),
):
    """Plots ROC (receiver operating characteristic) curve on given axes."""
    if y_true.ndim == 1:
        y_true = y_true.reshape((-1, 1))
    if y_pred.ndim == 1:
        y_pred = y_pred.reshape((-1, 1))
    # Make plots
    fpr_dm, tpr_dm, _ = roc_curve(
        y_true[:, node_num], y_pred[:, node_num], sample_weight=weights
    )
    fpr_dm = np.concatenate(([0], fpr_dm, [1]))
    tpr_dm = np.concatenate(([0], tpr_dm, [1]))
    ax.plot
    ax.plot(fpr_dm, tpr_dm, color=color, linestyle=linestyle)
    ax.set_title("ROC Curve")
    ax.set_xlabel("false positive rate")
    ax.set_ylabel("true positive rate")
    ax.set_ylim(ylim[0], ylim[-1])
    ax.set_yscale(yscal)
    # auc_value = my_roc_auc(y_true[:, node_num], y_pred[:, node_num], weights)
    auc_value = np.trapz(
        y=tpr_dm, x=fpr_dm
    )  # use Trapezium rule for simplicity
    return auc_value, fpr_dm, tpr_dm

# code from: https://github.com/SiLiKhon/my_roc_auc/blob/master/my_roc_auc.py
def my_roc_auc(
    classes: np.ndarray, predictions: np.ndarray, weights: np.ndarray = None
) -> float:
    """
    Calculating ROC AUC score as the probability of correct ordering
    """

    if weights is None:
        weights = np.ones_like(predictions)

    assert len(classes) == len(predictions) == len(weights)
    assert classes.ndim == predictions.ndim == weights.ndim == 1
    class0, class1 = sorted(np.unique(classes))

    data = np.empty(
        shape=len(classes),
        dtype=[
            ("c", classes.dtype),
            ("p", predictions.dtype),
            ("w", weights.dtype),
        ],
    )
    data["c"], data["p"], data["w"] = classes, predictions, weights

    data = data[np.argsort(data["c"])]
    data = data[
        np.argsort(data["p"], kind="mergesort")
    ]  # here we're relying on stability as we need class orders preserved

    correction = 0.0
    # mask1 - bool mask to highlight collision areas
    # mask2 - bool mask with collision areas' start points
    mask1 = np.empty(len(data), dtype=bool)
    mask2 = np.empty(len(data), dtype=bool)
    mask1[0] = mask2[-1] = False
    mask1[1:] = data["p"][1:] == data["p"][:-1]
    if mask1.any():
        mask2[:-1] = ~mask1[:-1] & mask1[1:]
        mask1[:-1] |= mask1[1:]
        (ids,) = mask2.nonzero()
        correction = (
            sum(
                [
                    ((dsplit["c"] == class0) * dsplit["w"] * msplit).sum()
                    * ((dsplit["c"] == class1) * dsplit["w"] * msplit).sum()
                    for dsplit, msplit in zip(
                        np.split(data, ids), np.split(mask1, ids)
                    )
                ]
            )
            * 0.5
        )

    weights_0 = data["w"] * (data["c"] == class0)
    weights_1 = data["w"] * (data["c"] == class1)
    cumsum_0 = weights_0.cumsum()

    return ((cumsum_0 * weights_1).sum() - correction) / (
        weights_1.sum() * cumsum_0[-1]
    )


text:
""" Generate random data for TimeSeries """

from pySPACE.missions.nodes.base_node import BaseNode
from pySPACE.missions.nodes.source.time_series_source import TimeSeriesSourceNode
from pySPACE.tools.memoize_generator import MemoizeGenerator
from pySPACE.resources.data_types.time_series import TimeSeries

import random
import numpy

class RandomTimeSeriesSourceNode(TimeSeriesSourceNode):
    """ Generate random data and act as a source for windowed TimeSeries

    This node acts as a source for windowed TimeSeries. The TimeSeries
    are generated randomly according to the given parameters and
    forwarded.

    The time series are generated according to the given generating function,
    and the class label by a uniform distribution according with a given threshold
    Only two classes are supported by now.

    **Parameters**

        :num_instances:
            The number of instances to be generated.

            (*optional, default: 20*)

        :generating_function_class_0:
            A function to generate data for class 0.
            Receives an index, which states the
            number of already generated samples.

            (*optional, default: lambda i: numpy.ones((2,2))*i*)

        :generating_function_class_1:
            A function to generate data for class 1.
            Receives an index, which states the
            number of already generated samples.

            (*optional, default: lambda i: numpy.ones((2,2))*i*)

        :channel_names: Channel names of the time series objects.

        :class_labels: The class labels of the generated time series.

        :choice_threshold:
            The threshold class assignment. The classes are
            generated randomly by generating a random number r
            between 0 and 1. If r < threshold, the class label is
            class_labels[0], and class_labels[1] otherwise.

        :sampling_frequency:
            Sampling frequency of the generated time series.

        :random:
            If true, the order of the data is randomly shuffled.

            (*optional, default: True*)

    **Exemplary Call**

    .. code-block:: yaml

        -
            node : RandomTimeSeriesSource

    :Author: Hendrik Woehrle (hendrik.woehrle@dfki.de)
    :Created: 2010/09/22
    """

    def __init__(self, num_instances = 20,
                        generating_function_class_0 = lambda i: numpy.ones((2,2))*i,
                        generating_function_class_1 = lambda i: numpy.ones((2,2))*i,
                        channel_names = ["X", "Y"],
                        class_labels = ['A','B'],
                        class_choice_function = random.random,
                        choice_threshold = 0.33,
                        sampling_frequency = 2,
                        **kwargs):
        super(RandomTimeSeriesSourceNode, self).__init__(**kwargs)
        # We have to create a dummy collection
        class DummyObject(object): pass
        collection = DummyObject()
        collection.meta_data = {'runs' : 1}
        collection.data = {}

        # only binary classification supported by now
        assert( len(class_labels) == 2)

        self.set_permanent_attributes(collection = collection,
                                        num_instances = num_instances,
                                        generating_function_class_0 = generating_function_class_0,
                                        generating_function_class_1 = generating_function_class_1,
                                        channel_names = channel_names,
                                        class_labels = class_labels,
                                        class_choice_function = class_choice_function,
                                        choice_threshold = choice_threshold,
                                        sampling_frequency = sampling_frequency)

    def generate_random_data(self):
        """ Method that is invoked by train and test data generation functions"""
        # invokes the given generating functions
        generated_data = []

        for i in range(self.num_instances):
            choice = self.class_choice_function()
            label = None

            if choice < self.choice_threshold:
                input_array = self.generating_function_class_0(i)
                label = self.class_labels[0]
            else:
                input_array = self.generating_function_class_1(i)
                label = self.class_labels[1]

            generated_data.append( (TimeSeries(input_array = input_array,
                                        channel_names = self.channel_names,
                                        sampling_frequency = self.sampling_frequency ),
                                        label))
        return generated_data

    def request_data_for_testing(self):
        """
        Returns the data that can be used for testing of subsequent nodes

        .. todo:: to document
        """
        # If we haven't read the data for testing yet
        if self.data_for_testing == None:

            generated_data = self.generate_random_data()

            # Create a generator that emits the windows
            test_data_generator = ((sample, label) \
                                     for (sample, label) in generated_data)

            self.data_for_testing = MemoizeGenerator(test_data_generator,
                                                     caching = True)

        # Return a fresh copy of the generator
        return self.data_for_testing.fresh()

    def request_data_for_training(self, use_test_data):
        """
        Returns the data that can be used for testing of subsequent nodes

        .. todo:: to document
        """
        if use_test_data:
            return self.request_data_for_testing()

        # If we haven't read the data for testing yet
        if self.data_for_training == None:

            generated_data = self.generate_random_data()

            # Create a generator that emits the windows
            train_data_generator = ((sample, label) \
                                     for (sample, label) in generated_data)

            self.data_for_training = MemoizeGenerator(train_data_generator,
                                                     caching = True)

        # Return a fresh copy of the generator
        return self.data_for_training.fresh()

    def get_metadata(self, key):
        """ This source node does not contain collection meta data. """
        return None

ext: py

max_stars_repo_name: pyspace/pyspace

max_issues_repo_name: pyspace/pyspace

max_forks_repo_name: pyspace/pyspace

reason: import numpy


text:
#!/usr/bin/python3
import RPi.GPIO as GPIO
import time
import sys
import numpy as np
import logging
from threading import Thread
from .pathcalc import *
import pigpio

position_dict = {}
position_dict["home"] = 8.0
position_dict["graveyard"] = 6.4 - 0.1
position_dict["filter"] = 9.6 + 0.1
#logging.basicConfig(filename="servo/servo.log",filemode="w",format='%(asctime)s -%(levelname)s -  %(message)s',datefmt='%d-%b-%y %H:%M:%S')
# logging.basicConfig(filename="servo/servo.log",level=logging.DEBUG)
class Servo(Exception):
	# brown = black GND
	# orange = red PWR
	# Yellow = white SIG
	sig_pin = 23
	p = 0 #what is p
	pos = position_dict["home"]
	is_moving = False
	round_val = 2
	action = ""
	def  __init__(self,tgt_sig_pin):
		self.sig_pin = tgt_sig_pin
		GPIO.setmode(GPIO.BCM)
		GPIO.setup( self.sig_pin, GPIO.OUT)
		self.p = GPIO.PWM(self.sig_pin,50)
		self.p.ChangeFrequency(50)
		self.p.start(position_dict["home"])
		self.p.ChangeDutyCycle(position_dict["home"])
		self.pos = position_dict["home"]
		self.is_moving = False
		# logging.info("started servo")
		self.start_servo()
	def cleanup(self):
		GPIO.setmode(GPIO.BCM)
		GPIO.setup(self.sig_pin, GPIO.OUT)
		GPIO.output(self.sig_pin, GPIO.LOW)
		GPIO.cleanup()
	def get_pos(self):
		return self.pos
	def set_pos_delay(self,tgt_pos,tgt_delay,tgt_wiggle):

		self.is_moving = True
		cur_diff = tgt_pos - self.pos
		# mov_list =  [round(curpos,self.round_val) for curpos in np.linspace(self.pos,tgt_pos,20)]
		mov_list = get_path(self.pos,tgt_pos,self.round_val)
		self.action = "moving to " + str(mov_list)
		for idx in mov_list:
			self.pos = idx
			self.start_servo()
			self.p.ChangeDutyCycle(self.pos)

			time.sleep(tgt_delay)
			self.stop_servo()
		# at position wiggle back and forth
		# if tgt_wiggle:
		# 	mov_list = get_wiggle_path(self.pos,tgt_pos,self.round_val)
		# 	self.action = "wiggle to " + str(mov_list)
		# 	for mov_pos in mov_list:
		# 		self.pos = mov_pos
		# 		self.p.ChangeDutyCycle(self.pos)
		# 		time.sleep(tgt_delay)
		self.pos = tgt_pos
		self.is_moving = False
		self.stop_servo()
	def move(self,tgt_pos,tgt_delay,tgt_wiggle):
		#same as set_pos_delay but adds threading
		x = Thread(target=self.set_pos_delay, args=(tgt_pos,tgt_delay,tgt_wiggle))
		x.start()
	def wiggle(self):
		# synchronous wiggle
		self.start_servo()
		mov_list = get_wiggle_path(1,self.pos,self.round_val)
		self.action = "wiggle sync " + str(mov_list)
		for idx in mov_list:
			# cur_diff = (idx-3)*0.1
			# tgt_pos = self.pos + cur_diff
			self.p.ChangeDutyCycle(self.pos)
			time.sleep(0.25)
		# self.stop_servo()
	def stop_servo(self):
		self.p.stop()
	def start_servo(self):
		self.p.ChangeFrequency(50)
		self.p.start(0)
		self.p.ChangeDutyCycle(self.pos)
	def set_pos(self,tgt_pos):
		self.start_servo()
		self.is_moving = True
		self.pos = round(tgt_pos,self.round_val)
		self.p.ChangeDutyCycle(self.pos)
		time.sleep(0.5)
		self.is_moving = False
		self.stop_servo()
	def run(self,tgt_cycle):
		self.is_moving = True
		cur_dir = 1
		if self.cur_cycle > tgt_cycle:
			cur_dir = -1
		print(self.cur_cycle,tgt_cycle,cur_dir)
		self.p.ChangeDutyCycle(tgt_cycle)
		time.sleep(0.1)
		self.pos = tgt_cycle
		self.is_moving = False
	def check_moving(self):
		return self.is_moving
	def get_action(self):
		return self.action

max_stars_repo_name: hackmanhattan/perlertools

max_issues_repo_name: hackmanhattan/perlertools

max_forks_repo_name: hackmanhattan/perlertools


text:
import numpy as np
import time
import cProfile, pstats
from subprocess import Popen, PIPE
from lmfit import minimize, Parameters,  __version__

def get_git_version():
    proc = Popen(['git', 'rev-parse', '--short', 'HEAD'], stdout=PIPE)
    return proc.communicate()[0].strip()

# define objective function: returns the array to be minimized
def fcn2min(params, x, data):
    """ model decaying sine wave, subtract data"""
    amp = params['amp']
    shift = params['shift']
    omega = params['omega']
    decay = params['decay']
    model = amp * np.sin(x * omega + shift) * np.exp(-x*x*decay)
    return model - data

def run_fit(nruns=100):
    # create data to be fitted
    x = np.linspace(0, 15, 601)
    np.random.seed(201)
    for i in range(nruns):
        data = (5. * np.sin(2 * x - 0.1) * np.exp(-x*x*0.025) +
                np.random.normal(size=len(x), scale=0.1) )
        params = Parameters()
        params.add('amp',   value= 1.0, min=0.0, max=100.0)
        params.add('decay', value= 0.0, min=-1.0, max=10.0)
        params.add('shift', value= 0.0, min=-np.pi/2., max=np.pi/2)
        params.add('omega', value= 1.0, min=-10.0, max=10.0)
        out = minimize(fcn2min, params, args=(x, data))
        # print out.params['amp']
        assert out.params['amp'].value < 5.5
        assert out.params['amp'].value > 4.5
        assert out.params['omega'].value < 2.25
        assert out.params['omega'].value > 1.75
        # print out.params['amp']

def show_profile(filename):
    stats = pstats.Stats(filename)
    stats.strip_dirs().sort_stats('tottime').print_stats(20)

def profile_command(command, filename=None):
    gitversion = get_git_version()
    if filename is None:
        filename = '%s.prof' % gitversion
    prof = cProfile.run(command, filename=filename)
    show_profile(filename)

profile_command('run_fit()')

max_stars_repo_name: tritemio/lmfit-py

max_issues_repo_name: tritemio/lmfit-py

max_forks_repo_name: tritemio/lmfit-py


text:
"""
Explanation:
https://blog.keras.io/a-ten-minute-introduction-to-sequence-to-sequence-learning-in-keras.html

Original Code:
https://github.com/fchollet/keras/blob/master/examples/lstm_seq2seq.py

Modified by: Chaunte Lacewell

Sequence to sequence example in Keras (character-level).

This script demonstrates how to implement a basic character-level
sequence-to-sequence model. We apply it to translating
short English sentences into short French sentences,
character-by-character. Note that it is fairly unusual to
do character-level machine translation, as word-level
models are more common in this domain.

# Summary of the algorithm:

- We start with input sequences from a domain (e.g. English sentences)
    and correspding target sequences from another domain
    (e.g. French sentences).
- An encoder LSTM turns input sequences to 2 state vectors
    (we keep the last LSTM state and discard the outputs).
- A decoder LSTM is trained to turn the target sequences into
    the same sequence but offset by one timestep in the future,
    a training process called "teacher forcing" in this context.
    Is uses as initial state the state vectors from the encoder.
    Effectively, the decoder learns to generate `targets[t+1...]`
    given `targets[...t]`, conditioned on the input sequence.
- In inference mode, when we want to decode unknown input sequences, we:
    - Encode the input sequence into state vectors
    - Start with a target sequence of size 1
        (just the start-of-sequence character)
    - Feed the state vectors and 1-char target sequence
        to the decoder to produce predictions for the next character
    - Sample the next character using these predictions
        (we simply use argmax).
    - Append the sampled character to the target sequence
    - Repeat until we generate the end-of-sequence character or we
        hit the character limit.

# Data download:

English to French sentence pairs.
http://www.manythings.org/anki/fra-eng.zip

Lots of neat sentence pairs datasets can be found at:
http://www.manythings.org/anki/

# References:

- Sequence to Sequence Learning with Neural Networks
    https://arxiv.org/abs/1409.3215
- Learning Phrase Representations using
    RNN Encoder-Decoder for Statistical Machine Translation
    https://arxiv.org/abs/1406.1078
"""
from __future__ import print_function
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.models import load_model
from tensorflow.keras.layers import Input, LSTM, Dense
from tensorflow.keras.callbacks import EarlyStopping
from nltk.translate.bleu_score import corpus_bleu
import numpy as np
import os
import argparse
import random
import io
from pathlib import Path

datalength = 145437

def vectorize_data(data, sample_indices):
    input_texts = []
    target_texts = []
    input_characters = set()
    target_characters = set()
    # lines = codecs.open(data_path, encoding='utf-8').read().split('\n')
    # for line in lines[: min(num_samples, len(lines) - 1)]:
    # lines = data[: min(num_samples, len(data) - 1)]
    lines = [data[i] for i in sample_indices]
    for line in lines:
        input_text, target_text = line.split('\t')
        # We use "tab" as the "start sequence" character
        # for the targets, and "\n" as "end sequence" character.
        target_text = '\t' + target_text + '\n'
        input_texts.append(input_text)
        target_texts.append(target_text)
        for char in input_text:
            if char not in input_characters:
                input_characters.add(char)
        for char in target_text:
            if char not in target_characters:
                target_characters.add(char)
    return input_texts, target_texts, input_characters, target_characters

def prepare_data(data_path, sample_indices, test_indices, shuffle=True):

    if shuffle:
        # Shuffle lines in data and save
        out_data_path = data_path.split('.txt')[0] + '_shuffled.txt'
        with io.open(data_path, 'r', encoding='utf-8') as f:
            data = [(random.random(), line.splitlines()[0]) for line in f]
            data.sort()
            data = [line[1] for line in data]
        with io.open(out_data_path, 'w', encoding='utf-8') as f:
            f.writelines('\n'.join(data))
    else:
        # Shuffle lines in data
        with io.open(data_path, 'r', encoding='utf-8') as f:
            data = [line.splitlines()[0] for line in f]

    # Vectorize the data.
    input_texts, target_texts, input_characters, target_characters = vectorize_data(data, sample_indices)
    input_texts_test, target_texts_test, input_characters_test, target_characters_test = vectorize_data(data, test_indices)

    input_characters = sorted(list(input_characters))
    target_characters = sorted(list(target_characters))
    input_characters_test = sorted(list(input_characters_test))
    target_characters_test = sorted(list(target_characters_test))
    all_input_characters = sorted(list(set(input_characters + input_characters_test)))
    all_target_characters = sorted(list(set(target_characters + target_characters_test)))
    all_input_texts = input_texts + input_texts_test
    all_target_texts = target_texts + target_texts_test
    num_encoder_tokens = len(all_input_characters)
    num_decoder_tokens = len(all_target_characters)
    max_encoder_seq_length = max([len(txt) for txt in all_input_texts])
    max_decoder_seq_length = max([len(txt) for txt in all_target_texts])

    print('Data path:', data_path)
    print('Number of samples:', len(input_texts))
    print('Number of unique input tokens:', num_encoder_tokens)
    print('Number of unique output tokens:', num_decoder_tokens)
    print('Max sequence length for inputs:', max_encoder_seq_length)
    print('Max sequence length for outputs:', max_decoder_seq_length)

    input_token_index = dict(
        [(char, i) for i, char in enumerate(all_input_characters)])
    target_token_index = dict(
        [(char, i) for i, char in enumerate(all_target_characters)])

    # Train data
    encoder_input_data = np.zeros(
        (len(input_texts), max_encoder_seq_length, num_encoder_tokens),
        dtype='float32')
    decoder_input_data = np.zeros(
        (len(input_texts), max_decoder_seq_length, num_decoder_tokens),
        dtype='float32')
    decoder_target_data = np.zeros(
        (len(input_texts), max_decoder_seq_length, num_decoder_tokens),
        dtype='float32')
    for i, (input_text, target_text) in enumerate(zip(input_texts, target_texts)):
        for t, char in enumerate(input_text):
            encoder_input_data[i, t, input_token_index[char]] = 1.
        for t, char in enumerate(target_text):
            # decoder_target_data is ahead of decoder_input_data by one timestep
            decoder_input_data[i, t, target_token_index[char]] = 1.
            if t > 0:
                # decoder_target_data will be ahead by one timestep
                # and will not include the start character.
                decoder_target_data[i, t - 1, target_token_index[char]] = 1.
    # Reverse-lookup token index to decode sequences back to
    # something readable.
    reverse_input_char_index = dict(
        (i, char) for char, i in input_token_index.items())
    reverse_target_char_index = dict(
        (i, char) for char, i in target_token_index.items())

    # Test data
    encoder_test_data = np.zeros(
        (len(input_texts_test), max_encoder_seq_length, num_encoder_tokens),
        dtype='float32')
    decoder_test_data = np.zeros(
        (len(input_texts_test), max_decoder_seq_length, num_decoder_tokens),
        dtype='float32')
    decoder_test_target_data = np.zeros(
        (len(input_texts_test), max_decoder_seq_length, num_decoder_tokens),
        dtype='float32')
    for i, (input_text_test, target_text_test) in enumerate(zip(input_texts_test, target_texts_test)):
        for t, char in enumerate(input_text_test):
            encoder_test_data[i, t, input_token_index[char]] = 1.
        for t, char in enumerate(target_text_test):
            # decoder_target_data is ahead of decoder_input_data by one timestep
            decoder_test_data[i, t, target_token_index[char]] = 1.
            if t > 0:
                # decoder_target_data will be ahead by one timestep
                # and will not include the start character.
                decoder_test_target_data[i, t - 1, target_token_index[char]] = 1.

    return input_texts, target_texts, num_encoder_tokens, num_decoder_tokens, max_encoder_seq_length,\
        max_decoder_seq_length, input_token_index, target_token_index, encoder_input_data,\
        decoder_input_data, decoder_target_data, reverse_input_char_index, reverse_target_char_index, \
        input_texts_test, target_texts_test, encoder_test_data, decoder_test_data, decoder_test_target_data

# returns train, inference_encoder and inference_decoder models
def define_models(n_input, n_output, n_units):
    # define training encoder
    encoder_inputs = Input(shape=(None, n_input))
    encoder = LSTM(n_units, return_state=True)
    encoder_outputs, state_h, state_c = encoder(encoder_inputs)
    encoder_states = [state_h, state_c]
    # define training decoder
    decoder_inputs = Input(shape=(None, n_output))
    decoder_lstm = LSTM(n_units, return_sequences=True, return_state=True)
    decoder_outputs, _, _ = decoder_lstm(decoder_inputs, initial_state=encoder_states)
    decoder_dense = Dense(n_output, activation='softmax')
    decoder_outputs = decoder_dense(decoder_outputs)
    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
    # define inference encoder
    encoder_model = Model(encoder_inputs, encoder_states)
    # define inference decoder
    decoder_state_input_h = Input(shape=(n_units,))
    decoder_state_input_c = Input(shape=(n_units,))
    decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]
    decoder_outputs, state_h, state_c = decoder_lstm(decoder_inputs, initial_state=decoder_states_inputs)
    decoder_states = [state_h, state_c]
    decoder_outputs = decoder_dense(decoder_outputs)
    decoder_model = Model([decoder_inputs] + decoder_states_inputs, [decoder_outputs] + decoder_states)

    # return all models
    return model, encoder_model, decoder_model

def load_defined_models(n_units):
    # Save all models
    model = load_model('models/seq2seq_model_{}units.h5'.format(n_units))
    encoder_model = load_model('models/seq2seq_encoder_model_{}units.h5'.format(n_units))
    decoder_model = load_model('models/seq2seq_decoder_model_{}units.h5'.format(n_units))

    # return all models
    return model, encoder_model, decoder_model

def predict_sequence(infenc, infdec, in_seq, n_output, target_token_ind,
                     reverse_target_char_ind, max_decoder_seq_len):
    # encode
    state = infenc.predict(in_seq)
    # start of sequence input
    target_seq = np.array([0.0 for _ in range(n_output)]).reshape(1, 1, n_output)
    # Populate the first character of target sequence with the start character.
    target_seq[0, 0, target_token_ind['\t']] = 1.
    # Sampling loop for a batch of sequences
    # (to simplify, here we assume a batch of size 1).
    stop_condition = False
    decoded_sentence = ''
    while not stop_condition:
        # predict next char
        yhat, h, c = infdec.predict([target_seq] + state)

        # Sample a token
        sampled_token_index = np.argmax(yhat[0, 0, :])
        sampled_char = reverse_target_char_ind[sampled_token_index]
        decoded_sentence += sampled_char

        # Exit condition: either hit max length
        # or find stop character.
        if (sampled_char == '\n' or
                len(decoded_sentence) > max_decoder_seq_len):
            stop_condition = True

        # Update the target sequence (of length 1).
        # target_seq = np.array([0.0 for _ in range(n_output)]).reshape(1, 1, n_output)
        # target_seq[0, 0, sampled_token_index] = 1.
        target_seq = yhat

        # Update states
        state = [h, c]

    return decoded_sentence

def main(input_params):
    current_mode = input_params.mode
    batch_size = input_params.batch_size  # Batch size for training.
    epochs = input_params.epochs  # Number of epochs to train for.
    num_units = input_params.num_units  # dimensionality of the output space
    num_samples_inf = input_params.num_inf_samples  # Number of samples for inferencing
    datafilename = input_params.datapath
    shuffle_data = input_params.shuffle_data

    # Path to the data txt file on disk.
    # data_path = os.path.realpath(os.getcwd()) + datafilename
    num_samples = input_params.num_trn_samples  # Number of samples to train on.
    train_indices = range(0, num_samples)
    test_indices = range(num_samples, min(num_samples + num_samples_inf, args.num_lines))
    input_texts, target_texts, num_encoder_tokens, num_decoder_tokens, max_encoder_seq_length,\
        max_decoder_seq_length, input_token_index, target_token_index, encoder_input_data,\
        decoder_input_data, decoder_target_data, reverse_input_char_index,\
        reverse_target_char_index, input_texts_test, target_texts_test, encoder_test_data, \
        decoder_test_data, decoder_test_target_data = prepare_data(datafilename, train_indices,
                                                                   test_indices, shuffle=shuffle_data)

    if current_mode == 'train':
        print("TRAINING====================================")
        model, encoder_model, decoder_model = define_models(num_encoder_tokens, num_decoder_tokens, num_units)

        # Run training
        early_stop = EarlyStopping(monitor='val_loss', verbose=1)
        model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])
        model.fit([encoder_input_data, decoder_input_data], decoder_target_data,
                  batch_size=batch_size, callbacks=[early_stop],
                  epochs=epochs,
                  validation_split=0.2, shuffle=True)

        # Save all models
        os.makedirs('models', exist_ok=True)
        model.save('models/seq2seq_model_{}units.h5'.format(num_units))
        encoder_model.save('models/seq2seq_encoder_model_{}units.h5'.format(num_units))
        decoder_model.save('models/seq2seq_decoder_model_{}units.h5'.format(num_units))

    elif current_mode == 'infer':
        print("INFERENCE====================================")
        model, encoder_model, decoder_model = load_defined_models(num_units)
        in_texts = []
        actual_texts = []
        predicted_texts = []
        for seq_index in range(len(test_indices)):
            # Take one sequence (part of the training test)
            # for trying out decoding.
            input_seq = encoder_test_data[seq_index: seq_index + 1]
            if len(input_seq) <= num_encoder_tokens:
                decoded_sentence = predict_sequence(encoder_model, decoder_model, input_seq, num_decoder_tokens,
                                                    target_token_index, reverse_target_char_index,
                                                    max_decoder_seq_length)

                actual_texts.append(target_texts_test[seq_index])
                in_texts.append(input_texts_test[seq_index])
                predicted_texts.append(decoded_sentence)
                print('INFERENCE SAMPLE {}===================================='.format(seq_index))
                print('Input sentence:', input_texts_test[seq_index])
                print('Decoded sentence:', decoded_sentence)
        # Bleu Scores
        print('Bleu-1: %f' % corpus_bleu(actual_texts, predicted_texts, weights=(1.0, 0, 0, 0)))
        print('Bleu-2: %f' % corpus_bleu(actual_texts, predicted_texts, weights=(0.5, 0.5, 0, 0)))
        print('Bleu-3: %f' % corpus_bleu(actual_texts, predicted_texts, weights=(0.3, 0.3, 0.3, 0)))
        print('Bleu-4: %f' % corpus_bleu(actual_texts, predicted_texts, weights=(0.25, 0.25, 0.25, 0.25)))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", type=str, dest='mode', default='infer',
                        choices=['train', 'infer'],
                        help='Either "train" for training or "infer" for inference')
    parser.add_argument("--batch", type=int, dest='batch_size', default=64,
                        help='Batch size')
    parser.add_argument("--epochs", type=int, dest='epochs', default=100,
                        help='Number of epochs')
    parser.add_argument("--units", type=int, dest='num_units', default=256,
                        help='Number of LSTM units')
    parser.add_argument("--datapath", type=Path, dest='datapath', default='data/fra_shuffled.txt',
                        help='Path to data [default: data/fra_shuffled.txt]')
    parser.add_argument("--shuffle_data", action='store_true',
                        help="Shuffle input data and save")
    parser.add_argument("--train_size", type=int, dest='num_trn_samples', default=10000,
                        help='Number of samples for training')
    parser.add_argument("--infer_size", type=int, dest='num_inf_samples', default=10000,
                        help='Number of samples for training')
    args = parser.parse_args()

    # Only shuffle data for training
    args.shuffle_data = args.shuffle_data if args.mode == 'train' else False
    args.datapath = os.path.join(os.path.realpath(os.getcwd()), str(args.datapath))
    args.num_lines = sum(1 for line in io.open(args.datapath, 'r', encoding='utf-8'))

    main(args)


text:
from __future__ import division
from __future__ import print_function
import numpy as np
import scipy.io.wavfile
import os
import pickle as pkl
import wave
import pyaudio
import argparse
from python_speech_features import mfcc
import sys

def read_tags(fname):
    if fname is None or not os.path.exists(fname):
        print('Input file not found, starting from no tags.')
        return {}
    return pkl.load(open(fname, 'rb'))

def get_correlation(sig1, sig2, rate=16000, cutoff=1.5):
    max_win = int(cutoff*rate)
    sig1_red = sig1[:max_win]
    sig2_red = sig2[:max_win]

    mf1 = mfcc(sig1_red, rate)
    mf2 = mfcc(sig2_red, rate)

    m1 = np.mean(mf1, axis=0)
    c1 = np.cov(mf1.T)
    m2 = np.mean(mf2, axis=0)
    c2 = np.cov(mf2.T)
    return kl_distance(m1, m2, c1, c2)

def kullback_liebler(m1, m2, c1, c2):
    k = m1.shape[0]
    c2_inv_term = np.matmul(np.linalg.inv(c2), c1)
    kl = np.matmul(np.linalg.inv(c2), (m2-m1))
    kl = np.matmul((m2-m1), kl)
    kl = np.matrix.trace(c2_inv_term)+kl-k-np.log(np.linalg.det(c2_inv_term))
    return 0.5*kl

def kl_distance(m1, m2, c1, c2):
    dist = kullback_liebler(m1, m2, c1, c2)
    dist += kullback_liebler(m2, m1, c2, c1)
    return dist

def match_audio(sig, tagged_data, rate=16000, kl_cutoff=35):
    for key, value in tagged_data.items():
        taged_sig = value['sig']
        corr = get_correlation(sig, taged_sig, rate=rate)
        if corr < kl_cutoff:
            return value['tag'], corr
    return None, None

def tag_audio(fname):
    wf = wave.open(fname, 'rb')
    p = pyaudio.PyAudio()

    def callback(in_data, frame_count, time_info, status):
        data = wf.readframes(frame_count)
        return (data, pyaudio.paContinue)

    stream = p.open(format=p.get_format_from_width(wf.getsampwidth()),
                    channels=wf.getnchannels(),
                    rate=wf.getframerate(),
                    output=True, stream_callback=callback)
    stream.start_stream()

    tag = input('Tag this audio (q! for early quitting):')
    if(tag == 'q!'):
        return -1

    stream.stop_stream()
    stream.close()
    wf.close()
    p.terminate()
    if(tag == ''):
        return 0
    return tag

def update_tags(fname, sig, tag, tagged_data, save_name, save=True):
    tagged_data[fname] = {}
    tagged_data[fname]['sig'] = sig
    tagged_data[fname]['tag'] = tag
    if(save):
        pkl.dump(tagged_data, open(save_name, 'wb'))
    return tagged_data

def main(args):
    if(args.directory is None):
        sys.exit('Must provide a directory with wav files')
    directory = args.directory
    fnames = [x for x in os.listdir(directory) if x[-3:] == 'wav']
    tagged_data = read_tags(args.input)
    for fname in fnames:
        fname = directory+fname
        rate, sig = scipy.io.wavfile.read(fname)
        tag, corr = match_audio(sig, tagged_data)
        if tag is None:
            tag = tag_audio(fname)
            if tag != 0 and tag != -1:
                tagged_data = update_tags(fname, sig, tag,
                                          tagged_data, args.output,
                                          args.save)
            if tag == -1:
                break
            print('\n\n')
        else:
            print('Closest tag is {} with a distance of {}'.format(tag,
                                                                   corr))
    pkl.dump(tagged_data, open(args.output, 'wb'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     epilog='')
    parser.add_argument('-i', action='store', dest='input', type=str,
                        help='Audio files already tagged to be used.',
                        default=None)
    parser.add_argument('-o', action='store', dest='output', type=str,
                        help='Output file name.',
                        default='ECHO_tags.pkl')
    parser.add_argument('-d', action='store', dest='directory', type=str,
                        help='Directory containing files to be tagged.',
                        default=None)
    parser.add_argument('-s', action='store_true', dest='save',
                        help='Save after every tag. Enables early quitting.')
    args = parser.parse_args()
    main(args)

max_stars_repo_name: ucl-pond/ESCAPE

max_issues_repo_name: ucl-pond/ESCAPE

max_forks_repo_name: ucl-pond/ESCAPE

reason: import numpy,import scipy


text:
#!/usr/bin/env python

'''
MO integrals in PBC code
'''

import numpy
from pyscf import ao2mo
from pyscf.pbc import gto, scf, tools

cell = gto.M(
    a = numpy.eye(3)*3.5668,
    atom = '''C     0.      0.      0.
              C     0.8917  0.8917  0.8917
              C     1.7834  1.7834  0.
              C     2.6751  2.6751  0.8917
              C     1.7834  0.      1.7834
              C     2.6751  0.8917  2.6751
              C     0.      1.7834  1.7834
              C     0.8917  2.6751  2.6751''',
    basis = 'gth-szv',
    pseudo = 'gth-pade',
    verbose = 4,
)

#
# MO integrals at Gamma point.
#
mf = scf.RHF(cell).run()
eri = mf.with_df.ao2mo(mf.mo_coeff)
#
# These integrals are real. They can be transformed under the 8-fold
# permutation symmetry.
#
eri = ao2mo.restore(8, eri, mf.mo_coeff.shape[1])

#
# MO integrals at arbitrary k-point.
#

nk = [2,2,2]
kpts = cell.make_kpts(nk)

kmf = scf.KRHF(cell, kpts).run()

nmo = kmf.mo_coeff[0].shape[1]
kconserv = tools.get_kconserv(cell, kpts)
nkpts = len(kpts)
for kp in range(nkpts):
    for kq in range(nkpts):
        for kr in range(nkpts):
            ks = kconserv[kp,kq,kr]

            eri_kpt = kmf.with_df.ao2mo([kmf.mo_coeff[i] for i in (kp,kq,kr,ks)],
                                        [kpts[i] for i in (kp,kq,kr,ks)])
            eri_kpt = eri_kpt.reshape([nmo]*4)

max_stars_repo_name: QuESt-Calculator/pyscf

max_issues_repo_name: QuESt-Calculator/pyscf

max_forks_repo_name: QuESt-Calculator/pyscf


text:
import pytest
import numpy as np

from doubleml import DoubleMLPLR, DoubleMLIRM, DoubleMLIIVM, DoubleMLPLIV
from doubleml.datasets import make_plr_CCDDHNR2018, make_irm_data, make_pliv_CHS2015, make_iivm_data

from sklearn.linear_model import Lasso, LogisticRegression

np.random.seed(3141)
dml_data_plr = make_plr_CCDDHNR2018(n_obs=100)
dml_data_pliv = make_pliv_CHS2015(n_obs=100, dim_z=1)
dml_data_irm = make_irm_data(n_obs=100)
dml_data_iivm = make_iivm_data(n_obs=100)

dml_plr = DoubleMLPLR(dml_data_plr, Lasso(), Lasso())
dml_plr.fit()
dml_pliv = DoubleMLPLIV(dml_data_pliv, Lasso(), Lasso(), Lasso())
dml_pliv.fit()
dml_irm = DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression())
dml_irm.fit()
dml_iivm = DoubleMLIIVM(dml_data_iivm, Lasso(), LogisticRegression(), LogisticRegression())
dml_iivm.fit()

# fit models with callable scores
plr_score = dml_plr._score_elements
dml_plr_callable_score = DoubleMLPLR(dml_data_plr, Lasso(), Lasso(),
                                     score=plr_score, draw_sample_splitting=False)
dml_plr_callable_score.set_sample_splitting(dml_plr.smpls)
dml_plr_callable_score.fit(store_predictions=True)

irm_score = dml_irm._score_elements
dml_irm_callable_score = DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression(),
                                     score=irm_score, draw_sample_splitting=False)
dml_irm_callable_score.set_sample_splitting(dml_irm.smpls)
dml_irm_callable_score.fit(store_predictions=True)

iivm_score = dml_iivm._score_elements
dml_iivm_callable_score = DoubleMLIIVM(dml_data_iivm, Lasso(), LogisticRegression(), LogisticRegression(),
                                       score=iivm_score, draw_sample_splitting=False)
dml_iivm_callable_score.set_sample_splitting(dml_iivm.smpls)
dml_iivm_callable_score.fit(store_predictions=True)

@pytest.mark.ci
@pytest.mark.parametrize('dml_obj',
                         [dml_plr, dml_pliv, dml_irm, dml_iivm])
def test_linear_score(dml_obj):
    assert np.allclose(dml_obj.psi,
                       dml_obj.psi_a * dml_obj.coef + dml_obj.psi_b,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_plr_callable_vs_str_score():
    assert np.allclose(dml_plr.psi,
                       dml_plr_callable_score.psi,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_plr.coef,
                       dml_plr_callable_score.coef,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_plr_callable_vs_pred_export():
    preds = dml_plr_callable_score.predictions
    g_hat = preds['ml_g'].squeeze()
    m_hat = preds['ml_m'].squeeze()
    psi_a, psi_b = plr_score(dml_data_plr.y, dml_data_plr.d,
                             g_hat, m_hat,
                             dml_plr_callable_score.smpls[0])
    assert np.allclose(dml_plr.psi_a.squeeze(),
                       psi_a,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_plr.psi_b.squeeze(),
                       psi_b,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_irm_callable_vs_str_score():
    assert np.allclose(dml_irm.psi,
                       dml_irm_callable_score.psi,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_irm.coef,
                       dml_irm_callable_score.coef,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_irm_callable_vs_pred_export():
    preds = dml_irm_callable_score.predictions
    g_hat0 = preds['ml_g0'].squeeze()
    g_hat1 = preds['ml_g1'].squeeze()
    m_hat = preds['ml_m'].squeeze()
    psi_a, psi_b = irm_score(dml_data_irm.y, dml_data_irm.d,
                             g_hat0, g_hat1, m_hat,
                             dml_irm_callable_score.smpls[0])
    assert np.allclose(dml_irm.psi_a.squeeze(),
                       psi_a,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_irm.psi_b.squeeze(),
                       psi_b,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_iivm_callable_vs_str_score():
    assert np.allclose(dml_iivm.psi,
                       dml_iivm_callable_score.psi,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_iivm.coef,
                       dml_iivm_callable_score.coef,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_iivm_callable_vs_pred_export():
    preds = dml_iivm_callable_score.predictions
    g_hat0 = preds['ml_g0'].squeeze()
    g_hat1 = preds['ml_g1'].squeeze()
    m_hat = preds['ml_m'].squeeze()
    r_hat0 = preds['ml_r0'].squeeze()
    r_hat1 = preds['ml_r1'].squeeze()
    psi_a, psi_b = iivm_score(dml_data_iivm.y, dml_data_iivm.z.squeeze(), dml_data_iivm.d,
                              g_hat0, g_hat1, m_hat, r_hat0, r_hat1,
                              dml_iivm_callable_score.smpls[0])
    assert np.allclose(dml_iivm.psi_a.squeeze(),
                       psi_a,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_iivm.psi_b.squeeze(),
                       psi_b,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_pliv_callable_vs_str_score():
    pliv_score = dml_pliv._score_elements
    dml_pliv_callable_score = DoubleMLPLIV(dml_data_pliv, Lasso(), Lasso(), Lasso(),
                                           score=pliv_score, draw_sample_splitting=False)
    dml_pliv_callable_score.set_sample_splitting(dml_pliv.smpls)
    dml_pliv_callable_score.fit()
    assert np.allclose(dml_pliv.psi,
                       dml_pliv_callable_score.psi,
                       rtol=1e-9, atol=1e-4)
    assert np.allclose(dml_pliv.coef,
                       dml_pliv_callable_score.coef,
                       rtol=1e-9, atol=1e-4)

@pytest.mark.ci
def test_pliv_callable_not_implemented():
    np.random.seed(3141)
    dml_data_pliv_2z = make_pliv_CHS2015(n_obs=100, dim_z=2)
    pliv_score = dml_pliv._score_elements

    dml_pliv_callable_score = DoubleMLPLIV._partialX(dml_data_pliv_2z, Lasso(), Lasso(), Lasso(),
                                                     score=pliv_score)
    msg = 'Callable score not implemented for DoubleMLPLIV.partialX with several instruments.'
    with pytest.raises(NotImplementedError, match=msg):
        dml_pliv_callable_score.fit()

    dml_pliv_callable_score = DoubleMLPLIV._partialZ(dml_data_pliv_2z, Lasso(),
                                                     score=pliv_score)
    msg = 'Callable score not implemented for DoubleMLPLIV.partialZ.'
    with pytest.raises(NotImplementedError, match=msg):
        dml_pliv_callable_score.fit()

    dml_pliv_callable_score = DoubleMLPLIV._partialXZ(dml_data_pliv_2z, Lasso(), Lasso(), Lasso(),
                                                      score=pliv_score)
    msg = 'Callable score not implemented for DoubleMLPLIV.partialXZ.'
    with pytest.raises(NotImplementedError, match=msg):
        dml_pliv_callable_score.fit()

max_stars_repo_name: SvenKlaassen/doubleml-for-py

max_issues_repo_name: SvenKlaassen/doubleml-for-py

max_forks_repo_name: SvenKlaassen/doubleml-for-py


text:
"""Functions to extract last year of deforestation or degradation from JRC data."""

import argparse
import re
from typing import Optional, Union

import numba
import numpy as np
import rioxarray as rxr
import xarray as xr
from tqdm.autonotebook import tqdm

from src.constants import JRC_PATH
from src.utils.logging import get_logger

logger = get_logger(__file__)

ANNUAL_CHANGE_PATH = JRC_PATH / "AnnualChange/tifs"
JRC_ANNUAL_CHANGE_UNDISTURBED = 1
JRC_ANNUAL_CHANGE_DEGRADATION = 2
JRC_ANNUAL_CHANGE_DEFORESTATION = 3
JRC_ANNUAL_CHANGE_RECOVERY = 4

@numba.njit
def last_of_value(
    array: np.ndarray,
    value: Union[int, float],
    fill_value: Union[int, float] = -9999,
    dtype: type = np.int16,
) -> np.ndarray:
    """
    Find index of last occurence of `value` along the first axis of the array.

    Args:
        array (np.ndarray): The 3D numpy array to investigate. The shape should be
            (obervations, rows, columns)
        value (Union[int, float]): The value for whichs last occurence should be noted
        fill_value (Union[int, float], optional): Fill value if `value` is not found
            along the first axis. Defaults to -9999.
        dtype (type, optional): The return datatype. Should be an integer type, since
            we deal with indices. Defaults to np.int16.

    Returns:
        np.array: A 2D numpy array of shape (rows, columns) and the value of the element
            [row, column] will be the last occurence of `value` in `array` (or
            `fill_value` if the value doesn't occur.)
    """
    assert len(array.shape) == 3, "This function assumes a 3d array"
    _, nrows, ncols = array.shape
    last_observation = np.ones_like(array[0], dtype=dtype) * fill_value
    for row in range(nrows):
        for col in range(ncols):
            value_observations = np.argwhere(array[:, row, col] == value)
            if len(value_observations) > 0:
                last_observation[row, col] = value_observations.max()

    return last_observation

@numba.njit
def first_of_value(
    array: np.ndarray,
    value: Union[int, float],
    fill_value: Union[int, float] = -9999,
    dtype: type = np.int16,
) -> np.ndarray:
    """
    Find index of first occurence of `value` along the first axis of the array.

    Args:
        array (np.ndarray): The 3D numpy array to investigate. The shape should be
            (obervations, rows, columns)
        value (Union[int, float]): The value for whichs first occurence should be noted
        fill_value (Union[int, float], optional): Fill value if `value` is not found
            along the first axis. Defaults to -9999.
        dtype (type, optional): The return datatype. Should be an integer type, since
            we deal with indices. Defaults to np.int16.

    Returns:
        np.array: A 2D numpy array of shape (rows, columns) and the value of the element
            [row, column] will be the first occurence of `value` in `array` (or
            `fill_value` if the value doesn't occur.)
    """
    assert len(array.shape) == 3, "This function assumes a 3d array"
    _, nrows, ncols = array.shape
    first_observation = np.ones_like(array[0], dtype=dtype) * fill_value
    for row in range(nrows):
        for col in range(ncols):
            value_observations = np.argwhere(array[:, row, col] == value)
            if len(value_observations) > 0:
                first_observation[row, col] = value_observations.min()

    return first_observation

def compute_last_observation(
    annual_change: xr.DataArray,
    jrc_class_value: int,
    first_observation: xr.DataArray = None,
    year_offset: int = 1990,
):

    if first_observation is not None:
        last_observation = xr.zeros_like(first_observation)
    else:
        last_observation = xr.zeros_like(annual_change[0])

    # Compute last year of observation (Note: Due to the size of the datasets this takes
    #  ~10 min, even with numba speedup)
    last_observation.data = last_of_value(annual_change.data, value=jrc_class_value)

    # Fix all negative/fill values and all values prior to 1990 with the first year of
    #  deforestation.
    if first_observation is not None:
        last_year = annual_change.year.data[-1]
        last_observation = xr.concat(
            [
                first_observation.where(first_observation <= last_year, other=0),
                last_observation + year_offset,
            ],
            dim="observations",
        ).max(dim="observations")
    else:
        last_observation = (year_offset + last_observation).where(
            last_observation > 0, other=0
        )

    return last_observation.astype(float)

def compute_recovery_period(
    annual_change: xr.DataArray,
    first_deforestation: Optional[xr.DataArray] = None,
    first_degradation: Optional[xr.DataArray] = None,
    as_startyear: bool = False,
):

    last_deforested = compute_last_observation(
        annual_change, JRC_ANNUAL_CHANGE_DEFORESTATION, first_deforestation
    )
    last_degraded = compute_last_observation(
        annual_change, JRC_ANNUAL_CHANGE_DEGRADATION, first_degradation
    )

    survey_year = annual_change.year.data[-1]
    deforested_before_survey = last_deforested < survey_year
    not_degraded_since_last_deforested = last_deforested > last_degraded
    recovering = annual_change.loc[survey_year] == JRC_ANNUAL_CHANGE_RECOVERY

    undisturbed_recovery = (
        recovering & deforested_before_survey & not_degraded_since_last_deforested
    )

    if as_startyear:
        return last_deforested.where(undisturbed_recovery)
    return survey_year - last_deforested.where(undisturbed_recovery)

# pylint: disable=redefined-outer-name
def calculate_jrc_last_observation_year(
    tile_identifier: str,
    jrc_class_value: int = JRC_ANNUAL_CHANGE_DEFORESTATION,  # 3 is deforestation class
    dataset: str = "DeforestationYear",
    overwrite: bool = False,
    final_year: int = 2019,
) -> bool:
    """
    Calculates the year of last observation of `jrc_class_value`in the JRC AnnualChange
    dataset.

    JRC dataset user guide:
    https://forobs.jrc.ec.europa.eu/TMF/download/TMF_DataUsersGuide.pdf

    Args:
        tile_identifier (str): The identifier of the JRC tile for which to compute the
            last observation of `jrc_class_value` for each pixel in the JRC AnnualChange
            dataset.
        jrc_class_value (int, optional): JRC class value. Class value `2` corresponds
            to degradation (in the AnnualChange datasett) and class value `3`
            corresponds to deforestation. Defaults to 3.
            More generally:
                1 - undistrubed
                2 - degraded
                3 - deforested
                4 - regrowth
                5 - permanent and seasonal water
                6 - other
        dataset (str, optional): Which dataset to use as a stand-in template for the
            results. This will usually be `DeforestationYear` or `DegradationYear`.
            Defaults to "DeforestationYear".
        overwrite (bool, optional): Whether to overwrite existing files. Default to
            False.

    Returns:
        bool: True, iff the calculation was successful and the output was saved.
    """

    first_observation_path = (
        JRC_PATH
        / dataset
        / f"JRC_TMF_{dataset}_v1_1982_{final_year}_{tile_identifier}.tif"
    )
    annual_change_paths = {
        year: (
            JRC_PATH
            / "AnnualChange"
            / "tifs"
            / f"JRC_TMF_AnnualChange_v1_{year}_{tile_identifier}.tif"
        )
        for year in range(1990, final_year + 1)
    }

    if not first_observation_path.exists():
        raise FileNotFoundError(
            f"No tile with identifier {tile_identifier} found at "
            f"{first_observation_path}."
        )

    for path in annual_change_paths.values():
        if not path.exists():
            raise FileNotFoundError(
                f"No tile with identifier {tile_identifier} found at {path}."
            )

    save_path = (
        JRC_PATH
        / f"Last{dataset}"
        / f"JRC_TMF_Last{dataset}_v1_1982_2019_{tile_identifier}.tif"
    )
    if save_path.exists() and not overwrite:
        raise FileExistsError(
            f"File {save_path} exists. To overwrite set overwrite to true."
        )

    # Creation of "last_year_of_observation" dataset:
    logger.info("Load Annual change maps for tile %s", tile_identifier)
    annual_change_maps = xr.concat(
        objs=[
            rxr.open_rasterio(path).squeeze().assign_coords({"year": year})
            for year, path in tqdm(annual_change_paths.items())
        ],
        dim="year",
    )

    # Get first year of observation map
    logger.info("Load year of %s dataset", dataset)
    first_observation = rxr.open_rasterio(first_observation_path).squeeze()

    last_observation = compute_last_observation(
        annual_change=annual_change_maps,
        jrc_class_value=jrc_class_value,
        first_observation=first_observation,
    )

    # Save result
    logger.info("Save result to compressed GeoTiff")
    save_path.parent.mkdir(parents=True, exist_ok=True)
    last_observation.rio.to_raster(
        save_path,
        compress="lzw",
    )

    return True

if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="GEDI order download script")
    parser.add_argument(
        "-v",
        "--jrc_class_value",
        help=(
            "JRC class value. Class value `2` corresponds to degradation (in the "
            "annual change data) and class value `3` corresponds to deforesttation. "
            "More generally: 1 - undistrubed, 2- degraded, 3 - deforested, 4 - "
            "regrowth, 5 - permanent and seasonal water, 6 - other. More here: "
            "https://forobs.jrc.ec.europa.eu/TMF/download/TMF_DataUsersGuide.pdf"
        ),
        type=int,
        default=3,
        nargs="?",  # Argument is optional
    )
    parser.add_argument(
        "-d",
        "--dataset",
        help=(
            "The dataset to look for a last observation for. Can be "
            "`DeforestationYear` or `DegradationYear`."
        ),
        type=str,
        default="DeforestationYear",
        nargs="?",  # Argument is optional
    )
    parser.add_argument(
        "-o",
        "--overwrite",
        help=(
            "If True, existing files are downloaded again and overwritten. "
            "Defaults to False."
        ),
        type=bool,
        default=False,
        nargs="?",  # Argument is optional
    )

    args = parser.parse_args()

    relevant_tiles = list((JRC_PATH / f"{args.dataset}").glob("*_SAM*"))

    for tile in tqdm(relevant_tiles):
        tile_identifier = re.findall(r"\w+_(SAM_[\w_]+)\.tif", tile.name)[0]
        print("Processing", tile_identifier)
        try:
            calculate_jrc_last_observation_year(
                tile_identifier,
                jrc_class_value=args.jrc_class_value,
                dataset=args.dataset,
                overwrite=args.overwrite,
            )
        except FileExistsError:
            print(f"{tile_identifier} exists.")
        except FileNotFoundError:
            print("Did not find files.")
        except Exception as e:  # pylint: disable=broad-except
            print(e)

max_stars_repo_name: Croydon-Brixton/gedi-biomass-mapping

max_issues_repo_name: Croydon-Brixton/gedi-biomass-mapping

max_forks_repo_name: Croydon-Brixton/gedi-biomass-mapping

reason: import numpy,import numba


text:
#!/usr/bin/python
''''''
"""
# ==============================
# flowlib.py
# library for optical flow processing
# Author: Ruoteng Li
# Date: 6th Aug 2016
# ==============================
"""
from . import png
import numpy as np
from PIL import Image
import cv2
import re

UNKNOWN_FLOW_THRESH = 1e7
SMALLFLOW = 0.0
LARGEFLOW = 1e8
"""
=============
Flow Section
=============
"""

def TODO():
    """
    TODO:
        update docstring in this page
    """

def read_flow(filename):
    """read optical flow data from flow file

    Args:
        filename (str): name of the flow file

    Returns:
        flow (array): optical flow data in numpy array (dtype: np.float32)
    """
    if filename.endswith('.flo'):
        flow = read_flo_file(filename)
    elif filename.endswith('.png'):
        flow = read_kitti_png_file(filename)
    elif filename.endswith('.pfm'):
        flow = read_pfm_file(filename)[:, :, :2].astype(np.float32)
    else:
        raise Exception('Invalid flow file format!')

    return flow

def write_flow(flow, filename):
    """write optical flow in Middlebury .flo format

    Args:
        flow (array, [HxWx2]): optical flow map
        filename (str): optical flow file path to be saved
    """
    f = open(filename, 'wb')
    magic = np.array([202021.25], dtype=np.float32)
    (height, width) = flow.shape[0:2]
    w = np.array([width], dtype=np.int32)
    h = np.array([height], dtype=np.int32)
    magic.tofile(f)
    w.tofile(f)
    h.tofile(f)
    flow.tofile(f)
    f.close()

def save_flow_image(flow, image_file):
    """save flow visualization into image file

    Args:
        flow (array, [HxWx2]): optical flow data
        image_file (str): image file path to be saved
    """
    # print flow.shape
    flow_img = flow_to_image(flow)
    img_out = Image.fromarray(flow_img)
    img_out.save(image_file)

def flowfile_to_imagefile(flow_file, image_file):
    """convert flowfile into image file

    Args:
        flow (str): optical flow file path
        image_file (str): image file path
    """
    flow = read_flow(flow_file)
    save_flow_image(flow, image_file)

def flow_error(tu, tv, u, v):
    """Calculate average end point error
    :param tu: ground-truth horizontal flow map
    :param tv: ground-truth vertical flow map
    :param u:  estimated horizontal flow map
    :param v:  estimated vertical flow map
    :return: End point error of the estimated flow
    """
    smallflow = 0.0
    '''
    stu = tu[bord+1:end-bord,bord+1:end-bord]
    stv = tv[bord+1:end-bord,bord+1:end-bord]
    su = u[bord+1:end-bord,bord+1:end-bord]
    sv = v[bord+1:end-bord,bord+1:end-bord]
    '''
    stu = tu[:]
    stv = tv[:]
    su = u[:]
    sv = v[:]

    idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (
        abs(stv) > UNKNOWN_FLOW_THRESH)
    stu[idxUnknow] = 0
    stv[idxUnknow] = 0
    su[idxUnknow] = 0
    sv[idxUnknow] = 0

    ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)]
    index_su = su[ind2]
    index_sv = sv[ind2]
    an = 1.0 / np.sqrt(index_su**2 + index_sv**2 + 1)

    index_stu = stu[ind2]
    index_stv = stv[ind2]
    tn = 1.0 / np.sqrt(index_stu**2 + index_stv**2 + 1)
    '''
    angle = un * tun + vn * tvn + (an * tn)
    index = [angle == 1.0]
    angle[index] = 0.999
    ang = np.arccos(angle)
    mang = np.mean(ang)
    mang = mang * 180 / np.pi
    '''

    epe = np.sqrt((stu - su)**2 + (stv - sv)**2)
    epe = epe[ind2]
    mepe = np.mean(epe)
    return mepe

def flow_kitti_error(tu, tv, u, v, mask):
    """
    Calculate average end point error
    :param tu: ground-truth horizontal flow map
    :param tv: ground-truth vertical flow map
    :param u:  estimated horizontal flow map
    :param v:  estimated vertical flow map
    :param mask: ground-truth mask
    :return: End point error of the estimated flow
    """
    tau = [3, 0.05]
    '''
    stu = tu[bord+1:end-bord,bord+1:end-bord]
    stv = tv[bord+1:end-bord,bord+1:end-bord]
    su = u[bord+1:end-bord,bord+1:end-bord]
    sv = v[bord+1:end-bord,bord+1:end-bord]
    '''
    stu = tu[:]
    stv = tv[:]
    su = u[:]
    sv = v[:]
    smask = mask[:]

    ind_valid = (smask != 0)
    n_total = np.sum(ind_valid)

    epe = np.sqrt((stu - su)**2 + (stv - sv)**2)
    mag = np.sqrt(stu**2 + stv**2) + 1e-5

    epe = epe[ind_valid]
    mag = mag[ind_valid]

    err = np.logical_and((epe > tau[0]), (epe / mag) > tau[1])
    n_err = np.sum(err)

    mean_epe = np.mean(epe)
    mean_acc = 1 - (float(n_err) / float(n_total))
    return (mean_epe, mean_acc)

def flow_to_image(flow, maxrad=-1):
    """Convert flow into middlebury color code image

    Args:
        flow (array, [HxWx2]): optical flow map

    Returns:
        img (array, [HxWx3]): optical flow image in middlebury color
    """
    u = flow[:, :, 0]
    v = flow[:, :, 1]

    maxu = -999.
    maxv = -999.
    minu = 999.
    minv = 999.

    idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
    u[idxUnknow] = 0
    v[idxUnknow] = 0

    if maxrad == -1:
        rad = np.sqrt(u**2 + v**2)
        maxrad = max(-1, np.max(rad))

    u = u / (maxrad + np.finfo(float).eps)
    v = v / (maxrad + np.finfo(float).eps)

    img = compute_color(u, v)

    idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
    img[idx] = 0

    return np.uint8(img)

def evaluate_flow_file(gt_file, pred_file):
    """
    evaluate the estimated optical flow end point error according to ground truth provided
    :param gt_file: ground truth file path
    :param pred_file: estimated optical flow file path
    :return: end point error, float32
    """
    # Read flow files and calculate the errors
    gt_flow = read_flow(gt_file)  # ground truth flow
    eva_flow = read_flow(pred_file)  # predicted flow
    # Calculate errors
    average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1],
                            eva_flow[:, :, 0], eva_flow[:, :, 1])
    return average_pe

def evaluate_flow(gt_flow, pred_flow):
    """
    gt: ground-truth flow
    pred: estimated flow
    """
    average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1],
                            pred_flow[:, :, 0], pred_flow[:, :, 1])
    return average_pe

def evaluate_kitti_flow(gt_flow, pred_flow, rigid_flow=None):
    if gt_flow.shape[2] == 2:
        gt_mask = np.ones((gt_flow.shape[0], gt_flow.shape[1]))
        epe, acc = flow_kitti_error(gt_flow[:, :, 0], gt_flow[:, :, 1],
                                    pred_flow[:, :, 0], pred_flow[:, :, 1],
                                    gt_mask)
    elif gt_flow.shape[2] == 3:
        epe, acc = flow_kitti_error(gt_flow[:, :, 0], gt_flow[:, :, 1],
                                    pred_flow[:, :, 0], pred_flow[:, :, 1],
                                    gt_flow[:, :, 2])
    return (epe, acc)

"""
==============
Disparity Section
==============
"""

def read_disp(file_name):
    # output: [H, W, 1 or 2]
    if file_name.endswith('.pfm'):
        disp = np.expand_dims(-read_pfm_file(file_name), axis=-1)
    elif file_name.endswith('.png'):
        disp = cv2.imread(file_name, -1)
        mask = np.float32(disp > 0)
        disp = np.float32(disp) / 256.
        disp = np.stack((disp, mask), axis=-1)
    else:
        raise Exception('Invalid disp file format!')

    return disp

def disp2flow(disp):
    padder = np.zeros((disp.shape[0], disp.shape[1]), dtype=np.float32)
    flow = np.stack((-disp[:, :, 0], padder), axis=-1)
    if disp.shape[2] > 1:
        flow = np.append(flow, disp[:, :, 1:], axis=-1)
    return flow

"""
==============
Others
==============
"""

def compute_color(u, v):
    """
    compute optical flow color map
    :param u: optical flow horizontal map
    :param v: optical flow vertical map
    :return: optical flow in color code
    """
    [h, w] = u.shape
    img = np.zeros([h, w, 3])
    nanIdx = np.isnan(u) | np.isnan(v)
    u[nanIdx] = 0
    v[nanIdx] = 0

    colorwheel = make_color_wheel()
    ncols = np.size(colorwheel, 0)

    rad = np.sqrt(u**2 + v**2)

    a = np.arctan2(-v, -u) / np.pi

    fk = (a + 1) / 2 * (ncols - 1) + 1

    k0 = np.floor(fk).astype(int)

    k1 = k0 + 1
    k1[k1 == ncols + 1] = 1
    f = fk - k0

    for i in range(0, np.size(colorwheel, 1)):
        tmp = colorwheel[:, i]
        col0 = tmp[k0 - 1] / 255
        col1 = tmp[k1 - 1] / 255
        col = (1 - f) * col0 + f * col1

        idx = rad <= 1
        col[idx] = 1 - rad[idx] * (1 - col[idx])
        notidx = np.logical_not(idx)

        col[notidx] *= 0.75
        img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx)))

    return img

def make_color_wheel():
    """
    Generate color wheel according Middlebury color code
    :return: Color wheel
    """
    RY = 15
    YG = 6
    GC = 4
    CB = 11
    BM = 13
    MR = 6

    ncols = RY + YG + GC + CB + BM + MR

    colorwheel = np.zeros([ncols, 3])

    col = 0

    # RY
    colorwheel[0:RY, 0] = 255
    colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))
    col += RY

    # YG
    colorwheel[col:col + YG, 0] = 255 - np.transpose(
        np.floor(255 * np.arange(0, YG) / YG))
    colorwheel[col:col + YG, 1] = 255
    col += YG

    # GC
    colorwheel[col:col + GC, 1] = 255
    colorwheel[col:col + GC, 2] = np.transpose(
        np.floor(255 * np.arange(0, GC) / GC))
    col += GC

    # CB
    colorwheel[col:col + CB, 1] = 255 - np.transpose(
        np.floor(255 * np.arange(0, CB) / CB))
    colorwheel[col:col + CB, 2] = 255
    col += CB

    # BM
    colorwheel[col:col + BM, 2] = 255
    colorwheel[col:col + BM, 0] = np.transpose(
        np.floor(255 * np.arange(0, BM) / BM))
    col += +BM

    # MR
    colorwheel[col:col + MR, 2] = 255 - np.transpose(
        np.floor(255 * np.arange(0, MR) / MR))
    colorwheel[col:col + MR, 0] = 255

    return colorwheel

def read_flo_file(filename):
    """
    Read from Middlebury .flo file
    :param flow_file: name of the flow file
    :return: optical flow data in matrix
    """
    f = open(filename, 'rb')
    magic = np.fromfile(f, np.float32, count=1)
    data2d = None

    if 202021.25 != magic:
        print('Magic number incorrect. Invalid .flo file')
    else:
        w = np.fromfile(f, np.int32, count=1)
        h = np.fromfile(f, np.int32, count=1)
        # print "Reading %d x %d flow file in .flo format" % (h, w)
        data2d = np.fromfile(f, np.float32, count=2 * int(w) * int(h))
        # reshape data into 3D array (columns, rows, channels)
        data2d = np.resize(data2d, (int(h), int(w), 2))
    f.close()
    return data2d

def read_png_file(flow_file):
    """
    Read from KITTI .png file
    :param flow_file: name of the flow file
    :return: optical flow data in matrix
    """
    flow_object = png.Reader(filename=flow_file)
    flow_direct = flow_object.asDirect()
    flow_data = list(flow_direct[2])
    (w, h) = flow_direct[3]['size']
    print("Reading %d x %d flow file in .png format" % (h, w))
    flow = np.zeros((h, w, 3), dtype=np.float64)
    for i in range(len(flow_data)):
        flow[i, :, 0] = flow_data[i][0::3]
        flow[i, :, 1] = flow_data[i][1::3]
        flow[i, :, 2] = flow_data[i][2::3]

    invalid_idx = (flow[:, :, 2] == 0)
    flow[:, :, 0:2] = (flow[:, :, 0:2] - 2**15) / 64.0
    flow[invalid_idx, 0] = 0
    flow[invalid_idx, 1] = 0
    return flow

def read_kitti_png_file(flow_file):
    flow_img = cv2.imread(flow_file, -1)
    flow_img = flow_img.astype(np.float32)
    flow_data = np.zeros(flow_img.shape, dtype=np.float32)
    flow_data[:, :, 0] = (flow_img[:, :, 2] - 2**15) / 64.0
    flow_data[:, :, 1] = (flow_img[:, :, 1] - 2**15) / 64.0
    flow_data[:, :, 2] = flow_img[:, :, 0]
    return flow_data

def read_pfm_file(file):
    file = open(file, 'rb')

    color = None
    width = None
    height = None
    scale = None
    endian = None

    header = file.readline().rstrip()
    if header.decode("ascii") == 'PF':
        color = True
    elif header.decode("ascii") == 'Pf':
        color = False
    else:
        raise Exception('Not a PFM file.')

    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode("ascii"))
    if dim_match:
        width, height = list(map(int, dim_match.groups()))
    else:
        raise Exception('Malformed PFM header.')

    scale = float(file.readline().decode("ascii").rstrip())
    if scale < 0:  # little-endian
        endian = '<'
        scale = -scale
    else:
        endian = '>'  # big-endian

    data = np.fromfile(file, endian + 'f')
    shape = (height, width, 3) if color else (height, width)

    data = np.reshape(data, shape)
    data = np.flipud(data)
    return data  #, scale

def resize_flow(flow, des_width, des_height, method='bilinear'):
    # improper for sparse flow
    src_height = flow.shape[0]
    src_width = flow.shape[1]
    if src_width == des_width and src_height == des_height:
        return flow
    ratio_height = float(des_height) / float(src_height)
    ratio_width = float(des_width) / float(src_width)
    if method == 'bilinear':
        flow = cv2.resize(
            flow, (des_width, des_height), interpolation=cv2.INTER_LINEAR)
    elif method == 'nearest':
        flow = cv2.resize(
            flow, (des_width, des_height), interpolation=cv2.INTER_NEAREST)
    else:
        raise Exception('Invalid resize flow method!')
    flow[:, :, 0] = flow[:, :, 0] * ratio_width
    flow[:, :, 1] = flow[:, :, 1] * ratio_height
    return flow

def horizontal_flip_flow(flow):
    flow = np.copy(np.fliplr(flow))
    flow[:, :, 0] *= -1
    return flow

def vertical_flip_flow(flow):
    flow = np.copy(np.flipud(flow))
    flow[:, :, 1] *= -1
    return flow

def remove_ambiguity_flow(flow_img, err_img, threshold_err=10.0):
    thre_flow = flow_img
    mask_img = np.ones(err_img.shape, dtype=np.uint8)
    mask_img[err_img > threshold_err] = 0.0
    thre_flow[err_img > threshold_err] = 0.0
    return (thre_flow, mask_img)

def write_kitti_png_file(flow_fn, flow_data, mask_data):
    flow_img = np.zeros((flow_data.shape[0], flow_data.shape[1], 3),
                        dtype=np.uint16)
    flow_img[:, :, 2] = flow_data[:, :, 0] * 64.0 + 2**15
    flow_img[:, :, 1] = flow_data[:, :, 1] * 64.0 + 2**15
    flow_img[:, :, 0] = mask_data[:, :]
    cv2.imwrite(flow_fn, flow_img)

def flow_kitti_mask_error(tu, tv, gt_mask, u, v, pd_mask):
    """
    Calculate average end point error
    :param tu: ground-truth horizontal flow map
    :param tv: ground-truth vertical flow map
    :param gt_mask: ground-truth mask

    :param u:  estimated horizontal flow map
    :param v:  estimated vertical flow map
    :param pd_mask: estimated flow mask
    :return: End point error of the estimated flow
    """
    tau = [3, 0.05]
    '''
    stu = tu[bord+1:end-bord,bord+1:end-bord]
    stv = tv[bord+1:end-bord,bord+1:end-bord]
    su = u[bord+1:end-bord,bord+1:end-bord]
    sv = v[bord+1:end-bord,bord+1:end-bord]
    '''
    stu = tu[:]
    stv = tv[:]
    su = u[:]
    sv = v[:]
    s_gt_mask = gt_mask[:]
    s_pd_mask = pd_mask[:]

    ind_valid = np.logical_and(s_gt_mask != 0, s_pd_mask != 0)
    n_total = np.sum(ind_valid)

    epe = np.sqrt((stu - su)**2 + (stv - sv)**2)
    mag = np.sqrt(stu**2 + stv**2) + 1e-5

    epe = epe[ind_valid]
    mag = mag[ind_valid]

    err = np.logical_and((epe > tau[0]), (epe / mag) > tau[1])
    n_err = np.sum(err)

    mean_epe = np.mean(epe)
    mean_acc = 1 - (float(n_err) / float(n_total))
    return (mean_epe, mean_acc)

max_stars_repo_name: SimOgaard/DF-VO

max_issues_repo_name: SimOgaard/DF-VO

max_forks_repo_name: SimOgaard/DF-VO