text:
from inspect import signature
from collections import namedtuple

import jax
import jax.numpy as jnp
import numpy as onp
import haiku as hk
from gym.spaces import Space, Discrete

from ..utils import safe_sample, batch_to_single, default_preprocessor
from ..proba_dists import ProbaDist
from .base_func import BaseFunc, ExampleData, Inputs, ArgsType1, ArgsType2, ModelTypes

__all__ = (
    'TransitionModel',
)

class TransitionModel(BaseFunc):
    r"""

    A deterministic transition function :math:`s'_\theta(s,a)`.

    Parameters
    ----------
    func : function

        A Haiku-style function that specifies the forward pass. The function signature must be the
        same as the example below.

    env : gym.Env

        The gym-style environment. This is used to validate the input/output structure of ``func``.

    observation_preprocessor : function, optional

        Turns a single observation into a batch of observations in a form that is convenient for
        feeding into :code:`func`. If left unspecified, this defaults to
        :attr:`proba_dist.preprocess_variate <coax.proba_dists.ProbaDist.preprocess_variate>`. The
        reason why the default is not :func:`coax.utils.default_preprocessor` is that we prefer
        consistence with :class:`coax.StochasticTransitionModel`.

    observation_postprocessor : function, optional

        Takes a batch of generated observations and makes sure that they are that are compatible
        with the original :code:`observation_space`. If left unspecified, this defaults to
        :attr:`proba_dist.postprocess_variate <coax.proba_dists.ProbaDist.postprocess_variate>`.

    action_preprocessor : function, optional

        Turns a single action into a batch of actions in a form that is convenient for feeding into
        :code:`func`. If left unspecified, this defaults
        :func:`default_preprocessor(env.action_space) <coax.utils.default_preprocessor>`.

    random_seed : int, optional

        Seed for pseudo-random number generators.

    """
    def __init__(
            self, func, env, observation_preprocessor=None, observation_postprocessor=None,
            action_preprocessor=None, random_seed=None):

        self.observation_preprocessor = observation_preprocessor
        self.observation_postprocessor = observation_postprocessor
        self.action_preprocessor = action_preprocessor

        # defaults
        if self.observation_preprocessor is None:
            self.observation_preprocessor = ProbaDist(env.observation_space).preprocess_variate
        if self.observation_postprocessor is None:
            self.observation_postprocessor = ProbaDist(env.observation_space).postprocess_variate
        if self.action_preprocessor is None:
            self.action_preprocessor = default_preprocessor(env.action_space)

        super().__init__(
            func,
            observation_space=env.observation_space,
            action_space=env.action_space,
            random_seed=random_seed)

    def __call__(self, s, a=None):
        r"""

        Evaluate the state-action function on a state observation :math:`s` or
        on a state-action pair :math:`(s, a)`.

        Parameters
        ----------
        s : state observation

            A single state observation :math:`s`.

        a : action

            A single action :math:`a`.

        Returns
        -------
        q_sa or q_s : ndarray

            Depending on whether :code:`a` is provided, this either returns a scalar representing
            :math:`q(s,a)\in\mathbb{R}` or a vector representing :math:`q(s,.)\in\mathbb{R}^n`,
            where :math:`n` is the number of discrete actions. Naturally, this only applies for
            discrete action spaces.

        """
        S = self.observation_preprocessor(self.rng, s)
        if a is None:
            S_next, _ = self.function_type2(self.params, self.function_state, self.rng, S, False)
            S_next = batch_to_single(S_next)  # (batch, num_actions, *) -> (num_actions, *)
            n = self.action_space.n
            s_next = [self.observation_postprocessor(self.rng, S_next, index=i) for i in range(n)]
        else:
            A = self.action_preprocessor(self.rng, a)
            S_next, _ = self.function_type1(self.params, self.function_state, self.rng, S, A, False)
            s_next = self.observation_postprocessor(self.rng, S_next)
        return s_next

    @property
    def function_type1(self):
        r"""

        Same as :attr:`function`, except that it ensures a type-1 function signature, regardless of
        the underlying :attr:`modeltype`.

        """
        if self.modeltype == 1:
            return self.function

        assert isinstance(self.action_space, Discrete)

        def project(A):
            assert A.ndim == 2, f"bad shape: {A.shape}"
            assert A.shape[1] == self.action_space.n, f"bad shape: {A.shape}"
            def func(leaf):  # noqa: E306
                assert isinstance(leaf, jnp.ndarray), f"leaf must be ndarray, got: {type(leaf)}"
                assert leaf.ndim >= 2, f"bad shape: {leaf.shape}"
                assert leaf.shape[0] == A.shape[0], \
                    f"batch_size (axis=0) mismatch: leaf.shape: {leaf.shape}, A.shape: {A.shape}"
                assert leaf.shape[1] == A.shape[1], \
                    f"num_actions (axis=1) mismatch: leaf.shape: {leaf.shape}, A.shape: {A.shape}"
                return jax.vmap(jnp.dot)(jnp.moveaxis(leaf, 1, -1), A)
            return func

        def type1_func(type2_params, type2_state, rng, S, A, is_training):
            S_next, state_new = self.function(type2_params, type2_state, rng, S, is_training)
            S_next = jax.tree_map(project(A), S_next)
            return S_next, state_new

        return type1_func

    @property
    def function_type2(self):
        r"""

        Same as :attr:`function`, except that it ensures a type-2 function signature, regardless of
        the underlying :attr:`modeltype`.

        """
        if self.modeltype == 2:
            return self.function

        if not isinstance(self.action_space, Discrete):
            raise ValueError(
                "input 'A' is required for type-1 dynamics model when action space is non-Discrete")

        n = self.action_space.n

        def reshape(leaf):
            # reshape from (batch * num_actions, *shape) -> (batch, *shape, num_actions)
            assert isinstance(leaf, jnp.ndarray), f"all leaves must be ndarray, got: {type(leaf)}"
            assert leaf.ndim >= 1, f"bad shape: {leaf.shape}"
            assert leaf.shape[0] % n == 0, \
                f"first axis size must be a multiple of num_actions, got shape: {leaf.shape}"
            leaf = jnp.reshape(leaf, (-1, n, *leaf.shape[1:]))  # (batch, num_actions, *shape)
            return leaf

        def type2_func(type1_params, type1_state, rng, S, is_training):
            rngs = hk.PRNGSequence(rng)
            batch_size = jax.tree_leaves(S)[0].shape[0]

            # example: let S = [7, 2, 5, 8] and num_actions = 3, then
            # S_rep = [7, 7, 7, 2, 2, 2, 5, 5, 5, 8, 8, 8]  # repeated
            # A_rep = [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]  # tiled
            S_rep = jax.tree_map(lambda x: jnp.repeat(x, n, axis=0), S)
            A_rep = jnp.tile(jnp.arange(n), batch_size)
            A_rep = self.action_preprocessor(next(rngs), A_rep)  # one-hot encoding

            # evaluate on replicas => output shape: (batch * num_actions, *shape)
            S_next_rep, state_new = self.function(
                type1_params, type1_state, next(rngs), S_rep, A_rep, is_training)
            S_next = jax.tree_map(reshape, S_next_rep)

            return S_next, state_new

        return type2_func

    @property
    def modeltype(self):
        r"""

        Specifier for how the transition function is modeled, i.e.

        .. math::

            (s,a)   &\mapsto s'(s,a)    &\qquad (\text{modeltype} &= 1) \\
            s       &\mapsto s'(s,.)    &\qquad (\text{modeltype} &= 2)

        Note that modeltype=2 is only well-defined if the action space is :class:`Discrete
        <gym.spaces.Discrete>`. Namely, :math:`n` is the number of discrete actions.

        """
        return self._modeltype

    @classmethod
    def example_data(
            cls, env, observation_preprocessor=None, action_preprocessor=None,
            batch_size=1, random_seed=None):

        if not isinstance(env.observation_space, Space):
            raise TypeError(
                "env.observation_space must be derived from gym.Space, "
                f"got: {type(env.observation_space)}")

        if not isinstance(env.action_space, Space):
            raise TypeError(
                f"env.action_space must be derived from gym.Space, got: {type(env.action_space)}")

        if observation_preprocessor is None:
            observation_preprocessor = ProbaDist(env.observation_space).preprocess_variate

        if action_preprocessor is None:
            action_preprocessor = default_preprocessor(env.action_space)

        rnd = onp.random.RandomState(random_seed)
        rngs = hk.PRNGSequence(rnd.randint(jnp.iinfo('int32').max))

        # input: state observations
        S = [safe_sample(env.observation_space, rnd) for _ in range(batch_size)]
        S = [observation_preprocessor(next(rngs), s) for s in S]
        S = jax.tree_multimap(lambda *x: jnp.concatenate(x, axis=0), *S)

        # input: actions
        A = [safe_sample(env.action_space, rnd) for _ in range(batch_size)]
        A = [action_preprocessor(next(rngs), a) for a in A]
        A = jax.tree_multimap(lambda *x: jnp.concatenate(x, axis=0), *A)

        # output: type1
        S_next_type1 = jax.tree_map(lambda x: jnp.asarray(rnd.randn(batch_size, *x.shape[1:])), S)
        q1_data = ExampleData(
            inputs=Inputs(args=ArgsType1(S=S, A=A, is_training=True), static_argnums=(2,)),
            output=S_next_type1)

        if not isinstance(env.action_space, Discrete):
            return ModelTypes(type1=q1_data, type2=None)

        # output: type2 (if actions are discrete)
        S_next_type2 = jax.tree_map(
            lambda x: jnp.asarray(rnd.randn(batch_size, env.action_space.n, *x.shape[1:])), S)
        q2_data = ExampleData(
            inputs=Inputs(args=ArgsType2(S=S, is_training=True), static_argnums=(1,)),
            output=S_next_type2)

        return ModelTypes(type1=q1_data, type2=q2_data)

    def _check_signature(self, func):
        sig_type1 = ('S', 'A', 'is_training')
        sig_type2 = ('S', 'is_training')
        sig = tuple(signature(func).parameters)

        if sig not in (sig_type1, sig_type2):
            sig = ', '.join(sig)
            alt = ' or func(S, is_training)' if isinstance(self.action_space, Discrete) else ''
            raise TypeError(
                f"func has bad signature; expected: func(S, A, is_training){alt}, got: func({sig})")

        if sig == sig_type2 and not isinstance(self.action_space, Discrete):
            raise TypeError("type-2 models are only well-defined for Discrete action spaces")

        Env = namedtuple('Env', ('observation_space', 'action_space'))
        example_data_per_modeltype = self.example_data(
            env=Env(self.observation_space, self.action_space),
            action_preprocessor=self.action_preprocessor,
            batch_size=1,
            random_seed=self.random_seed)

        if sig == sig_type1:
            self._modeltype = 1
            example_data = example_data_per_modeltype.type1
        else:
            self._modeltype = 2
            example_data = example_data_per_modeltype.type2

        return example_data

    def _check_output(self, actual, expected):
        expected_leaves, expected_structure = jax.tree_flatten(expected)
        actual_leaves, actual_structure = jax.tree_flatten(actual)
        assert all(isinstance(x, jnp.ndarray) for x in expected_leaves), "bad example_data"

        if actual_structure != expected_structure:
            raise TypeError(
                f"func has bad return tree_structure, expected: {expected_structure}, "
                f"got: {actual_structure}")

        if not all(isinstance(x, jnp.ndarray) for x in actual_leaves):
            bad_types = tuple(type(x) for x in actual_leaves if not isinstance(x, jnp.ndarray))
            raise TypeError(
                "all leaves of dist_params must be of type: jax.numpy.ndarray, "
                f"found leaves of type: {bad_types}")

        if not all(a.shape == b.shape for a, b in zip(actual_leaves, expected_leaves)):
            shapes_tree = jax.tree_multimap(
                lambda a, b: f"{a.shape} {'!=' if a.shape != b.shape else '=='} {b.shape}",
                actual, expected)
            raise TypeError(f"found leaves with unexpected shapes: {shapes_tree}")

max_stars_repo_name: coax-dev/coax

max_issues_repo_name: coax-dev/coax

max_forks_repo_name: coax-dev/coax

reason: import numpy,import jax


text:
import numpy as np

import chess.pgn
import chess

chess_dict = {
    'p':  [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    'n':  [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    'b':  [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    'r':  [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    'q':  [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
    'k':  [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
    'P':  [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
    'N':  [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
    'B':  [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
    'R':  [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
    'Q':  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
    'K':  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
    None: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
}

def board_to_matrix(board: chess.Board) -> np.array:
    matrix = [
        [
          None for _ in range(8)
        ] for _ in range(8)
    ]
    for i, square in enumerate(chess.SQUARES):
        rank = i // 8
        file = i % 8
        symbol = board.piece_at(square).symbol() if board.piece_at(square) else None
        matrix[rank][file] = chess_dict[symbol]
    return np.array([
        np.array([
            np.array(file, dtype=np.short) for file in rank
        ]) for rank in matrix
    ])

ext: py

reason: import numpy


text:
import os
from argparse import ArgumentParser

# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
# os.environ["CUDA_VISIBLE_DEVICES"]="0"
def build_parser():
    parser = ArgumentParser()
    parser.add_argument('--num_epochs', default=150, help="Number of training epochs (default: 150)", type=int)
    parser.add_argument('--batch_size', default=128, help="Batch size (default: 128)", type=int)
    parser.add_argument('--learning_rate', default=0.001, help="Learning rate for ADAM (default: 0.001)", type=float)
    parser.add_argument('--save_dir', default='tmp', help="checkpoint & summaries save dir name (default: tmp)")
    parser.add_argument('--gpu_num', default=0, help="CUDA visible device (default: 0)")
    parser.add_argument("--model_name", help="vggnet / vggnet2 / resnet / wide_resnet / inception", required=True)
    parser.add_argument("--augmentation_type", default="affine", help="none / affine / align (default: affine)")
    parser.add_argument("--resnet_layer_n", default=3, help="6n+2: {3, 5, 7, 9 ... 18} (default: 3)", type=int)
    parser.add_argument("--ignore_exist_model", default=False, help="Overwrite new model to exist model (default: false)", type=bool)
    parser.add_argument("--gpu_memory_fraction", default=0.3, help="If this value is 0.0, allow_growth option is on (default: 0.3)", type=float)

    return parser

parser = build_parser()
FLAGS = parser.parse_args()
print("\nParameters:")
for attr, value in sorted(vars(FLAGS).items()):
    print("{}={}".format(attr.upper(), value))
print("")
if FLAGS.save_dir == "tmp":
    print("========== save_dir is tmp!!! ==========")

os.environ["CUDA_VISIBLE_DEVICES"]=str(FLAGS.gpu_num)

import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
import vggnet, resnet, wide_resnet, inception_light
from solver import Solver
import time, datetime
import os, glob, shutil
import datagenerator
from utils import *

# tf.reset_default_graph()
SEED = 777
tf.set_random_seed(SEED)
np.random.seed(SEED)

mnist = input_data.read_data_sets('MNIST_data/', one_hot=True)
batch_size = FLAGS.batch_size
epoch_n = FLAGS.num_epochs
learning_rate = 0.001
N = mnist.train.num_examples

model = get_model(name=FLAGS.model_name, learning_rate=FLAGS.learning_rate, SEED=SEED, resnet_layer_n=FLAGS.resnet_layer_n)

SUMMARY_DIR = "./tmp/{}".format(FLAGS.save_dir)
CHECKPOINT_DIR = "./checkpoint/{}".format(FLAGS.save_dir)
if not os.path.exists(SUMMARY_DIR):
    os.makedirs(SUMMARY_DIR)
if not os.path.exists(CHECKPOINT_DIR):
    os.makedirs(CHECKPOINT_DIR)

def train():
    config = tf.ConfigProto()
    if FLAGS.gpu_memory_fraction > 0.0:
        if FLAGS.gpu_memory_fraction < 1.0:
            config.gpu_options.per_process_gpu_memory_fraction = FLAGS.gpu_memory_fraction
    else:
        config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    solver = Solver(sess, model)

    # saver
    saver = tf.train.Saver()
    if FLAGS.save_dir != "tmp":
        checkpoint = tf.train.get_checkpoint_state(CHECKPOINT_DIR)
        if checkpoint and checkpoint.model_checkpoint_path:
            print("model exist")
            if FLAGS.ignore_exist_model == False:
            	exit()
            # try:
            #     saver.restore(sess, checkpoint.model_checkpoint_path)
            #     print("Successfully loaded:", checkpoint.model_checkpoint_path)
            #     exit()
            # except:
            #     print("Error on loading old network weights")
        # else:
        #     print("Could not find old network weights")

    # reset summaries for our GPU
    for f in glob.glob(SUMMARY_DIR+"/*"):
        shutil.rmtree(f)

    # writers
    train_writer = tf.summary.FileWriter(SUMMARY_DIR + '/train', sess.graph)
    valid_writer = tf.summary.FileWriter(SUMMARY_DIR + '/valid')
    test_writer = tf.summary.FileWriter(SUMMARY_DIR + '/test')

    datagen = datagenerator.Generator(FLAGS.augmentation_type, mnist)

    max_train_acc = 0
    max_valid_acc = 0

    # train
    sess.run(tf.global_variables_initializer())

    start_time = time.time()
    print("========== training start ==========")
    print "batch_size=%d"%batch_size

    for epoch in range(epoch_n):
        # for _ in range(N // batch_size):
        i = 0
        for x, y in datagen.generate(batch_size=batch_size):
            print "start epoch No.%d, batch No.%d"%(epoch,i)
            _, train_loss = solver.train(x, y)
            i += 1

        train_loss, train_acc = solver.evaluate(mnist.train.images, mnist.train.labels, batch_size, writer=train_writer, step=epoch+1)
        valid_loss, valid_acc = solver.evaluate(mnist.validation.images, mnist.validation.labels, batch_size, writer=valid_writer, step=epoch+1)
        test_loss, test_acc = solver.evaluate(mnist.test.images, mnist.test.labels, batch_size, writer=test_writer, step=epoch+1)
        line = "[{}/{}] train: {:.4f}, {:.3%} / valid: {:.4f}, {:.2%} / test: {:.4f}, {:.2%}". \
        format(epoch+1, epoch_n, train_loss, train_acc, valid_loss, valid_acc, test_loss, test_acc)
        print line,
        elapsed_time = datetime.timedelta(seconds=int(time.time()-start_time))
        print("[{}]".format(elapsed_time))

        if train_acc > max_train_acc:
            max_train_acc = train_acc
            train_line = line
            saver.save(sess, CHECKPOINT_DIR + "/model-train")
        if valid_acc > max_valid_acc:
            max_valid_acc = valid_acc
            valid_line = line
            saver.save(sess, CHECKPOINT_DIR + "/model-valid")

    saver.save(sess, CHECKPOINT_DIR + "/model-final")

    print("[train max] {}".format(train_line))
    print("[valid max] {}".format(valid_line))
    print("[  final  ] {}".format(line))
    elapsed_time = time.time() - start_time
    formatted = datetime.timedelta(seconds=int(elapsed_time))
    print("=== training time elapsed: {}s ===".format(formatted))

# validation

# import validation
# validation.wrong_answer_check(solver, mnist.test.images, mnist.test.labels, filename="wrong_answers.png")

if __name__ == '__main__':
    train()


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

# Copyright 2016 Timothy Dozat
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections import Counter

import numpy as np

class KMeans(object):
    """
    Cluster sentences by their lengths
    """

    def __init__(self, k, len_cntr):
        """Do clustering

        Parameters
        ----------
        k : int
            number of clusters
        len_cntr : Counter
            length counter
        """
        # Error checking
        if len(len_cntr) < k:
            raise ValueError('Trying to sort %d data points into %d buckets' % (len(len_cntr), k))

        # Initialize variables
        self._k = k
        self._len_cntr = len_cntr
        self._lengths = sorted(self._len_cntr.keys())
        self._splits = []
        self._split2len_idx = {}
        self._len2split_idx = {}
        self._split_cntr = Counter()

        # Initialize the splits evenly
        lengths = []
        unique_length = []
        for length, count in list(self._len_cntr.items()):
            lengths.extend([length] * count)
            unique_length.append(length)
        lengths.sort()
        unique_length.sort()
        self._splits = [np.max(split) for split in np.array_split(lengths, self._k)]
        """Length of sentences"""

        i = len(self._splits) - 1
        while i > 0:
            while self._splits[i - 1] >= self._splits[i]:
                index = unique_length.index(self._splits[i - 1])
                if index == 0:
                    break
                self._splits[i - 1] = unique_length[index - 1]
            i -= 1

        unique_length.reverse()
        i = 1
        while i < len(self._splits) - 1:
            while self._splits[i] <= self._splits[i - 1]:
                index = unique_length.index(self._splits[i])
                if index == 0:
                    break
                self._splits[i] = unique_length[index - 1]
            i += 1

        # Reindex everything
        split_idx = 0
        split = self._splits[split_idx]
        for len_idx, length in enumerate(self._lengths):
            count = self._len_cntr[length]
            self._split_cntr[split] += count
            if length == split:
                self._split2len_idx[split] = len_idx
                split_idx += 1
                if split_idx < len(self._splits):
                    split = self._splits[split_idx]
                    self._split_cntr[split] = 0
            elif length > split:
                raise IndexError()

        # Iterate
        old_splits = None
        # print('0) Initial splits: %s; Initial mass: %d' % (self._splits, self.get_mass()))
        i = 0
        while self._splits != old_splits and i < 1000:
            old_splits = list(self._splits)
            self._recenter()
            i += 1
        # print('%d) Final splits: %s; Final mass: %d' % (i, self._splits, self.get_mass()))

        self._reindex()
        return

    def _recenter(self):
        """
        one iteration of k-means
        """
        for split_idx in range(len(self._splits)):
            split = self._splits[split_idx]
            len_idx = self._split2len_idx[split]
            if split == self._splits[-1]:
                continue
            right_split = self._splits[split_idx + 1]

            # Try shifting the centroid to the left
            if len_idx > 0 and self._lengths[len_idx - 1] not in self._split_cntr:
                new_split = self._lengths[len_idx - 1]
                left_delta = self._len_cntr[split] * (right_split - new_split) - self._split_cntr[split] * (
                        split - new_split)
                if left_delta < 0:
                    self._splits[split_idx] = new_split
                    self._split2len_idx[new_split] = len_idx - 1
                    del self._split2len_idx[split]
                    self._split_cntr[split] -= self._len_cntr[split]
                    self._split_cntr[right_split] += self._len_cntr[split]
                    self._split_cntr[new_split] = self._split_cntr[split]
                    del self._split_cntr[split]

            # Try shifting the centroid to the right
            elif len_idx < len(self._lengths) - 2 and self._lengths[len_idx + 1] not in self._split_cntr:
                new_split = self._lengths[len_idx + 1]
                right_delta = self._split_cntr[split] * (new_split - split) - self._len_cntr[split] * (
                        new_split - split)
                if right_delta <= 0:
                    self._splits[split_idx] = new_split
                    self._split2len_idx[new_split] = len_idx + 1
                    del self._split2len_idx[split]
                    self._split_cntr[split] += self._len_cntr[split]
                    self._split_cntr[right_split] -= self._len_cntr[split]
                    self._split_cntr[new_split] = self._split_cntr[split]
                    del self._split_cntr[split]

    def _reindex(self):
        """
        Index every sentence into a cluster
        """
        self._len2split_idx = {}
        last_split = -1
        for split_idx, split in enumerate(self._splits):
            self._len2split_idx.update(
                dict(list(zip(list(range(last_split + 1, split)), [split_idx] * (split - (last_split + 1))))))

    def __len__(self):
        return self._k

    def __iter__(self):
        return (split for split in self.splits)

    def __getitem__(self, key):
        return self._splits[key]

    @property
    def splits(self):
        """Get clusters

        Returns
        -------
        tuple
            (bucket, length) mapping
        """
        return self._splits

    @property
    def len2split_idx(self):
        """Get length to bucket mapping

        Returns
        -------
        tuple
             (length, bucket) mapping
        """
        return self._len2split_idx

def main():
    c = Counter({9: 3, 7: 2, 8: 2, 10: 2, 6: 2, 3: 2, 5: 2, 30: 1, 12: 1, 15: 1, 16: 1, 18: 1, 21: 1})
    k = 10
    print(KMeans(k, c).splits)

if __name__ == '__main__':
    main()


text:
""" In the first function, generate the normal distributed samples by summing up 12
 uniform distributed samples, as explained in the lecture
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
import math
import timeit

def sample_normal_distribution(a, b, sample):
    s = 0
    for i in range(12):
        t = np.random.uniform(a, b, sample)
        s += t
    return s/(np.sqrt(sample*(b-a)**2/12))

def sample_normal_twelve(mu, sigma):
    x = 0.5 * np.sum(np.random.uniform(-sigma, sigma, 12))
    return mu+x

def evaluate_sampling_dist(mu, sigma, n_samples, sample_function):
    n_bins = 100
    samples = []
    for i in range(n_samples):
        samples.append(sample_function(mu, sigma))
    print("%30s : mean = %.3f, std_dev = %.3f" % (sample_function.__name__, np.mean(samples), np.std(samples)))
    plt.figure()
    count, bins, ignored = plt.hist(samples, n_bins, density=True)
    plt.plot(bins, stats.norm(mu, sigma).pdf(bins), linewidth=2, color='r')
    plt.xlim([mu - 5*sigma, mu + 5*sigma])
    plt.title(sample_function.__name__)

def evaluage_sampling_time(mu, sigma, n_samples, sample_function):
    tic = timeit.default_timer()
    for i in range(n_samples):
        sample_function(mu, sigma)
    toc = timeit.default_timer()
    time_per_sample = (toc-tic)/n_samples*1e6
    print("%30s : %.3f us" % (sample_function.__name__, time_per_sample))

def sample_normal_rejection(mu, sigma):
    interval = 5*sigma
    max_density = stats.norm(mu, sigma).pdf(mu)

    # Rejection loop
    while True:
        x = np.random.uniform(mu-interval, mu+interval, 1)[0]
        y = np.random.uniform(0, max_density, 1)
        if y<= stats.norm(mu, sigma).pdf(x):
            break
    return x

def sample_normal_boxmuller(mu, sigma):
    u = np.random.uniform(0, 1, 2)

    x = math.cos(2*np.pi*u[0])*math.sqrt(-2*math.log(u[1]))
    return mu + sigma * x

if __name__ == "__main__":
    mu, sigma = 0, 1
    sample_functions =[
        sample_normal_twelve,
        sample_normal_rejection,
        sample_normal_boxmuller
    ]

    for fnc in sample_functions:
        evaluage_sampling_time(mu, sigma, 1000, fnc)

    n_samples = 10000
    print("evaluting sample distances with:")
    print(" mean :", mu)
    print("std_dev :", sigma)
    print("samples :", n_samples)

    for func in sample_functions:
       evaluate_sampling_dist(mu, sigma, n_samples, func)
    plt.show()

max_stars_repo_name: yubaoliu/AISLAM

max_issues_repo_name: yubaoliu/AISLAM

max_forks_repo_name: yubaoliu/AISLAM

reason: import numpy,from scipy


text:
# coding: utf-8

# In[2]:

import pandas as pd
import numpy as np
import cPickle
import keras.preprocessing.text as tokenizer
import gc
import string

# In[17]:

def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        return False

def convert_word_to_ix(data):
    result = []
    i=1
    x=set(wordtoix.keys())
    for sent in data:
        temp = []
        for w in sent:
            if w in x:
                temp.append(wordtoix[w])
            else:
                temp.append(1)
        temp.append(0)
        result.append(temp)
    return result

# In[36]:

pf_notes = pd.read_table('data/pratt_f.txt', sep='\n', header=None)[0]
pm_notes = pd.read_table('data/pratt_m.txt', sep='\n', header=None)[0]
tf_notes = pd.read_table('data/trinity_f.txt', sep='\n', header=None)[0]
tm_notes = pd.read_table('data/trinity_m.txt', sep='\n', header=None)[0]

notes = pf_notes.append(pm_notes, ignore_index=True).append(tf_notes, ignore_index=True).append(tm_notes, ignore_index=True)

# In[45]:

# stop_words = set([u'laughs', u'inaudible', u'laughter', u'pause', u'i', u'me', u'my', u'myself', u'we', u'our', u'ours', u'ourselves', u'you', u"you're", u"you've", u"you'll", u"you'd", u'your', u'yours', u'yourself', u'yourselves', u'he', u'him', u'his', u'himself', u'she', u"she's", u'her', u'hers', u'herself', u'it', u"it's", u'its', u'itself', u'they', u'them', u'their', u'theirs', u'themselves', u'what', u'which', u'who', u'whom', u'this', u'that', u"that'll", u'these', u'those', u'am', u'is', u'are', u'was', u'were', u'be', u'been', u'being', u'have', u'has', u'had', u'having', u'do', u'does', u'did', u'doing', u'a', u'an', u'the', u'and', u'but', u'if', u'or', u'because', u'as', u'until', u'while', u'of', u'at', u'by', u'for', u'with', u'about', u'against', u'between', u'into', u'through', u'during', u'before', u'after', u'above', u'below', u'to', u'from', u'up', u'down', u'in', u'out', u'on', u'off', u'over', u'under', u'again', u'further', u'then', u'once', u'here', u'there', u'when', u'where', u'why', u'how', u'all', u'any', u'both', u'each', u'few', u'more', u'most', u'other', u'some', u'such', u'no', u'nor', u'not', u'only', u'own', u'same', u'so', u'than', u'too', u'very', u's', u't', u'can', u'will', u'just', u'don', u"don't", u'should', u"should've", u'now', u'd', u'll', u'm', u'o', u're', u've', u'y', u'ain', u'aren', u"aren't", u'couldn', u"couldn't", u'didn', u"didn't", u'doesn', u"doesn't", u'hadn', u"hadn't", u'hasn', u"hasn't", u'haven', u"haven't", u'isn', u"isn't", u'ma', u'mightn', u"mightn't", u'mustn', u"mustn't", u'needn', u"needn't", u'shan', u"shan't", u'shouldn', u"shouldn't", u'wasn', u"wasn't", u'weren', u"weren't", u'won', u"won't", u'wouldn', u"wouldn't"])

sent = []
vocab = {}
x = set(vocab.keys())
for note in notes:
    note = note.replace('\xc2', '')
    note = note.replace('\xa0', '')
    note = note.replace('\xe2', '')
    note = note.replace('\x80', '')
    note = note.replace('\x91', '')

    temp = tokenizer.text_to_word_sequence(note, filters='!"#$%&()*+,-./:;<=>?@[\]^_`{|}~\n', lower=True, split=' ')
    temp = [j if not is_number(j) else '0' for j in temp]
    sent.append(temp)
    temp = set(temp)
    for word in temp:
        if word in x:
            vocab[word] += 1
        else:
            vocab[word] = 1

vocab = {x:y for x, y in vocab.items()}

ixtoword = {}
ixtoword[0] = 'END'
ixtoword[1] = 'UNK'
wordtoix = {}
wordtoix['END'] = 0
wordtoix['UNK'] = 1

ix = 2
for w in vocab:
    wordtoix[w] = ix
    ixtoword[ix] = w
    ix += 1

train_x = convert_word_to_ix(sent)

# In[48]:

cPickle.dump([vocab, wordtoix, ixtoword, train_x, sent], open('inputs_model.p', 'wb'), protocol=cPickle.HIGHEST_PROTOCOL)


text:
# https://atcoder.jp/contests/math-and-algorithm/tasks/math_and_algorithm_au

import numpy as np

MOD = 10 ** 9 + 7
n = int(input()) - 1
a = np.matrix([[2, 1], [1, 0]], dtype=np.uint64)
ans = np.identity(2, dtype=np.uint64)
while n > 0:
    if n % 2 != 0:
        ans = ans @ a
        ans %= MOD
    a = a @ a
    a %= MOD
    n >>= 1
print(int(ans[1, 0] + ans[1, 1]) % MOD)

max_stars_repo_name: silphire/training-with-books

max_issues_repo_name: silphire/training-with-books

max_forks_repo_name: silphire/training-with-books


text:
import numpy as np
from keras import backend as K

def categoricalCrossentropy(y_noisy, y_pred, the0, the1, the2, the3):
    '''
    Calculate the class-weighted categorical cross-entropy for the given
    predicted and true sets.

    y_true [in] The truth set to test against. This is a Tensor with a last
                dimension that contains a set of 1-of-N selections.
    y_pred [in] The predicted set to test against. This is a Tensor with a last
                dimension that contains a set of 1-of-N selections.
    returns     A Tensor function that will calculate the weighted categorical
                cross-entropy on the inputs.
    '''
    P = np.array([[1, the0, the0],
                  [the1,  1, (the0+the1)*.5],
                  [the1, (the0+the1)*.5, 1]])
    P = np.array([[1, 0, 0],
                  [0, 1, the0],
                  [the1, 0, 1]])
    if P is not None:
        # Wrap the loss weights in a tensor object.
        p = np.linalg.inv(P)
        p =  p / (np.sum(p))
        #assert (np.isclose(np.sum(p), 1) == True)
        theWeights =  K.constant(p, shape=p.shape)

        y_true = K.dot(y_noisy,theWeights)

    return  K.categorical_crossentropy(y_true, y_pred)

def noisy_loss(the0, the1, the2, the3):
    def loss(y_true, y_pred):
        return categoricalCrossentropy(y_true, y_pred,  the0, the1, the2, the3)
    return loss

max_stars_repo_name: BerlinRDT/roaddetection


text:
# -*- coding: utf-8 -*-
import sys
sys.path.insert(0,'..')

import numpy as np
np.random.seed(42)
import pytest
import sklearn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis

from ceml.sklearn import generate_counterfactual

def test_lda():
    # Load data
    X, y = load_iris(return_X_y=True)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=4242)

    # Create and fit model
    model = LinearDiscriminantAnalysis(store_covariance=True)
    model.fit(X_train, y_train)

    # Select data point for explaining its prediction
    x_orig = X_test[1:4][0,:]
    assert model.predict([x_orig]) == 2

    # Compute counterfactual
    features_whitelist = None

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", optimizer="mp", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 1, features_whitelist=features_whitelist, regularization="l1", optimizer="mp", return_as_dict=False)
    assert y_cf == 1
    assert model.predict(np.array([x_cf])) == 1

    cf = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", optimizer="mp", return_as_dict=True)
    assert cf["y_cf"] == 0
    assert model.predict(np.array([cf["x_cf"]])) == 0

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l2", optimizer="mp", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", C=1.0, optimizer="bfgs", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", C=1.0, optimizer="nelder-mead", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization=None, optimizer="bfgs", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization=None, optimizer="nelder-mead", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0

    features_whitelist = [0, 1, 2]
    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", optimizer="mp", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0
    assert all([True if i in features_whitelist else delta[i] <= 10e-5 for i in range(x_orig.shape[0])])

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l2", optimizer="mp", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0
    assert all([True if i in features_whitelist else delta[i] <= 10e-5 for i in range(x_orig.shape[0])])

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", C=1.0, optimizer="bfgs", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0
    assert all([True if i in features_whitelist else delta[i] == 0. for i in range(x_orig.shape[0])])

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization="l1", C=1.0, optimizer="nelder-mead", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0
    assert all([True if i in features_whitelist else delta[i] == 0. for i in range(x_orig.shape[0])])

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization=None, optimizer="bfgs", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0
    assert all([True if i in features_whitelist else delta[i] == 0. for i in range(x_orig.shape[0])])

    x_cf, y_cf, delta = generate_counterfactual(model, x_orig, 0, features_whitelist=features_whitelist, regularization=None, optimizer="nelder-mead", return_as_dict=False)
    assert y_cf == 0
    assert model.predict(np.array([x_cf])) == 0
    assert all([True if i in features_whitelist else delta[i] == 0. for i in range(x_orig.shape[0])])

    # Other stuff
    from ceml.sklearn import LdaCounterfactual
    with pytest.raises(TypeError):
        LdaCounterfactual(sklearn.linear_model.LogisticRegression())

    model = LinearDiscriminantAnalysis()
    model.fit(X_train, y_train)
    with pytest.raises(AttributeError):
        LdaCounterfactual(model)
    with pytest.raises(AttributeError):
        generate_counterfactual(model, x_orig, 0)

max_stars_repo_name: andreArtelt/ceml

max_issues_repo_name: andreArtelt/ceml

max_forks_repo_name: andreArtelt/ceml


text:
import numpy as np
from sklearn import linear_model
import decoders

# get k-fold splits
def get_test_train_splits(data, decisions, n_folds=5):
  """
  Returns a tuple of matched train sets and validation sets,
  rotating through the data.
  Used by multiple function wrappers as below.
  """
  fold_size = len(data)//n_folds

  training_X = [np.roll(data,fold_size*i, axis=0)[fold_size:] for i in range(n_folds)]
  val_X = [np.roll(data,fold_size*i, axis=0)[:fold_size] for i in range(n_folds)]

  training_Y = [np.roll(decisions,fold_size*i, axis=0)[fold_size:] for i in range(n_folds)]
  val_Y = [np.roll(decisions,fold_size*i, axis=0)[:fold_size] for i in range(n_folds)]

  return (training_X, training_Y), (val_X, val_Y)

# -----------------------------------------------------------------------------
# kfold val function for hardcoded logistic regression
# with customized C regularization parameter
def get_kfold_validation_score(data, decisions, C):
  '''
  gets mean validation accuracy with C reg without plotting or other shenanigans

  data: expects numpy array of shape (n_trials, n_units, n_bins)
  decisions: expects numpy array of shape (n_trials, 1)
    can be generalized to most categorical situations

  Output: returns mean validation accuracy of all kfold accuracies
  '''

  # use func from above to get splits
  (training_X, training_Y), (val_X, val_Y) = get_test_train_splits(data, decisions)

  scores = []

  for fold in range(5):
    training_Xi = training_X[fold]
    training_Yi = training_Y[fold]

    val_Xi = val_X[fold]
    val_Yi = val_Y[fold]

    # re-initialize each iteration
    logreg_model = linear_model.LogisticRegression(penalty='l2', solver = 'lbfgs',
                                                max_iter = 1000, C=C)
    # fit and score on training
    logreg_model.fit(training_Xi, training_Yi)
    logreg_model.score(training_Xi,training_Yi)

    # score on the validation data
    accuracy = logreg_model.score(val_Xi, val_Yi)
    scores.append(accuracy)

  return np.mean(scores)

# -----------------------------------------------------------------------------
# kfold training wrapper with helper function
def compute_accuracy(val_spikes, val_choices, model):
  predictions = model.predict(val_spikes)
  accuracy = np.sum(predictions == val_choices) / len(predictions)
  return accuracy

def kfold_train_wrapper(model_name, data, decisions,
                        n_folds, units, dropout, num_epochs, verbose):
  '''
  wrapper function to train and score on k-fold

  Inputs:
  ----------
  model_name: calls str model_name as class in decoders module
  data: expects numpy array of shape (n_trials, n_units, n_bins)
  decisions: expects numpy array of shape (n_trials, 1)
    can be generalized to most categorical situations

  Parameters:
  ----------
  n_folds: number of folds to use for k-fold cross validation
  units: Number of hidden units in each layer
  dropout: Proportion of units that get dropped out
  num_epochs: Number of epochs used for training
  verbose: binary, whether to show progress of the fit after each epoch

  Output: returns list of all kfold accuracies
  '''

  all_train_accuracy = []
  all_val_accuracy = []

  # use func from above to get splits
  (training_X, training_Y), (validation_X, validation_Y) = get_test_train_splits(data, decisions)

  for fold in range(n_folds):
      print(f"Fold {fold} of {n_folds}")
      train_Xi = training_X[fold]
      train_Yi = training_Y[fold]
      val_Xi = validation_X[fold]
      val_Yi = validation_Y[fold]

      method_to_call = getattr(decoders, model_name) # as class in separate module
      model = method_to_call(units=units, dropout=dropout, num_epochs=num_epochs, verbose=verbose)

      model.fit(train_Xi, train_Yi)
      train_accuracy = compute_accuracy(train_Xi, train_Yi, model)
      all_train_accuracy.append(train_accuracy)

      # score on the validation data
      val_accuracy = compute_accuracy(val_Xi, val_Yi, model)
      all_val_accuracy.append(val_accuracy)

      print("training fold accuracy: {}".format(train_accuracy))
      print("validation foldaccuracy: {}".format(val_accuracy))

  return all_train_accuracy, all_val_accuracy


text:
#          Copyright Rein Halbersma 2020-2021.
# Distributed under the Boost Software License, Version 1.0.
#    (See accompanying file LICENSE_1_0.txt or copy at
#          http://www.boost.org/LICENSE_1_0.txt)

import numpy as np

from ..enums import Action, Hand, nH

# Deterministic policy for a dealer who stands on 17.
stands_on_17 = np.full(nH, Action.HIT)

for _h in range(Hand.H17, Hand.H21 + 1):
    stands_on_17[_h] = Action.STAND

for _s in range(Hand.S17, Hand.BJ + 1):
    stands_on_17[_s] = Action.STAND

# Deterministic policy for a dealer who hits on soft 17.
hits_on_soft_17 = stands_on_17.copy()
hits_on_soft_17[Hand.S17] = Action.HIT

max_stars_repo_name: rhalbersma/blackjack


text:
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt

from numpy.fft import fft2, ifft2, fftshift, ifftshift

# %% read in the image and do transform
Q6_2 = np.array(Image.open("Q6_2.tif"))
row, col = Q6_2.shape
img_fourier = fftshift(fft2(Q6_2))

# %% generate degradation function
k = 0.0025
func = np.fromfunction(lambda mu, nu:
                       np.exp(-k * ((mu - row / 2) ** 2 + (nu - col / 2) ** 2) ** (5 / 6)),
                       (row, col))

# %% restore the image
def wiener(fourier=img_fourier, K=0.1):
    img_out_fourier = func ** 2 / func / (func ** 2 + K) * fourier
    img_out_view = np.log10(np.abs(fourier) + 1)
    img_out = np.abs(ifft2(ifftshift(img_out_fourier)))
    return img_out_fourier, img_out_view, img_out

_, img_out_view1, img_out1 = wiener(K=0.1)
_, img_out_view2, img_out2 = wiener(K=0.01)
_, img_out_view3, img_out3 = wiener(K=0.001)

plt.figure(figsize=(9.6, 6.4))
plt.subplot(231)
plt.imshow(img_out_view1, cmap='gray')
plt.title(
    'Fourier transform of\n'
    'the restored image\n'
    '(log transformed, $K = 0.1$)')
plt.subplot(234)
plt.imshow(img_out1, cmap='gray')
plt.title('Restored image')

plt.subplot(232)
plt.imshow(img_out_view2, cmap='gray')
plt.title(
    'Fourier transform of\n'
    'the restored image\n'
    '(log transformed, $K = 0.01$)')
plt.subplot(235)
plt.imshow(img_out2, cmap='gray')
plt.title('Restored image')

plt.subplot(233)
plt.imshow(img_out_view2, cmap='gray')
plt.title(
    'Fourier transform of\n'
    'the restored image\n'
    '(log transformed, $K = 0.001$)')
plt.subplot(236)
plt.imshow(img_out3, cmap='gray')
plt.title('Restored image')
plt.savefig('Q6_2_3.png')
plt.show()

max_stars_repo_name: kommunium/dip-lab

max_issues_repo_name: kommunium/dip-lab

max_forks_repo_name: kommunium/dip-lab

reason: import numpy,from numpy


text:
import numpy as np
from scipy.sparse import diags,bmat
from scipy.sparse.linalg import spsolve

class CurvedDuctFlowClass(object):
	"""This class provides a finite difference implementation for solving
	steady flow through a curved duct having a rectangular cross-section.
	Solutions may be computed both with and without using the Dean
	approximation (solutions in this are often referred to as Dean flow).

	The specific equations and non-dimensionalisation used within this class
	are described in my publication here doi.org/10.1017/S1446181118000287
	(noting, however, that this paper describes an alternative numerical method
	to solve these equations for a larger variety of cross-section shapes).
	Specifically, refer to the non-linear PDE given in equations (2.3).

	The solver uses Newton's method to solve the fully coupled non-linear problem.
	The linear system for each iteration is contructed in sparse matrix format
	and solved using sparse LU by default (although using bicgstab is an option).
	The discretisation is such that the solutions achieve second order
	convergence with respect to grid resolution.

    A note on correct non-dimensional use:
    You are free to choose any characteristic length X and velocity U that you wish.
    You need to pass the appropriate (non-dimensional) W,H along with epsilon=X/R
    and K=epsilon*Re^2 where Re=(rho/mu)*U*X. To obtain the correct flow you must
    now determine G such that the axial velocity of the non-dimensional solution
    respects your chosen velocity scale. For example, if your U is to be the mean axial
    velocity then G needs to be such that the mean of the returned axial velocity is 0.
    Finding the correct G is a non-linear optimisation problem (with non-zero K)
    and will require some form of iteration.

    A note on correct dimensional use:
    Pass your dimensional W,H values along with epsilon=1/R and K=epsilon*Re^2
    where Re=(rho/mu) in this instance (with appropriate units). You then need
    to determine G such that the returned axial velocity field meets some desired
    criterion (i.e. you might be aiming for a specifc flow rate). Again, this is a
    non-linear optimisation problem that you, the user, are responsible for solving.
	"""
	def __init__(self,m=65,n=65,W=2.0,H=2.0,epsilon=0.01,K=1.0,G=4.0):
		"""
		Initialise the class with the desired parameters.
		Some parameters can be changed later, excepting W,H (I would
		recommend using a separate class instance for different shapes).
		Similarly, changing m,n is loosely supported, but using different
		instances would probably better.

		Parameters
		----------
		m,n : int,int
			The number of points in the finite difference discretisation,
			m being the number of points across the width, and
			n being the number of points across the height.
			Note that array shape/dimensions will be ordered as (n,m).
			Defaults to 65,65
		W,H : float
			The width and height of the rectangular cross-section
			(in dimensionless coordinates).
			To be consistent with the non-dimensionalisation in the referenced
			paper these should generally be scaled such that min(W,H)=2.
			Defaults to 2.0,2.0
		epsilon: float
			The curvature of the duct, that is 1/R given a bend radius R
			in the appropriate dimensional or dimensionless measurements.
			Providing a value of 0 means the solution corresponds to 'Dean flow'.
			Defaults to 0.0
		K : float
			The Dean number used in the equations (specifically K=epsilon*Re^2).
			Observe this can be specified as non-zero even if epsilon=0.
			Defaults to 1.0
		G : float
			The pressure gradient driving the axial flow.
			Ideally it should be chosen such that max(u)=1.0 given K=0.0
			(e.g. 4.0/L^2 would be the appropriate choice for a circular
			 pipe and is generally close enough for our purposes).
			You should only change this if you understand the consequences on
			the scaling of both the axial velocity and stream-function.
			Defaults to 4.0
		"""
		self._shape = (n,m)
		self._size  = (H,W)
		self._epsilon = epsilon
		self._K = K
		self._G = G
		# Define other useful numbers
		self._L = min(0.5*W,0.5*H) # probably not really needed...
		self._N = n*m              # probably not really needed...
		self._ds = W/(m-1)
		self._dz = H/(n-1)
		# Construct arrays/meshgrids for the finite difference node coordinates (8*(2*m+n+3*n*m) bytes)
		self._s = np.linspace(-0.5*W,0.5*W,m)
		self._z = np.linspace(-0.5*H,0.5*H,n)
		self._S,self._Z = np.meshgrid(self._s,self._z) # Note: use default index order
		self._r = 1.0+epsilon*self._s # probably not really needed...
		self._R = 1.0+epsilon*self._S
		# Pre-allocate array for holding an initial guess (2*8*m*n bytes)
		self._u_old = np.zeros((n,m))
		self._Phi_old = np.zeros((n,m))
		self._u_old_is_zero = True # tracks whether an initial u iteration should be performed
		# Pre-allocate arrays of intermediate results (11*8*m*n bytes)
		self._duds = np.zeros((n,m))
		self._dudz = np.zeros((n,m))
		self._Lu = np.zeros((n,m))
		self._dPhids = np.zeros((n,m))
		self._dPhidz = np.zeros((n,m))
		self._d2Phidsdz = np.zeros((n,m))
		self._d2Phids2 = np.zeros((n,m))
		self._d2Phidz2 = np.zeros((n,m))
		#self._LPhi = np.zeros((n,m))   # don't need a persistent copy of this
		self._dLPhids = np.zeros((n,m))
		self._dLPhidz = np.zeros((n,m))
		self._L2Phi = np.zeros((n,m))
		# The following are 'advanced options' I don't expect a typical user to tweak.
		# They first two are related to the extrapolation of ghost cells used to
		# accurately enforce the Neumann boundary conditions.
		# The last relates to the construction of an initial guess if none is provided.
		self._use_first_order_extrap = False # toggle 1st order extrapolation of 2nd derivatives
		self._add_third_deriv_extrap = True  # toggle use of 3rd derivatives in extrapolation
		self._use_initial_u_iterate  = True  # toggle use of initial u iterate (when u_old=0)
		# Done
	def __del__(self):
		"""Destructor, deletes all internally stored arrays
		(noting memory may not be freed immediately: google 'python garbage collection')"""
		del self._s
		del self._z
		del self._r
		del self._S
		del self._Z
		del self._R
		del self._u_old
		del self._Phi_old
		del self._duds
		del self._dudz
		del self._Lu
		del self._dPhids
		del self._dPhidz
		del self._d2Phidsdz
		del self._d2Phids2
		del self._d2Phidz2
		del self._dLPhids
		del self._dLPhidz
		del self._L2Phi
		# Done
	def get_meshgrid(self):
		"""Returns the meshgrid of (s,z) coordinates for the finite
		difference points."""
		return self._S,self._Z
	def get_secondary_velocities(self,Phi=None):
		"""Returns the velocity fields associated with a stream-function
		of the secondary flow.
		Note the velocities will be given as exactly zero on the boundary."""
		if Phi is None:
			Phi = self._Phi_old
		else:
			assert Phi.shape==self._shape
		V = self._zeros_edge()
		W = self._zeros_edge()
		V[1:-1,1:-1] = -self._FD_dz(Phi)/self._R[1:-1,1:-1]
		W[1:-1,1:-1] =  self._FD_ds(Phi)/self._R[1:-1,1:-1]
		return V,W
	def set_initial_guess(self,u=None,Phi=None):
		"""Provide an initial guess for the solver.
		For example, this may be a previously computed solution with
		similar paramters, e.g. perhaps a nearby Dean number.
		With no arguments it resets initial guesses to zero.
		Note that generally the class will re-use the last computed solution
		as an initial guess unless one is provided/specified.
		"""
		if u is None:
			self._u_old.fill(0.0)
			self._u_old_is_zero = True
		else:
			assert u.shape==self._shape
			self._u_old[:,:] = u[:,:]  # copy into pre-allocated arrays
			self._u_old_is_zero = False
		if Phi is None:
			self._Phi_old.fill(0.0)
		else:
			assert Phi.shape==self._shape
			self._Phi_old[:,:] = Phi[:,:]
		return
	def set_K(self,K):
		"""Set a new 'Dean' number K
		(Recalling this specific non-dimensionalisation is such that
		 K is proportional to epsilon*Re^2)"""
		self._K = K
	def get_K(self):
		"""Get the current 'Dean' number K
		(Recalling this specific non-dimensionalisation is such that
		 K is proportional to epsilon*Re^2)"""
		return self._K
	def set_G(self,G):
		"""Set a (dimensionless) driving pressure G
		(Recall this specific non-dimensionalisation is such that
		 G is proportional to Dn=sqrt(K))"""
		self._G = G
	def get_G(self):
		"""Get the current (dimensionless) driving pressure G
		(Recall this specific non-dimensionalisation is such that
		 G is proportional to Dn=sqrt(K))"""
		return self._G
	def set_epsilon(self,epsilon):
		"""Set a new curvature ratio epsilon=L/R
		(R being the bend radius and L the characteristic duct size)
		Recall setting this to zero results in the 'Dean approximation'"""
		self._epsilon = epsilon
		self._r = 1.0+self._epsilon*self._s # probably not really needed...
		self._R = 1.0+self._epsilon*self._S # important to update this!!!
	def get_epsilon(self):
		"""Get the current curvature ratio epsilon=L/R
		(R being the bend radius and L the characteristic duct size)
		Recall setting this to zero results in the 'Dean approximation'"""
		return self._epsilon
	# TODO: consider adding functions for updating m,n,W,H...
	def _zeros_edge(self):
		"""Initialise an array of the correct size with zeros on the edge/boundary"""
		# Note: for small array sizes (essentially fitting in cache)
		# using np.zeros is quicker, however the folllowing scales better
		a=np.empty(self._shape);a[[0,-1],:]=0;a[1:-1,[0,-1]]=0;
		return a
	def _ones_edge(self):
		"""Initialise an array of the correct size with ones on the edge/boundary"""
		# Note: for small array sizes (essentially fitting in cache)
		# using np.ones is quicker, however the folllowing scales better
		a=np.empty(self._shape);a[[0,-1],:]=1;a[1:-1,[0,-1]]=1;
		return a
	def _initial_u_iterate(self):
		"""Construct and solve a Poisson problem to create an initial
		guess for the axial velocity u (which is stored in u_old)."""
		shape = self._shape
		m = shape[1]
		# Total memory for A,b is ~8*8*n*m bytes (including the 3 temporary arrays)
		A_p0_p0 = self._ones_edge()
		A_p0_p0[1:-1,1:-1] = -2.0/self._ds**2-2.0/self._dz**2
		A_pm1_p0 = self._zeros_edge()
		A_pm1_p0[1:-1,1:-1] = 1.0/self._ds**2
		A_p0_pm1 = self._zeros_edge()
		A_p0_pm1[1:-1,1:-1] = 1.0/self._dz**2
		b = self._zeros_edge()
		if self._epsilon==0.0:
			b[1:-1,1:-1] = -self._G
		else:
			b[1:-1,1:-1] = -self._G/self._R[1:-1,1:-1]
			# possibly also modify A_pm1_p0 with duds terms...
		A = diags([A_p0_p0.ravel(),\
				   A_pm1_p0.ravel()[:-1],A_pm1_p0.ravel()[1:],\
				   A_p0_pm1.ravel()[:-m],A_p0_pm1.ravel()[m:]],\
				  [0,1,-1,m,-m],format='csr')
		# It is difficult to predict how much memory spsolve will use
		self._u_old = spsolve(A,b.ravel()).reshape(shape)
		self._u_old_is_zero = False
		return
	def _FD_ds(self,array):
		"""Return centred finite difference estimate of the
		s (or x) derivative on the interior of the given array"""
		return (array[1:-1,2:]-array[1:-1,:-2])/(2.0*self._ds)
	def _FD_dz(self,array):
		"""Return centred finite difference estimate of the
		z (or y) derivative on the interior of the given array"""
		return (array[2:,1:-1]-array[:-2,1:-1])/(2.0*self._dz)
	def _FD_ds2(self,array):
		"""Return centred finite difference estimate of the second
		s (or x) derivative on the interior of the given array"""
		return (array[1:-1,2:]-2.0*array[1:-1,1:-1]+array[1:-1,:-2])/self._ds**2
	def _FD_dsdz(self,array):
		"""Return centred finite difference estimate of the mixed second
		s,z (or x,y) derivative on the interior of the given array"""
		return (array[2:,2:]-array[2:,:-2]-array[:-2,2:]+array[:-2,:-2])/(4.0*self._ds*self._dz)
	def _FD_dz2(self,array):
		"""Return centred finite difference estimate of the second
		z (or y) derivative on the interior of the given array"""
		return (array[2:,1:-1]-2.0*array[1:-1,1:-1]+array[:-2,1:-1])/self._dz**2
	def _FD_Laplace(self,array):
		"""Return centred finite difference estimate of the
		Laplacian on the interior of the given array"""
		return self._FD_ds2(array)+self._FD_dz2(array)
	def _create_ghost_array(self,array):
		"""Extrapolate a given array to provide one layer of ghost points
		around the exterior.
		The specific extrapolation here is such that the Neumann boundary
		conditions for Phi are satisfied (noting that the linear system
		construction must also be consistent with this).
		As such, the given array will generally be self._Phi_old"""
		shape = array.shape
		ghost = np.empty((shape[0]+2,shape[1]+2))
		ghost[1:-1,1:-1] = array
		if self._use_first_order_extrap:
			ghost[ 0,1:-1] = 1.5*ghost[ 1,1:-1]-1.0*ghost[ 2,1:-1]+0.5*ghost[ 3,1:-1]
			ghost[-1,1:-1] = 1.5*ghost[-2,1:-1]-1.0*ghost[-3,1:-1]+0.5*ghost[-4,1:-1]
			ghost[1:-1, 0] = 1.5*ghost[1:-1, 1]-1.0*ghost[1:-1, 2]+0.5*ghost[1:-1, 3]
			ghost[1:-1,-1] = 1.5*ghost[1:-1,-2]-1.0*ghost[1:-1,-3]+0.5*ghost[1:-1,-4]
		else:
			ghost[ 0,1:-1] =  2.0*ghost[ 1,1:-1]-2.5*ghost[ 2,1:-1] \
			                 +2.0*ghost[ 3,1:-1]-0.5*ghost[ 4,1:-1]
			ghost[-1,1:-1] =  2.0*ghost[-2,1:-1]-2.5*ghost[-3,1:-1] \
			                 +2.0*ghost[-4,1:-1]-0.5*ghost[-5,1:-1]
			ghost[1:-1, 0] =  2.0*ghost[1:-1, 1]-2.5*ghost[1:-1, 2] \
			                 +2.0*ghost[1:-1, 3]-0.5*ghost[1:-1, 4]
			ghost[1:-1,-1] =  2.0*ghost[1:-1,-2]-2.5*ghost[1:-1,-3] \
			                 +2.0*ghost[1:-1,-4]-0.5*ghost[1:-1,-5]
		if self._add_third_deriv_extrap:
			ghost[ 0,1:-1] -= (-ghost[ 1,1:-1]+3.0*ghost[ 2,1:-1] \
			                   -3.0*ghost[ 3,1:-1]+ghost[ 4,1:-1])/6.0
			ghost[-1,1:-1] += ( ghost[-2,1:-1]-3.0*ghost[-3,1:-1] \
							   +3.0*ghost[-4,1:-1]-ghost[-5,1:-1])/6.0
			ghost[1:-1, 0] -= (-ghost[1:-1, 1]+3.0*ghost[1:-1, 2] \
							   -3.0*ghost[1:-1, 3]+ghost[1:-1, 4])/6.0
			ghost[1:-1,-1] += ( ghost[1:-1,-2]-3.0*ghost[1:-1,-3] \
							   +3.0*ghost[1:-1,-4]-ghost[1:-1,-5])/6.0
		# Note that extrapolation of corner points is not needed
		return ghost
	def _update_intermediates(self,u,Phi):
		"""Update intermediate arrays containing finite difference
		estimates of the derivatives of solution vectors."""
		# TODO: determine which (if any) of the following can be avoided when epsilon==0
		self._duds[1:-1,1:-1] = self._FD_ds(u)
		self._dudz[1:-1,1:-1] = self._FD_dz(u)
		self._Lu[1:-1,1:-1] = self._FD_Laplace(u)
		self._dPhids[1:-1,1:-1] = self._FD_ds(Phi)
		self._dPhidz[1:-1,1:-1] = self._FD_dz(Phi)
		self._d2Phidsdz[1:-1,1:-1] = self._FD_dsdz(Phi)
		Phi_ghost = self._create_ghost_array(Phi)     # this is only used for intermediate calculation
		self._d2Phids2[:,:] = self._FD_ds2(Phi_ghost)
		self._d2Phidz2[:,:] = self._FD_dz2(Phi_ghost)
		LPhi = self._d2Phids2 + self._d2Phidz2        # this is only used for intermediate calculation
		self._dLPhids[1:-1,1:-1] = self._FD_ds(LPhi)
		self._dLPhidz[1:-1,1:-1] = self._FD_dz(LPhi)
		self._L2Phi[1:-1,1:-1] = self._FD_Laplace(LPhi)
		return
	def _generate_solver_residual(self,u=None,Phi=None,K=None,G=None):
		"""Generates the residual vector for the Newton iteration used by the solver.
		(Note this differs a little from the form written in my ANZIAM paper
		 by factors of R=1+epsilon*S)"""
		if u is None:
			u = self._u_old
		else:
			assert u.shape==self._shape
		if Phi is None:
			Phi = self._Phi_old
		else:
			assert Phi.shape==self._shape
		if K is None:
			K = self._K
		if G is None:
			G = self._G

		shape = self._shape
		eps = self._epsilon
		C_int = 1.0
		if eps!=0.0:
			C_int = 1.0/self._R[1:-1,1:-1]

		self._update_intermediates(u,Phi)
		duds_int = self._duds[1:-1,1:-1]
		dudz_int = self._dudz[1:-1,1:-1]
		Lu_int = self._Lu[1:-1,1:-1]
		dPhids_int = self._dPhids[1:-1,1:-1]
		dPhidz_int = self._dPhidz[1:-1,1:-1]
		d2Phids2_int = self._d2Phids2[1:-1,1:-1]
		d2Phidz2_int = self._d2Phidz2[1:-1,1:-1]
		L2Phi_int = self._L2Phi[1:-1,1:-1]
		d2Phidsdz_int = self._d2Phidsdz[1:-1,1:-1]
		dLPhids_int = self._dLPhids[1:-1,1:-1]
		dLPhidz_int = self._dLPhidz[1:-1,1:-1]

		b = np.empty((2,)+shape)
		b[:,[0,-1],:] = 0    # 0's on the edge/boundary
		b[:,1:-1,[0,-1]] = 0
		b[0,1:-1,1:-1] = -G*C_int-Lu_int \
						 -K*C_int*(dPhidz_int*duds_int-dPhids_int*dudz_int)
		b[1,1:-1,1:-1] =  2.0*u[1:-1,1:-1]*dudz_int-L2Phi_int \
						 +K*C_int*(dPhids_int*dLPhidz_int-dPhidz_int*dLPhids_int)
		if eps!=0.0:
			b[0,1:-1,1:-1] += eps*((eps-K*dPhidz_int)*u[1:-1,1:-1]*C_int-duds_int)*C_int
			b[1,1:-1,1:-1] +=  2.0*eps*dLPhids_int*C_int \
							  -3.0*eps**2*d2Phids2_int*C_int**2 \
							  +3.0*eps**3*dPhids_int*C_int**3 \
							  +eps*K*( dPhidz_int*(2.0*d2Phidz2_int+3.0*d2Phids2_int) \
									  -dPhids_int*d2Phidsdz_int)*C_int**2 \
							  -3.0*eps**2*K*dPhidz_int*dPhids_int*C_int**3

		return b.ravel() # Note: ravel gives a view, flatten creates a copy
	def _construct_Newton_system(self,u=None,Phi=None,K=None,G=None):
		"""Constructs the linear system corresponding to a Newton iteration
		of the finite difference discretisation of the problem.
		This is constructed into a sparse matrix format.
		The system is constructed using u_old,Phi_old."""
		if u is None:
			u = self._u_old
		else:
			assert u.shape==self._shape
		if Phi is None:
			Phi = self._Phi_old
		else:
			assert Phi.shape==self._shape
		if K is None:
			K = self._K
		if G is None:
			G = self._G

		m = self._shape[1]
		ds = self._ds
		dz = self._dz
		eps = self._epsilon
		C_int = 1.0
		if eps!=0.0:
			C_int = 1.0/self._R[1:-1,1:-1]

		b = self._generate_solver_residual(u,Phi,K,G)
		# Note the preceding line calls self._update_intermediates(u,Phi)

		duds_int = self._duds[1:-1,1:-1]
		dudz_int = self._dudz[1:-1,1:-1]
		Lu_int = self._Lu[1:-1,1:-1]
		dPhids_int = self._dPhids[1:-1,1:-1]
		dPhidz_int = self._dPhidz[1:-1,1:-1]
		d2Phids2_int = self._d2Phids2[1:-1,1:-1]
		d2Phidz2_int = self._d2Phidz2[1:-1,1:-1]
		L2Phi_int = self._L2Phi[1:-1,1:-1]
		d2Phidsdz_int = self._d2Phidsdz[1:-1,1:-1]
		dLPhids_int = self._dLPhids[1:-1,1:-1]
		dLPhidz_int = self._dLPhidz[1:-1,1:-1]

		# TODO: consider re-using 'temporary' arrays within here...

		A00_p0_p0 = self._ones_edge()
		A00_p0_p0[1:-1,1:-1] = -2.0/ds**2-2.0/dz**2
		A00_p1_p0 = self._zeros_edge()
		A00_p1_p0[1:-1,1:-1] = ( K/(2.0*ds))*C_int*dPhidz_int
		if eps!=0.0:
			A00_p0_p0[1:-1,1:-1] += (eps*K*dPhidz_int-eps**2)*C_int**2
			A00_p1_p0[1:-1,1:-1] += (eps/(2.0*ds))*C_int
		A00_m1_p0 = -A00_p1_p0
		A00_p1_p0[1:-1,1:-1] += 1.0/ds**2
		A00_m1_p0[1:-1,1:-1] += 1.0/ds**2
		A00_p0_p1 = self._zeros_edge()
		A00_p0_p1[1:-1,1:-1] = (-K/(2.0*dz))*C_int*dPhids_int
		A00_p0_m1 = -A00_p0_p1
		A00_p0_p1[1:-1,1:-1] += 1.0/dz**2
		A00_p0_m1[1:-1,1:-1] += 1.0/dz**2
		A00_block = diags([A00_p0_p0.ravel(),\
						   A00_p1_p0.ravel()[:-1],A00_m1_p0.ravel()[1:],\
						   A00_p0_p1.ravel()[:-m],A00_p0_m1.ravel()[m:]],\
						  [0,1,-1,m,-m],format='csr')

		A01_p1_p0 = self._zeros_edge()
		A01_p1_p0[1:-1,1:-1] = (-K/(2.0*ds))*C_int*dudz_int
		#A01_m1_p0 = -A01_p1_p0
		A01_p0_p1 = self._zeros_edge()
		A01_p0_p1[1:-1,1:-1] = ( K/(2.0*dz))*C_int*duds_int
		if eps!=0.0:
			A01_p0_p1[1:-1,1:-1] += (eps*K/(2.0*dz))*u[1:-1,1:-1]*C_int**2
		#A01_p0_m1 = -A01_p0_p1
		#A01_block = diags([A01_p1_p0.ravel()[:-1],A01_m1_p0.ravel()[1:],\
		#				   A01_p0_p1.ravel()[:-m],A01_p0_m1.ravel()[m:]], \
		#				  [1,-1,m,-m],format='csr')
		A01_block = diags([A01_p1_p0.ravel()[:-1],-A01_p1_p0.ravel()[1:],\
						   A01_p0_p1.ravel()[:-m],-A01_p0_p1.ravel()[m:]], \
						  [1,-1,m,-m],format='csr')

		A10_p0_p0 = self._zeros_edge()
		A10_p0_p1 = self._zeros_edge()
		A10_p0_p0[1:-1,1:-1] = -2.0*dudz_int
		A10_p0_p1[1:-1,1:-1] = (-1.0/dz)*u[1:-1,1:-1] # factor 2 cancels from 2*dz
		#A10_p0_m1 = -A10_p0_p1
		#A10_block = diags([A10_p0_p0.ravel(),\
		#                   A10_p0_p1.ravel()[:-m],A10_p0_m1.ravel()[m:]],\
		#                  [0,m,-m],format='csr')
		A10_block = diags([A10_p0_p0.ravel(),\
		                   A10_p0_p1.ravel()[:-m],-A10_p0_p1.ravel()[m:]],\
		                  [0,m,-m],format='csr')

		A11_p0_p0 = self._ones_edge()
		A11_p0_p0[1:-1,1:-1] =  6.0/ds**4+6.0/dz**4+8.0/(ds*dz)**2
		A11_p2_p0 = self._zeros_edge()
		A11_p1_p0 = self._zeros_edge()
		A11_p1_p0[1:-1,1:-1] = -4.0/ds**4-4.0/(ds*dz)**2 \
		                       -K/(2.0*ds)*C_int*dLPhidz_int \
		                       -K*(1.0/ds**3+1.0/(ds*dz**2))*C_int*dPhidz_int
		A11_m1_p0 = self._zeros_edge()
		A11_m1_p0[1:-1,1:-1] = -4.0/ds**4-4.0/(ds*dz)**2 \
		                       +K/(2.0*ds)*C_int*dLPhidz_int \
		                       +K*(1.0/ds**3+1.0/(ds*dz**2))*C_int*dPhidz_int
		A11_m2_p0 = self._zeros_edge()
		A11_p0_p2 = self._zeros_edge()
		A11_p0_p1 = self._zeros_edge()
		A11_p0_p1[1:-1,1:-1] = -4.0/dz**4-4.0/(ds*dz)**2 \
		                       +K/(2.0*dz)*C_int*dLPhids_int \
		                       +K*(1.0/dz**3+1.0/(dz*ds**2))*C_int*dPhids_int
		A11_p0_m1 = self._zeros_edge()
		A11_p0_m1[1:-1,1:-1] = -4.0/dz**4-4.0/(ds*dz)**2 \
		                       -K/(2.0*dz)*C_int*dLPhids_int \
		                       -K*(1.0/dz**3+1.0/(dz*ds**2))*C_int*dPhids_int
		A11_p0_m2 = self._zeros_edge()
		A11_p1_p1 = self._zeros_edge()
		A11_p1_p1[1:-1,1:-1] = 2.0/(ds*dz)**2+K/(2.0*ds*dz**2)*C_int*dPhidz_int \
											 -K/(2.0*dz*ds**2)*C_int*dPhids_int
		A11_m1_p1 = self._zeros_edge()
		A11_m1_p1[1:-1,1:-1] = 2.0/(ds*dz)**2-K/(2.0*ds*dz**2)*C_int*dPhidz_int \
											 -K/(2.0*dz*ds**2)*C_int*dPhids_int
		A11_p1_m1 = self._zeros_edge()
		A11_p1_m1[1:-1,1:-1] = 2.0/(ds*dz)**2+K/(2.0*ds*dz**2)*C_int*dPhidz_int \
											 +K/(2.0*dz*ds**2)*C_int*dPhids_int
		A11_m1_m1 = self._zeros_edge()
		A11_m1_m1[1:-1,1:-1] = 2.0/(ds*dz)**2-K/(2.0*ds*dz**2)*C_int*dPhidz_int \
											 +K/(2.0*dz*ds**2)*C_int*dPhids_int

		if eps!=0.0:
			A11_p0_p0[1:-1,1:-1] -= ( 6.0*(eps/ds)**2 \
			                         -(4.0*K*eps/dz**2+6.0*K*eps/ds**2)*dPhidz_int)*C_int**2
			A11_p1_p0[1:-1,1:-1] +=  2.0*(eps/(ds*dz**2)+eps/ds**3)*C_int \
			                        +( 3.0*(eps/ds)**2-3.0*K*eps/ds**2*dPhidz_int \
			                          +K*eps/(2.0*ds)*d2Phidsdz_int)*C_int**2 \
			                        +1.5*(-eps**3/ds+K*eps**2/ds*dPhidz_int)*C_int**3
			A11_m1_p0[1:-1,1:-1] += -2.0*(eps/(ds*dz**2)+eps/ds**3)*C_int \
									+( 3.0*(eps/ds)**2-3.0*K*eps/ds**2*dPhidz_int \
									  -K*eps/(2.0*ds)*d2Phidsdz_int)*C_int**2 \
									+1.5*( eps**3/ds-K*eps**2/ds*dPhidz_int)*C_int**3
			A11_p0_p1[1:-1,1:-1] -=  2.0*K*eps*C_int**2*( 0.5/dz*d2Phidz2_int+1.0/dz**2*dPhidz_int) \
									+1.5/dz*K*eps*C_int**2*d2Phids2_int \
									-1.5*K*eps**2/dz*C_int**3*dPhids_int
			A11_p0_m1[1:-1,1:-1] -=  2.0*K*eps*C_int**2*(-0.5/dz*d2Phidz2_int+1.0/dz**2*dPhidz_int) \
									-1.5/dz*K*eps*C_int**2*d2Phids2_int \
									+1.5*K*eps**2/dz*C_int**3*dPhids_int
			A11_p1_p1[1:-1,1:-1] -=  eps/(ds*dz**2)*C_int \
			                        -K*eps/(4.0*ds*dz)*C_int**2*dPhids_int
			A11_p1_m1[1:-1,1:-1] -=  eps/(ds*dz**2)*C_int \
			                        +K*eps/(4.0*ds*dz)*C_int**2*dPhids_int
			A11_m1_p1[1:-1,1:-1] +=  eps/(ds*dz**2)*C_int \
			                        -K*eps/(4.0*ds*dz)*C_int**2*dPhids_int
			A11_m1_m1[1:-1,1:-1] +=  eps/(ds*dz**2)*C_int \
			                        +K*eps/(4.0*ds*dz)*C_int**2*dPhids_int

		if eps==0.0:
			temp_p2_p0 = 1.0/ds**4+0.5*K/ds**3*dPhidz_int
			temp_m2_p0 = 1.0/ds**4-0.5*K/ds**3*dPhidz_int
		else:
			temp_p2_p0 = 1.0/ds**4+(0.5*K/ds**3*dPhidz_int-eps/ds**3)*C_int
			temp_m2_p0 = 1.0/ds**4-(0.5*K/ds**3*dPhidz_int-eps/ds**3)*C_int
		A11_p2_p0[1:-1,1:-2] = temp_p2_p0[:, :-1] # p2 part remains as is for left edge
		A11_m2_p0[1:-1,2:-1] = temp_m2_p0[:,1:  ] # m2 part remains as is for right edge
		A11_p2_p0[1:-1,-2] = 0.0
		A11_m2_p0[1:-1, 1] = 0.0
		# adjust left side appropriately
		if self._use_first_order_extrap: #+3/2,-1,+1/2
			A11_m1_p0[1:-1, 1] += 1.5*temp_m2_p0[:,0]
			A11_p0_p0[1:-1, 1] -= 1.0*temp_m2_p0[:,0]
			A11_p1_p0[1:-1, 1] += 0.5*temp_m2_p0[:,0]
		else: #+2,-5/2,+2,-1/2
			A11_m1_p0[1:-1, 1] += 2.0*temp_m2_p0[:,0]
			A11_p0_p0[1:-1, 1] -= 2.5*temp_m2_p0[:,0]
			A11_p1_p0[1:-1, 1] += 2.0*temp_m2_p0[:,0]
			A11_p2_p0[1:-1, 1] -= 0.5*temp_m2_p0[:,0]
		if self._add_third_deriv_extrap: #+1/6,-1/2,+1/2,-1/6
			A11_m1_p0[1:-1, 1] += 1.0/6.0*temp_m2_p0[:,0]
			A11_p0_p0[1:-1, 1] -= 0.5*temp_m2_p0[:,0]
			A11_p1_p0[1:-1, 1] += 0.5*temp_m2_p0[:,0]
			A11_p2_p0[1:-1, 1] -= 1.0/6.0*temp_m2_p0[:,0]
		# adjust right side appropriately
		if self._use_first_order_extrap: #+3/2,-1,+1/2
			A11_p1_p0[1:-1,-2] += 1.5*temp_p2_p0[:,-1]
			A11_p0_p0[1:-1,-2] -= 1.0*temp_p2_p0[:,-1]
			A11_m1_p0[1:-1,-2] += 0.5*temp_p2_p0[:,-1]
		else: #+2,-5/2,+2,-1/2
			A11_p1_p0[1:-1,-2] += 2.0*temp_p2_p0[:,-1]
			A11_p0_p0[1:-1,-2] -= 2.5*temp_p2_p0[:,-1]
			A11_m1_p0[1:-1,-2] += 2.0*temp_p2_p0[:,-1]
			A11_m2_p0[1:-1,-2] -= 0.5*temp_p2_p0[:,-1]
		if self._add_third_deriv_extrap: #+1/6,-1/2,+1/2,-1/6
			A11_p1_p0[1:-1,-2] += 1.0/6.0*temp_p2_p0[:,-1]
			A11_p0_p0[1:-1,-2] -= 0.5*temp_p2_p0[:,-1]
			A11_m1_p0[1:-1,-2] += 0.5*temp_p2_p0[:,-1]
			A11_m2_p0[1:-1,-2] -= 1.0/6.0*temp_p2_p0[:,-1]
		#
		temp_p0_p2 = 1.0/dz**4-K/(2.0*dz**3)*C_int*dPhids_int
		temp_p0_m2 = 1.0/dz**4+K/(2.0*dz**3)*C_int*dPhids_int
		A11_p0_p2[1:-2,1:-1] = temp_p0_p2[ :-1,:] # p2 part remains as is for bottom edge
		A11_p0_m2[2:-1,1:-1] = temp_p0_m2[1:  ,:] # m2 part remains as is for top edge
		A11_p0_p2[-2,1:-1] = 0.0
		A11_p0_m2[ 1,1:-1] = 0.0
		# adjust bottom side appropriately
		if self._use_first_order_extrap: #+3/2,-1,+1/2
			A11_p0_m1[ 1,1:-1] += 1.5*temp_p0_m2[0,:]
			A11_p0_p0[ 1,1:-1] -= 1.0*temp_p0_m2[0,:]
			A11_p0_p1[ 1,1:-1] += 0.5*temp_p0_m2[0,:]
		else: #+2,-5/2,+2,-1/2
			A11_p0_m1[ 1,1:-1] += 2.0*temp_p0_m2[0,:]
			A11_p0_p0[ 1,1:-1] -= 2.5*temp_p0_m2[0,:]
			A11_p0_p1[ 1,1:-1] += 2.0*temp_p0_m2[0,:]
			A11_p0_p2[ 1,1:-1] -= 0.5*temp_p0_m2[0,:]
		if self._add_third_deriv_extrap: #+1/6,-1/2,+1/2,-1/6
			A11_p0_m1[ 1,1:-1] += 1.0/6.0*temp_p0_m2[0,:]
			A11_p0_p0[ 1,1:-1] -= 0.5*temp_p0_m2[0,:]
			A11_p0_p1[ 1,1:-1] += 0.5*temp_p0_m2[0,:]
			A11_p0_p2[ 1,1:-1] -= 1.0/6.0*temp_p0_m2[0,:]
		# adjust top side appropriately
		if self._use_first_order_extrap: #+3/2,-1,+1/2
			A11_p0_p1[-2,1:-1] += 1.5*temp_p0_p2[-1,:]
			A11_p0_p0[-2,1:-1] -= 1.0*temp_p0_p2[-1,:]
			A11_p0_m1[-2,1:-1] += 0.5*temp_p0_p2[-1,:]
		else: #+2,-5/2,+2,-1/2
			A11_p0_p1[-2,1:-1] += 2.0*temp_p0_p2[-1,:]
			A11_p0_p0[-2,1:-1] -= 2.5*temp_p0_p2[-1,:]
			A11_p0_m1[-2,1:-1] += 2.0*temp_p0_p2[-1,:]
			A11_p0_m2[-2,1:-1] -= 0.5*temp_p0_p2[-1,:]
		if self._add_third_deriv_extrap: #+1/6,-1/2,+1/2,-1/6
			A11_p0_p1[-2,1:-1] += 1.0/6.0*temp_p0_p2[-1,:]
			A11_p0_p0[-2,1:-1] -= 0.5*temp_p0_p2[-1,:]
			A11_p0_m1[-2,1:-1] += 0.5*temp_p0_p2[-1,:]
			A11_p0_m2[-2,1:-1] -= 1.0/6.0*temp_p0_p2[-1,:]

		A11_block = diags([A11_p0_p0.ravel(),\
		                   A11_p2_p0.ravel()[:-2],A11_p1_p0.ravel()[:-1],\
						   A11_m1_p0.ravel()[1:],A11_m2_p0.ravel()[2:],\
						   A11_p0_p2.ravel()[:-2*m],A11_p0_p1.ravel()[:-m],\
						   A11_p0_m1.ravel()[m:],A11_p0_m2.ravel()[2*m:],\
						   A11_p1_p1.ravel()[:-1-m],A11_m1_p1.ravel()[:-m+1],\
						   A11_p1_m1.ravel()[m-1:],A11_m1_m1.ravel()[m+1:]],\
						  [0,2,1,-1,-2,2*m,m,-m,-2*m,1+m,-1+m,1-m,-1-m],format='csr')

		A = bmat([[A00_block,A01_block],[A10_block,A11_block]],format='csr')

		return A,b

	def get_PDE_residual(self,u,Phi,Phi_interior_only=False):
		"""Computes a finite difference estimate of the PDE residual given
		an approximation of u,Phi on a regular grid.
		Note that this residual is based on equation 3 in my ANZIAM paper,
		i.e. I have not multiplied or divided by R=1+epsilon*S as has been done for the solver.
		Values of m,n,W,H,epsilon,K,G are taken to be those specified within the class.

		Parameters
		----------
		u,Phi : np.array,np.array
			Approximations of the axial velocity u and stream-function of the secondary flow Phi
			Each must have the shape (n,m) and are assumed to be a regularly spaced full grid
		Phi_interior_only : bool
			If True then the residual for Phi is only calculated on the very interior,
			that is two points away from the boundary.
			(This can be useful for comparing the interior solution with that obtained
			 from other codes/implementations that do a poor job near the boundary.)

		Returns
		-------
		np.array,np.array
			The residuals of the u and Phi PDE equations respectively
		"""

		assert u.shape==self._shape
		assert Phi.shape==self._shape

		b = self._generate_solver_residual(u,Phi)
		ru   = b[:self._N].reshape(self._shape)
		rPhi = b[self._N:].reshape(self._shape)

		if self._epsilon!=0.0:
			ru   *= self._R
			rPhi /= self._R

		if Phi_interior_only:
			rPhi[[1,-2],1:-1] = 0
			rPhi[2:-2,[1,-2]] = 0

		return ru,rPhi
	def _Newton_solver_iterate(self,return_matrix=False):
		"""An implementation of a fully coupled Newton iteration in
		sparse matrix format for solving the curved duct flow equations.

		Parameters
		----------
		return_matrix : bool
			If true then the (spares) matrix and RHS vector will be returned
			(rather than the actual solution),
			this can be useful for debugging, or if a custom linear solver is to be used.

		Returns
		-------
		np.array,np.array
			An updated approximation of u,Phi respectively
		"""

		if self._u_old_is_zero and self._use_initial_u_iterate:
			self._initial_u_iterate()

		A,b = self._construct_Newton_system()

		if return_matrix:
			return A,b

		duPhi = spsolve(A,b).reshape((2,)+self._shape)

		# Note: we won't update self._u_old,self._Phi_old in this low level function

		return duPhi[0]+self._u_old,duPhi[1]+self._Phi_old

	def solve(self,K=None,K_ramp=None,tol=1.0E-10,ramp_tol=1.0E-3,max_its=20,iterative=False,verbose=False):
		"""
		This automates the full solution of the Dean flow problem.
		The default solversettings are desined to work well for the default
		values of H,G (and W>=H) in the class constructor. In other cases, or
		for very large K, the solver parameters may need tweaking.

		Parameters
		----------
		K : float (>=0)
			Specify the desired 'Dean number' K to solve for.
			If specified it will update the value stored within the class.
			If None then it will update the value stored within the class.
		K_ramp : [float (>=0),float (>=0)]
			If specified then it gives a starting K and a step size to be used
			in relation to solving for large K via progressively increasing K.
			This is particularly relevant if one provides a non-default value of
			any of H,W,G, or if a an initial guess is provided via the solution
			of another (possibly large) K. Otherwise the default is to ramp K from
			500 and increase in steps of 500 until the desired K is reached.
		tol : float
			Specifies a convergence tolerance
		ramp_tol : float
			Specifies a larger tolerance to use during K ramp iterations
		max_its : int
			Specifies a maximum iterations of inner solve (shouldn't need to be high
			when combined with appropriate ramping of K values)
		iterative : bool
			Toggle between iterative and direct solver
			(Not fully supported: I've not had much luck with iterative solvers at high resolution)
		verbose : bool
			Toggle the verbosity of the solver (mostly just prints convergence info)

		Returns
		-------
		np.array,np.array
			The flow solution u,Phi respectively
		"""
		if K is None:
			K = self._K
		else:
			self._K = K
		if K_ramp is None:
			# These default parameters are based on a default W,H,G
			# and are not guaranteed to work otherwise
			K_ramp = [min(500.0,K),500.0]
		else:
			assert len(K_ramp)==2

		# Extra setup in the case of an iterative solver...
		# (although this is not fully tested/supported...)
		if iterative:
			from scipy.sparse.linalg import gmres,bicgstab,spilu,LinearOperator
			iter_solver = bicgstab # bicgstab generally seems to be about twice as fast as gmres
			iter_tol = 1.0E-9
			preconditioner = "ilu" # ILU is the only viable option I've found...

		# Perform an initial u iterate if required...
		if self._u_old_is_zero and self._use_initial_u_iterate:
			self._initial_u_iterate()

		u   = self._u_old   # we avoid a copy here so that self._u_old is updated
		Phi = self._Phi_old # we avoid a copy here so that self._Phi_old is updated

		# TODO: consider if we want to implement more generally, e.g. for G,H,W,epsilon which are not the current class values (although I don't think it makes sense to do so...)

		# Generate list of K values to solve through
		if K<=K_ramp[0]:
			#K_list = [K] # this dis-allows backward ramping
			K_list = list(np.arange(K_ramp[0],K,K_ramp[1]))+[K] # this approach allows 'backward' ramping
		else:
			K_list = list(np.arange(K_ramp[0],K,K_ramp[1]))+[K]
		for k in K_list:
			its = 0
			err = 1.0
			temp_tol = ramp_tol
			if k==K:
				temp_tol = tol
			while err>temp_tol:
				# TODO: Consider modifying the following to avoid re-allocating temporary arrays every iteration
				A,b = self._construct_Newton_system(u,Phi,k)
				if iterative:
					res = 0
					if preconditioner=="ilu":
						B = spilu(A.tocsc(),drop_tol=1.0E-9,fill_factor=32.0) # drop_tol and fill_factor may need adjusting for different resolutions and values of K???
						M = LinearOperator(A.get_shape(),lambda x:B.solve(x))
						res = iter_solver(A,b.ravel(),M=M,tol=iter_tol)
					elif preconditioner=="diag":
						DI = 1.0/A.diagonal(0)
						M = LinearOperator(A.get_shape(),lambda x:x*DI)
						res = iter_solver(A,b.ravel(),M=M,tol=iter_tol)
					elif preconditioner=="jacobi":
						D = A.diagonal(0)
						M = LinearOperator(A.get_shape(),lambda x:(x-A.dot(x)+D*x)/D)
						res = iter_solver(A,b.ravel(),M=M,tol=iter_tol)
					else:
						res = iter_solver(A,b.ravel(),tol=iter_tol)
					if res[1]!=0:
						print("Error: itertive method was not successful... error code:",res[1])
					duPhi = res[0].reshape((2,)+self._shape)
					u   += duPhi[0]
					Phi += duPhi[1]
					err = max(np.linalg.norm(duPhi[0])/np.linalg.norm(u),\
							  np.linalg.norm(duPhi[1])/np.linalg.norm(Phi))
				else:
					duPhi = spsolve(A,b.ravel()).reshape((2,)+self._shape)
					u   += duPhi[0]
					Phi += duPhi[1]
					err = max(np.linalg.norm(duPhi[0])/np.linalg.norm(u),\
							  np.linalg.norm(duPhi[1])/np.linalg.norm(Phi))
				its += 1
				if its>=max_its:
					print("Warning: max_its reached, terminating early...")
					# I used this to see check that bifurcations are likely to be numerical artefacts
					# It is too slow/memory intensive to use for more extensive analysis though...
					#print("conditioning:",np.linalg.cond(A.todense()))
					break
				if verbose and False: # Optional next level debugging...
					print("Error on iteration ",its," is ",err)
			if verbose:
				print("Convergence for K=",k," reached in ",its," iterations (err,tol=",err,temp_tol,")")

		return u,Phi

	def _solver_test(self):
		"""A built-in test routine of the implementation above"""
		import matplotlib.pyplot as plt
		# Setup function for plotting
		S,Z = self.get_meshgrid()
		def plot_pair(u,Phi):
			plt.contourf(S,Z,u,32)
			plt.colorbar()
			plt.gca().set_aspect(1.0)
			plt.show()
			plt.contourf(S,Z,Phi,32)
			plt.colorbar()
			plt.gca().set_aspect(1.0)
			plt.show()
		# Set an initial K, solve the equations (with manual iterations), plot the solution
		self.set_K(500.0)
		print("Solving for K={:.1f}".format(self.get_K()))
		u0_s,Phi0_s = self._Newton_solver_iterate()
		if False: # Optionally view the result of the first iteration...
			plot_pair(u0_s,Phi0_s)
		for k in range(20):
			self.set_initial_guess(u0_s,Phi0_s)
			u1_s,Phi1_s = self._Newton_solver_iterate()
			u_conv = np.linalg.norm(u1_s-u0_s)/np.linalg.norm(u1_s)
			Phi_conv = np.linalg.norm(Phi1_s-Phi0_s)/np.linalg.norm(Phi1_s)
			print("    relative errors:",u_conv,Phi_conv)
			u0_s[:,:] = u1_s[:,:]
			Phi0_s[:,:] = Phi1_s[:,:]
			if max(u_conv,Phi_conv)<1.0E-10:
				print("    converged!")
				break
		print("Plotting solutions:")
		plot_pair(u1_s,Phi1_s)
		# Plot the PDE residual
		print("Plotting PDE residuals:")
		ru,rPhi = self.get_PDE_residual(u1_s,Phi1_s)
		plot_pair(ru,rPhi)
		# Try pushing K higher...
		self.set_initial_guess(u1_s,Phi1_s)
		self.set_K(5000.0)
		print("Solving for K={:.1f} via ramping".format(self.get_K()))
		u,Phi = self.solve(K_ramp=[1000.0,500.0],verbose=True)
		print("Plotting solutions:")
		plot_pair(u,Phi)
		# Plot the PDE residual
		print("Plotting PDE residuals:")
		ru,rPhi = self.get_PDE_residual(u,Phi)
		plot_pair(ru,rPhi)
		# Done
	# End of class

max_stars_repo_name: brendanharding/CurvedDuctFlow

max_issues_repo_name: brendanharding/CurvedDuctFlow

max_forks_repo_name: brendanharding/CurvedDuctFlow


text:
import sys
from scipy.stats import hypergeom

if len(sys.argv) < 3:
   exit("Usage: python feature_enrichment.py <feature association file> <genelist>")

gene_feature = {}
feature_dict = {}

association_file = sys.argv[1]

try:
   fassoc = open(association_file, "r")

   for line in fassoc:
     line = line.rstrip()
     id = line.split("\t")
     feature_id = id[1]
     num_fields = len(id)
     feature_dict[feature_id] = num_fields - 2

     for i in range(3, num_fields):
         gene_id = id[i]

         if gene_id not in gene_feature:
            feature_value = []
            feature_value.append(feature_id)
            gene_feature[gene_id] = feature_value
         else:
            gene_feature[gene_id].append(feature_id)
   fassoc.close()

except IOError:
       print ('cannot open', association_file)
       fassoc.close()

N = len(gene_feature.keys())

n = 0
featurefreq = {}

gene_file = sys.argv[2]

try:
   fgene = open(gene_file, "r")
   for gline in fgene:
     gline = gline.rstrip()
     n += 1
     geneids = gline.split(",")

     if geneids[0] in gene_feature:
        feature_list = gene_feature[geneids[0]]

        for feature in feature_list:
            if feature in featurefreq:
               featurefreq[feature] += 1
            else:
               featurefreq[feature] = 1

   for feature_key, frequency in featurefreq.items():
       k = frequency
       K = feature_dict[feature_key]
       prb = hypergeom.pmf(k, N, K, n)

       print ("Feature Id = " + feature_key + "\tN = " + str(N) + "\tK = " + str(K) + "\tn = " + str(n) + "\tk = " + str(k) + "\tSignificance = " +str(prb))
   fgene.close()

except IOError:
        print ('cannot open', gene_file)
        fgene.close()

reason: from scipy


text: from numpy.fft import *

reason: from numpy