text:
""" LELSD PyTorch model.
    From "Optimizing Latent Space Directions For GAN-based Local Image Editing"
    By Ehsan Pajouheshgar, Tong Zhang, and Sabine Susstrunk
    https://arxiv.org/abs/2111.12583
"""

import json
import os
from datetime import datetime

import numpy as np
import torch
import torchvision
from torch.autograd import Variable
from tqdm import tqdm

from utils.segmentation_utils import GANLinearSegmentation

try:
    from tensorboardX import SummaryWriter
except:
    from torch.utils.tensorboard import SummaryWriter

class LELSD:
    stylegan2_y_dims = {
        0: 512,  # 4 x 4
        1: 512,  # 8 x 8
        2: 512,  # 8 x 8
        3: 512,  # 16 x 16
        4: 512,  # 16 x 16
        5: 512,  # 32 x 32
        6: 512,  # 32 x 32
        7: 512,  # 64 x 64
        8: 512,  # 64 x 64
        9: 512,  # 128 x 128
        10: 256,  # 128 x 128
        11: 256,  # 256 x 256
        12: 128,  # 256 x 256
        13: 128,  # 512 x 512
        14: 64,  # 512 x 512
        15: 64,  # 1024 x 1024
        16: 32,  # 1024 x 1024
        17: 32,  # Unused
    }

    stylegan2_rgb_y_dims = {
        0: 512,  # 4 x 4
        1: 512,  # 8 x 8
        2: 512,  # 16 x 16
        3: 512,  # 32 x 32
        4: 512,  # 64 x 64
        5: 256,  # 128 x 128
        6: 128,  # 256 x 256
        7: 64,  # 512 x 512
        8: 32,  # 1024 x 1024
    }

    stylegan2_s_dims = {
        0: 512,  # 4 x 4 s1
        1: 512,  # 4 x 4 trgb
        2: 512,  # 8 x 8 s1
        3: 512,  # 8 x 8 s2
        4: 512,  # 8 x 8 trgb
        5: 512,  # 16 x 16 s1
        6: 512,  # 16 x 16 s2
        7: 512,  # 16 x 16 trgb
        8: 512,  # 32 x 32 s1
        9: 512,  # 32 x 32 s2
        10: 512,  # 32 x 32 trgb
        11: 512,  # 64 x 64 s1
        12: 512,  # 64 x 64 s2
        13: 512,  # 64 x 64 trgb
        14: 512,  # 128 x 128 s1
        15: 256,  # 128 x 128 s2
        16: 256,  # 128 x 128 trgb
        17: 256,  # 256 x 256 s1
        18: 128,  # 256 x 256 s2
        19: 128,  # 256 x 256 trgb
        20: 128,  # 512 x 512 s1
        21: 64,  # 512 x 512 s2
        22: 64,  # 512 x 512 trgb
        23: 64,  # 1024 x 1024 s1
        24: 32,  # 1024 x 1024 s2
        25: 32,  # 1024 x 1024 trgb
        26: 32  # 1024 x 1024 unused
    }

    """
    This class tries to find directions in the latent space of a pretrained GAN
    that if added to the initial latent vector, the change in the output image
    will be localized to a specific semantic part of the object.
    For example you can use this model to find latent directions that only change
    the mouth, or the nose of the generated face.
    For fitting this model, a pretrained GAN, and a pretrained segmentation model are required.
    Parameters
    ----------
    device: pytorch device
        Device used for performing the computation.
    localization_layers: list, required
        The layers that will be used in the Localization score.
        The Localozation Score tries to maximize the ratio of the
        changes in the desired part to changes outside
        the desired part. The mask determining the desired part will
        be bilinearly [up-down]-sampled to have the same shape as layer feature maps.
        Note that the layer indices should be ascendingly sorted.
    semantic_parts: list, required
        List of strings. This list will determine the semantic parts that determine the localized area
        of the output image we want to modify and edit. For fitting the model a segmentation object is
        required with an dict attribute named 'part_to_mask_idx' that will determine the mask idx corresponding
        to each semantic part in this list.
    loss_function: {'L1', 'L2', 'cos'}, default='L2'
        The loss function used to evaluate and optimize localization score
    mode: string, default='foreground'
        Either 'foreground' or 'background'
        Wether to treat the semantic parts as foreground or background.
    mask_aggregation: {'union', 'intersection', 'average'}, default='average'
        The method to combine two masks
    n_layers: int, default=18
        Number of layers in the generator network.
    num_latent_directions: int, default=3
        Number of distinct direction that will be find. Note that the model will try to
        find perpendicular directions by minimizing the correlation between the pair of directions.
    latent_dim: int, default=512
        Dimension of the latent space of the GAN
    batch_size: int, default=4
        Batch size used for fitting the model
    learning_rate: float, default=0.001,
        The learning rate used for optimizing the latent directions
    localization_layer_weights: list, default=None
        Weight of each layer in the Localization loss term.
        If not specified, the layers will have equal weights.
    gamma_correlation: float, default=1.0
        Coefficient of correlation loss
    unit_norm: bool, default=True,
        Setting this to True will force the optimization to normalize the latent
        directions to unit norm after each step of gradient descent.
        If this is set to false then (-1.0, 1.0) will be used as the alpha range
        with min_abs_alpha_value equal to 0.2
    latent_space: {'Z', 'W', 'W+', S*'}, Type of latent space used for extracting semantics.
        Z: Normally distributed latent codes
        W: StyleGAN intermediate Latent Space
        W+: StyleGAN Intermediate Latent Space separated for each layer
        S: StyleSpace defined in paper 'https://arxiv.org/abs/2011.12799'
            Examples
            S10: 10layer according to the paper
            S+: all layers
            S_conv: all conv layers (trgb styles not included)
    onehot_temperature: float, default=0.001
        Temperature of the RelaxedOneHotCategorical distribution used for generating
        weight of each latent direction for the linear interpolation.
        If this temperature is very small (close to zero) then  the coefficients in the
        linear interpolation will be almost onehot and if the temperature is to high, all
        of the coefficients will be equal to (1 / num_latent_dirs)
    min_alpha_value: float, default=-3.0
        The minimum alpha value used for adding the latent direction.
    min_alpha_value: float, default=-3.0
        The maximum alpha value used for adding the latent direction.
        Alpha will be sampled randomly in this range.
    min_abs_alpha_value: float, default=0.5
        If abs(alpha) for the sampled alpha is smaller than this value,
        alpha will be resampled. Notice that small values of alpha does not
        change the semantic noticeably and may cause unstable optimization.
    log_dir: str, default=log
        Directory used for saving the model, and optimization summaries
    random_seed: int, default=1234
        Random seed used for initializing the latent directions
    """

    def __init__(self, device, localization_layers, semantic_parts, loss_function="MSE",
                 mode='foreground', mask_aggregation='average', n_layers=18, num_latent_dirs=3, latent_dim=512,
                 batch_size=4,
                 learning_rate=0.001, localization_layer_weights=None, gamma_correlation=1.0,
                 unit_norm=True, latent_space='W',
                 onehot_temperature=0.001,
                 min_alpha_value=-3.0, max_alpha_value=3.0, min_abs_alpha_value=0.5,
                 log_dir='log', random_seed=1234):
        assert mask_aggregation in ['average', 'union', 'intersection']
        assert loss_function in ['L1', 'L2', 'cos']
        self.device = device
        self.localization_layers = localization_layers
        self.semantic_parts = semantic_parts
        self.loss_function = loss_function
        self.mode = mode
        self.mask_aggregation = mask_aggregation
        self.n_layers = n_layers
        self.num_latent_dirs = num_latent_dirs
        self.latent_dim = latent_dim
        self.batch_size = batch_size
        self.lr = learning_rate
        if localization_layer_weights is None:
            self.localization_layer_weights = np.array([1.0] * len(self.localization_layers))
        else:
            self.localization_layer_weights = np.array(localization_layer_weights)

        self.localization_layer_weights /= self.localization_layer_weights.sum()
        self.localization_layer_weights = list(self.localization_layer_weights)
        self.gamma_correlation = gamma_correlation

        self.unit_norm = unit_norm
        self.latent_space = latent_space

        self.min_alpha_value = min_alpha_value
        self.max_alpha_value = max_alpha_value
        self.min_abs_alpha_value = min_abs_alpha_value

        self.onehot_temperature = onehot_temperature
        assert self.onehot_temperature >= 0.001
        self.coefficient_sampling_dist = torch.distributions.relaxed_categorical.RelaxedOneHotCategorical(
            temperature=torch.tensor([self.onehot_temperature]),
            probs=torch.tensor([1 / self.num_latent_dirs] * self.num_latent_dirs))

        self.init_time_signature = datetime.now().strftime('%b%d_%H-%M-%S')
        self.log_dir = os.path.join(log_dir, self.init_time_signature)
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)
        self.random_seed = random_seed

        torch.manual_seed(self.random_seed)
        np.random.seed(self.random_seed)

        if self.latent_space == 'Z':
            self.latent_dirs = torch.randn(self.num_latent_dirs, self.latent_dim, device=self.device)
            self.latent_dirs = self.latent_dirs / torch.linalg.norm(self.latent_dirs, dim=1, keepdim=True)
        if self.latent_space == 'W':
            self.latent_dirs = torch.randn(self.num_latent_dirs, self.latent_dim, device=self.device)
            self.latent_dirs = self.latent_dirs / torch.linalg.norm(self.latent_dirs, dim=1, keepdim=True)
        if self.latent_space == 'W+':
            self.latent_dirs = torch.randn(self.num_latent_dirs, self.n_layers, self.latent_dim, device=self.device)
            self.latent_dirs = self.latent_dirs / torch.linalg.norm(self.latent_dirs, dim=2, keepdim=True)
        if self.latent_space.startswith("S"):
            # This feature is only implemented for StyleGAN2
            self.latent_dirs = []
            for layer in self.stylegan2_s_dims:
                dim = self.stylegan2_s_dims[layer]
                w = torch.randn(self.num_latent_dirs, dim, device=self.device)
                w = w / torch.linalg.norm(w, dim=1, keepdim=True)
                self.latent_dirs.append(w)
            if self.latent_space == 'S+':
                self.s_layers_to_apply = None
            elif self.latent_space == 'S_conv':
                self.s_layers_to_apply = [0, 2, 3, 5, 6, 8, 9, 11, 12, 14, 15, 17, 18, 20, 21, 23, 24, 26]
            else:
                self.s_layers_to_apply = [int(self.latent_space[1:])]

        if not self.latent_space.startswith("S"):
            self.latent_dirs = Variable(self.latent_dirs, requires_grad=True)
        else:
            #             pass
            self.latent_dirs = [Variable(w, requires_grad=True) for w in self.latent_dirs]

    def sample_alpha(self, min_value, max_value, min_abs_value=0.5):
        L = max_value - min_value
        alpha = L * np.random.rand() + min_value
        if abs(alpha) < min_abs_value:
            return self.sample_alpha(min_value, max_value, min_abs_value)
        return alpha

    def _move_latent_codes(self, latent_codes, latent_dir, alpha, layers_to_apply=None):
        if layers_to_apply:
            # if hasattr(self, "one_for_each_layer") and self.one_for_each_layer:
            if self.latent_space.startswith("S"):
                new_latent_codes = [x.clone() for x in latent_codes]
                for l in layers_to_apply:
                    if l >= len(new_latent_codes):
                        continue
                    new_latent_codes[l] += latent_dir[l] * alpha
            else:
                new_latent_codes = latent_codes.clone()
                if self.latent_space == 'W+':
                    new_latent_codes[:, layers_to_apply, :] += latent_dir[:, layers_to_apply, :] * alpha
                elif self.latent_space == 'W':
                    new_latent_codes[:, layers_to_apply, :] += latent_dir.unsqueeze(0) * alpha
                else:
                    # Z cannot be separately applied to different layers
                    new_latent_codes += latent_codes * alpha
        else:
            if self.latent_space.startswith("S"):
                new_latent_codes = [x.clone() for x in latent_codes]
                for l in range(len(new_latent_codes)):
                    new_latent_codes[l] += latent_dir[l] * alpha
            else:
                new_latent_codes = latent_codes.clone()
                if self.latent_space == 'W+':
                    new_latent_codes += latent_dir * alpha
                elif self.latent_space == 'W':
                    new_latent_codes += latent_dir.unsqueeze(0) * alpha
                else:
                    new_latent_codes += latent_dir * alpha

        return new_latent_codes

    def move_latent_codes(self, latent_codes, latent_dir_idx, alpha, layers_to_apply=None):
        """
        Moves the latent vector ws in the direction of self.latent_dirs[latent_dir_idx] by amount of alpha

        Parameters
        ----------
        latent_codes: If latent_space is 'Z' then Tensor with shape [batch_size, latent_dim]
            If latent_space is 'W' or 'W+ then Tensor with shape [batch_size, num_layers, latent_dim]
            If latent_space is 'S' then tuple of (ys, rgb_ys)
                where ys and rgb_ys are list of tensors of shape [batch_size, latent_dim]

        latent_dir_idx: int, Index of the latent direction
        alpha: float, Determines the magnitude of the vector that will be added to ws
        layers_to_apply: list, default=None
            list of layers that will use the updated w.
            If None then all of the layers will be updated
        """
        if not self.latent_space.startswith("S"):
            latent_dir = self.latent_dirs[latent_dir_idx: latent_dir_idx + 1]
        else:
            latent_dir = [x[latent_dir_idx: latent_dir_idx + 1] for x in self.latent_dirs]

        return self._move_latent_codes(latent_codes, latent_dir, alpha, layers_to_apply)

    @staticmethod
    def correlation_loss(ws):
        """
        Parameters
        ----------
        ws: Tensor of shape [num_ws, latent_dim]
            The rows of ws should have unit length.
        """
        covariance_matrix = torch.matmul(ws, ws.T)
        diagonal = torch.diag(covariance_matrix)
        normalizer = torch.sqrt(diagonal[:, None])
        normalizer = normalizer.T * normalizer
        correlation_matrix = covariance_matrix / normalizer
        non_diagonal_elements = correlation_matrix - torch.eye(len(ws), device=ws.device)
        return torch.mean(torch.abs(non_diagonal_elements))

    @staticmethod
    def combine_mask(mask1, mask2, method='average'):
        assert method in ['average', 'union', 'intersection']
        if method == 'average':
            return 0.5 * mask1 + 0.5 * mask2
        elif method == 'intersection':
            return mask1 * mask2
        else:
            return mask1 + mask2 - mask1 * mask2

    def fit(self, gan_sample_generator, segmentation_model, num_batches, num_lr_halvings=4, batch_size=None,
            pgbar=False, summary=True, snapshot_interval=200):
        """
        This method will optimize the latent directions using samples generated by gan_sample_generator
        that will be segmented by segmentation_model.
        Parameters
        ----------
        gan_sample_generator: SampleGenerator,
            This is a wrapper class on top of the pretrained generator network.
        segmentation_model: SegmentationModel
            This is a wrapper class on top of the pretrained segmentation network.
            This class should have an attribute named 'part_to_mask_idx' that will
            be used to translate each semantic part to the corresponding mask idx.
        num_batches: int
            Number of batches used for optimizing the latent directions
        num_lr_halvings: int, default=4
            Number of times that the learning rate will be halved.
        batch_size: int, default=None
            The batch size used for the optimization.
            if not specified the class attribute will be used.
        pgbar: bool, default=False
            Whether to show the progress bar or not
        summary: bool, default=True
            Whether to write summaries or not
        snapshot_interval: int, default=200
            The interval between saving image snapshots.
        """
        if batch_size is None:
            batch_size = self.batch_size
        assert hasattr(segmentation_model, 'part_to_mask_idx')
        if summary:
            with open(os.path.join(self.log_dir, "config.json"), "w") as f:
                config_dict = self.__dict__.copy()
                del config_dict['device']
                del config_dict['latent_dirs']
                del config_dict['coefficient_sampling_dist']
                json.dump(config_dict, f, indent=4, sort_keys=True, default=lambda x: x.__name__)

            try:
                summary_writer = SummaryWriter(logdir=self.log_dir)
            except:
                summary_writer = SummaryWriter(self.log_dir)

        part_ids = [segmentation_model.part_to_mask_idx[part_name] for part_name in self.semantic_parts]
        lr = self.lr
        lr_decay_steps = num_batches // (num_lr_halvings + 1)
        min_alpha_value = self.min_alpha_value
        max_alpha_value = self.max_alpha_value
        global_step = 0
        pbar = tqdm(range(54321, 54321 + num_batches, 1), disable=not pgbar, total=num_batches)
        pil_to_tensor = torchvision.transforms.ToTensor()

        if self.latent_space.startswith("S"):
            optimizer = torch.optim.Adam(self.latent_dirs, lr=lr)
        else:
            optimizer = torch.optim.Adam([self.latent_dirs], lr=lr)
        for i, seed in enumerate(pbar):
            if i % lr_decay_steps == (lr_decay_steps - 1):
                lr /= 2
                optimizer.param_groups[0]['lr'] = lr

            batch_data = gan_sample_generator.generate_batch(batch_size=batch_size, seed=seed, requires_grad=True,
                                                             return_all_layers=True, return_image=True)
            if isinstance(segmentation_model, GANLinearSegmentation):
                old_segmentation_output = segmentation_model.predict(batch_data, one_hot=False)
            else:
                old_segmentation_output = segmentation_model.predict(batch_data['image'], one_hot=False)
            segmentation_output_res = old_segmentation_output.shape[2]

            old_mask = 0.0
            for part_idx in part_ids:
                old_mask += 1.0 * (old_segmentation_output == part_idx)

            # for w_idx in range(self.num_ws):
            gan_sample_generator.zero_grad()
            segmentation_model.zero_grad()
            alpha = self.sample_alpha(min_alpha_value, max_alpha_value, min_abs_value=self.min_abs_alpha_value)
            latent_coefs = self.coefficient_sampling_dist.sample().unsqueeze(1).to(self.device)  # [num_ws, 1]
            if not self.latent_space.startswith("S"):
                if self.latent_space.startswith("W"):
                    old_latent = batch_data['ws'].detach()
                    if self.latent_space == "W+":
                        latent_coefs = latent_coefs.unsqueeze(2)
                    latent_dir = torch.sum(self.latent_dirs * latent_coefs, dim=0, keepdim=True)
                    new_latent = self._move_latent_codes(old_latent, latent_dir, alpha)
                    new_batch_data = gan_sample_generator.generate_batch_from_ws(new_latent, requires_grad=True,
                                                                                return_all_layers=True,
                                                                                return_image=True)
                else:
                    old_latent = batch_data['z'].detach()
                    latent_dir = torch.sum(self.latent_dirs * latent_coefs, dim=0, keepdim=True)
                    new_latent = self._move_latent_codes(old_latent, latent_dir, alpha)
                    new_batch_data = gan_sample_generator.generate_batch_from_z(new_latent, requires_grad=True,
                                                                                return_all_layers=True,
                                                                                return_image=True)

            else:
                old_latent = [tmp.detach() for tmp in batch_data['s']]
                latent_dir = [torch.sum(u * latent_coefs, dim=0, keepdim=True) for u in self.latent_dirs]
                new_latent = self._move_latent_codes(old_latent, latent_dir, alpha,
                                                     layers_to_apply=self.s_layers_to_apply)
                ys_tuple = gan_sample_generator.s_to_ys(new_latent)
                new_batch_data = gan_sample_generator.generate_batch_from_ys(ys_tuple, requires_grad=True,
                                                                             return_all_layers=True,
                                                                             return_image=True)

            if isinstance(segmentation_model, GANLinearSegmentation):
                new_segmentation_output = segmentation_model.predict(new_batch_data, one_hot=False)
            else:
                new_segmentation_output = segmentation_model.predict(new_batch_data['image'], one_hot=False)
            new_mask = 0.0
            for part_idx in part_ids:
                new_mask += 1.0 * (new_segmentation_output == part_idx)

            combined_mask = self.combine_mask(old_mask, new_mask, self.mask_aggregation)
            combined_mask = combined_mask.detach()
            # Cloning to avoid redundant computation
            mask = combined_mask.clone()

            last_layer_res = None
            localization_loss = 0
            # To maximize the Localization Score in localization layers
            for layer, layer_weight in zip(reversed(self.localization_layers),
                                           reversed(self.localization_layer_weights)):
                layer_res = gan_sample_generator.layer_to_resolution[layer]
                if last_layer_res != layer_res:
                    if layer_res != segmentation_output_res:
                        mask = torch.nn.functional.interpolate(mask, size=(layer_res, layer_res),
                                                               mode='bilinear',
                                                               align_corners=True)
                    else:
                        mask = combined_mask.clone()
                last_layer_res = layer_res
                if layer_weight == 0:
                    continue

                x1 = batch_data[f'layer_{layer}'].detach()
                x2 = new_batch_data[f'layer_{layer}']
                if self.loss_function == 'L1':
                    diff = torch.mean(torch.abs(x1 - x2), dim=1)
                elif self.loss_function == 'L2':
                    diff = torch.mean(torch.square(x1 - x2), dim=1)
                elif self.loss_function == 'cos':
                    diff = 1 - torch.nn.functional.cosine_similarity(x1, x2, dim=1, eps=1e-8)
                else:
                    diff = torch.mean(torch.square(x1 - x2), dim=1)
                indicator = mask[:, 0]
                if self.mode == 'background':
                    indicator = 1 - indicator

                localization_loss -= layer_weight * torch.sum(diff * indicator, dim=[1, 2]) / (
                        torch.sum(diff, dim=[1, 2]) + 1e-6)

            # -1.0 means perfect localization and 0 means poor localization
            localization_loss = torch.mean(localization_loss)
            if self.latent_space == "W+":
                correlation_loss = 0
                for layer in range(self.n_layers):
                    correlation_loss += self.correlation_loss(self.latent_dirs[:, layer, :])
                correlation_loss /= self.n_layers
            elif self.latent_space.startswith("S"):
                correlation_loss = 0
                if self.s_layers_to_apply is None:
                    for latent_dir in self.latent_dirs:
                        correlation_loss += self.correlation_loss(latent_dir)
                    correlation_loss /= len(self.latent_dirs)
                else:
                    for layer in self.s_layers_to_apply:
                        correlation_loss += self.correlation_loss(self.latent_dirs[layer])
                    correlation_loss /= len(self.s_layers_to_apply)
            else:
                correlation_loss = self.correlation_loss(self.latent_dirs)
            loss = localization_loss + self.gamma_correlation * correlation_loss
            loss.backward()
            if self.unit_norm:
                optimizer.step()
                self.latent_dirs.data = self.latent_dirs.data / torch.linalg.norm(self.latent_dirs.data, dim=1,
                                                                                keepdim=True)
                optimizer.zero_grad()
            else:
                pass
                optimizer.step()
                optimizer.zero_grad()

            if summary:
                if global_step % snapshot_interval == 0:
                    old_images = [pil_to_tensor(img.resize((256, 256))) for img in batch_data['image'][:4]]
                    old_mask_images = [
                        torch.nn.functional.interpolate(old_mask[b: b + 1], size=(256, 256), mode='bilinear',
                                                        align_corners=True)[0].repeat(3, 1, 1).detach().cpu() for b in
                        range(min(4, batch_size))]
                    new_images = [pil_to_tensor(img.resize((256, 256))) for img in new_batch_data['image'][:4]]
                    new_mask_images = [
                        torch.nn.functional.interpolate(new_mask[b: b + 1], size=(256, 256), mode='bilinear',
                                                        align_corners=True)[0].repeat(3, 1, 1).detach().cpu() for b in
                        range(min(4, batch_size))]
                    images = old_images + old_mask_images + new_images + new_mask_images
                    image_grid = torchvision.utils.make_grid(images, nrow=4)
                    torchvision.utils.save_image(image_grid,
                                                 os.path.join(self.log_dir, f"batch={i}-image_alpha={alpha:.2f}.jpg"))

                summary_writer.add_scalar("localization score", -localization_loss, global_step)
                summary_writer.add_scalar("correlation_loss", correlation_loss, global_step)
                summary_writer.add_scalar("loss", loss, global_step)
                summary_writer.add_scalar("alpha", alpha, global_step)
                summary_writer.add_scalar("lr", lr, global_step)
                global_step += 1

        if summary:
            summary_writer.close()
        if self.unit_norm:
            if not self.latent_space.startswith("S"):
                self.latent_dirs.data = self.latent_dirs.data / torch.linalg.norm(self.latent_dirs.data, dim=1,
                                                                                keepdim=True)
            else:
                for l in range(len(self.latent_dirs)):
                    self.latent_dirs[l].data = self.latent_dirs[l].data / torch.linalg.norm(self.latent_dirs[l].data,
                                                                                dim=1, keepdim=True)

    def edit_batch_data(self, gan_sample_generator, batch_data, latent_dir_idx, alpha, layers_to_apply=None):
        if not self.latent_space.startswith("S"):
            if self.latent_space.startswith("W"):
                old_latent = batch_data['ws'].detach()
                new_latent = self.move_latent_codes(old_latent, latent_dir_idx, alpha, layers_to_apply)
                new_batch_data = gan_sample_generator.generate_batch_from_ws(new_latent, return_style=True,
                                                                             return_image=True, return_all_layers=True)
                new_batch_data['ws'] = new_latent
            else:
                old_latent = batch_data['z'].detach()
                new_latent = self.move_latent_codes(old_latent, latent_dir_idx, alpha)
                new_batch_data = gan_sample_generator.generate_batch_from_z(new_latent, return_style=True,
                                                                            return_image=True, return_all_layers=True)
                new_batch_data['z'] = new_latent

        else:
            old_latent = [tmp.detach() for tmp in batch_data['s']]
            new_latent = self.move_latent_codes(old_latent, latent_dir_idx, alpha, layers_to_apply)
            ys_tuple = gan_sample_generator.s_to_ys(new_latent)
            new_batch_data = gan_sample_generator.generate_batch_from_ys(ys_tuple, return_image=True,
                                                                         return_all_layers=True)
            new_batch_data['s'] = new_latent

        return new_batch_data

    def save(self):
        torch.save(self, os.path.join(self.log_dir, "model.pth"))

    @staticmethod
    def load(model_path):
        return torch.load(model_path)

    def randomize_latent_dirs(self):
        if self.latent_space == 'Z':
            self.latent_dirs = torch.randn(self.num_latent_dirs, self.latent_dim, device=self.device)
            self.latent_dirs = self.latent_dirs / torch.linalg.norm(self.latent_dirs, dim=1, keepdim=True)
        if self.latent_space == 'W':
            self.latent_dirs = torch.randn(self.num_latent_dirs, self.latent_dim, device=self.device)
            self.latent_dirs = self.latent_dirs / torch.linalg.norm(self.latent_dirs, dim=1, keepdim=True)
        if self.latent_space == 'W+':
            self.latent_dirs = torch.randn(self.num_latent_dirs, self.n_layers, self.latent_dim, device=self.device)
            self.latent_dirs = self.latent_dirs / torch.linalg.norm(self.latent_dirs, dim=2, keepdim=True)
        if self.latent_space.startswith("S"):
            # This feature is only implemented for StyleGAN2
            self.latent_dirs = []
            for layer in self.stylegan2_s_dims:
                dim = self.stylegan2_s_dims[layer]
                w = torch.randn(self.num_latent_dirs, dim, device=self.device)
                w = w / torch.linalg.norm(w, dim=1, keepdim=True)
                self.latent_dirs.append(w)
            if self.latent_space == 'S+':
                self.s_layers_to_apply = None
            elif self.latent_space == 'S_conv':
                self.s_layers_to_apply = [0, 2, 3, 5, 6, 8, 9, 11, 12, 14, 15, 17, 18, 20, 21, 23, 24, 26]
            else:
                self.s_layers_to_apply = [int(self.latent_space[1:])]

        if not self.latent_space.startswith("S"):
            self.latent_dirs = Variable(self.latent_dirs, requires_grad=True)
        else:
            #             pass
            self.latent_dirs = [Variable(w, requires_grad=True) for w in self.latent_dirs]

        return self

    def load_from_InterFaceGAN(self, npy_path):
        if "_w_" in npy_path:
            latent_space = "W"
        else:
            latent_space = "Z"
        w = np.load(npy_path)
        self.latent_space = latent_space
        self.latent_dirs = torch.from_numpy(w).to(self.device)

    def load_from_StyleSpace(self, layer_idx, channel_idx):
        self.latent_space = f"S{layer_idx}"
        self.s_layers_to_apply = [layer_idx]
        self.latent_dirs = []
        for layer in self.stylegan2_s_dims:
            dim = self.stylegan2_s_dims[layer]
            w = torch.randn(1, dim, device=self.device) * 0.0
            if layer == layer_idx:
                w[:, channel_idx] = 1.0
            self.latent_dirs.append(w)

max_stars_repo_name: albertaillet/LELSD

max_issues_repo_name: albertaillet/LELSD

max_forks_repo_name: albertaillet/LELSD


text:
import pygad as pg
import pygad.plotting
import matplotlib.pyplot as plt
import matplotlib.patheffects as path_effects
import numpy as np
import utils
import glob
from scipy import stats
from multiprocessing import Pool

filename = __file__

def plot(args):
    halo = args[0]
    definition = args[1]

    path = '/ptmp/mpa/naab/REFINED/%s/SF_X/4x-2phase/out/snap_%s_4x_???' % (halo, halo)
    max = int(sorted(glob.glob(path))[-1][-3:])
    s, h, g = pg.prepare_zoom('/ptmp/mpa/naab/REFINED/%s/SF_X/4x-2phase/out/snap_%s_4x_%s' % (halo, halo, max), gas_trace='/u/mihac/data/%s/4x-2phase/gastrace_%s' % (halo, definition), star_form=None)
    s_m = s.gas[s.gas['num_recycled'] > -1]
    mask = s_m['mass_at_ejection'] > 0
    bins = np.linspace(0, s.cosmic_time(), 14*2)

    z_ej = s_m['metals_at_ejection'][mask] / s_m['mass_at_ejection'][mask]
    r_max = s_m['cycle_r_max'][mask]

    f, ax = plt.subplots(1, figsize=utils.figsize[::-1])
    _, _, _, cbar = pg.plotting.scatter_map(np.log10(r_max), np.log10(z_ej),
        qty=s_m['mass_at_ejection'][mask], logscale=True, clogscale=True,
        colors=s_m['T_at_ejection'][mask], colors_av=s_m['mass_at_ejection'][mask],
        extent=[[.5, 4], [-4, -.5]], bins=25, vlim=[10**6, 10**9.49], clim=[1e3, 10**7.9],
        ax=ax, zero_is_white=True)
    cbar_ax = cbar.ax
    cbar_ax.set_xlabel(r'$\log_{10}\left( T_\mathrm{ejection}\ \mathrm{[K]} \right)$', fontsize=utils.axes_labelsize*.8)
    plt.setp(cbar_ax.get_xticklabels(), fontsize=utils.tick_labelsize*.8)
    text = cbar_ax.xaxis.get_label().set_path_effects([path_effects.Stroke(linewidth=3, foreground='white'), path_effects.Normal()])
    cbar_ax.xaxis.set_label(text)
    ax.grid(False)
    ax.set_ylabel(r'$\log_{10}\left( Z_\mathrm{ejection} \right)$',
        fontsize=utils.axes_labelsize)
    ax.set_xlabel(r'$\log_{10}\left( r_\mathrm{max} \right)$',
        fontsize=utils.axes_labelsize)
    plt.setp(ax.get_xticklabels(), fontsize=utils.tick_labelsize)
    plt.setp(ax.get_yticklabels(), fontsize=utils.tick_labelsize)

    f.tight_layout()
    plt.subplots_adjust(top=0.93)
    f.suptitle('%s - %s' % (halo, definition), fontsize=32)

    plt.savefig(filename.split("/")[-1][:-3] + '_' + halo + '_' + definition + ".png")

p = Pool(7)
p.map(plot, utils.combinations)

ext: py

reason: import numpy,from scipy


text:
import typing
import sys
import numpy as np
import numba as nb

def solve(a, b, c) -> typing.NoReturn:
  cnt = 0
  q, b = divmod(b, 2)
  if c < q:
    cnt += c
    q -= c
    c = 0
    if a < 2 * q:
      cnt += a // 2
      return cnt
    cnt += q
    a -= q * 2
  else:
    cnt += q
    c -= q
  q, c = divmod(c, 2)
  if a < q:
    cnt += a
    return cnt
  cnt += q
  a -= q
  if a < c * 3:
    return cnt + a // 3
  cnt += c
  a -= c * 3
  q, a = divmod(a, 5)
  cnt += q
  return cnt

def main() -> typing.NoReturn:
  t = int(input())
  res = []
  for _ in range(t):
    a, b, c = map(int, input().split())
    res.append(solve(a, b, c))
  print(*res, sep='\n')

main()

max_stars_repo_name: kagemeka/competitive-programming

max_issues_repo_name: kagemeka/competitive-programming

max_forks_repo_name: kagemeka/competitive-programming

reason: import numpy,import numba


text:
import time
import numpy as np
from sklearn import linear_model, cross_validation, preprocessing, svm
from sklearn.grid_search import GridSearchCV

DESIGNS = [elm+str(num) for elm in ['a','b','c','l'] for num in range(0,5)]
morphed_DESIGNS = [elm+str(num) for elm in ['A','B','C','L'] for num in range(0,8)]
DESIGNS.extend(morphed_DESIGNS)

ATTRIBUTES = ['Active', 'Aggressive','Distinctive','Expressive','Innovative','Luxurious','Powerful','Sporty','Well Proportioned','Youthful']

features = np.zeros( [len(DESIGNS) ,len(ATTRIBUTES)])

for ind, attr in enumerate(ATTRIBUTES):
    features[:, ind] = np.loadtxt("../../data/processed_data/attribute_values/"+attr+"_full_rank.csv", delimiter=',')

baseline_features = features[0:20, :]

scaled_baseline_features = preprocessing.scale(baseline_features, axis=0)

num_data, num_features = np.shape(baseline_features)

a=[0]*5
b=[1]*5
c=[2]*5
l=[3]*5

a.extend(b)
a.extend(c)
a.extend(l)

brand_labels = np.array(a)

#---------------- Classifier Setup -------------------------------------------
clf_dict = { 0: 'Log Reg L1',
             1: 'Log Reg L2'}#,
#             3: 'SGD w/ ElasticNet',
#             4: 'Linear SVM'}
#             4: 'Decision Tree',
             #5: 'Random Forest',
             #6: 'Linear SVM'}

num_classifiers = len(clf_dict)
clf_array = np.empty((num_classifiers, ), dtype=object)
num_jobs = 1 # 12 if running on Foveros

#---------- L1 Logistic Regression -----------------------------------
lr_l1_tuned_parameters = [{'C': [0.01, 0.1, 1.0, 10, 100, 1000]}]
lr_l1_clf = GridSearchCV(estimator=linear_model.LogisticRegression(penalty='l1'),
                         param_grid=lr_l1_tuned_parameters, cv=3,
                         scoring='accuracy', refit=True, n_jobs=num_jobs)

#---------- L2 Logistic Regression -----------------------------------
lr_l2_tuned_parameters = [{'C': [0.01, 0.1, 1.0, 10, 100, 1000]}]
lr_l2_clf = GridSearchCV(estimator=linear_model.LogisticRegression(penalty='l2'),
                         param_grid=lr_l2_tuned_parameters, cv=3,
                         scoring='accuracy', refit=True, n_jobs=num_jobs)

huber_tuned_parameters = [{'alpha': [0.00001, 0.0001, 0.001, 0.01, 0.1, 1]}]
huber_clf = GridSearchCV(estimator=linear_model.SGDClassifier(penalty='l2', loss='huber'),
                       param_grid = huber_tuned_parameters, cv=3, scoring='accuracy',
                       refit=True, n_jobs=1)

#---------- Linear SVM -----------------------------------------------
svm_tuned_parameters = [{'C': [0.01, 0.1, 1.0, 10, 100, 1000]}]
svm_clf = GridSearchCV(estimator=svm.LinearSVC(loss='l2'), #dual=False
        param_grid = svm_tuned_parameters, cv=3, scoring='accuracy',
        refit=True, n_jobs=num_jobs)

clf_array[0] = lr_l1_clf
clf_array[1] = lr_l2_clf
#clf_array[2] = huber_clf
#clf_array[3] = svm_clf

num_experiments = 3
test_size_per_split = 0.2
seed = 0
#---------------- Define Error and Time Matrices -----------------------------
clf_error_matrix = np.empty((num_classifiers, num_experiments))
clf_time_matrix = np.empty((num_classifiers, num_experiments))

#---------------- Split Dataset into Train/ Test -------------------------
cv_iterator = cross_validation.ShuffleSplit(num_data,
        n_iter=num_experiments, test_size=test_size_per_split,
        random_state=seed)

#---------------- Begin Multiple Experiments for Each Classifier ---------
for exp_index, (train_index, test_index) in enumerate(cv_iterator):

    train_x = baseline_features[train_index, :]
    train_t = brand_labels[train_index]
    test_x = baseline_features[test_index, :]
    test_t = brand_labels[test_index]

    #---------- Train, Test, and Time ------------------------------------
    for clf_index, clf in enumerate(clf_array):
#            if reg_index ==4:
#                continue
        start_time = time.time()
        clf.fit(train_x, train_t)
        clf_time_matrix[clf_index, exp_index] = time.time() - start_time
        clf_error_matrix[clf_index, exp_index] = clf.score(test_x, test_t)

train_t

train_x.shape

#train_x

test_t

clf_error_matrix

clf_array[0].score(train_x,train_t)

np.mean(clf_error_matrix[0])

clf_array[0].score(test_x,test_t)

clf_time_matrix

omega=clf_array[0].best_estimator_.coef_

np.savetxt("omega_30percent_brand_recognition.csv", omega, delimiter=',')

omega

clf_array[0].best_estimator_.coef_

reason: import numpy


text:
"""
Copyright (C) 2015--2016 Jussi Leinonen

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to
do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
IN THE SOFTWARE.
"""

import numpy as np
from pyhdf import SD, HDF, VS

class EOSHDF(object):
    """
    Reader for HDF4 files, specifically those created by the NASA EOS data
    systems.

    Usage instructions:

    To open a file, initialize an EOSHDF object, passing a file name to the
    constructor. For example:
    import eoshdf
    eos = eoshdf.EOSHDF("example.hdf")

    To list the available datasets, use the list_datasets method:
    eos.list_datasets()

    To read a dataset, use the read_data method, passing the dataset name:
    data = eos.read_data("example_dataset")

    NOTE: No conversion factors or missing value masking are applied to read
    data. If you need these, you must apply them manually.
    """

    def __init__(self, file_name):
        self._file_name = file_name
        self._hdf = None
        self._vs = None
        self._sd = None

    def __del__(self):
        self._close_all()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._close_all()

    def _close_all(self):
        if self._vs is not None:
            self._vs.end()
            self._vs = None
            self._hdf.close()
            self._hdf = None
        if self._sd is not None:
            self._sd.end()
            self._sd = None

    def _open_vs(self):
        if self._vs is None:
            self._hdf = HDF.HDF(self._file_name)
            self._vs = self._hdf.vstart()
        return self._vs

    def _open_sd(self):
        if self._sd is None:
            self._sd = SD.SD(self._file_name)
        return self._sd

    def list_datasets(self):
        """
        Lists all available datasets. If you need to distinguish between
        the VS and SD interfaces, use the list_VS_datasets and
        list_SD_datasets methods.

        Returns:
            A list of the available datasets.

            This method might list some spurious datasets. In that case,
            refer to the data documentation to find which dataset to read.
        """

        vs_datasets = [v[0] for v in self.list_VS_datasets()]
        sd_datasets = list(self.list_SD_datasets().keys())
        return list(sorted(vs_datasets + sd_datasets))

    def read_data(self, ds_name, dtype=np.float64):
        """
        Reads a dataset. If you need to distinguish between
        the VS and SD interfaces, use the read_VS_data and
        read_SD_data methods.

        Args:
            ds_name: The dataset name.
            dtype: The datatype for the returned array.

        Returns:
            A numpy array containing the data.

            NOTE: No conversion factors or missing value masking are applied.
            If you need these, you must apply them manually.
        """
        try:
            data = self.read_SD_data(ds_name, dtype=dtype)
        except HDF.HDF4Error:
            try:
                data = self.read_VS_data(ds_name, dtype=dtype)
            except HDF.HDF4Error:
                raise IOError("Nonexistent data set.")
        return data

    def read_VS_data(self, ds_name, dtype=np.float64):
        vs = self._open_vs()
        vd = vs.attach(ds_name)
        arr = np.array(vd[:], dtype=dtype).ravel()
        vd.detach()
        return arr

    def list_VS_datasets(self):
        vs = self._open_vs()
        return vs.vdatainfo()

    def read_SD_data(self, ds_name, dtype=np.float64):
        sd = self._open_sd()
        sds = sd.select(ds_name)
        arr = np.array(sds[:], dtype=dtype)
        sds.endaccess()
        return arr

    def list_SD_datasets(self):
        sd = self._open_sd()
        return sd.datasets()

    def read_1D_data(self, var_name, dtype=np.float64):
        return self.read_VS_data(var_name, dtype=dtype)

    def read_2D_data(self, var_name, dtype=np.float64):
        return self.read_SD_data(var_name, dtype=dtype)

max_stars_repo_name: jleinonen/eoshdf

max_issues_repo_name: jleinonen/eoshdf

max_forks_repo_name: jleinonen/eoshdf


text:
import gym
import numpy as np
import torch as th
import time
import argparse

from sb3_contrib import DQNReg
from sb3_contrib import DQNClipped
from stable_baselines3.dqn import MlpPolicy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.callbacks import EvalCallback

# arguments
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--env', default='CartPole-v0')
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--total_timesteps', default=50000, type=int)
parser.add_argument('--loss_type', default=0, type=int)
parser.add_argument('--dqnreg_loss_weight', default=0.1, type=float)
args = parser.parse_args()
env_name = args.env
total_timesteps = args.total_timesteps
seed = args.seed
loss_type = args.loss_type
dqnreg_loss_weight = args.dqnreg_loss_weight

time_int = int(time.time())
model_save_name = "saved_models/dqnclippedreg_{}_{}_{}_{}".format(env_name, loss_type, seed, time_int)

env = gym.make(env_name)
env = Monitor(env)
policy_kwargs = dict(activation_fn=th.nn.ReLU, net_arch=[256, 256])
tensorboard_log = "tb_logs"

# hyperparams
buffer_size = max(total_timesteps // 100, 500)
learning_starts = max(total_timesteps // 1000, 100)
train_freq = 1
target_update_interval = 100
exploration_fraction = (learning_starts + 1000)/total_timesteps

# evaluations parameters
eval_env = gym.make(env_name)
eval_env = Monitor(eval_env)
eval_freq = max(1000, total_timesteps//20)
eval_log_path = "eval_logs/dqnclippedreg_{}_{}_{}_{}".format(env_name, loss_type, seed, time_int)
eval_callback = EvalCallback(eval_env, log_path=eval_log_path, eval_freq=eval_freq, deterministic=True, render=False, n_eval_episodes=25)

if env_name == 'MountainCar-v0':
    buffer_size = 10000  # max(total_timesteps // 100, 500)
    learning_starts = 1000  # max(total_timesteps // 1000, 100)
    learning_rate = 4e-3
    batch_size = 128
    gamma = 0.98
    train_freq = 16
    target_update_interval = 600
    gradient_steps = 8
    exploration_fraction = 0.2  # (learning_starts + 1000)/total_timesteps
    exploration_final_eps = 0.07

    if loss_type == 2:
        model = DQNClipped(MlpPolicy, env, policy_kwargs=policy_kwargs, target_update_interval=target_update_interval,
                    exploration_fraction=exploration_fraction,
                    buffer_size=buffer_size, train_freq=train_freq, learning_starts=learning_starts, seed=seed,
                    tensorboard_log=tensorboard_log, verbose=1,
                    batch_size=batch_size, learning_rate=learning_rate, gamma=gamma, gradient_steps=gradient_steps,
                    exploration_final_eps=exploration_final_eps)
    elif loss_type == 3:
        model = DQNReg(MlpPolicy, env, policy_kwargs=policy_kwargs, target_update_interval=target_update_interval,
                          exploration_fraction=exploration_fraction,
                          buffer_size=buffer_size, train_freq=train_freq, learning_starts=learning_starts, seed=seed,
                          tensorboard_log=tensorboard_log, verbose=1,
                          dqnreg_loss_weight=dqnreg_loss_weight,
                          batch_size=batch_size, learning_rate=learning_rate, gamma=gamma,
                          gradient_steps=gradient_steps,
                          exploration_final_eps=exploration_final_eps)

else:
    if loss_type == 2:
        model = DQNClipped(MlpPolicy, env, policy_kwargs=policy_kwargs, target_update_interval=target_update_interval,
                           exploration_fraction=exploration_fraction,
                    buffer_size=buffer_size, train_freq=train_freq, learning_starts=learning_starts, seed=seed,
                           tensorboard_log=tensorboard_log, verbose=1)

    elif loss_type == 3:
        model = DQNReg(MlpPolicy, env, policy_kwargs=policy_kwargs, target_update_interval=target_update_interval,
                       exploration_fraction=exploration_fraction,
                    buffer_size=buffer_size, train_freq=train_freq, learning_starts=learning_starts, seed=seed,
                       tensorboard_log=tensorboard_log, verbose=1,
                     dqnreg_loss_weight=dqnreg_loss_weight)

model.learn(total_timesteps=total_timesteps, log_interval=100, callback=eval_callback)

model.save(model_save_name)


text:
# coding=utf-8
"""
© 2013 LinkedIn Corp. All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
"""
import numpy
import os
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from mpl_toolkits.axes_grid1 import host_subplot
import mpl_toolkits.axisartist as AA
import logging
import naarad.naarad_constants as CONSTANTS

logger = logging.getLogger('naarad.graphing.matplotlib')

def convert_to_mdate(date_str):
  mdate = mdates.epoch2num(int(date_str)/1000)
  return mdate

# MPL-WA-07
# matplotlib does not rotate colors correctly when using multiple y axes. This method fills in that gap.
def get_current_color(index):
  return CONSTANTS.COLOR_PALETTE[index % len(CONSTANTS.COLOR_PALETTE)]

def get_graph_metadata(plots):
  height = 0
  width = 0
  title = ''
  for plot in plots:
    if plot.graph_height > height:
      height = plot.graph_height
    if plot.graph_width > width:
      width = plot.graph_width
    if title == '':
      title = plot.graph_title
    elif title != plot.graph_title:
      title = title + ',' + plot.graph_title
  return height/80, width/80, title

def curate_plot_list(plots):
  delete_nodes = []
  for plot in plots:
    if os.path.exists(plot.input_csv):
      if not os.path.getsize(plot.input_csv):
        logger.warning("%s file is empty. No plot corresponding to this file will be generated", plot.input_csv)
        delete_nodes.append(plot)
    else:
      logger.warning("%s file does not exist. No plot corresponding to this file will be generated", plot.input_csv)
      delete_nodes.append(plot)
  for node in delete_nodes:
    plots.remove(node)
  return plots

def highlight_region(plt, start_x, end_x):
  """
  Highlight a region on the chart between the specified start and end x-co-ordinates.
  param pyplot plt: matplotlibk pyplot which contains the charts to be highlighted
  param string start_x : epoch time millis
  param string end_x : epoch time millis
  """
  start_x = convert_to_mdate(start_x)
  end_x = convert_to_mdate(end_x)
  plt.axvspan(start_x, end_x, color=CONSTANTS.HIGHLIGHT_COLOR, alpha=CONSTANTS.HIGHLIGHT_ALPHA)

def graph_data(list_of_plots, output_directory, resource_path, output_filename):
  plots = curate_plot_list(list_of_plots)
  plot_count = len(plots)
  if plot_count == 0:
    return False, None
  graph_height, graph_width, graph_title = get_graph_metadata(list_of_plots)
  current_plot_count = 0
  fig, axis = plt.subplots()
  fig.set_size_inches(graph_width, graph_height)
  if plot_count < 2:
    fig.subplots_adjust(left=CONSTANTS.SUBPLOT_LEFT_OFFSET, bottom=CONSTANTS.SUBPLOT_BOTTOM_OFFSET, right=CONSTANTS.SUBPLOT_RIGHT_OFFSET)
  else:
    fig.subplots_adjust(left=CONSTANTS.SUBPLOT_LEFT_OFFSET, bottom=CONSTANTS.SUBPLOT_BOTTOM_OFFSET, right=CONSTANTS.SUBPLOT_RIGHT_OFFSET - CONSTANTS.Y_AXIS_OFFSET * (plot_count - 2))
  current_axis = axis
  for plot in plots:
    current_plot_count += 1
    logger.info('Processing: ' + plot.input_csv + ' [ ' + output_filename + ' ]')
    timestamp, yval = numpy.loadtxt(plot.input_csv, unpack=True, delimiter=',', converters={0: convert_to_mdate})
    maximum_yvalue = numpy.amax(yval) * (1.0 + CONSTANTS.ZOOM_FACTOR * current_plot_count)
    minimum_yvalue = numpy.amin(yval) * (1.0 - CONSTANTS.ZOOM_FACTOR * current_plot_count)

    if current_plot_count == 0:
      current_axis.yaxis.set_ticks_position('left')
    if current_plot_count > 1:
      current_axis = axis.twinx()
      current_axis.yaxis.grid(False)
      #Set right y-axis for additional plots
      current_axis.yaxis.set_ticks_position('right')
      #Offset the right y axis to avoid overlap
      current_axis.spines['right'].set_position(('axes', 1 + CONSTANTS.Y_AXIS_OFFSET * (current_plot_count-2)))
      current_axis.spines['right'].set_smart_bounds(False)
      current_axis.spines['right'].set_color(get_current_color(current_plot_count))
      current_axis.set_frame_on(True)
      current_axis.patch.set_visible(False)
    current_axis.set_ylabel(plot.y_label, color=get_current_color(current_plot_count), fontsize=CONSTANTS.Y_LABEL_FONTSIZE)
    current_axis.set_ylim([minimum_yvalue, maximum_yvalue])
    if plot.graph_type == 'line':
      current_axis.plot_date(x=timestamp, y=yval, linestyle='-', marker=None, color=get_current_color(current_plot_count))
    else:
      current_axis.plot_date(x=timestamp, y=yval, marker='.', color=get_current_color(current_plot_count))
    y_ticks = current_axis.get_yticklabels()
    for y_tick in y_ticks:
      y_tick.set_color(get_current_color(current_plot_count))
      y_tick.set_fontsize(CONSTANTS.Y_TICKS_FONTSIZE)
    for x_tick in current_axis.get_xticklabels():
      x_tick.set_fontsize(CONSTANTS.X_TICKS_FONTSIZE)
    if plot.highlight_regions is not None:
      for region in plot.highlight_regions:
        highlight_region(plt, str(region.start_timestamp), str(region.end_timestamp))
  axis.yaxis.grid(True)
  axis.xaxis.grid(True)
  axis.set_title(graph_title)
  axis.set_xlabel('Time')
  x_date_format = mdates.DateFormatter(CONSTANTS.X_TICKS_DATEFORMAT)
  axis.xaxis.set_major_formatter(x_date_format)
  plot_file_name = os.path.join(output_directory, output_filename + ".png")
  fig.savefig(plot_file_name)
  plt.close()
  #Create html fragment to be used for creation of the report
  with open(os.path.join(output_directory, output_filename + '.div'), 'w') as div_file:
    div_file.write('

' + os.path.basename(plot_file_name) + '

')
  return True, os.path.join(output_directory, output_filename + '.div')

def graph_data_on_the_same_graph(list_of_plots, output_directory, resource_path, output_filename):
  """
  graph_data_on_the_same_graph: put a list of plots on the same graph: currently it supports CDF
  """
  maximum_yvalue = -float('inf')
  minimum_yvalue = float('inf')
  plots = curate_plot_list(list_of_plots)
  plot_count = len(plots)
  if plot_count == 0:
    return False, None
  graph_height, graph_width, graph_title = get_graph_metadata(plots)
  current_plot_count = 0
  fig, axis = plt.subplots()
  fig.set_size_inches(graph_width, graph_height)
  if plot_count < 2:
    fig.subplots_adjust(left=CONSTANTS.SUBPLOT_LEFT_OFFSET, bottom=CONSTANTS.SUBPLOT_BOTTOM_OFFSET, right=CONSTANTS.SUBPLOT_RIGHT_OFFSET)
  else:
    fig.subplots_adjust(left=CONSTANTS.SUBPLOT_LEFT_OFFSET, bottom=CONSTANTS.SUBPLOT_BOTTOM_OFFSET, right=CONSTANTS.SUBPLOT_RIGHT_OFFSET - CONSTANTS.Y_AXIS_OFFSET * (plot_count - 2))
  # generate each plot on the graph
  for plot in plots:
    current_plot_count += 1
    logger.info('Processing: ' + plot.input_csv + ' [ ' + output_filename + ' ]')
    xval, yval = numpy.loadtxt(plot.input_csv, unpack=True, delimiter=',')
    axis.plot(xval, yval, linestyle='-', marker=None, color=get_current_color(current_plot_count), label=plot.plot_label)
    axis.legend()
    maximum_yvalue = max(maximum_yvalue, numpy.amax(yval) * (1.0 + CONSTANTS.ZOOM_FACTOR * current_plot_count))
    minimum_yvalue = min(minimum_yvalue, numpy.amin(yval) * (1.0 - CONSTANTS.ZOOM_FACTOR * current_plot_count))
  # set properties of the plots
  axis.yaxis.set_ticks_position('left')
  axis.set_xlabel(plots[0].x_label)
  axis.set_ylabel(plots[0].y_label, fontsize=CONSTANTS.Y_LABEL_FONTSIZE)
  axis.set_ylim([minimum_yvalue, maximum_yvalue])
  axis.yaxis.grid(True)
  axis.xaxis.grid(True)
  axis.set_title(graph_title)
  plot_file_name = os.path.join(output_directory, output_filename + ".png")
  fig.savefig(plot_file_name)
  plt.close()
  # Create html fragment to be used for creation of the report
  with open(os.path.join(output_directory, output_filename + '.div'), 'w') as div_file:
    div_file.write('

' + os.path.basename(plot_file_name) + '

')
  return True, os.path.join(output_directory, output_filename + '.div')

max_stars_repo_name: feng-tao/naarad

max_issues_repo_name: feng-tao/naarad

max_forks_repo_name: feng-tao/naarad


text:
import pandas as pd
import numpy as np
import sys

def profile_triple_SS(profile, models):
    SS = []
    for i in range(models.shape[0]):
        temp = models.iloc[i]
        try:
            pred = float(temp['Coef_p1'])*profile.loc[temp['Predictor1']] + \
                   float(temp['Coef_p2'])*profile.loc[temp['Predictor2']] + float(temp['Intercept'])
            SS.append((profile.loc[temp['Response']] - pred) ** 2)
        except KeyError:
            continue
    return(SS)

def profile_pair_SS(profile, models):
    SS = []
    for i in models.index:
        temp = models.iloc[i]
        try:
            pred = temp['Coef_p'] * profile.loc[temp.loc['Predictor']] + temp['Intercept']
            SS.append((profile.loc[temp['Response']] - pred) ** 2)
        except KeyError:
            continue
    return(SS)

def RMSE_test_profile(do_models_heal, tr_models_heal, do_models_ibd, tr_models_ibd, profile):

    heal_tr_ss = profile_triple_SS(profile, tr_models_heal)
    heal_do_ss = profile_pair_SS(profile, do_models_heal)
    heal_SS = list(heal_tr_ss) + list(heal_do_ss)
    heal_RMSE = (sum(heal_SS) / len(heal_SS)) ** (0.5)

    ibd_tr_ss = profile_triple_SS(profile, tr_models_ibd)
    ibd_do_ss = profile_pair_SS(profile, do_models_ibd)
    ibd_SS = list(ibd_tr_ss) + list(ibd_do_ss)
    ibd_RMSE = (sum(ibd_SS) / len(ibd_SS)) ** (0.5)

    # delta = heal_RMSE-(ibd_RMSE*multiplier)
    # delta = ['Healthy' if i <0 else 'IBD' for i in delta]
    return(heal_RMSE, ibd_RMSE)

def acc_for_multiplier(RMSEs, mult):
    heal_RMSE = RMSEs[0]
    ibd_RMSE = RMSEs[1]
    delta = heal_RMSE - (ibd_RMSE * mult)
    delta = ['Healthy' if i < 0 else 'IBD' for i in delta]
    return(delta)

def return_tax_names(ID_to_taxon, profiles):
    tax_code = {}
    with open(ID_to_taxon) as f:
        for line in f.read().splitlines():
            temp = line.split()
            # tax ID is in the 1st column tax name -- 2nd
            tax_code[temp[0]] = temp[1]
    profiles.index = [tax_code[x] for x in profiles.index]
    return(profiles)

# params = 'E:\\Lab\\FHM\\AGP\\Parallel\\py\\PARAM_FINAL_OTU_TABLE_ibd.txt'
# ibd_folders = 'E:\\Lab\\FHM\\AGP\\Parallel\\FINAL_OTU_TABLE_ibd'
# heal_folders = 'E:\\Lab\\FHM\\AGP\\Parallel\\FINAL_OTU_TABLE_heal'

params = sys.argv[1]
ibd_folders = sys.argv[2]
heal_folders = sys.argv[3]
out = sys.argv[4]
test_or_valid = sys.argv[5]

params = pd.read_table(params, header = 0, index_col = None)
print(len(params.index))
Pf_prev = -1
for i in params.index:

    Tt = str(params.loc[i,'Taxon_trim'])
    Sc = str(params.loc[i,'SparCC_cor'])
    Sp = str(params.loc[i,'SparCC_pval'])
    Ta = str(int(params.loc[i,'Tax_adjacency']))
    Pf = str(params.loc[i,'Pair_fstat'])

    if Pf_prev != Pf:

        Pf_prev = Pf

        template = '{}/trim_{}/pval_{}/dist_{}_cor_{}/fstat_{}'
        paths = {k:template.format(x, Tt, Sp, Ta, Sc, Pf)
                 for k,x in (('ibd',ibd_folders), ('heal',heal_folders))}
        print(paths)
        p_tax = {k:v for k,v in (
                                ('ibd',ibd_folders+'/trim_%s/tax_code.txt' % Tt),
                                ('heal',heal_folders+'/trim_%s/tax_code.txt' % Tt)
                                )}
        p_do_models = {k:v for k,v in (
                                        ('ibd', paths['ibd'] +'/pair_models_{}.txt'.format(Pf)),
                                        ('heal', paths['heal'] + '/pair_models_{}.txt'.format(Pf))
                                     )}
        p_tr_models = {k:v for k,v in (
                                        ('ibd', paths['ibd'] + '/triplet_models_{}.txt'.format(Pf)),
                                        ('heal',paths['heal'] + '/triplet_models_{}.txt'.format(Pf))
                                    )}
        p_test_profile = {k:v for k,v in (
                                        ('ibd', ibd_folders + '/trim_{}/{}.txt'.format(Tt,test_or_valid)),
                                        ('heal', heal_folders + '/trim_{}/{}.txt'.format(Tt,test_or_valid))
                                        )}
        try:
            with open(p_tr_models['ibd']) as f:
                if f.readline() == 'No pair models to base triple models on':
                    print('No pair models to base triple models on')
                    params.loc[i, 'Sensitivity'] = 'NA'
                    params.loc[i, 'Specificity'] = 'NA'
                    continue
            with open(p_tr_models['heal']) as f:
                if f.readline() == 'No pair models to base triple models on':
                    print('No pair models to base triple models on')
                    params.loc[i, 'Sensitivity'] = 'NA'
                    params.loc[i, 'Specificity'] = 'NA'
                    continue

            do_models = {k:pd.read_table(v, header=0, index_col=None, sep='\t') for k,v in p_do_models.items()}
            tr_models = {k:pd.read_table(v, header=0, index_col=None, sep='\t') for k, v in p_tr_models.items()}
            test_profile = {k:pd.read_table(v, header=0, index_col=0, sep='\t') for k,v in p_test_profile.items()}
            # tax = {k:pd.read_table(v, header=None, index_col=0, sep='\t') for k,v in tax.items()}

            for el in test_profile:
                test_profile[el] = return_tax_names(p_tax[el], test_profile[el])

            H = RMSE_test_profile(do_models['heal'], tr_models['heal'], do_models['ibd'], tr_models['ibd'],
                                  test_profile['heal'])
            I = RMSE_test_profile(do_models['heal'], tr_models['heal'], do_models['ibd'], tr_models['ibd'],
                                  test_profile['ibd'])

        except:
            print('File misssing')
            params.loc[i, 'Sensitivity'] = 'Miss'
            params.loc[i, 'Specificity'] = 'Miss'
            continue

    for file in (p_tr_models['ibd'], p_tr_models['heal']):
        try:
            with open(file) as f:
                if f.readline() == 'No pair models to base triple models on':
                    print('No pair models to base triple models on')
                    params.loc[i, 'Sensitivity'] = 'NA'
                    params.loc[i, 'Specificity'] = 'NA'
                    continue
        except:
            print('File misssing')
            params.loc[i, 'Sensitivity'] = 'Miss'
            params.loc[i, 'Specificity'] = 'Miss'
            continue

    multiplier = params.loc[i,'RMSE_sign']

    H_RMSE = acc_for_multiplier(H, multiplier)
    I_RMSE = acc_for_multiplier(I, multiplier)

    sens = I_RMSE.count('IBD') / float(len(I_RMSE))
    spec = H_RMSE.count('Healthy') / float(len(H_RMSE))
    params.loc[i, 'Sensitivity'] = sens
    params.loc[i, 'Specificity'] = spec
    print( str(i) + ' out of ' + str(params.shape[0]) )

params.to_csv(out, sep='\t', index=False)


text:
"""
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np

__all__ = ('pure_python_counts_in_cylinders', 'pure_python_idx_in_cylinders')

fixed_seed = 43

def pure_python_idx_in_cylinders(
        sample1, sample2, rp_max, pi_max, period=None):
    if period is None:
        xperiod, yperiod, zperiod = np.inf, np.inf, np.inf
    else:
        xperiod, yperiod, zperiod = period, period, period

    autocorr = False
    if sample2 is None:
        sample2, autocorr = sample1, True

    npts1, npts2 = len(sample1), len(sample2)

    indexes = []

    for i in range(npts1):
        for j in range(npts2):
            if i == j and autocorr:
                continue
            if _point_within_cylinder(sample1[i], sample2[j], rp_max[i], pi_max[i], xperiod, yperiod, zperiod):
                indexes.append((i, j))

    # https://github.com/numpy/numpy/issues/2407 for str("colname")
    return np.array(indexes, dtype=[(str("i1"), np.uint32), (str("i2"), np.uint32)])

def pure_python_counts_in_cylinders(
        sample1, sample2, rp_max, pi_max, period=None):
    """ Brute force pure python function calculating the distance
    between all pairs of points and storing the result into two matrices,
    one storing xy-distances, the other storing z-distances,
    account for possible periodicity of the box.
    """
    if period is None:
        xperiod, yperiod, zperiod = np.inf, np.inf, np.inf
    else:
        xperiod, yperiod, zperiod = period, period, period

    autocorr = False
    if sample2 is None:
        sample2, autocorr = sample1, True

    npts1, npts2 = len(sample1), len(sample2)

    counts = np.zeros(npts1, dtype=int)

    for i in range(npts1):
        for j in range(npts2):
            if i == j and autocorr:
                continue
            if _point_within_cylinder(sample1[i], sample2[j], rp_max[i], pi_max[i], xperiod, yperiod, zperiod):
                counts[i] += 1
    return counts

def _point_within_cylinder(p1, p2, rp_cylinder, pi_cylinder, xperiod, yperiod, zperiod):
    dx = p1[0] - p2[0]
    dy = p1[1] - p2[1]
    dz = p1[2] - p2[2]

    if dx > xperiod/2.:
        dx = xperiod - dx
    elif dx < -xperiod/2.:
        dx = -(xperiod + dx)

    if dy > yperiod/2.:
        dy = yperiod - dy
    elif dy < -yperiod/2.:
        dy = -(yperiod + dy)

    if dz > zperiod/2.:
        dz = zperiod - dz
    elif dz < -zperiod/2.:
        dz = -(zperiod + dz)

    rp = np.sqrt(dx*dx + dy*dy)
    dz = abs(dz)
    return (rp <= rp_cylinder) & (dz <= pi_cylinder)

max_stars_repo_name: pllim/halotools

max_issues_repo_name: pllim/halotools

max_forks_repo_name: pllim/halotools


text:
import os
import shutil
import numpy as np
import matplotlib.pyplot as plt

def generate(count):
    labels = np.zeros((count,)).astype(int)
    images = np.zeros((count, 3, 3)).astype(int)
    k = 0
    # primary images
    for i in range(3):
        for j in range(3):
            if np.absolute(i - j) < 2:
                images[k, i, :] = [0, 255, 255]
                images[k, j, 0] = 255
                labels[k] = 0
                k += 1
                images[k,:,:] = images[k-1,:,:].T
                labels[k] = 1
                k += 1
                if i != j:
                    images[k,:,:] = np.fliplr(images[k-2,:,:])
                    labels[k] = 0
                    k += 1
                    images[k,:,:] = images[k-1,:,:].T
                    labels[k] = 1
                    k += 1
    # random noising of pixels
    for i in range(count - 22):
        j = i % 22
        images[i+22, :, :] = images[j,:,:]
        for x in range(3):
            labels[i+22] = i % 2
            for y in range(3):
                images[i+22,x,y] = np.floor(images[i+22,x,y] * np.random.uniform(.5, 1))
    return images, labels

def flatten_save(images, labels, filename=""):
    data = np.zeros((labels.shape[0], 10)).astype(int)
    for i in range(labels.shape[0]):
        data[i, 0] = labels[i]
        data[i, 1:] = images[i,:,:].ravel()

    if len(filename) > 0:
        np.savetxt(filename, data, delimiter=",", fmt='%5d')

    return data

def image_save(images, labels, path):
    if os.path.exists(path):
        shutil.rmtree('%s' % path)
    os.makedirs(path)
    for i in range(labels.shape[0]):
        # if on a mac this could bork - should use os.path.join(...)
        if not os.path.exists('%s\%d' % (path, labels[i])):
            os.makedirs('%s\%d' % (path, labels[i]))
        plt.imsave('%s\%d\%08d.png' % (path, labels[i], i), images[i].astype(float), cmap='gray')

if __name__ == "__main__":
    images, labels = generate(1000)
    flatten_save(images, labels, 'fakes.csv')
    image_save(images, labels, 'data')


text:
import pandas as pd
import numpy as np
import math
import spacy
import os

import time

from utils import preprocessText
from MutantGeneration import MutantGeneration
from multiprocessing import Pool, Process, Queue, Manager
import multiprocessing

def compute_mut():
    '''for multiprocessing uaage'''
    while True:
        if not q.empty():
            label, text = q.get()
            text = preprocessText(text)
            mg = MutantGeneration(text)

            if len(mg.getMutants()) > 0:
                original = [text] * len(mg.getMutants())
                label = [label] * len(mg.getMutants())
                template = mg.getTemplates()
                mutant = mg.getMutants()
                gender = mg.getGenders()
                q_to_store.put((
                    original, label, template, mutant, gender
                ))
        else:
            print("Finished")
            return

df = pd.read_csv("../../asset/imdb/test.csv", names=["label", "sentence"], sep="\t")

# df = pd.read_csv("../../data/biasfinder/template_gender.csv")

# df = pd.read_csv("user_study_country-artha.csv")
# df = df[df["mod"] == 0]
# df = df[df["is_make_sense"] == "Yes"]

start = time.time()

originals = []
templates = []
mutants = []
labels = []
identifiers = []
types = []
genders = []
countries = []

n_template = 0

i = 0
counter = 0

manager = multiprocessing.Manager()

q = manager.Queue()
q_to_store = manager.Queue()

for index, row in df.iterrows():
    label = row["label"]
    text = row["sentence"]

    q.put((label, text))

numList = []
for i in range(8) :
    p = multiprocessing.Process(target=compute_mut, args=())
    numList.append(p)
    p.start()

for i in numList:
    i.join()

print("Generation Process finished.")

while not q_to_store.empty():
    original, label, template, mutant, gender = q_to_store.get()
    originals.extend(original)
    labels.extend(label)
    templates.extend(template)
    mutants.extend(mutant)
    genders.extend(gender)

end = time.time()
print("Execution Time: ", end-start)

dm = pd.DataFrame(data={"label": labels, "mutant": mutants, "template": templates, "original": originals, "gender": genders})
dm

dm["template"] = dm["template"].astype("category")
dm["template_id"] = dm["template"].cat.codes

# print(n_template)
# print(len(dm))

data_dir = "../../data/biasfinder/gender/"

if not os.path.exists(data_dir) :
    os.makedirs(data_dir)

dm.to_csv(data_dir + "test.csv", index=None, header=None, sep="\t")


text:
import numpy as np

#from pykalman import KalmanFilter
import pickle
from scipy import interpolate
import matplotlib.pylab as plt
import random
from scipy.optimize import leastsq

batteryTime=0
relayStatus=0
totalVoltage=0.0
totalCurrent=0.0
quitFlag=0

# chargeVoltage=np.array([3.023, 3.398, 3.490, 3.557, 3.618, 3.656, 3.721, 3.846, 3.943, 4.057, 4.184])
# chargeQuantity=np.array([74.2, 583, 1086.5, 1590.0, 2183.6, 2687.1, 3190.6, 3784.2, 4287.7, 4796.5, 5300.0])
# chargeSOC=np.array([0.014,0.11,0.205,0.30,0.412,0.507,0.602,0.714,0.809,0.905,1.0])

chargeQuantity=[10.0,74.2,159,243.8,328.6,413.4,498.2,583,662.5,747.3,832.1\
               ,916.9,1001.7,1086.5,1171.3,1256.1,1340.9,1425.7,1505.2\
               ,1590,1674.8,1759.6,1844.4,1929.2,2014,2098.8,2183.6,2268.4\
               ,2347.9,2432.7,2517.5,2602.3,2687.1,2771.9,2856.7,2941.5\
               ,3026.3,3111.1,3190.6,3275.4,3360.2,3445,3529.8,3614.6\
               ,3699.4,3784.2,3869,3953.8,4033.3,4118.1,4202.9,4287.7\
               ,4372.5,4457.3,4542.1,4626.9,4711.7,4796.5,4876,4960.8\
               ,5045.6,5130.4,5215.2,5300]

chargeVoltage=np.array([2.80,3.023,3.159,3.25,3.318,3.362,3.383,3.398,3.412,3.429,3.446\
              ,3.462,3.477,3.49,3.503,3.513,3.524,3.535,3.545,3.557,3.569\
              ,3.584,3.594,3.601,3.607,3.612,3.618,3.624,3.629,3.635,3.642\
              ,3.648,3.656,3.664,3.672,3.682,3.693,3.706,3.721,3.74,3.761\
              ,3.78,3.797,3.813,3.829,3.846,3.861,3.877,3.893,3.909,3.926\
              ,3.943,3.96,3.978,3.997,4.017,4.037,4.057,4.077,4.096,4.115\
              ,4.135,4.158,4.184])

chargeSOC=np.array([0.009,0.014,0.03,0.046,0.062,0.078,0.094,0.11\
          ,0.125,0.141,0.157,0.173,0.189,0.205,0.221,0.237\
          ,0.253,0.269,0.284,0.3,0.316,0.332,0.348,0.364\
          ,0.38,0.396,0.412,0.428,0.443,0.459,0.475,0.491\
          ,0.507,0.523,0.539,0.555,0.571,0.587,0.602,0.618\
          ,0.634,0.65,0.666,0.682,0.698,0.714,0.73,0.746,0.761\
          ,0.777,0.793,0.809,0.825,0.841,0.857,0.873,0.889\
          ,0.905,0.92,0.936,0.952,0.968,0.984,0.99])

def qvRealtion(voltage):
    global chargeVoltage, chargeQuantity
    vPoints=chargeVoltage
    qPoints=chargeQuantity
    #SOCPoints=chargeSOC
    qvCurveTck=interpolate.splrep(vPoints,qPoints,s=0.002)
    batteryQuantity=interpolate.splev(voltage,qvCurveTck,der=0)
    #SOCVCurveTck=interpolate.splrep(SOCPoints,vPoints,s=0.002)
    #batterySOC=interpolate.splev(voltage,SOCVCurveTck,der=0)
    #newX = np.arange(vPoints[0],vPoints[-1],0.1)
    # newY = interpolate.splev(newX,qvCurveTck,der=0)
    # plt.plot(vPoints,qPoints,'o',newX,newY,'-')
    # plt.title("V-Q curve")
    #newSOC=np.arange(SOCPoints[0],SOCPoints[-1],0.01)
    #newVoltage=interpolate.splev(newSOC,SOCVCurveTck,der=0)
    #plt.plot(SOCPoints,vPoints,'o',newSOC,newVoltage,'-')
    #plt.title("V-SOC curve")
    #plt.show()
    return batteryQuantity

#P0 =[ 3.346366    0.59163954 -0.00542907 -0.0092671  -0.06431112]
#real parameter[ 3.346366,0.59163954,0.00542907,-0.0092671,-0.06431112]
def fun(Zk,P):
    K0,K1,K2,K3,K4=P
    return K0+K1*Zk+K2/Zk+K3*np.log(Zk)+K4*np.log(1-Zk)
    #return 3.7638--0.00013*ZK-0.3650/ZK+0.2201*np.log(ZK)-0.0272*np.log(1-ZK)

def residuals(P,Yk,Zk):
    err = Yk-fun(Zk,P)
    return err

def plotfun(P):
    x=chargeSOC
    fig= plt.figure()
    fig.suptitle('Real SOC_OCV Curve and Fit Curve',fontsize=14,fontweight='bold')
    ax = fig.add_subplot(1,1,1)
    ax.set_xlabel('SOC(%)')
    ax.set_ylabel('Voltage(V)')
    y = fun(x,P)
    ax.plot(x*100,y,'b--',label='Estimated Curve')
    ax.plot(chargeSOC*100,chargeVoltage,'g',label='Real Curve')
    ax2 = ax.twinx()
    s2=(chargeVoltage-y)*100
    ax2.set_ylim(-10,90)
    ax2.plot(chargeSOC*100, s2, 'r--',label='Error')
    ax2.set_ylabel('Error(%)', color='r')
    ax.legend(loc="upper left")
    ax2.legend(loc="lower right")
    plt.show()

def calSOC():
    f=open('calSOC.txt','a')
    initQuantity=[0]*NumberOfCells*NumberOfStacks
    initSOC=[0]*NumberOfCells*NumberOfStacks
    dischargeQuantity=[0]*NumberOfCells*NumberOfStacks
    SOC=[0]*NumberOfCells*NumberOfStacks
    batteryTime=int(sendCommand('J000')) #sync the clock by reading clock in the battery simulator
    f.write(str(batteryTime)+'\t')
    if batteryTime == 0:
        for index in range(len(voltage)):
            initQuantity[index]=qvRealtion(voltage[index])
            initSOC[index]=initQuantity[index]/BATTERYCAPACITY
        totalQuantity=sum(initQuantity)
        totalSOC=totalQuantity/(BATTERYCAPACITY*NumberOfCells*NumberOfStacks)
        f.write(str(totalSOC)+'\t')
        f.write('\n')
        f.close()
        pickle.dump(initQuantity,open('data.pickle','w'))
    else:
        initialQuantity=pickle.load(open('data.pickle','r'))
        for index in range(len(initialQuantity)):
            dischargeQuantity[index]=initialQuantity[index]- totalCurrent * batteryTime/3600
            SOC[index]=dischargeQuantity[index]/BATTERYCAPACITY
        totalQuantity=sum(dischargeQuantity)
        totalSOC=totalQuantity/(BATTERYCAPACITY*NumberOfCells*NumberOfStacks)
        f.write(str(totalSOC)+'\t')
        f.write('\n')
        f.close()

if __name__ == '__main__':
    #qvRealtion(3.65)
    P0 = [1,1,1,1,1]
    r=leastsq(residuals,P0,args=(chargeVoltage,chargeSOC))
    print r[0]
    plotfun(r[0])
    #result:[ 3.346366    0.59163954 -0.00542907 -0.0092671  -0.06431112]

max_stars_repo_name: oeshine/batterySimulator

max_issues_repo_name: oeshine/batterySimulator

max_forks_repo_name: oeshine/batterySimulator


text:
"""
This module allows to handle the output of a microscopy experiment
analyzed with the software MoMA. It can parse existing files, help
to load data etc.
 """
# Author: Guillaume Witz, Biozentrum Basel, 2019
# License: MIT License

import os
import glob
import re
import struct
import numpy as np
import pandas as pd
import tifffile as tff

def find_firstfile(directory, file_regexp):
    """Returns the name and path of the first file found within a directory that
    conforms to a given regular expression

    Parameters
    ----------
    directory : string
        path of directory
    regepx: string
        regular expression

    Returns
    -------
    file : string
        first file found to conform within the tree
    filepath: string
        path of file
    """
    file = None
    filepath = None
    found =False
    for root, dirs, files in os.walk(directory,topdown=False):
        for f in files:
            if re.search(file_regexp,f):
                filepath = root
                file = f
                found = True
                break
        if found:
            break
    return file, filepath

class Momaobj:
    """Parsing of MicroManager metadata"""
    def __init__(self, version = None, data_folder = None, rawdata_folder = None,
                 moma_generic =None, tif_generic = None, gl = None, pos = None,
                 gl_list = None, pos_list = None, all_pos=None, all_gl = None, color = 1, time = 0,
                 col_nb = None, height = None, width = None, no_raw = False):
        """Standard __init__ method.

        Parameters
        ----------
        version : int
            moma version (way to save data)
        data_folder : str
            location of cut-out tiff data
        rawdata_folder : str
            location of raw MM images
        moma_generic : str
            generic name of a momoa file
        tif_generic : str
            generic name of a tif file
        gl : str
            current growth lane
        pos : str
            current position
        gl_list : list of str
            list of analysed growth lanes
        pos_list : list of str
            list of analysed positions
        all_gl : list of str
            list of growth lanes
        all_pos : list of str
            list of positions
        color : int
            current color
        time : int
            time point
        col_nb : int
            number of different colors
        height : int
            height of image
        width : int
            width of image
        no_raw : bool
            the original tiffs present in the folder
        """

        self.data_folder = data_folder
        self.rawdata_folder = rawdata_folder

        self.get_generics()
        self.get_version()

        if no_raw:
            self.get_pos_gl_noraw()

        else:
            self.get_pos_gl()
            self.get_pos_gl_analyzed()

        if self.moma_generic is None:
            self.gl = self.all_gl[0]
            self.pos = self.all_pos[0]
        else:
            self.gl = self.gl_list[0]
            self.pos = self.pos_list[0]

        self.col = color
        self.time = time
        self.col_nb = col_nb

        self.get_dims()

    #based on file name find if data are saved as stack or not
    def get_version(self):
        stack_or_not = len(re.findall('.*t\d*_c\d*.*',self.tif_generic))
        self.version = stack_or_not

    #going through all MoMA cropped data, find all available positions and growth lanes
    def get_pos_gl(self):
        all_pos_gl = []
        for root, dirs, files in os.walk(self.data_folder,topdown=False):
            if len(dirs)>0:
                for d in dirs:
                    pos_gl = re.findall('.*[pP]os(\d+).*_GL(\d+).*',d)
                    if len(pos_gl)>0:
                        all_pos_gl.append(pos_gl[0])
        self.all_pos = [x[0] for x in all_pos_gl]
        self.all_gl = [x[1] for x in all_pos_gl]
        return all_pos_gl

    #if image data are not available, position and growth lanes are taked from MoMA csv files
    def get_pos_gl_noraw(self):
        csvdir = os.path.dirname(self.moma_generic)
        pos_gl = [re.findall('.*[pP]os(\d+).*_GL(\d+).*',d)[0] for d in glob.glob(csvdir+'/*.csv')]
        self.all_pos = [x[0] for x in pos_gl]
        self.all_gl = [x[1] for x in pos_gl]

        self.pos_list = self.all_pos.copy()
        self.gl_list = self.all_gl.copy()

    #get generic path to MoMA csv and tif files
    def get_generics(self):
        csv_file, csv_path = find_firstfile(self.data_folder,'Exported.*csv')
        tif_file, tif_path = find_firstfile(self.data_folder,'.*tif')
        if csv_path is None:
            self.moma_generic = None
        else:
            self.moma_generic = csv_path+'/'+csv_file
        self.tif_generic = tif_path+'/'+tif_file

    #given a position and growth lane get the corresponding MoMA file name
    def get_momapath(self, pos = None, gl = None):
        if not pos:
            pos = self.pos
        if not gl:
            gl = self.gl

        new_csv = re.sub('(?<=[pP]os)\d*?(?=[/_])', pos,self.moma_generic)
        new_csv = re.sub('(?<=GL)\d*(?!=\d)', gl,new_csv)
        return new_csv

    #given a position and growth lane get the corresponding tif file name
    def get_tifpath(self, pos = None, gl = None):
        if not pos:
            pos = self.pos
        if not gl:
            gl = self.gl
        new_tif = re.sub('(?<=[pP]os)\d*?(?=[/_])', pos,self.tif_generic)
        new_tif = re.sub('(?<=GL)\d*(?!=\d)', gl,new_tif)
        if self.version == 1:
            new_tif = re.sub('(?<=t)\d*(?=_c)', (4-len(str(self.time+1)))*'0'+str(self.time+1),new_tif)
            new_tif = re.sub('(?<=t\d{4}_c)\d*(?=\.tif)', (4-len(str(self.col)))*'0'+str(self.col),new_tif)
        return new_tif

    #get a position and growth lane list of analyzed data (MoMA data)
    def get_pos_gl_analyzed(self):
        if self.moma_generic is None:
            self.pos_list = None
            self.gl_list = None
        else:
            all_pos = self.all_pos
            all_gl = self.all_gl
            pos_analyzed = [all_pos[x] for x in range(len(all_pos)) if os.path.exists(self.get_momapath(all_pos[x],all_gl[x]))]
            gl_analyzed = [all_gl[x] for x in range(len(all_gl)) if os.path.exists(self.get_momapath(all_pos[x],all_gl[x]))]
            #gl_analyzed = [x[1] for x in all_pos_gl if os.path.exists(self.get_momapath(x[0],x[1]))]
            self.pos_list = pos_analyzed
            self.gl_list = gl_analyzed

    #get maximum available time point in experiment
    def get_max_time(self):
        if self.version == 1:
            maxtime = np.max([int(re.findall('.*\_t(\d*)\_.*',os.path.basename(x))[0]) for x in glob.glob(os.path.dirname(self.tif_generic)+'/*.tif')])-1
        else:
            maxtime = int(len(tff.TiffFile(self.tif_generic).pages)/self.col_nb)-1
        return maxtime

    #setter for position and growth lane
    def set_pos_gl(self, pos, gl):
        self.pos = pos
        self.gl = gl
        self.get_dims()

    #get image dimensions
    def get_dims(self):
        image = self.load_moma_im()
        self.height = image.shape[0]
        self.width = image.shape[1]

    #load a tif file based on current status of object (color, time etc.)
    def load_moma_im(self):
        """Return image at current position, growth lane, color and time point"""

        tif_path = self.get_tifpath()
        if self.version == 1:
            #image = io.imread(tif_path)
            image = tff.imread(tif_path)
        else:
            '''img = Image.open(tif_path)
            img.seek((self.time)*self.col_nb+self.col-1)
            image = img.getdata()
            image = np.reshape(np.array(image),newshape=[img.height,img.width])'''
            image = tff.imread(tif_path,key=(self.time)*self.col_nb+self.col-1)
        return image

    #load a time series based on the times numpy array and the current status of object (color, time etc.)
    def load_moma_timeseries(self, times):
        """Return image at current position, growth lane, color and time point"""

        image=np.zeros((self.height, self.width, np.array(times).shape[0]))
        for idx, t in enumerate(times):
            self.time = t
            tif_path = self.get_tifpath()
            if self.version == 1:
                #image = io.imread(tif_path)
                image[:,:,idx] = tff.imread(tif_path)
            else:
                image[:,:,idx] = tff.imread(tif_path,key=(self.time)*self.col_nb+self.col-1)
        return image

max_stars_repo_name: guiwitz/DoubleAdderArticle

max_issues_repo_name: guiwitz/DoubleAdderArticle

max_forks_repo_name: guiwitz/DoubleAdderArticle


text:
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
import numpy as np
import scipy.linalg as linalg
import math

#This function pushes a matrix onto the stack that puts everything
#in the frame of a camera which is centered at position "P",
#is pointing towards "t", and has vector "r" to the right
#t - towards vector
#u - up vector
#r - right vector
#P - Camera center
def gotoCameraFrame(t, u, r, P):
    rotMat = np.array([ [r[0], u[0], -t[0], 0], [r[1], u[1], -t[1], 0], [r[2], u[2], -t[2], 0], [0, 0, 0, 1] ])
    rotMat = linalg.inv(rotMat)
    transMat = np.array([ [1, 0, 0, -P[0]], [0, 1, 0, -P[1]], [0, 0, 1, -P[2]], [0, 0, 0, 1] ])
    #Translate first then rotate
    mat = rotMat.dot(transMat)
    #OpenGL is column major and mine are row major so take transpose
    mat = mat.T
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    glMultMatrixd(mat.flatten())

class MousePolarCamera(object):
    #Coordinate system is defined as in OpenGL as a right
    #handed system with +z out of the screen, +x to the right,
    #and +y up
    #phi is CCW down from +y, theta is CCW away from +z
    def __init__(self, pixWidth, pixHeight, yfov = 0.75):
        self.pixWidth = pixWidth
        self.pixHeight = pixHeight
        self.yfov = yfov
        self.center = np.array([0, 0, 0])
        self.R = 1
        self.theta = 0
        self.phi = 0
        self.updateVecsFromPolar()

    def centerOnBBox(self, bbox, theta = -math.pi/2, phi = math.pi/2):
        self.center = np.mean(bbox, 0)
        diagLength = np.sqrt( (bbox[1, :] - bbox[0, :])**2 )
        self.R = diagLength*1.5
        self.theta = theta
        self.phi = phi
        self.updateVecsFromPolar()

    def centerOnPoints(self, X):
        bbox = np.array( [[np.max(X)], [np.min(X)]] )
        self.centerOnBBox(bbox)

    def updateVecsFromPolar(self):
        [sinT, cosT, sinP, cosP] = [math.sin(self.theta), math.cos(self.theta), math.sin(self.phi), math.cos(self.phi)]
        #Make the camera look inwards
        #i.e. towards is -dP(R, phi, theta)/dR, where P(R, phi, theta) is polar position
        self.towards = np.array([-sinP*cosT, -cosP, sinP*sinT])
        self.up = np.array([-cosP*cosT, sinP, cosP*sinT])
        self.eye = self.center - self.R*self.towards

    def gotoCameraFrame(self):
        gotoCameraFrame(self.towards, self.up, np.cross(self.towards, self.up), self.eye)

    def orbitUpDown(self, dP):
        dP = 1.5*dP/float(self.pixHeight)
        self.phi = self.phi+dP
        self.updateVecsFromPolar()

    def orbitLeftRight(self, dT):
        dT = 1.5*dT/float(self.pixWidth)
        self.theta = self.theta-dT
        self.updateVecsFromPolar()

    def zoom(self, rate):
        rate = rate / float(self.pixHeight)
        self.R = self.R*pow(4, rate)
        self.updateVecsFromPolar()

    def translate(self, dx, dy):
        length = np.sqrt((self.center-self.eye)**2)*math.tan(self.yfov);
        dx = length*dx / float(self.pixWidth)
        dy = length*dy / float(self.pixHeight)
        r = np.cross(self.towards, self.up)
        self.center = self.center - dx*r - dy*self.up
        self.updateVecsFromPolar()

max_stars_repo_name: ctralie/PublicationsCode

max_issues_repo_name: ybayle/PublicationsCode

max_forks_repo_name: ybayle/PublicationsCode

reason: import numpy,import scipy