<gh_stars>1-10
#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import pickle

from astropy.io import fits
from astropy.io import ascii
from argparse import ArgumentParser
from os import listdir
from os.path import isfile, join
from sklearn.decomposition import PCA
from scipy.stats import norm
from scipy.stats import lognorm

plt.rcParams["figure.figsize"] = (8, 6)
plt.rcParams["font.size"] = 14

''' Parsing arguments '''
parser = ArgumentParser()
parser.add_argument("-d", "--display", help="Plot spectra", dest='display', action='store_true')
parser.set_defaults(display=False)
parser.add_argument("-m", "--model", help="Model of PCA to generate the QSO", type=str, metavar="PATH_TO_MODEL", default="")
parser.add_argument("-o", "--output_file", help="Output file for the spectrum - Will overwrite existing file", type=str, metavar="OUTPUT_FILE", default="../output/spectrum.txt")
parser.add_argument("-p", "--set_p_values", help="Set the p_value at which to draw the coefficients of the components from the cumulative pdf (cdf) - no need to specify all the p_value of all the components", type=float, metavar="p_value_1 p_value_2 ... p_value_n", nargs='+')
parser.add_argument("-s", "--seed", help="Seed for random generator", type=int, metavar="SEED", default=0)
parser.add_argument("-z", "--redshift", help="Redshift of the generated spectrum", type=float, metavar="REDSHIFT")
arg = parser.parse_args()
model_file = arg.model
redshift = arg.redshift
output_file = arg.output_file
seed = arg.seed
display = arg.display
input_p_values = arg.set_p_values
np.random.seed(seed)
print(input_p_values)

''' Loading model '''
with open(model_file,"rb") as model:
	pca_redshifts = np.array(pickle.load(model,encoding='latin1'))
	pcas = pickle.load(model,encoding='latin1')
	lambda_centers = pickle.load(model,encoding='latin1')
	mu = pickle.load(model,encoding='latin1')
	sigma = pickle.load(model,encoding='latin1')
	lognormal = pickle.load(model,encoding='latin1')

''' Finding model for given redshift '''
delta_z = abs(pca_redshifts - redshift)
ind = np.where(delta_z == delta_z.min())[0][0]
print('Delta_z between spectrum and model:', delta_z.min())

''' Computing coeffs '''
bins_norm = np.linspace(-100,100,20000)
bins_lognorm = np.logspace(-1,7,20000)

p_values = np.ones(len(mu[ind])) * -99
p_values[:len(input_p_values)] = np.array(input_p_values)

coeffs = np.zeros(len(p_values))
for i,ip_value in enumerate(p_values):
	if ip_value==-99: # Draw random number from distrib
		if lognormal[ind][i]==0:
			coeffs[i] = np.random.normal(mu[ind][i],sigma[ind][i])
		else:
			coeffs[i] = np.random.lognormal(mu[ind][i],sigma[ind][i])
	else: # Estimate coefficient from number of sigmas
		if lognormal[ind][i]==0:
			cdf = lambda x: norm.cdf(x,mu[ind][i],sigma[ind][i])
			coeffs[i] = bins_norm[len(cdf(bins_norm)[cdf(bins_norm)<ip_value])]
		else:
			cdf = lambda x: lognorm.cdf(x,mu[ind][i],sigma[ind][i])
			coeffs[i] = bins_lognorm[max(len(cdf(bins_lognorm)[cdf(bins_lognorm)<=ip_value])-1,0)]
			print(bins_lognorm[max(len(cdf(bins_lognorm)[cdf(bins_lognorm)<=ip_value])-1,0)])

print(coeffs)
''' Saving spectra '''
lamb = lambda_centers[ind]*(1+redshift)
flux = pcas[ind].inverse_transform(coeffs)
ascii.write([lamb,flux],output_file,names=['lambda','flux'],overwrite=True)

''' Drawing random parameters '''
if display:
	plt.figure()
	plt.plot(lamb,flux)
	plt.xlabel(r'$\lambda$')
	plt.ylabel(r'$F - \bar{F}$')
	plt.grid(True,ls='--')
	plt.title("Simulated QSO")
	plt.show()


<filename>stonesoup/models/transition/tests/test_cv.py
# coding: utf-8
import datetime

from pytest import approx
import numpy as np
from scipy.stats import multivariate_normal

from ..linear import ConstantVelocity
from ....types.state import State

def test_cvmodel():
    """ ConstanVelocity Transition Model test """

    # State related variables
    state = State(np.array([[3.0], [1.0]]))
    old_timestamp = datetime.datetime.now()
    timediff = 1  # 1sec
    new_timestamp = old_timestamp + datetime.timedelta(seconds=timediff)
    time_interval = new_timestamp - old_timestamp

    # Model-related components
    noise_diff_coeff = 0.001  # m/s^2
    F = np.array([[1, timediff], [0, 1]])
    Q = np.array([[timediff**3 / 3,
                   timediff**2 / 2],
                  [timediff**2 / 2,
                   timediff]]) * noise_diff_coeff

    # Create and a Constant Velocity model object
    cv = ConstantVelocity(noise_diff_coeff=noise_diff_coeff)

    # Ensure ```cv.transfer_function(time_interval)``` returns F
    assert np.array_equal(F, cv.matrix(
        timestamp=new_timestamp, time_interval=time_interval))

    # Ensure ```cv.covar(time_interval)``` returns Q
    assert np.array_equal(Q, cv.covar(
        timestamp=new_timestamp, time_interval=time_interval))

    # Propagate a state vector through the model
    # (without noise)
    new_state_vec_wo_noise = cv.function(
        state,
        timestamp=new_timestamp,
        time_interval=time_interval)
    assert np.array_equal(new_state_vec_wo_noise, F@state.state_vector)

    # Evaluate the likelihood of the predicted state, given the prior
    # (without noise)
    prob = cv.pdf(State(new_state_vec_wo_noise),
                  state,
                  timestamp=new_timestamp,
                  time_interval=time_interval)
    assert approx(prob) == multivariate_normal.pdf(
        new_state_vec_wo_noise.T,
        mean=np.array(F@state.state_vector).ravel(),
        cov=Q)

    # Propagate a state vector throught the model
    # (with internal noise)
    new_state_vec_w_inoise = cv.function(
        state,
        noise=True,
        timestamp=new_timestamp,
        time_interval=time_interval)
    assert not np.array_equal(new_state_vec_w_inoise, F@state.state_vector)

    # Evaluate the likelihood of the predicted state, given the prior
    # (with noise)
    prob = cv.pdf(State(new_state_vec_w_inoise),
                  state,
                  timestamp=new_timestamp,
                  time_interval=time_interval)
    assert approx(prob) == multivariate_normal.pdf(
        new_state_vec_w_inoise.T,
        mean=np.array(F@state.state_vector).ravel(),
        cov=Q)

    # Propagate a state vector throught the model
    # (with external noise)
    noise = cv.rvs(timestamp=new_timestamp, time_interval=time_interval)
    new_state_vec_w_enoise = cv.function(
        state,
        timestamp=new_timestamp,
        time_interval=time_interval,
        noise=noise)
    assert np.array_equal(new_state_vec_w_enoise, F@state.state_vector+noise)

    # Evaluate the likelihood of the predicted state, given the prior
    # (with noise)
    prob = cv.pdf(State(new_state_vec_w_enoise), state,
                  timestamp=new_timestamp, time_interval=time_interval)
    assert approx(prob) == multivariate_normal.pdf(
        new_state_vec_w_enoise.T,
        mean=np.array(F@state.state_vector).ravel(),
        cov=Q)


<filename>bnpy/viz/PlotTrace.py<gh_stars>0
'''
PlotTrace.py

Executable for plotting trace stats of learning algorithm progress, including
* objective function (ELBO) vs laps thru data
* number of active components vs laps thru data
* hamming distance vs laps thru data

Usage (command-line)
-------
python -m bnpy.viz.PlotTrace dataName jobpattern [kwargs]
'''
import numpy as np
import argparse
import glob
import os
import scipy.io

from PlotUtil import pylab
from bnpy.ioutil import BNPYArgParser
from JobFilter import filterJobs

taskidsHelpMsg = "ids of trials/runs to plot from given job." + \
                 " Example: '4' or '1,2,3' or '2-6'."

Colors = [(0, 0, 0),  # black
          (0, 0, 1),  # blue
          (1, 0, 0),  # red
          (0, 1, 0.25),  # green (darker)
          (1, 0, 1),  # magenta
          (0, 1, 1),  # cyan
          (1, 0.6, 0),  # orange
          ]

LabelMap = dict(laps='num pass thru data',
                iters='num alg steps',
                times='elapsed time (sec)',
                K='num topics K',
                evidence='train objective',
                )
LabelMap['laps-saved-params'] = 'num pass thru data'
LabelMap['hamming-distance'] = 'Hamming dist.'
LabelMap['Keff'] = 'num topics K'

def plotJobsThatMatchKeywords(jpathPattern='/tmp/', **kwargs):
    ''' Create line plots for jobs matching pattern and provided kwargs
    '''
    if not jpathPattern.startswith(os.path.sep):
        jpathPattern = os.path.join(os.environ['BNPYOUTDIR'], jpathPattern)
    jpaths, legNames = filterJobs(jpathPattern, **kwargs)
    plotJobs(jpaths, legNames, **kwargs)

def plotJobs(jpaths, legNames, styles=None, density=2,
             xvar='laps', yvar='evidence', loc='upper right',
             xmin=None, xmax=None,
             taskids=None, savefilename=None, tickfontsize=None,
             bbox_to_anchor=None, **kwargs):
    ''' Create line plots for provided jobs.
    '''
    nLines = len(jpaths)
    if nLines == 0:
        raise ValueError('Empty job list. Nothing to plot.')

    nLeg = len(legNames)

    for lineID in xrange(nLines):
        if styles is None:
            curStyle = dict(colorID=lineID)
        else:
            curStyle = styles[lineID]

        plot_all_tasks_for_job(jpaths[lineID], legNames[lineID],
                               xvar=xvar, yvar=yvar,
                               taskids=taskids, density=density, **curStyle)

    # Y-axis limit determination
    # If we have "enough" data about the run beyond two full passes of dataset,
    # we zoom in on the region of data beyond lap 2
    if xvar == 'laps' and yvar == 'evidence':
        xmax = 0
        ymin = np.inf
        ymin2 = np.inf
        ymax = -np.inf
        allRunsHaveXBeyond1 = True
        for line in pylab.gca().get_lines():
            xd = line.get_xdata()
            yd = line.get_ydata()
            if xd.size < 3:
                allRunsHaveXBeyond1 = False
                continue
            posLap1 = np.searchsorted(xd, 1.0)
            posLap2 = np.searchsorted(xd, 2.0)
            if posLap1 < xd.size:
                ymin = np.minimum(ymin, yd[posLap1])
                ymax = np.maximum(ymax, yd[posLap1:].max())
            if posLap2 < xd.size:
                ymin2 = np.minimum(ymin2, yd[posLap2])
            xmax = np.maximum(xmax, xd.max())
            if xd.max() <= 1:
                allRunsHaveXBeyond1 = False
        if allRunsHaveXBeyond1 and xmax > 1.5:
            # If all relevant curves extend beyond x=1, only show that part
            xmin = 1.0 - 1e-5
        else:
            xmin = 0
        if allRunsHaveXBeyond1 and ymin2 < ymax:
            range1 = ymax - ymin
            range2 = ymax - ymin2
            if 10 * range2 < range1:
                # Y values jump from lap1 to lap2 is enormous,
                # so let's just show y values from lap2 onward...
                ymin = ymin2
        if (not np.allclose(ymax, ymin)) and allRunsHaveXBeyond1:
            pylab.ylim([ymin, ymax + 0.1 * (ymax - ymin)])
        pylab.xlim([xmin, xmax + .05 * (xmax - xmin)])

    if loc is not None and len(jpaths) > 1:
        pylab.legend(loc=loc, bbox_to_anchor=bbox_to_anchor)
    if tickfontsize is not None:
        pylab.tick_params(axis='both', which='major', labelsize=tickfontsize)

    if savefilename is not None:
        try:
            pylab.show(block=False)
        except TypeError:
            pass  # when using IPython notebook
        pylab.savefig(savefilename, bbox_inches='tight', pad_inches=0)
    else:
        try:
            pylab.show(block=True)
        except TypeError:
            pass  # when using IPython notebook

def plot_all_tasks_for_job(jobpath, label, taskids=None,
                           lineType='.-',
                           color=None,
                           colorID=0,
                           density=2,
                           yvar='evidence',
                           markersize=10,
                           linewidth=2,
                           xvar='laps', **kwargs):
    ''' Create line plot in current figure for each task/run of jobpath
    '''
    if not os.path.exists(jobpath):
        if not jobpath.startswith(os.path.sep):
            jobpath_tmp = os.path.join(os.environ['BNPYOUTDIR'], jobpath)
            if not os.path.exists(jobpath_tmp):
                raise ValueError("PATH NOT FOUND: %s" % (jobpath))
            jobpath = jobpath_tmp
    if color is None:
        color = Colors[colorID % len(Colors)]
    taskids = BNPYArgParser.parse_task_ids(jobpath, taskids)

    if yvar == 'hamming-distance':
        if xvar == 'laps':
            xvar = 'laps-saved-params'

    for tt, taskid in enumerate(taskids):
        try:
            var_ext = ''
            ytxtfile = os.path.join(jobpath, taskid, yvar + '.txt')
            if not os.path.isfile(ytxtfile):
                var_ext = '-saved-params'
                ytxtfile = os.path.join(
                    jobpath, taskid, yvar + var_ext + '.txt')
            ys = np.loadtxt(ytxtfile)

            xtxtfile = os.path.join(jobpath, taskid, xvar + var_ext + '.txt')
            if not os.path.isfile(xtxtfile):
                var_ext = '-saved-params'
                xtxtfile = os.path.join(
                    jobpath, taskid, xvar + var_ext + '.txt')
            xs = np.loadtxt(xtxtfile)

        except IOError as e:
            try:
                xs, ys = loadXYFromTopicModelFiles(jobpath, taskid)
            except ValueError:
                try:
                    xs, ys = loadXYFromTopicModelSummaryFiles(jobpath, taskid)
                except ValueError:
                    raise e

        if yvar == 'hamming-distance' or yvar == 'Keff':
            if xvar == 'laps-saved-params':
                # fix off-by-one error, if we save an extra dist on final lap
                if xs.size == ys.size - 1:
                    ys = ys[:-1]
                elif ys.size == xs.size - 1:
                    xs = xs[:-1]  # fix off-by-one error, if we quit early
            elif xs.size != ys.size:
                # Try to subsample both time series at laps where they
                # intersect
                laps_x = np.loadtxt(os.path.join(jobpath, taskid, 'laps.txt'))
                laps_y = np.loadtxt(os.path.join(jobpath, taskid,
                                                 'laps-saved-params.txt'))
                assert xs.size == laps_x.size
                if ys.size == laps_y.size - 1:
                    laps_y = laps_y[:-1]
                xs = xs[np.in1d(laps_x, laps_y)]
                ys = ys[np.in1d(laps_y, laps_x)]

        if xs.size != ys.size:
            raise ValueError('Dimension mismatch. len(xs)=%d, len(ys)=%d'
                             % (xs.size, ys.size))

        if xs.size < 3:
            continue

        # Cleanup laps data. Verify that it is sorted, with no collisions.
        if xvar == 'laps':
            diff = xs[1:] - xs[:-1]
            goodIDs = np.flatnonzero(diff >= 0)
            if len(goodIDs) < xs.size - 1:
                print 'WARNING: looks like multiple runs writing to this file!'
                print jobpath
                print 'Task: ', taskid
                print len(goodIDs), xs.size - 1
                xs = np.hstack([xs[goodIDs], xs[-1]])
                ys = np.hstack([ys[goodIDs], ys[-1]])

        if xvar == 'laps' and yvar == 'evidence':
            mask = xs >= 1.0
            xs = xs[mask]
            ys = ys[mask]

        # Force plot density (data points per lap) to desired specification
        # This avoids making plots that have huge file sizes,
        # due to too much content in the given display space
        if xvar == 'laps' and xs.size > 20 and np.sum(xs > 5) > 10:
            if (xs[-1] - xs[9]) != 0:
                curDensity = (xs.size - 10) / (xs[-1] - xs[9])
            else:
                curDensity = density
            while curDensity > density and xs.size > 11:
                # Thin xs and ys data by a factor of 2
                # while preserving the first 10 data points
                xs = np.hstack([xs[:10], xs[10::2]])
                ys = np.hstack([ys[:10], ys[10::2]])
                curDensity = (xs.size - 10) / (xs[-1] - xs[9])

        plotargs = dict(markersize=markersize, linewidth=linewidth, label=None,
                        color=color, markeredgecolor=color)
        plotargs.update(kwargs)
        if tt == 0:
            plotargs['label'] = label

        pylab.plot(xs, ys, lineType, **plotargs)

    pylab.xlabel(LabelMap[xvar])
    pylab.ylabel(LabelMap[yvar])

def loadXYFromTopicModelSummaryFiles(jobpath, taskid, xvar='laps', yvar='K'):
    ''' Load x and y variables for line plots from TopicModel files
    '''
    ypath = os.path.join(jobpath, taskid, 'predlik-' + yvar + '.txt')
    if not os.path.exists(ypath):
        raise ValueError('No TopicModel summary text files found')
    lappath = os.path.join(jobpath, taskid, 'predlik-lapTrain.txt')
    xs = np.loadtxt(lappath)
    ys = np.loadtxt(ypath)
    return xs, ys

def loadXYFromTopicModelFiles(jobpath, taskid, xvar='laps', yvar='K'):
    ''' Load x and y variables for line plots from TopicModel files
    '''
    tmpathList = glob.glob(os.path.join(jobpath, taskid, 'Lap*TopicModel.mat'))
    if len(tmpathList) < 1:
        raise ValueError('No TopicModel.mat files found')
    tmpathList.sort()  # ascending, from lap 0 to lap 1 to lap 100 to ...
    basenames = [x.split(os.path.sep)[-1] for x in tmpathList]
    laps = np.asarray([float(x[3:11]) for x in basenames])
    Ks = np.zeros_like(laps)
    for tt, tmpath in enumerate(tmpathList):
        if yvar == 'K':
            Q = scipy.io.loadmat(tmpath, variable_names=['K', 'probs'])
            try:
                Ks[tt] = Q['K']
            except KeyError:
                Ks[tt] = Q['probs'].size
        else:
            raise ValueError('Unknown yvar type for topic model: ' + yvar)
    return laps, Ks

def parse_args(xvar='laps', yvar='evidence'):
    ''' Returns Namespace of parsed arguments retrieved from command line
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument('dataName', type=str, default='AsteriskK8')
    parser.add_argument('jpath', type=str, default='demo*')

    parser.add_argument('--xvar', type=str, default=xvar,
                        choices=LabelMap.keys(),
                        help="name of x axis variable to plot.")

    parser.add_argument('--yvar', type=str, default=yvar,
                        choices=LabelMap.keys(),
                        help="name of y axis variable to plot.")

    helpMsg = "ids of trials/runs to plot from given job." + \
              " Example: '4' or '1,2,3' or '2-6'."
    parser.add_argument(
        '--taskids', type=str, default=None, help=helpMsg)
    parser.add_argument(
        '--savefilename', type=str, default=None,
        help="location where to save figure (absolute path directory)")

    args, unkList = parser.parse_known_args()

    argDict = BNPYArgParser.arglist_to_kwargs(unkList, doConvertFromStr=False)
    argDict.update(args.__dict__)
    argDict['jpathPattern'] = os.path.join(os.environ['BNPYOUTDIR'],
                                           args.dataName,
                                           args.jpath)
    del argDict['dataName']
    del argDict['jpath']
    return argDict

plotJobsThatMatch = plotJobsThatMatchKeywords

if __name__ == "__main__":
    argDict = parse_args('laps', 'evidence')
    plotJobsThatMatchKeywords(**argDict)


"""
This module creat data generators for the following datasets:
IIIT5K, IC13, SVHN
"""
import scipy.io
import mat73

len_IIIT5K = 500 # max 3000
len_IC13 = 500   # max 848
len_SVHN = 500   # max 13068

path_IIIT5K = 'data/IIIT5K/'
file_IIIT5K = path_IIIT5K + 'testdata.mat'

path_IC13 = 'data/Challenge2_Training_Task3_Images_GT/'
file_IC13 = path_IC13 + 'gt.txt'

path_SVHN = 'data/test/'
file_SVHN = path_SVHN + 'digitStruct.mat'

def get_label(content): # for SVHN

    label = ''

    if type(content) == list:

        for v in content:
            v = int(v.item())
            if v == 10: v = 0
            label += str(v)
    else:
        v = int(content.item())
        if v == 10: v = 0
        label += str(v)

    return label

def generator_IIIT5K(desc, len_):

    for i, row in enumerate(desc['testdata'][0]):
        if i >= len_:
            return

        yield path_IIIT5K + row[0][0], row[1][0]

def generator_IC13(desc, len_):

    for i, row in enumerate(desc):
        if i >= len_:
            return

        split = row.split(', ')

        yield path_IC13 + split[0], split[1][1:-2]

def generator_SVHN(desc, len_):

    for i, img_name in enumerate(desc['digitStruct']['name']):
        if i >= len_:
            return

        data = desc['digitStruct']['bbox'][i]['label']
        label = get_label(data)

        yield path_SVHN + img_name, label

def postprocessing_IIIT5K(predicted):

    punc = '''!()-[]{};:'"\, <>./?@#$%^&*_~'''

    for i, content in enumerate(predicted):

        result = ""

        for char in content:
            if char not in punc:
                result += char

        predicted[i] = result

    return predicted

def postprocessing_IC13(predicted):
    return predicted

def postprocessing_SVHN(predicted):
    return predicted

gt_IIIT5K = scipy.io.loadmat(file_IIIT5K)
gt_IC13 = open(file_IC13, 'r')
gt_SVHN = mat73.loadmat(file_SVHN, 'r')

gen_IIIT5K = generator_IIIT5K(gt_IIIT5K, len_IIIT5K)
gen_IC13 = generator_IC13(gt_IC13, len_IC13)
gen_SVHN = generator_SVHN(gt_SVHN, len_SVHN)

datasets = {
    "IIIT5K":{"gen":gen_IIIT5K,
              "len":len_IIIT5K,
              "postprocessing": postprocessing_IIIT5K},
    "IC13":  {"gen":gen_IC13,
              "len":len_IC13,
              "postprocessing": postprocessing_IC13},
    "SVHN":  {"gen":gen_SVHN,
              "len":len_SVHN,
              "postprocessing": postprocessing_SVHN}
}


<reponame>photonics-project/notebooks<filename>src/optics.py
# %%
# %matplotlib widget
import ipyvuetify as v
import ipywidgets as widgets
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import scipy.special as sc

from controls import (
    DetectorFormatControlPanel,
    OpticsControlPanel,
    WavelengthsControlPanel
)

parameters = {
    'diameter': 1,
    'focal_length': 2,
    'fnumber': 2,
    'Hdim': 1280,
    'Vdim': 720,
    'pitch': 20,
    'wavelengths': [3.0, 5.0],
}

output = widgets.Output()

class Figure():
    def __init__(self):
        (fig, [ax1, ax2]) = plt.subplots(ncols=2, constrained_layout=True)

        fig.canvas.header_visible = False
        fig.canvas.toolbar_visible = True
        fig.canvas.toolbar_position = 'right'

        self.fig = fig
        self.ax1 = ax1
        self.ax2 = ax2

        self.update()

    @property
    def canvas(self):
        return self.fig.canvas

    def update(self):
        Npixels = 9
        fnumber = parameters['fnumber']
        pitch = parameters['pitch']
        xlambda = np.array(parameters['wavelengths'])

        blur_spots = 2*sc.jn_zeros(1, 1)[0]/np.pi*xlambda*fnumber/pitch

        self.ax1.clear()
        self.ax2.clear()

        for ax in [self.ax1, self.ax2]:
            ax.set_aspect('equal')
            ax.set_xticklabels([])
            ax.set_xticks(np.arange(Npixels+1))
            ax.set_yticklabels([])
            ax.set_yticks(np.arange(Npixels+1))
            ax.set_axisbelow(True)
            ax.grid(True)

        self.ax1.set_title(f'Blur spot = {blur_spots[0]*pitch:g} µm', color='red')
        self.ax2.set_title(f'Blur spot = {blur_spots[1]*pitch:g} µm', color='blue')

        circle1 = mpl.patches.Circle((0.5, 0.5), blur_spots[0]/Npixels/2, color='red', alpha=0.5, transform=self.ax1.transAxes)
        circle2 = mpl.patches.Circle((0.5, 0.5), blur_spots[1]/Npixels/2, color='blue', alpha=0.5, transform=self.ax2.transAxes)

        self.ax1.add_patch(circle1)
        self.ax2.add_patch(circle2)

        self.ax1.grid(True)
        self.ax2.grid(True)

        self.fig.canvas.draw()

class Table():
    def __init__(self):
        self.widget = v.Html(tag='div', class_='d-flex flex-row', children=[])

        self.update()

    def update(self):
        fnumber = parameters['fnumber']
        Hdim = parameters['Hdim']
        Vdim = parameters['Vdim']
        pitch = parameters['pitch']
        xlambda = np.array(parameters['wavelengths'])

        blur_spots = 2*sc.jn_zeros(1, 1)[0]/np.pi*xlambda*fnumber/pitch

        table = v.DataTable(
            style_='width: 100%',
            hide_default_footer=True,
            disable_sort=True,
            headers=[
                {'text': 'Parameter', 'value': 'parameter'},
                {'text': 'Value', 'value': 'value'},
                {'text': 'Units', 'value': 'units'},
            ],
            items=[
                {
                    'parameter': 'Chip Width',
                    'value': f'{Hdim*(pitch/1e3):g}',
                    'units': 'mm',
                },
                {
                    'parameter': 'Chip Height',
                    'value': f'{Vdim*(pitch/1e3):g}',
                    'units': 'mm',
                },
                {
                    'parameter': 'Chip Area',
                    'value': f'{Hdim*Vdim*(pitch/1e3)**2:g}',
                    'units': 'mm²',
                },
                {
                    'parameter': 'Blur Spot (Wavelength #1)',
                    'value': f'{blur_spots[0]*pitch:g}',
                    'units': 'µm',
                },
                {
                    'parameter': 'Blur Spot (Wavelength #2)',
                    'value': f'{blur_spots[1]*pitch:g}',
                    'units': 'µm',
                },
            ],
        )
        self.widget.children = [table]

plt.ioff()

figure = Figure()

table = Table()

optics_control_panel = OpticsControlPanel(
    diameter=parameters['diameter'],
    focal_length=parameters['focal_length'],
)
detector_format_control_panel = DetectorFormatControlPanel(
    Hdim=parameters['Hdim'],
    Vdim=parameters['Vdim'],
    pitch=parameters['pitch'],
)
wavelengths_control_panel = WavelengthsControlPanel(xlambda=parameters['wavelengths'])

def update():
    figure.update()
    table.update()

def update_optics():
    parameters['diameter'] = optics_control_panel.diameter
    parameters['focal_length'] = optics_control_panel.focal_length
    parameters['fnumber'] = optics_control_panel.fnumber
    update()

def update_detector_format():
    parameters['Hdim'] = detector_format_control_panel.Hdim
    parameters['Vdim'] = detector_format_control_panel.Vdim
    parameters['pitch'] = float(detector_format_control_panel.pitch)
    update()

def update_wavelengths():
    parameters['wavelengths'] = wavelengths_control_panel.xlambda
    update()

optics_control_panel.on_change(update_optics)
detector_format_control_panel.on_change(update_detector_format)
wavelengths_control_panel.on_change(update_wavelengths)

v.Container(children=[
    v.Row(children=[
        v.Col(cols=12, md=6, children=[
            v.Card(
                class_='mb-4',
                outlined=True,
                children=[
                    v.CardTitle(children=['Optics']),
                    v.CardText(children=[
                        optics_control_panel.widget,
                    ]),
            ]),
            v.Card(
                class_='mb-4',
                outlined=True,
                children=[
                    v.CardTitle(children=['Wavelengths']),
                    v.CardText(children=[
                        wavelengths_control_panel.widget,
                    ]),
            ]),
            v.Card(
                class_='mb-4',
                outlined=True,
                children=[
                    v.CardTitle(children=['Detector']),
                    v.CardText(children=[
                        detector_format_control_panel.widget,
                    ]),
            ]),
            v.Card(
                class_='mb-4',
                outlined=True,
                children=[
                    v.CardTitle(children=['Results']),
                    v.CardText(children=[
                        table.widget,
                    ]),
            ]),
        ]),
        v.Col(cols=12, md=6, children=[
            v.Card(
                class_='mb-4',
                outlined=True,
                children=[
                    v.CardTitle(children=['Figure']),
                    v.CardText(children=[
                        figure.canvas,
                    ]),
            ]),
        ]),
    ]),
    # v.Row(children=[
    #     v.Col(cols=12, md=6, children=[
    #         output
    #     ]),
    # ]),
])

# %%


<filename>2021/day07/part1.py
from statistics import median

# Simple logging
from loguru import logger

# Personal utilities
from bsoyka_aoc_utils import get_data

CRABS = get_data(2021, 7, func=int, split=",")
logger.debug("Loaded crabs data")

# Finding the best alignment position can be done by getting the median
# of the crab positions.

align_position: int = int(median(CRABS))

fuel: int = sum(abs(crab - align_position) for crab in CRABS)

logger.success("Result: {}", fuel)


from network import *
import cv2, os, time, math
import glob
import scipy.io as io
from loss import *
from utils import *

is_training = False  # if test, set this 'False'
use_viewdirection = True  # use view direction
renew_input = True   # optimize input point features.
constant_initial = True  # use constant value for initialization.
use_RGB = True     # use RGB information for initialization.
random_crop = True  # crop image.

d = 32   # how many planes are used, identity with pre-processing.
h = 480  # image height, identity with pre-processing.
w = 640  # image width, identity with pre-processing.
top_left_v = 0  # top left position
top_left_u = 0  # top left position
h_croped = 240  # crop size height
w_croped = 320  # crop size width
forward_time = 4  # optimize input point features after cropping 4 times on one image.
overlap = 32  # size of overlap region of crops.

channels_i = int(8)  # dimension of input point features
channels_o = 3  # output image dimensions
channels_v = 3  # view direction dimensions

gpu_id = 3
num_epoch = 21
decrease_epoch = 7  # epochs, learning_rate_1 decreased.
learning_rate = 0.0001  # learning rate for network parameters optimization
learning_rate_1 = 0.01  # initial learning rate for input point features.

dataset = 'ScanNet'     # datasets
scene = 'scene0010_00'  # scene name
task = '%s_npcr_%s' % (dataset, scene)  # task name, also path of checkpoints file
dir1 = 'data/%s/%s/color/' % (dataset, scene)  # path of color image
dir2 = 'data/%s/%s/pose/' % (dataset, scene)  # path of camera poses.
dir3 = 'pre_processing_results/%s/%s/reproject_results_%s/' % (dataset, scene, d)  # voxelization information path.
dir4 = 'pre_processing_results/%s/%s/weight_%s/' % (dataset, scene, d)  # aggregation information path.
dir5 = 'pre_processing_results/%s/%s/point_clouds_simplified.ply' % (dataset, scene)  # point clouds file path

num_image = len(glob.glob(os.path.join(dir1, '*.jpg')))

image_names_train, index_names_train, camera_names_train, index_names_1_train,\
image_names_test, index_names_test, camera_names_test, index_names_1_test = prepare_data_ScanNet(dir1, dir2, dir3, dir4, num_image)

# load point clouds information
point_clouds, point_clouds_colors = loadfile(dir5)
num_points = point_clouds.shape[1]

# initial descriptor
descriptors = np.random.normal(0, 1, (1, num_points, channels_i))

if os.path.isfile('%s/descriptor.mat' % task):
    content = io.loadmat('%s/descriptor.mat' % task)
    descriptors = content['descriptors']
    print('loaded descriptors.')
else:
    if constant_initial:
        descriptors = np.ones((1, num_points, channels_i), dtype=np.float32) * 0.5

    if use_RGB:
        descriptors[0, :, 0:3] = np.transpose(point_clouds_colors) / 255.0

os.environ["CUDA_VISIBLE_DEVICES"] = "%s" % gpu_id
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
sess = tf.Session()

input1 = tf.placeholder(dtype=tf.float32, shape=[1, d, None, None, channels_i])
input2 = tf.placeholder(dtype=tf.float32, shape=[1, d, None, None, channels_v])
output = tf.placeholder(dtype=tf.float32, shape=[1, None, None, channels_o])

with tf.variable_scope(tf.get_variable_scope()):
    inputs = input1
    total_channels = channels_i

    if use_viewdirection:
        inputs = tf.concat([input1, input2], axis=4)
        total_channels = total_channels + channels_v

    color_layer, alpha, network = neural_render(input=inputs, reuse=False, use_dilation=True)
    loss, p0, p1, p2, p3, p4, p5 = VGG_loss(network, output, reuse=False)

    loss_all = loss

    # calculate gradient for aggregated point features.
    gradient = tf.gradients(loss_all, input1)

var_list_all = [var for var in tf.trainable_variables()]
opt = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss_all, var_list=var_list_all)

saver = tf.train.Saver(var_list=var_list_all, max_to_keep=1000)

sess.run(tf.global_variables_initializer())

ckpt = tf.train.get_checkpoint_state(task)
if ckpt:
    print('load ' + ckpt.model_checkpoint_path)
    saver.restore(sess, ckpt.model_checkpoint_path)

##############################################################################################
if is_training:
    print('begin training!')
    all = np.zeros(20000, dtype=float)
    cnt = 0

    for epoch in range(num_epoch):

        if epoch >= decrease_epoch:
            learning_rate_1 = 0.005

        if epoch >= decrease_epoch*2:
            learning_rate_1 = 0.001

        if os.path.isdir("%s/%04d" % (task, epoch)):
            continue

        for i in np.random.permutation(len(image_names_train)):
        # for i in range(4):
            st = time.time()
            image_descriptor = np.zeros([1, d, h, w, channels_i], dtype=np.float32)
            view_direction = np.zeros([1, d, h, w, channels_v], dtype=np.float32)
            input_gradient_all = np.zeros([1, d, h, w, channels_i], dtype=np.float32)
            count = np.zeros([1, d, h, w, 1], dtype=np.float32)
            camera_name = camera_names_train[i]
            index_name = index_names_train[i]
            image_name = image_names_train[i]
            index_name_1 = index_names_1_train[i]

            if not (os.path.isfile(camera_name) and os.path.isfile(image_name) and os.path.isfile(index_name) and os.path.isfile(index_name_1)):
                print("Missing file!")
                continue

            # we pre-process the voxelization and aggregation, in order to save time.
            npzfile = np.load(index_name)
            u = npzfile['u']  # u position on image plane
            v = npzfile['v']  # v position on image plane
            n = npzfile['d']  # indicates which plane
            select_index = npzfile['select_index']   # select index of all points.
            group_belongs = npzfile['group_belongs']  # points belong to which group/voxel
            index_in_each_group = npzfile['index_in_each_group']  # index in each group/voxel
            distance = npzfile['distance']  # distance to grid center
            each_split_max_num = npzfile['each_split_max_num']  # max num of points in one group/voxel in each plane.

            # load weight
            npzfile_weight = np.load(index_name_1)
            weight = npzfile_weight['weight_average']  # normalized weights for points aggregation.
            distance_to_depth_min = npzfile_weight['distance_to_depth_min']  # distance to minimum depth value in one group/voxel.

            # calculate update weight of each point feature
            descriptor_renew_weight = (1-distance)*(1/(1+distance_to_depth_min))

            extrinsic_matrix = CameraPoseRead(camera_name)  # camera to world
            camera_position = np.transpose(extrinsic_matrix[0:3, 3])

            max_num = np.max(each_split_max_num)  # max number of points in all group/voxel
            group_descriptor = np.zeros([(max(group_belongs+1)), max_num, channels_i], dtype=np.float32)
            group_descriptor[group_belongs, index_in_each_group, :] = descriptors[0, select_index, :] * np.expand_dims(weight, axis=1)

            image_descriptor[0, n, v, u, :] = np.sum(group_descriptor, axis=1)[group_belongs, :]

            view_direction[0, n, v, u, :] = np.transpose(point_clouds[0:3, select_index]) - camera_position
            view_direction[0, n, v, u, :] = view_direction[0, n, v, u, :] / (np.tile(np.linalg.norm(view_direction[0, n, v, u, :], axis=1, keepdims=True), (1, 3)) + 1e-10)

            image_output = np.expand_dims(cv2.resize(cv2.imread(image_name, -1), (w, h)), axis=0) / 255.0

            if random_crop:

                # limitation of memory etc, we crop the image.
                # Also, we hope crops almost cover the whole image to uniformly optimize point features.
                for j in np.random.permutation(forward_time):
                    movement_v = np.random.randint(0, overlap)
                    movement_u = np.random.randint(0, overlap)

                    if j==0:
                        top_left_u = 0 + movement_u
                        top_left_v = 0 + movement_v
                    if j==1:
                        top_left_u = w_croped - movement_u
                        top_left_v = 0 + movement_v
                    if j==2:
                        top_left_u = 0 + movement_u
                        top_left_v = h_croped - movement_v
                    if j==3:
                        top_left_u = w_croped - movement_u
                        top_left_v = h_croped - movement_v

                    [_, current_loss, l1, input_gradient] = sess.run([opt, loss_all, loss, gradient],
                                                                     feed_dict={input1: image_descriptor[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :],
                                                                                input2: view_direction[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :],
                                                                                output: image_output[:, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :]
                                                                                })

                    input_gradient_all[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :] = input_gradient[0] + input_gradient_all[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :]
                    count[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :] = count[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :] + 1
                    # print(np.max(count))

                if renew_input:
                    input_gradient_all = input_gradient_all/(count+1e-10)
                    descriptors[0, select_index, :] = descriptors[0, select_index, :] - learning_rate_1 * np.expand_dims(descriptor_renew_weight, axis=1) * input_gradient_all[0, n, v, u, :]

            else:

                [_, current_loss, l1, input_gradient] = sess.run([opt, loss_all, loss, gradient],
                                                                     feed_dict={input1: image_descriptor,
                                                                                input2: view_direction,
                                                                                output: image_output
                                                                                })

                if renew_input:
                    descriptors[0, select_index, :] = descriptors[0, select_index, :] - learning_rate_1 * np.expand_dims(descriptor_renew_weight, axis=1) * input_gradient[0][0, n, v, u, :]

            all[i] = current_loss * 255.0
            cnt = cnt + 1

            print('%s %s %s %.2f %.2f %s' % (epoch, i, cnt, current_loss, np.mean(all[np.where(all)]), time.time() - st))

        os.makedirs("%s/%04d" % (task, epoch))
        saver.save(sess, "%s/model.ckpt" % (task))
        io.savemat("%s/" % task + 'descriptor.mat', {'descriptors': descriptors})

        if epoch % 5 == 0:
            saver.save(sess, "%s/%04d/model.ckpt" % (task, epoch))
            io.savemat("%s/%04d/" % (task, epoch) + 'descriptor.mat', {'descriptors': descriptors})

        for id in range(len(image_names_test)):

            top_left_v = 120
            top_left_u = 160
            st = time.time()
            image_descriptor = np.zeros([1, d, h, w, channels_i])
            view_direction = np.zeros([1, d, h, w, channels_v])
            camera_name = camera_names_test[id]
            index_name = index_names_test[id]
            index_name_1 = index_names_1_test[id]

            if not (os.path.isfile(index_name) and os.path.isfile(camera_name) and os.path.isfile(index_name_1)):
                print('Missingg file 1!')
                continue

            npzfile = np.load(index_name)
            u = npzfile['u']
            v = npzfile['v']
            n = npzfile['d']
            select_index = npzfile['select_index']
            group_belongs = npzfile['group_belongs']
            index_in_each_group = npzfile['index_in_each_group']
            distance = npzfile['distance']
            each_split_max_num = npzfile['each_split_max_num']

            # load weight
            npzfile_weight = np.load(index_name_1)
            weight = npzfile_weight['weight_average']
            distance_to_depth_min = npzfile_weight['distance_to_depth_min']

            extrinsic_matrix = CameraPoseRead(camera_name)  # camera to world
            camera_position = np.transpose(extrinsic_matrix[0:3, 3])

            max_num = np.max(each_split_max_num)
            group_descriptor = np.zeros([(max(group_belongs + 1)), max_num, channels_i], dtype=np.float32)
            group_descriptor[group_belongs, index_in_each_group, :] = descriptors[0, select_index, :] * np.expand_dims(weight, axis=1)

            image_descriptor[0, n, v, u, :] = np.sum(group_descriptor, axis=1)[group_belongs, :]

            view_direction[0, n, v, u, :] = np.transpose(point_clouds[0:3, select_index]) - camera_position
            view_direction[0, n, v, u, :] = view_direction[0, n, v, u, :] / (np.tile(np.linalg.norm(view_direction[0, n, v, u, :], axis=1, keepdims=True), (1, 3)) + 1e-10)

            [result] = sess.run([network], feed_dict={input1: image_descriptor[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :],
                                                      input2: view_direction[:, :, top_left_v:(top_left_v + h_croped), top_left_u:(top_left_u + w_croped), :]})
            result = np.minimum(np.maximum(result, 0.0), 1.0) * 255.0
            cv2.imwrite('%s/%04d/%06d.png' % (task, epoch, id), np.uint8(result[0, :, :, :]))

            print(time.time() - st)
else:

    output_path = "%s/Test_Result/" % (task)
    if not os.path.isdir(output_path):
        os.makedirs(output_path)

    for id in range(len(camera_names_test)):

        st = time.time()
        image_descriptor = np.zeros([1, d, h, w, channels_i])
        view_direction = np.zeros([1, d, h, w, channels_v])
        camera_name = camera_names_test[id]
        index_name = index_names_test[id]
        index_name_1 = index_names_1_test[id]

        if not (os.path.isfile(index_name) and os.path.isfile(camera_name) and os.path.isfile(index_name_1)):
            print('Missingg file 1!')
            continue

        npzfile = np.load(index_name)
        u = npzfile['u']
        v = npzfile['v']
        n = npzfile['d']
        select_index = npzfile['select_index']
        group_belongs = npzfile['group_belongs']
        index_in_each_group = npzfile['index_in_each_group']
        distance = npzfile['distance']
        each_split_max_num = npzfile['each_split_max_num']

        # load weight
        npzfile_weight = np.load(index_name_1)
        weight = npzfile_weight['weight_average']
        distance_to_depth_min = npzfile_weight['distance_to_depth_min']

        extrinsic_matrix = CameraPoseRead(camera_name)  # camera to world
        camera_position = np.transpose(extrinsic_matrix[0:3, 3])

        max_num = np.max(each_split_max_num)
        group_descriptor = np.zeros([(max(group_belongs + 1)), max_num, channels_i], dtype=np.float32)
        group_descriptor[group_belongs, index_in_each_group, :] = descriptors[0, select_index, :] * np.expand_dims(weight, axis=1)

        image_descriptor[0, n, v, u, :] = np.sum(group_descriptor, axis=1)[group_belongs, :]

        view_direction[0, n, v, u, :] = np.transpose(point_clouds[0:3, select_index]) - camera_position
        view_direction[0, n, v, u, :] = view_direction[0, n, v, u, :] / (
        np.tile(np.linalg.norm(view_direction[0, n, v, u, :], axis=1, keepdims=True), (1, 3)) + 1e-10)

        [result] = sess.run([network], feed_dict={input1: image_descriptor, input2: view_direction})
        result = np.minimum(np.maximum(result, 0.0), 1.0) * 255.0
        cv2.imwrite(output_path + '%06d.png' % id, np.uint8(result[0, :, :, :]))

        print(time.time() - st)

if __name__ == '__main__':
    pass


import numpy as np
from scipy.stats import sem
from uncertainties import ufloat
import uncertainties.unumpy as unp
import matplotlib.pyplot as plt
import scipy.constants as sc
from scipy.optimize import curve_fit
#from pint import UnitRegistry
import operator

#u = UnitRegistry()
#Q_ = u.Quantity

Time_ges = 81591 #Bsp lambda_r = Q_(643.2, 'nanometer')
C_Start = 1444133
C_Stop = 4255

print("Gesamter Messzeitraum: ", Time_ges)
print("Gemessene Startimpulse: ", C_Start)
print("Gemessene Stopimpulse: ", C_Stop, '\n')
Time_Search = 10 *10**(-6)

C_Start_r = ufloat(C_Start, np.sqrt(C_Start))

Rate = C_Start_r/Time_ges
N_Search = Time_Search*Rate

W = N_Search * unp.exp(N_Search)
U_Search = W*C_Start/512

print("Rate pro Sekunde: ", Rate)
print("Wahrscheinlichkeit: ", W)
print("Untergrund nach Poisson: ", U_Search, '\n')

#Datenaufnahme etwas umständlich
Counts_M = np.genfromtxt('Messung.txt', unpack=True)
Channels = np.linspace(0, 511, 512)

Messung_Errors = np.sqrt(Counts_M)
Messung = unp.uarray(Counts_M, Messung_Errors)

#Kalibrierung_Errors = np.sqrt(Counts_K)
#Kalibrierung_Werte = np.array([ufloat(n, Kalibrierung_Errors[i]) for i,n in enumerate(Counts_K) ])
#Kalibrierung_y = np.array([Kalibrierung_Werte[i] for i,n in enumerate(Counts_K) if n!= 0])

Counts_K = np.genfromtxt('Kalibrierung.txt', unpack=True)
Kalibrierung_x = np.array([Channels[i] for i,n in enumerate(Counts_K) if n!= 0])
K_x = Kalibrierung_x
Kalibrierung_y = np.array([0.3,1,2,3,4,5,6,7,8,9,10])
Channels_real = np.array([K_x[0],K_x[2],K_x[3],K_x[4],K_x[6],K_x[7],K_x[8],K_x[9],K_x[10],K_x[12],K_x[13]])

Plateau_x = np.array([-24,-22,-20,-18,-16,-14,-12,-10,-8,-6,-4,-2,-1,0,1,2,4,6,8,10,12,14,16,18,20])

P1 = unp.uarray([154,186,186,173], np.sqrt([154,186,186,173]))
P2 = unp.uarray([176,189,164,172], np.sqrt([176,189,164,172]))
P4 = unp.uarray([183,167,178,194], np.sqrt([183,167,178,194]))
P6 = unp.uarray([178,193,165,180], np.sqrt([178,193,165,180]))
P8 = unp.uarray([186,168,147,168], np.sqrt([186,168,147,168]))

M0  = unp.uarray([167,184,180], np.sqrt([167,184,180]))
M1  = unp.uarray([177,189,178], np.sqrt([177,189,178]))
M2  = unp.uarray([169,180,171,171], np.sqrt([169,180,171,171]))
M4  = unp.uarray([177,168,183,158], np.sqrt([177,168,183,158]))
M6  = unp.uarray([167,164,179,153], np.sqrt([167,164,179,153]))
M8  = unp.uarray([176,144,151,135], np.sqrt([176,144,151,135]))
M10 = unp.uarray([123, 134], np.sqrt([123, 134]))

Plateau = unp.uarray([6,2,12,10,37,54,86,unp.nominal_values(M10.mean()),unp.nominal_values(M8.mean()),
unp.nominal_values(M6.mean()),unp.nominal_values(M4.mean()),unp.nominal_values(M2.mean()),
unp.nominal_values(M1.mean()),unp.nominal_values(M0.mean()),unp.nominal_values(P1.mean()),
unp.nominal_values(P2.mean()),unp.nominal_values(P4.mean()),unp.nominal_values(P6.mean()),
unp.nominal_values(P8.mean()),140,120,63,40,0,0] , [np.sqrt(6),np.sqrt(2),
np.sqrt(12),np.sqrt(10),np.sqrt(37),np.sqrt(54),np.sqrt(86),unp.std_devs(M10.mean()),
unp.std_devs(M8.mean()),unp.std_devs(M6.mean()),unp.std_devs(M4.mean()),unp.std_devs(M2.mean()),
unp.std_devs(M1.mean()),unp.std_devs(M0.mean()),unp.std_devs(P1.mean()),unp.std_devs(P2.mean()),
unp.std_devs(P4.mean()),unp.std_devs(P6.mean()),unp.std_devs(P8.mean()),np.sqrt(140),
np.sqrt(120),np.sqrt(63),np.sqrt(40),np.sqrt(0),np.sqrt(0)])

#print("Plateau: ", Plateau, '\n')

# nächste Abschnitt für das Plateau

fit_x = np.linspace(-25,25,1000)#Plateau_x[9],Plateau_x[18],100)

def PlateauFit(x, b):
    return 0*x+b

def Flankenfit(x, A, b):
    return A*x + b

LinksParams, LinksCovariance = curve_fit(Flankenfit, Plateau_x[3:10], unp.nominal_values(Plateau[3:10]))
RechtsParams, RechtsCovariance = curve_fit(Flankenfit, Plateau_x[19:24], unp.nominal_values(Plateau[19:24]))
Links_errors = np.sqrt(np.diag(LinksCovariance))
Rechts_errors = np.sqrt(np.diag(RechtsCovariance))

fitParams, fitCovariance = curve_fit(PlateauFit, Plateau_x[9:19], unp.nominal_values(Plateau[9:19]), sigma = unp.std_devs(Plateau[9:19]), absolute_sigma = True )
Plateau_errors = np.sqrt(np.diag(fitCovariance))
Höhe = ufloat(fitParams[0], Plateau_errors[0])

A_L = ufloat(LinksParams[0], Links_errors[0])
A_R = ufloat(RechtsParams[0], Rechts_errors[0])
B_L = ufloat(LinksParams[1], Links_errors[1])
B_R = ufloat(RechtsParams[1], Rechts_errors[1])

x_L = (Höhe/2-B_L)/A_L
x_R = (Höhe/2-B_R)/A_R

print("zeiten: ", x_L, x_R, '\n')
print("Links: ", Plateau_x[3:10], A_L, B_L, '\n')
print("rechts: ", Plateau_x[19:24], A_R, B_R, '\n')

plt.errorbar(Plateau_x, unp.nominal_values(Plateau), xerr= 0, yerr = unp.std_devs(Plateau), fmt = 'kx', label = r'Messwerte')
plt.plot(fit_x, PlateauFit(fit_x, *fitParams), 'r-', label = r'linearer Fit')
plt.axvline(unp.nominal_values(x_L), color = 'r', linestyle = '--', label = r'$T_{\mathrm{VZ,links}}$')
plt.axvline(unp.nominal_values(x_R), color = 'r', linestyle = '--', label = r'$T_{\mathrm{VZ,rechts}}$')
plt.ylabel(r'N(t)')
plt.xlabel(r'$T_{\mathrm{VZ}}$ in $\mathrm{ns}$')
plt.legend(loc = 'best')
plt.savefig('Plateau.pdf')
#plt.show()

print("Bestimmte Höhe des Plateaus: ", Höhe, '\n')

# Kalibrierung

def Kalibrierung(x, A, B):
    return A*x + B

KParams, KCovariance = curve_fit(Kalibrierung, Channels_real, Kalibrierung_y)
K_errors = np.sqrt(np.diag(KCovariance))
Steigung = ufloat(KParams[0], K_errors[0])
Abschnitt = ufloat(KParams[1], K_errors[1])
print("Abschnitt der Gerade: ", Abschnitt, '\n')
print("Steigung der Gerade: ", Steigung, '\n')

plt.clf()
plt.plot(Channels_real, Kalibrierung_y, 'kx', label = r'Messwerte')
plt.plot(Channels, Kalibrierung(Channels, *KParams), 'r-', label = r'linearer Fit')
plt.ylabel(r'$t$ in $\mathrm{\mu s}$')
plt.xlabel(r'Kanal')
plt.legend(loc = 'best')
#plt.savefig('Kalibrierung.pdf')
#plt.show()

# Auswertung der Zeiten

def Lebensdauer(x, N0, LB, U):
    return N0 * np.exp(-LB*x) + U

Zeiten = Kalibrierung(Channels, *KParams)

Zeiten_neu = np.linspace(0,0,512)
Counts_neu = np.linspace(0,0,512)

Zählstruktur = np.linspace(0,255,256)
Counts_neu_1 = np.array([Counts_M[2*i]+Counts_M[2*i+1] for i,n in enumerate(Zählstruktur)])
Zeiten_neu_1 = np.array([(Zeiten[2*i]+Zeiten[2*i+1])/2 for i,n in enumerate(Zählstruktur)])

for i,n in enumerate(Counts_neu_1):
    if n == 0:
        print("Index von 0:", i)

Counts_r = np.array([n for n in Counts_M])
Zeiten_r = np.array([n for n in Zeiten])

#print(Counts_M)

Counts_M[0] = 1
Counts_M[1] = 1
Counts_M[2] = 1

for i,n in enumerate(Counts_M):
    #print("Schlaufe", i, Counts_M[i-1:i+2])
    k = i
    #print(i)
    if Counts_M[i]!=0:
        Counts_neu[i] = Counts_M[i]
        Zeiten_neu[i] = Zeiten[i]
        #print(i,n,"--------------")
    elif Counts_M[i]==0:
        if Counts_M[i-1]!=0 and Counts_M[i+1]!=0:
            #print(i,0)
            Zeiten_neu[i-1] = (Zeiten[i]+Zeiten[i-1])/2
            Counts_neu[i] = Counts_M[i+1]
            Zeiten_neu[i] = Zeiten[i+1]
            #print("nächster", Counts_M[k+1])
            while k < 510:
                Counts_M[k] = Counts_M[k+1]
                Zeiten[k] = Zeiten[k+1]
                k=k+1
            #print(Counts_M[i:i+10])
        elif Counts_M[i-1]!=0 and Counts_M[i+1]==0 and Counts_M[i+2]!=0:
            #print(i, 00)
            Zeiten_neu[i-1] = (Zeiten[i]+Zeiten[i-1])/2
            Counts_neu[i] = Counts_M[i+2]
            Zeiten_neu[i] = (Zeiten[i+1]+Zeiten[i+2])/2
            while k <= 509:
                Counts_M[k] = Counts_M[k+2]
                Zeiten[k] = Zeiten[k+2]
                k=k+1
        elif Counts_M[i-1]==0:
            #print(i, i-1, i+1, "Abbruch")
            break
        else:
            #print(i,Counts_M[i-3:i+2], "Abbruch")
            break

for i,n in enumerate(Counts_neu):
    if n == 0:
        print("Fehler", i)
        break
#print(Counts_neu[0:378])

Counts_M = Counts_r
Zeiten = Zeiten_r

Counts_neu[0] = 0
Counts_neu[1] = 0
Counts_neu[2] = 0

#Counts_fusch = np.array([n for n in Counts_neu if n != 0]) # mach aus Counts_M Counts_neu
#Zeiten_fusch = np.array([Zeiten_neu[i] for i,n in enumerate(Counts_neu) if n != 0]) # mach aus Counts_M Counts_neu und Zeiten
#Errors_fusch = np.sqrt(Counts_fusch)

Counts_fusch = Counts_neu_1[1:184]
Zeiten_fusch = Zeiten_neu_1[1:184]
Errors_fusch = np.sqrt(Counts_fusch)

LD_Params, LD_Covariance = curve_fit(Lebensdauer, Zeiten_fusch, Counts_fusch,  sigma = Errors_fusch, absolute_sigma = True)
LD_Errors = np.sqrt(np.diag(LD_Covariance))

N0 = ufloat(LD_Params[0], LD_Errors[0])
LB = ufloat(LD_Params[1], LD_Errors[1])
LD = 1 /LB
U_fit = ufloat(LD_Params[2], LD_Errors[2])

print("Wert für N0: ", N0)
print("Zerfallskonstante: ", LB)
print("Lebensdauer: ", LD)
print("Untergrund durch Fit: ", U_fit, '\n')

Zeiten_x = np.linspace(-2, 12, 1000)

plt.clf()
plt.errorbar(Zeiten_fusch, Counts_fusch, xerr = 0, yerr = Errors_fusch, fmt = 'kx', label = r'Messwerte')
plt.plot(Zeiten_x, Lebensdauer(Zeiten_x, *LD_Params), 'r-', label = r'Exponentieller Fit')
plt.xlim(-0.5,11.5)
#plt.ylim(-2, 65)
plt.axvline(0, color = 'b', linestyle = '-', label = r'$t=0$')
plt.ylabel(r'N(t)')
plt.yscale('log')
plt.xlabel(r'$t$ in $\mathrm{ns}$')
plt.legend(loc = 'best')
plt.savefig('Spektrum_klein_log.pdf')
plt.show()

print("Verhältnisse Untergrund:", U_Search/U_fit)

LD_Lit = 2.20

print("Verhältniss Lebensdaier:", LD/LD_Lit)


<gh_stars>0
import sys
import os
import numpy
import scipy.linalg
import matplotlib.animation
import matplotlib.pyplot as plt

def phi(N, x):
    coef = tuple(0 if i < N else 1 for i in range(N+1))
    return numpy.polynomial.hermite.hermval(x, c=coef) * numpy.exp(-(x**2) / 2) / (numpy.power(numpy.pi, 0.25) * numpy.sqrt(2**N * numpy.math.factorial(N)))

def V1D(lam, x):
    return x*x / 2 + lam * numpy.power(x, 4)

def V2D(lam, x, y):
    return (x*x + y*y) / 2 + lam * x*x*y*y

def finite_difference_step(psi, tau, h, V):
    return psi + 1j * tau * ((numpy.roll(psi,1) + numpy.roll(psi,-1) - 2 * psi) / (2*h**2) - V * psi)

def finite_propagator_step(psi, tau, h, K, V):
    hbar_psi = psi
    state = psi
    for k in range(1, K+1):
        hbar_psi = - (1j * tau * ((numpy.roll(hbar_psi,1) + numpy.roll(hbar_psi,-1) - 2 * hbar_psi) / (2*h**2) - V * hbar_psi))
        state += (-1j*tau)**k / numpy.math.factorial(k) * hbar_psi
    return state

def implicit_scheme_step(psi, tau, h, V):
    upper = - 1j * tau / (4*h*h) * numpy.ones(len(psi))
    upper[0] = 0
    lower = - 1j * tau / (4*h*h) * numpy.ones(len(psi))
    lower[-1] = 0
    diag = 1j * tau / 2 * V + 1 + 1j * tau / (2*h*h)

    ab = numpy.array([upper, diag, lower], dtype=complex)
    b = psi + 1j * tau / 4 * ((numpy.roll(psi,1) + numpy.roll(psi,-1) - 2 * psi) / (h**2) - 2 * V * psi)

    return scipy.linalg.solve_banded((1,1), ab, b)

def time_evolution(N, lam, x, h, tau, niter, method, K=None):
    if method == 1 and K is None:
        print("Specify K for finite propagator")
        sys.exit()

    Vm = V1D(lam, x)
    states = numpy.zeros((niter+1, len(Vm)), dtype=complex)

    states[0] = phi(N, x)
    for n in range(niter):
        states[n][0] = 0
        states[n][-1] = 0
        if method == 0:
            states[n+1] = finite_difference_step(states[n], tau, h, Vm)
        elif method == 1:
            states[n+1] = finite_propagator_step(states[n], tau, h, K, Vm)
        elif method == 2:
            states[n+1] = implicit_scheme_step(states[n], tau, h, Vm)

    return states

def animate(x, states, filename, lam, h, tau, N, interval_base=200, interval_slow=500, range_slow=None, frames_slow=20):
    print("Rendering animation ...")
    fig, ax = plt.subplots()
    plt.suptitle("Time evolution of anharmonic oscilator wave function\n($\\lambda = {0}, h = {1}, \\tau = {2}, N = {3}$)".format(lam, h, tau, N))
    plt.xlabel(r"$x$")
    plt.ylabel(r"$|\psi|^2$")

    def animate_func(frame):
        ax.clear()
        ax.plot(x, numpy.abs(states[frame])**2)

    def animate_func2(frame):
        ax.clear()
        ax.plot(x, numpy.abs(states[frame+range_slow])**2)

    os.chdir("hw1_schroedinger-equation-spectrum-and-time-evolution/images/")

    animation = matplotlib.animation.FuncAnimation(fig, animate_func, frames=len(states), interval=interval_base)
    animation.save(filename+".html")

    if range_slow is not None:
        animation = matplotlib.animation.FuncAnimation(fig, animate_func2, frames=frames_slow, interval=interval_slow)
        animation.save(filename+"_slow.html")

def testing(method=2):
    if method == 0:  # tau < h**2
        N = 0
        lam = 0
        L = 10
        h = 0.1
        tau = 0.0099
        niter = 120
        x = numpy.linspace(-L, L, 2*L/h)

        states = time_evolution(N, lam, x, h, tau, niter, method)
        animate(x, states, "naloga1_finite_diff"+str(N)+","+str(lam)+","+str(h)+","+str(tau), lam, h, tau, N, range_slow=90)

    elif method == 1:  # tau < 2 * pi * h**2
        N = 0
        lam = 0
        L = 10
        h = 0.1
        tau = 0.0099 * 3.14
        niter = 400
        K = 10
        x = numpy.linspace(-L, L, 2*L/h)
        states = time_evolution(N, lam, x, h, tau, niter, method, K=K)
        animate(x, states, "naloga1_finite_propagator"+str(N)+","+str(lam)+","+str(h)+","+str(tau)+","+str(K), lam, h, tau, N, interval_base=20, range_slow=340, frames_slow=60)

    elif method == 2:
        N = 0
        lam = 0
        L = 10
        h = 0.1
        tau = 0.0099
        niter = 1000
        x = numpy.linspace(-L, L, 2*L/h)
        states = time_evolution(N, lam, x, h, tau, niter, method)
        animate(x, states, "naloga1_implicit_scheme"+str(N)+","+str(lam)+","+str(h)+","+str(tau), lam, h, tau, N, interval_base=5)

def task1():
    Ns = range(6)
    lams = numpy.linspace(0, 0.5, 3)
    L = 10
    h = 0.1
    tau = 0.0099
    niter = 1000
    x = numpy.linspace(-L, L, int(2*L/h) + 1)

    print("Generating states ...")
    N_to_lams_to_states = {}
    for N in Ns:
        N_to_lams_to_states[N] = {}
        for lam in lams:
            N_to_lams_to_states[N][lam] = time_evolution(N, lam, x, h, tau, niter, method=2)

    # # for lambda = 0 it's constant
    # fig, axes = plt.subplots(2, 3, figsize=(16,9))
    # fig.suptitle("Evolution for $\\lambda=0$ is time independent")
    # for N in Ns:
    #     i = N // 3
    #     j = N % 3
    #     axes[i, j].set_title("$N={}$".format(N))
    #     axes[i, j].set_ylim(bottom=0, top=1)
    #     axes[i, j].plot(x, numpy.abs(N_to_lams_to_states[N][0][0])**2)
    # fig.savefig("naloga1_N_lambda_0.pdf", bbox_inches="tight")
    # sys.exit()

    # plot each subplot different N
    print("Animating per N ...")
    fig, axes = plt.subplots(2, 3, figsize=(32,18))
    fig.suptitle("Time evolution")

    def animate_per_N(frame):
        print("\t{0} | {1}".format(niter, frame+1))
        for N in Ns:
            i = N // 3
            j = N % 3
            axes[i, j].clear()
            axes[i, j].set_title("$N = {}$".format(N))
            axes[i, j].set_ylim(bottom=0, top=1.5)
            for lam in lams:
                axes[i, j].plot(x, numpy.abs(N_to_lams_to_states[N][lam][frame])**2)

    animation = matplotlib.animation.FuncAnimation(fig, animate_per_N, frames=niter, interval=50)
    animation.save("naloga1_N.mp4")

    # plot each subplot different lambda
    print("Animating per lambda ...")
    fig, axes = plt.subplots(1, 3, figsize=(32,18))
    fig.suptitle("Time evolution")

    def animate_per_lam(frame):
        print("\t{0} | {1}".format(niter, frame+1))
        for lam in lams:
            i = int(lam * 4)
            axes[i].clear()
            axes[i].set_title("$\\lambda = {}$".format(lam))
            axes[i].set_ylim(bottom=0, top=1.5)
            for N in Ns:
                axes[i].plot(x, numpy.abs(N_to_lams_to_states[N][lam][frame])**2)

    animation = matplotlib.animation.FuncAnimation(fig, animate_per_lam, frames=niter, interval=50)
    animation.save("naloga1_lambda.mp4")

    # plot each subplot different N, lambda=0
    print("Animating different N, lambda = 0 ...")
    fig, axes = plt.subplots(2, 3, figsize=(32,18))
    fig.suptitle("Time evolution for $\\lambda=0$")

    def animate_per_N_fixed_lambda(frame):
        print("\t{0} | {1}".format(niter, frame+1))
        lam = 0
        for N in Ns:
            i = N // 3
            j = N % 3
            axes[i, j].clear()
            axes[i, j].set_title("$N = {}$".format(N))
            axes[i, j].set_ylim(bottom=0, top=1)
            axes[i, j].plot(x, numpy.abs(N_to_lams_to_states[N][lam][frame])**2)

    animation = matplotlib.animation.FuncAnimation(fig, animate_per_N_fixed_lambda, frames=niter, interval=5)
    animation.save("naloga1_N_lambda_0.mp4")

def time_evolution_changing_lambda(lams, x, h, tau, niter, L, a=5):
    interval = int(numpy.ceil(niter / len(lams)))

    print("Generating states ...")
    states = numpy.zeros((niter+1, len(x)), dtype=complex)
    states[0] = phi(0, x-a)
    for n in range(niter):
        Vm = V1D(lams[n//interval], x)
        states[n][0] = 0
        states[n][-1] = 0
        states[n+1] = implicit_scheme_step(states[n], tau, h, Vm)

    print("Animating TASK 2 ...")
    fig, ax = plt.subplots()
    fig.suptitle("Task 2: $N=0, a={0}, \\lambda \\in {1}$".format(a, lams))

    def animate(frame):
        print("Animating TASK 2\t{0} | {1}".format(niter, frame+1))
        ax.clear()
        ax.set_title("$\\lambda={}$".format(lams[frame//interval]))
        ax.set_ylim(bottom=0, top=1.5)
        ax.plot(x, numpy.abs(states[frame])**2)

    animation = matplotlib.animation.FuncAnimation(fig, animate, frames=niter, interval=5)
    animation.save("naloga2_L{0}_a{1}.mp4".format(L, a))

def task2(a):
    Ns = range(6)
    lams = numpy.linspace(0, 0.5, 11)
    lams = [0]
    L = 10
    h = 0.1
    tau = 0.0099
    niter = 10000
    x = numpy.linspace(-L, L, int(2*L/h) + 1)

    states = time_evolution_changing_lambda(lams, x, h, tau, niter, L, a=a)

def task2_extra():
    """
    Time of travel of Gaussian to 0 from a = L/2 with respect to L.
    """
    N = 0
    lam = 0
    Ls = numpy.array([2*L for L in range(1,23)])
    h = 0.01
    tau = 0.000099

    iterss = []

    for L in Ls:
        a = L // 2
        print(L)
        x = numpy.linspace(-L, L, int(2*L/h) + 1)
        # eps = int(0.1 * len(x))

        Vm = V1D(lam, x)
        state = phi(N, x-a)

        iters = 0
        while True:
            prob = numpy.abs(state)**2
            mid = int(2*L/h) // 2
            # if max(prob) in prob[mid-eps:mid+eps]:
            if numpy.argmax(prob) <= mid:
                print(iters)
                iterss.append(iters)
                break

            state[0] = 0
            state[-1] = 0
            state = implicit_scheme_step(state, tau, h, Vm)
            iters += 1

    fig = plt.figure()
    plt.title("Iterations of Gaussian travel to center")
    plt.xlabel("$L$")
    plt.ylabel("Time")
    plt.plot(Ls, tau*numpy.array(iterss))
    plt.show()
    fig.savefig("naloga2_iters_of_gaussian_travel.pdf", bbox_inches="tight")

def task2_extra2():
    """
    Time of travel of Gaussian to 0 for fixed L and different a.
    """
    N = 0
    lam = 0
    L = 10
    h = 0.001
    tau = 0.000099
    aa = numpy.array([0.25*a for a in range((L-1)*4)])
    x = numpy.linspace(-L, L, int(2*L/h) + 1)
    Vm = V1D(lam, x)
    # eps=int(0.1*len(x))

    iterss = []
    for a in aa:
        print(a)
        state = phi(N, x-a)

        iters = 0
        while True:
            prob = numpy.abs(state)**2
            mid = int(2*L/h) // 2
            # if max(prob) in prob[mid-eps:mid+eps]:
            if numpy.argmax(prob) <= mid:
                print(iters)
                iterss.append(iters)
                break

            state[0] = 0
            state[-1] = 0
            state = implicit_scheme_step(state, tau, h, Vm)
            iters += 1

    fig = plt.figure()
    plt.title("Iterations of Gaussian travel to center ($L={}$)".format(L))
    plt.xlabel("$a$")
    plt.ylabel("Time")
    plt.plot(aa, tau*numpy.array(iterss))
    plt.show()
    fig.savefig("naloga2_iters_of_gaussian_travel_fixedL={}.pdf".format(L), bbox_inches="tight")

def task2_extra3():
    N = 0
    lam = 0
    L = 10
    h = 0.01
    tau = 0.001
    aa = numpy.array([0, 0.1, 0.2, 0.3, 0.4] + [0.5*a for a in range(1,L-1)])
    x = numpy.linspace(-L, L, int(2*L/h) + 1)
    Vm = V1D(lam, x)
    niter = 20000

    errorss = []
    for a in aa:
        print(a)
        errors = []
        initial = numpy.abs(phi(N, x-a))**2
        initial[0] = 0
        initial[-1] = 0
        state = phi(N, x-a)
        for i in range(niter):
            diff = (numpy.abs(state)**2 - initial)
            errors.append(numpy.sqrt(diff.dot(diff)))
            state[0] = 0
            state[-1] = 0
            state = implicit_scheme_step(state, tau, h, Vm)
        errorss.append(errors)

    fig = plt.figure()
    plt.title("Distance between initial state and running state")
    plt.xlabel("Time")
    plt.ylabel("Distance ($||\\,|\psi(t=n\\tau)|^2 - |\\psi(0)|^2\\,||$)")
    plots = []
    for i in range(len(aa)):
        plots += plt.plot(tau*numpy.array(range(niter)), numpy.array(errorss[i]) / numpy.amax(errorss), linewidth=1)
    plt.legend(plots, aa, title="$a$", loc="upper right")
    fig.savefig("naloga2_distance.pdf", bbox_inches="tight")
    plt.show()

def task2_extra3a():
    N = 0
    lam = 0
    L = 10
    h = 0.01
    tau = 0.001
    a = 3
    x = numpy.linspace(-L, L, int(2*L/h) + 1)
    Vm = V1D(lam, x)
    niter = 20000

    states = numpy.zeros((niter+1, len(x)), dtype=complex)
    states[0] = phi(0, x-a)
    initial = numpy.abs(states[0])**2
    for n in range(niter):
        states[n][0] = 0
        states[n][-1] = 0
        states[n+1] = implicit_scheme_step(states[n], tau, h, Vm)

    print("Animating TASK 2 ...")
    fig, ax = plt.subplots()
    fig.suptitle("Time evolution ($a = 3$)")

    def animate(frame):
        print("Animating TASK 2\t{0} | {1}".format(niter, frame+1))
        diff = (numpy.abs(states[frame])**2 - initial)
        dist = numpy.sqrt(diff.dot(diff))
        ax.clear()
        ax.set_title("$d(\\psi(0), \\psi({0})) = {1}$".format(frame, numpy.round(dist, decimals=5)))
        ax.set_ylim(bottom=0, top=1)
        ax.plot(x, initial, "k-", linewidth=1)
        ax.plot(x, numpy.abs(states[frame])**2)

    animation = matplotlib.animation.FuncAnimation(fig, animate, frames=niter, interval=2)
    animation.save("naloga2_L{0}_a{1}.mp4".format(L, a))

def task3(a=5):
    Ns = range(6)
    lams = numpy.linspace(0, 0.5, 11)
    L = 10
    h = 0.1
    tau = 0.0099
    niter = 10000
    x = numpy.linspace(-L, L, int(2*L/h) + 1)

if __name__ == "__main__":
    os.chdir("hw1_schroedinger-equation-spectrum-and-time-evolution/images/")

    # task1()
    for a in [1, 3, 5]:
        task2(a)
    # task3()
    # task2_extra()
    # task2_extra2()
    task2_extra3a()


<filename>adapt/prepare/cluster.py<gh_stars>10-100
"""Functions for clustering sequences before input.

This includes computing a distance matrix using MinHash, and
clustering that matrix.
"""

from collections import defaultdict
from collections import OrderedDict
import logging
import operator

import numpy as np
from scipy.cluster import hierarchy

from adapt.utils import lsh

__author__ = '<NAME> <<EMAIL>>'

logger = logging.getLogger(__name__)

def make_signatures_with_minhash(family, seqs):
    """Construct a signature using MinHash for each sequence.

    Args:
        family: lsh.MinHashFamily object
        seqs: dict mapping sequence header to sequences

    Returns:
        dict mapping sequence header to signature
    """
    # Construct a single hash function; use the same for all sequences
    h = family.make_h()

    signatures = {}
    for name, seq in seqs.items():
        signatures[name] = h(seq)
    return signatures

def create_condensed_dist_matrix(n, dist_fn):
    """Construct a 1d condensed distance matrix for scipy.

    Args:
        n: number of elements whose pairwise distances to store in the
            matrix
        dist_fn: function such that dist_fn(i, j) gives the distance
            between i and j, for all i<j<n

    Returns:
        condensed 1d distance matrix for input to scipy functions
    """
    def idx(i, j):
        # Compute index in 1d vector for pair (i, j)
        return int((-1 * i*i)/2 + i*n - 3*i/2 + j - 1)

    dist_matrix_len = int(n*(n-1)/2)
    dist_matrix = np.zeros(dist_matrix_len)

    for j in range(n):
        for i in range(j):
            dist_matrix[idx(i, j)] = dist_fn(i, j)

    return dist_matrix

def cluster_from_dist_matrix(dist_matrix, threshold, num_clusters=None):
    """Use scipy to cluster a distance matrix.

    Args:
        dist_matrix: distance matrix, represented in scipy's 1d condensed form
        threshold: maximum inter-cluster distance to merge clusters (higher
            results in fewer clusters)
        num_clusters: if set, cluster such that this is the maximum number
            of clusters; threshold is ignored and must be None

    Returns:
        list c such that c[i] is a collection of all the observations
        (whose pairwise distances are indexed in dist) in the i'th
        cluster, in sorted order by descending cluster size
    """
    if threshold is None and num_clusters is None:
        raise ValueError(("threshold or num_clusters must be set"))
    elif threshold is not None and num_clusters is not None:
        raise ValueError(("Only one of threshold or num_clusters can be "
            "set"))

    linkage = hierarchy.linkage(dist_matrix, method='average')
    if threshold:
        clusters = hierarchy.fcluster(linkage, threshold, criterion='distance')
    elif num_clusters:
        clusters = hierarchy.fcluster(linkage, num_clusters,
            criterion='maxclust')

    # clusters are numbered starting at 1, but base the count on
    # first_clust_num just in case this changes
    first_clust_num = min(clusters)
    num_clusters_found = max(clusters) + 1 - first_clust_num
    elements_in_cluster = defaultdict(list)
    for i, clust_num in enumerate(clusters):
        elements_in_cluster[clust_num].append(i)
    cluster_sizes = {c: len(elements_in_cluster[c])
                     for c in range(first_clust_num,
                                    num_clusters_found + first_clust_num)}

    elements_in_cluster_sorted = []
    for clust_num, _ in sorted(cluster_sizes.items(),
            key=operator.itemgetter(1), reverse=True):
        elements_in_cluster_sorted += [elements_in_cluster[clust_num]]
    return elements_in_cluster_sorted

def cluster_with_minhash_signatures(seqs, k=12, N=100, threshold=0.1,
        num_clusters=None, return_dist_matrix_and_indices=False):
    """Cluster sequences based on their MinHash signatures.

    Args:
        seqs: dict mapping sequence header to sequences
        k: k-mer size to use for k-mer hashes (smaller is likely more
            sensitive for divergent genomes, but may lead to false positives
            in determining which genomes are close)
        N: number of hash values to use in a signature (higher is slower for
            clustering, but likely more sensitive for divergent genomes)
        threshold: maximum inter-cluster distance to merge clusters, in
            average nucleotide dissimilarity (1-ANI, where ANI is
            average nucleotide identity); higher results in fewer
            clusters
        num_clusters: if set, cluster such that this is the maximum number
            of clusters; threshold is ignored and must be None
        return_dist_matrix_and_indices: if set, return the pairwise distance
            matrix and sequences indices in each cluster; when used,
            seqs should be an OrderedDict

    Returns:
        if return_dist_matrix_and_indices:
            tuple (dm, c) such that dm is the pairwise distance matrix and
            c is a list such that c[i] gives a collection of sequence
            indices (corresponding to indices in dm) in the same cluster
            (note dm is a 1d condensed matrix in scipy's form); clusters in
            c are sorted in descending order of size
        else:
            list c such that c[i] gives a collection of sequence headers
            in the same cluster, and the clusters in c are sorted
            in descending order of size
    """
    if len(seqs) == 1:
        # Simply return one cluster
        if return_dist_matrix_and_indices:
            return None, [[0]]
        else:
            return [[list(seqs.keys())[0]]]

    family = lsh.MinHashFamily(k, N=N)
    signatures_map = make_signatures_with_minhash(family, seqs)

    # Map each sequence header to an index (0-based), and get
    # the signature for the corresponding index
    num_seqs = len(seqs)
    seq_headers = []
    signatures = []
    for name, seq in seqs.items():
        seq_headers += [name]
        signatures += [signatures_map[name]]

    # Eq. 4 of the Mash paper (Ondov et al. 2016) shows that the
    # Mash distance, which is shown to be closely related to 1-ANI, is:
    #  D = (-1/k) * ln(2*j/(1+j))
    # where j is a Jaccard similarity. Solving for j:
    #  j = 1/(2*exp(k*D) - 1)
    # So, for a desired distance D in terms of 1-ANI, the corresponding
    # Jaccard distance is:
    #  1.0 - 1/(2*exp(k*D) - 1)
    # We can use this to calculate a clustering threshold in terms of
    # Jaccard distance
    if threshold is not None:
        jaccard_dist_threshold = 1.0 - 1.0/(2.0*np.exp(k*threshold) - 1)
    else:
        # Ignore inter-cluster distance; use num_clusters instead
        jaccard_dist_threshold = None

    def jaccard_dist(i, j):
        # Return estimated Jaccard dist between signatures at
        # index i and index j
        return family.estimate_jaccard_dist(
            signatures[i], signatures[j])

    dist_matrix = create_condensed_dist_matrix(num_seqs, jaccard_dist)
    clusters = cluster_from_dist_matrix(dist_matrix,
        jaccard_dist_threshold, num_clusters=num_clusters)

    if return_dist_matrix_and_indices:
        return dist_matrix, clusters
    else:
        seqs_in_cluster = []
        for cluster_idxs in clusters:
            seqs_in_cluster += [[seq_headers[i] for i in cluster_idxs]]
        return seqs_in_cluster

def find_representative_sequences(seqs, k=12, N=100, threshold=0.1,
        num_clusters=None, frac_to_cover=1.0):
    """Find a set of representative sequences.

    This clusters seqs, and then determines a medoid for each cluster.
    It returns the medoids.

    This will not return representative sequences with ambiguity or NNNs.

    Args:
        seqs, k, N, threshold, num_clusters: see cluster_with_minhash_signatures()
        frac_to_cover: return medoids from clusters that collectively
            account for at least this fraction of all sequences; this
            allows ignoring representative sequences for outlier
            clusters

    Returns:
        tuple (set of sequence headers representing cluster medoids,
        fraction of all sequences contained in cluster)
    """
    seqs = OrderedDict(seqs)
    dist_matrix, clusters = cluster_with_minhash_signatures(
            seqs, k=k, N=N, threshold=threshold, num_clusters=num_clusters,
            return_dist_matrix_and_indices=True)

    seqs_items = list(seqs.items())
    n = len(seqs)
    def idx(i, j):
        # Compute index in 1d vector for pair (i, j)
        if i > j:
            i, j = j, i
        return int((-1 * i*i)/2 + i*n - 3*i/2 + j - 1)

    rep_seqs = []
    rep_seqs_frac = []
    num_seqs_accounted_for = 0
    for cluster_idxs in clusters:
        # Stop if we have already accounted for frac_to_cover of the
        # sequences
        # Note that clusters should be sorted in descending order of
        # size, so any clusters after this one will be the same size
        # or smaller
        if float(num_seqs_accounted_for) / len(seqs) >= frac_to_cover:
            break

        # Find the medoid of this cluster
        # Simply look over all pairs in the cluster (there are faster
        # algorithms, though not linear)
        curr_medoid = None
        curr_medoid_dist_total = None
        for i in cluster_idxs:
            # Only allow i to be the medoid if it does not have ambiguity
            seq = seqs_items[i][1]
            if sum(seq.count(b) for b in ('A','C','G','T')) != len(seq):
                # Has ambiguity or NNNs; skip
                continue

            # Compute the total distance to all other sequences in this
            # cluster, and check if this is the medoid
            dist_total = 0
            for j in cluster_idxs:
                if i == j:
                    continue
                dist_total += dist_matrix[idx(i,j)]
            if curr_medoid is None or dist_total < curr_medoid_dist_total:
                curr_medoid = i
                curr_medoid_dist_total = dist_total
        if curr_medoid is not None:
            rep_seqs += [curr_medoid]
            rep_seqs_frac += [float(len(cluster_idxs)) / len(seqs)]
            num_seqs_accounted_for += len(cluster_idxs)
        else:
            # All sequences have ambiguity or NNNs; raise a warning and
            # skip this cluster
            logger.warning(("Cannot find medoid for cluster of size %d "
                "because all sequences have ambiguity or NNNs; skipping "
                "this cluster"),
                len(cluster_idxs))

    return ([seqs_items[i][0] for i in rep_seqs], rep_seqs_frac)