<commit_before>from itertools import product, chain

import numpy as np
from scipy.sparse import issparse, csc_matrix
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from metal.analysis import (
    plot_probabilities_histogram,
    lf_summary,
    confusion_matrix,
)
from metal.classifier import Classifier
from metal.label_model.lm_defaults import lm_model_defaults
from metal.utils import recursive_merge_dicts
from metal.label_model.graph_utils import get_clique_tree

class LabelModel(Classifier):
    def __init__(self, m, k, p=None, deps=[], **kwargs):
        """
        Args:
            seed: int: Random state seed
            deps: list: A list of LF dependencies as tuples of LF indices
            kwargs: TBD
        """
        self.config = recursive_merge_dicts(lm_model_defaults, kwargs)
        super().__init__()
        self.k = k
        self.m = m

        # Class balance- assume uniform if not provided
        if p is None:
            self.p = (1/self.k) * np.ones(self.k)
        else:
            self.p = p
        self.P = torch.diag(torch.from_numpy(self.p)).float()

        # Dependencies
        self.deps = deps
        self.c_tree = get_clique_tree(range(self.m), self.deps)

        # Whether to take the simple conditionally independent approach, or the
        # "inverse form" approach for handling dependencies
        # This flag allows us to eg test the latter even with no deps present
        self.inv_form = (len(self.deps) > 0)

    def config_set(self, update_dict):
        """Updates self.config with the values in a given update dictionary"""
        self.config = recursive_merge_dicts(self.config, update_dict)

    def _get_augmented_label_matrix(self, L, offset=0, higher_order=False):
        """Returns an augmented version of L where each column is an indicator
        for whether a certain source or clique of sources voted in a certain
        pattern.

        Args:
            - L: A dense n x m numpy array, where n is the number of data points
                and m is the number of sources, with values in {0,1,...,k}
            - offset: Create indicators for values {offset,...,k}
        """
        n, m = L.shape
        km = self.k + 1 - offset

        # Create a helper data structure which maps cliques (as tuples of member
        # sources) --> {start_index, end_index, maximal_cliques}, where
        # the last value is a set of indices in this data structure
        self.c_data = {}
        for i in range(m):
            self.c_data[i] = {
                'start_index': i*km,
                'end_index': (i+1)*km,
                'max_cliques': set([j for j in self.c_tree.nodes()
                    if i in self.c_tree.node[j]['members']])
            }

        # Form the columns corresponding to unary source labels
        L_aug = np.zeros((n, m * km))
        for y in range(offset, self.k+1):
            L_aug[:, y-offset::km] = np.where(L == y, 1, 0)

        # Get the higher-order clique statistics based on the clique tree
        # First, iterate over the maximal cliques (nodes of c_tree) and
        # separator sets (edges of c_tree)
        if higher_order:
            for item in chain(self.c_tree.nodes(), self.c_tree.edges()):
                if isinstance(item, int):
                    C = self.c_tree.node[item]
                    C_type = 'node'
                elif isinstance(item, tuple):
                    C = self.c_tree[item[0]][item[1]]
                    C_type = 'edge'
                else:
                    raise ValueError(item)
                members = list(C['members'])
                nc = len(members)

                # If a unary maximal clique, just store its existing index
                if nc == 1:
                    C['start_index'] = members[0] * km
                    C['end_index'] = (members[0]+1) * km

                # Else add one column for each possible value
                else:
                    L_C = np.ones((n, km ** nc))
                    for i, vals in enumerate(product(range(km), repeat=nc)):
                        for j, v in enumerate(vals):
                            L_C[:,i] *= L_aug[:, members[j]*km + v]

                    # Add to L_aug and store the indices
                    C['start_index'] = L_aug.shape[1]
                    C['end_index'] = L_aug.shape[1] + L_C.shape[1]
                    L_aug = np.hstack([L_aug, L_C])

                # Add to self.c_data as well
                self.c_data[tuple(members)] = {
                    'start_index': C['start_index'],
                    'end_index': C['end_index'],
                    'max_cliques': set([item]) if C_type=='node' else set(item)
                }
        return L_aug

    def _generate_O(self, L):
        """Form the overlaps matrix, which is just all the different observed
        combinations of values of pairs of sources

        Note that we only include the k non-abstain values of each source,
        otherwise the model not minimal --> leads to singular matrix
        """
        self.n, self.m = L.shape
        L_aug = self._get_augmented_label_matrix(L, offset=1)
        self.d = L_aug.shape[1]
        self.O = torch.from_numpy( L_aug.T @ L_aug / self.n ).float()

    def _generate_O_inv(self, L):
        """Form the *inverse* overlaps matrix"""
        self._generate_O(L)
        self.O_inv = torch.from_numpy(np.linalg.inv(self.O.numpy())).float()

    def _init_params(self):
        """Initialize the learned params

        - \mu is the primary learned parameter, where each row corresponds to
        the probability of a clique C emitting a specific combination of labels,
        conditioned on different values of Y (for each column); that is:

            self.mu[i*self.k + j, y] = P(\lambda_i = j | Y = y)

        and similarly for higher-order cliques.
        - Z is the inverse form version of \mu.
        - mask is the mask applied to O^{-1}, O for the matrix approx constraint
        """
        # Initialize mu so as to break basic reflective symmetry
        # TODO: Update for higher-order cliques!
        self.mu_init = torch.zeros(self.d, self.k)
        for i in range(self.m):
            for y in range(self.k):
                self.mu_init[i*self.k + y, y] += np.random.random()
        self.mu = nn.Parameter(self.mu_init.clone()).float()

        if self.inv_form:
            self.Z = nn.Parameter(torch.randn(self.d, self.k)).float()

        self.mask = torch.ones(self.d, self.d).byte()
        for ci in self.c_data.values():
            si, ei = ci['start_index'], ci['end_index']
            for cj in self.c_data.values():
                sj, ej = cj['start_index'], cj['end_index']

                # Check if ci and cj are part of the same maximal clique
                # If so, mask out their corresponding blocks in O^{-1}
                if len(ci['max_cliques'].intersection(cj['max_cliques'])) > 0:
                    self.mask[si:ei, sj:ej] = 0
                    self.mask[sj:ej, si:ei] = 0

    def get_conditional_probs(self, source=None):
        """Returns the full conditional probabilities table as a numpy array,
        where row i*(k+1) + ly is the conditional probabilities of source i
        emmiting label ly (including abstains 0), conditioned on different
        values of Y, i.e.:

            c_probs[i*(k+1) + ly, y] = P(\lambda_i = ly | Y = y)

        Note that this simply involves inferring the kth row by law of total
        probability and adding in to mu.

        If `source` is not None, returns only the corresponding block.
        """
        c_probs = np.zeros((self.m * (self.k+1), self.k))
        mu = self.mu.detach().clone().numpy()

        for i in range(self.m):
            # si = self.c_data[(i,)]['start_index']
            # ei = self.c_data[(i,)]['end_index']
            # mu_i = mu[si:ei, :]
            mu_i = mu[i*self.k:(i+1)*self.k, :]
            c_probs[i*(self.k+1) + 1:(i+1)*(self.k+1), :] = mu_i

            # The 0th row (corresponding to abstains) is the difference between
            # the sums of the other rows and one, by law of total prob
            c_probs[i*(self.k+1), :] = 1 - mu_i.sum(axis=0)
        c_probs = np.clip(c_probs, 0.01, 0.99)

        if source is not None:
            return c_probs[source*(self.k+1):(source+1)*(self.k+1)]
        else:
            return c_probs

    def get_label_probs(self, L):
        """Returns the n x k matrix of label probabilities P(Y | \lambda)"""
        L_aug = self._get_augmented_label_matrix(L, offset=1)
        mu = np.clip(self.mu.detach().clone().numpy(), 0.01, 0.99)

        # Create a "junction tree mask" over the columns of L_aug / mu
        if len(self.deps) > 0:
            jtm = np.zeros(L_aug.shape[1])

            # All maximal cliques are +1
            for i in self.c_tree.nodes():
                node = self.c_tree.node[i]
                jtm[node['start_index']:node['end_index']] = 1

            # All separator sets are -1
            for i, j in self.c_tree.edges():
                edge = self.c_tree[i][j]
                jtm[edge['start_index']:edge['end_index']] = 1
        else:
            jtm = np.ones(L_aug.shape[1])

        # Note: We omit abstains, effectively assuming uniform distribution here
        X = np.exp( L_aug @ np.diag(jtm) @ np.log(mu) + np.log(self.p) )
        Z = np.tile(X.sum(axis=1).reshape(-1,1), self.k)
        return X / Z

    def loss_inv_Z(self, l2=0.0):
        return torch.norm((self.O_inv + self.Z @ self.Z.t())[self.mask])**2

    def get_Q(self):
        """Get the model's estimate of Q = \mu P \mu^T

        We can then separately extract \mu subject to additional constraints,
        e.g. \mu P 1 = diag(O).
        """
        Z = self.Z.detach().clone().numpy()
        O = self.O.numpy()
        I_k = np.eye(self.k)
        return O @ Z @ np.linalg.inv(I_k + Z.T @ O @ Z) @ Z.T @ O

    def loss_inv_mu(self, l2=0.0):
        loss_1 = torch.norm(self.Q - self.mu @ self.P @ self.mu.t())**2
        loss_2 = torch.norm(
            torch.sum(self.mu @ self.P, 1) - torch.diag(self.O))**2
        return loss_1 + loss_2

    def loss_mu(self, l2=0.0):
        loss_1 = torch.norm(
            (self.O - self.mu @ self.P @ self.mu.t())[self.mask])**2
        loss_2 = torch.norm(
            torch.sum(self.mu @ self.P, 1) - torch.diag(self.O))**2
        # loss_l2 = torch.norm( self.mu - self.mu_init )**2
        loss_l2 = 0
        return loss_1 + loss_2 + l2 * loss_l2

    def train(self, L, **kwargs):
        """Train the model (i.e. estimate mu) in one of two ways, depending on
        whether source dependencies are provided or not:

        (1) No dependencies (conditionally independent sources): Estimate mu
        subject to constraints:
            (1a) O_{B(i,j)} - (mu P mu.T)_{B(i,j)} = 0, for i != j, where B(i,j)
                is the block of entries corresponding to sources i,j
            (1b) np.sum( mu P, 1 ) = diag(O)

        (2) Source dependencies:
            - First, estimate Z subject to the inverse form
            constraint:
                (2a) O_\Omega + (ZZ.T)_\Omega = 0, \Omega is the deps mask
            - Then, compute Q = mu P mu.T
            - Finally, estimate mu subject to mu P mu.T = Q and (1b)
        """
        self.config = recursive_merge_dicts(self.config, kwargs,
            misses='ignore')

        if self.inv_form:
            # Compute O, O^{-1}, and initialize params
            if self.config['verbose']:
                print("Computing O^{-1}...")
            self._generate_O_inv(L)
            self._init_params()

            # Estimate Z, compute Q = \mu P \mu^T
            if self.config['verbose']:
                print("Estimating Z...")
            self._train(self.loss_inv_Z)
            self.Q = torch.from_numpy(self.get_Q()).float()

            # Estimate \mu
            if self.config['verbose']:
                print("Estimating \mu...")
            self._train(self.loss_inv_mu)
        else:
            # Compute O and initialize params
            if self.config['verbose']:
                print("Computing O...")
            self._generate_O(L)
            self._init_params()

            # Estimate \mu
            if self.config['verbose']:
                print("Estimating \mu...")
            self._train(self.loss_mu)

    def _train(self, loss_fn):
        """Train model (self.parameters()) by optimizing the provided loss fn"""
        train_config = self.config['train_config']

        # Set optimizer as SGD w/ momentum
        optimizer_config = self.config['train_config']['optimizer_config']
        optimizer = optim.SGD(
            self.parameters(),
            **optimizer_config['optimizer_common'],
            **optimizer_config['sgd_config']
        )

        # Train model
        for epoch in range(train_config['n_epochs']):
            optimizer.zero_grad()

            # Compute gradient and take a step
            # Note that since this uses all N training points this is an epoch!
            loss = loss_fn(l2=train_config['l2'])
            if torch.isnan(loss):
                raise Exception("Loss is NaN. Consider reducing learning rate.")

            loss.backward()
            optimizer.step()

            # Print loss every print_every steps
            if (self.config['verbose'] and
                (epoch % train_config['print_every'] == 0
                or epoch == train_config['n_epochs'] - 1)):
                msg = f"[Epoch {epoch}] Loss: {loss.item():0.6f}"
                print(msg)<commit_msg>Add task_graph to init<commit_after>from itertools import product, chain

import numpy as np
from scipy.sparse import issparse, csc_matrix
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from metal.analysis import (
    plot_probabilities_histogram,
    lf_summary,
    confusion_matrix,
)
from metal.classifier import Classifier
from metal.label_model.lm_defaults import lm_model_defaults
from metal.utils import recursive_merge_dicts
from metal.label_model.graph_utils import get_clique_tree

class LabelModel(Classifier):
    def __init__(self, m, k=2, task_graph=None, p=None, deps=[], **kwargs):
        """
        Args:
            m: int: Number of sources
            k: int: Number of true classes
            task_graph: TaskGraph: A TaskGraph which defines a feasible set of
                task label vectors; note this overrides k
            p: np.array: Class balance
            deps: list: A list of source dependencies as tuples of indices
            kwargs:
                - seed: int: Random state seed
        """
        self.config = recursive_merge_dicts(lm_model_defaults, kwargs)
        super().__init__()
        self.k = k
        self.m = m

        # TaskGraph; note overrides k if present
        self.task_graph = task_graph
        if self.task_graph is not None:
            self.k = len(self.task_graph)

        # Class balance- assume uniform if not provided
        if p is None:
            self.p = (1/self.k) * np.ones(self.k)
        else:
            self.p = p
        self.P = torch.diag(torch.from_numpy(self.p)).float()

        # Dependencies
        self.deps = deps
        self.c_tree = get_clique_tree(range(self.m), self.deps)

        # Whether to take the simple conditionally independent approach, or the
        # "inverse form" approach for handling dependencies
        # This flag allows us to eg test the latter even with no deps present
        self.inv_form = (len(self.deps) > 0)

    def config_set(self, update_dict):
        """Updates self.config with the values in a given update dictionary"""
        self.config = recursive_merge_dicts(self.config, update_dict)

    def _get_augmented_label_matrix(self, L, offset=0, higher_order=False):
        """Returns an augmented version of L where each column is an indicator
        for whether a certain source or clique of sources voted in a certain
        pattern.

        Args:
            - L: A dense n x m numpy array, where n is the number of data points
                and m is the number of sources, with values in {0,1,...,k}
            - offset: Create indicators for values {offset,...,k}
        """
        n, m = L.shape
        km = self.k + 1 - offset

        # Create a helper data structure which maps cliques (as tuples of member
        # sources) --> {start_index, end_index, maximal_cliques}, where
        # the last value is a set of indices in this data structure
        self.c_data = {}
        for i in range(m):
            self.c_data[i] = {
                'start_index': i*km,
                'end_index': (i+1)*km,
                'max_cliques': set([j for j in self.c_tree.nodes()
                    if i in self.c_tree.node[j]['members']])
            }

        # Form the columns corresponding to unary source labels
        L_aug = np.zeros((n, m * km))
        for y in range(offset, self.k+1):
            L_aug[:, y-offset::km] = np.where(L == y, 1, 0)

        # Get the higher-order clique statistics based on the clique tree
        # First, iterate over the maximal cliques (nodes of c_tree) and
        # separator sets (edges of c_tree)
        if higher_order:
            for item in chain(self.c_tree.nodes(), self.c_tree.edges()):
                if isinstance(item, int):
                    C = self.c_tree.node[item]
                    C_type = 'node'
                elif isinstance(item, tuple):
                    C = self.c_tree[item[0]][item[1]]
                    C_type = 'edge'
                else:
                    raise ValueError(item)
                members = list(C['members'])
                nc = len(members)

                # If a unary maximal clique, just store its existing index
                if nc == 1:
                    C['start_index'] = members[0] * km
                    C['end_index'] = (members[0]+1) * km

                # Else add one column for each possible value
                else:
                    L_C = np.ones((n, km ** nc))
                    for i, vals in enumerate(product(range(km), repeat=nc)):
                        for j, v in enumerate(vals):
                            L_C[:,i] *= L_aug[:, members[j]*km + v]

                    # Add to L_aug and store the indices
                    C['start_index'] = L_aug.shape[1]
                    C['end_index'] = L_aug.shape[1] + L_C.shape[1]
                    L_aug = np.hstack([L_aug, L_C])

                # Add to self.c_data as well
                self.c_data[tuple(members)] = {
                    'start_index': C['start_index'],
                    'end_index': C['end_index'],
                    'max_cliques': set([item]) if C_type=='node' else set(item)
                }
        return L_aug

    def _generate_O(self, L):
        """Form the overlaps matrix, which is just all the different observed
        combinations of values of pairs of sources

        Note that we only include the k non-abstain values of each source,
        otherwise the model not minimal --> leads to singular matrix
        """
        self.n, self.m = L.shape
        L_aug = self._get_augmented_label_matrix(L, offset=1)
        self.d = L_aug.shape[1]
        self.O = torch.from_numpy( L_aug.T @ L_aug / self.n ).float()

    def _generate_O_inv(self, L):
        """Form the *inverse* overlaps matrix"""
        self._generate_O(L)
        self.O_inv = torch.from_numpy(np.linalg.inv(self.O.numpy())).float()

    def _init_params(self):
        """Initialize the learned params

        - \mu is the primary learned parameter, where each row corresponds to
        the probability of a clique C emitting a specific combination of labels,
        conditioned on different values of Y (for each column); that is:

            self.mu[i*self.k + j, y] = P(\lambda_i = j | Y = y)

        and similarly for higher-order cliques.
        - Z is the inverse form version of \mu.
        - mask is the mask applied to O^{-1}, O for the matrix approx constraint
        """
        # Initialize mu so as to break basic reflective symmetry
        # TODO: Update for higher-order cliques!
        self.mu_init = torch.zeros(self.d, self.k)
        for i in range(self.m):
            for y in range(self.k):
                self.mu_init[i*self.k + y, y] += np.random.random()
        self.mu = nn.Parameter(self.mu_init.clone()).float()

        if self.inv_form:
            self.Z = nn.Parameter(torch.randn(self.d, self.k)).float()

        self.mask = torch.ones(self.d, self.d).byte()
        for ci in self.c_data.values():
            si, ei = ci['start_index'], ci['end_index']
            for cj in self.c_data.values():
                sj, ej = cj['start_index'], cj['end_index']

                # Check if ci and cj are part of the same maximal clique
                # If so, mask out their corresponding blocks in O^{-1}
                if len(ci['max_cliques'].intersection(cj['max_cliques'])) > 0:
                    self.mask[si:ei, sj:ej] = 0
                    self.mask[sj:ej, si:ei] = 0

    def get_conditional_probs(self, source=None):
        """Returns the full conditional probabilities table as a numpy array,
        where row i*(k+1) + ly is the conditional probabilities of source i
        emmiting label ly (including abstains 0), conditioned on different
        values of Y, i.e.:

            c_probs[i*(k+1) + ly, y] = P(\lambda_i = ly | Y = y)

        Note that this simply involves inferring the kth row by law of total
        probability and adding in to mu.

        If `source` is not None, returns only the corresponding block.
        """
        c_probs = np.zeros((self.m * (self.k+1), self.k))
        mu = self.mu.detach().clone().numpy()

        for i in range(self.m):
            # si = self.c_data[(i,)]['start_index']
            # ei = self.c_data[(i,)]['end_index']
            # mu_i = mu[si:ei, :]
            mu_i = mu[i*self.k:(i+1)*self.k, :]
            c_probs[i*(self.k+1) + 1:(i+1)*(self.k+1), :] = mu_i

            # The 0th row (corresponding to abstains) is the difference between
            # the sums of the other rows and one, by law of total prob
            c_probs[i*(self.k+1), :] = 1 - mu_i.sum(axis=0)
        c_probs = np.clip(c_probs, 0.01, 0.99)

        if source is not None:
            return c_probs[source*(self.k+1):(source+1)*(self.k+1)]
        else:
            return c_probs

    def get_label_probs(self, L):
        """Returns the n x k matrix of label probabilities P(Y | \lambda)"""
        L_aug = self._get_augmented_label_matrix(L, offset=1)
        mu = np.clip(self.mu.detach().clone().numpy(), 0.01, 0.99)

        # Create a "junction tree mask" over the columns of L_aug / mu
        if len(self.deps) > 0:
            jtm = np.zeros(L_aug.shape[1])

            # All maximal cliques are +1
            for i in self.c_tree.nodes():
                node = self.c_tree.node[i]
                jtm[node['start_index']:node['end_index']] = 1

            # All separator sets are -1
            for i, j in self.c_tree.edges():
                edge = self.c_tree[i][j]
                jtm[edge['start_index']:edge['end_index']] = 1
        else:
            jtm = np.ones(L_aug.shape[1])

        # Note: We omit abstains, effectively assuming uniform distribution here
        X = np.exp( L_aug @ np.diag(jtm) @ np.log(mu) + np.log(self.p) )
        Z = np.tile(X.sum(axis=1).reshape(-1,1), self.k)
        return X / Z

    def loss_inv_Z(self, l2=0.0):
        return torch.norm((self.O_inv + self.Z @ self.Z.t())[self.mask])**2

    def get_Q(self):
        """Get the model's estimate of Q = \mu P \mu^T

        We can then separately extract \mu subject to additional constraints,
        e.g. \mu P 1 = diag(O).
        """
        Z = self.Z.detach().clone().numpy()
        O = self.O.numpy()
        I_k = np.eye(self.k)
        return O @ Z @ np.linalg.inv(I_k + Z.T @ O @ Z) @ Z.T @ O

    def loss_inv_mu(self, l2=0.0):
        loss_1 = torch.norm(self.Q - self.mu @ self.P @ self.mu.t())**2
        loss_2 = torch.norm(
            torch.sum(self.mu @ self.P, 1) - torch.diag(self.O))**2
        return loss_1 + loss_2

    def loss_mu(self, l2=0.0):
        loss_1 = torch.norm(
            (self.O - self.mu @ self.P @ self.mu.t())[self.mask])**2
        loss_2 = torch.norm(
            torch.sum(self.mu @ self.P, 1) - torch.diag(self.O))**2
        # loss_l2 = torch.norm( self.mu - self.mu_init )**2
        loss_l2 = 0
        return loss_1 + loss_2 + l2 * loss_l2

    def train(self, L, **kwargs):
        """Train the model (i.e. estimate mu) in one of two ways, depending on
        whether source dependencies are provided or not:

        (1) No dependencies (conditionally independent sources): Estimate mu
        subject to constraints:
            (1a) O_{B(i,j)} - (mu P mu.T)_{B(i,j)} = 0, for i != j, where B(i,j)
                is the block of entries corresponding to sources i,j
            (1b) np.sum( mu P, 1 ) = diag(O)

        (2) Source dependencies:
            - First, estimate Z subject to the inverse form
            constraint:
                (2a) O_\Omega + (ZZ.T)_\Omega = 0, \Omega is the deps mask
            - Then, compute Q = mu P mu.T
            - Finally, estimate mu subject to mu P mu.T = Q and (1b)
        """
        self.config = recursive_merge_dicts(self.config, kwargs,
            misses='ignore')

        if self.inv_form:
            # Compute O, O^{-1}, and initialize params
            if self.config['verbose']:
                print("Computing O^{-1}...")
            self._generate_O_inv(L)
            self._init_params()

            # Estimate Z, compute Q = \mu P \mu^T
            if self.config['verbose']:
                print("Estimating Z...")
            self._train(self.loss_inv_Z)
            self.Q = torch.from_numpy(self.get_Q()).float()

            # Estimate \mu
            if self.config['verbose']:
                print("Estimating \mu...")
            self._train(self.loss_inv_mu)
        else:
            # Compute O and initialize params
            if self.config['verbose']:
                print("Computing O...")
            self._generate_O(L)
            self._init_params()

            # Estimate \mu
            if self.config['verbose']:
                print("Estimating \mu...")
            self._train(self.loss_mu)

    def _train(self, loss_fn):
        """Train model (self.parameters()) by optimizing the provided loss fn"""
        train_config = self.config['train_config']

        # Set optimizer as SGD w/ momentum
        optimizer_config = self.config['train_config']['optimizer_config']
        optimizer = optim.SGD(
            self.parameters(),
            **optimizer_config['optimizer_common'],
            **optimizer_config['sgd_config']
        )

        # Train model
        for epoch in range(train_config['n_epochs']):
            optimizer.zero_grad()

            # Compute gradient and take a step
            # Note that since this uses all N training points this is an epoch!
            loss = loss_fn(l2=train_config['l2'])
            if torch.isnan(loss):
                raise Exception("Loss is NaN. Consider reducing learning rate.")

            loss.backward()
            optimizer.step()

            # Print loss every print_every steps
            if (self.config['verbose'] and
                (epoch % train_config['print_every'] == 0
                or epoch == train_config['n_epochs'] - 1)):
                msg = f"[Epoch {epoch}] Loss: {loss.item():0.6f}"
                print(msg)<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf8 -*-

""" Take CRIRES fits file and generate a tapas xml request to submit on tapas

"""

from astropy.io import fits
from astropy import units as u
from astropy.coordinates import Angle

def _parser():
    """Take care of all the argparse stuff.

    :returns: the args
    """
    parser = argparse.ArgumentParser(description='Telluric Removal')
    parser.add_argument('fname', help='Input fits file')
    parser.add_argument("-l", "--listpectra", help="If fname points to list of spectra for observation", default=False)
    parser.add_argument("-r", "--resolution", help="Spcify Instrument Resolution", default=False)
    parser.add_argument("-s", "--sampling", help="Sampling ratio", default=10)
    parser.add_argument("-i", "--ifuntion", help="Instrument function-Gaussian(1) or None(-1)", default="gaussian")
    parser.add_argument("-u", "--unit", help="Spectra Unit - Options: Air, Vacuum, Wavenumber", default="Air")
    parser.add_argument("-b", "--berv", help="Have BERV RV correction applied to the Tapas spectra", default=False)
    parser.add_argument("-f", "--tapas_format", help="Tapas file format Options: FITS,ASCII,NETCDF,VO", default="ASCII")
    # name  with a flag to specify if the name is to a listspectra.txt file to do whole observation!

# Many flags for tapas settings
# Try get most from fits files
    #parser.add_argument('-x', '--export', default=False,
 #                       help='Export result to fits file True/False')
    args = parser.parse_args()
    return args

def main(fname="/home/jneal/Phd/data/Crires/BDs-DRACS/HD30501-1/Combined_Nods/CRIRE.2012-04-07T00:08:29.976_1.nod.ms.norm.sum.wavecal.fits", listspectra=False, resolution=False, unit="air", ifunction="gaussian", sampling=10, berv=False, tapas_format="ASCII" )

    output_file = "/home/jneal/Phd/Codes/UsefulModules/Tapas_xml_request_file.xml"

    # Observation Settings
    if listspectra:
        pass
#        date =
        #target_ra = header["RA"]
        #target_dec = header["DEC"]
    else:
        header = fits.getheader(fname)

        date = header["DATE-OBS"][:-1]+"Z"
        print(header["DATE-OBS"])
        target_ra = header["RA"]
        target_dec = header["DEC"]
        wl_min = header["HIERARCH ESO INS WLEN MIN"])
        wl_max = header["HIERARCH ESO INS WLEN MAX"])
        instrument = header["INSTRUME"]
        telescope = header["TELESCOP"]
        slit_width = header["HIERARCH ESO INS SLIT1 WID"]
   # Observation Specifications
    # Date
    #MJD-OBS =       56024.00590250 / Obs start 2012-04-07T00:08:29.976
    #DATE-OBS= '2012-04-07T00:08:29.9764' / Observing date
    #EXPTIME =          180.0000000 / Integration time

        print("TODO !!! - Time correct for tapas timing issue")

   # Observatory settings
    if "VLT" in telscope:
    	obs_name = "ESO Paranal Chile (VLT)"
    else:
    	print("Don't have TAPAS telecope location nam (Using the Obervation value)")
    	obs_name = telescope

    obs_long = header["HIERARCH ESO TEL GEOLON"]
    obs_lat = header["HIERARCH ESO TEL GEOLAT"]
    obs_alt = header["HIERARCH ESO TEL GEOELEV"]

    # Target Settings
    ra_angle = Angle(target_ra, u.degree)
    ra_j2000 = str(ra_angle.to_string(unit=u.hour, sep=':',precision=0, pad=True))  # Extra str to get rid of u"string" which failed in template
    dec_angle = Angle(target_dec, u.deg)
    dec_j2000 = str(dec_angle.to_string(unit=u.degree, sep=':', precision=0, pad=True)) # Extra str to get rid of u"string" which failed in template
    print("RA Angle =",ra_angle, "RA deg =", target_ra, "ra_2000 h:m:s=", ra_j2000)
    print("Dec Angle =",dec_angle, "DEC deg =", target_dec, "dec_2000 d:m:s=", dec_j2000)

    #spectral_range = "{0} {1}".format(min_range, max_range)
    spec_range_min = round(wl_min - 10)
    spec_range_max = round(wl_max + 10)

    if resolution:
        resolving_power = resolution
    else:
        if "CRIRES" in instrument:
            if slit_width == 0.400:
                print("Slit width was 0.4, \nSetting resolving_power = 50000")
                resolving_power = 50000    # if Crires take the two slit width/Resolution values
            elif slit_width == 0.200:
                print("Slit width was 0.2, \nSetting resolving_power = 100000")
                resolving_power = 100000    # if Crires take the two slit width/Resolution values
            else:
        	    print("Slitwidth of CRIRES does not match 2 fixed settings")
        	    print("Is this older CRIRES data?")
        else:
            print("Resolving power not defined")

# Things want out of observations files
#["DATE-OBS"]
#TELESCOP= 'ESO-VLT-U1'
#RA      =            71.409028          / 04:45:38.1 RA (J2000) pointing (deg)
#DEC     =            -50.07734
#HIERARCH ESO INS SLIT1 WID   =        0.400
#INSTRUME= 'CRIRES  '
#HIERARCH ESO TEL GEOELEV     =         2648. / Elevation above sea level (m)
#HIERARCH ESO TEL GEOLAT      =     -24.6276 / Tel geo latitute (+=North) (deg)
#HIERARCH ESO TEL GEOLON      =     -70.4051 / Tel geo longitude (+=East) (deg)
#HIERARCH ESO INS WLEN MAX    =     2160.205 / Wavelength Limit + [nm].
#HIERARCH ESO INS WLEN MIN    =     2121.056 / Wavelength Limit - [nm].

#Test_file =
    #Template_xml_file = "/home/jneal/Phd/Codes/UsefulModules/tapas_xml_template.txt"

    # Tapas Specifications
    if tapas_format in ["ASCII","FITS","NETCDF","VO"]:
          tapas_format = tapas_format
    else:
        print("TAPAS format was not correct, using default of ASCII")
        tapas_format = "ASCII"

# open save file, find request id, add 1
    try:
        with open(output_file, "r") as f:
            for line in f:
                if line.startswith("<request_Id"):
                    ID_num = line.split('"')[1]
                    print("ID_number from last file", ID_num)
    else:
        Request_ID = "10000"    # Iterate on previous ID if found

    ray = "YES"
    h20 = "NO"   # alternate yes no for h20 for scaling
    o3 = "YES"
    o2 = "YES"
    co2 = "YES"
    ch4 = "YES"
    n2o = "YES"

    if "air" in unit.lower():
        spectral_choice = "Standard Wavelength (nm)"
    if "vac" in unit.lower():
        spectral_choice = "Vacuum Wavelength (nm)"
    if "wave" in unit.lower():
        spectral_choice = "Wavenumber (cm-1)"

    ifunction = str(ifunction).lower()
    if ifunction == "-1" or "none" in ifunction:
        ilsf_choice = -1    # -1, 0, 1
    elif ifunction == "1" or "gaussian" in ifunction:
        ilsf_choice = 1    # -1, 0, 1
    else:
        print("Instrumnet function not specifid correctly\n Valid choices are none(-1) or gaussian(1), The defualt is gaussian.")

    sampling_ratio = sampling # defualt 10
    if berv:
        apply_berv = "YES"
    else:
        apply_berv = "NO"

    #arletty_file = "canr_2013111700.arl"
    #ecmwf_file = "canr_2013111700_qo3.txt"
    #date = "2013-11-17T00:41:00.000Z"
    arletty_file = "canr_"+ date[0:4] + date[5:7] + date[8:10] + date[11:13] + ".arl"
    ecmwf_file = "canr_" + date[0:4] + date[5:7] + date[8:10] + date[11:13] + "_qo3.txt"

    print("arletty_file", arletty_file)
    print("ecmwf_file", ecmwf_file)

    d = {"Request_ID":Request_ID, "Format":Format, "ray":ray, "h20":h20, "o3":o3, "o2":o2, "co2":co2, "ch4":ch4, "n2o":n2o, "date":date, "obs_name":obs_name, "obs_long":obs_long, "obs_lat":obs_lat, "obs_alt":obs_alt, "ra_j2000":ra_j2000, "dec_j2000":dec_j2000, "spectral_choice":spectral_choice, "spec_range_min":spec_range_min, "spec_range_max":spec_range_max, "ilsf_choice":ilsf_choice, "resolving_power":resolving_power, "sampling_ratio":sampling_ratio, "apply_berv":apply_berv, "arletty_file":arletty_file, "ecmwf_file":ecmwf_file}

    template = """‹?xml version="1.0" encoding="UTF-8"?›
<tapas Id="Ether_TAPAS_999"›
<request Id="$Request_ID"›
<preferences>
<format valid="VO,ASCII,FITS,NETCDF">$Format</format>
<rayleigh_extinction valid="YES,NO">$ray</rayleigh_extinction>
<h2o_extinction valid="YES,NO">$h20</h2o_extinction>
<o3_extinction valid="YES,NO">$o3</o3_extinction>
<o2_extinction valid="YES,NO">$o3</o2_extinction>
<co2_extinction valid="YES,NO">$co2</co2_extinction>
<ch4_extinction valid="YES,NO">$ch4</ch4_extinction>
<n2o_extinction valid="YES,NO">$n2o</n2o_extinction>
</preferences>
<observation>
<date>$date</date>
<observatory>
<name>$obs_name</name>
<longitude min="-180" max="180">$obs_long</longitude>
<latitude min="-90" max="90">$obs_lat</latitude>
<altitude min="0" max="10000">$obs_alt</altitude>
</observatory>
<los>
<ra_j2000>$ra_j2000</ra_j2000>
<dec_j2000>$dec_j2000</dec_j2000>
</los>
<instrument>
<spectral_choice valid="Vacuum Wavelength (nm),Standard Wavelength (nm),Wavenumber (cm-1)">$spectral_choice</spectral_choice>
<spectral_range>$spec_range_min $spec_range_max</spectral_range>
<ilsf_choice valid="-1,0,1">$ilsf_choice</ilsf_choice>
<resolving_power min="0">$resolving_power</resolving_power>
<sampling_ratio min="0">$sampling_ratio</sampling_ratio>
<appli_berv valid="YES,NO">$apply_berv</appli_berv>
</instrument>
</observation>
<atmosphere>
<reference valid="0,1,2,3,4,5,6">0</reference>
<arletty_file>/data/tapas///arletty/$arletty_file</arletty_file>
<ecmwf_file>/data/tapas///ecmwf/$ecmwf_file</ecmwf_file>
</atmosphere>
</request>
</tapas> """

# ECMWF = ARLETTY
#[{"text":"ECMWF","value":0},{"text":"TROPICAL","value":1},{"text":"MEDIUM_LATITUDE_SUMMER","value":2},{"text":"MEDIUM_LATITUDE_WINTER","value":3},{"text":"SUBARCTIC_SUMMER","value":4},{"text":"SUBARCTIC_WINTER","value":5},{"text":"US_STANDARD_1976","value":6}]

#######################TODO - Missing all in one transmission

#tapasTexts["NM_STANDARD"] = "Standard Wavelength (nm)";
#            tapasTexts["NM_VACUUM"] = "Vacuum Wavelength (nm)";
#            tapasTexts["CM"] = "Wavenumber (cm-1)";
#
#            tapasTexts["NONE"] = "None";
#            tapasTexts["GAUSSIAN"] = "Gaussian";
#            tapasTexts["ECMWF"] = "ARLETTY";
#            tapasTexts["TROPICAL"] = "Tropical";
#            tapasTexts["MEDIUM_LATITUDE_SUMMER"] = "Average latitude summer";
#            tapasTexts["MEDIUM_LATITUDE_WINTER"] = "Average latitude winter";
#            tapasTexts["SUBARCTIC_SUMMER"] = "Subarctic summer";
#            tapasTexts["SUBARCTIC_WINTER"] = "Subarctic winter";
#            tapasTexts["US_STANDARD_1976"] = "Standard US 1976";

    from string import Template
    s = Template(template)

    print(s)

    sub = s.substitute(d)
    print(sub)

    # Save xml ouptut for copying to tapas
    #TOTRY in future - submit straight to tapas

    #output_file = "/home/jneal/Phd/data/Tapas/Tapas_xml_request_file.xml"
    output_file = "/home/jneal/Phd/Codes/UsefulModules/Tapas_xml_request_file.xml"
    with open(output_file, "w") as out:
    	out.write(sub)
    print("Saved tapas xml request to copy at \t {}".format(output_file))

if __name__ == "__main__":
#    args = vars(_parser())
#    fname = args.pop('fname')
#    opts = {k: args[k] for k in args}

#    main(fname, **opts)
    main()
<commit_msg>Species selection mapping<commit_after>#!/usr/bin/env python3
# -*- coding: utf8 -*-

""" Take CRIRES fits file and generate a tapas xml request to submit on tapas

"""

from astropy.io import fits
from astropy import units as u
from astropy.coordinates import Angle

def _parser():
    """Take care of all the argparse stuff.

    :returns: the args
    """
    parser = argparse.ArgumentParser(description='Telluric Removal')
    parser.add_argument('fname', help='Input fits file')
    parser.add_argument("-l", "--listpectra", help="If fname points to list of spectra for observation", default=False)
    parser.add_argument("-r", "--resolution", help="Spcify Instrument Resolution", default=False)
    parser.add_argument("-s", "--sampling", help="Sampling ratio", default=10)
    parser.add_argument("-i", "--ifuntion", help="Instrument function-Gaussian(1) or None(-1)", default="gaussian")
    parser.add_argument("-u", "--unit", help="Spectra Unit - Options: Air, Vacuum, Wavenumber", default="Air")
    parser.add_argument("-b", "--berv", help="Have BERV RV correction applied to the Tapas spectra", default=False)
    parser.add_argument("-f", "--tapas_format", help="Tapas file format Options: FITS,ASCII,NETCDF,VO", default="ASCII")
    # name  with a flag to specify if the name is to a listspectra.txt file to do whole observation!

# Many flags for tapas settings
# Try get most from fits files
    #parser.add_argument('-x', '--export', default=False,
 #                       help='Export result to fits file True/False')
    args = parser.parse_args()
    return args

def main(fname="/home/jneal/Phd/data/Crires/BDs-DRACS/HD30501-1/Combined_Nods/CRIRE.2012-04-07T00:08:29.976_1.nod.ms.norm.sum.wavecal.fits", listspectra=False, resolution=False, unit="air", ifunction="gaussian", sampling=10, berv=False, tapas_format="ASCII"):

    output_file = "/home/jneal/Phd/Codes/UsefulModules/Tapas_xml_request_file.xml"

    # Observation Settings
    if listspectra:
        pass
#        date =
        #target_ra = header["RA"]
        #target_dec = header["DEC"]
    else:
        header = fits.getheader(fname)

        date = header["DATE-OBS"][:-1]+"Z"
        print(header["DATE-OBS"])
        target_ra = header["RA"]
        target_dec = header["DEC"]
        wl_min = header["HIERARCH ESO INS WLEN MIN"])
        wl_max = header["HIERARCH ESO INS WLEN MAX"])
        instrument = header["INSTRUME"]
        telescope = header["TELESCOP"]
        slit_width = header["HIERARCH ESO INS SLIT1 WID"]
   # Observation Specifications
    #MJD-OBS =       56024.00590250 / Obs start 2012-04-07T00:08:29.976
    #DATE-OBS= '2012-04-07T00:08:29.9764' / Observing date
    #EXPTIME =          180.0000000 / Integration time

        print("TODO !!! - Time correct for tapas timing issue")

   # Observatory settings
    if "VLT" in telscope:
    	obs_name = "ESO Paranal Chile (VLT)"
    else:
    	print("Don't have TAPAS telecope location nam (Using the Obervation value)")
    	obs_name = telescope

    obs_long = header["HIERARCH ESO TEL GEOLON"]
    obs_lat = header["HIERARCH ESO TEL GEOLAT"]
    obs_alt = header["HIERARCH ESO TEL GEOELEV"]

    # Target Settings
    ra_angle = Angle(target_ra, u.degree)
    ra_j2000 = str(ra_angle.to_string(unit=u.hour, sep=':',precision=0, pad=True))  # Extra str to get rid of u"string" which failed in template
    dec_angle = Angle(target_dec, u.deg)
    dec_j2000 = str(dec_angle.to_string(unit=u.degree, sep=':', precision=0, pad=True)) # Extra str to get rid of u"string" which failed in template
    print("RA Angle =",ra_angle, "RA deg =", target_ra, "ra_2000 h:m:s=", ra_j2000)
    print("Dec Angle =",dec_angle, "DEC deg =", target_dec, "dec_2000 d:m:s=", dec_j2000)

    #spectral_range = "{0} {1}".format(min_range, max_range)
    spec_range_min = round(wl_min - 10)
    spec_range_max = round(wl_max + 10)

    if resolution:
        resolving_power = resolution
    else:
        if "CRIRES" in instrument:
            if slit_width == 0.400:
                print("Slit width was 0.4, \nSetting resolving_power = 50000")
                resolving_power = 50000    # if Crires take the two slit width/Resolution values
            elif slit_width == 0.200:
                print("Slit width was 0.2, \nSetting resolving_power = 100000")
                resolving_power = 100000    # if Crires take the two slit width/Resolution values
            else:
        	    print("Slitwidth of CRIRES does not match 2 fixed settings")
        	    print("Is this older CRIRES data?")
        else:
            print("Resolving power not defined")

    # Tapas Specifications
    if tapas_format in ["ASCII","FITS","NETCDF","VO"]:
          tapas_format = tapas_format
    else:
        print("TAPAS format was not correct, using default of ASCII")
        tapas_format = "ASCII"

# open save file, find request id, add 1
    try:
        with open(output_file, "r") as f:
            for line in f:
                if line.startswith("<request_Id"):
                    ID_num = line.split('"')[1]
                    print("ID_number from last file", ID_num)
    else:
        Request_ID = "10000"    # Iterate on previous ID if found

# Specify different tapas spectra
    if species == "all":
        species_map = [1,1,1,1,1,1]
    elif species == "ray":
        species_map = [1,0,0,0,0,0]
    elif species == "h2o":
        species_map = [0,1,0,0,0,0]
    elif species == "o2":
        species_map = [0,0,1,0,0,0]
    elif species == "o3":
        species_map = [0,0,0,1,0,0]
    elif species == "co2":
3       species_map = [0,0,0,0,1,0]
    elif species == "not_h2o":
        species_map = [1,0,1,1,1,1]

    if species_map[0]:
        ray = "YES"
    else:
        ray = "NO"
    if species_map[1]:
        h20 = "YES"   # alternate yes no for h20 for scaling
    else:
        h20 = "NO"   # alternate yes no for h20 for scaling
    if species_map[2]:
        o3 = "YES"
    else:
        o3 = "NO"
    if species_map[3]:
        o2 = "YES"
    else:
        o2 = "NO"
    if species_map[4]:
        co2 = "YES"
    else:
        co2 = "NO"
    if species_map[5]:
        ch4 = "YES"
    else:
        ch4 = "NO"

    if "air" in unit.lower():
        spectral_choice = "Standard Wavelength (nm)"
    if "vac" in unit.lower():
        spectral_choice = "Vacuum Wavelength (nm)"
    if "wave" in unit.lower():
        spectral_choice = "Wavenumber (cm-1)"

    ifunction = str(ifunction).lower()
    if ifunction == "-1" or "none" in ifunction:
        ilsf_choice = -1    # -1, 0, 1
    elif ifunction == "1" or "gaussian" in ifunction:
        ilsf_choice = 1    # -1, 0, 1
    else:
        print("Instrumnet function not specifid correctly\n Valid choices are none(-1) or gaussian(1), The defualt is gaussian.")

    sampling_ratio = sampling # defualt 10
    if berv:
        apply_berv = "YES"
    else:
        apply_berv = "NO"

    #arletty_file = "canr_2013111700.arl"
    #ecmwf_file = "canr_2013111700_qo3.txt"
    #date = "2013-11-17T00:41:00.000Z"
    arletty_file = "canr_"+ date[0:4] + date[5:7] + date[8:10] + date[11:13] + ".arl"
    ecmwf_file = "canr_" + date[0:4] + date[5:7] + date[8:10] + date[11:13] + "_qo3.txt"

    print("arletty_file", arletty_file)
    print("ecmwf_file", ecmwf_file)

    d = {"Request_ID":Request_ID, "Format":Format, "ray":ray, "h20":h20, "o3":o3, "o2":o2, "co2":co2, "ch4":ch4, "n2o":n2o, "date":date, "obs_name":obs_name, "obs_long":obs_long, "obs_lat":obs_lat, "obs_alt":obs_alt, "ra_j2000":ra_j2000, "dec_j2000":dec_j2000, "spectral_choice":spectral_choice, "spec_range_min":spec_range_min, "spec_range_max":spec_range_max, "ilsf_choice":ilsf_choice, "resolving_power":resolving_power, "sampling_ratio":sampling_ratio, "apply_berv":apply_berv, "arletty_file":arletty_file, "ecmwf_file":ecmwf_file}

    template = """‹?xml version="1.0" encoding="UTF-8"?›
<tapas Id="Ether_TAPAS_999"›
<request Id="$Request_ID"›
<preferences>
<format valid="VO,ASCII,FITS,NETCDF">$Format</format>
<rayleigh_extinction valid="YES,NO">$ray</rayleigh_extinction>
<h2o_extinction valid="YES,NO">$h20</h2o_extinction>
<o3_extinction valid="YES,NO">$o3</o3_extinction>
<o2_extinction valid="YES,NO">$o3</o2_extinction>
<co2_extinction valid="YES,NO">$co2</co2_extinction>
<ch4_extinction valid="YES,NO">$ch4</ch4_extinction>
<n2o_extinction valid="YES,NO">$n2o</n2o_extinction>
</preferences>
<observation>
<date>$date</date>
<observatory>
<name>$obs_name</name>
<longitude min="-180" max="180">$obs_long</longitude>
<latitude min="-90" max="90">$obs_lat</latitude>
<altitude min="0" max="10000">$obs_alt</altitude>
</observatory>
<los>
<ra_j2000>$ra_j2000</ra_j2000>
<dec_j2000>$dec_j2000</dec_j2000>
</los>
<instrument>
<spectral_choice valid="Vacuum Wavelength (nm),Standard Wavelength (nm),Wavenumber (cm-1)">$spectral_choice</spectral_choice>
<spectral_range>$spec_range_min $spec_range_max</spectral_range>
<ilsf_choice valid="-1,0,1">$ilsf_choice</ilsf_choice>
<resolving_power min="0">$resolving_power</resolving_power>
<sampling_ratio min="0">$sampling_ratio</sampling_ratio>
<appli_berv valid="YES,NO">$apply_berv</appli_berv>
</instrument>
</observation>
<atmosphere>
<reference valid="0,1,2,3,4,5,6">0</reference>
<arletty_file>/data/tapas///arletty/$arletty_file</arletty_file>
<ecmwf_file>/data/tapas///ecmwf/$ecmwf_file</ecmwf_file>
</atmosphere>
</request>
</tapas> """

# ECMWF = ARLETTY
#[{"text":"ECMWF","value":0},{"text":"TROPICAL","value":1},{"text":"MEDIUM_LATITUDE_SUMMER","value":2},{"text":"MEDIUM_LATITUDE_WINTER","value":3},{"text":"SUBARCTIC_SUMMER","value":4},{"text":"SUBARCTIC_WINTER","value":5},{"text":"US_STANDARD_1976","value":6}]

#######################TODO - Missing all in one transmission

#tapasTexts["NM_STANDARD"] = "Standard Wavelength (nm)";
#            tapasTexts["NM_VACUUM"] = "Vacuum Wavelength (nm)";
#            tapasTexts["CM"] = "Wavenumber (cm-1)";
#
#            tapasTexts["NONE"] = "None";
#            tapasTexts["GAUSSIAN"] = "Gaussian";
#            tapasTexts["ECMWF"] = "ARLETTY";
#            tapasTexts["TROPICAL"] = "Tropical";
#            tapasTexts["MEDIUM_LATITUDE_SUMMER"] = "Average latitude summer";
#            tapasTexts["MEDIUM_LATITUDE_WINTER"] = "Average latitude winter";
#            tapasTexts["SUBARCTIC_SUMMER"] = "Subarctic summer";
#            tapasTexts["SUBARCTIC_WINTER"] = "Subarctic winter";
#            tapasTexts["US_STANDARD_1976"] = "Standard US 1976";

    from string import Template
    s = Template(template)

    print(s)

    sub = s.substitute(d)
    print(sub)

    # Save xml ouptut for copying to tapas
    #TOTRY in future - submit straight to tapas

    #output_file = "/home/jneal/Phd/data/Tapas/Tapas_xml_request_file.xml"
    output_file = "/home/jneal/Phd/Codes/UsefulModules/Tapas_xml_request_file.xml"
    with open(output_file, "w") as out:
    	out.write(sub)
    print("Saved tapas xml request to copy at \t {}".format(output_file))

if __name__ == "__main__":
#    args = vars(_parser())
#    fname = args.pop('fname')
#    opts = {k: args[k] for k in args}

#    main(fname, **opts)
    main()
<|endoftext|>


<commit_before>import sys
import math
import bencode
from hashlib import sha1
from random import sample
from constants import bsize
from piece import Piece, FinalPiece
from protocol import PeerProtocolFactory

DEBUG = True

class Torrent(object):
    """Container for parsed metadata from a .torrent file."""

    def __init__(self, filename):
        self.filename = filename
        self.parse_metainfo()

    def parse_metainfo(self):
        torrent_dict = bencode.bdecode(open(self.filename).read())

        self.announce_url = torrent_dict['announce']
        info = torrent_dict['info']
        self.info_hash = sha1(bencode.bencode(info)).digest()
        self.piece_length = info['piece length']

        if not 'files' in info:
            self.length = info['length']
            self.names_length = [[info['name'], self.length]]
        else:
            self.length = sum(f['length'] for f in info['files'])
            self.names_length = [(f['path'][0], f['length'])
                                    for f in info['files']]

        #break string pieces up into a list of those pieces
        pieces_string = info['pieces']
        hashes = [pieces_string[i: i+20] for i in xrange(0, len(pieces_string), 20)]
        num_pieces = len(hashes)
        blocks = self.piece_length / bsize

        #calculate size of last piece
        leftover = self.length - ((num_pieces - 1) * self.piece_length)
        final_blocks = int(math.ceil(float(leftover) / self.piece_length))

        self.pieces = [Piece(hashes[i], blocks) for i in xrange(num_pieces-1)]
        self.pieces.append(FinalPiece(hashes[-1], final_blocks, leftover))

class ActiveTorrent(Torrent):
    """Represents a torrent that is in the process of being downloaded."""

    def __init__(self, client, filename):
        super(ActiveTorrent, self).__init__(filename)
        self.client = client
        self.uploaded = 0
        self.downloaded = 0
        self.factory = PeerProtocolFactory(client, self)
        self.outfile = open('temp_file', 'w')
        self.to_dl = set(range(len(self.pieces)))

    def add_block(self, index, offset, block):
        if index not in self.to_dl: return #multiple requests for same piece

        piece = self.pieces[index]
        piece.add(offset, block)
        if piece.is_full:
            if piece.check_hash():
                self.write_piece(index)
            else:
                raise ValueError('Shit, hash didn\'t match')

    def write_piece(self, index):
        data = self.pieces[index].full_data
        self.outfile.seek(self.piece_length * index)
        self.outfile.write(data)
        self.downloaded += len(data)
        self.clean_up(index)

    def clean_up(self, index):
        self.to_dl.remove(index)
        if not self.left: #torrent is finished downloading
            self.finish()

    def get_block(self):
        if not self.to_dl: return
        index, = sample(self.to_dl, 1)
        offset_index = self.pieces[index].first_nothave()
        return index, offset_index

    def finish(self):
        self.outfile.close()
        del self.factory
        self.client.delete_torrent(self)

    def connect_to_peer(self, (host, port)):
        from twisted.internet import reactor
        reactor.connectTCP(host, port, self.factory)

    @property
    def left(self):
        return self.length - self.downloaded

if __name__ == '__main__':
    from client import TorrentClient
    client = TorrentClient(ActiveTorrent(sys.argv[1]))
<commit_msg>writing out the multiple files that are downloaded and then delete the temp<commit_after>import os
import sys
import math
import bencode
from hashlib import sha1
from random import sample
from constants import bsize
from piece import Piece, FinalPiece
from protocol import PeerProtocolFactory

DEBUG = True

class Torrent(object):
    """Container for parsed metadata from a .torrent file."""

    def __init__(self, filename):
        self.filename = filename
        self.parse_metainfo()

    def parse_metainfo(self):
        torrent_dict = bencode.bdecode(open(self.filename).read())

        self.announce_url = torrent_dict['announce']
        info = torrent_dict['info']
        self.info_hash = sha1(bencode.bencode(info)).digest()
        self.piece_length = info['piece length']

        if not 'files' in info:
            self.length = info['length']
            self.names_length = [[info['name'], self.length]]
        else:
            self.length = sum(f['length'] for f in info['files'])
            self.names_length = [(f['path'][0], f['length'])
                                    for f in info['files']]

        #break string pieces up into a list of those pieces
        pieces_string = info['pieces']
        hashes = [pieces_string[i: i+20] for i in xrange(0, len(pieces_string), 20)]
        num_pieces = len(hashes)
        blocks = self.piece_length / bsize

        #calculate size of last piece
        leftover = self.length - ((num_pieces - 1) * self.piece_length)
        final_blocks = int(math.ceil(float(leftover) / self.piece_length))

        self.pieces = [Piece(hashes[i], blocks) for i in xrange(num_pieces-1)]
        self.pieces.append(FinalPiece(hashes[-1], final_blocks, leftover))

class ActiveTorrent(Torrent):
    """Represents a torrent that is in the process of being downloaded."""

    def __init__(self, client, filename):
        super(ActiveTorrent, self).__init__(filename)
        self.client = client
        self.uploaded = 0
        self.downloaded = 0
        self.factory = PeerProtocolFactory(client, self)
        self.outfile = open('temp_file', 'w')
        self.to_dl = set(range(len(self.pieces)))

    def add_block(self, index, offset, block):
        if index not in self.to_dl: return #multiple requests for same piece

        piece = self.pieces[index]
        piece.add(offset, block)
        if piece.is_full:
            if piece.check_hash():
                self.write_piece(index)
            else:
                raise ValueError('Shit, hash didn\'t match')

    def write_piece(self, index):
        data = self.pieces[index].full_data
        self.outfile.seek(self.piece_length * index)
        self.outfile.write(data)
        self.downloaded += len(data)
        self.clean_up(index)

    def clean_up(self, index):
        self.to_dl.remove(index)
        if not self.left: #torrent is finished downloading
            self.finish()

    def get_block(self):
        if not self.to_dl: return
        index, = sample(self.to_dl, 1)
        offset_index = self.pieces[index].first_nothave()
        return index, offset_index

    def finish(self):
        self.outfile.close()
        self.write_files()
        del self.factory
        self.client.delete_torrent(self)

    def connect_to_peer(self, (host, port)):
        from twisted.internet import reactor
        reactor.connectTCP(host, port, self.factory)

    def write_files(self):
        with open(self.outfile) as out:
            for fname, size in self.names_length:
                with open(fname, 'w') as cur:
                    cur.write(out.read(size))
        os.remove(self.outfile)

    @property
    def left(self):
        return self.length - self.downloaded

if __name__ == '__main__':
    from client import TorrentClient
    client = TorrentClient(ActiveTorrent(sys.argv[1]))
<|endoftext|>


<commit_before># Original _suncless.py downloaded from
# http://github.com/toymachine/concurrence/raw/master/lib/concurrence/_stackless.py
# at Thu Jan  7 22:59:54 CET 2010

# Copyright (C) 2009, Hyves (Startphone Ltd.)
#
# This module is part of the Concurrence Framework and is released under
# the New BSD License: http://www.opensource.org/licenses/bsd-license.php

"""This module implements the stackless API on top of py.magic greenlet API
This way it is possible to run concurrence applications on top of normal python
using the greenlet module.
Because the greenlet module uses only 'hard' switching as opposed to stackless 'soft' switching
it is a bit slower (about 35%), but very usefull because you don't need to install stackless.
Note that this does not aim to be a complete implementation of stackless on top of greenlets,
just enough of the stackless API to make concurrence run.
This code was inspired by:
http://aigamedev.com/programming-tips/round-robin-multi-tasking and
also by the pypy implementation of the same thing (buggy, not being maintained?) at
https://codespeak.net/viewvc/pypy/dist/pypy/lib/stackless.py?view=markup
"""

try:
    from py.magic import greenlet #as of version 1.0 of py, it does not supply greenlets anymore
except ImportError:
    from greenlet import greenlet #there is an older package containing just the greenlet lib

from collections import deque

class TaskletExit(SystemExit):pass

import __builtin__
__builtin__.TaskletExit = TaskletExit

class bomb(object):
    """used as a result value for sending exceptions trough a channel"""
    def __init__(self, exc_type = None, exc_value = None, exc_traceback = None):
        self.type = exc_type
        self.value = exc_value
        self.traceback = exc_traceback

    def raise_(self):
        raise self.type, self.value, self.traceback

class channel(object):
    """implementation of stackless's channel object"""
    def __init__(self):
        self.balance = 0
        self.preference = -1
        self.queue = deque()

    def receive(self):
        return _scheduler._receive(self, self.preference)

    def send(self, data):
        return _scheduler._send(self, data, self.preference)

    def send_exception(self, exp_type, *args):
        self.send(bomb(exp_type, exp_type(*args)))

    def send_sequence(self, iterable):
        for item in iterable:
            self.send(item)

class tasklet(object):
    """implementation of stackless's tasklet object"""

    def __init__(self, f = None, greenlet = None, alive = False):
        self.greenlet = greenlet
        self.func = f
        self.alive = alive
        self.blocked = False
        self.data = None

    def bind(self, func):
        if not callable(func):
            raise TypeError('tasklet function must be a callable')
        self.func = func

    def __call__(self, *args, **kwargs):
        """this is where the new task starts to run, e.g. it is where the greenlet is created
        and the 'task' is first scheduled to run"""
        if self.func is None:
            raise TypeError('tasklet function must be a callable')

        def _func(*_args, **_kwargs):
            try:
                self.func(*args, **kwargs)
            except TaskletExit:
                pass #let it pass silently
            except:
                import logging
                logging.exception('unhandled exception in greenlet')
                #don't propagate to parent
            finally:
                assert _scheduler.current == self
                _scheduler.remove(self)
                if _scheduler._runnable: #there are more tasklets scheduled to run next
                    #this make sure that flow will continue in the correct greenlet, e.g. the next in the schedule
                    self.greenlet.parent = _scheduler._runnable[0].greenlet
                self.alive = False
                del self.greenlet
                del self.func
                del self.data

        self.greenlet = greenlet(_func)
        self.alive = True
        _scheduler.append(self)
        return self

    def kill(self):
        _scheduler.throw(self, TaskletExit)

    def raise_exception(self, *args):
        _scheduler.throw(self, *args)

    def __str__(self):
        return repr(self)

    def __repr__(self):
        if hasattr(self, 'name'):
            _id = self.name
        else:
            _id = str(self.func)
        return '<tasklet %s at %0x>' % (_id, id(self))

    def remove(self):
        """Remove self from the main scheduler queue.

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        if self.blocked:
            raise RuntimeError('You cannot remove a blocked tasklet.')
        i = 0
        for tasklet_obj in _scheduler._runnable:
            if tasklet_obj is self:
                del _scheduler._runnable[i]
                return self
            i += 1
        return self

    def insert(self):
        """Add self to the end of the scheduler queue, unless already in.

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        if not self.alive:
            raise RuntimeError('You cannot run an unbound(dead) tasklet')
        if self.blocked:
            raise RuntimeError('You cannot run a blocked tasklet')
        if self not in _scheduler._runnable:
            _scheduler._runnable.append(self)

    @property
    def next(self):
        """Return the next tasklet in the doubly-linked list.

        Stackless implements this method for all tasklets. This implementation
        raises a NotImplementedError unless self is scheduler.getcurrent() or
        self is the last runnable tasklet.
        """
        runnable = _scheduler._runnable
        if self is runnable[0]:
            return runnable[len(runnable) > 1]
        elif self is runnable[-1]:
            return runnable[0]
        else:
            raise NotImplementedError('tasklet.next for not current or last')

    @property
    def prev(self):
        """Return the next tasklet in the doubly-linked list.

        Stackless implements this method for all tasklets. This implementation
        raises a NotImplementedError unless self is scheduler.getcurrent() or
        self is the last runnable tasklet.
        """
        runnable = _scheduler._runnable
        if self is runnable[0]:
            return runnable[-1]
        elif self is runnable[-1]:
            return runnable[-2]
        else:
            raise NotImplementedError('tasklet.prev for not current or last')

    def run(self):
        """Switch execution to self, make it stackless.getcurrent().

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        if self.blocked:
            raise RuntimeError('You cannot run a blocked tasklet')
        if not self.alive:
            raise RuntimeError('You cannot run an unbound(dead) tasklet')
        runnable = _scheduler._runnable
        i = 0
        for tasklet_obj in runnable:
            if tasklet_obj is self:
                if i:
                    runnable.rotate(-i)
                    self.greenlet.switch()
                return
            i += 1
        runnable.appendleft(self)
        self.greenlet.switch()

class scheduler(object):
    def __init__(self):
        self._main_task = tasklet(greenlet = greenlet.getcurrent(), alive = True)
        #all non blocked tast are in this queue
        #all tasks are only onces in this queue
        #the current task is the first item in the queue
        self._runnable = deque([self._main_task])

    def schedule(self):
        """schedules the next tasks and puts the current task back at the queue of runnables"""
        self._runnable.rotate(-1)
        next_task = self._runnable[0]
        next_task.greenlet.switch()

    def schedule_block(self):
        """blocks the current task and schedules next"""
        self._runnable.popleft()
        next_task = self._runnable[0]
        next_task.greenlet.switch()

    def throw(self, task, *args):
        if not task.alive: return #this is what stackless does

        assert task.blocked or task in self._runnable

        task.greenlet.parent = self._runnable[0].greenlet
        if task.blocked:
            self._runnable.appendleft(task)
        else:
            self._runnable.remove(task)
            self._runnable.appendleft(task)

        task.greenlet.throw(*args)

    def _receive(self, channel, preference):
        #Receiving 1):
        #A tasklet wants to receive and there is
        #a queued sending tasklet. The receiver takes
        #its data from the sender, unblocks it,
        #and inserts it at the end of the runnables.
        #The receiver continues with no switch.
        #Receiving 2):
        #A tasklet wants to receive and there is
        #no queued sending tasklet.
        #The receiver will become blocked and inserted
        #into the queue. The next sender will
        #handle the rest through "Sending 1)".
        if channel.balance > 0: #some sender
            channel.balance -= 1
            sender = channel.queue.popleft()
            sender.blocked = False
            data, sender.data = sender.data, None
            if preference >= 0:
                #sender preference
                self._runnable.rotate(-1)
                self._runnable.appendleft(sender)
                self._runnable.rotate(1)
                self.schedule()
            else:
                #receiver preference
                self._runnable.append(sender)
        else: #no sender
            current = self._runnable[0]
            channel.queue.append(current)
            channel.balance -= 1
            current.blocked = True
            try:
                self.schedule_block()
            except:
                channel.queue.remove(current)
                channel.balance += 1
                current.blocked = False
                raise

            data, current.data = current.data, None

        if isinstance(data, bomb):
            data.raise_()
        else:
            return data

    def _send(self, channel, data, preference):
        #  Sending 1):
        #    A tasklet wants to send and there is
        #    a queued receiving tasklet. The sender puts
        #    its data into the receiver, unblocks it,
        #    and inserts it at the top of the runnables.
        #    The receiver is scheduled.
        #  Sending 2):
        #    A tasklet wants to send and there is
        #    no queued receiving tasklet.
        #    The sender will become blocked and inserted
        #    into the queue. The next receiver will
        #    handle the rest through "Receiving 1)".
        #print 'send q', channel.queue
        if channel.balance < 0: #some receiver
            channel.balance += 1
            receiver = channel.queue.popleft()
            receiver.data = data
            receiver.blocked = False
            #put receiver just after current task in runnable and schedule (which will pick it up)
            if preference < 0: #receiver pref
                self._runnable.rotate(-1)
                self._runnable.appendleft(receiver)
                self._runnable.rotate(1)
                self.schedule()
            else: #sender pref
                self._runnable.append(receiver)
        else: #no receiver
            current = self.current
            channel.queue.append(current)
            channel.balance += 1
            current.data = data
            current.blocked = True
            try:
                self.schedule_block()
            except:
                channel.queue.remove(current)
                channel.balance -= 1
                current.data = None
                current.blocked = False
                raise

    def remove(self, task):
        assert task.blocked or task in self._runnable
        if task in self._runnable:
            self._runnable.remove(task)

    def append(self, task):
        assert task not in self._runnable
        self._runnable.append(task)

    @property
    def runcount(self):
        return len(self._runnable)

    @property
    def current(self):
        return self._runnable[0]

#there is only 1 scheduler, this is it:
_scheduler = scheduler()

def getruncount():
    return _scheduler.runcount

def getcurrent():
    return _scheduler.current

def schedule():
    return _scheduler.schedule()

<commit_msg>added stackless.schedule_remove(); fixed tasklet.kill()<commit_after># Original _suncless.py downloaded from
# http://github.com/toymachine/concurrence/raw/master/lib/concurrence/_stackless.py
# at Thu Jan  7 22:59:54 CET 2010

# Copyright (C) 2009, Hyves (Startphone Ltd.)
#
# This module is part of the Concurrence Framework and is released under
# the New BSD License: http://www.opensource.org/licenses/bsd-license.php

"""This module implements the stackless API on top of py.magic greenlet API
This way it is possible to run concurrence applications on top of normal python
using the greenlet module.
Because the greenlet module uses only 'hard' switching as opposed to stackless 'soft' switching
it is a bit slower (about 35%), but very usefull because you don't need to install stackless.
Note that this does not aim to be a complete implementation of stackless on top of greenlets,
just enough of the stackless API to make concurrence run.
This code was inspired by:
http://aigamedev.com/programming-tips/round-robin-multi-tasking and
also by the pypy implementation of the same thing (buggy, not being maintained?) at
https://codespeak.net/viewvc/pypy/dist/pypy/lib/stackless.py?view=markup
"""

try:
    from py.magic import greenlet #as of version 1.0 of py, it does not supply greenlets anymore
except ImportError:
    from greenlet import greenlet #there is an older package containing just the greenlet lib

assert hasattr(greenlet, 'throw'), (
    'wrong version of greenlet loaded; please get greenlet from svn co '
    'http://codespeak.net/svn/py/release/0.9.x/py/c-extension/greenlet')

from collections import deque

class TaskletExit(SystemExit):pass

import __builtin__
__builtin__.TaskletExit = TaskletExit

class bomb(object):
    """used as a result value for sending exceptions trough a channel"""
    def __init__(self, exc_type = None, exc_value = None, exc_traceback = None):
        self.type = exc_type
        self.value = exc_value
        self.traceback = exc_traceback

    def raise_(self):
        raise self.type, self.value, self.traceback

class channel(object):
    """implementation of stackless's channel object"""
    def __init__(self):
        self.balance = 0
        self.preference = -1
        self.queue = deque()

    def receive(self):
        return _scheduler._receive(self, self.preference)

    def send(self, data):
        return _scheduler._send(self, data, self.preference)

    def send_exception(self, exp_type, *args):
        self.send(bomb(exp_type, exp_type(*args)))

    def send_sequence(self, iterable):
        for item in iterable:
            self.send(item)

class tasklet(object):
    """implementation of stackless's tasklet object"""

    def __init__(self, f = None, greenlet = None, alive = False):
        self.greenlet = greenlet
        self.func = f
        self.alive = alive
        self.blocked = False
        self.data = None

    def bind(self, func):
        if not callable(func):
            raise TypeError('tasklet function must be a callable')
        self.func = func

    def __call__(self, *args, **kwargs):
        """this is where the new task starts to run, e.g. it is where the greenlet is created
        and the 'task' is first scheduled to run"""
        if self.func is None:
            raise TypeError('tasklet function must be a callable')

        def _func(*_args, **_kwargs):
            try:
                self.func(*args, **kwargs)
            except TaskletExit:
                pass #let it pass silently
            except:
                import logging
                logging.exception('unhandled exception in greenlet')
                #don't propagate to parent
            finally:
                assert _scheduler.current == self
                _scheduler.remove(self)
                if _scheduler._runnable: #there are more tasklets scheduled to run next
                    #this make sure that flow will continue in the correct greenlet, e.g. the next in the schedule
                    self.greenlet.parent = _scheduler._runnable[0].greenlet
                self.alive = False
                del self.greenlet
                del self.func
                del self.data

        self.greenlet = greenlet(_func)
        self.alive = True
        _scheduler.append(self)
        return self

    def kill(self):
        _scheduler.throw(self, TaskletExit)

    def raise_exception(self, *args):
        _scheduler.throw(self, *args)

    def __str__(self):
        return repr(self)

    def __repr__(self):
        if hasattr(self, 'name'):
            _id = self.name
        else:
            _id = str(self.func)
        return '<tasklet %s at %0x>' % (_id, id(self))

    def remove(self):
        """Remove self from the main scheduler queue.

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        if self.blocked:
            raise RuntimeError('You cannot remove a blocked tasklet.')
        i = 0
        for task in _scheduler._runnable:
            if task is self:
                del _scheduler._runnable[i]
                return self
            i += 1
        return self

    def insert(self):
        """Add self to the end of the scheduler queue, unless already in.

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        if not self.alive:
            raise RuntimeError('You cannot run an unbound(dead) tasklet')
        if self.blocked:
            raise RuntimeError('You cannot run a blocked tasklet')
        if self not in _scheduler._runnable:
            _scheduler._runnable.append(self)

    @property
    def next(self):
        """Return the next tasklet in the doubly-linked list.

        Stackless implements this method for all tasklets. This implementation
        raises a NotImplementedError unless self is scheduler.getcurrent() or
        self is the last runnable tasklet.
        """
        runnable = _scheduler._runnable
        if self is runnable[0]:
            return runnable[len(runnable) > 1]
        elif self is runnable[-1]:
            return runnable[0]
        else:
            raise NotImplementedError('tasklet.next for not current or last')

    @property
    def prev(self):
        """Return the next tasklet in the doubly-linked list.

        Stackless implements this method for all tasklets. This implementation
        raises a NotImplementedError unless self is scheduler.getcurrent() or
        self is the last runnable tasklet.
        """
        runnable = _scheduler._runnable
        if self is runnable[0]:
            return runnable[-1]
        elif self is runnable[-1]:
            return runnable[-2]
        else:
            raise NotImplementedError('tasklet.prev for not current or last')

    def run(self):
        """Switch execution to self, make it stackless.getcurrent().

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        if self.blocked:
            raise RuntimeError('You cannot run a blocked tasklet')
        if not self.alive:
            raise RuntimeError('You cannot run an unbound(dead) tasklet')
        runnable = _scheduler._runnable
        i = 0
        for task in runnable:
            if task is self:
                if i:
                    runnable.rotate(-i)
                    self.greenlet.switch()
                return
            i += 1
        runnable.appendleft(self)
        self.greenlet.switch()

class scheduler(object):
    def __init__(self):
        self._main_task = tasklet(greenlet = greenlet.getcurrent(), alive = True)
        #all non blocked tast are in this queue
        #all tasks are only onces in this queue
        #the current task is the first item in the queue
        self._runnable = deque([self._main_task])

    def schedule(self):
        """schedules the next tasks and puts the current task back at the queue of runnables"""
        self._runnable.rotate(-1)
        self._runnable[0].greenlet.switch()

    def schedule_remove(self):
        """makes stackless.getcurrent() not runnable, schedules next tasks"""
        runnable = self._runnable
        if len(runnable) > 1:
            runnable.popleft()
            self._runnable[0].greenlet.switch()

    def schedule_block(self):
        """blocks the current task and schedules next"""
        self._runnable.popleft()
        next_task = self._runnable[0]
        next_task.greenlet.switch()

    def throw(self, task, *args):
        """Raise an exception in the tasklet, and run it.

        Please note that this implementation has O(r) complexity, where r is
        the number or runnable (non-blocked) tasklets. The implementation in
        Stackless has O(1) complexity.
        """
        assert args
        if not task.alive: return #this is what stackless does
        runnable = self._runnable
        task.greenlet.parent = runnable[0].greenlet
        if not task.blocked:
            i = 0
            for task2 in runnable:
                if task2 is task:
                    if i:
                        runnable.rotate(-i)
                        task.greenlet.throw(*args)
                        return
                    else:
                        exc_class = args.pop(0)
                        raise exc_class(*args)
                i += 1
        runnable.appendleft(task)
        task.greenlet.throw(*args)

    def _receive(self, channel, preference):
        #Receiving 1):
        #A tasklet wants to receive and there is
        #a queued sending tasklet. The receiver takes
        #its data from the sender, unblocks it,
        #and inserts it at the end of the runnables.
        #The receiver continues with no switch.
        #Receiving 2):
        #A tasklet wants to receive and there is
        #no queued sending tasklet.
        #The receiver will become blocked and inserted
        #into the queue. The next sender will
        #handle the rest through "Sending 1)".
        if channel.balance > 0: #some sender
            channel.balance -= 1
            sender = channel.queue.popleft()
            sender.blocked = False
            data, sender.data = sender.data, None
            if preference >= 0:
                #sender preference
                self._runnable.rotate(-1)
                self._runnable.appendleft(sender)
                self._runnable.rotate(1)
                self.schedule()
            else:
                #receiver preference
                self._runnable.append(sender)
        else: #no sender
            current = self._runnable[0]
            channel.queue.append(current)
            channel.balance -= 1
            current.blocked = True
            try:
                self.schedule_block()
            except:
                channel.queue.remove(current)
                channel.balance += 1
                current.blocked = False
                raise

            data, current.data = current.data, None

        if isinstance(data, bomb):
            data.raise_()
        else:
            return data

    def _send(self, channel, data, preference):
        #  Sending 1):
        #    A tasklet wants to send and there is
        #    a queued receiving tasklet. The sender puts
        #    its data into the receiver, unblocks it,
        #    and inserts it at the top of the runnables.
        #    The receiver is scheduled.
        #  Sending 2):
        #    A tasklet wants to send and there is
        #    no queued receiving tasklet.
        #    The sender will become blocked and inserted
        #    into the queue. The next receiver will
        #    handle the rest through "Receiving 1)".
        #print 'send q', channel.queue
        if channel.balance < 0: #some receiver
            channel.balance += 1
            receiver = channel.queue.popleft()
            receiver.data = data
            receiver.blocked = False
            #put receiver just after current task in runnable and schedule (which will pick it up)
            if preference < 0: #receiver pref
                self._runnable.rotate(-1)
                self._runnable.appendleft(receiver)
                self._runnable.rotate(1)
                self.schedule()
            else: #sender pref
                self._runnable.append(receiver)
        else: #no receiver
            current = self.current
            channel.queue.append(current)
            channel.balance += 1
            current.data = data
            current.blocked = True
            try:
                self.schedule_block()
            except:
                channel.queue.remove(current)
                channel.balance -= 1
                current.data = None
                current.blocked = False
                raise

    def remove(self, task):
        assert task.blocked or task in self._runnable
        if task in self._runnable:
            self._runnable.remove(task)

    def append(self, task):
        assert task not in self._runnable
        self._runnable.append(task)

    @property
    def runcount(self):
        return len(self._runnable)

    @property
    def current(self):
        return self._runnable[0]

#there is only 1 scheduler, this is it:
_scheduler = scheduler()

def getruncount():
    return _scheduler.runcount

def getcurrent():
    return _scheduler.current

def schedule():
    return _scheduler.schedule()

def schedule_remove():
    return _scheduler.schedule_remove()
<|endoftext|>


<commit_before>import torch
from torch.utils.data import DistributedSampler as _DistributedSampler

class DistributedSampler(_DistributedSampler):

    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
        super().__init__(dataset, num_replicas=num_replicas, rank=rank)
        self.shuffle = shuffle

    def __iter__(self):
        # deterministically shuffle based on epoch
        if self.shuffle:
            g = torch.Generator()
            g.manual_seed(self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()
        else:
            indices = torch.arange(len(self.dataset)).tolist()

        # add extra samples to make it evenly divisible
        indices += indices[:(self.total_size - len(indices))]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)
<commit_msg>Fix bug in distributed sampler (#4257)<commit_after>import math

import torch
from torch.utils.data import DistributedSampler as _DistributedSampler

class DistributedSampler(_DistributedSampler):

    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
        super().__init__(dataset, num_replicas=num_replicas, rank=rank)
        self.shuffle = shuffle

    def __iter__(self):
        # deterministically shuffle based on epoch
        if self.shuffle:
            g = torch.Generator()
            g.manual_seed(self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()
        else:
            indices = torch.arange(len(self.dataset)).tolist()

        # add extra samples to make it evenly divisible
        # in case that indices is shorter than half of total_size
        indices = (indices *
                   math.ceil(self.total_size / len(indices)))[:self.total_size]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)
<|endoftext|>


<commit_before>from os.path import dirname, join

import numpy as np
import pandas.io.sql as psql
import sqlite3 as sql

from bokeh.plotting import Figure
from bokeh.models import ColumnDataSource, HoverTool, HBox, VBoxForm
from bokeh.models.widgets import Slider, Select, TextInput
from bokeh.io import curdoc
from bokeh.sampledata.movies_data import movie_path

conn = sql.connect(movie_path)
query = open(join(dirname(__file__), 'query.sql')).read()
movies = psql.read_sql(query, conn)

movies["color"] = np.where(movies["Oscars"] > 0, "orange", "grey")
movies["alpha"] = np.where(movies["Oscars"] > 0, 0.9, 0.25)
movies.fillna(0, inplace=True)  # just replace missing values with zero
movies["revenue"] = movies.BoxOffice.apply(lambda x: '{:,d}'.format(int(x)))

with open(join(dirname(__file__), "razzies-clean.csv")) as f:
    razzies = f.read().splitlines()
movies.loc[movies.imdbID.isin(razzies), "color"] = "purple"
movies.loc[movies.imdbID.isin(razzies), "alpha"] = 0.9

axis_map = {
    "Tomato Meter": "Meter",
    "Numeric Rating": "numericRating",
    "Number of Reviews": "Reviews",
    "Box Office (dollars)": "BoxOffice",
    "Length (minutes)": "Runtime",
    "Year": "Year",
}

# Create Input controls
reviews = Slider(title="Minimum number of reviews", value=80, start=10, end=300, step=10)
min_year = Slider(title="Year released", start=1940, end=2014, value=1970, step=1)
max_year = Slider(title="End Year released", start=1940, end=2014, value=2014, step=1)
oscars = Slider(title="Minimum number of Oscar wins", start=0, end=4, value=0, step=1)
boxoffice = Slider(title="Dollars at Box Office (millions)", start=0, end=800, value=0, step=1)
genre = Select(title="Genre", value="All",
               options=open(join(dirname(__file__), 'genres.txt')).read().split())
director = TextInput(title="Director name contains")
cast = TextInput(title="Cast names contains")
x_axis = Select(title="X Axis", options=sorted(axis_map.keys()), value="Tomato Meter")
y_axis = Select(title="Y Axis", options=sorted(axis_map.keys()), value="Number of Reviews")

# Create Column Data Source that will be used by the plot
source = ColumnDataSource(data=dict(x=[], y=[], color=[], title=[], year=[], revenue=[]))

hover = HoverTool(tooltips=[
    ("Title","@title"),
    ("Year", "@year"),
    ("$", "@revenue")
])

p = Figure(plot_height=600, plot_width=800, title="", toolbar_location=None, tools=[hover])
p.circle(x="x", y="y", source=source, size=7, color="color", line_color=None, fill_alpha="alpha")

def select_movies():
    genre_val = genre.value
    director_val = director.value.strip()
    cast_val = cast.value.strip()
    selected = movies[
        (movies.Reviews >= reviews.value) &
        (movies.BoxOffice >= (boxoffice.value * 1e6)) &
        (movies.Year >= min_year.value) &
        (movies.Year <= max_year.value) &
        (movies.Oscars >= oscars.value)
    ]
    if (genre_val != "All"):
        selected = selected[selected.Genre.str.contains(genre_val)==True]
    if (director_val != ""):
        selected = selected[selected.Director.str.contains(director_val)==True]
    if (cast_val != ""):
        selected = selected[selected.Cast.str.contains(cast_val)==True]
    return selected

def update(attrname, old, new):
    df = select_movies()
    x_name = axis_map[x_axis.value]
    y_name = axis_map[y_axis.value]

    p.xaxis.axis_label = x_axis.value
    p.yaxis.axis_label = y_axis.value
    p.title = "%d movies selected" % len(df)
    source.data = dict(
        x=df[x_name],
        y=df[y_name],
        color=df["color"],
        title=df["Title"],
        year=df["Year"],
        revenue=df["revenue"],
        alpha=df["alpha"],
    )

controls = [reviews, boxoffice, genre, min_year, max_year, oscars, director, cast, x_axis, y_axis]
for control in controls:
    control.on_change('value', update)

inputs = HBox(VBoxForm(*controls), width=300)

update(None, None, None) # initial load of the data

curdoc().add_root(HBox(inputs, p, width=1100))
<commit_msg>New layout does better when not specifying sizes for boxes (#4544)<commit_after>from os.path import dirname, join

import numpy as np
import pandas.io.sql as psql
import sqlite3 as sql

from bokeh.plotting import Figure
from bokeh.models import ColumnDataSource, HoverTool, HBox, VBoxForm
from bokeh.models.widgets import Slider, Select, TextInput
from bokeh.io import curdoc
from bokeh.sampledata.movies_data import movie_path

conn = sql.connect(movie_path)
query = open(join(dirname(__file__), 'query.sql')).read()
movies = psql.read_sql(query, conn)

movies["color"] = np.where(movies["Oscars"] > 0, "orange", "grey")
movies["alpha"] = np.where(movies["Oscars"] > 0, 0.9, 0.25)
movies.fillna(0, inplace=True)  # just replace missing values with zero
movies["revenue"] = movies.BoxOffice.apply(lambda x: '{:,d}'.format(int(x)))

with open(join(dirname(__file__), "razzies-clean.csv")) as f:
    razzies = f.read().splitlines()
movies.loc[movies.imdbID.isin(razzies), "color"] = "purple"
movies.loc[movies.imdbID.isin(razzies), "alpha"] = 0.9

axis_map = {
    "Tomato Meter": "Meter",
    "Numeric Rating": "numericRating",
    "Number of Reviews": "Reviews",
    "Box Office (dollars)": "BoxOffice",
    "Length (minutes)": "Runtime",
    "Year": "Year",
}

# Create Input controls
reviews = Slider(title="Minimum number of reviews", value=80, start=10, end=300, step=10)
min_year = Slider(title="Year released", start=1940, end=2014, value=1970, step=1)
max_year = Slider(title="End Year released", start=1940, end=2014, value=2014, step=1)
oscars = Slider(title="Minimum number of Oscar wins", start=0, end=4, value=0, step=1)
boxoffice = Slider(title="Dollars at Box Office (millions)", start=0, end=800, value=0, step=1)
genre = Select(title="Genre", value="All",
               options=open(join(dirname(__file__), 'genres.txt')).read().split())
director = TextInput(title="Director name contains")
cast = TextInput(title="Cast names contains")
x_axis = Select(title="X Axis", options=sorted(axis_map.keys()), value="Tomato Meter")
y_axis = Select(title="Y Axis", options=sorted(axis_map.keys()), value="Number of Reviews")

# Create Column Data Source that will be used by the plot
source = ColumnDataSource(data=dict(x=[], y=[], color=[], title=[], year=[], revenue=[]))

hover = HoverTool(tooltips=[
    ("Title","@title"),
    ("Year", "@year"),
    ("$", "@revenue")
])

p = Figure(plot_height=600, plot_width=800, title="", toolbar_location=None, tools=[hover])
p.circle(x="x", y="y", source=source, size=7, color="color", line_color=None, fill_alpha="alpha")

def select_movies():
    genre_val = genre.value
    director_val = director.value.strip()
    cast_val = cast.value.strip()
    selected = movies[
        (movies.Reviews >= reviews.value) &
        (movies.BoxOffice >= (boxoffice.value * 1e6)) &
        (movies.Year >= min_year.value) &
        (movies.Year <= max_year.value) &
        (movies.Oscars >= oscars.value)
    ]
    if (genre_val != "All"):
        selected = selected[selected.Genre.str.contains(genre_val)==True]
    if (director_val != ""):
        selected = selected[selected.Director.str.contains(director_val)==True]
    if (cast_val != ""):
        selected = selected[selected.Cast.str.contains(cast_val)==True]
    return selected

def update(attrname, old, new):
    df = select_movies()
    x_name = axis_map[x_axis.value]
    y_name = axis_map[y_axis.value]

    p.xaxis.axis_label = x_axis.value
    p.yaxis.axis_label = y_axis.value
    p.title = "%d movies selected" % len(df)
    source.data = dict(
        x=df[x_name],
        y=df[y_name],
        color=df["color"],
        title=df["Title"],
        year=df["Year"],
        revenue=df["revenue"],
        alpha=df["alpha"],
    )

controls = [reviews, boxoffice, genre, min_year, max_year, oscars, director, cast, x_axis, y_axis]
for control in controls:
    control.on_change('value', update)

inputs = HBox(VBoxForm(*controls))

update(None, None, None) # initial load of the data

curdoc().add_root(HBox(inputs, p))
<|endoftext|>