text:
from typing import Optional

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.path import Path as MplPath
import matplotlib.patches as patches
import networkx as nx
import numpy as np

def normalize_vector(vector: np.array, normalize_to: float) -> np.array:
    """Make `vector` norm equal to `normalize_to`

    vector: np.array
        Vector with 2 coordinates
    normalize_to: float
        A norm of the new vector

    Returns
    -------
    Vector with the same direction, but with length normalized to `normalize_to`
    """
    vector_norm = np.linalg.norm(vector)

    return vector * normalize_to / vector_norm

def orthogonal_vector(
    point: np.array, width: float, normalize_to: Optional[float] = None
) -> np.array:
    """Get orthogonal vector to a `point`

    point: np.array
        Vector with x and y coordinates of a point
    width: float
        Distance of the x-coordinate of the new vector from the `point` (in orthogonal direction)
    normalize_to: Optional[float] = None
        If a number is provided, normalize a new vector length to this number

    Returns
    -------
    Array with x and y coordinates of the vector, which is orthogonal to the vector from (0, 0) to `point`
    """
    EPSILON = 0.000001

    x = width
    y = -x * point[0] / (point[1] + EPSILON)

    ort_vector = np.array([x, y])

    if normalize_to is not None:
        ort_vector = normalize_vector(ort_vector, normalize_to)

    return ort_vector

def draw_self_loop(
    point: np.array,
    ax: Optional[plt.Axes] = None,
    padding: float = 1.5,
    width: float = 0.3,
    plot_size: int = 10,
    linewidth: float = 0.2,
    color: str = "pink",
) -> plt.Axes:
    """Draw a loop from `point` to itself

    !Important! By "center" we assume a (0, 0) point. If your data is centered around a different points,
    it is strongly recommended to center it around zero. Otherwise, you will probably get ugly plots

    Parameters
    ----------
    point: np.array
        1D array with 2 coordinates of the point. Loop will be drawn from and to these coordinates.
    padding: float = 1.5
        Controls how the distance of the loop from the center. If `padding` > 1, the loop will be
        from the outside of the `point`. If `padding` < 1, the loop will be closer to the center
    width: float = 0.3
        Controls the width of the loop
    linewidth: float = 0.2
        Width of the line of the loop
    ax: Optional[matplotlib.pyplot.Axes]:
        Axis on which to draw a plot. If None, a new Axis is generated
    plot_size: int = 7
        Size of the plot sides in inches. Ignored if `ax` is provided
    color: str = "pink"
        Color of the arrow

    Returns
    -------
    Matplotlib axes with the self-loop drawn
    """

    if ax is None:
        fig, ax = plt.subplots(figsize=(plot_size, plot_size))

    point_with_padding = padding * point
    ort_vector = orthogonal_vector(point, width, normalize_to=width)

    first_anchor = ort_vector + point_with_padding
    second_anchor = -ort_vector + point_with_padding

    verts = [point, first_anchor, second_anchor, point]
    codes = [MplPath.MOVETO, MplPath.CURVE4, MplPath.CURVE4, MplPath.CURVE4]

    path = MplPath(verts, codes)

    patch = patches.FancyArrowPatch(
        path=path,
        lw=linewidth,
        arrowstyle="-|>",
        color=color,
        alpha=0.5,
        mutation_scale=30,  # arrowsize in nx.draw_networkx_edges()
    )
    ax.add_patch(patch)

    return ax

def draw_graph_edge(
    graph: nx.Graph,
    pos: dict,
    edge: tuple,
    edge_weight: float,
    ax: plt.Axes,
    color: str,
    arc_radius: float = 0.2,
) -> plt.Axes:
    """Draw the given edge of the network

    Parameters
    ----------
    graph: nx.Graph
        A graph, edge of which you want to draw
    pos: dict
        A dictionary with the positions of the nodes. Keys are nodes of the `graph`,
        and values are their positions
    edge: tuple
        Edge in the `graph`. Tuple must contain 2 elements: the start and the end of the edge
    edge_weight: float
        Number describing how much weight does the edge have. It is used for the visualisation
    ax: plt.Axes
        Axis on which to draw a plot
    color: str
        Color of the node
    arc_radius: float = 0.2
        Edges are drawn as arcs, not as straight lines. This parameter controls the curvature of
        the `edge` arc

    Returns
    -------
    Matplotlib axes with the edge drawn
    """

    nx.draw_networkx_edges(
        graph,
        pos=pos,
        width=edge_weight,
        edgelist=[edge],
        alpha=0.5,
        edge_color=color,
        ax=ax,
        arrowsize=30,
        connectionstyle=f"arc3,rad={arc_radius}",
        node_size=1000
    )

    return ax

def graph_edges_weights(graph):
    return {edge: graph.edges[edge]["weight"] for edge in graph.edges}

def draw_graph_edges(graph: nx.Graph, pos: dict, ax: plt.Axes):
    """Draw graph edges so that edges in the opposite directions look differently

    Parameters
    ----------
    graph: nx.Graph
        Graph, edges of which you want to draw
    pos: dict
        Dictionary, where keys are nodes and values are their positions. Can be obtained
        through networkx layout algorithms (e. g. nx.circular_layout())
    ax: plt.Axes
        Axis on which the edges are drawn

    Returns
    -------
    Axis with the edges drawn
    """

    edge_weights = graph_edges_weights(graph)
    edges_to_draw = set(graph.edges)

    for edge in graph.edges:
        if edge not in edges_to_draw:
            continue

        if edge[0] == edge[1]:  # By default, networkx doesn't draw self loops correctly
            draw_self_loop(point=pos[edge[0]], ax=ax, linewidth=edge_weights[edge])
            continue

        draw_graph_edge(
            graph, pos, edge, edge_weight=edge_weights[edge], ax=ax, color="pink"
        )
        edges_to_draw.remove(edge)

        # Edges between the same vertices look confusing, if they have the same style
        # So we draw such edges with different colors and curvature
        reverse_edge = (edge[1], edge[0])

        if reverse_edge in graph.edges and reverse_edge in edges_to_draw:
            draw_graph_edge(
                graph,
                pos,
                reverse_edge,
                edge_weight=edge_weights[edge],
                ax=ax,
                color="lightblue",
                arc_radius=0.1
            )
            edges_to_draw.remove(reverse_edge)

    return ax

def draw_chord_diagram(
    graph: Optional[nx.Graph] = None,
    matrix: Optional[pd.DataFrame] = None,
    ax=None,
    plot_size=10
) -> plt.Axes:
    """Draw weighted connections between the nodes of the `graph`"""

    if graph is None and matrix is None:
        raise ValueError("Please provide either `graph` or `matrix`")

    elif graph is None:
        graph = nx.DiGraph(matrix)

    if ax is None:
        fig, ax = plt.subplots(figsize=(plot_size, plot_size))

    pos = nx.circular_layout(graph, center=(0, 0))

    nx.draw_networkx_nodes(
        graph, pos=pos, node_color="springgreen", ax=ax, alpha=0.7, node_size=1000
    )

    draw_graph_edges(graph, pos, ax)

    nx.draw_networkx_labels(graph, pos=pos, font_weight="bold", ax=ax)

    ax.margins(0.1)  # Give some space for labels

    return ax

def draw_dotplot(df, figsize=(12, 6)):
    df["radius"] = np.sqrt(df["count"]) / 3
    fig, ax = plt.subplots(figsize=figsize)

    ax.scatter(df["index"], df["variable"], s=0)

    # plot data row-wise as text with circle radius according to count
    for _, row in df.iterrows():
        bbox_props = {
            "boxstyle": f'circle, pad = {row["radius"]}',
            "fc": "w",
            "ec": "r",
            "lw": 2
        }

        if row["count"]:
            ax.annotate(
                str(row["count"]),
                xy=(row["index"], row["variable"]),
                bbox=bbox_props,
                ha="center",
                va="center",
                zorder=2,
                clip_on=True
            )

    ax.grid(ls="--", zorder=1)
    fig.tight_layout()
    ax.set_title("Interactions counts", fontsize=18, fontweight="bold")

    return ax

max_forks_repo_name: VladimirShitov/communiquer

reason: import numpy,import networkx


text:
"""
Pytest for confirming model import assessment works using Telescope data
retrieval from LogAnalyticsDataClient. Usage:

Train a brain using model import's instructions, then run

pytest tests/test_model_import.py -s \
    --brain_name <BRAIN_NAME> \
    --log_analy_workspace <LAW_WORKSPACE_ID> \
    --custom_assess_name <CUSTOM_ASSESSMENT_NAME> \
    --model_file_path "./Machine-Teaching-Examples/model_import/state_transform_deep.zip"

    or
    --model_file_path "./Machine-Teaching-Examples/model_import/state_transform_deep.onnx"
"""

__author__ = "Journey McDowell"
__copyright__ = "Copyright 2021, Microsoft Corp."

import os
import pytest
import glob
import pandas as pd
import numpy as np
from azure.loganalytics import LogAnalyticsDataClient
from azure.common.credentials import get_azure_cli_credentials
from azure.loganalytics.models import QueryBody
import ast
import json
import matplotlib.pyplot as plt
import time

# Allowing optional flags to replace defaults for pytest from tests/conftest.py
@pytest.fixture()
def brain_name(pytestconfig):
    return pytestconfig.getoption("brain_name")

@pytest.fixture()
def brain_version(pytestconfig):
    return pytestconfig.getoption("brain_version")

@pytest.fixture()
def concept_name(pytestconfig):
    return pytestconfig.getoption("concept_name")

# json file with assessment configs for input
@pytest.fixture()
def file_name(pytestconfig):
    return pytestconfig.getoption("file_name")

@pytest.fixture()
def simulator_package_name(pytestconfig):
    return pytestconfig.getoption("simulator_package_name")

@pytest.fixture()
def inkling_fname(pytestconfig):
    return pytestconfig.getoption("inkling_fname")

@pytest.fixture()
def instance_count(pytestconfig):
    return pytestconfig.getoption("instance_count")

@pytest.fixture()
def custom_assess_name(pytestconfig):
    return pytestconfig.getoption("custom_assess_name")

@pytest.fixture()
def log_analy_workspace(pytestconfig):
    return pytestconfig.getoption("log_analy_workspace")

@pytest.fixture()
def import_name(pytestconfig):
    return pytestconfig.getoption("import_name")

@pytest.fixture()
def model_file_path(pytestconfig):
    return pytestconfig.getoption("model_file_path")

@pytest.fixture()
def episode_iteration_limit(pytestconfig):
    return pytestconfig.getoption("episode_iteration_limit")

# Use CLI to import a ML model as .onnx or tf
def test_model_import(import_name, model_file_path):
    os.system('bonsai importedmodel create --name "{}" --modelfilepath {}'.format(
        import_name,
        model_file_path,
    ))

    os.system('bonsai importedmodel show --name "{}" -o json > status.json'.format(
        import_name,
    ))

    # Confirm model import succeeded
    with open('status.json') as fname:
        status = json.load(fname)

    assert status['Status'] == 'Succeeded'

# Use CLI to create, upload inkling, train, and wait til complete
def test_train_brain(brain_name, brain_version, inkling_fname, simulator_package_name):
    os.system('bonsai brain create -n {}'.format(
        brain_name,
    ))
    os.system('bonsai brain version update-inkling -n {} --version {} -f {}'.format(
        brain_name,
        brain_version,
        inkling_fname
    ))
    concept_names = ['ImportedConcept', 'MoveToCenter']
    for concept in concept_names:
        time.sleep(20)
        if concept == 'ImportedConcept':
            os.system('bonsai brain version start-training -n {} --version {} -c {}'.format(
                brain_name,
                brain_version,
                concept
            ))
        else:
            os.system('bonsai brain version start-training -n {} --version {} --simulator-package-name {} -c {}'.format(
                brain_name,
                brain_version,
                simulator_package_name,
                concept
            ))

        # Do not continue until training is complete
        running = True
        while running:
            time.sleep(20)
            os.system('bonsai brain version show --name {} --version {} -o json > status.json'.format(
                brain_name,
                brain_version,
            ))
            with open('status.json') as fname:
                status = json.load(fname)
            if status['trainingState'] == 'Active':
                pass
            else:
                running = False
                time.sleep(300)
                print('Training complete...')

        # Final check concepts are complete
        with open('status.json') as fname:
            status = json.load(fname)
        assert status['status'] == 'Succeeded'
    print('All Concepts trained')

# Main test function for
# 1. running custom assessment using bonsai-cli
# 2. retrieving data from Log Analytics Workspace using LogAnalyticsDataClient
# 3. flattening states, actions, and configs
# 4. making plots for episode metrics
# 5. qualifying pass/fail
def test_assessment_brain(brain_name, brain_version, concept_name, file_name, simulator_package_name, instance_count, custom_assess_name, log_analy_workspace, episode_iteration_limit):
    # Run custom assessment
    os.system('bonsai brain version assessment start --brain-name {} --brain-version {} --concept-name {} --file {} --simulator-package-name {} --instance-count {} --name {}  --episode-iteration-limit {}'.format(
        brain_name,
        brain_version,
        concept_name,
        file_name,
        simulator_package_name,
        instance_count,
        custom_assess_name,
        episode_iteration_limit,
    ))

    # Do not continue until assessment is complete and waited 5 minutes
    running = True
    while running:
        time.sleep(10)
        os.system('bonsai brain version assessment show --brain-name {} --brain-version {} --name {} -o json > status.json'.format(
            brain_name,
            brain_version,
            custom_assess_name
        ))
        with open('status.json') as fname:
            status = json.load(fname)
        if status['status'] == 'Completed':
            running = False
            print('Assessment complete, waiting 5 min for data to appear in LAW...')
            for i in range(5):
                print('{} min...'.format(5-i))
                time.sleep(60)

    # Extract telescope from LAW using workspace ID and return flattened
    df = extract_telescope(log_analy_workspace, brain_name, brain_version, custom_assess_name)

    df = df.reset_index(drop=True)

    # Make plots
    # EpisodeIndex based on unique EpisodeIds
    df['EpisodeIndex'] = np.zeros(len(df))
    k = 1
    for i in list(set(df['EpisodeId'])):
        for j in range(len(df)):
            if df['EpisodeId'].iloc[j] == i:
                df.at[j,'EpisodeIndex'] = k
        k += 1

    # Manipulate
    df['distance_to_center'] = np.sqrt(df['ball_x'] ** 2 + df['ball_y'] ** 2)
    df['velocity_magnitude'] = np.sqrt(df['ball_vel_x'] ** 2 + df['ball_vel_y'] ** 2)

    # Create dataframe consisting of episode finish metrics
    df_last = pd.DataFrame({})
    mse_dist_list = []
    mse_vel_list = []
    for ep in range(1, int(df['EpisodeIndex'].max())+1):
        last_iter = df[(df['EpisodeIndex']==ep) & (df['IterationIndex']==len(df[df['EpisodeIndex']==ep]))]
        print(last_iter.head())

        mse_dist_list.append(np.square(np.subtract(0, df[(df['EpisodeIndex']==ep)]['distance_to_center'])).mean())
        mse_vel_list.append(np.square(np.subtract(0, df[(df['EpisodeIndex']==ep)]['velocity_magnitude'])).mean())

        df_last = pd.concat([df_last, last_iter], sort=False)

    df_last['mse_dist'] = mse_dist_list
    df_last['mse_vel'] = mse_vel_list

    # Create dataframe consisting of summary info
    df_summary = {}
    df_summary['percentage_full_episodes'] = (len(df[df['IterationIndex']==251]) / df['EpisodeIndex'].max()) * 100
    df_summary['avg_final_distance_to_center'] = np.mean(df_last['distance_to_center'])
    df_summary['avg_final_velocity_magnitude'] = np.mean(df_last['velocity_magnitude'])
    df_summary['mse_dist_total'] = np.square(np.subtract(0, df['distance_to_center'])).mean()
    df_summary['mse_vel_total'] = np.square(np.subtract(0, df['velocity_magnitude'])).mean()

    # Save Summary values as json
    with open('brain_summary.json', 'w') as outfile:
        json.dump(df_summary, outfile)

    # Plot Final Values
    fig, ax = plt.subplots(1, 1, figsize=(16, 8))
    episodes = [i for i in range(1, int(df['EpisodeIndex'].max()+1))]
    ax.plot(episodes, df_last['distance_to_center'])
    ax.plot(episodes, df_last['velocity_magnitude'])
    ax.legend(['Final Distance to Center', 'Final Velocity Mag'])
    ax.set_xlabel('Custom Assessment Episodes')
    fig.suptitle('brain \n Percentage of Full Episodes: {:0.2f}% \n Average Final Ball Distance To Center: {:0.4f} \n Average Final Ball Velocity Mag: {:0.4f} \n'.format(df_summary['percentage_full_episodes'], df_summary['avg_final_distance_to_center'], df_summary['avg_final_velocity_magnitude']), fontsize=14)

    # Plot MSE Values
    figa, axa = plt.subplots(1, 1, figsize=(16, 8))
    axa.plot(episodes, df_last['mse_dist'])
    axa.plot(episodes, df_last['mse_vel'])
    axa.legend(['MSE Distance', 'MSE Vel Mag'])
    axa.set_xlabel('Custom Assessment Episodes')
    figa.suptitle('brain \n Percentage of Full Episodes: {:0.2f}% \n Average MSE Ball Distance To Center: {:0.4f} \n Average MSE Ball Velocity Mag: {:0.4f} \n'.format(df_summary['percentage_full_episodes'], df_summary['mse_dist_total'], df_summary['mse_vel_total']), fontsize=14)

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    for key, val in df_summary.items():
        print(key, val)
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    #plt.show()

    # Assert tests for qualification
    assert df_summary['percentage_full_episodes'] >= 50
    assert df_summary['avg_final_distance_to_center'] <= 0.1
    assert df_summary['avg_final_velocity_magnitude'] <= 0.1
    assert df_summary['mse_dist_total'] <= 0.008
    assert df_summary['mse_vel_total'] <= 0.008

# Extract telescope data using query
@pytest.mark.skip(reason="helper")
def extract_telescope(log_analy_workspace_id, brain_name, brain_version, assessment_name):
    creds, _ = get_azure_cli_credentials(resource="https://api.loganalytics.io")
    log_client = LogAnalyticsDataClient(creds)
    myWorkSpaceId = log_analy_workspace_id
    result = log_client.query(myWorkSpaceId, QueryBody(**{
        'query': (
            'EpisodeLog_CL'
            '| where BrainName_s == "{}" and BrainVersion_d == "{}" and AssessmentName_s == "{}"'
            '| where  TimeGenerated > ago(30d)'
            '| join kind=inner ('
            'IterationLog_CL'
            '| sort by Timestamp_t desc'
            ') on EpisodeId_g'
            '| project AssessmentName = AssessmentName_s, EpisodeId = EpisodeId_g, IterationIndex = IterationIndex_d, Timestamp = Timestamp_t, SimState = parse_json(SimState_s), SimAction = parse_json(SimAction_s), Reward = Reward_d, CumulativeReward = CumulativeReward_d, Terminal = Terminal_b, LessonIndex = LessonIndex_d, SimConfig = parse_json(SimConfig_s), GoalMetrics = parse_json(GoalMetrics_s), EpisodeType = EpisodeType_s, FinishReason = FinishReason_s'
            '| order by EpisodeId asc, IterationIndex asc'
        ).format(brain_name, str(brain_version), assessment_name)
    }))

    df = pd.DataFrame(result.tables[0].rows, columns=[result.tables[0].columns[i].name for i in range(len(result.tables[0].columns))])

    df_flattened = format_kql_logs(df)
    df_flattened.to_csv('flattened_telescope.csv')
    return df_flattened

# Flatten data
@pytest.mark.skip(reason="helper")
def format_kql_logs(df: pd.DataFrame) -> pd.DataFrame:
    ''' Function to format a dataframe obtained from KQL query.
        Output format: keeps only selected columns, and flatten nested columns [SimAction, SimState, SimConfig]
        Parameters
        ----------
        df : DataFrame
            dataframe obtained from running KQL query then exporting `_kql_raw_result_.to_dataframe()`
    '''
    selected_columns = ["Timestamp", "IterationIndex", "Reward", "CumulativeReward", "Terminal", "SimState", "SimAction", "SimConfig", "EpisodeId"]
    nested_columns =  ["SimState", "SimAction", "SimConfig"]
    df_selected_columns = df[selected_columns]
    series_lst = []
    ordered_columns = ["EpisodeId", "IterationIndex", "Reward", "Terminal"]
    for i in nested_columns:
        try:
            new_series = df_selected_columns[i].apply(ast.literal_eval).apply(pd.Series)
        except:
            df_selected_columns[i].fillna(value=str({key: str(np.nan) for key, value in ast.literal_eval(df_selected_columns[i][1]).items()}), inplace=True)
            new_series = df_selected_columns[i].apply(ast.literal_eval).apply(pd.Series)
        column_names = new_series.columns.values.tolist()
        series_lst.append(new_series)
        if len(column_names) > 0:
            ordered_columns.extend(column_names)
        del(df_selected_columns[i])

    series_lst.append(df_selected_columns)
    formated_df = pd.concat(series_lst, axis=1)
    formated_df = formated_df.sort_values(by='Timestamp',ascending=True) # reorder df based on Timestamp
    formated_df.index = range(len(formated_df)) # re-index
    formated_df['Timestamp']=pd.to_datetime(formated_df['Timestamp']) # convert Timestamp to datetime

    formated_df = formated_df[ordered_columns]
    return formated_df.sort_values(by=["EpisodeId", "IterationIndex"])

ext: py

max_stars_repo_name: microsoft/moabsim-py

max_issues_repo_name: microsoft/moabsim-py

max_forks_repo_name: microsoft/moabsim-py

reason: import numpy


text:
"""The :mod:`genome` module defines the ``Genome`` type and provides genome translation, spawning and simplification.

A ``Genome`` is a persistent collection of gene Atoms (any `Atom` that isn't a ``CodeBlock``). It can be translated
into a ``CodeBlock``.

The ``GeneSpawner`` is a factory capable of generating random genes and random genomes. It is used for initializing a
population as well producing new genes used by variation operators (ie. mutation).

The genome simplification process is useful for removing superfluous genes from a genome without negatively impacting
the behavior of the program produced by the genome. This process has many benefits including: improving generalization,
shrinking the size of the serialized solution, and in some cases making the program easier to explain.

"""
from __future__ import annotations

from enum import Enum
from typing import Sequence, Union, Any, Callable, Tuple

import numpy as np
from pyrsistent import PRecord, field, CheckedPVector, l

from pyshgp.gp.evaluation import Evaluator
from pyshgp.push.instruction_set import InstructionSet
from pyshgp.push.program import ProgramSignature, Program
from pyshgp.push.atoms import Atom, CodeBlock, Closer, Literal, InstructionMeta, Input
from pyshgp.push.type_library import infer_literal
from pyshgp.tap import tap
from pyshgp.utils import DiscreteProbDistrib

class Opener(PRecord):
    """Marks the start of one or more CodeBlock."""

    count = field(type=int, mandatory=True)

    def dec(self) -> Opener:
        """Create an ``Opener`` with ``count`` decremented."""
        return Opener(count=self.count - 1)

def _has_opener(seq: Sequence) -> bool:
    for el in seq:
        if isinstance(el, Opener):
            return True
    return False

class Genome(CheckedPVector):
    """A linear sequence of genes (aka any atom that isn't a ``CodeBlock``).

    PyshGP uses the Plushy genome representation.

    See: http://gpbib.cs.ucl.ac.uk/gp-html/Spector_2019_GPTP.html

    """

    __type__ = Atom
    __invariant__ = lambda a: (not isinstance(a, CodeBlock), 'CodeBlock')

def genome_to_code(genome: Genome) -> CodeBlock:
    """Translate into nested CodeBlocks.

    These CodeBlocks can be considered the Push program representation of
    the Genome which can be executed by a PushInterpreter and evaluated
    by an Evaluator.

    """
    plushy_buffer = l()
    for atom in genome[::-1]:
        if isinstance(atom, InstructionMeta) and atom.code_blocks > 0:
            plushy_buffer = plushy_buffer.cons(Opener(count=atom.code_blocks))
        plushy_buffer = plushy_buffer.cons(atom)

    push_buffer = []
    while True:
        # If done with plush but unclosed opens, recur with one more close.
        if len(plushy_buffer) == 0 and _has_opener(push_buffer):
            plushy_buffer = plushy_buffer.cons(Closer())
        # If done with plush and all opens closed, return push.
        elif len(plushy_buffer) == 0:
            return CodeBlock(push_buffer)
        else:
            atom = plushy_buffer.first
            plushy_buffer = plushy_buffer.rest
            # If next instruction is a close, and there is an open.
            if isinstance(atom, Closer) and _has_opener(push_buffer):
                ndx, opener = [(ndx, el) for ndx, el in enumerate(push_buffer) if isinstance(el, Opener)][-1]
                post_open = push_buffer[ndx + 1:]
                pre_open = push_buffer[:ndx]
                push_buffer = pre_open + [CodeBlock(post_open)]
                if opener.count > 1:
                    opener = opener.dec()
                    push_buffer.append(opener)
            # If next instruction is a close, and there is no open, ignore it.
            # If next instruction is not a close.
            elif not isinstance(atom, Closer):
                push_buffer.append(atom)

class GeneTypes(Enum):
    """An ``Enum`` denoting the different types of genes that can appear in a Genome."""

    INPUT = 1
    INSTRUCTION = 2
    CLOSE = 3
    LITERAL = 4
    ERC = 5

class GeneSpawner:
    """A factory of random Genes (Atoms) and Genomes.

    When  spawning a random gene, the result can be one of three types of Atoms.
    An Instruction, a Closer, or a Literal. If the Atom is a Literal, it may
    be one of the supplied Literals, or it may be the result of running one of
    the Ephemeral Random Constant generators.

    Reference for ERCs:
    "A field guide to genetic programming", Section 3.1
    Riccardo Poli and William B. Langdon and Nicholas Freitag McPhee,
    http://www.gp-field-guide.org.uk/

    Attributes
    ----------
    n_input : int
        Number of input instructions that could appear the genomes.
    instruction_set : pyshgp.push.instruction_set.InstructionSet
        InstructionSet containing instructions to use when spawning genes and
        genomes.
    literals : Sequence[pyshgp.push.instruction_set.atoms.Literal]
        A list of Literal objects to pull from when spawning genes and genomes.
    erc_generator : Sequence[Callable]
        A list of functions (aka Ephemeral Random Constant generators). When one
        of these functions is called, the output is placed in a Literal and
        returned as the spawned gene.
    distribution : pyshgp.utils.DiscreteProbDistrib
        A probability distribution describing how frequently to produce
        Instructions, Closers, Literals, and ERCs.

    """

    def __init__(self,
                 n_inputs: int,
                 instruction_set: Union[InstructionSet, str],
                 literals: Sequence[Any],
                 erc_generators: Sequence[Callable],
                 distribution: DiscreteProbDistrib = "proportional"):
        self.n_inputs = n_inputs
        self.erc_generators = erc_generators

        self.instruction_set = instruction_set
        if self.instruction_set == "core":
            self.instruction_set = InstructionSet(register_core=True)
        self.type_library = self.instruction_set.type_library
        self.literals = [lit if isinstance(lit, Literal) else infer_literal(lit, self.type_library) for lit in literals]

        if distribution == "proportional":
            self.distribution = (
                DiscreteProbDistrib()
                .add(GeneTypes.INPUT, self.n_inputs)
                .add(GeneTypes.INSTRUCTION, len(self.instruction_set))
                .add(GeneTypes.CLOSE, sum([i.code_blocks for i in self.instruction_set.values()]))
                .add(GeneTypes.LITERAL, len(literals))
                .add(GeneTypes.ERC, len(erc_generators))
            )
        else:
            self.distribution = distribution

    def random_input(self) -> Input:
        """Return a random ``Input``.

        Returns
        -------
        pyshgp.push.atoms.Input

        """
        return Input(input_index=np.random.randint(self.n_inputs))

    def random_instruction(self) -> InstructionMeta:
        """Return a random Instruction from the InstructionSet.

        Returns
        -------
        pyshgp.push.atoms.InstructionMeta
            A randomly selected Literal.

        """
        i = np.random.choice(list(self.instruction_set.values()))
        return InstructionMeta(name=i.name, code_blocks=i.code_blocks)

    def random_literal(self) -> Literal:
        """Return a random Literal from the set of Literals.

        Returns
        -------
        pyshgp.push.atoms.Literal
            A randomly selected Literal.

        """
        lit = np.random.choice(self.literals)
        if not isinstance(lit, Literal):
            lit = infer_literal(lit, self.type_library)
        return lit

    def random_erc(self) -> Literal:
        """Materialize a random ERC generator into a Literal and return it.

        Returns
        -------
        pyshgp.push.atoms.Literal
            A Literal whose value comes from running a ERC generator function.

        """
        erc_value = np.random.choice(self.erc_generators)()
        if not isinstance(erc_value, Literal):
            erc_value = infer_literal(erc_value, self.type_library)
        return erc_value

    def random_gene(self) -> Atom:
        """Return a random Atom based on the GenomeSpawner's distribution.

        Returns
        -------
        pyshgp.push.atoms.Atom
            An random Atom. Either an Instruction, Closer, or Literal.

        """
        atom_type = self.distribution.sample()
        if atom_type is GeneTypes.INPUT:
            return self.random_input()
        elif atom_type is GeneTypes.INSTRUCTION:
            return self.random_instruction()
        elif atom_type is GeneTypes.CLOSE:
            return Closer()
        elif atom_type is GeneTypes.LITERAL:
            return self.random_literal()
        elif atom_type is GeneTypes.ERC:
            return self.random_erc()
        else:
            raise ValueError("GenomeSpawner distribution bad atom type {t}".format(t=str(atom_type)))

    def spawn_genome(self, size: Union[int, Sequence[int]]) -> Genome:
        """Return a random Genome based on the GenomeSpawner's distribution.

        The genome will contain the specified number of Atoms if size is an
        integer. If size is a pair of integers, the genome will be of a random
        size in the range of the two integers.

        Parameters
        ----------
        size
            The resulting genome will contain this many Atoms if size is an
            integer. If size is a pair of integers, the genome will be of a random
            size in the range of the two integers.

        Returns
        -------
        pyshgp.gp.genome.Genome
            A Genome with random contents of a given size.

        """
        if isinstance(size, Sequence):
            size = np.random.randint(size[0], size[1]) + 1
        return Genome([self.random_gene() for _ in range(size)])

class GenomeSimplifier:
    """Simplifies a genome while preserving, or improving, its error.

    Genomes, and Push programs, can contain superfluous Push code. This extra
    code often has no effect on the program behavior, but occasionally it can
    introduce subtle errors or behaviors that is not covered by the training
    cases. Removing the superfluous code makes genomes (and thus programs)
    smaller and easier to understand. More importantly, simplification can
    improve the generalization of the given genome/program.

    The process of genome simplification is iterative and closely resembles
    simple hill climbing. For each iteration, the simplifier will randomly
    select a small number of random genes to remove. The Genome is re-evaluated
    and if its error gets worse, the change is reverted. After repeating this
    for some number of steps, the resulting genome will be the same size or
    smaller while containing the same (or better) error value.

    Reference:
    "Improving generalization of evolved programs through automatic simplification"
    Thomas Helmuth, Nicholas Freitag McPhee, Edward Pantridge, and Lee Spector. 2017.
    In Proceedings of the Genetic and Evolutionary Computation Conference (GECCO '17).
    ACM, New York, NY, USA, 937-944. DOI: https://doi.org/10.1145/3071178.3071330

    See: https://dl.acm.org/citation.cfm?id=3071178.3071330

    """

    # @TODO: Add noop swaps to simplification.

    def __init__(self,
                 evaluator: Evaluator,
                 program_signature: ProgramSignature):
        self.evaluator = evaluator
        self.program_signature = program_signature

    def _remove_rand_genes(self, genome: Genome) -> Genome:
        # @todo DRY with deletion variation operator.
        gn = genome
        n_genes_to_remove = min(np.random.randint(1, 4), len(genome) - 1)
        ndx_of_genes_to_remove = np.random.choice(np.arange(len(gn)), n_genes_to_remove, replace=False)
        ndx_of_genes_to_remove[::-1].sort()
        for ndx in ndx_of_genes_to_remove:
            gn = gn.delete(ndx)
        return gn

    def _errors_of_genome(self, genome: Genome) -> np.ndarray:
        cb = genome_to_code(genome)
        program = Program(code=cb, signature=self.program_signature)
        return self.evaluator.evaluate(program)

    @tap
    def _step(self, genome: Genome, errors_to_beat: np.ndarray) -> Tuple[Genome, np.ndarray]:
        new_gn = self._remove_rand_genes(genome)
        new_errs = self._errors_of_genome(new_gn)
        if np.sum(new_errs) <= np.sum(errors_to_beat):
            return new_gn, new_errs
        return genome, errors_to_beat

    @tap
    def simplify(self,
                 genome: Genome,
                 original_errors: np.ndarray,
                 steps: int = 2000) -> Tuple[Genome, np.ndarray]:
        """Simplify the given genome while maintaining error.

        Parameters
        ----------
        genome
            The Genome to simplify.
        original_errors
            Error vector of the genome to simplify.
        steps
            Number of simplification iterations to perform. Default is 2000.

        Returns
        -------
        pyshgp.gp.genome.Genome
            The shorter Genome that expresses the same computation.

        """
        gn = genome
        errs = original_errors
        for step in range(steps):
            gn, errs = self._step(gn, errs)
            if len(gn) == 1:
                break
        return gn, errs

max_stars_repo_name: RedBeansAndRice/pyshgp

max_issues_repo_name: RedBeansAndRice/pyshgp

max_forks_repo_name: RedBeansAndRice/pyshgp


text:
# test different matching algorithms

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

from dhdt.generic.handler_im import \
    nan_resistant_conv2, nan_resistant_diff2
from dhdt.generic.filtering_statistical import mad_filtering
from dhdt.generic.test_tools import create_sample_image_pair, \
    construct_phase_plane, signal_to_noise, test_phase_plane_localization, \
    test_normalize_power_spectrum, test_subpixel_localization, \
    construct_correlation_peak, \
    create_sheared_image_pair, create_scaled_image_pair
from dhdt.preprocessing.image_transforms import \
    histogram_equalization
from dhdt.processing.matching_tools import \
    get_integer_peak_location
from dhdt.processing.matching_tools_frequency_affine import \
    scaling_through_power_summation
from dhdt.processing.matching_tools_frequency_correlators import \
    phase_corr, phase_only_corr, robust_corr, \
    orientation_corr, windrose_corr, binary_orientation_corr, \
    cosi_corr, sign_only_corr, cosine_corr, cross_corr, \
    symmetric_phase_corr, amplitude_comp_corr, upsampled_cross_corr, \
    projected_phase_corr
from dhdt.processing.matching_tools_frequency_subpixel import \
    phase_svd, phase_pca, phase_tpss, phase_difference, phase_ransac, \
    phase_lsq, phase_radon, phase_hough, phase_gradient_descend,\
    phase_weighted_pca, phase_secant #, phase_jac # jac:mag er later wel uit
from dhdt.processing.matching_tools_frequency_filters import \
    normalize_power_spectrum, perdecomp, low_pass_circle, raised_cosine, \
    local_coherence, thresh_masking, adaptive_masking, gaussian_mask, \
    low_pass_ellipse

from dhdt.processing.matching_tools_differential import \
    simple_optical_flow, hough_optical_flow, affine_optical_flow

from dhdt.processing.matching_tools_spatial_correlators import \
    normalized_cross_corr, cumulative_cross_corr, sum_sq_diff, \
    sum_sad_diff, weighted_normalized_cross_correlation
from dhdt.processing.matching_tools_spatial_subpixel import \
    get_top_moment, get_top_gaussian, get_top_parabolic, \
    get_top_birchfield, get_top_2d_gaussian, get_top_paraboloid, \
    get_top_blue

im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**4, max_range=1, integer=False)
print('di:{:+.4f}'.format(di)+' dj:{:+.4f}'.format(dj))
mu = np.mean(im2.ravel())

W1 = np.random.rand(2**5,2**5)<.5
W2 = np.random.rand(2**9,2**9)<.5
#W1, W2 = im2>mu, im1_big>mu
wncc = weighted_normalized_cross_correlation(im2,im1_big,
                                             W1, W2)

di, dj = 2*np.random.random()-1, 2*np.random.random()-1
#di, dj = 10*di, 10*dj
C_tilde = construct_correlation_peak(np.zeros((2**7,2**7)), di, dj)

ddi,ddj,ii,jj = get_top_blue(C_tilde, ds=1)

from dhdt.presentation.image_io import make_image

from dhdt.preprocessing.shadow_filters import anistropic_diffusion_scalar

im1_ani = anistropic_diffusion_scalar(im1_big)

m,n = 20,16
nan_elem = 40
I = np.random.random((m,n))
I_clean = np.copy(I)
nan_idx = np.random.randint(m*n, size=nan_elem)
I[np.unravel_index(nan_idx, I.shape, 'C')] = np.nan

#dy = np.outer(np.array([2,4,2]),np.array([1,2,1]))
dy = np.outer(np.array([-1,0,+1]),np.array([1,2,1]))
I_dy = nan_resistant_diff2(I,dy)

di_hof,dj_hof,sc_hof = hough_optical_flow(im1,im2,param_resol=40, sample_fraction=1)
print('di:{:+.4f}'.format(di_hof[0])+' dj:{:+.4f}'.format(dj_hof[0]))

im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**4, max_range=1, integer=False)
C_cos = cosine_corr(np.pad(im1,((4,16),(10,10))),\
                             np.pad(im1,((0,20),(0,20))))

C_cos = cosine_corr(im1,im2)

# di_cos, dj_cos

from dhdt.preprocessing.shadow_filters import anistropic_diffusion_scalar

im1_ani = anistropic_diffusion_scalar(im1_big)
wncc = weighted_normalized_cross_correlation(im2,im1_big[200:-200,200:-200])

im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**5, max_range=1, integer=False)

beta_1,beta_2 = 0.35, 0.5
W_1,W_2 = raised_cosine(im1,beta_1), raised_cosine(im2,beta_2)
im1,_ = perdecomp(im1)
im2,_ = perdecomp(im2)
S_1,S_2 = np.fft.fft2(im1), np.fft.fft2(im2)
Q_hat = (np.multiply(W_1,S_1))*np.conjugate(np.multiply(W_2,S_2))

C_hat = np.real(np.fft.ifft2(Q_hat))

Q_tilde = construct_phase_plane(im1, di, dj, indexing='ij')
W = low_pass_ellipse(Q_tilde, r1=.1, r2=.3)

C_tilde = np.real(np.fft.fftshift(np.fft.ifft2(W*Q_tilde)))

np.fft.fftshift(C_hat)

im1, im2, di, dj, im1_big = create_sheared_image_pair(d=2**5,
                                                      sh_i=0.00, sh_j=0.04,
                                                      max_range=1)
print(di, dj)
u, v, A, snr = affine_optical_flow(im1, im2, model='affine', iteration=10)

im1, im2, di, dj, im1_big = create_sheared_image_pair(d=2**5,
                                                      sh_i=0.00, sh_j=0.00,
                                                      max_range=1)

u,v = hough_optical_flow(im1, im2, num_estimates=1)

print('test phase plane localizators')
#tries = 2
#d_stack = np.zeros((tries,4))
#for i in np.arange(tries):

im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**5, max_range=1, integer=False)
di_hof,dj_hof = hough_optical_flow(im1,im2,param_resol=100, sample_fraction=1)

    # load data
#im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**7, max_range=100, integer=True)
#Q_tilde = construct_phase_plane(im1, di, dj)
#W_tilde = np.ones_like(Q_tilde, dtype=np.bool_)
#print('di:{:+.4f}'.format(di)+' dj:{:+.4f}'.format(dj))

#im1,im2,di,dj,_ = create_sheared_image_pair(d=2**7,sh_x=-0.1, sh_y=0.2, max_range=1)

im1,im2,di,dj,_ = create_scaled_image_pair(d=2**7,sc_x=1.3, sc_y=1.0, max_range=0)

im1,_ = perdecomp(im1)
im2,_ = perdecomp(im2)
S_1,S_2 = np.fft.fft2(im1), np.fft.fft2(im2)

a,b = scaling_through_power_summation(S_1,S_2)

di_soc, dj_soc = sign_only_corr(im1, im2)

di_prj, dj_prj = projected_phase_corr(im1, im2)

di_cos, dj_cos = cosine_corr(im1, im2)

beta_1,beta_2 = 0.35, 0.5
W_1,W_2 = raised_cosine(im1,beta_1), raised_cosine(im2,beta_2)
im1,_ = perdecomp(im1)
im2,_ = perdecomp(im2)
S_1,S_2 = np.fft.fft2(im1), np.fft.fft2(im2)
Q_hat = (np.multiply(W_1,S_1))*np.conjugate(np.multiply(W_2,S_2))

# test Jacobian
#dQd1,dQd2 = phase_jac(Q_tilde, np.array([di+.3, dj-.4]))

di_sec, dj_sec = phase_secant(Q_tilde)
print('di:{:+.4f}'.format(di_sec)+' dj:{:+.4f}'.format(dj_sec))
#di_gd, dj_gd = phase_gradient_descend(Q_tilde)

W_hat = adaptive_masking(Q_hat)
di_sec, dj_sec = phase_secant(Q_hat, W=W_hat)
(m, snr) = phase_tpss(Q_hat, W_hat, np.array([0.0, 0.0]), j=1e2)

#OK di_svd, dj_svd = phase_svd(Q_tilde, W_tilde, rad=0.4)
#OK di_lsq, dj_lsq = phase_lsq(Q_tilde, W_tilde)
di_pca, dj_pca = phase_pca(Q_tilde, W_tilde)

plt.imshow(np.fft.fftshift(np.angle(Q_tilde)), vmin=-np.pi, vmax=+np.pi),
plt.colorbar(), plt.show()
plt.imshow(np.fft.fftshift(np.angle(Q_hat)), vmin=-np.pi, vmax=+np.pi),
plt.colorbar(), plt.show()

W_gauss = gaussian_mask(Q_tilde)

di_pcw, dj_pcw = phase_weighted_pca(Q_tilde, W_gauss)
#    d_stack[i,:] = np.array([di, dj, di_pca, dj_pca])

#import matplotlib.pyplot as plt
#plt.scatter(d_stack[:,0],d_stack[:,2]), plt.show()

#plt.scatter(d_stack[:,1],d_stack[:,3]), plt.show()

W_hat = adaptive_masking(Q_hat)

(m, snr) = phase_tpss(Q_hat, W_hat, np.array([0.0, 0.0]), j=5e2)
(ddi,ddj) = test_subpixel_localization(m[0],m[1], di,dj, tolerance=.1)
print('phase_tpss is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

di_gd, dj_gd = phase_gradient_descend(Q_hat,W_hat)
di_pca, dj_pca = phase_pca(Q_hat, W_hat)

#di_hat, dj_hat = phase_ransac(Q_tilde, precision_threshold=.05)
#(ddi,ddj) = test_subpixel_localization(di_hat,dj_hat,di,dj, tolerance=.1)
#print('phase_ransac is OK, maximum difference is '+\
#      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

print('di:{:.4f}'.format(di)+' dj:{:.4f}'.format(dj))
m_close = np.array([np.round(1e1*di)/1e1, \
                    np.round(1e1*dj)/1e1])
#(m, snr) = phase_newton(Q_tilde, np.ones_like(Q_tilde,dtype=np.float32), \
#                      m_close, j=5)

# load data
im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**7, max_range=10)

print('testing (spatial) correlators')
C_hat = normalized_cross_corr(im2, im1_big)

#C_hat = cumulative_cross_corr(im2, im1_big)
#di_hat,dj_hat,snr,corr = get_integer_peak_location(C_hat)
#assert( np.isclose(np.round(di), di_hat) and \
#       np.isclose(np.round(dj), dj_hat) )
#ddi,ddj = di-di_hat, dj-dj_hat
#print('sum_sq_diff correlator is OK, maximum difference is '+\
#      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

C_hat = -1*sum_sq_diff(im2, im1_big)
di_hat,dj_hat,snr,corr = get_integer_peak_location(C_hat)
assert( np.isclose(np.round(di), di_hat) and \
       np.isclose(np.round(dj), dj_hat) )
ddi,ddj = di-di_hat, dj-dj_hat
print('sum_sq_diff correlator is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

temp_size = 2**6
im1,im2,di,dj,im1_big = create_sample_image_pair(d=temp_size, max_range=1)
print('{:.4f}'.format(di))
print('{:.4f}'.format(dj))

(di_hat,dj_hat,di_lamb,dj_lamb) = simple_optical_flow(im1, im2, temp_size, \
                                                      temp_size//2, \
                                                      temp_size//2,
                                                      sigma=2)

S1, S2 = np.fft.fft2(im1), np.fft.fft2(im2)
di_02,dj_02 = upsampled_cross_corr(S1,S2, upsampling=2) #todo
di_04,dj_04 = upsampled_cross_corr(S1,S2, upsampling=4)
di_08,dj_08 = upsampled_cross_corr(S1,S2, upsampling=8)
di_16,dj_16 = upsampled_cross_corr(S1,S2, upsampling=16)
di_32,dj_32 = upsampled_cross_corr(S1,S2, upsampling=32)

C_hat = cosine_corr(im1, im2)
di_hat,dj_hat,snr,corr = get_integer_peak_location(C_hat)

print('test phase cycle localizators')

im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**4, max_range=10)
Q_tilde = construct_phase_plane(im1, di, dj)

di_tilde, dj_tilde = phase_hough(Q_tilde, sample_fraction=1)

Q_hat = phase_only_corr(im1, im2)
Q_hat = robust_corr(im1, im2)
di_hat, dj_hat = phase_hough(Q_hat, sample_fraction=.01)

di_hat,dj_hat,snr,corr = get_integer_peak_location(C_hat)
assert( np.isclose(np.round(di), di_hat) and \
       np.isclose(np.round(dj), dj_hat) )
ddi,ddj = di-di_hat, dj-dj_hat
print('normalized cross correlator is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

print('testing (spatial) subpixel estimators')

di, dj = 2*np.random.random()-1, 2*np.random.random()-1
di, dj = 10*di, 10*dj
C_tilde = construct_correlation_peak(np.zeros((2**7,2**7)), di, dj)

# test test-function
di_tilde,dj_tilde,_,_ = get_integer_peak_location(C_tilde)
assert np.round(di)==di_tilde
assert np.round(di)==di_tilde

ddi,ddj,ii,jj = get_top_moment(C_tilde, ds=1)
assert np.isclose(di, ii+ddi, atol=2e-1)
assert np.isclose(dj, jj+ddj, atol=2e-1)
dddi,dddj = (di-(ii+ddi)), (dj-(jj+ddj))
print('moment locator is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(dddi), abs(dddj))) + ' pixels')

di, dj = 2*np.random.random()-1, 2*np.random.random()-1
di, dj = 10*di, 10*dj
C_tilde = construct_correlation_peak(np.zeros((2**7,2**7)), di, dj)

ddi,ddj,ii,jj = get_top_gaussian(C_tilde)
assert np.isclose(di, ii+ddi, atol=5e-1)
assert np.isclose(dj, jj+ddj, atol=5e-1)
dddi,dddj = (di-(ii+ddi)), (dj-(jj+ddj))
print('gaussian locator is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(dddi), abs(dddj))) + ' pixels')

ddi,ddj,ii,jj = get_top_parabolic(C_tilde)
assert np.isclose(di, ii+ddi, atol=5e-1)
assert np.isclose(dj, jj+ddj, atol=5e-1)
dddi,dddj = (di-(ii+ddi)), (dj-(jj+ddj))
print('parabolic locator is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(dddi), abs(dddj))) + ' pixels')

di, dj = 2*np.random.random()-1, 2*np.random.random()-1
di, dj = 10*di, 10*dj
C_tilde = construct_correlation_peak(np.zeros((2**7,2**7)), di, dj, fwhm=3)
print('{:.4f}'.format(di))
print('{:.4f}'.format(dj))

ddi,ddj,ii,jj = get_top_paraboloid(C_tilde)
(dj, (jj+ddj))
(di, (ii+ddi))

ddi,ddj,ii,jj = get_top_2d_gaussian(C_tilde)

ddi,ddj,ii,jj = get_top_birchfield(C_tilde)
(dj, (jj+ddj))
(di, (ii+ddi))
assert np.isclose(di, ii+ddi, atol=5e-1)
assert np.isclose(dj, jj+ddj, atol=5e-1)
dddi,dddj = (di-(ii+ddi)), (dj-(jj+ddj))
print('birchfield locator is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(dddi), abs(dddj))) + ' pixels')

Q_hat = phase_only_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('phase_only_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

print('test phase cycle localizators')

im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**7, max_range=10)
Q_tilde = construct_phase_plane(im1, di, dj)
print('di:{:.4f}'.format(di)+' dj:{:.4f}'.format(dj))

di_hat, dj_hat = phase_difference(Q_tilde)
(ddi,ddj) = test_subpixel_localization(di_hat,dj_hat,di,dj, tolerance=.1)
print('phase_difference is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

di_hat, dj_hat = phase_radon(Q_tilde, coord_system='ij')
(ddi,ddj) = test_subpixel_localization(di_hat,dj_hat,di,dj, tolerance=.1)
print('phase_difference is OK, maximum difference is '+\
      '{:.4f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

# load data
im1,im2,di,dj,im1_big = create_sample_image_pair(d=2**7, max_range=10)
Q_tilde = construct_phase_plane(im1, di, dj)

(im1,_), (im2,_) = perdecomp(im1), perdecomp(im2)

print('testing (frequency) correlators')

Q_hat = phase_only_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('phase_only_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

Q_hat = symmetric_phase_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('symmetric_phase_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

Q_hat = amplitude_comp_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('amplitude_comp_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

Q_hat = phase_corr(im1, im2)
_ = test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('phase_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

Q_hat = orientation_corr(im1, im2)
_ = test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('orientation_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

#Q_hat = windrose_corr(im1, im2)
#_ = test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
#snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
#print('windrose_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

Q_hat = cross_corr(im1, im2)
_ = test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('cross_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

print('binary_orientation_corr does not seem to do a good job')
# Q_hat = binary_orientation_corr(im1, im2)
# _ = test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
# snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
# print('binary_orientation_corr correlator is OK, snr is '+'{:.2f}'.format(snr))

(Q_hat,W,m0) = cosi_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('cosi_corr is correlator OK, snr is '+'{:.2f}'.format(snr))

Q_hat = robust_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=1)
snr = signal_to_noise(Q_hat, Q_tilde, norm=1)
print('robust_corr is correlator OK, snr is '+'{:.2f}'.format(snr))

(m, n) = Q_tilde.shape
fi = np.flip( (np.arange(0,m)-(m/2)) /m)
fj = np.flip( (np.arange(0,n)-(n/2)) /n)

Fj = np.repeat(fj[np.newaxis,:],m,axis=0)
Fi = np.repeat(fi[:,np.newaxis],n,axis=1)
Fj = np.fft.fftshift(Fj)
Fi = np.fft.fftshift(Fi)

di_hat, dj_hat = phase_lsq(Fi.flatten(), Fj.flatten(), Q_tilde.flatten())
(ddi,ddj) = test_subpixel_localization(di_hat,dj_hat,di,dj, tolerance=.1)
print('phase_lsq is OK, maximum difference is '+\
      '{:.2f}'.format(np.maximum(abs(ddi), abs(ddj))) + ' pixels')

C_hat = sign_only_corr(im1, im2)
test_phase_plane_localization(Q_hat, di, dj, tolerance=.5)
snr = signal_to_noise(Q_tilde, Q_hat, norm=1)
print('cosine_corr is OK, snr is '+'{:.2f}'.format(snr))

test_phase_plane_localization

#Q_hat = phase_corr(im1, im2) # cross-power spectrum
#C_hat = sign_only_corr(im1, im2)
Q_hat = cosine_corr(im1, im2)
#Q_hat = phase_only_corr(im1, im2)
#Q_hat = robust_corr(im1, im2)
#(Q_hat,W,m0) = cosi_corr(im1, im2)

Q_hat = normalize_power_spectrum(Q_hat)
Q = np.fft.fftshift(Q_hat)

C = local_coherence(Q_hat, ds=1)

# look at signal-to-noise
snr = signal_to_noise(Q_hat,Q_tilde,norm=2)
print(snr)

(m,n) = Q.shape
fy = np.flip((np.arange(0,m)-(m/2)) /m)
fx = np.flip((np.arange(0,n)-(n/2)) /n)

Fx = np.repeat(fx[np.newaxis,:],m,axis=0)
Fy = np.repeat(fy[:,np.newaxis],n,axis=1)

#IN = Q!=0
IN = np.fft.fftshift(low_pass_circle(Q, r =.35))

Q_dummy = np.fft.fftshift(Q_tilde)

di = phase_ransac(Q)

di_hat = phase_difference_1d(Q)

di_hat,dj_hat = phase_pca(Fy[IN], Fx[IN], Q_dummy[IN])
di_hat,dj_hat = phase_pca(Fy[IN], Fx[IN], Q[IN])

# Fr = np.cos(np.radians(theta))*Fx \
#     - np.sin(np.radians(theta))*Fy

snr_spat = test_phase_plane_localization(Q_hat, di, dj, tolerance=0.5)

imShow= False
if imShow:
    ax = plt.subplot()
    im = ax.imshow(np.angle(Q_hat))
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    plt.show()

    ax = plt.subplot()
    im = ax.imshow(np.angle(Q_tilde))
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    plt.show()

#estimate period
Q_shift = np.roll(Q, (1,1))

Q_di = np.multiply(np.conj(Q), np.roll(Q, (1,0)))
Q_dj = np.multiply(np.conj(Q), np.roll(Q, (0,1)))

ang_Q_di, ang_Q_dj = np.angle(Q_di), np.angle(Q_dj)

ang_dQ = np.arctan2(ang_Q_di, ang_Q_dj)

if imShow:
    ax = plt.subplot()
    im = ax.imshow(np.angle(Q_di))
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    plt.show()

    ax = plt.subplot()
    im = ax.imshow(np.angle(Q_tilde))
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    plt.show()

max_stars_repo_name: GO-Eratosthenes/dhdt

max_issues_repo_name: GO-Eratosthenes/dhdt

max_forks_repo_name: GO-Eratosthenes/dhdt


text:
# Some norm functions with the weights for MM algorithms
#
# Copyright 2020 Robin Scheibler
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import numpy as np

eps = 1e-15

def lp_norm(E, p=1):
    assert p > 0 and p < 2
    weights = p / np.maximum(eps, 2.0 * np.abs(E) ** (2 - p))
    return weights

def lpq_norm(E, p=1, q=2, axis=1):
    assert p > 0 and q >= p and q <= 2.0

    rn = np.sum(np.abs(E) ** q, axis=axis, keepdims=True) ** (1 - p / q)
    qfn = np.abs(E) ** (2 - q)
    weights = p / np.maximum(eps, 2.0 * rn * qfn)
    return weights


text:
# -*- coding: utf-8 -*-
"""

Author
------
Bo Zhang

Email
-----
bozhang@nao.cas.cn

Created on
----------
- Fri Feb  26 21:44:00 2016     convolution

Modifications
-------------
- Sun Mar   6 21:00:00 2016     convolve from arbitrary R_hi to arbitrary R_lo

Aims
----
- degrade high resolution spectra to low resolution spectra

"""

# ###########################################################################
# Enlightening example:
#
# for one pixel:
# R_hi = 2000 @5000A    FWHM_hi =  2.5A
# R_lo =  500 @5000A    FWHM_lo = 10.0A
# FWHM_GK = sqrt(10.0**2 - 2.5**2) = 9.68A
#
# for next pixel:
# R_hi = 1000 @5000A    FWHM_hi =  5.0A
# R_lo =  500 @5000A    FWHM_lo = 10.0A
# FWHM_GK = sqrt(10.0**2 - 5.0**2) = 8.66A
#
# Therefore, to keep the number of pixels of Gaussian Kernels the same,
# we have to adjust the FWHM_interp (and thus the R_interp).
#
# For pixels 2000->500, the FWHM_GK = 9.68 (wider GK),
# while for pixels 1000-> 500, FWHM_GK = 8.66 (narrower GK),
# thus the R_interp should be 9.68/8.66=1.12 times higher,
# i.e., the delta_lambda should be  1./1.12 times smaller.
# ###########################################################################

from __future__ import print_function

import datetime

import numpy as np
from scipy.interpolate import PchipInterpolator

from .spec import spec_quick_init

import astropy.constants as const
import astropy.units as u
from inspect import isfunction

# deprecated constants ########################################################
# OVER_SAMPLING = 10.
# # threshold for R_hi_spec/R_hi
# # if R_hi_spec > OVER_SAMPLING*R_hi, then R_interp = R_hi_spec
# # else R_interp = OVER_SAMPLING * np.max([R_hi_spec, R_hi])
#
# KERNEL_LENGTH_FWHM = 4.24
# # kernel array length is 4.24 times FWHM, i.e., 10 sigma
# #############################################################################

# ########################################################################### #
# ################  transformation between R and FWHM  ###################### #
# ########################################################################### #

def resolution2fwhm(R, wave=5000.):
    # assert R is positive
    if np.isscalar(R):
        assert R > 0.
    else:
        assert np.all(R > 0.)

    # assert wave is correct in shape
    assert np.isscalar(wave) or len(wave) == len(R)

    return wave / R

def fwhm2resolution(fwhm, wave=5000.):
    # assert fwhm is positive
    if np.isscalar(fwhm):
        assert fwhm > 0.
    else:
        assert np.all(fwhm > 0.)

    # assert wave is correct in shape
    assert np.isscalar(wave) or len(wave) == len(fwhm)

    return wave / fwhm

# ########################################################################### #
# ######################  generate wave array given R  ###################### #
# ########################################################################### #

def _generate_wave_array_R_fixed(wave_start, wave_stop, R=2000.,
                                over_sample=1.):
    """ generate wave array given a fixed R """
    R_ = over_sample * R - .5
    # determine wave_step_min
    wave_step_min = wave_start / R_
    # wave guess
    wave_step_guess = np.zeros((wave_stop-wave_start)/wave_step_min)
    wave_guess = np.zeros_like(wave_step_guess)
    wave_step_guess[0] = wave_step_min
    wave_guess[0] = wave_start
    # iterate for real
    for i in np.arange(1, len(wave_guess)):
        wave_guess[i] = wave_guess[i-1] + wave_step_guess[i-1]
        wave_step_guess[i] = wave_guess[i] / R_
    return wave_guess[
        np.logical_and(wave_guess >= wave_start, wave_guess <= wave_stop)]

def _generate_wave_array_R_func(wave_start, wave_stop, R=(lambda x: x),
                                over_sample=1., wave_test_step=1.):
    """ generate wave array given R as a function """
    # determine wave_step_min
    wave_test = np.arange(wave_start, wave_stop, wave_test_step)
    R_test = over_sample * R(wave_test)
    wave_step_min = np.min(wave_test / R_test)
    # wave guess
    wave_guess = np.zeros(np.int(np.ceil((wave_stop-wave_start)/wave_step_min)))
    wave_guess[0] = wave_start
    # iterate for real # single side R !!!
    for i in np.arange(1, len(wave_guess)):
        wave_guess[i] = wave_guess[i-1] + \
                        wave_guess[i-1] / (over_sample * R(wave_guess[i-1]))
    return wave_guess[
        np.logical_and(wave_guess >= wave_start, wave_guess <= wave_stop)]

def generate_wave_array_R(wave_start, wave_stop, R=2000.,
                          over_sample=1., wave_test_step=1.):
    """ generate a wavelength array matching the given R

    Parameters
    ----------
    wave_start: float
        start from this wavelength
    wave_stop: float
        stop at this wavelength
    R: float or function
        specify a fixed R or specify R as a function of wavelength
    over_sample: float
        over-sampling rate, default is 1.
    wave_test_step:
        used to determing the wave_step_min

    Returns
    -------
    wave: array
        an array matching the given R

    Example
    -------
    >>> def R(x): return 0.2*x
    >>> wave_array_R = generate_wave_array_R(4000., 5000., R)

    """
    if np.isscalar(R):
        # if R is scalar
        return _generate_wave_array_R_fixed(
            wave_start, wave_stop, R=R,
            over_sample=over_sample)
    else:
        # if R is a function / Interpolator
        return _generate_wave_array_R_func(
            wave_start, wave_stop, R=R,
            over_sample=over_sample, wave_test_step=wave_test_step)

# ########################################################################### #
# ################  generate wave array given delta_lambda  ################# #
# ########################################################################### #

def _generate_wave_array_delta_lambda_fixed(wave_start, wave_stop,
                                            delta_lambda,
                                            over_sample=1.):
    """ generate a wavelength array matching the given delta_lambda (fixed) """
    return np.arange(wave_start, wave_stop, delta_lambda / over_sample)

def _generate_wave_array_delta_lambda_func(wave_start, wave_stop,
                                           delta_lambda=(lambda x: 1.),
                                           over_sample=1.,
                                           wave_test_step=1.):
    """ generate a wavelength array matching the given delta_lambda
        (specified as a function of wavelength) """
    wave_test = np.arange(wave_start, wave_stop, wave_test_step)
    delta_lambda_min = np.min(delta_lambda(wave_test))
    wave_guess = np.arange(wave_start, wave_stop, delta_lambda_min)
    for i in xrange(1, len(wave_guess)):
        wave_guess[i] = \
            wave_guess[i-1] + delta_lambda(wave_guess[i-1]) / over_sample
    return wave_guess[
        np.logical_and(wave_guess >= wave_start, wave_guess <= wave_stop)]

def generate_wave_array_delta_lambda(wave_start, wave_stop,
                                     delta_lambda=(lambda x: 1.),
                                     over_sample=1.,
                                     wave_test_step=1.):
    """ generate a wavelength array matching the given delta_lambda
        (delta_lambda given as a fixed number or a function)
    Parameters
    ----------
    wave_start: float
        where the wavelength starts
    wave_stop: float
        where the wavelength stops
    delta_lambda: float or function
        specifies the delta_lambda as a fixed number or a function of wavelength
    over_sample: float
        over-sampling
    wave_test_step: float
        tests for the smallest wave_guess step

    Returns
    -------
    wave_guess: array

    Example
    -------
    >>> def dl(x): return 0.002*x
    >>> wave_array_dl = generate_wave_array_delta_lambda(4000., 5000., dl)

    """
    if np.isscalar(delta_lambda):
        # if delta_lambda is scalar
        return _generate_wave_array_delta_lambda_fixed(
            wave_start, wave_stop, delta_lambda=delta_lambda,
            over_sample=over_sample)
    else:
        # if delta_lambda is a function / Interpolator
        return _generate_wave_array_delta_lambda_func(
            wave_start, wave_stop, delta_lambda=delta_lambda,
            over_sample=over_sample, wave_test_step=wave_test_step)

# ########################################################################### #
# ############## find spectral R (FWHM) and R_max (FWHM_min) ################ #
# ########################################################################### #

def find_R_for_wave_array(wave):
    """ find the R of wavelength array (sampling resolution array) """
    wave = wave.flatten()
    wave_diff = np.diff(wave)
    wave_diff_ = (wave_diff[1:] + wave_diff[:-1]) / 2.
    return np.hstack((
        wave[0]/wave_diff[0], wave[1:-1]/wave_diff_, wave[-1]/wave_diff[-1]))

def find_R_max_for_wave_array(wave):
    """ find the maximum sampling resolution of a given wavelength array """
    return np.max(find_R_for_wave_array(wave))

def find_delta_lambda_for_wave_array(wave):
    """ find the delta_lambda of wavelength array (delta_lambda array) """
    wave = wave.flatten()
    wave_diff = np.diff(wave)
    wave_diff_ = (wave_diff[1:] + wave_diff[:-1]) / 2.
    return np.hstack((wave_diff[0], wave_diff_, wave[-1]))

def find_delta_lambda_min_for_wave_array(wave):
    """ find the minimum delta_lambda of a given wavelength array """
    return np.min(find_delta_lambda_for_wave_array(wave))

# ########################################################################### #
# ###############################  find Rgk  ################################ #
# ########################################################################### #

def find_Rgk(R_hi=2000., R_lo=500., over_sample=1.):
    """ find Rgk as a function of wavelength

    Parameters
    ----------
    R_hi: float or funtion
        higher resolution (as a function of wavelength)
    R_lo: float or funtion
        lower resolution (as a function of wavelength)
    over_sample: float
        over-sampled resolution, default is 1.

    Returns
    -------
    Rgk: function
        Gaussian Kernel resolution as a function of wavelength

    """
    if np.isscalar(R_hi):
        # if R_hi is a fixed number
        R_hi_ = lambda x: R_hi
    else:
        R_hi_ = R_hi

    if np.isscalar(R_lo):
        # if R_lo is a fixed number
        R_lo_ = lambda x: R_lo
    else:
        R_lo_ = R_lo

    Rgk = lambda x: \
        (over_sample * x / np.sqrt((x/R_lo_(x))**2. - (x/R_hi_(x))**2.))
    return Rgk

# def find_appropriate_R_interp_for_convolution(R_hi, R_hi_spec):
#
#     R_hi_spec_R_hi = R_hi_spec / R_hi
#
#     if R_hi_spec_R_hi > OVER_SAMPLING:
#         # already over sampled
#         R_interp = R_hi_spec
#     elif R_hi_spec_R_hi >= 1.:
#         # already over sampled
#         R_interp = OVER_SAMPLING * R_hi_spec
#     else:
#         R_interp = OVER_SAMPLING * R_hi
#
#     return R_interp

# def find_gaussian_kernel_fwhm(R_hi, R_lo, return_type='fwhm'):
#     assert R_hi > R_lo
#     fwhm_hi = resolution2fwhm(R_hi)
#     fwhm_lo = resolution2fwhm(R_lo)
#     fwhm = np.sqrt(fwhm_lo**2. - fwhm_hi**2.)
#     if return_type == 'fwhm':
#         return fwhm
#     elif return_type == 'R':
#         return fwhm2resolution(fwhm)

# ########################################################################### #
# #########################  Gaussian Kernel ################################ #
# ########################################################################### #

def fwhm2sigma(fwhm):
    return fwhm / (2. * np.sqrt(2. * np.log(2.)))

def sigma2fwhm(sigma):
    return 2. * np.sqrt(2. * np.log(2.)) * sigma

def normalized_gaussian_array(x, b=0., c=1.):
    # a = 1. / (np.sqrt(2*np.pi) * c)
    ngs_arr = np.exp(- (x-b)**2. / (2.*c**2.))
    return ngs_arr / np.sum(ngs_arr)

def generate_gaussian_kernel_array(over_sample_Rgk, sigma_num):
    """ generate gaussian kernel array according to over_sample_Rgk

    Parameters
    ----------
    over_sample_Rgk: float
        over_sample rate
    sigma_num: float
        1 sigma of the Gaussian = sigma_num pixels

    Returns
    -------
    normalized gaussian array

    """
    sigma_pixel_num = fwhm2sigma(over_sample_Rgk)

    array_length = np.fix(sigma_num * sigma_pixel_num)
    if array_length % 2 == 0:
        array_length += 1
    array_length_half = (array_length-1) / 2.

    xgs = np.arange(- array_length_half, array_length_half + 1)
    return normalized_gaussian_array(xgs, b=0., c=sigma_pixel_num)

def conv_spec(wave,
              flux,
              R_hi=2000.,
              R_lo=500.,
              over_sample_additional=3.,
              gaussian_kernel_sigma_num=8.,
              wave_new=None,
              wave_new_oversample=5.,
              verbose=True,
              return_type='array'):
    """ to convolve high-R spectrum to low-R spectrum

    Parameters
    ----------
    wave: array
        wavelength
    flux: array
        flux array
    R_hi: float or function
        higher resolution
    R_lo: float or function
        lower R
    over_sample_additional: float
        additional over-sample rate
    gaussian_kernel_sigma_num: float
        the gaussian kernel width in terms of sigma
    wave_new: None or float or array
        if None: wave_new auto-generated using wave_new_oversample
        if float: this specifies the over-sample rate
        if voctor: this specifies the new wave_new array
    verbose: bool
        if True, print the details on the screen
    return_type: string
        if 'array': return wave and flux as array
        if 'spec': retrun spec object

    Returns
    -------
    spec: bopy.spec.spec.Spec
        Spec, based on astropy.table.Table class

    """
    if verbose:
        start = datetime.datetime.now()
        print('--------------------------------------------------------------')
        print('@Cham: Welcome to the spectral convolution code developed by Bo Zhang (@NAOC) ...')

    # 1. re-format R_hi & R_lo
    assert R_hi is not None and R_lo is not None

    if np.isscalar(R_hi):
        R_hi_ = lambda x: R_hi
    else:
        R_hi_ = R_hi

    if np.isscalar(R_lo):
        R_lo_ = lambda x: R_lo
    else:
        R_lo_ = R_lo

    # 2. find Rgk
    Rgk = find_Rgk(R_hi_, R_lo_, over_sample=1.)

    # 3. find appropriate over_sample
    R_hi_specmax = find_R_max_for_wave_array(wave)
    R_hi_max = np.max(R_hi_(wave))
    over_sample = over_sample_additional * np.fix(np.max([
        R_hi_specmax/Rgk(wave), R_hi_max/Rgk(wave)]))

    # 4. find wave_interp & flux_interp
    if verbose:
        print('@Cham: interpolating orignal spectrum to wave_interp ...')
    wave_max = np.max(wave)
    wave_min = np.min(wave)
    wave_interp = generate_wave_array_R(wave_min, wave_max,
                                        Rgk, over_sample=over_sample)
    P = PchipInterpolator(wave, flux, extrapolate=None)
    flux_interp = P(wave_interp)
    assert not np.any(np.isnan(flux_interp))

    # 5. generate Gaussian Kernel array
    if verbose:
        print('@Cham: generating gaussian kernel array ...')
    gk_array = generate_gaussian_kernel_array(over_sample,
                                              gaussian_kernel_sigma_num)
    gk_len = len(gk_array)
    gk_len_half = np.int((gk_len - 1) / 2.)

    # 6. convolution
    if verbose:
        print('@Cham: convolution ...')
        print('@Cham: estimated convolution time: %.2f seconds ...'
              % (len(flux_interp) * gk_len / 55408. / 657. * 0.05))
    convolved_flux = np.convolve(flux_interp, gk_array)[gk_len_half:-gk_len_half]

    # 7. find new wave array
    if wave_new is None:
        # wave_new is None
        # default: 5 times over-sample
        if verbose:
            print('@Cham: using default 5 times over-sample wave array ...')
        wave_new = generate_wave_array_R(wave_interp[0], wave_interp[-1],
                                         R_lo, wave_new_oversample)
    elif np.isscalar(wave_new):
        # wave_new specifies the new wave array over_sampling_lo rate
        # default is 5. times over-sample
        if verbose:
            print('@Cham: using user-specified %.2f times over-sample wave array ...' % wave_new)
        wave_new = generate_wave_array_R(wave_interp[0], wave_interp[-1],
                                         R_lo, wave_new)
    else:
        # wave_new specified
        if verbose:
            print('@Cham: using user-specified wave array ...')

    # 8. interpolate convolved flux to new wave array
    if verbose:
        print('@Cham: interpolating convolved spectrum to new wave array ...')
    P = PchipInterpolator(wave_interp, convolved_flux, extrapolate=False)
    flux_new = P(wave_new)
    if verbose:
        stop = datetime.datetime.now()
        print('@Cham: total time spent: %.2f seconds' % (stop-start).total_seconds())
        print('---------------------------------------------------------------')

    if return_type == 'array':
        return wave_new, flux_new
    elif return_type == 'spec':
        return spec_quick_init(wave_new, flux_new)

def test_bc03_degrade_to_R500():
    # test a BC03 population spectrum
    from bopy.spec.spec import Spec
    # 1. read spectrum
    fp = '/home/cham/PycharmProjects/bopy/bopy/data/model_bc03/bc2003_hr_m42_chab_ssp_020.spec'
    data = np.loadtxt(fp)
    spec = Spec(data, names=['wave', 'flux'])
    print(spec)
    spec = spec.extract_chunk_wave_interval([[4000., 8000.]])[0]

    # 2.convolve spectum
    spec_ = conv_spec(spec['wave'], spec['flux'],
                      lambda x: 0.2*x, lambda x: 0.1*x,
                      over_sample_additional=3.,
                      gaussian_kernel_sigma_num=6.,
                      wave_new=None,
                      wave_new_oversample=5.,
                      verbose=False)
    spec_.pprint()

    print(find_R_for_wave_array(spec_['wave']))
    # 3.plot results
    # fig = plt.figure()
    # plt.plot(wave, flux)
    # plt.plot(spec_['wave'], spec_['flux'], 'r')
    # fig.show()
    # fig.savefig(''')
    print('@Cham: test OK ...')

if __name__ == '__main__':
    test_bc03_degrade_to_R500()

max_stars_repo_name: hypergravity/bopy

max_issues_repo_name: hypergravity/bopy

max_forks_repo_name: hypergravity/bopy

reason: import numpy,from scipy,import astropy


text:
#!/usr/bin/env python3
#
#  Copyright (c) 2015 - 2022, Intel Corporation
#  SPDX-License-Identifier: BSD-3-Clause
#

'''
Creates a model relating power-limit to both runtime and energy
consumption. Offers power-governor policy recommendations for
minimum energy use subject to a runtime degradation constraint.
'''

import argparse
import math
import sys

import pandas
import numpy as np
from numpy.polynomial.polynomial import Polynomial

import geopmpy.io

from experiment import common_args
from experiment import machine

class PowerLimitModel(object):
    "Abstract class for predicting a statistic based on a power limit."
    def __init__(self):
        "Default constructor."
        raise NotImplementedError

    def train(self, df, key):
        """
        Train this model to fit the data of column key contained in the
        dataframe (df), which is indexed by power limit. This modifies
        the instance it is called on, preparing it for calls of the
        methods test and batch_test."""
        raise NotImplementedError

    def evaluate(self, PL):
        """
        Evaluate a trained model on a single power limit (PL). Returns
        the model's prediction for this power limit."""
        raise NotImplementedError

    def batch_evaluate(self, PL_list):
        """
        Evaluate a trained model on a list-like container of power limits
        (PL_list). Returns a list of values."""
        return [self.evaluate(PL) for PL in PL_list]

    def serialize(self):
        "Convert model to JSON."
        raise NotImplementedError

    def __str__(self):
        return "<%s>" % "PowerLimitModel"

    def repr(self):
        return self.__str__()

class PolynomialFitPowerModel(PowerLimitModel):
    """
    Implementation of a polynomial model of fixed degree that can be a
    trained to predict a statistic from a power limit."""
    def __init__(self, degree):
        """
        Default constructor; degree (a positive integer) is the degree
        of the polynomial being fit to."""
        self._degree = degree
        self._model = None

    def train(self, df, key):
        new_series, [resid, rank, sv, rcond] = Polynomial.fit(df.index,
                                                              df[key].values,
                                                              deg=self._degree,
                                                              full=True)
        self._model = new_series.convert()
        self._resid = (resid[0] / len(df[key].values))**0.5

    def evaluate(self, PL):
        return sum(self._model.coef * PL ** np.arange(self._degree + 1))

    def residual(self):
        return self._resid

    def __str__(self):
        if self._model is not None:
            terms = ["", " PL"] + \
                    [" PL ** {}".format(i) for i in range(2, self._degree + 1)]
            coeffs = list(self._model.coef)
            expression = ""
            terms.reverse()
            coeffs.reverse()
            firstTerm = True
            for tt, cc in zip(terms, coeffs):
                if not firstTerm:
                    if cc < 0:
                        expression += " - "
                        cc *= -1
                    else:
                        expression += " + "
                expression += str(cc) + tt
                firstTerm = False
            return "<{} {}>".format(self.__class__.__name__, expression)
        else:
            return "<{} (untrained)>".format(self.__class__.__name__)

class CubicPowerModel(PolynomialFitPowerModel):
    """
    Implementation of a cubic model that can be a trained to predict a
    statistic from a power limit."""
    def __init__(self):
        "Simple constructor."
        super(CubicPowerModel, self).__init__(degree=3)

class CrossValidationModel(PowerLimitModel):
    """
    Implementation of a Cross Validation model. This creates an ensemble of
    sub-models, gives each one a fraction of the data, and deduces confidence
    values from the distribution of fit values."""
    def __init__(self, base_gen, frac=0.8, size=100):
        """
        Simple constructor; base_gen is a callable that returns a
        PowerLimitModel instance (or a derived class), frac is the proportion
        of training data to give to each instance during training (default 0.8),
        and size is the number of sub-models to instantiate for the ensemble
        (default 100)."""

        self._model = [base_gen() for _ in range(size)]
        self._frac = frac

    def train(self, df, key):
        for mm in self._model:
            mm.train(df.sample(frac=self._frac), key)

    def evaluate(self, PL, full=False):
        """
        Evaluate a trained model on a single power limit (PL). Returns
        the model's prediction for this power limit. If the parameter
        full is True, also return the standard deviation and residuals
        (default False)."""

        vals = [mm.evaluate(PL) for mm in self._model]
        mean = sum(vals)/len(vals)
        dev = (sum([val**2 for val in vals])/len(vals) - mean**2)**0.5
        res = (sum([mm.residual() for mm in self._model])/len(self._model))
        if full:
            return {'mean': mean, 'dev': [dev, res]}
        return mean

def extract_columns(df, region_filter = None):
    """
    Extract the columns of interest from the full report collection
    dataframe. This returns a dataframe indexed by the power limit
    and columns 'runtime' and 'energy'. region_filter (if provided)
    is a container that specifies which regions to include (by default,
    include all of them)."""
    df_filtered = df
    if region_filter and region_filter != 'Epoch':
        df_filtered = df[df['region'].isin(region_filter.split(','))]

    # these are the only columns we need
    df_cols = df_filtered[['POWER_PACKAGE_LIMIT_TOTAL',
                           'host', 'Profile',
                           'runtime (s)',
                           'package-energy (J)']]

    df_cols = df_cols.rename({'POWER_PACKAGE_LIMIT_TOTAL': 'power_limit',
                              'runtime (s)': 'runtime',
                              'package-energy (J)': 'energy'
                             }, axis = 1)

    df_cols = df_cols.groupby(['power_limit', 'host', 'Profile']).sum().reset_index()
    df_cols = df_cols[['power_limit', 'runtime', 'energy']].set_index('power_limit')

    return df_cols

def dump_stats_summary(df, fname):
    """
    Write mean runtime and energy and the standard deviation of
    runtime and energy for each power limit in CSV format to the
    file fname."""
    means = df.groupby(level=0).mean().rename(lambda x: x+"_mean", axis=1)
    stds = df.groupby(level=0).std().rename(lambda x: x+"_stddev", axis=1)
    pandas.concat([means, stds], axis=1).to_csv(fname)

def policy_min_energy(plrange, enmodel, rtmodel = None, pltdp = None,
                      max_degradation = None):
    """
    Find the power limit over the range plrange (list-like) that
    has the minimum predicted energy usage (according to the energy model
    enmodel, an instance of PowerLimitModel), subject to the constraint
    that its runtime does not exceed the runtime at power limit pltdp
    by more than a factor of (1 + max_degradation), if max_degradation is
    specified. Returns a dictionary with keys power, runtime, and energy,
    and values the optimal power limit and the predicted runtime and energy
    at that limit, respectively."""

    if pltdp is None:
        pltdp = max(plrange)

    en_predictions = enmodel.batch_evaluate(plrange)
    if max_degradation is None:
        # we don't need the runtime model in this case
        best_energy, best_pl = min(zip(en_predictions, plrange))
        if rtmodel:
            best_runtime = rtmodel.evaluate(best_pl)
        else:
            best_runtime = None
        return {'power': best_pl,
                'runtime': best_runtime,
                'energy': best_energy}
    else:
        rt_predictions = rtmodel.batch_evaluate(plrange)
        rt_at_tdp = rtmodel.evaluate(pltdp)
        constrained_values = [(energy, runtime, pl)
                              for pl, runtime, energy
                              in zip(plrange, rt_predictions, en_predictions)
                              if runtime <= (1 + max_degradation) * rt_at_tdp]
        best_energy, best_runtime, best_pl = min(constrained_values)
        return {'power': best_pl,
                'runtime': best_runtime,
                'energy': best_energy}

def rootssq(ll):
    """Return the root of the sum of the squares of the list of values. If a
    non-iterable is passed in, just return it."""
    try:
        return sum([ss**2  for ss in ll])**0.5
    except TypeError:
        pass  # ll is not iterable
    return ll

def normal_comparison(dist1, dist2):
    """Returns the probability that a sample drawn from the normal distribution
    described by dist1 (a dictionary with keys 'mean' and 'dev' and values the
    mean and standard deviation of the distribution, respectively; the standard
    deviation may be a simple float or may be a list of error magnitudes in the
    same dimensions as the mean which will be combined by rootssq) will be
    larger than a sample drawn from a normal distribution drawn from the normal
    distribution described by dist2."""
    mu1, sigma1 = dist1['mean'], dist1['dev']
    mu2, sigma2 = dist2['mean'], dist2['dev']
    sigma1 = rootssq(sigma1)
    sigma2 = rootssq(sigma2)
    return 0.5 * (1 + math.erf((mu1 - mu2) / (2 * (sigma1 ** 2 + sigma2 ** 2)**0.5)))

def policy_confident_energy(plrange, enmodel, rtmodel = None, pltdp = None,
                            max_degradation = None, confidence = 0.9):
    """
    Find the power limit over the range plrange (list-like) that
    has the minimum predicted energy usage (according to the energy model
    enmodel, an instance of PowerLimitModel), subject to the constraint
    that its runtime does not exceed the runtime at power limit pltdp
    by more than a factor of (1 + max_degradation), if max_degradation is
    specified, with at least a confidence (default 0.9) probability. Returns a
    dictionary with keys power, runtime, and energy, and values the optimal
    power limit, the predicted runtime and energy at that limit,
    respectively, and the standard deviations of these values."""
    if pltdp is None:
        pltdp = max(plrange)

    en_predictions = [enmodel.evaluate(pl, True) for pl in plrange]
    if max_degradation is None:
        # we don't need the runtime model in this case
        best_energy, best_pl = min(zip(en_predictions, plrange))
        if rtmodel:
            best_runtime = rtmodel.evaluate(best_pl, True)
        else:
            best_runtime = None
        return {'power': best_pl,
                'runtime': best_runtime['mean'],
                'runtimedev': rootssq(best_runtime['dev']),
                'energy': best_energy['mean'],
                'energydev': rootssq(best_energy['dev'])}
    else:
        rt_predictions = [rtmodel.evaluate(pl, True) for pl in plrange]
        rt_tdp_dist = rtmodel.evaluate(pltdp, True)
        rt_at_tdp, [dev_tdp, res_tdp] = rt_tdp_dist['mean'], rt_tdp_dist['dev']
        degraded_rt = {'mean': rt_at_tdp * (1 + max_degradation),\
                'dev': [dev_tdp * (1 + max_degradation), res_tdp * (1 + max_degradation)]}
        constrained_values = [(energy['mean'], energy, runtime, pl)
                              for pl, runtime, energy
                              in zip(plrange, rt_predictions, en_predictions)
                              if normal_comparison(degraded_rt, runtime) > confidence]
        if len(constrained_values) == 0:
            # this means that we don't ever have enough confidence to guarantee this
            return None
        _, best_energy, best_runtime, best_pl = min(constrained_values)
        return {'power': best_pl,
                'runtime': best_runtime['mean'],
                'runtimedev': rootssq(best_runtime['dev']),
                'energy': best_energy['mean'],
                'energydev': rootssq(best_energy['dev'])}

def main(full_df, region_filter, dump_prefix, min_pl, max_pl, tdp, max_degradation, cross_validation=True):
    """
    The main function. full_df is a report collection dataframe, region_filter
    is a list of regions to include, dump_prefix a filename prefix for
    debugging output (if specified), min_pl, max_pl, tdp are the minimum,
    maximum, and reference power limits, respectively, max_degradation
    specifies what maximum runtime degradation is accepted. If max_degradation
    is None, then just find the minimum energy configuration, ignoring
    runtime."""
    df = extract_columns(full_df, region_filter)
    if dump_prefix:
        df.to_csv("{}.dat".format(dump_prefix))
        dump_stats_summary(df, "{}.stats".format(dump_prefix))

    if cross_validation:
        runtime_model = CrossValidationModel(CubicPowerModel)
        energy_model = CrossValidationModel(CubicPowerModel)
    else:
        runtime_model = CubicPowerModel()
        energy_model = CubicPowerModel()

    runtime_model.train(df, key='runtime')
    energy_model.train(df, key='energy')

    plrange = [min_pl + i for i in range(int(max_pl - min_pl))] + [max_pl]
    if cross_validation:
        best_policy = policy_confident_energy(plrange, energy_model, runtime_model, tdp,
                                              max_degradation = max_degradation)
        tdprt, tdpen = runtime_model.evaluate(tdp, True), energy_model.evaluate(tdp, True)
        tdp_values = {'power': tdp,
                      'runtime': tdprt['mean'],
                      'runtimedev': rootssq(tdprt['dev']),
                      'energy': tdpen['mean'],
                      'energydev': rootssq(tdpen['dev'])}
    else:
        best_policy = policy_min_energy(plrange, energy_model, runtime_model, tdp,
                                        max_degradation = max_degradation)
        tdprt, tdpen = runtime_model.evaluate(tdp), energy_model.evaluate(tdp)
        tdp_values = {'power': tdp, 'runtime': tdprt, 'energy': tdpen}

    return {'tdp': tdp_values, 'best': best_policy}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    common_args.add_min_power(parser)
    common_args.add_max_power(parser)
    parser.add_argument('--path', required=True,
                        help='path containing reports and machine.json')
    parser.add_argument('--region-filter', default=None, dest='region_filter',
                        help='comma-separated list of regions to include, '
                             'default to include all regions, use \'Epoch\' '
                             'here to analyze just the Epoch data')
    parser.add_argument('--dump-prefix', dest='dump_prefix',
                        help='prefix to dump statistics to, empty to not dump '
                             'stats')
    parser.add_argument('--tdp', default=None, type=float,
                        help='tdp (sticker power limit), default behavior is '
                             'to use max power')
    parser.add_argument('--max-degradation',
                        default=0.1, type=float, dest='max_degradation',
                        help='maximum allowed runtime degradation, default is '
                             '0.1 (i.e., 10%%)')
    parser.add_argument('--min_energy', action='store_true',
                        help='ignore max degradation, just give the minimum '
                             'energy possible')
    parser.add_argument('--disable-confidence', dest='confidence',
                        action='store_false',
                        help='Ignore uncertainty when giving a recommendation.')

    args = parser.parse_args()

    try:
        if args.region_filter == 'Epoch':
            df = geopmpy.io.RawReportCollection('*report', dir_name=args.path).get_epoch_df()
        else:
            df = geopmpy.io.RawReportCollection('*report', dir_name=args.path).get_df()
    except RuntimeError:
        sys.stderr.write('<geopm> Error: No report data found in ' + path + \
                         '; run a power sweep before using this analysis.\n')
        sys.exit(1)

    min_pl, tdp, max_pl = None, None, None
    if None in [args.min_power, args.tdp, args.max_power] or \
            min(args.min_power, args.tdp, args.max_power) < 0:
        try:
            machine_info = machine.get_machine(args.path)
            min_pl = machine_info.power_package_min()
            tdp = machine_info.power_package_tdp()
            max_pl = machine_info.power_package_tdp()
        except RuntimeError:
            sys.stderr.write('Warning: couldn\'t open machine.json. Falling '
                             'back to default values.\n')
            min_pl = 150
            tdp = 280
            max_pl = 280

    # user-provided arguments override
    if args.min_power and args.min_power > 0:
        min_pl = args.min_power
    if args.tdp and args.tdp > 0:
        tdp = args.tdp
    if args.max_power and args.max_power > 0:
        max_pl = args.max_power

    output = main(df, args.region_filter, args.dump_prefix, min_pl, max_pl, tdp,
                  None if args.min_energy else args.max_degradation,
                  args.confidence)

    if args.confidence:
        sys.stdout.write('AT TDP = {power:.0f}W, '
                         'RUNTIME = {runtime:.0f} s +/- {runtimedev:.0f}, '
                         'ENERGY = {energy:.0f} J +/- {energydev:.0f}\n'.format(**output['tdp']))
        if output['best']:
            sys.stdout.write('AT PL  = {power:.0f}W, '
                             'RUNTIME = {runtime:.0f} s +/- {runtimedev:.0f}, '
                             'ENERGY = {energy:.0f} J +/- {energydev:.0f}\n'.format(**output['best']))
        else:
            sys.stdout.write("NO SUITABLE POLICY WAS FOUND.\n")
    else:
        sys.stdout.write('AT TDP = {power:.0f}W, '
                         'RUNTIME = {runtime:.0f} s, '
                         'ENERGY = {energy:.0f} J\n'.format(**output['tdp']))
        sys.stdout.write('AT PL  = {power:.0f}W, '
                         'RUNTIME = {runtime:.0f} s, '
                         'ENERGY = {energy:.0f} J\n'.format(**output['best']))

    relative_delta = lambda new, old: 100 * (new - old) / old

    if output['best']:
        sys.stdout.write('DELTA          RUNTIME = {:.1f} %,  ENERGY = {:.1f} %\n'
                         .format(relative_delta(output['best']['runtime'], output['tdp']['runtime']),
                                 relative_delta(output['best']['energy'], output['tdp']['energy'])))

max_stars_repo_name: avilcheslopez/geopm

max_issues_repo_name: avilcheslopez/geopm

max_forks_repo_name: avilcheslopez/geopm

reason: import numpy,from numpy


text:
'''
Wenchang Yang (wenchang@princeton.edu)
'''
import xarray as xr
from scipy.signal import freqz
from ..butter import _lowpass_ba, _highpass_ba, _bandpass_ba

def lowpass_response(cutoff=0.25, order=2, fs=1.0):
    '''calculate the lowpass response given the lowcut and order of the Butterworth
    filter.'''
    b, a = _lowpass_ba(cutoff=cutoff, order=order, fs=fs)
    f, h = freqz(b=b, a=a, whole=False, fs=fs)
    # the actual order is 2x the given order since a forward-backward filter is appled
    da = xr.DataArray( ( h*h.conj() ).astype('float'),
        dims='freq', coords=[f,],
        name='response')
    da.attrs['info'] = f'Butterworth {order}-order forward-backward lowpass response'
    da['freq'].attrs['units'] = 'sample freq'

    return da

def highpass_response(cutoff=0.25, order=2, fs=1.0):
    '''calculate the highpass response given the lowcut and order of the Butterworth
    filter.'''
    b, a = _highpass_ba(cutoff=cutoff, order=order, fs=fs)
    f, h = freqz(b=b, a=a, whole=False, fs=fs)
    # the actual order is 2x the given order since a forward-backward filter is appled
    da = xr.DataArray( ( h*h.conj() ).astype('float'),
        dims='freq', coords=[f,],
        name='response')
    da.attrs['info'] = f'Butterworth {order}-order forward-backward highpass response'
    da['freq'].attrs['units'] = 'sample freq'

    return da

def bandpass_response(cutoff=(0.125, 0.375), order=2, fs=1.0):
    '''calculate the bandpass response given the cutoff and order of the Butterworth
    filter.'''
    b, a = _bandpass_ba(cutoff=cutoff, order=order, fs=fs)
    f, h = freqz(b=b, a=a, whole=False, fs=fs)
    # the actual order is 2x the given order since a forward-backward filter is appled
    da = xr.DataArray( ( h*h.conj() ).astype('float'),
        dims='freq', coords=[f,],
        name='response')
    da.attrs['info'] = f'Butterworth {order}-order forward-backward bandpass response'
    da['freq'].attrs['units'] = 'sample freq'

    return da

reason: from scipy