text:
# Author: Sudeep Pillai <spillai@csail.mit.edu>
# License: MIT
from __future__ import print_function

import time
import logging
import os.path

import six
from six.moves import cPickle
from collections import defaultdict

import tables as tb
import numpy as np

from pybot.utils.misc import progressbar
from pybot.utils.io_utils import create_path_if_not_exists
from pybot.utils.itertools_recipes import izip, imap, islice, grouper, chunks

# =============================================================================
# Pytables helpers
# =============================================================================

# Avoids printing natural name warning when writing
# dictionaries with ints as keys
import warnings
warnings.filterwarnings('ignore', category=tb.NaturalNameWarning)

def load_pickled_dict(fn):
    return cPickle.load(open(os.path.expanduser(fn), 'rb'))

def save_pickled_dict(fn, d):
    cPickle.dump(d, open(os.path.expanduser(fn), 'w'))

def load_json_dict(fn):
    import json
    return json.load(open(os.path.expanduser(fn), 'r'))

def save_json_dict(fn, d, pretty=False):
    import json
    with open(os.path.expanduser(fn), 'w') as fp:
        if pretty:
            json.dump(d, fp, sort_keys=True, indent=4, separators=(',', ':'))
        else:
            json.dump(d, fp)

def load_yaml_dict(fn):
    import yaml
    return yaml.load(open(os.path.expanduser(fn), 'r'))

def save_yaml_dict(fn, d):
    import yaml
    with open(os.path.expanduser(fn), 'w') as fp:
        fp.write(yaml.dump(d, default_flow_style=False))

def load_mat(fn):
    import scipy.io as io
    return io.loadmat(os.path.expanduser(fn))

def save_mat(fn, d):
    import scipy.io as io
    io.savemat(os.path.expanduser(fn), d)

def read_pytable(h5f, group=None):
    if group is None:
        group = h5f.root

    data = AttrDict()
    for child in h5f.list_nodes(group):
        item = None
        try:
            if isinstance(child, tb.group.Group):
                item = read_pytable(h5f, child)
            else:
                item = child.read()
                if isinstance(item, str) and item.startswith('OBJ_'):
                    item = cPickle.loads(item[4:])
            data[child._v_name] = item
        except tb.NoSuchNodeError:
            warnings.warn('No such node: "%s", skipping...' % repr(child))
            pass

    return data

def load_pytable(fn):
    try:
        h5f = tb.open_file(os.path.expanduser(
            fn), mode='r', title='Title: %s' % fn)
        data = read_pytable(h5f, group=h5f.root)
        h5f.close()
    except Exception as e:
        raise RuntimeError('{}'.format(e))
        data = AttrDict()
    return data

def get_node(g, k):
    if g._v_pathname.endswith('/'):
        return ''.join([g._v_pathname, k])
    else:
        return ''.join([g._v_pathname, '/', k])

def flush_pytable(h5f, data=None, group=None, table=None, force=True):
    # if data is None: data = self.data
    # if table is None: table = self.tables
    # if group is None: group = self.groups

    # print 'Keys: ', data.keys()
    for k, v in six.iteritems(data):
        # print 'key,val', k,v, type(v)

        try:
            k = str(k)
        except:
            print('Cannot save to DB, key is not string %s ' % k)
            continue

        # if not isinstance(k, str):
        #     continue

        # Clean up before writing
        if force:
            try:
                h5f.remove_node(get_node(group._gp, k), recursive=True)
            except tb.NoSuchNodeError:
                pass

        # print 'In Group: ', group, k, v
        if isinstance(v, dict):
            # self.log.debug('Attempting to save dict type')
            # assert(k not in table);
            table[k] = AttrDict()
            group[k] = AttrDict()
            group[k]._gp = h5f.create_group(group._gp, k)
            h5f.flush()
            # self.log.debug('Out Group: %s' % group[k])
            flush_pytable(h5f, data=v, group=group[k], table=table[k])
        elif isinstance(v, np.ndarray):
            # self.log.debug('Attempting to save ndarray %s' % type(v))
            table[k] = h5f.create_array(group._gp, k, v)
            # self.log.debug('Out Table: %s' % table[k])
        else:
            # print 'Attempting to save arbitrary type %s' % type(v), k, group._gp
            try:
                assert v is not None
                table[k] = h5f.create_carray(group._gp, k, obj=v)
            except (TypeError, ValueError, AssertionError):
                v = 'OBJ_' + cPickle.dumps(v, -1)
                table[k] = h5f.create_array(group._gp, k, v)
                # print 'TypeError', v
            finally:
                h5f.flush()
    return

def save_pytable(fn, d):
    create_path_if_not_exists(fn)
    h5f = tb.open_file(os.path.expanduser(fn), mode='w', title='%s' % fn)

    tables = AttrDict()
    groups = AttrDict()
    groups._gp = h5f.root

    flush_pytable(h5f, data=d, group=groups, table=tables)
    h5f.close()

# =============================================================================
# AttrDict
# =============================================================================

class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)

    def __getitem__(self, attr):
        return super(AttrDict, self).__getitem__(attr)

    def __setitem__(self, attr, value):
        super(AttrDict, self).__setitem__(attr, value)

    def __getattr__(self, attr):
        return super(AttrDict, self).__getitem__(attr)

    def __setattr__(self, attr, value):
        super(AttrDict, self).__setitem__(attr, value)

    def __getstate__(self):
        pass

    def __setstate__(self):
        pass

    def to_dict(self):
        return dict(self)

    @staticmethod
    def load_dict(fn):
        return AttrDict(load_pickled_dict(fn))

    def save_dict(self, fn):
        save_pickled_dict(fn, self)

    @staticmethod
    def load_json(fn):
        return AttrDict(load_json_dict(fn))

    def save_json(self, fn):
        save_json_dict(fn, self)

    @staticmethod
    def load_yaml(fn):
        return AttrDict(load_yaml_dict(fn))

    def save_yaml(self, fn):
        print(self.to_dict())
        save_yaml_dict(fn, self.to_dict())

    @staticmethod
    def load_mat(fn):
        return AttrDict(load_mat(fn))

    def save_mat(self, fn):
        save_mat(fn, self)

    @staticmethod
    def load(fn):
        return load_pytable(fn)

    def save(self, fn):
        fn = os.path.expanduser(fn)
        print('Saving ', fn)
        create_path_if_not_exists(fn)
        return save_pytable(fn, self)

    def pprint(self):
        import json
        print(json.dumps(self.to_dict(),
                         indent=4,
                         sort_keys=True))

class IterDB(object):
    def __init__(self, filename, mode):
        """
        An iterable database that should theoretically allow
        scalable reading/writing of datasets.
           batch_size: length of list

        Notes:
           meta_file should contain all the related meta data
        including keys, their corresponding value lengths,
        overall file size etc
        """
        fn = os.path.expanduser(filename)
        if mode == 'w' or mode == 'a':
            print('{}::{} at {}'.format(
                'Writing' if mode == 'w' else 'Appending',
                self.__class__.__name__, fn))
            self.h5f_ = tb.open_file(fn, mode=mode, title='%s' % fn)
            self.data_ = {}
        elif mode == 'r':
            if not os.path.exists(fn):
                import sys
                raise IOError('Failed to open file, {}'.format(fn))

            self.h5f_ = tb.open_file(fn, mode=mode, title='%s' % fn)
            print('{}::Loaded with fields: {}'.format(
                self.__class__.__name__, self.keys))
        else:
            raise RuntimeError('Unknown mode %s' % mode)

    def __del__(self):
        self.close()

    @property
    def keys(self):
        return [child._v_name for child in self.h5f_.list_nodes(self.h5f_.root)]

    @property
    def filename(self):
        return self.filename_

    def append(self, key, item):
        if key not in self.data_:
            filters = tb.Filters(complevel=5, complib='blosc')
            if isinstance(item, np.ndarray):
                atom = tb.Atom.from_type(item.dtype.name, item.shape[1:])
            else:
                atom = tb.VLStringAtom()
            self.data_[key] = self.h5f_.create_vlarray(self.h5f_.root, key,
                                                       atom, filters=filters)
            print('Creating VLArray, and appending to key {}'.format(key))
            print(self.data_[key])

        self.data_[key].append(self.pack(item))

    def pack(self, item):
        if isinstance(item, np.ndarray):
            return item
        else:
            return 'OBJ_' + cPickle.dumps(item, -1)

    def unpack(self, item):
        if isinstance(item, np.ndarray):
            return item
        elif isinstance(item, str) and item.startswith('OBJ_'):
            return cPickle.loads(item[4:])
        else:
            raise ValueError('Unknown type written to pytables')

    def extend(self, key, items):
        for item in items:
            self.data_[key].append(item)

    def node_str(self, key):
        return ''.join(['/', key])

    def get_node(self, key):
        # _f_get_child(key)
        return self.h5f_.get_node(self.node_str(key))

    def get_node_value(self, key):
        return list(self.get_node(key))

    def length(self, key):
        return self.get_node(key).nrows

    def itervalues_for_key(self, key, inds=None, every_k=None, verbose=False):
        if key not in self.keys:
            raise RuntimeError(
                'Key %s not found in dataset. keys: %s' % (key, self.keys))
        return imap(self.unpack,
                    islice(self.get_node(key).iterrows(), None, None, every_k))

    def itervalues_for_keys(self, keys, inds=None, verbose=False):
        for key in keys:
            if key not in self.keys:
                raise RuntimeError(
                    'Key %s not found in dataset. keys: %s' % (key, self.keys))

        items = (self.itervalues_for_key(key) for key in keys)
        return izip(*items)

    def iterchunks(self, key, batch_size=10, verbose=False):
        if key not in self.keys:
            raise RuntimeError(
                'Key %s not found in dataset. keys: %s' % (key, self.keys))

        itchunks = chunks(self.itervalues_for_key(key), batch_size)
        for chunk in itchunks:
            yield chunk

    def iterchunks_keys(self, keys, batch_size=10, verbose=False):
        """
        Iterate in chunks and izip specific keys
        """
        for key in keys:
            if key not in self.keys:
                raise RuntimeError(
                    'Key %s not found in dataset. keys: %s' % (key, self.keys))

        iterables = (self.iterchunks(key, batch_size=batch_size,
                                     verbose=verbose) for key in keys)
        return izip(*iterables)

    def close(self):
        self.h5f_.close()

class IterDBDeprecated(object):
    def __init__(self, filename, mode, fields=[], batch_size=5):
        """
        An iterable database that should theoretically allow
        scalable reading/writing of datasets.
           batch_size: length of list

        Notes:
           meta_file should contain all the related meta data
        including keys, their corresponding value lengths,
        overall file size etc
        """

        # Setup filenames, meta filenames, and chunk idx
        self.setup(filename)

        if mode == 'w':
            print('{}::Writing with fields: {}, batch size: {}'.format(
                self.__class__.__name__, fields, batch_size))

            # Keep data_.keys(), and keys_ consistent
            self.data_ = AttrDict()
            self.keys_ = set()
            for field in fields:
                self.data_[field] = []
                self.keys_.add(field)
            self.meta_file_ = AttrDict(chunks=[], keys=[])
            self.batch_size_ = batch_size
            self.lengths_ = defaultdict(lambda: 0)
        elif mode == 'r':
            # Load first chunk, and keep keys_ consistent
            self.meta_file_ = AttrDict.load(self.meta_filename_)
            self.keys_ = self.meta_file_.keynames
            self.lengths_ = {}
            for key in self.keys_:
                l = 0
                for chunk in self.meta_file_.chunks:
                    l += chunk[key]
                self.lengths_[key] = l

            print('{}::Loaded with fields: {}, metafiles: {}'.format(
                self.__class__.__name__, self.keys_, len(self.meta_file_)))

            # For the time-being, dynamically disattach append, extend functionality
            self.append, self.extend, self.add_fields = None, None, None
            self.flush, self.save = None, None

        else:
            raise RuntimeError('Unknown mode %s' % mode)

    def setup(self, filename):
        self.chunk_idx_ = 0
        self.folder_, self.filename_ = os.path.split(filename)
        self.meta_filename_ = os.path.join(
            self.folder_, self.filename_, 'meta_data.h5')
        self.get_chunk_filename = lambda idx: os.path.join(
            self.folder_, self.filename_, 'chunks', 'chunk_%04i.h5' % (idx))

    # def __del__(self):
    #     if len(self.data_):
    #         print 'Seems like you have data stored away, but has not been flushed!'
    #         self.flush()

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

    def __setitem__(self, key, value):
        self.meta_file_[key] = value

    def add_fields(self, fields):
        for field in fields:
            if field in self.data_:
                raise RuntimeError(
                    'Field %s is already in data, cannot update/replace fields!' % field)
            self.data_[field] = []
            self.keys_.add(field)

    def append(self, key, item):
        if len(self.data_[key]) >= self.batch_size_:
            self.flush()
        self.data_[key].append(item)
        # print [(k, len(v)) for k,v in self.data_.iteritems()]

    def extend(self, key, items):
        self.data_[key].extend(item)

    def length(self, key):
        return self.lengths_[key]

    def itervalues(self, key, inds=None, verbose=False):
        if key not in self.keys_:
            raise RuntimeError(
                'Key %s not found in dataset. keys: %s' % (key, self.keys_))

        idx, ii, done = 0, 0, False
        total_chunks = len(self.meta_file_.chunks)

        if inds is not None:
            import warnings
            warnings.warn('{}::itervalues indices provided must be sorted'.format(
                self.__class__.__name__))
            inds = np.sort(inds)

        for chunk_idx, chunk in enumerate(self.meta_file_.chunks):
            data = AttrDict.load(self.get_chunk_filename(chunk_idx))
            if inds is None:
                for item in data[key]:
                    yield item
                    idx += 1
            else:
                for i, item in enumerate(data[key]):
                    if ii >= len(inds):
                        done = True
                        break
                    if idx == inds[ii]:
                        # print i, ii, idx
                        yield item
                        ii += 1
                    idx += 1

                if done:
                    break

    def keys(self):
        return self.keys_

    def iter_keys_values(self, keys, inds=None, verbose=False):
        for key in keys:
            if key not in self.keys_:
                raise RuntimeError(
                    'Key %s not found in dataset. keys: %s' % (key, self.keys_))

        idx, ii = 0, 0
        total_chunks = len(self.meta_file_.chunks)
        inds = np.sort(inds) if inds is not None else None

        for chunk_idx, chunk in enumerate(self.meta_file_.chunks):
            data = AttrDict.load(self.get_chunk_filename(chunk_idx))

            # if inds is None:
            items = (data[key] for key in keys)
            for item in izip(*items):
                yield item
            # else:
            #     for i, item in enumerate(data[key]):
            #         if inds[ii] == idx + i:
            #             yield item
            #             ii += 1
            #             if ii >= len(inds): break
            #     idx += len(data[key])

    def iterchunks(self, key, batch_size=10, verbose=False):
        if key not in self.keys_:
            raise RuntimeError(
                'Key %s not found in dataset. keys: %s' % (key, self.keys_))

        idx, ii = 0, 0
        total_chunks = len(self.meta_file_.chunks)
        batch_chunks = chunks(range(len(self.meta_file_.chunks)), batch_size)

        for chunk_group in progressbar(batch_chunks, size=total_chunks / batch_size, verbose=verbose):
            items = []
            # print key, chunk_group
            for chunk_idx in chunk_group:
                # grouper will fill chunks with default none values
                if chunk_idx is None:
                    continue
                # Load chunk
                data = AttrDict.load(self.get_chunk_filename(chunk_idx))
                for item in data[key]:
                    items.append(item)
            yield items

    def iterchunks_keys(self, keys, batch_size=10, verbose=False):
        """
        Iterate in chunks and izip specific keys
        """
        for key in keys:
            if key not in self.keys_:
                raise RuntimeError(
                    'Key %s not found in dataset. keys: %s' % (key, self.keys_))

        iterables = [self.iterchunks(
            key, batch_size=batch_size, verbose=verbose) for key in keys]
        return izip(*iterables)

    def flush(self):
        """
        Save dictionary with metadata, move to new chunk,
        and prepare it for writing.
        """
        self.meta_file_.keynames = list(self.data_.keys())
        self.meta_file_.chunks.append(
            AttrDict((k, len(v)) for k, v in six.items(self.data_)))
        self.meta_file_.save(self.meta_filename_)
        for k, v in six.items(self.data_):
            self.lengths_[k] += len(v)

        self.data_.save(self.get_chunk_filename(self.chunk_idx_))
        self._next_chunk_to_write()
        print('Flushing', self.meta_file_.chunks[-1])

    def finalize(self):
        self.flush()

    def _next_chunk_to_write(self):
        for k, v in six.items(self.data_):
            self.data_[k] = []
        self.chunk_idx_ += 1

class AttrDictDB(object):
    # Set up tables first and then flush
    def __init__(self, filename='', data=AttrDict(), mode='r',
                 force=False, recursive=True, maxlen=100, ext='.h5'):
        self.log = logging.getLogger(self.__class__.__name__)

        # Set up output file
        if ext in filename:
            output_filename = "%s" % filename
        else:
            output_filename = "%s%s" % (filename, ext)
        self.log.info('DictDB: Opening file %s' % output_filename)

        # Open the db, no append support just yet
        self.h5f = tb.open_file(
            output_filename, mode=mode, title='%s' % filename)

        # Init the dict
        self.data = data
        self.tables = AttrDict()
        self.groups = AttrDict()
        self.groups._gp = self.h5f.root
        self.maxlen = maxlen

        # Read the db based on the mode
        if mode == 'r' or mode == 'a':
            self.log.info('Reading DB to DictDB')
            self.data = self.read(group=self.h5f.root)

            self.__getattr__ = self.data.__getitem__
            self.__getitem__ = self.data.__getitem__

        self.__setattr__ = self.data.__setitem__
        self.__setitem__ = self.data.__setitem__

    def __del__(self):
        self.h5f.close()

    def close(self):
        try:
            self.h5f.flush()
            self.h5f.close()
        except tb.exceptions.ClosedFileError as e:
            print('AttrDictDB already closed')

    def read(self, group=None):
        if group is None:
            group = self.h5f.root

        data = AttrDict()
        for child in self.h5f.list_nodes(group):
            item = None
            try:
                if isinstance(child, tb.group.Group):
                    item = self.read(child)
                else:
                    item = child.read()
                    if isinstance(item, str) and item.startswith('OBJ_'):
                        item = cPickle.loads(item[4:])
                data[child._v_name] = item
            except tb.NoSuchNodeError:
                warnings.warn('No such node: "%s", skipping...' % repr(child))
                pass
        return data

    def get_node(self, g, k):
        if g._v_pathname.endswith('/'):
            return ''.join([g._v_pathname, k])
        else:
            return ''.join([g._v_pathname, '/', k])

    def flush(self, data=None, group=None, table=None, force=True):
        if data is None:
            data = self.data
        if table is None:
            table = self.tables
        if group is None:
            group = self.groups

        # print 'Keys: ', data.keys()
        for k, v in six.items(data):
            # print 'key,val', k,v, type(v)

            if not isinstance(k, str):
                self.log.debug('Cannot save to DB, key is not string %s ' % k)
                continue

            # Clean up before writing
            if force:
                try:
                    self.h5f.remove_node(self.get_node(
                        group._gp, k), recursive=True)
                except tb.NoSuchNodeError:
                    pass

            # print 'In Group: ', group, k, v
            if isinstance(v, dict):
                self.log.debug('Attempting to save dict type')
                # assert(k not in table);
                table[k] = AttrDict()
                group[k] = AttrDict()
                group[k]._gp = self.h5f.create_group(group._gp, k)
                self.h5f.flush()
                self.log.debug('Out Group: %s' % group[k])
                self.flush(data=v, group=group[k], table=table[k])
            elif isinstance(v, np.ndarray):
                self.log.debug('Attempting to save ndarray %s' % type(v))
                table[k] = self.h5f.create_array(group._gp, k, v)
                self.log.debug('Out Table: %s' % table[k])
            # elif isinstance(v,io_utils.TableWriter):
            #     self.log.debug('Attempting to save with custom writer')
            #     table[k] = self.h5f.createTable(group._gp, name=k,
            #                                     description=v.description,
            #                                     title='%s-data' % (k) )
            #     v.write(table[k])
            #     # print 'Creating table with group:%s name:%s desc:%s' % (group._gp, k, writer.description)
            #     # print 'Out Table: ', table[k]
            else:
                self.log.debug(
                    'Attempting to save arbitrary type %s' % type(v))
                try:
                    assert v is not None
                    table[k] = self.h5f.create_carray(group._gp, k, obj=v)
                except (TypeError, ValueError, AssertionError):
                    v = 'OBJ_' + cPickle.dumps(v, -1)
                    table[k] = self.h5f.create_array(group._gp, k, v)
                    # print 'TypeError', v
                finally:
                    self.h5f.flush()
        return

if __name__ == "__main__":
    # print '\nTesting AttrDict()'
    # a = AttrDict()
    # a.a = 1
    # a.b = 'string'
    # a.c = np.arange(10)
    # a.d = np.arange(36).reshape((6,6))
    # a['test'] = 'test'

    # a.e = AttrDict()
    # a.e.a = 2
    # a.e.b= 'string2'
    # a.e.c= np.arange(10)
    # print a

    # print '\nTesting DictDB() write'
    # db = AttrDictDB(filename='test.h5', mode='w')
    # db.data.a = 1
    # db.data.b = 'string'
    # db.data.c = np.arange(10)

    # db.data.e = AttrDict()
    # db.data.e.a = 2
    # db.data.e.b = 'string2'
    # db.data.e.c = np.arange(10)
    # db.data.e.d = ('this','hat')
    # # db.data.e.f = io_utils.Feature3DWriter(data=[])
    # print 'Write DB: ', db.data
    # wkeys = db.data.keys()
    # db.flush()
    # db.close()
    # print 'OK'

    # print '\nTesting DictDB() read'
    # rdb = AttrDictDB(filename='test', mode='r')
    # print 'Read DB: ', rdb.data
    # rkeys = rdb.data.keys()
    # print 'Rkeys: ', rkeys
    # print 'Wkeys: ', wkeys
    # print 'Diff: ', set(rkeys).difference(set(wkeys))
    # rdb.close()
    # print 'OK'

    # print '\nTesting DictDB() read and then write'
    # rwdb = AttrDictDB(filename='test', mode='a')
    # print 'Read DB: ', rwdb.data
    # print 'Keys: ', rwdb.data.keys()
    # add = AttrDict()
    # add.desc = np.ones((64,2), dtype=np.float32)
    # add.desc2 = np.eye(4)
    # add.string = 'string'
    # add.tup = ('this', 'that')
    # rwdb.flush(data=add)
    # rwdb.close()

    # print '\nTesting DictDB() re-read'
    # rdb = AttrDictDB(filename='test', mode='r')
    # # print 'Read DB: ', rdb.data
    # rkeys = rdb.data.keys()
    # print 'keys: ', rkeys
    # rdb.close()
    # print 'OK'

    print('Testing IterDB')
    from pybot.geometry import RigidTransform
    p = RigidTransform.identity()
    A = [np.random.rand(400, 1000, 3) for j in range(3)]

    # print('Writing to IterDB a,b,c')
    # db = IterDB(filename='iterdb_test.h5', mode='w', batch_size=10000)
    # for j in range(100):
    #     db.append('a', A[0])
    #     db.append('b', A[1])
    #     db.append('c', p)
    #     print j
    # print('a: {}'.format(db.length('a')))
    # print('b: {}'.format(db.length('b')))
    # print('c: {}'.format(db.length('c')))
    # db.close()
    # print('OK')

    print('Reading from IterDB a,c')
    db = IterDB(filename='iterdb_test.h5', mode='r')
    iter_a = db.itervalues_for_key('a')
    for item in db.itervalues_for_key('c'):
        print(item)
    print(db.length('c'))
    db.close()
    print('OK')

    # print('Appending to IterDB d,e,f')
    # db = IterDB(filename='iterdb_test.h5', mode='a')
    # for j in range(300):
    #     db.append('d', A[0])
    #     db.append('e', A[1])
    #     db.append('f', p)
    # print('d: {}'.format(db.length('d')))
    # print('e: {}'.format(db.length('e')))
    # print('f: {}'.format(db.length('f')))
    # db.close()
    # print('OK')

    # print('Reading c and f')
    # db = IterDB(filename='iterdb_test.h5', mode='r')
    # for (c,f) in db.itervalues_for_keys(['c','f']):
    #     print (c,f)
    # db.close()
    # print('OK')

max_stars_repo_name: spillai/pybot

max_issues_repo_name: spillai/pybot

max_forks_repo_name: spillai/pybot

reason: import numpy,import scipy


text:
import pickle
import numpy as np
import matplotlib.pyplot as plt
import sys
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt

from utils.evalResults import readData, reductionProcedures

from scipy.io import loadmat
from widerface_evaluate.bbox import bbox_overlaps
from IPython import embed

file=open("/content/drive/My Drive/RetinaFace/Pytorch_Retinaface/errorAnal/fp.pickle","rb")
a=pickle.load(file)
file.close()
errorList=["Blur","Small","Dark","Tree","occulusion","in group","dif object","non annotated face","less iou"]
# errorList=["Small","Dark","in group","not predicted","less iou"]

dtype=[(x,int) for x in errorList]
setList={setname: set() for setname in errorList}

totalWrongPredBoxes=0
for filename in a:
    nparray=a[filename]
    totalWrongPredBoxes+=len(np.where(np.sum(nparray,axis=1)>0)[0])
    for i , setname in enumerate(setList):
        ok=(np.where(nparray[:,i]==1)[0])
        for j in ok:
            setList[setname].add(filename+str(j))

evalDataFolder="/content/drive/My Drive/RetinaFace/Pytorch_Retinaface/evalData/"
fileName=evalDataFolder+"Resnet50_epoch_28_noGrad_FT_Adam_lre3"+"/outResults_val.pickle"
preds=readData(fileName)

b=[]
# for x in setList["Small"]:
for x in setList["non annotated face"]:
    imageName,no=x.split(".jpg")
    imageName=imageName+".jpg"
    no=int(no)
    # print(imageName,no)
    pred_data=preds[imageName]
    dets,predbox=reductionProcedures(pred_data,nms_threshold=0.5,confidence_threshold=0.055)
    b.append((dets[no][2]*dets[no][3]))

print(len(b))

ext: py

reason: import numpy,from scipy


text:
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
import MixMatch_pytorch.models.bottom_model_plus as models
from 即将删除的文件.liver_models import BottomModel
import numpy as np
import os

parser = argparse.ArgumentParser(description='liver framework settings')
parser.add_argument('--epochs', default=150, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('-b', '--batch-size', default=32, type=int,
                    metavar='N', help='mini-batch size (default: 128)')
parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.5, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=5e-4, type=float,
                    metavar='W', help='weight decay (default: 5e-4)')
parser.add_argument('--use-cuda', default=True, type=bool, help='use cuda device or not')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                    help='evaluate model on validation set')
parser.add_argument('--save-model-dir', dest='save_model_dir',
                    help='The directory used to save the trained models',
                    default='./saved_models/liver_saved_models', type=str)
parser.add_argument('--save-csv-dir', dest='save_csv_dir',
                    help='The directory used to save the csv files(outputs a tsne embedding)',
                    default='./csv_files/liver_csv_files', type=str)
parser.add_argument('--save-every', dest='save_every',
                    help='Saves checkpoints at every specified number of epochs',
                    type=int, default=50)
parser.add_argument('--if-cluster-outputsA', help='if_cluster_outputsA',
                    type=bool, default=True)
parser.add_argument('--path-dataset', help='path_dataset',
                    type=str, default='./datasets/HCV-Egy-Data/HCV-Egy-Data.csv')
parser.add_argument('--use-mal-optim', help='use mal optim or not',
                    type=bool, default=False)
parser.add_argument('--half', help='number of features owned by                                    the adversarial participant (int)',
                    type=int, default=11)

# 4 possible defenses on/off
parser.add_argument('--ppdl', help='turn_on_privacy_preserving_deep_learning',
                    type=bool, default=False)
parser.add_argument('--gc', help='turn_on_gradient_compression',
                    type=bool, default=False)
parser.add_argument('--lap-noise', help='turn_on_lap_noise',
                    type=bool, default=False)
parser.add_argument('--sign-sgd', help='turn_on_sign_sgd',
                    type=bool, default=False)
# setting about possible defenses
parser.add_argument('--ppdl-theta-u', help='theta-u parameter for defense privacy-preserving deep learning',
                    type=float, default=0.75)
parser.add_argument('--gc-preserved-percent', help='preserved-percent parameter for defense gradient compression',
                    type=float, default=0.75)
parser.add_argument('--noise-scale', help='noise-scale parameter for defense noisy gradients',
                    type=float, default=1e-1)

def train_per_batch(x_a, target, model, optimizer, loss=nn.CrossEntropyLoss()):
    optimizer.zero_grad()
    output = model(x_a)
    loss_value = loss(output, target)
    loss_value.backward()
    optimizer.step()
    return loss_value

def test_per_epoch(test_loader, model, loss=nn.CrossEntropyLoss()):
    test_loss = 0
    correct = 0

    right_samples_num = 0
    TP_samples_num = 0
    TN_samples_num = 0
    FP_samples_num = 0
    FN_samples_num = 0
    wrong_samples_num = 0

    # set all sub-models to eval mode.
    model.eval()

    for data, target in test_loader:
        data = data.float()
        target = target.long()
        if args.use_cuda:
            data = data.cuda()
            target = target.cuda()
        # run forward process of the whole framework
        x_a = data[:, :args.half]  # (batch size，features)
        output = model(x_a)

        # sum up batch loss
        test_loss += loss(output, target).data.item()
        # get the index of the max log-probability
        pred = output.data.max(1, keepdim=True)[1]
        # print(pred)
        target_data = target.data.view_as(pred)
        # print("target_data:", target_data)
        # total number of correctly predicted samples
        correct += pred.eq(target_data).cpu().sum()
        # number of correctly predicted positive samples
        if args.use_cuda:
            target_data = target_data.cpu()
            pred = pred.cpu()
        y_true = np.array(target_data)
        y_pred = np.array(pred)
        for i in range(len(y_true)):
            if y_true[i] == y_pred[i]:
                if y_true[i] == 1.:
                    TP_samples_num += 1
                else:
                    TN_samples_num += 1
                right_samples_num += 1
            else:
                if y_pred[i] == 1.:
                    FP_samples_num += 1
                else:
                    FN_samples_num += 1
                wrong_samples_num += 1

    total_samples_num = right_samples_num + wrong_samples_num
    print('Accuracy: ', right_samples_num / total_samples_num)
    if (TP_samples_num + FP_samples_num) != 0:
        precision = TP_samples_num / (TP_samples_num + FP_samples_num)
    else:
        precision = 0
    if (TP_samples_num + FN_samples_num) != 0:
        recall = TP_samples_num / (TP_samples_num + FN_samples_num)
    else:
        recall = 0
    if (precision + recall) != 0:
        f1 = 2 * (precision * recall) / (precision + recall)
    else:
        f1 = 0
    print('Precision:', precision, end='')
    print('  Recall:', recall, end='')
    print("  F1", f1)
    print('TP Samples Num: ', TP_samples_num)
    print('TN Samples Num: ', TN_samples_num)
    print('FP Samples Num: ', FP_samples_num)
    print('FN Samples Num: ', FN_samples_num)
    print('Right Samples Num: ', right_samples_num)
    print('Wrong Samples Num: ', wrong_samples_num)
    print('Total Samples Num: ', total_samples_num)
    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        float(correct) * 100 / len(test_loader.dataset)))
    model.train()

def main():
    global args
    args = parser.parse_args()

    # write experiment setting into file name
    setting_str = ""
    setting_str += "_"
    setting_str += "lr="
    setting_str += str(args.lr)
    if args.use_mal_optim:
        setting_str += "_"
        setting_str += "mal"
    else:
        setting_str += "_"
        setting_str += "normal"
    if args.ppdl:
        setting_str += "_"
        setting_str += "ppdl-theta_u="
        setting_str += str(args.ppdl_theta_u)
    if args.gc:
        setting_str += "_"
        setting_str += "gc-preserved_ percent="
        setting_str += str(args.gc_preserved_percent)
    if args.lap_noise:
        setting_str += "_"
        setting_str += "lap_noise-scale="
        setting_str += str(args.noise_scale)
    setting_str += "_"
    setting_str += "num_adv_feat="
    setting_str += str(args.half)
    print("settings:", setting_str)

    # Check the save_dir exists or not
    if not os.path.exists(args.save_model_dir):
        os.makedirs(args.save_model_dir)
    if not os.path.exists(args.save_csv_dir):
        os.makedirs(args.save_csv_dir)

    model = models.BottomModelPlus(4, 4)
    model.bottom_model = BottomModel(half=args.half, is_adversary=True)
    if args.use_cuda:
        model = model.cuda()

    # liver Dataset

    train_loader = liver_helper.LiverDataset(args.path_dataset, args.batch_size).get_trainloader()
    test_loader = liver_helper.LiverDataset(args.path_dataset, args.batch_size).get_testloader()
    # Data Loader (Input Pipeline)
    print("len train_loader:", len(train_loader))
    print("len test_loader:", len(test_loader))
    optimizer = optim.SGD(model.parameters(), lr=args.lr,
                          momentum=args.momentum,
                          weight_decay=args.weight_decay)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)
    for epoch in range(1, args.epochs + 1):
        scheduler.step()
        for batch_idx, (data, target) in enumerate(train_loader):
            # print("batch_idx:", batch_idx)
            data = data.float()
            target = target.long()
            if args.use_cuda:
                data = data.cuda()
                target = target.cuda()
            x_a = data[:, :args.half]
            # print("datasets:", datasets)
            loss = train_per_batch(x_a, target, model, optimizer)
            if batch_idx % 100 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                           100. * batch_idx / len(train_loader), loss.data.item()))
                with torch.no_grad():
                    test_per_epoch(test_loader=test_loader, model=model)
        with torch.no_grad():
            test_per_epoch(test_loader=test_loader, model=model)

if __name__ == '__main__':
    main()

max_stars_repo_name: jiahuanluo/label-inference-attacks

max_issues_repo_name: jiahuanluo/label-inference-attacks

max_forks_repo_name: jiahuanluo/label-inference-attacks

reason: import numpy


text:
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from matplotlib import cm

import torch
import pyro
from pyro.distributions import MultivariateNormal, Normal
from typing import List, Set, Dict, Tuple, Optional, Iterable, Mapping, Union
Dist = torch.distributions.Distribution

class GMM():
    """Mixture of Gaussians (MoG), aka. Gaussian Mixture Model (GMM) Implementation.
    This implementation handles the mode collapsing that can happen while learning *multiple* gaussian
    components under the maximum-likelihood criterion.
    Here we use the heuristic to set the mean to a random number and cov to be some large det. when a component's
    mean becomes  exactly the same as one of the data points.
    Althernatively, one can train the GMM under the maximum-a-posteioriory criterion to avoid the mode collaping

    Model parameters:
    - pi (vector of length K): [pi_1, ... pi_K]
    - means (matrix of shape (K, dim_x): Each kth row is a vector of length dim_x that represents a center of
     the kth Gaussian component distribution
    - covs (3dim tensor of shape (K, dim_x, dim_x): Each kth matrix is a covariance matrix for the kth gaussian component
    distribution

    Learning algorithm: EM method applied to Guassian Mixture Model. See PRML Ch 9.2
    - Estep: choose the posterior probability function of Z|X; \theta to be
    the one that lower-bounds the ELBO objective function tightest, given the observed
    data X and current iteration's model parameter values.

    - Mstep: Compute the new values for the model parameters s.t. the new values maximizes
    the ELBO, given the observed data and the posterior probabilities (ie. the soft mixture assignment
    for each data point)

    """

    def __init__(self, obs: Union[torch.Tensor, np.ndarray], n_mixtures: int, dtype=None,
                 device='cpu'):
        self.K = n_mixtures
        self.N, self.D = obs.shape
        self.dtype = torch.float32 if dtype is None else dtype
        self.obs = obs
        self.device = device

        # Initialize the model parameters as random initial values
        self.pi = torch.ones(self.K, dtype=self.dtype).div(self.K)  # Uniform prior on mixing coefficients
        self.means = torch.randn(self.K, self.D, dtype=self.dtype)
        self.covs = torch.zeros(self.K, self.D, self.D, dtype=self.dtype)
        for k in range(self.K):
            self.covs[k] = torch.diagflat(torch.ones(self.D, dtype=self.dtype))  # initialize with identity mtx

        # Register model parameters as in a dictionary
        # todo: as a global dictionary. Note - this is why pyro maintains a similar global dict
        #         self.params = {'pi': self.pi,
        #                        'means': self.means,
        #                        'covs': self.covs}

        # Initialize posterior probabilities for each data point: R
        self.R = torch.zeros(self.N, self.K, dtype=self.dtype)
        self.update_posterior()

    @property
    def params(self):
        return {'pi': self.pi,
                'means': self.means,
                'covs': self.covs}

    @property
    def gaussians(self):
        return make_gaussians(self.means, self.covs)

    def log_prob(self, X: Union[torch.Tensor, np.ndarray] = None, reduction: str = 'sum') -> float:
        """
        Compute the loglikelihood of the observation data X.

        Args:
        - X: Observed data
            - If it's not given, evaluate the loglikelihood of the observation data
            at initialization of this object
        - reduction (str): similar to torch.nn.MSELoss's `reduction` keyword.
            - 'mean': the llh is averaged over each llh element in the batch,
            ie. we divide the sum of loglikelihood at all observation points by number of observations

            - 'sum': the returned llh is sumed over each llh element in the batch
        """

        if X is None:
            X = self.obs

        llh = torch.zeros(self.N)
        for k, dist in enumerate(self.gaussians):
            llh += self.pi[k] * dist.log_prob(X)

        if reduction == 'mean':
            return llh.mean().item()
        elif reduction == 'sum':
            return llh.sum().item()
        else:
            raise ValueError(f"reduction must be either mean or sum: {reduction}")

    def update_posterior(self):
        """Estep: Evaluate the posterior probabilities (p(z_n|X,\theta)) given the observed data
        and the current values of the model parameters"""
        # Zero out R
        self.R = torch.zeros(self.N, self.K, dtype=self.dtype)
        normalizer = torch.zeros(self.N, 1, dtype=self.dtype)

        for k, dist in enumerate(self.gaussians):
            self.R[:, k] = self.pi[k] * dist.log_prob(self.obs)
            normalizer += self.R[:, [k]]
        self.R.div_(normalizer)

    def update_params(self):
        """Mstep: Re-estimate the model parameters using the current posterior probabilities (and
        observed data)"""
        counts = self.R.sum(dim=0)
        assert counts.shape[-1] == self.K

        # re-estimate means and covariances
        for k in range(self.K):
            # Mean
            self.means[k] = (self.R[:, [k]] * self.obs).sum(dim=0).div(counts[k])
            #             print(k, self.means[k].shape)

            # Covariance
            centered = self.obs - self.means[[k]]
            cov = torch.zeros(self.D, self.D, dtype=self.dtype, device=self.device)
            assert cov.shape == (self.D, self.D)
            for n in range(len(self.obs)):
                centered = self.obs[[n]] - self.means[[k]]
                cov += self.R[n, k] * torch.matmul(centered.T, centered)
            self.covs[k] = cov.div(counts[k])

        # Re-estimate mixture proportions
        self.pi = counts.div(self.N)

    def step(self):
        self.update_posterior()
        self.update_params()

    def info(self):
        for name, param in self.params.items():
            print(f"==={name}===")
            print(param.numpy())
        print(f"===posterior probs===")
        print(self.R)

    def visualize_model(self, title: str = 'GMM', padx: float = 5., pady: float = 5., cmap: str = 'Set1'):
        """Note:Works for 2dim gaussians only.
        Visualize the component gaussian (conditional) distributions as ellipse (w/ each dist's center
        and 1 std)

        For matplotlib.cm, refer to this example:
        - https://matplotlib.org/tutorials/colors/colormaps.html#codecell8:~:text=%23%20Indices%20to%20step%20through%20colormap,-x
        """
        if self.D != 2:
            print("Only 2Dim gaussians are supported.")
            return
        xmin, ymin = self.means.min(dim=0).values
        xmax, ymax = self.means.max(dim=0).values
        xlim = (xmin.item() - padx, xmax.item() + padx)
        ylim = (ymin.item() - pady, ymax.item() + pady)

        f, ax = plt.subplots()
        f.suptitle(title)
        # Indices to step through colormap
        cmap_inds = np.linspace(0.0, 1.0, self.K)
        colors = cm.get_cmap(plt.get_cmap(cmap))(cmap_inds)[:, :3]
        for k in range(self.K):
            #             print(k, colors[k])
            ax.scatter(self.means[k, 0], self.means[k, 1], label=f'{k}', s=10, c=colors[[k]])
            e = get_cov_ellipse(self.means[k], self.covs[k], n_std=1, fc=colors[k], alpha=0.3)
            ax.add_artist(e)
        ax.set_aspect('equal')
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        ax.legend()

        return f

    def visualize_assignments(self, title: str = 'Assignments', cmap: str = 'Set1'):
        # Indices to step through colormap
        cmap_inds = np.linspace(0.0, 1.0, self.K)
        colors = cm.get_cmap(plt.get_cmap(cmap))(cmap_inds)[:, :3]

        f, ax = plt.subplots()
        f.suptitle(title)
        for n in range(self.N):
            x = self.obs[n]
            # posterior probs {r_k}
            ratios = self.R[n]
            assert torch.isclose(ratios.sum(), torch.tensor(1.0))

            for k, ratio in enumerate(ratios):
                ax.scatter(*x, s=10, c=colors[[k]], alpha=ratio.item())
        return f

    def visualize(self, title: str = "GMM", cmap: str = 'Set1', fig_kw: Optional[Dict] = None):
        """Note:Works for 2dim gaussians only.
        Visualizes current GMM model and posterior probs (aka. soft assignments of all data)
        """
        if self.D != 2:
            print("Only 2Dim gaussians are supported.")
            return

        f, ax = plt.subplots() if fig_kw is None else plt.subplots(**fig_kw)
        f.suptitle(title)
        # Indices to step through colormap
        cmap_inds = np.linspace(0.0, 1.0, self.K)
        colors = cm.get_cmap(plt.get_cmap(cmap))(cmap_inds)[:, :3]

        # Draw gaussian dists as ellipses
        for k in range(self.K):
            ax.scatter(self.means[k, 0], self.means[k, 1], label=f'{k}', s=10, c=colors[[k]])
            e = get_cov_ellipse(self.means[k], self.covs[k], n_std=1, fc=colors[k], alpha=0.3)
            ax.add_artist(e)

        # Show each data point's soft assignmenets
        for n in range(self.N):
            x = self.obs[n]
            # posterior probs {r_k}
            ratios = self.R[n]
            assert torch.isclose(ratios.sum(), torch.tensor(1.0))

            for k, ratio in enumerate(ratios):
                ax.scatter(*x, s=10, c=colors[[k]], alpha=ratio.item())

        ax.set_aspect('equal', 'datalim')
        ax.legend()
        return f

def make_guassian(mean: torch.Tensor, cov: torch.Tensor) -> Dist:
    return MultivariateNormal(mean, cov)

def make_gaussians(means: torch.Tensor, covs: torch.Tensor) -> Iterable[Dist]:
    ds = []
    K = means.shape[-1]
    for k in range(K):
        ds.append(MultivariateNormal(means[k], covs[k]))
    return ds

def get_cov_ellipse(mean, cov, n_std, **kwargs):
    """
    Return a matplotlib Ellipse patch representing the covariance matrix
    cov centred at centre and scaled by the factor n_std.

    Usage:
    e = get_cov_ellipse(cov, (height_mean, mass_mean), 3,
                        fc=colours[gender], alpha=0.4)
    ax.add_artist(e)
    """
    if isinstance(mean, np.ndarray) and mean.ndim == 2:
        mean = mean.flatten()

    # Find and sort eigenvalues and eigenvectors into descending order
    eigvals, eigvecs = np.linalg.eigh(cov)
    order = eigvals.argsort()[::-1]
    eigvals, eigvecs = eigvals[order], eigvecs[:, order]

    # The anti-clockwise angle to rotate our ellipse by
    vx, vy = eigvecs[:, 0][0], eigvecs[:, 0][1]
    theta = np.arctan2(vy, vx)

    # Width and height of ellipse to draw
    width, height = 2 * n_std * np.sqrt(eigvals)
    return Ellipse(xy=mean, width=width, height=height,
                   angle=np.degrees(theta), **kwargs)

max_stars_repo_name: cocoaaa/ReprLearn

max_issues_repo_name: cocoaaa/ReprLearn

max_forks_repo_name: cocoaaa/ReprLearn


text:
import numpy as np

def fpudl(initial=32,damp=0.9,a=0,b=0,dt=0.001,show=True):
    if type(initial) == int:
        initial = np.zeros(initial)
    if type(initial) == list:
        initial = np.array(initial)

    x = initial
    v = np.zeros(x.shape)
    if show:
        import matplotlib.animation as animation
        from matplotlib import pyplot as plt
        fig, ax = plt.subplots()
        xp, = ax.plot(range(len(x)),x,'o')
        vp, = ax.plot(range(len(x)),v,'o')
        ax.set_ylim(-1.3,1.3)

    crt = 2**(1/3)
    den = (2*(2-crt))
    c = [1/den,(1-crt)/den,(1-crt)/den,1/den]
    d = [2/den,-2*crt/den,2/den,0]

    if show:
        def anim(i,x=x,v=v):
            xp.set_ydata(x)
            vp.set_ydata(v)
            return xp,vp,
        ani = animation.FuncAnimation(fig,anim,interval=1)
        dgc = (ani,anim)
    else:
        dgc = None
    if show:
        plt.show(block=0)
    def fpu_delay_line(inp,dt=dt,damp=damp,x=x,v=v,a=a,b=b,c=c,d=d,dgc=dgc):
        #step with symplectic integration
        def accel():
            shifted_p = np.roll(x,1)
            shifted_p[0] = inp.real
            shifted_n = np.roll(x,-1)
            shifted_n[-1] = 0
            d_p = shifted_p - x
            d_n = shifted_n - x
            sq = [d_p*d_p,d_n*d_n]
            return d_p+d_n + a*(sq[0]-sq[1]) +\
                b*(sq[0]*d_p+sq[1]*d_n)
        #now can velocity verlet
        for i in range(len(c)):
            x += v*dt*c[i]
            v += accel()*dt*d[i]
        v[-1] *= damp
        out = (lambda n: n+a*n*n+b*n*n*n)(x[-1])
        return out

    return fpu_delay_line
def fpu(initial=32,h=1,a=0,b=0,dt=0.001,simrate=0.001,show=True):
    if type(initial) == int:
        initial = np.sin(h*(np.arange(initial)+1)/(initial+1) * np.pi)
    if type(initial) == list:
        initial = np.array(initial)

    x = initial
    v = np.zeros(x.shape)
    if show:
        import matplotlib.animation as animation
        from matplotlib import pyplot as plt
        fig, ax = plt.subplots()
        xp, = ax.plot(range(len(x)),x,'o')
        vp, = ax.plot(range(len(x)),v,'o')
        ax.set_ylim(-1.3,1.3)

    crt = 2**(1/3)
    den = (2*(2-crt))
    c = [1/den,(1-crt)/den,(1-crt)/den,1/den]
    d = [2/den,-2*crt/den,2/den,0]
    def update(dt=dt,x=x,v=v,a=a,b=b,c=c,d=d):
        #step with symplectic integration
        def accel():
            shifted_p = np.roll(x,1)
            shifted_p[0] = 0
            shifted_n = np.roll(x,-1)
            shifted_n[-1] = 0
            d_p = shifted_p - x
            d_n = shifted_n - x
            sq = [d_p*d_p,d_n*d_n]
            return d_p+d_n + a*(sq[0]-sq[1]) +\
                b*(sq[0]*d_p+sq[1]*d_n)
        #now can velocity verlet
        for i in range(len(c)):
            x += v*dt*c[i]
            v += accel()*dt*d[i]
    if show:
        def anim(i,x=x,v=v):
            xp.set_ydata(x)
            vp.set_ydata(v)
            return xp,vp,
        ani = animation.FuncAnimation(fig,anim,interval=1)
        dgc = (ani,anim)
    else:
        dgc = None
    def gen(dontgc = dgc,simrate=simrate,dt=dt):
        time = 0
        while 1:
            yield 0
            for i in range(len(x)):
                yield x[i]+1j*v[i]
            time += simrate/2
            while time > dt:
                update()
                time -= dt
            yield 0
            for i in range(len(x)-1,-1,-1):
                yield -x[i]-1j*v[i]
            time += simrate/2
            while time > dt:
                update()
                time -= dt
    if show:
        plt.show(block=0)
    return gen()

def fpuv(initial=32,h=1,a=0,b=0,dt=0.001):
    if type(initial) == int:
        initial = np.sin(h*(np.arange(initial)+1)/(initial+1) * np.pi)
    if type(initial) == list:
        initial = np.array(initial)

    x = initial
    v = np.zeros(x.shape)
    import matplotlib.animation as animation
    from matplotlib import pyplot as plt
    fig, ax = plt.subplots()
    xp, = ax.plot(range(len(x)),x,'o')
    vp, = ax.plot(range(len(x)),v,'o')
    ax.set_ylim(-1.3,1.3)

    def anim(i,x=x,v=v,dt=dt,a=a,b=b):
        #step with velocity verlet
        def accel():
            shifted_p = np.roll(x,1)
            shifted_p[0] = 0
            shifted_n = np.roll(x,-1)
            shifted_n[-1] = 0
            d_p = shifted_p - x
            d_n = shifted_n - x
            sq = [d_p*d_p,d_n*d_n]
            return d_p+d_n + a*(sq[0]-sq[1]) +\
                b*(sq[0]*d_p+sq[1]*d_n)
        #now can velocity verlet
        a = accel()
        x += v*dt+.5*dt*dt*a
        v += .5*(a+accel())

        xp.set_ydata(x)
        vp.set_ydata(v)
        return xp,vp,

    ani = animation.FuncAnimation(fig,anim,interval=1)
    plt.show(block=0)
    return ani,anim


text:
import numpy as np
import os
import matplotlib.pyplot as plt
import math
import AnalysisTools as at

## Reading probes into the file
folder = "output"
os.chdir(folder)
probes = 'record_rectangle1.csv'

datalist = at.readProbeFile(probes)
probeType = datalist[0]
time = datalist[1]
data = datalist[2]
rotq_e3 = data[:,7]

alpha = []
for i in range(0,len(rotq_e3)):
    alpha.append(2*math.asin(rotq_e3[i]))
    alpha[i] *= 360/(2*math.pi)
alpha = np.array(alpha)

it = np.where(time>2.5)[0][0]
period = at.zeroCrossing(time[:it],alpha[:it],up=False)[0]

period_ref = 0.93
err = 100*abs(period_ref-period)/abs(period_ref)
val = open('validation_FCC.txt', 'w')
val.write('Period for the rotation angle'+'\n')
val.write('Theory'+'\t'+'Simulation'+'\t'+'\t'+'Error (%)'+'\n')
val.write(str(period_ref)+'\t'+str(period)+'\t'+str(err))
val.close()

plt.plot(time,alpha)
plt.xlabel('time [sec]')
plt.ylabel('rotation angle [deg]')
plt.suptitle('Rotation angle for the floating caisson 2D')
plt.xlim((0.,4.5))
plt.grid(True)
plt.savefig('rotation_angle.png')
plt.show()

max_stars_repo_name: AntoninoScala/air-water-vv

max_issues_repo_name: AntoninoScala/air-water-vv

max_forks_repo_name: AntoninoScala/air-water-vv


text:
import os
import numpy as np
import astropy.io.fits as fits

from .base import interpolate_data, interpolate_param
from ..utils.download import get_file

class _YaPSI(object):
    '''YaPSI stellar evolution tracks.

    '''
    _y_nodes   = [0.25, 0.28, 0.31, 0.34, 0.37]
    _feh_nodes = [-1.5, -1.0, -0.5, 0.0, +0.3]
    _xz_nodes  = (
        ((0.749455,0.000545),(0.719477,0.000523),(0.689499,0.000501),
         (0.659520,0.000480),(0.629542,0.000458)), # [Fe/H] = -1.5
        ((0.748279,0.001721),(0.718348,0.001652),(0.688417,0.001583),
         (0.658485,0.001515),(0.628554,0.001446)), # [Fe/H] = -1.0
        ((0.744584,0.005416),(0.714801,0.005199),(0.685018,0.004982),
         (0.655234,0.004766),(0.625451,0.004549)), # [Fe/H] = -0.5
        ((0.733138,0.016862),(0.703812,0.016188),(0.674487,0.015513),
         (0.645161,0.014839),(0.615836,0.014164)), # [Fe/H] = +0.0
        ((0.717092,0.032908),(0.688408,0.031592),(0.659725,0.030275),
         (0.631041,0.028959),(0.602357,0.027643)), # [Fe/H] = +0.3
    )
    _mass_nodes = np.concatenate((
        np.arange(0.15, 0.40, 0.01),
        np.arange(0.40, 0.90, 0.02),
        np.arange(0.90, 1.80, 0.05),
        np.arange(1.80, 3.00, 0.10),
        np.arange(3.00, 5.01, 0.20),
    ))

    # amlt: alpha mixing length
    _amlt1, _amlt2 = 1.82126, 1.91804

    def __init__(self):
        self._track_data = {}

    def _load_tracks(self):
        '''Read evoution tracks.
        '''

    def _get_trackid(self, y, feh, mass, amlt):
        '''Get Track ID.

        Args:
            y (float): Initial helium content.
            feh (float): Metallcity.
            mass (float): Stellar mass.
            amlt (float): Mixing length.
        Returns:
            tuple: A tuple of four integers (`y_id`, `feh_id`, `mass_id`,
                `alpha_id`) as the track ID.
        '''
        y_id    = int(round(y*100))
        feh_id  = int(round(feh*10))
        mass_id = int(round(mass*100))
        if abs(amlt-self._amlt1)<1e-5:
            amlt_id = 1
        elif abs(amlt-self._amlt2)<1e-5:
            amlt_id = 2
        else:
            print('Alpha ID does not exist')
            raise ValueError

        return (y_id, feh_id, mass_id, amlt_id)

    def get_track(self, y, feh, mass, n=0, minage=0):
        '''
        Get the evolution track for given *Y*, [Fe/H] and *M*.

        Args:
            y (float):
            feh (float): Metallicity.
            mass (float): Stellar mass.
            n (integer): Number of points in evolution track.
            minage (float): Minimum age returned in the evolution track.

        Returns:
            tuple: A tuple containing:

        '''

        if y in self._y_nodes:
            track = self._get_track_of_feh(y, feh, mass, n, minage)
        else:
            track_lst = []
            iy = self._get_inodes(self._y_nodes, y)
            y_lst = self._y_nodes[iy:iy+4]
            for _y in y_lst:
                track = self._get_track_of_feh(_y, feh, mass, n, minage)
                track_lst.append(track)
            track = interpolate_param(track_lst, y_lst, y)

        return track

    def _get_inodes(self, nodes, value):
        '''Get the begining index of the 4-points interpolation.

        Args:
            nodes (list): Input node list.
            value (intger or float): Input value.
        Returns:
            integer: Beginning index of 4-points interpolation.
        '''
        i0 = np.searchsorted(nodes, value)
        i = i0-2
        i = max(i, 0)
        i = min(i, len(nodes)-4)
        return i

    def _load_track(self, y, feh, mass, amlt):
        '''
        Load the track of given *y* and *feh* values and pack it into the cache.

        Args:
            y (float): Helium content (*Y*).
            feh (float): Metallicity ([Fe/H]).
            mass (float): Stellar mass (*M*).
            amlt (float): Mixing-length (*alpha*).

        Notes:
            The track data is a tuple containing four arrays
            (log\ *T*:sub:`eff`, log\ *L*, age, log\ *g*).

        '''
        iy   = self._y_nodes.index(y)
        ifeh = self._feh_nodes.index(feh)
        x,z  = self._xz_nodes[ifeh][iy]

        data_path  = 'thirdpartydata/yapsi'
        folder = 'X{:8.6f}_Z{:8.6f}'.format(x, z).replace('.','p')
        fname = 'M{:4.2f}_X{:8.6f}_Z{:8.6f}_A{:7.5f}'.format(
                mass, x, z, amlt).replace('.', 'p')+'.trk'
        filepath = os.path.join(data_path, folder, fname)
        filename = get_file(filepath)

        # load data
        data = []
        file1 = open(filename)
        for row in file1:
            if row[0]=='#':
                continue
            col = row.split()
            age     = float(col[2])
            logL    = float(col[6])
            logR    = float(col[7])
            logg    = float(col[8])
            logTeff = float(col[9])
            item = (logTeff, logL, age, logR, logg)
            data.append(item)
        file1.close()

        # repack parameters
        logTeff_lst = np.array([row[0] for row in data])
        logL_lst    = np.array([row[1] for row in data])
        age_lst     = np.array([row[2] for row in data])
        logR_lst    = np.array([row[3] for row in data])
        logg_lst    = np.array([row[4] for row in data])

        track = (logTeff_lst, logL_lst, age_lst, logg_lst, logR_lst)
        return track

    def _get_track_of_mass(self, y, feh, mass, n=0, minage=0):
        if y not in self._y_nodes:
            print('Error: Y = {} not in y_nodes'.format(y))
            raise ValueError
        if feh not in self._feh_nodes:
            print('Error: [Fe/H] = {} not in feh_nodes'.format(feh))
            raise ValueError

        # check if input mass is in mass nodes
        m = np.abs(self._mass_nodes - mass)<1e-3
        if m.sum()>0:
            amlt = (self._amlt1, self._amlt2)[mass<=1.1]
            trackid = self._get_trackid(y, feh, mass, amlt)
            if trackid in self._track_data:
                track = self._track_data[trackid]
            else:
                track = self._load_track(y, feh, mass, amlt)
                self._track_data[trackid] = track

            if minage!=0:
                m = track[2] > minage
                track = (track[0][m], track[1][m], track[2][m],
                         track[3][m], track[4][m])
            if n>0:
                track = interpolate_data(track, n)
        else:
            track_lst = []
            imass = self._get_inodes(self._mass_nodes, mass)
            mass_lst = self._mass_nodes[imass:imass+4]
            for _mass in mass_lst:
                amlt = (self._amlt1, self._amlt2)[_mass<=1.1]

                trackid = self._get_trackid(y, feh, _mass, amlt)
                if trackid in self._track_data:
                    track = self._track_data[trackid]
                else:
                    track = self._load_track(y, feh, _mass, amlt)
                    self._track_data[trackid] = track

                if minage!=0:
                    m = track[2] > minage
                    track = (track[0][m], track[1][m], track[2][m],
                             track[3][m], track[4][m])
                if n>0:
                    track = interpolate_data(track, n)
                track_lst.append(track)
            track = interpolate_param(track_lst, mass_lst, mass)
        return track

    def _get_track_of_feh(self, y, feh, mass, n=0, minage=0):
        if y not in self._y_nodes:
            print('Error: Y = %g not in y_nodes')
            raise ValueError

        if feh in self._feh_nodes:
            # load missing tracks
            #trackid0 = self._get_trackid(y, feh, 1.0, self._amlt1)
            #if trackid0 not in self._track_data:
            #    self._load_track(y, feh, mass, amlt)
            track = self._get_track_of_mass(y, feh, mass, n, minage)
        else:
            # need interpolation over feh
            track_lst = []
            ifeh = self._get_inodes(self._feh_nodes, feh)
            feh_lst = self._feh_nodes[ifeh:ifeh+4]
            for _feh in feh_lst:
                # load missing tracks
                #trackid0 = self._get_trackid(y, _feh, 1.0, self._amlt1)
                #if trackid0 not in self._track_data:
                #    self._load_track(y, _feh)
                track = self._get_track_of_mass(y, _feh, mass, n, minage)
                track_lst.append(track)
            track = interpolate_param(track_lst, feh_lst, feh)
        return track

YaPSI = _YaPSI()

max_stars_repo_name: wangleon/stella

max_issues_repo_name: wangleon/stella

max_forks_repo_name: wangleon/stella

reason: import numpy,import astropy


text:
"""
Test ants_image.py

nptest.assert_allclose
self.assertEqual
self.assertTrue

"""

import os
import unittest
from common import run_tests

from tempfile import mktemp

import numpy as np
import numpy.testing as nptest
import pandas as pd

import ants

class TestClass_LabelImage(unittest.TestCase):
    """
    Test ants.ANTsImage class
    """
    def setUp(self):
        square2d = np.ones((20,20))
        square2d[:10,:] += 1
        square2d[:,:10] += 2
        img2d = ants.from_numpy(square2d).astype('uint8')
        info2d = []
        for i in img2d.unique():
            info2d.append(['Value=%i'%i, 'ModTwo=%i'%int(i%2), 'ModThree=%i'%int(i%3)])
        info2d = pd.DataFrame(info2d, index=img2d.unique(), columns=['Value','ModTwo','ModThree'])

        self.img2d = img2d
        self.info2d = info2d

        square3d = np.ones((20,20,20))
        square3d[:10,...] += 1
        square3d[:,:10,:] += 2
        square3d[...,:10] += 3
        img3d = ants.from_numpy(square3d).astype('uint8')

        info3d = []
        for i in img3d.unique():
            info3d.append(['Value=%i'%i, 'ModTwo=%i'%int(i%2), 'ModThree=%i'%int(i%3)])
        info3d = pd.DataFrame(info2d, index=img2d.unique(), columns=['Value','ModTwo','ModThree'])

        self.img3d = img3d
        self.info3d = info3d

    def tearDown(self):
        pass

    def test_init(self):
        lbl_img2d = ants.LabelImage(label_image=self.img2d, label_info=self.info2d)
        lbl_img3d = ants.LabelImage(label_image=self.img2d, label_info=self.info2d)

    def test_init_with_template(self):
        arr2d = np.abs(np.random.randn(*self.img2d.shape))
        template2d = ants.from_numpy(arr2d)
        lbl_img2d = ants.LabelImage(label_image=self.img2d, label_info=self.info2d,
                                    template=template2d)

        arr3d = np.abs(np.random.randn(*self.img3d.shape))
        template3d = ants.from_numpy(arr3d)
        lbl_img3d = ants.LabelImage(label_image=self.img3d, label_info=self.info3d,
                                    template=template3d)

    def test__repr__(self):
        lbl_img2d = ants.LabelImage(label_image=self.img2d, label_info=self.info2d)
        lbl_img3d = ants.LabelImage(label_image=self.img2d, label_info=self.info2d)

        r1 = lbl_img2d.__repr__()
        r2 = lbl_img3d.__repr__()
        self.assertTrue(isinstance(r1,str))
        self.assertTrue(isinstance(r2,str))

if __name__ == '__main__':
    run_tests()

max_stars_repo_name: xemio/ANTsPy

max_issues_repo_name: xemio/ANTsPy

max_forks_repo_name: xemio/ANTsPy


text:
import numpy as np

def ode0(z,UJet,a):

	# first order formulation of ODE

	# z = [u0,u1,ud0,ud1]
	u = z[:2]
	ud = z[2:]

	U = np.zeros((2,2))

	U[1,0] = UJet[1][0](u)
	U[0,1] = UJet[0][1](u)

	r = np.array([U[1,0],U[0,1]])

	return np.concatenate([ud,r])

def ode2(z,UJet,a,h):

	# first order formulation of modified ODE 2

	u = z[:2]
	ud = z[2:]

	U = np.zeros((4,4))

	U[1,0] = UJet[1][0](u)
	U[0,1] = UJet[0][1](u)

	U[0,2] = UJet[0][2](u)
	U[0,3] = UJet[0][3](u)

	U[1,1] = UJet[1][1](u)
	U[1,2] = UJet[1][2](u)

	U[2,1] = UJet[2][1](u)

	U[3,0] = UJet[3][0](u)

	xd1 = ud[0]
	xd2 = ud[1]

	rhs = np.array([(3*a[1,2]*(U[0,1]*U[0,2] + xd2**2*U[0,3] + U[1,0]*U[1,1] + 2*xd1*xd2*U[1,2] + xd1**2*U[2,1]) +
	(-4 + 3*a[1,1])*(U[0,1]*U[1,1] + xd2**2*U[1,2] + U[1,0]*U[2,0] + 2*xd1*xd2*U[2,1] + xd1**2*U[3,0]))/12.,
	(-4*(U[0,1]*U[0,2] + xd2**2*U[0,3] + U[1,0]*U[1,1] + 2*xd1*xd2*U[1,2] + xd1**2*U[2,1]) +
	3*a[2,2]*(U[0,1]*U[0,2] + xd2**2*U[0,3] + U[1,0]*U[1,1] + 2*xd1*xd2*U[1,2] + xd1**2*U[2,1]) +
	3*a[1,2]*(U[0,1]*U[1,1] + xd2**2*U[1,2] + U[1,0]*U[2,0] + 2*xd1*xd2*U[2,1] + xd1**2*U[3,0]))/12.])

	r = ode0(z,UJet,a)[2:] + h**2*rhs

	return np.concatenate([ud,r])

def ode4(z,UJet,a,h):

	# first order formulation of modified ODE 4

	u = z[:2]
	ud = z[2:]

	U = np.zeros((6,6))

	U[0,1] = UJet[0][1](u)
	U[0,2] = UJet[0][2](u)
	U[0,3] = UJet[0][3](u)
	U[0,4] = UJet[0][4](u)
	U[0,5] = UJet[0][5](u)

	U[1,0] = UJet[1][0](u)
	U[1,1] = UJet[1][1](u)
	U[1,2] = UJet[1][2](u)
	U[1,3] = UJet[1][3](u)
	U[1,4] = UJet[1][4](u)

	U[2,0] = UJet[2][0](u)
	U[2,1] = UJet[2][1](u)
	U[2,2] = UJet[2][2](u)
	U[2,3] = UJet[2][3](u)

	U[3,0] = UJet[3][0](u)
	U[3,1] = UJet[3][1](u)
	U[3,2] = UJet[3][2](u)

	U[4,0] = UJet[4][0](u)
	U[4,1] = UJet[4][1](u)

	U[5,0] = UJet[5][0](u)

	xd1 = ud[0]
	xd2 = ud[1]

	rhs =     np.array([(128*xd2**2*U[0,3]*U[1,1] - 150*xd2**2*a[1,1]*U[0,3]*U[1,1] + 45*xd2**2*a[1,1]**2*U[0,3]*U[1,1] - 60*xd2**2*a[2,2]*U[0,3]*U[1,1] +
	45*xd2**2*a[1,1]*a[2,2]*U[0,3]*U[1,1] + 128*U[1,0]*U[1,1]**2 - 150*a[1,1]*U[1,0]*U[1,1]**2 + 45*a[1,1]**2*U[1,0]*U[1,1]**2 - 60*a[2,2]*U[1,0]*U[1,1]**2 +
	45*a[1,1]*a[2,2]*U[1,0]*U[1,1]**2 + 9*(16 - 30*a[1,1] + 15*a[1,1]**2)*U[0,1]**2*U[1,2] + 192*xd2**2*U[0,2]*U[1,2] - 360*xd2**2*a[1,1]*U[0,2]*U[1,2] +
	180*xd2**2*a[1,1]**2*U[0,2]*U[1,2] + 448*xd1*xd2*U[1,1]*U[1,2] - 660*xd1*xd2*a[1,1]*U[1,1]*U[1,2] + 270*xd1*xd2*a[1,1]**2*U[1,1]*U[1,2] -
	120*xd1*xd2*a[2,2]*U[1,1]*U[1,2] + 90*xd1*xd2*a[1,1]*a[2,2]*U[1,1]*U[1,2] + 48*xd2**4*U[1,4] - 90*xd2**4*a[1,1]*U[1,4] + 45*xd2**4*a[1,1]**2*U[1,4] +
	128*xd2**2*U[1,2]*U[2,0] - 210*xd2**2*a[1,1]*U[1,2]*U[2,0] + 90*xd2**2*a[1,1]**2*U[1,2]*U[2,0] + 128*U[1,0]*U[2,0]**2 - 210*a[1,1]*U[1,0]*U[2,0]**2 +
	90*a[1,1]**2*U[1,0]*U[2,0]**2 + 192*xd1*xd2*U[0,2]*U[2,1] - 360*xd1*xd2*a[1,1]*U[0,2]*U[2,1] + 180*xd1*xd2*a[1,1]**2*U[0,2]*U[2,1] +
	320*xd1**2*U[1,1]*U[2,1] + 192*xd2**2*U[1,1]*U[2,1] - 510*xd1**2*a[1,1]*U[1,1]*U[2,1] - 360*xd2**2*a[1,1]*U[1,1]*U[2,1] +
	225*xd1**2*a[1,1]**2*U[1,1]*U[2,1] + 180*xd2**2*a[1,1]**2*U[1,1]*U[2,1] - 60*xd1**2*a[2,2]*U[1,1]*U[2,1] + 45*xd1**2*a[1,1]*a[2,2]*U[1,1]*U[2,1] +
	448*xd1*xd2*U[2,0]*U[2,1] - 780*xd1*xd2*a[1,1]*U[2,0]*U[2,1] + 360*xd1*xd2*a[1,1]**2*U[2,0]*U[2,1] + 288*xd2**2*U[1,0]*U[2,2] -
	540*xd2**2*a[1,1]*U[1,0]*U[2,2] + 270*xd2**2*a[1,1]**2*U[1,0]*U[2,2] + 192*xd1*xd2**3*U[2,3] - 360*xd1*xd2**3*a[1,1]*U[2,3] +
	180*xd1*xd2**3*a[1,1]**2*U[2,3] + 144*U[1,0]**2*U[3,0] - 270*a[1,1]*U[1,0]**2*U[3,0] + 135*a[1,1]**2*U[1,0]**2*U[3,0] + 192*xd1*xd2*U[1,1]*U[3,0] -
	360*xd1*xd2*a[1,1]*U[1,1]*U[3,0] + 180*xd1*xd2*a[1,1]**2*U[1,1]*U[3,0] + 320*xd1**2*U[2,0]*U[3,0] - 570*xd1**2*a[1,1]*U[2,0]*U[3,0] +
	270*xd1**2*a[1,1]**2*U[2,0]*U[3,0] + 576*xd1*xd2*U[1,0]*U[3,1] - 1080*xd1*xd2*a[1,1]*U[1,0]*U[3,1] + 540*xd1*xd2*a[1,1]**2*U[1,0]*U[3,1] +
	U[0,1]*((128 + 45*a[1,1]**2 - 60*a[2,2] + 15*a[1,1]*(-10 + 3*a[2,2]))*U[0,2]*U[1,1] +
	2*(9*xd2**2*(16 - 30*a[1,1] + 15*a[1,1]**2)*U[1,3] + (64 - 105*a[1,1] + 45*a[1,1]**2)*U[1,1]*U[2,0] +
	18*xd1*xd2*(16 - 30*a[1,1] + 15*a[1,1]**2)*U[2,2] + 9*(16 - 30*a[1,1] + 15*a[1,1]**2)*(U[1,0]*U[2,1] + xd1**2*U[3,1]))) + 288*xd1**2*xd2**2*U[3,2] -
	540*xd1**2*xd2**2*a[1,1]*U[3,2] + 270*xd1**2*xd2**2*a[1,1]**2*U[3,2] + 288*xd1**2*U[1,0]*U[4,0] - 540*xd1**2*a[1,1]*U[1,0]*U[4,0] +
	270*xd1**2*a[1,1]**2*U[1,0]*U[4,0] + 192*xd1**3*xd2*U[4,1] - 360*xd1**3*xd2*a[1,1]*U[4,1] + 180*xd1**3*xd2*a[1,1]**2*U[4,1] +
	15*a[1,2]*(9*(-2 + a[1,1] + a[2,2])*U[0,1]**2*U[0,3] + 3*xd2**4*(-2 + a[1,1] + a[2,2])*U[0,5] - 10*U[0,2]*U[1,0]*U[1,1] + 3*a[1,1]*U[0,2]*U[1,0]*U[1,1] +
	6*a[2,2]*U[0,2]*U[1,0]*U[1,1] - 24*xd1**2*U[1,1]*U[1,2] + 12*xd1**2*a[1,1]*U[1,1]*U[1,2] + 12*xd1**2*a[2,2]*U[1,1]*U[1,2] +
	12*xd1*xd2**3*(-2 + a[1,1] + a[2,2])*U[1,4] - 18*U[1,0]*U[1,1]*U[2,0] + 12*a[1,1]*U[1,0]*U[1,1]*U[2,0] + 3*a[2,2]*U[1,0]*U[1,1]*U[2,0] -
	10*xd1**2*U[0,2]*U[2,1] + 3*xd1**2*a[1,1]*U[0,2]*U[2,1] + 6*xd1**2*a[2,2]*U[0,2]*U[2,1] - 18*U[1,0]**2*U[2,1] + 9*a[1,1]*U[1,0]**2*U[2,1] +
	9*a[2,2]*U[1,0]**2*U[2,1] - 28*xd1**2*U[2,0]*U[2,1] + 15*xd1**2*a[1,1]*U[2,0]*U[2,1] + 12*xd1**2*a[2,2]*U[2,0]*U[2,1] +
	U[0,1]*((-10 + 3*a[1,1] + 6*a[2,2])*U[0,2]**2 + 18*xd2**2*(-2 + a[1,1] + a[2,2])*U[0,4] - 14*U[1,1]**2 + 9*a[1,1]*U[1,1]**2 + 3*a[2,2]*U[1,1]**2 -
	36*U[1,0]*U[1,2] + 18*a[1,1]*U[1,0]*U[1,2] + 18*a[2,2]*U[1,0]*U[1,2] + 36*xd1*xd2*(-2 + a[1,1] + a[2,2])*U[1,3] + (-4 + 3*a[1,1])*U[0,2]*U[2,0] -
	36*xd1**2*U[2,2] + 18*xd1**2*a[1,1]*U[2,2] + 18*xd1**2*a[2,2]*U[2,2]) +
	xd2**2*((-34 + 15*a[1,1] + 18*a[2,2])*U[0,2]*U[0,3] + (-38 + 21*a[1,1] + 15*a[2,2])*U[1,1]*U[1,2] - 36*U[1,0]*U[1,3] + 18*a[1,1]*U[1,0]*U[1,3] +
	18*a[2,2]*U[1,0]*U[1,3] - 4*U[0,3]*U[2,0] + 3*a[1,1]*U[0,3]*U[2,0] - 36*xd1**2*U[2,3] + 18*xd1**2*a[1,1]*U[2,3] + 18*xd1**2*a[2,2]*U[2,3]) -
	14*xd1**2*U[1,1]*U[3,0] + 9*xd1**2*a[1,1]*U[1,1]*U[3,0] + 3*xd1**2*a[2,2]*U[1,1]*U[3,0] - 36*xd1**2*U[1,0]*U[3,1] + 18*xd1**2*a[1,1]*U[1,0]*U[3,1] +
	18*xd1**2*a[2,2]*U[1,0]*U[3,1] + 2*xd1*xd2*(6*(-2 + a[1,1] + a[2,2])*U[0,3]*U[1,1] + (-22 + 9*a[1,1] + 12*a[2,2])*U[0,2]*U[1,2] - 16*U[1,2]*U[2,0] +
	9*a[1,1]*U[1,2]*U[2,0] + 6*a[2,2]*U[1,2]*U[2,0] - 26*U[1,1]*U[2,1] + 15*a[1,1]*U[1,1]*U[2,1] + 9*a[2,2]*U[1,1]*U[2,1] - 36*U[1,0]*U[2,2] +
	18*a[1,1]*U[1,0]*U[2,2] + 18*a[2,2]*U[1,0]*U[2,2] - 12*xd1**2*U[3,2] + 6*xd1**2*a[1,1]*U[3,2] + 6*xd1**2*a[2,2]*U[3,2]) - 6*xd1**4*U[4,1] +
	3*xd1**4*a[1,1]*U[4,1] + 3*xd1**4*a[2,2]*U[4,1]) + 48*xd1**4*U[5,0] - 90*xd1**4*a[1,1]*U[5,0] + 45*xd1**4*a[1,1]**2*U[5,0] +
	45*a[1,2]**2*(2*U[1,0]*U[1,1]**2 + 3*U[0,1]**2*U[1,2] + xd2**4*U[1,4] + U[0,2]*U[1,0]*U[2,0] + U[1,0]*U[2,0]**2 + 6*xd1**2*U[1,1]*U[2,1] +
	4*xd1*xd2**3*U[2,3] + xd1**2*U[0,2]*U[3,0] + 3*U[1,0]**2*U[3,0] + 5*xd1**2*U[2,0]*U[3,0] +
	U[0,1]*(3*U[0,2]*U[1,1] + 6*xd2**2*U[1,3] + U[1,1]*U[2,0] + 6*U[1,0]*U[2,1] + 12*xd1*xd2*U[2,2] + 6*xd1**2*U[3,1]) +
	xd2**2*(2*U[0,3]*U[1,1] + 5*U[0,2]*U[1,2] + U[1,2]*U[2,0] + 4*U[1,1]*U[2,1] + 6*U[1,0]*U[2,2] + 6*xd1**2*U[3,2]) + 6*xd1**2*U[1,0]*U[4,0] +
	2*xd1*xd2*(3*U[0,2]*U[2,1] + 3*U[2,0]*U[2,1] + 2*U[1,1]*(2*U[1,2] + U[3,0]) + 6*U[1,0]*U[3,1] + 2*xd1**2*U[4,1]) + xd1**4*U[5,0]))/720.,
	(128*U[0,2]*U[1,0]*U[1,1] + 45*a[1,2]**2*U[0,2]*U[1,0]*U[1,1] - 210*a[2,2]*U[0,2]*U[1,0]*U[1,1] + 90*a[2,2]**2*U[0,2]*U[1,0]*U[1,1] -
	210*a[1,2]*U[1,0]*U[1,1]**2 + 45*a[1,1]*a[1,2]*U[1,0]*U[1,1]**2 + 135*a[1,2]*a[2,2]*U[1,0]*U[1,1]**2 + 192*xd1**2*U[1,1]*U[1,2] +
	180*xd1**2*a[1,2]**2*U[1,1]*U[1,2] - 360*xd1**2*a[2,2]*U[1,1]*U[1,2] + 180*xd1**2*a[2,2]**2*U[1,1]*U[1,2] +
	9*U[0,1]**2*((16 + 15*a[1,2]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[0,3] + 15*a[1,2]*(-2 + a[1,1] + a[2,2])*U[1,2]) +
	xd2**4*((48 + 45*a[1,2]**2 - 90*a[2,2] + 45*a[2,2]**2)*U[0,5] + 45*a[1,2]*(-2 + a[1,1] + a[2,2])*U[1,4]) - 60*a[1,2]*U[0,2]*U[1,0]*U[2,0] +
	45*a[1,2]*a[2,2]*U[0,2]*U[1,0]*U[2,0] + 128*U[1,0]*U[1,1]*U[2,0] - 60*a[1,1]*U[1,0]*U[1,1]*U[2,0] + 135*a[1,2]**2*U[1,0]*U[1,1]*U[2,0] -
	150*a[2,2]*U[1,0]*U[1,1]*U[2,0] + 45*a[1,1]*a[2,2]*U[1,0]*U[1,1]*U[2,0] + 45*a[2,2]**2*U[1,0]*U[1,1]*U[2,0] - 150*a[1,2]*U[1,0]*U[2,0]**2 +
	90*a[1,1]*a[1,2]*U[1,0]*U[2,0]**2 + 45*a[1,2]*a[2,2]*U[1,0]*U[2,0]**2 + 128*xd1**2*U[0,2]*U[2,1] + 45*xd1**2*a[1,2]**2*U[0,2]*U[2,1] -
	210*xd1**2*a[2,2]*U[0,2]*U[2,1] + 90*xd1**2*a[2,2]**2*U[0,2]*U[2,1] + 144*U[1,0]**2*U[2,1] + 135*a[1,2]**2*U[1,0]**2*U[2,1] - 270*a[2,2]*U[1,0]**2*U[2,1] +
	135*a[2,2]**2*U[1,0]**2*U[2,1] - 570*xd1**2*a[1,2]*U[1,1]*U[2,1] + 225*xd1**2*a[1,1]*a[1,2]*U[1,1]*U[2,1] + 315*xd1**2*a[1,2]*a[2,2]*U[1,1]*U[2,1] +
	192*xd1**2*U[2,0]*U[2,1] + 225*xd1**2*a[1,2]**2*U[2,0]*U[2,1] - 360*xd1**2*a[2,2]*U[2,0]*U[2,1] + 180*xd1**2*a[2,2]**2*U[2,0]*U[2,1] +
	12*xd1*xd2**3*((16 + 15*a[1,2]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[1,4] + 15*a[1,2]*(-2 + a[1,1] + a[2,2])*U[2,3]) - 60*xd1**2*a[1,2]*U[0,2]*U[3,0] +
	45*xd1**2*a[1,2]*a[2,2]*U[0,2]*U[3,0] - 270*a[1,2]*U[1,0]**2*U[3,0] + 135*a[1,1]*a[1,2]*U[1,0]**2*U[3,0] + 135*a[1,2]*a[2,2]*U[1,0]**2*U[3,0] +
	128*xd1**2*U[1,1]*U[3,0] - 60*xd1**2*a[1,1]*U[1,1]*U[3,0] + 90*xd1**2*a[1,2]**2*U[1,1]*U[3,0] - 150*xd1**2*a[2,2]*U[1,1]*U[3,0] +
	45*xd1**2*a[1,1]*a[2,2]*U[1,1]*U[3,0] + 45*xd1**2*a[2,2]**2*U[1,1]*U[3,0] - 510*xd1**2*a[1,2]*U[2,0]*U[3,0] + 270*xd1**2*a[1,1]*a[1,2]*U[2,0]*U[3,0] +
	225*xd1**2*a[1,2]*a[2,2]*U[2,0]*U[3,0] + 288*xd1**2*U[1,0]*U[3,1] + 270*xd1**2*a[1,2]**2*U[1,0]*U[3,1] - 540*xd1**2*a[2,2]*U[1,0]*U[3,1] +
	270*xd1**2*a[2,2]**2*U[1,0]*U[3,1] + U[0,1]*((128 + 45*a[1,2]**2 - 210*a[2,2] + 90*a[2,2]**2)*U[0,2]**2 + 128*U[1,1]**2 - 60*a[1,1]*U[1,1]**2 +
	90*a[1,2]**2*U[1,1]**2 - 150*a[2,2]*U[1,1]**2 + 45*a[1,1]*a[2,2]*U[1,1]**2 + 45*a[2,2]**2*U[1,1]**2 + 288*U[1,0]*U[1,2] + 270*a[1,2]**2*U[1,0]*U[1,2] -
	540*a[2,2]*U[1,0]*U[1,2] + 270*a[2,2]**2*U[1,0]*U[1,2] +
	18*xd2**2*((16 + 15*a[1,2]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[0,4] + 15*a[1,2]*(-2 + a[1,1] + a[2,2])*U[1,3]) - 150*a[1,2]*U[1,1]*U[2,0] +
	90*a[1,1]*a[1,2]*U[1,1]*U[2,0] + 45*a[1,2]*a[2,2]*U[1,1]*U[2,0] + 45*a[1,2]*U[0,2]*((-6 + a[1,1] + 4*a[2,2])*U[1,1] + a[1,2]*U[2,0]) -
	540*a[1,2]*U[1,0]*U[2,1] + 270*a[1,1]*a[1,2]*U[1,0]*U[2,1] + 270*a[1,2]*a[2,2]*U[1,0]*U[2,1] + 288*xd1**2*U[2,2] + 270*xd1**2*a[1,2]**2*U[2,2] -
	540*xd1**2*a[2,2]*U[2,2] + 270*xd1**2*a[2,2]**2*U[2,2] +
	36*xd1*xd2*((16 + 15*a[1,2]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[1,3] + 15*a[1,2]*(-2 + a[1,1] + a[2,2])*U[2,2]) - 540*xd1**2*a[1,2]*U[3,1] +
	270*xd1**2*a[1,1]*a[1,2]*U[3,1] + 270*xd1**2*a[1,2]*a[2,2]*U[3,1]) +
	xd2**2*(320*U[1,1]*U[1,2] - 60*a[1,1]*U[1,1]*U[1,2] - 510*a[2,2]*U[1,1]*U[1,2] + 45*a[1,1]*a[2,2]*U[1,1]*U[1,2] + 225*a[2,2]**2*U[1,1]*U[1,2] +
	5*U[0,2]*((64 + 45*a[1,2]**2 - 114*a[2,2] + 54*a[2,2]**2)*U[0,3] + 3*a[1,2]*(-28 + 12*a[1,1] + 15*a[2,2])*U[1,2]) + 288*U[1,0]*U[1,3] -
	540*a[2,2]*U[1,0]*U[1,3] + 270*a[2,2]**2*U[1,0]*U[1,3] + 288*xd1**2*U[2,3] - 540*xd1**2*a[2,2]*U[2,3] + 270*xd1**2*a[2,2]**2*U[2,3] +
	45*a[1,2]**2*(6*U[1,1]*U[1,2] + 6*U[1,0]*U[1,3] + U[0,3]*U[2,0] + 6*xd1**2*U[2,3]) +
	15*a[1,2]*((-14 + 3*a[1,1] + 9*a[2,2])*U[0,3]*U[1,1] + (-10 + 6*a[1,1] + 3*a[2,2])*U[1,2]*U[2,0] +
	6*(-2 + a[1,1] + a[2,2])*(2*U[1,1]*U[2,1] + 3*U[1,0]*U[2,2] + 3*xd1**2*U[3,2]))) - 540*xd1**2*a[1,2]*U[1,0]*U[4,0] +
	270*xd1**2*a[1,1]*a[1,2]*U[1,0]*U[4,0] + 270*xd1**2*a[1,2]*a[2,2]*U[1,0]*U[4,0] + 48*xd1**4*U[4,1] + 45*xd1**4*a[1,2]**2*U[4,1] - 90*xd1**4*a[2,2]*U[4,1] +
	45*xd1**4*a[2,2]**2*U[4,1] + 2*xd1*xd2*(6*(16 + 15*a[1,2]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[0,3]*U[1,1] - 390*a[1,2]*U[1,1]*U[1,2] +
	135*a[1,1]*a[1,2]*U[1,1]*U[1,2] + 225*a[1,2]*a[2,2]*U[1,1]*U[1,2] + 96*U[1,2]*U[2,0] + 135*a[1,2]**2*U[1,2]*U[2,0] - 180*a[2,2]*U[1,2]*U[2,0] +
	90*a[2,2]**2*U[1,2]*U[2,0] + 224*U[1,1]*U[2,1] - 60*a[1,1]*U[1,1]*U[2,1] + 180*a[1,2]**2*U[1,1]*U[2,1] - 330*a[2,2]*U[1,1]*U[2,1] +
	45*a[1,1]*a[2,2]*U[1,1]*U[2,1] + 135*a[2,2]**2*U[1,1]*U[2,1] - 330*a[1,2]*U[2,0]*U[2,1] + 180*a[1,1]*a[1,2]*U[2,0]*U[2,1] +
	135*a[1,2]*a[2,2]*U[2,0]*U[2,1] + U[0,2]*((224 + 135*a[1,2]**2 - 390*a[2,2] + 180*a[2,2]**2)*U[1,2] + 15*a[1,2]*(-16 + 6*a[1,1] + 9*a[2,2])*U[2,1]) +
	288*U[1,0]*U[2,2] + 270*a[1,2]**2*U[1,0]*U[2,2] - 540*a[2,2]*U[1,0]*U[2,2] + 270*a[2,2]**2*U[1,0]*U[2,2] - 180*a[1,2]*U[1,1]*U[3,0] +
	90*a[1,1]*a[1,2]*U[1,1]*U[3,0] + 90*a[1,2]*a[2,2]*U[1,1]*U[3,0] - 540*a[1,2]*U[1,0]*U[3,1] + 270*a[1,1]*a[1,2]*U[1,0]*U[3,1] +
	270*a[1,2]*a[2,2]*U[1,0]*U[3,1] + 96*xd1**2*U[3,2] + 90*xd1**2*a[1,2]**2*U[3,2] - 180*xd1**2*a[2,2]*U[3,2] + 90*xd1**2*a[2,2]**2*U[3,2] -
	180*xd1**2*a[1,2]*U[4,1] + 90*xd1**2*a[1,1]*a[1,2]*U[4,1] + 90*xd1**2*a[1,2]*a[2,2]*U[4,1]) - 90*xd1**4*a[1,2]*U[5,0] + 45*xd1**4*a[1,1]*a[1,2]*U[5,0] +
	45*xd1**4*a[1,2]*a[2,2]*U[5,0])/720.])

	r = ode2(z,UJet,a,h)[2:] + h**4*rhs

	return np.concatenate([ud,r])

def Hmod(x,UJet,a,h):

	dx = (-x[4:]+8*x[3:-1]-8*x[1:-3]+x[:-4])/(12*h) 	# fourth oder finite differences
	dnrm2 = np.sum(dx**2,1)
	nrm2 = np.sum(x[2:-2]**2,1)

	u = np.transpose(x[2:-2])

	U = np.zeros((6,6,u.shape[1]))

	U[0,0] = UJet[0][0](u)

	U[0,1] = UJet[0][1](u)
	U[0,2] = UJet[0][2](u)
	U[0,3] = UJet[0][3](u)
	U[0,4] = UJet[0][4](u)
	U[0,5] = UJet[0][5](u)

	U[1,0] = UJet[1][0](u)
	U[1,1] = UJet[1][1](u)
	U[1,2] = UJet[1][2](u)
	U[1,3] = UJet[1][3](u)
	U[1,4] = UJet[1][4](u)

	U[2,0] = UJet[2][0](u)
	U[2,1] = UJet[2][1](u)
	U[2,2] = UJet[2][2](u)
	U[2,3] = UJet[2][3](u)

	U[3,0] = UJet[3][0](u)
	U[3,1] = UJet[3][1](u)
	U[3,2] = UJet[3][2](u)

	U[4,0] = UJet[4][0](u)
	U[4,1] = UJet[4][1](u)

	U[5,0] = UJet[5][0](u)

	xd1 = dx[:,0]
	xd2 = dx[:,1]

	H0 = 1/2*dnrm2 - U[0,0]

	H2diff=((-4 + 3*a[2,2])*U[0,1]**2 + 8*xd2**2*U[0,2] - 6*xd2**2*a[2,2]*U[0,2] + 6*a[1,2]*U[0,1]*U[1,0] - 4*U[1,0]**2 + 3*a[1,1]*U[1,0]**2 - 6*xd2**2*a[1,2]*U[1,1] +
	xd1**2*(-6*a[1,2]*U[1,1] + 2*(4 - 3*a[1,1])*U[2,0]) - 2*xd1*xd2*(3*a[1,2]*U[0,2] + (-8 + 3*a[1,1] + 3*a[2,2])*U[1,1] + 3*a[1,2]*U[2,0]))/24.

	H2 = H0+h**2*H2diff

	H4diff=(90*xd1**2*a[1,2]*U[0,2]*U[1,1] - 45*xd1**2*a[1,1]*a[1,2]*U[0,2]*U[1,1] - 45*xd1**2*a[1,2]*a[2,2]*U[0,2]*U[1,1] - 90*a[1,2]*U[1,0]**2*U[1,1] +
	45*a[1,1]*a[1,2]*U[1,0]**2*U[1,1] + 45*a[1,2]*a[2,2]*U[1,0]**2*U[1,1] - 32*xd1**2*U[1,1]**2 + 90*xd1**2*a[1,1]*U[1,1]**2 - 45*xd1**2*a[1,1]**2*U[1,1]**2 -
	45*xd1**2*a[1,2]**2*U[1,1]**2 - 15*xd1**2*a[2,2]*U[1,1]**2 +
	U[0,1]**2*((48 + 45*a[1,2]**2 - 90*a[2,2] + 45*a[2,2]**2)*U[0,2] + 45*a[1,2]*(-2 + a[1,1] + a[2,2])*U[1,1]) -
	3*xd2**4*((16 + 15*a[1,2]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[0,4] + 15*a[1,2]*(-2 + a[1,1] + a[2,2])*U[1,3]) + 48*U[1,0]**2*U[2,0] -
	90*a[1,1]*U[1,0]**2*U[2,0] + 45*a[1,1]**2*U[1,0]**2*U[2,0] + 45*a[1,2]**2*U[1,0]**2*U[2,0] + 60*xd1**2*a[1,2]*U[1,1]*U[2,0] -
	45*xd1**2*a[1,1]*a[1,2]*U[1,1]*U[2,0] - 45*xd1**2*a[1,2]*a[2,2]*U[1,1]*U[2,0] - 32*xd1**2*U[2,0]**2 + 75*xd1**2*a[1,1]*U[2,0]**2 -
	45*xd1**2*a[1,1]**2*U[2,0]**2 - 45*xd1**2*a[1,2]**2*U[2,0]**2 + 180*xd1**2*a[1,2]*U[1,0]*U[2,1] - 90*xd1**2*a[1,1]*a[1,2]*U[1,0]*U[2,1] -
	90*xd1**2*a[1,2]*a[2,2]*U[1,0]*U[2,1] - 3*xd1*xd2**3*(60*a[1,2]**2*U[1,3] + (64 - 30*a[1,1] + 15*a[1,1]**2 - 90*a[2,2] + 45*a[2,2]**2)*U[1,3] +
	15*a[1,2]*(-2 + a[1,1] + a[2,2])*(U[0,4] + 3*U[2,2])) - 96*xd1**2*U[1,0]*U[3,0] + 180*xd1**2*a[1,1]*U[1,0]*U[3,0] - 90*xd1**2*a[1,1]**2*U[1,0]*U[3,0] -
	90*xd1**2*a[1,2]**2*U[1,0]*U[3,0] - 3*U[0,1]*(-32*U[1,0]*U[1,1] + 30*a[1,1]*U[1,0]*U[1,1] - 15*a[1,1]**2*U[1,0]*U[1,1] + 30*a[2,2]*U[1,0]*U[1,1] -
	15*a[2,2]**2*U[1,0]*U[1,1] + xd2**2*((32 + 30*a[1,2]**2 - 60*a[2,2] + 30*a[2,2]**2)*U[0,3] + 30*a[1,2]*(-2 + a[1,1] + a[2,2])*U[1,2]) +
	32*xd1**2*U[2,1] - 90*xd1**2*a[1,1]*U[2,1] + 45*xd1**2*a[1,1]**2*U[2,1] + 30*xd1**2*a[2,2]*U[2,1] - 15*xd1**2*a[2,2]**2*U[2,1] -
	30*a[1,2]**2*(U[1,0]*U[1,1] - xd1**2*U[2,1]) + xd1*xd2*
	(60*a[1,2]**2*U[1,2] + (64 - 90*a[1,1] + 45*a[1,1]**2 - 30*a[2,2] + 15*a[2,2]**2)*U[1,2] + 15*a[1,2]*(-2 + a[1,1] + a[2,2])*(3*U[0,3] + U[2,1])) -
	15*a[1,2]*(-2 + a[1,1] + a[2,2])*(U[0,2]*U[1,0] + U[1,0]*U[2,0] + xd1**2*(-3*U[1,2] + U[3,0]))) + 90*xd1**4*a[1,2]*U[3,1] -
	45*xd1**4*a[1,1]*a[1,2]*U[3,1] - 45*xd1**4*a[1,2]*a[2,2]*U[3,1] -
	xd2**2*((32 + 45*a[1,2]**2 - 75*a[2,2] + 45*a[2,2]**2)*U[0,2]**2 + 15*a[1,2]*(-4 + 3*a[1,1] + 3*a[2,2])*U[0,2]*U[1,1] + 32*U[1,1]**2 + 15*a[1,1]*U[1,1]**2 -
	90*a[2,2]*U[1,1]**2 + 45*a[2,2]**2*U[1,1]**2 + 96*U[1,0]*U[1,2] + 90*a[1,1]*U[1,0]*U[1,2] - 45*a[1,1]**2*U[1,0]*U[1,2] - 270*a[2,2]*U[1,0]*U[1,2] +
	135*a[2,2]**2*U[1,0]*U[1,2] + 288*xd1**2*U[2,2] - 270*xd1**2*a[1,1]*U[2,2] + 135*xd1**2*a[1,1]**2*U[2,2] - 270*xd1**2*a[2,2]*U[2,2] +
	135*xd1**2*a[2,2]**2*U[2,2] + 45*a[1,2]**2*(U[1,1]**2 + 2*U[1,0]*U[1,2] + 6*xd1**2*U[2,2]) -
	45*a[1,2]*(-2 + a[1,1] + a[2,2])*(U[0,3]*U[1,0] - U[1,1]*U[2,0] - 3*U[1,0]*U[2,1] - 3*xd1**2*(U[1,3] + U[3,1]))) - 48*xd1**4*U[4,0] +
	90*xd1**4*a[1,1]*U[4,0] - 45*xd1**4*a[1,1]**2*U[4,0] - 45*xd1**4*a[1,2]**2*U[4,0] -
	xd1*xd2*(64*U[0,2]*U[1,1] - 90*a[1,1]*U[0,2]*U[1,1] + 45*a[1,1]**2*U[0,2]*U[1,1] - 60*a[2,2]*U[0,2]*U[1,1] + 45*a[2,2]**2*U[0,2]*U[1,1] + 64*U[1,1]*U[2,0] -
	60*a[1,1]*U[1,1]*U[2,0] + 45*a[1,1]**2*U[1,1]*U[2,0] - 90*a[2,2]*U[1,1]*U[2,0] + 45*a[2,2]**2*U[1,1]*U[2,0] + 192*U[1,0]*U[2,1] -
	90*a[1,1]*U[1,0]*U[2,1] + 45*a[1,1]**2*U[1,0]*U[2,1] - 270*a[2,2]*U[1,0]*U[2,1] + 135*a[2,2]**2*U[1,0]*U[2,1] + 192*xd1**2*U[3,1] -
	270*xd1**2*a[1,1]*U[3,1] + 135*xd1**2*a[1,1]**2*U[3,1] - 90*xd1**2*a[2,2]*U[3,1] + 45*xd1**2*a[2,2]**2*U[3,1] +
	90*a[1,2]**2*(U[0,2]*U[1,1] + U[1,1]*U[2,0] + 2*U[1,0]*U[2,1] + 2*xd1**2*U[3,1]) +
	15*a[1,2]*(3*(-2 + a[1,1] + a[2,2])*U[0,2]**2 + 2*(-5 + 3*a[1,1] + 3*a[2,2])*U[1,1]**2 + 2*U[0,2]*U[2,0] +
	3*(-2 + a[1,1] + a[2,2])*(U[2,0]**2 + U[1,0]*(U[1,2] + 3*U[3,0]) + xd1**2*(3*U[2,2] + U[4,0])))))/720

	H4 = H2+h**4*H4diff

	return H0, H2, H4

def Imod(x,VJet,a,h):

	# computation of modified angular momentum if U=1/2V(#1,#2) and the matrix coefficients A commute with a rotation matrix

	dx = (-x[4:]+8*x[3:-1]-8*x[1:-3]+x[:-4])/(12*h) 	# fourth oder finite differences
	nrm2 = np.sum(x[2:-2]**2,1)

	V1 = VJet[1](nrm2)
	V2 = VJet[2](nrm2)
	V3 = VJet[3](nrm2)
	V4 = VJet[4](nrm2)

	xd1 = dx[:,0]
	xd2 = dx[:,1]

	x1 = x[2:-2,0]
	x2 = x[2:-2,1]

	# I0

	I0 = -(x2*xd1) + x1*xd2

	# I2

	I2diff= (x1*(-12*xd1*a[1,2] + xd2*(8 + 3*a[1,1] - 9*a[2,2]))*V1 - 6*x1**3*xd1*a[1,2]*V2 + 6*x2**3*xd2*a[1,2]*V2 +
	6*x1*x2**2*(xd1*a[1,2] + xd2*(a[1,1] - a[2,2]))*V2 + x2*
	((12*xd2*a[1,2] + xd1*(-8 + 9*a[1,1] - 3*a[2,2]))*V1 + 6*x1**2*(-(xd2*a[1,2]) + xd1*(a[1,1] - a[2,2]))*V2))/12.

	I2 = I0+h**2*I2diff

	# I4

	I4diff = (180*x2**5*xd2*a[1,2]*(-2 + a[1,1] + a[2,2])*(V2**2 + 3*V1*V3 + 2*xd2**2*V4) +
	180*x1*x2**4*(-2 + a[1,1] + a[2,2])*(xd2*(a[1,1] - a[2,2])*(V2**2 + 3*V1*V3 + 2*xd2**2*V4) + xd1*a[1,2]*(V2**2 + 3*V1*V3 + 6*xd2**2*V4)) +
	6*x1*x2**2*(2*(45*xd1**3*a[1,2]*(-2 + a[1,1] + a[2,2]) + 5*xd1*xd2**2*a[1,2]*(-34 + 21*a[1,1] + 21*a[2,2]) +
	2*xd2**3*(-32 + 30*a[1,1]**2 + 95*a[2,2] - 45*a[2,2]**2 - 5*a[1,1]*(7 + 3*a[2,2])) +
	xd1**2*xd2*(128 + 135*a[1,1]**2 + 50*a[2,2] - 75*a[2,2]**2 + 10*a[1,1]*(-29 + 6*a[2,2])))*V3 +
	V1*(60*xd1*a[1,2]*(-1 + a[1,1] + a[2,2])*V2 + xd2*(-64 + 135*a[1,1]**2 + 240*a[2,2] - 165*a[2,2]**2 - 30*a[1,1]*(4 + a[2,2]))*V2 +
	90*x1**2*xd2*(-2*a[1,1] + a[1,1]**2 - (-2 + a[2,2])*a[2,2])*V3) +
	30*x1**2*(-2 + a[1,1] + a[2,2])*(2*xd1**3*a[1,2]*V4 - 6*xd1*xd2**2*a[1,2]*V4 + xd2*(a[1,1] - a[2,2])*(V2**2 + 6*xd1**2*V4))) +
	6*x2**3*(2*xd2*(45*xd1**2*a[1,2]*(-2 + a[1,1] + a[2,2]) + 5*xd2**2*a[1,2]*(-34 + 21*a[1,1] + 21*a[2,2]) +
	xd1*xd2*(64 + 45*a[1,1]**2 - 20*a[2,2] - 15*a[2,2]**2 + 10*a[1,1]*(-10 + 3*a[2,2])))*V3 +
	V1*(20*xd2*a[1,2]*(-11 + 9*a[1,1] + 9*a[2,2])*V2 + xd1*(64 + 45*a[1,1]**2 - 20*a[2,2] - 15*a[2,2]**2 + 10*a[1,1]*(-10 + 3*a[2,2]))*V2 +
	90*x1**2*xd1*(-2*a[1,1] + a[1,1]**2 - (-2 + a[2,2])*a[2,2])*V3) +
	30*x1**2*(-2 + a[1,1] + a[2,2])*(6*xd1**2*xd2*a[1,2]*V4 - 2*xd2**3*a[1,2]*V4 + xd1*(a[1,1] - a[2,2])*(V2**2 + 6*xd2**2*V4))) -
	x1*(-6*xd2**3*(-64 + 15*a[1,1]**2 + a[1,1]*(20 - 30*a[2,2]) + 100*a[2,2] - 45*a[2,2]**2)*V2 +
	60*xd1*xd2**2*a[1,2]*((-10 + 3*a[1,1] + 3*a[2,2])*V2 + 9*x1**2*(-2 + a[1,1] + a[2,2])*V3) +
	xd2*((64 - 45*a[1,1]**2 - 90*a[1,2]**2 + 45*a[2,2] + 45*a[2,2]**2 + 15*a[1,1]*(-13 + 6*a[2,2]))*V1**2 -
	6*x1**2*(-64 + 15*a[1,1]**2 + a[1,1]*(20 - 30*a[2,2]) + 100*a[2,2] - 45*a[2,2]**2)*V1*V2 -
	6*xd1**2*((-64 + 45*a[1,1]**2 + a[1,1]*(40 - 30*a[2,2]) + 80*a[2,2] - 75*a[2,2]**2)*V2 +
	2*x1**2*(-64 + 15*a[1,1]**2 + a[1,1]*(20 - 30*a[2,2]) + 100*a[2,2] - 45*a[2,2]**2)*V3)) +
	30*xd1*a[1,2]*((8 + 3*a[1,1] + 3*a[2,2])*V1**2 + 6*x1**4*(-2 + a[1,1] + a[2,2])*V2**2 +
	2*V1*(2*x1**2*(-11 + 9*a[1,1] + 9*a[2,2])*V2 + 9*x1**4*(-2 + a[1,1] + a[2,2])*V3) +
	2*xd1**2*((-6 + 9*a[1,1] + 9*a[2,2])*V2 + x1**2*(-34 + 21*a[1,1] + 21*a[2,2])*V3 + 6*x1**4*(-2 + a[1,1] + a[2,2])*V4))) +
	x2*((30*xd2*a[1,2]*(8 + 3*a[1,1] + 3*a[2,2]) + xd1*(64 + 45*a[1,1]**2 - 90*a[1,2]**2 - 195*a[2,2] - 45*a[2,2]**2 + 45*a[1,1]*(1 + 2*a[2,2])))*V1**2 +
	6*x1**2*V1*(-30*xd2*a[1,2]*(2*(-1 + a[1,1] + a[2,2])*V2 + 3*x1**2*(-2 + a[1,1] + a[2,2])*V3) +
	xd1*((64 + 165*a[1,1]**2 + 30*a[1,1]*(-8 + a[2,2]) + 120*a[2,2] - 135*a[2,2]**2)*V2 + 90*x1**2*(-2*a[1,1] + a[1,1]**2 - (-2 + a[2,2])*a[2,2])*V3))
	+ 6*(-30*xd2*a[1,2]*(x1**4*(-2 + a[1,1] + a[2,2])*V2**2 + xd2**2*((2 - 3*a[1,1] - 3*a[2,2])*V2 + 3*x1**2*(-2 + a[1,1] + a[2,2])*V3)) +
	xd1*(30*x1**4*(-2*a[1,1] + a[1,1]**2 - (-2 + a[2,2])*a[2,2])*V2**2 +
	xd2**2*((64 + 75*a[1,1]**2 - 40*a[2,2] - 45*a[2,2]**2 + 10*a[1,1]*(-8 + 3*a[2,2]))*V2 +
	2*x1**2*(-128 + 75*a[1,1]**2 + 290*a[2,2] - 135*a[2,2]**2 - 10*a[1,1]*(5 + 6*a[2,2]))*V3)) +
	10*xd1**2*xd2*a[1,2]*((-10 + 3*a[1,1] + 3*a[2,2])*V2 - x1**2*(-34 + 21*a[1,1] + 21*a[2,2])*V3 - 18*x1**4*(-2 + a[1,1] + a[2,2])*V4) +
	xd1**3*((64 + 45*a[1,1]**2 - 20*a[2,2] - 15*a[2,2]**2 + 10*a[1,1]*(-10 + 3*a[2,2]))*V2 +
	4*x1**2*(32 + 45*a[1,1]**2 + 35*a[2,2] - 30*a[2,2]**2 + 5*a[1,1]*(-19 + 3*a[2,2]))*V3 +
	60*x1**4*(-2*a[1,1] + a[1,1]**2 - (-2 + a[2,2])*a[2,2])*V4))))/720.

	I4 = I2+h**4*I4diff

	return I0, I2, I4

max_stars_repo_name: Christian-Offen/BEAConjugateSymplectic

max_issues_repo_name: Christian-Offen/BEAConjugateSymplectic

max_forks_repo_name: Christian-Offen/BEAConjugateSymplectic


text:
import os

import torch
import git
import numpy as np
from PIL import Image
from torchvision import datasets, transforms
from sklearn.preprocessing import OneHotEncoder

from strimadec.datasets.utils import get_git_root

class FashionMNIST(datasets.FashionMNIST):

    """standard FashionMNIST dataset with one-hot vectors as labels instead of class labels"""

    def __init__(self, train):
        git_root_path = get_git_root(os.getcwd())
        store_dir = os.path.join(git_root_path, "strimadec/datasets/data")
        transform = transforms.ToTensor()
        target_transform, download = None, True
        super().__init__(
            root=store_dir,
            train=train,
            transform=transform,
            target_transform=target_transform,
            download=download,
        )
        return

    def __getitem__(self, idx):
        img, target = super().__getitem__(idx)
        # convert target to one hot vector (there are 10 classes)
        target = torch.eye(10)[target]
        return img, target

class FullMNIST(datasets.MNIST):
    """standard MNIST dataset with one-hot vectors as labels instead of class labels"""

    def __init__(self, train):
        git_root_path = get_git_root(os.getcwd())
        store_dir = os.path.join(git_root_path, "strimadec/datasets/data")
        transform = transforms.ToTensor()
        target_transform, download = None, True
        super(FullMNIST, self).__init__(store_dir, train, transform, target_transform, download)
        return

    def __getitem__(self, idx):
        img, target = super().__getitem__(idx)
        # convert target to one hot vector (there are 10 classes)
        target = torch.eye(10)[target]
        return img, target

class Letters(datasets.EMNIST):
    """26 letters dataset with one-hot vectors as labels"""

    def __init__(self) -> None:
        git_root_path = get_git_root(os.getcwd())
        store_dir = os.path.join(git_root_path, "strimadec/datasets/data")
        transform = transforms.ToTensor()
        target_transform, download = None, True
        split = "letters"
        super().__init__(
            root=store_dir,
            split=split,
            transform=transform,
            target_transform=target_transform,
            download=download,
        )
        return

    def __getitem__(self, idx):
        img, target = super().__getitem__(idx)
        # convert target to one hot vector (there are 26 classes)
        target = torch.eye(26)[target - 1]  # target goes from 1 to 26
        return img, target

class SimplifiedMNIST(torch.utils.data.TensorDataset):
    """contains only a subset of MNIST digits with one-hot vectors as labels"""

    def __init__(self, train, digits):
        git_root_path = get_git_root(os.getcwd())
        store_dir = os.path.join(git_root_path, "strimadec/datasets/data")
        transform = transforms.ToTensor()
        target_transform, download = None, True
        MNIST_dataset = datasets.MNIST(store_dir, train, transform, target_transform, download)
        # select only specific digits
        data = []
        labels = []
        for digit in digits:
            indices_digits = MNIST_dataset.targets == digit
            torch_imgs = [
                transforms.ToTensor()(Image.fromarray(img.numpy(), mode="L"))
                for img in MNIST_dataset.data[indices_digits]
            ]
            data.append(torch.vstack(torch_imgs))
            labels.extend([digit] * sum(indices_digits))
        # vertical stack torch tensors within data list
        data = torch.vstack(data).unsqueeze(1)
        # create one-hot encoded labels
        labels = OneHotEncoder().fit_transform(np.array(labels).reshape(-1, 1)).toarray()
        # define tensor dataset
        super(SimplifiedMNIST, self).__init__(data, torch.from_numpy(labels))
        return


text:
# NOTE:
# To test on 'real' bigdata:
# xtc_dir = "/reg/d/psdm/xpp/xpptut15/scratch/mona/test"
# >bsub -n 64 -q psfehq -o log.txt mpirun python user_loops.py

# cpo found this on the web as a way to get mpirun to exit when
# one of the ranks has an exception
import sys
# Global error handler
def global_except_hook(exctype, value, traceback):
    sys.stderr.write("except_hook. Calling MPI_Abort().\n")
    # NOTE: mpi4py must be imported inside exception handler, not globally.
    # In chainermn, mpi4py import is carefully delayed, because
    # mpi4py automatically call MPI_Init() and cause a crash on Infiniband environment.
    import mpi4py.MPI
    mpi4py.MPI.COMM_WORLD.Abort(1)
    sys.__excepthook__(exctype, value, traceback)
sys.excepthook = global_except_hook

#import logging
#logger = logging.getLogger('psana.psexp.node')
#logger.setLevel(logging.DEBUG)
#ch = logging.StreamHandler()
#ch.setLevel(logging.DEBUG)
#formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
#ch.setFormatter(formatter)
#logger.addHandler(ch)

import os
import vals
import numpy as np
from psana import DataSource

def filter_fn(evt):
    # Create a detector within filter is possible
    run = evt.run()
    det = run.Detector('xppcspad')

    # Access step event
    step_evt = run.step(evt)
    return True

xtc_dir = os.path.join(os.environ.get('TEST_XTC_DIR', os.getcwd()),'.tmp')

# Usecase 1a : two iterators with filter function
ds = DataSource(exp='xpptut13', run=1, dir=xtc_dir, filter=filter_fn, monitor=True)
#beginJobCode
for run in ds.runs():
    #beginRunCode
    # Detector interface identified by detector name
    det = run.Detector('xppcspad')

    # Environment values are accessed also through detector interface
    edet = run.Detector('HX2:DVD:GCC:01:PMON')
    sdet = run.Detector('motor2')

    for evt in run.events():
        padarray = vals.padarray
        # 4 segments, two per file
        assert(np.array_equal(det.raw.calib(evt),np.stack((padarray,padarray,padarray,padarray))))
        assert edet.dtype == float
        assert sdet.dtype == float
        assert edet(evt) is None or edet(evt) == 41.0
        assert sdet(evt) == 42.0
        assert run.expt == 'xpptut15' # this is from xtc file
        assert run.runnum == 14
        assert run.timestamp == 4294967297
    #endRunCode
#endJobCode

# Usecase#2 looping through steps
ds = DataSource(exp='xpptut13', run=1, dir=xtc_dir, batch_size=10)
for run in ds.runs():
    det = run.Detector('xppcspad')
    for step in run.steps():
        for evt in step.events():
            padarray = vals.padarray
            assert(np.array_equal(det.raw.calib(evt),np.stack((padarray,padarray,padarray,padarray))))

# Usecase#3: singlefile ds
ds = DataSource(files=os.path.join(xtc_dir,'data-r0001-s00.xtc2'))
for run in ds.runs():
    det = run.Detector('xppcspad')
    edet = run.Detector('HX2:DVD:GCC:01:PMON')
    sdet = run.Detector('motor2')
    for step in run.steps():
        for evt in step.events():
            calib = det.raw.calib(evt)
            assert calib.shape == (2,3,6)
    assert run.expt == 'xpptut15'
    assert run.runnum == 14


text:
import uuid

import numpy as np

class BColors:
	HEADER = '\033[95m'
	OKBLUE = '\033[94m'
	OKGREEN = '\033[92m'
	WARNING = '\033[93m'
	FAIL = '\033[91m'
	ENDC = '\033[0m'
	BOLD = '\033[1m'
	UNDERLINE = '\033[4m'

class Indv:
	def __init__(self, indv_size):
		self.indv_size = indv_size
		self.weights = np.random.uniform(-4.0, 4.0, indv_size)
		self.uuid = uuid.uuid4()

	def __str__(self):
		return 'Indv weights: {self.weights} uuid: {self.uuid}'.format(self=self)

	def __repr__(self):
		return self.weights.__str__()

	def __len__(self):
		return self.indv_size

class Population:
	def __init__(self, pop_size, indv_size):
		self.population = []
		self.pop_size = pop_size
		self.indv_size = indv_size
		self.uuid = uuid.uuid4()
		for indv_idx in range(self.pop_size):
			self.population.append(Indv(indv_size))

	def __str__(self):
		return 'Population of {self.pop_size} indviduals with uuid: {self.uuid}\n{self.population}'.format(self=self)

	def __repr__(self):
		return self.population.__str__()

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

def fitness_offsprings(mutated_offsprings, func_weights):
	fitness_array = []
	for indv in mutated_offsprings:
		fitness_array.append(fitness_max_function(indv.weights, func_weights))
	print("Actual offsprings fitness values:")
	print(fitness_array)
	return fitness_array

def fitness_population(population, func_weights):
	fitness_array = []
	for indv in population.population:
		fitness_array.append(fitness_max_function(indv.weights,
												  func_weights))  # Calculate the fitness value for each indv of the population
	print("Actual fitness values:")
	print(fitness_array)
	return fitness_array

def fitness_max_function(indv_weights, func_weights):
	return np.dot(indv_weights, func_weights)

def select_parents(population, fitness_values):
	num_parents = int(len(population) / 2)  # We want the best 50% of the parents
	parents = []
	for parent_num in range(num_parents):  # For each parent
		max_fitness_idx = np.where(fitness_values == np.max(fitness_values))  # Get the index of the max fitness value
		max_fitness_idx = max_fitness_idx[0][0]
		parents.append(population.population[max_fitness_idx])
		fitness_values[max_fitness_idx] = -99999999999  # Change his fitness to a minimum
	print("Selected parents:")
	print(parents)
	return parents


text:
# coding=utf-8
# Creation date: 28 окт. 2020
# Creation time: 19:16
# Creator: SteamPeKa

from typing import Collection, Hashable, Union

import numpy
import pytest

def assert_equal_tensors(expected: numpy.ndarray, actual: numpy.ndarray, additional_string: Union[None, str] = None):
    assert expected.shape == actual.shape, f"Tensors have different shapes: {expected.shape} and {actual.shape}"
    abs_difference = numpy.max(numpy.abs(expected - actual))
    fail_message = (
        f"Tensors are not equal with maximum abs difference of {abs_difference}.\n"
        f"Expected:\n"
        f"{expected}\n"
        f"Actual:\n"
        f"{actual}\n"
        f"Difference:\n"
        f"{expected - actual}"
    )
    if additional_string is not None:
        fail_message += "\n####\n" + additional_string + "\n####"
    assert abs_difference == pytest.approx(0), fail_message

def assert_collections_equal_as_sets(expected: Collection[Hashable], actual: Collection[Hashable]):
    expected_set = set(expected)
    actual_set = set(actual)
    sym_diff = expected_set.symmetric_difference(actual_set)
    assert len(sym_diff) == 0, (f"Collections are not equal.\n"
                                f"Expected collection-set: {expected_set} (of size {len(expected_set)})\n"
                                f"Actual collection-set: {actual_set} (of size {len(actual_set)})\n"
                                f"Symmetric difference: {sym_diff} (of length {len(sym_diff)} != 0)")