text:
import numpy as np
import theano as th
from theano import tensor as T
from theano.tensor import switch, nnet, tanh

import netutils as nu

try:
    import matplotlib.pyplot as plt
except:
    plt = None

__doc__ = \
    """File with activation functions and other general tools required for building a network."""

# Make a callable subclass of dictionary
cdict = type('cdict', (dict,), {'__call__': lambda self, inp: self['function'](inp)})

# Activation Functions
# Rectified Linear
def relu(alpha=0):  # Gradient Tested and Working
    return lambda x: switch(x > 0, x, alpha * x)
    # switch(x < 0, 0, x)

# Truncated Rectified Linear
def trelu(threshold=1, alpha=0):
    return lambda x: switch(x > threshold, x, alpha * x)

# Sigmoid
def sigmoid(gain=1, spread=1, mode='hard'):  # Possible modes: 'fast', 'full', 'hard'.
    if mode == 'fast':
        return lambda x: gain * nnet.ultra_fast_sigmoid(spread * x)
    elif mode == 'full':
        return lambda x: gain * nnet.sigmoid(spread * x)
    elif mode == 'hard':
        return lambda x: gain * nnet.hard_sigmoid(spread * x)
    else:
        return lambda x: gain * nnet.ultra_fast_sigmoid(spread * x)

# Linear (duh)
def linear(gain=1.):
    return lambda x: gain * x

# Truncated Linear
def trlinear(threshold=1, gain=1):
    return lambda x: switch(T.abs_(x) > threshold, gain*x, 0)

# Tanh of ReLU: Saturates at gain for large positive values, gives zero otherwise.
def relutanh(gain=10, spread=0.1):
    return lambda x: gain * tanh(spread * relu()(x))

# Symmetric Linear Unit (SyLU)
def sylu(gain=10, spread=0.1):
    return lambda x: switch(T.ge(x, (1 / spread)), gain, 0) + \
                     switch(T.and_(T.gt((1 / spread), x), T.gt(x, -(1 / spread))), gain * spread * x, 0) + \
                     switch(T.le(x, -(1 / spread)), -gain, 0)

# Exponential Linear Unit
def elu(alpha=1):
    return lambda x: switch(x > 0., x, alpha * (T.exp(x) - 1))

# Make activation functions learnable
def learnable(func, fmapsin=None, a=None, b=None, c=None, d=None):
    """
    Make the activation function `func` learnable, such that func --> c * f(a*x + b) + d

    :type func: callable
    :param func: Activation function (to make learnable)

    :type fmapsin: int
    :param fmapsin: Number of input feature maps. Required if the function is to have independent parameters for every
                   feature map.

    :type a: callable
    :param a: a of func --> c * f(a*x + b) + d.

    :type b: callable
    :param b: b of func --> c * f(a*x + b) + d

    :type c: callable
    :param c: c of func --> c * f(a*x + b) + d

    :type d: callable
    :param d: d of func --> c * f(a*x + b) + d

    :return: c * f(a*x + b) + d with a, b, c, d learnable (if not None)
    """

    # Fmapsin is only required if the function is to have learnable parameters for every feature map. Otherwise, the
    # parameters are broadcasted over the entire stack of feature maps.
    fmapsin = (fmapsin, ) if isinstance(fmapsin, int) else ()

    # Instantiate a list of trainable parameters
    trainables = []

    if a is not None:
        a = nu.getshared(shape=fmapsin, value=a, name='func-a')
        trainables.append(a)
    else:
        a = nu.getshared(shape=fmapsin, value=1.)

    if b is not None:
        b = nu.getshared(shape=fmapsin, value=b, name='func-b')
        trainables.append(b)
    else:
        b = nu.getshared(shape=fmapsin, value=0.)

    if c is not None:
        c = nu.getshared(shape=fmapsin, value=c, name='func-c')
        trainables.append(c)
    else:
        c = nu.getshared(shape=fmapsin, value=1.)

    if d is not None:
        d = nu.getshared(shape=fmapsin, value=d, name='func-d')
        trainables.append(d)
    else:
        d = nu.getshared(shape=fmapsin, value=0.)

    # Define function
    def _func(inp):
        if inp.ndim == 4:
            # Broadcast variables accordingly
            bcst = lambda *variables: [var.dimshuffle('x', 0, 'x', 'x') if var.ndim == 1 else var for var in variables]

        elif inp.ndim == 5:
            # Broadcast variables accordingly
            bcst = lambda *variables: [var.dimshuffle('x', 'x', 0, 'x', 'x') if var.ndim == 1 else var
                                       for var in variables]

        else:
            raise NotImplementedError("Invalid input dimensionality.")

        # Get broadcasted variables
        ba, bb, bc, bd = bcst(a, b, c, d)
        # Make function
        return bc * func(ba * inp + bb) + bd

    # Return a callable dictionary
    return cdict({"function": _func, "trainables": trainables})

# Initializers
# Identity kernel (can be used for Net2DeeperNet)
def idkernel(shape, shared=True):
    # Init a new zero variable of shape shape
    idk = np.zeros(shape=shape, dtype=th.config.floatX)
    # Check if shape corresponds to a 2D or 3D kernel
    if len(shape) == 4:
        # Check if the kernel size is odd or even, and generate appropriate slice objects
        slicelist = [slice(kerdim/2, kerdim/2 + 1) if kerdim % 2 == 1 else slice(kerdim/2 - 1, kerdim/2)
                     for kerdim in shape[2:]]
        # Make identity kernel
        idk[..., slicelist[0], slicelist[1]] = 1.
    elif len(shape) == 5:
        slicelist = [slice(kerdim/2, kerdim/2 + 1) if kerdim % 2 == 1 else slice(kerdim/2 - 1, kerdim/2)
                     for kerdim in shape[1] + shape[3:]]
        # Make identity kernel
        idk[:, slicelist[0], :, slicelist[1], slicelist[2]] = 1

    # Convert to shared variable if requested
    if shared:
        return th.shared(value=idk)
    else:
        return idk

# Initialization schemes
# Glorot Initialization (Normal)
def xavier(shape, dtype=th.config.floatX):
    if len(shape) == 4:
        # 2D convolutions
        # Determine the number of output maps
        fmapsout = shape[0]
        fmapsin = shape[1]
        fov = np.prod(shape[2:])
    elif len(shape) == 5:
        # 3D convolutions
        fmapsout = shape[0]
        fmapsin = shape[2]
        fov = np.prod(shape[3:]) * shape[1]
    else:
        raise NotImplementedError

    # Compute variance for Xavier init
    var = 2./((fmapsin + fmapsout) * fov)
    # Build kernel
    ker = np.random.normal(loc=0., scale=np.sqrt(var), size=tuple(shape)).astype(dtype)
    return ker

# He Initialization
def he(shape, gain=1., dtype=th.config.floatX):
    if len(shape) == 4:
        # 2D convolutions
        fmapsin = shape[1]
        fov = np.prod(shape[2:])
    elif len(shape) == 5:
        # 3D convolutions
        fmapsin = shape[2]
        fov = np.prod(shape[3:]) * shape[1]
    else:
        raise NotImplementedError

    # Parse gain
    if isinstance(gain, str):
        if gain.lower() == 'relu':
            gain = 2.
        elif gain.lower() in ['sigmoid', 'linear', 'tanh']:
            gain = 1.

    # Compute variance for He init
    var = gain/(fmapsin * fov)
    # Build kernel
    ker = np.random.normal(loc=0., scale=np.sqrt(var), size=tuple(shape)).astype(dtype)
    return ker

# Training Monitors
# Batch Number
# Loss
# Cost
# Training Error
# Validation Error
# Gradient Norm Monitor
# Monitor to check if the model is exploring previously unexplored parameter space. Returns n, where n is the number of
# dimensions in which the model has explored new param range.
# Monitor for update norm

if plt is not None:
    # Plotter Function
    pass

reason: import numpy,import theano,from theano


text:
import numpy as np
from sklearn.datasets import load_breast_cancer

import deeprob.spn.structure as spn
from deeprob.spn.models.sklearn import SPNEstimator, SPNClassifier

if __name__ == '__main__':
    # Load the dataset and set the features distributions
    data, target = load_breast_cancer(return_X_y=True)
    _, n_features = data.shape
    distributions = [spn.Gaussian] * n_features

    # =================================================================================================================

    # Instantiate and fit a SPN density estimator
    clf = SPNEstimator(
        distributions,
        learn_leaf='mle',     # Learn leaf distributions by MLE
        split_rows='kmeans',  # Use K-Means for splitting rows
        split_cols='rdc',     # Use RDC for splitting columns
        min_rows_slice=64,    # The minimum number of rows required to split furthermore
        random_state=42       # The random state, used for reproducibility
    )
    clf.fit(data)

    # Compute the average log-likelihood and two standard deviations
    score = clf.score(data)
    print('Train - Mean LL: {} - Stddev LL: {}'.format(score['mean_ll'], score['stddev_ll']))

    # Sample some data and compute the average log-likelihood and two standard deviations
    samples = clf.sample(n=100)
    score = clf.score(samples)
    print('Sampling - Mean LL: {} - Stddev LL: {}'.format(score['mean_ll'], score['stddev_ll']))

    # =================================================================================================================

    # Instantiate and fit a SPN classifier
    clf = SPNClassifier(
        distributions,
        learn_leaf='mle',     # Learn leaf distributions by MLE
        split_rows='kmeans',  # Use K-Means for splitting rows
        split_cols='rdc',     # Use RDC for splitting columns
        min_rows_slice=64,    # The minimum number of rows required to split furthermore
        random_state=42       # The random state, used for reproducibility
    )
    clf.fit(data, target)

    # Compute the accuracy score
    print('Train - Accuracy: {}'.format(clf.score(data, target)))

    # Sample some data from the conditional distribution and compute the accuracy score
    classes = np.array([1, 0, 0, 1, 0, 1, 1, 0, 1, 1])
    samples = clf.sample(y=classes)
    print('Sampling - Accuracy: {}'.format(clf.score(samples[:, :-1], classes)))

ext: py

reason: import numpy


text:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#
# sugarcane.py
# projects
#
# Created by Haibao Tang on 12/02/19
# Copyright © 2019 Haibao Tang. All rights reserved.
#
import logging
import os.path as op
import sys
from random import random, sample
from itertools import groupby
from collections import Counter
import numpy as np
import matplotlib.pyplot as plt

from jcvi.apps.base import OptionParser, ActionDispatcher, mkdir
from jcvi.graphics.base import adjust_spines, markup, normalize_axes, savefig
from jcvi.utils.validator import validate_in_choices

SoColor = "#7436a4"  # Purple
SsColor = "#5a8340"  # Green

# Computed using prepare(), corrected with real sizes
ChrSizes = {
    "SO-chr01": 148750011,
    "SO-chr02": 119865146,
    "SO-chr03": 103845728,
    "SO-chr04": 104559946,
    "SO-chr05": 93134056,
    "SO-chr06": 74422021,
    "SO-chr07": 81308893,
    "SO-chr08": 71010813,
    "SO-chr09": 86380266,
    "SO-chr10": 73923121,
    "SS-chr01": 114519418,
    "SS-chr02": 119157314,
    "SS-chr03": 85009228,
    "SS-chr04": 79762909,
    "SS-chr05": 90584537,
    "SS-chr06": 95848354,
    "SS-chr07": 83589369,
    "SS-chr08": 64028871,
}

# Simulate genome composition
class Genome:
    def __init__(self, name, prefix, ploidy, haploid_chromosome_count):
        """
        Simulate a genome with given ploidy and haploid_chromosome_count. Example:

        >>> print(Genome("t", "pf", 2, 3))
        test: pf-chr01_a,pf-chr01_b,pf-chr02_a,pf-chr02_b,pf-chr03_a,pf-chr03_b
        """
        self.name = name
        chromosomes = []
        for i in range(haploid_chromosome_count):
            chromosomes += [
                f"{prefix}-chr{i + 1:02d}_{chr(ord('a') + j)}" for j in range(ploidy)
            ]
        self.chromosomes = chromosomes

    def __len__(self):
        return len(self.chromosomes)

    @classmethod
    def make(cls, name, chromosomes):
        genome = Genome(name, "", 0, 0)
        genome.chromosomes = chromosomes
        return genome

    @property
    def gamete(self):
        """Randomly generate a gamete from current genome that"""
        self.chromosomes.sort()
        gamete_chromosomes = []

        # Check for any chromosome that have 2 identical copies, if so, we will assume disomic
        # inheritance for that chromosome and always keep one and only copy
        duplicate_chromosomes = []
        singleton_chromosomes = []
        for chromosome, chromosomes in groupby(self.chromosomes):
            chromosomes = list(chromosomes)
            ncopies = len(chromosomes)
            duplicate_chromosomes += [chromosome] * (ncopies // 2)
            if ncopies % 2 == 1:
                singleton_chromosomes.append(chromosome)

        # Get one copy of each duplicate chromosome first
        gamete_chromosomes += duplicate_chromosomes

        def prefix(x):
            return x.split("_", 1)[0]

        # Randomly assign the rest, singleton chromosomes
        for group, chromosomes in groupby(singleton_chromosomes, key=prefix):
            chromosomes = list(chromosomes)
            halfn = len(chromosomes) // 2
            # Odd number, e.g. 5, equal chance to be 2 or 3
            if len(chromosomes) % 2 != 0 and random() < 0.5:
                halfn += 1
            gamete_chromosomes += sorted(sample(chromosomes, halfn))
        return Genome.make(self.name + " gamete", gamete_chromosomes)

    def mate_nplusn(self, name, other_genome, verbose=True):
        if verbose:
            print(
                f"Crossing '{self.name}' x '{other_genome.name}' (n+n)", file=sys.stderr
            )
        f1_chromosomes = sorted(
            self.gamete.chromosomes + other_genome.gamete.chromosomes
        )
        return Genome.make(name, f1_chromosomes)

    def mate_nx2plusn(self, name, other_genome, verbose=True):
        if verbose:
            print(
                f"Crossing '{self.name}' x '{other_genome.name}' (2xn+n)",
                file=sys.stderr,
            )
        f1_chromosomes = sorted(
            2 * self.gamete.chromosomes + other_genome.gamete.chromosomes
        )
        return Genome.make(name, f1_chromosomes)

    def mate_2nplusn(self, name, other_genome, verbose=True):
        if verbose:
            print(
                f"Crossing '{self.name}' x '{other_genome.name}' (2n+n)",
                file=sys.stderr,
            )
        f1_chromosomes = sorted(self.chromosomes + other_genome.gamete.chromosomes)
        return Genome.make(name, f1_chromosomes)

    def __str__(self):
        return self.name + ": " + ",".join(self.chromosomes)

    @property
    def summary(self):
        def prefix(x, sep="-"):
            return x.split(sep, 1)[0]

        def size(chromosomes):
            return sum(ChrSizes[prefix(x, sep="_")] for x in chromosomes)

        # Chromosome count
        total_count = 0
        total_unique = 0
        total_size = 0
        total_so_size = 0
        ans = []
        for group, chromosomes in groupby(self.chromosomes, prefix):
            chromosomes = list(chromosomes)
            uniq_chromosomes = set(chromosomes)
            group_count = len(chromosomes)
            group_unique = len(uniq_chromosomes)
            group_so_size = size({x for x in uniq_chromosomes if x[:2] == "SO"})
            group_size = size(uniq_chromosomes)
            total_count += group_count
            total_unique += group_unique
            total_so_size += group_so_size
            total_size += group_size
            ans.append((group, group_count, group_unique, group_so_size, group_size))
        ans.append(("Total", total_count, total_unique, total_so_size, total_size))
        return ans

    def print_summary(self):
        print("[SUMMARY]")
        for group, group_count, group_unique in self.summary:
            print(f"{group}: count={group_count}, unique={group_unique}")

class GenomeSummary:
    def __init__(self, SO_data, SS_data, percent_SO_data):
        self.SO_data = SO_data
        self.SS_data = SS_data
        self.percent_SO_data = percent_SO_data
        self.percent_SS_data = [100 - x for x in percent_SO_data]

    def _summary(self, a, tag, precision=0):
        mean, min, max = (
            round(np.mean(a), precision),
            round(np.min(a), precision),
            round(np.max(a), precision),
        )
        s = f"*{tag}* chr: {mean:.0f}"
        if min == mean and max == mean:
            return s
        return s + f" ({min:.0f}-{max:.0f})"

    def _percent_summary(self, a, tag, precision=1):
        mean, min, max = (
            round(np.mean(a), precision),
            round(np.min(a), precision),
            round(np.max(a), precision),
        )
        s = f"*{tag}*%: {mean:.1f}%"
        print(s)
        if min == mean and max == mean:
            return s
        return s + f" ({min:.1f}-{max:.1f}%)"

    @property
    def percent_SO_summary(self):
        return self._percent_summary(self.percent_SO_data, "So")

    @property
    def percent_SS_summary(self):
        return self._percent_summary(self.percent_SS_data, "Ss")

    @property
    def SO_summary(self):
        return self._summary(self.SO_data, "So")

    @property
    def SS_summary(self):
        return self._summary(self.SS_data, "Ss")

def simulate_F1(SO, SS, mode="nx2+n", verbose=False):
    SO_SS_F1 = (
        SO.mate_nx2plusn("SOxSS F1", SS, verbose=verbose)
        if mode == "nx2+n"
        else SO.mate_2nplusn("SOxSS F1", SS, verbose=verbose)
    )
    if verbose:
        SO_SS_F1.print_summary()
    return SO_SS_F1

def simulate_F2(SO, SS, mode="nx2+n", verbose=False):
    SO_SS_F1 = simulate_F1(SO, SS, mode=mode, verbose=verbose)
    SO_SS_F2_nplusn = SO_SS_F1.mate_nplusn("SOxSS F2", SO_SS_F1, verbose=verbose)
    if verbose:
        SO_SS_F2_nplusn.print_summary()
    return SO_SS_F2_nplusn

def simulate_F1intercross(SO, SS, mode="nx2+n", verbose=False):
    SO_SS_F1_1 = simulate_F1(SO, SS, mode=mode, verbose=verbose)
    SO_SS_F1_2 = simulate_F1(SO, SS, mode=mode, verbose=verbose)
    SO_SS_F1intercross_nplusn = SO_SS_F1_1.mate_nplusn(
        "SOxSS F1 intercross", SO_SS_F1_2, verbose=verbose
    )
    return SO_SS_F1intercross_nplusn

def simulate_BCn(n, SO, SS, mode="nx2+n", verbose=False):
    SS_SO_F1 = simulate_F1(SO, SS, mode=mode, verbose=verbose)
    SS_SO_BC1, SS_SO_BC2_nplusn, SS_SO_BC3_nplusn, SS_SO_BC4_nplusn = (
        None,
        None,
        None,
        None,
    )
    # BC1
    if n >= 1:
        SS_SO_BC1 = (
            # Expecting one more round of female restitution in BC1
            SO.mate_nx2plusn("SSxSO BC1", SS_SO_F1, verbose=verbose)
            if mode == "nx2+n"
            else SO.mate_2nplusn("SSxSO BC1", SS_SO_F1, verbose=verbose)
        )
    # BC2
    if n >= 2:
        SS_SO_BC2_nplusn = SO.mate_nplusn("SSxSO BC2", SS_SO_BC1, verbose=verbose)
    # BC3
    if n >= 3:
        SS_SO_BC3_nplusn = SO.mate_nplusn(
            "SSxSO BC3", SS_SO_BC2_nplusn, verbose=verbose
        )
    # BC4
    if n >= 4:
        SS_SO_BC4_nplusn = SO.mate_nplusn(
            "SSxSO BC4", SS_SO_BC3_nplusn, verbose=verbose
        )
    return [
        None,
        SS_SO_BC1,
        SS_SO_BC2_nplusn,
        SS_SO_BC3_nplusn,
        SS_SO_BC4_nplusn,
    ][n]

def plot_summary(ax, samples: list[Genome]) -> GenomeSummary:
    """Plot the distribution of chromosome numbers given simulated samples.

    Args:
        ax (Axes): Matplotlib axes.
        samples (list[Genome]): Summarized genomes.

    Returns:
        GenomeSummary: Summary statistics of simulated genomes.
    """
    SO_data = []
    SS_data = []
    percent_SO_data = []
    for sample in samples:
        summary = sample.summary
        try:
            _, _, group_unique, _, _ = [x for x in summary if x[0] == "SO"][0]
        except:
            group_unique = 0
        SO_data.append(group_unique)
        try:
            _, _, group_unique, _, _ = [x for x in summary if x[0] == "SS"][0]
        except:
            group_unique = 0
        SS_data.append(group_unique)
        total_tag, _, _, total_so_size, total_size = summary[-1]
        assert total_tag == "Total"
        percent_SO = total_so_size * 100.0 / total_size
        percent_SO_data.append(percent_SO)
    # Avoid overlapping bars
    SS_counter, SO_counter = Counter(SS_data), Counter(SO_data)
    overlaps = SS_counter.keys() & SO_counter.keys()
    shift = 0.5  # used to offset bars a bit to avoid cluttering
    if overlaps:
        for overlap in overlaps:
            logging.debug(f"Modify bar offsets at {overlap} due to SS and SO overlaps")
            SS_counter[overlap - shift] = SS_counter[overlap]
            del SS_counter[overlap]
            SO_counter[overlap + shift] = SO_counter[overlap]
            del SO_counter[overlap]

    def modify_range_end(d: dict, value: int):
        if value not in d:
            return
        # Has data at the range end, but no adjacent data points (i.e. isolated bar)
        if value in d and (value - 1 in d or value + 1 in d):
            return
        logging.debug(f"Modify bar offsets at {value} due to end of range ends")
        d[value - shift if value else value + shift] = d[80]
        del d[value]

    modify_range_end(SS_counter, 0)
    modify_range_end(SS_counter, 80)
    modify_range_end(SO_counter, 0)
    modify_range_end(SO_counter, 80)

    x, y = zip(*sorted(SS_counter.items()))
    ax.bar(np.array(x), y, color=SsColor, ec=SsColor)
    x, y = zip(*sorted(SO_counter.items()))
    ax.bar(np.array(x), y, color=SoColor, ec=SoColor)
    ax.set_xlim(80, 0)
    ax.set_ylim(0, len(samples) / 2)
    ax.set_yticks([])
    summary = GenomeSummary(SO_data, SS_data, percent_SO_data)

    # Write the stats summary within the plot
    summary_style = dict(
        size=9,
        ha="center",
        va="center",
        transform=ax.transAxes,
    )
    ax.text(0.75, 0.85, markup(summary.SS_summary), color=SsColor, **summary_style)
    ax.text(
        0.75, 0.65, markup(summary.percent_SS_summary), color=SsColor, **summary_style
    )
    ax.text(0.25, 0.85, markup(summary.SO_summary), color=SoColor, **summary_style)
    ax.text(
        0.25, 0.65, markup(summary.percent_SO_summary), color=SoColor, **summary_style
    )

    return summary

def write_chromosomes(genomes: list[Genome], filename: str):
    """Write simulated chromosomes to file

    Args:
        genomes (list[Genome]): List of simulated genomes.
        filename (str): File path to write to.
    """
    print(f"Write chromosomes to `{filename}`", file=sys.stderr)
    with open(filename, "w") as fw:
        for genome in genomes:
            print(genome, file=fw)

def write_SO_percent(summary: GenomeSummary, filename: str):
    """Write SO % to file

    Args:
        summary (GenomeSummary): List of simulated genomes.
        filename (str): File path to write to.
    """
    print(f"Write SO percent to `{filename}`", file=sys.stderr)
    with open(filename, "w") as fw:
        print("\n".join(str(x) for x in sorted(summary.percent_SO_data)), file=fw)

def simulate(args):
    """
    %prog simulate [2n+n|nx2+n]

    Run simulation on female restitution. There are two modes:
    - 2n+n: merger between a somatic and a germline
    - nx2+n: merger between a doubled germline and a germline

    These two modes would impact the sequence diversity in the progeny
    genome in F1, F2, BCn ... the goal of this simulation, is thus to
    understand the mode and the spread of such diversity in the hybrid
    progenies.
    """
    import seaborn as sns

    sns.set_style("darkgrid")

    p = OptionParser(simulate.__doc__)
    p.add_option(
        "--verbose",
        default=False,
        action="store_true",
        help="Verbose logging during simulation",
    )
    p.add_option("-N", default=10000, type="int", help="Number of simulated samples")
    opts, args, iopts = p.set_image_options(args, figsize="6x6")
    if len(args) != 1:
        sys.exit(not p.print_help())

    (mode,) = args
    validate_in_choices(mode, ["2n+n", "nx2+n"], "Mode")
    logging.info(f"Transmission: {mode}")

    # Construct a composite figure with 6 tracks
    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])
    rows = 6
    ypad = 0.05
    yinterval = (1 - 2 * ypad) / (rows + 1)
    yy = 1 - ypad
    xpad = 0.2
    xwidth = 0.7

    # Axes are vertically stacked, and share x-axis
    axes = []
    yy_positions = []  # Save yy positions so we can show details to the right laterr
    for idx in range(rows):
        yy_positions.append(yy)
        yy -= yinterval
        ax = fig.add_axes([xpad, yy, xwidth, yinterval * 0.85])
        if idx != rows - 1:
            plt.setp(ax.get_xticklabels(), visible=False)
        axes.append(ax)
    ax1, ax2, ax3, ax4, ax5, ax6 = axes

    # Prepare the simulated data
    # Simulate two parents
    SS = Genome("SS", "SS", 10, 8)
    SO = Genome("SO", "SO", 8, 10)

    verbose = opts.verbose
    N = opts.N
    all_F1s = [simulate_F1(SO, SS, mode=mode, verbose=verbose) for _ in range(N)]
    all_F2s = [simulate_F2(SO, SS, mode=mode, verbose=verbose) for _ in range(N)]
    all_BC1s = [simulate_BCn(1, SO, SS, mode=mode, verbose=verbose) for _ in range(N)]
    all_BC2s = [simulate_BCn(2, SO, SS, mode=mode, verbose=verbose) for _ in range(N)]
    all_BC3s = [simulate_BCn(3, SO, SS, mode=mode, verbose=verbose) for _ in range(N)]
    all_BC4s = [simulate_BCn(4, SO, SS, mode=mode, verbose=verbose) for _ in range(N)]

    # Plotting
    all_F1s_summary = plot_summary(ax1, all_F1s)
    all_F2s_summary = plot_summary(ax2, all_F2s)
    plot_summary(ax3, all_BC1s)
    plot_summary(ax4, all_BC2s)
    plot_summary(ax5, all_BC3s)
    plot_summary(ax6, all_BC4s)

    # Show title to the left
    xx = xpad / 2
    for (title, subtitle), yy in zip(
        (
            (r"$\mathrm{F_1}$", None),
            (r"$\mathrm{F_2}$", None),
            (r"$\mathrm{BC_1}$", None),
            (r"$\mathrm{BC_2}$", None),
            (r"$\mathrm{BC_3}$", None),
            (r"$\mathrm{BC_4}$", None),
        ),
        yy_positions,
    ):
        if subtitle:
            yy -= 0.06
        else:
            yy -= 0.07
        root.text(
            xx,
            yy,
            title,
            color="darkslategray",
            ha="center",
            va="center",
            fontweight="semibold",
        )
        if subtitle:
            yy -= 0.02
            root.text(
                xx, yy, subtitle, color="lightslategray", ha="center", va="center"
            )

    axes[-1].set_xlabel("Number of unique chromosomes")
    adjust_spines(axes[-1], ["bottom"], outward=True)
    normalize_axes(root)

    # Title
    mode_title = r"$n_1\times2 + n_2$" if mode == "nx2+n" else r"$2n + n$"
    root.text(0.5, 0.95, f"Transmission: {mode_title}", ha="center")

    savefig(f"{mode}.pdf", dpi=120)

    outdir = f"simulations_{mode}"
    mkdir(outdir)
    # Write chromosomes to disk
    for genomes, filename in (
        (all_F1s, "all_F1s"),
        (all_F2s, "all_F2s"),
        (all_BC1s, "all_BC1s"),
        (all_BC2s, "all_BC2s"),
        (all_BC3s, "all_BC3s"),
        (all_BC4s, "all_BC4s"),
    ):
        write_chromosomes(genomes, op.join(outdir, filename))

    # Write the SO percent in simulated samples so that we can compute P-value
    for summary, SO_percent_filename in (
        (all_F1s_summary, "all_F1s_SO_percent"),
        (all_F2s_summary, "all_F2s_SO_percent"),
    ):
        write_SO_percent(summary, op.join(outdir, SO_percent_filename))

def _get_sizes(filename, prefix_length, tag, target_size=None):
    """Returns a dictionary of chromome lengths from a given file.

    Args:
        filename ([str]): Path to the input file. Input file is 2-column file
        with rows `seqid length`.
        prefix_length (int): Extract first N characters.
        tag (str): Prepend `tag-` to the seqid.
        target_size (int): Expected genome size. Defaults to None.
    """
    from collections import defaultdict

    sizes_list = defaultdict(list)
    with open(filename) as fp:
        for row in fp:
            if not row.startswith("Chr"):
                continue
            name, size = row.split()
            idx = int(name[3:prefix_length])
            size = int(size)
            name = f"{tag}-chr{idx:02d}"
            sizes_list[name].append(size)

    # Get the average length
    sizes = dict(
        (name, int(round(np.mean(size_list)))) for name, size_list in sizes_list.items()
    )
    print(sizes)
    if target_size is None:
        return sizes

    total_size = sum(sizes.values())
    correction_factor = target_size / total_size
    print(
        f"{tag} total:{total_size} target:{target_size} correction:{correction_factor:.2f}x"
    )
    return dict(
        (name, int(round(correction_factor * size))) for name, size in sizes.items()
    )

def prepare(args):
    """
    %prog SoChrLen.txt SsChrLen.txt

    Calculate lengths from real sugarcane data.
    """
    p = OptionParser(prepare.__doc__)
    opts, args = p.parse_args(args)
    if len(args) != 2:
        sys.exit(not p.print_help())

    solist, sslist = args
    # The haploid set of LA Purple is 957.2 Mb and haploid set of US56-14-4 is 732.5 Mb
    sizes = _get_sizes(solist, 5, "SO", target_size=int(957.2 * 1e6))
    sizes.update(_get_sizes(sslist, 4, "SS", target_size=int(732.5 * 1e6)))
    print(sizes)

def main():

    actions = (
        ("prepare", "Calculate lengths from real sugarcane data"),
        ("simulate", "Run simulation on female restitution"),
    )
    p = ActionDispatcher(actions)
    p.dispatch(globals())

if __name__ == "__main__":
    main()

max_stars_repo_name: inspirewind/jcvi

max_issues_repo_name: xuzhougeng/jcvi

max_forks_repo_name: xuzhougeng/jcvi


text:
import numpy as np
from copy import deepcopy
import scipy.signal

def calc_spikes(trace, threshold=1.0, mode='local_max'):
    trace = np.asarray(trace)

    if mode=='local_max':
        trace_above_th = find_local_max(trace, threshold)
    else:
        trace_above_th = trace > threshold
        trace_above_th = replace_1bySpike(trace_above_th)
    trace_above_th = np.asarray(trace_above_th, dtype=int)

    return trace_above_th

def find_local_max(V, threshold):
    time_diff = 0.1 # ms
    wsize = int(0.5 / time_diff)
    voltage = np.zeros_like(V)
    V = scipy.signal.savgol_filter(V, wsize, 3)

    spike_inds = np.where((V[1:-1] > threshold) & (np.diff(V[:-1]) >= 0) & (np.diff(V[1:]) <= 0))[0]
    voltage[spike_inds] = 1.0
    return voltage

def replace_1bySpike(binary_trace):
    spike_list = np.zeros_like(binary_trace)

    current = 0
    counter = 0
    for num in binary_trace:
        if num:
            current += 1
        else:
            if current > 0:
                spike_list[int(counter - current/2)] = 1
            current = 0
        counter += 1
    return spike_list

def subtract_PDfromPY(PY_spikes, PD_spikes, vicinity=10, sampling_frequency=10000):
    counter = 0
    PYminusPD = deepcopy(PY_spikes)
    for s_PY in PY_spikes:
        PDinVicitity = False
        if s_PY > 0:
            PDinVicitity = check_spike_vicinity(counter, PD_spikes, vicinity=vicinity, sampling_frequency=sampling_frequency)
        if PDinVicitity:
            PYminusPD[counter] = 0.0
        counter += 1
    return PYminusPD

def check_spike_vicinity(counter, PD_spikes, vicinity=10, sampling_frequency=10000):
    if np.any(PD_spikes[int(counter-vicinity*sampling_frequency/1000):int(counter+vicinity*sampling_frequency/1000)]): return True
    else: return False

reason: import numpy,import scipy


text:
"""
use VO tables of photometric filters. Put filters XML files in
'_dir'. All proper filters will be automatically loaded and made
available to all functions in the module.

ref: http://svo.cab.inta-csic.es/theory/fps3/index.php?mode=browse

http://svo2.cab.inta-csic.es/theory/fps3/fps.php?ID=Generic/Bessell.U

"""
import os, sys
import numpy as np
import urllib2

_dir = '../SPIPS/DATA/FILTERS' # will automatically loads all filters from there
if not os.path.exists(_dir):
    _dir = os.path.join(filter(lambda d: d.endswith('/SPIPS'), sys.path)[0], 'DATA/FILTERS/')

if not os.path.isdir(_dir):
    import sys
    rotastarDir =filter(lambda x: 'SPIPS' in x, sys.path)
    if len(rotastarDir)==1:
        _dir = os.path.join(rotastarDir[0], _dir)
        print '\033[43m', _dir, '\033[0m'

# ------------

def downloadFilters(directory=_dir):
    """
    Attempt to make an automatic dowload a set of commonly used filters from the Spanish
    Virtual Observatory
    """
    files = {}
    # build tree of files
    files['Generic']=[]
    for f in ['U','B','V','R','I']:
        files['Generic'].append('Bessell.'+f)
    for f in ['U','B','V','R','I','J','M']:
        files['Generic'].append('Johnson.'+f)
    for f in ['u','v','b','y']:
        files['Generic'].append('Stromgren.'+f)
    files['OSN']=[]
    for f in ['I', 'R']:
        files['OSN'].append('Circ.Cousins_'+f)
    files['IRAS']=[]
    for f in ['12','25','60','100']:
        files['IRAS'].append('IRAS.'+f+'mu')
    files['2MASS']=[]
    for f in ['J','H','Ks']:
        files['2MASS'].append('2MASS.'+f)
    files['SLOAN']=[]
    for f in ['u','g','r','i','z']:
        files['SLOAN'].append('SDSS.'+f)
    files['CTIO']=[]
    for f in ['J','H','K']:
        files['CTIO'].append('ANDICAM.'+f)
    # load the files and save them
    for k in files.keys():
        for f in files[k]:
            if not os.path.exists(os.path.join(directory, f+'.xml')):
                url = 'http://svo2.cab.inta-csic.es/theory/fps3/fps.php?ID='+k+'/'+f
                print url
                l = urllib2.urlopen(url)
                f = open(os.path.join(directory, f+'.xml'), 'w')
                f.write(l.read())
                f.close()
            else:
                pass
                #print k+'/'+f+'.xml already exists'
    return

def wavelRange(filtname):
    global _data
    return _data[filtname]['profile'][0].min(),\
           _data[filtname]['profile'][0].max()

def sliding_avg_interp1d(x,y,x0,dx0=None):
    """
    use to interpret y(x) for x0+-dx0/2.

    use if y is over-resolved compared to dx0 (same shape as x0)

    assumes x and x0 sorted!!!
    """
    if not isinstance(x, np.ndarray):
        x = np.array(x)
    if not isinstance(y, np.ndarray):
        y = np.array(y)
    if not isinstance(x0, np.ndarray):
        x0 = np.array(x0)
    if dx0==None:
        dx0 = np.gradient(x0)
    y0 = np.zeros(len(x0))

    for i in range(len(x0)):
        # -- slow :(
        #y0[i] = np.mean(y[(x>=x0[i]-dx0[i]/2)*(x<x0[i]+dx0[i]/2)])
        # -- faster
        y0[i] = np.mean(y[np.searchsorted(x, x0[i]-dx0[i]/2., side='left'):
                          np.searchsorted(x, x0[i]+dx0[i]/2., side='right')])

    if not all(np.isfinite(y0)):
        wf = np.where(np.isfinite(y0))
        wi = np.where(1-np.isfinite(y0))
        y0[wi] = np.interp(x0[wi], x0[wf], y0[wf])
    return y0

def sliding_avg_interp1d_WEAVE(x,y,x0,dx0=None):
    """
    use to interpret y(x) for x0+-dx0/2.

    use if y is over-resolved compared to dx0 (same shape as x0)

    assumes x and x0 sorted!!!
    """
    from scipy import weave
    from scipy.weave import converters

    if not isinstance(x, np.ndarray):
        x = np.array(x)
    if not isinstance(y, np.ndarray):
        y = np.array(y)
    if not isinstance(x0, np.ndarray):
        x0 = np.array(x0)
    if dx0==None:
        dx0 = np.gradient(x0)

    LX = len(x)
    LX0 = len(x0)
    y0 = np.zeros(len(x0))
    # using WEAVE
    code =\
    """
    double n;
    int j;
    int i;
    int j0;
    int inloop;
    j0=1;

    for (i=0; i<LX0; i++){
        n=0.0;
        j=j0;
        inloop=0;
        while ((x(j)<=x0(i)+dx0(i)/2) && (j<LX-1) ) {
            if ((x(j)>=x0(i)-dx0(i)/2.) && (x(j)<=x0(i)+dx0(i)/2.)) {
               n=n+(x(j+1)-x(j-1))/2.0;
               y0(i)=y0(i)+y(j)*(x(j+1)-x(j-1))/2.0;
               inloop=1;
            }
            if (inloop==0) {
                j0=j; // memorize where we started before
                }
            j++;
        }
        y0(i) /= n;
    }
    """
    err = weave.inline(code,
                        ['x', 'y', 'x0', 'dx0', 'LX', 'LX0','y0'],
                       type_converters=converters.blitz,
                       compiler = 'gcc')
    # try to catch 'nan' and 'inf'
    if not all(np.isfinite(y0)):
        wf = np.where(np.isfinite(y0))
        wi = np.where(1-np.isfinite(y0))
        y0[wi] = np.interp(x0[wi], x0[wf], y0[wf])
    return y0

def _atmoTrans(wl, expo=1.5):
    """
    wl in um
    return T**expo where T is the transmission of the atmosphere between
    0.3 and 5um in wavelength,
    for 1.0mm of water vapor if expo=1.0
    """
    global _atmoData
    try:
        n = len(_atmoData)
    except:
        _ad = '../SPIPS/DATA/ATMO/'
        if not os.path.exists(_ad):
            _ad = os.path.join(filter(lambda d: d.endswith('/SPIPS'), sys.path)[0], 'DATA/ATMO/')

        f = open(os.path.join(_ad, 'transmission_300_5000_pwv10.txt'))
        _atmoData = []
        for l in f.readlines():
            _atmoData.append([float(l.split()[0])/1000., float(l.split()[1])])
        f.close()
        _atmoData = np.array(_atmoData)
    #  expo to crudly decrease water vapor, 0 to eliminate it
    return np.interp(wl, _atmoData[:,0], _atmoData[:,1]**expo)

def _ccd(wl):
    # -- http://www.faculty.virginia.edu/rwoclass/astr511/im/ESO-ccdQEs.gif
    # -- "EEV Average"
    _wl = [.3, .325, .35, .375, .4, .45, .5, .55, .6, .65, .7, .75, .8, .85, .9, .95, 1., 1.1]
    _T =  [0., .45, .55, .7,  .8,  .9,  .92, .9, .9, .85, .8, .7,  .57, .45, .3, .15, .5, 0]
    return np.interp(wl, _wl, _T)

def Transmission(filtname, withAtmo='auto'):
    global _data
    if withAtmo == 'auto':
        # -- facilities in space or which already took into account atmo
        space = ['IRAS', 'Spitzer', 'Hipparcos', '2MASS', 'DENIS', 'TYCHO',
                'TESS', 'GAIA', 'MSX', 'AKARI', 'DIRBE', 'HST', 'JWST']
        withAtmo = not any([s.lower() in filtname.lower() for s in space])

    if withAtmo:
        return lambda x: np.interp(x,
                               _data[filtname]['profile'][0],
                               _data[filtname]['profile'][1],
                               left=0, right=0)*_atmoTrans(x)
    else:
        return lambda x: np.interp(x,
                               _data[filtname]['profile'][0],
                               _data[filtname]['profile'][1],
                               left=0, right=0)

def Tag(l, tag):
    return l.split(tag+'="')[1].split('"')[0]

def LoadFilters(directory=_dir):
    """
    quick and dirty reader for .xml files

    for more download VOTable files in directory '_dir' from
    http://svo.cab.inta-csic.es/theory/fps/index.php?mode=browse
    """
    res={}
    download=False
    try:
        filz = os.listdir(directory)
        filz = filter(lambda x: x.endswith('.xml'), filz)
    except:
        pass

    if download or len(filz)==0:
        if not os.path.isdir(directory):
            print '  -> creating', directory
            os.makedirs(directory)
        downloadFilters(directory) # only download the required set of filters
        filz = os.listdir(directory)
        filz = filter(lambda x: x.endswith('.xml'), filz)
    filz.sort()
    for fil in filz:
        f = open(os.path.join(directory, fil))
        lines = f.readlines()
        tmp = {}
        prof = []
        #print '-'*80
        for l in lines:
            if '<PARAM' in l and not 'ProfileReference' in l:
                #print Tag(l, 'name')
                try:
                    tmp[Tag(l, 'name')] = float(Tag(l,'value'))
                except:
                    tmp[Tag(l, 'name')] = Tag(l,'value')
            if l.strip().startswith(''):
                prof.append(float(l.split('>')[1].split('<')[0]))
        # -- compatibility:
        if not 'ID' in tmp.keys() and 'filterID' in tmp.keys():
            tmp['ID'] = tmp['filterID']
        prof[1] = 0 # force first value in profile to 0
        prof[-1] = 0 # force last value of filter to be 0
        # -- use um for the wavelength, instead of Angstrom
        tmp['profile']=(np.array(prof[::2])*1e-4,np.array(prof[1::2]))
        # -- build a more common name that the 'ID': 'band_system'
        band = tmp['ID'].split('.')[1]

        if tmp['ID'].split('/')[0]=='Generic' or \
               tmp['ID'].split('/')[0]=='Paranal':
            typ = tmp['ID'].split('/')[1].split('.')[0]
        else:
            typ = tmp['ID'].split('/')[0]

        # -- name of filter == name of file
        k = fil.split('.')[:-1] # remove .xml
        if 'Generic' in k:
            k.remove('Generic')
        if len(k)>2 and k[-2]==k[-3]:
            k.pop(-2)
        k = k[-1]+'_'+'_'.join(k[:-1])
        res[k] = tmp
    #__WalravenFilters()
    return res

def writeFilter(filt):
    """
    filt is the dictionarry read from an xml file (or built manually)

    assumes profile's wavelength are in um, will be converted in Angstrom
    """
    f = open('filterTemplate.xml')
    filename = filt['ID'].replace(' ', '').replace('/','_')+'.xml'
    g = open(filt['ID'].replace(' ', '').replace('/','_')+'.xml', 'w')
    for l in f.readlines():
        for k in filt.keys():
            if '__'+k in l:
                #print '>>> found', k
                if k!= 'profile':
                    l = l.replace('__'+k, str(filt[k]))
                else:
                    s = ["
\n %f	\n%f	\n
\n"%(10000*filt['profile'][0][i], filt['profile'][1][i])
                        for i in range(len(filt['profile'][0]))]
                    s = ''.join(s)
                    l = l.replace('__'+k, s)
                #print l
        g.write(l)
    f.close()
    g.close()
    return

def _DIRBE(flambda=None):
    """
    import atlas9
    # == crude Vega Model:
    diam = 3.1
    c = np.pi*(diam/2.0)**2 # angular surface, in mas**2
    c *= (np.pi/(180.0*3600*1000))**2 # conversion in rad**2
    flambda = lambda x: c*atlas9.flambda(x, 9800., 3.0)*10e-9

    """
    data = {'wl':[]}
    bands = [1,2,3,4,5,6,7,8,9,10]
    for b in bands:
        data[b] = []
    f = open(os.path.join(_dir, 'dirbe.txt'))
    for l in f.readlines():
        if not l.strip().startswith('#'):
            tmp = np.float_(l.split())
            data['wl'].append(tmp[0]*1e4) # um -> A
            for b in bands:
                data[b].append(tmp[b])
    f.close()

    for b in bands:
        F = {'Description': 'DIRBE '+str(b),
            'ID':'Generic/DIRBE.'+str(b),
            'FWHM': 110,
            'PhotSystem': 'Vega',
            'WavelengthPeak': -1,
            'WavelengthMin': -1,
            'WavelengthMax': -1,
            'WavelengthMean': -1,
            'WavelengthEff': -1,
            'WidthEff': -1,
            }
        _wl = np.array(data['wl'])
        _f = np.array(data[b])
        _i = np.arange(len(_f))
        i_min = min(_i[_f>0])-2
        i_max = max(_i[_f>0])+2
        profile = np.transpose(np.array([_wl[i_min:i_max], _f[i_min:i_max]]))
        F['WavelengthMean'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
        F['WavelengthEff'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
        F['WavelengthMin'] = np.min(profile[:,0])
        F['WavelengthMax'] = np.max(profile[:,0])
        F['WavelengthPeak'] = profile[np.argmax(profile[:,1]),0]
        #print F
        F['profile'] = (profile[:,0]*1e-4, profile[:,1])
        if flambda is None:
            F['ZeroPoint'] = 1.0 # fake
        else:
            tmp = np.sum(flambda(profile[:,0]*1e-4)*profile[:,1])/np.sum(profile[:,1])
            F['ZeroPoint'] = convert_Wm2um_to_Jy(tmp, F['WavelengthEff']*1e-4)
            print tmp, 'W/m2/um', F['ZeroPoint'], 'Jy', F['WavelengthEff']*1e-4
            writeFilter(F)

    return data

def _OGLE():
    f = open(os.path.join(_dir, 'QEOGLE.txt'))
    QE = []
    for l in f.readlines():
        QE.append([float(l.split(',')[0]), float(l.split(',')[1])])
    f.close
    QE = np.array(QE)

    V = {'Description': 'OGLE V',
        'ID':'OGLE/SPIPS.V',
        'FWHM': 110,
        'PhotSystem': 'Vega',
        'WavelengthPeak': -1,
        'WavelengthMin': -1,
        'WavelengthMax': -1,
        'WavelengthMean': -1,
        'WavelengthEff': -1,
        'WidthEff': -1,
         }
    f = open(os.path.join(_dir, 'VOGLE.txt'))
    profile = []
    for l in f.readlines():
        profile.append([float(l.split(',')[0]), float(l.split(',')[1])])
    f.close
    profile = np.array(profile)
    profile[0,1] = 0 # 0 at the edge
    profile[-1,1] = 0 # 0 at the edge
    V['WavelengthMean'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
    V['WavelengthEff'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
    V['WavelengthMin'] = np.min(profile[:,0])
    V['WavelengthMax'] = np.max(profile[:,0])
    V['WavelengthPeak'] = profile[np.argmax(profile[:,1]),0]
    V['profile'] = (profile[:,0]/10000, profile[:,1]*np.interp(profile[:,0],
                                                        QE[:,0], QE[:,1]))
    V['ZeroPoint'] = convert_Wm2um_to_Jy(3.85263580174e-08, V['WavelengthEff']/1e4)
    writeFilter(V)

    I = {'Description': 'OGLE I',
        'ID':'OGLE/SPIPS.I',
        'FWHM': 110,
        'PhotSystem': 'Vega',
        'WavelengthPeak': -1,
        'WavelengthMin': -1,
        'WavelengthMax': -1,
        'WavelengthMean': -1,
        'WavelengthEff': -1,
        'WidthEff': -1,
         }
    f = open(os.path.join(_dir, 'IOGLE.txt'))
    profile = []
    for l in f.readlines():
        profile.append([float(l.split(',')[0]), float(l.split(',')[1])])
    f.close
    profile = np.array(profile)
    profile[0,1] = 0 # 0 at the edge
    profile[-1,1] = 0 # 0 at the edge
    I['WavelengthMean'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
    I['WavelengthEff'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
    I['WavelengthMin'] = np.min(profile[:,0])
    I['WavelengthMax'] = np.max(profile[:,0])
    I['WavelengthPeak'] = profile[np.argmax(profile[:,1]),0]
    I['profile'] = (profile[:,0]/10000, profile[:,1]*np.interp(profile[:,0],
                                                        QE[:,0], QE[:,1]))
    I['ZeroPoint'] = convert_Wm2um_to_Jy(1.14521005341e-08 , I['WavelengthEff']/1e4)
    writeFilter(I)
    return

def _OGLE_MG():
    """
    calibrated from Martin Groenewegen, email from Alex 2017/
    """
    zp_jy = {'u':2012.02405, 'b':3861.51953, 'v':3625.37329, 'i':2369.19482}
    for band in ['u','b','v','i']:
        T = {'Description': 'OGLE %s calibrated by M. Groenewegen'%band.upper(),
            'ID':'OGLE/MG.'+band.upper(),
            'FWHM': 110,
            'PhotSystem': 'Vega',
            'WavelengthPeak': -1,
            'WavelengthMin': -1,
            'WavelengthMax': -1,
            'WavelengthMean': -1,
            'WavelengthEff': -1,
            'WidthEff': -1,
             }
        f = open(os.path.join(_dir, band+'_ogle.dat'))
        profile = []
        for l in f.readlines():
            #print '"'+l+'"'
            if len(l.strip())>5:
                profile.append([float(l.split()[0]), float(l.split()[1])])
        f.close()
        profile = np.array(profile)
        profile[0,1] = 0 # 0 at the edge
        profile[-1,1] = 0 # 0 at the edge
        T['WavelengthMean'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
        T['WavelengthEff'] = np.sum(profile[:,1]*profile[:,0])/np.sum(profile[:,1])
        T['WavelengthMin'] = np.min(profile[:,0])
        T['WavelengthMax'] = np.max(profile[:,0])
        T['WavelengthPeak'] = profile[np.argmax(profile[:,1]),0]
        T['profile'] = (profile[:,0]/10000, profile[:,1])
        T['ZeroPoint'] = zp_jy[band]
        writeFilter(T)
    pass

def _HipparcosBessell2000():
    """
    http://ulisse.pd.astro.it/Astro/ADPS/ADPS2/FileHtml/index_f135.html
    """
    H = {'Description': 'Hp Hipparcos Bessell 2000 version',
    'ID':'B2000/Hipparcos.Hp',
    'FWHM': 134.0,
    'PhotSystem': 'Hipparcos',
    'WavelengthPeak': 4365.0,
    'WavelengthMin': 3400.0,
    'WavelengthMax': 8800.0,
    'WavelengthMean': 4208.0,
    'WavelengthEff': 4208.0,
    'WidthEff': -1,
     }
    profile = [(3400,  0.000),
            (3500,  0.015),
            (3600,  0.032),
            (3800,  0.103),
            (3900,  0.155),
            (4000,  0.227),
            (4100,  0.300),
            (4200,  0.400),
            (4300,  0.530),
            (4400,  0.700),
            (4500,  0.845),
            (4600,  0.928),
            (4700,  0.970),
            (4800,  1.000),
            (4900,  0.997),
            (5000,  0.988),
            (5100,  0.973),
            (5200,  0.956),
            (5300,  0.935),
            (5400,  0.911),
            (5500,  0.887),
            (5600,  0.858),
            (5700,  0.825),
            (5800,  0.789),
            (5900,  0.752),
            (6000,  0.714),
            (6100,  0.675),
            (6200,  0.637),
            (6300,  0.599),
            (6400,  0.562),
            (6500,  0.525),
            (6600,  0.487),
            (6700,  0.449),
            (6800,  0.411),
            (6900,  0.372),
            (7000,  0.335),
            (7100,  0.299),
            (7200,  0.264),
            (7300,  0.232),
            (7400,  0.203),
            (7500,  0.176),
            (8000,  0.074),
            (8500,  0.020),
            (8800,  0.000),]
    H['profile'] = (np.array(profile)[:,0]/10000, np.array(profile)[:,1])
    # -- fitted to delta cep
    H['ZeroPoint'] = convert_Wm2um_to_Jy(3.752e-8, H['WavelengthEff']/1e4)
    writeFilter(H)
    return

def _WalravenFilters():
    """
    see http://ulisse.pd.astro.it/Astro/ADPS/Systems/Sys_021/index_021.html
    """
    # -- unit conversion: http://www.stsci.edu/hst/nicmos/tools/conversion_form.html
    # -- erg/cm2/s/A -> W/m2/um
    c = 10.

    W = {'Description': '-2.5 Walraven filter W',
    'ID':'Walraven.W',
    'FWHM': 134.0,
    'PhotSystem': 'Walraven',
    'WavelengthPeak': 3270.0,
    'WavelengthMin': 3110.0,
    'WavelengthMax': 3410.0,
    'WavelengthMean': 3254.0,
    'WavelengthEff': 3254.0,
    'WidthEff': -1,
     }
    profile = [[3100, 0.000], [3110, 0.000], [3120, 0.025], [3130, 0.061],
               [3140, 0.130], [3150, 0.215], [3160, 0.300], [3170, 0.388],
               [3180, 0.475], [3190, 0.565], [3200, 0.652], [3210, 0.736],
               [3220, 0.817], [3230, 0.893], [3240, 0.956], [3250, 0.997],
               [3260, 0.995], [3270, 0.953], [3280, 0.889], [3290, 0.814],
               [3300, 0.725], [3310, 0.628], [3320, 0.534], [3330, 0.444],
               [3340, 0.359], [3350, 0.285], [3360, 0.211], [3370, 0.144],
               [3380, 0.080], [3390, 0.028], [3400, 0.005], [3410, 0.000]]
    W['profile'] = (np.array(profile)[:,0]/10000, np.array(profile)[:,1])
    # -- http://ulisse.pd.astro.it/Astro/ADPS/Systems/Sys_021/fig_021_gif.html
    W['ZeroPoint'] = convert_Wm2um_to_Jy(2.12e-11*c, W['WavelengthEff']/1e4)
    writeFilter(W)

    U = {'Description': '-2.5 Walraven filter U',
    'ID':'Walraven.U',
    'FWHM': 260.,
    'PhotSystem': 'Walraven',
    'WavelengthPeak': 3620.0,
    'WavelengthMin': 3380.0,
    'WavelengthMax': 3920.0,
    'WavelengthMean': 3633.0,
    'WavelengthEff': 3633.0,
    'WidthEff': -1,
     }
    profile = [[3380, 0.000],[3390, 0.004],[3400, 0.010],[3410, 0.020],
               [3420, 0.033],[3430, 0.058],[3440, 0.097],[3450, 0.142],
               [3460, 0.192],[3470, 0.245],[3480, 0.300],[3490, 0.356],
               [3500, 0.415],[3510, 0.475],[3520, 0.535],[3530, 0.596],
               [3540, 0.657],[3550, 0.721],[3560, 0.784],[3570, 0.846],
               [3580, 0.901],[3590, 0.948],[3600, 0.982],[3610, 0.995],
               [3620, 1.000],[3630, 0.993],[3640, 0.978],[3650, 0.956],
               [3660, 0.929],[3670, 0.898],[3680, 0.865],[3690, 0.827],
               [3700, 0.787],[3710, 0.742],[3720, 0.691],[3730, 0.636],
               [3740, 0.575],[3750, 0.512],[3760, 0.444],[3770, 0.373],
               [3780, 0.306],[3790, 0.249],[3800, 0.200],[3810, 0.159],
               [3820, 0.127],[3830, 0.100],[3840, 0.076],[3850, 0.058],
               [3860, 0.043],[3870, 0.031],[3880, 0.023],[3890, 0.015],
               [3900, 0.007],[3910, 0.002],[3920, 0.000],]

    U['profile'] = (np.array(profile)[:,0]/10000, np.array(profile)[:,1])
    # -- http://ulisse.pd.astro.it/Astro/ADPS/Systems/Sys_021/fig_021_gif.html
    U['ZeroPoint'] = convert_Wm2um_to_Jy(1.61e-11*c, U['WavelengthEff']/1e4)
    writeFilter(U)

    L = {'Description': '-2.5 Walraven filter L',
    'ID':'Walraven.L',
    'FWHM': 140.,
    'PhotSystem': 'Walraven',
    'WavelengthPeak': 3900.0,
    'WavelengthMin': 3590.0,
    'WavelengthMax': 4105.0,
    'WavelengthMean': 3838.0,
    'WavelengthEff': 3838.0,
    'WidthEff': -1,
     }
    profile = [[3590, 0.000],[3600, 0.003],[3610, 0.006],[3620, 0.011],
               [3630, 0.020],[3640, 0.035],[3650, 0.064],[3660, 0.108],
               [3670, 0.157],[3680, 0.215],[3690, 0.274],[3700, 0.336],
               [3710, 0.401],[3720, 0.468],[3730, 0.538],[3740, 0.609],
               [3750, 0.684],[3760, 0.759],[3770, 0.832],[3780, 0.900],
               [3790, 0.945],[3800, 0.974],[3810, 0.990],[3820, 0.999],
               [3830, 0.998],[3840, 0.985],[3850, 0.961],[3860, 0.929],
               [3870, 0.891],[3880, 0.855],[3890, 0.815],[3900, 0.775],
               [3910, 0.727],[3920, 0.676],[3930, 0.621],[3940, 0.562],
               [3950, 0.496],[3960, 0.425],[3970, 0.356],[3980, 0.299],
               [3990, 0.246],[4000, 0.198],[4015, 0.138],[4030, 0.090],
               [4045, 0.053],[4060, 0.028],[4075, 0.012],[4090, 0.004],
               [4105, 0.000],]

    L['profile'] = (np.array(profile)[:,0]/10000, np.array(profile)[:,1])
    # -- http://ulisse.pd.astro.it/Astro/ADPS/Systems/Sys_021/fig_021_gif.html
    L['ZeroPoint'] = convert_Wm2um_to_Jy(1.52e-11*c, L['WavelengthEff']/1e4)
    writeFilter(L)

    V = {'Description': '-2.5 Walraven filter V',
    'ID':'Walraven.V',
    'FWHM': 706.,
    'PhotSystem': 'Walraven',
    'WavelengthPeak': 5400.0,
    'WavelengthMin': 4800.0,
    'WavelengthMax': 6480.0,
    'WavelengthMean': 5467.0,
    'WavelengthEff': 5467.0,
    'WidthEff': -1,}
    # -- http://ulisse.pd.astro.it/Astro/ADPS/Systems/Sys_021/fig_021_gif.html
    V['ZeroPoint'] = convert_Wm2um_to_Jy(6.73e-12*c, V['WavelengthEff']/1e4)
    profile = [[4775, 0.000],[4800, 0.001],[4825, 0.009],[4850, 0.024],
               [4875, 0.045],[4900, 0.076],[4925, 0.115],[4950, 0.161],
               [4975, 0.209],[5000, 0.264],[5025, 0.320],[5050, 0.380],
               [5075, 0.441],[5100, 0.504],[5125, 0.568],[5150, 0.634],
               [5175, 0.699],[5200, 0.759],[5225, 0.811],[5250, 0.860],
               [5275, 0.902],[5300, 0.938],[5325, 0.967],[5350, 0.986],
               [5375, 0.998],[5400, 0.999],[5425, 0.992],[5450, 0.975],
               [5475, 0.955],[5500, 0.930],[5525, 0.904],[5550, 0.876],
               [5575, 0.847],[5600, 0.812],[5625, 0.779],[5650, 0.746],
               [5675, 0.709],[5700, 0.671],[5725, 0.633],[5750, 0.594],
               [5775, 0.554],[5800, 0.514],[5840, 0.451],[5880, 0.388],
               [5920, 0.321],[5960, 0.258],[6000, 0.201],[6040, 0.152],
               [6080, 0.113],[6120, 0.084],[6160, 0.058],[6200, 0.039],
               [6240, 0.029],[6280, 0.022],[6320, 0.016],[6360, 0.012],
               [6400, 0.008],[6440, 0.004],[6480, 0.000],]
    V['profile'] = (np.array(profile)[:,0]/10000, np.array(profile)[:,1])
    writeFilter(V)

    B = {'Description': '-2.5 Walraven filter B',
        'ID':'Walraven.B',
        'FWHM': 420.,
        'PhotSystem': 'Walraven',
        'WavelengthPeak': .0,
        'WavelengthMin': 3870,
        'WavelengthMax': 4825,
        'WavelengthMean': 4325.0,
        'WavelengthEff': 4325.0,
        'WidthEff': -1,}
    # -- http://ulisse.pd.astro.it/Astro/ADPS/Systems/Sys_021/fig_021_gif.html
    B['ZeroPoint'] = convert_Wm2um_to_Jy(1.23e-11*c, B['WavelengthEff']/1e4)
    profile = [[3870,  0.000],[3880, 0.001],[3890, 0.004],[3900, 0.008],
               [3910, 0.014],[3920, 0.023],[3930, 0.035],[3940, 0.051],
               [3950, 0.071],[3960, 0.092],[3970, 0.112],[3980, 0.136],
               [3990, 0.161],[4000, 0.186],[4015, 0.228],[4030, 0.273],
               [4045, 0.320],[4060, 0.368],[4075, 0.419],[4090, 0.471],
               [4105, 0.523],[4120, 0.575],[4135, 0.627],[4150, 0.682],
               [4165, 0.733],[4180, 0.781],[4195, 0.824],[4210, 0.863],
               [4225, 0.898],[4240, 0.929],[4255, 0.954],[4270, 0.976],
               [4285, 0.990],[4300, 0.998],[4320, 0.997],[4340, 0.985],
               [4360, 0.960],[4380, 0.929],[4400, 0.893],[4420, 0.850],
               [4440, 0.800],[4460, 0.744],[4480, 0.683],[4500, 0.621],
               [4520, 0.562],[4540, 0.504],[4560, 0.445],[4580, 0.390],
               [4600, 0.335],[4620, 0.282],[4640, 0.230],[4660, 0.181],
               [4680, 0.138],[4700, 0.101],[4725, 0.063],[4750, 0.034],
               [4775, 0.015],[4800, 0.003],[4825 , 0.000]]
    B['profile'] = (np.array(profile)[:,0]/10000, np.array(profile)[:,1])
    writeFilter(B)
    return

def _TESS():
    """
    VERY APPROXIMATIVE!
    """
    T = {'Description': 'TESS space mission',
    'ID':'TESS.T',
    'FWHM': 4000.0,
    'PhotSystem': 'TESS',
    'WavelengthPeak': 8800,
    'WavelengthMin': 5800.,
    'WavelengthMax': 11050,
    'WavelengthMean': 8000,
    'WavelengthEff': 8000,
    'WidthEff': -1,
     }
    # http://astronomicaltelescopes.spiedigitallibrary.org/data/Journals/JATIAG/931021/JATIS_1_1_014003_f001.png
    profile = [[580, 0], [600, 0.9], [860, 1.0], [900, .95],
                [920, 0.9], [930, 0.85], [940, 0.8], [960, 0.7],
                [980, 0.6], [1000, 0.45], [1030, 0.2], [1060, 0.1],
                [1100, 0.08], [1130, 0]]

    T['profile'] = (np.array(profile)[:,0]/1000., np.array(profile)[:,1])
    # -- from rotastar.VegaFlux('T_TESS.T')
    T['ZeroPoint'] = convert_Wm2um_to_Jy(1.31837883589e-08, T['WavelengthEff']/1e4)
    writeFilter(T)

def effWavelength_um(filtname):
    global _data
    if isinstance(filtname, str):
        return _data[filtname.strip()]['WavelengthEff']*1e-4
        #return _data[filtname.strip()]['WavelengthMean']*1e-4
    elif isinstance(filtname, list):
        return [effWavelength_um(f) for f in filtname]
    elif isinstance(filtname, np.ndarray):
        return np.array([effWavelength_um(f) for f in filtname])
    else:
        return None

def zeroPoint_Wm2um(filtname):
    global _data
    kz = filtname.strip()
    wl0 = effWavelength_um(filtname)
    return convert_Jy_to_Wm2um(_data[kz]['ZeroPoint'],wl0)

def convert_Jy_to_Wm2um(f_jy, wl_um):
    if isinstance(wl_um, str):
        wl_um = effWavelength_um(wl_um)
    return f_jy/(wl_um**2/2.9979e-12)

def convert_Wm2um_to_Jy(f, wl_um):
    if isinstance(wl_um, str):
        wl_um = effWavelength_um(wl_um)
    return f*(wl_um**2/2.9979e-12)

def convert_mag_to_Wm2um(mag, filtname):
    zp = zeroPoint_Wm2um(filtname)
    return zp*10.0**(-mag/2.5)

def convert_Wm2um_to_mag(flux, filtname):
    zp = zeroPoint_Wm2um(filtname)
    return -2.5*np.log10(flux/zp)

def convert_mag_to_Jy(mag, filtname):
    return convert_Wm2um_to_Jy(convert_mag_to_Wm2um(mag, filtname),
                    effWavelength_um(filtname))

def convert_Jy_to_mag(f, filtname):
    return convert_Wm2um_to_mag(convert_Jy_to_Wm2um(f, filtname), filtname)

_data = LoadFilters()

print ' | available filters from "photfilt2":'

print '   ',
kz = _data.keys()
kz.sort()
for k in kz:
    print '   ', k, '('+_data[k]['ID']+')', _data[k]['Description']
print ' | for more download VOTable files (XML) in', _dir, 'from'
print ' | http://svo.cab.inta-csic.es/theory/fps3/index.php?mode=browse'
print ''

max_stars_repo_name: amerand/SPIPS

max_issues_repo_name: amerand/SPIPS

max_forks_repo_name: amerand/SPIPS

reason: import numpy,from scipy


text:
# coding=utf-8

from __future__ import print_function
from gen.gen_captcha import gen_dataset, load_templates, candidates
from model.nn import load_model_nn
from model.common import find_model_ckpt
import os
import tensorflow as tf
from gen.utils import vec2str
import numpy as np
from PIL import Image

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
trainer_dir = os.path.dirname(os.path.abspath(__file__))
home_dir = os.path.dirname(trainer_dir)
graph_log_dir = os.path.join(trainer_dir, 'logs')

def show_im(dataset):
    data = np.uint8(dataset[0]).reshape((30, 96)) * 255
    im = Image.fromarray(data)
    im.show()

def test_model():
    templates = load_templates()

    model = load_model_nn()
    x = model['x']
    keep_prob = model['keep_prob']
    saver = model['saver']
    prediction = model['prediction']
    graph = model['graph']
    model_ckpt_path, _ = find_model_ckpt()
    print("Used the model:", model_ckpt_path)

    with tf.Session(graph=graph) as session:
        tf.global_variables_initializer().run()
        saver.restore(session, model_ckpt_path)

        dataset, labels = gen_dataset(1, templates)  # generate one image

        label = prediction.eval(feed_dict={x: dataset, keep_prob: 1.0}, session=session)[0]
        string = ''
        for i in range(4):
            string += chr(label[i] + ord('0'))
        print("predict label:", string)
        show_im(dataset)
        print("actual label:", vec2str(labels[0]))

if __name__ == "__main__":
    test_model()


text:
import numpy, itertools, pickle, warnings as _builtin_warnings
from nutils import *
from nutils.testing import *
_ = numpy.newaxis

@parametrize
class check(TestCase):

  def setUp(self):
    super().setUp()
    domain = mesh.rectilinear([1]*self.ndim, periodic=[0])[0]
    self.sample = domain.sample('uniform', 2)
    param = rootcoords = evaluable.ApplyTransforms(evaluable.PopHead(self.ndim, evaluable.SelectChain(0)), evaluable.Points(self.sample.npoints, self.sample.ndims))
    if self.ndim == 1:
      self.geom = param**2
      poly = numpy.array([[1,-2,1],[0,2,-2],[0,0,1]]) # 2nd order bernstein
    elif self.ndim == 2:
      r, theta = param[:,0], param[:,1]
      self.geom = r[:,_] * evaluable.stack([evaluable.cos(theta), evaluable.sin(theta)], axis=1)
      poly = numeric.poly_outer_product([[1,-2,1],[0,2,-2],[0,0,1]], [[1,-1],[0,1]]) # 2nd x 1st order bernstein
    else:
      raise Exception('invalid ndim {!r}'.format(self.ndim))
    numpy.random.seed(0)
    self.args = [evaluable.Guard(evaluable.Polyval(numeric.dot(numpy.random.uniform(size=shape+poly.shape[:1], low=self.low, high=self.high), poly), rootcoords)) for shape in self.shapes]
    if self.pass_geom:
        self.args += [self.geom]
    self.n_op_argsfun = self.n_op(*self.sample_eval(self.args))
    self.op_args = self.op(*self.args)
    self.shapearg = numpy.random.uniform(size=self.n_op_argsfun.shape, low=self.low, high=self.high)
    self.pairs = [(i, j) for i in range(self.op_args.ndim-1) for j in range(i+1, self.op_args.ndim) if self.op_args.shape[i] == self.op_args.shape[j]]
    _builtin_warnings.simplefilter('ignore', evaluable.ExpensiveEvaluationWarning)

  def assertArrayAlmostEqual(self, actual, desired, decimal):
    if actual.shape != desired.shape:
      self.fail('shapes of actual {} and desired {} are incompatible.'.format(actual.shape, desired.shape))
    error = actual - desired if not actual.dtype.kind == desired.dtype.kind == 'b' else actual ^ desired
    approx = error.dtype.kind in 'fc'
    indices = tuple(zip(*(numpy.greater_equal(abs(error), 1.5 * 10**-decimal) if approx else error).nonzero()))
    if not indices:
      return
    lines = ['arrays are not equal']
    if approx:
      lines.append(' up to {} decimals'.format(decimal))
    lines.append(' in {}/{} entries:'.format(len(indices), error.size))
    n = 5
    lines.extend('\n  {} actual={} desired={} difference={}'.format(index, actual[index], desired[index], error[index]) for index in indices[:n])
    if len(indices) > 2*n:
      lines.append('\n  ...')
      n = -n
    lines.extend('\n  {} actual={} desired={} difference={}'.format(index, actual[index], desired[index], error[index]) for index in indices[n:])
    self.fail(''.join(lines))

  def assertFunctionAlmostEqual(self, actual, desired, decimal):
    evalargs = dict(_transforms=[trans[0] for trans in self.sample.transforms], _points=self.sample.points[0])
    with self.subTest('vanilla'):
      self.assertArrayAlmostEqual(actual.eval(**evalargs), desired, decimal)
    with self.subTest('simplified'):
      self.assertArrayAlmostEqual(actual.simplified.eval(**evalargs), desired, decimal)
    with self.subTest('optimized'):
      self.assertArrayAlmostEqual(actual.optimized_for_numpy.eval(**evalargs), desired, decimal)
    with self.subTest('sample'):
      self.assertArrayAlmostEqual(self.sample_eval(actual), desired, decimal)

  @util.positional_only
  @util.single_or_multiple
  def sample_eval(self, funcs, arguments=...):
    datas = self.sample._eval(funcs, arguments)
    return tuple(map(sparse.toarray, datas))

  def test_evalconst(self):
    constargs = [numpy.random.uniform(size=(self.sample.npoints, *shape)) for shape in self.shapes]
    if self.pass_geom:
      constargs += [numpy.random.uniform(size=self.geom.shape)]
    self.assertFunctionAlmostEqual(decimal=15,
      desired=self.n_op(*[constarg for constarg in constargs]),
      actual=self.op(*constargs))

  def test_eval(self):
    self.assertFunctionAlmostEqual(decimal=15,
      actual=self.op_args,
      desired=self.n_op_argsfun)

  def test_getitem(self):
    for idim in range(self.op_args.ndim):
      for item in range(self.op_args.shape[idim]):
        s = (Ellipsis,) + (slice(None),)*idim + (item,) + (slice(None),)*(self.op_args.ndim-idim-1)
        self.assertFunctionAlmostEqual(decimal=15,
          desired=self.n_op_argsfun[s],
          actual=self.op_args[s])

  def test_transpose(self):
    trans = numpy.arange(self.op_args.ndim,0,-1) % self.op_args.ndim
    self.assertFunctionAlmostEqual(decimal=15,
      desired=numpy.transpose(self.n_op_argsfun, trans),
      actual=evaluable.transpose(self.op_args, trans))

  def test_insertaxis(self):
    for axis in range(self.op_args.ndim+1):
      with self.subTest(axis=axis):
        self.assertFunctionAlmostEqual(decimal=15,
          desired=numpy.repeat(numpy.expand_dims(self.n_op_argsfun, axis), 2, axis),
          actual=evaluable.insertaxis(self.op_args, axis, 2))

  def test_takediag(self):
    for ax1, ax2 in self.pairs:
      self.assertFunctionAlmostEqual(decimal=15,
        desired=numeric.takediag(self.n_op_argsfun, ax1, ax2),
        actual=evaluable.takediag(self.op_args, ax1, ax2))

  def test_eig(self):
    if self.op_args.dtype == float:
      for ax1, ax2 in self.pairs:
        A = self.sample_eval(self.op_args)
        L, V = self.sample_eval(list(evaluable.eig(self.op_args, axes=(ax1,ax2))))
        self.assertArrayAlmostEqual(decimal=11,
          actual=(numpy.expand_dims(V,ax2) * numpy.expand_dims(L,ax2+1).swapaxes(ax1,ax2+1)).sum(ax2+1),
          desired=(numpy.expand_dims(A,ax2) * numpy.expand_dims(V,ax2+1).swapaxes(ax1,ax2+1)).sum(ax2+1))

  def test_inv(self):
    for ax1, ax2 in self.pairs:
      trans = [i for i in range(self.n_op_argsfun.ndim) if i not in (ax1,ax2)] + [ax1,ax2]
      invtrans = list(map(trans.index, range(len(trans))))
      self.assertFunctionAlmostEqual(decimal=10,
        desired=numeric.inv(self.n_op_argsfun.transpose(trans)).transpose(invtrans),
        actual=evaluable.inverse(self.op_args, axes=(ax1,ax2)))

  def test_determinant(self):
    for ax1, ax2 in self.pairs:
      self.assertFunctionAlmostEqual(decimal=11,
        desired=numpy.linalg.det(self.n_op_argsfun.transpose([i for i in range(self.n_op_argsfun.ndim) if i not in (ax1,ax2)] + [ax1,ax2])),
        actual=evaluable.determinant(self.op_args, axes=(ax1,ax2)))

  def test_take(self):
    indices = [0,-1]
    for iax, sh in enumerate(self.op_args.shape):
      if sh >= 2:
        self.assertFunctionAlmostEqual(decimal=15,
          desired=numpy.take(self.n_op_argsfun, indices, axis=iax),
          actual=evaluable.take(self.op_args, indices, axis=iax))

  def test_take_nomask(self):
    for iax, sh in enumerate(self.op_args.shape):
      if sh >= 2:
        indices = [0,sh-1]
        self.assertFunctionAlmostEqual(decimal=15,
          desired=numpy.take(self.n_op_argsfun, indices, axis=iax),
          actual=evaluable.take(self.op_args, evaluable.Guard(evaluable.asarray(indices)), axis=iax))

  def test_take_reversed(self):
    indices = [-1,0]
    for iax, sh in enumerate(self.op_args.shape):
      if sh >= 2:
        self.assertFunctionAlmostEqual(decimal=15,
          desired=numpy.take(self.n_op_argsfun, indices, axis=iax),
          actual=evaluable.take(self.op_args, indices, axis=iax))

  def test_inflate(self):
    for iax, sh in enumerate(self.op_args.shape):
      dofmap = evaluable.Constant(numpy.arange(sh) * 2)
      desired = numpy.zeros(self.n_op_argsfun.shape[:iax] + (sh*2-1,) + self.n_op_argsfun.shape[iax+1:], dtype=self.n_op_argsfun.dtype)
      desired[(slice(None),)*iax+(slice(None,None,2),)] = self.n_op_argsfun
      self.assertFunctionAlmostEqual(decimal=15,
        desired=desired,
        actual=evaluable._inflate(self.op_args, dofmap=dofmap, length=sh*2-1, axis=iax))

  def test_diagonalize(self):
    for axis in range(self.op_args.ndim):
      for newaxis in range(axis+1, self.op_args.ndim+1):
        self.assertFunctionAlmostEqual(decimal=15,
          desired=numeric.diagonalize(self.n_op_argsfun, axis, newaxis),
          actual=evaluable.diagonalize(self.op_args, axis, newaxis))

  def test_product(self):
    for iax in range(self.op_args.ndim):
      self.assertFunctionAlmostEqual(decimal=15,
        desired=numpy.product(self.n_op_argsfun, axis=iax),
        actual=evaluable.product(self.op_args, axis=iax))

  def test_getslice(self):
    for idim in range(self.op_args.ndim):
      if self.op_args.shape[idim] == 1:
        continue
      s = (Ellipsis,) + (slice(None),)*idim + (slice(0,self.op_args.shape[idim]-1),) + (slice(None),)*(self.op_args.ndim-idim-1)
      self.assertFunctionAlmostEqual(decimal=15,
        desired=self.n_op_argsfun[s],
        actual=self.op_args[s])

  def test_sumaxis(self):
    for idim in range(self.op_args.ndim):
      self.assertFunctionAlmostEqual(decimal=14,
        desired=self.n_op_argsfun.sum(idim),
        actual=self.op_args.sum(idim))

  def test_add(self):
    self.assertFunctionAlmostEqual(decimal=15,
      desired=self.n_op_argsfun + self.shapearg,
      actual=(self.op_args + self.shapearg))

  def test_multiply(self):
    self.assertFunctionAlmostEqual(decimal=15,
      desired=self.n_op_argsfun * self.shapearg,
      actual=(self.op_args * self.shapearg))

  def test_dot(self):
    for iax in range(self.op_args.ndim):
      self.assertFunctionAlmostEqual(decimal=13,
        desired=numeric.contract(self.n_op_argsfun, self.shapearg, axis=iax),
        actual=evaluable.dot(self.op_args, self.shapearg, axes=iax))

  def test_pointwise(self):
    self.assertFunctionAlmostEqual(decimal=15,
      desired=numpy.sin(self.n_op_argsfun).astype(float), # "astype" necessary for boolean operations (float16->float64)
      actual=evaluable.sin(self.op_args))

  def test_power(self):
    self.assertFunctionAlmostEqual(decimal=13,
      desired=self.n_op_argsfun**3,
      actual=(self.op_args**3))

  def test_power0(self):
    power = (numpy.arange(self.op_args.size) % 2).reshape(self.op_args.shape)
    self.assertFunctionAlmostEqual(decimal=13,
      desired=self.n_op_argsfun**power,
      actual=self.op_args**power)

  def test_sign(self):
    if self.n_op_argsfun.dtype.kind != 'b':
      self.assertFunctionAlmostEqual(decimal=15,
        desired=numpy.sign(self.n_op_argsfun),
        actual=evaluable.sign(self.op_args))

  def test_mask(self):
    for idim in range(self.op_args.ndim):
      if self.op_args.shape[idim] <= 1:
        continue
      mask = numpy.ones(self.op_args.shape[idim], dtype=bool)
      mask[0] = False
      if self.op_args.shape[idim] > 2:
        mask[-1] = False
      self.assertFunctionAlmostEqual(decimal=15,
        desired=self.n_op_argsfun[(slice(None,),)*idim+(mask,)],
        actual=evaluable.mask(self.op_args, mask, axis=idim))

  def test_ravel(self):
    for idim in range(self.op_args.ndim-1):
      self.assertFunctionAlmostEqual(decimal=15,
        desired=self.n_op_argsfun.reshape(self.n_op_argsfun.shape[:idim]+(-1,)+self.n_op_argsfun.shape[idim+2:]),
        actual=evaluable.ravel(self.op_args, axis=idim))

  def test_unravel(self):
    for idim in range(self.op_args.ndim):
      length = self.n_op_argsfun.shape[idim]
      unravelshape = (length//3,3) if (length%3==0) else (length//2,2) if (length%2==0) else (length,1)
      self.assertFunctionAlmostEqual(decimal=15,
        desired=self.n_op_argsfun.reshape(self.n_op_argsfun.shape[:idim]+unravelshape+self.n_op_argsfun.shape[idim+1:]),
        actual=evaluable.unravel(self.op_args, axis=idim, shape=unravelshape))

  def test_desparsify(self):
    args = []
    for arg in self.args:
      for i in range(arg.ndim):
        arg = evaluable._inflate(arg, evaluable.Guard(numpy.arange(arg.shape[i])), arg.shape[i], i)
      args.append(arg)
    op_args = self.op(*args).simplified
    _transforms = [trans[0] for trans in self.sample.transforms]
    for axis, prop in enumerate(op_args._axes):
      if isinstance(prop, evaluable.Sparse):
        actual = numpy.zeros_like(self.n_op_argsfun)
        for ind, f in op_args._desparsify(axis):
          _ind = ind.eval(_transforms=_transforms)
          actual[(slice(None),)*(axis)+(_ind,)] += f.eval(_transforms=_transforms, _points=self.sample.points[0])
        self.assertArrayAlmostEqual(actual, self.n_op_argsfun, decimal=15)

  def find(self, target, xi0):
    elemtrans, = self.sample.transforms[0]
    ndim, = self.geom.shape
    J = evaluable.localgradient(self.geom, ndim)
    Jinv = evaluable.inverse(J)
    countdown = 5
    iiter = 0
    self.assertEqual(target.shape[-1:], self.geom.shape)
    if xi0.shape != target.shape:
      tmp = numpy.empty_like(target)
      tmp[...] = xi0
      xi0 = tmp
    target = target.reshape(-1, target.shape[-1])
    xi = xi0.reshape(-1, xi0.shape[-1])
    while countdown:
      err = target - self.geom.eval(_transforms=[elemtrans], _points=points.CoordsPoints(xi))
      if numpy.less(numpy.abs(err), 1e-12).all():
        countdown -= 1
      dxi_root = (Jinv.eval(_transforms=[elemtrans], _points=xi) * err[...,_,:]).sum(-1)
      #xi = xi + numpy.dot(dxi_root, self.elem.inv_root_transform.T)
      xi = xi + dxi_root
      iiter += 1
      self.assertLess(iiter, 100, 'failed to converge in 100 iterations')
    return xi.reshape(xi0.shape)

  @parametrize.enable_if(lambda hasgrad, **kwargs: hasgrad)
  def test_derivative(self):
    eps = 1e-4
    fddeltas = numpy.array([1,2,3])
    fdfactors = numpy.linalg.solve(2*fddeltas**numpy.arange(1,1+2*len(fddeltas),2)[:,None], [1]+[0]*(len(fddeltas)-1))
    for iarg, shape in enumerate(self.shapes):
      shape = (self.sample.npoints, *shape)
      x0 = self.sample_eval(self.args[iarg])
      dx = numpy.random.normal(size=shape)
      x = evaluable.Argument('x', shape)
      f = self.op(*self.args[:iarg]+[x]+self.args[iarg+1:])
      exact = numeric.contract(self.sample_eval(evaluable.derivative(f, x), x=x0), dx, range(f.ndim, f.ndim+dx.ndim))
      if exact.dtype.kind in 'bi' or self.zerograd:
        approx = numpy.zeros_like(exact)
        scale = 1
      else:
        fdvals = numpy.stack([self.sample_eval(f, x=x0+eps*n*dx) for n in (*-fddeltas, *fddeltas)], axis=0)
        if fdvals.dtype.kind == 'i':
          fdvals = fdvals.astype(float)
        fdvals = fdvals.reshape(2, len(fddeltas), *fdvals.shape[1:])
        approx = ((fdvals[1] - fdvals[0]).T @ fdfactors).T / eps
        scale = numpy.linalg.norm(self.sample_eval(f, x=x0).ravel()) or 1
      self.assertArrayAlmostEqual(exact / scale, approx / scale, decimal=10)

def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, pass_geom=False, ndim=2, low=-1, high=1):
  check(name, op=op, n_op=n_op, shapes=shapes, hasgrad=hasgrad, zerograd=zerograd, pass_geom=pass_geom, ndim=ndim, low=low, high=high)

_check('identity', lambda f: evaluable.asarray(f), lambda a: a, [(2,4,2)])
_check('const', lambda f: evaluable.asarray(numpy.arange(16, dtype=float).reshape(2,4,2)), lambda a: numpy.arange(16, dtype=float).reshape(2,4,2), [(2,4,2)])
_check('zeros', lambda f: evaluable.zeros([1,4,3,4]), lambda a: numpy.zeros([1,4,3,4]), [(4,3,4)])
_check('ones', lambda f: evaluable.ones([1,4,3,4]), lambda a: numpy.ones([1,4,3,4]), [(4,3,4)])
_check('range', lambda f: evaluable.Range(4, offset=2)[numpy.newaxis], lambda a: numpy.arange(2,6)[numpy.newaxis], [(4,)])
_check('sin', evaluable.sin, numpy.sin, [(4,)])
_check('cos', evaluable.cos, numpy.cos, [(4,)])
_check('tan', evaluable.tan, numpy.tan, [(4,)])
_check('sqrt', evaluable.sqrt, numpy.sqrt, [(4,)], low=0)
_check('log', evaluable.ln, numpy.log, [(4,)], low=0)
_check('log2', evaluable.log2, numpy.log2, [(4,)], low=0)
_check('log10', evaluable.log10, numpy.log10, [(4,)], low=0)
_check('exp', evaluable.exp, numpy.exp, [(4,)])
_check('arctanh', evaluable.arctanh, numpy.arctanh, [(4,)])
_check('tanh', evaluable.tanh, numpy.tanh, [(4,)])
_check('cosh', evaluable.cosh, numpy.cosh, [(4,)])
_check('sinh', evaluable.sinh, numpy.sinh, [(4,)])
_check('abs', evaluable.abs, numpy.abs, [(4,)])
_check('sign', evaluable.sign, numpy.sign, [(4,4)], zerograd=True)
_check('power', evaluable.power, numpy.power, [(4,1),(1,4)], low=0)
_check('negative', evaluable.negative, numpy.negative, [(4,)])
_check('reciprocal', evaluable.reciprocal, numpy.reciprocal, [(4,)], low=-2, high=-1)
_check('arcsin', evaluable.arcsin, numpy.arcsin, [(4,)])
_check('arccos', evaluable.arccos, numpy.arccos, [(4,)])
_check('arctan', evaluable.arctan, numpy.arctan, [(4,)])
_check('ln', evaluable.ln, numpy.log, [(4,)], low=0)
_check('product', lambda a: evaluable.product(a,2), lambda a: numpy.product(a,-2), [(4,3,4)])
_check('sum', lambda a: evaluable.sum(a,2), lambda a: a.sum(-2), [(4,3,4)])
_check('transpose1', lambda a: evaluable.transpose(a,[0,1,3,2]), lambda a: a.transpose([0,1,3,2]), [(4,4,4)], low=0, high=20)
_check('transpose2', lambda a: evaluable.transpose(a,[0,2,3,1]), lambda a: a.transpose([0,2,3,1]), [(4,4,4)])
_check('insertaxis', lambda a: evaluable.insertaxis(a,2,3), lambda a: numpy.repeat(numpy.expand_dims(a,2), 3, 2), [(2,4)])
_check('get', lambda a: evaluable.get(a,2,1), lambda a: a[...,1,:], [(4,3,4)])
_check('takediag141', lambda a: evaluable.takediag(a,1,3), lambda a: numeric.takediag(a,1,3), [(1,4,1)])
_check('takediag434', lambda a: evaluable.takediag(a,1,3), lambda a: numeric.takediag(a,1,3), [(4,3,4)])
_check('takediag343', lambda a: evaluable.takediag(a,1,3), lambda a: numeric.takediag(a,1,3), [(3,4,3)])
_check('determinant141', lambda a: evaluable.determinant(a,(1,3)), lambda a: numpy.linalg.det(a.swapaxes(-3,-2)), [(1,4,1)])
_check('determinant434', lambda a: evaluable.determinant(a,(1,3)), lambda a: numpy.linalg.det(a.swapaxes(-3,-2)), [(4,3,4)])
_check('determinant4433', lambda a: evaluable.determinant(a,(3,4)), lambda a: numpy.linalg.det(a), [(4,4,3,3)])
_check('determinant200', lambda a: evaluable.determinant(a,(2,3)), lambda a: numpy.linalg.det(a) if a.shape[-1] else numpy.ones(a.shape[:-2], float), [(2,0,0)], zerograd=True)
_check('inverse141', lambda a: evaluable.inverse(a+evaluable.Diagonalize(evaluable.ones([1]))[None,:,None],(1,3)), lambda a: numpy.linalg.inv(a.swapaxes(-3,-2)+numpy.eye(1)).swapaxes(-3,-2), [(1,4,1)])
_check('inverse434', lambda a: evaluable.inverse(a+evaluable.Diagonalize(evaluable.ones([4]))[None,:,None],(1,3)), lambda a: numpy.linalg.inv(a.swapaxes(-3,-2)+numpy.eye(4)).swapaxes(-3,-2), [(4,3,4)])
_check('inverse4422', lambda a: evaluable.inverse(a+evaluable.Diagonalize(evaluable.ones([2]))[None]), lambda a: numpy.linalg.inv(a+numpy.eye(2)), [(4,4,2,2)])
_check('repeat', lambda a: evaluable.repeat(a,3,2), lambda a: numpy.repeat(a,3,-2), [(4,1,4)])
_check('diagonalize', lambda a: evaluable.diagonalize(a,1,3), lambda a: numeric.diagonalize(a,1,3), [(4,4,4,4)])
_check('multiply', evaluable.multiply, numpy.multiply, [(4,1),(4,4)])
_check('dot', lambda a,b: evaluable.dot(a,b,axes=2), lambda a,b: (a*b).sum(2), [(4,2,4),(4,2,4)])
_check('divide', evaluable.divide, lambda a, b: a * b**-1, [(4,4),(1,4)], low=-2, high=-1)
_check('divide2', lambda a: evaluable.asarray(a)/2, lambda a: a/2, [(4,1)])
_check('add', evaluable.add, numpy.add, [(4,1),(1,4)])
_check('subtract', evaluable.subtract, numpy.subtract, [(4,1),(1,4)])
_check('dot2', lambda a,b: evaluable.multiply(a,b).sum(-2), lambda a,b: (a*b).sum(-2), [(4,2,4),(1,2,4)])
_check('min', lambda a,b: evaluable.Minimum(a,b), numpy.minimum, [(4,4),(4,4)])
_check('max', lambda a,b: evaluable.Maximum(a,b), numpy.maximum, [(4,4),(4,4)])
_check('equal', lambda a,b: evaluable.Equal(*evaluable._numpy_align(a,b)), numpy.equal, [(4,1),(1,4)], zerograd=True)
_check('greater', lambda a,b: evaluable.Greater(*evaluable._numpy_align(a,b)), numpy.greater, [(4,1),(1,4)], zerograd=True)
_check('less', lambda a,b: evaluable.Less(*evaluable._numpy_align(a,b)), numpy.less, [(4,1),(1,4)], zerograd=True)
_check('arctan2', evaluable.arctan2, numpy.arctan2, [(4,1),(1,4)])
_check('stack', lambda a,b: evaluable.stack([a,b], 1), lambda a,b: numpy.concatenate([a[...,_,:],b[...,_,:]], axis=-2), [(4,),(4,)])
_check('eig', lambda a: evaluable.eig(a+a.swapaxes(1,2),symmetric=True)[1], lambda a: numpy.linalg.eigh(a+a.swapaxes(1,2))[1], [(4,4)], hasgrad=False)
_check('trignormal', lambda a: evaluable.TrigNormal(a), lambda a: numpy.array([numpy.cos(a), numpy.sin(a)]).T, [()])
_check('trigtangent', lambda a: evaluable.TrigTangent(a), lambda a: numpy.array([-numpy.sin(a), numpy.cos(a)]).T, [()])
_check('mod', lambda a,b: evaluable.mod(a,b), lambda a,b: numpy.mod(a,b), [(4,),(4,)], hasgrad=False)
_check('mask', lambda f: evaluable.mask(f,numpy.array([True,False,True,False,True,False,True]),axis=2), lambda a: a[:,:,::2], [(4,7,4)])
_check('ravel', lambda f: evaluable.ravel(f,axis=2), lambda a: a.reshape(-1,4,4,4,4), [(4,2,2,4,4)])
_check('unravel', lambda f: evaluable.unravel(f,axis=2,shape=[2,2]), lambda a: a.reshape(-1,4,2,2,4,4), [(4,4,4,4)])
_check('inflate', lambda f: evaluable._inflate(f,dofmap=evaluable.Guard([0,3]),length=4,axis=2), lambda a: numpy.concatenate([a[:,:,:1], numpy.zeros_like(a), a[:,:,1:]], axis=2), [(4,2,4)])
_check('inflate-constant', lambda f: evaluable._inflate(f,dofmap=[0,3],length=4,axis=2), lambda a: numpy.concatenate([a[:,:,:1], numpy.zeros_like(a), a[:,:,1:]], axis=2), [(4,2,4)])
_check('choose', lambda a, b, c: evaluable.Choose(evaluable.appendaxes(evaluable.Int(a)%2, (3,3)), [b,c]), lambda a, b, c: numpy.choose(a[:,_,_].astype(int)%2, [b,c]), [(), (3,3), (3,3)])
_check('slice', lambda a: evaluable.asarray(a)[:,::2], lambda a: a[:,::2], [(5,3)])

_polyval_mask = lambda shape, ndim: 1 if ndim == 0 else numpy.array([sum(i[-ndim:]) < shape[-1] for i in numpy.ndindex(shape)], dtype=int).reshape(shape)
_polyval_desired = lambda c, x: sum(c[(...,*i)]*(x[(slice(None),*[None]*(c.ndim-x.shape[1]))]**i).prod(-1) for i in itertools.product(*[range(c.shape[-1])]*x.shape[1]) if sum(i) < c.shape[-1])
_check('polyval_1d_p0', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(1,)], pass_geom=True, ndim=1)
_check('polyval_1d_p1', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(2,)], pass_geom=True, ndim=1)
_check('polyval_1d_p2', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(3,)], pass_geom=True, ndim=1)
_check('polyval_2d_p0', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(1,1)], pass_geom=True, ndim=2)
_check('polyval_2d_p1', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(2,2)], pass_geom=True, ndim=2)
_check('polyval_2d_p2', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(3,3)], pass_geom=True, ndim=2)
_check('polyval_2d_p1_23', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), evaluable.asarray(x)), _polyval_desired, [(2,3,2,2)], pass_geom=True, ndim=2)

class blocks(TestCase):

  def setUp(self):
    super().setUp()
    _builtin_warnings.simplefilter('ignore', evaluable.ExpensiveEvaluationWarning)

  def test_multiply_equal(self):
    ((i,), f), = evaluable.multiply(evaluable._inflate([1,2], dofmap=[0,2], length=3, axis=0), evaluable._inflate([3,4], dofmap=[0,2], length=3, axis=0)).blocks
    self.assertAllEqual(i.eval(), [0,2])
    self.assertAllEqual(f.eval(), [1*3,2*4])

  def test_multiply_embedded(self):
    ((i,), f), = evaluable.multiply([1,2,3], evaluable._inflate([4,5], dofmap=[0,2], length=3, axis=0)).blocks
    self.assertAllEqual(i.eval(), [0,2])
    self.assertAllEqual(f.eval(), [1*4,3*5])

  def test_multiply_overlapping(self):
    ((i,), f), = evaluable.multiply(evaluable._inflate([1,2], dofmap=[0,1], length=3, axis=0), evaluable._inflate([3,4], dofmap=[1,2], length=3, axis=0)).blocks
    self.assertAllEqual(i.eval(), [1])
    self.assertAllEqual(f.eval(), [2*3])

  def test_multiply_disjoint(self):
    blocks = evaluable.multiply(evaluable._inflate([1,2], dofmap=[0,2], length=4, axis=0), evaluable._inflate([3,4], dofmap=[1,3], length=4, axis=0)).blocks
    self.assertEqual(blocks, ())

  def test_multiply_overlap(self):
    ((i,), f), = evaluable.multiply(evaluable._inflate([1,2], dofmap=evaluable.Guard([0,1]), length=3, axis=0), evaluable._inflate([3,4], dofmap=evaluable.Guard([1,2]), length=3, axis=0)).blocks
    self.assertAllEqual(i.eval(), [1])
    self.assertAllEqual(f.eval(), [2*3])

  def test_takediag(self):
    ((i,), f), = evaluable.takediag([[1,2,3],[4,5,6],[7,8,9]]).blocks
    self.assertAllEqual(i.eval(), [0,1,2])
    self.assertAllEqual(f.eval(), [1,5,9])

  def test_takediag_embedded_axis(self):
    ((i,), f), = evaluable.takediag(evaluable._inflate([[1,2,3],[4,5,6]], dofmap=[0,2], length=3, axis=0)).blocks
    self.assertAllEqual(i.eval(), [0,2])
    self.assertAllEqual(f.eval(), [1,6])

  def test_takediag_embedded_rmaxis(self):
    ((i,), f), = evaluable.takediag(evaluable._inflate([[1,2],[3,4],[5,6]], dofmap=[0,2], length=3, axis=1)).blocks
    self.assertAllEqual(i.eval(), [0,2])
    self.assertAllEqual(f.eval(), [1,6])

  def test_takediag_overlapping(self):
    ((i,), f), = evaluable.takediag(evaluable._inflate(evaluable._inflate([[1,2],[3,4]], dofmap=[0,1], length=3, axis=0), dofmap=[1,2], length=3, axis=1)).blocks
    self.assertAllEqual(i.eval(), [1])
    self.assertAllEqual(f.eval(), [3])

  def test_takediag_disjoint(self):
    blocks = evaluable.takediag(evaluable._inflate(evaluable._inflate([[1,2],[3,4]], dofmap=[0,2], length=4, axis=0), dofmap=[1,3], length=4, axis=1)).blocks
    self.assertEqual(blocks, ())

  def test_takediag_overlap(self):
    ((i,), f), = evaluable.takediag(evaluable._inflate(evaluable._inflate([[1,2],[3,4]], dofmap=evaluable.Guard([0,1]), length=3, axis=0), dofmap=evaluable.Guard([1,2]), length=3, axis=1)).blocks
    self.assertAllEqual(i.eval(), [1])
    self.assertAllEqual(f.eval(), [3])

class commutativity(TestCase):

  def setUp(self):
    super().setUp()
    numpy.random.seed(0)
    self.A = evaluable.asarray(numpy.random.uniform(size=[2,3]))
    self.B = evaluable.asarray(numpy.random.uniform(size=[2,3]))

  def test_add(self):
    self.assertEqual(evaluable.add(self.A, self.B), evaluable.add(self.B, self.A))

  def test_multiply(self):
    self.assertEqual(evaluable.multiply(self.A, self.B), evaluable.multiply(self.B, self.A))

  def test_dot(self):
    self.assertEqual(evaluable.dot(self.A, self.B, axes=[0]), evaluable.dot(self.B, self.A, axes=[0]))

  def test_combined(self):
    self.assertEqual(evaluable.add(self.A, self.B) * evaluable.dot(self.A, self.B, axes=[0]), evaluable.dot(self.B, self.A, axes=[0]) * evaluable.add(self.B, self.A))

class sampled(TestCase):

  def test_match(self):
    f = evaluable.Sampled(numpy.array([[1,2], [3,4]]), numpy.array([[1,2], [3,4]]))
    self.assertAllEqual(f.eval(), numpy.eye(2))

  def test_no_match(self):
    f = evaluable.Sampled(numpy.array([[1,2], [3,4]]), numpy.array([[3,4], [1,2]]))
    with self.assertRaises(Exception):
      f.eval()

@parametrize
class piecewise(TestCase):

  def setUp(self):
    super().setUp()
    self.domain, self.geom = mesh.rectilinear([1])
    x, = self.geom
    if self.partition:
      left, mid, right = function.partition(x, .2, .8)
      self.f = left + function.sin(x) * mid + x**2 * right
    else:
      self.f = function.piecewise(x, [.2,.8], 1, function.sin(x), x**2)

  def test_evalf(self):
    f_ = self.domain.sample('uniform', 4).eval(self.f) # x=.125, .375, .625, .875
    assert numpy.equal(f_, [1, numpy.sin(.375), numpy.sin(.625), .875**2]).all()

  def test_deriv(self):
    g_ = self.domain.sample('uniform', 4).eval(function.grad(self.f, self.geom)) # x=.125, .375, .625, .875
    assert numpy.equal(g_, [[0], [numpy.cos(.375)], [numpy.cos(.625)], [2*.875]]).all()

piecewise(partition=False)
piecewise(partition=True)

class elemwise(TestCase):

  def setUp(self):
    super().setUp()
    self.domain, geom = mesh.rectilinear([5])
    self.index = self.domain.f_index.prepare_eval(ndims=self.domain.ndims)
    self.data = tuple(map(types.frozenarray, (
      numpy.arange(1, dtype=float).reshape(1,1),
      numpy.arange(2, dtype=float).reshape(1,2),
      numpy.arange(3, dtype=float).reshape(3,1),
      numpy.arange(4, dtype=float).reshape(2,2),
      numpy.arange(6, dtype=float).reshape(3,2),
    )))
    self.func = evaluable.Elemwise(self.data, self.index, float)

  def test_evalf(self):
    for i, (trans, points) in enumerate(zip(self.domain.transforms, self.domain.sample('gauss', 1).points)):
      with self.subTest(i=i):
        numpy.testing.assert_array_almost_equal(self.func.eval(_transforms=(trans,), _points=points), self.data[i])

  def test_shape(self):
    for i, (trans, points) in enumerate(zip(self.domain.transforms, self.domain.sample('gauss', 1).points)):
      with self.subTest(i=i):
        self.assertEqual(self.func.size.eval(_transforms=(trans,), _points=points), self.data[i].size)

  def test_derivative(self):
    self.assertTrue(evaluable.iszero(evaluable.localgradient(self.func, self.domain.ndims).simplified))

  def test_shape_derivative(self):
    self.assertEqual(evaluable.localgradient(self.func, self.domain.ndims).shape, self.func.shape+(self.domain.ndims,))

class jacobian(TestCase):

  def setUp(self):
    super().setUp()
    self.domain, self.geom = mesh.unitsquare(1, 'square')
    self.basis = self.domain.basis('std', degree=1)
    arg = function.Argument('dofs', [4])
    self.v = self.basis.dot(arg)
    self.X = (self.geom[numpy.newaxis,:] * [[0,1],[-self.v,0]]).sum(-1) # X_i = <x_1, -2 x_0>_i
    self.J = function.J(self.X, 2)
    self.dJ = function.derivative(self.J, arg)

  def test_shape(self):
    self.assertEqual(self.J.shape, ())
    self.assertEqual(self.dJ.shape, (4,))

  def test_value(self):
    values = self.domain.sample('uniform', 2).eval(self.J, dofs=[2]*4)
    numpy.testing.assert_almost_equal(values, [2]*4)
    values1, values2 = self.domain.sample('uniform', 2).eval([self.J,
      self.v + self.v.grad(self.geom)[0] * self.geom[0]], dofs=[1,2,3,10])
    numpy.testing.assert_almost_equal(values1, values2)

  def test_derivative(self):
    values1, values2 = self.domain.sample('uniform', 2).eval([self.dJ,
      self.basis + self.basis.grad(self.geom)[:,0] * self.geom[0]], dofs=[1,2,3,10])
    numpy.testing.assert_almost_equal(values1, values2)

  def test_zeroderivative(self):
    otherarg = function.Argument('otherdofs', (10,))
    values = self.domain.sample('uniform', 2).eval(function.derivative(self.dJ, otherarg))
    self.assertEqual(values.shape[1:], self.dJ.shape + otherarg.shape)
    self.assertAllEqual(values, 0)

class grad(TestCase):

  def assertEvalAlmostEqual(self, topo, factual, fdesired):
    actual, desired = topo.sample('uniform', 2).eval([function.asarray(factual), function.asarray(fdesired)])
    self.assertAllAlmostEqual(actual, desired)

  def test_0d(self):
    domain, (x,) = mesh.rectilinear([1])
    self.assertEvalAlmostEqual(domain, function.grad(x**2, x), 2*x)

  def test_1d(self):
    domain, x = mesh.rectilinear([1]*2)
    self.assertEvalAlmostEqual(domain, function.grad([x[0]**2, x[1]**2], x), [[2*x[0], 0], [0, 2*x[1]]])

  def test_2d(self):
    domain, x = mesh.rectilinear([1]*4)
    x = function.unravel(x, 0, (2, 2))
    self.assertEvalAlmostEqual(domain, function.grad(x, x), numpy.eye(4, 4).reshape(2, 2, 2, 2))

  def test_3d(self):
    domain, x = mesh.rectilinear([1]*4)
    x = function.unravel(function.unravel(x, 0, (2, 2)), 0, (2, 1))
    self.assertEvalAlmostEqual(domain, function.grad(x, x), numpy.eye(4, 4).reshape(2, 1, 2, 2, 1, 2))

class normal(TestCase):

  def assertEvalAlmostEqual(self, topo, factual, fdesired):
    actual, desired = topo.sample('uniform', 2).eval([function.asarray(factual), function.asarray(fdesired)])
    self.assertAllAlmostEqual(actual, desired)

  def test_0d(self):
    domain, (x,) = mesh.rectilinear([1])
    self.assertEvalAlmostEqual(domain.boundary['right'], function.normal(x), 1)
    self.assertEvalAlmostEqual(domain.boundary['left'], function.normal(x), -1)

  def test_1d(self):
    domain, x = mesh.rectilinear([1]*2)
    for bnd, n in ('right', [1, 0]), ('left', [-1, 0]), ('top', [0, 1]), ('bottom', [0, -1]):
      self.assertEvalAlmostEqual(domain.boundary[bnd], function.normal(x), n)

  def test_2d(self):
    domain, x = mesh.rectilinear([1]*2)
    x = function.unravel(x, 0, [2, 1])
    for bnd, n in ('right', [1, 0]), ('left', [-1, 0]), ('top', [0, 1]), ('bottom', [0, -1]):
      self.assertEvalAlmostEqual(domain.boundary[bnd], function.normal(x), numpy.array(n)[:,_])

  def test_3d(self):
    domain, x = mesh.rectilinear([1]*2)
    x = function.unravel(function.unravel(x, 0, [2, 1]), 0, [1, 2])
    for bnd, n in ('right', [1, 0]), ('left', [-1, 0]), ('top', [0, 1]), ('bottom', [0, -1]):
      self.assertEvalAlmostEqual(domain.boundary[bnd], function.normal(x), numpy.array(n)[_,:,_])

class asciitree(TestCase):

  @unittest.skipIf(sys.version_info < (3, 6), 'test requires dicts maintaining insertion order')
  def test_plain(self):
    f = evaluable.Sin(evaluable.Inflate(1, 1, 2)**evaluable.Diagonalize(evaluable.Argument('arg', (2,))))
    self.assertEqual(f.asciitree(),
                     '%9 = Sin(a2,a2)\n'
                     ': %8 = Power(a2,a2)\n'
                     '  : %7 = Transpose(i2,s2)\n'
                     '  | : %6 = InsertAxis(s2,i2)\n'
                     '  |   : %4 = Inflate(s2)\n'
                     '  |   | : %1 = Constant()\n'
                     '  |   | : %1\n'
                     '  |   | : %2 = Constant()\n'
                     '  |   : %2\n'
                     '  : %5 = Diagonalize(d2,d2)\n'
                     '    : %3 = Argument(a2)\n'
                     '      : %0 = Evaluable')

  @unittest.skipIf(sys.version_info < (3, 6), 'test requires dicts maintaining insertion order')
  def test_rich(self):
    f = evaluable.Sin(evaluable.Inflate(1, 1, 2)**evaluable.Diagonalize(evaluable.Argument('arg', (2,))))
    self.assertEqual(f.asciitree(richoutput=True),
                     '%9 = Sin(a2,a2)\n'
                     '└ %8 = Power(a2,a2)\n'
                     '  ├ %7 = Transpose(i2,s2)\n'
                     '  │ └ %6 = InsertAxis(s2,i2)\n'
                     '  │   ├ %4 = Inflate(s2)\n'
                     '  │   │ ├ %1 = Constant()\n'
                     '  │   │ ├ %1\n'
                     '  │   │ └ %2 = Constant()\n'
                     '  │   └ %2\n'
                     '  └ %5 = Diagonalize(d2,d2)\n'
                     '    └ %3 = Argument(a2)\n'
                     '      └ %0 = Evaluable')

class simplify(TestCase):

  def test_multiply_transpose(self):
    dummy = evaluable.Argument('dummy', shape=[2,2,2], dtype=float)
    f = evaluable.multiply(dummy,
          evaluable.Transpose(evaluable.multiply(dummy,
            evaluable.Transpose(dummy, (2,0,1))), (2,0,1)))
    # The test below is not only to verify that no simplifications are
    # performed, but also to make sure that simplified does not get stuck in a
    # circular dependence. This used to be the case prior to adding the
    # isinstance(other_trans, Transpose) restriction in Transpose._multiply.
    self.assertEqual(f.simplified, f)

max_stars_repo_name: joostvanzwieten/nutils

max_issues_repo_name: joostvanzwieten/nutils

max_forks_repo_name: joostvanzwieten/nutils