text:
from app import app, es, db, flask_bcrypt, login_manager
from app import User, Organization
from flask import jsonify
from flask import render_template
from flask import url_for
from flask import redirect
from flask import Response
from flask import request
from flask import session
from flask import abort
from flask import flash
from flask import send_file
from flask.ext.login import (current_user, login_required, login_user,
        logout_user, confirm_login, fresh_login_required)
from flask import Blueprint

from werkzeug import secure_filename
import flask
import tempfile
import json
from elasticsearch_dsl import Search, Q
import io
import re
import magic
import requests
from nltk import word_tokenize
from nltk.corpus import stopwords
from collections import Counter
import os,sys
import subprocess

from bulk import bulk_download, bulk_search
from config import tmp_dir
from util.network import make_graph, document_graph
import time

import sklearn
from sklearn.feature_extraction.text import CountVectorizer
from sklearn import decomposition
import corex as ce
import numpy as np
import phonenumbers
from phonenumbers import geocoder

parent = os.path.dirname(os.path.realpath(__file__))
sys.path.append('/home/gmueller/geodict') #change path to MITIE top level

import geodict_lib

DEFAULT_INDEX = 'dossiers'
uni = Blueprint('unicorn', __name__, url_prefix='/unicorn')

@uni.route('/_bulk_search', methods=['POST'])
@login_required
def bulk_search_route():

    search_results = request.form['searches']
    searches = search_results.split('\n')

    data = bulk_search(searches)
    return send_file(io.BytesIO(data.xls), as_attachment=True,
            attachment_filename='bulk_{}.xls'.format(time.time()))

    # Bulk search all of these queries
    # Bundle results into excel

@uni.route('/bulk_download')
@login_required
def bulk_download_route():
    if 'last_query' not in session:
        return abort(404)

    last_query = session['last_query']

    ids = last_query['ids']
    query = last_query['query']

    data = bulk_download(ids)

    return send_file(io.BytesIO(data.xls), as_attachment=True,
            attachment_filename=query + '.xls')

@uni.route('/<doc_id>/debug')
@login_required
def request_doc(doc_id):
    q = {
            "query" : {
                "match" : {
                    "_id" : doc_id
                    }
                },
            }
    response = es.search(body=q, index=DEFAULT_INDEX)
    return response

def get_file(doc_id):
    ''' Render base64 encoded contents of a given file by its doc_id '''
    response = request_doc(doc_id)
    try:
        base64 = response['hits']['hits'][0]['_source']['file']
        fn = response['hits']['hits'][0]['_source']['title']
    except KeyError, IndexError:
        abort(404)

    return base64, fn

@uni.route('/<doc_id>/entities')
@login_required
def get_entities(doc_id):
    response = request_doc(doc_id)
    try:
        entities = response['hits']['hits'][0]['_source']['entities']
    except KeyError, IndexError:
        return jsonify([])

    return jsonify({'entities': entities})

@uni.route('/view/<doc_id>')
@login_required
def view_doc(doc_id):
    ''' In-depth view of a particular document.
    Displays pdf version of document, extracted entities,
    as well as other analytics. '''

    if is_owner_of_doc(doc_id):
        return render_template('doc-view.html', doc_id=doc_id)
    else:
        return abort(403)

@uni.route('/pdf/<doc_id>')
@login_required
def pdf_endpoint(doc_id):
    base64, fn = get_file(doc_id)
    b = base64.decode('base64')
    mimetype = magic.from_buffer(b, mime=True)

    if mimetype == 'application/pdf':
        return send_file(io.BytesIO(b), as_attachment=True,
                attachment_filename=fn)

    else:
        fd, fname = tempfile.mkstemp(prefix=tmp_dir)
        stream = os.fdopen(fd, 'wb')
        stream.write(b)
        out_fname = fname + '.out'
        stream.close()

        os.chmod(fname, 0777)
        try:
            subprocess.check_output(['unoconv', '-o', out_fname, fname],
                    stderr=subprocess.STDOUT)
            with open(out_fname, 'rb') as converted_stream:
                out = send_file(out_fname, as_attachment=True,
                        attachment_filename=fn + '.pdf')
        except subprocess.CalledProcessError as e:
            print e.output
            # Return error pdf
            out = "no pdf available"

        os.remove(fname)
        os.remove(out_fname)

        return out

@uni.route('/topics/<doc_id>')
def get_topics(doc_id):
    topics=json.loads(open('/home/gmueller/unicorn/topics.json').read())
    return json.dumps(topics['documents'][doc_id])

@uni.route('/all_topics')
def alltopics():
    topics=json.loads(open('/home/gmueller/unicorn/topics.json').read())
    docs=len(topics['documents'])
    dist={}
    for idx,x in enumerate(np.bincount([np.argmax(item[1]) for item in topics['documents'].items()])):
        dist['topic'+str(idx+1)]=float(x)/docs
    topics['dist']=dist
    return json.dumps(topics)

@uni.route('/download/<doc_id>')
@login_required
def download_endpoint(doc_id):

    base64, fn = get_file(doc_id)
    f = io.BytesIO(base64.decode('base64'))
    return send_file(f,
            as_attachment=True,
            attachment_filename=fn)

@uni.route('/search/<query>/<page>/preserve')
@login_required
def search_preserve(query, page):
    return search_endpoint(query, page, box_only=True)

@uni.route('/search')
@uni.route('/search/<query>')
@uni.route('/search/<query>/<page>')
@login_required
def search_endpoint(query=None, page=None, box_only=False):
    if not query and not page:
        last_query = session.get('last_query', None)
        if last_query:
            query, page = last_query['query'], last_query['page']
        else:
            # better error
            return abort(404)

    if not page:
        page = 1

    session['last_query'] = {'query': query, 'page': page, 'ids': []}
    # convert pages to records for ES
    start = int(page)
    if start > 1:
        start *= 10

    q = {
            "fields": ["title", "highlight", "entities", "owner"],
            "from": start,
            "query" : {
                "match" : {
                    "file" : query
                    }
                },

            "highlight": { "fields": { "file": { } },
                "pre_tags" : [""],
                "post_tags" : [""]
                }
            }
    raw_response = es.search(body=q, index=DEFAULT_INDEX,
            df="file",
            size=10)

    hits = []

    for resp in raw_response['hits']['hits']:
        # Store returned ids
        session['last_query']['ids'].append(resp['_id'])

        if is_owner(resp['fields']['owner'][0]):
            # Flatten structure for individual hits
            hits.append({'id': resp['_id'],
                'title': resp['fields']['title'][0],
                'highlight': resp['highlight']['file'][0],
                'permissions': True
                })
        else:
            hits.append({'id': resp['_id'],
                'title': resp['fields']['title'][0],
                'permissions': False
                })

    results = {
            'hits': hits,
            'took': float(raw_response['took'])/1000,
            'total': "{:,}".format(raw_response['hits']['total']),
            'total_int': int(raw_response['hits']['total']),
            'query': query,
            'from': int(page)
            }

    if box_only:
        return render_template('search-results-box.html', results=results)

    return render_template('search-template.html', results=results)

@uni.route('/')
@login_required
def root():
    user_struct = {
            'notifs': 0
            }
    return render_template('index-dash.html', user=user_struct)

@uni.route('/user')
@login_required
def user_page():
    # If current user is admin

    return render_template('user-invite.html')

@uni.route('/upload/')
@login_required
def upload_form():
    return render_template('upload.html')

@uni.route('/_upload-documents', methods=['POST'])
@login_required
def upload_endpoint():
    files = request.files.getlist('file[]')
    d = {}
    for f in files:
        sf = secure_filename(f.filename)

        es_dict = {
                'file': f.read().encode('base64'),
                'title': sf,
                'owner': current_owner.organization.organization
                }
        es.index(index=DEFAULT_INDEX, doc_type='attachment', body=es_dict)
        f.close()

    return redirect(url_for('.root'))

@uni.route('/viz')
@login_required
def viz_all():
    q = {
        "fields" : ["entities","title"],
        "query" : {
            "match_all" : {}
            },
        "size": 100
        }
    r = es.search(body=q, index=DEFAULT_INDEX)
    graph = document_graph(r['hits']['hits'])

    return json.dumps(graph)

@uni.route('/geo')
@login_required
def geo_endpoint():
    query=session['last_query']['query']
    url='http://localhost:9200/dossiers/_search'
    q = {
        "fields" : ["file"],
        "query" : {
            "term" : { "file" : query }
            }
        }
    #r=requests.post(url,data=json.dumps(q))
    r=es.search(body=q,index=DEFAULT_INDEX)
    data=r
    locations=[]
    for hit in data['hits']['hits']:
        for location in geodict_lib.find_locations_in_text(re.sub('\s', ' ', str(hit['fields']['file']))):
            for token in location['found_tokens']:
                locations.append({'lat':token['lat'],'lon':token['lon'],'name':token['matched_string']})

    #geo=map(lambda x: x['found_tokens'])
    return json.dumps(locations)

@uni.route('/viz/<query>')
@login_required
def viz_endpoint(query):
    url='http://localhost:9200/dossiers/_search'
    q = {
        "_source": ["entity"],
        "fields" : ["entities","title"],
        "query" : {
            "term" : { "file" : query }
            },
        "size": 100
        }
    #r=requests.post(url,data=json.dumps(q))
    r=es.search(body=q,index=DEFAULT_INDEX)
    data=r
    #graph = make_graph(data)
    graph = document_graph(data['hits']['hits'])
    return json.dumps(graph)

@uni.route('/viz_latest')
@login_required
def viz_latest():
    return viz_endpoint(session['last_query']['query'])

@uni.route('/wc_latest')
@login_required
def wc_latest():
    return wc(session['last_query']['query'])

@uni.route('/url_list')
@uni.route('/url_list/<query>')
@login_required
def url_fetch(query=""):
    #query="http"
    if not query:
        query=session['last_query']['query']
    stopset=set(stopwords.words('english'))
    url='http://localhost:9200/dossiers/_search'
    q = {
        "fields" : ["file"],
        "query" : {
            "term" : { "file" : query }
            }
        }
    #r=requests.post(url,data=json.dumps(q))
    r=es.search(body=q,index=DEFAULT_INDEX)
    data=r['hits']['hits']
    urls=[]
    pn=[]
    for doc in data:
        urls.append(re.findall(r'(https?://[^\s]+)', doc['fields']['file'][0]))
        try:
            for match in phonenumbers.PhoneNumberMatcher(doc['fields']['file'][0], region=None):
                    pn.append({'number':phonenumbers.format_number(match.number, phonenumbers.PhoneNumberFormat.E164),'location':geocoder.description_for_number(match.number,"en")})
        except KeyError:
            pass
    urls=filter(lambda x: x!=[],urls)
    #urls_flat=reduce(lambda x,y: x.extend(y),urls)
    urls_flat=[item for sublist in urls for item in sublist]
    return json.dumps({'urls':dict(Counter(urls_flat)), 'pn':pn})

@uni.route('/wordcloud/<query>')
@login_required
def wc(query):
    stopset=set(stopwords.words('english'))
    url='http://localhost:9200/dossiers/_search'
    q = {
        "fields" : ["file"],
        "query" : {
            "term" : { "file" : query }
            }
        }
    #r=requests.post(url,data=json.dumps(q))
    r=es.search(body=q,index=DEFAULT_INDEX)
    #r=requests.post(url,data=json.dumps(q))
    data=r
    frequency=[]
    documents=[]
    for hit in data['hits']['hits']:
        text=hit['fields']['file'][0]
        nowhite=re.sub('\s', ' ', text)
        nowhite=re.sub(r'[^\w\s]', '',text)
        wt=word_tokenize(nowhite)
        documents.append(wt)

    docflat=[item for sublist in documents for item in sublist]
    wc=dict(Counter(docflat))
    for k,v in wc.iteritems():
                frequency.append(dict({"text":k,"size":v*3}))
    frequency=filter(lambda x:x['size']>6 and x['text'].lower() not in stopset,frequency)
    return json.dumps(frequency)

@uni.route('/topicmodel')
@uni.route('/topicmodel/<query>')
@login_required
def tm(query):
    #count_vectorizer.fit_transform(train_set)
    #print "Vocabulary:", count_vectorizer.vocabulary
    # Vocabulary: {'blue': 0, 'sun': 1, 'bright': 2, 'sky': 3}
    #freq_term_matrix = count_vectorizer.transform(test_set)
    #print freq_term_matrix.todense()
    stopset=set(stopwords.words('english'))
    url='http://localhost:9200/dossiers/_search'
    q = {
        "fields" : ["file"],
        "query" : {
            "term" : { "file" : query }
            }
        }

    return json.dumps(topic_words[0])

@uni.route('/<doc_id>/related')
@login_required
def more_like_this(doc_id):
    ''' Returns similar documents '''
    q = {
      "fields": ["title"],
        "query": {
            "more_like_this" : {
            "docs" : [
            {
                "_index" : "dossiers",
                "_type" : "attachment",
                "_id" : doc_id
            }]
          }
        },
        "size": 10
    }

    response = es.search(body=q, index=DEFAULT_INDEX)
    results = {'results': []}
    try:
        for r in response['hits']['hits']:
            results['results'].append({
                'id': r['_id'],
                'name': r['fields']['title'][0]
                })
    except KeyError, IndexError:
        pass

    return jsonify(results)

@uni.route('/index.html')
@login_required
def reroute_index():
    return redirect(url_for('.root'))

@uni.route('/register', methods=['POST'])
@login_required
def handle_registration():
    if not current_user.moderator:
        return redirect(url_for('.root'))

    email = request.form['email']
    password = request.form['password']

    # Change this to: current user's group
    organization = current_user.organization

    pw_hashed = flask_bcrypt.generate_password_hash(password)
    new_user = User(email=email,
                    password=pw_hashed,
                    organization=organization)

    db.session.add(new_user)
    db.session.commit()

    return redirect(url_for('.root'))

######################################
# Registration blueprint:
######################################

@uni.route("/login", methods=["GET", "POST"])
def login():
    if request.method == 'GET':
        return render_template("/user-login.html")

    email = request.form["email"]
    password = request.form["password"]
    user = User.query.filter_by(email=email).first()
    if user and flask_bcrypt.check_password_hash(user.password,
            password):
            if login_user(user):
                return redirect(url_for('.root'))
            else:
                flash("Invalid email or password")

    return render_template("/user-login.html")

@uni.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('.login'))

@login_manager.unauthorized_handler
def login_redirect():
    return redirect(url_for('.login'))

@login_manager.user_loader
def load_user(userid):
    return User.query.get(userid)

def is_owner_of_doc(doc):
    owner = es.get(index=DEFAULT_INDEX, doc_type='attachment', id=doc,
            fields='owner')['fields']['owner'][0]
    return is_owner(owner)

def is_owner(org):
    return current_user.organization.organization == 'admins' or \
            current_user.organization.organization == org

@app.errorhandler(403)
def permission_denied(e):
    return render_template('permission-denied.html'), 403


text:
import logging
import os
import pickle
import sys
from pathlib import Path

import click
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import zsampler
from dotenv import load_dotenv, find_dotenv

from src.inference.context_geo import GridContextGeo
from src.inference.hmc import HMCSampler
from src.inference.priors import BetaPriorWithIntercept
from src.experiment.model_management import build_block_mixture_flat_uid
from src.experiment.visualize import plot_traceplots

class FlatRegionMixture:

    def __init__(self, *, uid=None,
                 grid_context=None,
                 K=1,
                 thinning=5,
                 n_info_interval=1000,
                 block_type="msoa"):
        self.uid = uid
        self.context = grid_context
        self.K = K
        self.NN = self.context.mask.shape[0]
        self.thinning = thinning
        self.n_info_interval = n_info_interval
        N = grid_context.counts.shape[0]

        # do a random assignment to mixtures
        initial_Z = np.zeros((N, K), dtype=int)
        initial_Z[np.arange(N), np.random.choice(K, N)] = 1

        # Create an (N x 1) vector which gives the corresponding block for each cell.
        if block_type == "lad":
            block_assignment = np.asarray(grid_context.lads)
        elif block_type == "msoa":
            block_assignment = np.asarray(grid_context.msoas)
        elif block_type == "ward":
            block_assignment = np.asarray(grid_context.wards)
        elif block_type == "lsoa":
            block_assignment = np.asarray(grid_context.lsoas)
        else:
            block_assignment = np.repeat(1, N) # a single block

        unique_block_labels = np.unique(block_assignment)
        self.block_assignment_numeric = np.zeros(block_assignment.shape[0], dtype=np.int)
        for idx_cell, block_label in enumerate(block_assignment):
            self.block_assignment_numeric[idx_cell] = np.where(unique_block_labels == block_label)[0]
        self.block_assignment = block_assignment

        self.B = np.max(self.block_assignment_numeric) + 1

        # Create B x K matrix which counts number of cells in a block b with assignment k.
        self.block_label_counts = np.zeros(shape=(self.B, K), dtype=np.int64)
        for i in range(N):
            self.block_label_counts[self.block_assignment_numeric[i], :] += initial_Z[i, :]

        # Beta prior
        self.beta_prior = BetaPriorWithIntercept(a=1, b=0.01)
        self.Z = initial_Z

        self.Z_samples = []
        self.alpha_samples = []
        self.beta_samples = []

        self.logger = logging.getLogger(__name__)

    def loglik(self, estimand):
        """ Compute log p(y | beta, Z) + log p(beta)"""
        J = self.context.J
        K = self.K
        covariates = self.context.covariates
        counts = self.context.counts

        beta = estimand[:(J * K)]
        beta_matrix = beta.reshape((J, K), order='F')

        fixed_effects = np.sum(np.multiply(self.Z, np.dot(covariates, beta_matrix)), axis=1)

        poisson_part = np.sum(np.multiply(counts, fixed_effects) - np.exp(fixed_effects))
        self.logger.debug(f"Poisson part: {poisson_part}")
        beta_part = self.beta_prior.log_pdf(beta, J)

        output = poisson_part + beta_part
        return output

    def nabla_loglik(self, estimand):
        """ Compute the gradient of log p(y | beta, Z) + log p(beta) with respect to beta"""

        J = self.context.J
        K = self.K
        covariates = self.context.covariates
        counts = self.context.counts

        beta = estimand[:(J * K)]
        beta_matrix = beta.reshape((J, K), order='F')
        fixed_effects = np.sum(np.multiply(self.Z, np.dot(covariates, beta_matrix)), axis=1)

        # nabla beta
        nabla_beta_matrix = np.zeros(beta_matrix.shape)
        nabla_beta_matrix += np.dot(covariates.T, self.Z * counts[:, np.newaxis])

        temp = np.exp(fixed_effects)
        nabla_beta_matrix += (- np.dot(covariates.T, self.Z * temp[:, np.newaxis]))
        nabla_beta = nabla_beta_matrix.flatten('F')
        nabla_beta += self.beta_prior.nabla_beta_log_pdf(beta, J)

        output = nabla_beta
        return output

    def plot_traces(self, hmc_samples):

        samples_array = np.asarray(hmc_samples)

        S = samples_array.shape[0]
        J = self.context.J
        N = self.context.N
        K = self.K

        # discard irrelevant samples
        self.Z_samples = self.Z_samples[(-S):]
        Z_samples_array = np.asarray(self.Z_samples)

        mixture_allocation = np.zeros((S, N, K))
        mixture_allocation[np.repeat(range(S), N), np.tile(range(N), S), Z_samples_array.flatten(order='C')] = 1
        average_alloc = np.mean(mixture_allocation, axis=0)

        for k in range(self.K):
            plt.figure()
            self.context.plot_realisations(average_alloc[:, k], 111)
            plt.show()

            beta_k_samples = samples_array[:, (k * J):((k + 1) * J)]
            plot_traceplots(beta_k_samples, self.context.covariates_names)
            plt.show()

            k_component_indices = np.where(average_alloc[:, k] > (1 / K))[0]

            # Fitted surface
            fitted_surface_map = np.multiply(average_alloc[:, k],
                                             np.dot(self.context.covariates, np.mean(beta_k_samples, axis=0)))
            plt.figure()
            self.context.plot_realisations(fitted_surface_map, 111, plot_title="Log-intensity fitted")
            plt.show()

            # Correlation Matrix
            crime_surface_k = np.dot(self.context.covariates[k_component_indices, :],
                                     np.mean(beta_k_samples, axis=0))
            surface_vars = np.concatenate((np.log(1 + self.context.counts[k_component_indices]).reshape(-1, 1),
                                           crime_surface_k.reshape(-1, 1),
                                           self.context.covariates[k_component_indices, 1:]), axis=1)
            surface_vars_df = pd.DataFrame(surface_vars)
            surface_vars_df.columns = ['log-y', 'log-fitted'] + self.context.covariates_names[1:]
            corr = surface_vars_df.corr()

            f, ax = plt.subplots(figsize=(15, 10))
            sns.heatmap(corr, annot=True, linewidths=.5, ax=ax)
            plt.show()

    def load_samples_snapshot(self, iteration_no):
        beta_filepath = Path(os.getcwd()) / "models" / "snapshots" / f"beta-samples--{self.uid}--{iteration_no}.npy"
        z_filepath = Path(os.getcwd()) / "models" / "snapshots" / f"Z-samples--{self.uid}--{iteration_no}.npy"
        alpha_filepath = Path(os.getcwd()) / "models" / "snapshots" / f"alpha-samples--{self.uid}--{iteration_no}.npy"
        beta_samples = np.load(beta_filepath)
        z_samples = np.load(z_filepath)
        try:
            alpha_samples = np.load(alpha_filepath)
            return beta_samples, z_samples, alpha_samples
        except FileNotFoundError:
            alpha_samples = np.zeros((1, self.K))
            return beta_samples, z_samples, alpha_samples

    def get_initial_estimand(self):
        beta = np.random.normal(0, 1, self.context.J * self.K)
        return beta

    def get_mass_matrix_diag(self):
        beta_m_diag = 5e2 * np.ones(self.context.J * self.K)
        return beta_m_diag

    def __save_output(self, iteration):

        folder_name = Path(os.getcwd()) / "models" / "snapshots"

        if not os.path.exists(folder_name):
            os.makedirs(folder_name)

        Z_full_path = folder_name / f"Z-samples--{self.uid}--{iteration}"
        Z_samples_array = np.asarray(self.Z_samples)
        if Z_samples_array.shape[0] > 0:
            np.save(Z_full_path, Z_samples_array[::self.thinning, :])

        alpha_full_path = folder_name / f"alpha-samples--{self.uid}--{iteration}"
        alpha_array = np.asarray(self.alpha_samples)
        if alpha_array.shape[0] > 0:
            np.save(alpha_full_path, alpha_array[::self.thinning, :])

        beta_full_path = folder_name / f"beta-samples--{self.uid}--{iteration}"
        beta_array = np.asarray(self.beta_samples)
        if beta_array.shape[0] > 0:
            np.save(beta_full_path, beta_array[::self.thinning, :])

    def run_sampling(self, beta_sampler, number_of_iterations):
        N = self.context.N
        J = self.context.J
        K = self.K
        B = self.block_label_counts.shape[0]

        alpha = np.repeat(1/K, K)
        covariates = self.context.covariates
        counts = self.context.counts

        iteration = 0
        while iteration < number_of_iterations:

            ##########################################################################################
            # BOOKKEEPING
            ##########################################################################################
            # The HMC sampler is adaptive and therefore will discard samples during adaptive phase.
            if len(beta_sampler.samples) < len(self.Z_samples):
                num_current_samples = len(beta_sampler.samples)
                self.Z_samples = self.Z_samples[(-num_current_samples):]
                self.alpha_samples = self.alpha_samples[(-num_current_samples):]
                self.beta_samples = beta_sampler.samples

            if (iteration + 1) % self.n_info_interval == 0:
                self.__save_output(iteration)

            ##########################################################################################
            # SAMPLE ALPHA
            ##########################################################################################
            # We fix alpha to 1/K so there is no need for sampling. Should we choose to treat alpha
            # as random, the samples would have to be saved so we keep the code below for that eventuality.
            self.alpha_samples.append(alpha)
            self.logger.debug(f"Alpha: {alpha[0]}")

            ##########################################################################################
            # SAMPLE BETA
            ##########################################################################################
            beta_sampler.sample_one()

            ##########################################################################################
            # SAMPLE Z
            ##########################################################################################
            current_hmc_estimand = beta_sampler.estimand
            current_beta = current_hmc_estimand[:(J * K)].reshape((J, K), order='F')

            Z_new = zsampler.sample_region(self.Z.astype(np.int64),
                                           counts.astype(np.int64),
                                           covariates.astype(np.float64),
                                           current_beta.astype(np.float64),
                                           alpha.astype(np.float64),
                                           self.block_assignment_numeric.astype(np.int64),
                                           self.block_label_counts)

            self.Z_samples.append(np.where(Z_new > 0)[1])
            self.Z = Z_new

            iteration += 1

        self.logger.info("Sampling completed - saving model.")
        self.beta_samples = beta_sampler.samples
        self.__save_output(iteration)

@click.command()
@click.option('--year', '-y', type=str, default='12015-122015')
@click.option('--type', '-t', default='burglary')
@click.option('--resolution', '-r', type=int, default=400)
@click.option('--model_name', '-m', type=str, default='burglary_raw_0')
@click.option('--interpolation', '-i', type=str, default='weighted')
@click.option('--num_mixtures', '-K', type=int, default=3)
@click.option('--uid', type=str, default=None)
@click.option('--verbose', is_flag=True)
@click.option('--block_type', type=str, default="msoa")
@click.option('--collection_unit', type=str, default="lsoa")
def main(year, type, resolution, model_name, interpolation, num_mixtures, uid, verbose,
         block_type, collection_unit):

    if uid is None:
        uid = build_block_mixture_flat_uid(prefix="BLOCK-MIXTURE-FLAT", chain_no=1, block_scheme=block_type,
                                           c_type=type, t_period=year, model_spec=model_name,
                                           cov_interpolation=interpolation, resolution=resolution, K=num_mixtures)

    log_fmt = '[%(levelname)s] [%(asctime)s] [%(name)s] %(message)s'
    datefmt = '%H:%M:%S'
    if verbose:
        logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, format=log_fmt)
    else:
        logging.basicConfig(filename=Path('models') / f"log-{uid}.log",
                            filemode='a',
                            format=log_fmt,
                            datefmt=datefmt,
                            level=logging.DEBUG)
    logger = logging.getLogger(__name__)

    logger.info("Building the context.")
    grid_context = GridContextGeo(interpolation=interpolation,
                                  year=year,
                                  resolution=resolution,
                                  crime_type=type,
                                  model_name=model_name,
                                  cov_collection_unit=collection_unit,
                                  covariates_type='raw')

    logger.info("Writing sampling context into a file.")
    context_filename = Path(os.getcwd()) / "models" / f"context--{uid}.pickle"
    with open(context_filename, 'wb') as context_file:
        context_info = {
            'context': grid_context,
            'K': num_mixtures
        }
        pickle.dump(context_info, context_file)

    logger.info("Initialising the model with estimand and mass matrix diagonal")
    model = FlatRegionMixture(uid=uid,
                              grid_context=grid_context,
                              K=num_mixtures,
                              n_info_interval=50_000,
                              thinning=10,
                              block_type=block_type)
    init_estimand = model.get_initial_estimand()
    mass_matrix_diag = model.get_mass_matrix_diag()

    logger.info("Launching HMC sampler.")
    hmc_all_iterations = 120_000
    sampler = HMCSampler(func_lpdf=model.loglik,
                         func_nabla_lpdf=model.nabla_loglik,
                         func_plot=model.plot_traces if verbose else None,
                         init_estimand=init_estimand,
                         init_M_diag=mass_matrix_diag,
                         init_L=10,
                         init_epsilon=5.0e-2,
                         n_burnin=30_000,
                         n_calib=60_000,
                         S=hmc_all_iterations,
                         n_info_interval=model.n_info_interval,
                         thinning=model.thinning,
                         unique_estimation_id=uid,
                         adaptive=True)
    if verbose:
        plt.figure()
        grid_context.plot_realisations(np.log(model.block_assignment_numeric + 1), 111)
        plt.show()

    model.run_sampling(sampler, hmc_all_iterations)
    logger.info("Procedure finished.")

if __name__ == "__main__":
    load_dotenv(find_dotenv())
    main()

ext: py

reason: import numpy


text:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import json
import os
from abc import abstractmethod

import numpy as np
import torch

import nni
from nni.nas.pytorch.callbacks import ArchitectureCheckpoint
from callbacks import *
from datasets import build_dataset
from losses import build_loss_fn
from mutator import build_mutator
from networks import build_model
from utils import (MyLogger, flops_size_counter, generate_optimizer,
                           generate_scheduler, metrics,
                           parse_cfg_for_scheduler)

from .base_trainer import BaseTrainer

__all__ = [
    'TorchTensorEncoder',
    'Trainer'
]

class TorchTensorEncoder(json.JSONEncoder):
    def default(self, o):  # pylint: disable=method-hidden
        if isinstance(o, torch.Tensor):
            olist = o.tolist()
            if "bool" not in o.type().lower() and all(map(lambda d: d == 0 or d == 1, olist)):
                print("Every element in %s is either 0 or 1. "
                                "You might consider convert it into bool.", olist)
            return olist
        return super().default(o)

class Trainer(BaseTrainer):
    def __init__(self, cfg, **kwargs):
        """
        Trainer initialization.

            Parameters
            ----------
                model : nn.Module
                    Model with mutables.
                mutator : BaseMutator
                    A mutator object that has been initialized with the model.
                loss : callable
                    Called with logits and targets. Returns a loss tensor.
                metrics : callable
                    Returns a dict that maps metrics keys to metrics data.
                optimizer : Optimizer
                    Optimizer that optimizes the model.
                num_epochs : int
                    Number of epochs of training.
                dataset_train : torch.utils.data.Dataset
                    Dataset of training.
                dataset_valid : torch.utils.data.Dataset
                    Dataset of validation/testing.
                batch_size : int
                    Batch size.
                workers : int
                    Number of workers used in data preprocessing.
                device : torch.device
                    Device object. Either ``torch.device("cuda")`` or ``torch.device("cpu")``. When ``None``, trainer will
                    automatic detects GPU and selects GPU first.
                log_frequency : int
                    Number of mini-batches to log metrics.
                callbacks : list of Callback
                    Callbacks to plug into the trainer. See Callbacks.
        """
        self.cfg = cfg
        self.logger = MyLogger(__name__, cfg).getlogger()
        self.set_up()

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.mutator.to(self.device)
        self.loss.to(self.device)

        self.callbacks = self.generate_callbacks()
        for callback in self.callbacks:
            callback.build(self.model, self.mutator, self)

    def set_up(self):
        # model
        self.model = build_model(self.cfg)
        self.logger.info(f"Building model {self.cfg.model.name} ...")

        # mutator
        # self.logger.info('Cell choices: {}'.format(model.layers[0].nodes[0].cell_x.op_choice.choices))
        self.mutator = build_mutator(self.model, self.cfg)
        for x in self.mutator.mutables:
            if isinstance(x, nni.nas.pytorch.mutables.LayerChoice):
                self.logger.info('Cell choices: {}'.format(x.choices))
                break

        self.logger.info(f"Building mutator {self.cfg.mutator.name} ...")

        # dataset
        self.batch_size = self.cfg.dataset.batch_size
        self.workers = self.cfg.dataset.workers
        self.dataset_train, self.dataset_valid = build_dataset(self.cfg)
        self.logger.info(f"Building dataset {self.cfg.dataset.name} ...")

        # loss
        self.loss = build_loss_fn(self.cfg)
        self.logger.info(f"Building loss function {self.cfg.loss.name} ...")

        # optimizer
        self.optimizer = generate_optimizer(
            model=self.model,
            optim_name=self.cfg.optim.name,
            lr=self.cfg.optim.base_lr,
            momentum=self.cfg.optim.momentum,
            weight_decay=self.cfg.optim.weight_decay)
        self.logger.info(f"Building optimizer {self.cfg.optim.name} ...")

        # scheduler
        self.scheduler_params = parse_cfg_for_scheduler(self.cfg, self.cfg.optim.scheduler.name)
        self.lr_scheduler = generate_scheduler(
            self.optimizer,
            self.cfg.optim.scheduler.name,
            **self.scheduler_params)
        self.logger.info(f"Building optimizer scheduler {self.cfg.optim.scheduler.name} ...")

        # miscellaneous
        self.num_epochs = self.cfg.trainer.num_epochs
        self.log_frequency = self.cfg.logger.log_frequency
        self.start_epoch = 0

    @abstractmethod
    def train_one_epoch(self, epoch):
        pass

    @abstractmethod
    def validate_one_epoch(self, epoch):
        pass

    @abstractmethod
    def test_one_epoch(self, epoch):
        pass

    def train(self, validate=True):
        self.resume()
        self.train_meters = None
        self.valid_meters = None
        for epoch in range(self.start_epoch, self.num_epochs):
            for callback in self.callbacks:
                callback.on_epoch_begin(epoch)

            # training
            self.logger.info("Epoch {} Training".format(epoch))
            self.train_meters = self.train_one_epoch(epoch)
            self.logger.info("Final training metric: {}".format(self.train_meters))

            if validate:
                # validation
                self.logger.info("Epoch {} Validatin".format(epoch))
                self.valid_meters = self.validate_one_epoch(epoch)
                self.logger.info("Final test metric: {}".format(self.valid_meters))

            for callback in self.callbacks:
                if isinstance(callback, CheckpointCallback):
                    if self.valid_meters:
                        meters = self.valid_meters
                    else:
                        meters = self.train_meters
                    callback.update_best_metric(meters.meters['save_metric'].avg)
                callback.on_epoch_end(epoch)

    def validate(self):
        self.validate_one_epoch(-1)

    def test(self):
        return self.test_one_epoch(-1)

    def export(self, file):
        """
        Call ``mutator.export()`` and dump the architecture to ``file``.

        Parameters
        ----------
        file : str
            A file path. Expected to be a JSON.
        """
        mutator_export = self.mutator.export()
        with open(file, "w") as f:
            json.dump(mutator_export, f, indent=2, sort_keys=True, cls=TorchTensorEncoder)

    def checkpoint(self):
        """
        Return trainer checkpoint.
        """
        raise NotImplementedError("Not implemented yet")

    def enable_visualization(self):
        """
        Enable visualization. Write graph and training log to folder ``logs/<timestamp>``.
        """
        sample = None
        for x, _ in self.train_loader:
            sample = x.to(self.device)[:2]
            break
        if sample is None:
            self.logger.warning("Sample is %s.", sample)
        self.logger.info("Creating graph json, writing to %s. Visualization enabled.", self.cfg.logger.path)
        with open(os.path.join(self.cfg.logger.path, "graph.json"), "w") as f:
            json.dump(self.mutator.graph(sample), f)
        self.visualization_enabled = True

    def _write_graph_status(self):
        if hasattr(self, "visualization_enabled") and self.visualization_enabled:
            print(json.dumps(self.mutator.status()), file=self.status_writer, flush=True)

    def model_size(self, name='size'):
        assert name in ['size', 'flops']
        size = self.cfg.input.size
        if self.cfg.dataset.is_3d:
            input_size = (1,1,self.cfg.dataset.slice_num,*size)
        else:
            input_size = (1,3,*size)
        return flops_size_counter(self.model, input_size)[name]

    def generate_callbacks(self):
        '''
        Args：
            func: a function to generate other callbacks, must return a list
        Return:
            a list of callbacks.
        '''
        self.ckpt_callback = CheckpointCallback(
            checkpoint_dir=self.cfg.logger.path,
            name='best_search.pth',
            mode=self.cfg.callback.checkpoint.mode)
        self.arch_callback = ArchitectureCheckpoint(self.cfg.logger.path)
        self.relevance_callback = RelevanceCallback(
            save_path=self.cfg.logger.path,
            filename=self.cfg.callback.relevance.filename
        )
        callbacks = [self.ckpt_callback, self.arch_callback,self.relevance_callback]
        return callbacks

    def metrics(self, *args, **kwargs):
        return metrics( *args, **kwargs)

    def resume(self):
        self.best_metric = -999
        path = self.cfg.model.resume_path
        if path:
            assert os.path.exists(path), "{} does not exist".format(path)
            ckpt = torch.load(path)
            self.start_epoch = ckpt['epoch'] + 1
            self.model.load_state_dict(ckpt['model_state_dict'])
            self.mutator.load_state_dict(ckpt['mutator_state_dict'])
            self.optimizer.load_state_dict(ckpt['optimizer_state_dict'])
            self.lr_scheduler.load_state_dict(ckpt['lr_scheduler_state_dict'])
            self.best_metric = ckpt['best_metric']
            self.logger.info('Resuming training from epoch {}'.format(self.start_epoch))

        for callback in self.callbacks:
            if isinstance(callback, CheckpointCallback):
                if self.best_metric == -999:
                    self.best_metric = callback.best_metric
                else:
                    callback.best_metric = self.best_metric

        if len(self.cfg.trainer.device_ids) > 1:
            device_ids = self.cfg.trainer.device_ids
            num_gpus_available = torch.cuda.device_count()
            assert num_gpus_available >= len(device_ids), "you can only use {} device(s)".format(num_gpus_available)
            self.model = torch.nn.DataParallel(self.model, device_ids=device_ids)
            # self.mutator = torch.nn.DataParallel(self.mutator, device_ids=device_ids) # mutator doesn't support dataparallel yet.

    def save_cfg(self):
        cfg_file = self.cfg.logger.cfg_file
        with open(cfg_file, 'w') as f:
            f.write(str(self.cfg))
        print(f'Saving config file to {cfg_file}')


text:
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
import numpy as np
import pandas as pd
from scipy import optimize
import math
import csv
import imageio
import sys as sys
sys.path.insert(1, './Tools')

from interp import interp
from finite_differences import fd_d1_o4
from Gaussian_fit import gaussian_fit
from Gaussian_fit import gaussian_fit_GUI
from max_pedestal_finder import find_pedestal_from_data
from max_pedestal_finder import find_pedestal
from read_profiles import read_profile_file
from read_profiles import read_geom_file
from read_EFIT_file import get_geom_pars
from write_iterdb import output_iterdb
from max_stat_tool import Poly_fit
from max_stat_tool import coeff_to_Poly
from DispersionRelationDeterminantFullConductivityZeff import VectorFinder_auto
from DispersionRelationDeterminantFullConductivityZeff import VectorFinder_manual
from DispersionRelationDeterminantFullConductivityZeff import VectorFinder_auto_Extensive
#Created by Max T. Curie  09/17/2021
#Last edited by Max Curie 09/17/2021
#Supported by scripts in IFS

class mode_finder:
    profile_type=('pfile','ITERDB')
    geomfile_type=('gfile','GENE_tracor')

    def __init__(self,profile_type,profile_name,geomfile_type,geomfile_name,\
            outputpath,inputpath,x0_min,x0_max,zeff_manual=-1,suffix='.dat',\
            mref=2.,Impurity_charge=6.):
        n0=1.
        Z=float(Impurity_charge)
        print("suffix"+str(suffix))
        self.profile_name=profile_name
        self.geomfile_name=geomfile_name
        self.outputpath=outputpath
        self.inputpath=inputpath
        self.x0_min=x0_min
        self.x0_max=x0_max

        if profile_type not in mode_finder.profile_type:
            raise ValueError(f'{profile_type} is not a valid profile type, need to be pfile or ITERDB')
        else:
            self.profile_type  = profile_type

        if geomfile_type not in mode_finder.geomfile_type:
            raise ValueError(f'{geomfile_type} is not a valid geomfile_type, need to be gfile or GENE_tracor')
        else:
            self.geomfile_type  = geomfile_type

        #*************Loading the data******************************************
        if profile_type=="ITERDB":
            rhot0, rhop0, te0, ti0, ne0, ni0, nz0, vrot0 = read_profile_file(profile_type,profile_name,geomfile_name,suffix)
        elif profile_type=="pfile":
            rhot0, rhop0, te0, ti0, ne0, ni0, nz0, vrot0 = read_profile_file(profile_type,profile_name,geomfile_name,suffix)

        if geomfile_type=="gfile":
            xgrid, q, R_ref= read_geom_file(geomfile_type,geomfile_name,suffix)
        elif geomfile_type=="GENE_tracor":
            xgrid, q, Lref, R_ref, Bref, x0_from_para = read_geom_file(geomfile_type,geomfile_name,suffix)

        if geomfile_type=="GENE_tracor" and profile_type!="profile":
            rhot0_range_min=np.argmin(abs(rhot0-xgrid[0]))
            rhot0_range_max=np.argmin(abs(rhot0-xgrid[-1]))
            rhot0=rhot0[rhot0_range_min:rhot0_range_max]
            rhop0=rhop0[rhot0_range_min:rhot0_range_max]
            te0=te0[rhot0_range_min:rhot0_range_max]
            ti0=ti0[rhot0_range_min:rhot0_range_max]
            ne0=ne0[rhot0_range_min:rhot0_range_max]
            ni0=ni0[rhot0_range_min:rhot0_range_max]
            vrot0=vrot0[rhot0_range_min:rhot0_range_max]

        uni_rhot = np.linspace(min(rhot0),max(rhot0),2000)

        uni_rhop = interp(rhot0,rhop0,uni_rhot)
        ti_u = interp(rhot0,ti0,uni_rhot)
        ni_u = interp(rhot0,ni0,uni_rhot)
        te_u = interp(rhot0,te0,uni_rhot)
        ne_u = interp(rhot0,ne0,uni_rhot)
        nz_u = interp(rhot0,nz0,uni_rhot)
        vrot_u = interp(rhot0,vrot0,uni_rhot)
        q      = interp(xgrid,q,uni_rhot)
        tprime_e = -fd_d1_o4(te_u,uni_rhot)/te_u
        nprime_e = -fd_d1_o4(ne_u,uni_rhot)/ne_u
        qprime = fd_d1_o4(q,uni_rhot)/q

        print(str(x0_min)+', '+str(x0_max))

        midped, topped=find_pedestal(file_name=geomfile_name, path_name='', plot=False)
        x0_center=midped
        print('mid pedestal is at r/a = '+str(x0_center))

        if geomfile_type=="gfile":
            Lref, Bref, R_major, q0, shat0=get_geom_pars(geomfile_name,x0_center)

        print("Lref="+str(Lref))

        index_begin=np.argmin(abs(uni_rhot-x0_min))
        index_end  =np.argmin(abs(uni_rhot-x0_max))

        te_u = te_u[index_begin:index_end]
        ne_u = ne_u[index_begin:index_end]
        ni_u = ni_u[index_begin:index_end]
        nz_u = nz_u[index_begin:index_end]
        vrot_u = vrot_u[index_begin:index_end]
        q      = q[index_begin:index_end]
        tprime_e = tprime_e[index_begin:index_end]
        nprime_e = nprime_e[index_begin:index_end]
        qprime   = qprime[index_begin:index_end]
        uni_rhot = uni_rhot[index_begin:index_end]

        center_index = np.argmin(abs(uni_rhot-x0_center))

        #*************End of loading the data******************************************

        #****************Start setting up ******************
        ne=ne_u/(10.**19.)      # in 10^19 /m^3
        ni=ni_u/(10.**19.)      # in 10^19 /m^3
        nz=nz_u/(10.**19.)      # in 10^19 /m^3
        te=te_u/1000.           # in keV
        ti=ti_u/1000.           # in keV
        m_SI = mref *1.6726*10**(-27)
        me_SI = 9.11*10**(-31)
        c  = 1.
        qref = 1.6*10**(-19)
        #refes to GENE manual
        coll_c=2.3031*10**(-5)*Lref*ne/(te)**2*(24-np.log(np.sqrt(ne*10**13)/(te*1000)))
        coll_ei=4.*coll_c*np.sqrt(te*1000.*qref/me_SI)/Lref
        nuei=coll_ei
        beta=403.*10**(-5)*ne*te/Bref**2.

        nref=ne_u[center_index]
        te_mid=te_u[center_index]
        Tref=te_u[center_index] * qref

        cref = np.sqrt(Tref / m_SI)
        Omegaref = qref * Bref / m_SI / c
        rhoref = cref / Omegaref
        rhoref_temp = rhoref * np.sqrt(te_u/te_mid)
        kymin=n0*q0*rhoref/(Lref*x0_center)
        kyGENE =kymin * (q/q0) * np.sqrt(te_u/te_mid) * (x0_center/uni_rhot) #Add the effect of the q varying
        #from mtm_doppler
        omMTM = kyGENE*(tprime_e+nprime_e)
        gyroFreq = 9.79E3/np.sqrt(mref)*np.sqrt(te_u)/Lref
        mtmFreq = omMTM*gyroFreq/(2.*np.pi*1000.)
        omegaDoppler = abs(vrot_u*n0/(2.*np.pi*1000.))
        omega=mtmFreq + omegaDoppler

        if zeff_manual==-1:
            zeff = (  ( ni+nz*(float(Z)**2.) )/ne  )[center_index]
        else:
            zeff=zeff_manual
        print('********zeff*********')
        print('zeff='+str(zeff))
        print('********zeff*********')

        omega_n_GENE=kyGENE*(nprime_e)       #in cs/a
        omega_n=omega_n_GENE*gyroFreq/(2.*np.pi*1000.)  #in kHz

        coll_ei=coll_ei/(2.*np.pi*1000.)  #in kHz

        Lt=1./tprime_e
        Ln=1./nprime_e
        Lq=1./(Lref/(R_ref*q)*qprime)

        shat=Ln/Lq
        eta=Ln/Lt
        ky=kyGENE*np.sqrt(2.)
        nu=(coll_ei)/(np.max(omega_n))

        self.mref=mref
        self.Z=Z
        self.Bref=Bref

        self.ne=ne      # in 10^19 /m^3
        self.ni=ni      # in 10^19 /m^3
        self.nz=nz      # in 10^19 /m^3
        self.te=te      # in keV
        self.ti=ti      # in keV
        self.nref=nref
        self.r_sigma=(1./shat)*( (me_SI/m_SI)**0.5 )
        self.R_ref=R_ref
        self.cs_to_kHz=gyroFreq/(2.*np.pi*1000.)
        self.omn=omega_n    #omega_n in kHz
        self.cs=cref
        self.rho_s=rhoref*np.sqrt(2.)
        self.Lref=Lref
        self.x=uni_rhot
        self.rhop=uni_rhop
        self.shat=shat
        self.eta=eta
        self.ky=ky
        self.nu=nu
        self.zeff=zeff
        self.beta=beta
        self.q=q
        self.ome=mtmFreq
        self.Doppler=omegaDoppler
        self.coll_ei=coll_ei
        self.gyroFreq=gyroFreq

        self.x_nominal=uni_rhot
        self.nref_nominal=nref
        self.ne_nominal=self.ne
        self.te_nominal=self.te
        self.r_sigma_nominal=self.r_sigma
        self.cs_to_kHz_nominal=self.cs_to_kHz
        self.omn_nominal=self.omn
        self.cs_nominal=self.cs
        self.rho_s_nominal=self.rho_s
        self.shat_nominal=self.shat
        self.eta_nominal=self.eta
        self.ky_nominal=self.ky
        self.nu_nominal=self.nu
        self.zeff_nominal=self.zeff
        self.beta_nominal=self.beta
        self.q_nominal=self.q
        self.ome_nominal=self.ome
        self.Doppler_nominal=self.Doppler
        self.coll_ei_nominal=self.coll_ei
        self.gyroFreq_nominal=self.gyroFreq

    def __str__(self):
        try:
            self.x_peak
        except:
            self.ome_peak_range(0.1)
        index=np.argmin(abs(self.x-self.x_peak))

        return f'Parameter at peak:\n\
        shat={self.shat[index]},\n\
        eta={self.eta[index]},\n\
        ky={self.ky[index]},\n\
        nu={self.nu[index]},\n\
        zeff={self.zeff},\n\
        beta={self.beta[index]},\n\
        r_sigma={self.r_sigma[index]},\n\
        xstar={self.xstar}'

    def set_xstar(self,xstar):
        self.xstar=xstar
        #setter for xstar

    def find_pedestal(self,x,p):
        dp=np.gradient(p,x)  #Second order of pressure
        ddp   = np.gradient(dp,x)  #Second order of pressure
        midped = x[np.argmin(dp)]
        topped = x[np.argmin(ddp)]
        return topped,midped

    def inside_freq_band_check(self,f,freq_min_list,freq_max_list):
        for freq_min, freq_max in zip(freq_min_list,freq_max_list):
            if freq_min<=f and f<=freq_max:
                return True
                break

    def profile_fit(self,show_plot=False):
        x=self.x
        ne=self.ne
        te=self.te

        if show_plot==True:
            fig, ax=plt.subplots(nrows=3,ncols=1,sharex=True)
            ax[0].plot(x,ne,label='original')
            ax[0].set_ylabel(r'$n_e(10^19/m^3)$',fontsize=20)
            ax[1].plot(x,te,label='original')
            ax[1].set_ylabel(r'$T_e(keV)$',fontsize=20)
            ax[1].set_xlabel(r'$\rho_{tor}$',fontsize=20)
            plt.tight_layout()
            plt.show()

        ne_popt, pcov = optimize.curve_fit(tanh_func, x, ne)
        ne_mid_ped, ne_ped_width, ne_height,ne_edge=ne_popt[0], ne_popt[1], ne_popt[2], ne_popt[3]

        te_popt, pcov = optimize.curve_fit(tanh_func, x, te)
        te_mid_ped, te_ped_width, te_height,te_edge=te_popt[0], te_popt[1], te_popt[2], te_popt[3]

        if show_plot==True:
            ne_fit=tanh_func(x, ne_popt[0], ne_popt[1], ne_popt[2], ne_popt[3])
            te_fit=tanh_func(x, te_popt[0], te_popt[1], te_popt[2], te_popt[3])

            fig, ax=plt.subplots(nrows=2,ncols=1,sharex=True)
            ax[0].plot(x,ne,label='original')
            ax[0].plot(x,ne_fit,label='fit')
            ax[0].legend()
            ax[0].set_ylabel(r'$n_e(10^19/m^3)$',fontsize=10)
            ax[1].plot(x,te,label='original')
            ax[1].plot(x,te_fit,label='fit')
            ax[1].set_ylabel(r'$T_e(keV)$',fontsize=10)
            ax[1].set_xlabel(r'$\rho_{tor}$',fontsize=10)
            plt.tight_layout()
            plt.show()

        self.ne_mid_ped=ne_mid_ped
        self.ne_ped_width=ne_ped_width
        self.ne_height=ne_height
        self.ne_edge=ne_edge
        self.te_mid_ped=te_mid_ped
        self.te_ped_width=te_ped_width
        self.te_height=te_height
        self.te_edge=te_edge

        return ne_mid_ped, ne_ped_width, ne_height,ne_edge,\
                te_mid_ped, te_ped_width, te_height,te_edge

    def q_fit(self,q_order=5,show_plot=False):
        dq= np.gradient(self.q,self.x)
        x_dq_min=x[np.argmin(abs(dq))]
        q_coeff=Poly_fit(self.x-x_dq_min, self.q, q_order, show=False)
        q_fit=coeff_to_Poly(self.x, x_dq_min, q_coeff, show=False)
        if show_plot==Ture:
            plt.clf()
            plt.plot(self.x,self.q,label='orginal')
            plt.plot(self.x,q_fit,label='fit')
            plt.legend()
            plt.show()

    def modify_profile(self,\
                    q_scale=1.,q_shift=0.,\
                    shat_scale=1.,\
                    ne_scale=1.,te_scale=1.,\
                    ne_shift=0.,te_shift=0.,\
                    Doppler_scale=1.,\
                    show_plot=False):

        #************start of modification***********
        #********************************************

        #**********************************
        #*************modify ne************
        try:
            self.ne_weight
        except:
            ne_top_ped,ne_mid_ped=self.find_pedestal(self.x,self.ne)
            self.ne_mid_ped=ne_mid_ped
            ne_width=ne_mid_ped-ne_top_ped

            self.ne_weight = 0.5+np.tanh((self.x-ne_top_ped)/ne_width)/2

        index=np.argmin(abs(self.x-self.ne_mid_ped))
        ne_mod=self.ne*(1.+(ne_scale-1.)*self.ne_weight)+(ne_shift)*self.ne[index]
        #*************modify ne************
        #**********************************

        #**********************************
        #*************modify Te************
        try:
            self.te_weight
        except:
            te_top_ped,te_mid_ped=self.find_pedestal(self.x,self.te)
            self.te_mid_ped=te_mid_ped
            te_width=te_mid_ped-te_top_ped
            self.te_weight = 0.5+np.tanh((self.x-te_top_ped)/te_width)/2

        index=np.argmin(abs(self.x-self.te_mid_ped))

        te_mod=self.te*(1.+(te_scale-1.)*self.te_weight)+(te_shift)*self.te[index]
        #*************modify Te************
        #**********************************

        mid_ped,top_ped=self.find_pedestal(self.x,self.ome)

        width=(mid_ped-top_ped)*10.
        try:
            self.shat_weight
        except:
            self.shat_weight = -1.*np.tanh((self.x-mid_ped)/width)/2

        q_mod=self.q*q_scale*(1.+(shat_scale-1.)*self.shat_weight)+q_shift

        Doppler_mod=self.Doppler*Doppler_scale

        #show_plot=True

        if show_plot==True:
            fig, ax=plt.subplots(nrows=4,ncols=1,sharex=True)
            ax[0].plot(self.x,self.ne,label='original')
            ax[0].plot(self.x,ne_mod,label='modified')
            ax[0].legend()
            ax[0].set_ylabel(r'$n_e(10^19/m^3)$',fontsize=10)
            ax[1].plot(self.x,self.te)
            ax[1].plot(self.x,te_mod)
            ax[1].set_ylabel(r'$T_e(keV)$',fontsize=10)
            ax[2].plot(self.x,self.q)
            ax[2].plot(self.x,q_mod)
            ax[2].set_ylabel(r'$q$',fontsize=10)
            ax[3].plot(self.x,self.Doppler)
            ax[3].plot(self.x,Doppler_mod)
            ax[3].set_ylabel(r'$f_{Doppler}(kHz)$',fontsize=10)
            ax[3].set_xlabel(r'$\rho_{tor}$',fontsize=10)
            plt.tight_layout()
            plt.show()

        #*************end of modification************
        #********************************************

        tprime_e = -fd_d1_o4(te_mod,self.x)/te_mod
        nprime_e = -fd_d1_o4(ne_mod,self.x)/ne_mod
        qprime = fd_d1_o4(q_mod,self.x)/q_mod

        center_index=np.argmin(abs(self.x-mid_ped))

        Tref_mod=te_mod[center_index]

        #*********start of calculation******
        #****************Start setting up ******************
        #get the data from attribute
        ni=self.ni
        nz=self.nz

        te_u=te_mod*1000.
        ne_u=ne_mod*10.**19.
        ni_u=ni*10.**19.
        nz_u=nz*10.**19.
        mref=self.mref
        Lref=self.Lref
        Bref=self.Bref
        uni_rhot=self.x
        Z=self.Z
        R_ref=self.R_ref

        x0_center=mid_ped
        ne=ne_mod
        te=te_mod
        q=q_mod

        n0=1.
        m_SI = mref *1.6726*10**(-27)
        me_SI = 9.11*10**(-31)
        c  = 1.
        qref = 1.6*10**(-19)
        #refes to GENE manual
        coll_c=2.3031*10**(-5)*Lref*ne/(te)**2*(24-np.log(np.sqrt(ne*10**13)/(te*1000)))
        coll_ei=4.*coll_c*np.sqrt(te*1000.*qref/me_SI)/Lref
        nuei=coll_ei
        beta=403.*10**(-5)*ne*te/Bref**2.

        nref=ne_u[center_index]
        te_mid=te_u[center_index]
        Tref=te_u[center_index]*qref
        q0=q_mod[center_index]

        cref = np.sqrt(Tref / m_SI)
        Omegaref = qref * Bref / m_SI / c
        rhoref = cref / Omegaref
        rhoref_temp = rhoref * np.sqrt(te_u/te_mid)
        kymin=n0*q0*rhoref/(Lref*x0_center)
        kyGENE =kymin * (q/q0) * np.sqrt(te_u/te_mid) * (x0_center/uni_rhot) #Add the effect of the q varying
        #from mtm_doppler
        omMTM = kyGENE*(tprime_e+nprime_e)
        gyroFreq = 9.79E3/np.sqrt(mref)*np.sqrt(te_u)/Lref
        mtmFreq = omMTM*gyroFreq/(2.*np.pi*1000.)
        omegaDoppler = Doppler_mod
        omega=mtmFreq + omegaDoppler

        zeff=self.zeff*nref/self.nref

        omega_n_GENE=kyGENE*(nprime_e)       #in cs/a
        omega_n=omega_n_GENE*gyroFreq/(2.*np.pi*1000.)  #in kHz

        coll_ei=coll_ei/(2.*np.pi*1000.)  #in kHz

        Lt=1./tprime_e
        Ln=1./nprime_e
        Lq=1./(Lref/(R_ref*q)*qprime)

        shat=Ln/Lq
        eta=Ln/Lt
        ky=kyGENE*np.sqrt(2.)
        nu=(coll_ei)/(np.max(omega_n))

        self.x=self.x[2:-2]
        self.nref=nref
        self.ne=ne[2:-2]      # in 10^19 /m^3
        self.ni=ni[2:-2]      # in 10^19 /m^3
        self.nz=nz[2:-2]      # in 10^19 /m^3
        self.te=te[2:-2]      # in keV
        self.r_sigma=(1./shat[:-2])*( (me_SI/m_SI)**0.5 )
        self.cs_to_kHz=gyroFreq[:-2]/(2.*np.pi*1000.)
        self.omn=omega_n[2:-2]    #omega_n in kHz
        self.cs=cref
        self.rho_s=rhoref*np.sqrt(2.)
        self.shat=shat[2:-2]
        self.eta=eta[2:-2]
        self.ky=ky[2:-2]
        self.nu=nu[2:-2]
        self.zeff=zeff
        self.beta=beta[2:-2]
        self.q=q[2:-2]
        self.ome=mtmFreq[2:-2]
        self.Doppler=omegaDoppler[2:-2]
        self.coll_ei=coll_ei[2:-2]
        self.gyroFreq=gyroFreq[2:-2]
        #**********end of calculation*******
        if show_plot==True:
            fig, ax=plt.subplots(nrows=7,\
            ncols=1,sharex=True)

            ax[0].plot(self.x_nominal,self.shat_nominal,label='nominal')
            ax[0].plot(self.x,self.shat,label='modified')
            ax[0].legend()
            ax[0].set_ylabel(r'$L_{ne}/L_{q}$')
            ax[1].plot(self.x_nominal,self.eta_nominal)
            ax[1].plot(self.x,self.eta)
            ax[1].set_ylabel(r'$L_{ne}/L_{Te}$')
            ax[2].plot(self.x_nominal,self.ky_nominal)
            ax[2].plot(self.x,self.ky)
            ax[2].set_ylabel(r'$k_y\rho_s$')
            ax[3].plot(self.x_nominal,self.nu_nominal)
            ax[3].plot(self.x,self.nu)
            ax[3].set_ylabel(r'$\nu_{ei}/\omega_{*ne}$')
            ax[4].plot(self.x_nominal,self.beta_nominal)
            ax[4].plot(self.x,self.beta)
            ax[4].set_ylabel(r'$\beta$')
            ax[5].plot(self.x_nominal,self.q_nominal)
            ax[5].plot(self.x,self.q)
            ax[5].set_ylabel(r'$q$')
            ax[6].plot(self.x_nominal,self.ome_nominal)
            ax[6].plot(self.x,self.ome)
            ax[6].set_ylabel(r'$\omega_{*e}(kHz)$')
            ax[6].set_xlabel(r'$\rho_{tor}$')
            plt.subplots_adjust(wspace=0, hspace=0)
            plt.show()

    def reset_profile(self):
        self.x=self.x_nominal
        self.ne=self.ne_nominal
        self.te=self.te_nominal
        self.r_sigma=self.r_sigma_nominal
        self.cs_to_kHz=self.cs_to_kHz_nominal
        self.omn=self.omn_nominal
        self.cs=self.cs_nominal
        self.rho_s=self.rho_s_nominal
        self.shat=self.shat_nominal
        self.eta=self.eta_nominal
        self.ky=self.ky_nominal
        self.nu=self.nu_nominal
        self.zeff=self.zeff_nominal
        self.beta=self.beta_nominal
        self.q=self.q_nominal
        self.ome=self.ome_nominal
        self.Doppler=self.Doppler_nominal
        self.coll_ei=self.coll_ei_nominal
        self.gyroFreq=self.gyroFreq_nominal
        self.nref=self.nref_nominal

    def q_modify(self,q_scale,q_shift,shat_scale=1.):
        mid_ped,top_ped=self.find_pedestal(self.x,self.te)

        width=(mid_ped-top_ped)*10.
        try:
            self.shat_weight
        except:
            self.shat_weight = -1.*np.tanh((self.x-mid_ped)/width)/2

        q_mod=self.q*q_scale*(1.+(shat_scale-1.)*self.shat_weight)+q_shift

        self.ome= self.ome*q_mod/self.q
        self.omn= self.omn*q_mod/self.q
        self.shat= self.shat*( (q_scale*self.q)/q_mod )
        self.ky=self.ky*q_mod/self.q

        self.q=q_mod

        return self.q

    def q_back_to_nominal(self):
        self.q=self.q_nominal
        self.ome= self.ome_nominal
        self.omn= self.omn_nominal
        self.shat= self.shat_nominal
        self.ky=self.ky_nominal

    def omega_gaussian_fit_GUI(self,root,x,data,rhoref,Lref):
        amplitude,mean,stddev=gaussian_fit_GUI(root,x,data)
        #print(f'amplitude,mean,stddev={amplitude},{mean},{stddev}')
        mean_rho=mean*Lref/rhoref         #normalized to rhoi
        xstar=abs(stddev*Lref/rhoref)

        popt=[0]*3
        popt[0] = amplitude
        popt[1] = mean
        popt[2] = stddev

        #print(popt)
        #print(mean_rho,xstar)

        return mean_rho,xstar

    def omega_gaussian_fit(self,manual=False):
        x=self.x
        data=self.ome
        rhoref=self.rho_s
        Lref=self.Lref

        amplitude,mean,stddev=gaussian_fit(x,data,manual)
        #print(f'amplitude,mean,stddev={amplitude},{mean},{stddev}')
        mean_rho=mean*Lref/rhoref         #normalized to rhoi
        xstar=abs(stddev*Lref/rhoref)

        popt=[0]*3
        popt[0] = amplitude
        popt[1] = mean
        popt[2] = stddev

        #print(popt)
        #print(mean_rho,xstar)

        return mean_rho,xstar

    def ome_peak_range(self,peak_percent=0.01):
        y1=self.ome
        x=self.x
        peak_index=np.argmax(self.ome)

        left_x=x[:peak_index]
        right_x=x[peak_index:]
        left_y1 =y1[:peak_index]
        right_y1=y1[peak_index:]

        left_index=np.argmin(abs(left_y1  - y1[peak_index]*(1.-peak_percent)))
        right_index=np.argmin(abs(right_y1  - y1[peak_index]*(1.-peak_percent)))

        self.x_peak=self.x[peak_index]
        self.x_min=left_x[left_index]
        self.x_max=right_x[right_index]

        return self.x_peak,self.x_min,self.x_max

    def Rational_surface(self,n0):
        q=self.q
        uni_rhot=self.x

        x_list=[]
        m_list=[]

        q_prime=(fd_d1_o4(q,uni_rhot)/q)[1:-1]
        peak_index=np.argmin(abs(q_prime))

        qmin = np.min(q)
        qmax = np.max(q)
        m_min = math.ceil(qmin*n0)
        m_max = math.floor(qmax*n0)
        mnums = np.arange(m_min,m_max+1)
        dq=np.max(abs(q[:-1]-q[1:]))

        if np.min(q_prime)<-0.001: #for non-monotonic increasing q profiles
            for m in mnums:
                #print(m)
                x_list_temp=[]
                index_list_temp=[]
                q0=float(m)/float(n0)
                for index0 in range(len(q)):
                    if abs(q[index0]-q0)<dq:
                        x_list_temp.append(uni_rhot[index0])
                        index_list_temp.append(index0)
                index_diff=np.array(index_list_temp[:-1],dtype=int)-np.array(index_list_temp[1:],dtype=int)
                #print(index_diff)
                #print(x_list_temp)
                #print(index_list_temp)
                if len(x_list_temp)>0:
                    x_list.append(x_list_temp[0])
                    m_list.append(m)
                for i in range(len(index_diff)):
                    if index_diff[i]==-1:
                        pass
                    else:
                        x_list.append(x_list_temp[i])
                        m_list.append(m)

            #print(x_list)
            #print(m_list)
            return x_list, m_list
        else: #for monotonic increasing q profiles
            for m in mnums:
                q0=float(m)/float(n0)
                index0=np.argmin(abs(q-q0))
                if abs(q[index0]-q0)<0.1:
                    x_list.append(uni_rhot[index0])
                    m_list.append(m)

            #print(x_list)
            #print(m_list)
            return x_list, m_list

    def Rational_surface_peak_surface(self,n0):
        try:
            x_peak=self.x_peak
        except:
            x_peak,x_min,x_max=self.ome_peak_range(peak_percent=0.1)

        x_list, m_list=self.Rational_surface(n0)
        index=np.argmin(abs(x_list-x_peak))
        x_surface_near_peak=x_list[index]
        m_surface_near_peak=m_list[index]

        return x_surface_near_peak, m_surface_near_peak

    def Rational_surface_top_surfaces(self,n0,top=1):
        try:
            x_peak=self.x_peak
        except:
            x_peak,x_min,x_max=self.ome_peak_range(peak_percent=0.1)

        x_list, m_list=self.Rational_surface(n0)
        temp_ndarray=(np.array([x_list, m_list,abs(x_list-x_peak)])).transpose()
        surface_df = pd.DataFrame(temp_ndarray,\
                    columns = ['x_list','m_list','peak_distance'])
        #print(surface_df)
        surface_df_sort=surface_df.sort_values(by='peak_distance')
        if len(surface_df_sort['x_list'])<top:
            x_surface_near_peak_list=surface_df_sort['x_list']
            m_surface_near_peak_list=surface_df_sort['m_list']
        else:
            x_surface_near_peak_list=surface_df_sort['x_list'][:top]
            m_surface_near_peak_list=surface_df_sort['m_list'][:top]
        x_surface_near_peak_list=np.array(x_surface_near_peak_list)
        m_surface_near_peak_list=np.array(m_surface_near_peak_list)
        return x_surface_near_peak_list, m_surface_near_peak_list

    def parameter_for_dispersion(self,x0,n):
        index=np.argmin(abs(x0-self.x))

        factor_temp=np.sqrt(\
                (float(n)*self.q[index]/self.Lref)**2.\
                +(float(n)/self.R_ref)**2.)\
                *self.Lref/self.q[index]

        self.factor=factor_temp

        nu=self.nu[index]/factor_temp
        zeff=self.zeff
        eta=self.eta[index]
        shat=self.shat[index]
        beta=self.beta[index]
        ky=self.ky[index]*factor_temp

        try:
            x_peak=self.x_peak
        except:
            x_peak,x_min,x_max=self.ome_peak_range(peak_percent=0.1)

        try:
            xstar=self.xstar
        except:
            xstar=0.
        mu=(x0-x_peak)*self.Lref/(self.rho_s)
        return nu,zeff,eta,shat,beta,ky,mu,xstar

    def Dispersion(self,nu,zeff,eta,shat,beta,ky,ModIndex,mu,xstar,manual=False):
        if manual==True:
            w0=VectorFinder_manual(nu,zeff,eta,shat,beta,ky,ModIndex,abs(mu),xstar)
        elif manual==5:
            w0=VectorFinder_auto_Extensive(nu,zeff,eta,shat,beta,ky,ModIndex,abs(mu),xstar)
        else:
            w0=VectorFinder_auto(nu,zeff,eta,shat,beta,ky,ModIndex,abs(mu),xstar)
        return w0

    def plot_parameter(self):
        fig, ax=plt.subplots(nrows=7,\
            ncols=1,sharex=True)

        ax[0].plot(self.x,self.shat)
        ax[0].set_ylabel(r'$L_{ne}/L_{q}$')
        ax[1].plot(self.x,self.eta)
        ax[1].set_ylabel(r'$L_{ne}/L_{Te}$')
        ax[2].plot(self.x,self.ky)
        ax[2].set_ylabel(r'$k_y\rho_s$')
        ax[3].plot(self.x,self.nu)
        ax[3].set_ylabel(r'$\nu_{ei}/\omega_{*ne}$')
        ax[4].plot(self.x,self.beta)
        ax[4].set_ylabel(r'$\beta$')
        ax[5].plot(self.x,self.q,label=r'$q$')
        ax[5].set_ylabel(r'$q$')
        ax[6].plot(self.x,self.ome)
        ax[6].set_ylabel(r'$\omega_{*e}(kHz)$')

        ax[6].set_xlabel(r'$\rho_{tor}$')
        plt.subplots_adjust(wspace=0, hspace=0)
        plt.show()

    def Plot_q_scale_rational_surfaces_colored(self,peak_percent,\
            q_scale,q_shift,q_uncertainty,n_list,unstable_list,color_list):
        marker_size=5
        x_peak,x_stability_boundary_min,x_stability_boundary_max=self.ome_peak_range(peak_percent)

        x_min_index=np.argmin(abs(self.x-x_stability_boundary_min))
        x_max_index=np.argmin(abs(self.x-x_stability_boundary_max))

        q_range=self.q[x_min_index:x_max_index]

        x_list=[]
        y_list=[]
        x_list_error=[]
        y_list_error=[]

        fig, ax = plt.subplots()

        ax.fill_between(self.x, self.q*(1.-q_uncertainty), self.q*(1.+q_uncertainty), color='blue', alpha=.0)
        ax.axvline(x_stability_boundary_min,color='orange',alpha=0.6)
        ax.axvline(x_stability_boundary_max,color='orange',alpha=0.6)
        ax.plot(self.x,self.q,label=r'Safety factor $q_0$')

        for i in range(len(n_list)):
            n=n_list[i]
            x1=[]
            y1=[]
            x1_error=[]
            y1_error=[]

            qmin = np.min(self.q*(1.-q_uncertainty))
            qmax = np.max(self.q*(1.+q_uncertainty))

            m_min = math.ceil(qmin*n)
            m_max = math.floor(qmax*n)
            m_list=np.arange(m_min,m_max+1)

            for m in m_list:
                surface=float(m)/float(n)
                if np.min(q_range)*(1.-q_uncertainty)<surface and surface<np.max(q_range)*(1.+q_uncertainty):
                    x1.append(0.5*(x_stability_boundary_max+x_stability_boundary_min))
                    y1.append(surface)
                    x1_error.append(0.5*(x_stability_boundary_max-x_stability_boundary_min))
                    y1_error.append(0)
            if color_list==-1:
                ax.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',linestyle='none',ms=marker_size,label='n='+str(n))
            else:
                ax.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',color=color_list[i],linestyle='none',ms=marker_size,label='n='+str(n))

            x_list.append(x1)
            y_list.append(y1)
            x_list_error.append(x1_error)
            y_list_error.append(y1_error)

        ax.plot(self.x,self.q*q_scale+q_shift,color='orange',label=r'Modified $q=$'+str(q_scale)+r'$*q_0$')
        ax.set_xlim(self.x0_min,self.x0_max)
        ax.set_xlabel(r'$\rho_{tor}$')
        ax.set_ylabel('Safety factor')
        plt.legend(loc='upper left')
        plt.show()

    #color_lsit=-1 for auto color assignment
    def Plot_q_scale_rational_surfaces_colored_obj(self,peak_percent,\
            q_scale,q_shift,q_uncertainty,n_list,unstable_list,color_list):
        marker_size=5
        x_peak,x_stability_boundary_min,x_stability_boundary_max=self.ome_peak_range(peak_percent)

        x_min_index=np.argmin(abs(self.x-x_stability_boundary_min))
        x_max_index=np.argmin(abs(self.x-x_stability_boundary_max))

        q_range=self.q[x_min_index:x_max_index]

        x_list=[]
        y_list=[]
        x_list_error=[]
        y_list_error=[]
        fig = Figure(figsize = (5, 5),\
                        dpi = 100)
        plot1 = fig.add_subplot(111)

        plot1.fill_between(self.x, self.q*(1.-q_uncertainty), self.q*(1.+q_uncertainty), color='blue', alpha=.0)
        plot1.axvline(x_stability_boundary_min,color='orange',alpha=0.6)
        plot1.axvline(x_stability_boundary_max,color='orange',alpha=0.6)
        plot1.plot(self.x,self.q,label=r'Safety factor $q_0$')

        for i in range(len(n_list)):
            n=n_list[i]
            x1=[]
            y1=[]
            x1_error=[]
            y1_error=[]

            qmin = np.min(self.q*(1.-q_uncertainty))
            qmax = np.max(self.q*(1.+q_uncertainty))

            m_min = math.ceil(qmin*n)
            m_max = math.floor(qmax*n)
            m_list=np.arange(m_min,m_max+1)

            for m in m_list:
                surface=float(m)/float(n)
                if np.min(q_range)*(1.-q_uncertainty)<surface and surface<np.max(q_range)*(1.+q_uncertainty):
                    x1.append(0.5*(x_stability_boundary_max+x_stability_boundary_min))
                    y1.append(surface)
                    x1_error.append(0.5*(x_stability_boundary_max-x_stability_boundary_min))
                    y1_error.append(0)
            if color_list==-1:
                plot1.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',linestyle='none',ms=marker_size,label='n='+str(n))
            else:
                plot1.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',color=color_list[i],linestyle='none',ms=marker_size,label='n='+str(n))

            x_list.append(x1)
            y_list.append(y1)
            x_list_error.append(x1_error)
            y_list_error.append(y1_error)

        plot1.plot(self.x,self.q*q_scale+q_shift,color='orange',label=r'Modified $q=$'+str(q_scale)+r'$*q_0$')
        plot1.set_xlim(self.x0_min,self.x0_max)
        plot1.set_xlabel(r'$\rho_{tor}$')
        plot1.set_ylabel('Safety factor')
        plot1.legend(loc='upper left')

        return fig

    def Plot_q_scale_rational_surfaces_red_and_green(self,peak_percent,\
            q_scale,q_shift,q_uncertainty,n_list,unstable_list):
        marker_size=5
        #calculate the radial stability boundary
        x_peak,x_stability_boundary_min,x_stability_boundary_max=self.ome_peak_range(peak_percent)

        x_min_index=np.argmin(abs(self.x-x_stability_boundary_min))
        x_max_index=np.argmin(abs(self.x-x_stability_boundary_max))

        q_range=self.q[x_min_index:x_max_index]

        x_list=[]
        y_list=[]
        x_list_error=[]
        y_list_error=[]
        fig, ax = plt.subplots()
        ax.fill_between(self.x, self.q*(1.-q_uncertainty), self.q*(1.+q_uncertainty), color='blue', alpha=.2)
        ax.axvline(x_stability_boundary_min,color='orange',alpha=0.6)
        ax.axvline(x_stability_boundary_max,color='orange',alpha=0.6)
        ax.plot(self.x,self.q,color='blue',label=r'Safety factor $q_0$')

        stable_counter=0
        unstable_counter=0
        for i in range(len(n_list)):
            n=n_list[i]

            x1=[]
            y1=[]
            x1_error=[]
            y1_error=[]

            qmin = np.min(self.q*(1.-q_uncertainty))
            qmax = np.max(self.q*(1.+q_uncertainty))

            m_min = math.ceil(qmin*n)
            m_max = math.floor(qmax*n)
            m_list=np.arange(m_min,m_max+1)

            for m in m_list:
                surface=float(m)/float(n)
                if np.min(q_range)*(1.-q_uncertainty)<surface and surface<np.max(q_range)*(1.+q_uncertainty):
                    x1.append(0.5*(x_stability_boundary_max+x_stability_boundary_min))
                    y1.append(surface)
                    x1_error.append(0.5*(x_stability_boundary_max-x_stability_boundary_min))
                    y1_error.append(0)

            if unstable_list[i]==1:#for unstable case
                if unstable_counter==0:
                    ax.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',linestyle='none',color='red',ms=marker_size,label='Unstable')
                else:
                    ax.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',linestyle='none',color='red',ms=marker_size)
                unstable_counter=unstable_counter+1
            elif unstable_list[i]==0:#for stable case
                if stable_counter==0:
                    ax.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',linestyle='none',color='green',ms=marker_size,label='Stable')
                else:
                    ax.errorbar(x1,y1,xerr=x1_error,yerr=y1_error,marker='s',linestyle='none',color='green',ms=marker_size)
                stable_counter=stable_counter+1
            x_list.append(x1)
            y_list.append(y1)
            x_list_error.append(x1_error)
            y_list_error.append(y1_error)

        ax.plot(self.x,self.q*q_scale+q_shift,color='orange',label=r'Modified $q=$'+str(q_scale)+r'$*q_0$')
        ax.set_xlim(self.x0_min,self.x0_max)
        ax.set_xlabel(r'$\rho_{tor}$')
        ax.set_ylabel('Safety factor')
        plt.legend(loc='upper left')
        plt.show()

    #this attribute only plot the nearest rational surface for the given toroidal mode number
    def Plot_ome_q_surface_demo(self,peak_percent,n_min,n_max,f_min,f_max):
        n_list=np.array(np.arange(n_min,n_max+1),dtype=int)
        x_peak,x_min,x_max=self.ome_peak_range(peak_percent)
        x_peak_plot,x_min_plot,x_max_plot=self.ome_peak_range(peak_percent*10.)
        fig, host = plt.subplots()
        fig.subplots_adjust(right=0.75)

        #par1 = host.twinx()
        #p1, = par1.plot(self.x, self.q, "b-", label='Safety Factor')

        x_fill=[x_min,x_min,x_max,x_max]
        y_fill=[f_min,f_max,f_max,f_min]
        #matplotlib.patches.Rectangle((0,total_trans-trans_error),10,2.*trans_error,alpha=0.4)
        p1, = host.fill(x_fill,y_fill,color='blue',alpha=0.3,label='Unstable area')

        for n in n_list:
            Unstabel_surface_counter=0
            try:
                x0,m=self.Rational_surface_peak_surface(n)
            except:
                continue

            if x_min_plot<x0 and x0<x_max_plot:
                if x_min<x0 and x0<x_max:
                    host.axvline(x0,color='orange',alpha=1)
                    host.scatter([x0],[float(n)*(self.ome+self.Doppler)[np.argmin(abs(self.x-x0))]],s=30,color='blue')
                    Unstabel_surface_counter=Unstabel_surface_counter+1
                else:
                    host.axvline(x0,color='orange',alpha=0.3)
                    host.scatter([x0],[float(n)*(self.ome+self.Doppler)[np.argmin(abs(self.x-x0))]],s=30,color='blue')

            if Unstabel_surface_counter>0:
                p2, = host.plot(self.x,(self.ome+self.Doppler)*float(n), "r-", label=r'Unstable $\omega_{*e}$')
            else:
                p2, = host.plot(self.x,(self.ome+self.Doppler)*float(n), "k-", label=r'Stable $\omega_{*e}$')

        host.set_xlim(np.min(self.x),np.max(self.x))
        host.set_ylim(0, np.max(f_max)*1.2)
        #par1.set_ylim(np.min(self.q)*0.8,np.max(self.q)*1.2)

        host.set_xlabel(r"$\rho_{tor}$")
        host.set_ylabel(r"$\omega_{*e}$(kHz)")
        #par1.set_ylabel("Safety factor")

        host.yaxis.label.set_color('black')
        #par1.yaxis.label.set_color('black')

        tkw = dict(size=4, width=1.5)
        host.tick_params(axis='y', colors='black', **tkw)
        #par1.tick_params(axis='y', colors='black', **tkw)
        #par1.tick_params(axis='y', colors=p3.get_color(), **tkw)
        host.tick_params(axis='x', **tkw)

        lines = [p1,p2]

        #host.legend(lines, ['Unstable area',r'Unstable $\omega_{*e}$',r'Stable $\omega_{*e}$'])
        host.legend(lines, [l.get_label() for l in lines])
        plt.show()

    def Plot_ome_q_surface_frequency_list(self,peak_percent,n_min,n_max,\
                Frequency_list,Frequency_error=0.2,save_imag=False,\
                image_name='./plot.jpg'):
        n_list=np.array(np.arange(n_min,n_max+1),dtype=int)
        x_peak,x_min,x_max=self.ome_peak_range(peak_percent)
        x_peak_plot,x_min_plot,x_max_plot=self.ome_peak_range(peak_percent*10.)

        fig, host = plt.subplots()
        fig.subplots_adjust(right=0.75)

        #par1 = host.twinx()
        #p1, = par1.plot(self.x, self.q, "b-", label='Safety Factor')
        for f in Frequency_list:
            f_min=f*(1-Frequency_error)
            f_max=f*(1+Frequency_error)
            x_fill=[x_min,x_min,x_max,x_max]
            y_fill=[f_min,f_max,f_max,f_min]
            #matplotlib.patches.Rectangle((0,total_trans-trans_error),10,2.*trans_error,alpha=0.4)
            p1, = host.fill(x_fill,y_fill,color='blue',alpha=0.3,label='Unstable area')

        for n in n_list:
            Unstabel_surface_counter=0
            try:
                x0,m=self.Rational_surface_peak_surface(n)
            except:
                continue
            if x_min_plot<x0 and x0<x_max_plot:
                if x_min<x0 and x0<x_max:
                    host.axvline(x0,color='orange',alpha=1)
                    host.scatter([x0],[float(n)*(self.ome+self.Doppler)[np.argmin(abs(self.x-x0))]],s=30,color='blue')
                    Unstabel_surface_counter=Unstabel_surface_counter+1
                else:
                    host.axvline(x0,color='orange',alpha=0.3)
                    host.scatter([x0],[float(n)*(self.ome+self.Doppler)[np.argmin(abs(self.x-x0))]],s=30,color='blue')

            if Unstabel_surface_counter>0:
                p2, = host.plot(self.x,(self.ome+self.Doppler)*float(n), "r-", label=r'Unstable $\omega_{*e}$')
            else:
                p2, = host.plot(self.x,(self.ome+self.Doppler)*float(n), "k-", label=r'Stable $\omega_{*e}$')

        f_max=np.max(Frequency_list)
        host.set_xlim(np.min(self.x),np.max(self.x))
        host.set_ylim(0, np.max(f_max)*(1+Frequency_error)*1.2)
        #par1.set_ylim(np.min(self.q)*0.8,np.max(self.q)*1.2)

        host.set_xlabel(r"$\rho_{tor}$")
        host.set_ylabel(r"$\omega_{*e}$(kHz)")
        #par1.set_ylabel("Safety factor")

        host.yaxis.label.set_color('black')
        #par1.yaxis.label.set_color('black')

        tkw = dict(size=4, width=1.5)
        host.tick_params(axis='y', colors='black', **tkw)
        #par1.tick_params(axis='y', colors='black', **tkw)
        #par1.tick_params(axis='y', colors=p3.get_color(), **tkw)
        host.tick_params(axis='x', **tkw)

        lines = [p1,p2]

        #host.legend(lines, ['Unstable area',r'Unstable $\omega_{*e}$',r'Stable $\omega_{*e}$'])
        host.legend(lines, [l.get_label() for l in lines])
        if save_imag==True:
            plt.savefig(image_name)#save the
        else:
            plt.show()

        plt.close(fig)

    def Plot_ome_q_stability_boundary(self,peak_percent):
        fig, host = plt.subplots()
        fig.subplots_adjust(right=0.75)

        x_peak,x_min,x_max=self.ome_peak_range(peak_percent)
        par1 = host.twinx()
        x=self.x
        y1=self.ome
        y3=self.ome+self.Doppler
        print(self.Doppler)
        y2=self.q
        p1, = host.plot(x, y1, "b-", label=r'$\omega_{*e}$')
        p2, = host.plot(x, y3, "g-", label=r'$\omega_{*e}$ in lab frame')
        p3, = par1.plot(x, y2, "r-", label=r'safety factor')
        p4, = par1.plot([x_min]*1000,np.arange(0,100,(100./1000.)),color='orange',alpha=0.6, label='Unstable range')
        par1.axvline(x_max,color='orange',alpha=0.6)

        host.set_xlim(self.x0_min, self.x0_max)
        host.set_ylim(0., np.max(self.ome+self.Doppler)*1.2)
        par1.set_ylim(np.min(self.q)*0.7, np.max(self.q)*1.2)

        host.set_xlabel(r"$\rho_{tor}$")
        host.set_ylabel(r"$\omega_{*e}$(kHz)")
        par1.set_ylabel("Safety factor")

        host.yaxis.label.set_color(p1.get_color())
        par1.yaxis.label.set_color(p3.get_color())

        tkw = dict(size=4, width=1.5)
        host.tick_params(axis='y', colors=p1.get_color(), **tkw)
        par1.tick_params(axis='y', colors=p3.get_color(), **tkw)
        #par1.tick_params(axis='y', colors=p3.get_color(), **tkw)
        host.tick_params(axis='x', **tkw)

        lines = [p1, p2, p3, p4]

        host.legend(lines, [l.get_label() for l in lines], loc='upper left')
        plt.show()

    def Plot_highlight_top_percent_ome(self,peak_percent,n_min,n_max,f_min,f_max):
        n_list=np.array(np.arange(n_min,n_max+1),dtype=int)
        x_peak,x_min,x_max=self.ome_peak_range(peak_percent)
        plt.clf()
        fig, host = plt.subplots()
        fig.subplots_adjust(right=0.75)
        fig.set_size_inches(8, 6)

        x_min=x_min
        x_max=x_max
        y_min=0
        y_max=10000
        x_fill=[x_min,x_min,x_max,x_max]
        y_fill=[y_min,y_max,y_max,y_min]

        p1,=host.plot(self.x,(self.ome+self.Doppler)*float(1), "b-",alpha=0.5, label=r'$\omega_{*e}$')
        p2,=host.fill(x_fill,y_fill,color='purple',alpha=0.1,label=r'$\mu/x_{*} \leq 0.3$')

        host.set_xlim(np.min(self.x),np.max(self.x))
        host.set_ylim(0, np.max(self.ome+self.Doppler)*1.2)

        host.set_xlabel(r"$\rho_{tor}$")
        host.set_ylabel(r"$\omega_{*e}$(kHz)")

        host.yaxis.label.set_color('black')

        tkw = dict(size=4, width=1.5)
        host.tick_params(axis='y', colors='black', **tkw)
        host.tick_params(axis='x', **tkw)

        lines = [ p1,p2]

        host.legend(lines, [l.get_label() for l in lines])
        plt.savefig('./define_mu.png')
        plt.show()

reason: import numpy,from scipy


text:
import numpy as np
import ot
from scipy import stats
from scipy.sparse import random
from bregman import sinkhorn_scaling

def random_gamma_init(p,q, **kwargs):
    """ Returns random coupling matrix with marginal p,q
    """
    rvs=stats.beta(1e-1,1e-1).rvs
    S=random(len(p), len(q), density=1, data_rvs=rvs)
    return sinkhorn_scaling(p,q,S.A, **kwargs)

def init_matrix_np(X1, X2, v1, v2):
    """Return loss matrices and tensors for COOT fast computation
    Returns the value of |X1-X2|^{2} \otimes T as done in [1] based on [2] for the Gromov-Wasserstein distance.
    Where :
        - X1 : The source dataset of shape (n,d)
        - X2 : The target dataset of shape (n',d')
        - v1 ,v2 : weights (histograms) on the columns of resp. X1 and X2
        - T : Coupling matrix of shape (n,n')
    Parameters
    ----------
    X1 : numpy array, shape (n, d)
         Source dataset
    X2 : numpy array, shape (n', d')
         Target dataset
    v1 : numpy array, shape (d,)
        Weight (histogram) on the features of X1.
    v2 : numpy array, shape (d',)
        Weight (histogram) on the features of X2.

    Returns
    -------
    constC : ndarray, shape (n, n')
        Constant C matrix (see paragraph 1.2 of supplementary material in [1])
    hC1 : ndarray, shape (n, d)
        h1(X1) matrix (see paragraph 1.2 of supplementary material in [1])
    hC2 : ndarray, shape (n', d')
        h2(X2) matrix (see paragraph 1.2 of supplementary material in [1])
    References
    ----------
    .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas
          "CO-Optimal Transport"
    .. [2] Peyré, Gabriel, Marco Cuturi, and Justin Solomon,
        "Gromov-Wasserstein averaging of kernel and distance matrices."
        International Conference on Machine Learning (ICML). 2016.
    """
    def f1(a):
        return (a ** 2)

    def f2(b):
        return (b ** 2)

    def h1(a):
        return a

    def h2(b):
        return 2 * b

    constC1 = np.dot(np.dot(f1(X1), v1.reshape(-1, 1)),
                     np.ones(f1(X2).shape[0]).reshape(1, -1))
    constC2 = np.dot(np.ones(f1(X1).shape[0]).reshape(-1, 1),
                     np.dot(v2.reshape(1, -1), f2(X2).T))

    constC = constC1 + constC2
    hX1 = h1(X1)
    hX2 = h2(X2)

    return constC, hX1, hX2

def cot_numpy(X1, X2, w1 = None, w2 = None, v1 = None, v2 = None,
              niter=10, algo='emd', reg=0,algo2='emd',
              reg2=0, verbose=True, log=False, random_init=False, C_lin=None):

    """ Returns COOT between two datasets X1,X2 (see [1])

    The function solves the following optimization problem:
    .. math::
        COOT = \min_{Ts,Tv} \sum_{i,j,k,l} |X1_{i,k}-X2_{j,l}|^{2}*Ts_{i,j}*Tv_{k,l}

    Where :
    - X1 : The source dataset
    - X2 : The target dataset
    - w1,w2  : weights (histograms) on the samples (rows) of resp. X1 and X2
    - v1,v2  : weights (histograms) on the features (columns) of resp. X1 and X2

    Parameters
    ----------
    X1 : numpy array, shape (n, d)
         Source dataset
    X2 : numpy array, shape (n', d')
         Target dataset
    w1 : numpy array, shape (n,)
        Weight (histogram) on the samples of X1. If None uniform distribution is considered.
    w2 : numpy array, shape (n',)
        Weight (histogram) on the samples of X2. If None uniform distribution is considered.
    v1 : numpy array, shape (d,)
        Weight (histogram) on the features of X1. If None uniform distribution is considered.
    v2 : numpy array, shape (d',)
        Weight (histogram) on the features of X2. If None uniform distribution is considered.
    niter : integer
            Number max of iterations of the BCD for solving COOT.
    algo : string
            Choice of algorithm for solving OT problems on samples each iteration. Choice ['emd','sinkhorn'].
            If 'emd' returns sparse solution
            If 'sinkhorn' returns regularized solution
    algo2 : string
            Choice of algorithm for solving OT problems on features each iteration. Choice ['emd','sinkhorn'].
            If 'emd' returns sparse solution
            If 'sinkhorn' returns regularized solution
    reg : float
            Regularization parameter for samples coupling matrix. Ignored if algo='emd'
    reg2 : float
            Regularization parameter for features coupling matrix. Ignored if algo='emd'
    eps : float
        Threshold for the convergence
    random_init : bool
            Wether to use random initialization for the coupling matrices. If false identity couplings are considered.
    log : bool, optional
         record log if True
    C_lin : numpy array, shape (n, n')
            Prior on the sample correspondences. Added to the cost for the samples transport

    Returns
    -------
    Ts : numpy array, shape (n,n')
           Optimal Transport coupling between the samples
    Tv : numpy array, shape (d,d')
           Optimal Transport coupling between the features
    cost : float
            Optimization value after convergence
    log : dict
        convergence information and coupling marices
    References
    ----------
    .. [1] Redko Ievgen, Vayer Titouan, Flamary R{\'e}mi and Courty Nicolas
          "CO-Optimal Transport"
    Example
    ----------
    import numpy as np
    from cot import cot_numpy

    n_samples=300
    Xs=np.random.rand(n_samples,2)
    Xt=np.random.rand(n_samples,1)
    cot_numpy(Xs,Xt)
    """
    if v1 is None:
       v1 = np.ones(X1.shape[1]) / X1.shape[1]  # is (d,)
    if v2 is None:
       v2 = np.ones(X2.shape[1]) / X2.shape[1]  # is (d',)
    if w1 is None:
       w1 = np.ones(X1.shape[0]) / X1.shape[0]  # is (n',)
    if w2 is None:
       w2 = np.ones(X2.shape[0]) / X2.shape[0]  # is (n,)

    if not random_init:
        Ts = np.ones((X1.shape[0], X2.shape[0])) / (X1.shape[0] * X2.shape[0])  # is (n,n')
        Tv = np.ones((X1.shape[1], X2.shape[1])) / (X1.shape[1] * X2.shape[1])  # is (d,d')
    else:
       Ts=random_gamma_init(w1,w2)
       Tv=random_gamma_init(v1,v2)

    constC_s, hC1_s, hC2_s = init_matrix_np(X1, X2, v1, v2)

    constC_v, hC1_v, hC2_v = init_matrix_np(X1.T, X2.T, w1, w2)
    cost = np.inf

    log_out ={}
    log_out['cost'] = []

    for i in range(niter):
        Tsold = Ts
        Tvold = Tv
        costold = cost

        M = constC_s - np.dot(hC1_s, Tv).dot(hC2_s.T)
        if C_lin is not None:
            M=M+C_lin
        if algo == 'emd':
            Ts = ot.emd(w1, w2, M, numItermax=1e7)
        elif algo == 'sinkhorn':
            Ts = ot.sinkhorn(w1, w2, M, reg)

        M = constC_v - np.dot(hC1_v, Ts).dot(hC2_v.T)

        if algo2 == 'emd':
            Tv = ot.emd(v1, v2, M, numItermax=1e7)
        elif algo2 == 'sinkhorn':
            Tv = ot.sinkhorn(v1,v2, M, reg2)

        delta = np.linalg.norm(Ts - Tsold) + np.linalg.norm(Tv - Tvold)
        cost = np.sum(M * Tv)

        if log:
            log_out['cost'].append(cost)

        if verbose:
            print('Delta: {0}  Loss: {1}'.format(delta, cost))

        if delta < 1e-16 or np.abs(costold - cost) < 1e-7:
            if verbose:
                print('converged at iter ', i)
            break
    if log:
        return Ts, Tv, cost, log_out
    else:
        return Ts, Tv, cost

max_stars_repo_name: ievred/COOT

max_issues_repo_name: ievred/COOT

max_forks_repo_name: ievred/COOT