text:
import torch.nn as nn
import torch
import os
import time
import numpy as np
import pandas as pd
import warnings
import utils
from collections import OrderedDict
from tabnet import TabNetModel
from torch.utils.tensorboard import SummaryWriter

class TabNet(object):

    default_train_params = {
        "batch_size": 8192,
        "validation_batch_size": 1024,
        "run_self_supervised_training": False,
        "run_supervised_training": True,
        "early_stopping": True,
        "early_stopping_min_delta_pct": 0,
        "early_stopping_patience": 20,
        "max_epochs_supervised": 500,
        "max_epochs_self_supervised": 500,
        "epoch_save_frequency": 100,
        "train_generator_shuffle": True,
        "train_generator_n_workers": 0,
        "epsilon": 1e-7,
        "learning_rate": 0.01,
        "learning_rate_decay_factor": 0.95,
        "learning_rate_decay_step_rate": 1000,
        "weight_decay": 0.001,
        "sparsity_regularization": 0.0001,
        "p_mask": 0.2,
    }
    default_save_params = {
        "model_name": "forest_cover",
        "tensorboard_folder": "../runs/",
        "save_folder": "../runs/model_backups/",
    }
    default_model_params = {
        "n_steps": 5,
        "n_dims_d": 16,
        "n_dims_a": 16,
        "batch_norm_momentum": 0.85,
        "dropout_p": 0.3,
        "categorical_variables": [],
        "categorical_config": {},
        "discrete_target_mapping": {},
        "embedding_dim": 2,
        "discrete_outputs": False,
        "gamma": 1.5,
    }
    model_save_path = None
    model = None

    def __init__(self, model_params={}, use_cuda=True, save_file=None):
        self.__configure_device(use_cuda)

        c_model_save = self.__load_model(save_file)
        if c_model_save:
            self.model_params, self.model = c_model_save
        else:
            self.model_params = self.default_model_params
            self.model_params.update(model_params)

    def __get_tensorboard_writer(self, X_test_cont, X_test_cat):
        """Creates a tensorboard writer object, and pre-populates the model
        graph on Tensorboard with a test_batch.

        Parameters
        ----------
        X_test_cont : a batch of the continuous features to pre-populate the Tensorboard graph
        X_test_cat : a batch of the categorical features to pre-populate the Tensorboard graph

        Returns
        -------
        writer : torch.utils.tensorboard.SummaryWriter object
        """
        if not os.path.isdir(self.save_params["tensorboard_folder"]):
            os.mkdir(self.save_params["tensorboard_folder"])

        writer = SummaryWriter(
            "{}/{}_{}.run".format(
                self.save_params["tensorboard_folder"],
                int(time.time()),
                self.save_params["model_name"],
            )
        )
        # Populate with dummy data to bypass torch/tensorboard restrictions
        if len(X_test_cat) == 0:
            X_test_cat = {-1: torch.zeros(X_test_cont.size()[0])}
        writer.add_graph(
            self.model,
            [
                X_test_cont,
                X_test_cat,
                self.__generate_model_mask(0, batch_size=X_test_cont.size()[0]),
            ],
        )
        return writer

    def __load_model(self, save_file):
        """Loads model from a local file.

        Parameters
        ----------
        save_file : the path to the model save file

        Returns
        -------
        model_params : Model parameters
        model : Model initialized with saved weights
        """
        if save_file is None:
            return
        if not os.path.isfile(save_file):
            print(
                "File {} not found. Run `train` to configure a new model".format(
                    save_file
                )
            )
            return
        try:
            model_params, model_state_dict = torch.load(
                save_file, map_location=self.device
            )
            model = TabNetModel(**model_params)
            model.load_state_dict(model_state_dict)
            model.to(self.device)
            return model_params, model
        except:
            print(
                "File {} not correctly formatted. Run `train` to configure a new model".format(
                    save_file
                )
            )
            return

    def __save_model(self, save_identifier=None, self_supervised=False):
        """Saves current model parameters to a local file.

        Parameters
        ----------
        save_identifier : the model identifier to be used in the savefile name
        self_supervised : flag to identify whether the model has only been
        trained with self-supervision

        Returns
        -------
        None : Model is saved to disk
        """
        if not os.path.isdir(self.save_params["save_folder"]):
            os.mkdir(self.save_params["save_folder"])

        model_stub = "predictive_model"
        if self_supervised:
            model_stub = "self_supervised_model"

        self.model_save_path = "{}/{}_{}_{}_{}.pt".format(
            self.save_params["save_folder"],
            int(time.time()),
            self.save_params["model_name"],
            model_stub,
            save_identifier,
        )
        print("Saving model to: {}".format(self.model_save_path))
        torch.save((self.model_params, self.model.state_dict()), self.model_save_path)

    def __configure_device(self, use_cuda):
        """Configures torch device, based on configured properties and
        available devices."""
        use_cuda_available = torch.cuda.is_available() and use_cuda
        if use_cuda and use_cuda != use_cuda_available:
            print("Device configuration: Cuda not available - check GPU configuration.")
        self.device = torch.device("cuda:0" if use_cuda_available else "cpu")
        print(
            "Device configuration: Using {} for training/inference".format(self.device)
        )
        torch.backends.cudnn.benchmark = True  # Cudnn configuration

    def __generate_model_mask(self, p, batch_size):
        """Generates a [1,0] mask as input into the model, with the masking
        probability defined by `p`.

        Note: This function handles categorical encodings, to ensure masks
        apply to the entire embedding as opposed to just one dimension.

        Parameters
        ----------
        p : Masking probability. Setting p to 1 will return all zeros
        batch_size : Batch_size
        """
        mask = torch.bernoulli(
            torch.ones(batch_size, self.model_params["n_original_input_dims"]) * (1 - p)
        )
        idx_slices = set(
            [
                self.model_params["categorical_config"][i]["idx"]
                for i in self.model_params["categorical_config"]
            ]
        )
        mask_keep_idx = []
        mask_arr = []
        for c_idx in range(self.model_params["n_original_input_dims"]):
            if c_idx not in idx_slices:
                mask_keep_idx.append(c_idx)
            else:
                c_mask = (
                    mask[:, c_idx]
                    .unsqueeze(-1)
                    .repeat(1, self.model_params["embedding_dim"])
                )
                mask_arr.append(c_mask)
        return torch.cat([mask[:, mask_keep_idx]] + mask_arr, -1).to(self.device)

    def __get_reconstruction_loss(self, x, x_reconstruction, feature_mask):
        """Compute reconstruction loss as per guidance in paper"""
        std = torch.std(x.data, dim=0)
        std[std == 0] = 1  # Correct for cases where std_dev along dimension = 0
        return torch.norm(
            (torch.ones_like(x) - feature_mask) * (x_reconstruction - x) / std
        ).sum()

    def __train(
        self,
        train_generator,
        val_generator=None,
        epochs=None,
        self_supervised=False,
        step_offset=0,
    ):
        """Internal function to fit a TabNet model to an input dataset, with a
        provided number of epochs. Supports early stopping (configured through
        model training parameters at the top of this file).

        If no validation dataset is provided, no validation metrics will be logged
        and early stopping functionality will not be available.

        Parameters
        ----------
        train_generator : torch.utils.DataLoader generator for the training dataset
        val_generator : torch.utils.DataLoader generator for the validation dataset
        epochs : the maximum number of training epochs
        self_supervised : flag to use a self-supervision training objective, as
        opposed to a prediction training objective
        step_offset : the step offset to start logging from

        Returns
        -------
        step : The number of steps the model was trained for
        """
        # Define optimizer
        optimizer = torch.optim.Adam(
            self.model.parameters(),
            lr=self.train_params["learning_rate"],
            weight_decay=self.train_params["weight_decay"],
        )
        lambda_scale = lambda epoch: self.train_params["learning_rate_decay_factor"]
        scheduler = torch.optim.lr_scheduler.MultiplicativeLR(
            optimizer, lr_lambda=lambda_scale
        )

        # Define predictive criterion
        criterion = None
        if self.model_params["discrete_outputs"]:
            criterion = torch.nn.CrossEntropyLoss()
        else:
            criterion = torch.nn.MSELoss()

        # Define logging criterion
        print_stub_name = "Predictive"
        if self_supervised:
            print_stub_name = "Self-supervised"

        # Define early stopping
        es_tracker = None
        criterion_val_loss = None
        if self.train_params["early_stopping"]:
            es_tracker = utils.EarlyStopping(
                min_delta=self.train_params["early_stopping_min_delta_pct"],
                patience=self.train_params["early_stopping_patience"],
                percentage=True,
            )

        # Running max
        model_max_state_dict = None
        model_max_criteria = torch.tensor(float("NaN")).to(self.device)

        # Training
        step = step_offset
        for c_epoch in range(epochs):
            loss_avg = []
            for batch_idx, (x_batch_cont, x_batch_cat, y_batch) in enumerate(
                train_generator
            ):
                step += 1

                criterion_loss, reconstruction_loss, sparsity_loss, masks = (
                    None,
                    None,
                    None,
                    None,
                )

                ones_mask = self.__generate_model_mask(0, x_batch_cont.size()[0])
                if self_supervised:
                    self_supervised_mask = self.__generate_model_mask(
                        self.train_params["p_mask"], x_batch_cont.size()[0]
                    )

                    x_embedded, y_pred_logits, x_reconstruct_batch, masks = self.model(
                        x_batch_cont, x_batch_cat, self_supervised_mask, mask_input=True
                    )

                    reconstruction_loss = self.__get_reconstruction_loss(
                        x_embedded, x_reconstruct_batch, self_supervised_mask
                    )

                else:
                    x_embedded, y_pred_logits, x_reconstruct_batch, masks = self.model(
                        x_batch_cont, x_batch_cat, ones_mask, mask_input=False
                    )

                    # Define criterion loss
                    if self.model_params["discrete_outputs"]:
                        y_batch = torch.squeeze(y_batch)
                    criterion_loss = criterion(y_pred_logits, y_batch)

                # Define sparsity loss
                sparsity_loss = (
                    -1
                    * torch.stack(
                        [
                            (c_mask * torch.log(c_mask + self.train_params["epsilon"]))
                            .sum(dim=-1)
                            .mean()
                            for c_mask in masks
                        ]
                    ).mean()
                )
                if reconstruction_loss:
                    self.tensorboard_writer.add_scalar(
                        "{}/Reconstruction loss".format(print_stub_name),
                        reconstruction_loss,
                        step,
                    )
                if criterion_loss:
                    self.tensorboard_writer.add_scalar(
                        "{}/Criterion loss".format(print_stub_name),
                        criterion_loss,
                        step,
                    )
                self.tensorboard_writer.add_scalar(
                    "{}/Sparsity loss".format(print_stub_name),
                    sparsity_loss,
                    step,
                )

                loss = self.train_params["sparsity_regularization"] * sparsity_loss
                if self_supervised:
                    loss += reconstruction_loss  # Only optimise reconstruction in self-supervised regime
                else:
                    loss += criterion_loss  # Only optimise target criterion in prediction phase

                # Store in loss buffer
                loss_avg.append(loss)

                self.tensorboard_writer.add_scalar(
                    "{}/Total loss".format(print_stub_name),
                    loss,
                    step,
                )

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                if step % self.train_params["learning_rate_decay_step_rate"] == 0:
                    scheduler.step()
                    print(
                        "Decaying learning rate. Revised learning rate: {}".format(
                            scheduler.get_last_lr()
                        )
                    )

            if (c_epoch > 0) and (
                c_epoch % self.train_params["epoch_save_frequency"] == 0
            ):
                self.__save_model("epoch_{}".format(c_epoch), self_supervised)

            if val_generator is not None:
                if self_supervised:
                    reconstruction_val_loss = self.__validation_reconstruct_loss(
                        val_generator
                    )
                    self.tensorboard_writer.add_scalar(
                        "{}/Validation reconstruction loss".format(print_stub_name),
                        reconstruction_val_loss,
                        step,
                    )
                    print(
                        "{} - Epoch: {}, Step: {}, Loss: {}, {}: {}".format(
                            print_stub_name,
                            c_epoch + 1,
                            step,
                            torch.mean(torch.stack(loss_avg)),
                            "Validation reconstruction loss",
                            np.round(reconstruction_val_loss.item(), 4),
                        )
                    )

                    if reconstruction_val_loss == torch.min(
                        reconstruction_val_loss, model_max_criteria
                    ) or torch.isnan(model_max_criteria):
                        model_max_state_dict = self.model.state_dict()
                        model_max_criteria = reconstruction_val_loss

                    if self.train_params["early_stopping"] and es_tracker.step(
                        reconstruction_val_loss
                    ):
                        self.model.load_state_dict(model_max_state_dict)
                        print(
                            "Early stopping criterion met - ending training, and using best weights (val reconstruction loss: {})".format(
                                np.round(model_max_criteria.item(), 4)
                            )
                        )
                        break
                else:
                    y_val_pred, y_val_pred_logits, y_val = self.__validation_predict(
                        val_generator
                    )
                    metric_value, metric_name = None, None
                    if self.model_params["discrete_outputs"]:
                        metric_name = "Validation accuracy"
                        metric_value = torch.true_divide(
                            (y_val_pred == y_val).sum(), y_val.size(0)
                        )
                    else:
                        metric_name = "Validation MSE"
                        metric_value = torch.square(y_val_pred - y_val).mean()

                    criterion_val_loss = criterion(y_val_pred_logits, y_val.squeeze())

                    self.tensorboard_writer.add_scalar(
                        "{}/{}".format(print_stub_name, metric_name),
                        metric_value,
                        step,
                    )
                    self.tensorboard_writer.add_scalar(
                        "{}/Validation criterion loss".format(print_stub_name),
                        criterion_val_loss,
                        step,
                    )
                    print(
                        "{} - Epoch: {}, Step: {}, Total train loss: {}, {}: {}, {}: {}".format(
                            print_stub_name,
                            c_epoch + 1,
                            step,
                            np.round(torch.mean(torch.stack(loss_avg)).item(), 4),
                            "Validation criterion loss",
                            np.round(criterion_val_loss.item(), 4),
                            metric_name,
                            np.round(metric_value.item(), 4),
                        )
                    )
                    if criterion_val_loss == torch.min(
                        criterion_val_loss, model_max_criteria
                    ) or torch.isnan(model_max_criteria):
                        model_max_state_dict = self.model.state_dict()
                        model_max_criteria = criterion_val_loss

                    if self.train_params["early_stopping"] and es_tracker.step(
                        criterion_val_loss
                    ):
                        self.model.load_state_dict(model_max_state_dict)
                        print(
                            "Early stopping criterion met - ending training, and using best weights (val criterion loss: {})".format(
                                np.round(model_max_criteria.item(), 4)
                            )
                        )
                        break
            else:
                print(
                    "{} - Epoch: {}, Step: {}, Total train loss: {}".format(
                        print_stub_name,
                        c_epoch + 1,
                        step,
                        np.round(torch.mean(torch.stack(loss_avg)).item(), 4),
                    )
                )
        self.__save_model("final", self_supervised)
        return step

    def __validation_reconstruct_loss(self, generator):
        """Returns model reconstruction loss for an input dataset.

        Parameters
        ----------
        generator : torch.utils.DataLoader generator for the dataset

        Returns
        -------
        loss : Reconstruction loss across the full input dataset
        """
        self.model.eval()
        with torch.no_grad():
            out_loss = []
            for batch_idx, (x_batch_cont, x_batch_cat, y_batch) in enumerate(generator):
                ones_mask = self.__generate_model_mask(0, x_batch_cont.size()[0])
                self_supervised_mask = self.__generate_model_mask(
                    self.train_params["p_mask"], x_batch_cont.size()[0]
                )
                x_embedded, y_pred_logits, x_reconstruct_batch, masks = self.model(
                    x_batch_cont, x_batch_cat, self_supervised_mask, mask_input=True
                )

                out_loss.append(
                    self.__get_reconstruction_loss(
                        x_embedded, x_reconstruct_batch, self_supervised_mask
                    )
                )

        self.model.train()
        return torch.stack(out_loss).mean()

    def __validation_predict(self, generator):
        """Returns prediction probabilities, logits, and targets for an input
        dataset. Typically used to produce predictions for the validation set.

        Parameters
        ----------
        generator : torch.utils.DataLoader generator for the dataset

        Returns
        -------
        y_pred_agg : Model predictions probabilities for the input dataset
        y_pred_logits_agg: Model prediction logits for the input dataset
        """
        self.model.eval()
        with torch.no_grad():
            out_y = []
            out_y_pred_logits = []
            out_y_pred = []
            for batch_idx, (x_batch_cont, x_batch_cat, y_batch) in enumerate(generator):
                ones_mask = self.__generate_model_mask(0, x_batch_cont.size()[0])
                x_embedded, y_pred_logits, x_reconstruct_batch, masks = self.model(
                    x_batch_cont, x_batch_cat, ones_mask
                )
                if self.model_params["discrete_outputs"]:
                    y_val_pred = torch.argmax(
                        nn.functional.softmax(y_pred_logits, dim=-1), dim=-1
                    )
                else:
                    y_val_pred = y_pred_logits.squeeze()
                out_y_pred_logits.append(y_pred_logits)
                out_y_pred.append(y_val_pred)
                out_y.append(y_batch)
            y_pred_agg = torch.cat(out_y_pred, dim=0).squeeze()
            y_pred_logits_agg = torch.cat(out_y_pred_logits, dim=0).squeeze()
            y_agg = torch.cat(out_y, dim=0).squeeze()
        self.model.train()
        return (y_pred_agg, y_pred_logits_agg, y_agg)

    def fit(
        self, X_train, y_train, X_val=None, y_val=None, train_params={}, save_params={}
    ):
        """Prepares and runs full train loop, for a given set of training and
        saving parameters. If a validation set is not provided, validation
        metrics will not be provided during training, and early stopping
        functionality will not be available.

        Parameters
        ----------
        X_train : Pandas / Numpy array of training features
        y_train : Pandas / Numpy array of training targets
        X_val : (Optional) Pandas / Numpy array of validation features
        y_val : (Optional) Pandas / Numpy array of validation targets
        train_params : A dictionary of training parameters
        save_params: A dictionary of save params

        Returns
        -------
        None
        """
        self.save_params = self.default_save_params
        self.save_params.update(save_params)

        self.train_params = self.default_train_params
        self.train_params.update(train_params)

        if self.model is None and self.model_params["discrete_outputs"]:
            self.model_params[
                "discrete_target_mapping"
            ] = utils.generate_categorical_to_ordinal_map(y_train)

        if (type(X_train) is not type(X_val)) or (type(y_train) is not type(y_val)):
            raise ValueError("Training and validation datasets should have same type")

        data_columns = None
        if isinstance(X_train, pd.DataFrame):
            data_columns = dict(zip(X_train.columns, range(X_train.shape[1])))
        else:
            data_columns = dict(zip(range(X_train.shape[1]), range(X_train.shape[1])))
        if (
            (X_val is not None)
            and isinstance(X_val, pd.DataFrame)
            and not (X_val.columns == X_train.columns).all()
        ):
            raise ValueError("X_train and X_val have differing columns!")
        if (
            (X_val is not None)
            and isinstance(X_val, np.ndarray)
            and X_val.shape[1] != X_train.shape[1]
        ):
            raise ValueError(
                "Training and validation datasets have differing number of columns"
            )

        # Configure categorical variables
        if len(self.model_params["categorical_variables"]) > 0:
            config_dict = {}
            if isinstance(X_train, np.ndarray):
                # TODO: See below
                # Check whether categorical_variables is an int
                # Check whether max/min falls within range of 0, dims - 1
                pass
            if isinstance(X_train, pd.DataFrame):
                # TODO: See below
                # If string: check whether categorical variables is missing from training dataset
                # If int: check whether min/max falls within range of 0, dims - 1
                pass

            for col in self.model_params["categorical_variables"]:
                cast_map = utils.generate_categorical_to_ordinal_map(X_train[col])
                config_dict[col] = {
                    "map": cast_map,
                    "n_dims": len(cast_map.keys()),
                    "idx": data_columns[col],
                    "identifier": col,
                }
            self.model_params["categorical_config"] = config_dict

        # Cast DataFrames to numpy arrays
        if isinstance(X_train, pd.DataFrame):
            X_train = X_train.values
        if isinstance(y_train, pd.DataFrame) or isinstance(y_train, pd.Series):
            y_train = y_train.values
        if isinstance(X_val, pd.DataFrame):
            X_val = X_val.values
        if isinstance(y_val, pd.DataFrame) or isinstance(y_val, pd.Series):
            y_val = y_val.values

        # Build generators for train / validation
        train_data = utils.TrainingDataset(
            X_train,
            y_train,
            self.model_params["discrete_target_mapping"],
            self.model_params["categorical_config"],
            columns=data_columns,
            device=self.device,
        )
        train_generator = torch.utils.data.DataLoader(
            train_data,
            **{
                "batch_size": self.train_params["batch_size"],
                "shuffle": self.train_params["train_generator_shuffle"],
                "num_workers": self.train_params["train_generator_n_workers"],
            }
        )
        val_data = None
        val_generator = None
        if X_val is not None:
            val_data = utils.TrainingDataset(
                X_val,
                y_val,
                output_mapping=self.model_params["discrete_target_mapping"],
                categorical_mapping=self.model_params["categorical_config"],
                columns=data_columns,
                device=self.device,
            )
            val_generator = torch.utils.data.DataLoader(
                val_data,
                **{
                    "batch_size": self.train_params["validation_batch_size"],
                    "shuffle": False,
                }
            )

        # Adjust train_data input dims based on categorical embeddings
        n_input_dims = X_train.shape[1] + (
            len(self.model_params["categorical_config"])
            * (self.model_params["embedding_dim"] - 1)
        )
        n_continuous_dims = X_train.shape[1] - len(
            self.model_params["categorical_config"]
        )

        # Update with correct dimensions
        if self.model is None:
            self.model_params.update(
                {
                    "n_input_dims": n_input_dims,
                    "n_original_input_dims": X_train.shape[1],
                    "n_continuous_input_dims": n_continuous_dims,
                    "n_output_dims": train_data.n_output_dims,
                    "column_index_map": data_columns,
                }
            )
            self.model = TabNetModel(**self.model_params)
            self.model.to(self.device)

        X_test_batch_cont, X_test_batch_cat, _ = train_data.random_batch(
            self.train_params["batch_size"]
        )
        print("Setting-up Tensorboard instance...")
        self.tensorboard_writer = self.__get_tensorboard_writer(
            X_test_batch_cont, X_test_batch_cat
        )
        self.model.train()  # Enable training mode
        print("Starting training...")

        if (
            self.train_params["run_self_supervised_training"] == False
            and self.train_params["run_supervised_training"] == False
        ):
            raise ValueError(
                "No training scheme defined: set `run_self_supervised_training` or `run_supervised_training` to True"
            )

        step = 0
        if self.train_params["run_self_supervised_training"]:
            print("Training model with self-supervision objective")
            step = self.__train(
                train_generator=train_generator,
                val_generator=val_generator,
                epochs=self.train_params["max_epochs_self_supervised"],
                self_supervised=True,
                step_offset=step,
            )
        if self.train_params["run_supervised_training"]:
            print("Training model with predictive objective")
            step = self.__train(
                train_generator=train_generator,
                val_generator=val_generator,
                epochs=self.train_params["max_epochs_supervised"],
                self_supervised=False,
                step_offset=step,
            )

        if self.tensorboard_writer is not None:
            self.tensorboard_writer.flush()
            self.tensorboard_writer.close()

    def __predict(self, X, batch_size=1024):
        """Internal method to produce predictions for provided input and
        batch_size features."""
        if not self.model:
            raise ValueError(
                "Model not yet initialized. Run `train` to fit a model to an input dataset"
            )

        self.model.eval()
        data_columns = None
        if isinstance(X, pd.DataFrame):
            data_columns = dict(zip(X.columns, range(X.shape[1])))
        else:
            data_columns = dict(zip(range(X.shape[1]), range(X.shape[1])))
        if isinstance(X, pd.DataFrame):
            X = X.values
        pred_data = utils.InferenceDataset(
            X,
            self.model_params["categorical_config"],
            columns=data_columns,
            device=self.device,
        )
        with torch.no_grad():

            pred_generator = torch.utils.data.DataLoader(
                pred_data, **{"batch_size": batch_size, "shuffle": False}
            )

            out_y_pred = []
            for batch_idx, (x_batch_cont, x_batch_cat) in enumerate(pred_generator):
                ones_mask = self.__generate_model_mask(0, x_batch_cont.size()[0])
                y_val_pred = None
                x_embedded, y_pred_logits, x_reconstruct_batch, masks = self.model(
                    x_batch_cont, x_batch_cat, ones_mask
                )
                if self.model_params["discrete_outputs"]:
                    y_val_pred = nn.functional.softmax(y_pred_logits, dim=-1)
                else:
                    y_val_pred = y_pred_logits.squeeze()
                out_y_pred.append(y_val_pred)
            y_pred_agg = torch.cat(out_y_pred, dim=0).squeeze()
        self.model.train()  # Enable training mode
        return y_pred_agg

    def predict_proba(self, X, batch_size=1024):
        """Uses model to predict for a given set of input features.

        Returns probability distribution for classification tasks.
        Raises ValueError if `discrete_outputs` is set to False.

        Parameters
        ----------
        X : Numpy array of features to predict on

        Returns
        -------
        y : Pandas array of predictions for input features, with column
        names corresponding to associated class names

        Raises
        ------
        ValueError: If model is not yet trained, or model is not a Classifier
        """
        if not self.model_params["discrete_outputs"]:
            raise ValueError("`Predict_proba` not available for regression models")
        else:
            ret_data = pd.DataFrame(self.__predict(X, batch_size=batch_size).numpy())
            ret_data.columns = utils.map_ordinals_to_categoricals(
                np.array(ret_data.columns), self.model_params["discrete_target_mapping"]
            )
            return ret_data

    def predict(self, X, batch_size=1024):
        """Uses model to predict for a given set of input features.

        Returns probability distribution for classification tasks.

        Parameters
        ----------
        X : Numpy array of features to predict on

        Returns
        -------
        y : List of predictions for input features
        """
        if self.model_params["discrete_outputs"]:
            return utils.map_ordinals_to_categoricals(
                torch.argmax(self.__predict(X, batch_size=batch_size), dim=-1),
                self.model_params["discrete_target_mapping"],
            )
        else:
            return self.__predict(X, batch_size=batch_size).numpy()


text:
import unittest

import numpy as np

from quantiphyse.test.widget_test import WidgetTest

from .widget import PcaWidget

NUM_PCA = 5
NAME = "pca"

class PcaWidgetTest(WidgetTest):

    def widget_class(self):
        return PcaWidget

    def testNoData(self):
        if self.w._run.runBtn.isEnabled():
            self.w._run.runBtn.clicked.emit()
        self.assertFalse(self.error)

    def test3dData(self):
        """ 4D data required """
        self.ivm.add(self.data_3d, grid=self.grid, name="data_3d")
        self.assertFalse(bool(self.w._options.option("data").value))

    def test4dData(self):
        """ 4d clustering """
        self.ivm.add(self.data_4d, grid=self.grid, name="data_4d")
        self.processEvents()
        self.w._options.option("data").value = "data_4d"
        self.processEvents()
        self.assertTrue(self.w._run.runBtn.isEnabled())
        self.assertEquals(self.w._options.option("output-name").value, "data_4d_pca")

        self.w._options.option("n-components").value = NUM_PCA
        self.w._options.option("output-name").value = NAME
        self.w._run.runBtn.clicked.emit()
        self.processEvents()

        for mode in range(NUM_PCA):
            self.assertTrue("%s%i" % (NAME, mode) in self.ivm.data)
        self.assertFalse("%s%i" % (NAME, NUM_PCA+1) in self.ivm.data)

        self.assertEquals(self.ivm.current_data.name, "%s0" % NAME)
        self.assertFalse(self.error)

    def test4dDataWithRoi(self):
        """ 4d clustering within an ROI"""
        self.ivm.add(self.data_4d, grid=self.grid, name="data_4d")
        self.ivm.add(self.mask, grid=self.grid, roi=True, name="mask")
        self.w._options.option("data").value = "data_4d"
        self.w._options.option("roi").value = "mask"

        self.assertTrue(self.w._run.runBtn.isEnabled())
        self.assertEquals(self.w._options.option("output-name").value, "data_4d_pca")

        self.w._options.option("n-components").value = NUM_PCA
        self.w._options.option("output-name").value = NAME
        self.w._run.runBtn.clicked.emit()
        self.processEvents()

        for mode in range(NUM_PCA):
            name = "%s%i" % (NAME, mode)
            self.assertTrue(name in self.ivm.data)
            # Data value is always zero outside the ROI
            data = self.ivm.data[name].raw()
            self.assertTrue(np.all(data[self.mask == 0] == 0))
            self.assertFalse(self.error)

        self.assertFalse("%s%i" % (NAME, NUM_PCA+1) in self.ivm.data)
        self.assertEquals(self.ivm.current_data.name, "%s0" % NAME)
        self.assertFalse(self.error)

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

ext: py

max_stars_repo_name: physimals/quantiphyse

max_issues_repo_name: ibme-qubic/quantiphyse

max_forks_repo_name: physimals/quantiphyse

reason: import numpy


text:
import numpy as np

def simulate_RR_uniform_flavor(policy, T=3600):
    # policy: a list of choice probability, same policy for all restaurants
    # reward function: currently uniform across time and all restaurants
    policies = np.tile(policy, 4).reshape((4, 4), order='C')
    reward = 0
    probs = [0, 0.2, 0.8, 1]
    wait_time = [7, 5, 3, 1]
    reward_time = 2
    steps = 0
    accepts = []
    while T < 3600:
        tone = np.random.randint(4)
        restaurant = steps % 4
        # choice
        if np.random.random() < policies[restaurant][tone]:
            # accept, outcome
            T -= wait_time[tone]
            accepts.append(steps)
            if np.random.random() < probs[tone]:
                reward += 1
                T -= reward_time
        steps += 1
    return reward

def policy_optimize(T, k_reps=10):
    policy = [0.5] * 4
    objective = lambda plc: sum(simulate_RR_uniform_flavor(plc, T) for i in range(k_reps)) / k_reps
    # scipy optimize objective function with the argument policy
    optimal = None
    return optimal


text:
from mlspeclib import MLSchema, MLObject
from random import randint, random

results_ml_object = MLObject()

results_ml_object.set_type(
    schema_type=result_ml_object_schema_type,
    schema_version=result_ml_object_schema_version,
)

# Execute code below. Examples:

"""
1) Execute on the worker node this step is running on

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn import metrics

dataset = pd.read_csv('Weather.csv') # Example dataset

X = dataset['MinTemp'].values.reshape(-1,1)
y = dataset['MaxTemp'].values.reshape(-1,1)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)
regressor = LinearRegression()
regressor.fit(X_train, y_train) # Training the algorithm

y_pred = regressor.predict(X_test)
print('Mean Absolute Error:', metrics.mean_absolute_error(y_test, y_pred))
print('Mean Squared Error:', metrics.mean_squared_error(y_test, y_pred))
print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test, y_pred)))

"""

"""
2) Execute a long running job on an external service

from pyspark import SparkContext
sc = SparkContext("local", "Log Processing App")
numAs = logData.filter(lambda s: 'a' in s).count()
numBs = logData.filter(lambda s: 'b' in s).count()
print "Lines with a: %i, lines with b: %i" % (numAs, numBs)
"""

"""
3) Execute a job on a hosted service

from azureml.core import Workspace, Experiment
from azureml.core.authentication import ServicePrincipalAuthentication

azure_credentials = os.environ.get("INPUT_AZURE_CREDENTIALS")
sp_auth = ServicePrincipalAuthentication(
    tenant_id=azure_credentials.get("tenantId"),
    service_principal_id=azure_credentials.get("clientId"),
    service_principal_password=azure_credentials.get("clientSecret",)
)
ws = Workspace.from_config(
    path=config_file_path,
    _file_name=config_file_name,
    auth=sp_auth
)

experiment = Experiment(workspace=ws, "this_experiment_name")
run_config = load_runconfig_yaml(runconfig_yaml_file="run_config.yml")
run = experiment.submit(config=run_config, tags=parameters.get("tags", {}))
"""

# Execution metrics
results_ml_object.execution_profile.system_memory_utilization = random()
results_ml_object.execution_profile.network_traffic_in_bytes = randint(7e9, 9e10)
results_ml_object.execution_profile.gpu_temperature = randint(70, 130)
results_ml_object.execution_profile.disk_io_utilization = random()
results_ml_object.execution_profile.gpu_percent_of_time_accessing_memory = random()
results_ml_object.execution_profile.cpu_utilization = random()
results_ml_object.execution_profile.gpu_utilization = random()
results_ml_object.execution_profile.gpu_memory_allocation = random()

max_stars_repo_name: mlspec/mlspeclib-action-docker

max_issues_repo_name: mlspec/mlspeclib-action-docker

max_forks_repo_name: mlspec/mlspeclib-action-docker


text:
# %% Import

import numpy as np
import glob
import pickle

# %%

root = 'data/'

scenarios_dict = {
    'standard':[
        'uniaxial',
        'biaxial',
        'planar',
        'volumetric',
        'shear'
        ],
    # 'test':['test1', 'test2', 'test3'],
    'advanced': [
        'biaxial2',
        'biaxial3',
        'uniaxial-shear'
        ]
    }
scenarios_dict['all'] = scenarios_dict['standard'] + scenarios_dict['advanced']

# %%

def get_p(folder):
    with open(f'{folder}meta.pkl', 'rb') as file:
        meta = pickle.load(file)
    p = [meta['E']] + meta['t']
    return p

def get_ps_folders():
    folders = glob.glob(root+'*/')
    ps = np.zeros([len(folders), 4])
    for i in range(len(folders)):
        with open(f'{folders[i]}meta.pkl', 'rb') as file:
            meta = pickle.load(file)
        ps[i] = [meta['E']] + meta['t']
    return ps, folders

def get_folder(t=[1, 1, 1], E=1):
    p = np.array([E] + list(t), dtype=np.float32)
    ps, folders = get_ps_folders()
    found = False
    c = 0
    while (not found) and (c < ps.shape[0]):
        p_ = ps[c]
        if np.linalg.norm(p - p_) < 1e-6:
            found = True
            break
        c += 1
    if found:
        return folders[c]
    else:
        raise Exception('Meta not available.')

def get_case(
        t=[1, 1, 1],
        E=1,
        folder='',
        scenario='uniaxial',
        tc='t', #t or c
        ):
    if len(folder) == 0:
        folder = get_folder(E=E, t=t)
        p = np.array([E] + list(t), dtype=np.float32)
    else:
        p = get_p(folder)
    p = np.array(p, dtype=np.float32)
    Fs = []
    ps = []
    Ws = []
    Ps = []
    if tc == 'c':
        case = f'{scenario}_compression'
    else:
        case = f'{scenario}_tension'
    file = glob.glob(f'{folder}*{case}*')[0]
    data = np.loadtxt(file, skiprows=1).astype('float32')
    if tc == 'c':
        data = data = np.flip(data, 0)
    Fs = data[:, 1:1+9].reshape([-1, 3, 3])
    ps = np.ones([data.shape[0], 4], dtype='float32')*p
    Ws = 1e6*data[:, 21].reshape([-1, 1])
    Ps = 1e6*data[:, 11:11+9].reshape([-1, 3, 3])
    return Fs, ps, Ws, Ps

def get_all(
        concat_all=False
        ):
    cases = []
    for scenario in scenarios_dict['all']:
        cases.append([scenario, 't'])
        cases.append([scenario, 'c'])
    _, folders = get_ps_folders()
    Fs = np.zeros([len(folders), len(cases), 151, 3, 3], dtype=np.float32)
    ps = np.zeros([len(folders), len(cases), 151, 4], dtype=np.float32)
    Ws = np.zeros([len(folders), len(cases), 151, 1], dtype=np.float32)
    Ps = np.zeros([len(folders), len(cases), 151, 3, 3], dtype=np.float32)
    for i_f in range(len(folders)):
        for i_c in range(len(cases)):
            Fs[i_f, i_c], ps[i_f, i_c], Ws[i_f, i_c], Ps[i_f, i_c] = get_case(
                folder=folders[i_f],
                scenario=cases[i_c][0],
                tc=cases[i_c][1]
                )
    if concat_all:
        Fs = np.reshape(Fs, [-1, 3, 3])
        ps = np.reshape(ps, [-1, 4])
        Ws = np.reshape(Ws, [-1, 1])
        Ps = np.reshape(Ps, [-1, 3, 3])
    return Fs, ps, Ws, Ps

max_stars_repo_name: CPShub/sim-data

max_issues_repo_name: CPShub/sim-data

max_forks_repo_name: CPShub/sim-data


text:
# Calling syntax: nodectrl.py [MQTT server address]
from __future__ import division, print_function
import numpy as np
import sys
import csv
from obnpy.obnnode import *

class Controller(OBNNode):
    """Implements a controller for a DC motor."""
    def __init__(self, workspace, server):
        OBNNode.__init__(self, "ctrl", workspace, server)

        self.create_input("v", "scalar", "double")  # Create velocity input
        self.create_input("sp", "scalar", "double") # Create setpoint input
        self.create_output("u", "scalar", "double") # Create command output

        MAINBLOCK = 0                   # There is only one block in this node
        self.on_init(self.initCallback) # assign callback to initialize the node
        self.on_term(self.termCallback) # simple callback for node's termination
        self.on_block_output(self.ctrlOutput, MAINBLOCK) # callback to send output
        self.on_block_state(self.updateState, MAINBLOCK) # callback to update the controller's states

        # Initialize the controller's parameters
        self.Kp = 46.0
        self.Ki = 1.5
        self.Kd = 0.9

    def initCallback(self):
        # Initialize the controller's states
        self.integrator = 0.0
        self.derivator = 0.0
        self.command = 0.0
        self.output_ports["u"].set(self.command) # Initialize the output (not necessary)
        self.dumpfile = open("controller.txt", "wb")
        self.dump = csv.writer(self.dumpfile, delimiter="\t")
        print("At {} simulation started.".format(self.sim_time()))

    def termCallback(self):
        self.dumpfile.close()
        print("Controller node terminated.")

    def ctrlOutput(self):
        self.sp = self.input_ports["sp"].get() # Get the setpoint
        self.v = self.input_ports["v"].get() # Get the velocity
        error = self.sp - self.v
        P_val = self.Kp * error
        D_val = self.Kd * (error - self.derivator)
        self.derivator = error
        self.integrator += error
        I_val = self.Ki * self.integrator
        self.command = P_val + I_val + D_val
        self.output_ports["u"].set(self.command)

    def updateState(self):
        # At this point, the inputs are all up-to-date, so we can dump log data
        self.dump.writerow([self.sim_time(), self.sp, self.v, self.command, self.integrator])

def main():
    if len(sys.argv) < 2:
        server = 'tcp://localhost:1883'
    else:
        server = sys.argv[1]

    node = Controller("test2", server)
    print("Ready to run the controller node; please start all other nodes ...")
    status = node.run(60)
    print("Simulation stopped with status = {}".format(status))

if __name__ == '__main__':
    main()


text:
import numba

def i4_0d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.int32)
    d.update(kvs)
    return d

def i4_1d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.int32[:])
    d.update(kvs)
    return d

def i8_0d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.int64)
    d.update(kvs)
    return d

def i8_1d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.int64[:])
    d.update(kvs)
    return d

def f4_0d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.float32)
    d.update(kvs)
    return d

def f4_1d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.float32[:])
    d.update(kvs)
    return d

def f8_0d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.float64)
    d.update(kvs)
    return d

def f8_1d_dict(**kvs):
    d = numba.typed.Dict.empty(key_type=numba.types.unicode_type,
                               value_type=numba.types.float64[:])
    d.update(kvs)
    return d

i4_dict = i4_0d_dict
i8_dict = i8_0d_dict
f4_dict = f4_0d_dict
f8_dict = f8_0d_dict

reason: import numba


text:
"""
Put script documenation here.
"""

import logging
log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())

from pymeasure.instruments.rohdeschwarz import RohdeSMB100A, RohdeFSQ
from pymeasure.instruments.agilent import Agilent4156, AgilentU2040X
from pymeasure.log import console_log, file_log

import numpy as np
import pandas as pd

class HarmonicTest(object):
    """
    Class to define all the calibration and measurement sequences to measure harmonics and power handling in RF switch and amplifier devices.
    """

    def __init__(self, sweep_values, **kwargs):
        self.smu = Agilent4156(
            "GPIB0::25", read_termination='\n', write_termination='\n', timeout=None)
        log.info('Connected to {}'.format(self.smu.id))

        self.pm = AgilentU2040X(
            'USB0::0x2A8D::0x1E01::MY56360005::0::INSTR', timeout=None)
        log.info('Connected to {}'.format(self.pm.id))

        self.sig = RohdeSMB100A(
            "GPIB0::29", read_termination='\n', write_termination='\n', timeout=None)
        log.info('Connected to {}'.format(self.sig.id))

        self.sa = RohdeFSQ("GPIB0::20", read_termination='\n',
                           write_termination='\n', timeout=None)
        log.info('Connected to {}'.format(self.sa.id))

        self.DEV_TYPE = 'SWITCH'    # or AMPLIFIER
        self.FUND_FREQ = 983  # MHz
        self.CAL_POWER = -50  # dBm
        self.RES_BW = 100  # Hz
        self.FREQ_SPAN = 1  # kHz
        self.SWP_AVG = 10
        self.SWEEP_VALUES = sweep_values
        self.NHARMONICS = 3

    def setup_instruments(self):
        """ Apply instrument settings """

        # signal generator
        self.sig.reset()
        self.sig.freq_unit = 'MHz'
        self.sig.power_unit = 'DBM'
        self.sig.fixed_freq = self.FUND_FREQ
        self.sig.power_level = self.CAL_POWER
        log.info('Finished setting up signal generator.')

        # signal analyzer
        self.sa.reset()
        self.sa.freq_unit = 'MHz'
        self.sa.power_unit = 'DBM'
        self.sa.center_freq = self.FUND_FREQ
        self.sa.freq_unit = 'kHz'
        self.sa.freq_span = self.FREQ_SPAN
        self.sa.freq_unit = 'Hz'
        self.sa.res_bw = self.RES_BW
        self.sa.video_bw = self.RES_BW
        self.sa.sweep_count = self.SWP_AVG
        self.sa.continuous_mode = 'OFF'
        self.sa.all_markers_off()
        self.sa.freq_counter = 'ON'
        self.freq_unit = 'MHz'
        log.info('Finished setting up signal analyzer.')

        # setup power meter
        self.pm.reset()
        self.pm.freq_unit = 'MHZ'
        self.pm.freq = self.FUND_FREQ
        self.pm.continuous_mode = 'OFF'
        log.info('Finished setting up power meter.')

    def input_calibration(self):
        """ Determine input loss offset via power calibration """
        input(
            'Connect power meter to probe end of the input RF cable. Press any Enter to continue...')
        self.sig.output = 'ON'
        self.sig.power_level = self.CAL_POWER
        self.pm.init()
        self.INPUT_OFFSET = self.pm.read
        self.sig.output = 'OFF'

max_stars_repo_name: vik-s/rfscripts

max_issues_repo_name: vik-s/rfscripts

max_forks_repo_name: vik-s/rfscripts


text:
import numpy as np

from .exceptions import LineSearchFailError
from .geodesic import geodesic_move
from .newton import newton_step

def minimize(y0, fun, grEu, hessEu=None, maxit=50, tol=1e-8, ltol=1e-3,
             alpha=0.2, beta=0.5):
    '''
    Minimizing a function over a Gr(n,p), using Newton's method if a
    Hessian is provided, and the steepest descent method otherwise.

    Parameters
    ----------
    y0 : (n, p) ndarray
        Starting point at the Grassmann manifold.
    fun : function
        Function to be minimized. `fun` should take a (n, p) ndarray as
        input and return a scalar.
    grEu : function
        Gradient of `fun` in Euclidean space. `grEu` should take a
        (n, p) ndarray as input and return a (n, p) ndarray.
    hessEu : function
        Hessian of `fun` in Euclidean space. `hessEu` should take a
        (n, p) ndarray as input and return a (np, np) ndarray.
    maxit: integer
        Maximum number of iterations.
    tol : double
      Tolerated value for `(grad f)'B(grad f)` (Newton) or `||grad f||^2`
      (steepest descent) at the optimal point.
    ltol : double
        Smallest eigenvalue of the Hessian that is tolerated. If the
        Hessian has an eigenvalue smaller than this, a multiple of the
        identity matrix is added to the Hessian so that the resulting
        matrix has `ltol` as its smallest eigenvalue.
    alpha : double
        Parameter for line search.
    beta : double
        Parameter for line search.

    Returns
    -------
    Dictionary with items
        value : double
            Value of \code{fun} at the found solution.
        y : list of (n, p) ndarrays
            List with the points on the Grassmann manifold (matrices)
            that has been gone through during the iterations.
        lam2 : list of doubles
            List with the values of `(grad f)'B(grad f)` (Newton) or
            `||grad f||^2` (steepest descent) at the points in `y`.

    Raises
    ------
    LineSearchFailError
        If line search fails.

    Notes
    -----
    The Newton step is `v = Pi(B^-1 grad f)`, where `Pi` is the
    projection onto the tangent plane, `B = Hess(f) + cI`, `Hess(f)` is
    a matrix interpretation of the Hessian of `fun` on `Gr(n,p)`, `c` is
    chosen such that the smallest eigenvalue of `B` is `ltol` and
    `grad f` is the gradient of `fun` along `Gr(n,p)`.

    For the steepest descent method, the corresponding step is `-grad f`,
    i.e. the negative of the gradient along `Gr(n,p)`.

    For both methods, backtracking line search is used.

    References
    ----------
    Edelman, A., Arias, T. A., and Smith S. T. (1998) The geometry of
    algorithms with orthogonality constraints.
    SIAM J. Matrix Anal. Appl., 20(2), 303-353.
    '''
    steepest_descent = hessEu is None

    n = y0.shape[0]
    y = [y0]
    lam2 = -np.ones(maxit)
    yy_next = y0

    for it_main in range(maxit):
        yy = np.copy(yy_next)
        grEuy = grEu(yy)
        gr = np.dot((np.eye(n) - np.dot(yy, yy.T)), grEuy)
        v = -gr if steepest_descent else newton_step(
            yy, fun(yy), grEuy, hessEu(yy), ltol=ltol)
        lam2[it_main] = -np.sum(v*gr)

        if lam2[it_main] < 2*tol:
            break

        r = 1
        it_linesearch = 0
        fyy = fun(yy)
        yy_next = geodesic_move(yy, v, r)

        while fun(yy_next) >= fyy - alpha*r*lam2[it_main]:
            r *= beta
            yy_next = geodesic_move(yy, v, r)
            it_linesearch += 1
            if it_linesearch > 50:
                raise LineSearchFailError

        y.append(yy_next)

    return({'value': fun(yy_next), 'y': y, 'lam2': lam2})

max_stars_repo_name: kjohnsson/grassopt

max_issues_repo_name: kjohnsson/grassopt

max_forks_repo_name: kjohnsson/grassopt


text:
#!/usr/bin/env python

"""optim.py: optimize experiments python program for Doppler shift"""

__author__ = "Chakraborty, S."
__copyright__ = "Copyright 2020, SuperDARN@VT"
__credits__ = []
__license__ = "MIT"
__version__ = "1.0."
__maintainer__ = "Chakraborty, S."
__email__ = "shibaji7@vt.edu"
__status__ = "Research"

import os
import sys
sys.path.append("sd/")
import datetime as dt
import pandas as pd
import argparse
from dateutil import parser as dparser
import numpy as np
from netCDF4 import Dataset
import time

from superdarn import SuperDARN

_I_, _J_, _K_, _L_ = 10, 10, 10, 10

if __name__ == "__main__":
        parser = argparse.ArgumentParser()
        parser.add_argument("-r", "--rad", default="bks", help="Radar code (default bks)")
        parser.add_argument("-ev", "--event", default=dt.datetime(2015,5,5,22,10), help="Start date (default 2015-05-05T22:10)",
                type=dparser.isoparse)
        parser.add_argument("-b", "--bmnum", type=int, default=0, help="Beam num starting from 0 in center and +, - on right-left")
        parser.add_argument("-dra", "--d_ratio", type=float, default=100., help="D region e-denisity peak times (100)")
        parser.add_argument("-ds", "--d_start", type=float, default=10., help="D region start (60.)")
        parser.add_argument("-de", "--d_end", type=float, default=35., help="D region end (80.)")
        parser.add_argument("-drt", "--d_rtime", type=float, default=2., help="D region e-denisity rise time (2.)")
        parser.add_argument("-era", "--e_ratio", type=float, default=1., help="E region e-denisity peak order (1.)")
        parser.add_argument("-fra", "--f_ratio", type=float, default=1.01, help="F region e-denisity peak order")
        parser.add_argument("-fs", "--f_start", type=float, default=130., help="F region start (180.)")
        parser.add_argument("-fe", "--f_end", type=float, default=190, help="F region end (240.)")
        parser.add_argument("-frt", "--f_rtime", type=float, default=0.5, help="F region e-denisity rise time (0.5)")
        parser.add_argument("-rd", "--save_radar", action="store_false", help="Save riometer data (default True)")
        parser.add_argument("-ps", "--plot_summary", action="store_true", help="Plot summary report (default False)")
        parser.add_argument("-plt", "--plot", action="store_false", help="Plot summary report (default True)")
        parser.add_argument("-sr", "--save_result", action="store_false", help="Save results (default True)")
        parser.add_argument("-c", "--clear", action="store_true", help="Clear pervious stored files (default False)")
        parser.add_argument("-v", "--verbose", action="store_false", help="Increase output verbosity (default True)")
        parser.add_argument("-fr", "--frequency", type=float, default=12., help="Frequency of oprrations in MHz (default 12 MHz)")
        parser.add_argument("-mr", "--mrange", type=float, default=1990., help="Max ground range in km (default 1000 km)")
        parser.add_argument("-nmr", "--nmrange", type=int, default=200, help="Number of steps in ground range (default 101)")
        parser.add_argument("-sh", "--sheight", type=float, default=50., help="Start height in km (default 50 km)")
        parser.add_argument("-nmh", "--eheight", type=float, default=350., help="End height in km (default 350 km)")
        parser.add_argument("-hinc", "--hinc", type=float, default=1., help="Step in height in km (default 1 km)")
        parser.add_argument("-es", "--selev", type=float, default=30., help="Start elevation angle deg")
        parser.add_argument("-ee", "--eelev", type=float, default=35., help="End elevation angle deg")
        parser.add_argument("-ei", "--ielev", type=float, default=1., help="Inc of elevation angle deg (default 20*)")
        parser.add_argument("-nhops", "--nhops", type=float, default=1, help="Number of hops (default 1)")
        args = parser.parse_args()
        if args.verbose:
            print("\n Parameter list for simulation ")
            for k in vars(args).keys():
                print("     " , k , "->" , str(vars(args)[k]))
        if args.plot:
            dat = Dataset("config/optim.nc")
            d_ratios = dat.variables["d_ratio"][:]
            d_rtimes = dat.variables["d_rtime"][:]
            f_ratios = dat.variables["f_ratio"][:]
            f_rtimes = dat.variables["f_rtime"][:]
            v_d, v_f = dat.variables["vd"][:], dat.variables["vf"][:]
            u = []
            for i, dr in enumerate(d_ratios):
                for j, dtr in enumerate(d_rtimes):
                    for k, fr in enumerate(f_ratios):
                        for l, ftr in enumerate(f_rtimes):
                            x = i+j+k+l
                            u.append({"dratio": d_ratios[i], "drtime": d_rtimes[j], "fratio": f_ratios[k], "frtime": f_rtimes[l],
                                "vd": v_d[i,j,k,l], "vf": v_f[i,j,k,l]})
            df = pd.DataFrame.from_records(u)
            import statsmodels.api as sm
            from statsmodels.formula.api import ols
            models = {}
            models["m1"] = ols("vf~dratio+fratio+drtime+frtime", data=df).fit()
            models["m2"] = ols("vf~dratio*fratio+drtime+frtime", data=df).fit()
            models["m3"] = ols("vf~dratio+fratio+drtime*frtime", data=df).fit()
            models["m4"] = ols("vf~dratio*fratio*drtime*frtime", data=df).fit()
            tables = {}
            for m in models.keys():
                tables[m] = sm.stats.anova_lm(models[m], typ=2)
                print(tables[m])
        else:
            d_ratios = np.linspace(2., 400., _I_)
            d_rtimes = np.linspace(0.5, 3.5, _J_)
            f_ratios = np.linspace(1., 1.2, _K_)
            f_rtimes = np.linspace(0.5, 3., _L_)
            os.system("rm data/sim/optim.nc")
            rootgrp = Dataset("data/sim/optim.nc", "w", format="NETCDF4")
            rootgrp.description = """Optimization based on 4 parameters."""
            rootgrp.history = "Created " + time.ctime(time.time())
            rootgrp.source = "SuperDARN Sudden Phase Anomaly: Doppler Flash"
            rootgrp.createDimension("ndratio", len(d_ratios))
            rootgrp.createDimension("ndrtime", len(d_rtimes))
            rootgrp.createDimension("nfratio", len(f_ratios))
            rootgrp.createDimension("nfrtime", len(f_rtimes))
            x = rootgrp.createVariable("d_ratio","f4",("ndratio",))
            x.description = "Fractor by which D-region density goes up."
            x[:] = d_ratios
            x = rootgrp.createVariable("f_ratio","f4",("nfratio",))
            x.description = "Fractor by which F-region density goes up."
            x[:] = f_ratios
            x = rootgrp.createVariable("d_rtime","f4",("ndrtime",))
            x.description = "Time taken by the D region e-denisty rise. (in mins)"
            x[:] = d_rtimes
            x = rootgrp.createVariable("f_rtime","f4",("nfrtime",))
            x.description = "Time taken by the F region e-denisty rise. (in mins)"
            x[:] = f_rtimes
            vd, vf = np.zeros((len(d_ratios),len(d_rtimes),len(f_ratios),len(f_rtimes))),\
                    np.zeros((len(d_ratios),len(d_rtimes),len(f_ratios),len(f_rtimes)))
            vd[:], vf[:] = np.nan, np.nan
            for i, dr in enumerate(d_ratios):
                for j, dtr in enumerate(d_rtimes):
                    for k, fr in enumerate(f_ratios):
                        for l, ftr in enumerate(f_rtimes):
                            args.d_ratio, args.d_rtime, args.f_ratio, args.f_rtime = dr, dtr, fr, ftr
                            vd[i,j,k,l], vf[i,j,k,l], _ = SuperDARN(args)._exe_()
            x = rootgrp.createVariable("vd","f8",("ndratio","ndrtime","nfratio","nfrtime"))
            x.description = "Velocity estimated from D-region density change. (in m/s)"
            x[:] = vd
            x = rootgrp.createVariable("vf","f8",("ndratio","ndrtime","nfratio","nfrtime"))
            x.description = "Velocity estimated from F-region density change. (in m/s)"
            x[:] = vf
            print("")
            rootgrp.close()
            os.system("cp data/sim/optim.nc config/")
        if os.path.exists("sd/__pycache__/"):
            os.system("rm -rf sd/__pycache__/")
            os.system("rm -rf py*.log")

reason: import numpy,import statsmodels,from statsmodels


text:
from tflearn.layers.core import input_data, dropout, fully_connected
from tflearn.layers.estimator import regression
from tflearn.optimizers import Adam
from tflearn.layers.normalization import batch_normalization
from tflearn import DNN
from tflearn import activations
import tensorflow as tf
import os

import numpy as np

class Network:
    """
    This is a wrapper class for the neural network structure
    """
    def __init__(self, name, base_name):
        self.base_name = base_name
        self.name = name
        self.model = None
        self.net = None

    #Initializes the network
    def init_self_net(self ,hidden_layers, hidden_neurons,drop=0.8,beta1=0.99,lr=0.001, activation="relu", n_features=12, n_classes=11, regularizer="", regularization_penalty=0.001):
        self.net = input_data(shape=[None, n_features], name='input')
        for i in range(hidden_layers):
            self.net = fully_connected(self.net, hidden_neurons, regularizer=regularizer, activation=activation, weight_decay=regularization_penalty)
            self.net = batch_normalization(self.net)
            self.net = dropout(self.net, drop)

        self.net = fully_connected(self.net, n_classes, activation='softmax')

        # Define the optimizer
        o = Adam(learning_rate=lr, beta1=beta1)
        self.net = regression(self.net, optimizer=o, loss='categorical_crossentropy', name='targets')

        return self.net
    # Trains the netowk and stores all the outputs and logs in the respective directories
    def train(self,X_train,Y_train,X_val,Y_val,epochs=5, batch_size = 64, save=False, file_path="models"):
        if self.net is None:
            raise EmptyNetError("No self.net to be trained!")
        self.model = DNN(self.net, tensorboard_verbose=3, tensorboard_dir="Neural_Network/logs/{}".format(self.base_name))
        self.model.fit({'input': X_train}, {'targets': Y_train}, n_epoch=epochs, validation_set=({'input': X_val}, {'targets': Y_val}),
        snapshot_step=50, show_metric=True, batch_size=batch_size, run_id=self.name)
        if save:
            if not os.path.isdir("Neural_Network/"+file_path+"/"+self.base_name+"/"+self.name):
                os.makedirs("Neural_Network/"+file_path+"/"+self.base_name+"/"+self.name)
            self.save_net("Neural_Network/{}/{}/{}/{}".format(file_path, self.base_name, self.name, self.name))

    def evaluate_model(self, X,Y):
        if X.size == 0 or Y.size == 0:
            raise EmptyDataError("No data to train on!")
        if self.net is None or self.model is None:
            raise EmptyNetError("No self.net to be trained!")
        return self.model.evaluate(X,Y, batch_size=X.shape[0])

    def predict(self,X):
        if X.size == 0:
            raise EmptyDataError("No data to train on!")
        if self.net is None or self.model is None:
            raise EmptyNetError("No self.net to be trained!")
        labels = []
        for feature_set in X:
            labels.append(np.argmax(self.model.predict([feature_set])))
        return np.array(labels)

    def load_net(self, file_path):
        if self.net is None:
            raise EmptyNetError("No self.net to be loaded!")
        self.model = DNN(self.net)
        self.model.load(file_path)

    def save_net(self, file_path):
        if self.net is None or self.model is None:
            raise EmptyNetError("No self.net to be saved!")
        self.model.save(file_path)


text:
#
#
# This example shows the diffraction by a Si 111 crystal calculated in its simplest implementation:
#
#
#    - calculate_simple_diffraction()
#      Uses a crystal setup and calculates the complex transmitivity and reflectivity
#
#
import numpy

from crystalpy.diffraction.GeometryType import BraggDiffraction
from crystalpy.diffraction.DiffractionSetup import DiffractionSetup
from crystalpy.diffraction.DiffractionSetupDabax import DiffractionSetupDabax
from crystalpy.diffraction.Diffraction import Diffraction

from crystalpy.util.Vector import Vector
from crystalpy.util.Photon import Photon
from crystalpy.diffraction.PerfectCrystalDiffraction import PerfectCrystalDiffraction

from dabax.dabax_xraylib import DabaxXraylib

import time

def calculate_simple_diffraction_angular_scan_accelerated():

    # Create a diffraction setup.

    diffraction_setup_dabax = DiffractionSetupDabax(geometry_type          = BraggDiffraction(),  # GeometryType object
                                               crystal_name           = "YB66",                             # string
                                               thickness              = 7e-3,                             # meters
                                               miller_h               = 4,                                # int
                                               miller_k               = 0,                                # int
                                               miller_l               = 0,                                # int
                                               asymmetry_angle        = 0,#10.0*numpy.pi/180.,                              # radians
                                               azimuthal_angle        = 0.0,
                                               dabax=DabaxXraylib())                              # radians

    energy                 = 8040.0                           # eV
    angle_deviation_min    = 20e-6                          # radians
    angle_deviation_max    = 80e-6                           # radians
    angle_deviation_points = 200

    angle_step = (angle_deviation_max-angle_deviation_min)/angle_deviation_points

    #
    # gets Bragg angle needed to create deviation's scan
    #
    bragg_angle = diffraction_setup_dabax.angleBragg(energy)

    print("Bragg angle for E=%f eV is %f deg"%(energy,bragg_angle*180.0/numpy.pi))

    # Create a Diffraction object (the calculator)
    diffraction = Diffraction()
    diffraction_dabax = Diffraction()

    # initialize arrays for storing outputs
    deviations = numpy.zeros(angle_deviation_points)
    intensityS = numpy.zeros(angle_deviation_points)
    intensityP = numpy.zeros(angle_deviation_points)
    intensityS_dabax = numpy.zeros(angle_deviation_points)
    intensityP_dabax = numpy.zeros(angle_deviation_points)

    for ia in range(angle_deviation_points):
        deviation = angle_deviation_min + ia * angle_step
        angle = deviation  + bragg_angle

        # calculate the components of the unitary vector of the incident photon scan
        # Note that diffraction plane is YZ
        yy = numpy.cos(angle)
        zz = - numpy.abs(numpy.sin(angle))
        photon = Photon(energy_in_ev=energy,direction_vector=Vector(0.0,yy,zz))

        # perform the calculation
        coeffs = diffraction.calculateDiffractedComplexAmplitudes(diffraction_setup_dabax, photon)

        # store results
        deviations[ia] = deviation
        intensityS[ia] = coeffs['S'].intensity()
        intensityP[ia] = coeffs['P'].intensity()

    psi_0, psi_H, psi_H_bar = diffraction_setup_dabax.psiAll(energy)

    for ia in range(angle_deviation_points):
        deviation = angle_deviation_min + ia * angle_step
        angle = deviation  + bragg_angle

        # calculate the components of the unitary vector of the incident photon scan
        # Note that diffraction plane is YZ
        yy = numpy.cos(angle)
        zz = - numpy.abs(numpy.sin(angle))
        photon = Photon(energy_in_ev=energy,direction_vector=Vector(0.0,yy,zz))

        # perform the calculation
        # coeffs_dabax = diffraction_dabax.calculateDiffractedComplexAmplitudes(diffraction_setup_dabax, photon)
        #
        # # store results
        # deviations[ia] = deviation
        # intensityS_dabax[ia] = coeffs_dabax['S'].intensity()
        # intensityP_dabax[ia] = coeffs_dabax['P'].intensity()

        # Create PerfectCrystalDiffraction instance.
        perfect_crystal = PerfectCrystalDiffraction(geometry_type=diffraction_setup_dabax.geometryType(),
                                                    bragg_normal=diffraction_setup_dabax.normalBragg(),
                                                    surface_normal=diffraction_setup_dabax.normalSurface(),
                                                    bragg_angle=diffraction_setup_dabax.angleBragg(energy),
                                                    psi_0=psi_0,
                                                    psi_H=psi_H,
                                                    psi_H_bar=psi_H_bar,
                                                    thickness=diffraction_setup_dabax.thickness(),
                                                    d_spacing=diffraction_setup_dabax.dSpacing() * 1e-10)

        complex_amplitudes = perfect_crystal.calculateDiffraction(photon)

        deviations[ia] = deviation
        intensityS_dabax[ia] = complex_amplitudes['S'].intensity()  # 0.0 # coeffs_dabax['S'].intensity()
        intensityP_dabax[ia] = complex_amplitudes['P'].intensity()  # 0.0 # coeffs_dabax['P'].intensity()

    # plot results
    import matplotlib.pylab as plt
    plt.plot(1e6*deviations,intensityS_dabax)
    plt.plot(1e6*deviations,intensityP_dabax)
    plt.xlabel("deviation angle [urad]")
    plt.ylabel("Reflectivity")
    plt.legend(["Sigma-polarization DABAX","Pi-polarization DABAX"])
    plt.show()

#
# main
#
if __name__ == "__main__":

    calculate_simple_diffraction_angular_scan_accelerated()

max_stars_repo_name: oasys-kit/crystalpy

max_issues_repo_name: oasys-kit/crystalpy

max_forks_repo_name: oasys-kit/crystalpy


text:
"""Core data-science utilities."""
import csv
import math
import os
import warnings
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

import matplotlib
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from PIL import Image, ImageChops

from .pathlib import pathify

# Silence warning due to pandas using deprecated matplotlib API.
# See: https://github.com/pandas-dev/pandas/pull/32444
warnings.filterwarnings("ignore", category=matplotlib.cbook.mplDeprecation)

DD_T = Dict[str, Any]

class PyExec(object):
    """Helper class to execute valid Python codes and get all the newly created symbols.

    Typical use-case: to implement a data dictionary that can mix-in Python code to construct certain variables.
    """

    @staticmethod
    def from_file(path: Union[str, Path, os.PathLike], valid_symbols: Iterable[str] = []) -> DD_T:
        """Execute the Python file and export the symbols as a dictionary."""
        with pathify(path).open("r") as f:
            return PyExec.from_str(f.read(), valid_symbols)

    @staticmethod
    def from_str(s, valid_symbols: Iterable[str] = []) -> DD_T:
        """Execute the Python codes and export the symbols as a dictionary."""
        dd_raw: DD_T = {}
        exec(s, dd_raw)
        dd = {k: dd_raw[k] for k in dd_raw}
        return dd

class DFBuilder(object):
    """A helper class to build a dataframe incremental, row-by-row."""

    def __init__(self, dtype=None, columns=None):
        """Construct dataframe row by row.

        Args:
            dtype ([type], optional): see pandas.DataFrame.from_dict. Defaults to None.
            columns ([type], optional): see pandas.DataFrame.from_dict. Defaults to None.
        """
        self.dtype = dtype
        self.columns = columns
        self.rows = []

    def __add__(self, other: Iterable[Any]) -> "DFBuilder":
        """Add a new row to this instance (and modify this instance in-place)."""
        self.rows.append(other)
        return self

    @property
    def df(self):
        """Return the dataframe representation of this instance."""
        return pd.DataFrame.from_dict({i: row for i, row in enumerate(self.rows)}, orient="index", columns=self.columns)

def json_np_converter(o: Union[np.int64, np.float64]) -> Union[int, float]:
    """Convert numpy values to JSON-compliant data type.

    This only works for values. Unfortunately, {json.dump()} still need keys to be strictly Python data types.

    Args:
        o (np.int64 or np.float64): JSON value

    Raises:
        TypeError: if o is not np.int64 or np.float64

    Returns:
        int or float: Python data type that corresponds to o

    Examples:
        >>> import json
        >>> json.dumps(data, default=json_np_converter)
    """
    if isinstance(o, np.int64):
        return int(o)
    if isinstance(o, np.float64):
        return float(o)
    raise TypeError(f"Unknown instance: {type(o)}")

class SimpleMatrixPlotter(object):
    """A simple helper class to fill-in subplot one after another.

    Sample usage using add():

    >>> import pandas as pd
    >>> df = pd.DataFrame({'a': [1,1,1,2,2,2,3,3,3,4,4]})
    >>> gb = df.groupby(by=['a'])
    >>>
    >>> smp = SimpleMatrixPlotter(gb.ngroups)
    >>> for group_name, df_group in gb:
    >>>     ax, _ = smp.add(df_group.plot)
    >>>     assert ax == _
    >>>     ax.set_title(f"Item={group_name}")
    >>> # smp.trim(); plt.show()
    >>> smp.savefig("/tmp/testfigure.png")  # After this, figure & axes are gone.

    Alternative usage using pop():

    >>> smp = SimpleMatrixPlotter(gb.ngroups)
    >>> for group_name, df_group in gb:
    >>>     ax = smp.pop()
    >>>     df_group.plot(ax=ax, title=f"Item={group_name}")
    >>> # smp.trim(); plt.show()
    >>> smp.savefig("/tmp/testfigure.png")  # After this, figure & axes are gone.

    Attributes:
        i (int): Index of the currently free subplot
    """

    def __init__(self, ncols: Optional[int] = None, init_figcount: int = 5, figsize=(6.4, 4.8), dpi=100, **kwargs):
        """Initialize a ``SimpleMatrixPlotter`` instance.

        Args:
            ncols (int, optional): Number of columns. Passing None means to set to sqrt(init_figcount) clipped at
                5 and 20. Defaults to None.
            init_figcount (int, optional): Total number of subplots. Defaults to 5.
            figsize (tuple, optional): size per subplot, see ``figsize`` for matplotlib. Defaults to (6.4, 4.8).
            dpi (int, optional): dot per inch, see ``dpi`` in matplotlib. Defaults to 100.
            kwargs (optional): Keyword argumetns for plt.subplots, but these are ignored and will be overriden:
                ``ncols``, ``nrows``, ``figsize``, ``dpi``.
        """
        # Initialize subplots
        if ncols is None:
            ncols = min(max(5, int(math.sqrt(init_figcount))), 20)
        nrows = init_figcount // ncols + (init_figcount % ncols > 0)

        kwargs = {k: v for k, v in kwargs.items() if k not in {"nrows", "ncols", "figsize", "dpi"}}
        self.fig, _ = plt.subplots(
            nrows=nrows, ncols=ncols, figsize=(figsize[0] * ncols, figsize[1] * nrows), dpi=dpi, **kwargs
        )
        self.axes = self.fig.axes  # Cache list of axes returned by self.fig.axes
        self.fig.subplots_adjust(hspace=0.35)
        self._i = 0  # Index of the current free subplot

        # Warn if initial pixels exceed matplotlib limit.
        pixels = np.ceil(self.fig.get_size_inches() * self.fig.dpi).astype("int")
        if (pixels > 2 ** 16).any():
            warnings.warn(f"Initial figure is {pixels} pixels, and at least one dimension exceeds 65536 pixels.")

    @property
    def i(self):
        """:int: Index of the earliest unused subplot."""
        return self._i

    @property
    def ncols(self):
        """Return the number of columns."""
        return self.axes[0].get_geometry()[1] if len(self.axes) > 0 else 0

    @property
    def nrows(self):
        """Return the number of rows."""
        return self.axes[0].get_geometry()[0] if len(self.axes) > 0 else 0

    @property
    def shape(self):
        """Return a tuple of (rows, cols)."""
        return (self.nrows, self.ncols)

    def add(self, plot_fun, *args, **kwargs) -> Tuple[plt.Axes, Any]:
        """Fill the current free subplot using `plot_fun()`, and set the axes and figure as the current ones.

        Args:
            plot_fun (callable): A function that must accept `ax` keyword argument.

        Returns:
            (plt.Axes, Any): a tuple of (axes, return value of plot_fun).
        """
        ax = self.pop()
        retval = plot_fun(*args, ax=ax, **kwargs)
        return ax, retval

    def pop(self) -> plt.Axes:
        """Get the next axes in this subplot, and set it and its figure as the current axes and figure, respectively.

        Returns:
            plt.Axes: the next axes
        """
        # TODO: extend with new subplots:
        # http://matplotlib.1069221.n5.nabble.com/dynamically-add-subplots-to-figure-td23571.html#a23572
        ax = self.axes[self._i]
        plt.sca(ax)
        plt.figure(self.fig.number)
        self._i += 1
        return ax

    def trim(self):
        """Delete unused subplots."""
        for ax in self.axes[self._i :]:
            self.fig.delaxes(ax)
        self.axes = self.axes[: self._i]

    def savefig(self, *args, **kwargs):
        """Save plotted subplots, then destroy the underlying figure.

        Subsequent operations are undefined and may raise errors.
        """
        self.trim()
        kwargs["bbox_inches"] = "tight"
        self.fig.savefig(*args, **kwargs)
        # Whatever possible ways to release figure
        self.fig.clf()
        plt.close(self.fig)
        del self.fig
        self.fig = None

# NOTE: to simplify the pager implementation, just destroy-old-and-create-new SimpleMatrixPlotter instances.
#       Should it be clear that the overhead of this approach is not acceptable, then reset-and-reuse shall be
#       considered.
#
# As of now, using pager does caps memory usage (in addition to making sure not to hit matplotlib limit of 2^16 pixels
# per figure dimension). The following benchmark to render 10 montages at 100 subplots/montage tops at 392MB RSS, when
# measured on MBP early 2015 model, Mojave 10.14.6, python-3.7.6.
#
# from smallmatter.ds import MontagePager
# import pandas as pd
# mp = MontagePager()
# for i in range(1000):
#     title = f"chart-{i:04d}"
#     pd.Series(rand(6)).plot(ax=mp.pop(title), title=title)
# mp.savefig()
class MontagePager(object):
    """A pager to group and save subplots into multiple montage image files."""

    def __init__(
        self,
        path: Path = Path("."),
        prefix: str = "montage",
        page_size: int = 100,
        savefig_kwargs: Dict[str, Any] = {},
        **kwargs,
    ):
        """Render plots to one or more montages.

        Each montage has at most ``page_size`` subplots. This pager automatically saves an existing montage on overflow,
        which occurs when the montage is full and an attempt was made to add a new subplot to it. After the existing
        montage is saved, a new blank montage is created, and the new subplot will be added to it. Callers are expected
        to explicitly save the last montage.

        >>> import pandas as pd
        >>>
        >>> mp = MontagePlotter('output', savefig_kwargs=dict(transparent=TrueFalse))
        >>> for i in range(128):
        >>>     title = f"chart-{i:04d}"
        >>>     pd.Series(rand(6)).plot(ax=mp.pop(title), title=title)
        >>> mp.savefig()  # Save the last montage which may be partially filled.

        >>> mp = MontagePlotter('output', savefig_kwargs=dict(transparent=True))
        >>> ...

        Args:
            prefix (str, optional): Prefix of output filenames. Defaults to "montage".
            page_size (int, optional): Number of subplots per montage. Defaults to 100.
            savefig_kwargs (dict, optional): Keyword arguments to SimpleMatrixPlotter.savefig(), but ``fname`` will be
                overriden by MontagePager.
            kwargs: Keyword arguments to instantiate each montage (i.e., SimpleMatrixPlotter.__init__()).
        """
        self.path = path
        self.montage_path = path / "montages"
        self.individual_path = path / "individuals"
        self.prefix = prefix
        self.page_size = page_size
        self.smp_kwargs = kwargs
        self.smp_kwargs["init_figcount"] = page_size
        self.savefig_kwargs = savefig_kwargs
        self.smp = SimpleMatrixPlotter(**self.smp_kwargs)
        self._i = 0
        self.itemid: List[Any] = []
        self.csvwriter = csv.writer((path / "mappings.csv").open("w"))
        self.csvwriter.writerow(["individual", "title", "montage", "subplot", "row", "col"])

    @property
    def i(self):
        """:int: Sequence number of the current montage (zero-based)."""
        return self._i

    @property
    def filename(self):
        return f"{self.prefix}-{self._i:04d}.png"

    def pop(self, subplot_id: Any = "", **kwargs):
        """Return the next axes, and associate the returned axes with `subplot_id`."""
        if self.smp.i >= self.page_size:
            self.savefig()
            self.smp = SimpleMatrixPlotter(**self.smp_kwargs)
            self._i += 1
        self.itemid.append(subplot_id)
        return self.smp.pop()

    def savefig(self):
        """Save the current montage to a file."""
        # These must be done before smp.savefig() which destroys the underlying figure.
        subplot_cnt = self.smp.i
        bg_rgb = tuple((int(255 * channel) for channel in self.smp.fig.get_facecolor()[:3]))

        if subplot_cnt < 1:
            return

        with BytesIO() as buf:
            # Get the image buffer of the bbox-transformed canvas -- print_figure() in matplotlib/backend_bases.py.
            #
            # NOTE: methods described in https://stackoverflow.com/questions/4325733/save-a-subplot-in-matplotlib) was
            # tested and found not robust. The best outcome was using `exent = ax.get_tightbbox(...)`, which was still
            # not good enought as the next-row title still creeps into individual subplot image (tested with matplotlib
            # 3.2.1 and 3.3.0).
            self.smp.savefig(buf, format="png", **self.savefig_kwargs)
            buf.seek(0)
            im = Image.open(buf)
            im.load()

        im.save(self.montage_path / f"{self.prefix}-{self._i:04d}.png")
        self._save_pieces(im, subplot_cnt, bg_rgb)
        self._save_csv()

    def _save_csv(self):
        """Write row ["individual", "title", "montage-fname", "subplot-idx", "row", "col"]."""
        ncols = self.smp.ncols
        mtg_i = self._i
        mtg_fname = self.filename
        for i, itemid in enumerate(self.itemid):
            row, col = divmod(i, ncols)
            s = str(itemid).encode("unicode-escape").decode("utf-8")
            self.csvwriter.writerow((f"{mtg_i:04d}-{i:02d}.png", s, mtg_fname, i, row, col))
        self.itemid.clear()

    def _save_pieces(
        self,
        im: Image.Image,
        subplot_cnt: int,
        bg_rgb: Tuple[float, float, float] = (255, 255, 255),
        debug: bool = False,
    ):
        """Chop to pieces and save, row-wise."""

        def subplot_size():
            true_nrows = (subplot_cnt // self.smp.nrows) + ((subplot_cnt % self.smp.ncols) != 0)
            true_ncols = min(self.smp.ncols, subplot_cnt)
            h = im.height // true_nrows
            w = im.width // true_ncols
            if debug:
                #    print(f"{im.height=} {im.width=} {subplot_cnt=} {true_nrows=} {true_ncols=} {h=} {w=}")
                print(
                    f"im.height={im.height} im.width={im.width} im.subplot_cnt={subplot_cnt}"
                    f"true_nrows={true_nrows} true_ncols={true_ncols} h={h} w={w}"
                )
            return h, w

        subplot_h, subplot_w = subplot_size()

        def fixed_bbox(i):
            """To crop by fix size, but may produce excess border & plot not centered."""
            row, col = (i // self.smp.nrows), (i % self.smp.ncols)
            up = row * subplot_h
            left = col * subplot_w
            right = left + subplot_w
            bottom = up + subplot_h
            if debug:
                #    print(f"fixed_bbox: {i=} {row=} {col=} {left=} {up=} {right=} {bottom=}")
                print(f"fixed_bbox: i={i} row={row} col={col} left={left} up={up} right={right} bottom={bottom}")
            return left, up, right, bottom

        def tighten(im):
            if self.savefig_kwargs.get("transparent", False):
                bbox = im.getbbox()
            else:
                bg = Image.new("RGB", im.size, bg_rgb)
                diff = ImageChops.difference(im.convert("RGB"), bg)
                bbox = diff.getbbox()

            # Tight crop, but with small pads for a more pleasant view
            cropped = im.crop(pad(*bbox))
            return cropped

        def pad(left: float, up: float, right: float, bottom: float, pixels=4) -> Tuple[float, float, float, float]:
            left = min(left - pixels, 0)
            up = min(up - pixels, 0)
            right = min(right + pixels, im.width)
            bottom = min(bottom + pixels, im.height)
            return (left, up, right, bottom)

        subplot_h, subplot_w = subplot_size()
        for i in range(subplot_cnt):
            cropped_fixed = im.crop(fixed_bbox(i))
            cropped_tight = tighten(cropped_fixed)
            cropped_tight.save(self.individual_path / f"{self._i:04d}-{i:02d}.png")

def plot_binpat(
    img, figsize: Optional[Tuple[float, float]] = None, title: str = "", xlabel: str = "", ylabel: str = ""
):
    """Plot an binary-heatmap image using matplotlib.

    Args:
        img (array-like or PIL image): The image data (see matplotlib.pyplot.imshow())
        figsize (Optional[Tuple[float, float]], optional): Figure size. Defaults to None.
        title (str, optional): Title of the plot. Defaults to "".
        xlabel (str, optional): x-label of the plot. Defaults to "".
        ylabel (str, optional): y-label of the plot. Defaults to "".
    """
    # FIXME: to tidy up.
    plt.figure(figsize=figsize)
    plt.title(title)
    plt.ylabel(ylabel)
    plt.xlabel(xlabel)
    plt.gca().set_yticks([])
    plt.gca().set_xticks([])
    plt.imshow(img, cmap="gray", interpolation="none")

max_stars_repo_name: Georgi-Petkov/smallmatter-package

max_issues_repo_name: Georgi-Petkov/smallmatter-package

max_forks_repo_name: Georgi-Petkov/smallmatter-package


text:
import pickle
from os import path as osp
from torch.utils.data import Dataset
from torchvision.datasets.folder import default_loader
from augmentation import transform_train, transform_aug
from torchvision import transforms
import torch
import numpy as np

class CUB200_2011Dataset(Dataset):
    def __init__(self, root, train=True, transform=transform_train, nnIndex=None, augment=transform_train, target_transform=None):
        self.root = root
        self.dataset_dir = osp.join(self.root, self.__class__.__name__)
        self.transform = transform
        self.augment = augment
        self.target_transform = target_transform
        self.ae = False
        self.train = train
        self.loader = default_loader
        self.cls_to_idx = {}
        self.ps_label = []
        self._paths, self._targets = self.get_processed_data()
        self.nnIndex = nnIndex

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        tmp = 'train' if self.train is True else 'test'
        fmt_str += '    Split: {}\n'.format(tmp)
        fmt_str += '    Root Location: {}\n'.format(self.root)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        return fmt_str

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

    def get_processed_data(self):
        self.gen_cls2idx()
        if self.train:
            data_file = 'train.txt'
        else:
            data_file = 'test.txt'
        paths = []
        targets = []
        with open(osp.join(self.dataset_dir, data_file)) as f:
            for line in f:
                filename = line.strip()
                target = self.cls_to_idx[osp.split(filename)[0]]
                filepath = osp.join(self.dataset_dir, 'images', filename)
                paths.append(filepath)
                targets.append(target)
        return paths, targets

    def gen_cls2idx(self):
        if osp.exists(osp.join(self.dataset_dir, 'cls2idx.dict')):
            self.cls_to_idx = pickle.load(open(osp.join(self.dataset_dir, 'cls2idx.dict'), 'rb'))
        else:

            with open(osp.join(self.dataset_dir, 'classes.txt'), 'r') as f:
                for line in f:
                    self.cls_to_idx[line.strip()] = len(self.cls_to_idx)
            pickle.dump(self.cls_to_idx, open(osp.join(self.dataset_dir, 'cls2idx.dict'), 'wb'))

    def __getitem__(self, idx):
        if self.ae:
            img = default_loader(self.paths[idx])
            img = self.augment(img)
            train_img = self.transform(img)
            target_img = self.target_transform(img)
            target = self.targets[idx]
            return target_img, train_img, target, idx
        if self.nnIndex is not None:
            img1 = default_loader(self.paths[idx])
            img2 = default_loader(self.paths[self.nnIndex[idx]])
            target = self.targets[idx]
            img1 = self.transform(img1)
            img2 = self.augment(img2)
            return img1, img2, target, idx
        else:
            img = default_loader(self.paths[idx])
            img = self.transform(img)
            target = self.targets[idx]
            return img, target, idx

    def update_ps_label(self, labels):
        self.ps_label += labels

    @property
    def paths(self):
        return self._paths

    @property
    def targets(self):
        return self._targets

def handle_cub200(root, train=True):
    data_dir = osp.join(root, 'CUB200_2011Dataset')
    cls2idx = pickle.load(open(osp.join(data_dir, 'cls2idx.dict'), 'rb'))
    if train:
        data_file = 'train.txt'
    else:
        data_file = 'test.txt'
    paths = []
    targets = []
    with open(osp.join(data_dir, data_file)) as f:
        for line in f:
            filename = line.strip()
            target = cls2idx[osp.split(filename)[0]]
            filepath = osp.join(data_dir, 'images', filename)
            paths.append(filepath)
            targets.append(target)
    return paths, targets

def handle_dataset_alpha(root, train=True, dataset_name='cub200'):
    print("Loading", dataset_name)
    data_dir = osp.join(root, dataset_name)
    if train:
        img_data = np.load(data_dir + '/{}_{}_256resized_img.npy'.format('training', dataset_name))
        img_label = np.load(data_dir + '/{}_{}_256resized_label.npy'.format('training', dataset_name))
    else:
        img_data = np.load(data_dir + '/{}_{}_256resized_img.npy'.format('validation', dataset_name))
        img_label = np.load(data_dir + '/{}_{}_256resized_label.npy'.format('validation', dataset_name))
    return img_data, img_label

class MetricLearningDataset(Dataset):
    def __init__(self, root, train=True, dataset_name='cub200', transform=transform_train, augmentation=transform_train, nn_index=None):
        super(MetricLearningDataset, self).__init__()
        self.data, self.targets = handle_dataset_alpha(root, train, dataset_name)
        # nearest neighbor index
        self.nnIndex = nn_index
        if not self.nnIndex:
            self.nnIndex = np.arange(len(self.data))
        self.transform = transform
        self.augment = augmentation

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

    def __getitem__(self, idx):
        img1 = self.data[idx]
        img2 = self.data[self.nnIndex[idx]]
        target = self.targets[idx]
        img1 = self.transform(img1)
        img2 = self.augment(img2)
        return img1, img2, target

class AutoencoderDataset(Dataset):
    def __init__(self, root, train=True, dataset_name='cub200', input_sie=224, target_size=224):
        self.data, self.targets = handle_dataset_alpha(root, train, dataset_name)

        normalize = transforms.Compose([
            transforms.Lambda(lambda x: x * 255.0),
            transforms.Normalize(mean=[122.7717, 115.9465, 102.9801], std=[1, 1, 1]),
            transforms.Lambda(lambda x: x[[2, 1, 0], ...]),
            # transforms.Lambda(lambda x: x / 255.0),

        ])
        if train:
            self.augment = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((input_sie, input_sie)),
                transforms.RandomHorizontalFlip(),
                transforms.RandomRotation(30),
            ])
        else:
            self.augment = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((input_sie, input_sie)),
            ])

        self.transform_target = transforms.Compose([
            transforms.Resize((target_size, target_size)),
            transforms.ToTensor(),
            normalize
        ])

        self.transform_train = transforms.Compose([
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x + torch.zeros(x.size()).normal_(0, 0.09)),
            normalize
        ])

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

    def __getitem__(self, idx):
        img = self.data[idx]
        img = self.augment(img)
        train_img = self.transform_train(img)
        target_img = self.transform_target(img)
        return train_img, target_img

class MLDataInstance(Dataset):
    """Metric Learning Dataset.
    """

    def __init__(self, src_dir, dataset_name, train=True, transform=None, target_transform=None, nnIndex=None,
                 augmenter=None):

        data_dir = src_dir + '/' + dataset_name + '/'
        if train:
            img_data = np.load(data_dir + '{}_{}_256resized_img.npy'.format('training', dataset_name))
            img_label = np.load(data_dir + '{}_{}_256resized_label.npy'.format('training', dataset_name))
        else:
            img_data = np.load(data_dir + '{}_{}_256resized_img.npy'.format('validation', dataset_name))
            img_label = np.load(data_dir + '{}_{}_256resized_label.npy'.format('validation', dataset_name))

        self.img_data = img_data
        self.targets = img_label
        self.transform = transform
        self.target_transform = target_transform
        self.nnIndex = nnIndex
        self.augmenter = augmenter

    def __getitem__(self, index):

        if self.nnIndex is not None:

            img1, img2, target = self.img_data[index], self.img_data[self.nnIndex[index]], self.targets[index]

            img1 = self.transform(img1)
            if self.augmenter:
                img2 = self.augmenter(img2)
            else:
                img2 = self.transform(img2)
            if self.target_transform is not None:
                target = self.target_transform(target)
            return img1, img2, target

        else:
            img, target = self.img_data[index], self.targets[index]
            img = self.transform(img)
            if self.target_transform is not None:
                target = self.target_transform(target)
            return img, img, target

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