<filename>harmonizome/utility_functions.py
# Adapted from code created by <NAME>

import datetime
import os
import zipfile

import numpy as np
import pandas as pd
import scipy.spatial.distance as dist
import scipy.sparse as sp
from statsmodels.distributions.empirical_distribution import ECDF

from tqdm import tqdm

def remove_impute(df):
    '''
    Removes rows and columns that have more than 95% of their data missing,
    or 0. Replacing any missing data leftover after removal with
    the means of the rows.
    '''
    r, c = df.shape
    df.loc[np.sum(np.logical_or(np.isnan(df), df == 0), axis=1) < 0.05 * r,
           np.sum(np.logical_or(np.isnan(df), df == 0), axis=0) < 0.05 * c]

    return df.fillna(df.mean(axis=0))

def merge(df, axis):
    '''
    Merges duplicate rows or columns, depending on the axis specified. The
    final values of the merged rows or columns is determined by the method.
    '''
    if axis == 'column':
        return df.groupby(df.columns, axis=1).mean()
    elif axis == 'row':
        return df.groupby(level=0, axis=0).mean()

def quantile_normalize(df):
    '''
    Performs quantile normalization on the input DataFrame.
    '''
    # from ayhan on StackOverflow
    rank_mean = df.stack().groupby(
        df.rank(method='first').stack().astype(int)).mean()
    return df.rank(method='min').stack().astype(int).map(
        rank_mean).unstack()

def zscore(df, epsilon=0):
    '''
    Calculates the modified z-score of df according to the specified axis.

    Parameters:
        axis - the axis on which to calculate the z-scores. Either 'row' or 'column'
        epsilon - small adjustment in the case of divide by 0 errors.
    '''
    np.seterr(divide='ignore', invalid='ignore')
    median_y = np.median(df, axis=1)[:, np.newaxis]
    abs_dev = np.abs(df - median_y)
    median_dev = np.median(abs_dev, axis=1)
    mean_dev = np.mean(abs_dev, axis=1)
    median_abs_dev = np.broadcast_to(median_dev[:, np.newaxis], df.shape)
    mean_abs_dev = np.broadcast_to(mean_dev[:, np.newaxis], df.shape)
    modified_z_scores = np.where(median_abs_dev != 0,
                                 0.6745 * (df - median_y) / median_abs_dev,
                                 (df - median_y) / (1.253314 * mean_abs_dev + epsilon))

    return pd.DataFrame(data=modified_z_scores, index=df.index,
                        columns=df.columns)

def log2(df):
    '''
    Returns a dataframe with the adjusted log2 values of the input.
    '''
    return np.log2(df + 1)

def map_symbols(df, symbol_lookup, remove_duplicates=False):
    '''
    Replaces the index of the df, which are gene names, with
    corresponding approved gene symbols according to the given symbol_lookup
    dictionary. If any gene names are not in the mapping, they are discarded
    from the DataFrame.
    '''
    tqdm.pandas()
    df = df.reset_index()

    df.iloc[:, 0] = df.iloc[:, 0].progress_map(
        lambda x: symbol_lookup.get(x, np.nan))

    df = df.dropna(subset=[df.columns[0]])
    if remove_duplicates:
        df = df.drop_duplicates()
    df = df.set_index(df.columns[0])
    return df

def binary_matrix(df):
    '''
    Creates an adjacency matrix from df, which is a gene-attribute edge
    list.
    '''
    matrix = pd.crosstab(df.index, df.iloc[:, 0]) > 0
    matrix.index.name = df.index.name
    matrix.columns.name = df.columns[0]
    return matrix

def ternary_matrix(df):
    '''
    Returns the input matrix with all significant values, greater than 0.95
    or less than -0.95, mapped to 1 or -1, respectively. All other values
    are mapped to 0.
    '''
    def mapter(x):
        if x >= 0.95:
            return 1
        elif x <= -0.95:
            return -1
        else:
            return 0

    return df.applymap(mapter)

def save_setlib(df, lib, direction, path, name):
    '''
    If lib = 'gene', this creates a file which lists all attributes and the
    genes that are correlated in the direction given with that attribute.

    If lib = 'attribute', this creates a file which lists all genes and the
    attributes that are correlated in the direction given with that gene.
    The year and month are added at the end of the name. The path the file is
    saved to is thus
        path + name + '_<year>_<month>.gmt'
    '''
    filenameGMT = file_name(path, name, 'gmt')
    direction = {'up': 1, 'down': -1}[direction]

    if not (lib == 'gene' or lib == 'attribute'):
        return
    if lib == 'attribute':
        df = df.T

    with open(filenameGMT, 'w') as f:
        arr = df.reset_index(drop=True).to_numpy(dtype=np.int_)
        attributes = df.columns

        w, h = arr.shape
        for i in tqdm(range(h)):
            print(attributes[i], *df.index[arr[:, i] == direction],
                  sep='\t', end='\n', file=f)

def similarity_matrix(df, metric, dtype=None, sparse=False):
    '''
    Creates a similarity matrix between the rows of the df based on
    the metric specified. The resulting matrix has both rows and columns labeled
    by the index of df.
    '''
    if sparse and metric == 'jaccard':
        # from na-o-ys on Github
        sparse = sp.csr_matrix(df.to_numpy(dtype=np.bool).astype(np.int))
        cols_sum = sparse.getnnz(axis=1)
        ab = sparse * sparse.T
        denom = np.repeat(cols_sum, ab.getnnz(axis=1)) + \
            cols_sum[ab.indices] - ab.data
        ab.data = ab.data / denom
        similarity_matrix = ab.todense()
        np.fill_diagonal(similarity_matrix, 1)

    else:
        similarity_matrix = dist.pdist(df.to_numpy(dtype=dtype), metric)
        similarity_matrix = dist.squareform(similarity_matrix)
        similarity_matrix = 1 - similarity_matrix

    similarity_df = pd.DataFrame(
        data=similarity_matrix, index=df.index, columns=df.index)
    similarity_df.index.name = None
    similarity_df.columns.name = None
    return similarity_df

def gene_list(df, geneid_lookup):
    '''
    Creates a list of genes and the corresponding Entrez Gene IDs(supplied by
    the NCBI)

    Note: this differs from the previous function in its behavior with dealing
    with genes that do not have an ID. This function will set the id of the gene
    to -1, whereas the previous script will set them to np.nan.
    '''
    gene_ids = np.array([geneid_lookup.get(x, -1)
                         if np.isfinite(geneid_lookup.get(x, -1))
                         else -1 for x in tqdm(df.index)], dtype=np.int_)
    df = pd.DataFrame(gene_ids, index=df.index,
                      columns=['Gene ID'])
    return df

def attribute_list(df, metaData=None):
    '''
    Creates a list of attributes in the form of a DataFrame, with the attributes
    as the indices. If metaData is specified, it returns appends the attributes
    of df onto the metaData DataFrame.
    '''
    if metaData is not None:
        attribute_list = metaData.reindex(df.columns)
        attribute_list.index.name = df.columns.name
    else:
        attribute_list = pd.DataFrame(index=df.columns)
    return attribute_list

def standardized_matrix(df):
    '''
    Creates a standardized matrix by using an emperical CDF for each row.
    Each row in the df should represent a single gene.

    Requires:
    Indices of the DataFrame are unique.
    '''
    arr = df.to_numpy(copy=True)

    def process(array):
        ourECDF = ECDF(array)
        array = ourECDF(array)
        mean = np.mean(array)
        array = 2 * (array - mean)
        return array

    for i in tqdm(range(arr.shape[0])):
        arr[i, :] = process(arr[i, :])

    values = arr.flatten()
    ourECDF = ECDF(values)
    ourECDF = ourECDF(values).reshape(arr.shape)

    mean = np.mean(ourECDF)
    ourECDF = 2 * (ourECDF - mean)
    newDF = pd.DataFrame(data=ourECDF, index=df.index,
                         columns=df.columns)
    return newDF

def edge_list(df):
    '''
    Creates the gene-attribute edge list from the given input DataFrame,
    attribute and gene lists. The year and month are added at the
    end of the name. The path the file is saved to is thus
        path + name + '_<year>_<month>.gmt'
    Also prints the number of cells in df that are statistically
    significant, i.e. > 0.95 confidence.
    Requires:
        attributelist and genelist were generated from running
        createAttributeList and createGeneList on df, respectively.
    '''
    count = np.sum(np.sum(df >= 0.95) + np.sum(df <= -0.95))
    df = df.stack()
    df.name = 'Weight'
    print('The number of statisticaly relevent gene-attribute associations is: %d' % count)
    return df

def file_name(path, name, ext):
    '''
    Returns the file name by taking the path and name, adding the year and month
    and then the extension. The final string returned is thus
        '<path>/<name>_<year>_<month>.ext'
    '''
    date = str(datetime.date.today())[0:7].replace('-', '_')
    filename = ''.join([name, '_', date, '.', ext])
    return os.path.join(path, filename)

def save_data(df, path, name, compression=None, ext='tsv',
              symmetric=False, dtype=None, **kwargs):
    '''
    Save df according to the compression method given.
    compression can take these values:
        None or 'gmt' - defaults to pandas to_csv() function.
        'gzip' - uses the gzip compression method of the pandas to_csv() function
        'npz' - converts the DataFrame to a numpy array, and saves the array.
                The array is stored as 'axes[0]_axes[1]'. If symmetric is true,
                it is stored as 'axes[0]_axes[1]_symmetric' instead.
    ext is only used if compression is None or 'gzip'. The extension of the file
    will be .ext, or .ext.gz if 'gzip' is specified.
    axes must only be specified if compression is 'npz'. It is a string tuple
    that describes the index and columns df, i.e. (x, y) where x, y =
    'gene' or 'attribute'.
    symmetric is only used if compression is 'npz', and indicates if df
    is symmetric and can be stored as such.
    dtype is only used if compression is 'npz', and indicates a dtype that the
    array can be cast to before storing.

    The year and month are added at the end of the name. The path the file is
    saved to is thus
        path + name + '_<year>_<month>.ext'
    where ext is .ext, .ext.gz, or .npz depending on the compression method.
    '''

    if compression is None:
        name = file_name(path, name, ext)
        df.to_csv(name, sep='\t', **kwargs)
    elif compression == 'gzip':
        name = file_name(path, name, ext + '.gz')
        df.to_csv(name, sep='\t', compression='gzip', **kwargs)
    elif compression == 'npz':
        name = file_name(path, name, 'npz')

        data = df.to_numpy(dtype=dtype)
        index = np.array(df.index)
        columns = np.array(df.columns)

        if symmetric:
            data = np.triu(data)
            np.savez_compressed(name, symmetric=data, index=index)
        else:
            np.savez_compressed(name, nonsymmetric=data,
                                index=index, columns=columns)

def load_data(filename):
    '''
    Loads a pandas DataFrame stored in a .npz data numpy array format.
    '''
    with np.load(filename, allow_pickle=True) as data_load:
        arrays = data_load.files
        if arrays[0] == 'symmetric':
            data = data_load['symmetric']
            index = data_load['index']
            data = data + data.T - np.diag(data.diagonal())
            df = pd.DataFrame(data=data, index=index, columns=index)
            return df
        elif arrays[0] == 'nonsymmetric':
            data = data_load['nonsymmetric']
            index = data_load['index']
            columns = data_load['columns']
            df = pd.DataFrame(data=data, index=index, columns=columns)
            return df

def archive(path):
    with zipfile.ZipFile('output_archive.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(path):
            for f in files:
                zipf.write(os.path.join(root, f))


"""
Utilities

<NAME> <<EMAIL>>

The Australian National University
Australian Centre for Robotic Vision
"""

import os
import torch
import pickle
import numpy as np
import scipy.io as sio

from tqdm import tqdm
from collections import defaultdict
from torch.utils.data import Dataset

from hicodet.hicodet import HICODet

import pocket
from pocket.core import DistributedLearningEngine
from pocket.utils import DetectionAPMeter, BoxPairAssociation

import sys
sys.path.append('detr')
import datasets.transforms as T

def get_iou(bb1, bb2):
    assert bb1['x1'] < bb1['x2']
    assert bb1['y1'] < bb1['y2']
    assert bb2['x1'] < bb2['x2']
    assert bb2['y1'] < bb2['y2']

    # determine the coordinates of the intersection rectangle
    x_left = max(bb1['x1'], bb2['x1'])
    y_top = max(bb1['y1'], bb2['y1'])
    x_right = min(bb1['x2'], bb2['x2'])
    y_bottom = min(bb1['y2'], bb2['y2'])

    if x_right < x_left or y_bottom < y_top:
        return 0.0

    # The intersection of two axis-aligned bounding boxes is always an
    # axis-aligned bounding box
    intersection_area = (x_right - x_left) * (y_bottom - y_top)

    # compute the area of both AABBs
    bb1_area = (bb1['x2'] - bb1['x1']) * (bb1['y2'] - bb1['y1'])
    bb2_area = (bb2['x2'] - bb2['x1']) * (bb2['y2'] - bb2['y1'])

    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
    assert iou >= 0.0
    assert iou <= 1.0
    return iou

def custom_collate(batch):
    images = []
    targets = []
    for im, tar in batch:
        images.append(im)
        targets.append(tar)
    return images, targets

class DataFactory(Dataset):
    def __init__(self, name, partition, data_root):
        if name not in ['hicodet', 'vcoco']:
            raise ValueError("Unknown dataset ", name)

        if name == 'hicodet':
            assert partition in ['train2015', 'test2015'], \
                "Unknown HICO-DET partition " + partition
            self.dataset = HICODet(
                root=os.path.join(data_root, 'hico_20160224_det/images', partition),
                anno_file=os.path.join(data_root, 'instances_{}.json'.format(partition)),
                target_transform=pocket.ops.ToTensor(input_format='dict')
            )
        else:
            print("ERROR!!!")
            exit(0)

        # Prepare dataset transforms
        normalize = T.Compose([
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
        if partition.startswith('train'):
            self.transforms = T.Compose([
                T.RandomHorizontalFlip(),
                T.ColorJitter(.4, .4, .4),
                T.RandomSelect(
                    T.RandomResize(scales, max_size=1333),
                    T.Compose([
                        T.RandomResize([400, 500, 600]),
                        T.RandomSizeCrop(384, 600),
                        T.RandomResize(scales, max_size=1333),
                    ])
                ), normalize,
        ])
        else:
            self.transforms = T.Compose([
                T.RandomResize([800], max_size=1333),
                normalize,
            ])

        self.name = name

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

    def __getitem__(self, i):
        image, target = self.dataset[i]
        if self.name == 'hicodet':
            target['labels'] = target['verb']
            # Convert ground truth boxes to zero-based index and the
            # representation from pixel indices to coordinates
            target['boxes_h'][:, :2] -= 1
            target['boxes_o'][:, :2] -= 1
        else:
            target['labels'] = target['actions']
            target['object'] = target.pop('objects')

        image, target = self.transforms(image, target)

        return image, target

class CacheTemplate(defaultdict):
    """A template for VCOCO cached results """
    def __init__(self, **kwargs):
        super().__init__()
        for k, v in kwargs.items():
            self[k] = v
    def __missing__(self, k):
        seg = k.split('_')
        # Assign zero score to missing actions
        if seg[-1] == 'agent':
            return 0.
        # Assign zero score and a tiny box to missing <action,role> pairs
        else:
            return [0., 0., .1, .1, 0.]

class CustomisedDLE(DistributedLearningEngine):
    def __init__(self, net, dataloader, max_norm=0, num_classes=117, **kwargs):
        super().__init__(net, None, dataloader, **kwargs)
        self.max_norm = max_norm
        self.num_classes = num_classes

    def _on_each_iteration(self):
        loss_dict = self._state.net(
            *self._state.inputs, targets=self._state.targets)

        if loss_dict['interaction_loss'].isnan():
            raise ValueError(f"The HOI loss is NaN for rank {self._rank}")

        self._state.loss = sum(loss for loss in loss_dict.values())
        self._state.optimizer.zero_grad(set_to_none=True)
        self._state.loss.backward()
        if self.max_norm > 0:
            torch.nn.utils.clip_grad_norm_(self._state.net.parameters(), self.max_norm)
        self._state.optimizer.step()

    @torch.no_grad()
    def test_hico(self, dataloader):
        net = self._state.net
        net.eval()

        all_correct = []
        verb_correct_obj_wrong = []
        obj_correct_verb_wrong = []
        obj_wrong_verb_wrong = []
        all_wrong = []
        missed = []

        dataset = dataloader.dataset.dataset
        associate = BoxPairAssociation(min_iou=0.5)
        conversion = torch.from_numpy(np.asarray(
            dataset.object_n_verb_to_interaction, dtype=float
        ))

        #test_anno = [0 for _ in range(len(dataset.anno_interaction))]
        #test_anno[30] = 2
        meter = DetectionAPMeter(
            91 * 2, nproc=1,
            #num_gt=test_anno,
            num_gt=dataset.anno_interaction,
            algorithm='11P'
        )

        for batch_idx, batch in tqdm(enumerate(dataloader)):
            if batch_idx < 4:
                continue

            inputs = pocket.ops.relocate_to_cuda(batch[0])
            #print("----")
            output = net(inputs)

            # Skip images without detections
            if output is None or len(output) == 0:
                continue
            # Batch size is fixed as 1 for inference
            assert len(output) == 1, f"Batch size is not 1 but {len(output)}."
            output = pocket.ops.relocate_to_cpu(output[0], ignore=True)
            output.pop("attn_maps")
            target = batch[-1][0]
            # Format detections
            boxes = output['boxes']
            boxes_h, boxes_o = boxes[output['pairing']].unbind(0)
            objects = output['objects']
            scores = output['scores']
            verbs = output['labels']
            interactions = conversion[objects, verbs]
            gt_bx_h = net.module.recover_boxes(target['boxes_h'], target['size'])
            print(gt_bx_h)
            gt_bx_o = net.module.recover_boxes(target['boxes_o'], target['size'])
            print(gt_bx_o)
            labels = torch.zeros_like(scores)
            unique_hoi = interactions.unique()

            for hoi_idx in unique_hoi:
                gt_idx = torch.nonzero(target['hoi'] == hoi_idx).squeeze(1)
                det_idx = torch.nonzero(interactions == hoi_idx).squeeze(1)
                if len(gt_idx):
                    labels[det_idx] = associate(
                        (gt_bx_h[gt_idx].view(-1, 4),
                        gt_bx_o[gt_idx].view(-1, 4)),
                        (boxes_h[det_idx].view(-1, 4),
                        boxes_o[det_idx].view(-1, 4)),
                        scores[det_idx].view(-1)
                    )
            print(".....")
            print(scores)
            print(interactions)
            print(labels)
            meter.append(scores, interactions, labels)

            inx = np.array([i for i in range(len(boxes_h)) if i % 2 == 0])
            boxes_h_filter = boxes_h[inx]
            boxes_o_filter = boxes_o[inx]
            objects_filter = objects[inx]
            scores_reshape = scores.reshape(-1, 2)

            pred_hbox = []
            pred_obox = []
            pred_obj = []
            pre_verb = []
            pred_verb_score = []
            for hbox, obox, score, obj in zip(boxes_h_filter, boxes_o_filter, scores_reshape, objects_filter.reshape(-1, 1)):
                max_score, max_idx = torch.max(score, 0)
                if max_score.item() > 0.1:
                    pred_hbox.append(hbox)
                    pred_obox.append(obox)
                    pred_obj.append(obj.item())
                    pre_verb.append(max_idx.item())
                    pred_verb_score.append(score)

            if len(pre_verb) == 0:
                max_score, max_idx = torch.max(scores, 0)
                pred_hbox.append(boxes_h[max_idx])
                pred_obox.append(boxes_o[max_idx])
                pred_obj.append(objects[max_idx].item())
                pre_verb.append(max_idx.item() % 2)
                pred_verb_score.append(scores[max_idx])
            # Recover target box scale
            all_correct.append(0)
            verb_correct_obj_wrong.append(0)
            obj_correct_verb_wrong.append(0)
            obj_wrong_verb_wrong.append(0)
            all_wrong.append(0)
            for hbox, obox, verb, obj in zip(pred_hbox, pred_obox, pre_verb, pred_obj):
                print(str(hbox) + " - " + str(obox) + " :" + str(obj) + " - " + str(verb))
                found = False
                for ghbox, gobox, gverb, gobj in zip(gt_bx_h, gt_bx_o, target["verb"], target["object"]):
                    hbox_overlap = get_iou({"x1": hbox[0].item(), "x2": hbox[2].item(), "y1": hbox[1].item(), "y2": hbox[3].item()},
                                           {"x1": ghbox[0].item(), "x2": ghbox[2].item(), "y1": ghbox[1].item(),
                                            "y2": ghbox[3].item()})
                    obox_overlap = get_iou({"x1": obox[0].item(), "x2": obox[2].item(), "y1": obox[1].item(), "y2": obox[3].item()},
                                           {"x1": gobox[0].item(), "x2": gobox[2].item(), "y1": gobox[1].item(),
                                            "y2": gobox[3].item()})

                    if hbox_overlap > 0.5 and obox_overlap > 0.5:
                        found = True
                        if verb == gverb.item() and obj == gobj.item():
                            all_correct[-1] += 1
                        elif verb == gverb:
                            verb_correct_obj_wrong[-1] += 1
                        elif obj == gobj:
                            obj_correct_verb_wrong[-1] += 1
                        else:
                            obj_wrong_verb_wrong[-1] += 1
                        break
                if not found:
                    all_wrong[-1] += 1

            missed.append(len(target["verb"]) - sum([all_correct[-1], verb_correct_obj_wrong[-1], obj_correct_verb_wrong[-1], obj_wrong_verb_wrong[-1]]))
            break
        return meter.eval(), {"all_correct": sum(all_correct), "verb_correct_obj_wrong": sum(verb_correct_obj_wrong), "obj_correct_verb_wrong": sum(obj_correct_verb_wrong),
                              "obj_wrong_verb_wrong": sum(obj_wrong_verb_wrong), "all_wrong": sum(all_wrong), "missed": sum(missed)}

    @torch.no_grad()
    def cache_hico(self, dataloader, cache_dir='matlab'):
        net = self._state.net
        net.eval()

        dataset = dataloader.dataset.dataset
        conversion = torch.from_numpy(np.asarray(
            dataset.object_n_verb_to_interaction, dtype=float
        ))
        object2int = dataset.object_to_interaction

        # Include empty images when counting
        nimages = len(dataset.annotations)
        all_results = np.empty((600, nimages), dtype=object)

        for i, batch in enumerate(tqdm(dataloader)):
            inputs = pocket.ops.relocate_to_cuda(batch[0])
            output = net(inputs)

            # Skip images without detections
            if output is None or len(output) == 0:
                continue
            # Batch size is fixed as 1 for inference
            assert len(output) == 1, f"Batch size is not 1 but {len(output)}."
            output = pocket.ops.relocate_to_cpu(output[0], ignore=True)
            # NOTE Index i is the intra-index amongst images excluding those
            # without ground truth box pairs
            image_idx = dataset._idx[i]
            # Format detections
            boxes = output['boxes']
            boxes_h, boxes_o = boxes[output['pairing']].unbind(0)
            objects = output['objects']
            scores = output['scores']
            verbs = output['labels']
            interactions = conversion[objects, verbs]
            # Rescale the boxes to original image size
            ow, oh = dataset.image_size(i)
            h, w = output['size']
            scale_fct = torch.as_tensor([
                ow / w, oh / h, ow / w, oh / h
            ]).unsqueeze(0)
            boxes_h *= scale_fct
            boxes_o *= scale_fct

            # Convert box representation to pixel indices
            boxes_h[:, 2:] -= 1
            boxes_o[:, 2:] -= 1

            # Group box pairs with the same predicted class
            permutation = interactions.argsort()
            boxes_h = boxes_h[permutation]
            boxes_o = boxes_o[permutation]
            interactions = interactions[permutation]
            scores = scores[permutation]

            # Store results
            unique_class, counts = interactions.unique(return_counts=True)
            n = 0
            for cls_id, cls_num in zip(unique_class, counts):
                all_results[cls_id.long(), image_idx] = torch.cat([
                    boxes_h[n: n + cls_num],
                    boxes_o[n: n + cls_num],
                    scores[n: n + cls_num, None]
                ], dim=1).numpy()
                n += cls_num

        # Replace None with size (0,0) arrays
        for i in range(600):
            for j in range(nimages):
                if all_results[i, j] is None:
                    all_results[i, j] = np.zeros((0, 0))
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
        # Cache results
        for object_idx in range(80):
            interaction_idx = object2int[object_idx]
            sio.savemat(
                os.path.join(cache_dir, f'detections_{(object_idx + 1):02d}.mat'),
                dict(all_boxes=all_results[interaction_idx])
            )

    @torch.no_grad()
    def cache_vcoco(self, dataloader, cache_dir='vcoco_cache'):
        net = self._state.net
        net.eval()

        dataset = dataloader.dataset.dataset
        all_results = []
        for i, batch in enumerate(tqdm(dataloader)):
            inputs = pocket.ops.relocate_to_cuda(batch[0])
            output = net(inputs)

            # Skip images without detections
            if output is None or len(output) == 0:
                continue
            # Batch size is fixed as 1 for inference
            assert len(output) == 1, f"Batch size is not 1 but {len(output)}."
            output = pocket.ops.relocate_to_cpu(output[0], ignore=True)
            # NOTE Index i is the intra-index amongst images excluding those
            # without ground truth box pairs
            image_id = dataset.image_id(i)
            # Format detections
            boxes = output['boxes']
            boxes_h, boxes_o = boxes[output['pairing']].unbind(0)
            scores = output['scores']
            actions = output['labels']
            # Rescale the boxes to original image size
            ow, oh = dataset.image_size(i)
            h, w = output['size']
            scale_fct = torch.as_tensor([
                ow / w, oh / h, ow / w, oh / h
            ]).unsqueeze(0)
            boxes_h *= scale_fct
            boxes_o *= scale_fct

            for bh, bo, s, a in zip(boxes_h, boxes_o, scores, actions):
                a_name = dataset.actions[a].split()
                result = CacheTemplate(image_id=image_id, person_box=bh.tolist())
                result[a_name[0] + '_agent'] = s.item()
                result['_'.join(a_name)] = bo.tolist() + [s.item()]
                all_results.append(result)

        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
        with open(os.path.join(cache_dir, 'cache.pkl'), 'wb') as f:
            # Use protocol 2 for compatibility with Python2
            pickle.dump(all_results, f, 2)


<filename>code/lr_train.py
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Reddit
from torch_geometric.data import ClusterData, ClusterLoader, Data
from torch_geometric.nn import SAGEConv, GCNConv

import argparse, time, datetime, shutil
import sys, os, glob, json

import warnings
warnings.filterwarnings("ignore")
import numpy as np
import pandas as pd
from sklearn.metrics import classification_report
from sklearn.metrics import f1_score, recall_score, precision_score, accuracy_score

from torch.utils.tensorboard import SummaryWriter

import nltk
nltk.download('punkt')
from torch.autograd import Variable
import torch.nn as nn
from statistics import stdev
import sys
sys.path.append("..")

from models.model import *
from utils.utils import *
from utils.data_utils_gnn import *
from utils.data_utils_txt import *
from utils.data_utils_hygnn import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class LR_model(torch.nn.Module):
    def __init__(self, config):
        super(LR_model, self).__init__()
        self.text_dim = 1024 if config['text_encoder'] == 'roberta' else 384
        self.graph_dim = config['graph_dim']

        if config['mode'] == 'gnn':
            self.in_dim = self.graph_dim*3 if config['model_name'] in ['gat', 'rgat'] else self.graph_dim
        elif config['mode'] == 'text':
            self.in_dim = self.text_dim
        else:
            self.in_dim = self.graph_dim*3 if config['model_name'] == 'gat' else self.graph_dim
            self.in_dim+=self.text_dim
        self.classifier = nn.Linear(self.in_dim, config['n_classes'])

    def forward(self, x):
        out = self.classifier(x)
        return out

class LR_Learner():
    def __init__(self, config):
        self.best_val_acc, self.best_val_f1, self.best_val_recall, self.best_val_precision = 0, 0, 0, 0
        self.actual_best_f1 = 0
        self.preds_list, self.labels_list = [] , []
        self.train_f1, self.train_precision, self.train_recall, self.train_accuracy = 0,0,0,0
        self.train_loss = []
        self.threshold = 0
        self.prev_val_loss, self.not_improved  = 0, 0
        self.best_val_loss = 1e4
        self.total_iters = 0
        self.terminate_training = False
        self.start_epoch, self.iters = 1, 0
        self.preds, self.loss = 0, 0
        self.start = time.time()

        self.model = LR_model(config)
        if config['optimizer'] == 'Adam':
            self.optimizer = torch.optim.AdamW(self.model.parameters(), lr = config['lr'], weight_decay = config['weight_decay'])
        else:
            self.optimizer = torch.optim.SGD(self.model.parameters(), lr = config['lr'], momentum = config['momentum'], weight_decay = config['weight_decay'])
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=config['lr_decay_step'], gamma = config['lr_decay_factor'])
        self.criterion = nn.BCELoss() if config['loss_func'] == 'bce' else nn.CrossEntropyLoss()

    def eval_lr(self, test = False):
        self.model.eval()
        preds_list, labels_list = [], []
        docs_list = []
        eval_loss = []

        loader = test_loader if test else val_loader
        with torch.no_grad():
            for iters, (batch_x, batch_y, doc) in enumerate(loader):
                batch_x = Variable(batch_x)
                batch_y = Variable(batch_y)
                preds = self.model(batch_x)
                if config['loss_func'] == 'bce':
                    preds = F.sigmoid(preds)
                labels = batch_y.float()
                # labels = batch_y.float()
                preds = torch.where(torch.isnan(preds), torch.zeros_like(preds), preds)
                loss = self.criterion(preds.to(device), labels.to(device))
                eval_loss.append(loss.detach().item())
                if config['loss_func'] == 'ce':
                    preds = F.softmax(preds, dim=1)
                    preds = torch.argmax(preds, dim=1)
                else:
                    preds = (F.sigmoid(preds)>0.5).type(torch.FloatTensor) if config['loss_func'] == 'bce_logits' else (preds>0.5).type(torch.FloatTensor)
                preds_list.append(preds.cpu().detach().numpy())
                labels_list.append(labels.cpu().detach().numpy())
                docs_list.append(doc)

            preds_list = [pred for batch_pred in preds_list for pred in batch_pred]
            labels_list = [label for batch_labels in labels_list for label in batch_labels]
            docs_list = [doc for docs in docs_list for doc in docs]
            # if test:
            #     self.save_correct_preds(docs_list, preds_list, labels_list)
            eval_loss = sum(eval_loss)/len(eval_loss)

            if test:
                print(classification_report(np.array(labels_list), np.array(preds_list)))

            if not test:
                eval_f1, eval_macro_f1, eval_recall, eval_precision, eval_accuracy = evaluation_measures(config, np.array(preds_list), np.array(labels_list))
                return eval_f1, eval_macro_f1, eval_precision, eval_recall, eval_accuracy, eval_loss

            else:
                eval_f1, eval_macro_f1, eval_recall, eval_precision, eval_accuracy = evaluation_measures(config, np.array(preds_list), np.array(labels_list))
                return eval_f1, eval_macro_f1, eval_precision, eval_recall, eval_accuracy, eval_loss

    def save_correct_preds(self, docs, preds, labels):
        print(len(preds), len(labels), len(docs))
        idx_correct_preds = [idx for idx in range(len(preds)) if preds[idx]==labels[idx]]
        correct_pred_docs = [docs[idx] for idx in range(len(docs)) if idx in idx_correct_preds]
        print(len(correct_pred_docs))
        temp_dict = {'correct_preds' : list(correct_pred_docs)}
        model = 'roberta' if config['mode'] == 'text' else config['model_name']
        dataname = config['data_name']
        correct_doc_file = os.path.join('data', 'complete_data', dataname, 'cached_embeds', 'correct_docs_{}.json'.format(model))

        print("Saving the list of correct test preds in : ", correct_doc_file)
        with open(correct_doc_file, 'w+') as j:
            json.dump(temp_dict, j)

    def train_epoch_step(self):
        self.model.train()
        lr = self.scheduler.get_lr()[0]
        self.total_iters += self.iters
        self.preds_list = [pred for batch_pred in self.preds_list for pred in batch_pred]
        self.labels_list = [label for batch_labels in self.labels_list for label in batch_labels]

        self.train_f1, self.train_macro_f1, self.train_recall, self.train_precision, self.train_accuracy = evaluation_measures(config, np.array(self.preds_list), np.array(self.labels_list))

        # Evaluate on dev set
        self.eval_f1, self.eval_macro_f1, self.eval_precision, self.eval_recall, self.eval_accuracy, self.eval_loss = self.eval_lr()

        # print stats
        print_stats(config, self.epoch, self.train_loss, self.train_accuracy, self.train_f1, self.train_precision, self.train_recall,
                    self.eval_loss, self.eval_accuracy, self.eval_f1, self.eval_precision, self.eval_recall, self.start, lr)

        if self.eval_f1 > self.best_val_f1:
            print("New High Score! Saving model...")
            self.best_val_f1 = self.eval_f1
            self.best_val_acc = self.eval_accuracy
            self.best_val_recall = self.eval_recall
            self.best_val_precision = self.eval_precision
            best_model = self.model.state_dict()
            # Save the state and the vocabulary
            torch.save({
                'epoch': self.epoch,
                'best_val_f1' : self.best_val_f1,
                'model_state_dict': best_model,
                # 'model_classif_state_dict': model.classifier.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
            }, os.path.join(model_path, config['model_save_name']))

        if self.epoch==1:
            print("Saving model !")
            best_model = self.model.state_dict()
            # Save the state and the vocabulary
            torch.save({
                'epoch': self.epoch,
                'best_val_f1' : self.best_val_f1,
                'model_state_dict': best_model,
                # 'model_classif_state_dict': model.classifier.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
            }, os.path.join(model_path, config['model_save_name']))

        self.scheduler.step()

        if self.eval_f1 - self.best_val_f1!=0 and self.eval_f1 - self.best_val_f1 < 1e-3:
            self.not_improved+=1
            print(self.not_improved)
            if self.not_improved >= config['patience']:
                self.terminate_training= True
        else:
            self.not_improved = 0

        if self.eval_f1 > self.best_val_f1 and self.eval_f1 - self.best_val_f1 > 1e-3:
            self.best_val_f1 = self.eval_f1
            self.not_improved=0

        # if self.best_val_loss - self.eval_loss < 1e-3:
        #     self.not_improved+=1
        #     print(self.not_improved)
        #     if self.not_improved >= config['patience']:
        #         self.terminate_training= True
        # else:
        #     self.not_improved = 0

        # if self.eval_loss < self.best_val_loss and self.best_val_loss - self.eval_loss > 1e-3:
        #     self.best_val_loss = self.eval_loss
        #     self.not_improved = 0

        self.preds_list = []
        self.labels_list = []

    def train_iters_step(self):
        if config['loss_func'] == 'bce':
            self.preds = F.sigmoid(self.preds)

        self.batch_labels = self.batch_y.float()
        self.loss = self.criterion(self.preds.to(device),  self.batch_labels.to(device))

        self.optimizer.zero_grad()
        self.loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5)
        self.optimizer.step()

        if config['loss_func'] == 'ce':
            self.preds = F.softmax(self.preds, dim=1)
            self.preds = torch.argmax(self.preds, dim=1)
        elif config['loss_func'] == 'bce':
            self.preds = (self.preds>0.5).type(torch.FloatTensor)
        elif config['loss_func'] == 'bce_logits':
            self.preds = F.sigmoid(self.preds)
            self.preds = (self.preds>self.threshold).type(torch.FloatTensor).squeeze(1)

        self.preds_list.append(self.preds.cpu().detach().numpy())
        self.labels_list.append(self.batch_labels.cpu().detach().numpy())

        self.train_loss.append(self.loss.detach().item())

    def train_main(self):
        print("\n\n"+ "="*100 + "\n\t\t\t\t\t Training Network\n" + "="*100)

        # Seeds for reproduceable runs
        torch.manual_seed(config['seed'])
        torch.cuda.manual_seed(config['seed'])
        np.random.seed(config['seed'])
        random.seed(config['seed'])
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

        print("\nBeginning training at:  {} \n".format(datetime.datetime.now()))

        for self.epoch in range(1, config['max_epoch']+1):
            for self.iters, (batch_x, batch_y, doc) in enumerate(train_loader):
                self.batch_x = Variable(batch_x)
                self.batch_y = Variable(batch_y)
                self.preds = self.model(self.batch_x)
                # print(self.preds)
                self.train_iters_step()
            self.train_epoch_step()
            if self.terminate_training:
                    break

        # Termination message
        if self.terminate_training:
            print("\n" + "-"*100 + "\nTraining terminated early because the Validation loss did not improve for   {}   epochs" .format(config['patience']))
        else:
            print("\n" + "-"*100 + "\nMaximum epochs reached. Finished training !!")

        print("\n" + "-"*50 + "\n\t\tEvaluating on test set\n" + "-"*50)

        if os.path.isfile(os.path.join(model_path, config['model_save_name'])):
            checkpoint = torch.load(os.path.join(model_path, config['model_save_name']))
            self.model.load_state_dict(checkpoint['model_state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        else:
            raise ValueError("No Saved model state_dict found for the chosen model...!!! \nAborting evaluation on test set...".format(config['model_name']))

        # Evaluate on dev set
        test_f1, test_macro_f1, test_precision, test_recall, test_accuracy, test_loss = self.eval_lr(test=True)

        print_test_stats(test_accuracy, test_precision, test_recall, test_f1, test_macro_f1, self.best_val_acc, self.best_val_precision, self.best_val_recall, self.best_val_f1)
        return test_f1, test_macro_f1, test_accuracy

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # Required Paths
    parser.add_argument('--data_path', type = str, default = './data/complete_data',
                          help='path to dataset folder that contains the adj and feat matrices, etc')
    parser.add_argument('--model_checkpoint_path', type = str, default = './model_checkpoints_lr',
                          help='Directory for saving trained model checkpoints')
    parser.add_argument("--model_save_name", type=str, default= 'best_model_lr.pt',
                       help = 'saved model name')

    #### Training Params ####

    # Named params
    parser.add_argument('--data_name', type = str, default = 'gossipcop',
                          help='dataset name: politifact / gossipcop / pheme / HealthRelease / HealthStory')
    parser.add_argument('--model_name', type = str, default = 'HGCN',
                          help='model name: gcn / graph_sage / graph_conv / gat / rgcn / HGCN')
    parser.add_argument('--text_encoder', type = str, default = 'roberta',
                          help='text encoder: cnn / roberta')
    parser.add_argument('--mode', type=str, default='gnn+text',
                        help='what features to use for classification: gnn / text / gnn+text')
    parser.add_argument('--loss_func', type = str, default = 'bce',
                        help = 'Loss function to use for optimization: bce / bce_logits / ce')
    parser.add_argument('--scheduler', type = str, default = 'step',
                        help = 'The type of lr scheduler to use anneal learning rate: step/multi_step')
    parser.add_argument('--optimizer', type = str, default = 'SGD',
                        help = 'Optimizer to use for training')

    # Dimensions/sizes params
    parser.add_argument('--batch_size', type = int, default = 16,
                          help='batch size for training"')
    parser.add_argument('--graph_dim', type = int, default = 512,
                          help='batch size for training"')

    # Numerical params
    parser.add_argument('--lr', type = float, default = 1e-3,
                          help='Learning rate for training')
    parser.add_argument('--weight_decay', type = float, default = 2e-3,
                        help = 'weight decay for optimizer')
    parser.add_argument('--momentum', type = float, default = 0.8,
                        help = 'Momentum for optimizer')
    parser.add_argument('--max_epoch', type = int, default = 50,
                        help = 'Max epochs to train for')
    parser.add_argument('--lr_decay_step', type = float, default = 3,
                        help = 'No. of epochs after which learning rate should be decreased')
    parser.add_argument('--lr_decay_factor', type = float, default = 0.8,
                        help = 'Decay the learning rate of the optimizer by this multiplicative amount')
    parser.add_argument('--patience', type = float, default = 10,
                        help = 'Patience no. of epochs for early stopping')
    parser.add_argument('--seed', type=int, default=168,
                        help='set seed for reproducability')

    # Options params
    parser.add_argument('--shuffle', type=bool, default=True,
                        help='whether to shuffle batches')

    args, unparsed = parser.parse_known_args()
    config = args.__dict__

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    config['device'] = device

    config['n_classes'] = 1

    # Check all provided paths:
    model_path = os.path.join(config['model_checkpoint_path'], config['data_name'], config['model_name'])

    if not os.path.exists(config['data_path']):
        raise ValueError("[!] ERROR: Dataset path does not exist")
    else:
        print("\nData path checked..")
    if not os.path.exists(model_path):
        print("\nCreating checkpoint path for saved models at:  {}\n".format(model_path))
        os.makedirs(model_path)
    else:
        print("\nModel save path checked..")

    # Print args
    print("\n" + "x"*50 + "\n\nRunning training with the following parameters: \n")
    for key, value in config.items():
        print(key + ' : ' + str(value))
    print("\n" + "x"*50)

    # # Prepare dataset and iterators for training
    # train_loader, val_loader, test_loader = prepare_lr_training(config)

    seeds = [3, 21, 42, 84, 168]
    # seeds = [3]
    f1_list = []
    macro_f1_list = []
    acc_list = []
    for seed in seeds:
        config['seed']= 21
        print("\nseed= ", seed)
        # Prepare dataset and iterators for training
        prep_data = Prepare_Dataset(config)
        train_loader, val_loader, test_loader = prep_data.prepare_lr_training(config, seed)
        lr_model = LR_Learner(config)
        f1, macro_f1, acc = lr_model.train_main()
        f1_list.append(f1)
        macro_f1_list.append(macro_f1)
        acc_list.append(acc)
        # sys.exit()

    print(f1_list)
    print("\nmean accuracy= ", sum(acc_list)/len(acc_list))
    print("std accuracy = ", stdev(acc_list))
    print("\nmean f1= ", sum(f1_list)/len(f1_list))
    print("std f1= ", stdev(f1_list))
    print("\nmean macro-f1= ", sum(macro_f1_list)/len(macro_f1_list))
    print("std macro-f1= ", stdev(macro_f1_list))

    # try:
    #     lr_model = LR_Learner(config)
    #     lr_model.train_main()
    # except KeyboardInterrupt:
    #     print("Keyboard interrupt by user detected...\nClosing the tensorboard writer!")
    #     print("Best val f1 = ", lr_model.best_val_f1)


#!/usr/bin/python

print(f'Loading {__file__}')

### Andrei - importing channel access function - sprry for that DAMA :) 12 Dec 2019

from bluesky.plan_stubs import one_1d_step, abs_set, wait, sleep
import time
from collections import ChainMap
import bluesky.plans as bp
import matplotlib.ticker as mtick
get_fields = db.get_fields
get_images = db.get_images
get_table = db.get_table

from lmfit import  Model
from lmfit import minimize, Parameters, Parameter, report_fit
from scipy.special import erf

# TODO: create a conda package for it and include to collection profiles
import peakutils

def get_scan(scan_id, debug=False):
    """Get scan from databroker using provided scan id.
from Maksim
    :param scan_id: scan id from bluesky.
    :param debug: a debug flag.
    :return: a tuple of scan and timestamp values.
    """
    scan = db[scan_id]
    #t = datetime.datetime.fromtimestamp(scan['start']['time']).strftime('%Y-%m-%d %H:%M:%S')
    #t = dtt.datetime.fromtimestamp(scan['start']['time']).strftime('%Y-%m-%d %H:%M:%S')
    t='N.A. conflicting with other macro'
    if debug:
        print(scan)
    print('Scan ID: {}  Timestamp: {}'.format(scan_id, t))
    return scan, t

def get_data(scan_id, field='ivu_gap', intensity_field='elm_sum_all', det=None, debug=False):
    """Get data from the scan stored in the table.
from Maksim
    :param scan_id: scan id from bluesky.
    :param field: visualize the intensity vs. this field.
    :param intensity_field: the name of the intensity field.
    :param det: the name of the detector.
    :param debug: a debug flag.
    :return: a tuple of X, Y and timestamp values.
    """
    scan, t = get_scan(scan_id)
    if det:
        imgs = get_images(scan, det)
        im = imgs[-1]
        if debug:
            print(im)

    table = get_table(scan)
    fields = get_fields(scan)

    if debug:
        print(table)
        print(fields)
    x = table[field]
    y = table[intensity_field]

    return x, y, t

def ps(uid='-1',det='default',suffix='default',shift=.5,logplot='off', der  = False ):
    '''
    YG Copied from CHX beamline@March 18, 2018
    function to determine statistic on line profile (assumes either peak or erf-profile)
    calling sequence: uid='-1',det='default',suffix='default',shift=.5)
    det='default' -> get detector from metadata, otherwise: specify, e.g. det='eiger4m_single'
    suffix='default' -> _stats1_total / _sum_all, otherwise: specify, e.g. suffix='_stats2_total'
    shift: scale for peak presence (0.5 -> peak has to be taller factor 2 above background)
    '''
    #import datetime
    #import time
    #import numpy as np
    #from PIL import Image
    #from databroker import db, get_fields, get_images, get_table
    #from matplotlib import pyplot as pltfrom
    #from lmfit import  Model
    #from lmfit import minimize, Parameters, Parameter, report_fit
    #from scipy.special import erf

    # get the scan information:
    if uid == '-1':
        uid=-1
    if det == 'default':
        if db[uid].start.detectors[0] == 'elm' and suffix=='default':
            intensity_field='elm_sum_all'
        elif db[uid].start.detectors[0] == 'elm':
            intensity_field='elm'+suffix
        elif suffix == 'default':
            intensity_field= db[uid].start.detectors[0]+'_stats4_total'
        else:
            intensity_field= db[uid].start.detectors[0]+suffix
    else:
        if det=='elm' and suffix == 'default':
            intensity_field='elm_sum_all'
        elif det=='elm':
            intensity_field = 'elm'+suffix
        elif suffix == 'default':
            intensity_field=det+'_stats4_total'
        else:
            intensity_field=det+suffix

    field = db[uid].start.motors[0]

    #field='dcm_b';intensity_field='elm_sum_all'
    [x,y,t]=get_data(uid,field=field, intensity_field=intensity_field, det=None, debug=False)  #need to re-write way to get data
    x=np.array(x)
    y=np.array(y)
    #print(t)
    if der:
        y = np.diff( y )
        x = x[1:]

    PEAK=x[np.argmax(y)]
    PEAK_y=np.max(y)
    COM=np.sum(x * y) / np.sum(y)

    ### from Maksim: assume this is a peak profile:
    def is_positive(num):
        return True if num > 0 else False

    # Normalize values first:
    ym = (y - np.min(y)) / (np.max(y) - np.min(y)) - shift  # roots are at Y=0

    positive = is_positive(ym[0])
    list_of_roots = []
    for i in range(len(y)):
        current_positive = is_positive(ym[i])
        if current_positive != positive:
            list_of_roots.append(x[i - 1] + (x[i] - x[i - 1]) / (abs(ym[i]) + abs(ym[i - 1])) * abs(ym[i - 1]))
            positive = not positive
    if len(list_of_roots) >= 2:
        FWHM=abs(list_of_roots[-1] - list_of_roots[0])
        CEN=list_of_roots[0]+0.5*(list_of_roots[1]-list_of_roots[0])
        ps.fwhm=FWHM
        ps.cen=CEN
        #return {
        #    'fwhm': abs(list_of_roots[-1] - list_of_roots[0]),
        #    'x_range': list_of_roots,
       #}
    else:    # ok, maybe it's a step function..
        print('no peak...trying step function...')
        ym = ym + shift
        def err_func(x, x0, k=2, A=1,  base=0 ):     #### erf fit from Yugang
            return base - A * erf(k*(x-x0))
        mod = Model(  err_func )
        ### estimate starting values:
        x0=np.mean(x)
        #k=0.1*(np.max(x)-np.m getattr(quadem, f"current{i}").mean_value.kind = "hinted"in(x))
        pars  = mod.make_params( x0=x0, k=200,  A = 1., base = 0. )
        result = mod.fit(ym, pars, x = x )
        CEN=result.best_values['x0']
        FWHM = result.best_values['k']
        ps.cen = CEN
        ps.fwhm = FWHM

    ### re-plot results:
    if logplot=='on':
        plt.close(999)
        plt.figure(999)
        plt.semilogy([PEAK,PEAK],[np.min(y),np.max(y)],'k--',label='PEAK')
        #plt.hold(True)
        plt.semilogy([CEN,CEN],[np.min(y),np.max(y)],'r-.',label='CEN')
        plt.semilogy([COM,COM],[np.min(y),np.max(y)],'g.-.',label='COM')
        plt.semilogy(x,y,'bo-')
        plt.xlabel(field);plt.ylabel(intensity_field)
        plt.legend()
        plt.title('uid: '+str(uid)+' @ '+str(t)+'\nPEAK: '+str(PEAK_y)[:8]+' @ '+str(PEAK)[:8]+'   COM @ '+str(COM)[:8]+ '\n FWHM: '+str(FWHM)[:8]+' @ CEN: '+str(CEN)[:8],size=9)
        plt.show()
    else:
        plt.close(999)
        plt.figure(999)
        plt.plot([PEAK,PEAK],[np.min(y),np.max(y)],'k--',label='PEAK')
        #plt.hold(True)
        plt.plot([CEN,CEN],[np.min(y),np.max(y)],'r-.',label='CEN')
        plt.plot([COM,COM],[np.min(y),np.max(y)],'g.-.',label='COM')
        plt.plot(x,y,'bo-')
        plt.xlabel(field);plt.ylabel(intensity_field)
        plt.legend()
        plt.title('uid: '+str(uid)+' @ '+str(t)+'\nPEAK: '+str(PEAK_y)[:8]+' @ '+str(PEAK)[:8]+'   COM @ '+str(COM)[:8]+ '\n FWHM: '+str(FWHM)[:8]+' @ CEN: '+str(CEN)[:8],size=9)
        plt.show()

    ### assign values of interest as function attributes:
    ps.peak=PEAK
    ps.com=COM
    #return x, y

def set_abs_value(pv_prefix, abs_value):
    """
    Use an absolute value for a PV
    Input
    ---
    pv_prefix:string, the prefix of a pv, e.g., 'XF:12ID1-ES{XtalDfl-Ax:IH}' for XtalDfl IH
    abs_value, float, the absolute value to be set

    Example:
    set_abs_value( 'XF:12ID1-ES{XtalDfl-Ax:IH}', 0 ) #set diff.yv abolute value to 0
    """
    pv_set = EpicsSignal(pv_prefix + 'Mtr.VAL', name="pv_set")
    pv_use_button = EpicsSignal(pv_prefix + 'Mtr.SET', name="pv_use_button")

    yield from bps.mv(pv_use_button, 'Set')

    old_val = pv_set.get()
    yield from bps.mv(pv_set, abs_value)
    yield from bps.mv(pv_use_button, 'Use')

    print('The absolute value of %s was changed from %s to %s.'%(pv_set, old_val, abs_value))


import config
import numpy as np
import pandas as pd
from math import sqrt
from scipy.io import FortranFile

h = list(map(lambda r1, r2: sqrt(r2-r1), config.RADIUS1, config.RADIUS2))

def body_matrix(fileprefix):
    table = pd.read_csv(f"{config.MATPATH}/{fileprefix}.table", delim_whitespace=True)
    M = table.shape[0]
    y = table['residual']
    error = table['error']
    cor1 = table['ellipcor']
    cor2 = table['crustcor']

    K = np.ndarray(shape=(M,config.N))
    f = FortranFile(f"{config.MATPATH}/{fileprefix}.swp.mat", 'r')

    for i in range(M):
        mh = f.read_record(dtype='i4, a4, i4, i4')
        (nrow, sta, yr, dat) = mh[0]
        rin = f.read_record(dtype=np.float32)
        kmat = f.read_record(dtype=np.int32)

        row = np.zeros(config.N)
        for j in range(nrow):
            row[kmat[j]-1] = rin[j] / error.values[i]
            # if is_thickness_to_weight: row[kmat[j]-1] /= h[(kmat[j]-1)//2578]
        K[i][:] = row
        # if i == 1:
            # print(csr_matrix(row).nonzero())

    residuals = (y.values[0:M]+cor1.values[0:M]+cor2.values[0:M])/error.values[0:M]
    residuals = np.reshape(residuals,(len(residuals),1))
    return residuals, K

def surface_matrix(filenames):
    K = np.ndarray(shape=(2578*len(filenames), config.N))
    residuals = []

    for i in range(len(filenames)):
        f = FortranFile(f"{config.MATPATH}/{filenames[i]}", 'r')
        count = 0
        while(1):
            if count == 2578: break
            mh = f.read_record(dtype='f4, f4, i4')
            if not mh: break
            (rem, res, nrow) = mh[0]

            rin = f.read_record(dtype=np.float32)
            kmat = f.read_record(dtype=np.int32)

            residuals.append(rem/(2*res))

            row = np.zeros(config.N)
            for j in range(len(kmat)):
            # for j in range(nrow):
                row[kmat[j]-1] = rin[j] / (2*res)
                # if is_thickness_to_weight: row[kmat[j]-1] /= h[(kmat[j]-1)//2578]
            K[count + 2578*i] = row

            count += 1

    residuals = np.reshape(residuals,(len(residuals),1))
    return residuals, K

def revert_h_weighting(m):
    newm = [ m[j]/h[j//2578] for j in range(len(m))]
    return np.array(newm)

def revert_c_weighting(m, c):
    return m[:,0]/c if len(m.shape)==2 else m[:]/c

def createInputFile(A, filename):

    f = FortranFile(filename, 'w')
    if len(A.shape) == 2:
        (M, N) = A.shape
    else:
        (M, N) = (len(A), 1)
    f.write_record(int(M))
    f.write_record(int(N))
    # f.write_record(np.array([M,N], dtype=int))
    # for val in A.flatten(order='F'):
    #     f.write_record(float(val))
    for i in range(N):
        this_column = A[:,i] if N>1 else A[:]
        this_column = np.array(this_column, dtype=np.float64)
        f.write_record(this_column)

def readOutputFile(filename):
    data = np.fromfile(filename, dtype='i4')
    (M, N) = (data[0], data[1])
    del data
    data = np.fromfile(filename, dtype=np.float64)[1:].reshape(M, N, order='F')
    return data


"""
Created by anonymous author on 9/16/20

Adapted from https://github.com/vlukiyanov/pt-dec
"""

import numpy as np
from scipy.optimize import linear_sum_assignment

import torch
import torch.nn as nn
from torch.nn import Parameter
from typing import Optional, Tuple
from cross_loss_influence.data.scripts.generate_mog_data import MOGDataset
from sklearn.cluster import KMeans

class ClusterAssignment(nn.Module):
    def __init__(
        self,
        cluster_number: int,
        embedding_dimension: int,
        alpha: float = 1.0,
        cluster_centers: Optional[torch.Tensor] = None,
    ) -> None:
        """
        Module to handle the soft assignment, for a description see in 3.1.1. in Xie/Girshick/Farhadi,
        where the Student's t-distribution is used measure similarity between feature vector and each
        cluster centroid.
        :param cluster_number: number of clusters
        :param embedding_dimension: embedding dimension of feature vectors
        :param alpha: parameter representing the degrees of freedom in the t-distribution, default 1.0
        :param cluster_centers: clusters centers to initialise, if None then use Xavier uniform
        """
        super(ClusterAssignment, self).__init__()
        self.embedding_dimension = embedding_dimension
        self.cluster_number = cluster_number
        self.alpha = alpha
        if cluster_centers is None:
            initial_cluster_centers = torch.zeros(
                self.cluster_number, self.embedding_dimension, dtype=torch.float
            )
            # initial_cluster_centers = torch.rand((self.cluster_number, self.embedding_dimension), dtype=torch.float)
            nn.init.xavier_uniform_(initial_cluster_centers)
        else:
            initial_cluster_centers = cluster_centers
        self.cluster_centers = Parameter(initial_cluster_centers)

    def forward(self, batch: torch.Tensor) -> torch.Tensor:
        """
        Compute the soft assignment for a batch of feature vectors, returning a batch of assignments
        for each cluster.
        :param batch: FloatTensor of [batch size, embedding dimension]
        :return: FloatTensor [batch size, number of clusters]
        """
        norm_squared = torch.sum((batch.unsqueeze(1) - self.cluster_centers) ** 2, 2)
        numerator = 1.0 / (1.0 + (norm_squared / self.alpha))
        power = float(self.alpha + 1) / 2
        numerator = numerator ** power
        return numerator / torch.sum(numerator, dim=1, keepdim=True)

class DEC(nn.Module):
    def __init__(
        self,
        input_dim: int,
        cluster_number: int,
        hidden_dimension: int,
        alpha: float = 1.0,
    ):
        """
        Module which holds all the moving parts of the DEC algorithm, as described in
        Xie/Girshick/Farhadi; this includes the AutoEncoder stage and the ClusterAssignment stage.
        :param cluster_number: number of clusters
        :param hidden_dimension: hidden dimension, output of the encoder
        :param encoder: encoder to use
        :param alpha: parameter representing the degrees of freedom in the t-distribution, default 1.0
        """
        super(DEC, self).__init__()
        self.embedding_layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dimension),
            nn.ReLU(),
            nn.Linear(hidden_dimension, hidden_dimension)
        )
        self.hidden_dimension = hidden_dimension
        self.cluster_number = cluster_number
        self.alpha = alpha
        self.assignment = ClusterAssignment(
            cluster_number, self.hidden_dimension, alpha
        )

    def encode(self, data_in: torch.Tensor) -> torch.Tensor:
        return self.embedding_layers(data_in)

    def forward(self, data_in: torch.Tensor) -> torch.Tensor:
        """
        Compute the cluster assignment using the ClusterAssignment after running the batch
        through the encoder part of the associated AutoEncoder module.
        :param batch: [batch size, embedding dimension] FloatTensor
        :return: [batch size, number of clusters] FloatTensor
        """
        embedded_data = self.embedding_layers(data_in)
        return self.assignment(embedded_data)

def init_model(model_in: torch.nn.Module,
               dataset_in: Tuple[torch.Tensor, torch.Tensor] = None,
               device: str = 'cpu',
               save: bool = True,
               save_dir=None):
    import os
    from cross_loss_influence.config import MODEL_SAVE_DIR
    if save_dir is not None:
        model_save_path = save_dir
    else:
        model_save_path = MODEL_SAVE_DIR
    model_fn = os.path.join(model_save_path, 'mog_model_init.pth.tar')
    if save:
        ds = MOGDataset(dataset_in)
        num_clusters = len(np.unique(dataset_in[1]))

        kmeans = KMeans(n_clusters=model_in.cluster_number, n_init=num_clusters)
        model_in.train()
        features = []
        actual = []
        # form initial cluster centres
        for data, label in zip(dataset_in[0], dataset_in[1]):
            if ds.kl:
                label = torch.argmax(label)
            actual.append(label.item())
            data = data.to(device)
            features.append(model_in.encode(data).detach().cpu())
        predicted = kmeans.fit_predict(torch.stack(features).numpy())
        cluster_centers = torch.tensor(
            kmeans.cluster_centers_, dtype=torch.float, requires_grad=True, device=device
        )
        cluster_centers = cluster_centers + torch.randn(cluster_centers.size()).to(device)/2.0
        with torch.no_grad():
            model_in.state_dict()["assignment.cluster_centers"].copy_(cluster_centers)
        torch.save({
            'model':model_in.state_dict(),
        }, model_fn)
    else:
        chk = torch.load(model_fn, map_location='cpu')
        model_in.load_state_dict(chk['model'])
    return model_in

def cluster_accuracy(y_true, y_predicted, cluster_number: Optional[int] = None):
    """
    Calculate clustering accuracy after using the linear_sum_assignment function in SciPy to
    determine reassignments.
    :param y_true: list of true cluster numbers, an integer array 0-indexed
    :param y_predicted: list  of predicted cluster numbers, an integer array 0-indexed
    :param cluster_number: number of clusters, if None then calculated from input
    :return: reassignment dictionary, clustering accuracy
    """
    if cluster_number is None:
        cluster_number = (
            max(y_predicted.max(), y_true.max()) + 1
        )  # assume labels are 0-indexed
    count_matrix = np.zeros((cluster_number, cluster_number), dtype=np.int64)
    for i in range(y_predicted.size):
        count_matrix[y_predicted[i], y_true[i]] += 1

    row_ind, col_ind = linear_sum_assignment(count_matrix.max() - count_matrix)
    reassignment = dict(zip(row_ind, col_ind))
    accuracy = count_matrix[row_ind, col_ind].sum() / y_predicted.size
    return reassignment, accuracy

def target_distribution(batch: torch.Tensor) -> torch.Tensor:
    """
    Compute the target distribution p_ij, given the batch (q_ij), as in 3.1.3 Equation 3 of
    Xie/Girshick/Farhadi; this is used the KL-divergence loss function.
    :param batch: [batch size, number of clusters] Tensor of dtype float
    :return: [batch size, number of clusters] Tensor of dtype float
    """
    weight = (batch ** 2) / torch.sum(batch, 0)
    return (weight.t() / torch.sum(weight, 1)).t()


<reponame>omangin/multimodal<filename>multimodal/lib/sklearn_utils.py
# Authors: <NAME> <<EMAIL>>
#          <NAME> <<EMAIL>>
#          <NAME> <<EMAIL>>
#          <NAME> <<EMAIL>>
#          <NAME>
#          <NAME>
#          <NAME>

# These files where extracted from the scikit-learn project
# (http://scikit-learn.org/).

# Scikit-learn is a Python module for machine learning built on top
# of SciPy and distributed under the 3-Clause BSD license.

# New BSD License
#
# Copyright (c) 2007--2013 The scikit-learn developers.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   a. Redistributions of source code must retain the above copyright notice,
#      this list of conditions and the following disclaimer.
#   b. Redistributions in binary form must reproduce the above copyright
#      notice, this list of conditions and the following disclaimer in the
#      documentation and/or other materials provided with the distribution.
#   c. Neither the name of the Scikit-learn Developers  nor the names of
#      its contributors may be used to endorse or promote products
#      derived from this software without specific prior written
#      permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
# DAMAGE.

import numpy as np
from scipy import sparse

# Source utils/fixes.py

def safe_copy(X):
        # Copy, but keep the order
        return np.copy(X, order='K')

# Source utils/validation.py

def assert_all_finite(X):
    """Throw a ValueError if X contains NaN or infinity.

    Input MUST be an np.ndarray instance or a scipy.sparse matrix."""

    # First try an O(n) time, O(1) space solution for the common case that
    # there everything is finite; fall back to O(n) space np.isfinite to
    # prevent false positives from overflow in sum method.
    if X.dtype.char in np.typecodes['AllFloat'] and not np.isfinite(X.sum()) \
      and not np.isfinite(X.data if sparse.issparse(X) else X).all():
            raise ValueError("array contains NaN or infinity")

def array2d(X, dtype=None, order=None, copy=False):
    """Returns at least 2-d array with data from X"""
    if sparse.issparse(X):
        raise TypeError('A sparse matrix was passed, but dense data '
                        'is required. Use X.todense() to convert to dense.')
    X_2d = np.asarray(np.atleast_2d(X), dtype=dtype, order=order)
    if X is X_2d and copy:
        X_2d = safe_copy(X_2d)
    return X_2d

def atleast2d_or_csr(X, dtype=None, order=None, copy=False):
    """Like numpy.atleast_2d, but converts sparse matrices to CSR format

    Also, converts np.matrix to np.ndarray.
    """
    if sparse.issparse(X):
        # Note: order is ignored because CSR matrices hold data in 1-d arrays
        if dtype is None or X.dtype == dtype:
            X = X.tocsr()
        else:
            X = sparse.csr_matrix(X, dtype=dtype)
    else:
        X = array2d(X, dtype=dtype, order=order, copy=copy)
    assert_all_finite(X)
    return X

# Source utils/extmath.py

def safe_sparse_dot(a, b, dense_output=False):
    """Dot product that handle the sparse matrix case correctly"""
    if sparse.issparse(a) or sparse.issparse(b):
        ret = a * b
        if dense_output and hasattr(ret, "toarray"):
            ret = ret.toarray()
        return ret
    else:
        return np.dot(a, b)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 27 09:23:27 2021

@author: gianni
"""

#for testing with pytest (just type "pytest" in the terminal)

import ionisation_photodissociation as ip
import numpy as np
import pytest
from scipy import constants
import os
import itertools

class Test_ISF_scalings():

    fields = [ip.DraineISF,ip.HabingField]

    def test_scaling(self):
        scalings = [lambda wavelength: 2, lambda wavelength: wavelength**2]
        for field in self.fields:
            unscaled_field = field()
            for scaling in scalings:
                scaled_field = field(scaling=scaling)
                wavelength = unscaled_field.lambda_grid
                expected_flux = unscaled_field.flux(wavelength=wavelength)\
                                              *scaling(wavelength=wavelength)
                assert np.all(expected_flux == scaled_field.flux(wavelength=wavelength))

class Test_DraineISF():

    def test_lambda_check(self):
        isf = ip.DraineISF()
        too_small_lambda = np.array((0.1*isf.lambda_min,0.5*isf.lambda_min))
        too_large_lambda = np.array((1.01*isf.lambda_max,2*isf.lambda_max))
        for lamb in (too_small_lambda,too_large_lambda):
            assert np.all(isf.flux(lamb) == 0)

class Test_Habing():

    def test_interpolation(self):
        field = ip.HabingField()
        out_of_range_lambda = np.array([90,2000])*constants.nano
        for lamb in out_of_range_lambda:
            assert field.flux(lamb) == 0
        test_lamb = 650.2*constants.nano
        test_energy = constants.h*constants.c/test_lamb
        expected_flux = 3.90754e+06*test_energy/constants.centi**2/constants.nano
        assert np.isclose(field.flux(650.2*constants.nano),expected_flux,rtol=1e-6,atol=0)

def generate_test_atm():
    atm = ip.StellarAtmosphere()
    atm.lambda_grid = np.array([1,2.2,3.3,4.6])
    atm.modelflux = np.array((10.2,20,40.4,90))
    atm.ref_distance = 2
    return atm

class Test_StellarAtmosphere():

    atm = generate_test_atm()

    def test_lamb_limits(self):
        too_large_lamb = np.array((5,7))
        too_small_lamb = np.array((0.1,0.9))
        for lamb in (too_large_lamb,too_small_lamb):
            assert np.all(self.atm.flux(wavelength=lamb,distance=1)==0)

    def test_flux_interpolation(self):
        assert self.atm.flux(wavelength=self.atm.lambda_grid[0],
                             distance=self.atm.ref_distance)\
                                                   == self.atm.modelflux[0]
        assert self.atm.flux(wavelength=self.atm.lambda_grid[0],
                             distance=2*self.atm.ref_distance)\
                                                   == self.atm.modelflux[0]/4

    def test_luminosity(self):
        expected_lum = np.trapz(self.atm.modelflux,self.atm.lambda_grid)\
                                *4*np.pi*self.atm.ref_distance**2
        assert self.atm.luminosity() == expected_lum

    def test_scaling(self):
        scalings = [lambda wavelength: 3.5,
                    lambda wavelength: wavelength**2]
        for scaling in scalings:
            test_atm = generate_test_atm()
            original_model_flux = test_atm.modelflux.copy()
            unscaled_flux = test_atm.flux(wavelength=test_atm.lambda_grid,
                                          distance=test_atm.ref_distance)
            test_atm._scale_spectrum(scaling=scaling)
            scaling_factors = scaling(wavelength=test_atm.lambda_grid)
            assert np.all(original_model_flux*scaling_factors==test_atm.modelflux)
            scaled_flux = test_atm.flux(wavelength=test_atm.lambda_grid,
                                        distance=test_atm.ref_distance)
            assert np.all(unscaled_flux*scaling_factors==scaled_flux)
            expected_luminosity = np.trapz(unscaled_flux*scaling_factors,
                                           test_atm.lambda_grid)\
                                 *4*np.pi*test_atm.ref_distance**2
            assert expected_luminosity == test_atm.luminosity()

    def test_plot_model(self):
        self.atm.plot_model()

    def test_file_writing(self):
        test_filepath = 'test.npz'
        dist = 1*constants.au
        self.atm.write_modelflux_to_file(filepath=test_filepath,distance=dist)
        data = np.load(test_filepath)
        assert np.all(data['wavelength']==self.atm.lambda_grid)
        assert np.all(data['flux']
                        ==self.atm.flux(wavelength=data['wavelength'],distance=dist))
        os.remove(test_filepath)

def general_test_flux_scaling(unscaled_atm,atm_cls,atm_kwargs):
    scalings = [lambda wavelength: 3,lambda wavelength: wavelength/3]
    dist = 1*constants.au
    for scaling in scalings:
        #make a copy in order not to modify atm_kwargs
        scaled_atm_kwargs = atm_kwargs.copy()
        scaled_atm_kwargs['scaling'] = scaling
        scaled_atm = atm_cls(**scaled_atm_kwargs)
        scaling_factors = scaling(wavelength=scaled_atm.lambda_grid)
        assert np.all(unscaled_atm.modelflux*scaling_factors == scaled_atm.modelflux)
        flux_kwargs = {'wavelength':scaled_atm.lambda_grid,'distance':dist}
        unscaled_flux = unscaled_atm.flux(**flux_kwargs)
        scaled_flux = scaled_atm.flux(**flux_kwargs)
        assert np.allclose(scaling_factors*unscaled_flux,scaled_flux,rtol=1e-6,atol=0)
        expected_lum = np.trapz(unscaled_flux*scaling_factors,unscaled_atm.lambda_grid)\
                                                               *4*np.pi*dist**2
        assert np.isclose(expected_lum,scaled_atm.luminosity(),rtol=1e-6,atol=0)

class Test_ATLAS():

    test_grid = np.array((1,2,3,4))
    R_Sun = 6.955e8
    template_init_kwargs = {'Teff':5780,'metallicity':-0.6,'logg':4.3,'Rstar':R_Sun,
                            'calibration_spec':None,'verbose':True}
    template_atm = ip.ATLASModelAtmosphere(**template_init_kwargs)

    def test_grid_checking(self):
        for value in (0,5):
            with pytest.raises(AssertionError):
                ip.ATLASModelAtmosphere.assert_within_grid(value=value,
                                                           grid=self.test_grid)

    def test_grid_searching(self):
        values = np.array((0,2,2.9,7))
        expected_grid_values = np.array((1,2,3,4))
        for v,expec_v in zip(values,expected_grid_values):
            grid_value = ip.ATLASModelAtmosphere.get_closest_grid_value(
                                              value=v,grid=self.test_grid)
            assert grid_value == expec_v

    def test_bad_grid_values(self):
        bad_init_kwargs = {'Teff':[2000,1e5],'metallicity':[-3,3],'logg':[-1,7]}
        for kw,bad_values in bad_init_kwargs.items():
            for v in bad_values:
                kwargs = self.template_init_kwargs.copy()
                kwargs[kw] = v
                with pytest.raises(AssertionError):
                    ip.ATLASModelAtmosphere(**kwargs)

    def test_initialisation(self):
        #test all None:
        kwargs = self.template_init_kwargs.copy()
        kwargs['Rstar'] = None
        with pytest.raises(ValueError):
            ip.ATLASModelAtmosphere(**kwargs)
        #test 2x not None
        for key in ('calibration_spec','obs_luminosity'):
            kwargs = self.template_init_kwargs.copy()
            assert kwargs['Rstar'] is not None
            kwargs[key] = 1
            with pytest.raises(AssertionError):
                ip.ATLASModelAtmosphere(**kwargs)
        #test 3x not None
        kwargs = self.template_init_kwargs.copy()
        assert kwargs['Rstar'] is not None
        kwargs['calibration_spec'] = 1
        kwargs['obs_luminosity'] = 1
        with pytest.raises(AssertionError):
            ip.ATLASModelAtmosphere(**kwargs)

    def test_adopted_grid_values(self):
        assert self.template_atm.Teff == 5750
        assert self.template_atm.metallicity == -0.5
        assert self.template_atm.logg == 4.5

    def test_positive_metallicity(self):
        kwargs = self.template_init_kwargs.copy()
        kwargs['metallicity'] = 0.3
        ip.ATLASModelAtmosphere(**kwargs)

    def test_distance_scaling(self):
        test_lamb = 700*constants.nano
        assert self.template_atm.flux(wavelength=test_lamb,distance=1*constants.au)\
                == 4*self.template_atm.flux(wavelength=test_lamb,distance=2*constants.au)

    def test_RJ_extrapolation(self):
        atm = ip.ATLASModelAtmosphere(**self.template_init_kwargs)
        test_lamb = 2*constants.milli
        dist = 1*constants.au
        RJ_flux = atm.flux(wavelength=test_lamb,distance=dist)
        expected_flux = atm.modelflux[-1]*(atm.lambda_grid[-1]/test_lamb)**4\
                         *(atm.ref_distance/dist)**2
        assert np.isclose(RJ_flux,expected_flux,rtol=1e-3,atol=0)

    def test_calibration_with_luminosity(self):
        test_lum = 3
        test_kwargs = self.template_init_kwargs.copy()
        test_kwargs["Rstar"] = None
        test_kwargs['obs_luminosity'] = test_lum
        test_atm = ip.ATLASModelAtmosphere(**test_kwargs)
        assert test_lum == test_atm.luminosity()

    def test_calibration_with_spec(self):
        ref_atm = ip.ATLASModelAtmosphere(**self.template_init_kwargs)
        calibration_spec = {'wave':ref_atm.lambda_grid,'flux':ref_atm.modelflux/2,
                            'ref_distance':ref_atm.ref_distance}
        kwargs = self.template_init_kwargs.copy()
        kwargs['Rstar'] = None
        kwargs['calibration_spec'] = calibration_spec
        atm = ip.ATLASModelAtmosphere(**kwargs)
        expected_scaling = 0.5
        assert np.isclose(atm.spec_calibration_scaling,expected_scaling,rtol=1e-4,atol=0)

    def test_flux_scaling(self):
        general_test_flux_scaling(unscaled_atm=self.template_atm,
                                  atm_cls=ip.ATLASModelAtmosphere,
                                  atm_kwargs=self.template_init_kwargs)

    def test_plot_model(self):
        atm = ip.ATLASModelAtmosphere(**self.template_init_kwargs)
        atm.plot_model()

class Test_betaPic():

    betaPic = ip.betaPicObsSpectrum()

    def test_flux_reading(self):
        test_lamb = 1327.75*constants.angstrom
        expected_flux = 10**(1.65467e-001)*constants.erg/constants.centi**2/constants.angstrom
        flux = self.betaPic.flux(wavelength=test_lamb,distance=1*constants.au)
        assert expected_flux == flux

    def test_dilution(self):
        betaPic_diluted = ip.betaPicObsSpectrum(dilution=3)
        assert self.betaPic.lambda_grid.size > betaPic_diluted.lambda_grid.size
        betaPic_int = np.trapz(self.betaPic.modelflux,self.betaPic.lambda_grid)
        betaPic_diluted_int = np.trapz(betaPic_diluted.modelflux,
                                       betaPic_diluted.lambda_grid)
        assert np.isclose(betaPic_int,betaPic_diluted_int,rtol=1e-3,atol=0)

    def test_scaling(self):
        general_test_flux_scaling(
                     unscaled_atm=self.betaPic,atm_cls=ip.betaPicObsSpectrum,
                     atm_kwargs={})

    def test_plot_model(self):
        self.betaPic.plot_model()

class Test_pd_cross_section():

    cs = ip.PhotodissociationCrossSection('CO.hdf5')

    def test_crosssection(self):
        too_small_lamb = 5e-9
        too_large_lamb = 1.8e-7
        for lamb in (too_small_lamb,too_large_lamb):
            assert self.cs.crosssection(lamb) == 0

    def test_plot(self):
        self.cs.plot()

def test_Osterbrock_ionisation_and_recombination():
    for element in ip.ionisation_potential.keys():
        ionisation = ip.OsterbrockIonisationCrossSection(element=element)
        nu_min = ip.ionisation_potential[element]/constants.h
        lamb_max = constants.c/nu_min
        assert ionisation.crosssection(wavelength=10*lamb_max) == 0
        recomb = ip.OsterbrockRecombination(element=element)
        too_small_T, too_large_T = 7000, 16000
        for T in (too_small_T,too_large_T):
            with pytest.raises(AssertionError):
                recomb.recombination_coeff(T)

def test_Nahar_ionisation_and_recombination():
    for element,io_potential in ip.ionisation_potential.items():
        ionisation = ip.NaharIonisationCrossSection(element=element)
        non_ionising_wavelength = 1.1*constants.h*constants.c/io_potential
        assert ionisation.crosssection(non_ionising_wavelength) == 0
        recomb = ip.NaharRecombination(element=element)
        too_small_T, too_large_T = 1, 10e10
        for T in (too_small_T,too_large_T):
            with pytest.raises(AssertionError):
                recomb.recombination_coeff(T)

solar_atmosphere = ip.ATLASModelAtmosphere(
                          Teff=5780,metallicity=0.01,logg=4.43,Rstar=6.955e8)

def pd_crosssections_iterator():
    for mol in ('CO','H2O','OH'):
        yield ip.PhotodissociationCrossSection(f'{mol}.hdf5')
def io_crosssections_iterator():
    for element in ('C','O'):
        for cs in (ip.NaharIonisationCrossSection(element=element),
                   ip.OsterbrockIonisationCrossSection(element=element)):
            yield cs
isf = ip.DraineISF()

class Test_rate():

    def rate_iterator(self):
        for crosssection in itertools.chain(pd_crosssections_iterator(),
                                            io_crosssections_iterator()):
            for atm in (solar_atmosphere,None):
                rate = ip.Rate(stellar_atmosphere=atm,crosssection=crosssection)
                yield crosssection,atm,rate

    def test_lambda_grid_construction(self):
        for crosssection,atm,rate in self.rate_iterator():
            assert np.all(np.diff(rate.lambda_grid)>0)
            assert np.all(np.unique(rate.lambda_grid) == rate.lambda_grid)
            assert np.all(np.isin(isf.lambda_grid,rate.lambda_grid))
            if hasattr(crosssection,'lambda_grid'):
                assert np.all(np.isin(crosssection.lambda_grid,rate.lambda_grid))
            if atm is not None:
                assert np.all(np.isin(atm.lambda_grid,rate.lambda_grid))
            if not hasattr(crosssection,'lambda_grid') and atm is None:
                assert np.all(rate.lambda_grid==isf.lambda_grid)

    def test_stellar_rate_distance_scaling(self):
        for crosssection,atm,rate in self.rate_iterator():
            if atm is not None:
                assert rate.stellar_rate(distance=1*constants.au)\
                              == 4*rate.stellar_rate(distance=2*constants.au)

    def test_stellar_rate_flux_scaling(self):
        dist = 1*constants.au
        scaling = lambda wavelength: 3.5
        atm_kwargs = {'Teff':6000,'metallicity':0.01,'logg':4.2,'Rstar':7e8}
        for crosssection in itertools.chain(pd_crosssections_iterator(),
                                            io_crosssections_iterator()):
            atm = ip.ATLASModelAtmosphere(**atm_kwargs)
            original_rate = ip.Rate(stellar_atmosphere=atm,crosssection=crosssection)
            original_stellar_rate = original_rate.stellar_rate(distance=dist)
            scaled_atm = ip.ATLASModelAtmosphere(scaling=scaling,**atm_kwargs)
            scaled_rate = ip.Rate(stellar_atmosphere=scaled_atm,crosssection=crosssection)
            scaled_stellar_rate = scaled_rate.stellar_rate(distance=dist)
            scaling_factors = scaling(wavelength=atm.lambda_grid)
            assert np.isclose(scaled_stellar_rate,scaling_factors*original_stellar_rate,
                              rtol=1e-6,atol=0)

    def test_ISF_scaling(self):
        scalar_scaling = 6.7
        ISF_scaling = lambda wavelength: scalar_scaling
        for crosssection,atm,rate in self.rate_iterator():
            unscaled_isf_rate = rate.isf_rate
            scaled_rate = ip.Rate(stellar_atmosphere=atm,crosssection=crosssection,
                                  ISF_scaling=ISF_scaling)
            assert np.isclose(scalar_scaling*unscaled_isf_rate,scaled_rate.isf_rate,
                              rtol=1e-6,atol=0)

    def test_rate(self):
        for crosssection,atm,rate in self.rate_iterator():
            rate.isf_rate
            if atm is not None:
                rate.stellar_rate(distance=1*constants.au)

    def test_plot_print(self):
        for crosssection,atm,rate in self.rate_iterator():
            rate = ip.Rate(stellar_atmosphere=solar_atmosphere,
                           crosssection=crosssection)
            rate.plot_ISF_rate()
            dist = 1*constants.au
            rate.plot_stellar_rate(distance=dist)
            rate.print_rates(distance=dist)

class Test_pd_and_io_rate():

    def rate_iterator(self):
        for crosssection in pd_crosssections_iterator():
            for rate_cls in (ip.PhotodissociationRate,ip.IonisationRate):
                rate = rate_cls(stellar_atmosphere=solar_atmosphere,
                                crosssection=crosssection)
            yield crosssection,rate

    def test_total_rate(self):
        for crosssection,rate in self.rate_iterator():
            dist = 1*constants.au
            stellar_rate = rate.stellar_rate(distance=dist)
            if isinstance(rate,ip.PhotodissociationRate):
                CR_rate = 0
            elif isinstance(rate,ip.IonisationRate):
                CR_rate = ip.IonisationRate.CR_ionisation_rate
            tot_rate = rate.isf_rate + stellar_rate + CR_rate
            assert rate.total_rate(distance=dist) == tot_rate

    def test_print(self):
        for crosssection,rate in self.rate_iterator():
            rate = ip.PhotodissociationRate(stellar_atmosphere=solar_atmosphere,
                                            crosssection=crosssection)
            rate.print_rates(distance=1*constants.au)

class Test_ionisation_balance():

    elements = ('C','O')
    recombinations = [ip.NaharRecombination,ip.OsterbrockRecombination]
    io_crosssections = [ip.NaharIonisationCrossSection,
                        ip.OsterbrockIonisationCrossSection]

    def io_balance_iterator(self):
        for element in self.elements:
            for recomb,io_cs in itertools.product(self.recombinations,
                                                  self.io_crosssections):
                recombination = recomb(element=element)
                io_rate = ip.IonisationRate(crosssection=io_cs(element=element),
                                            ISF_scaling=lambda wavelength: 1,
                                            stellar_atmosphere=solar_atmosphere)
                balance = ip.IonisationBalance(ionisation_rate=io_rate,
                                               recombination=recombination)
                yield balance

    def test_balance_computations(self):
        #don't know how to test this properly, so mostly just call all functions
        #to make sure that they at least execute without throwing errors
        n = 500/constants.centi**3
        n_ion = 300/constants.centi**3
        n_neutral = n-n_ion
        n_e = 100/constants.centi**3
        distance = 100*constants.au
        T = 8000
        for balance in self.io_balance_iterator():
            balance.get_ionisation_rate(distance=distance)
            balance.get_recomb_coeff(T=T)
            balance.determine_n_neutral(n_ion=n_ion,n_e=n_e,distance=distance,T=T)
            n_neu_e_from_ion = balance.determine_n_neutral_e_from_ion(
                                 n_ion=n_ion,distance=distance,T=T)
            n_neu_e_from_ion_check = balance.determine_n_neutral(
                                       n_ion=n_ion,n_e=n_ion,distance=distance,T=T)
            assert n_neu_e_from_ion == n_neu_e_from_ion_check
            balance.determine_n_ion(n_neutral=n_neutral,n_e=n_e,distance=distance,T=T)
            n_ion_e_from_ion = balance.determine_n_ion_e_from_ion(
                                  n_neutral=n_neutral,distance=distance,T=T)
            n_ion_e_from_ion_check = balance.determine_n_ion(
                                      n_neutral=n_neutral,n_e=n_ion_e_from_ion,
                                      distance=distance,T=T)
            assert np.isclose(n_ion_e_from_ion,n_ion_e_from_ion_check,rtol=1e-6,atol=0)
            balance.determine_ionisation_balance(n=n,n_e=n_e,distance=distance,T=T)
            balance.determine_ionisation_balance_e_from_ion(n=n,distance=distance,T=T)

    def test_ionisation_fraction(self):
        for balance in self.io_balance_iterator():
            assert balance.ionisation_fraction(n_neutral=0.3,n_ion=0.7) == 0.7


from pyqm import *
import os, sys
import time

### Parameters:

nx = 500 #mesh size
sigma = .4 #gaussian standard deviation
dx = .01 #spatial resolution
X = create1Dmesh(dx,nx)
X = X.reshape(len(X),-1)
T = 10 #number of timesteps
dt = 1/(2/(dx**2)) #.08
#python toolbox nibot crank-nicolson  ilikear

## Hamiltonian
def potential(x): return renormalize(x**2)
Hamiltonian = -laplacematrix(dx,nx) #+ X*X

### Time Evolution
def propagate(H,dt):
    '''H must be a matrix in K basis'''
    #return (1 - 1j * dt * H)
    from scipy.linalg import expm
    return matrix(expm(-1j*H))#* dt * H)
    #return inv((1 + .5j * dt * H)) * (1 - .5j * dt * H)

def evolve(initial, t):
    U = propagate(Hamiltonian,dt)
    allvec = [initial]
    for i in range(t):
        #allvec.append(renormalize(U*allvec[i]))
        allvec.append(U*allvec[i])
    return allvec
    #return [initial, U*initial, U*U*initial, U*U*U*initial]

### Actions

Initial = renormalize(gaussian(X,sigma))
#Initial = renormalize(sin(X))

ion()
amplitudes = map(abs,evolve(Initial,T))
animate1D(X,amplitudes)
#waitforbuttonpress()

#createvideo(amplitudes,plotsurface)