# -*- coding: utf-8 -*-
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations
import warnings

import numpy as np
from scipy import stats
from sklearn.utils.validation import check_is_fitted, check_consistent_length, check_array
from tqdm.auto import tqdm

from .base import HubnessReduction

class MutualProximity(HubnessReduction):
    """ Hubness reduction with Mutual Proximity [1]_.

    Parameters
    ----------
    method: 'normal' or 'empiric', default = 'normal'
        Model distance distribution with 'method'.

        - 'normal' or 'gaussi' model distance distributions with independent Gaussians (fast)
        - 'empiric' or 'exact' model distances with the empiric distributions (slow)

    verbose: int, default = 0
        If verbose > 0, show progress bar.

    References
    ----------
    .. [1] <NAME>., <NAME>., <NAME>., & <NAME>. (2012).
           Local and global scaling reduce hubs in space. The Journal of Machine
           Learning Research, 13(1), 2871–2902.
    """

    def __init__(self, method: str = 'normal', verbose: int = 0, **kwargs):
        super().__init__(**kwargs)
        self.method = method
        self.verbose = verbose

    def fit(self, neigh_dist, neigh_ind, X=None, assume_sorted=None, *args, **kwargs) -> MutualProximity:
        """ Fit the model using neigh_dist and neigh_ind as training data.

        Parameters
        ----------
        neigh_dist: np.ndarray, shape (n_samples, n_neighbors)
            Distance matrix of training objects (rows) against their
            individual k nearest neighbors (columns).

        neigh_ind: np.ndarray, shape (n_samples, n_neighbors)
            Neighbor indices corresponding to the values in neigh_dist.

        X: ignored

        assume_sorted: ignored
        """
        # Check equal number of rows and columns
        check_consistent_length(neigh_ind, neigh_dist)
        check_consistent_length(neigh_ind.T, neigh_dist.T)
        check_array(neigh_dist, force_all_finite=False)
        check_array(neigh_ind)

        self.n_train = neigh_dist.shape[0]

        if self.method in ['exact', 'empiric']:
            self.method = 'empiric'
            self.neigh_dist_train_ = neigh_dist
            self.neigh_ind_train_ = neigh_ind
        elif self.method in ['normal', 'gaussi']:
            self.method = 'normal'
            self.mu_train_ = np.nanmean(neigh_dist, axis=1)
            self.sd_train_ = np.nanstd(neigh_dist, axis=1, ddof=0)
        else:
            raise ValueError(f'Mutual proximity method "{self.method}" not recognized. Try "normal" or "empiric".')

        return self

    def transform(self, neigh_dist, neigh_ind, X=None, assume_sorted=None, *args, **kwargs):
        """ Transform distance between test and training data with Mutual Proximity.

        Parameters
        ----------
        neigh_dist: np.ndarray
            Distance matrix of test objects (rows) against their individual
            k nearest neighbors among the training data (columns).

        neigh_ind: np.ndarray
            Neighbor indices corresponding to the values in neigh_dist

        X: ignored

        assume_sorted: ignored

        Returns
        -------
        hub_reduced_dist, neigh_ind
            Mutual Proximity distances, and corresponding neighbor indices

        Notes
        -----
        The returned distances are NOT sorted! If you use this class directly,
        you will need to sort the returned matrices according to hub_reduced_dist.
        Classes from :mod:`skhubness.neighbors` do this automatically.
        """
        check_is_fitted(self, ['mu_train_', 'sd_train_', 'neigh_dist_train_', 'neigh_ind_train_'], all_or_any=any)
        check_array(neigh_dist, force_all_finite='allow-nan')
        check_array(neigh_ind)

        n_test, n_indexed = neigh_dist.shape

        if n_indexed == 1:
            warnings.warn(f'Cannot perform hubness reduction with a single neighbor per query. '
                          f'Skipping hubness reduction, and returning untransformed distances.')
            return neigh_dist, neigh_ind

        hub_reduced_dist = np.empty_like(neigh_dist)

        # Show progress in hubness reduction loop
        disable_tqdm = False if self.verbose else True
        range_n_test = tqdm(range(n_test),
                            desc=f'MP ({self.method})',
                            disable=disable_tqdm,
                            )

        # Calculate MP with independent Gaussians
        if self.method == 'normal':
            mu_train = self.mu_train_
            sd_train = self.sd_train_
            for i in range_n_test:
                j_mom = neigh_ind[i]
                mu = np.nanmean(neigh_dist[i])
                sd = np.nanstd(neigh_dist[i], ddof=0)
                p1 = stats.norm.sf(neigh_dist[i, :], mu, sd)
                p2 = stats.norm.sf(neigh_dist[i, :], mu_train[j_mom], sd_train[j_mom])
                hub_reduced_dist[i, :] = (1 - p1 * p2).ravel()
        # Calculate MP empiric (slow)
        elif self.method == 'empiric':
            max_ind = self.neigh_ind_train_.max()
            for i in range_n_test:
                dI = neigh_dist[i, :][np.newaxis, :]  # broadcasted afterwards
                dJ = np.zeros((dI.size, n_indexed))
                for j in range(n_indexed):
                    tmp = np.zeros(max_ind + 1) + (self.neigh_dist_train_[neigh_ind[i, j], -1] + 1e-6)
                    tmp[self.neigh_ind_train_[neigh_ind[i, j]]] = self.neigh_dist_train_[neigh_ind[i, j]]
                    dJ[j, :] = tmp[neigh_ind[i]]
                # dJ = self.neigh_dist_train_[neigh_ind[i], :n_indexed]
                d = dI.T
                hub_reduced_dist[i, :] = 1. - (np.sum((dI > d) & (dJ > d), axis=1) / n_indexed)
        else:
            raise ValueError(f"Internal: Invalid method {self.method}.")

        # Return the hubness reduced distances
        # These must be sorted downstream
        return hub_reduced_dist, neigh_ind


import numpy as np
import open3d as o3d
from transformations import *
import os,sys,yaml,copy,pickle,time,cv2,socket,argparse,inspect,trimesh,operator,gzip,re,random,torch
import resource
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))
from scipy.spatial import cKDTree
code_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(code_dir)
sys.path.append("{}/../".format(code_dir))
sys.path.append("{}/ss-pybullet".format(code_dir))
from dexnet.grasping.grasp import ParallelJawPtGrasp3D
from autolab_core import YamlConfig
from dexnet.grasping.grasp_sampler import PointConeGraspSampler,NocsTransferGraspSampler
from PIL import Image
from Utils import *
from data_reader import *
from pointnet2 import *
from aligning import *
import PointGroup.data.dataset_seg as dataset_seg
from PointGroup.model.pointgroup.pointgroup import PointGroup
import PointGroup.lib.pointgroup_ops.functions.pointgroup_ops as pointgroup_ops
import PointGroup.util.config as config_pg
import spconv
from spconv.modules import SparseModule
from dataset_nunocs import NunocsIsolatedDataset
from dataset_grasp import GraspDataset
from functools import partial
from sklearn.cluster import DBSCAN,MeanShift
torch.multiprocessing.set_sharing_strategy('file_system')
from multiprocessing import Pool
import multiprocessing
from functools import partial
from itertools import repeat

class GraspPredicter:
  def __init__(self,class_name):
    self.class_name_to_artifact_id = {
      'nut': 47,
      'hnm': 51,
      'screw': 50,
    }
    artifact_id = self.class_name_to_artifact_id[class_name]
    code_dir = os.path.dirname(os.path.realpath(__file__))
    artifact_dir = f"{code_dir}/artifacts/artifacts-{artifact_id}"
    print('GraspPredicter artifact_dir',artifact_dir)
    with open(f"{artifact_dir}/config_grasp.yml",'r') as ff:
      self.cfg = yaml.safe_load(ff)

    normalizer_dir = '{}/normalizer.pkl'.format(artifact_dir)
    if os.path.exists(normalizer_dir):
      with open(normalizer_dir,'rb') as ff:
        tmp = pickle.load(ff)
      self.cfg['mean'] = tmp['mean']
      self.cfg['std'] = tmp['std']

    self.dataset = GraspDataset(self.cfg,phase='test',class_name=class_name)

    self.model = PointNetCls(n_in=self.cfg['input_channel'],n_out=len(self.cfg['classes'])-1)
    self.model = load_model(self.model,ckpt_dir='{}/best_val.pth.tar'.format(artifact_dir))
    self.model.cuda().eval()

  def predict_batch(self,data,grasp_poses):
    with torch.no_grad():
      batch_size = 200
      input_datas = []
      for i in range(len(grasp_poses)):
        data_transformed = self.dataset.transform(copy.deepcopy(data),grasp_poses[i])
        input_data = torch.from_numpy(data_transformed['input'])
        input_datas.append(input_data)

      input_datas = torch.stack(input_datas,dim=0)
      n_split = int(np.ceil(len(input_datas)/batch_size))
      ids = np.arange(len(input_datas))
      ids_split = np.array_split(ids,n_split)

      out = []
      for i in range(len(ids_split)):
        ids = ids_split[i]
        input_data = input_datas[ids].cuda().float()
        pred = self.model(input_data)[0]
        pred = pred.softmax(dim=1).data.cpu().numpy()
        for b in range(len(pred)):
          cur_pred = pred[b]
          pred_label = cur_pred.argmax()
          confidence = cur_pred[pred_label]
          out.append([pred_label,confidence,cur_pred])
    torch.cuda.empty_cache()

    return out

class NunocsPredicter:
  def __init__(self,class_name):
    self.class_name = class_name
    self.class_name_to_artifact_id = {
      'nut': 78,
      'hnm': 73,
      'screw': 76
    }
    if self.class_name=='nut':
      self.min_scale = [0.005,0.005,0.001]
      self.max_scale = [0.05,0.05,0.05]
    elif self.class_name=='hnm':
      self.min_scale = [0.005,0.005,0.005]
      self.max_scale = [0.15,0.05,0.05]
    elif self.class_name=='screw':
      self.min_scale = [0.005,0.005,0.005]
      self.max_scale = [0.15,0.05,0.05]

    artifact_id = self.class_name_to_artifact_id[class_name]
    code_dir = os.path.dirname(os.path.realpath(__file__))
    artifact_dir = f"{code_dir}/artifacts/artifacts-{artifact_id}"
    print('NunocsPredicter artifact_dir',artifact_dir)
    with open(f"{artifact_dir}/config_nunocs.yml",'r') as ff:
      self.cfg = yaml.safe_load(ff)
    if os.path.exists('{}/normalizer.pkl'.format(artifact_dir)):
      with open('{}/normalizer.pkl'.format(artifact_dir),'rb') as ff:
        tmp = pickle.load(ff)
      self.cfg['mean'] = tmp['mean']
      self.cfg['std'] = tmp['std']
    self.dataset = NunocsIsolatedDataset(self.cfg,phase='test')

    self.model = PointNetSeg(n_in=self.cfg['input_channel'],n_out=3*self.cfg['ce_loss_bins'])

    self.model = load_model(self.model,ckpt_dir='{}/best_val.pth.tar'.format(artifact_dir))
    self.model.cuda().eval()

  def predict(self,data):
    with torch.no_grad():
      data['cloud_nocs'] = np.zeros(data['cloud_xyz'].shape)
      data['cloud_rgb'] = np.zeros(data['cloud_xyz'].shape)
      data_transformed = self.dataset.transform(copy.deepcopy(data))
      self.data_transformed = data_transformed
      ori_cloud = data_transformed['cloud_xyz_original']
      input_data = torch.from_numpy(data_transformed['input']).cuda().float().unsqueeze(0)

      pred = self.model(input_data)[0].reshape(-1,3,self.cfg['ce_loss_bins'])
      bin_resolution = 1/self.cfg['ce_loss_bins']
      pred_coords = pred.argmax(dim=-1).float()*bin_resolution
      probs = pred.softmax(dim=-1)
      confidence_z = torch.gather(probs[:,2,:],dim=-1,index=pred[:,2,:].argmax(dim=-1).unsqueeze(-1)).data.cpu().numpy().reshape(-1)
      conf_color = array_to_heatmap_rgb(confidence_z)
      nocs_cloud = pred_coords.data.cpu().numpy()-0.5

      nocs_cloud_down = copy.deepcopy(nocs_cloud)
      ori_cloud_down = copy.deepcopy(ori_cloud)

      best_ratio = 0
      best_transform = None
      best_nocs_cloud = None
      best_symmetry_tf = None
      for symmetry_tf in [np.eye(4)]:
        tmp_nocs_cloud_down = (symmetry_tf@to_homo(nocs_cloud_down).T).T[:,:3]
        for thres in [0.003,0.005]:
          use_kdtree_for_eval = False
          kdtree_eval_resolution = 0.003
          transform, inliers = estimate9DTransform(source=tmp_nocs_cloud_down,target=ori_cloud_down,PassThreshold=thres,max_iter=10000,use_kdtree_for_eval=use_kdtree_for_eval,kdtree_eval_resolution=kdtree_eval_resolution,max_scale=self.max_scale,min_scale=self.min_scale,max_dimensions=np.array([1.2,1.2,1.2]))
          if transform is None:
            continue

          if np.linalg.det(transform[:3,:3])<0:
            continue
          scales = np.linalg.norm(transform[:3,:3],axis=0)
          print("thres",thres)
          print("estimated scales",scales)
          print("transform:\n",transform)
          transformed = (transform@to_homo(tmp_nocs_cloud_down).T).T[:,:3]
          err_thres = 0.003

          cloud_at_canonical = (np.linalg.inv(transform)@to_homo(ori_cloud_down).T).T[:,:3]
          dimensions = cloud_at_canonical.max(axis=0)-cloud_at_canonical.min(axis=0)
          print("estimated canonical dimensions",dimensions)

          errs = np.linalg.norm(transformed-ori_cloud_down, axis=1)
          ratio = np.sum(errs<=err_thres)/len(errs)
          inliers = np.where(errs<=err_thres)[0]

          print("inlier ratio",ratio)

          if ratio>best_ratio:
            best_ratio = ratio
            best_symmetry_tf = symmetry_tf
            best_transform = transform.copy()
            best_nocs_cloud = copy.deepcopy(tmp_nocs_cloud_down)

      if best_transform is None:
        return None,None

      print(f"nocs predictor best_ratio={best_ratio}, scales={np.linalg.norm(best_transform[:3,:3],axis=0)}")
      print("nocs pose\n",best_transform)
      self.best_ratio = best_ratio
      transform = best_transform
      self.nocs_pose = transform.copy()
      nocs_cloud = (best_symmetry_tf@to_homo(nocs_cloud).T).T[:,:3]

      return nocs_cloud, transform

class PointGroupPredictor:
  def __init__(self,class_name):
    self.class_name_to_artifact_id = {
      'nut': 40,
      'hnm': 68,
      'screw': 77,
    }
    self.class_name = class_name
    artifact_id = self.class_name_to_artifact_id[class_name]
    code_dir = os.path.dirname(os.path.realpath(__file__))
    artifact_dir = f"{code_dir}/artifacts/artifacts-{artifact_id}"
    print('PointGroupPredictor artifact_dir',artifact_dir)
    config_dir = f"{artifact_dir}/config_pointgroup.yaml"
    self.cfg_pg = config_pg.get_parser(config_dir=config_dir)
    with open(config_dir,'r') as ff:
      self.cfg = yaml.safe_load(ff)

    self.dataset = dataset_seg.Dataset(cfg=self.cfg,cfg_pg=self.cfg_pg,phase='test')

    self.model = PointGroup(self.cfg_pg)
    self.model = load_model(self.model,ckpt_dir='{}/best_val.pth.tar'.format(artifact_dir))
    self.model.cuda().eval()

    self.n_slice_per_side = 1

  def predict(self,data):
    with torch.no_grad():
      xmax = data['cloud_xyz'][:,0].max()
      xmin = data['cloud_xyz'][:,0].min()
      ymax = data['cloud_xyz'][:,1].max()
      ymin = data['cloud_xyz'][:,1].min()
      xlen = (xmax-xmin)/self.n_slice_per_side
      ylen = (ymax-ymin)/self.n_slice_per_side

      batch_offsets = [0]
      locs = []
      xyz_original_all = []
      feats = []
      colors = []

      for ix in range(self.n_slice_per_side):
        for iy in range(self.n_slice_per_side):
          xstart = xmin+ix*xlen
          ystart = ymin+iy*ylen
          keep_mask = (data['cloud_xyz'][:,0]>=xstart) & (data['cloud_xyz'][:,0]<=xstart+xlen) & (data['cloud_xyz'][:,1]>=ystart) & (data['cloud_xyz'][:,1]<=ystart+ylen)
          xyz_origin = data['cloud_xyz'][keep_mask]
          normals = data['cloud_normal'][keep_mask]
          color = data['cloud_rgb'][keep_mask]

          pcd = toOpen3dCloud(xyz_origin)
          pcd = pcd.voxel_down_sample(voxel_size=self.cfg['downsample_size'])
          pts = np.asarray(pcd.points).copy()
          kdtree = cKDTree(xyz_origin)
          dists,indices = kdtree.query(pts)
          xyz_origin = xyz_origin[indices]
          normals = normals[indices]
          color = color[indices]

          xyz = xyz_origin * self.dataset.scale
          xyz -= xyz.min(0)
          batch_offsets.append(batch_offsets[-1] + xyz.shape[0])

          i = ix+iy*self.n_slice_per_side
          locs.append(torch.cat([torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long()], 1))

          xyz_original_all.append(torch.from_numpy(xyz_origin))
          feats.append(torch.from_numpy(normals))
          colors.append(torch.from_numpy(color))

      batchsize = len(batch_offsets)-1
      batch_offsets = torch.tensor(batch_offsets, dtype=torch.int)
      locs = torch.cat(locs, 0)
      xyz_original_all = torch.cat(xyz_original_all, 0).to(torch.float32)
      feats = torch.cat(feats, 0)
      colors = torch.cat(colors, 0)

      spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.dataset.full_scale[0], None)

      voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(locs, len(batch_offsets)-1, self.dataset.mode)

      coords = locs.cuda()
      voxel_coords = voxel_locs.cuda()
      p2v_map = p2v_map.cuda()
      v2p_map = v2p_map.cuda()

      coords_float = xyz_original_all.cuda().float()
      feats = feats.cuda().float()

      batch_offsets = batch_offsets.cuda()

      if self.cfg_pg.use_coords:
          feats = torch.cat((feats, coords_float), 1)
      voxel_feats = pointgroup_ops.voxelization(feats, v2p_map, self.cfg_pg.mode)

      input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, self.cfg_pg.batch_size)

      ret = self.model(input_, p2v_map, coords_float, coords[:, 0].int(), batch_offsets, epoch=self.model.prepare_epochs-1)
      offsets = ret['pt_offsets'].data.cpu().numpy()
      xyz_original_all = xyz_original_all.data.cpu().numpy()

      pcd = toOpen3dCloud(xyz_original_all)
      pcd = pcd.voxel_down_sample(voxel_size=0.002)
      xyz_down = np.asarray(pcd.points).copy()
      kdtree = cKDTree(xyz_original_all)
      dists,indices = kdtree.query(xyz_down)
      xyz_down = xyz_original_all[indices]
      xyz_shifted = xyz_down+offsets[indices]
      self.xyz_shifted = xyz_shifted

      if self.class_name=='hnm':
        eps = 0.003
        min_samples = 20
        bandwidth = 0.005
      elif self.class_name=='nut':
        eps = 0.003
        min_samples = 5
        bandwidth = 0.007
      elif self.class_name=='screw':
        eps = 0.003
        min_samples = 5
        bandwidth = 0.009
      else:
        raise NotImplemented

      labels = MeanShift(bandwidth=bandwidth,cluster_all=True,n_jobs=-1,seeds=None).fit_predict(xyz_shifted)

      kdtree = cKDTree(xyz_down)
      dists,indices = kdtree.query(data['cloud_xyz'])
      labels_all = labels[indices]

      return labels_all


import sys
sys.path.append('./')

import pandas as pd
import os
from collections import namedtuple
import numpy as np
import cv2
from shutil import copyfile
from tqdm import tqdm
from scipy import stats
import time
import random
import warnings
from shutil import copyfile
import pyexcel
import yaml

try:
    dj.__version__
except: # ? onlu import tables if we havent already
    try:
        import datajoint as dj
        from database.dj_config import start_connection
        dbname, _ = start_connection()
    except Exception as e:
        print("Could not connect to database: ", e)
    else:
        # print("Importing tables")
        from database.TablesDefinitionsV4 import *

from Utilities.matplotlib_config_figures import *

from Utilities.file_io.files_load_save import *
from Utilities.video_and_plotting.video_editing import Editor
from Utilities.maths.math_utils import *
from Utilities.maths.distributions import *
from Utilities.maths.filtering import *
from Utilities.constants import *

from Plotting.utils.plotting_utils import *
from Plotting.utils.plot_distributions import *
from Plotting.utils.colors import *

mpl.rc('text', usetex=False)
params = {
    'text.latex.preamble': ['\\usepackage{gensymb}'],
    # 'image.origin': 'lower',
    'image.interpolation': 'nearest',
    'image.cmap': 'gray',
    'axes.grid': False,
    'savefig.dpi': 150,  # to adjust notebook inline plot size
    'axes.labelsize': 12, # fontsize for x and y labels (was 10)
    'axes.titlesize': 12,
    'font.size': 12, # was 10
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'text.usetex': False,        # ! <----- use TEX
    'legend.fontsize': 14,
}
mpl.rcParams.update(params)


<reponame>PlaytikaResearch/abexp
# MIT License
#
# Copyright (c) 2021 Playtika Ltd.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
from scipy.stats import ttest_ind
from statsmodels.stats.power import NormalIndPower, TTestIndPower

from abexp.statistics.stats_metrics import cohens_d, cohens_h

class SampleSize:
    """
    This class provides some utils to be used before running A/B test experiments. It includes minimum sample size
    determination, power calculation and effect size estimation. It handles both the case of means comparison and
    proportions comparison. Results are computed via power analysis with closed-form solution or simulation under the
    assumption that sample data are normally distributed.
    """
    def __init__(self):
        pass

    @staticmethod
    def ssd_mean(mean_contr, mean_treat, std_contr, alpha=0.05, power=0.8):
        """
        Sample size determination (SDD) to compare means. Compute the minimum sample size needed to run A/B test
        experiments. The result is computed via power analysis with closed-form solution t-test. Effect size estimation
        is calculated with cohen's d coefficient.

        Parameters
        ----------
        mean_contr : float
            Mean of the control group.
        mean_treat : float
            Mean of the treatment group.
        std_contr : float > 0
            Standard deviation of the control group. It assumes that the standard deviation of the control group is
            equal to the standard deviation of the treatment group.
        alpha : float in interval (0,1)
            Significance level, default 0.05. It is the probability of a type I error, that is wrong rejections if the
            null hypothesis is true.
        power : float in interval (0,1)
            Statistical power of the test, default 0.8. It is one minus the probability of a type II error. Power is
            the probability that the test correctly rejects the null hypothesis if the alternative hypothesis is true.

        Returns
        -------
        sample_size : int
            Minimum sample size per each group
        """

        # If the means are equals the function returns infinity
        if mean_contr == mean_treat:
            return np.Inf

        # Compute effect size as Cohen's d
        effect_size = cohens_d(mu_1=mean_contr, mu_2=mean_treat, std=std_contr)

        # Compute t-test to solve sample size
        analysis = TTestIndPower()
        sample_size = analysis.solve_power(effect_size=effect_size, power=power, alpha=alpha, nobs1=None, ratio=1.,
                                           alternative='two-sided')

        return round(sample_size)

    @staticmethod
    def ssd_mean_sim(mean_contr, mean_treat, std_contr, alpha=0.05, power=0.8, sims=1000, start_size=100, step_size=0,
                     max_size=10000):
        """
        Sample size determination (SDD) to compare means with simulation. Compute the minimum sample size needed to run
        A/B test experiments. The result is computed via power analysis with simulation through t-test.

        Parameters
        ----------
        mean_contr : float
            Mean of the control group.
        mean_treat : float
            Mean of the treatment group.
        std_contr : float > 0
            Standard deviation of the control group. It assumes that the standard deviation of the control group is
            equal to the standard deviation of the treatment group.
        alpha : float in interval (0,1)
            Significance level, default 0.05. It is the probability of a type I error, that is wrong rejections if the
            Null Hypothesis is true.
        power : float in interval (0,1)
            Statistical Power of the test, default 0.8. It is one minus the probability of a type II error. Power is the
            probability that the test correctly rejects the Null Hypothesis if the Alternative Hypothesis is true.
        sims : int
            Number simulations, default 1000.
        start_size : int
            Initial sample size, default 100, used for the first iteration.
        step_size : int
            Spacing between samples size, default 50. This is the distance between two adjacent sample size,
            sample_size[i+1] - sample_size[i].
        max_size : int
            Maximum sample size, default 10000. The function returns this value if the desired power is not reached via
            simulation.

        Returns
        -------
        sample_size : int
            Minimum sample size per each group
        """

        # If the means are equals the function returns infinity
        if mean_contr == mean_treat:
            return np.Inf

        # Initialize simulated power and sample size, sample size
        sim_power, sample_size = 0, start_size - step_size

        while sim_power < power and sample_size <= max_size - step_size:

            # Keep incrementing sample size by step_size till the required power is reached
            sample_size += step_size

            # Model the variable for the two groups as normal distributions
            obs_contr = np.random.normal(loc=mean_contr, scale=std_contr, size=(sample_size, sims))
            obs_treat = np.random.normal(loc=mean_treat, scale=std_contr, size=(sample_size, sims))

            # Compute t-test for control and treatment groups
            _, p_value = ttest_ind(obs_treat, obs_contr)

            # Power is the fraction of times in the simulation when the p-value was less than 0.05
            sim_power = (p_value < alpha).sum() / sims

        return round(sample_size)

    @staticmethod
    def ssd_prop(prop_contr, prop_treat, alpha=0.05, power=0.8):
        """
        Sample size determination (SDD) to compare proportions. Compute the minimum sample size needed to run A/B test
        experiments. The result is computed via power analysis with closed-form solution z-test. Effect size
        estimation is calculated with cohen's h coefficient.

        Parameters
        ----------
        prop_contr : float in interval (0,1)
            Proportion in the control group.
        prop_treat : float in interval (0,1)
            Proportion in the treatment group.
        alpha : float in interval (0,1)
            Significance level, default 0.05. It is the probability of a type I error, that is wrong rejections if the
            Null Hypothesis is true.
        power : float in interval (0,1)
            Statistical Power of the test, default 0.8. It is one minus the probability of a type II error. Power is
            the probability that the test correctly rejects the Null Hypothesis if the Alternative Hypothesis is true.

        Returns
        -------
        sample_size : int
            Minimum sample size per each group
        """

        # If proportions are equals the function returns infinity
        if prop_contr == prop_treat:
            return np.Inf

        # Compute effect size as Cohen's h
        effect_size = cohens_h(p1=prop_contr, p2=prop_treat)

        # Compute t-test to solve sample size
        analysis = NormalIndPower()
        sample_size = analysis.solve_power(effect_size=effect_size, nobs1=None, alpha=alpha, power=power, ratio=1.,
                                           alternative='two-sided')
        return round(sample_size)


from os import environ
from scipy.optimize import curve_fit
import numpy as np
import matplotlib.pyplot as plt
import os
import time
from random import randint
import subprocess
from itertools import combinations
import matplotlib

matplotlib.use('Agg')

pp = "./build/release/examples/"

graphs = [
    "graph_alexnet",
    "graph_googlenet",
    "graph_mobilenet",
    "graph_squeezenet",
    "graph_resnet50"
]

cp = ["--target=NEON", "--threads=4"]
gp = ["--target=CL"]
sm = ["taskset", "-c", "0-3"]
bg = ["taskset", "-c", "4-7"]

targets = {
    "cpu_big": lambda g: bg + [pp + g] + cp,
    "cpu_small": lambda g: sm + [pp + g] + cp,
    "gpu": lambda g: [pp + g] + gp
}

T = 5
dt = 0.01

def temp_func(t, b, c): return b * (1 - np.exp(- t / c))

def get_temp():
    with open('/sys/class/thermal/thermal_zone0/temp') as f:
        temp = int(f.readline())
        return temp

def plot_temp(cmd_combi, x, y, popt, i, t0):
    fig = plt.figure()
    plt.plot(x, [t0+yy for yy in y])
    yp = [t0+temp_func(xx, *popt) for xx in x]
    plt.plot(x, yp,
        'r-',
        label='fit: a=%5.3f, b=%5.3f' % tuple(popt))
    plt.axhline(80000, linestyle='--', color='r')
    plt.xlabel('Time')
    plt.ylabel('Temperature')
    combi_name = "-".join([target for target, _ in cmd_combi])
    plt.title('Execution ' + graph + " " + combi_name)
    plt.legend(loc='best')
    name ='temp_profile_v4/' + graph + "_" + combi_name + "_" + str(i) + ".png"
    print(name)
    plt.savefig(name)
    plt.close(fig)

if __name__ == "__main__":
    env = dict(os.environ)
    env['LD_LIBRARY_PATH'] = './build/release'
    for graph in graphs:
        print("graph = "+graph)
        RCs = {}
        cmd_combis = []
        cmds = list((target, targets[target](graph)) for target in targets)
        for i in range(1, len(cmds) + 1):
            cmd_combis += list(combinations(cmds, i))
        for cmd_combi in cmd_combis:
            print('> ', "-".join([target for target, _ in cmd_combi]))
            bp = 30000
            cp = 5
            bss = []
            css = []
            for i in range(10):
                time.sleep(randint(0, T))
                t=0
                running = []
                for target_cmd in cmd_combi:
                    running.append(subprocess.Popen(target_cmd[1], env=env))
                x=[]
                y = []
                while None in (p.poll() for p in running):
                    x.append(t)
                    y.append(get_temp())
                    time.sleep(dt)
                    t+=dt
                y.append(get_temp())
                x.append(t)
                t0 = y[0]
                mx = t0
                y = [yy-t0 for yy in y]
                try:
                    popt, _ = curve_fit(temp_func, x, y, bounds=(0, np.inf), p0=(bp, cp))
                except:
                    continue
                [b, c] = popt
                bp = b
                cp = c
                bss.append(b)
                css.append(c)
                print('b=%5.3f, c=%5.3f' % (b, c))
                plot_temp(cmd_combi, x, y, popt, i, t0)
            bs = sum(bss)/len(bss)
            cs  = sum(css)/len(css)
            print('Average: b=%5.3f, c=%5.3f' % (bs, cs))


<reponame>waterjump/alignedCelebFaces
# Just disables the warning, doesn't enable AVX/FMA
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf
import numpy as np
from scipy import misc
from tensorflow.examples.tutorials.mnist import input_data
import random
import math

def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)

def bias_variable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)

IMAGE_WIDTH = 64
IMAGE_HEIGHT = 64

learning_rate = 0.0000  # Used to be 0.001
inputs_ = tf.placeholder(tf.float32, (None, IMAGE_HEIGHT, IMAGE_WIDTH, 3), name='inputs')
targets_ = tf.placeholder(tf.float32, (None, IMAGE_HEIGHT, IMAGE_WIDTH, 3), name='targets')

### Encoder
conv0 = tf.layers.conv2d(inputs=inputs_, filters=120, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 64x64x25
maxpool0 = tf.layers.max_pooling2d(conv0, pool_size=(2,2), strides=(2,2), padding='same')
# Now 32x32x25
conv1 = tf.layers.conv2d(inputs=maxpool0, filters=160, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 32x32x40
maxpool1 = tf.layers.max_pooling2d(conv1, pool_size=(2,2), strides=(2,2), padding='same')
# Now 16x16x40
conv2 = tf.layers.conv2d(inputs=maxpool1, filters=200, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 16x16x60
maxpool2 = tf.layers.max_pooling2d(conv2, pool_size=(2,2), strides=(2,2), padding='same')
# Now 8x8x60
conv3 = tf.layers.conv2d(inputs=maxpool2, filters=240, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 8x8x80
maxpool3 = tf.layers.max_pooling2d(conv3, pool_size=(2,2), strides=(2,2), padding='same')
# Now 4x4x80

maxpool3_flat = tf.reshape(maxpool3, [-1,4*4*240])

W_fc1 = weight_variable([4*4*240, 300])
b_fc1 = bias_variable([300])
tesy = tf.matmul(maxpool3_flat, W_fc1)
encoded = tf.nn.relu(tf.matmul(maxpool3_flat, W_fc1) + b_fc1)

W_fc2 = weight_variable([300, 4*4*240])
b_fc2 = bias_variable([4*4*240])
predecoded_flat = tf.nn.relu(tf.matmul(encoded, W_fc2) + b_fc2)

predecoded = tf.reshape(predecoded_flat, [-1,4,4,240])

### Decoder
upsample1 = tf.image.resize_images(predecoded, size=(8,8), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
# Now 8x8x80
conv4 = tf.layers.conv2d(inputs=upsample1, filters=200, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 8x8x60
upsample2 = tf.image.resize_images(conv4, size=(16,16), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
# Now 16x16x60
conv5 = tf.layers.conv2d(inputs=upsample2, filters=160, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 16x16x40
upsample3 = tf.image.resize_images(conv5, size=(32,32), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
# Now 32x32x40
conv6 = tf.layers.conv2d(inputs=upsample3, filters=120, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 32x32x25
upsample4 = tf.image.resize_images(conv6, size=(64,64), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
# Now 64x64x25
conv7 = tf.layers.conv2d(inputs=upsample4, filters=15, kernel_size=(3,3), padding='same', activation=tf.nn.relu)
# Now 64x64x10

logits = tf.layers.conv2d(inputs=conv7, filters=3, kernel_size=(3,3), padding='same', activation=None)
#Now 64x64x1

# Pass logits through sigmoid to get reconstructed image
decoded = tf.nn.sigmoid(logits)

# Pass logits through sigmoid and calculate the cross-entropy loss
loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=targets_, logits=logits)

# Get cost and define the optimizer
cost = tf.reduce_mean(loss)
opt = tf.train.AdamOptimizer(learning_rate).minimize(cost)

print("made it here! :D")
IMAGE_COUNT = 13016
DENSE_SIZE = 300
sess = tf.Session()
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())
saver.restore(sess,  "/media/rob/Ma Book1/alignedCelebFaces/models/model27674.ckpt")

denseArray = np.zeros((IMAGE_COUNT,DENSE_SIZE))
for imageIndex in range(IMAGE_COUNT):
    imagio = misc.imread('/media/rob/Ma Book1/alignedCelebFaces/data/dataFace'+str(imageIndex)+'.png')
    imagio = imagio[:,:,0:3].reshape((1,64,64,3))/255.0
    denseRep = sess.run([encoded], feed_dict={inputs_: imagio,targets_: imagio})
    denseArray[imageIndex] = np.array(denseRep)
    if imageIndex%20 == 0:
        print(imageIndex)
np.save("denseArray27K.npy",denseArray)


"""
Direct NumPy port of tfgan.eval.classifier_metrics
"""
# CREDITS: https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/tpu_utils/classifier_metrics_numpy.py

import numpy as np
import scipy.special

def log_softmax(x, axis):
    return x - scipy.special.logsumexp(x, axis=axis, keepdims=True)

def kl_divergence(p, p_logits, q):
    assert len(p.shape) == len(p_logits.shape) == 2
    assert len(q.shape) == 1
    return np.sum(p * (log_softmax(p_logits, axis=1) - np.log(q)[None, :]), axis=1)

def _symmetric_matrix_square_root(mat, eps=1e-10):
    """Compute square root of a symmetric matrix.
    Note that this is different from an elementwise square root. We want to
    compute M' where M' = sqrt(mat) such that M' * M' = mat.
    Also note that this method **only** works for symmetric matrices.
    Args:
      mat: Matrix to take the square root of.
      eps: Small epsilon such that any element less than eps will not be square
        rooted to guard against numerical instability.
    Returns:
      Matrix square root of mat.
    """
    u, s, vt = np.linalg.svd(mat)
    # sqrt is unstable around 0, just use 0 in such case
    si = np.where(s < eps, s, np.sqrt(s))
    return u.dot(np.diag(si)).dot(vt)

def trace_sqrt_product(sigma, sigma_v):
    """Find the trace of the positive sqrt of product of covariance matrices.
    '_symmetric_matrix_square_root' only works for symmetric matrices, so we
    cannot just take _symmetric_matrix_square_root(sigma * sigma_v).
    ('sigma' and 'sigma_v' are symmetric, but their product is not necessarily).
    Let sigma = A A so A = sqrt(sigma), and sigma_v = B B.
    We want to find trace(sqrt(sigma sigma_v)) = trace(sqrt(A A B B))
    Note the following properties:
    (i) forall M1, M2: eigenvalues(M1 M2) = eigenvalues(M2 M1)
       => eigenvalues(A A B B) = eigenvalues (A B B A)
    (ii) if M1 = sqrt(M2), then eigenvalues(M1) = sqrt(eigenvalues(M2))
       => eigenvalues(sqrt(sigma sigma_v)) = sqrt(eigenvalues(A B B A))
    (iii) forall M: trace(M) = sum(eigenvalues(M))
       => trace(sqrt(sigma sigma_v)) = sum(eigenvalues(sqrt(sigma sigma_v)))
                                     = sum(sqrt(eigenvalues(A B B A)))
                                     = sum(eigenvalues(sqrt(A B B A)))
                                     = trace(sqrt(A B B A))
                                     = trace(sqrt(A sigma_v A))
    A = sqrt(sigma). Both sigma and A sigma_v A are symmetric, so we **can**
    use the _symmetric_matrix_square_root function to find the roots of these
    matrices.
    Args:
      sigma: a square, symmetric, real, positive semi-definite covariance matrix
      sigma_v: same as sigma
    Returns:
      The trace of the positive square root of sigma*sigma_v
    """

    # Note sqrt_sigma is called "A" in the proof above
    sqrt_sigma = _symmetric_matrix_square_root(sigma)

    # This is sqrt(A sigma_v A) above
    sqrt_a_sigmav_a = sqrt_sigma.dot(sigma_v.dot(sqrt_sigma))

    return np.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a))

def classifier_score_from_logits(logits):
    """Classifier score for evaluating a generative model from logits.
    This method computes the classifier score for a set of logits. This can be
    used independently of the classifier_score() method, especially in the case
    of using large batches during evaluation where we would like precompute all
    of the logits before computing the classifier score.
    This technique is described in detail in https://arxiv.org/abs/1606.03498. In
    summary, this function calculates:
    exp( E[ KL(p(y|x) || p(y)) ] )
    which captures how different the network's classification prediction is from
    the prior distribution over classes.
    Args:
      logits: Precomputed 2D tensor of logits that will be used to compute the
        classifier score.
    Returns:
      The classifier score. A floating-point scalar of the same type as the output
      of `logits`.
    """
    assert len(logits.shape) == 2

    # Use maximum precision for best results.
    logits_dtype = logits.dtype
    if logits_dtype != np.float64:
        logits = logits.astype(np.float64)

    p = scipy.special.softmax(logits, axis=1)
    q = np.mean(p, axis=0)
    kl = kl_divergence(p, logits, q)
    assert len(kl.shape) == 1
    log_score = np.mean(kl)
    final_score = np.exp(log_score)

    if logits_dtype != np.float64:
        final_score = final_score.astype(logits_dtype)

    return final_score

def frechet_classifier_distance_from_activations(
    real_activations, generated_activations
):
    """Classifier distance for evaluating a generative model.
    This methods computes the Frechet classifier distance from activations of
    real images and generated images. This can be used independently of the
    frechet_classifier_distance() method, especially in the case of using large
    batches during evaluation where we would like precompute all of the
    activations before computing the classifier distance.
    This technique is described in detail in https://arxiv.org/abs/1706.08500.
    Given two Gaussian distribution with means m and m_w and covariance matrices
    C and C_w, this function calculates
                  |m - m_w|^2 + Tr(C + C_w - 2(C * C_w)^(1/2))
    which captures how different the distributions of real images and generated
    images (or more accurately, their visual features) are. Note that unlike the
    Inception score, this is a true distance and utilizes information about real
    world images.
    Note that when computed using sample means and sample covariance matrices,
    Frechet distance is biased. It is more biased for small sample sizes. (e.g.
    even if the two distributions are the same, for a small sample size, the
    expected Frechet distance is large). It is important to use the same
    sample size to compute frechet classifier distance when comparing two
    generative models.
    Args:
      real_activations: 2D Tensor containing activations of real data. Shape is
        [batch_size, activation_size].
      generated_activations: 2D Tensor containing activations of generated data.
        Shape is [batch_size, activation_size].
    Returns:
     The Frechet Inception distance. A floating-point scalar of the same type
     as the output of the activations.
    """
    assert len(real_activations.shape) == len(generated_activations.shape) == 2

    activations_dtype = real_activations.dtype
    if activations_dtype != np.float64:
        real_activations = real_activations.astype(np.float64)
        generated_activations = generated_activations.astype(np.float64)

    # Compute mean and covariance matrices of activations.
    m = np.mean(real_activations, 0)
    m_w = np.mean(generated_activations, 0)
    num_examples_real = float(real_activations.shape[0])
    num_examples_generated = float(generated_activations.shape[0])

    # sigma = (1 / (n - 1)) * (X - mu) (X - mu)^T
    real_centered = real_activations - m
    sigma = real_centered.T.dot(real_centered) / (num_examples_real - 1)

    gen_centered = generated_activations - m_w
    sigma_w = gen_centered.T.dot(gen_centered) / (num_examples_generated - 1)

    # Find the Tr(sqrt(sigma sigma_w)) component of FID
    sqrt_trace_component = trace_sqrt_product(sigma, sigma_w)

    # Compute the two components of FID.

    # First the covariance component.
    # Here, note that trace(A + B) = trace(A) + trace(B)
    trace = np.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component

    # Next the distance between means.
    mean = np.sum(np.square(m - m_w))  # Equivalent to L2 but more stable.
    fid = trace + mean
    if activations_dtype != np.float64:
        fid = fid.astype(activations_dtype)

    return fid


import numpy as np
from numpy import ndarray
from typing import Callable, Tuple
from pynverse import inversefunc
from scipy.stats import chi2
from basic.types import vector, matrix, Configuration, elemental, TP_HOTELLING
from basic.decorators import document, type_checker
import basic.docfunc as doc
import copy

def _mvn_params(x: matrix) -> Tuple[vector, matrix]:
    return x.mean(axis=0), np.cov(x.T, bias=True)

def _upper_domain(func: Callable, heuristic_step: float = 10) -> int:
    upper = heuristic_step
    while True:
        try:
            res = inversefunc(func, y_values=1, domain=[0, upper])
        except ValueError:
            upper += heuristic_step  # heuristically increase upper of domain
        else:
            return res

def hotelling_threshold(df: int, level: float = 0.05) -> float:
    """calculate the threshold of chi2 distribution, by a given test level"""
    dis = chi2(df=df)
    upper = _upper_domain(dis.cdf, 10)
    return inversefunc(dis.cdf, y_values=1-level, domain=[0, upper])

def a(x: matrix, miu: vector, sigma: matrix) -> ndarray:
    res = []
    m_inv = np.array(np.matrix(sigma).I)  # inverse matrix
    for i in range(len(x)):
        _ = x[i] - miu
        res.append(np.matmul(np.matmul(_, m_inv), _.T))
    return np.array(res)

@document(doc.en_Hotelling)
class Hotelling:

    settings: Configuration = {
        'model_import': np.array([]),  # ndarray, matrix-like
        'level': 0.05,  # float, 0~1
        'data_import': np.array([]),  # ndarray, matrix-like
    }

    @type_checker(in_class=True, kwargs_types=TP_HOTELLING, elemental_types=elemental)
    def __init__(self, **settings: Configuration):
        assert np.all([k in settings.keys() for k in ['model_import']]) == 1, 'missing required arg model_import.'
        self.settings.update({k: v for k, v in settings.items()})
        self.model = self.settings.get('model_import')
        self.mean, self.sigma = _mvn_params(self.model)
        self.threshold = hotelling_threshold(self.model.shape[1], self.settings.get('level'))

    @type_checker(in_class=True, kwargs_types=TP_HOTELLING, elemental_types=elemental)
    @document(doc.en_Hotelling_predict)
    def predict(self, **settings: Configuration):
        assert np.all([k in settings.keys() for k in ['data_import']]) == 1, 'missing required arg data_import.'
        _settings = copy.deepcopy(self.settings)
        _settings.update({k: v for k, v in settings.items()})
        return a(_settings.get('data_import'), self.mean, self.sigma) <= self.threshold

if __name__ == '__main__':
    pass


<filename>IAM_GPU_improved_lib.py<gh_stars>0
import matplotlib
matplotlib.use('Agg')

from numba import cuda
from timeit import default_timer as timer
from matplotlib import pylab
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

from IAM_lib import *

import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
import scipy.io as sio
import skimage.morphology as skimorph
import skimage.filters as skifilters
import scipy.ndimage.morphology as scimorph

import math, numba, cv2, csv, gc
import os, errno, sys, shutil

import code
from PIL import Image

# Turn interactive plotting off
plt.ioff()

def iam_lots_gpu_compute(csv_filename="", patch_size=[1,2,4,8],
                         blending_weights=[0.65,0.2,0.1,0.05], num_sample=[512],
                         alpha=0.5, thrsh_patches = True, bin_tresh=0.5, save_jpeg=True,
                         delete_intermediary=False, nawm_preprocessing=False):
    '''
    FUNCTION'S SUMMARY:

    Main function of the LOTS-IAM-GPU algorithm. This function produces (i.e. saving)
    age maps that indicate level of irregularity of voxels in brain FLAIR MRI. This
    function reads a list of FLAIR MR image (NifTI), ICV mask (NifTI), CSF mask (NifTI),
    NAWM mask (NifTI), and Cortical mask (NifTI) to produce the corresponding age maps
    from a CSV file. Please note that this version only accept NifTI (.nii/.nii.gz) files.

    NOTE: NAWM and Cortical masks are optional. They will be used if they are included
    in the CSV file.

    Format of the CSV input file (NOTE: spaces are used to make the format clearer):

        path_to_mri_codebase_folder, mri_code_name, path_FLAIR, path_ICV, path_CSF,
        path_NAWM (optional), path_Cortical (optional)

    Example (NOTE: spaces are used to make the format clearer):

        /dir/MRI_DB/, MRI001, /dir/MRI_DB/MRI001/FLAIR.nii.gz, /dir/MRI_DB/MRI001/ICV.nii.gz,
        /dir/MRI_DB/MRI001/CSF.nii.gz, /dir/MRI_DB/MRI001/NAWM.nii.gz (optional),
        /dir/MRI_DB/MRI001/Cortex.nii.gz (optional)

    By default, the age maps are calculated by using four different sizes of source/target
    patches (i.e. 1x1, 2x2, 4x4, and 8x8) and 64 target samples. Furthermore, all intermediary
    files are saved in .mat (Matlab) and JPEG files.

    INPUT PARAMETERS:

    This function's behavior can be set by using input parameters below.

        1. output_filedir   : Path of directory for saving all results. Format of the path:
                              "output_path/name_of_experiment"

        2. csv_filename     : Name of a CSV input file which contains list all files to be
                              processed by the LOTS-IAM-GPU. Example: "input.csv"

        3. patch_size       : Size of source/target patches for IAM's computation. Default:
                              [1,2,4,8] to calculate age maps from four different sizes of
                              source/target patches i.e. 1x1, 2x2, 4x4, and 8x8. The sizes
                              of source/target patches must be in the form of python's list.

        4. blending_weights : Weights used for blending age maps produced by different size of
                              source/target patches. The weights must be in the form of python's
                              list, summed to 1, and its length must be the same as `patch_size`.

        5. num_sample       : A list of numbers used for randomly sampling target patches to be
                              used in the LOTS-IAM-GPU calculation. Default: [512]. Available
                              values: [64, 128, 256, 512, 1024, 2048]. Some important notes:

                                a. Smaller number will make computation faster.
                                b. Input the numbers as a list to automatically produce
                                   age maps by using all different numbers of target patches.
                                   The software will automatically create different output
                                   folders for different number of target samples.
                                c. For this version, only 64, 128, 256, 512, 1024, and 2048
                                   can be used as input numbers (error will be raised if other
                                   numbers are used).

        6. alpha            : Weight of distance function to blend maximum difference and
                              average difference between source and target patches. Default:
                              0.5. Input value should be between 0 and 1 (i.e. floating points).
                              The current distance function being used is:

                                  d = (alpha . |max(s - t)|) + ((1 - alpha) . |mean(s - t)|)

                              where d is distance value, s is source patch, and t is target patch.

        7. bin_tresh        : Threshold value for cutting of probability values of brain masks,
                              if probability masks are given instead of binary masks.

        8. save_jpeg        : True  --> Save all JPEG files for visualisation.
                              False --> Do not save the JPEG files.

        9. delete_intermediary : False --> Save all intermediary files (i.e. JPEG/.mat files).
                                 True  --> Delete all intermediary files, saving some spaces in
                                           the hard disk drive.

    OUTPUT:

    The software will automatically create a new folder provided in "output_filedir" variable.
    Please make sure that the directory is accessible and writable.

    Inside the experiment’s folder, each patient/MRI mri_code will have its own folder. In default,
    there are 6 sub-folders which are:
    1. 1: Contains age maps of each slice generated by using 1x1 patch.
    2. 2: Contains age maps of each slice generated by using 2x2 patch.
    3. 4: Contains age maps of each slice generated by using 4x4 patch.
    4. 8: Contains age maps of each slice generated by using 8x8 patch.
    5. IAM_combined_python: Contains two sub-folders:
        a. Patch: contains visualisation of age maps of each slices in JPEG files, and
        b. Combined: contains visualisation of the final output of LOTS-IAM-GPU’s computation.
    6. IAM_GPU_nifti_python: Contains one Matlab (.mat) file and three NIfTI files (.nii.gz):
        a. all_slice_dat.mat: processed mri_code of all slices in Matlab file,
        b. IAM_GPU_COMBINED.nii.gz: the original age map values,
        c. IAM_GPU_GN.nii.gz: the final age map values (i.e. GN and penalty), and
        d. IAM_GPU_GN_postprocessed.nii.gz: the final age map values plus post-processing
           (only if NAWM mask is provided).

    Note: If parameter value of `delete_intermediary` is `True`, then all folders listed above
    will be deleted, except for folder `IAM_GPU_nifti_python` and its contents.

    MORE HELP:

    Please read README.md file provided in:
        https://github.com/febrianrachmadi/lots-iam-gpu

    VERSION (dd/mm/yyyy):
    - 31/05/2018b: NAWM and Cortical brain masks are now optional input (will be used if available).
    - 31/05/2018a: Fix header information of the LOTS-IAM-GPU's result.
    - 08/05/2018 : Add lines to cutting off probability mask and deleting intermediary folders.
    - 07/05/2018 : Initial release code.
    '''

    ## Check availability of input files and output path
    if csv_filename == "":
        raise ValueError("Please set output folder's name and CSV mri_code filename. See: help(iam_lots_gpu)")
        return 0

    ## Check compatibility between 'patch_size' and 'blending_weights'
    if len(patch_size) != len(blending_weights):
        raise ValueError("Lengths of 'patch_size' and 'blending_weights' variables are not the same. Length of 'patch_size' is " + str(len(patch_size)) + ", while 'blending_weights' is " + str(len(blending_weights)) + ".")
        return 0

    ## If intermediary files to be deleted, don't even try to save JPEGs
    if delete_intermediary:
        save_jpeg = False

    ''' Set number of mean samples automatically '''
    ''' num_samples_all = [64, 128, 256, 512, 1024, 2048] '''
    ''' num_mean_samples_all = [16, 32, 32, 64, 128, 128] '''
    num_samples_all = num_sample
    num_mean_samples_all = []
    for sample in num_samples_all:
        if sample == 64:
            num_mean_samples_all.append(16)
        elif sample == 128:
            num_mean_samples_all.append(32)
        elif sample == 256:
            num_mean_samples_all.append(32)
        elif sample == 512:
            num_mean_samples_all.append(64)
        elif sample == 1024:
            num_mean_samples_all.append(128)
        elif sample == 2048:
            num_mean_samples_all.append(128)
        else:
            raise ValueError("Number of samples must be either 64, 128, 256, 512, 1024 or 2048!")
            return 0

    print("--- PARAMETERS - CHECKED ---")
    print('CSV mri_code filename: ' + csv_filename)
    print('Patch size(s): ' + str(patch_size))
    print('Number of samples (all): ' + str(num_samples_all))
    print('Number of mean samples (all): ' + str(num_mean_samples_all))
    print('Save JPEGs? ' + str(save_jpeg))
    print("--- PARAMETERS - CHECKED ---\n")

    for ii_s in range(0, len(num_samples_all)):
        num_samples = num_samples_all[ii_s]
        num_mean_samples = num_mean_samples_all[ii_s]
        print('Number of samples for IAM: ' + str(num_samples))
        print('Number of mean samples for IAM: ' + str(num_mean_samples))

        with open(csv_filename, newline='') as csv_file:
            num_subjects = len(csv_file.readlines())
            print('Number of subject(s): ' + str(num_subjects))

        with open(csv_filename, newline='', encoding="utf-8-sig") as csv_file:
            reader = csv.reader(csv_file)

            timer_idx = 0
            elapsed_times_all = np.zeros((num_subjects))
            elapsed_times_patch_all = np.zeros((num_subjects, len(patch_size)))
            for row in reader:
                mri_code = row[2]

                dirOutput = row[1]
                print('Output dir: ' + dirOutput + '\n--')

                try:
                    os.makedirs(dirOutput)
                except OSError as e:
                    if e.errno != errno.EEXIST:
                        raise

                print('--\nNow processing mri_code: ' + mri_code)

                inputSubjectDir = row[0]
                print('Input filename (full path): ' + inputSubjectDir)

                ''' Create output folder(s) '''
                dirOutData = dirOutput + '/' + mri_code
                dirOutDataCom = dirOutput + '/' + mri_code + '/IAM_combined_python/'
                dirOutDataPatch = dirOutput + '/' + mri_code + '/IAM_combined_python/Patch/'
                dirOutDataCombined = dirOutput + '/' + mri_code + '/IAM_combined_python/Combined/'
                try:
                    print(dirOutDataCom)
                    os.makedirs(dirOutDataCom)
                    os.makedirs(dirOutDataPatch)
                    os.makedirs(dirOutDataCombined)
                except OSError as e:
                    if e.errno != errno.EEXIST:
                        raise

                mri_data = sio.loadmat(row[0])     # Loading FLAIR
                mri_data = mri_data["flair"]
                [x_len, y_len, z_len] = mri_data.shape

                one_mri_data = timer()
                for xy in range(0, len(patch_size)):
                    print('>>> Processing patch-size: ' + str(patch_size[xy]) + ' <<<\n')

                    try:
                        os.makedirs(dirOutData + '/' + str(patch_size[xy]))
                    except OSError as e:
                        if e.errno != errno.EEXIST:
                            raise

                    one_patch = timer()
                    for zz in range(0, mri_data.shape[2]):
                        print('---> Slice number: ' + str(zz) + ' <---')

                        '''
                        KEY POINT: PRE-PROCESSING P.2 - START
                        -------------------------------------
                        This version still does per slice operation for extracting brain tissues.
                        Two important variables used in the next part of the code are:
                        1. mask_slice --->  Combination of ICV & CSF masks. It is used to find valid source patches
                                            for LOTS-IAM-GPU computation (i.e. brain tissues' source patches).
                        2. brain_slice -->  Brain tissues' information from FLAIR slice.
                        '''

                        mask_slice = np.nan_to_num(mri_data[:, :, zz])
                        mask_slice[mask_slice > 0] = 1

                        brain_slice = np.nan_to_num(mri_data[:, :, zz])

                        '''
                        -----------------------------------
                        KEY POINT: PRE-PROCESSING P.2 - END
                        '''

                        ## Show brain slice to be used for computation
                        #fig, ax = plt.subplots()
                        #cax = ax.imshow(icv_slice, cmap="jet")
                        #cbar = fig.colorbar(cax)
                        #fig.show()
                        #plt.savefig("plot.jpg")

                        # Vol distance threshold
                        vol_slice = np.count_nonzero(brain_slice) / (x_len * y_len)                         ## Proportion of brain slice compared to full image
                        print('DEBUG-Patch: brain_slice - ' + str(np.count_nonzero(brain_slice)) +
                              ', x_len * y_len - ' + str(x_len * y_len) + ', vol: ' + str(vol_slice))       ## x_len/y_len = 512 here

                        # Patch's sampling number treshold
                        TRSH = 0.50
                        if patch_size[xy] == 1:
                            if vol_slice < 0.010: TRSH = 0
                            elif vol_slice < 0.035: TRSH = 0.15
                            elif vol_slice < 0.070 and vol_slice >= 0.035: TRSH = 0.60
                            elif vol_slice >= 0.070: TRSH = 0.80
                        elif patch_size[xy] == 2:
                            if vol_slice < 0.010: TRSH = 0
                            elif vol_slice < 0.035: TRSH = 0.15
                            elif vol_slice < 0.070 and vol_slice >= 0.035: TRSH = 0.60
                            elif vol_slice >= 0.070: TRSH = 0.80
                        elif patch_size[xy] == 4 or patch_size[xy] == 8:
                            if vol_slice < 0.035: TRSH = 0

                        print('DEBUG-Patch: Size - ' + str(patch_size[xy]) + ', slice - ' + str(zz) +
                              ', vol: ' + str(vol_slice) + ', TRSH: ' + str(TRSH))

                        counter_y = int(y_len / patch_size[xy])                                             ## counter_y = 512 if patch of size 1 and image of size 512x512
                        counter_x = int(x_len / patch_size[xy])
                        source_patch_len = counter_x * counter_y                                            ## How many source patches are neede (e.g. for 1, we need one for each pixel)
                        age_values_all = np.zeros(source_patch_len)                                         ## Age Map that will be filled with the actual values

                        valid = 0
                        if ((vol_slice >= 0.008 and vol_slice < 0.035) and (patch_size[xy] == 1 or patch_size[xy] == 2)) or \
                            ((vol_slice >= 0.035 and vol_slice < 0.065) and (patch_size[xy] == 1 or patch_size[xy] == 2 or \
                             patch_size[xy] == 4)) or (vol_slice > 0.065):
                            valid = 1

                            ## Creating grid-patch 'xy-by-xy'
                            #  -- Column
                            y_c = np.ceil(patch_size[xy] / 2)
                            y_c_sources = np.zeros(int(y_len / patch_size[xy]))
                            for iy in range(0, int(y_len / patch_size[xy])):
                                y_c_sources[iy] = (iy * patch_size[xy]) + y_c - 1

                            #  -- Row
                            x_c = np.ceil(patch_size[xy] / 2)
                            x_c_sources = np.zeros(int(x_len / patch_size[xy]))
                            for ix in range(0, int(x_len / patch_size[xy])):
                                x_c_sources[ix] = (ix * patch_size[xy]) + x_c - 1

                            ''' Extracting Source Patches '''
                            area_source_patch = np.zeros([1,patch_size[xy],patch_size[xy]])
                            center_source_patch = np.zeros([1,2])
                            icv_source_flag = np.zeros([source_patch_len])
                            icv_source_flag_valid = np.ones([source_patch_len])
                            index_mapping = np.ones([source_patch_len]) * -1

                            flag = 1
                            index = 0
                            index_source= 0

                            if patch_size[xy] == 1:
                                area_source_patch = brain_slice[mask_slice == 1]
                                area_source_patch = area_source_patch.reshape([area_source_patch.shape[0], 1, 1])
                                index = source_patch_len
                                index_source = area_source_patch.shape[0]
                                icv_source_flag = mask_slice.flatten()
                                positive_indices = (np.where(brain_slice.flatten() > 0))[0]
                                index = 0
                                for i in positive_indices:
                                    index_mapping[i] = index
                                    index += 1

                            else:
                                area_source_patch = []
                                for isc in range(0, counter_x):
                                    for jsc in range(0, counter_y):
                                            icv_source_flag[index] = mask_slice[int(x_c_sources[isc]), int(y_c_sources[jsc])]
                                            if icv_source_flag[index] == 1:
                                                temp = get_area(x_c_sources[isc], y_c_sources[jsc],
                                                                patch_size[xy], patch_size[xy], brain_slice)
                                                area_source_patch.append(temp.tolist())
                                                index_mapping[index] = index_source
                                                index_source += 1

                                            index += 1
                                area_source_patch = np.asarray(area_source_patch)

                            icv_source_flag_valid = icv_source_flag_valid[0:index_source]
                            age_values_valid = np.zeros(index_source)

                            """ TO DELETE, IT'S JUST FOR DISSERTATION
                            for i in range(area_source_patch.shape[2]):
                                plt.imshow(area_source_patch[i]) #Needs to be in row,col order
                                plt.savefig("test.jpg")
                            """

                            ''' Extracting Target Patches '''
                            target_patches = []
                            index_debug = 0
                            random_array = np.random.randint(10, size=(x_len, y_len))
                            index_possible = np.zeros(brain_slice.shape)
                            index_possible[(mask_slice != 0) & (random_array > TRSH*10)] = 1
                            index_possible = np.argwhere(index_possible)

                            for index_chosen in index_possible:
                                x, y = index_chosen
                                area = get_area(x, y, patch_size[xy], patch_size[xy], brain_slice)
                                if area.size == patch_size[xy] * patch_size[xy]:
                                    if np.random.randint(low=1, high=10)/10 < (100/(x*y)) * num_samples:
                                        pass
                                    target_patches.append(area)
                                    index_debug += 1

                            target_patches_np = get_shuffled_patches(target_patches, num_samples)
                            target_patches_np = target_patches_np[0:num_samples,:,:]
                            print('Sampling finished: ' + ' with: ' + str(index_debug) + ' samples from: ' + str(x_len * y_len))
                            area = []

                            ''''''
                            ''' Reshaping array mri_code '''
                            area_source_patch_cuda_all = np.reshape(area_source_patch,(area_source_patch.shape[0],
                                                            area_source_patch.shape[1] * area_source_patch.shape[2]))
                            target_patches_np_cuda_all = np.reshape(target_patches_np, (target_patches_np.shape[0],
                                                            target_patches_np.shape[1] * target_patches_np.shape[2]))

                            #if patch_size[xy] == 2:
                            #    code.interact(local=dict(globals(), **locals()))

                            melvin = timer()
                            source_len = icv_source_flag_valid.shape[0]
                            loop_len = 512 # def: 512
                            loop_num = int(np.ceil(source_len / loop_len))
                            print('\nLoop Information:')
                            print('Total number of source patches: ' + str(source_len))
                            print('Number of voxels processed in one loop: ' + str(loop_len))
                            print('Number of loop needed: ' + str(loop_num))
                            print('Check GPU memory: ' + str(cuda.current_context().get_memory_info()))

                            for il in range(0, loop_num):
                                ''' Debug purposed printing '''
                                print('.', end='')
                                if np.remainder(il+1, 32) == 0:
                                    print(' ' + str(il+1) + '/' + str(loop_num)) # Print newline

                                ''' Only process sub-array '''
                                source_patches_loop = area_source_patch_cuda_all[il*loop_len:(il*loop_len)+loop_len,:]

                                '''  SUBTRACTION '''
                                sub_result_gm = cuda.device_array((source_patches_loop.shape[0],
                                                                   target_patches_np_cuda_all.shape[0],
                                                                   target_patches_np_cuda_all.shape[1]))
                                TPB = (4,256)
                                BPGx = int(math.ceil(source_patches_loop.shape[0] / TPB[0]))
                                BPGy = int(math.ceil(target_patches_np_cuda_all.shape[0] / TPB[1]))
                                BPGxy = (BPGx,BPGy)
                                cu_sub_st[BPGxy,TPB](source_patches_loop, target_patches_np_cuda_all, sub_result_gm)

                                '''  MAX-MEAN-ABS '''
                                sub_max_mean_result = cuda.device_array((source_patches_loop.shape[0],
                                                                         target_patches_np_cuda_all.shape[0],2))
                                cu_max_mean_abs[BPGxy,TPB](sub_result_gm, sub_max_mean_result)
                                sub_result_gm = 0  # Free memory

                                '''  DISTANCE '''
                                distances_result = cuda.device_array((source_patches_loop.shape[0],
                                                                      target_patches_np_cuda_all.shape[0]))
                                cu_distances[BPGxy,TPB](sub_max_mean_result,
                                                        icv_source_flag_valid[il*loop_len:(il*loop_len)+loop_len],
                                                        distances_result, alpha)
                                sub_max_mean_result = 0  # Free memory

                                ''' SORT '''
                                TPB = 256
                                BPG = int(math.ceil(distances_result.shape[0] / TPB))
                                cu_sort_distance[BPG,TPB](distances_result)

                                ''' MEAN (AGE-VALUE) '''
                                idx_start = 8 # Starting index of mean calculation (to avoid bad example)
                                distances_result_for_age = distances_result[:,idx_start:idx_start+num_mean_samples]
                                distances_result = 0  # Free memory
                                cu_age_value[BPG,TPB](distances_result_for_age,
                                                      age_values_valid[il*loop_len:(il*loop_len)+loop_len])
                                distances_result_for_age = 0  # Free memory
                                del source_patches_loop  # Free memory
                                #code.interact(local=dict(globals(), **locals()))
                            print(' - Finished!\n')
                            print(timer() - melvin)
                            raise Exception()
                        ''' Mapping from age_value_valid to age value_all '''
                        if valid == 1:
                            index = 0
                            for idx_val in index_mapping:
                                if idx_val != -1:
                                    age_values_all[index] = age_values_valid[int(idx_val)]
                                index += 1

                        ''' Normalisation to probabilistic map (0...1) '''
                        if (np.max(age_values_all) - np.min(age_values_all)) == 0:
                            all_mean_distance_normed = age_values_all
                        else:
                            all_mean_distance_normed = np.divide((age_values_all - np.min(age_values_all)),
                                (np.max(age_values_all) - np.min(age_values_all)))

                        ''' SAVE Result (JPG) '''
                        slice_age_map = np.zeros([counter_x,counter_y])
                        index = 0
                        for ix in range(0, counter_x):
                            for iy in range(0, counter_y):
                                slice_age_map[ix,iy] = all_mean_distance_normed[index]
                                index += 1

                        ## Save mri_data
                        sio.savemat(dirOutData + '/' + str(patch_size[xy]) + '/' + str(zz) + '_dat.mat',
                                    {'slice_age_map':slice_age_map})

                        print('Check GPU memory: ' + str(cuda.current_context().get_memory_info()))
                        print('GPU flushing..\n--\n')
                        numba.cuda.profile_stop()
                    elapsed_times_patch_all[timer_idx,xy] = timer() - one_patch
                    print('IAM for MRI ID: ' + mri_code + ' with patch size: ' + str(patch_size[xy])
                          + ' elapsed for: ' + str(elapsed_times_patch_all[timer_idx,xy]))

                elapsed_times_all[timer_idx] = timer() - one_mri_data
                print('IAM for MRI ID: ' + mri_code + ' elapsed for: ' + str(elapsed_times_all[timer_idx]))
                timer_idx += 1

                ''' Save all elapsed times '''
                sio.savemat(dirOutput + '/elapsed_times_all_' + str(num_samples) + 's' + str(num_mean_samples) + 'm.mat',
                            {'elapsed_times_all':elapsed_times_all})
                sio.savemat(dirOutput + '/elapsed_times_patch_all_' + str(num_samples) + 's' + str(num_mean_samples) + 'm.mat',
                            {'elapsed_times_patch_all':elapsed_times_patch_all})
                ''' IAM's (GPU Part) Computation ENDS here '''

                '''
                KEY POINT: IAM's Combination, Penalisation, and Post-processing - START
                -----------------------------------------------------------------------
                Part 0 - Saving output results in .mat and JPEG files.
                Part 1 - Combination of multiple age maps.
                Part 2 - Global normalisation and penalisation of age maps based on brain tissues.
                Part 3 - Post-processing.

                Hint: You can search the keys of Part 0/1/2/3.
                '''
                combined_age_map_mri = np.zeros((x_len, y_len, z_len))
                combined_age_map_mri_mult = np.zeros((x_len, y_len, z_len))
                combined_age_map_mri_mult_normed = np.zeros((x_len, y_len, z_len))
                for zz in range(0, mri_data.shape[2]):
                    mri_slice = mri_data[:,:,zz]
                    mask_slice = np.nan_to_num(mri_slice)
                    mask_slice[mask_slice > 0] = 1
                    penalty_slice = np.nan_to_num(mri_slice)   ### PENALTY

                    slice_age_map_all = np.zeros((len(patch_size), x_len, y_len))

                    dirOutData = dirOutput + '/' + mri_code
                    for xy in range(0, len(patch_size)):
                        mat_contents = sio.loadmat(dirOutData + '/' + str(patch_size[xy]) + '/' + str(zz) + '_dat.mat')
                        slice_age_map = mat_contents['slice_age_map']
                        slice_age_map_res = cv2.resize(slice_age_map, None, fx=patch_size[xy],
                                                       fy=patch_size[xy], interpolation=cv2.INTER_CUBIC)
                        slice_age_map_res = skifilters.gaussian(slice_age_map_res,sigma=0.5,truncate=2.0)
                        #if zz== 20:
                        #    code.interact(local=dict(globals(), **locals()))
                        slice_age_map_res = np.multiply(mask_slice, slice_age_map_res)
                        slice_age_map_all[xy,:,:] = slice_age_map_res
                    slice_age_map_all = np.nan_to_num(slice_age_map_all)

                    if save_jpeg:
                        ''' >>> Part 0 <<<'''
                        ''' Show all age maps based on patch's size and saving the mri_data '''
                        fig, axes = plt.subplots(2, 2, sharex=True, sharey=True)
                        fig.set_size_inches(10, 10)
                        fig.suptitle('All Patches Gaussian Filtered', fontsize=16)

                        axes[0,0].set_title('Patch 1 x 1')
                        im1 = axes[0,0].imshow(np.rot90(slice_age_map_all[0,:,:]), cmap="jet", vmin=0, vmax=1)
                        divider1 = make_axes_locatable(axes[0,0])
                        cax1 = divider1.append_axes("right", size="7%", pad=0.05)
                        cbar1 = plt.colorbar(im1, ticks=[0, 0.5, 1], cax=cax1)

                        if len(patch_size) > 1:
                            axes[0,1].set_title('Patch 2 x 2')
                            im2 = axes[0,1].imshow(np.rot90(slice_age_map_all[1,:,:]), cmap="jet", vmin=0, vmax=1)
                            divider2 = make_axes_locatable(axes[0,1])
                            cax2 = divider2.append_axes("right", size="7%", pad=0.05)
                            cbar2 = plt.colorbar(im2, ticks=[0, 0.5, 1], cax=cax2)

                            if len(patch_size) > 2:
                                axes[1,0].set_title('Patch 4 x 4')
                                im3 = axes[1,0].imshow(np.rot90(slice_age_map_all[2,:,:]), cmap="jet", vmin=0, vmax=1)
                                divider3 = make_axes_locatable(axes[1,0])
                                cax3 = divider3.append_axes("right", size="7%", pad=0.05)
                                cbar3 = plt.colorbar(im3, ticks=[0, 0.5, 1], cax=cax3)

                                if len(patch_size) > 3:
                                    axes[1,1].set_title('Patch 8 x 8')
                                    im4 = axes[1,1].imshow(np.rot90(slice_age_map_all[3,:,:]), cmap="jet", vmin=0, vmax=1)
                                    divider4 = make_axes_locatable(axes[1,1])
                                    cax4 = divider4.append_axes("right", size="7%", pad=0.05)
                                    cbar4 = plt.colorbar(im4, ticks=[0, 0.5, 1], cax=cax4)

                        plt.tight_layout()
                        plt.subplots_adjust(top=0.95)

                        ''' >>> Part 0 <<<'''
                        ''' Save mri_data in *_all.jpg '''
                        dirOutData = dirOutput + '/' + mri_code + '/IAM_combined_python/Patch/'
                        fig.savefig(dirOutData + str(zz) + '_all.jpg', dpi=100)
                        print('Saving files: ' + dirOutData + str(zz) + '_all.jpg')
                        plt.close()

                    ''' >>> Part 1 <<< '''
                    ''' Combined all patches age map information '''
                    combined_age_map = 0
                    for bi in range(len(patch_size)):
                        combined_age_map += np.multiply(blending_weights[bi],slice_age_map_all[bi,:,:])
                    combined_age_map_mri[:,:,zz] = combined_age_map

                    ''' Global Normalisation - saving needed mri_data '''
                    combined_age_map_mri_mult[:,:,zz] = np.multiply(np.multiply(combined_age_map, penalty_slice), mask_slice)  ### PENALTY
                    normed_only = np.divide((combined_age_map_mri[:,:,zz] - np.min(combined_age_map_mri[:,:,zz])),\
                                            (np.max(combined_age_map_mri[:,:,zz]) - np.min(combined_age_map_mri[:,:,zz])))
                    normed_mult = np.multiply(np.multiply(normed_only, penalty_slice), mask_slice)  ### PENALTY
                    normed_mult_normed = np.divide((normed_mult - np.min(normed_mult)), \
                                            (np.max(normed_mult) - np.min(normed_mult)))
                    combined_age_map_mri_mult_normed[:,:,zz] = normed_mult_normed

                    ''' Save mri_data in *.mat '''
                    dirOutData = dirOutput + '/' + mri_code + '/IAM_combined_python/Patch/'
                    print('Saving files: ' + dirOutData + 'c' + str(zz) + '_combined.mat\n')
                    sio.savemat(dirOutData + 'c' + str(zz) + '_combined.mat', {'slice_age_map_all':slice_age_map_all,
                                                                'combined_age_map':normed_only,
                                                                'mri_slice_mul_normed':normed_mult_normed,
                                                                'combined_mult':combined_age_map_mri_mult[:,:,zz]})

                ''' >>> Part 2 <<< '''
                ''' Penalty + Global Normalisation (GN) '''
                combined_age_map_mri_normed = np.divide((combined_age_map_mri - np.min(combined_age_map_mri)),\
                                            (np.max(combined_age_map_mri) - np.min(combined_age_map_mri)))
                combined_age_map_mri_mult_normed = np.divide((combined_age_map_mri_mult - np.min(combined_age_map_mri_mult)),\
                                            (np.max(combined_age_map_mri_mult) - np.min(combined_age_map_mri_mult)))

                if save_jpeg:
                    for zz in range(0, mri_data.shape[2]):
                        fig2, axes2 = plt.subplots(1, 3)
                        fig2.set_size_inches(16,5)

                        axes2[0].set_title('Combined and normalised')
                        im1 = axes2[0].imshow(np.rot90(np.nan_to_num(combined_age_map_mri_normed[:,:,zz])), cmap="jet", vmin=0, vmax=1)
                        divider1 = make_axes_locatable(axes2[0])
                        cax1 = divider1.append_axes("right", size="7%", pad=0.05)
                        cbar1 = plt.colorbar(im1, ticks=[0, 0.5, 1], cax=cax1)

                        axes2[1].set_title('Combined, penalised and normalised')
                        im2 = axes2[1].imshow(np.rot90(np.nan_to_num(combined_age_map_mri_mult_normed[:,:,zz])), cmap="jet", vmin=0, vmax=1)
                        divider2 = make_axes_locatable(axes2[1])
                        cax2 = divider2.append_axes("right", size="7%", pad=0.05)
                        cbar2 = plt.colorbar(im2, ticks=[0, 0.5, 1], cax=cax2)

                        axes2[2].set_title('Original MRI slice')
                        im3 = axes2[2].imshow(np.rot90(np.nan_to_num(mri_data[:,:,zz])), cmap="gray")
                        divider3 = make_axes_locatable(axes2[2])
                        cax3 = divider3.append_axes("right", size="7%", pad=0.05)
                        cbar3 = plt.colorbar(im3, cax=cax3)

                        plt.tight_layout()
                        # Make space for title
                        plt.subplots_adjust(top=0.95)

                        ''' Save mri_data in *_combined.jpg '''
                        dirOutData = dirOutput + '/' + mri_code + '/IAM_combined_python/Combined/'
                        fig2.savefig(dirOutData + str(zz) + '_combined.jpg', dpi=100)
                        print('Saving files: ' + dirOutData + str(zz) + '_combined.jpg')
                        plt.close()

                ''' Save mri_data in *.mat '''
                sio.savemat(dirOutDataCom + '/all_slice_dat.mat', {'combined_age_map_all_slice':combined_age_map_mri,
                                                   'mri_slice_mul_all_slice':combined_age_map_mri_mult,
                                                   'combined_age_map_mri_normed':combined_age_map_mri_normed,
                                                   'combined_age_map_mri_mult_normed':combined_age_map_mri_mult_normed})

                '''
                combined_age_map_mri_img = nib.Nifti1Image(combined_age_map_mri_normed, mri_nii.affine)
                nib.save(combined_age_map_mri_img, str(dirOutDataFin + '/IAM_GPU_COMBINED.nii.gz'))

                combined_age_map_mri_GN_img = nib.Nifti1Image(combined_age_map_mri_mult_normed, mri_nii.affine)
                nib.save(combined_age_map_mri_GN_img, str(dirOutDataFin + '/IAM_GPU_GN.nii.gz'))
                '''

                ''' >>> Part 3 <<< '''
                ''' Post-processing '''
                ''' COMMENTED OUT BECAUSE NOT AVAILABLE
                if nawm_available and ~nawm_preprocessing:
                    combined_age_map_mri_mult_normed = np.multiply(combined_age_map_mri_mult_normed,nawm_mri_code)
                    combined_age_map_mri_GN_img = nib.Nifti1Image(combined_age_map_mri_mult_normed, mri_nii.affine)
                    nib.save(combined_age_map_mri_GN_img, str(dirOutDataFin + '/IAM_GPU_GN_postprocessed.nii.gz'))
                '''
                '''
                ---------------------------------------------------------------------
                KEY POINT: IAM's Combination, Penalisation, and Post-processing - END
                '''

                if delete_intermediary:
                    shutil.rmtree(dirOutDataCom, ignore_errors=True)
                    for xy in range(0, len(patch_size)):
                        shutil.rmtree(dirOutput + '/' + mri_code + '/' + str(patch_size[xy]), ignore_errors=True)

                del temp
                del center_source_patch, icv_source_flag
                del icv_source_flag_valid, index_mapping
                del area_source_patch, target_patches_np   # Free memory
                del area_source_patch_cuda_all, target_patches_np_cuda_all   # Free memory
                gc.collect()

        ## Print the elapsed time information
        print('\n--\nSpeed statistics of this run..')
        print('mean elapsed time  : ' + str(np.mean(elapsed_times_all)) + ' seconds')
        print('std elapsed time   : ' + str(np.std(elapsed_times_all)) + ' seconds')
        print('median elapsed time : ' + str(np.median(elapsed_times_all)) + ' seconds')
        print('min elapsed time   : ' + str(np.min(elapsed_times_all)) + ' seconds')
        print('max elapsed time   : ' + str(np.max(elapsed_times_all)) + ' seconds')


<gh_stars>0
"""A library for 1D nodal-DG numerics."""

import numpy as np
import scipy.special as ss

def jacobi_gq(alpha, beta, n):
    """Computes the n'th order Gauss quadrature points and weights."""
    p = ss.jacobi(n + 1, alpha, beta)
    r = np.sort(p.r)
    i = np.argsort(p.r)
    w = p.weights[i]
    return r, w

def jacobi_gl(alpha, beta, n):
    """Computes the n'th order Gauss lobatto quadrature points."""
    if n == 1:
        x = np.zeros((2,))
        x[0] = -1.0
        x[1] = 1.0
        return x

    x, _ = jacobi_gq(alpha + 1, beta + 1, n - 2)
    return np.concatenate([[-1.0], x, [1.0]])

def jacobi_p(x, alpha, beta, n):
    """Computes the Jacobi Polynomial.

    Args:
      x: Nodal points where the polynomial is computed.
      alpha: Parameter of the Jacobi polynomial, must be greeater than -1.
      beta: Parameter of the Jacobi polynomial, must be greeater than -1.
      n: Degree of the polynomial.

    Returns:
      The Jacobi polynomial at points `r` for order `n`.
    """
    p = np.zeros((len(x), 2))

    # Initial values of p_0(x) and p_1(x).
    gamma_0 = 2**(alpha + beta + 1) / (alpha + beta + 1) * \
            ss.gamma(alpha + 1) * ss.gamma(beta + 1) / \
            ss.gamma(alpha + beta + 1)
    p[:, 0] = 1.0 / np.sqrt(gamma_0)
    if n == 0:
        return p[:, 0]
    gamma_1 = (alpha + 1) * (beta + 1) / (alpha + beta + 3) * gamma_0
    p[:, 1] = ((alpha + beta + 2) * x / 2 + (alpha - beta) / 2) / \
            np.sqrt(gamma_1)
    if n == 1:
        return p[:, 1]

    a_old = 2.0 / (2.0 + alpha + beta) * \
            np.sqrt((alpha + 1) * (beta + 1) / (alpha + beta + 3))

    for i in range(1, n):
        h_1 = 2 * i + alpha + beta
        a_new = 2.0 / (h_1 + 2) * np.sqrt((i + 1) * \
                (i + 1 + alpha + beta) * (i + 1 + alpha) * \
                (i + 1 + beta) / (h_1 + 1) / (h_1 + 3))
        b_new = -(alpha**2 - beta**2) / h_1 / (h_1 + 2)
        p_new = 1.0 / a_new * (-a_old * p[:, 0] + (x - b_new) * p[:, 1])
        p[:, 0] = p[:, 1]
        p[:, 1] = p_new
        a_old = a_new

    return p[:, 1]

def grad_jacobi_p(r, alpha, beta, n):
    """Computes the derivative of the Jacobi Polynomial.

    Args:
      r: Nodal points where the derivative is computed.
      alpha: Parameter of the Jacobi polynomial, must be greeater than -1.
      beta: Parameter of the Jacobi polynomial, must be greeater than -1.
      n: Degree of the polynomial.

    Returns:
      The derivative of the Jacobi polynomial at points `r` for order `n`.
    """
    if n < 0:
        raise ValueError(
            'Degree of the polynomial must be non-negative. '
            '{} is given.'.format(n))
    dp = np.zeros((len(r), 1))
    if n > 0:
        dp = (np.sqrt(n * (n + alpha + beta + 1.0)) *
                jacobi_p(r, alpha + 1, beta + 1, n - 1))
    return dp

def vandermonde_1d(n, r):
    """Initializes the 1D Vandermonde matrix, V_{ij} = phi_j(r_i)."""
    v = np.zeros((len(r), n + 1))
    for j in range(n + 1):
        v[:, j] = jacobi_p(r, 0, 0, j)
    return v

def grad_vandermonde_1d(n, r):
    """Initializes the gradient of the modal basis at `r` at order `n`."""
    d_vr = np.zeros((len(r), n + 1))
    for i in range(n + 1):
        d_vr[:, i] = np.squeeze(grad_jacobi_p(r, 0, 0, i))
    return d_vr

def d_matrix_1d(n, r, v):
    """Initializes the differentiation matrices on the interval.

    Args:
      n: The order of the polynomial.
      r: The nodal points.
      v: The Vandemonde matrix.

    Returns:
      The gradient matrix D.
    """
    vr = grad_vandermonde_1d(n, r)
    return np.linalg.lstsq(v.T, vr.T, rcond=None)[0].T

def minmod(v):
    """Computes the minmod function."""
    m, n = v.shape
    s = np.sum(np.sign(v), axis=0) / m
    return np.where(np.abs(s) == 1, s * np.min(np.abs(v), axis=0), 0)

def minmodb(v, max_dd2, h):
    """Computes the TVB modified minmod function.

    Args:
      v: A 2D array with shape (n, k) representing k vectors of length n.
      max_dd2: The upper bound on the second derivative at the local
        extrema.
      h: The grid spacing.

    Returns:
      The modified minmod function.
    """
    m_fn = v[0, :]
    ids = np.where(np.abs(m_fn) > max_dd2 * h**2)

    if len(ids[0]) > 0:
        m_fn[ids] = minmod(v[:, ids])

    return m_fn

def slope_limit_lin(u_x, x_l, v_m1, v_0, v_p1):
    """Applies slope limiter to a linear function u.

    Args:
      u_x: The gradient of a linear polynomial of size (n_p, 1).
      x_l: The collocation points of the linear element. The size of x_l is
        (n_p, 1).
      v_m1: The average of the left element.
      v_0: The average of the current element.
      v_p1: The average of the right element.

    Returns:
      The slope limited polynomial in the current element.
    """
    n_p, k = u_x.shape
    if x_l.shape[0] != n_p or x_l.shape[1] != k:
        raise ValueError('The shape of input polynomial and collocation '
                'points mismatches: u_l ({}), x_l({})'.format( \
                        u_x.shape, x_l.shape))
    # Computes the grid size.
    h = x_l[-1, :] - x_l[0, :]
    h_n = np.tile(h, (n_p, 1))

    # Computes the coordinates at the middle point of elements.
    x_0 = np.tile(x_l[0, :] + 0.5 * h, (n_p, 1))

    # Computes the limit function.
    return np.tile(v_0, (n_p, 1)) + (x_l - x_0) * ( \
            np.tile( \
            minmod(np.array( \
            [u_x[0, :], (v_p1 - v_0) / h, (v_0 - v_m1) / h])), \
            (n_p, 1)))


"""
    Test program to cac laulte the relation between focal length and
    the accomodation parameter
"""

from poptics.lens import Eye
import poptics.wavelength as wl
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit

def fit(x,a,b,c,d):
    """
    the fit function for accomodation against focal length
    """
    return a*x**3 + b*x**2 + c*x + d

def main():

    #     Mahe default eye
    eye = Eye(0.0)
    print("Original Focalength : " + str(eye.backFocalLength()))

    #      Make a range of accomodation parameters, Note 1.0 is no accomodatiomn
    acc = np.linspace(1.0,1.5,50)
    focal = np.zeros(acc.size)

    #      Go round loops finding focal length for each accomodation value
    for i,a in enumerate(acc):
        eye.accommodation(a)
        focal[i] = eye.backFocalLength(wl.PhotopicPeak)

    #     Do a fit
    popt,pvar = curve_fit(fit,focal,acc)
    print(str(popt))           # print out fite values

    #     No a plot to make sure it is sensible.
    plt.plot(focal,acc)
    plt.plot(focal,fit(focal,*popt),"x")
    plt.show()

main()