text:
import numpy as np
import torch
import torch.nn.functional as F
from numpy import clip, exp
from scipy.signal import convolve2d
import matplotlib.pyplot as plt

def expand(x, r):
    return np.repeat(np.repeat(x, r, axis = 0), r, axis = 1)

def show_tensor(tensor, cmap='magma', scale=False):
    im = tensor.numpy()
    if scale:
        im = im / im.max()
    else:
        im = clip(im, 0, 1)

    if im.shape[0] == 1:
        plt.imshow(im[0], cmap=cmap)
    else:
        plt.imshow(im.transpose((1, 2, 0)))

def tensor_to_numpy(x):
    x = x.detach().cpu().numpy()
    if x.ndim == 4:
        x = x[0]
    if x.ndim == 2:
        return x

    if x.shape[0] == 1:
        return x[0]
    elif x.shape[0] == 3:
        return x.transpose((1, 2, 0))
    else:
        raise

def plot_tensors(tensor_list, titles=None, **kwargs):

    color = True if tensor_list[0].shape[1] == 3 else False
    image_list = [tensor_to_numpy(tensor) for tensor in tensor_list]
    width = len(image_list)
    fig, ax = plt.subplots(1, width, sharex='col', sharey='row', figsize=(width * 4, 4))

    for i in range(width):
        if image_list[i].ndim == 2:
            if 'cmap' not in kwargs:
                kwargs['cmap'] = 'Greys_r'
            ax[i].imshow(image_list[i], **kwargs)
        else:
            ax[i].imshow(image_list[i], **kwargs)
        if titles:
            ax[i].set_title(titles[i])
        ax[i].get_xaxis().set_ticks([])
        ax[i].get_yaxis().set_ticks([])
    fig

def show_data(datapt):
    # For datasets of the form (noise1, noise2, ground truth), shows all three concatenated
    show_tensor(torch.cat((datapt[0], datapt[1], datapt[2]), dim=2))

def scale_tensor(x):
    return (x - x.min()) / (x.max() - x.min())

def plot_grid(images, height, width, **kwargs):

    if not isinstance(images, np.ndarray):
        images = np.concatenate([im[np.newaxis] for im in images])
    assert images.shape[0] >= width * height

    if 'cmap' not in kwargs:
        kwargs['cmap'] = 'Greys_r'

    images = images[:width * height]
    fig, ax = plt.subplots(height, width, sharex='col', sharey='row', figsize=(width * 4, height * 4))
    image_grid = images.reshape(height, width, images.shape[1], images.shape[2])

    # axes are in a two-dimensional array, indexed by [row, col]
    for i in range(height):
        for j in range(width):
            if (height > 1):
                ax[i, j].imshow(image_grid[i, j], **kwargs)
                ax[i, j].get_xaxis().set_ticks([])
                ax[i, j].get_yaxis().set_ticks([])
            else:
                ax[j].imshow(image_grid[i, j], **kwargs)
                ax[j].get_xaxis().set_ticks([])
                ax[j].get_yaxis().set_ticks([])
    fig

def show(image, **kwargs):
    import matplotlib.pyplot as plt
    plt.imshow(image, cmap=plt.cm.gray, **kwargs)
    plt.gca().get_xaxis().set_ticks([])
    plt.gca().get_yaxis().set_ticks([])

def plot_images(image_list, **kwargs):
    images = np.concatenate([im[np.newaxis] for im in image_list])
    plot_grid(images, 1, len(image_list), **kwargs)

def clamp_tensor(x):
    return torch.clamp(x, 0, 1)

def random_noise(img, params):
    """Parameters for random noise include the mode and the type.

    mode: gaussian, poisson, or gaussian_poisson noise type
    std: std of gaussian
    photons_at_max: at image with intensity 1 has this many photons on average
    clamp: clamp result to [0,1]
    """

    noisy = img

    if params['mode'] == 'poisson' or params['mode'] == 'gaussian_poisson':
        noisy = torch.poisson(noisy * params['photons_at_max']) / params['photons_at_max']

    if params['mode'] == 'gaussian' or params['mode'] == 'gaussian_poisson':
        noise = torch.randn(img.size()).to(img.device) * params['std']
        noisy = noise + noisy

    if params['mode'] == 'bernoulli':
        noisy = noisy * torch.bernoulli(torch.ones(noisy.shape) * params['p'])

    if 'clamp' in params and params['clamp']:
        noisy = torch.clamp(noisy, 0, 1)

    return noisy

def test_bernoulli_noise():
    torch.manual_seed(2018)
    p = 0.2
    shape = (10, 1, 100, 100)
    n = 10 * 100 * 100
    img = torch.ones(shape)
    noisy = random_noise(img, {'mode': 'bernoulli', 'p': p})

    var = n * p * (1 - p)

    assert torch.abs(noisy.sum() - p * img.sum()) < 3 * (var ** 0.5)

def psnr(x, x_true, max_intensity=1.0, pad=None, rescale=False):
    '''A function computing the PSNR of a noisy tensor x approximating a tensor x_true.

    It vectorizes over the batch.

    PSNR := 10*log10 (MAX^2/MSE)

    where the MSE is the averaged squared error over all pixels and channels.
    '''

    return 10 * torch.log10((max_intensity ** 2) / mse(x, x_true, pad=pad, rescale=rescale))

def test_psnr():
    std = 0.1
    noise = torch.randn(10, 3, 100, 100) * std
    x_true = torch.ones(10, 3, 100, 100) / 2
    x = x_true + noise
    # MSE should be 0.01. PSNR should be 20.
    assert (torch.abs(psnr(x, x_true) - 20) < 0.1).all()

    x = 256 * x
    x_true = 256 * x_true
    assert (torch.abs(psnr(x, x_true, 256) - 20) < 0.2).all()

def test_mse_rescale():
    y = torch.randn(10, 3, 10, 10)
    x = 10 * y + 7
    assert (mse(x, y, rescale=True) < 1e-5).all()

    # Normalized values are (1, 1, 0, -2) and (1, 1, -1, -1)
    y = torch.Tensor([3, 3, 2, 0]).unsqueeze(0).unsqueeze(0).unsqueeze(0)
    x = torch.Tensor([5, 5, 0, 0]).unsqueeze(0).unsqueeze(0).unsqueeze(0)

    assert mse(x, y, rescale=True).sum() == 0.5

def mse(x, y, pad=None, rescale=False):
    if pad:
        x = x[:, :, pad:-pad, pad:-pad]
        y = y[:, :, pad:-pad, pad:-pad]

    def batchwise_mean(z):
        return z.reshape(z.shape[0], -1).mean(dim=1).reshape(-1, 1, 1, 1)

    if rescale:
        x = x - batchwise_mean(x)
        y = y - batchwise_mean(y)
        a = batchwise_mean(x * y) / batchwise_mean(x * x)
        x = a * x

    return batchwise_mean((x - y) ** 2).reshape(-1)

def gaussian(window_size, sigma):
    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
    return gauss / gauss.sum()

def create_window(window_size, channel):
    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
    return window

def _ssim(img1, img2, window, window_size, channel, size_average=True):
    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)

    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1 * mu2

    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2

    C1 = 0.01 ** 2
    C2 = 0.03 ** 2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))

    if size_average:
        return ssim_map.mean()
    else:
        return ssim_map.mean(1).mean(1).mean(1)

def ssim(img1, img2, window_size=11, size_average=True):
    (_, channel, _, _) = img1.size()
    window = create_window(window_size, channel)

    if img1.is_cuda:
        window = window.cuda(img1.get_device())
    window = window.type_as(img1)

    return _ssim(img1, img2, window, window_size, channel, size_average)

def smooth(tensor):
    kernel = np.array([[0.5, 1.0, 0.5], [1.0, 2.0, 1.0], (0.5, 1.0, 0.5)])
    kernel = kernel[np.newaxis, np.newaxis, :, :]
    kernel = torch.Tensor(kernel).to(tensor.device)
    kernel = kernel / kernel.sum()

    filtered_tensor = torch.nn.functional.conv2d(tensor, kernel, stride=1, padding=1)
    return filtered_tensor

def normalize(x, pmin=3, pmax=99.8, axis=None, clip=False, eps=1e-20, dtype=np.float32):
    """Percentile-based image normalization."""

    mi = np.percentile(x, pmin, axis=axis, keepdims=True)
    ma = np.percentile(x, pmax, axis=axis, keepdims=True)
    return normalize_mi_ma(x, mi, ma, clip=clip, eps=eps, dtype=dtype)

def normalize_mi_ma(x, mi, ma, clip=False, eps=1e-20, dtype=np.float32):
    if dtype is not None:
        x = x.astype(dtype, copy=False)
        mi = dtype(mi) if np.isscalar(mi) else mi.astype(dtype, copy=False)
        ma = dtype(ma) if np.isscalar(ma) else ma.astype(dtype, copy=False)
        eps = dtype(eps)

    try:
        import numexpr
        x = numexpr.evaluate("(x - mi) / ( ma - mi + eps )")
    except ImportError:
        x = (x - mi) / (ma - mi + eps)

    if clip:
        x = np.clip(x, 0, 1)

    return x

class PercentileNormalizer():
    """Percentile-based image normalization.
    Parameters
    ----------
    pmin : float
        Low percentile.
    pmax : float
        High percentile.
    dtype : type
        Data type after normalization.
    kwargs : dict
        Keyword arguments for :func:`csbdeep.utils.normalize_mi_ma`.
    """

    def __init__(self, pmin=2, pmax=99.8, dtype=np.float32, **kwargs):
        if not (np.isscalar(pmin) and np.isscalar(pmax) and 0 <= pmin < pmax <= 100):
            raise ValueError
        self.pmin = pmin
        self.pmax = pmax
        self.dtype = dtype
        self.kwargs = kwargs

        self.mi = None
        self.ma = None

    def normalize(self, img, channel=1):
        """Percentile-based normalization of raw input image.
        Note that percentiles are computed individually for each channel (if present in `axes`).
        """
        axes = tuple((d for d in range(img.ndim) if d != channel))
        self.mi = np.percentile(img, self.pmin, axis=axes, keepdims=True).astype(self.dtype, copy=False)
        self.ma = np.percentile(img, self.pmax, axis=axes, keepdims=True).astype(self.dtype, copy=False)
        return normalize_mi_ma(img, self.mi, self.ma, dtype=self.dtype, **self.kwargs)

    def denormalize(self, mean):
        """Undo percentile-based normalization to map restored image to similar range as input image.
        """
        alpha = self.ma - self.mi
        beta = self.mi
        return alpha * mean + beta

def test_percentile_normalizer():
    a = np.arange(1000).reshape(10, 1, 10, 10).astype(np.uint16)

    norm = PercentileNormalizer(pmin=0, pmax=100, dtype=np.float32, clip=False)
    assert norm.normalize(a).min() == 0 and norm.normalize(a).max() == 1

    # gap between 10th and 90th percentile is 100 to 900, so
    # the transform is (x - 100)/800
    norm = PercentileNormalizer(pmin=10, pmax=90, dtype=np.float32, clip=False)
    assert norm.normalize(a).max() == 1.125

    norm = PercentileNormalizer(pmin=2, pmax=99.8, dtype=np.float32, clip=True)
    assert norm.normalize(a).max() == 1.0

def gpuinfo(gpuid):
    import subprocess
    sp = subprocess.Popen(['nvidia-smi', '-q', '-i', str(gpuid), '-d', 'MEMORY'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out_str = sp.communicate()
    out_list = out_str[0].decode("utf-8").split('BAR1', 1)[0].split('\n')
    out_dict = {}
    for item in out_list:
        try:
            key, val = item.split(':')
            key, val = key.strip(), val.strip()
            out_dict[key] = val
        except:
            pass
    return out_dict

def getfreegpumem(id):
    return int(gpuinfo(id)['Free'].replace('MiB', '').strip())

def getbestgpu():
    freememlist = []
    for id in range(4):
        freemem = getfreegpumem(id)
        print("GPU device %d has %d MiB left." % (id, freemem))
        freememlist.append(freemem)
    idbest = freememlist.index(max(freememlist))
    print("--> GPU device %d was chosen" % idbest)
    return idbest

def get_args():
    global args
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("config_files",
                        help="configuration file for experiment.",
                        type=str,
                        nargs='+')
    parser.add_argument("--device",
                        help="cuda device",
                        type=str,
                        required=True)
    args = parser.parse_args()

# Function that takes in stacked 2 channel data and returns images
# whose difference in max values is below a threshold
def clean_two_channel_data(data, thresh=5000):
  max_differences = np.abs((data[:,0].max(-1).max(-1)) - (data[:,1].max(-1).max(-1)))
  idxs = max_differences < thresh
  return data[idxs]

# Function that takes in stacked 3 channel data and returns images
# whose difference in max values is below a threshold
def clean_three_channel_data(data, thresh=5000):
  max_differences_1_2 = np.abs((data[:,0].max(-1).max(-1)) - (data[:,1].max(-1).max(-1)))
  max_differences_1_3 = np.abs((data[:,0].max(-1).max(-1)) - (data[:,2].max(-1).max(-1)))
  max_differences_2_3 = np.abs((data[:,1].max(-1).max(-1)) - (data[:,2].max(-1).max(-1)))

  idxs_1_2 = (max_differences_1_2 < thresh).reshape(-1)
  idxs_1_3 = (max_differences_1_3 < thresh).reshape(-1)
  idxs_2_3 = (max_differences_2_3 < thresh).reshape(-1)

  all_idxs = np.stack([idxs_1_2, idxs_1_3, idxs_2_3])
  idxs = (np.sum(all_idxs, axis=0) == 3)

  return data[idxs]

def clamp_data(data, min=0, max=1500):
  return np.clip(data[:], 0, max).reshape(len(data), -1, 128, 128)

reason: import numpy,from numpy,from scipy,import numexpr


text:
"""Utility functions for ray computation."""
import math
import numpy as np
from scipy.spatial.transform import Rotation as R
import box_utils
import torch

def apply_batched_transformations(inputs, transformations):
    """Batched transformation of inputs.

    Args:
        inputs: List of [R, S, 3]
        transformations: [R, 4, 4]

    Returns:
        transformed_inputs: List of [R, S, 3]
    """
    # if rotation_only:
    #    transformations[:, :3, 3] = torch.zeros((3,), dtype=torch.float)

    transformed_inputs = []
    for x in inputs:
        N_samples = x.size()[1]
        homog_transformations = transformations.unsqueeze(1)  # [R, 1, 4, 4]
        homog_transformations = torch.tile(homog_transformations, (1, N_samples, 1, 1))  # [R, S, 4, 4]
        homog_component = torch.ones_like(x)[..., 0:1]  # [R, S, 1]
        homog_x = torch.cat((x, homog_component), axis=-1)  # [R, S, 4]
        homog_x = homog_x.unsqueeze(2)
        transformed_x = torch.matmul(
            homog_x,
            torch.transpose(homog_transformations, 2, 3))  # [R, S, 1, 4]
        transformed_x = transformed_x[..., 0, :3]  # [R, S, 3]
        transformed_inputs.append(transformed_x)
    return transformed_inputs

def get_transformation_from_params(params):
    translation, rotation = [0, 0, 0], [0, 0, 0]
    if 'translation' in params:
        translation = params['translation']
    if 'rotation' in params:
        rotation = params['rotation']
    translation = torch.tensor(translation, dtype=torch.float)
    rotmat = torch.tensor(R.from_euler('xyz', rotation, degrees=True).as_matrix(), dtype=torch.float)
    return translation, rotmat

def rotate_dirs(dirs, rotmat):
    """
    Args:
        dirs: [R, 3] float tensor.
        rotmat: [3, 3]
    """
    if type(dirs) == np.ndarray:
        dirs = torch.tensor(dirs).float()
    #rotmat = rotmat.unsqueeze(0)
    rotmat = torch.broadcast_to(rotmat, (dirs.shape[0], 3, 3))  # [R, 3, 3]
    dirs_obj = torch.matmul(dirs.unsqueeze(1), torch.transpose(rotmat, 1, 2))  # [R, 1, 3]
    dirs_obj = dirs_obj.squeeze(1)  # [R, 3]
    return dirs_obj

def transform_dirs(dirs, params, inverse=False):
    _, rotmat = get_transformation_from_params(params)  # [3,], [3, 3]
    if inverse:
        rotmat = torch.transpose(rotmat, 0, 1)  # [3, 3]
    dirs_transformed = rotate_dirs(dirs, rotmat)
    return dirs_transformed

def transform_rays(ray_batch, params, use_viewdirs, inverse=False):
    """Transform rays into object coordinate frame given o2w transformation params.

    Note: do not assume viewdirs is always the normalized version of rays_d (e.g., in staticcam case).

    Args:
        ray_batch: [R, M] float tensor. Batch of rays.
        params: Dictionary containing transformation parameters:
            'translation': List of 3 elements. xyz translation.
            'rotation': List of 3 euler angles in xyz.
        use_viewdirs: bool. Whether to we are using viewdirs.
        inverse: bool. Whether to apply inverse of the transformations provided in 'params'.

    Returns:
        ray_batch_obj: [R, M] float tensor. The ray batch, in object coordinate frame.
    """
    rays_o, rays_d = ray_batch[:, 0:3], ray_batch[:, 3:6]
    translation, rotmat = get_transformation_from_params(params)  # [3,], [3, 3]

    if inverse:
        translation = -1 * translation  # [3,]
        rotmat = torch.transpose(rotmat, 1, 0)  # [3, 3]

    translation_inverse = -1 * translation
    rotmat_inverse = torch.transpose(rotmat, 1, 0)

    # Transform the ray origin.
    rays_o_obj, _ = box_utils.ray_to_box_coordinate_frame_pairwise(
        box_center=translation_inverse,
        box_rotation_matrix=rotmat_inverse,
        rays_start_point=rays_o,
        rays_end_point=rays_d)

    # Only apply rotation to rays_d.
    rays_d_obj = rotate_dirs(rays_d, rotmat)

    ray_batch_obj = update_ray_batch_slice(ray_batch, rays_o_obj, 0, 3)
    ray_batch_obj = update_ray_batch_slice(ray_batch_obj, rays_d_obj, 3, 6)
    if use_viewdirs:
        # Grab viewdirs from the ray batch itself. Because it may be different from rays_d
        # (as in the staticcam case).
        viewdirs = ray_batch[:, 8:11]
        viewdirs_obj = rotate_dirs(viewdirs, rotmat)
        ray_batch_obj = update_ray_batch_slice(ray_batch_obj, viewdirs_obj, 8, 11)
    return ray_batch_obj

def transform_points_into_world_coordinate_frame(pts, params, check_numerics=False):
    translation, rotmat = get_transformation_from_params(params)  # [3,], [3, 3]

    # pts_flat = pts.view(-1, 3)  # [RS, 3]
    # num_examples = pts_flat.size()[0]  # RS

    # translation = translation.unsqueeze(0)
    # translation = torch.tile(translation, (num_examples, 1))  # [RS, 3]
    # rotmat = rotmat.unsqueeze(0)
    # rotmat = torch.tile(rotmat, (num_examples, 1, 1))

    # # pts_flat_transformed = torch.matmul(pts_flat[:, None, :], torch.transpose(rotmat, 2, 1))  # [RS, 1, 3]
    # pts_flat_transformed = pts_flat[:, None, :]  # [RS, 1, 3]
    # pts_flat_transformed += translation[:, None, :]  # [RS, 1, 3]
    # pts_transformed = pts_flat_transformed.view(pts.size())  # [R, S, 3]
    chunk = 256
    # Check batch transformations works without rotation.
    if check_numerics:
        transformations = np.eye(4)
        transformations[:3, 3] = translation
        transformations = torch.tensor(transformations, dtype=torch.float)  # [4, 4]
        transformations = torch.tile(transformations[None, ...], (pts.size()[0], 1, 1))  # [R, 4, 4]
        pts_transformed1 = []
        for i in range(0, pts.size()[0], chunk):
            pts_transformed1_chunk = apply_batched_transformations(
                inputs=[pts[i:i+chunk]], transformations=transformations[i:i+chunk])[0]
            pts_transformed1.append(pts_transformed1_chunk)
        pts_transformed1 = torch.cat(pts_transformed1, dim=0)

        pts_transformed2 = pts + translation[None, None, :]

    # Now add rotation
    transformations = np.eye(4)
    transformations = torch.tensor(transformations, dtype=torch.float)
    transformations[:3, :3] = rotmat
    transformations[:3, 3] = translation
    #transformations = torch.tensor(transformations, dtype=torch.float)  # [4, 4]
    transformations = torch.tile(transformations[None, ...], (pts.size()[0], 1, 1))  # [R, 4, 4]
    pts_transformed = []
    for i in range(0, pts.size()[0], chunk):
        pts_transformed_chunk = apply_batched_transformations(
            inputs=[pts[i:i+chunk]], transformations=transformations[i:i+chunk])[0]
        pts_transformed.append(pts_transformed_chunk)
    pts_transformed = torch.cat(pts_transformed, dim=0)
    return pts_transformed

# def transform_rays(ray_batch, translation, use_viewdirs):
#    """Apply transformation to rays.

#    Args:
#        ray_batch: [R, M] float tensor. All information necessary
#            for sampling along a ray, including: ray origin, ray direction, min
#            dist, max dist, and unit-magnitude viewing direction.
#        translation: [3,] float tensor. The (x, y, z) translation to apply.
#        use_viewdirs: Whether to use view directions.

#    Returns:
#        ray_batch: [R, M] float tensor. Transformed ray batch.
#    """
#   assert translation.size()[0] == 3, "translation.size()[0] must be 3..."

#    # Since we are only supporting translation for now, only ray origins need to be
#    # modified. Ray directions do not need to change.
#    rays_o = ray_batch[:, 0:3] + translation
#    rays_remaining = ray_batch[:, 3:]
#    ray_batch = torch.cat((rays_o, rays_remaining), dim=1)
#    return ray_batch

def compute_rays_length(rays_d):
    """Compute ray length.

    Args:
        rays_d: [R, 3] float tensor. Ray directions.

    Returns:
        rays_length: [R, 1] float tensor. Ray lengths.
    """
    rays_length = torch.norm(rays_d, dim=-1, keepdim=True)  # [N_rays, 1]
    return rays_length

def normalize_rays(rays):
    """Normalize ray directions.

    Args:
        rays: [R, 3] float tensor. Ray directions.

    Returns:
        normalized_rays: [R, 3] float tensor. Normalized ray directions.
    """
    normalized_rays = rays / compute_rays_length(rays_d=rays)
    return normalized_rays

def compute_ray_dirs_and_length(rays_o, rays_dst):
    """Compute ray directions.

    Args:
        rays_o: [R, 3] float tensor. Ray origins.
        rays_dst: [R, 3] float tensor. Ray destinations.

    Returns:
        rays_d: [R, 3] float tensor. Normalized ray directions.
    """
    # The ray directions are the difference between the ray destinations and the
    # ray origins.
    rays_d = rays_dst - rays_o  # [R, 3]  # Direction out of light source

    # Compute the length of the rays.
    rays_length = compute_rays_length(rays_d=rays_d)

    # Normalized the ray directions.
    rays_d = rays_d / rays_length  # [R, 3]  # Normalize direction
    return rays_d, rays_length

def update_ray_batch_slice(ray_batch, x, start, end):
    left = ray_batch[:, :start]  # [R, ?]
    right = ray_batch[:, end:]  # [R, ?]
    updated_ray_batch = torch.cat((left, x, right), dim=-1)
    return updated_ray_batch

def update_ray_batch_bounds(ray_batch, bounds):
    updated_ray_batch = update_ray_batch_slice(ray_batch=ray_batch, x=bounds,
                                               start=6, end=8)
    return updated_ray_batch

def create_ray_batch(
    rays_o, rays_dst, rays_i, use_viewdirs, rays_near=None, rays_far=None, epsilon=1e-10):
    # Compute the ray directions.
    rays_d = rays_dst - rays_o  # [R,3]  # Direction out of light source
    rays_length = compute_rays_length(rays_d=rays_d)  # [R, 1]
    rays_d = rays_d / rays_length  # [R, 3]  # Normalize direction
    viewdirs = rays_d  # [R, 3]

    # If bounds are not provided, set the beginning and end of ray as sampling bounds.
    if rays_near is None:
        rays_near = torch.zeros((rays_o.size()[0], 1), dtype=torch.float) + epsilon  # [R, 1]
    if rays_far is None:
        rays_far = rays_length  # [R, 1]

    ray_batch = torch.cat((rays_o, rays_d, rays_near, rays_far), dim=-1)
    if use_viewdirs:
        ray_batch = torch.cat((ray_batch, viewdirs), dim=-1)
    ray_batch = torch.cat((ray_batch, rays_i), dim=-1)
    return ray_batch

def sample_random_lightdirs(num_rays, num_samples, upper_only=False):
    """Randomly sample directions in the unit sphere.

    Args:
        num_rays: int or tensor shape dimension. Number of rays.
        num_samples: int or tensor shape dimension. Number of samples per ray.
        upper_only: bool. Whether to sample only on the upper hemisphere.

    Returns:
        lightdirs: [R, S, 3] float tensor. Random light directions sampled from the unit
            sphere for each sampled point.
    """
    if upper_only:
        min_z = 0
    else:
        min_z = -1

    phi = torch.rand(num_rays, num_samples) * (2 * math.pi)  # [R, S]
    cos_theta = torch.rand(num_rays, num_samples) * (1 - min_z) + min_z  # [R, S]
    theta = torch.acos(cos_theta)  # [R, S]

    x = torch.sin(theta) * torch.cos(phi)
    y = torch.sin(theta) * torch.sin(phi)
    z = torch.cos(theta)

    lightdirs = torch.cat((x[..., None], y[..., None], z[..., None]), dim=-1)  # [R, S, 3]
    return lightdirs

def get_light_positions(rays_i, img_light_pos):
    """Extracts light positions given scene IDs.

    Args:
        rays_i: [R, 1] float tensor. Per-ray image IDs.
        img_light_pos: [N, 3] float tensor. Per-image light positions.

    Returns:
        rays_light_pos: [R, 3] float tensor. Per-ray light positions.
    """
    rays_light_pos = img_light_pos[rays_i.long()].squeeze()  # [R, 3]
    return rays_light_pos

def get_lightdirs(lightdirs_method, num_rays=None, num_samples=None, rays_i=None,
    metadata=None, ray_batch=None, use_viewdirs=False, normalize=False):
    """Compute lightdirs.

    Args:
        lightdirs_method: str. Method to use for computing lightdirs.
        num_rays: int or tensor shape dimension. Number of rays.
        num_samples: int or tensor shape dimension. Number of samples per ray.
        rays_i: [R, 1] float tensor. Ray image IDs.
        metadata: [N, 3] float tensor. Metadata about each image. Currently only light
            position is provided.
        ray_batch: [R, M] float tensor. Ray batch.
        use_viewdirs: bool. Whether to use viewdirs.
        normalize: bool. Whether to normalize lightdirs.

    Returns;
        lightdirs: [R, S, 3] float tensor. Light directions for each sample.
    """
    if lightdirs_method == 'viewdirs':
        raise NotImplementedError
        assert use_viewdirs
        lightdirs = ray_batch[:, 8:11]  # [R, 3]
        lightdirs *= 1.5
        lightdirs = torch.tile(lightdirs[:, None, :], (1, num_samples, 1))
    elif lightdirs_method == 'metadata':
        lightdirs = get_light_positions(rays_i, metadata)  # [R, 3]
        lightdirs = torch.tile(lightdirs[:, None, :], (1, num_samples, 1))  # [R, S, 3]
    elif lightdirs_method == 'random':
        lightdirs = sample_random_lightdirs(num_rays, num_samples)  # [R, S, 3]
    elif lightdirs_method == 'random_upper':
        lightdirs = sample_random_lightdirs(num_rays, num_samples, upper_only=True)  # [R, S, 3]
    else:
        raise ValueError(f'Invalid lightdirs_method: {lightdirs_method}.')
    if normalize:
        lightdirs_flat = lightdirs.view(-1, 3)  # [RS, 3]
        lightdirs_flat = normalize_rays(lightdirs_flat)  # [RS, 3]
        lightdirs = lightdirs_flat.view(lightdirs.size())  # [R, S, 3]
    return lightdirs

ext: py

max_stars_repo_name: rhgao/ObjectFolder

max_issues_repo_name: rhgao/ObjectFolder

max_forks_repo_name: rhgao/ObjectFolder

reason: import numpy,from scipy


text:
import os
import pickle

import sys; sys.path.append('.')
import numpy as np
import matplotlib.pylab as plt

# parameters to include in the plot
SAVE_PLOT = True
################################################################################
#                             Make plot                                        #
################################################################################

# 0. build figures
fig = plt.figure(1, figsize=(6, 6))

# 1. load performances and experiment parameters
try:
    experiment_tag = sys.argv[1]
except:
    raise ValueError('Please specify a valid experiment tag.')

print('\nCalculating learning capacity for the Random Sequence Task...')

experiment_folder = 'RandomSequenceTask_{}'.format(experiment_tag)
experiment_path = 'backup/{}/'.format(experiment_folder)
experiment_n = len(os.listdir(experiment_path))
performance = np.zeros(experiment_n)
l_list = np.zeros(experiment_n, dtype=np.int)
n_list = np.zeros(experiment_n, dtype=np.int)

for exp, exp_name in enumerate(os.listdir(experiment_path)):
    exp_n = [int(s) for s in exp_name.split('_') if s.isdigit()]
    stats = pickle.load(open('{}/stats.p'.format(experiment_path+exp_name), 'rb'))
    params = pickle.load(open('{}/init_params.p'.format(experiment_path+exp_name), 'rb'))
    performance[exp] = stats.performance
    l_list[exp] = int(params.L)
    n_list[exp] = int(params.N_e)

print(n_list)
print(performance)

performance_mean = np.zeros(len(np.unique(l_list)))
for i, l in enumerate(np.unique(l_list)):
    performance_mean[i] = performance[np.where(l_list == l)].mean()
plt.plot(np.unique(l_list), 1-performance_mean, 'o', label=experiment_tag[5:])

# 4. adjust figure parameters and save
fig_lettersize = 15
plt.legend(loc=(0.05, 0.55), frameon=False, fontsize=fig_lettersize, title='Network size')
plt.xlabel(r'$L$', fontsize=fig_lettersize)
plt.ylabel(r'error', fontsize=fig_lettersize)
# plt.xlim([0, 8000])
plt.ylim([0., 0.9])
plt.xscale('log')

if SAVE_PLOT:
    plots_dir = 'plots/RandomSequenceTask/'
    if not os.path.exists(plots_dir):
        os.makedirs(plots_dir)
    plt.savefig('{}LC.pdf'.format(plots_dir), format='pdf')
plt.show()

max_stars_repo_name: carolinscholl/SORN

max_issues_repo_name: carolinscholl/SORN

max_forks_repo_name: carolinscholl/SORN

reason: import numpy


text:
"""
Contains utility functions for accessing NUPACK 4 and ViennaRNA energy calculation algorithms.

The main functions are
:meth:`pfunc` (for calculating complex free energy with NUPACK, along with its helper functions
:meth:`secondary_structure_single_strand` and :meth:`binding`),
:meth:`nupack_complex_base_pair_probabilities` (for calculating base pair probabilities with NUPACK),
:meth:`rna_duplex_multiple` (for calculating an approximation to two-strand complex free energy
that is much faster than calling :meth:`pfunc` on the same pair of strands).
"""  # noqa

import collections
import math
import itertools
import os
import logging
import random
import subprocess as sub
import sys
from functools import lru_cache
from multiprocessing.pool import ThreadPool
from typing import Sequence, Union, Tuple, List, Dict, Optional, cast, Deque

import numpy as np

import dsd.constraints as dc

os_is_windows = sys.platform == 'win32'

parameter_set_directory = 'nupack_viennaRNA_parameter_files'

default_vienna_rna_parameter_filename = 'dna_mathews1999.par'  # closer to nupack than dna_mathews2004.par

default_temperature = 37.0
default_magnesium = 0.0125
default_sodium = 0.05

_cached_nupack_models = {}

def calculate_strand_association_penalty(temperature: float, num_seqs: int) -> float:
    """
    Additive adjustment factor to convert NUPACK's mole fraction units to molarity.

    For details on why this is needed for multi-stranded complexes, see Section S1.1 of
    http://www.nupack.org/downloads/serve_public_file/fornace20_supp.pdf?type=pdf and Figure 2 of
    http://www.nupack.org/downloads/serve_public_file/nupack_user_guide_3.2.2.pdf?type=pdf

    :param temperature:
        temperature in Celsius
    :param num_seqs:
        number of sequences
    :return:
        Additive adjustment factor to convert NUPACK's mole fraction units to molar.
    """
    r = 0.0019872041  # Boltzmann's constant in kcal/mol/K (value on Wikipedia under Molar Gas Constant:
    # https://en.wikipedia.org/wiki/Gas_constant)
    # r = 0.001985875  # Boltzmann's constant in kcal/mol/K (value on Wikipedia under Boltzman's Constant:
    # https://en.wikipedia.org/wiki/Boltzmann_constant#Value_in_different_units, not sure why different
    # from the Molar Gas Constant, but luckily it's not until the fourth non-zero digit)
    water_conc = 55.14  # molar concentration of water at 37 C; ignore temperature dependence, ~5%
    temperature_kelvin = temperature + 273.15  # Kelvin
    # converts from NUPACK mole fraction units to molar units, per association
    adjust = r * temperature_kelvin * math.log(water_conc)
    return adjust * (num_seqs - 1)

@lru_cache(maxsize=10_000)
def pfunc(seqs: Union[str, Tuple[str, ...]],
          temperature: float = default_temperature,
          sodium: float = default_sodium,
          magnesium: float = default_magnesium,
          strand_association_penalty: bool = True,
          ) -> float:
    """
    Calls pfunc from NUPACK 4 (http://www.nupack.org/) on a complex consisting of the unique strands in
    seqs, returns energy ("delta G"), i.e., generally a negative number.

    By default, a strand association penalty is applied that is not applied by NUPACK's pfunc.
    See `strand_association_penalty` parameter documentation for details.

    NUPACK 4 must be installed. Installation instructions can be found at
    https://piercelab-caltech.github.io/nupack-docs/start/.

    :param seqs:
        DNA sequences (tuple, or a single DNA sequence),
        whose order indicates a cyclic permutation of the complex.
        For one or two sequences, there is only one cyclic permutation,
        so the order doesn't matter in such cases.
    :param temperature:
        temperature in Celsius
    :param sodium:
        molarity of sodium in moles per liter
    :param magnesium:
        molarity of magnesium in moles per liter
    :param strand_association_penalty:
        Add strand association penalty for a complex, related to converting NUPACK's mole fraction units
        to molarity. The quantity added is that returned by :meth:`calculate_strand_association_penalty`
        with parameters `temperature` and `len(seqs)`.
        For most constraints, which involve only one size of complex, this factor won't matter other than
        to adjust the energy threshold by the same factor. The factor depends only on the number of strands
        in `seqs`, but not on their sequences. However, this factor is needed for a meaningful comparison
        of energies between complexes of different sizes, e.g., to calculate equilibrium concentrations
        of complexes of various sizes.
        For details on why this is needed for multi-stranded complexes, see Section S1.1 of
        http://www.nupack.org/downloads/serve_public_file/fornace20_supp.pdf?type=pdf and Figure 2 of
        http://www.nupack.org/downloads/serve_public_file/nupack_user_guide_3.2.2.pdf?type=pdf
    :return:
        complex free energy ("delta G") of ordered complex with strands in given cyclic permutation
    """
    if isinstance(seqs, str):
        seqs = (seqs,)

    try:
        from nupack import pfunc as nupack_pfunc  # type: ignore
        from nupack import Model  # type: ignore
    except ModuleNotFoundError:
        raise ImportError(
            'NUPACK 4 must be installed to use pfunc. Installation instructions can be found at '
            'https://piercelab-caltech.github.io/nupack-docs/start/.')

    # expensive to create a Model, so don't create the same one twice
    param = (temperature, sodium, magnesium)
    if param not in _cached_nupack_models:
        model = Model(celsius=temperature, sodium=sodium, magnesium=magnesium, material='dna')
        _cached_nupack_models[param] = model
    else:
        model = _cached_nupack_models[param]
    (_, dg) = nupack_pfunc(strands=seqs, model=model)

    if strand_association_penalty and len(seqs) > 1:
        dg += calculate_strand_association_penalty(temperature, len(seqs))

    return dg

def nupack_complex_base_pair_probabilities(strand_complex: 'dc.Complex', # circular import causes problems
                                           temperature: float = default_temperature,
                                           sodium: float = default_sodium,
                                           magnesium: float = default_magnesium) -> np.ndarray:
    """
    Calculates base-pair probabilities according to NUPACK 4.

    :param strand_complex:
        Ordered tuple of strands in complex (specifying a particular circular ordering, which is
        imposed on all considered secondary structures)
    :param temperature:
        temperature in Celsius
    :param sodium:
        molarity of sodium in moles per liter
    :param magnesium:
        molarity of magnesium in moles per liter
    :return:
        2D Numpy array of floats, with `result[i1][i2]` giving the base-pair probability of base at position
        `i1` with base at position `i2` (if `i1` != `i2`), where `i1` and `i2` are the absolute positions
        of the bases in the entire ordered list of strands. For example, with strands AAAA and TTTTT,
        there are nine indices 0,1,2,3,4,5,6,7,8, with positions 0,1,2,3 on the first strand AAAA,
        and positions 4,5,6,7,8 on the second strand TTTTT.
        If `i1` == `i2`, then `result[i1][i1]` is the probability that the base at position `i1` is
        *unpaired*.
    """
    try:
        from nupack import Complex as NupackComplex
        from nupack import Model as NupackModel
        from nupack import ComplexSet as NupackComplexSet
        from nupack import Strand as NupackStrand
        from nupack import SetSpec as NupackSetSpec
        from nupack import complex_analysis as nupack_complex_analysis
        from nupack import PairsMatrix as NupackPairsMatrix
        from nupack import Model  # type: ignore
    except ModuleNotFoundError:
        raise ImportError(
            'NUPACK 4 must be installed to use nupack_complex_base_pair_probabilities. '
            'Installation instructions can be found at '
            'https://piercelab-caltech.github.io/nupack-docs/start/.')

    param = (temperature, sodium, magnesium)
    if param not in _cached_nupack_models:
        model = Model(celsius=temperature, sodium=sodium, magnesium=magnesium, material='dna')
        _cached_nupack_models[param] = model
    else:
        model = _cached_nupack_models[param]

    nupack_strands = [NupackStrand(strand_.sequence(), name=strand_.name) for strand_ in strand_complex]
    nupack_complex: NupackComplex = NupackComplex(nupack_strands)
    nupack_complex_set = NupackComplexSet(
        nupack_strands, complexes=NupackSetSpec(max_size=0, include=(nupack_complex,)))
    nupack_complex_analysis_result = nupack_complex_analysis(
        nupack_complex_set, compute=['pairs'], model=model)
    pairs: NupackPairsMatrix = nupack_complex_analysis_result[nupack_complex].pairs
    nupack_complex_result: np.ndarray = pairs.to_array()
    return nupack_complex_result

def call_subprocess(command_strs: List[str], user_input: str) -> Tuple[str, str]:
    """
    Calls system command through a subprocess. Assumes running on a POSIX operating system.

    If running on Windows, automatically appends "wsl -e" to start of command to call command
    through Windows subsystem for Linux, so wsl must be installed for this to work:
    `https://docs.microsoft.com/en-us/windows/wsl/install-win10 <https://docs.microsoft.com/en-us/windows/wsl/install-win10>`_

    :param command_strs:
        List of command and command line arguments, i.e., to call ``ls -l -a``,
        `command_strs` should be the list ['ls', '-l', '-a'].
    :param user_input:
        Input to give once program is running (i.e., what would user type).
    :return:
        pair of strings (output, error), giving the strings written to stdout and stderr, respectively.
    """
    # When porting the code from python2 to python3 we found an issue with sub.Popen().
    # Passing either of the keyword arguments universal_newlines=True or encoding='utf8'
    # solves the problem for python3.6. For python3.7 (but not 3.6) one can use text=True
    # XXX: Then why are none of those keyword arguments being used here??
    process: Optional[sub.Popen] = None
    command_strs = (['wsl.exe', '-e'] if os_is_windows else []) + command_strs
    try:
        with sub.Popen(command_strs, stdin=sub.PIPE, stdout=sub.PIPE, stderr=sub.PIPE) as process:
            output, stderr = process.communicate(user_input.encode())
    except BaseException as error:
        if process is not None:
            process.kill()
        raise error
    output_decoded = output.decode()
    stderr_decoded = stderr.decode()
    return output_decoded, stderr_decoded

# https://github.com/python/cpython/blob/42336def77f53861284336b3335098a1b9b8cab2/Lib/functools.py#L485
_sentinel = object()
_max_size = 1_000_000
_rna_duplex_cache: Dict[Tuple[Tuple[str, str], float, str], float] = {}
_rna_duplex_queue: Deque[Tuple[Tuple[str, str], float, str]] = collections.deque(maxlen=_max_size)

def _rna_duplex_single_cacheable(seq_pair: Tuple[str, str], temperature: float,
                                 parameters_filename: str) -> Optional[float]:
    # return None if arguments are not in _rna_duplex_cache,
    # otherwise return cached energy in _rna_duplex_cache
    key = (seq_pair, temperature, parameters_filename)
    result: Union[object, float] = _rna_duplex_cache.get(key, _sentinel)
    if result is _sentinel:
        return None
    else:
        result_float: float = cast(float, result)
        return result_float

def rna_duplex_multiple(seq_pairs: Sequence[Tuple[str, str]],
                        logger: logging.Logger = logging.root,
                        temperature: float = default_temperature,
                        parameters_filename: str = default_vienna_rna_parameter_filename,
                        cache: bool = True,
                        # cache: bool = False,  # off until I implement LRU queue to bound cache size
                        ) -> List[float]:
    """
    Calls RNAduplex (from ViennaRNA package: https://www.tbi.univie.ac.at/RNA/)
    on a list of pairs, specifically:
    [ (seq1, seq2), (seq2, seq3), (seq4, seq5), ... ]
    where seqi is a string over {A,C,T,G}. Temperature is in Celsius.
    Returns a list (in the same order as seqpairs) of free energies.

    :param seq_pairs:
        sequence (list or tuple) of pairs of DNA sequences
    :param logger:
        logger to use for printing error messages
    :param temperature:
        temperature in Celsius
    :param parameters_filename:
        name of parameters file for NUPACK
    :param cache:
        Whether to cache results to save on number of sequences to give to RNAduplex.
    :return:
        list of free energies, in the same order as `seq_pairs`
    """
    # print(f'rna_duplex_multiple.lru_cache = {rna_duplex_multiple.cache_info()}')

    # NB: the string NA_parameter_set needs to be exactly the intended filename;
    # e.g. any extra whitespace characters cause RNAduplex to default to RNA parameter set
    # without warning the user!

    # Note that loading parameter set dna_mathews2004.par throws a warning encoded in that parameter set:
    # WARNING: stacking enthalpies not symmetric

    # https://stackoverflow.com/questions/10174211/how-to-make-an-always-relative-to-current-module-file-path

    # fill in cached energies and determine which indices still need to have their energies calculated
    if cache:
        energies = [_rna_duplex_single_cacheable(seq_pair, temperature, parameters_filename)
                    for seq_pair in seq_pairs]
    else:
        energies = [None] * len(seq_pairs)
    idxs_to_calculate = [i for i, energy in enumerate(energies) if energy is None]
    idxs_to_calculate_set = set(idxs_to_calculate)
    seq_pairs_to_calculate = [seq_pair for i, seq_pair in enumerate(seq_pairs) if i in idxs_to_calculate_set]

    full_parameters_filename = os.path.join(os.path.dirname(__file__),
                                            parameter_set_directory, parameters_filename)

    if os_is_windows:
        full_parameters_filename = _fix_filename_windows(full_parameters_filename)

    command_strs: List[str] = ['RNAduplex', '-P', full_parameters_filename, '-T', str(temperature),
                               '--noGU', '−−noconv']

    # DNA sequences to type after RNAduplex starts up
    user_input = '\n'.join(f'{seq_pair[0]}\n{seq_pair[1]}' for seq_pair in seq_pairs_to_calculate) + '\n@\n'

    output, error = call_subprocess(command_strs, user_input)

    if error.strip() != '':
        logger.warning('error from RNAduplex: ', error)
        if error.split('\n')[0] != 'WARNING: stacking enthalpies not symmetric':
            raise ValueError('I will ignore errors about "stacking enthalpies not symmetric", but this '
                             'is a different error that I don\'t know how to handle. Exiting...'
                             f'\nerror:\n{error}')

    energies_to_calculate: List[float] = []
    lines = output.split('\n')
    for line in lines[:-1]:
        energies_to_calculate.append(float(line.split(':')[1].split('(')[1].split(')')[0]))
    if len(lines) - 1 != len(seq_pairs_to_calculate):
        raise ValueError(f'lengths do not match: #lines:{len(lines) - 1} #seqpairs:{len(seq_pairs)}')

    assert len(energies_to_calculate) == len(seq_pairs_to_calculate)

    # put calculated energies into list to return alongside cached energies
    assert len(idxs_to_calculate) == len(energies_to_calculate)
    for i, energy, seq_pair in zip(idxs_to_calculate, energies_to_calculate, seq_pairs_to_calculate):
        energies[i] = energy
        if cache:
            key = (seq_pair, temperature, parameters_filename)
            _rna_duplex_cache[key] = energy

            # clear out oldest cache key if _rna_duplex_queue is full
            if len(_rna_duplex_queue) == _rna_duplex_queue.maxlen:
                lru_item = _rna_duplex_queue[0]
                del _rna_duplex_cache[lru_item]

            _rna_duplex_queue.append(key)

    return energies

def _fix_filename_windows(parameters_filename: str) -> str:
    # FIXME: Here's the story: when developing in Windows, os.path will construct the path with Windows
    #  absolute paths. But we need to pass off the computation to wsl.exe (Windows Subsystem for Linux),
    #  which expects Linux-style paths (and has no idea what to do with 'C:\'). So we manually translate
    #  the absolute path. But this is fugly, and we should be not using absolute paths in this way.
    for drive in ['C', 'c', 'D', 'd', 'E', 'e', 'F', 'f']:
        parameters_filename = parameters_filename.replace(f'{drive}:\\', f'/mnt/{drive.lower()}/')
    parameters_filename = parameters_filename.replace('\\', '/')
    return parameters_filename

def rna_cofold_multiple(seq_pairs: Sequence[Tuple[str, str]],
                        logger: logging.Logger = logging.root,
                        temperature: float = default_temperature,
                        parameters_filename: str = default_vienna_rna_parameter_filename,
                        ) -> List[float]:
    """
    Calls RNAcofold (from ViennaRNA package: https://www.tbi.univie.ac.at/RNA/)
    on a list of pairs, specifically:
    [ (seq1, seq2), (seq2, seq3), (seq4, seq5), ... ]
    where seqi is a string over {A,C,T,G}. Temperature is in Celsius.
    Returns a list (in the same order as seqpairs) of free energies.

    :param seq_pairs:
        sequence (list or tuple) of pairs of DNA sequences
    :param logger:
        logger to use for printing error messages
    :param temperature:
        temperature in Celsius
    :param parameters_filename:
        name of NUPACK parameters file
    :return:
        list of free energies, in the same order as `seq_pairs`
    """

    # NB: the string NA_parameter_set needs to be exactly the intended filename;
    # e.g. any extra whitespace characters cause RNAduplex to default to RNA parameter set
    # without warning the user!

    # Note that loading parameter set dna_mathews2004.par throws a warning encoded in that parameter set:
    # WARNING: stacking enthalpies not symmetric

    # https://stackoverflow.com/questions/10174211/how-to-make-an-always-relative-to-current-module-file-path
    full_parameters_filename = os.path.join(os.path.dirname(__file__),
                                            parameter_set_directory, parameters_filename)

    if os_is_windows:
        full_parameters_filename = _fix_filename_windows(full_parameters_filename)

    # DNA sequences to type after RNAcofold starts up
    user_input = '\n'.join(seqpair[0] + '&' + seqpair[1] for seqpair in seq_pairs) + '\n@\n'

    command_strs: List[str] = ['RNAcofold', '-P', full_parameters_filename, '-T', str(temperature),
                               '--noGU', '−−noconv', '-p']

    output, stderr = call_subprocess(command_strs, user_input)

    if stderr.strip() != '':
        logger.warning('error from RNAduplex: ', stderr)
        if stderr.split('\n')[0] != 'WARNING: stacking enthalpies not symmetric':
            raise ValueError('I will ignore errors about "stacking enthalpies not symmetric", but this '
                             'is a different error that I don\'t know how to handled. Exiting.')

    lines = output.split('\n')
    dg_list: List[float] = []
    for line in lines[:-1]:
        dg_list.append(-float(line.split(':')[1].split('(')[1].split(')')[0]))
    if len(lines) - 1 != len(seq_pairs):
        raise AssertionError(f'lengths do not match: #lines:{len(lines) - 1} #seqpairs:{len(seq_pairs)}')

    return dg_list

_wctable = str.maketrans('ACGTacgt', 'TGCAtgca')

def wc(seq: str) -> str:
    """Return reverse Watson-Crick complement of `seq`."""
    return seq.translate(_wctable)[::-1]

def complex_free_energy_single_strand(
        seq: str, temperature: float = default_temperature, sodium: float = default_sodium,
        magnesium: float = default_magnesium) -> float:
    """Computes the complex free energy of a single strand.

    NUPACK 4 must be installed. Installation instructions can be found at
    https://piercelab-caltech.github.io/nupack-docs/start/.
    """
    return pfunc((seq,), temperature, sodium, magnesium)

def binding_complement(seq: str, temperature: float = default_temperature, sodium: float = default_sodium,
                       magnesium: float = default_magnesium, subtract_indv: bool = True) -> float:
    """Computes the complex free energy of a strand with its perfect Watson-Crick complement.

    NUPACK 4 must be installed. Installation instructions can be found at
    https://piercelab-caltech.github.io/nupack-docs/start/.
    """
    seq1 = seq
    seq2 = wc(seq)
    # this is a hack to save time since (seq1,seq2) and (seq2,seq1) are
    #   considered different tuples hence are cached differently by lrucache;
    #   but pfunc is a symmetric function with only two sequences, so it's safe to swap the order
    if seq1 > seq2:
        seq1, seq2 = seq2, seq1
    association_energy = pfunc((seq1, seq2), temperature, sodium, magnesium)
    if subtract_indv:
        # ddG_reaction = dG(products) - dG(reactants)
        association_energy -= (pfunc(seq1, temperature, sodium, magnesium) +
                               pfunc(seq2, temperature, sodium, magnesium))
    return association_energy

def binding(seq1: str, seq2: str, *, temperature: float = default_temperature,
            sodium: float = default_sodium, magnesium: float = default_magnesium) -> float:
    """Computes the complex free energy of association between two strands.

    NUPACK 4 must be installed. Installation instructions can be found at
    https://piercelab-caltech.github.io/nupack-docs/start/.
    """
    # this is a hack to save time since (seq1,seq2) and (seq2,seq1) are
    #   considered different tuples hence are cached differently by lrucache;
    #   but pfunc is a symmetric function so it's safe to swap the order
    if seq1 > seq2:
        seq1, seq2 = seq2, seq1
    return pfunc((seq1, seq2), temperature, sodium, magnesium) - (
            pfunc(seq1, temperature, sodium, magnesium) + pfunc(seq2, temperature, sodium, magnesium))

def random_dna_seq(length: int, bases: Sequence = 'ACTG') -> str:
    """Chooses a random DNA sequence."""
    return ''.join(random.choices(population=bases, k=length))

LOG_ENERGY = False

def log_energy(energy: float) -> None:
    if LOG_ENERGY:
        print(f'{energy:.1f}')

global_thread_pool = ThreadPool()

def domain_orthogonal(seq: str, seqs: Sequence[str], temperature: float, sodium: float,
                      magnesium: float, orthogonality: float,
                      orthogonality_ave: float = -1, parallel: bool = False) -> bool:
    """test orthogonality of domain with all others and their wc complements

    NUPACK 4 must be installed. Installation instructions can be found at
    https://piercelab-caltech.github.io/nupack-docs/start/.
    """
    if parallel:
        raise NotImplementedError()

    def binding_callback(s1: str, s2: str) -> float:
        return binding(s1, s2, temperature=temperature, sodium=sodium, magnesium=magnesium)

    if parallel:
        results = [
            global_thread_pool.apply_async(binding_callback, args=(s, s))
            for s in (seq, wc(seq))]
        energies = [result.get() for result in results]
        if max(energies) > orthogonality:
            return False
    else:
        ss = binding(seq, seq, temperature=temperature, sodium=sodium, magnesium=magnesium)
        log_energy(ss)
        if ss > orthogonality:
            return False
        wsws = binding(wc(seq), wc(seq), temperature=temperature, sodium=sodium, magnesium=magnesium)
        log_energy(wsws)
        if wsws > orthogonality:
            return False
    energy_sum = 0.0
    for altseq in seqs:
        if parallel:
            results = [
                global_thread_pool.apply_async(binding_callback,
                                               args=(seq1, seq2, temperature, sodium, magnesium))
                for seq1, seq2 in itertools.product((seq, wc(seq)), (altseq, wc(altseq)))]
            energies = [result.get() for result in results]
            if max(energies) > orthogonality:
                return False
            energy_sum += sum(energies)
        else:
            sa = binding(seq, altseq, temperature=temperature, sodium=sodium, magnesium=magnesium)
            log_energy(sa)
            if sa > orthogonality:
                return False
            sw = binding(seq, wc(altseq), temperature=temperature, sodium=sodium, magnesium=magnesium)
            log_energy(sw)
            if sw > orthogonality:
                return False
            wa = binding(wc(seq), altseq, temperature=temperature, sodium=sodium, magnesium=magnesium)
            log_energy(wa)
            if wa > orthogonality:
                return False
            ww = binding(wc(seq), wc(altseq), temperature=temperature, sodium=sodium, magnesium=magnesium)
            log_energy(ww)
            if ww > orthogonality:
                return False
            energy_sum += sa + sw + wa + ww
    if orthogonality_ave > 0:
        energy_ave = energy_sum / (4 * len(seqs)) if len(seqs) > 0 else 0.0
        return energy_ave <= orthogonality_ave
    else:
        return True

def domain_pairwise_concatenated_no_sec_struct(seq: str, seqs: Sequence[str], temperature: float,
                                               sodium: float, magnesium: float,
                                               concat: float, concat_ave: float = -1,
                                               parallel: bool = False) -> bool:
    """test lack of secondary structure in concatenated domains

    NUPACK 4 must be installed. Installation instructions can be found at
    https://piercelab-caltech.github.io/nupack-docs/start/.
    """

    if parallel:
        raise NotImplementedError()

    energy_sum = 0.0
    for altseq in seqs:
        wc_seq = wc(seq)
        wc_altseq = wc(altseq)
        if parallel:
            results = [global_thread_pool.apply_async(complex_free_energy_single_strand,
                                                      args=(seq1 + seq2, temperature, sodium, magnesium)) for
                       (seq1, seq2) in
                       [(seq, altseq),
                        (seq, wc_altseq),
                        (wc_seq, altseq),
                        (wc_seq, wc_altseq),
                        (altseq, seq),
                        (wc_altseq, seq),
                        (altseq, wc_seq),
                        (wc_altseq, wc_seq)]]
            energies = [result.get() for result in results]
            #             print len(results)
            #             print 'pair: %s' % [round(e,1) for e in energies]
            if max(energies) > concat:
                return False
            energy_sum += sum(energies)
        else:
            seq_alt = complex_free_energy_single_strand(seq + altseq, temperature, sodium, magnesium)
            if seq_alt > concat:
                return False
            seq_wcalt = complex_free_energy_single_strand(seq + wc_altseq, temperature, sodium, magnesium)
            if seq_wcalt > concat:
                return False
            wcseq_alt = complex_free_energy_single_strand(wc_seq + altseq, temperature, sodium, magnesium)
            if wcseq_alt > concat:
                return False
            wcseq_wcalt = complex_free_energy_single_strand(wc_seq + wc_altseq, temperature, sodium,
                                                            magnesium)
            if wcseq_wcalt > concat:
                return False
            alt_seq = complex_free_energy_single_strand(altseq + seq, temperature, sodium, magnesium)
            if alt_seq > concat:
                return False
            alt_wcseq = complex_free_energy_single_strand(altseq + wc_seq, temperature, sodium, magnesium)
            if alt_wcseq > concat:
                return False
            wcalt_seq = complex_free_energy_single_strand(wc_altseq + seq, temperature, sodium, magnesium)
            if wcalt_seq > concat:
                return False
            wcalt_wcseq = complex_free_energy_single_strand(wc_altseq + wc_seq, temperature, sodium,
                                                            magnesium)
            if wcalt_wcseq > concat:
                return False
            energy_sum += (seq_alt + seq_wcalt + wcseq_alt + wcseq_wcalt +
                           alt_seq + alt_wcseq + wcalt_seq + wcalt_wcseq)
    if concat_ave > 0:
        energy_ave = energy_sum / (8 * len(seqs)) if len(seqs) > 0 else 0.0
        return energy_ave <= concat_ave
    else:
        return True

_binaryGCTable = str.maketrans('ACTG', '0101')

def domain_concatenated_no_4gc(seq: str, seqs: Sequence[str]) -> bool:
    """prevent {G,C}^4 under concatenation"""
    for altseq in seqs:
        catseq = altseq + seq + altseq
        strength = catseq.translate(_binaryGCTable)
        if '1111' in strength:
            return False
    return True

def domain_no_4gc(seq: str) -> bool:
    """prevent {G,C}^4"""
    return '1111' not in seq.translate(_binaryGCTable)

def domain_concatenated_no_4g_or_4c(seq: str, seqs: Sequence[str]) -> bool:
    """prevent G^4 and C^4 under concatenation"""
    for altseq in seqs:
        catseq = altseq + seq + altseq
        if 'GGGG' in catseq:
            #             print '|GGGG# seq: %s altseq: %s|' % (seq,altseq)
            return False
        if 'CCCC' in catseq:
            #             print '|CCCC# seq: %s altseq: %s|' % (seq,altseq)
            return False
    return True

max_stars_repo_name: dsd-test-nupack-in-github-workflow/dsd

max_issues_repo_name: dsd-test-nupack-in-github-workflow/dsd

max_forks_repo_name: dsd-test-nupack-in-github-workflow/dsd


text:
import numpy as np

import Symbol

class Concordance:
    def __init__(self):
        pass

    def calculate_concordance(self,X):

        return 0

class ConcordanceAndSmallestCoveringSet(Concordance):
    '''
    This class is related to the paper of "Concordance and the Smallest Covering Set of Preference Orderings"
    by Dr. Zhiwe Lin  at Ulster University
       Prof. Hui Wang at Ulster University
       Prof. Cees H. Elzinga at VU University Amsterdam and NIDI

    The contribution of the work includes:
        1) to measure concordance of preference orderings
        2) to extract meaningful descriptive patterns from preference orderings
	 Paper available at http://arxiv.org/pdf/1609.04722v1.pdf

     If you have any questions or want to report bugs,
     please contact Dr. Zhiwei Lin at http://scm.ulster.ac.uk/zhiwei.lin/
    '''

    def Upsilon(self, x, T_matrix,varPhi,u, i,omega):
        '''
        Equation (38)

        :param x:
        :type x:
        :param T_matrix:
        :type T_matrix:
        :param varPhi:
        :type varPhi:
        :param u:
        :type u:
        :param i:
        :type i:
        :param omega:
        :type omega:
        :return:
        :rtype:
        '''
        if omega[i]==0:
            return []
        Q=[j for j in range(i+1,len(x)) if  T_matrix[j][i]==1]

        if len(Q)==0:
            return [u]
        S=[]
        for j in Q:
            v=list(u)
            v.append(x[j])
            S.extend(self.Upsilon(x,T_matrix,varPhi,v,j, omega))

        return S

    def check_occurrence(self,x,T_matrix,LCSs):
            no_occurrences=[]
            for (i,sym) in enumerate(x):
                if T_matrix[i][i]==0:
                    continue
                occurrence=False
                for lcs in LCSs:
                    if sym in lcs:
                        occurrence=True
                        break
                if not occurrence:
                    no_occurrences.append(i)
            return no_occurrences

    def extractSmallestCoveringSet(self,x, T_matrix,varPhi, LCSs,omega):
        '''
        Extract the smallest covering set from X
        :param x:
        :type x:
        :param T_matrix:
        :type T_matrix:
        :param varPhi:
        :type varPhi:
        :param LCSs:
        :type LCSs:
        :param omega:
        :type omega:
        :return:
        :rtype:
        '''
        A=list(LCSs)

        no_occurrences=self.check_occurrence(x,T_matrix,A)

        while len(no_occurrences)!=0:
            i=no_occurrences[0]

            if omega[i]==0:
                continue

            C=self.Theta_function(x,T_matrix,varPhi,[x[i]],i,omega)
            for c in C:
                B=self.Upsilon(x,T_matrix,varPhi, c, i, omega)

            omega[i]=0

            if len(B)>0:
                A.extend(B)
            no_occurrences=self.check_occurrence(x,T_matrix,A)

        return A

    def Theta_function(self, x,T_matrix, varPhi, u, i,omega):
        '''
        Equation (34)
        :param x:
        :type x:
        :param T_matrix:
        :type T_matrix:
        :param varPhi:
        :type varPhi:
        :param u:
        :type u:
        :param i:
        :type i:
        :param omega:
        :type omega:
        :return:
        :rtype:
        '''

        if omega[i]==0:
            return []

        P=[j for j in range(i) if varPhi[j]+1==varPhi[i] and T_matrix[i][j]==1 ]

        if len(P)==0:
            return [u]
        S=[]
        for j in P:
            v=[x[j]]
            v.extend(u)
            S.extend(self.Theta_function(x,T_matrix,varPhi, v,j,omega))

        return S

    def extract_all_lcs(self, x, T_matrix, varPhi, omega, llcs):
        '''
        To return a set of  all longest common subsequences when
        calculate_concordance is called.
        '''

        R=[i for i in range(len(varPhi)) if varPhi[i]==llcs]
        LCSs=[]

        for i in R:
            LCSs.extend(self.Theta_function(x,T_matrix,varPhi,[x[i]],i,omega))

        return LCSs

    def calculate_concordance(self,X):
        '''
        Given a set X of ordering sequences, this will calculate the concordance  $\kappa(X)$ of X
        shown in the paper.

        :param X:
        :type X:
        :return $\kappa(X)$, $\ell(X)$, the set of the longest common subsequences and the smallest covering set :
        :rtype:
        '''
        x=X[0]
        I=[]

        #Initialization
        for y in X:
            I_k=[]
            for x_j in x:

                found=False
                for i in range(len(y)):
                    if y[i]==x_j:
                        I_k.append(i)
                        found=True
                        break

                if not found:
                    I_k.append(float('inf'))
            I.append(I_k)

        T_matrix=np.zeros((len(x),len(x)))

        for j in range(len(x)):
            for i in range(j+1):
                T_matrix[j][i]=1
                for k in range(len(X)):
                    if I[k][i]> I[k][j] or I[k][j]==float('inf'):
                        T_matrix[j][i]=0
                        break;
        psi=np.zeros(len(x)+1)
        varPhi=np.zeros(len(x))

        omega=[1 for i in x]
        psi[0]=1
        for j in range(len(x)):
            length=0
            psi[j+1]+=T_matrix[j][j]

            for i in range(0,j):
                psi[j+1]+=psi[(i+1)%(j+1)]*T_matrix[j][i]
                length=np.maximum(length,varPhi[i]*T_matrix[j][i])
            varPhi[j]=(length+1)*T_matrix[j][j]

        concordance=np.sum(psi)
        llcs=np.max(varPhi)
        LCSs=self.extract_all_lcs(x, T_matrix, varPhi, omega,llcs)
        SCS=self.extractSmallestCoveringSet(x, T_matrix,varPhi, LCSs,omega)

        return concordance,llcs,LCSs,SCS

if __name__=="__main__":
    def print_results(X,con_results):
        print("================================================")
        print("********************************")
        print("The Set X of the Ordering Sequences:")
        for x in X:
            print (x)
        print("********************************")
        print("k(X)= {}   --The concordance of X".format(con_results[0]))
        print("l(x)= {}   --The length of the longest common subsequences".format(con_results[1]))
        print("All LCSs: {} ".format(con_results[2]))
        print("The  Smallest Covering Set C(X): {} ".format(con_results[3]))
        print("================================================\n\n")

    X=[]
    X.append([Symbol.SymbolItem(i) for i in ['a', 'b', 'c', 'd','e', 'f']])
    X.append([Symbol.SymbolItem(i) for i in [ 'a', 'c', 'f', 'b','d','e']])
    X.append([Symbol.SymbolItem(i) for i in ['a',  'b', 'd','c', 'f', 'e']])
    con=ConcordanceAndSmallestCoveringSet()
    print_results(X,con.calculate_concordance(X))

    X=[]
    X.append([Symbol.SymbolItem(i) for i in ['a', 'b', 'c', 'd','f', 'e']])
    X.append([Symbol.SymbolItem(i) for i in [ 'a', 'd', 'b', 'c','f','e']])
    X.append([Symbol.SymbolItem(i) for i in ['a',  'd', 'f','b', 'e', 'c']])
    con=ConcordanceAndSmallestCoveringSet()
    print_results(X,con.calculate_concordance(X))

    X=[]
    X.append([Symbol.SymbolItem(i) for i in ['f', 'a', 'b', 'c', 'd', 'e']])
    X.append([Symbol.SymbolItem(i) for i in ['e', 'a', 'f', 'd', 'b', 'c']])
    X.append([Symbol.SymbolItem(i) for i in ['a', 'f', 'e', 'd', 'b', 'c']])
    con=ConcordanceAndSmallestCoveringSet()
    print_results(X,con.calculate_concordance(X))

    X=[]
    X.append([Symbol.SymbolItem(i) for i in ['a', 'b', 'c', 'd', 'e']])
    X.append([Symbol.SymbolItem(i) for i in ['a', 'b', 'd', 'c', 'e']])
    X.append([Symbol.SymbolItem(i) for i in ['b', 'd', 'c', 'e']])
    con=ConcordanceAndSmallestCoveringSet()
    print_results(X, con.calculate_concordance(X))

    X=[]
    X.append([Symbol.SymbolItem(i) for i in ['f', 'a', 'b', 'c', 'd', 'e']])
    X.append([Symbol.SymbolItem(i) for i in ['f', 'a', 'b', 'c', 'd', 'e']])
    X.append([Symbol.SymbolItem(i) for i in ['f', 'a', 'b', 'c', 'd', 'e']])
    con=ConcordanceAndSmallestCoveringSet()
    print_results(X, con.calculate_concordance(X))

    X=[]
    X.append([Symbol.SymbolItem(i) for i in ['a', 'b', 'c', 'd', 'e']])
    X.append([Symbol.SymbolItem(i) for i in ['e', 'a', 'd', 'b', 'c']])
    X.append([Symbol.SymbolItem(i) for i in ['a', 'e', 'd', 'b', 'c']])
    con=ConcordanceAndSmallestCoveringSet()
    print_results(X,con.calculate_concordance(X))

max_stars_repo_name: zhiweiuu/secs

max_issues_repo_name: zhiweiuu/secs

max_forks_repo_name: zhiweiuu/secs


text:
'''
Created on Sep 4, 2012

@author: rch
'''
from traits.api import \
    Int, Instance, Property, List, \
    cached_property, Event, on_trait_change

from traitsui.api import \
    View, Item, Group, HGroup, RangeEditor

from bmcs_beam.mxn.mxn_tree_node import \
    MxNTreeNode

from bmcs_beam.mxn.cross_section import \
    CrossSection

from bmcs_beam.mxn.reinf_layout import \
    RLCTexUniform

import numpy as np

from bmcs_beam.mxn.matrix_cross_section import \
    MatrixCrossSection, MCSGeoRect

class MxNDiagram(MxNTreeNode):

    modified = Event
    def set_modified(self):
        self.modified = True

    tree_node_list = List(Instance(CrossSection))
    def _tree_node_list_default(self):
        return [CrossSection(notify_change_ext=self.set_modified,
                            reinf=[RLCTexUniform(
                                      n_layers=12,
                                      material='default_fabric')],
                            matrix_cs=MatrixCrossSection(geo=MCSGeoRect(
                                      width=0.2, height=0.06), n_cj=20,
                                      material='default_mixture',
                                      material_law='constant'))]

    @on_trait_change('tree_node_list')
    def cs_changed(self):
        self.modified = True

    cs = Property(depends_on='tree_node_list')
    def _get_cs(self):
        val = self.tree_node_list[0]
        val.notify_change_ext = self.set_modified
        return val
    def _set_cs(self, val):
        self.tree_node_list = [val]

    eps_cu = Property()
    def _get_eps_cu(self):
        return -self.cs.matrix_cs_with_state.material_law_.eps_c_u

    strain_env_reinf = Property()
    '''Contains the reinforcement element that determines the strain envelope.
    '''
    def _get_strain_env_reinf(self):
        eps_lo_arr = np.array([[r, r.convert_eps_u_2_lo(eps_up=self.eps_cu)]
                               for r in self.cs.reinf_components_with_state])
        try:
            min_index = np.argmin(eps_lo_arr[:, 1])
        except:
            # Returns None in case there is no reinforcement.
            return None
        return eps_lo_arr[min_index][0]

    n_eps = Int(20, auto_set=False, enter_set=True)
    eps_range = Property(depends_on='n_eps,modified')
    @cached_property
    def _get_eps_range(self):
        eps_cu = self.eps_cu
        eps_ccu = 0.8 * eps_cu
        env_reinf = self.strain_env_reinf

        if env_reinf:
            '''Strain envelope for reinforced cross section
            '''
            eps_t_lo = env_reinf.convert_eps_u_2_lo(eps_up=eps_cu)
            eps_t_lo_0 = env_reinf.convert_eps_u_2_lo(eps_up=0.)
            eps_t_u = env_reinf.material_law_.eps_u

            # Strain arrays for the lower rim
            eps_cc_0 = np.linspace(eps_ccu, 0., self.n_eps)
            eps_0_tlo = np.linspace(0., eps_t_lo, self.n_eps)
            eps_tlo_tlo0 = np.linspace(eps_t_lo, eps_t_lo_0, self.n_eps)
            eps_tlo0_tu = np.linspace(eps_t_lo_0, eps_t_u, self.n_eps)

            # Strain arrays for the upper rim
            eps_cc_c = np.linspace(eps_ccu, eps_cu, self.n_eps)
            eps_c_const = self.eps_cu * np.ones_like(eps_cc_0)
            eps_c_0 = np.linspace(eps_cu, 0., self.n_eps)
            eps_0_tu = np.linspace(0., eps_t_u, self.n_eps)

            eps1 = np.vstack([eps_cc_0, eps_cc_c])
            eps2 = np.vstack([eps_0_tlo, eps_c_const])
            eps3 = np.vstack([eps_tlo_tlo0, eps_c_0])
            eps4 = np.vstack([eps_tlo0_tu, eps_0_tu])

            return np.hstack([eps1, eps2, eps3, eps4])
        else:
            '''Strain envelope for cross section without reinforcement
            '''
            # Strain arrays for the lower rim
            eps_cc_0 = np.linspace(eps_ccu, 0., self.n_eps)
            eps_00 = np.zeros_like(eps_cc_0)

            # Strain arrays for the upper rim
            eps_cc_c = np.linspace(eps_ccu, eps_cu, self.n_eps)
            eps_c_0 = np.linspace(eps_cu, 0., self.n_eps)

            eps1 = np.vstack([eps_cc_0, eps_cc_c])
            eps2 = np.vstack([eps_00, eps_c_0])

            return np.hstack([eps1, eps2])

    n_eps_range = Property(depends_on='n_eps,modified')
    @cached_property
    def _get_n_eps_range(self):
        return self.eps_range.shape[1]

    #===========================================================================
    # MN Diagram
    #===========================================================================

    def _get_MN_fn(self, eps_lo, eps_up):
        self.cs.set(eps_lo=eps_lo,
                    eps_up=eps_up)
        return (self.cs.M, self.cs.N)

    MN_vct = Property()
    def _get_MN_vct(self):
        return np.vectorize(self._get_MN_fn)

    MN_arr = Property(depends_on='modified,n_eps')
    @cached_property
    def _get_MN_arr(self):
        return self.MN_vct(self.eps_range[0, :], self.eps_range[1, :])

    #===========================================================================
    # f_eps Diagram
    #===========================================================================

    current_eps_idx = Int(1)  # , auto_set = False, enter_set = True)

    current_eps = Property(depends_on='current_eps_idx,modified')
    @cached_property
    def _get_current_eps(self):
        return self.eps_range[(0, 1), self.current_eps_idx - 1]

    current_MN = Property(depends_on='current_eps_idx,modified')
    @cached_property
    def _get_current_MN(self):
        return self._get_MN_fn(*self.current_eps)

    def plot_eps(self, ax):
        ax.plot(-self.eps_range,
                [0, self.cs.matrix_cs_with_state.geo.height],
                color='black')
        ax.plot(-self.current_eps,
                [0, self.cs.matrix_cs_with_state.geo.height],
                lw=3, color='red')

        ax.spines['left'].set_position('zero')
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        ax.spines['left'].set_smart_bounds(True)
        ax.spines['bottom'].set_smart_bounds(True)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')

    def plot_MN(self, ax):
        ax.plot(self.MN_arr[0], -self.MN_arr[1], lw=2, color='blue')
        ax.plot(self.current_MN[0], -self.current_MN[1], 'g.',
                markersize=20.0, color='red')

        ax.spines['left'].set_position('zero')
        ax.spines['bottom'].set_position('zero')
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        ax.spines['left'].set_smart_bounds(True)
        ax.spines['bottom'].set_smart_bounds(True)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')
        ax.grid(b=None, which='major')
        x1, x2, z1, z2 = ax.axis()
        ax.axis([0, x2, z1, z2])

    def plot_custom(self, ax, color='blue', linestyle='-',
                    linewidth=2, label='<unnamed>'):
        M = self.MN_arr[0]
        N = -self.MN_arr[1]
        ax.plot(M, N, lw=linewidth, color=color,
                ls=linestyle, label=label)

        ax.spines['left'].set_position('zero')
        ax.spines['bottom'].set_position('zero')
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        ax.spines['left'].set_smart_bounds(True)
        ax.spines['bottom'].set_smart_bounds(True)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')
        ax.grid(b=None, which='major')
        x1, x2, z1, z2 = ax.axis()
        z_min = np.min(np.hstack([[z1], N]))
        z_max = np.max(np.hstack([[z2], N]))
        x_max = np.max(np.hstack([[x2], M]))
        ax.axis([0, x_max, z_min, z_max])

    #===========================================================================
    # Visualisation related attributes
    #===========================================================================

    def plot(self, fig):
        ax1 = fig.add_subplot(1, 2, 1)
        self.plot_eps(ax1)
        ax2 = fig.add_subplot(1, 2, 2)
        self.plot_MN(ax2)

    node_name = 'MxN diagram'

    tree_view = View(Group(
                HGroup(
                Group(Item('n_eps', springy=True),
                      Item('current_eps_idx', editor=RangeEditor(low=1,
                                   high_name='n_eps_range',
                                   format='(%s)',
                                   mode='slider',
                                   auto_set=False,
                                   enter_set=False,
                                   ),
                           show_label=False,
                           ),
                      label='Discretization',
                      springy=True
                      ),
                springy=True,
                ),
                scrollable=True,
                ),
                width=1.0,
                height=0.8,
                resizable=True,
                buttons=['OK', 'Cancel'])

if __name__ == '__main__':
    pass


text:
import unittest
import blur
import numpy as np

class DiffOfDiffCheck(unittest.TestCase):

    def setUp(self):
        self.arr = np.array([0,1,2,3,4])

    def testDiffIndex0(self):
        self.assertEqual(blur._diffofdiff(self.arr, 0), 2)

    def testDiffIndex1(self):
        self.assertEqual(blur._diffofdiff(self.arr, 1), 1)

    def testDiff(self):
        self.assertEqual(blur._diffofdiff(self.arr, 2), 0)

    def testDiffPenultimateIndex(self):
        self.assertEqual(blur._diffofdiff(self.arr, 3), 1)

    def testUltimateIndex(self):
        self.assertEqual(blur._diffofdiff(self.arr, 4), 2)

    def testIndexErrorUnderflow(self):
        self.assertRaises(IndexError, blur._diffofdiff, [0], -1)

    def testIndexErrorOverflow(self):
        self.assertRaises(IndexError, blur._diffofdiff, [0], 1)

class ContrastCheck(unittest.TestCase):

    def setUp(self):
        self.arr = [50,20,30,40,10]

    def testContrastIndex0(self):
        self.assertEqual(blur._contrast(self.arr, 0), 0)

    def testContrastIndex1(self):
        self.assertEqual(blur._contrast(self.arr, 1), 30)

    def testContrastIndex2(self):
        self.assertEqual(blur._contrast(self.arr, 2), 10)

    def testContrastIndex3(self):
        self.assertEqual(blur._contrast(self.arr, 3), 10)

    def testContrastIndex4(self):
        self.assertEqual(blur._contrast(self.arr, 4), 30)

    def testIndexErrorUnderflow(self):
        self.assertRaises(IndexError, blur._contrast, [0], -1)

    def testIndexErrorOverflow(self):
        self.assertRaises(IndexError, blur._contrast, [0], 1)

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: buruzaemon/moojoos

max_issues_repo_name: buruzaemon/moojoos

max_forks_repo_name: buruzaemon/moojoos


text:
import concurrent.futures
import numpy.random as npr
from utils.parallel_utils import *
import os
import sys
import time
import copy
import collections
from tqdm import tqdm
from utils.visualization_utils import plot_action

sys.path.append('./')
sys.path.append('..')
from scripts.create_agent import createAgent
from utils.parameters import *
from storage.buffer import QLearningBufferExpert
import rospy
from src.envs.dual_bin_front_rear import DualBinFrontRear
from utils.logger import Logger
from utils.torch_utils import augmentData2Buffer

ExpertTransition = collections.namedtuple('ExpertTransition',
                                          'state obs action reward next_state next_obs done step_left expert')

def set_seed(s):
    np.random.seed(s)
    torch.manual_seed(s)
    torch.cuda.manual_seed(s)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def getCurrentObs(in_hand, obs):
    obss = []
    for i, o in enumerate(obs):
        obss.append((o.squeeze(), in_hand[i].squeeze()))
    return obss

def saveModelAndInfo(logger, agent):
    logger.saveModel(logger.num_episodes, env, agent, create_dir=True)
    logger.saveLossCurve(100)
    logger.saveRewards()
    logger.saveSGDtime()
    logger.saveLosses()
    logger.saveTdErrors()
    logger.saveStepLeftCurve(1000)
    logger.saveExpertSampleCurve(100)
    logger.saveLearningCurve(learning_curve_avg_window)
    logger.saveLearningCurve2(learning_curve_avg_window)

def train_step(agent, replay_buffer, logger):
    batch = replay_buffer.sample(batch_size, onpolicydata=sample_onpolicydata, onlyfailure=onlyfailure)
    loss, td_error = agent.update(batch)
    logger.trainingBookkeeping(loss, td_error.mean().item())
    logger.num_training_steps += 1

def main():
    if not no_bar:
        pbar = tqdm(total=max_episode)
        print('\n')
        pbar.set_description('Steps:0; Avg Reward:0.0; Reward:0.0; Explore:0.0; Loss:0.0; Time:0.0')
    timer_start = time.time()

    envs.ur5.moveToHome()
    envs.p_reset()
    envs.p_move_reward(return_reward=False)
    rospy.sleep(1)
    while logger.num_steps < max_episode:
        p_agent.eps = init_eps if logger.num_episodes < step_eps else final_eps
        action = Action.get_var('main')
        IsRobotReady.get_var('main')
        # print('getting action')
        envs.p_picking(action)
        # print('picked')
        t_move_reward_start = time.time()
        reward = envs.p_move_reward()
        t_move_reward = time.time() - t_move_reward_start
        rospy.sleep(max(1 - t_move_reward, 0.1))
        # print('checked reward')
        logger.num_steps += 1
        envs.p_place_move_center(is_request=(logger.num_steps != max_episode))
        IsSGDFinished.get_var('main')
        # print('placed')
        if not no_bar:
            timer_final = time.time()
            r_avg = np.array(logger.rewards).mean() if is_test \
                else logger.getCurrentAvgReward(learning_curve_avg_window)
            description = 'Grasp{}: r_avg={:.03f}; r={:}; eps={:.01f}; ' \
                          'loss={:.03f}; t_episode={:.02f}'.format(
                logger.num_steps, r_avg,
                reward, p_agent.eps, float(logger.getCurrentLoss()), timer_final - timer_start)
            pbar.set_description(description)
            timer_start = timer_final
            pbar.update(logger.num_steps - pbar.n)
        if logger.num_steps % (max_episode // num_saves) == 0 or logger.num_steps == max_episode:
            saveModelAndInfo(logger, agent)
            logger.saveCheckPoint(args, envs, agent, replay_buffer, save_envs=False)
        # print('SGD finished')
        # print('--------finished step', logger.num_steps, '---------')

    if is_test:
        print('test success rate: {:.2f} over {} grasps'.format(np.array(logger.rewards).mean(), max_episode))
    envs.ur5.moveToHome()
    State.set_var('main', SENTINEL)
    Action.set_var('main', SENTINEL)
    Reward.set_var('main', SENTINEL)
    Request.set_var('main', SENTINEL)
    print('training finished')

class AgentWrapper:
    """
    Wrap the agent to perform parallel thread training/ testing
    """

    def __init__(self):
        self.all_state = None  # all_state include (grasping_state, in_hands, obs)
        self.network = threading.Lock()
        self.obs = None
        self.in_hand_obs = None
        self.actions_star = None
        self.q_value_maps = None
        self.actions_star_idx = None
        self.eps = 1.
        self.last_grasp_succeed = -1
        if is_test:
            self.plot_logger = logger
            self.tau = test_tau
        else:
            self.plot_logger = None
            self.tau = train_tau
        primative_idx, x_idx, y_idx, z_idx, self.rot_idx = map(lambda a: action_sequence.find(a),
                                                               ['p', 'x', 'y', 'z', 'r'])

    def get_action(self):
        """
        Infer action from the q maps
        Triggered by the state s
        """
        while True:
            state = State.get_var('get_action')
            # print('got state')
            if state is SENTINEL:
                break
            # print('get_action getting network')
            self.network.acquire()
            # print('get_action got network')
            self.all_state = state
            states, in_hands, obs = self.all_state
            q_value_maps, actions_star_idx, actions_star, in_hand_obs = \
                agent.getBoltzmannActions(states, in_hands, obs, temperature=self.tau, eps=self.eps, return_patch=True)
            self.network.release()
            # print('got action')
            action = [*list(map(lambda x: x.item(), actions_star[0])), 0]
            # print('agent s action', actions_star_idx)
            self.obs = obs
            self.in_hand_obs = in_hand_obs
            self.actions_star = actions_star
            self.actions_star_idx = actions_star_idx
            self.q_value_maps = q_value_maps
            Action.set_var('get_action', action)
        print('get_action killed')

    def store_transition_SGD(self):
        """
        Store the transition to the buffer and perform SGD step
        Triggered by reward
        """
        while True:
            reward = Reward.get_var('store_transition_SGD')
            if reward is SENTINEL or self.all_state is SENTINEL:
                break

            states, in_hands, obs = self.all_state
            buffer_obs = getCurrentObs(in_hands, obs)
            steps_lefts = envs.getStepLeft()
            dones = torch.tensor(False, dtype=torch.float32).view(1)

            for i in range(num_processes):
                is_expert = False
                data = ExpertTransition(states[i], buffer_obs[i], self.actions_star_idx[i], reward[i], states[i],
                                        buffer_obs[i], dones[i], steps_lefts[i], torch.tensor(is_expert))
                augmentData2Buffer(replay_buffer, data, agent.rzs,
                                   onpolicy_data_aug_n, onpolicy_data_aug_rotate, onpolicy_data_aug_flip)

            logger.stepBookkeeping(reward.numpy(), steps_lefts.numpy(), dones.numpy())
            # print('transition augmented')
            # if logger.num_episodes >= training_offset:
            # print('store transition SGD getting network')
            self.network.acquire()
            # print('store transition SGD got network')
            if (not is_test and logger.num_episodes >= training_offset) \
                    or (is_test and is_baseline and not reward) \
                    or (is_test and not is_baseline):
                for training_iter in range(training_iters):
                    SGD_start = time.time()
                    train_step(agent, replay_buffer, logger)
                    logger.SGD_time.append(time.time() - SGD_start)
            if is_baseline and (self.last_grasp_succeed == 0 and reward == 1):
                agent.loadModel(load_model_pre)
            self.network.release()
            # print('trained a SGD step')
            # print('training')
            self.last_grasp_succeed = reward
            IsSGDFinished.set_var('store_transition_SGD', True)
            # if render and not reward:
            if render:
                plot_action(obs, agent, self.actions_star, self.actions_star_idx, self.q_value_maps, num_rotations,
                            patch_size, None, self.in_hand_obs, action_sequence, logger=self.plot_logger)
        print('store_transition_SGD killed')

if __name__ == '__main__':
    format_ = "%(asctime)s: %(message)s"
    logging.basicConfig(format=format_, level=logging.INFO,
                        datefmt="%H:%M:%S")
    logging.getLogger().setLevel(logging.DEBUG)
    rospy.init_node('robot_exp')

    start_time = time.time()
    if seed is not None:
        set_seed(seed)

    # setup environment
    envs = DualBinFrontRear(ws_center=robot_ws_center, ws_x=real_workspace_size, ws_y=real_workspace_size,
                            cam_size=(cam_size, cam_size), action_sequence=action_sequence,
                            in_hand_mode='raw', pick_offset=pick_place_offset,
                            place_offset=pick_place_offset, in_hand_size=patch_size,
                            obs_source=obs_source, safe_z_region=safe_z_region, bin_size=workspace_size,
                            place_open_pos=place_open_pos, bin_size_pixel=112, z_heuristic=env_config['z_heuristic'])
    envs.State = State
    envs.Action = Action
    envs.Reward = Reward
    envs.IsRobotReady = IsRobotReady
    envs.SENTINEL = SENTINEL
    env_config['render'] = True

    # setup agent
    agent = createAgent()
    agent.train()
    if load_model_pre:
        agent.loadModel(load_model_pre)

    # setup the parallel agent
    p_agent = AgentWrapper()
    p_agent.eps = init_eps
    replay_buffer = QLearningBufferExpert(buffer_size)

    # logging
    simulator_str = copy.copy(simulator)
    if simulator == 'pybullet':
        simulator_str += ('_' + robot)
    log_dir = os.path.join(log_pre, '{}_{}_{}_{}_{}'.format(alg, model, simulator_str, num_objects, max_episode_steps))
    if note:
        log_dir += '_'
        log_dir += note
    logger = Logger(log_dir, env, 'train', num_processes, max_episode, log_sub)
    logger.saveParameters(hyper_parameters)
    if load_sub:
        logger.loadCheckPoint(os.path.join(log_dir, load_sub, 'checkpoint'), env, agent, replay_buffer)

    # starting parallel training
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        executor.submit(p_agent.get_action)
        executor.submit(p_agent.store_transition_SGD)
        executor.submit(main)


text:
import numpy as np
import pickle as pk
import multiprocessing
import time

class data():
    def __init__(self, batch_size, batch_length, pc_size):
        self.batch_size=batch_size
        self.batch_length=batch_length
        self.pc_size=pc_size

        self.pquat_sample_list=[0,1,2,4,5,16,17,18,19]
        self.trans=np.load('./trans.mcl.20p8a.dat').reshape((20, 8, 3000, 3))
        self.pquat=np.load('./pmat.mcl.20p8a.dat')[...,self.pquat_sample_list,:,:].reshape((20, 8, 3000, 9, 3, 3))
        self.betas=np.load('./betas.mcl.20p8a.dat') #(20,10)
        self.gender=np.load('./gender.mcl.20p8a.dat') #(20,1)
        self.pc=pk.load(open('./pc.mcl.20p8a.dat', 'rb'))

        print('load trans:', self.trans.shape)
        print('load pquat:', self.pquat.shape)
        print('load betas:', self.betas.shape)
        print('load gender:', self.gender.shape)
        print('load pc:', len(self.pc), len(self.pc[0]), len(self.pc[0][0]))

        assert self.trans.shape[0]==self.pquat.shape[0]
        assert self.trans.shape[1]==self.pquat.shape[1]
        assert self.trans.shape[2]==self.pquat.shape[2]
        assert self.trans.shape[0]==self.betas.shape[0]
        assert self.trans.shape[0]==self.gender.shape[0]
        assert self.trans.shape[0]==len(self.pc)
        assert self.trans.shape[1]==len(self.pc[0])
        assert self.trans.shape[2]==len(self.pc[0][0])

        self.name_size=self.trans.shape[0] #8
        self.act_size=self.trans.shape[1] #7
        self.total_length_size=self.trans.shape[2] #3000
        self.train_length_size=2400
        self.test_length_size=self.total_length_size-self.train_length_size

        # train
        self.train_trans=self.trans[:,:,:self.train_length_size]
        self.train_pquat=self.pquat[:,:,:self.train_length_size]
        self.train_pc=[[x[:self.train_length_size] for x in y] for y in self.pc]

        # test
        self.test_trans=self.trans[:,:,self.train_length_size:]
        self.test_pquat=self.pquat[:,:,self.train_length_size:]
        self.test_pc=[[x[self.train_length_size:] for x in y] for y in self.pc]

        # multi-processing
        self.next_buffer_size=128
        self.test_buffer_size=2
        self.next_extra_process_size=16

        self.q_next_batch=multiprocessing.Manager().Queue()
        self.q_test=multiprocessing.Manager().Queue()
        self.v_flag=multiprocessing.Manager().Value('b', True)

        self.task_test=multiprocessing.Process(target=self.maintain_test_pc)
        self.task_next=multiprocessing.Process(target=self.maintain_next_batch, args=(True,))
        self.task_list=[]
        for i in range(self.next_extra_process_size):
            self.task_list.append(multiprocessing.Process(target=self.maintain_next_batch, args=(False,)))

        self.task_next.start()
        self.task_test.start()
        for i in range(self.next_extra_process_size):
            self.task_list[i].start()

    def close(self):
        self.v_flag.value=False
        print('set to close', self.v_flag)
        time.sleep(10)
        while not self.q_next_batch.empty():
            self.q_next_batch.get()
        while not self.q_test.empty():
            self.q_test.get()
        self.task_next.join()
        self.task_test.join()
        for i in range(self.next_extra_process_size):
            self.task_list[i].join()

    def maintain_test_pc(self):
        while self.v_flag.value:
            if self.q_test.qsize()<self.test_buffer_size:
                self.q_test.put(self.prepare_test_pc())
            else:
                time.sleep(1)
        print('test pc closed')

    def maintain_next_batch(self, print_flag):
        while self.v_flag.value:
            if self.q_next_batch.qsize()<self.next_buffer_size:
                self.q_next_batch.put(self.prepare_next_batch())
            else:
                time.sleep(0.2)
        print('next batch closed')

    def prepare_test_pc(self):
        pc_test_list=[]
        for name_no in range(self.name_size):
            pc_name_list=[]
            for act_no in range(self.act_size):
                pc_tmp_list=[]
                for j in range(self.test_length_size):
                    pc_tmp=np.zeros((self.pc_size, 6), dtype=np.float32)
                    pc_no=self.test_pc[name_no][act_no][j].shape[0]
                    if pc_no<self.pc_size:
                        fill_list=np.random.choice(self.pc_size, size=pc_no, replace=False)
                        fill_set=set(fill_list)
                        pc_tmp[fill_list]=self.test_pc[name_no][act_no][j]
                        dupl_list=[x for x in range(self.pc_size) if x not in fill_set]
                        dupl_pc=np.random.choice(pc_no, size=len(dupl_list), replace=True)
                        pc_tmp[dupl_list]=self.test_pc[name_no][act_no][j][dupl_pc]
                    else:
                        pc_list=np.random.choice(pc_no, size=self.pc_size, replace=False)
                        pc_tmp=self.test_pc[name_no][act_no][j][pc_list]
                    pc_tmp_list.append(pc_tmp)
                pc_name_list.append(pc_tmp_list)
            pc_test_list.append(pc_name_list)
        return np.asarray(pc_test_list)

    def prepare_next_batch(self):
        batch_pc_list=[]
        batch_trans_list=[]
        batch_pquat_list=[]
        batch_betas_list=[]
        batch_gender_list=[]
        length_list=np.random.choice(self.train_length_size-self.batch_length+1, size=self.batch_size, replace=True)
        name_list=np.random.choice(self.name_size, size=self.batch_size, replace=True)
        act_list=np.random.choice(self.act_size, size=self.batch_size, replace=True)
        for i in range(self.batch_size):
            length_begin=length_list[i]
            name_choice=name_list[i]
            act_choice=act_list[i]
            batch_trans_list.append(self.train_trans[name_choice, act_choice, length_begin: length_begin+self.batch_length])
            batch_pquat_list.append(self.train_pquat[name_choice, act_choice, length_begin: length_begin+self.batch_length])
            batch_betas_list.append(np.tile([self.betas[name_choice]], (self.batch_length, 1)))
            batch_gender_list.append(np.tile([self.gender[name_choice]], (self.batch_length, 1)))
            pc_tmp_list=[]
            for j in range(length_begin, length_begin+self.batch_length):
                pc_tmp=np.zeros((self.pc_size, 6), dtype=np.float32)
                pc_no=self.train_pc[name_choice][act_choice][j].shape[0]
                if pc_no<self.pc_size:
                    fill_list=np.random.choice(self.pc_size, size=pc_no, replace=False)
                    fill_set=set(fill_list)
                    pc_tmp[fill_list]=self.train_pc[name_choice][act_choice][j]
                    dupl_list=[x for x in range(self.pc_size) if x not in fill_set]
                    dupl_pc=np.random.choice(pc_no, size=len(dupl_list), replace=True)
                    pc_tmp[dupl_list]=self.train_pc[name_choice][act_choice][j][dupl_pc]
                else:
                    pc_list=np.random.choice(pc_no, size=self.pc_size, replace=False)
                    pc_tmp=self.train_pc[name_choice][act_choice][j][pc_list]
                pc_tmp_list.append(pc_tmp)
            batch_pc_list.append(pc_tmp_list)
        return np.asarray(batch_pc_list), np.asarray(batch_pquat_list), np.asarray(batch_trans_list), np.asarray(batch_betas_list), np.asarray(batch_gender_list)

    def get_test_pc(self):
        while self.q_test.empty():
            time.sleep(0.1)
        return self.q_test.get()

    def next_batch(self):
        while self.q_next_batch.empty():
            time.sleep(0.1)
        return self.q_next_batch.get()

if __name__=='__main__':
    d=data(3, 11, 180)
    print('total:', len(d.pc), len(d.pc[0]), len(d.pc[0][0]), d.pquat.shape, d.trans.shape, d.betas.shape, d.gender.shape)
    print('train:', len(d.train_pc), len(d.train_pc[0]), len(d.train_pc[0][0]), d.train_pquat.shape, d.train_trans.shape)
    print('test:', len(d.test_pc), len(d.test_pc[0]), len(d.test_pc[0][0]), d.test_pquat.shape, d.test_trans.shape)
    p, q, t, b, g = d.next_batch()
    print('Next batch:', p.shape, q.shape, t.shape, b.shape, g.shape)
    p=d.get_test_pc()
    print('Test pc:', p.shape)
    time.sleep(10)
    d.close()

max_stars_repo_name: HavocFiXer/mmMesh

max_issues_repo_name: HavocFiXer/mmMesh

max_forks_repo_name: HavocFiXer/mmMesh