text:
import torch
import numpy as np
import gym
import os
import pickle
from absl import app
from absl import flags

# from utils import _windows_enable_ANSI
from model import PNet, VNet
import ralamb
from trajectory import Buffer, Trajectory
from agent import Agent, Discriminator

FLAGS = flags.FLAGS
flags.DEFINE_string("env", "Taxi-v3", "gym env name")
flags.DEFINE_bool("use_gpu", True, "using GPU")
flags.DEFINE_integer("hid_num", 16, "hidden layer unit size")
flags.DEFINE_float("lr", 1e-4, "learning rate")
flags.DEFINE_float("final_lr", 1e-2, "final learning rate for AdaBound")
flags.DEFINE_float("weight_decay", 1e-6, "weight decay")
flags.DEFINE_float("gamma", 0.99, "reward discount factor")
flags.DEFINE_float("lam", 0.95, "GAE lambda")
flags.DEFINE_float("eps", 1e-6, "small value for preventing 0 division in Advantage")
flags.DEFINE_integer("batch_size", 16, "batch size for each training")
flags.DEFINE_float("clip_param", 0.2, "PPO clipping liklihood ratio")
flags.DEFINE_float("ent_c", 0.2, "entropy coefficient")
flags.DEFINE_float("v_loss_c", 0.9, "value loss coefficient")
flags.DEFINE_float("max_grad_norm", 0.5, "maximum gradient norm")
flags.DEFINE_bool("use_discrim", False, "using AIRL Discriminator")

flags.DEFINE_integer("n_sample_episodes", 10, "sampling episode in each training epoch")
flags.DEFINE_integer("n_train_epochs", 10000, "training epoch")
flags.DEFINE_integer("n_opt_epochs", 3, "update epoch in each iteration")
flags.DEFINE_integer("n_test_epochs", 10, "test epoch")
flags.DEFINE_bool("render_env", True, "env render in test epoch")

flags.DEFINE_string("data_dir", "./data", "data save directory")
# flags.DEFINE_string("data_dir", "/content/drive/My Drive/data", "data save directory")    # Google Drive
flags.DEFINE_bool("resume", True, "load model parameters")

def create_directory(dir_name):
    if not os.path.exists(dir_name):
        print("Create directory: {0}".format(dir_name))
        os.mkdir(dir_name)

def get_device(use_gpu):
    # サポート対象のGPUがあれば使う
    if use_gpu:
        print("Check GPU available")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device("cpu")
    if device == "cuda":
        torch.backends.cudnn.benchmark = True
    return device

def train(env, agent, max_rew, model_filename_base, device, discrim, discrim_filename_base, expert_traj):
    print("Train start")
    try:
        for epoch in range(FLAGS.n_train_epochs):
            # 評価モード
            agent.eval()
            traj = Trajectory()
            episode_rewards = []
            if discrim is not None:
                real_episode_rewards = []
            for episode in range(FLAGS.n_sample_episodes):
                ob = env.reset()
                done = False
                buf = Buffer()
                while not done:
                    ob_tsr = torch.from_numpy(np.array(ob)).to(device=device)  # テンソル化
                    action, value, llh, _ = agent.get_action(ob_tsr)
                    next_ob, reward, done, info = env.step(action)
                    buf_data = dict(obs=ob, actions=action, rewards=reward, values=value, llhs=llh)

                    if discrim is not None:
                        pseudo_rew = discrim.get_pseudo_reward(ob_tsr)
                        buf_data["real_rewards"] = buf_data["rewards"]
                        buf_data["rewards"] = pseudo_rew
                        buf_data["next_obs"] = next_ob
                        buf_data["dones"] = done

                    buf.append(buf_data)
                    ob = next_ob

                buf.finish_path(FLAGS.gamma, FLAGS.lam)
                traj.append(buf.data_map)
                episode_rewards.append(buf.data_map["rewards"].sum())

                if discrim is not None:
                    real_episode_rewards.append(buf.data_map["real_rewards"].sum())

            # 学習モード
            agent.train()
            traj.finish_path(FLAGS.eps, device)
            pol_loss, v_loss, entropy = agent.update(
                traj,
                FLAGS.batch_size,
                FLAGS.n_opt_epochs,
                FLAGS.clip_param,
                FLAGS.v_loss_c,
                FLAGS.ent_c,
                FLAGS.max_grad_norm,
            )
            if discrim is not None:
                discrim.train()
                discrim_a_loss, discrim_e_loss = discrim.update(traj, expert_traj, FLAGS.batch_size, FLAGS.gamma, agent)

            mean_rew = np.mean(episode_rewards)
            print("Epoch[{:5d}/{:5d}], ".format(epoch + 1, FLAGS.n_train_epochs), end="")
            if discrim is not None:
                print(
                    "Loss(P/V/E/Da/De)[{:+.6f}/{:+.6f}/{:+.6f}/{:+.6f}/{:+.6f}], ".format(
                        pol_loss, v_loss, entropy, discrim_a_loss, discrim_e_loss
                    ),
                    end="",
                )
            else:
                print("Loss(P/V/E)[{:+.6f}/{:+.6f}/{:+.6f}], ".format(pol_loss, v_loss, entropy), end="")
            reward_format = "Reward(Mean/Min/Max)[{:+3.3f}/{:+3d}/{:+3d}]"
            print(reward_format.format(mean_rew, int(np.min(episode_rewards)), int(np.max(episode_rewards))), end="")
            if discrim is not None:
                mean_real_rew = np.mean(real_episode_rewards)
                reward_format = ", Reward(Real)(Mean/Min/Max)[{:+3.3f}/{:+3.3f}/{:+3.3f}]"
                print(
                    reward_format.format(mean_real_rew, np.min(real_episode_rewards), np.max(real_episode_rewards)),
                    end="",
                )
            print()

            if (((epoch + 1) % 10) == 0) and (mean_rew > max_rew):
                print("Reward Update: Old[{:+3.3f}] -> New[{:+3.3f}]".format(max_rew, mean_rew))
                max_rew = mean_rew
                print("Save best model")
                save_info = dict(max_rew=max_rew)
                agent.save_model(model_filename_base, save_info, "best")

                if discrim is not None:
                    discrim.save_model(discrim_filename_base, "best")

            if ((epoch + 1) % 100) == 0:
                print("Save checkpoint model")
                save_info = dict(max_rew=max_rew)
                agent.save_model(model_filename_base, save_info, "checkpoint")

                if discrim is not None:
                    discrim.save_model(discrim_filename_base, "checkpoint")

    except KeyboardInterrupt:
        print("Keyboard interrupt")
    print("Train complete")

    print("Save last model")
    save_info = dict(max_rew=max_rew)
    agent.save_model(model_filename_base, save_info, "last")

    if discrim is not None:
        discrim.save_model(discrim_filename_base, "last")

def test(env, agent, device):
    print("Test start")
    try:
        # 評価モード
        agent.eval()
        test_reward = []
        for epoch in range(FLAGS.n_test_epochs):
            ob = env.reset()
            done = False
            step = 0
            episode_reward = 0
            if FLAGS.render_env:
                print("-- env[{0:02d}:{1:03d}] ".format(epoch + 1, step) + "-" * 30)
                env.render()

            while not done:
                ob_tsr = torch.from_numpy(np.array(ob)).to(device=device)
                action, value, _, pi = agent.get_action(ob_tsr)
                next_ob, reward, done, info = env.step(action)
                episode_reward += reward
                ob = next_ob

                step += 1
                if FLAGS.render_env:
                    print("-- env[{0:02d}:{1:03d}] ".format(epoch + 1, step) + "-" * 30)
                    env.render()
                    print(info, (pi * 1000).astype(int), "{:.6f}".format(value))

            episode_reward = np.mean(episode_reward)
            test_reward.append(episode_reward)
            print("Epoch[{:3d}], ".format(epoch + 1), end="")
            print("Reward[{0}]".format(episode_reward), end="")
            print()
        print("Mean reward[{0}]".format(np.mean(test_reward)))
    except KeyboardInterrupt:
        print("Keyboard interrupt")
    print("Test complete")

def main(_):
    device = get_device(FLAGS.use_gpu)
    print("Use device: {}".format(device))

    # モデル保存用フォルダ生成
    data_dir = FLAGS.data_dir
    create_directory(data_dir)
    create_directory(os.path.join(data_dir, "models"))

    env = gym.make(FLAGS.env)

    p_net = PNet(env.observation_space, env.action_space, FLAGS.hid_num)
    v_net = VNet(env.observation_space, FLAGS.hid_num)
    print(p_net)
    print(v_net)
    p_net.to(device)
    v_net.to(device)
    optim_p = ralamb.Ralamb(p_net.parameters(), lr=FLAGS.lr, weight_decay=FLAGS.weight_decay)
    optim_v = ralamb.Ralamb(v_net.parameters(), lr=FLAGS.lr, weight_decay=FLAGS.weight_decay)
    agent = Agent(p_net, v_net, optim_p, optim_v, device)

    if FLAGS.use_discrim:
        expert_filename = os.path.join(FLAGS.data_dir, "expert_data", "taxi_expert.pkl")
        print("Load expert data: ", expert_filename)
        with open(expert_filename, "rb") as f:
            expert_traj = Trajectory()
            expert_epis = pickle.load(f)
            for epi in expert_epis:
                epi["next_obs"] = np.append(epi["obs"][1:], epi["obs"][0])
                expert_traj.append(epi)
            expert_traj.to_tensor(device)

        pseudo_rew_net = VNet(env.observation_space, FLAGS.hid_num)
        shaping_val_net = VNet(env.observation_space, FLAGS.hid_num)
        print(pseudo_rew_net)
        print(shaping_val_net)
        pseudo_rew_net.to(device)
        shaping_val_net.to(device)
        optim_discrim = ralamb.Ralamb(
            list(pseudo_rew_net.parameters()) + list(shaping_val_net.parameters()),
            lr=FLAGS.lr,
            weight_decay=FLAGS.weight_decay,
        )
        discrim = Discriminator(pseudo_rew_net, shaping_val_net, optim_discrim, device)
    else:
        discrim = None
        expert_traj = None

    # モデルの途中状態を読込む
    max_rew = -1e6
    model_filename_base = os.path.join(FLAGS.data_dir, "models", "model_" + FLAGS.env + "_PPO_H" + str(FLAGS.hid_num))
    discrim_filename_base = None
    if FLAGS.resume:
        print("Load last model")
        load_info = agent.load_model(model_filename_base, "last")
        if load_info:
            max_rew = load_info["max_rew"]
            print("Max reward: {0}".format(max_rew))
        else:
            print("Model file not found")

        if FLAGS.use_discrim:
            discrim_filename_base = os.path.join(
                FLAGS.data_dir, "models", "discrim_" + FLAGS.env + "_AIRL_H" + str(FLAGS.hid_num)
            )
            discrim.load_model(discrim_filename_base, "last")

    train(env, agent, max_rew, model_filename_base, device, discrim, discrim_filename_base, expert_traj)
    test(env, agent, device)

if __name__ == "__main__":
    # _windows_enable_ANSI()
    app.run(main)


text:
# -*- coding: utf-8 -*-
"""
Utilities to help in the DTI pre-processing
"""
from nipype.algorithms.rapidart import ArtifactDetect
from src.env import DATA, ATLAS_TYPES
from src.postproc.utils import execute
import os
from os.path import join as opj
import nibabel as nib
import numpy as np

def nlmeans_denoise(in_file, mask_file, out_file='', N=12):
    """ Filepath interface to the nlmeans_denoise_img in pypes.preproc."""
    import os.path as op
    import nibabel as nib

    from src.preproc import nlmeans_denoise_img
    from src.utils import rename

    den = nlmeans_denoise_img(nib.load(in_file), mask=nib.load(mask_file), N=N)

    if not out_file:
        out_file = rename(in_file, '_denoised')

    den.to_filename(out_file)

    return op.abspath(out_file)

def reslice(in_file, new_zooms=None, order=3, out_file=''):
    """
    Performs regridding of an image to set isotropic voxel sizes using dipy.
    """
    import os.path as op

    import nibabel as nib

    from src.preproc import reslice_img
    from src.utils import rename

    img = reslice_img(nib.load(in_file), new_zooms=new_zooms, order=order)

    if not out_file:
        out_file = rename(in_file, '_resliced')

    img.to_filename(out_file)

    return op.abspath(out_file)

def dti_acquisition_parameters(in_file, epi_factor=128):
    """ Return the path to the acqp_file and index file necessary for the new Eddy tool.
    This works when `in_file` is a nifti file obtained from a recent version from dcm2nii.

    This assumes a standard Siemens acquisition if the 'descrip.phaseDir' field is not
    in the nifti header.

    Please check the InPlanePhaseEncodingDirection field in the DICOM files, if its values is
    'COL', you are good to go.

    Notes
    -----
    # Comments on the `eddy` tool from FSL FDT.

    A description of the tool:
    http://fsl.fmrib.ox.ac.uk/fsl/fslwiki/Eddy/UsersGuide

    Our problem to run this tool instead of the good-old `eddy_correct` is the `--acqp` argument, an
    acquisitions parameters file.

    A detailed description of the --acpq input file is here:
    http://fsl.fmrib.ox.ac.uk/fsl/fslwiki/eddy/Faq#How_do_I_know_what_to_put_into_my_--acqp_file

    In the following subsections I describe each of the fields needed to check and build the acquisitions parameters
    file.

    ## Phase Encoding Direction
    Dicom header field: (0018,1312) InPlanePhaseEncodingDirection
    The phase encoding direction is the OPPOSITE of frequency encoding direction:
    - 'COL' = A/P (freqDir L/R),
    - 'ROW' = L/R (freqDir A/P)

    Nifti header field: "descrip.phaseDir:'+'" is for 'COL' in the DICOM InPlanePhaseEncodingDirection value.
    So if you have only one phase encoding oriendation and a '+' in the previous header field,
    the first 3 columns for the `--acqp` parameter file should be:
    0 1 0

    indicating that the scan was acquired with phase-encoding in the anterior-posterior direction.

    For more info:
    http://web.stanford.edu/group/vista/cgi-bin/wiki/index.php/DTI_Preprocessing_User_Manual#Frequency_vs_Phase_Encode_Direction
    https://xwiki.nbirn.org:8443/xwiki/bin/view/Function-BIRN/PhaseEncodeDirectionIssues

    ## Effective Echo Spacing (aka dwell time)

    Effective Echo Spacing (s) = 1/(BandwidthPerPixelPhaseEncode * MatrixSizePhase)

    effective echo spacing = 1 / [(0019,1028) * (0051,100b component #1)] (from the archives)
    https://www.jiscmail.ac.uk/cgi-bin/webadmin?A3=ind1303&L=FSL&E=quoted-printable&P=29358&B=--B_3444933351_15386849&T=text%2Fhtml;%20charset=ISO-8859-1&pending=

    The dwell time is in the nifti header `descrip.dwell`, in seconds (or look at the field `time_units`).
    http://www.mit.edu/~satra/nipype-nightly/interfaces/generated/nipype.interfaces.fsl.epi.html

    More info:
    http://lcni.uoregon.edu/kb-articles/kb-0003

    ## EPI factor

    The EPI factor is not included in the nifti header.
    You can read it using the Grassroots DICOM tool called `gdcmdump`, for example:
    >>> gdcmdump -C IM-0126-0001.dcm | grep 'EPIFactor'
    sFastImaging.lEPIFactor                  = 128

    More info:
    http://dicomlookup.com/default.htm

    # The fourth element of the acquisitions parameter file

    The fourth element in each row is the time (in seconds) between reading the center of the first echo and reading the
    center of the last echo.
    It is the "dwell time" multiplied by "number of PE steps - 1" and it is also the reciprocal of the PE bandwidth/pixel.

    Total readout time (FSL) = (number of echoes - 1) * echo spacing

    Total Readout Time (SPM) = 1/(BandwidthPerPixelPhaseEncode)
    Since the Bandwidth Per Pixel Phase Encode is in Hz, this will give the readout time in seconds

    # The `---index` argument

    The index argument is a text file with a row of numbers. Each number
    indicates what line (starting from 1) in the `acqp` file corresponds to
    each volume in the DTI acquisition.

    # What to do now with the `dcm2nii` files?

    I see two options to calculate the `acqp` lines with these files.

    1. We already have the `dwell` but we don't have the EPI factor.
    We know that the standard in Siemens is 128 and we could stick to that.

    2. Use the `slice_duration * 0.001` which is very near the calculated value.

    # Summary

    So, for example, if we had these acquisition parameters:

    ```
    Phase enc. dir. P >> A
    Echo spacing 0.75 [ms]
    EPI factor 128
    ```

    We should put in the `acqp` file this line:
    0 1 0 0.095
    """
    # the imports must be inside if you want them to work in a nipype.Function node.
    import warnings
    import os.path as op
    import nibabel as nib

    acqp_file = "diff.acqp"
    index_file = "diff.index"

    image = nib.load(in_file)
    n_volumes = image.shape[-1]

    # read the value of the descrip field in the nifti header
    descrip_field = image.header["descrip"].astype(str)[()]
    descrip = dict([item.split("=", 1) for item in descrip_field.split(";")])

    # assume by default phaseDir = '+'
    phase_dir = descrip.get("phase", '+')
    pe_axis   = "0 1 0"
    if phase_dir == "-":
        pe_axis = "0 -1 0"

    if 'phase' not in descrip:
        warnings.warn("Could not find or understand the value for phaseDir: {}. "
                      "Using default PE axis {}.".format(descrip.get("phase"), pe_axis),
                      category=RuntimeWarning, stacklevel=2)

    # (number of phase-encode steps - 1) *
    # (echo spacing time in milliseconds) *
    # (seconds per millisecond)
    total_readout_time = (epi_factor - 1) * float(descrip["dwell"]) * 1e-3

    with open(acqp_file, "wt") as fout:
        fout.write("{} {}\n".format(pe_axis, total_readout_time))

    with open(index_file, "wt") as fout2:
        fout2.write("{}\n".format(" ".join(n_volumes * ["1"])))

    return op.abspath(acqp_file), op.abspath(index_file)

def correct_dwi_space_atlas(subject_list, session_list):
    """
    Function to correct atlas after resampling (looses some rois),
    this function recovers those lost rois
    """
    sub_ses_comb = [[subject, session] for subject in subject_list
                    for session in session_list]

    for sub, ses in sub_ses_comb:
        for atlas in ATLAS_TYPES:

            atlas_old = opj(DATA, 'external',
                            'bha_' + atlas + '_1mm_mni09c.nii.gz')
            atlas_new = opj(DATA, 'processed', 'diff', '_session_id_' +
                            ses + '_subject_id_' + sub,
                            'r' + sub + '_' + ses + '_' + atlas + '.nii')

            atlas_new_img = nib.load(atlas_new)
            m = atlas_new_img.affine[:3, :3]

            atlas_old_data = nib.load(atlas_old).get_data()
            atlas_old_data_rois = np.unique(atlas_old_data)
            atlas_new_data = atlas_new_img.get_data()
            atlas_new_data_rois = np.unique(atlas_new_data)

            diff_rois = np.setdiff1d(atlas_old_data_rois, atlas_new_data_rois)

            for roi in diff_rois:
                p = np.argwhere(atlas_old_data == roi)[0]
                x, y, z = (np.round(np.diag(np.divide(p, m)))).astype(int)
                atlas_new_data[x, y, z] = roi

            atlas_new_data_img_corrected = nib.Nifti1Image(atlas_new_data,
                                                           affine=atlas_new_img.affine)
            os.remove(atlas_new)
            nib.save(atlas_new_data_img_corrected,
                     atlas_new)

def compress_file_gz(path_file):
    from subprocess import check_call

    check_call(['gzip', path_file])
    check_call(['rm', path_file])

    return path_file + '.gz'

def get_con_matrix_matlab(subject_list, session_list):
    """
    Function to get atlas-wise connectivity matrix. End-to-end fiber counting.
    Camino's workflow is not calculating end-to-end
    """
    sub_ses_comb = [[subject, session] for subject in subject_list
                    for session in session_list]

    for sub, ses in sub_ses_comb:
        for atlas in ATLAS_TYPES:

            num_nodes = atlas[-4:]
            out_folder = opj(DATA, 'processed', 'tract',
                             '_session_id_' + ses + '_subject_id_' + sub)
            tracts_file = opj(out_folder, 'tracts.Bfloat_' + num_nodes)
            reference_file = opj(DATA, 'processed', 'diff',
                                 '_session_id_' + ses + '_subject_id_' + sub,
                                 'eddy_corrected_avg_b0.nii.gz')
            atlas_file = opj(DATA, 'processed', 'diff',
                             '_session_id_' + ses + '_subject_id_' + sub,
                             'r' + sub + '_' + ses + '_' + atlas + '.nii')

            try:
                atlas_file = compress_file_gz(atlas_file)
            except:
                atlas_file = atlas_file + '.gz'

            matlab_path = opj(os.getcwd(), 'src', 'matlab')

            # Extract brain from subject space
            command = ["matlab",
                       "-nodisplay",
                       "-nojvm",
                       "-r",
                       "addpath(\'" + matlab_path + "\');" +
                       "calc_cm(" +
                       "\'" + tracts_file + "\'," +
                       "\'" + reference_file + "\'," +
                       "\'" + atlas_file + "\'," +
                       "\'" + out_folder + "\'," +
                       num_nodes +
                       "); exit;"
                       ]
            for output in execute(command):
                print(output)

ext: py

max_stars_repo_name: erramuzpe/ruber

max_issues_repo_name: erramuzpe/ruber

max_forks_repo_name: erramuzpe/ruber

reason: import numpy


text:
# Multiple landmark detection in 3D ultrasound images of fetal head
# Network inference with evaluations
#
# Reference
# Fast Multiple Landmark Localisation Using a Patch-based Iterative Network
# https://arxiv.org/abs/1806.06987
#
# Code EDITED BY: Xingchen Xiao
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import scipy.ndimage
import time
import torch
from utils import input_data_torch, shape_model_func, patch, save, visual
from network_torch import VisitNet
np.random.seed(0)

class Config(object):
    """Inference configurations."""
    # File paths
    data_dir = './data_2d/Images'
    label_dir = '../landmark_detection_PIN/data_2d/Landmarks_old'
    # label_dir = '../landmark_detection_PIN/data_2d/landmarks'
    train_list_file = './data_2d/train_260/list_train.txt'
    test_list_file = './data_2d/train_260/list_test.txt'
    # train_list_file = './data_2d/list_train.txt'
    # test_list_file = './data_2d/list_test.txt'
    model_dir = './cnn_model'
    # Shape model parameters
    shape_model_file = './shape_model/shape_model_260_3/ShapeModel.mat'
    shape_model_file = './shape_model/shape_model_260_3/ShapeModel.mat'
    eigvec_per = 0.995      # Percentage of eigenvectors to keep
    sd = 3.0                # Standard deviation of shape parameters
    landmark_count = 3   # Number of landmarks
    landmark_unwant = []     # list of unwanted landmark indices
    # Testing parameters
    box_size = 121          # patch size (odd number)
    max_test_steps = 5     # Number of inference steps
    num_random_init = 3    # Number of random initialisations used
    predict_mode = 1        # How the new patch position is computed.
                            # 0: Classification and regression. Hard classification
                            # 1: Classification and regression. Soft classification. Multiply classification probabilities with regressed distances
                            # 2: Regression only
                            # 3: Classification only
    # Visualisation parameters
    visual = True           # Whether to save visualisation
    dropout = 0.8

def main():
    config = Config()

    # Load shape model
    shape_model = shape_model_func.load_shape_model(config.shape_model_file, config.eigvec_per)
    num_cnn_output_c = 2 * shape_model['Evectors'].shape[1]
    num_cnn_output_r = shape_model['Evectors'].shape[1]

    # Load images and landmarks
    data = input_data_torch.read_data_sets(config.data_dir,
                                     config.label_dir,
                                     config.train_list_file,
                                     config.test_list_file,
                                     config.landmark_count,
                                     config.landmark_unwant,
                                     shape_model)

    print("Start inference...")
    model = VisitNet(num_cnn_output_c,num_cnn_output_r,config.dropout)

    model.load_state_dict(torch.load('../landmark_detection_PIN/cnn_model/model2_train_260_99999.pt'))
    model.eval()

    with torch.no_grad():
        # # Evaluation on test-set
        predict(data[1], config, shape_model, False,
                model)

        # # # Evaluation on train-set
        # predict(data[0], config, shape_model, True,
        #         model)

def predict(data, config, shape_model, train,
            model):
    """Find the path of the landmark iteratively, and evaluate the results.

    Args:
      data: dataset
      config: testing parameters
      shape_model: structure containing shape model
      train: True: train or False: test dataset
      sess, x, action_ind, yc, yr, keep_prob: tensorflow nodes required for inference

    """

    images = data.images
    landmarks_gt = data.labels
    names = data.names
    pix_dim = data.pix_dim
    num_landmarks = config.landmark_count
    img_count = len(images)
    max_test_steps = config.max_test_steps
    num_examples = config.num_random_init + 1

    landmarks_all_steps = np.zeros((img_count, max_test_steps + 1, num_examples, num_landmarks, 3))
    landmarks_mean = np.zeros((img_count, num_landmarks, 3), dtype=np.float32)
    landmarks_mean_unscale = np.zeros((img_count, num_landmarks, 3), dtype=np.float32)
    landmarks_gt_unscale = np.zeros((img_count, num_landmarks, 3), dtype=np.float32)
    images_unscale = []
    time_elapsed = np.zeros(img_count)

    for i in range(img_count):
        # Predict landmarks iteratively
        start_time = time.time()
        landmarks_all_steps[i], landmarks_mean[i] = predict_landmarks(images[i], config, shape_model,
                      model)
        end_time = time.time()
        time_elapsed[i] = end_time - start_time

        # Convert the scaling back to that of the original image.
        landmarks_mean_unscale[i] = landmarks_mean[i] * pix_dim[i, :]
        landmarks_gt_unscale[i] = landmarks_gt[i] * pix_dim[i, :]
        images_unscale.append(scipy.ndimage.zoom(images[i][:, :, :, 0], pix_dim[i]))

        print("Image {}/{}: {}, time = {:.10f}s".format(i+1, img_count, names[i], time_elapsed[i]))

    # Time
    time_elapsed_mean = np.mean(time_elapsed)
    print("Mean running time = {:.10f}s\n".format(time_elapsed_mean))

    # Evaluate distance error
    err, err_mm = compute_err(landmarks_mean, landmarks_gt, pix_dim)

    # Save distance error to txt file
    save.save_err('./results/dist_err', train, names, err, err_mm)

    # Save predicted landmarks as txt files. Landmarks are in voxel coordinates. Not in CNN coordinates.
    save.save_landmarks('./results/landmarks', train, names, landmarks_mean_unscale, config.landmark_unwant)

    # Visualisation
    # if config.visual:
    #     print("Show visualisation...")
    #     for i in range(img_count):
    #         print("Processing visualisation {}/{}: {}".format(i+1, img_count, names[i]))
    #         # visual.plot_landmarks_2d('./results/landmarks_visual2D', train, names[i], images_unscale[i],
    #         #                          landmarks_mean_unscale[i], landmarks_gt_unscale[i])
    #         visual.plot_landmarks_3d('./results/landmarks_visual3D', train, names[i], landmarks_mean[i],
    #                                  landmarks_gt[i], images[i].shape)
    #         visual.plot_landmarks_path('./results/landmark_path', train, names[i], landmarks_all_steps[i],
    #                                    landmarks_gt[i], images[i].shape)

def predict_landmarks(image, config, shape_model,
                      model):
    """Predict landmarks iteratively.

    Args:
      image: image with dimensions=[width, height, depth, channel].
      config: test parameters
      shape_model: structure containing shape model
      sess, x, action_ind, yc, yr, keep_prob: tensorflow nodes required for inference

    Returns:
      landmarks_all_steps: predicted landmarks. [max_test_steps + 1, num_examples, num_landmarks, 3]
      landmarks_mean: mean predicted landmarks across all num_examples. [num_landmarks, 3]

    """

    num_landmarks = config.landmark_count
    max_test_steps = config.max_test_steps
    box_size = config.box_size
    box_r = int((box_size-1)/2)

    # Initialise shape parameters, b=0 and landmarks, x
    b = shape_model_func.init_shape_params(config.num_random_init, None, config.sd, shape_model)     # b=[num_examples, num_shape_params]
    landmarks = shape_model_func.b2landmarks(b, shape_model)    # landmarks=[num_examples, num_landmarks, 3]
    num_examples = b.shape[0]

    # Extract patches from landmarks
    patches = np.zeros((num_examples, box_size, box_size, 3*num_landmarks))
    for j in range(num_examples):
        patches[j] = patch.extract_patch_all_landmarks(image, landmarks[j], box_r)       # patches=[num_examples, box_size, box_size, 3*num_landmarks]

    landmarks_all_steps = np.zeros((max_test_steps + 1, num_examples, num_landmarks, 3))
    landmarks_all_steps[0] = landmarks

    patches = torch.from_numpy(patches).permute(0, 3, 1, 2).to(torch.float32)

    for j in range(max_test_steps):    # find path of landmark iteratively

        # Predict CNN outputs
        yc_val, yr_val,_ = model(patches)

        yc_val = yc_val.detach().numpy()
        yr_val= yr_val.detach().numpy()

        # Compute classification probabilities
        action_prob = np.exp(yc_val - np.expand_dims(np.amax(yc_val, axis=1), 1))
        action_prob = action_prob / np.expand_dims(np.sum(action_prob, axis=1), 1)      # action_prob=[num_examples, 2*num_shape_params]

        # Update b values by combining classification and regression outputs
        b = update_b(b, action_prob, yr_val, config.predict_mode)

        # Convert b to landmarks
        landmarks = shape_model_func.b2landmarks(b, shape_model)    # landmarks=[num_examples, num_landmarks, 3]
        landmarks_all_steps[j+1] = landmarks

        # patches = patches.permute(0, 1, 2, 3)
        # Extract patches from landmarks
        for k in range(num_examples):
            patches[k] = torch.from_numpy(patch.extract_patch_all_landmarks(image, landmarks[k], box_r)).permute(2,0,1)
             # patches=[num_examples, box_size, box_size, 3*num_landmarks]

    # Compute mean of all initialisations
    landmarks_mean = np.mean(landmarks_all_steps[-1, :, :, :], axis=0)  # landmarks_mean=[num_landmarks, 3]

    return landmarks_all_steps, landmarks_mean

def update_b(b, action_prob, yr_val, predict_mode):
    """Update new shape parameters b using the regression and classification output.

    Args:
      b: current shape parameters values. [num_examples, num_shape_params].
      action_prob: classification output. [num_actions]=[num_examples, 2*num_shape_params]
      yr_val: values of db to regress. yr=b-b_gt. [num_examples, num_shape_params]
      predict_mode: 0: Hard classification. Move regressed distance only in the direction with maximum probability.
                    1: Soft classification. Multiply classification probabilities with regressed distances.
                    2: Regression only.
                    3: Classification only.

    Returns:
      b_new: new b after update. [num_examples, num_shape_params]

    """
    if predict_mode == 0:
        # Hard classification. Move regressed distance only in the direction with maximum probability.
        ind = np.argmax(np.amax(np.reshape(action_prob, (b.shape[0], b.shape[1], 2)), axis=2), axis=1)  # ind = [num_examples]
        row_ind = np.arange(b.shape[0])
        b[row_ind, ind] = b[row_ind, ind] - yr_val[row_ind, ind]
    elif predict_mode == 1:
        # Soft classification. Multiply classification probabilities with regressed distances.
        b = b - yr_val * np.amax(np.reshape(action_prob, (b.shape[0], b.shape[1], 2)), axis=2)
    elif predict_mode == 2:
        # Regression only.
        b = b - yr_val
    elif predict_mode == 3:
        # Classification only
        step = 1
        action_prob_reshape = np.reshape(action_prob, (b.shape[0], b.shape[1], 2))
        ind = np.argmax(np.amax(action_prob_reshape, axis=2), axis=1)   # ind=[num_examples]
        row_ind = np.arange(b.shape[0])
        is_negative = np.argmax(action_prob_reshape[row_ind, ind], axix=1)    # is_negative=[num_examples]
        # Move b in either positive or negative direction
        b[row_ind[is_negative], ind[is_negative]] = b[row_ind[is_negative], ind[is_negative]] + step
        b[row_ind[np.logical_not(is_negative)], ind[np.logical_not(is_negative)]] = b[row_ind[np.logical_not(is_negative)], ind[np.logical_not(is_negative)]] - step

    return b

def compute_err(landmarks, landmarks_gt, pix_dim):
    """Compute the distance error between predicted and GT landmarks.

    Args:
      landmarks: Predicted landmarks [img_count, num_landmarks, 3].
      landmarks_gt: Ground truth landmarks. [img_count, num_landmarks, 3]
      pix_dim: Pixel spacing. [img_count, 3]

    Returns:
      err: distance error in voxel. [img_count, num_landmarks]
      err_mm: distance error in mm. [img_count, num_landmarks]

    """
    num_landmarks = landmarks.shape[1]
    err = np.sqrt(np.sum(np.square(landmarks - landmarks_gt), axis=-1))
    err_mm = np.sqrt(np.sum(np.square((landmarks - landmarks_gt) * pix_dim[:, np.newaxis, :]), axis=-1))
    err_mm_landmark_mean = np.mean(err_mm, axis=0)
    err_mm_landmark_std = np.std(err_mm, axis=0)
    err_mm_mean = np.mean(err_mm)
    err_mm_std = np.std(err_mm)
    str = "Mean distance error (mm): "
    for j in range(num_landmarks):
        str += ("{:.10f} ".format(err_mm_landmark_mean[j]))
    print("{}".format(str))
    str = "Std distance error (mm): "
    for j in range(num_landmarks):
        str += ("{:.10f} ".format(err_mm_landmark_std[j]))
    print("{}".format(str))
    print("Mean distance error (mm) = {:.10f} \nStd distance error (mm) = {:.10f}\n".format(err_mm_mean, err_mm_std))
    return err, err_mm

if __name__ == '__main__':
    main()

reason: import numpy,import scipy


text:
from datetime import datetime
from newspaper import Article
from micromort.resources.configs.mongodbconfig import mongodb_config
from micromort.data_stores.mysql import db, cursor
from micromort.data_stores.mongodb import getConnection
from multiprocessing.dummy import Pool as ThreadPool
from micromort.utils.logger import logger
from dateutil import parser
import sys
import numpy as np

class Newspaper_scraper:
    def __init__(self, classify=False):
        self.db = db
        self.cursor = cursor
        self.classify = classify

        mongo_db = mongodb_config['dbs']['news_websites']['db']
        collection =  "newstweets_categorized_news"
        #mongodb_config['dbs']['news_websites']['collection']
        logger.debug("Creating mongo connection with db: " + mongo_db + " collection: " + collection)
        self.mongo_connection = getConnection(mongo_db, collection)

        rss_mongo_db = mongodb_config['dbs']['rss']['db']
        rss_mongo_collection =  mongodb_config['dbs']['rss']['collection']
        logger.debug("Creating mongo connection with db: " + rss_mongo_db + " collection: " + rss_mongo_collection)
        self.rss_connection = getConnection(rss_mongo_db, rss_mongo_collection)

        if classify:
            from micromort.models.trained_models.svm_mean_embeddings import Classifier, MeanEmbeddingVectorizer
            self.classifier = Classifier()

        pass

    def getRssData(self, url):
        return self.rss_connection.find_one({"link" : url})

    def scrape(self, url, useRssData=True):
        try:
            article = Article(url)
            article.download()
            article.parse()
        except Exception as ex:
            logger.exception("Failed to load content for: " + url)
            return -1

        #logging for testing:
        logger.debug("scrapped title: " + article.title)

        #check the content
        #Well if we dont have article text, then it is useless !
        if len(article.text) < 1:
            return -1

        #Get Rss data from mongo (For summary and published date)
        summary = ""
        published = ""
        title = article.title
        if useRssData:
            rssData = self.getRssData(url)
            try:
                summary = rssData.get("summary", "")
                title = rssData.get("title", article.title)
                published = parser.parse(rssData.get("published", ""))
            except Exception:
                summary = ""
                published = ""

        ob = {
            "url" : article.url,
            "original_url" : url,
            "title" : title,
            "text" : article.text,
            "images" : article.images,
            "summary" : summary,
            "published" : published,
            "top_image" : article.top_image,
            "movies" :  article.movies,
            "meta": {
                "updated_at" :  datetime.utcnow(),
            }
        }
        return ob

    def storeInMongo(self,  collection,  item):
        url = item["url"]
        collection.update({
            "url"  : url
            }
        ,  dict(item),  True)
        logger.debug("Article added to MongoDB database!")

    def getUrls(self):
        fromDate = "2017-12-01"
        toDate = "2018-03-01"
        self.cursor.execute(
                    #"""SELECT url FROM article_urls WHERE url like \"%www.asiaone.com%\" or url like \"%abc%\" """,
                    """SELECT url FROM article_urls where created_at between %s and %s""",
                    [fromDate, toDate]
                )
        urls = cursor.fetchall()
        return urls

    """
        @urls : Array of urls to worked upon
        @store : Boolean flag to decide if the scraped data to be stored in mongo
        @useRssData : Boolean flag to decide if (summary, published_data and title)
        is to be used from micromort.rss mongo collection. This should be set
        to true only for the urls for which RSS fed service is collecting the data

        Returns the array of json with article's details and predicted labels
    """

    def main(self, urls, store=True, useRssData=False):
        print "coming here"
        try:
            #data = []
            count=0
            for _url in urls:
                #url = _url[0].strip()
                url = _url.strip()[:1024]
                count=count+1
                if(count%10000 == 0):
                    print "Done:::::::", count
                logger.info("scrapping :" + url)
                item = self.scrape(url, useRssData)
                if item != -1:
                    #data.append(item)
                    if self.classify:
                        item["labels"] = self.classifier.predict_single(item["title"] + " " + item["text"], True)
                        logger.info("Predicted labels:" + str(item["labels"]))
                    if store:
                        self.storeInMongo(self.mongo_connection, item)
        except Exception as e:
            print e
        return 1

def batch(iterable, n=1):
    l = len(iterable)
    for ndx in range(0, l, n):
        yield iterable[ndx:min(ndx + n, l)]

if __name__ == "__main__":
    ob = Newspaper_scraper(classify=False)
    urls = []
    file = sys.argv[1]
    with open('/home/mannu/code/work/micromort/data/urls/' + file, 'r') as fp:
        read_lines = fp.readlines()
        for line in read_lines:
            urls.append(line.strip())
    ob.main(urls)
    # batch_len = len(urls)
    # n_threads = 10
    # urls_batch = np.array_split(np.array(urls), n_threads)
    # pool = ThreadPool(n_threads)
    # #print "test", np.asarray(urls_batch)
    # try:
    #     print "coming here?"
    #     results = pool.map(ob.main, np.asarray(urls_batch))
    # except Exception as e:
    #     print e
    # #d = ob.main(urls, store=True, useRssData=False)

    # for _ in d:
    #     print _["url"], ",", _["title"]


text:
#! /usr/bin/env python

from __future__ import division
from builtins import range
from LLC_Membranes.llclib import file_rw, transform, topology
import mdtraj as md
import numpy as np
import matplotlib.path as mplPath
import mdtraj as md
from random import randint
import tqdm

class region:
    """
    Define a region as an extrusion of a polygon in the z direction
    """

    def __init__(self, sides):
        """
        :param sides: number of sides making up the region in the xy plane
        :return: region
        """
        self.sides = sides

    def xyregion(self, corners):
        """
        :param corners: points defining the corners of the polygon making up the xy region
        :return: a region defined by corners
        """
        path = mplPath.Path(corners)

def thickness(filename, ref_atoms, grid, *traj, **kwargs):
    """
    :param filename: name of .gro file
    :param ref_atoms: atoms which thickness will be based on
    :param traj: trajectory of positions
    :return: trajectory of thicknesses or single thickness based on max/min z coordinate of reference atoms
    """

    if traj:

        traj = np.asarray(traj)[0]  # optional arguments of the form *args need to be convert back to numpy arrays
        nT = traj.shape[0]  # number of trajectory points

        thick = np.zeros([nT])
        z_max = np.zeros([nT])
        z_min = np.zeros([nT])
        thick_std = np.zeros([nT])
        for t in range(nT):
            z_max_t = max(traj[t, :, 2])
            z_min_t = min(traj[t, :, 2])
            thick[t] = z_max_t - z_min_t
            z_max[t] = z_max_t
            z_min[t] = z_min_t

    else:
        f = open(filename, "r")  # .gro file whose positions of Na ions will be read

        a = []  # list to hold lines of file
        for line in f:
            a.append(line)

        line = 0
        while a[line].count('HII') == 0:
            line += 1

        if grid:

            t = md.load(filename)
            pos = t.xyz[0, :, :]  # positions of all atoms

            if kwargs['exclude']:
                keep = [a.index for a in t.topology.atoms if a.residue.name != kwargs['exclude']]
                pos = t.atom_slice(keep).xyz[0, :, :]

            # define boundaries of each grid area
            grid_res = kwargs['grid_res']

            nregions = (grid_res - 1) ** 2
            g = np.zeros([2, grid_res, grid_res])
            dims = a[-1].split()
            xbox = np.linalg.norm([dims[0], dims[3], dims[4]])
            ybox = np.linalg.norm([dims[1], dims[5], dims[6]])
            yangle = np.arctan(float(dims[1])/abs(float(dims[5])))

            for i in range(grid_res):
                g[0, i, :] = np.linspace(0, xbox, grid_res) + (float(i) / grid_res)*float(dims[5])
                g[1, :, i] = np.linspace(0, ybox, grid_res)*np.sin(yangle)

            corners = np.zeros([nregions, 4, 2])
            zmaxes = np.zeros([nregions])
            zmins = np.zeros([nregions])
            thicks = np.zeros([nregions])

            for i in range(grid_res - 1):
                for j in range(grid_res - 1):
                    # define corners of grid region
                    r = i*(grid_res - 1) + j
                    corners[r, 0, :] = g[:, i, j]
                    corners[r, 1, :] = g[:, i + 1, j]
                    corners[r, 2, :] = g[:, i + 1, j + 1]
                    corners[r, 3, :] = g[:, i, j + 1]

                    # create a region using the corners (corners need to be traced in order)
                    path = mplPath.Path(corners[r, :, :])
                    contained = path.contains_points(pos[:, :2])  # check whether each point is in the region
                    z = pos[np.where(contained), 2]  # get the z position of all atoms contained in the region
                    zmaxes[r] = np.max(z)
                    zmins[r] = np.min(z)
                    thicks[r] = zmaxes[r] - zmins[r]

            # bootstrap to get statistics
            nboot = 2000
            vmax = np.zeros([nboot])
            vmin = np.zeros([nboot])

            for i in range(nboot):
                imax = randint(0, nregions - 1)
                imin = randint(0, nregions - 1)
                vmax[i] = zmaxes[imax]
                vmin[i] = zmins[imin]

            z_max = np.mean(vmax)
            z_min = np.mean(vmin)
            thick = np.mean(vmax - vmin)
            thick_std = np.std(vmax - vmin)

        else:
            z = []  # list to hold z positions of all atoms

            while a[line].count('HII') != 0:
                if str.strip(a[line][11:15]) in ref_atoms:
                    z.append(float(a[line][36:44]))
                line += 1

            z_max = max(z)
            z_min = min(z)
            thick = z_max - z_min
            thick_std = 0

    return thick, z_max, z_min, thick_std

def conc(t, comp, b):
    """
    Calculate the concentration of the specified component
    :param t: mdtraj trajectory object for system being studied
    :param comp: component which you want the concentration of
    :param b: buffer. distance into membrane to go before starting calculation
    :return: concentration
    """

    from pymbar import timeseries

    box = t.unitcell_vectors
    equil = timeseries.detectEquilibration(box[:, 2, 2])[0]
    thick = np.mean(box[equil:, 2, 2])

    z_max = thick
    z_min = 0
    buffer = thick*b
    z_max -= buffer
    z_min += buffer
    thick = z_max - z_min

    # Calculate concentration (an average of all frames)
    keep = [a.index for a in t.topology.atoms if a.name == comp]
    t_comp = t.atom_slice(keep)

    pos = t_comp.xyz
    ncomp = pos.shape[1]  # number of components in the simulation which you want the concentration of
    nT = pos.shape[0]

    if b > 0:
        count = np.zeros([nT])
        box_vol = np.zeros([nT])
        cross = np.zeros([nT])
        for t in range(nT):
            x_dim = np.linalg.norm(box[t, 0, :])
            y_dim = np.linalg.norm(box[t, 1, :])
            cross[t] = x_dim*y_dim
            box_vol[t] = x_dim*y_dim*thick
            for c in range(ncomp):
                if z_max >= pos[t, c, 2] >= z_min:
                    count[t] += 1
    else:
        count = ncomp*np.ones([nT])
        box_vol = np.zeros([nT])
        cross = np.zeros([nT])
        for t in range(nT):
            x_dim = np.linalg.norm(box[t, 0, :])
            y_dim = np.linalg.norm(box[t, 1, :])
            cross[t] = x_dim*y_dim
            box_vol[t] = x_dim*y_dim*thick

    factor = 1 / (1*10**-27)  # convert from ions/nm^3 to ions/m^3
    conc = np.zeros([nT])
    for c in range(nT):
        conc[c] = (count[c] / box_vol[c]) * factor

    avg_conc = np.mean(conc)
    std = np.std(conc)
    avg_cross = np.mean(cross)

    return avg_conc, std, avg_cross, thick, z_max, z_min

def avg_pore_loc(npores, pos, box, buffer=0, spline=False, npts=20, progress=False, bins=False, spline_name='spline.pl'):
    """ Calculate average pore location for each pore at each frame

    :param no_pores: the number of pores in the unit cell
    :param pos: the coordinates of the component(s) which you are using to locate the pore centers
    :param box: box vectors (t.unitcell_vectors when trajectory is load with mdtraj)
    :param buffer: fraction (of membrane thickness) of top and bottom of membrane to exclude from p2p calculations
    :param spline: trace pore centers with a spline
    :param npts: number of points making up the spline in each pore
    :param progress: show progress bar while constructing splines
    :param bins: return the bin centers of each spline for plotting purposes
    :param spline_name: name of spline. Include absolute path if not in same directory where script was run

    :type no_pores: int
    :type pos: numpy.ndarray, shape(ncomponents, 3) or numpy.ndarray, shape(nframes, ncomponents, 3)
    :type buffer: float
    :type spline: bool
    :type box: numpy.ndarray, shape(nframes, 3, 3)
    :type npts: int
    :type progress: bool
    :type bins: bool
    :type spline_name: str

    :return: numpy array containing the x, y coordinates of the center of each pore at each frame
    """

    # Find the average location of the pores w.r.t. x and y

    if spline:
        if box is None:
            print('You must supply box vectors if you are to trace the pores with a spline')
            exit()
        else:

            print('Calculating pore spline...')
            centers, bin_centers = trace_pores(pos, box, npts, npores=4, progress=progress, savename=spline_name)

            if bins:
                return centers, bin_centers
            else:
                return centers
    else:

        if len(pos.shape) == 3:  # multiple frames

            nT = np.shape(pos)[0]
            comp_ppore = np.shape(pos)[1] // npores

            p_center = np.zeros([nT, npores, 2])

            for i in range(nT):

                positions = wrap_box(pos[i, ...], box[i, ...])

                if buffer > 0:

                    include = np.full(pos.shape[1], True)

                    include[np.where(pos[i, :, 2] > box[i, 2, 2] + buffer)] = False
                    include[np.where(pos[i, :, 2] < buffer)] = False

                    for j in range(npores):
                        p_center[i, j, :] = positions[comp_ppore * j:comp_ppore * (j + 1), :2].mean(axis=0)
                        count = 0
                        for k in range(comp_ppore * j, comp_ppore * (j + 1)):
                            if include[k]:
                                p_center[i, j, :] += positions[k, :2]
                                count += 1
                        p_center[i, j, :] /= count  # take the average

                else:

                    for j in range(npores):
                        p_center[i, j, :] = positions[comp_ppore*j:comp_ppore*(j + 1), :2].mean(axis=0)

        elif len(pos.shape) == 2:  # single frame

            comp_ppore = pos.shape[0] // npores

            p_center = np.zeros([npores, 2])

            for j in range(npores):
                for k in range(comp_ppore*j, comp_ppore*(j + 1)):
                    p_center[j, :] += pos[k, :2]
                p_center[j, :] /= comp_ppore

        else:
            return 'Please use a position array with valid dimensions'
            exit()

        return p_center

def p2p(p_centers, distances):
    """
    :param p_centers: the x, y locations of the pore centers in the format return from avg_pore_loc()
    :param distances: the number of distinct distances between pores
    :return: all of the pore to pore distances
    """
    nT = np.shape(p_centers)[2]
    p2ps = np.zeros([distances, nT])  # distances in the order 1-2, 1-3, 1-4, 2-3, 2-4, 3-4
    for i in range(nT):
        # So ugly ... sadness :(
        p2ps[0, i] = np.linalg.norm(p_centers[:, 0, i] - p_centers[:, 1, i])
        p2ps[1, i] = np.linalg.norm(p_centers[:, 0, i] - p_centers[:, 2, i])
        p2ps[2, i] = np.linalg.norm(p_centers[:, 0, i] - p_centers[:, 3, i])
        p2ps[3, i] = np.linalg.norm(p_centers[:, 1, i] - p_centers[:, 2, i])
        p2ps[4, i] = np.linalg.norm(p_centers[:, 1, i] - p_centers[:, 3, i])
        p2ps[5, i] = np.linalg.norm(p_centers[:, 2, i] - p_centers[:, 3, i])

    return p2ps

def limits(pos, pcenters):
    """
    Estimate the pore 'radius' based on the position of some component and it's maximum deviation from the pore center
    :param: pos: the positions of all atoms included in making the estimate
    :param: pcenters: the x,y positions of the pore centers for each frame
    :return: an approximate pore radius. Beyond which, we have entered the alkane region
    """

    nT = pcenters.shape[0]
    npores = pcenters.shape[1]
    natoms = pos.shape[1]
    atom_ppore = natoms // npores

    deviation = np.zeros([nT, npores, atom_ppore])
    for f in tqdm.tqdm(range(nT)):
        for i in range(atom_ppore):
            for j in range(npores):
                deviation[f, j, i] = np.linalg.norm(pos[f, j*atom_ppore + i, :2] - pcenters[f, j, :])

    #deviation = np.reshape(deviation, (nT, natoms))

    fr = np.zeros([nT])
    frstd = np.zeros([nT])
    #
    # for i in range(nT):
    #     fr[i] = np.mean(deviation[i, :])  # + np.std(deviation[i, :]) # maybe?
    #     frstd[i] = np.std(deviation[i, :])

    radii = np.zeros([nT, npores])

    for t in range(nT):
        for p in range(npores):
            radii[t, p] = np.mean(deviation[t, p, :])

    return radii

def put_in_box(pt, x_box, y_box, m, angle):
    """
    :param pt: The point to place back in the box
    :param x_box: length of box in x dimension
    :param y_box: length of box in y dimension
    :param m: slope of box vector
    :param angle: angle between x axis and y box vector
    :return: coordinate shifted into box
    """

    b = - m * x_box  # y intercept of box vector that does not pass through origin (right side of box)

    if pt[1] < 0:
        pt[:2] += [np.cos(angle)*x_box, np.sin(angle)*x_box]  # if the point is under the box
    if pt[1] > y_box:
        pt[:2] -= [np.cos(angle)*x_box, np.sin(angle)*x_box]
    if pt[1] > m*pt[0]:  # if the point is on the left side of the box
        pt[0] += x_box
    if pt[1] < (m*pt[0] + b):  # if the point is on the right side of the box
        pt[0] -= x_box

    return pt

def trace_pores(pos, box, npoints, npores=4, progress=True, save=True, savename='spline.pl'):
    """
    Find the line which traces through the center of the pores
    :param pos: positions of atoms used to define pore location (args.ref) [natoms, 3]
    :param box: xy box vectors, [2, 2], mdtraj format (t.unitcell_vectors)
    :param npoints: number of points for spline in each pore
    :param npores: number of pores in unit cell (assumed that atoms are number sequentially by pore. i.e. pore 1 atom
    numbers all precede those in pore 2)
    :param progress: set to True if you want a progress bar to be shown
    :param save: save spline as pickled object
    :param savename: path to spline. If absolute path is not provided, will look in current directory

    :type pos: np.ndarray
    :type box: np.ndarray
    :type npoints: int
    :type npores: int
    :type progress: bool
    :type save: bool
    :type savename: str

    :return: points which trace the pore center
    """

    try:
        print('Attempting to load spline ... ', end='', flush=True)
        spline = file_rw.load_object(savename)
        print('Success!')

        return spline[0], spline[1]

    except FileNotFoundError:

        print('%s not found ... Calculating spline' % savename)

        single_frame = False
        if np.shape(pos.shape)[0] == 2:
            pos = pos[np.newaxis, ...]  # add a new axis if we are looking at a single frame
            box = box[np.newaxis, ...]
            single_frame = True

        nframes = pos.shape[0]
        atoms_p_pore = int(pos.shape[1] / npores)  # atoms in each pore

        v = np.zeros([nframes, 4, 2])  # vertices of unitcell box
        bounds = []

        v[:, 0, :] = [0, 0]
        v[:, 1, 0] = box[:, 0, 0]
        v[:, 3, :] = np.vstack((box[:, 1, 0], box[:, 1, 1])).T
        v[:, 2, :] = v[:, 3, :] + np.vstack((box[:, 0, 0], np.zeros([nframes]))).T
        center = np.vstack((np.mean(v[..., 0], axis=1), np.mean(v[..., 1], axis=1), np.zeros(nframes))).T

        for t in range(nframes):
            bounds.append(mplPath.Path(v[t, ...]))  # create a path tracing the vertices, v

        angle = np.arcsin(box[:, 1, 1]/box[:, 0, 0])  # specific to case where magnitude of x and y box lengths are equal
        angle = np.where(box[:, 1, 0] < 0, angle + np.pi / 2, angle)  # haven't tested this well yet

        m = (v[:, 3, 1] - v[:, 0, 1]) / (v[:, 3, 0] - v[:, 0, 0])  # slope from points connecting first and third vertices

        centers = np.zeros([nframes, npores, npoints, 3])
        bin_centers = np.zeros([nframes, npores, npoints])

        for t in tqdm.tqdm(range(nframes), disable=(not progress)):
            for p in range(npores):

                pore = pos[t, p*atoms_p_pore:(p+1)*atoms_p_pore, :]  # coordinates for atoms belonging to a single pore

                while np.min(pore[:, 2]) < 0 or np.max(pore[:, 2]) > box[t, 2, 2]:  # because cross-linked configurations can extend very far up and down

                    pore[:, 2] = np.where(pore[:, 2] < 0, pore[:, 2] + box[t, 2, 2], pore[:, 2])
                    pore[:, 2] = np.where(pore[:, 2] > box[t, 2, 2], pore[:, 2] - box[t, 2, 2], pore[:, 2])

                _, bins = np.histogram(pore[:, 2], bins=npoints)  # bin z-positions

                section_indices = np.digitize(pore[:, 2], bins)  # list that tells which bin each atom belongs to
                bin_centers[t, p, :] = [(bins[i] + bins[i + 1])/2 for i in range(npoints)]

                for l in range(1, npoints + 1):

                    atom_indices = np.where(section_indices == l)[0]

                    before = pore[atom_indices[0], :]  # choose the first atom as a reference

                    shift = transform.translate(pore[atom_indices, :], before, center[t, :])  # shift everything to towards the center

                    for i in range(shift.shape[0]):  # check if the points are within the bounds of the unitcell
                        while not bounds[t].contains_point(shift[i, :2]):
                            shift[i, :] = put_in_box(shift[i, :], box[t, 0, 0], box[t, 1, 1], m[t], angle[t])  # if its not in the unitcell, shift it so it is

                    c = [np.mean(shift, axis=0)]

                    centers[t, p, l - 1, :] = transform.translate(c, center[t, :], before)  # move everything back to where it was

                    while not bounds[t].contains_point(centers[t, p, l - 1, :]):  # make sure everything is in the box again
                        centers[t, p, l - 1, :] = put_in_box(centers[t, p, l - 1, :], box[t, 0, 0], box[t, 1, 1], m[t], angle[t])

        if single_frame:
            return centers[0, ...]  # doesn't return bin center yet

        else:

            if save:
                file_rw.save_object((centers, bin_centers), savename)

            return centers, bin_centers

def center_of_mass(pos, mass_atoms):
    """ Calculate center of mass of residues over a trajectory

    :param pos: xyz coordinates of atoms
    :param mass_atoms : mass of atoms in order they appear in pos

    :type pos: np.array (nframes, natoms, 3)
    :type mass_atoms: list

    :return: center of mass of each residue at each frame
    """

    nframes = pos.shape[0]
    natoms = len(mass_atoms)

    com = np.zeros([nframes, pos.shape[1] // natoms, 3])  # track the center of mass of each residue

    for f in range(nframes):
        for i in range(com.shape[1]):
            w = (pos[f, i * natoms:(i + 1) * natoms, :].T * mass_atoms).T  # weight each atom in the residue by its mass
            com[f, i, :] = np.sum(w, axis=0) / sum(mass_atoms)  # sum the coordinates and divide by the mass of the residue

    return com

def residue_center_of_mass(t, res):
    """ Calculate the center of mass versus time of a residue in an MD trajectory

    :param t: mdtraj trajectory object
    :param res: name of residue to track

    :type t: object
    :type res: str

    :return: center of mass of residue versus time
    """

    residue = topology.Residue(res)  # get resiude attributes

    ndx = [a.index for a in t.topology.atoms if a.residue.name == res]  # index of all residue atoms
    names = [a.name for a in t.topology.atoms if a.residue.name == res][:residue.natoms]  # names of atoms in one residue
    mass = [residue.mass[x] for x in names]  # mass of atoms in order that they appear in file
    print('Calculating center of mass trajectories of residue %s' % residue.name)

    return center_of_mass(t.xyz[:, ndx, :], mass)  # determine center of mass trajectories

def compdensity(coord, pore_centers, box, cut=1.5, nbins=50, spline=False):
    """ Measure the density of a component as a function of the distance from the pore centers.

    :param coord: the coordinates of the component(s) which you want a radial distribution of at each frame
    :param pore_centers: a numpy array of the locations of each pore center at each trajectory frame
    :param cut: cutoff distance for distance calculations. Will not count anything further than cut from the pore center
    :param nbins: number of bins in r direction
    :param spline: calculate RDF with respect to spline

    :type coord: numpy.ndarray
    :type pore_centers: numpy.ndarray
    :type cut: float
    :type nbins: int
    :type spline: bool

    :return: Radial distance from pore center r, and the density of a species, whose positions are defined by
    `coordinates`, as a function the distance from the pore center.
    """

    nT = coord.shape[0]
    pores = pore_centers.shape[1]
    density = np.zeros([nT, nbins])  # number / nm^3

    for t in tqdm.tqdm(range(nT), unit=' Frames'):
        for p in range(pores):

            if spline:
                distances = radial_distance_spline(pore_centers[t, p, ...], coord[t, ...], box[t, ...])
            else:
                distances = np.linalg.norm(coord[t, :, :2] - pore_centers[t, p, :], axis=1)

            hist, bin_edges = np.histogram(distances, bins=nbins, range=(0, cut))

            density[t, :] += hist

        density[t, :] /= (pores * box[t, 2, 2])  # normalize by z-dimension

    # normalize based on volume of anulus where bin is located (just need to divide by area since height done above)
    r = np.zeros([nbins])
    for i in range(nbins):
        density[:, i] /= (np.pi * (bin_edges[i + 1] ** 2 - bin_edges[i] ** 2))
        r[i] = (bin_edges[i + 1] + bin_edges[i]) / 2  # center of bins

    return r, density

def distance_from_pore_center(coord, pore_centers, box, spline=False):
    """ Measure the density of a component as a function of the distance from the pore centers.

    :param coord: the coordinates of the component(s) which you want a radial distribution of at each frame
    :param pore_centers: a numpy array of the locations of each pore center at each trajectory frame
    :param cut: cutoff distance for distance calculations. Will not count anything further than cut from the pore center
    :param

    :type coord: numpy.ndarray
    :type pore_centers: numpy.ndarray
    :type cut: float

    :return: Radial distance of each individual solute/component, defined by coords, as a function of time
    """

    nT = coord.shape[0]
    pores = pore_centers.shape[1]
    nsolute = coord.shape[1]

    r_distances = np.zeros([nT, nsolute])

    for t in tqdm.tqdm(range(nT), unit=' Frames'):
        rd = np.zeros([nsolute, pores])
        for p in range(pores):

            if spline:
                rd[:, p] = radial_distance_spline(pore_centers[t, p, ...], coord[t, ...], box[t, ...])
            else:
                rd[:, p] = np.linalg.norm(coord[t, :, :2] - pore_centers[t, p, :], axis=1)

        # Move the minimum solute--pore-center distance for each solute to the first index of rd
        # This removes any assumption that there is a constant number of solutes per pore and that the solute
        # stays in the same pore.
        for i, r in enumerate(rd):  # there is probably a vectorized way to do this with argsort
            rd[i, :] = r[np.argsort(r)]

        r_distances[t, :] = rd[:, 0]

    return r_distances

def radial_distance_spline(spline, com, box):
    """ Calculate radial distance from pore center based on distance from center of mass to closest z point in spline

    :param spline: coordinates of spline for a single pore and frame
    :param com: atomic center of mass z-coordinates
    :param zbox: z box dimension (nm)

    :type spline: np.ndarray [npts_spline, 3]
    :type com: np.ndarray [n_com, 3]
    :type zbox: float

    :return: array of distances from pore center
    """

    edges = np.zeros([spline.shape[0] + 1])
    edges[1:-1] = ((spline[1:, 2] - spline[:-1, 2]) / 2) + spline[:-1, 2]
    edges[-1] = box[2, 2]

    com = wrap_box(com, box)
    # while np.min(com[:, 2]) < 0 or np.max(com[:, 2]) > zbox:  # because cross-linked configurations can extend very far up and down
    #     com[:, 2] = np.where(com[:, 2] < 0, com[:, 2] + zbox, com[:, 2])
    #     com[:, 2] = np.where(com[:, 2] > zbox, com[:, 2] - zbox, com[:, 2])

    zbins = np.digitize(com[:, 2], edges)

    # handle niche case where coordinate lies exactly on the upper or lower bound
    zbins = np.where(zbins == 0, zbins + 1, zbins)
    zbins = np.where(zbins == edges.size, zbins - 1, zbins)

    return np.linalg.norm(com[:, :2] - spline[zbins - 1, :2], axis=1)

def minimum_image_distance(dist, box):
    """ Calculate minimum image distances from a vector of distances. This assumes a monoclinic unit cell where the x
    box vector is fixed along the x-axis, the z-box vector is perpendicular to the xy plane, and the y-box vector makes
    an angle, theta, with the x-axis.

    :param d: a vector of distances (n, 3) where n is number of points
    :param box: box vectors meant to enclose d, mdtraj format: (3, 3)

    :return:
    """

    x_box = box[0, 0]  # length of x-box vector
    y_box = box[1, 1]  # perpendicular distance from x-axis to top of box in y-direction
    z_box = box[2, 2]  # length of z-box vector
    d = np.copy(dist)
    angle = np.arcsin(y_box / x_box)  # angle between y-box vector and x-box vector in radians

    # check x coordinates
    while np.max(np.abs(d[:, 0])) > 0.5*x_box:  # iterate in case subtracting/adding box vector length once isn't enough
        d[:, 0] = np.where(d[:, 0] > 0.5*x_box, d[:, 0] - x_box, d[:, 0])
        d[:, 0] = np.where(d[:, 0] < -0.5*x_box, d[:, 0] + x_box, d[:, 0])

    # check y coordinates
    while np.amax(np.abs(d[:, 1])) > 0.5*y_box:  # written differently because np.where didn't know how to handle 2 axes
        d[np.where(d[:, 1] > 0.5*y_box)[0], :2] -= [x_box*np.cos(angle), y_box]
        d[np.where(d[:, 1] < -0.5*y_box)[0], :2] += [x_box*np.cos(angle), y_box]

    # check z coordinates
    while np.max(np.abs(d[:, 2])) > 0.5*z_box:
        d[:, 2] = np.where(d[:, 2] > 0.5*z_box, d[:, 2] - z_box, d[:, 2])
        d[:, 2] = np.where(d[:, 2] < -0.5*z_box, d[:, 2] + z_box, d[:, 2])

    return d

def partition(com, pore_centers, r, buffer=0, unitcell=None, npores=4, spline=False, spline_range=None):
    """ Partition residue center of masses into tail and pore region

    :param com: positions of centers of mass of particle whose partition we are calculating
    :param pore_centers: positions of pore centers
    :param r: pore radius, outside of which atoms will be considered in the tail region
    :param buffer: z distance (nm) to cut out from top and bottom of membrane (in cases where there is a water gap)
    :param unitcell: unitcell vectors in mdtraj format (t.unitcell_vectors). Only needed if buffer and/or spline is used
    :param npores: number of pores
    :param spline: calculate partition with respect to pore spline
    :param spline_range: range of frames to use in spline. Provide a tuple with the first and last (non-inclusive)
    frame that should be included.

    :type com: numpy.ndarray (nT, ncom, 3)
    :type pore_centers: numpy.ndarray (nT, npores, 2) or (nT, npores, 3) or (nT, npores, npts, 3) if spline=True where
    npts=number of points in spline
    :type r: float
    :type buffer: float
    :type unitcell: numpy.ndarray (nT, 3, 3)
    :type npores: int
    :type spline: bool
    :type spline_path: str
    :type spline_range: NoneType or tuple

    :return part: boolean numpy array with shape (nT, com.shape[1]) where True indicates a center of mass that is
    inside the inner region (i.e. < r)
    """

    nT = com.shape[0]

    if spline:
        npts = pore_centers.shape[2]  # number of points in each spline

    if nT < pore_centers.shape[0]:

        if spline_range is not None:

            start, end = spline_range
            pore_centers = pore_centers[start:end]

        else:

            print('The number of frames in the trajectory is less than the number frames in the spline. I will assume '
                  'that the difference has been chopped off from the front of the full trajectory and truncate the '
                  'spline accordingly. If this seems wrong, check out partition() in LLC_Membranes.llclib.physical')

            diff = pore_centers.shape[0] - nT

            pore_centers = pore_centers[diff:, ...]  # assumes trajectory h

    part = np.zeros([nT, com.shape[1]], dtype=bool)  # Will be changed to True if solute in pores

    print('Calculating solute partition...')
    for i in tqdm.tqdm(range(nT)):

        if buffer > 0:
            xy_positions = com[i, (com[i, :, 2] > buffer) & (com[i, :, 2] < unitcell[i, 2, 2] - buffer), :2]
        else:
            xy_positions = com[i, :, :2]

        if spline:

            z = com[i, :, 2]  # extract z-coordinates for this frame
            zbox = unitcell[i, 2, 2]  # z-box vector for this frame

            # make sure z-component of every particle in the box
            while np.max(z) > zbox or np.min(z) < 0:  # might need to do this multiple times
                z = np.where(z > zbox, z - zbox, z)
                z = np.where(z < 0, z + zbox, z)

            zbins = np.digitize(z, np.linspace(0, zbox, npts + 1))

            # handle niche case where coordinate lies exactly on the upper or lower bound
            zbins = np.where(zbins == 0, zbins + 1, zbins)
            zbins = np.where(zbins == npts + 1, zbins - 1, zbins)
            zbins -= 1  # digitize numbers bins starting at 1 (0 is below the bottom bin)

            pore = []

            for p in range(npores):
                d = np.linalg.norm(xy_positions - pore_centers[i, p, zbins, :2], axis=1)
                pore += np.where(d <= r)[0].tolist()

        else:

            pore = []

            for p in range(npores):
                d = np.linalg.norm(xy_positions - pore_centers[i, p, :], axis=1)
                pore += np.where(d <= r)[0].tolist()

        part[i, pore] = True

    return part

def wrap_box(positions, box, tol=1e-6):
    """ Put all atoms in box

    :param positions: xyz atomic position [n_atoms, 3]
    :param box: box vectors [3, 3] (as obtained from mdtraj t.unitcell_vectors)

    :type positions: np.ndarray
    :type box: np.ndarray

    :return: positions moved into box
    """

    xy = positions[:, :2]  # xy coordinates have dependent changes so this makes things neater below
    z = positions[:, 2]

    xbox, ybox, zbox = box[0, 0], box[1, 1], box[2, 2]

    angle = np.arcsin(ybox / xbox)  # angle between y-box vector and x-box vector in radians
    m = np.tan(angle)
    b = - m * xbox  # y intercept of box vector that does not pass through origin (right side of box)

    while max(xy[:, 1]) > ybox or min(xy[:, 1]) < 0:
        xy[np.where(xy[:, 1] > ybox)[0], :2] -= [xbox*np.cos(angle), ybox]
        xy[np.where(xy[:, 1] < 0)[0], :2] += [xbox * np.cos(angle), ybox]

    # added tolerance for corner case
    while len(np.where(xy[:, 0] - (xy[:, 1] / m) < -tol)[0]) > 0 or \
            len(np.where(xy[:, 0] - ((xy[:, 1] - b) / m) > 0)[0]) > 0:
        xy[np.where(xy[:, 0] < (xy[:, 1] / m))[0], 0] += xbox
        xy[np.where(xy[:, 0] > ((xy[:, 1] - b) / m))[0], 0] -= xbox

    # check z coordinates
    while np.max(z) > zbox or np.min(z) < 0:  # might need to do this multiple times
        z = np.where(z > zbox, z - zbox, z)
        z = np.where(z < 0, z + zbox, z)

    return np.concatenate((xy, z[:, np.newaxis]), axis=1)

def fft_3D_monoclinic(xyz, box_vectors, bins):
    """ Calculate 3D discrete fourier transform for each frame of a trajectory of coordinates in monoclinic unit cell

    :param xyz: (nframes, natoms, 3) coordinate array of atoms whose 3D DFT we want to calculate
    :param box_vectors: matrix of box vectors of shape (nT, 3, 3)
    :param bins: number of bins in each dimension

    :type xyz: np.ndarray
    :type box_vectors: np.ndarray
    :type bins: list of ints

    :return: 3D structure factor
    :type: np.ndarray
    """

    nT = xyz.shape[0]  # number of frames
    L = np.linalg.norm(box_vectors, axis=2)

    locations, L = rescale(xyz, L)  # make unit cell constant size

    # define histograme bin edges
    x = np.linspace(0, L[0], int(bins[0]))
    y = np.linspace(0, L[1], int(bins[1]))
    z = np.linspace(0, L[2], int(bins[2]))

    zv = [0.0, 0.0, 0.0]  # zero vector

    # put all atoms inside box - works for single frame and multiframe
    for it in range(locations.shape[0]):  # looped to save memory
        locations[it, ...] = np.where(locations[it, ...] < L, locations[it, ...], locations[it, ...] - L)  # get positions in periodic cell
        locations[it, ...] = np.where(locations[it, ...] > zv, locations[it, ...], locations[it, ...] + L)

    # fourier transform loop
    sf = np.zeros([x.size - 1, y.size - 1, z.size - 1])
    for frame in tqdm.tqdm(range(nT), unit=' Frames'):
        H, edges = np.histogramdd(locations[frame, ...], bins=(x, y, z))
        if sf:
            fft = np.fft.fftn(H - H.mean())
            sf += (fft * fft.conjugate()).real
        else:
            sf += np.abs(np.fft.fftn(H)) ** 2

    sf /= nT  # average of all frames

    return sf


text:
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Conv MLP module."""

import numpy as np
from ravens.models import MlpModel
from ravens.models import ResNet43_8s
import tensorflow as tf
tf.compat.v1.enable_eager_execution()
from tensorflow.keras import layers  # pylint: disable=g-import-not-at-top
import tensorflow_hub as hub

def compute_spatial_soft_argmax(x, batch_size, H=149, W=69, C=64):  # pylint: disable=invalid-name
  """Parameter-less, extract coordinates for each channel.

  Args:
    x: shape: (batch_size, ~H, ~W, C)
    batch_size: int
    H: size related to original image H size
    W: size related to original image W size
    C: channels

  Returns:
    shape: (batch_size, C, 2)
  """
  # see: https://github.com/tensorflow/tensorflow/issues/6271
  x = tf.transpose(x, [0, 3, 1, 2])
  x = tf.reshape(x, [batch_size * C, H * W])
  softmax = tf.nn.softmax(x)
  softmax = tf.reshape(softmax, [batch_size, C, H, W])
  softmax = tf.transpose(softmax, [0, 2, 3, 1])

  posx, posy = tf.meshgrid(
      tf.linspace(-1., 1., num=H), tf.linspace(-1., 1., num=W), indexing="ij")

  image_coords = tf.stack((posx, posy), axis=2)  # (H, W, 2)
  # Convert softmax to shape [B, H, W, C, 1]
  softmax = tf.expand_dims(softmax, -1)
  # Convert image coords to shape [H, W, 1, 2]
  image_coords = tf.expand_dims(image_coords, 2)
  # Multiply (with broadcasting) and reduce over image dimensions to get the
  # result of shape [B, C, 2].
  spatial_soft_argmax = tf.reduce_sum(softmax * image_coords, axis=[1, 2])
  return spatial_soft_argmax

class ConvMLP(tf.keras.Model):
  """Conv MLP module."""

  def __init__(self, d_action, use_mdn, pretrained=True):
    super(ConvMLP, self).__init__()

    if pretrained:
      inception = hub.KerasLayer(
          "https://tfhub.dev/google/imagenet/inception_v1/feature_vector/4",
          trainable=True)
      for i in inception.weights:
        if "Conv2d_1a_7x7/weights" in i.name:
          conv1weights = i
          break

    self.d_action = d_action

    input_shape = (None, 320, 160, 3)

    # filters = [64, 32, 16]
    filters = [64, 64, 64]

    if pretrained:
      self.conv1 = layers.Conv2D(
          filters=filters[0],
          kernel_size=(7, 7),
          strides=(2, 2),
          weights=[conv1weights.numpy(), tf.zeros(64)],
          input_shape=input_shape)

    else:
      self.conv1 = layers.Conv2D(
          filters=filters[0],
          kernel_size=(7, 7),
          strides=(2, 2),
          input_shape=input_shape)

    self.convd = layers.Conv2D(
        filters=filters[0],
        kernel_size=(7, 7),
        strides=(2, 2),
        input_shape=input_shape)
    self.bnd = layers.BatchNormalization()
    self.relud = layers.ReLU()

    self.bn1 = layers.BatchNormalization()
    self.relu1 = layers.ReLU()

    self.conv2 = layers.Conv2D(
        filters=filters[1], kernel_size=(5, 5), strides=(1, 1))
    self.bn2 = layers.BatchNormalization()
    self.relu2 = layers.ReLU()

    self.conv3 = layers.Conv2D(
        filters=filters[2], kernel_size=(5, 5), strides=(1, 1))
    self.bn3 = layers.BatchNormalization()
    self.relu3 = layers.ReLU()

    self.flatten = layers.Flatten()

    self.mlp = MlpModel(
        None,
        filters[-1] * 2,
        d_action,
        "relu",
        use_mdn,
        dropout=0.0,
        use_sinusoid=False)
    # note: no dropout on top of conv
    # the conv layers seem to help regularize the mlp layers

  def set_batch_size(self, batch_size):
    self.batch_size = batch_size

  def call(self, x):
    """FPROP through module.

    Args:
      x: shape: (batch_size, H, W, C)

    Returns:
      shape: (batch_size, self.d_action)
    """
    rgb = self.relu1(self.bn1(self.conv1(
        x[:, :, :, :3])))  # only rgb channels through pre-trained
    depth = self.relud(self.bnd(self.convd(x[:, :, :, 3:])))  # for depth
    x = tf.concat((rgb, depth), axis=-1)

    x = self.relu2(self.bn2(self.conv2(x)))
    x = self.relu3(self.bn3(self.conv3(x)))
    # shape (B, ~H, ~W, C=16)

    x = compute_spatial_soft_argmax(x, self.batch_size)  # shape (B, C, 2)

    x = self.flatten(x)  # shape (B, C*2)

    return self.mlp(x)

class DeepConvMLP:
  """Deep conv MLP module."""

  def __init__(self, input_shape, d_action, use_mdn):
    del use_mdn

    self.batch_size = 4
    self.input_shape = input_shape
    self.d_action = d_action

    channel_depth_dim = 16
    in0, out0 = ResNet43_8s(
        self.input_shape,
        channel_depth_dim,
        prefix="s0_",
        cutoff_early=False,
        include_batchnorm=True)

    # out0 = tf.nn.avg_pool(out0, ksize=(1,4,4,1), strides=(1,4,4,1),
    #                       padding="SAME", data_format="NHWC")

    out0 = compute_spatial_soft_argmax(out0, self.batch_size, 320, 160,
                                       16)  # shape (B, C, 2)

    out0 = tf.keras.layers.Flatten()(out0)
    out0 = tf.keras.layers.Dense(
        128,
        kernel_initializer="normal",
        bias_initializer="normal",
        activation="relu")(
            out0)
    out0 = tf.keras.layers.Dense(
        128,
        kernel_initializer="normal",
        bias_initializer="normal",
        activation="relu")(
            out0)
    out0 = tf.keras.layers.Dense(
        self.d_action, kernel_initializer="normal", bias_initializer="normal")(
            out0)

    self.model = tf.keras.Model(inputs=[in0], outputs=[out0])

  def set_batch_size(self, x):
    pass

  @property
  def trainable_variables(self):
    return self.model.trainable_variables

  def __call__(self, x):
    return self.model(x)

def main():
  cfg = tf.config.experimental
  gpus = cfg.list_physical_devices("GPU")
  mem_limit = 1024 * 4
  dev_cfg = [cfg.VirtualDeviceConfiguration(memory_limit=mem_limit)]
  cfg.set_virtual_device_configuration(gpus[0], dev_cfg)

  conv_mlp = ConvMLP(d_action=3, use_mdn=None)

  img = np.random.randn(7, 320, 160, 3)
  out = conv_mlp(img)
  print(out.shape)

if __name__ == "__main__":
  main()

max_stars_repo_name: egonrian/google-research

max_issues_repo_name: egonrian/google-research

max_forks_repo_name: egonrian/google-research


text:
# SPDX-FileCopyrightText: 2022 UChicago Argonne, LLC
# SPDX-License-Identifier: MIT

from .plugin import *
from .plugin_openmc import *
from .plugin_moose import *
from .plugin_pyarc import *
from .plugin_sas import *
from .plugin_relap5 import *
from .template import *
from .parameters import *
from .database import *

# This allows a user to write watts.Quantity
from astropy.units import Quantity

__version__ = '0.2.0'

reason: from astropy


text:
# -*- coding: utf-8 -*-
"""
Created on Fri Oct 24 16:52:46 2014

Mix of various routines around Zernike modes

@author: tristanbuey
"""

# Load libraries
import numpy as np;

def gammX(n, i0):
    """
    gammX(n, i0)
    Computes Noll matrix for derivatives of Zernike Gx.
    The matrix applies on a Zernike vector z, and produces
    the Zernike decomposition z' of x-derivative :
    z' = Gx . z

    n = number of zernike coefficients on the input vector z.
    i0 = zernike index of the first coefficient of z (1=piston, 2=tip, ...)

    It results that Gx is a matrix with a size (n+i0-1, n).
    """
    gg = np.zeros((i0+n-1,n));
    # variable i will span Zernike indexes, starting at i0.
    for i in range (i0, i0+n):
        # variable j spans Zernike indexes, starting at piston
        # and stopping at i
        for j in range(1, i+1):
            gg[j-1,i-i0] = gamX(i,j);
    return gg;

def gammY(n, i0):
    """
    gammY(n, i0)
    Computes Noll matrix for derivatives of Zernike Gy.
    The matrix applies on a Zernike vector z, and produces
    the Zernike decomposition z' of y-derivative :
    z' = Gy . z

    n = number of zernike coefficients on the input vector z.
    i0 = zernike index of the first coefficient of z (1=piston, 2=tip, ...)

    It results that Gy is a matrix with a size (n+i0-1, n).
    """
    gg = np.zeros((i0+n-1,n));
    # variable i will span Zernike indexes, starting at i0.
    for i in range(i0, i0+n):
        # variable j spans Zernike indexes, starting at piston
        # and stopping at i
        for j in range(1, i+1):
            gg[j-1,i-i0] = gamY(i,j);
    return gg;

"""
A lot of sub-functions to calculate the Noll matrix
"""

def pair(number):
    return number % 2 == 0;

def impair(num):
   return num % 2 != 0;

def nm(i):
    """
    For a given Zernike mode of index <i>, returns the
    radial and azimutal orders (n,m)
    """
    n = int( (-1.+np.sqrt(8*(i-1)+1))/2.);
    p = (i-(n*(n+1))/2);
    k = n%2;
    m = int((p+k)/2)*2 - k;
    return (n,m);

def gamY(i,j):
    """
    Input arguments:
    2 scalar int i and j, that are indexes of Zernike modes.

    Returns the coefficient of the derivative matrix of (Noll R.J., 1976)
    The algorithm coded below is a python translation of the series of
    rules that Noll has enounced in his article of 1976, for derivating
    Zernike.
    Warning: Unfortunately Noll had made a little error in his rules,
    that has been corrected in this program.
    """
    # determine radial and azimutal orders of Zernike number i
    ni,mi = nm(i);
    # idem for j
    n,m = nm(j);

    # Noll's rules :
    if(mi==(m-1) or mi==(m+1)):
        if(m==0 or mi==0):
            if((m==0 and impair(i)) or (mi==0 and impair(j))):
                return np.sqrt(2*(n+1)*(ni+1));
            else:
                return 0.00;
        else:
            if(impair(i+j)):
                if((mi==m+1 and impair(j)) or (mi==m-1 and pair(j))):
                    return -np.sqrt((n+1)*(ni+1));
                else:
                    return np.sqrt((n+1)*(ni+1));
            else:
                return 0.0;
    else:
        return 0.0;

    return;

def gamX(i,j):
    """
    Input arguments:
    2 scalar int i and j, that are indexes of Zernike modes.

    Returns the coefficient of the derivative matrix of (Noll R.J., 1976)
    The algorithm coded below is a python translation of the series of
    rules that Noll has enounced in his article of 1976, for derivating
    Zernike.
    Warning: Unfortunately Noll had made a little error in his rules,
    that has been corrected in this program.
    """
    # determine radial and azimutal orders of Zernike number i
    ni,mi = nm(i);
    # idem for j
    n,m = nm(j);

    # Noll's rules :
    if(mi==m-1 or mi==m+1):
        if(m==0 or mi==0):
            if((m==0 and pair(i)) or (mi==0 and pair(j))):
                return np.sqrt(2*(n+1)*(ni+1));
            else:
                return 0.00;
        else:
            if( (j+i)%2==0 ):
                return np.sqrt((n+1)*(ni+1));
            else:
                return 0.00;
    else:
        return 0.0;

    return;

def polyfute(m,n):
    """
    Les coefs des poly de zer sont des K_mn(s).
    Le coeff K_mn(s) pondère r^(n-2s)
    Il y a la relation de recurrence
    K_mn(s+1) =  K_mn(s) * ((n+m)/2-s)*((n-m)/2-s)/(s+1)/(n-s)
    Il y a aussi
    K_mn(0) = n! / ((n+m)/2)! / ((n-m)/2)!
    """

    a = np.zeros(n+1)

    # Calcul de K_mn(0)
    st = 2                           # start index for dividing by ((n-m)/2)!
    coef = 1.00
    for i in range((n+m)/2+1, n+1):
        if( st<=((n-m)/2) and i%st==0 ) :
            j = i/st
            st = st+1
            coef = coef*j
        else:
            coef = coef*i

    # division by ((n-m)/2)! (has already been partially done)
    for i in range(st,(n-m)/2+1):
        coef = coef / i

    a[n] = round(coef);   # pour K_nm(0)

    for i in range(1,(n-m)/2+1):
        coef = -coef * ((n+m)/2-i+1)*((n-m)/2-i+1);
        coef = coef / i;
        coef = coef / (n-i+1);
        a[n-2*i] = round(coef)

    return a

def evaluate_poly(n,m,a,r):
    """
     evaluate_poly(n,m,a,r)
     n is the radial order
     m is the azimutal order
     a[] is the list of coefficient, with a(i+1) the coeff of r^i
     r is the variable of the polynomial
    """
    if n>1 :
        r2 = r*r

    p = a[n]
    for i in range(n-2,m-1,-2):
        p = p*r2 + a[i]

    if(m==0): return p
    elif(m==1): p*=r
    elif(m==2): p*=r2
    else: p = p * r**m

    return p

def zer(r,t,i):
    """
    Computes Zernike polynom of index i, at point (r,t)

    The algo is using
    1) a recursive function to compute coefficients of the polynom, so that
    there is no factorial function of large numbers to invoke (risk of
    roundoff errors, plus takes more exec time)

    2) a smarter way to compute polynomial expressions such as
     ax^3+bx^2+cx+d = x(x(ax+b)+c)+d
    to avoid roundoff errors and minimize number of operations

    """

    if(i==1):
        return np.ones_like(r+t)

    # calcul de n et m a partir de i
    n = int( (-1.+np.sqrt(8*(i-1)+1))/2.)
    p = (i-(n*(n+1))/2);
    k = n%2;
    m = int((p+k)/2)*2 - k;

    a = polyfute(m,n)

    Z = evaluate_poly(n,m,a,r) * np.sqrt(n+1);
    if( m!=0 ):
        Z *= np.sqrt(2);
        if( i%2 ):
            Z *= np.sin(m*t)
        else:
            Z *= np.cos(m*t)
    return Z

def mkxy(npt, center, xy=0):
    """
    Defines a meshgrid of npt X npt points, and express
    it in polar coordinates.
    Returns a tuple (r,theta).

    When center==1, the center of the x,y map is 'Fourier-centered'.
    When optional argument xy==1, the returned tuple is (x, y)
    When optional argument xy==2, the returned tuple is (r, theta, x, y)
    """
    # generate an array of coordinates
    if center==1:
        x = np.linspace(-1,1,npt+1)[0:npt]
    else:
        x = np.linspace(-1,1,npt)
    # Warning: meshgrid swaps the 2 axis x and y.
    # That's why it's required to re-swap them here
    y,x = np.meshgrid(x,x)
    if( xy==1 ):
        return (x,y)
    # generates a map of the distance of subapertures to pupil center
    r = np.sqrt(x**2 + y**2)
    # generates a map of the azimut angle of subapertures
    theta = np.arctan2(y,x)
    if( xy==2 ):
        return r,theta,x,y
    else:
        return r,theta

max_stars_repo_name: soylentdeen/CIAO-commissioning-tools

max_issues_repo_name: soylentdeen/CIAO-commissioning-tools

max_forks_repo_name: soylentdeen/CIAO-commissioning-tools


text:
#tie stuff together around the hzpt object
import numpy as np
import warnings
from gzpt.zel import *
from scipy.interpolate import InterpolatedUnivariateSpline as ius
import os.path as path
import warnings

class hzpt:
    """
    Parent object for the extended HZPT model for Power spectrum and 2PCF for dark matter and tracers.
    Each hzpt instance corresponds to one cosmology and one (effective) redshift.

    Parameters
    ----------
    cosmo: (cosmology object)
    z: redshift, float

    Methods
    ----------
    __init__

    Example
    --------
    >>> k,plin = np.loadtxt("my_linear_power_file.txt",unpack=True)
    >>> z = 0.5
    >>> model = gzpt.hzpt(k,plin)
    >>> mm = matter.Correlator(params, model)
    >>> gg = tracers.AutoCorrelator(params,model)
    >>> gm = tracers.CrossCorrelator(params,model)

    #3D stats
    r = np.logspace(-1,3)
    ximm = mm.Xi()(r)
    k = np.logspace(-3,1)
    Pgg = gg.Power()(k)

    """
    #For now will only support single z, but can come back to this later
    def __init__(self,klin,plin,config=True):
        self.plin = loginterp(klin,plin) #interpolator

        #compute ZA
        self.cleft = CLEFT(klin,plin)
        self.cleft.make_ptable()
        self.kza,pza = self.cleft.pktable.T #evaluated at klin
        self.P_zel = loginterp(self.kza,pza) #callable
        self.config = config
        if(self.config):
            self.rxi = np.logspace(-1,3,4000) #matching up with the SBT in zel
            xiza = self.cleft.compute_xi_real(self.rxi)
            self.Xi_zel = loginterp(self.rxi,xiza) #callable
        else:
            warnings.warn("Configuration space models will not work if config flag is off!",
                           RuntimeWarning)

    def update_redshift(self,Dz):
        self.cleft.make_ptable(Dz=Dz)
        self.P_zel = loginterp(*self.cleft.pktable.T)
        if(self.config):
            self.Xi_zel = loginterp(self.rxi,self.cleft.compute_xi_real(self.rxi))

max_stars_repo_name: jmsull/gzpt

max_issues_repo_name: jmsull/gzpt

max_forks_repo_name: jmsull/gzpt

reason: import numpy,from scipy


text:
#deneme 24 - HSV (2/3)
#Trackbar ile parametrelerle oynayarak HSV degerleri bulunacak
#resimlerdeki cisimler filtre yapilarak digerlerinden ayrilacak

import cv2
import numpy as np

def nothing(x) :
    pass

cv2.namedWindow('tracking')

cv2.createTrackbar('LH', 'tracking', 0, 255, nothing)
cv2.createTrackbar('LS', 'tracking', 0, 255, nothing)
cv2.createTrackbar('LV', 'tracking', 0, 255, nothing)

cv2.createTrackbar('UH', 'tracking', 255, 255, nothing)
cv2.createTrackbar('US', 'tracking', 255, 255, nothing)
cv2.createTrackbar('UV', 'tracking', 255, 255, nothing)

while True :
    img = cv2.imread('smarties.png')
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    l_h = cv2.getTrackbarPos('LH', 'tracking')
    l_s = cv2.getTrackbarPos('LS', 'tracking')
    l_v = cv2.getTrackbarPos('LV', 'tracking')
    u_h = cv2.getTrackbarPos('UH', 'tracking')
    u_s = cv2.getTrackbarPos('US', 'tracking')
    u_v = cv2.getTrackbarPos('UV', 'tracking')

    l_b = np.array([l_h, l_s, l_v])
    u_b = np.array([u_h, u_s, u_v])

    mask = cv2.inRange(hsv, l_b, u_b)

    result = cv2.bitwise_and(img, img, mask=mask)

    cv2.imshow('original', img)
    cv2.imshow('mask', mask)
    cv2.imshow('result', result)

    key = cv2.waitKey(1) & 0xFF
    if key == 27 :
        break

cv2.destroyAllWindows()

max_stars_repo_name: EnesMercan/Computer-Vision-Python

max_issues_repo_name: EnesMercan/Computer-Vision-Python

max_forks_repo_name: EnesMercan/Computer-Vision-Python


text:
import cv2
import base64
import numpy as np
import random
from datetime import datetime

def findSliderX(bigImg_base64,smallImg_base64):
    # print(len(smallImg_base64))
    # print(len(bigImg_base64))
    smallImg_base64 = smallImg_base64[22:]
    bigImg_base64 = bigImg_base64[22:]

    try:
        bigImg_bytes=base64.urlsafe_b64decode(bigImg_base64 + '=' * (4 - len(bigImg_base64) % 4))
        smallImg_bytes=base64.urlsafe_b64decode(smallImg_base64 + '=' * (4 - len(smallImg_base64) % 4))
        bigImg_array = np.fromstring(bigImg_bytes,np.uint8)
        smallImg_array = np.fromstring(smallImg_bytes,np.uint8)
        bigImg = cv2.imdecode(bigImg_array, cv2.IMREAD_COLOR)
        smallImg = cv2.imdecode(smallImg_array, cv2.IMREAD_COLOR)
        # cv2.imshow('bigImg', bigImg)
        # cv2.imshow('smallImg', smallImg)
    except Exception:
        pass

    # 分析小图片，找到纵坐标
    # img = cv2.imread('slider2-1.png', 0)
    img = smallImg
    blur = cv2.GaussianBlur(img, (3, 3), 0)  # 用高斯滤波处理原图像降噪
    canny = cv2.Canny(blur, 150, 200)  # canny算子边缘检测 170是最小阈值,200是最大阈值
    # print(len(canny))
    # print(len(canny[0]))
    # cv2.imshow('smallImg_canny', canny)
    YFromTop = 0
    for i in range(len(canny)):
        if(canny[i][5]) > 0:
            YFromTop = i
            break
    # print(YFromTop)

    # 分析大图片，找到横坐标
    # img = cv2.imread('slider2-2.png', 0)
    img = bigImg
    blur = cv2.GaussianBlur(img, (3, 3), 0)
    canny = cv2.Canny(blur, 150, 200)
    # print(len(canny))
    # print(len(canny[0]))
    # cv2.imshow('bigImg_canny', canny)
    XFromLeft = 0
    # 根据前面得到的纵坐标开始分析,向下浮动5个像素点
    XFromLeft = findLineEdge(canny, YFromTop, 0)
    if XFromLeft > 0:
        # 暂不启用
        pass
        # 向右约75个像素是缺口的右边沿
        # 向右70个像素，浮动10个像素点
        # for x_i in range(10):
        #     result = findLineEdge(canny, YFromTop, XFromLeft+70+x_i)
        #     if result > 0:
        #         break
        # else:
        #     print("未检测到边沿2 " +str(datetime.now()))
    else:
        print("未检测到边沿1 " +str(datetime.now()))
        XFromLeft = random.randint(20, len(canny[0])-20)

    # 已废弃
    # for y_i in range(5):
    #     for i in range(len(canny[0])):
    #         if(canny[YFromTop+y_i][i]) > 0:
    #             # 向下延伸10个像素，如果都是边缘则认为是缺口的左边沿
    #             for j in range(10):
    #                 if(canny[YFromTop+j+y_i][i] == 0):
    #                     break
    #             else:
    #                 # 向右约73个像素是缺口的右边沿
    #                 # 向右71个像素，向右浮动5个像素点
    #                 for x_i in range(5):
    #                     # 向下延伸10个像素，如果都是边缘则认为是缺口的右边沿
    #                     for j in range(10):
    #                         if(canny[YFromTop+j+y_i][i+x_i] == 0):
    #                             # 不是边沿
    #                             break
    #                     else:
    #                         # 是边沿
    #                         XFromLeft = i
    #                         break
    #                 else:
    #                     # 不是边沿
    #                     continue
    #                 # 是边沿
    #                 break
    #     else:
    #         # 不是边沿
    #         continue
    #     # 是边沿
    #     break
    # print(XFromLeft)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    XFromLeft_scale = XFromLeft/len(canny[0])
    return XFromLeft_scale

# 检测滑块纵向边沿
def findLineEdge(Img, y_base, x_base):
    XFromLeft = -1
    # 向下浮动5个像素点
    for y_i in range(5):
        # 横向，从左至右分析
        for i in range(x_base, len(Img[0])):
            if(Img[y_base+y_i][i]) > 0:
                # 向下延伸10个像素，如果都是边缘则认为是缺口的左边沿
                for j in range(10):
                    if(Img[y_base+j+y_i][i] == 0):
                        # 不是边沿
                        break
                else:
                    # 是边沿
                    XFromLeft = i
                    break
        else:
            # 不是边沿
            continue
        # 是边沿
        break
    return XFromLeft


text:
# licensed under ieee project
#Auhor Rahul Kumar Chauhan

import re
import nltk

import pandas as pd
import numpy as np

from bs4 import BeautifulSoup
from nltk.corpus import stopwords

class Word2VecUtility(object):

    @staticmethod
    def review_to_wordlist( review, remove_stopwords=False ):

        review_text = BeautifulSoup(review).get_text()

        review_text = re.sub("[^a-zA-Z]"," ", review_text)

        words = review_text.lower().split()

        if remove_stopwords:
            stops = set(stopwords.words("english"))
            words = [w for w in words if not w in stops]

        return(words)

    @staticmethod
    def review_to_sentences( review, tokenizer, remove_stopwords=False ):

        raw_sentences = tokenizer.tokenize(review.decode('utf8').strip())

        sentences = []
        for raw_sentence in raw_sentences:

            if len(raw_sentence) > 0:

                sentences.append( Word2VecUtility.review_to_wordlist( raw_sentence, \
                  remove_stopwords ))

        return sentences