text:
"""
Copyright (c) Facebook, Inc. and its affiliates.

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.

---

This code borrows heavily from the DimeNet implementation as part of
pytorch-geometric: https://github.com/rusty1s/pytorch_geometric. License:

---

Copyright (c) 2020 Matthias Fey <matthias.fey@tu-dortmund.de>

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

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

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

import torch
from torch import nn
from torch_geometric.nn import radius_graph
from torch_geometric.nn.acts import swish
from torch_geometric.nn.inits import glorot_orthogonal
from torch_geometric.nn.models.dimenet import (
    BesselBasisLayer,
    EmbeddingBlock,
    Envelope,
    ResidualLayer,
    SphericalBasisLayer,
)
from torch_geometric.utils import degree
from torch_scatter import scatter
from torch_sparse import SparseTensor

from ocpmodels.common.registry import registry
from ocpmodels.common.utils import (
    conditional_grad,
    get_pbc_distances,
    radius_graph_pbc,
)
#from ocpmodels.models.utils.pna import AGGREGATORS, SCALERS

try:
    import sympy as sym
except ImportError:
    sym = None

class InteractionPPBlock(torch.nn.Module):
    def __init__(
        self,
        hidden_channels,
        int_emb_size,
        basis_emb_size,
        num_spherical,
        num_radial,
        num_before_skip,
        num_after_skip,
        act=swish,
    ):
        super(InteractionPPBlock, self).__init__()
        self.act = act

        # Transformations of Bessel and spherical basis representations.
        self.lin_rbf1 = nn.Linear(num_radial, basis_emb_size, bias=False)
        self.lin_rbf2 = nn.Linear(basis_emb_size, hidden_channels, bias=False)
        self.lin_sbf1 = nn.Linear(
            num_spherical * num_radial, basis_emb_size, bias=False
        )
        self.lin_sbf2 = nn.Linear(basis_emb_size, int_emb_size, bias=False)

        # Dense transformations of input messages.
        self.lin_kj = nn.Linear(hidden_channels, hidden_channels)
        self.lin_ji = nn.Linear(hidden_channels, hidden_channels)

        # Embedding projections for interaction triplets.
        self.lin_down = nn.Linear(hidden_channels, int_emb_size, bias=False)
        self.lin_up = nn.Linear(int_emb_size, hidden_channels, bias=False)

        # Residual layers before and after skip connection.
        self.layers_before_skip = torch.nn.ModuleList(
            [
                ResidualLayer(hidden_channels, act)
                for _ in range(num_before_skip)
            ]
        )
        self.lin = nn.Linear(hidden_channels, hidden_channels)
        self.layers_after_skip = torch.nn.ModuleList(
            [
                ResidualLayer(hidden_channels, act)
                for _ in range(num_after_skip)
            ]
        )

        self.reset_parameters()

    def reset_parameters(self):
        glorot_orthogonal(self.lin_rbf1.weight, scale=2.0)
        glorot_orthogonal(self.lin_rbf2.weight, scale=2.0)
        glorot_orthogonal(self.lin_sbf1.weight, scale=2.0)
        glorot_orthogonal(self.lin_sbf2.weight, scale=2.0)

        glorot_orthogonal(self.lin_kj.weight, scale=2.0)
        self.lin_kj.bias.data.fill_(0)
        glorot_orthogonal(self.lin_ji.weight, scale=2.0)
        self.lin_ji.bias.data.fill_(0)

        glorot_orthogonal(self.lin_down.weight, scale=2.0)
        glorot_orthogonal(self.lin_up.weight, scale=2.0)

        for res_layer in self.layers_before_skip:
            res_layer.reset_parameters()
        glorot_orthogonal(self.lin.weight, scale=2.0)
        self.lin.bias.data.fill_(0)
        for res_layer in self.layers_after_skip:
            res_layer.reset_parameters()

    def forward(self, x, rbf, sbf, idx_kj, idx_ji):
        # Initial transformations.
        x_ji = self.act(self.lin_ji(x))
        x_kj = self.act(self.lin_kj(x))

        # Transformation via Bessel basis.
        rbf = self.lin_rbf1(rbf)
        rbf = self.lin_rbf2(rbf)
        x_kj = x_kj * rbf

        # Down-project embeddings and generate interaction triplet embeddings.
        x_kj = self.act(self.lin_down(x_kj))

        # Transform via 2D spherical basis.
        sbf = self.lin_sbf1(sbf)
        sbf = self.lin_sbf2(sbf)
        x_kj = x_kj[idx_kj] * sbf

        # Aggregate interactions and up-project embeddings.
        x_kj = scatter(x_kj, idx_ji, dim=0, dim_size=x.size(0))
        x_kj = self.act(self.lin_up(x_kj))

        h = x_ji + x_kj
        for layer in self.layers_before_skip:
            h = layer(h)
        h = self.act(self.lin(h)) + x
        for layer in self.layers_after_skip:
            h = layer(h)

        return h

class OutputPPBlock(torch.nn.Module):
    def __init__(
        self,
        num_radial,
        hidden_channels,
        out_emb_channels,
        out_channels,
        num_layers,
        act,
        aggregators,
        scalers,
    ):
        super(OutputPPBlock, self).__init__()
        self.act = act
        self.pna = False
        if aggregators and scalers:
            self.pna = True
            self.aggregators = [AGGREGATORS[aggr] for aggr in aggregators]
            self.scalers = [SCALERS[scale] for scale in scalers]
            # average degree computed from data
            self.avg_deg = {"lin": 33.53, "log": 3.47}
            self.aggr_down = nn.Linear(
                len(aggregators) * len(scalers) * hidden_channels,
                out_emb_channels,
                bias=True,
            )
        else:
            self.lin_up = nn.Linear(
                hidden_channels, out_emb_channels, bias=True
            )

        self.lin_rbf = nn.Linear(num_radial, hidden_channels, bias=False)
        self.lins = torch.nn.ModuleList()
        for _ in range(num_layers):
            self.lins.append(nn.Linear(out_emb_channels, out_emb_channels))
        self.lin = nn.Linear(out_emb_channels, out_channels, bias=False)

        self.reset_parameters()

    def reset_parameters(self):
        glorot_orthogonal(self.lin_rbf.weight, scale=2.0)
        if not self.pna:
            glorot_orthogonal(self.lin_up.weight, scale=2.0)
        for lin in self.lins:
            glorot_orthogonal(lin.weight, scale=2.0)
            lin.bias.data.fill_(0)
        self.lin.weight.data.fill_(0)

    def forward(self, x, rbf, i, num_nodes=None):
        x = self.lin_rbf(rbf) * x
        if self.pna:
            outs = [
                aggr(x, i, dim_size=num_nodes) for aggr in self.aggregators
            ]
            out = torch.cat(outs, dim=-1)
            deg = degree(i, num_nodes, dtype=x.dtype).view(-1, 1)
            outs = [scaler(out, deg, self.avg_deg) for scaler in self.scalers]
            x = torch.cat(outs, dim=-1)
            x = self.aggr_down(x)
        else:
            x = scatter(x, i, dim=0, dim_size=num_nodes)
            x = self.lin_up(x)
        for lin in self.lins:
            x = self.act(lin(x))
        return self.lin(x)

class DimeNetPlusPlus(torch.nn.Module):
    r"""DimeNet++ implementation based on https://github.com/klicperajo/dimenet.

    Args:
        hidden_channels (int): Hidden embedding size.
        out_channels (int): Size of each output sample.
        num_blocks (int): Number of building blocks.
        int_emb_size (int): Embedding size used for interaction triplets
        basis_emb_size (int): Embedding size used in the basis transformation
        out_emb_channels(int): Embedding size used for atoms in the output block
        num_spherical (int): Number of spherical harmonics.
        num_radial (int): Number of radial basis functions.
        cutoff: (float, optional): Cutoff distance for interatomic
            interactions. (default: :obj:`5.0`)
        envelope_exponent (int, optional): Shape of the smooth cutoff.
            (default: :obj:`5`)
        num_before_skip: (int, optional): Number of residual layers in the
            interaction blocks before the skip connection. (default: :obj:`1`)
        num_after_skip: (int, optional): Number of residual layers in the
            interaction blocks after the skip connection. (default: :obj:`2`)
        num_output_layers: (int, optional): Number of linear layers for the
            output blocks. (default: :obj:`3`)
        act: (function, optional): The activation funtion.
            (default: :obj:`swish`)
    """

    url = "https://github.com/klicperajo/dimenet/raw/master/pretrained"

    def __init__(
        self,
        hidden_channels,
        out_channels,
        num_blocks,
        int_emb_size,
        basis_emb_size,
        out_emb_channels,
        num_spherical,
        num_radial,
        cutoff=5.0,
        envelope_exponent=5,
        num_before_skip=1,
        num_after_skip=2,
        num_output_layers=3,
        act=swish,
        aggregators=None,
        scalers=None,
    ):
        super(DimeNetPlusPlus, self).__init__()

        self.cutoff = cutoff

        if sym is None:
            raise ImportError("Package `sympy` could not be found.")

        self.num_blocks = num_blocks

        self.rbf = BesselBasisLayer(num_radial, cutoff, envelope_exponent)
        self.sbf = SphericalBasisLayer(
            num_spherical, num_radial, cutoff, envelope_exponent
        )

        self.emb = EmbeddingBlock(num_radial, hidden_channels, act)

        self.output_blocks = torch.nn.ModuleList(
            [
                OutputPPBlock(
                    num_radial,
                    hidden_channels,
                    out_emb_channels,
                    out_channels,
                    num_output_layers,
                    act,
                    aggregators,
                    scalers,
                )
                for _ in range(num_blocks + 1)
            ]
        )

        self.interaction_blocks = torch.nn.ModuleList(
            [
                InteractionPPBlock(
                    hidden_channels,
                    int_emb_size,
                    basis_emb_size,
                    num_spherical,
                    num_radial,
                    num_before_skip,
                    num_after_skip,
                    act,
                )
                for _ in range(num_blocks)
            ]
        )

        self.reset_parameters()

    def reset_parameters(self):
        self.rbf.reset_parameters()
        self.emb.reset_parameters()
        for out in self.output_blocks:
            out.reset_parameters()
        for interaction in self.interaction_blocks:
            interaction.reset_parameters()

    def triplets(self, edge_index, num_nodes):
        row, col = edge_index  # j->i

        value = torch.arange(row.size(0), device=row.device)
        adj_t = SparseTensor(
            row=col, col=row, value=value, sparse_sizes=(num_nodes, num_nodes)
        )
        adj_t_row = adj_t[row]
        num_triplets = adj_t_row.set_value(None).sum(dim=1).to(torch.long)

        # Node indices (k->j->i) for triplets.
        idx_i = col.repeat_interleave(num_triplets)
        idx_j = row.repeat_interleave(num_triplets)
        idx_k = adj_t_row.storage.col()
        mask = idx_i != idx_k  # Remove i == k triplets.
        idx_i, idx_j, idx_k = idx_i[mask], idx_j[mask], idx_k[mask]

        # Edge indices (k-j, j->i) for triplets.
        idx_kj = adj_t_row.storage.value()[mask]
        idx_ji = adj_t_row.storage.row()[mask]

        return col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji

    def forward(self, z, pos, batch=None):
        """"""
        raise NotImplementedError

@registry.register_model("dimenetplusplus")
class DimeNetPlusPlusWrap(DimeNetPlusPlus):
    def __init__(
        self,
        num_atoms,
        bond_feat_dim,  # not used
        num_targets,
        use_pbc=True,
        regress_forces=True,
        regress_position=False,
        hidden_channels=128,
        num_blocks=4,
        int_emb_size=64,
        basis_emb_size=8,
        out_emb_channels=256,
        num_spherical=7,
        num_radial=6,
        otf_graph=False,
        cutoff=10.0,
        envelope_exponent=5,
        num_before_skip=1,
        num_after_skip=2,
        num_output_layers=3,
        aggregators=None,
        scalers=None,
    ):
        self.regress_position = regress_position

        if self.regress_position:
            self.num_targets = 3
        else:
            self.num_targets = num_targets

        self.regress_forces = regress_forces
        self.use_pbc = use_pbc
        self.cutoff = cutoff
        self.otf_graph = otf_graph

        super(DimeNetPlusPlusWrap, self).__init__(
            hidden_channels=hidden_channels,
            out_channels=self.num_targets,
            num_blocks=num_blocks,
            int_emb_size=int_emb_size,
            basis_emb_size=basis_emb_size,
            out_emb_channels=out_emb_channels,
            num_spherical=num_spherical,
            num_radial=num_radial,
            cutoff=cutoff,
            envelope_exponent=envelope_exponent,
            num_before_skip=num_before_skip,
            num_after_skip=num_after_skip,
            num_output_layers=num_output_layers,
            aggregators=aggregators,
            scalers=scalers,
        )

    @conditional_grad(torch.enable_grad())
    def _forward(self, data):
        pos = data.pos
        batch = data.batch

        if self.otf_graph:
            edge_index, cell_offsets, neighbors = radius_graph_pbc(
                data, self.cutoff, 50, data.pos.device
            )
            data.edge_index = edge_index
            data.cell_offsets = cell_offsets
            data.neighbors = neighbors

        if self.use_pbc:
            out = get_pbc_distances(
                pos,
                data.edge_index,
                data.cell,
                data.cell_offsets,
                data.neighbors,
                return_offsets=True,
            )

            edge_index = out["edge_index"]
            dist = out["distances"]
            offsets = out["offsets"]

            j, i = edge_index
        else:
            edge_index = radius_graph(pos, r=self.cutoff, batch=batch)
            j, i = edge_index
            dist = (pos[i] - pos[j]).pow(2).sum(dim=-1).sqrt()

        _, _, idx_i, idx_j, idx_k, idx_kj, idx_ji = self.triplets(
            edge_index, num_nodes=data.atomic_numbers.size(0)
        )

        # Calculate angles.
        pos_i = pos[idx_i].detach()
        pos_j = pos[idx_j].detach()
        if self.use_pbc:
            pos_ji, pos_kj = (
                pos[idx_j].detach() - pos_i + offsets[idx_ji],
                pos[idx_k].detach() - pos_j + offsets[idx_kj],
            )
        else:
            pos_ji, pos_kj = (
                pos[idx_j].detach() - pos_i,
                pos[idx_k].detach() - pos_j,
            )

        a = (pos_ji * pos_kj).sum(dim=-1)
        b = torch.cross(pos_ji, pos_kj).norm(dim=-1)
        angle = torch.atan2(b, a)

        rbf = self.rbf(dist)
        sbf = self.sbf(dist, angle, idx_kj)

        # Embedding block.
        x = self.emb(data.atomic_numbers.long(), rbf, i, j)
        P = self.output_blocks[0](x, rbf, i, num_nodes=pos.size(0))

        # Interaction blocks.
        for interaction_block, output_block in zip(
            self.interaction_blocks, self.output_blocks[1:]
        ):
            x = interaction_block(x, rbf, sbf, idx_kj, idx_ji)
            P += output_block(x, rbf, i, num_nodes=pos.size(0))

        energy = (
            P.sum(dim=0)
            if data.batch is None
            else scatter(P, data.batch, dim=0)
        )

        return P, energy

    def forward(self, data):
        if self.regress_forces:
            data.pos.requires_grad_(True)

        P, energy = self._forward(data)

        if self.regress_forces:
            forces = -1 * (
                torch.autograd.grad(
                    energy,
                    data.pos,
                    grad_outputs=torch.ones_like(energy),
                    create_graph=True,
                )[0]
            )
            return energy, forces
        elif self.regress_position:
            return P
        else:
            return energy

    @property
    def num_params(self):
        return sum(p.numel() for p in self.parameters())

max_stars_repo_name: AdeeshKolluru/ocp

max_issues_repo_name: AdeeshKolluru/ocp

max_forks_repo_name: AdeeshKolluru/ocp

reason: import sympy


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

"""
The ``driver`` module implements Numba's interface to the dpctl library, a
lightweight Python and C wrapper to a subset of DPC++'s runtime API. The
``driver`` module includes:

- LLVM IR builders for dpctl C API functions to be called directly from a Numba
  generated LLVM module.
- Functions to lauch kernels on the dpctl "current queue".

"""
from numba_dppy.driver.dpctl_capi_fn_builder import DpctlCAPIFnBuilder
from numba_dppy.driver.kernel_launch_ops import KernelLaunchOps
from numba_dppy.driver.usm_ndarray_type import USMNdArrayType

__all__ = [
    DpctlCAPIFnBuilder,
    KernelLaunchOps,
    USMNdArrayType,
]

ext: py

max_stars_repo_name: reazulhoque/numba-dppy

max_issues_repo_name: reazulhoque/numba-dppy

max_forks_repo_name: reazulhoque/numba-dppy

reason: from numba


text:
# Copyright (c) 2020 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""An algorithm for adjusting entropy regularization strength."""
from absl import logging
import gin
import math
import numpy as np
import torch

import alf
from alf.algorithms.algorithm import Algorithm
from alf.data_structures import namedtuple, AlgStep, LossInfo, StepType
from alf.summary import should_record_summaries
from alf.utils.averager import ScalarWindowAverager
from alf.utils.dist_utils import calc_default_target_entropy, entropy_with_fallback

EntropyTargetLossInfo = namedtuple("EntropyTargetLossInfo", ["neg_entropy"])
EntropyTargetInfo = namedtuple("EntropyTargetInfo", ["loss"])

@gin.configurable
class EntropyTargetAlgorithm(Algorithm):
    """Algorithm for adjusting entropy regularization.

    It tries to adjust the entropy regularization (i.e. alpha) so that the
    the entropy is not smaller than ``target_entropy``.

    The algorithm has three stages:

    0. init stage. This is an optional stage. If the initial entropy is already
       below ``max_entropy``, then this stage is skipped. Otherwise, the alpha will
       be slowly decreased so that the entropy will land at ``max_entropy`` to
       trigger the next ``free_stage``. Basically, this stage let the user to choose
       an arbitrary large init alpha without considering every specific case.
    1. free stage. During this stage, the alpha is not changed. It transitions
       to adjust_stage once entropy drops below ``target_entropy``.
    2. adjust stage. During this stage, ``log_alpha`` is adjusted using this formula:

       .. code-block:: python

            ((below + 0.5 * above) * decreasing - (above + 0.5 * below) * increasing) * update_rate

       Note that ``log_alpha`` will always be decreased if entropy is increasing
       even when the entropy is below the target entropy. This is to prevent
       overshooting ``log_alpha`` to a too big value. Same reason for always
       increasing ``log_alpha`` even when the entropy is above the target entropy.
       ``update_rate`` is initialized to ``fast_update_rate`` and is reduced by a
       factor of 0.9 whenever the entropy crosses ``target_entropy``. ``udpate_rate``
       is reset to ``fast_update_rate`` if entropy drops too much below
       ``target_entropy`` (i.e., ``fast_stage_thresh`` in the code, which is the half
       of ``target_entropy`` if it is positive, and twice of ``target_entropy`` if
       it is negative.

    ``EntropyTargetAlgorithm`` can be used to approximately reproduce the learning
    of temperature in `Soft Actor-Critic Algorithms and Applications <https://arxiv.org/abs/1812.05905>`_.
    To do so, you need to use the same ``target_entropy``, set ``skip_free_stage``
    to True, and  set ``slow_update_rate`` and ``fast_update_rate`` to the 4
    times of the learning rate for temperature.
    """

    def __init__(self,
                 action_spec,
                 initial_alpha=0.1,
                 skip_free_stage=False,
                 max_entropy=None,
                 target_entropy=None,
                 very_slow_update_rate=0.001,
                 slow_update_rate=0.01,
                 fast_update_rate=np.log(2),
                 min_alpha=1e-4,
                 average_window=2,
                 debug_summaries=False,
                 name="EntropyTargetAlgorithm"):
        """
        Args:
            action_spec (nested BoundedTensorSpec): representing the actions.
            initial_alpha (float): initial value for alpha; make sure that it's
                large enough for initial meaningful exploration
            skip_free_stage (bool): If True, directly goes to the adjust stage.
            max_entropy (float): the upper bound of the entropy. If not provided,
                ``min(initial_entropy * 0.8, initial_entropy / 0.8)`` is used.
                initial_entropy is estimated from the first ``average_window``
                steps. 0.8 is to ensure that we can get a policy a less random
                as the initial policy before starting the free stage.
            target_entropy (float): the lower bound of the entropy. If not
                provided, a default value proportional to the action dimension
                is used. This value should be less or equal than ``max_entropy``.
            very_slow_update_rate (float): a tiny update rate for ``log_alpha``;
                used in stage 0.
            slow_update_rate (float): minimal update rate for ``log_alpha``; used
                in stage 2.
            fast_update_rate (float): maximum update rate for ``log_alpha``; used
                in state 2.
            min_alpha (float): the minimal value of alpha. If <=0, :math:`e^{-100}`
                is used.
            average_window (int): window size for averaging past entropies.
            debug_summaries (bool): True if debug summaries should be created.
        """
        super().__init__(debug_summaries=debug_summaries, name=name)

        self.register_buffer(
            '_log_alpha',
            torch.tensor(np.log(initial_alpha), dtype=torch.float32))
        self.register_buffer('_stage', torch.tensor(-2, dtype=torch.int32))
        self._avg_entropy = ScalarWindowAverager(average_window)
        self.register_buffer(
            "_update_rate", torch.tensor(
                fast_update_rate, dtype=torch.float32))
        self._action_spec = action_spec
        self._min_log_alpha = -100.
        if min_alpha >= 0.:
            self._min_log_alpha = np.log(min_alpha)
        self._min_log_alpha = torch.tensor(self._min_log_alpha)

        flat_action_spec = alf.nest.flatten(self._action_spec)
        if target_entropy is None:
            target_entropy = np.sum(
                list(map(calc_default_target_entropy, flat_action_spec)))
            logging.info("target_entropy=%s" % target_entropy)

        if max_entropy is None:
            # max_entropy will be estimated in the first `average_window` steps.
            max_entropy = 0.
            self._stage.fill_(-2 - average_window)
        else:
            assert target_entropy <= max_entropy, (
                "Target entropy %s should be less or equal than max entropy %s!"
                % (target_entropy, max_entropy))
        self.register_buffer("_max_entropy",
                             torch.tensor(max_entropy, dtype=torch.float32))

        if skip_free_stage:
            self._stage.fill_(1)

        if target_entropy > 0:
            self._fast_stage_thresh = 0.5 * target_entropy
        else:
            self._fast_stage_thresh = 2.0 * target_entropy
        self._target_entropy = target_entropy
        self._very_slow_update_rate = very_slow_update_rate
        self._slow_update_rate = torch.tensor(slow_update_rate)
        self._fast_update_rate = torch.tensor(fast_update_rate)

    def rollout_step(self, distribution, step_type, on_policy_training):
        """Rollout step.

        Args:
            distribution (nested Distribution): action distribution from the
                policy.
            step_type (StepType): the step type for the distributions.
            on_policy_training (bool): If False, this step does nothing.

        Returns:
            AlgStep: ``info`` field is ``LossInfo``, other fields are empty. All
            fields are empty If ``on_policy_training=False``.
        """
        if on_policy_training:
            return self.train_step(distribution, step_type)
        else:
            return AlgStep()

    def train_step(self, distribution, step_type):
        """Train step.

        Args:
            distribution (nested Distribution): action distribution from the
                policy.
            step_type (StepType): the step type for the distributions.
        Returns:
            AlgStep: ``info`` field is ``LossInfo``, other fields are empty.
        """
        entropy, entropy_for_gradient = entropy_with_fallback(distribution)
        return AlgStep(
            output=(),
            state=(),
            info=EntropyTargetInfo(
                loss=LossInfo(
                    loss=-entropy_for_gradient,
                    extra=EntropyTargetLossInfo(neg_entropy=-entropy))))

    def calc_loss(self, experience, info: EntropyTargetInfo, valid_mask=None):
        """Calculate loss.

        Args:
            experience (Experience): experience for gradient update
            info (EntropyTargetInfo): for computing loss.
            valid_mask (tensor): valid mask to be applied on time steps.

        Returns:
            LossInfo:
        """
        loss_info = info.loss
        mask = (experience.step_type != StepType.LAST).type(torch.float32)
        if valid_mask:
            mask = mask * (valid_mask).type(torch.float32)
        entropy = -loss_info.extra.neg_entropy * mask
        num = torch.sum(mask)
        not_empty = num > 0
        num = max(num, 1)
        entropy2 = torch.sum(entropy**2) / num
        entropy = torch.sum(entropy) / num
        entropy_std = torch.sqrt(
            torch.max(torch.tensor(0.0), entropy2 - entropy * entropy))

        if not_empty:
            self.adjust_alpha(entropy)
            if self._debug_summaries and should_record_summaries():
                with alf.summary.scope(self.name):
                    alf.summary.scalar("entropy_std", entropy_std)

        alpha = torch.exp(self._log_alpha)
        return loss_info._replace(loss=loss_info.loss * alpha)

    def adjust_alpha(self, entropy):
        """Adjust alpha according to the current entropy.

        Args:
            entropy (scalar Tensor): the current entropy.
        Returns:
            adjusted entropy regularization
        """
        prev_avg_entropy = self._avg_entropy.get()
        avg_entropy = self._avg_entropy.average(entropy)

        def _init_entropy():
            self._max_entropy.fill_(
                torch.min(0.8 * avg_entropy, avg_entropy / 0.8))
            self._stage.add_(1)

        def _init():
            below = avg_entropy < self._max_entropy
            decreasing = (avg_entropy < prev_avg_entropy).type(torch.float32)
            # -1 * (1 - decreasing) + 0.5 * decreasing
            update_rate = (-1 + 1.5 * decreasing) * self._very_slow_update_rate
            self._stage.add_(below.type(torch.int32))
            self._log_alpha.fill_(
                torch.max(self._log_alpha + update_rate, self._min_log_alpha))

        def _free():
            crossing = avg_entropy < self._target_entropy
            self._stage.add_(crossing.type(torch.int32))

        def _adjust():
            previous_above = self._stage.type(torch.bool)
            above = avg_entropy > self._target_entropy
            self._stage.fill_(above.type(torch.int32))
            crossing = above != previous_above
            update_rate = self._update_rate
            update_rate = torch.where(crossing, 0.9 * update_rate, update_rate)
            update_rate = torch.max(update_rate, self._slow_update_rate)
            update_rate = torch.where(entropy < self._fast_stage_thresh,
                                      self._fast_update_rate, update_rate)
            self._update_rate.fill_(update_rate)
            above = above.type(torch.float32)
            below = 1 - above
            decreasing = (avg_entropy < prev_avg_entropy).type(torch.float32)
            increasing = 1 - decreasing
            log_alpha = self._log_alpha + (
                (below + 0.5 * above) * decreasing -
                (above + 0.5 * below) * increasing) * update_rate
            log_alpha = torch.max(log_alpha, self._min_log_alpha)
            self._log_alpha.fill_(log_alpha)

        if self._stage < -2:
            _init_entropy()
        if self._stage == -2:
            _init()
        if self._stage == -1:
            _free()
        if self._stage >= 0:
            _adjust()
        alpha = torch.exp(self._log_alpha)

        if self._debug_summaries and should_record_summaries():
            with alf.summary.scope(self.name):
                alf.summary.scalar("alpha", alpha)
                alf.summary.scalar("avg_entropy", avg_entropy)
                alf.summary.scalar("stage", self._stage)
                alf.summary.scalar("update_rate", self._update_rate)

        return alpha

reason: import numpy


text:
# Copyright 2019-2020 ETH Zurich and the DaCe authors. All rights reserved.
import numpy as np
import dace
from dace.memlet import Memlet

# Create nested SDFG
nsdfg = dace.SDFG('nested')
nsdfg.add_stream('nout', dace.int32)
state = nsdfg.add_state()
t = state.add_tasklet('task', set(), {'o'}, 'o = 2')
w = state.add_write('nout')
state.add_edge(t, 'o', w, None, Memlet.simple('nout', '0'))

# Create SDFG
sdfg = dace.SDFG('nested_stream_test')
state = sdfg.add_state('a')

# Nodes
sdfg.add_stream('SB', dace.int32, transient=True)
sdfg.add_array('B', (2, ), dace.int32)
SB = state.add_access('SB')
B = state.add_write('B')
n = state.add_nested_sdfg(nsdfg, None, set(), {'nout'})
state.add_edge(n, 'nout', SB, None, Memlet.simple('SB', '0'))
state.add_nedge(SB, B, Memlet.simple('B', '0'))

if __name__ == '__main__':
    print('Nested stream test')

    Bdata = np.zeros([2], np.int32)
    sdfg(B=Bdata)

    B_regression = np.array([2, 0], dtype=np.int32)

    diff = np.linalg.norm(B_regression - Bdata)
    print("Difference:", diff)
    print("==== Program end ====")
    exit(0 if diff == 0 else 1)

max_stars_repo_name: targetsm/dace

max_issues_repo_name: targetsm/dace

max_forks_repo_name: targetsm/dace


text:
import matplotlib as mplt
mplt.use('Agg')

import matplotlib.pyplot as plt
from matplotlib import cm

from netCDF4 import Dataset

import sys, argparse
import numpy as np

def ext(data):
    s = data.shape
    ndata = np.zeros((s[0], s[1]+1))
    ndata[:, 0:-1] = data
    ndata[:, -1] = data[:, 0]
    return ndata

def ext_axis(lon):
    return np.append(lon, 360)

parser = argparse.ArgumentParser()
parser.add_argument('--data-file')
parser.add_argument('--domain-file')
parser.add_argument('--lev-file', type=str, default="")
parser.add_argument('--lev-varname', type=str, default="lev")
parser.add_argument('--output-dir')
parser.add_argument('--casename')

parser.add_argument('--varname-mean')
parser.add_argument('--varname-std')
parser.add_argument('--title', default="")
parser.add_argument('--subtitles', default="")
parser.add_argument('--colormap-mean', default="bwr")
parser.add_argument('--colormap-std', default="hot_r")
parser.add_argument('--auto-clevs', action="store_true", default=False)
parser.add_argument('--cmin-mean', type=float)
parser.add_argument('--cmax-mean', type=float)
parser.add_argument('--cmax-std', type=float)
parser.add_argument('--clevs', type=int)
parser.add_argument('--tick-levs-mean', type=int, default=-1)
parser.add_argument('--tick-levs-std', type=int, default=-1)
parser.add_argument('--clabel-mean', default="")
parser.add_argument('--clabel-std', default="")
parser.add_argument('--offset', type=float, default=0.0)
parser.add_argument('--scale', default="1.0")
parser.add_argument('--idx-t', type=int, default=-1)
parser.add_argument('--idx-x', type=int, default=-1)
parser.add_argument('--extra-filename', default="")
parser.add_argument('--land-transparent', action="store_true", default=False)
parser.add_argument('--central-longitude', type=float, default=180.0)
parser.add_argument('--figsize', type=str, default="20,8")

args = parser.parse_args()

figsize   = tuple(map(float, args.figsize.split(',')))
subtitles = tuple(map(str, args.subtitles.split(',')))

if args.lev_file == "":
    args.lev_file = args.data_file

f = Dataset(args.data_file, "r")
g = Dataset(args.domain_file, "r")
h = Dataset(args.lev_file, "r")

lon = g.variables["xc"][1, :]                   #-- read clon
lat = g.variables["yc"][:, 1]                   #-- read clat
lev = h.variables[args.lev_varname][:]          #-- read level

g.close()
h.close()

args.scale = eval(args.scale)

var_mean = f.variables[args.varname_mean]
var_std  = f.variables[args.varname_std]

if args.idx_x == -1:
    data_mean = var_mean[:, :, :]
    data_std  = var_std[:, :, :]
else:
    data_mean = var_mean[:, :, :, args.idx_x]
    data_std  = var_std[:, :, :, args.idx_x]

if data_mean.shape[0] != data_std.shape[0]:
    raise Exception("Data lengths of mean and std in time must match.")

data_N, Nz, Ny = data_mean.shape

if args.tick_levs_mean == -1:
    args.tick_levs_mean = args.clevs

if args.tick_levs_std == -1:
    args.tick_levs_std = args.clevs

"""

DOM = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])

for s in range(4):
    N = 0
    for m in range(s*3, (s+1)*3):
        N += DOM[m]
        _data_mean[s, :, :] +=  data_mean[m, :, :] * DOM[m]
        _data_var[s, :, :]  +=  ( data_mean[m, :, :]**2.0 + data_var[m, :, :] ) * DOM[m]

    _data_mean[s, :, :] /= N
    _data_var[s, :, :] = _data_var[s, :, :] / N - _data_mean[s, :, :] ** 2.0
    _data_std[s, :, :] = np.sqrt(_data_var[s, :, :])

_data_mean -= args.offset
_data_mean /= args.scale
_data_std  /= args.scale
"""

f.close()

clevels_mean = np.linspace(args.cmin_mean, args.cmax_mean, args.clevs+1)
clevels_std = np.linspace(0, args.cmax_std, args.clevs+1)
cmap_mean = cm.get_cmap(args.colormap_mean)
cmap_std  = cm.get_cmap(args.colormap_std)

tick_levels_mean = np.linspace(args.cmin_mean, args.cmax_mean, args.tick_levs_mean+1)
tick_levels_std = np.linspace(0, args.cmax_std, args.tick_levs_std+1)

fig, ax = plt.subplots(nrows=2, ncols=data_N, figsize=figsize)

fig.suptitle(args.title)

for s in range(data_N):

    ax0 = ax[0, s]
    ax1 = ax[1, s]

    _mean = data_mean[s, :, :]
    _std  = data_std[s, :, :]

    mappable_mean = ax0.contourf(lat, lev, _mean, clevels_mean, cmap=cmap_mean, extend="both")
    mappable_std  = ax1.contourf(lat, lev, _std,  clevels_std,  cmap=cmap_std,  extend="max")

    ax0.set_title(subtitles[s])
    ax1.set_xlabel("Latitude [deg]")
    ax0.set_xlim([-90, 90])
    ax1.set_xlim([-90, 90])
    ax0.set_xticks([-90, -60, -30, 0, 30, 60, 90])
    ax1.set_xticks([-90, -60, -30, 0, 30, 60, 90])

    for a in ax[:, s]:
        a.invert_yaxis()
        #a.set_global()
        #a.set_aspect('auto')

cb_mean = fig.colorbar(mappable_mean, ax=ax[0, :], orientation="vertical", ticks=tick_levels_mean)
cb_std  = fig.colorbar(mappable_std, ax=ax[1, :], orientation="vertical", ticks=tick_levels_std)

cb_mean.ax.set_ylabel(args.clabel_mean, rotation=90)
cb_std.ax.set_ylabel(args.clabel_std, rotation=90)

filename = "%s/%s_atm_meridional_mean_std_%s.png" % (args.output_dir, args.casename, args.extra_filename)
fig.savefig(filename, dpi=200)
print("Output %s" % (filename,))

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['ProgressParallel', 'unpack_body_models', 'fast_amass_unpack', 'viable_slice', 'global_index_map',
           'load_npz', 'AMASS']

# Cell
import argparse
import os
from shutil import unpack_archive
import joblib
from tqdm.auto import tqdm

class ProgressParallel(joblib.Parallel):
    def __call__(self, *args, **kwargs):
        with tqdm(total=kwargs["total"]) as self._pbar:
            del kwargs["total"]
            return joblib.Parallel.__call__(self, *args, **kwargs)

    def print_progress(self):
        self._pbar.total = self.n_dispatched_tasks
        self._pbar.n = self.n_completed_tasks
        self._pbar.refresh()

def unpack_body_models(tardir, outdir, n_jobs=1):
    tar_root, _, tarfiles = [x for x in os.walk(tardir)][0]
    tarfiles = [x for x in tarfiles if "tar" in x.split(".")]
    tarpaths = [os.path.join(tar_root, tar) for tar in tarfiles]
    for tarpath in tarpaths:
        print(f"{tarpath} extracting to {outdir}")
    ProgressParallel(n_jobs=n_jobs)(
        (joblib.delayed(unpack_archive)(tarpath, outdir) for tarpath in tarpaths),
        total=len(tarpaths),
    )

def fast_amass_unpack():
    parser = argparse.ArgumentParser(
        description="Unpack all the body model tar files in a directory to a target directory"
    )
    parser.add_argument(
        "tardir",
        type=str,
        help="Directory containing tar.bz2 body model files",
    )
    parser.add_argument(
        "outdir",
        type=str,
        help="Output directory",
    )
    parser.add_argument(
        "-n",
        default=1,
        type=int,
        help="Number of jobs to run the tar unpacking with",
    )
    args = parser.parse_args()
    unpack_body_models(args.tardir, args.outdir, n_jobs=args.n)

# Cell
import math
import numpy as np
import torch
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

# Cell
def viable_slice(cdata, keep):
    """
    Inspects a dictionary loaded from `.npz` numpy dumps
    and creates a slice of the viable indexes.
    args:

        - `cdata`: dictionary containing keys:
            ['poses', 'gender', 'mocap_framerate', 'betas',
             'marker_data', 'dmpls', 'marker_labels', 'trans']
        - `keep`: ratio of the file to keep, between zero and 1.,
            drops leading and trailing ends of the arrays

    returns:

        - viable: slice that can access frames in the arrays:
            cdata['poses'], cdata['marker_data'], cdata['dmpls'], cdata['trans']
    """
    assert (
        keep > 0.0 and keep <= 1.0
    ), "Proportion of array to keep must be between zero and one"
    n = cdata["poses"].shape[0]
    drop = (1.0 - keep) / 2.0
    return slice(int(n * drop), int(n * keep + n * drop))

# Cell
def global_index_map(npz_directory, overlapping, clip_length, keep=0.8):
    """
    args:
        - `npz_directory`: Directory containing `.npz` files
        - `overlapping`: Whether clips can overlap
        - `clip_length`:
    returns:
        - map from global index to corresponding file and array indexes
    """
    for r, d, f in os.walk(npz_directory):
        npz_files = [x for x in f if "npz" in x.split(".")]
        npz_paths = [os.path.join(npz_directory, r, x) for x in npz_files]
    # array slices for each file
    viable_slices = {
        npz_path: viable_slice(np.load(npz_path), keep=keep) for npz_path in npz_paths
    }
    # clip index -> array index
    def clip_to_array_index(i, array_slice):
        if not overlapping:
            i = i * clip_length
        return [i + array_slice.start + j for j in range(clip_length)]

    # length of a slice
    def lenslice(s):
        if overlapping:
            return (s.stop - s.start) - (clip_length - 1)
        else:
            return math.floor((s.stop - s.start) / clip_length)

    # global index -> file, relative index
    def find_array(i):
        global_index, j = 0, 0
        for npz_path in viable_slices:
            global_index += lenslice(viable_slices[npz_path])
            if i < global_index:
                return npz_path, i - j
            j = global_index

    # how many examples are there in this dataset, total
    n_examples = sum(lenslice(viable_slices[npz_path]) for npz_path in viable_slices)
    # create map function
    def global_to_array(i):
        npz_path, j = find_array(i)
        return npz_path, clip_to_array_index(j, viable_slices[npz_path])

    return global_to_array, n_examples

# Cell
def load_npz(npz_path, indexes):
    cdata = np.load(npz_path)

    # unpack and enforce data type
    poses = cdata["poses"][indexes].astype(np.float32)
    dmpls = cdata["dmpls"][indexes].astype(np.float32)
    trans = cdata["trans"][indexes].astype(np.float32)
    betas = np.repeat(
        cdata["betas"][np.newaxis].astype(np.float32), repeats=len(indexes), axis=0
    )

    def gender_to_int(g):
        # casting gender to integer will raise a warning in future
        g = str(g.astype(str))
        return {"male": -1, "neutral": 0, "female": 1}[g]

    gender = np.array([gender_to_int(cdata["gender"]) for _ in indexes])

    return dict(poses=poses, dmpls=dmpls, trans=trans, betas=betas, gender=gender)

# Cell
class AMASS(Dataset):
    def __init__(self, unpacked_directory, clip_length, overlapping, transform=None):
        self.global_to_array, self.n_examples = global_index_map(
            unpacked_directory, overlapping=overlapping, clip_length=clip_length
        )
        self.transform = transform

    def __len__(self):
        return self.n_examples

    def __getitem__(self, i):
        data = load_npz(*self.global_to_array(i))
        return {k: self.transform(data[k]) for k in data}


text:
"""
This script generates sample datasets that can be used to
compare analyses between different statistics platforms.
"""

import numpy as np
import pathlib
PATH_HERE = pathlib.Path(__file__).parent.resolve()
np.set_printoptions(suppress=True)

def generateNormalSet1():
    np.random.seed(0)
    vals = np.random.standard_normal(1000) * 12 + 34
    csvPath = PATH_HERE.joinpath("norm_1000_12_34.csv")
    np.savetxt(csvPath, vals, "%.05f")

if __name__ == "__main__":
    generateNormalSet1()

max_stars_repo_name: EFeru/ScottPlot

max_issues_repo_name: EFeru/ScottPlot

max_forks_repo_name: EFeru/ScottPlot


text:
'''
Max-Planck-Gesellschaft zur Foerderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights on this computer program.
Using this computer program means that you agree to the terms in the LICENSE file (https://flame.is.tue.mpg.de/modellicense) included
with the FLAME model. Any use not explicitly granted by the LICENSE is prohibited.

Copyright 2020 Max-Planck-Gesellschaft zur Foerderung der Wissenschaften e.V. (MPG). acting on behalf of its
Max Planck Institute for Intelligent Systems. All rights reserved.

More information about FLAME is available at http://flame.is.tue.mpg.de.
For comments or questions, please email us at flame@tue.mpg.de
'''

import numpy as np
import chumpy as ch
from os.path import join

from psbody.mesh import Mesh
from smpl_webuser.serialization import load_model
from sbody.mesh_distance import ScanToMesh
from sbody.robustifiers import GMOf
from sbody.alignment.objectives import sample_from_mesh
from fitting.landmarks import load_embedding, landmark_error_3d, mesh_points_by_barycentric_coordinates
from fitting.util import load_binary_pickle, write_simple_obj, safe_mkdir, get_unit_factor

# -----------------------------------------------------------------------------

def compute_approx_scale(lmk_3d, model, lmk_face_idx, lmk_b_coords, opt_options=None):
    """ function: compute approximate scale to align scan and model

    input:
        lmk_3d: input landmark 3d, in shape (N,3)
        model: FLAME face model
        lmk_face_idx, lmk_b_coords: landmark embedding, in face indices and barycentric coordinates
        opt_options: optimizaton options

    output:
        model.r: fitted result vertices
        model.f: fitted result triangulations (fixed in this code)
        parms: fitted model parameters

    """

    scale = ch.ones(1)
    scan_lmks = scale*ch.array(lmk_3d)
    model_lmks = mesh_points_by_barycentric_coordinates( model, model.f, lmk_face_idx, lmk_b_coords )
    lmk_err = scan_lmks-model_lmks

    # options
    if opt_options is None:
        print("fit_lmk3d(): no 'opt_options' provided, use default settings.")
        import scipy.sparse as sp
        opt_options = {}
        opt_options['disp']    = 1
        opt_options['delta_0'] = 0.1
        opt_options['e_3']     = 1e-4
        opt_options['maxiter'] = 2000
        sparse_solver = lambda A, x: sp.linalg.cg(A, x, maxiter=opt_options['maxiter'])[0]
        opt_options['sparse_solver'] = sparse_solver

    # on_step callback
    def on_step(_):
        pass

    ch.minimize( fun      = lmk_err,
                 x0       = [ scale, model.trans, model.pose[:3] ],
                 method   = 'dogleg',
                 callback = on_step,
                 options  = opt_options )
    return scale.r

# -----------------------------------------------------------------------------

def fit_scan(  scan,                        # input scan
               lmk_3d,                      # input scan landmarks
               model,                       # model
               lmk_face_idx, lmk_b_coords,  # landmark embedding
               weights,                     # weights for the objectives
               gmo_sigma,                   # weight of the robustifier
               shape_num=300, expr_num=100, opt_options=None ):

    """ function: fit FLAME model to a 3D scan

    input:
        scan: input scan
        lmk_3d: input landmark 3d, in shape (N,3)
        model: FLAME face model
        lmk_face_idx, lmk_b_coords: landmark embedding, in face indices and barycentric coordinates
        weights: weights for each objective
        shape_num, expr_num: numbers of shape and expression compoenents used
        opt_options: optimizaton options

    output:
        model.r: fitted result vertices
        model.f: fitted result triangulations (fixed in this code)
        parms: fitted model parameters

    """

    # variables
    shape_idx      = np.arange( 0, min(300,shape_num) )        # valid shape component range in "betas": 0-299
    expr_idx       = np.arange( 300, 300+min(100,expr_num) )   # valid expression component range in "betas": 300-399
    used_idx       = np.union1d( shape_idx, expr_idx )
    model.betas[:] = np.random.rand( model.betas.size ) * 0.0  # initialized to zero
    model.pose[:]  = np.random.rand( model.pose.size ) * 0.0   # initialized to zero
    free_variables = [model.trans, model.pose, model.betas[used_idx]]

    # weights
    print("fit_scan(): use the following weights:")
    for kk in weights.keys():
        print("fit_scan(): weights['%s'] = %f" % ( kk, weights[kk] ))

    # objectives
    # landmark error
    lmk_err = landmark_error_3d(mesh_verts=model, mesh_faces=model.f,  lmk_3d=lmk_3d, lmk_face_idx=lmk_face_idx, lmk_b_coords=lmk_b_coords)

    # scan-to-mesh distance, measuring the distance between scan vertices and the closest points in the model's surface
    sampler = sample_from_mesh(scan, sample_type='vertices')
    s2m = ScanToMesh(scan, model, model.f, scan_sampler=sampler, rho=lambda x: GMOf(x, sigma=gmo_sigma))

    # regularizer
    shape_err = weights['shape'] * model.betas[shape_idx]
    expr_err  = weights['expr']  * model.betas[expr_idx]
    pose_err  = weights['pose']  * model.pose[3:] # exclude global rotation

    objectives = {'s2m': weights['s2m']*s2m, 'lmk': weights['lmk']*lmk_err, 'shape': shape_err, 'expr': expr_err, 'pose': pose_err}

    # options
    if opt_options is None:
        print("fit_lmk3d(): no 'opt_options' provided, use default settings.")
        import scipy.sparse as sp
        opt_options = {}
        opt_options['disp']    = 1
        opt_options['delta_0'] = 0.1
        opt_options['e_3']     = 1e-4
        opt_options['maxiter'] = 2000
        sparse_solver = lambda A, x: sp.linalg.cg(A, x, maxiter=opt_options['maxiter'])[0]
        opt_options['sparse_solver'] = sparse_solver

    # on_step callback
    def on_step(_):
        pass

    # optimize
    # step 1: rigid alignment
    from time import time
    timer_start = time()
    print("\nstep 1: start rigid fitting...")
    ch.minimize( fun      = lmk_err,
                 x0       = [ model.trans, model.pose[:3] ],
                 method   = 'dogleg',
                 callback = on_step,
                 options  = opt_options )
    timer_end = time()
    print("step 1: fitting done, in %f sec\n" % ( timer_end - timer_start ))

    # step 2: non-rigid alignment
    timer_start = time()
    print("step 2: start non-rigid fitting...")
    ch.minimize( fun      = objectives,
                 x0       = free_variables,
                 method   = 'dogleg',
                 callback = on_step,
                 options  = opt_options )
    timer_end = time()
    print("step 2: fitting done, in %f sec\n" % ( timer_end - timer_start ))

    # return results
    parms = { 'trans': model.trans.r, 'pose': model.pose.r, 'betas': model.betas.r }
    return model.r, model.f, parms

def run_fitting():
    # input scan
    scan_path = './data/scan.obj'

    # landmarks of the scan
    scan_lmk_path = './data/scan_lmks.npy'

    # measurement unit of landmarks ['m', 'cm', 'mm', 'NA']
    # When using option 'NA', the scale of the scan will be estimated by rigidly aligning model and scan landmarks
    scan_unit = 'm'

    scan = Mesh(filename=scan_path)
    print("loaded scan from:", scan_path)

    lmk_3d = np.load(scan_lmk_path)
    print("loaded scan landmark from:", scan_lmk_path)

    # model
    model_path = './models/generic_model.pkl' # change to 'female_model.pkl' or 'male_model.pkl', if gender is known
    model = load_model(model_path)       # the loaded model object is a 'chumpy' object, check https://github.com/mattloper/chumpy for details
    print("loaded model from:", model_path)

    # landmark embedding
    lmk_emb_path = './models/flame_static_embedding.pkl'
    lmk_face_idx, lmk_b_coords = load_embedding(lmk_emb_path)
    print("loaded lmk embedding")

    # scale scans and scan landmarks to be in the same local coordinate systems as the FLAME model
    if scan_unit.lower() == 'na':
        print('No scale specifiec - compute approximate scale based on the landmarks')
        scale_factor = compute_approx_scale(lmk_3d, model, lmk_face_idx, lmk_b_coords)
        print('Scale factor: %f' % scale_factor)
    else:
        scale_factor = get_unit_factor('m') / get_unit_factor(scan_unit)
        print('Scale factor: %f' % scale_factor)

    scan.v[:] *= scale_factor
    lmk_3d[:] *= scale_factor

    # output
    output_dir = './output'
    safe_mkdir(output_dir)

    # weights
    weights = {}
    # scan vertex to model surface distance term
    weights['s2m']   = 2.0
    # landmark term
    weights['lmk']   = 1e-2
    # shape regularizer (weight higher to regularize face shape more towards the mean)
    weights['shape'] = 1e-4
    # expression regularizer (weight higher to regularize facial expression more towards the mean)
    weights['expr']  = 1e-4
    # regularization of head rotation around the neck and jaw opening (weight higher for more regularization)
    weights['pose']  = 1e-3
    # Parameter of the Geman-McClure robustifier (higher weight for a larger bassin of attraction which makes it less robust to outliers)
    gmo_sigma = 1e-4

    # optimization options
    import scipy.sparse as sp
    opt_options = {}
    opt_options['disp']    = 1
    opt_options['delta_0'] = 0.1
    opt_options['e_3']     = 1e-4
    opt_options['maxiter'] = 2000
    sparse_solver = lambda A, x: sp.linalg.cg(A, x, maxiter=opt_options['maxiter'])[0]
    opt_options['sparse_solver'] = sparse_solver

    # run fitting
    mesh_v, mesh_f, parms = fit_scan(   scan=scan,                                             # input scan
                                        lmk_3d=lmk_3d,                                         # input landmark 3d
                                        model=model,                                           # model
                                        lmk_face_idx=lmk_face_idx, lmk_b_coords=lmk_b_coords,  # landmark embedding
                                        weights=weights,                                       # weights for the objectives
                                        gmo_sigma=gmo_sigma,                                   # parameter of the regularizer
                                        shape_num=300, expr_num=100, opt_options=opt_options ) # options

    # write result
    output_path = join( output_dir, 'fit_scan_result.obj' )
    write_simple_obj( mesh_v=mesh_v, mesh_f=mesh_f, filepath=output_path, verbose=False )
    print('output mesh saved to: ', output_path)

    # output scaled scan for reference (output scan fit and the scan should be spatially aligned)
    output_path = join( output_dir, 'scan_scaled.obj' )
    write_simple_obj( mesh_v=scan.v, mesh_f=scan.f, filepath=output_path, verbose=False )

# -----------------------------------------------------------------------------

if __name__ == '__main__':
    run_fitting()

max_stars_repo_name: Rubikplayer/flame-fitting

max_issues_repo_name: Rubikplayer/flame-fitting

max_forks_repo_name: Rubikplayer/flame-fitting

reason: import numpy,import scipy


text:
"""
Tests for psi4wrapper.py
"""

import SuperCoolFast as scf
import numpy as np
import os

def test_opt():
    """
    We pick the physicist's water molecule in a relatively large
    basis set (aug-cc-pVDZ), whose uhf and rhf solutions are the
    same. And we test that energies of the following four cases
    are identical:
        RHF + DIIS
        RHF + ODA
        UHF + DIIS
        UHF + ODA
    Note for DIIS, we in addition require the convergence is achieved
    in 30 iterations.
    """
    # get integrals from psi4
    ao_ints, scf_params, e_ZZ_repul = \
        scf.init(os.path.dirname(__file__) + '/test_opt.yml')
    # RHF, DIIS
    scf_params['max_iter'] = 30
    eps, C, D, F = scf.scf(ao_ints, scf_params, e_ZZ_repul)
    E_rhf_diis = scf.get_SCF_energy(ao_ints['H'], F, D, False) + e_ZZ_repul
    # RHF, ODA
    scf_params['max_iter'] = 300
    scf_params['opt'] = "oda"
    eps, C, D, F = scf.scf(ao_ints, scf_params, e_ZZ_repul)
    E_rhf_oda = scf.get_SCF_energy(ao_ints['H'], F, D, False) + e_ZZ_repul
    # UHF, DIIS
    scf_params['max_iter'] = 30
    scf_params['opt'] = "diis"
    scf_params['unrestricted'] = True
    eps, C, D, F = scf.scf(ao_ints, scf_params, e_ZZ_repul)
    E_uhf_diis = scf.get_SCF_energy(ao_ints['H'], F, D, True) + e_ZZ_repul
    # UHF, ODA
    scf_params['max_iter'] = 300
    scf_params['opt'] = "oda"
    eps, C, D, F = scf.scf(ao_ints, scf_params, e_ZZ_repul)
    E_uhf_oda = scf.get_SCF_energy(ao_ints['H'], F, D, True) + e_ZZ_repul
    # RHF, psi4
    import psi4
    mol = psi4.geometry(scf_params['geometry'])
    mol.update_geometry()
    ctrl_string = "SCF/" + scf_params['basis']
    psi4.set_options({"scf_type": "pk"})
    psi4_E_rhf = psi4.energy(ctrl_string, molecule = mol)
    # check
    assert(np.allclose(E_rhf_diis, E_rhf_oda) == True)
    assert(np.allclose(E_rhf_diis, E_uhf_diis) == True)
    assert(np.allclose(E_rhf_diis, E_uhf_oda) == True)
    assert(np.allclose(E_rhf_diis, psi4_E_rhf) == True)

max_stars_repo_name: hongzhouye/sigma-SCF

max_issues_repo_name: hongzhouye/sigma-SCF

max_forks_repo_name: hongzhouye/sigma-SCF


text:
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Dec 21 12:01:40 2017

@author: www.github.com/GustavZ
"""
import numpy as np
import tensorflow as tf
import os
import sys

# Root directory of the project
ROOT_DIR = os.path.abspath("../src/realtime_object_detection")
sys.path.append(ROOT_DIR)

from helper import FPS, WebcamVideoStream, SessionWorker, conv_detect2track, conv_track2detect
from model import Model
from config import Config
from vis_utils import visualize_objectdetection
from tf_utils import reframe_box_masks_to_image_masks

def detection(model,config):
    # Tracker
    if config.USE_TRACKER:
        import sys
        sys.path.append(ROOT_DIR+'/kcf')
        import KCF
        tracker = KCF.kcftracker(False, True, False, False)
        tracker_counter = 0
        track = False

    print("> Building Graph")
    # tf Session Config
    tf_config = model.tf_config
    detection_graph = model.detection_graph
    category_index = model.category_index
    with detection_graph.as_default():
        with tf.Session(graph=detection_graph,config=tf_config) as sess:
            # start Videostream
            vs = WebcamVideoStream(config.CV_INPUT,config.WIDTH,config.HEIGHT).start()
            # Define Input and Ouput tensors
            tensor_dict = model.get_tensordict(['num_detections', 'detection_boxes', 'detection_scores','detection_classes', 'detection_masks'])
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            # Mask Transformations
            if 'detection_masks' in tensor_dict:
                # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
                detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
                detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
                real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
                detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
                detection_masks_reframed = reframe_box_masks_to_image_masks(
                                            detection_masks, detection_boxes, vs.real_height, vs.real_width)
                detection_masks_reframed = tf.cast(tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                # Follow the convention by adding back the batch dimension
                tensor_dict['detection_masks'] = tf.expand_dims(detection_masks_reframed, 0)
            if config.SPLIT_MODEL:
                score_out = detection_graph.get_tensor_by_name('Postprocessor/convert_scores:0')
                expand_out = detection_graph.get_tensor_by_name('Postprocessor/ExpandDims_1:0')
                score_in = detection_graph.get_tensor_by_name('Postprocessor/convert_scores_1:0')
                expand_in = detection_graph.get_tensor_by_name('Postprocessor/ExpandDims_1_1:0')
                # Threading
                score = model.score
                expand = model.expand
                gpu_worker = SessionWorker("GPU",detection_graph,tf_config)
                cpu_worker = SessionWorker("CPU",detection_graph,tf_config)
                gpu_opts = [score_out, expand_out]
                cpu_opts = [tensor_dict['detection_boxes'], tensor_dict['detection_scores'], tensor_dict['detection_classes'], tensor_dict['num_detections']]

            fps = FPS(config.FPS_INTERVAL).start()
            masks = None
            print('> Starting Detection')
            while vs.isActive():
                # Detection
                if not (config.USE_TRACKER and track):
                    if config.SPLIT_MODEL:
                        # split model in seperate gpu and cpu session threads
                        if gpu_worker.is_sess_empty():
                            # read video frame, expand dimensions and convert to rgb
                            frame = vs.read()
                            # put new queue
                            gpu_feeds = {image_tensor: vs.expanded()}
                            if config.VISUALIZE:
                                gpu_extras = frame # for visualization frame
                            else:
                                gpu_extras = None
                            gpu_worker.put_sess_queue(gpu_opts,gpu_feeds,gpu_extras)
                        g = gpu_worker.get_result_queue()
                        if g is None:
                            # gpu thread has no output queue. ok skip, let's check cpu thread.
                            pass
                        else:
                            # gpu thread has output queue.
                            score,expand,frame = g["results"][0],g["results"][1],g["extras"]

                            if cpu_worker.is_sess_empty():
                                # When cpu thread has no next queue, put new queue.
                                # else, drop gpu queue.
                                cpu_feeds = {score_in: score, expand_in: expand}
                                cpu_extras = frame
                                cpu_worker.put_sess_queue(cpu_opts,cpu_feeds,cpu_extras)
                        c = cpu_worker.get_result_queue()
                        if c is None:
                            # cpu thread has no output queue. ok, nothing to do. continue
                            continue # If CPU RESULT has not been set yet, no fps update
                        else:
                            boxes, scores, classes, num, frame = c["results"][0],c["results"][1],c["results"][2],c["results"][3],c["extras"]
                    else:
                        # default session
                        frame = vs.read()
                        output_dict = sess.run(tensor_dict, feed_dict={image_tensor: vs.expanded()})
                        num = output_dict['num_detections'][0]
                        classes = output_dict['detection_classes'][0]
                        boxes = output_dict['detection_boxes'][0]
                        scores = output_dict['detection_scores'][0]
                        if 'detection_masks' in output_dict:
                            masks = output_dict['detection_masks'][0]

                    # reformat detection
                    num = int(num)
                    boxes = np.squeeze(boxes)
                    classes = np.squeeze(classes).astype(np.uint8)
                    scores = np.squeeze(scores)

                    # Visualization
                    vis = visualize_objectdetection(frame,boxes,classes,scores,masks,category_index,fps._glob_numFrames,
                                                    config.MAX_FRAMES,fps.fps_local(),config.PRINT_INTERVAL,config.PRINT_TH,
                                                    config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                    if not vis:
                        break

                    # Activate Tracker
                    if config.USE_TRACKER and num <= config.NUM_TRACKERS:
                        tracker_frame = frame
                        track = True
                        first_track = True

                # Tracking
                else:
                    frame = vs.read()
                    if first_track:
                        trackers = []
                        tracker_boxes = boxes
                        for box in boxes[~np.all(boxes == 0, axis=1)]:
                                tracker.init(conv_detect2track(box,vs.real_width, vs.real_height), tracker_frame)
                                trackers.append(tracker)
                        first_track = False

                    for idx,tracker in enumerate(trackers):
                        tracker_box = tracker.update(frame)
                        tracker_boxes[idx,:] = conv_track2detect(tracker_box, vs.real_width, vs.real_height)
                    vis = visualize_objectdetection(frame,tracker_boxes,classes,scores,masks,category_index,fps._glob_numFrames,
                                                    config.MAX_FRAMES,fps.fps_local(),config.PRINT_INTERVAL,config.PRINT_TH,
                                                    config.OD_MODEL_NAME+config._DEV+config._OPT,config.VISUALIZE)
                    if not vis:
                        break

                    tracker_counter += 1
                    #tracker_frame = frame
                    if tracker_counter >= config.TRACKER_FRAMES:
                        track = False
                        tracker_counter = 0

                fps.update()

    # End everything
    vs.stop()
    fps.stop()
    if config.SPLIT_MODEL:
        gpu_worker.stop()
        cpu_worker.stop()

def main():
    config = Config()
    model = Model('od',config.OD_MODEL_NAME,config.OD_MODEL_PATH,config.LABEL_PATH,
                config.NUM_CLASSES,config.SPLIT_MODEL, config.SSD_SHAPE).prepare_od_model()
    detection(model, config)

if __name__ == '__main__':
    main()


text:
# Copyright (c) 2020, Huawei Technologies.All rights reserved.
#
# Licensed under the BSD 3-Clause License  (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import numpy as np
import sys
from common_utils import TestCase, run_tests
from common_device_type import dtypes, instantiate_device_type_tests
from util_test import create_common_tensor

def create_all_one_tensor(item):
    dtype = item[0]
    format = item[1]
    shape = item[2]
    input1 = np.ones(shape).astype(dtype)
    cpu_input = torch.from_numpy(input1)
    npu_input = torch.from_numpy(input1).to("npu")
    if format != -1:
        npu_input = npu_input.npu_format_cast(format)
    return cpu_input, npu_input

class TestAllclose(TestCase):
    def cpu_op_exec(self, input_x, input_y):
        output = torch.allclose(input_x, input_y)
        return output

    def npu_op_exec(self, input_x, input_y):
        output = torch.allclose(input_x, input_y)
        return output

    def test_allclose_random(self, device):
        test_cases = [
            [[np.float32, -1, (1, 2)], [np.float32, -1, (1, 2)]],
            [[np.float32, -1, (1234, 2234)], [np.float32, -1, (1234, 2234)]],
            [[np.float32, -1, (321, 421, 521)], [np.float32, -1, (421, 521)]],
            [[np.float32, -1, (1, 600)], [np.float32, -1, (400, 200, 1)]],
            [[np.float32, -1, (20, 30, 40, 1)], [np.float32, -1, (30, 40, 50)]],
            [[np.float16, -1, (1, 2)], [np.float16, -1, (1, 2)]],
            [[np.float16, -1, (1234, 2234)], [np.float16, -1, (1234, 2234)]],
            [[np.float16, -1, (321, 421, 521)], [np.float16, -1, (421, 521)]],
            [[np.float16, -1, (1, 600)], [np.float16, -1, (400, 200, 1)]],
            [[np.float16, -1, (20, 30, 40, 1)], [np.float16, -1, (30, 40, 50)]],
            [[np.int8, -1, (1, 2)], [np.int8, -1, (1, 2)]],
            [[np.int8, -1, (1234, 2234)], [np.int8, -1, (1234, 2234)]],
            [[np.int8, -1, (321, 421, 521)], [np.int8, -1, (421, 521)]],
            [[np.int8, -1, (1, 600)], [np.int8, -1, (400, 200, 1)]],
            [[np.int8, -1, (20, 30, 40, 1)], [np.int8, -1, (30, 40, 50)]],
            [[np.uint8, -1, (1, 2)], [np.uint8, -1, (1, 2)]],
            [[np.uint8, -1, (1234, 2234)], [np.uint8, -1, (1234, 2234)]],
            [[np.int32, -1, (1, 2)], [np.int32, -1, (1, 2)]],
            [[np.int32, -1, (1234, 2234)], [np.int32, -1, (1234, 2234)]],
            [[np.int32, -1, (321, 421, 521)], [np.int32, -1, (421, 521)]],
            [[np.int32, -1, (1, 600)], [np.int32, -1, (400, 200, 1)]],
            [[np.int32, -1, (20, 30, 40, 1)], [np.int32, -1, (30, 40, 50)]],

        ]
        for item in test_cases:
            cpu_input_x, npu_input_x = create_common_tensor(item[0], 0, 100)
            cpu_input_y, npu_input_y = create_common_tensor(item[1], 0, 100)

            if cpu_input_x.dtype == torch.float16:
                cpu_input_x = cpu_input_x.to(torch.float32)

            if cpu_input_y.dtype == torch.float16:
                cpu_input_y = cpu_input_y.to(torch.float32)

            cpu_output = np.array(self.cpu_op_exec(cpu_input_x, cpu_input_y))
            npu_output = np.array(self.npu_op_exec(npu_input_x, npu_input_y))
            self.assertRtolEqual(cpu_output, npu_output)

    def test_allclose_x_equal_y(self, device):
        test_cases = [
            [[np.float32, -1, (1, 2)], [np.float32, -1, (1, 2)]],
            [[np.float32, -1, (1234, 2234)], [np.float32, -1, (1234, 2234)]],
            [[np.float32, -1, (321, 421, 521)], [np.float32, -1, (421, 521)]],
            [[np.float32, -1, (1, 600)], [np.float32, -1, (400, 200, 1)]],
            [[np.float32, -1, (20, 30, 40, 1)], [np.float32, -1, (30, 40, 50)]],
            [[np.float16, -1, (1, 2)], [np.float16, -1, (1, 2)]],
            [[np.float16, -1, (1234, 2234)], [np.float16, -1, (1234, 2234)]],
            [[np.float16, -1, (321, 421, 521)], [np.float16, -1, (421, 521)]],
            [[np.float16, -1, (1, 600)], [np.float16, -1, (400, 200, 1)]],
            [[np.float16, -1, (20, 30, 40, 1)], [np.float16, -1, (30, 40, 50)]],
            [[np.int8, -1, (1, 2)], [np.int8, -1, (1, 2)]],
            [[np.int8, -1, (1234, 2234)], [np.int8, -1, (1234, 2234)]],
            [[np.int8, -1, (321, 421, 521)], [np.int8, -1, (421, 521)]],
            [[np.int8, -1, (1, 600)], [np.int8, -1, (400, 200, 1)]],
            [[np.int8, -1, (20, 30, 40, 1)], [np.int8, -1, (30, 40, 50)]],
            [[np.uint8, -1, (1, 2)], [np.uint8, -1, (1, 2)]],
            [[np.uint8, -1, (1234, 2234)], [np.uint8, -1, (1234, 2234)]],
            [[np.int32, -1, (1, 2)], [np.int32, -1, (1, 2)]],
            [[np.int32, -1, (1234, 2234)], [np.int32, -1, (1234, 2234)]],
            [[np.int32, -1, (321, 421, 521)], [np.int32, -1, (421, 521)]],
            [[np.int32, -1, (1, 600)], [np.int32, -1, (400, 200, 1)]],
            [[np.int32, -1, (20, 30, 40, 1)], [np.int32, -1, (30, 40, 50)]],
        ]
        for item in test_cases:
            cpu_input_x, npu_input_x = create_all_one_tensor(item[0])
            cpu_input_y, npu_input_y = create_all_one_tensor(item[1])

            if cpu_input_x.dtype == torch.float16:
                cpu_input_x = cpu_input_x.to(torch.float32)

            if cpu_input_y.dtype == torch.float16:
                cpu_input_y = cpu_input_y.to(torch.float32)

            cpu_output = np.array(self.cpu_op_exec(cpu_input_x, cpu_input_y))
            npu_output = np.array(self.npu_op_exec(npu_input_x, npu_input_y))
            self.assertRtolEqual(cpu_output, npu_output)

    def test_allclose_scalar_1(self, device):
        input_x = np.array([1e-08]).astype(np.float32)
        input_y = np.array([1e-09]).astype(np.float32)
        cpu_input_x = torch.from_numpy(input_x)
        npu_input_x = torch.from_numpy(input_x).to("npu")
        cpu_input_y = torch.from_numpy(input_y)
        npu_input_y = torch.from_numpy(input_y).to("npu")

        cpu_output = np.array(self.cpu_op_exec(cpu_input_x, cpu_input_y))
        npu_output = np.array(self.npu_op_exec(npu_input_x, npu_input_y))
        self.assertRtolEqual(cpu_output, npu_output)

    def test_allclose_scalar_2(self, device):
        input_x = np.array([1e-07]).astype(np.float32)
        input_y = np.array([1e-08]).astype(np.float32)
        cpu_input_x = torch.from_numpy(input_x)
        npu_input_x = torch.from_numpy(input_x).to("npu")
        cpu_input_y = torch.from_numpy(input_y)
        npu_input_y = torch.from_numpy(input_y).to("npu")

        cpu_output = np.array(self.cpu_op_exec(cpu_input_x, cpu_input_y))
        npu_output = np.array(self.npu_op_exec(npu_input_x, npu_input_y))
        self.assertRtolEqual(cpu_output, npu_output)

instantiate_device_type_tests(TestAllclose, globals(), except_for='cpu')
if __name__ == "__main__":
    torch.npu.set_device("npu:3")
    run_tests()

max_stars_repo_name: Ascend/pytorch

max_issues_repo_name: Ascend/pytorch

max_forks_repo_name: Ascend/pytorch


text:
import torch
from torch import nn
import time
import numpy as np
import torch.nn.functional as F
import copy
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from evaluate import score_BIO

cuda_flag = True and torch.cuda.is_available()

def train(model, train_iter, dev_iter, test_iter, args):
    # Train:
    # HyperParameter
    n_epochs = args.EPOCHS
    learning_rate = args.lr

    n_gpu = torch.cuda.device_count()
    print("CUDA: " +str(cuda_flag))
    if cuda_flag:
        model = model.cuda()
    # rnn = TextCNN(300, output_size, max_length)
    if args.optimizer == "SGD":
        optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    if args.optimizer == "Adam":
        optimizer = torch.optim.Adam(model.parameters())
    if args.optimizer == "Adadelta":
        optimizer = torch.optim.Adadelta(model.parameters())
    if args.freeze:
        for param in model.word_rep.word_embed.parameters():
            param.requires_grad = False
    print_every = 100
    plot_every = 30
    # Keep track of losses for plotting
    current_loss = []
    all_losses = []
    test_acc = []
    acc_index = 0
    es_len = 0

    start = time.time()
    best_score = 0

    # scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
    # np.random.seed([3, 1415])

    for epoch in range(int(n_epochs)):
        loss_sum = 0
        tr_loss = 0
        tr_loss_avi = 0
        nb_tr_examples, nb_tr_steps = 0, 0
        for step, batch in enumerate(train_iter):

            model.train()
            output_p = model(batch)
            loss = 0
            if model.use_crf:
                loss = model.crf._neg_log_likelihood(output_p, batch.label)
            else:
                # for output, label in zip(output_p, batch.label):
                #     loss += F.nll_loss(output, label, ignore_index=-1)
                loss = F.nll_loss(output_p.view(-1, 3), batch.label.view(-1), ignore_index=-1)
            if n_gpu > 1:
                loss = loss.mean()
            loss.backward()
            optimizer.step()
            model.zero_grad()

            current_loss.append(loss.item())
            loss_sum += loss.item()

            tr_loss += loss.item()
            tr_loss_avi += loss.item() * batch.batch_size
            nb_tr_examples += batch.batch_size
            nb_tr_steps += 1

            if (step+1) % plot_every == 0:
                # print(batch_epoch)
                all_losses.append(sum(current_loss) / len(current_loss))
                current_loss = []

        with torch.no_grad():
            eval_dict = eval(model, dev_iter, args)
        print("Epoch:%d" % epoch)
        print("DEV: p:%.4f, r:%.4f, f:%.4f" % (eval_dict['precision'], eval_dict['recall'], eval_dict['f1']))
        with torch.no_grad():
            test_dict = eval(model, test_iter, args)
        print("TEST: p:%.4f, r:%.4f, f:%.4f" % (test_dict['precision'], test_dict['recall'], test_dict['f1']))
        # print("OPINION: p:%.4f, r:%.4f, f:%.4f" % (pred_acc_t2[0], pred_acc_t2[1], pred_acc_t2[2]))

        if eval_dict['main_metric'] > best_score:
            best_score = eval_dict['main_metric']
            max_print = ("Epoch%d\n" % epoch
                         + "DEV: p:%.4f, r:%.4f, f:%.4f\n" % (eval_dict['precision'], eval_dict['recall'], eval_dict['f1'])
                         + "TEST: p:%.4f, r:%.4f, f:%.4f\n" % (test_dict['precision'], test_dict['recall'], test_dict['f1']))
            best_model = "backup/%s_%.4f_%.4f.pt" % (args.model, best_score, test_dict['main_metric'])
            best_dict = copy.deepcopy(model)

        print("Epoch: %d, loss: %.4f" % (epoch, loss_sum))

    torch.save(best_dict, best_model)
    print("Best Result:")
    print(max_print)
    # plt.figure()
    # plt.plot(all_losses)
    time_stamp = time.asctime().replace(':', '_').split()
    print(time_stamp)
    # plt.savefig("fig/foor_%s.png" % '_'.join(time_stamp))
    # plt.show()
    return best_dict

def category_from_output(output):

    top_n, top_i = output.topk(1) # Tensor out of Variable with .data
    # print(top_i)
    category_i = top_i.view(output.size()[0], -1).detach().cpu().numpy().tolist()
    return category_i

def eval(model, dev_iter, args):
    model.eval()

    logit_list = []
    labels_eval_list = []
    predicted_result = []
    for eval_batch in dev_iter:
        output = model(eval_batch)
        if args.use_crf:
            _, category_i_p = model.crf._viterbi_decode(output)
        else:
            category_i_p = category_from_output(output)
        # print(eval_batch.target)
        # print(output)
        # print(category_i_p)
        predicted_result.extend(category_i_p)
        label_ids = eval_batch.label.to('cpu').numpy()
        labels_eval_list.extend(label_ids.tolist())
        logits = output.detach().cpu().numpy()
        logit_list.extend(logits.tolist())

    eval_dict = {}

    score_dict = score_BIO(predicted_result, labels_eval_list, ignore_index=-1)

    eval_dict.update(score_dict)
    eval_dict['main_metric'] = score_dict['f1']
    return eval_dict


text:
import qiskit
import numpy as np
import sys
from math import pi
import qiskit.providers.aer.noise as noise
from deconstruct import *
from hqca.hamiltonian import *
from hqca.instructions import *
from hqca.processes import *
from hqca.acse import *
from hqca.core import *
from hqca.core.primitives import *
from pyscf import gto
from hqca.transforms import *
from functools import partial
from hqca.tools import *
from hqca.state_tomography import *
from tests._generic import *
from delayed_assert import delayed_assert as da
np.set_printoptions(precision=3)

class Ins(Instructions):
    def __init__(self,operator,
            Nq=4,propagate=False,**kw):
        try:
            operator = operator.op_form()
        except AttributeError as Exception:
            pass
        if Nq==2:
            self.use = []
            self._gates = []
            a1 = ['XY','YX']
            for op in operator:
                if op.s==a1[0]:
                    self.use.append('a1')
                if len(self.use)==3:
                    break
            get_funcs = {
                    'a1':self._type_A1,
                    }
            for u in self.use[:3]:
                self._gates.append(
                        [(operator,),get_funcs[u]]
                        )
        if propagate:
            self._applyH(Nq=Nq,**kw)

    @property
    def gates(self):
        return self._gates

    @gates.setter
    def gates(self,a):
        self._gates = a

    def _type_A1(self,Q,op):
        '''
        connectvitiy?

        0-1-2-3 (old)
        ---->

        1-2-3-0
        '''
        paulis = ['YX','XY']
        para = []
        for p in paulis:
            for o in op:
                if o.s==p:
                    para.append(o.c.imag)
        Q.s(0)
        Q.s(1)
        Q.h(1)
        Q.Cx(1,0)
        Q.Ry(1,para[0])
        Q.Ry(0,-para[1])
        Q.Cx(1,0)
        Q.h(1)
        Q.si(1)
        Q.si(0)

    def _applyH(self,
            HamiltonianOperator,
            scaleH=0.5,Nq=2,**kw):
        self.hss = scaleH
        if Nq==2:
            self.use_H = ['h0']
            get_funcs = {
                    'h0':self._type_H0,
                    }
            for u in self.use_H[:]:
                self._gates.append(
                        [(HamiltonianOperator,),get_funcs[u]]
                        )

    def _type_H0(self,Q,op):
        h0 = ['ZI','IZ','ZZ','XX','YY']
        para = []
        for p in h0:
            for o in op:
                if o.s==p:
                    para.append(o.c.real*self.hss)
        #
        Q.Rz(0,para[0])
        Q.Rz(1,para[1])

        Q.Cx(0,1)
        Q.Rz(1,para[2])
        Q.Cx(0,1)

        Q.h(1)
        Q.Cx(1,0)
        Q.Ry(1,para[3])
        Q.Ry(0,para[4])
        Q.Cx(1,0)
        Q.h(1)
        Q.si(1)
        Q.z(0)

def test_noisy():
    ham = generic_molecular_hamiltonian()
    mol = generic_mol()
    s = Stabilizer(ham._qubOp,verbose=True)
    s.gaussian_elimination()
    s.find_symmetry_generators()
    Tr,iTr = get_transform_from_symmetries(
            Transform=JordanWigner,
            symmetries=['ZIIZ','ZIZ'],
            qubits=[3,2],
            eigvals=[-1,+1],
            )
    ham = MolecularHamiltonian(mol,transform=Tr)

    #
    #fermi = Operator([
    #    FermiString(coeff=0.5,ops='+-p',indices=[2,3,1],N=4),
    #    FermiString(coeff=0.5,ops='+-p',indices=[3,2,1],N=4),
    #    ])
    #print(fermi.transform(Tr))
    #sys.exit()

    st = StorageACSE(ham)
    qs = QuantumStorage()
    pr = StabilizerProcess('filter_diagonal')
    Nq = 2
    Ns = 8192
    error = 1.0
    qs.set_algorithm(st)
    qs.set_backend(
            backend='qasm_simulator',
            Nq=Nq,
            num_shots=Ns,
            provider='Aer')
    nm = model_v2(scaling=error)
    qs.set_noise_model(custom=True,
            noise_model=nm)
    qs.initial_transform=iTr
    proc = StabilizerProcess('filter_diagonal')
    qs.set_error_mitigation('ansatz_shift',coeff=1.0,protocol='current')
    # run 1
    ins = Line
    tomoRe = ReducedTomography(qs,verbose=False)
    tomoRe.generate(real=True,imag=True,
            simplify=True,transform=Tr,
            method='gt',strategy='lf')
    tomoIm = ReducedTomography(qs,verbose=False)
    tomoIm.generate(real=True,imag=True,
            simplify=True,transform=Tr,
            method='gt',strategy='lf')
    #st = StorageACSE(ham,closed_ansatz=-2)
    st = StorageACSE(ham,closed_ansatz=True)
    acse = RunACSE(
            st,qs,ins,processor=proc,
            method='newton',
            update='quantum',
            opt_thresh=1e-10,
            S_thresh_rel=1e-6,
            S_min=1e-6,
            use_trust_region=True,
            convergence_type='trust',
            split_ansatz=True,

            hamiltonian_step_size=0.5,
            max_iter=5,
            initial_trust_region=2.0,
            newton_step=-1,
            restrict_S_size=0.5,
            tomo_S = tomoIm,
            tomo_Psi = tomoRe,
            verbose=True,
            )
    acse.build()
    da.expect(abs(acse.e0+0.783792654277353)<=1e-10)
    acse.run()
    #da.expect(abs(acse.e0+0.846147736093)<=1e-8)
    da.assert_expectations()
    #   #
    #  #
    # #
    ##
test_noisy()

max_stars_repo_name: damazz/HQCA

max_issues_repo_name: damazz/HQCA

max_forks_repo_name: damazz/HQCA


text:
"""
Residue-related classes --- :mod:`prolif.residue`
=================================================
"""
import re
from typing import Optional, List
from collections import UserDict
import numpy as np
from .rdkitmol import BaseRDKitMol

_RE_RESID = re.compile(r'([A-Z]{,3})?(\d*)\.?(\w)?')
NoneType = type(None)

class ResidueId:
    """A unique residue identifier

    Parameters
    ----------
    name : str
        3-letter residue name
    number : int
        residue number
    chain : str or None, optionnal
        1-letter protein chain
    """
    def __init__(self,
                 name: str = "UNK",
                 number: int = 0,
                 chain: Optional[str] = None):
        self.name = name or "UNK"
        self.number = number or 0
        self.chain = chain or None

    def __repr__(self):
        return f"ResidueId({self.name}, {self.number}, {self.chain})"

    def __str__(self):
        resid = f"{self.name}{self.number}"
        if self.chain:
            resid += f".{self.chain}"
        return resid

    def __hash__(self):
        return hash((self.name, self.number, self.chain))

    def __eq__(self, other):
        return hash(self) == hash(other)

    def __lt__(self, other):
        return (self.chain, self.number) < (other.chain, other.number)

    @classmethod
    def from_atom(cls, atom):
        """Creates a ResidueId from an RDKit atom

        Parameters
        ----------
        atom : rdkit.Chem.rdchem.Atom
            An atom that contains an RDKit :class:`~rdkit.Chem.rdchem.AtomMonomerInfo`
        """
        mi = atom.GetMonomerInfo()
        if mi:
            name = mi.GetResidueName()
            number = mi.GetResidueNumber()
            chain = mi.GetChainId()
            return cls(name, number, chain)
        return cls()

    @classmethod
    def from_string(cls, resid_str):
        """Creates a ResidueId from a string

        Parameters
        ----------
        resid_str : str
            A string in the format ``<3-letter code><residue number>.<chain>``
            All arguments are optionnal, and the dot should be present only if
            the chain identifier is also present

        Examples
        --------

        +-----------+----------------------------------+
        | string    | Corresponding ResidueId          |
        +===========+==================================+
        | "ALA10.A" | ``ResidueId("ALA", 10, "A")``    |
        +-----------+----------------------------------+
        | "GLU33"   | ``ResidueId("GLU", 33, None)``   |
        +-----------+----------------------------------+
        | "LYS.B"   | ``ResidueId("LYS", 0, "B")``     |
        +-----------+----------------------------------+
        | "ARG"     | ``ResidueId("ARG", 0, None)``    |
        +-----------+----------------------------------+
        | "5.C"     | ``ResidueId("UNK", 5, "C")``     |
        +-----------+----------------------------------+
        | "42"      | ``ResidueId("UNK", 42, None)``   |
        +-----------+----------------------------------+
        | ".D"      | ``ResidueId("UNK", 0, "D")``     |
        +-----------+----------------------------------+
        | ""        | ``ResidueId("UNK", 0, None)``    |
        +-----------+----------------------------------+

        """
        matches = _RE_RESID.search(resid_str)
        name, number, chain = matches.groups()
        number = int(number) if number else 0
        return cls(name, number, chain)

class Residue(BaseRDKitMol):
    """A class for residues as RDKit molecules

    Parameters
    ----------
    mol : rdkit.Chem.rdchem.Mol
        The residue as an RDKit molecule

    Attributes
    ----------
    resid : prolif.residue.ResidueId
        The residue identifier

    Notes
    -----
    The name of the residue can be converted to a string by using
    ``str(Residue)``
    """
    def __init__(self, mol):
        super().__init__(mol)
        self.resid = ResidueId.from_atom(mol.GetAtomWithIdx(0))

    def __repr__(self): # pragma: no cover
        name = ".".join([self.__class__.__module__, self.__class__.__name__])
        return f"<{name} {self.resid} at {id(self):#x}>"

    def __str__(self):
        return str(self.resid)

class ResidueGroup(UserDict):
    """A container to store and retrieve Residue instances easily

    Parameters
    ----------
    residues : list
        A list of :class:`~prolif.residue.Residue`

    Attributes
    ----------
    n_residues : int
        Number of residues in the ResidueGroup

    Notes
    -----
    Residues in the group can be accessed by :class:`ResidueId`, string, or
    index. See the :class:`~prolif.molecule.Molecule` class for an example.
    You can also use the :meth:`~prolif.residue.ResidueGroup.select` method to
    access a subset of a ResidueGroup.
    """
    def __init__(self, residues: List[Residue]):
        self._residues = np.asarray(residues, dtype=object)
        resinfo = [(r.resid.name, r.resid.number, r.resid.chain)
                   for r in self._residues]
        try:
            name, number, chain = zip(*resinfo)
        except ValueError:
            self.name = np.array([], dtype=object)
            self.number = np.array([], dtype=np.uint8)
            self.chain = np.array([], dtype=object)
        else:
            self.name = np.asarray(name, dtype=object)
            self.number = np.asarray(number, dtype=np.uint16)
            self.chain = np.asarray(chain, dtype=object)
        super().__init__([(r.resid, r) for r in self._residues])

    def __getitem__(self, key):
        # bool is a subclass of int but shouldn't be used here
        if isinstance(key, bool):
            raise KeyError("Expected a ResidueId, int, or str, "
                           f"got {type(key).__name__!r} instead")
        if isinstance(key, int):
            return self._residues[key]
        elif isinstance(key, str):
            key = ResidueId.from_string(key)
            return self.data[key]
        elif isinstance(key, ResidueId):
            return self.data[key]
        raise KeyError("Expected a ResidueId, int, or str, "
                       f"got {type(key).__name__!r} instead")

    def select(self, mask):
        """Locate a subset of a ResidueGroup based on a boolean mask

        Parameters
        ----------
        mask : numpy.ndarray
            A 1D array of ``dtype=bool`` with the same length as the number of
            residues in the ResidueGroup. The mask should be constructed by
            using conditions on the "name", "number", and "chain" residue
            attributes as defined in the :class:`~prolif.residue.ResidueId`
            class

        Returns
        -------
        rg : prolif.residue.ResidueGroup
            A subset of the original ResidueGroup

        Examples
        --------
        ::

            >>> rg
            <prolif.residue.ResidueGroup with 200 residues at 0x7f9a68719ac0>
            >>> rg.select(rg.chain == "A")
            <prolif.residue.ResidueGroup with 42 residues at 0x7fe3fdb86ca0>
            >>> rg.select((10 <= rg.number) & (rg.number < 30))
            <prolif.residue.ResidueGroup with 20 residues at 0x7f5f3c69aaf0>
            >>> rg.select((rg.chain == "B") & (np.isin(rg.name, ["ASP", "GLU"])))
            <prolif.residue.ResidueGroup with 3 residues at 0x7f5f3c510c70>

        As seen in these examples, you can combine masks with different
        operators, similarly to numpy boolean indexing or pandas
        :meth:`~pandas.DataFrame.loc` method

            * AND --> ``&``
            * OR --> ``|``
            * XOR --> ``^``
            * NOT --> ``~``

        """
        return ResidueGroup(self._residues[mask])

    def __repr__(self): # pragma: no cover
        name = ".".join([self.__class__.__module__, self.__class__.__name__])
        return f"<{name} with {self.n_residues} residues at {id(self):#x}>"

    @property
    def n_residues(self):
        return len(self)

max_stars_repo_name: liudongliangHI/ProLIF

max_issues_repo_name: liudongliangHI/ProLIF

max_forks_repo_name: liudongliangHI/ProLIF


text:
from __future__ import print_function, division
import sys
import itertools
try:
    import pathlib
except ImportError:
    import pathlib2 as pathlib
from skimage.exposure import adjust_gamma
import json
import numpy as np
import pytiff
from PIL import Image
from matplotlib import colors
import csv
import json
import yaml
from collections import OrderedDict

def composite_channel(target, image, color, range_min, range_max):
    ''' Render _image_ in pseudocolor and composite into _target_
    Args:
        target: Numpy float32 array containing composition target image
        image: Numpy uint16 array of image to render and composite
        color: Color as r, g, b float array, 0-1
        range_min: Threshhold range minimum, 0-65535
        range_max: Threshhold range maximum, 0-65535
    '''
    f_image = (image.astype('float32') - range_min) / (range_max - range_min)
    f_image = f_image.clip(0,1, out=f_image)
    for i, component in enumerate(color):
        target[:, :, i] += f_image * component

def render_color_tiles(input_file, output_dir, tile_size, num_channels, config_rows):

    EXT = 'jpg'

    print('Processing:', str(input_file))

    output_path = pathlib.Path(output_dir)
    tiff = pytiff.Tiff(str(input_file), encoding='utf-8')

    assert tiff.number_of_pages % num_channels == 0, "Pyramid/channel mismatch"

    num_levels = tiff.number_of_pages // num_channels

    render_groups = OrderedDict()
    for i in config_rows:
        if i['Group'] not in render_groups.keys():
            render_groups[i['Group']] = {
                'Channel Number': [],
                'Low': [],
                'High': [],
                'Color': [],
                'Marker Name': []
            }
        for key in render_groups[i['Group']]:
            render_groups[i['Group']][key].append(i[key])

    for level in range(num_levels):

        page_base = level * num_channels
        tiff.set_page(page_base)
        ny = int(np.ceil(tiff.shape[0] / tile_size))
        nx = int(np.ceil(tiff.shape[1] / tile_size))
        print('    level {} ({} x {})'.format(level, ny, nx))

        for ty, tx in itertools.product(range(0, ny), range(0, nx)):

            iy = ty * tile_size
            ix = tx * tile_size
            filename = '{}_{}_{}.{}'.format(level, tx, ty, EXT)

            for name, settings in render_groups.items():

                channels = '--'.join(
                    settings['Channel Number'][i] + '__' + settings['Marker Name'][i]
                    for i in range(len(settings['Color']))
                )

                group_dir = name.replace(' ', '-') + '_' + channels
                if not (output_path / group_dir).exists():
                    (output_path / group_dir).mkdir(parents=True)
                for i, (marker, color, start, end) in enumerate(zip(
                    settings['Channel Number'], settings['Color'],
                    settings['Low'], settings['High']
                )):
                    tiff.set_page(page_base + int(marker))
                    tile = tiff[iy:iy+tile_size, ix:ix+tile_size]
                    tile = adjust_gamma(tile, 1/2.2)
                    if i == 0:
                        target = np.zeros(tile.shape + (3,), np.float32)
                    composite_channel(
                        target, tile, colors.to_rgb(color), float(start), float(end)
                    )
                np.clip(target, 0, 1, out=target)
                target_u8 = (target * 255).astype(np.uint8)
                img = Image.frombytes('RGB', target.T.shape[1:], target_u8.tobytes())
                img.save(str(output_path / group_dir / filename), quality=85)

max_stars_repo_name: thejohnhoffer/minerva-author

max_issues_repo_name: thejohnhoffer/minerva-author

max_forks_repo_name: thejohnhoffer/minerva-author