text:
from datetime import datetime

import numpy as np
import pandas as pd

from create_sector_grid_data import SectorProcessor
from hagelslag.processing.ObjectMatcher import centroid_distance, time_distance
from hagelslag.processing.ObjectMatcher import shifted_centroid_distance

# date_index = pd.DatetimeIndex([pd.Timestamp.utcnow().strftime("%Y%m%d")])
work_path = "/ai-hail/aburke/2018_HREFv2_data/"
scratch_path = "/hail/aburke/testing_weights/figuring_out_files/"
ensemble_members = ['nam_00', 'nam_12', 'arw_00', 'arw_12', 'nssl_00', 'nssl_12', 'nmmb_00', 'nmmb_12']

config = dict(dates=None,
              start_hour=12,
              end_hour=36,
              watershed_variable="MAXUVV",
              ensemble_name="HREFv2",
              ensemble_members=ensemble_members,
              model_path=work_path + 'symbolic_hrefv2_data',
              model_watershed_params=(8, 1, 80, 100, 60),
              size_filter=12,
              gaussian_window=2,
              mrms_path=work_path + "MRMS/",
              mrms_variable="MESH_Max_60min_00.50",
              mrms_watershed_params=(19, 1, 100, 100, 75),
              object_matcher_params=([shifted_centroid_distance], np.array([1.0]),
                                     np.array([24000])),

              track_matcher_params=([centroid_distance, time_distance],
                                    np.array([80000, 2])),

              storm_variables=['MAXUVV', 'Storm relative helicity_3000', 'Storm relative helicity_1000',
                               'MAXREF', 'MXUPHL_5000', 'MAXDVV'],

              potential_variables=['Precipitable water_0', 'Temperature_1000', 'Dew point temperature_1000',
                                   'Geopotential Height_500', 'Temperature_500',
                                   'Dew point temperature_500', 'U component of wind_500', 'V component of wind_500',
                                   'Geopotential Height_700', 'Temperature_700', 'Dew point temperature_700',
                                   'U component of wind_700',
                                   'V component of wind_700', 'Geopotential Height_850', 'Temperature_850',
                                   'Dew point temperature_850',
                                   'U component of wind_850', 'V component of wind_850', 'MAXUW', 'MAXVW',
                                   'Surface lifted index', 'Convective available potential energy_0',
                                   'Convective inhibition_0'],

              tendency_variables=[],
              shape_variables=["area", "eccentricity", "major_axis_length", "minor_axis_length", "orientation",
                               "extent"],
              variable_statistics=["mean", "max", "min", "std", "skew",
                                   "percentile_10", "percentile_50", "percentile_90"],
              csv_path=scratch_path + "track_data_spring2018_MAXUVV_closest_csv/",
              geojson_path=scratch_path + "track_data_spring2018_MAXUVV_closest_json/",
              nc_path=scratch_path + "track_data_spring2018_MAXUVV_patch_nc/",
              unique_matches=True,
              patch_radius=16,
              closest_matches=True,
              match_steps=True,
              train=False,
              single_step=True,
              label_type="gamma",
              model_map_file="/hail/aburke/hagelslag/mapfiles/hrefv2_2018_map.txt",
              mask_file="/hail/aburke/hagelslag/mapfiles/hrefv2_us_mask.nc",
              run_date_format="%Y%m%d-%H%M",
              sector_csv_outpath=scratch_path + 'track_data_spring2018_MAXUVV_closest_csv/',
              sector_nc_outpath=scratch_path + 'track_data_spring2018_MAXUVV_patch_nc/',
              sector_mapfile="/hail/aburke/hagelslag/mapfiles/hrefv2_sectors/C_2018_map.txt"
              )

sector = 'E'

date_index = pd.DatetimeIndex(start="2018-05-01T00:00", end="2018-06-02T00:00", freq="1D")

sector_mapfile = "/hail/aburke/hagelslag/mapfiles/hrefv2_sectors/{0}_2018_map.txt".format(sector)
ensemble_name = "HREFv2"
member = 'arw_00'
run_date = datetime(2018, 7, 1)
run_date_format = "%Y%m%d-%H%M"
csv_path = '/hail/aburke/HREF_Scripts_and_Data/hwt_2018/track_data_spring2018_MAXUVV_closest_csv/'
nc_path = '/hail/aburke/HREF_Scripts_and_Data/hwt_2018/track_data_spring2018_MAXUVV_patch_nc/'
sector_csv_outpath = "/hail/aburke/testing_weights/sector_trained/{0}/track_data_spring2018_MAXUVV_closest_csv/".format(
    sector)
sector_nc_outpath = "/hail/aburke/testing_weights/sector_trained/{0}/track_data_spring2018_MAXUVV_patch_nc/".format(
    sector)
patch_radius = 16

for run_date in date_index:
    for member in ensemble_members:
        sector = SectorProcessor(sector_mapfile,
                                 ensemble_name, member,
                                 run_date, run_date_format)

        sector.output_sector_netcdf(nc_path, sector_nc_outpath, patch_radius, config)

        # for key in ['track_step', 'track_total']:
        #    sector.output_sector_csv(csv_path,key,sector_csv_outpath)

max_stars_repo_name: djgagne/hagelslag

max_issues_repo_name: djgagne/hagelslag

max_forks_repo_name: djgagne/hagelslag


text:
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# 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.
"""Some generally useful tensorflow operations."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v1 as tf
# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.data.ops import dataset_ops

# pylint: enable=g-direct-tensorflow-import

def affine(x, output_size, weight_name, bias_name=None, weight_init=None):
  """Affine transformation of the input `x`.

  Args:
    x: <float32>[..., x_dim]
    output_size: size of the last output dimension
    weight_name: Name of the weight variable to use
    bias_name: Name for the bias variable, if one should be used
    weight_init: Initializer of the weight variable

  Returns:
    transformed <float32>[..., `output_size`]
  """
  dim = x.shape.as_list()[-1]
  w = tf.get_variable(
      weight_name, (dim, output_size), tf.float32, initializer=weight_init)
  out = tf.tensordot(x, w, [[len(x.shape) - 1], [0]])
  if bias_name:
    b = tf.get_variable(
        bias_name, (output_size,),
        tf.float32,
        initializer=tf.zeros_initializer())
    for _ in range(len(out.shape) - 1):
      b = tf.expand_dims(b, 0)
    out += b
  return out

def last_dim_weighted_sum(x, weight_name, weight_init=None, keepdims=False):
  """Computes a weighted sum of the last dimension of `x`.

  Args:
    x: <float32>[..., x_dim]
    weight_name: Name of the weight variable to use
    weight_init: Initializer of the weight variable
    keepdims: Whether the output should hav an ending size one dim

  Returns:
    summed: <float32>[...] or <float32>[..., 1] iff `keepdims`
  """

  dim = x.shape.as_list()[-1]
  w = tf.get_variable(weight_name, dim, initializer=weight_init)
  out = tf.tensordot(x, w, [[len(x.shape) - 1], [0]])
  if keepdims:
    return tf.expand_dims(out, len(out.shape))
  else:
    return out

def mask_logits(vec, mask):
  """Mask `vec` in log-space.

  Elements in `vec` that are not in `mask` will be set to be very
  negative, so that when no longer in log-space (e.i., after `tf.exp(vec)`)
  their values will be very close to zero.

  Args:
    vec: <float32>[...] tensor to mask
    mask: Either a None (in which case this is a no-op), a boolean or 0/1 float
      mask that matches all, or all but the last, dimensions of `vec`, or 1-D
      integer length mask

  Raises:
    ValueError: If `mask` cannot be matched to `vec`
  Returns:
    masked: vec:<float32>[...]
  """
  if mask is None:
    return vec

  if mask.dtype == tf.int32:
    # Assume `mask` holds sequence lengths
    if len(vec.shape) not in [2, 3]:
      raise ValueError("Can't use a length mask on tensor of rank>3")
    mask = tf.sequence_mask(mask, tf.shape(vec)[1], tf.float32)
  else:
    mask = tf.to_float(mask)

  if len(mask.shape) == (len(vec.shape) - 1):
    mask = tf.expand_dims(mask, len(vec.shape) - 1)

  return vec * mask - (1 - mask) * 1E20

def _lowercase(x):
  # Specify `np.object` to avoid incorrectly returning a np.string type arr
  return np.array([w.lower() for w in x], np.object)

def lowercase_op(string_tensor):
  """Lowercase an arbitrarily sized string tensor."""

  shape = tf.shape(string_tensor)
  lc = tf.py_func(_lowercase, [tf.reshape(string_tensor, [-1])], tf.string,
                  False)
  return tf.reshape(lc, shape)

def bucket_by_quantiles(len_fn, batch_size, n_buckets, hist_bounds):
  """Dynamically bucket a `tf.data.Dataset` based on the element's length.

  Keeps track of a histogram of the input elements lengths, and yields batches
  of examples that belong to the same length quantile.

  Useful in cases where you want to bucket data, but don't know what the
  optimal bucketing ranges should be

  Args:
    len_fn: Function mapping elements in the dataset to an integer length
    batch_size: Maximum size of the output batches
    n_buckets: Number of quantiles to break the data into
    hist_bounds: List of integer bounds to use when building the histograms,
      should cover a range so that at most a single quantile of elements are
      lower/higher then the bucket range. More fine-grained buckets will make
      the histogram more precise, but adds to the computational overhead

  Raises:
    ValueError: If `hist_bounds` or `len_fn` are invalid
  Returns:
    A function that can be used with tf.data.Dataset.apply to batch a dataset
  """
  n_hist_binds = len(hist_bounds)

  if n_hist_binds < n_buckets:
    raise ValueError("Requested %d buckets, but only have %d histogram bins" %
                     (n_buckets, n_hist_binds))
  if any(hist_bounds[i] >= hist_bounds[i + 1] for i in range(n_hist_binds - 1)):
    raise ValueError("Bins must be descending")

  # Need to use `use_resource = True` to make this work correctly
  # within tf.data.Dataset
  hist_counts = tf.get_local_variable(
      "hist-counts",
      n_hist_binds + 1,
      tf.int64,
      tf.zeros_initializer(),
      use_resource=True)
  hist_bounds = tf.constant(hist_bounds, tf.int64)

  def bucket_fn(x):
    """Compute the element bucket and update the histogram."""
    ix = len_fn(x)
    if ix.dtype == tf.int32:
      ix = tf.to_int64(ix)
    elif ix.dtype != tf.int64:
      raise ValueError("Len function returned a non-int")

    adds_to_bins = tf.to_int64(tf.greater(hist_bounds, ix))
    # pad with a 1 for the "larger than all" bin
    adds_to_bins = tf.pad(adds_to_bins, [[0, 1]], constant_values=1)
    new_counts = tf.assign_add(hist_counts, adds_to_bins)
    bin_ix = n_hist_binds - tf.reduce_sum(adds_to_bins)
    # Computes the quantile based on the counts of the exammple's bucket
    bucket_ix = tf.floordiv(((n_buckets - 1) * new_counts[bin_ix]),
                            new_counts[-1])
    return bucket_ix

  def reduce_fn(_, x):
    return x.padded_batch(batch_size, dataset_ops.get_legacy_output_shapes(x))

  return tf.data.experimental.group_by_window(bucket_fn, reduce_fn, batch_size)

ext: py

max_stars_repo_name: naveenjafer/language

max_issues_repo_name: naveenjafer/language

max_forks_repo_name: naveenjafer/language

reason: import numpy


text:
# Copyright 2018-2020 Xanadu Quantum Technologies Inc.

# 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.
"""
Unit tests for the :mod:`pennylane.collection` submodule.
"""
from collections.abc import Sequence

import pytest
from pennylane import numpy as np

import pennylane as qml

try:
    import torch
except ImportError as e:
    pass

try:
    import tensorflow as tf

    if tf.__version__[0] == "1":
        print(tf.__version__)
        import tensorflow.contrib.eager as tfe

        tf.enable_eager_execution()
        Variable = tfe.Variable
    else:
        from tensorflow import Variable
except ImportError as e:
    pass

@pytest.fixture
def qnodes(interface, tf_support, torch_support):
    """fixture returning some QNodes"""
    if interface == "torch" and not torch_support:
        pytest.skip("Skipped, no torch support")

    if interface == "tf" and not tf_support:
        pytest.skip("Skipped, no tf support")

    dev1 = qml.device("default.qubit", wires=2)
    dev2 = qml.device("default.qubit", wires=2)

    @qml.qnode(dev1, interface=interface)
    def qnode1(x):
        qml.RX(x[0], wires=0)
        qml.RY(x[1], wires=0)
        qml.CNOT(wires=[0, 1])
        return qml.var(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))

    @qml.qnode(dev2, interface=interface)
    def qnode2(x):
        qml.Hadamard(wires=0)
        qml.RX(x[0], wires=0)
        qml.RY(x[1], wires=1)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1))

    return qnode1, qnode2

class TestMap:
    """Test for mapping ansatz over observables or devices,
    to return a QNode collection"""

    def test_template_not_callable(self):
        """Test that an exception is correctly called if a
        template is not callable"""
        with pytest.raises(ValueError, match="template is not a callable"):
            qml.map(5, 0, 0)

    def test_mapping_over_observables(self):
        """Test that mapping over a list of observables produces
        a QNodeCollection with the correct QNodes, with a single
        device broadcast."""
        dev = qml.device("default.qubit", wires=1)
        obs_list = [qml.PauliX(0), qml.PauliY(0)]
        template = lambda x, wires: qml.RX(x, wires=0)

        qc = qml.map(template, obs_list, dev)

        assert len(qc) == 2

        # evaluate collection so that queue is populated
        qc(1)

        assert len(qc[0].ops) == 2
        assert qc[0].ops[0].name == "RX"
        assert qc[0].ops[1].name == "PauliX"

        assert len(qc[1].ops) == 2
        assert qc[1].ops[0].name == "RX"
        assert qc[1].ops[1].name == "PauliY"

        # test that device is broadcast
        assert qc[0].device is qc[1].device

    def test_mapping_over_observables_as_tuples(self):
        """Test that mapping over a tuple of observables produces
        a QNodeCollection with the correct QNodes, with a single
        device broadcast."""
        dev = qml.device("default.qubit", wires=1)
        obs_list = (qml.PauliX(0), qml.PauliY(0))
        template = lambda x, wires: qml.RX(x, wires=0)

        qc = qml.map(template, obs_list, dev)

        assert len(qc) == 2

        # evaluate collection so that queue is populated
        qc(1)

        assert len(qc[0].ops) == 2
        assert qc[0].ops[0].name == "RX"
        assert qc[0].ops[1].name == "PauliX"

        assert len(qc[1].ops) == 2
        assert qc[1].ops[0].name == "RX"
        assert qc[1].ops[1].name == "PauliY"

        # test that device is broadcast
        assert qc[0].device is qc[1].device

    def test_mapping_over_devices(self):
        """Test that mapping over a list of devices produces
        a QNodeCollection with the correct QNodes"""
        dev_list = [qml.device("default.qubit", wires=1), qml.device("default.qubit", wires=1)]

        obs_list = [qml.PauliX(0), qml.PauliY(0)]
        template = lambda x, wires: qml.RX(x, wires=0)

        qc = qml.map(template, obs_list, dev_list)

        assert len(qc) == 2

        # evaluate collection so that queue is populated
        qc(1)

        assert len(qc[0].ops) == 2
        assert qc[0].ops[0].name == "RX"
        assert qc[0].ops[1].name == "PauliX"

        assert len(qc[1].ops) == 2
        assert qc[1].ops[0].name == "RX"
        assert qc[1].ops[1].name == "PauliY"

        # test that device is not broadcast
        assert qc[0].device is not qc[1].device
        assert qc[0].device is dev_list[0]
        assert qc[1].device is dev_list[1]

    def test_mapping_over_measurements(self):
        """Test that mapping over a list of measurement types produces
        a QNodeCollection with the correct QNodes"""
        dev = qml.device("default.qubit", wires=1)

        obs_list = [qml.PauliX(0), qml.PauliY(0)]
        template = lambda x, wires: qml.RX(x, wires=0)

        qc = qml.map(template, obs_list, dev, measure=["expval", "var"])

        assert len(qc) == 2

        # evaluate collection so that queue is populated
        qc(1)

        assert len(qc[0].ops) == 2
        assert qc[0].ops[0].name == "RX"
        assert qc[0].ops[1].name == "PauliX"
        assert qc[0].ops[1].return_type == qml.operation.Expectation

        assert len(qc[1].ops) == 2
        assert qc[1].ops[0].name == "RX"
        assert qc[1].ops[1].name == "PauliY"
        assert qc[1].ops[1].return_type == qml.operation.Variance

    def test_invalid_obserable(self):
        """Test that an invalid observable raises an exception"""
        dev = qml.device("default.qubit", wires=1)

        obs_list = [qml.PauliX(0), qml.S(wires=0)]
        template = lambda x, wires: qml.RX(x, wires=0)

        with pytest.raises(ValueError, match="Some or all observables are not valid"):
            qml.map(template, obs_list, dev, measure=["expval", "var"])

class TestApply:
    """Tests for the apply function"""

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf"])
    def test_apply_summation(self, qnodes, interface, tf_support, torch_support, tol):
        """Test that summation can be applied using all interfaces"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])

        if interface == "tf":
            sfn = tf.reduce_sum
        elif interface == "torch":
            sfn = torch.sum
        else:
            sfn = np.sum

        cost = qml.apply(sfn, qc)

        params = [0.5643, -0.45]
        res = cost(params)
        expected = sfn(qc(params))

        assert np.allclose(res, expected, atol=tol, rtol=0)

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf"])
    def test_nested_apply(self, qnodes, interface, tf_support, torch_support, tol):
        """Test that nested apply can be done using all interfaces"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])

        if interface == "tf":
            sinfn = tf.sin
            sfn = tf.reduce_sum
        elif interface == "torch":
            sinfn = torch.sin
            sfn = torch.sum
        else:
            sinfn = np.sin
            sfn = np.sum

        cost = qml.apply(sfn, qml.apply(sinfn, qc))

        params = [0.5643, -0.45]
        res = cost(params)
        expected = sfn(sinfn(qc(params)))

        assert np.allclose(res, expected, atol=tol, rtol=0)

class TestSum:
    """Tests for the sum function"""

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf", None])
    def test_apply_summation(self, qnodes, interface, tf_support, torch_support, tol):
        """Test that summation can be applied using all interfaces"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])
        cost = qml.sum(qc)

        params = [0.5643, -0.45]
        res = cost(params)
        expected = sum(qc[0](params) + qc[1](params))

        if interface in ("tf", "torch"):
            res = res.numpy()
            expected = expected.numpy()

        assert np.allclose(res, expected, atol=tol, rtol=0)

    def test_unknown_interface(self, monkeypatch):
        """Test exception raised if the interface is unknown"""
        monkeypatch.setattr(qml.QNodeCollection, "interface", "invalid")
        dev = qml.device("default.qubit", wires=1)

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        qnodes = [qml.QNode(circuit, dev) for i in range(4)]
        qc = qml.QNodeCollection(qnodes)
        with pytest.raises(ValueError, match="Unknown interface invalid"):
            qml.sum(qc)

class TestDot:
    """Tests for the sum function"""

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf", None])
    def test_dot_product_tensor_qnodes(self, qnodes, interface, tf_support, torch_support):
        """Test that the dot product of tensor.qnodes can be applied using all interfaces"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])
        coeffs = [0.5, -0.1]

        if interface == "torch":
            coeffs = torch.tensor(coeffs, dtype=torch.float64)

        if interface == "tf":
            coeffs = tf.cast(coeffs, dtype=tf.float64)

        # test the dot product of tensor, qnodes
        cost = qml.dot(coeffs, qc)

        params = [0.5643, -0.45]
        res = cost(params)

        qcval = qc(params)

        if interface in ("tf", "torch"):
            res = res.numpy()
            qcval = qcval.numpy()
            coeffs = coeffs.numpy()

        expected = np.dot(coeffs, qcval)
        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf", None])
    def test_dot_product_qnodes_qnodes(self, qnodes, interface, tf_support, torch_support):
        """Test that the dot product of qnodes.qnodes can be applied using all interfaces"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qnode1, qnode2 = qnodes
        qc1 = qml.QNodeCollection([qnode1, qnode2])
        qc2 = qml.QNodeCollection([qnode1, qnode2])

        # test the dot product of qnodes, qnodes
        cost = qml.dot(qc1, qc2)

        params = [0.5643, -0.45]
        res = cost(params)

        qc1val = qc1(params)
        qc2val = qc2(params)

        if interface in ("tf", "torch"):
            res = res.numpy()
            qc1val = qc1val.numpy()
            qc2val = qc2val.numpy()

        expected = np.dot(qc1val, qc2val)
        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf", None])
    def test_dot_product_qnodes_tensor(self, qnodes, interface, tf_support, torch_support):
        """Test that the dot product of qnodes.tensor can be applied using all interfaces"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qnode1, _ = qnodes
        qc = qml.QNodeCollection([qnode1])
        coeffs = [0.5, -0.1]

        if interface == "torch":
            coeffs = torch.tensor(coeffs, dtype=torch.float64)

        if interface == "tf":
            coeffs = tf.cast(coeffs, dtype=tf.float64)

        # test the dot product of qnodes, tensor
        cost = qml.dot(qc, coeffs)

        params = [0.5643, -0.45]
        res = cost(params)
        qcval = qc(params)

        if interface in ("tf", "torch"):
            res = res.numpy()
            qcval = qcval.numpy()
            coeffs = coeffs.numpy()

        expected = np.dot(qcval, coeffs)
        assert np.all(res == expected)

    def test_unknown_interface(self, monkeypatch):
        """Test exception raised if the interface is unknown"""
        monkeypatch.setattr(qml.QNodeCollection, "interface", "invalid")
        dev = qml.device("default.qubit", wires=1)

        @qml.qnode(dev)
        def circuit1(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        @qml.qnode(dev)
        def circuit2(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        qc = qml.QNodeCollection([circuit1, circuit2])
        with pytest.raises(ValueError, match="Unknown interface invalid"):
            qml.dot([1, 2], qc)

    def test_no_qnodes(self):
        """Test exception raised if no qnodes are provided as arguments"""
        with pytest.raises(ValueError, match="At least one argument must be a QNodeCollection"):
            qml.dot([1, 2], [3, 4])

class TestQNodeCollection:
    """Tests for the QNodeCollection class"""

    def test_empty_init(self):
        """Test that an empty QNode collection can be initialized"""
        qc = qml.QNodeCollection()
        assert qc.qnodes == []
        assert len(qc) == 0

    def test_init_with_qnodes(self):
        """Test that a QNode collection can be initialized with QNodes"""
        dev = qml.device("default.qubit", wires=1)

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        qnodes = [qml.QNode(circuit, dev) for i in range(4)]
        qc = qml.QNodeCollection(qnodes)

        assert qc.qnodes == qnodes
        assert len(qc) == 4

    @pytest.mark.parametrize("interface", ["autograd", "numpy", "torch", "tf"])
    def test_interface_property(self, interface, tf_support, torch_support):
        """Test that the interface property correctly
        resolves interfaces from the internal QNodes"""
        if interface == "torch" and not torch_support:
            pytest.skip("Skipped, no torch support")

        if interface == "tf" and not tf_support:
            pytest.skip("Skipped, no tf support")

        qc = qml.QNodeCollection()
        assert qc.interface is None

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        dev = qml.device("default.qubit", wires=1)
        qnodes = [qml.QNode(circuit, dev, interface=interface) for i in range(4)]
        qc = qml.QNodeCollection(qnodes)

        if interface == "numpy":
            # Note: the "numpy" interface is deprecated, and
            # now resolves to "autograd"
            interface = "autograd"

        assert qc.interface == interface

    def test_append_qnode(self):
        """Test that a QNode is correctly appended"""
        qc = qml.QNodeCollection()
        assert qc.qnodes == []

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        dev = qml.device("default.qubit", wires=1)
        qnode = qml.QNode(circuit, dev)
        qc.append(qnode)

        assert qc.qnodes == [qnode]

    def test_extend_qnodes(self):
        """Test that a list of QNodes is correctly appended"""
        qc = qml.QNodeCollection()
        assert qc.qnodes == []

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        dev = qml.device("default.qubit", wires=1)
        qnodes = [qml.QNode(circuit, dev) for i in range(4)]
        qc.extend(qnodes)

        assert qc.qnodes == [] + qnodes

    def test_extend_multiple_interface_qnodes(self):
        """Test that an error is returned if QNodes with differing
        interfaces are attempted to be added to a QNodeCollection"""
        qc = qml.QNodeCollection()

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        dev = qml.device("default.qubit", wires=1)
        qnodes = [
            qml.QNode(circuit, dev, interface="autograd"),
            qml.QNode(circuit, dev, interface=None),
        ]

        with pytest.raises(ValueError, match="do not all use the same interface"):
            qc.extend(qnodes)

    def test_extend_interface_mistmatch(self):
        """Test that an error is returned if QNodes with a differing
        interface to the QNode collection are appended"""
        qc = qml.QNodeCollection()

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        dev = qml.device("default.qubit", wires=1)
        qnode1 = qml.QNode(circuit, dev, interface="autograd")
        qnode2 = qml.QNode(circuit, dev, interface=None)

        qc.extend([qnode1])

        with pytest.raises(ValueError, match="Interface mismatch"):
            qc.extend([qnode2])

    def test_indexing(self):
        """Test that indexing into the QNodeCollection correctly works"""

        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        dev = qml.device("default.qubit", wires=1)
        qnodes = [qml.QNode(circuit, dev) for i in range(4)]

        qc = qml.QNodeCollection(qnodes)
        assert qc[2] == qnodes[2]

    def test_sequence(self):
        """Test that the QNodeCollection is a sequence type"""
        qc = qml.QNodeCollection()
        assert isinstance(qc, Sequence)

    @pytest.mark.parametrize("interface", ["autograd", "numpy"])
    def test_eval_autograd(self, qnodes):
        """Test correct evaluation of the QNodeCollection using
        the Autograd interface"""
        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])
        params = [0.5643, -0.45]

        res = qc(params)
        expected = np.vstack([qnode1(params), qnode2(params)])
        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["autograd", "numpy"])
    def test_grad_autograd(self, qnodes):
        """Test correct gradient of the QNodeCollection using
        the Autograd interface"""
        qnode1, qnode2 = qnodes

        params = [0.5643, -0.45]
        qc = qml.QNodeCollection([qnode1, qnode2])
        cost_qc = lambda params: np.sum(qc(params))
        grad_qc = qml.grad(cost_qc, argnum=0)

        cost_expected = lambda params: np.sum(qnode1(params) + qnode2(params))
        grad_expected = qml.grad(cost_expected, argnum=0)

        res = grad_qc(params)
        expected = grad_expected(params)

        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["torch"])
    def test_eval_torch(self, qnodes, skip_if_no_torch_support):
        """Test correct evaluation of the QNodeCollection using
        the torch interface"""
        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])
        params = [0.5643, -0.45]

        res = qc(params).numpy()
        expected = np.vstack([qnode1(params), qnode2(params)])
        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["torch"])
    def test_grad_torch(self, qnodes, skip_if_no_torch_support):
        """Test correct gradient of the QNodeCollection using
        the torch interface"""
        qnode1, qnode2 = qnodes

        # calculate the gradient of the collection using pytorch
        params = torch.autograd.Variable(torch.tensor([0.5643, -0.45]), requires_grad=True)
        qc = qml.QNodeCollection([qnode1, qnode2])
        cost = torch.sum(qc(params))
        cost.backward()
        res = params.grad.numpy()

        # calculate the gradient of the QNodes individually using pytorch
        params = torch.autograd.Variable(torch.tensor([0.5643, -0.45]), requires_grad=True)
        cost = torch.sum(qnode1(params) + qnode2(params))
        cost.backward()
        expected = params.grad.numpy()

        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["tf"])
    def test_eval_tf(self, qnodes, skip_if_no_tf_support):
        """Test correct evaluation of the QNodeCollection using
        the tf interface"""
        qnode1, qnode2 = qnodes
        qc = qml.QNodeCollection([qnode1, qnode2])
        params = [0.5643, -0.45]

        res = qc(params).numpy()
        expected = np.vstack([qnode1(params), qnode2(params)])
        assert np.all(res == expected)

    @pytest.mark.parametrize("interface", ["tf"])
    def test_grad_tf(self, qnodes, skip_if_no_tf_support):
        """Test correct gradient of the QNodeCollection using
        the tf interface"""
        qnode1, qnode2 = qnodes

        # calculate the gradient of the collection using tf
        params = Variable([0.5643, -0.45])
        qc = qml.QNodeCollection([qnode1, qnode2])

        with tf.GradientTape() as tape:
            tape.watch(params)
            cost = sum(qc(params))
            res = tape.gradient(cost, params).numpy()

        # calculate the gradient of the QNodes individually using tf
        params = Variable([0.5643, -0.45])

        with tf.GradientTape() as tape:
            tape.watch(params)
            cost = sum(qnode1(params) + qnode2(params))
            expected = tape.gradient(cost, params).numpy()

        assert np.all(res == expected)

max_stars_repo_name: zaqqwerty/pennylane

max_issues_repo_name: zaqqwerty/pennylane

max_forks_repo_name: zaqqwerty/pennylane


text:
#!/usr/bin/env python3
"""
@author: MD.Nazmuddoha Ansary
"""
from __future__ import print_function
from termcolor import colored

from progressbar import ProgressBar
import time
import os
import numpy as np
from glob import glob
from sklearn.model_selection import train_test_split
from coreLib.utils import DataSet,readJson,LOG_INFO,readh5,saveh5,to_tfrecord
#-----------------------------------------------------Load Config----------------------------------------------------------
config_data=readJson('config.json')

class FLAGS:
    DATA_DIR     = config_data['FLAGS']['DATA_DIR']
    CLASSES      = config_data['FLAGS']['CLASSES']
    OUTPUT_DIR   = config_data['FLAGS']['OUTPUT_DIR']
    IMAGE_DIM    = config_data['FLAGS']['IMAGE_DIM']
    ROT_ANGLES   = config_data['FLAGS']['ROT_ANGLES']
    EVAL_SPLIT   = config_data['FLAGS']['EVAL_SPLIT']
    TEST_SPLIT   = config_data['FLAGS']['TEST_SPLIT']
    BATCH_SIZE   = config_data['FLAGS']['BATCH_SIZE']

#-----------------------------------------------------------------------------------------------------------------------------------
def create_dataset(FLAGS):
    DS=DataSet(FLAGS)
    DS.create()
    return DS

def get_batched_len(batch_size,data_len):
    new_len= (data_len//batch_size)*batch_size
    return new_len

def create_split(DS,FLAGS):
    LOG_INFO('Data Split')
    # paths for saving data
    X_TEST_PATH     = os.path.join(DS.test_dir,'X_test.h5')
    Y_TEST_PATH     = os.path.join(DS.test_dir,'Y_test.h5')
    X_TRAIN_PATH    = os.path.join(DS.train_dir,'X_train.h5')
    Y_TRAIN_PATH    = os.path.join(DS.train_dir,'Y_train.h5')
    X_EVAL_PATH     = os.path.join(DS.train_dir,'X_eval.h5')
    Y_EVAL_PATH     = os.path.join(DS.train_dir,'Y_eval.h5')
    # complete data
    X=readh5(DS.X_PATH)
    Y=readh5(DS.Y_PATH)
    # split data
    X_comb,X_test,Y_comb,Y_test = train_test_split(X,Y,test_size=FLAGS.TEST_SPLIT)
    X_train,X_eval,Y_train,Y_eval = train_test_split(X_comb,Y_comb,test_size=FLAGS.EVAL_SPLIT)
    # batch data
    train_len=get_batched_len(FLAGS.BATCH_SIZE,X_train.shape[0])
    eval_len =get_batched_len(FLAGS.BATCH_SIZE,X_eval.shape[0])
    test_len =get_batched_len(FLAGS.BATCH_SIZE,X_test.shape[0])
    # finalize data
    X_train=X_train[:train_len]
    Y_train=Y_train[:train_len]
    X_eval =X_eval[:eval_len]
    Y_eval =Y_eval[:eval_len]
    X_test =X_test[:test_len]
    Y_test =Y_test[:test_len]

    # save data
    # test
    saveh5(X_TEST_PATH,X_test)
    saveh5(Y_TEST_PATH,Y_test)
    # train
    saveh5(X_TRAIN_PATH,X_train)
    saveh5(Y_TRAIN_PATH,Y_train)
    # eval
    saveh5(X_EVAL_PATH,X_eval)
    saveh5(Y_EVAL_PATH,Y_eval)

    LOG_INFO('Train Data Size:{}'.format(X_train.shape[0]))
    LOG_INFO('Test Data Size:{}'.format(X_test.shape[0]))
    LOG_INFO('Eval Data Size:{}'.format(X_eval.shape[0]))
    os.remove(DS.X_PATH)
    os.remove(DS.Y_PATH)

def create_tfrecords(DS,FLAGS):
    # complete image paths
    images_path_list=glob(os.path.join(DS.image_dir,'*.png'))
    # test split
    total_data_len=len(images_path_list)
    test_len=int(FLAGS.TEST_SPLIT*total_data_len)
    test_image_path_list=images_path_list[:test_len]
    test_image_path_list=test_image_path_list[:get_batched_len(FLAGS.BATCH_SIZE,len(test_image_path_list))]
    images_path_list=images_path_list[test_len:]
    # eval and train split
    total_data_len=len(images_path_list)
    eval_len=int(FLAGS.EVAL_SPLIT*total_data_len)
    eval_image_path_list=images_path_list[:eval_len]
    eval_image_path_list = eval_image_path_list[:get_batched_len(FLAGS.BATCH_SIZE,len( eval_image_path_list))]
    train_image_path_list=images_path_list[eval_len:]
    train_image_path_list=train_image_path_list[:get_batched_len(FLAGS.BATCH_SIZE,len(train_image_path_list))]
    # tfrecords
    LOG_INFO('Creating TFRECORDS')
    to_tfrecord(train_image_path_list,'Train',DS.ds_dir,FLAGS.CLASSES)
    to_tfrecord( eval_image_path_list,'Eval' ,DS.ds_dir,FLAGS.CLASSES)
    to_tfrecord( test_image_path_list,'Test' ,DS.ds_dir,FLAGS.CLASSES)

#-----------------------------------------------------------------------------------------------------------------------------------
def main(FLAGS):
    start_time=time.time()
    LOG_INFO('Creating Dataset')
    DS=create_dataset(FLAGS)
    create_split(DS,FLAGS)
    #create_tfrecords(DS,FLAGS)
    LOG_INFO('Time Taken:{} s'.format(time.time()-start_time),p_color='yellow')

if __name__ == "__main__":
    main(FLAGS)

max_stars_repo_name: shakiritis/pySKIND

max_issues_repo_name: shakiritis/pySKIND

max_forks_repo_name: shakiritis/pySKIND


text:
from mmelemental.models import forcefield as ff
import numpy
import pytest
import os

natoms = 10
nbonds = natoms - 1
nangles = natoms - 2
ndihedrals = natoms - 3

# Data
rand_npts = numpy.random.rand(natoms)
rand_nbond = numpy.random.rand(nbonds)
rand_nangle = numpy.random.rand(nangles)
rand_ndih = numpy.random.rand(ndihedrals)
rand_int_ndih = numpy.random.randint(ndihedrals, size=ndihedrals)

# Pytest parametrization args
nb_pytest_params = "form,params", [
    ("LennardJones", {"epsilon": rand_npts, "sigma": rand_npts})
]
bonds_pytest_params = "form,lengths,params", [
    ("Harmonic", rand_nbond, {"spring": rand_nbond}),
    ("Gromos96", rand_nbond, {"spring": rand_nbond}),
]
angles_pytest_params = "form,lengths,params", [
    ("Harmonic", rand_nangle, {"spring": rand_nangle})
]
dihedrals_pytest_params = "form,lengths,params", [
    (
        "Charmm",
        rand_ndih,
        {"energy": rand_ndih, "periodicity": rand_int_ndih, "phase": rand_ndih},
    ),
    (
        "CharmmMulti",
        rand_ndih,
        {
            "energy": (rand_ndih, rand_ndih),
            "periodicity": (rand_int_ndih, rand_int_ndih),
            "phase": (rand_ndih, rand_ndih),
        },
    ),
]
dihedrals_improper_pytest_params = "form,lengths,params", [
    (
        "Charmm",
        rand_ndih,
        {"energy": rand_ndih, "periodicity": rand_int_ndih, "phase": rand_ndih},
    ),
    (
        "CharmmMulti",
        rand_ndih,
        {
            "energy": (rand_ndih, rand_ndih),
            "periodicity": (rand_int_ndih, rand_int_ndih),
            "phase": (rand_ndih, rand_ndih),
        },
    ),
]

def rewrite(filename):
    import json

    with open(filename, "r") as fp:
        data = json.load(fp)

    with open(filename, "w") as fp:
        json.dump(data, fp, indent=4)

def rand_angles(npts):
    return rand_data * numpy.pi

###########################################
####### NON-BONDED ########################
@pytest.mark.parametrize(*nb_pytest_params)
def test_nonbonded(form, params):
    potential = getattr(ff.nonbonded.potentials, form)(**params)
    nonbonded = ff.nonbonded.NonBonded(params=potential, form=form)
    assert nonbonded.form == form
    return nonbonded

############################################
######## BONDS #############################
@pytest.mark.parametrize(*bonds_pytest_params)
def test_bonds(form, lengths, params):
    indices = [(i, i + 1, 1.0) for i in range(nbonds)]
    potential = getattr(ff.bonded.bonds.potentials, form)(**params)
    bonds = ff.bonded.Bonds(
        params=potential,
        lengths=lengths,
        form=form,
        connectivity=indices,
    )
    assert bonds.form == form
    return bonds

###########################################
####### ANGLES ############################
@pytest.mark.parametrize(*angles_pytest_params)
def test_angles(form, lengths, params):
    indices = [(i, i + 1, i + 2) for i in range(nangles)]
    potential = getattr(ff.bonded.angles.potentials, form)(**params)
    angles = ff.bonded.Angles(
        params=potential,
        angles=lengths,
        angles_units="radians",
        connectivity=indices,
        form=form,
    )
    assert angles.form == form
    return angles

###########################################
####### DIHEDRALS #########################
@pytest.mark.parametrize(*dihedrals_pytest_params)
def test_dihedrals(form, lengths, params):
    indices = [(i, i + 1, i + 2, i + 3) for i in range(ndihedrals)]
    potential = getattr(ff.bonded.dihedrals.potentials, form)(**params)
    dihedrals = ff.bonded.Dihedrals(
        params=potential,
        angles=lengths,
        angles_units="radians",
        connectivity=indices,
        form=form,
    )
    assert dihedrals.form == form
    return dihedrals

###########################################
####### IMPROPER DIHEDRALS ################
@pytest.mark.parametrize(*dihedrals_pytest_params)
def test_improper_dihedrals(form, lengths, params):
    indices = [(i, i + 1, i + 2, i + 3) for i in range(ndihedrals)]
    potential = getattr(ff.bonded.dihedrals_improper.potentials, form)(**params)
    dihedrals_improper = ff.bonded.DihedralsImproper(
        params=potential,
        angles=lengths,
        angles_units="radians",
        connectivity=indices,
        form=form,
    )
    assert dihedrals_improper.form == form
    return dihedrals_improper

@pytest.mark.parametrize(*dihedrals_pytest_params)
def test_forcefield(form, lengths, params):
    nonbonded = test_nonbonded(
        form="LennardJones", params={"epsilon": rand_npts, "sigma": rand_npts}
    )
    bonds = test_bonds(
        form="Harmonic", lengths=rand_nbond, params={"spring": rand_nbond}
    )
    angles = test_angles(
        form="Harmonic", lengths=rand_nangle, params={"spring": rand_nangle}
    )
    dihedrals = test_dihedrals(form, lengths, params)
    dihedrals_improper = test_improper_dihedrals(form, lengths, params)
    mm_ff = ff.ForceField(
        nonbonded=nonbonded,
        bonds=bonds,
        angles=angles,
        dihedrals=dihedrals,
        dihedrals_improper=dihedrals_improper,
        charges=numpy.random.rand(natoms),
        symbols=["H" for _ in range(natoms)],
    )
    mm_ff.to_file("forcefield-single.json")
    rewrite("forcefield-single.json")
    os.remove("forcefield-single.json")

@pytest.mark.skip("Cannot run this now")
def test_forcefield_hybrid():
    nonbonded = test_nonbonded()
    bonds = [test_bonds_linear(), test_bonds_gromos96()]
    angles = test_angles()
    dihedrals = test_dihedrals()
    mm_ff = ff.ForceField(
        nonbonded=nonbonded,
        bonds=bonds,
        angles=angles,
        dihedrals=dihedrals,
        charges=numpy.random.rand(natoms),
        symbols=["H" for _ in range(natoms)],
    )
    mm_ff.to_file("forcefield.json")
    rewrite("forcefield.json")
    os.remove("forcefield.json")

max_stars_repo_name: MolSSI/MMElemental

max_issues_repo_name: MolSSI/MMElemental

max_forks_repo_name: MolSSI/mmelemental


text:
# -*- coding: utf-8 -*-
import numpy as np
import ot as pot
import scipy.sparse
import scipy.sparse
import scipy.stats
import sklearn.decomposition
import sklearn.metrics

def compute_growth_scores(proliferation, apoptosis, beta_max=1.7, beta_center=0.25, delta_max=1.7, delta_min=0.3,
                          beta_min=0.3):
    birth = __beta(proliferation, beta_max=beta_max, center=beta_center, beta_min=beta_min)
    death = __delta(apoptosis, delta_max=delta_max, delta_min=delta_min)
    return np.exp(birth - death)

def __gen_logistic(p, beta_max, beta_min, pmax, pmin, center, width):
    return beta_min + __logistic(p, L=beta_max - beta_min, k=4 / width,
                                 x0=center)

def __beta(p, beta_max=1.7, beta_min=0, pmax=1.0, pmin=-0.5, center=0.25):
    # map proliferation p into interval beta_max, beta_min
    return __gen_logistic(p, beta_max, beta_min, pmax, pmin, center, width=0.5)

def __delta(a, delta_max=1.7, delta_min=0.3, amax=0.5, amin=-0.4, center=0.1):
    return __gen_logistic(a, delta_max, delta_min, amax, amin, center,
                          width=0.2)

def __logistic(x, L=1, k=1, x0=0):
    f = L / (1 + np.exp(-k * (x - x0)))
    return f

def sample_randomly(exp1, exp2, tm, g, npairs):
    # if args.npairs is None or args.npairs <= 0:
    #     l = tm / tm.sum(axis=0)
    #     l = l.flatten()
    #     qq = np.where(l > args.tm_thresh)
    #     n = len(qq[0])
    # else:
    n = npairs

    # column_sums = np.sum(tm, axis=0)
    # row_sums = np.sum(tm, axis=1)
    # p = np.outer(row_sums, column_sums)

    p = g
    q = np.ones(exp2.shape[0]) * np.average(g)
    p = np.outer(p, q)
    p = p.flatten()
    probabilties = p / p.sum()
    # pairs = np.random.multinomial(n, p, size=1)
    # pairs = np.nonzero(pairs.reshape(exp1.shape[0], exp2.shape[0]))

    s = np.random.multinomial(n, probabilties, size=1)
    reshaped_s = s.reshape(exp1.shape[0], exp2.shape[0])
    pairs = np.nonzero(reshaped_s)
    weights = reshaped_s[pairs]
    return {'pc0': exp1[pairs[0]], 'pc1': exp2[pairs[1]],
            'indices0': pairs[0],
            'indices1': pairs[1],
            'weights': weights}

def sample_uniformly(exp1, exp2, tm, npairs):
    # if args.npairs is None or args.npairs <= 0:
    #     l = tm / tm.sum(axis=0)
    #     l = l.flatten()
    #     qq = np.where(l > args.tm_thresh)
    #     n = len(qq[0])
    # else:
    n = npairs

    p = np.ones(exp1.shape[0] * exp2.shape[0])
    p = p / p.sum()
    pairs = np.random.multinomial(n, p, size=1)
    pairs = np.nonzero(pairs.reshape(exp1.shape[0], exp2.shape[0]))
    return exp1[pairs[0]], exp2[pairs[1]]

def sample_from_transport_map(exp1, exp2, tm, npairs, t_interpolate):
    probabilties = tm / np.power(tm.sum(axis=0), 1.0 - t_interpolate)
    probabilties = probabilties.flatten()
    # if args.npairs is None or args.npairs <= 0:
    #     # l = l / l.sum()
    #     reshaped_tmap = probabilties.reshape(exp1.shape[0], exp2.shape[0])
    #     q = np.where(reshaped_tmap > args.tm_thresh)
    #     qq = np.where(probabilties > args.tm_thresh)
    #     return exp1[q[0]], exp2[q[1]], probabilties[qq]
    # else:
    probabilties = probabilties / probabilties.sum()
    # pairs = np.random.multinomial(args.npairs, probabilties, size=1)
    # pairs = np.nonzero(pairs.reshape(exp1.shape[0], exp2.shape[0]))
    #  return {'pc0': exp1[pairs[0]], 'pc1': exp2[pairs[1]], 'indices0': pairs[0], 'indices1': pairs[1], 'weights': None}
    s = np.random.multinomial(npairs, probabilties, size=1)
    reshaped_s = s.reshape(exp1.shape[0], exp2.shape[0])
    pairs = np.nonzero(reshaped_s)
    weights = reshaped_s[pairs]
    return {'pc0': exp1[pairs[0]], 'pc1': exp2[pairs[1]],
            'indices0': pairs[0],
            'indices1': pairs[1],
            'weights': weights}

def interpolate_with_ot(p0, p1, tmap, interp_frac, size):
    """
    Interpolate between p0 and p1 at fraction t_interpolate knowing a transport map from p0 to p1

    Parameters
    ----------
    p0 : 2-D array
        The genes of each cell in the source population
    p1 : 2-D array
        The genes of each cell in the destination population
    tmap : 2-D array
        A transport map from p0 to p1
    t_interpolate : float
        The fraction at which to interpolate
    size : int
        The number of cells in the interpolated population

    Returns
    -------
    p05 : 2-D array
        An interpolated population of 'size' cells
    """
    p0 = p0.toarray() if scipy.sparse.isspmatrix(p0) else p0
    p1 = p1.toarray() if scipy.sparse.isspmatrix(p1) else p1
    p0 = np.asarray(p0, dtype=np.float64)
    p1 = np.asarray(p1, dtype=np.float64)
    tmap = np.asarray(tmap, dtype=np.float64)
    if p0.shape[1] != p1.shape[1]:
        raise ValueError("Unable to interpolate. Number of genes do not match")
    if p0.shape[0] != tmap.shape[0] or p1.shape[0] != tmap.shape[1]:
        raise ValueError("Unable to interpolate. Tmap size is {}, expected {}"
                         .format(tmap.shape, (len(p0), len(p1))))
    I = len(p0);
    J = len(p1)
    # Assume growth is exponential and retrieve growth rate at t_interpolate
    p = tmap / np.power(tmap.sum(axis=0), 1. - interp_frac)
    p = p.flatten(order='C')
    p = p / p.sum()
    choices = np.random.choice(I * J, p=p, size=size)
    return np.asarray([p0[i // J] * (1 - interp_frac) + p1[i % J] * interp_frac for i in choices], dtype=np.float64)

def interpolate_randomly(p0, p1, t_interpolate, size):
    """
    Interpolate between p0 and p1 at fraction t_interpolate

    Parameters
    ----------
    p0 : 2-D array
        The genes of each cell in the source population
    p1 : 2-D array
        The genes of each cell in the destination population
    t_interpolate : float
        The fraction at which to interpolate
    size : int
        The number of cells in the interpolated population

    Returns
    -------
    p05 : 2-D array
        An interpolated population of 'size' cells
    """
    t = t_interpolate
    p0 = p0.toarray() if scipy.sparse.isspmatrix(p0) else p0
    p1 = p1.toarray() if scipy.sparse.isspmatrix(p1) else p1
    p0 = np.asarray(p0, dtype=np.float64)
    p1 = np.asarray(p1, dtype=np.float64)
    if p0.shape[1] != p1.shape[1]:
        raise ValueError("Unable to interpolate. Number of genes do not match")
    I = len(p0)
    J = len(p1)
    choices = np.random.choice(I * J, size=size)
    return np.asarray([p0[i // J] * (1 - t) + p1[i % J] * t for i in choices], dtype=np.float64)

def interpolate_randomly_with_growth(p0, p1, t, size, g):
    p0 = p0.toarray() if scipy.sparse.isspmatrix(p0) else p0
    p1 = p1.toarray() if scipy.sparse.isspmatrix(p1) else p1
    p0 = np.asarray(p0, dtype=np.float64)
    p1 = np.asarray(p1, dtype=np.float64)

    p = g
    q = np.ones(p1.shape[0]) * np.average(g)

    p = np.outer(p, q)
    p = p.flatten(order='C')
    p = p / p.sum()
    I = len(p0)
    J = len(p1)
    choices = np.random.choice(I * J, p=p, size=size)
    return np.asarray([p0[i // J] * (1 - t) + p1[i % J] * t for i in choices], dtype=np.float64)

def earth_mover_distance(cloud1, cloud2, eigenvals=None, weights1=None, weights2=None):
    """
    Returns the earth mover's distance between two point clouds

    Parameters
    ----------
    cloud1 : 2-D array
        First point cloud
    cloud2 : 2-D array
        Second point cloud

    Returns
    -------
    distance : float
        The distance between the two point clouds
    """
    cloud1 = cloud1.toarray() if scipy.sparse.isspmatrix(cloud1) else cloud1
    cloud2 = cloud2.toarray() if scipy.sparse.isspmatrix(cloud2) else cloud2
    if eigenvals is not None:
        cloud1 = cloud1.dot(eigenvals)
        cloud2 = cloud2.dot(eigenvals)
    if weights1 is None:
        p = np.ones(len(cloud1)) / len(cloud1)
    else:
        weights1 = weights1.astype('float64')
        p = weights1 / weights1.sum()

    if weights2 is None:
        q = np.ones(len(cloud2)) / len(cloud2)
    else:
        weights2 = weights2.astype('float64')
        q = weights2 / weights2.sum()

    pairwise_dist = np.ascontiguousarray(sklearn.metrics.pairwise.pairwise_distances(
        cloud1, Y=cloud2, metric='sqeuclidean', n_jobs=-1))
    return np.sqrt(pot.emd2(p, q, pairwise_dist, numItermax=1e7))

def compute_pca(m1, m2, n_components):
    matrices = list()
    matrices.append(m1 if not scipy.sparse.isspmatrix(m1) else m1.toarray())
    matrices.append(m2 if not scipy.sparse.isspmatrix(m2) else m2.toarray())
    x = np.vstack(matrices)
    mean_shift = x.mean(axis=0)
    x = x - mean_shift
    n_components = min(n_components, x.shape[0])  # n_components must be <= ncells
    pca = sklearn.decomposition.PCA(n_components=n_components, random_state=58951)
    pca.fit(x.T)
    comp = pca.components_.T
    m1_len = m1.shape[0]
    m2_len = m2.shape[0]
    pca_1 = comp[0:m1_len]
    pca_2 = comp[m1_len:(m1_len + m2_len)]
    return pca_1, pca_2, pca, mean_shift

def get_pca(dim, *args):
    """
    Get a PCA projector for the arguments.

    Parameters
    ----------
    dim : int
        The number of components to use for PCA, i.e. number of dimensions of the resulting space.
    *args : ndarray
        The points to compute dimensionality reduction for. Can be several sets of points.

    Returns
    -------
    pca : sklearn.decomposition.PCA
        A PCA projector. Use `pca.transform(x)` to get the PCA-space projection of x.

    Example
    -------
    >>> pca = get_pca(30, p0_x)
    >>> pca.transform(p0_x)
    >>> # -> project p0_x to PCA space
    >>> pca = get_pca(30, p0_x, p1_x)
    >>> p0, p05, p1 = [ pca.transform(x) for x in [p0_x, p05_x, p1_x] ]
    >>> # -> project all three sets of points to PCA space computed from p0_x and p1_x
    """
    args = [a.toarray() if scipy.sparse.isspmatrix(a) else a for a in args]
    x = np.vstack(args)
    mean = x.mean(axis=0)
    x = x - mean
    pca = sklearn.decomposition.PCA(n_components=dim)
    return pca.fit(x), mean

def pca_transform(pca, mean, arr):
    """
    Apply a PCA transformation to argument

    Parameters
    ----------
    pca : sklearn.decomposition.PCA
        A PCA projector. See wot.ot.get_pca
    arr : numpy ndarray or scipy.sparse matrix
        The array to project.

    Returns
    -------
    result : ndarray
    """
    ndarr = arr.toarray() if scipy.sparse.isspmatrix(arr) else arr
    ndarr = ndarr - mean
    return pca.transform(ndarr)

max_stars_repo_name: matthieuheitz/wot

max_issues_repo_name: matthieuheitz/wot

max_forks_repo_name: matthieuheitz/wot

reason: import numpy,import scipy


text:
"""
Programmer: Chris Tralie

Purpose: To simulate a 2D environment and a robot inside of it with
an omnidirectional range scanner
"""
import numpy as np
import matplotlib.pyplot as plt
from skimage.measure import find_contours
import skimage
import skimage.io

def clamp(arrs, boolexpr):
    """
    Clamp numpy arrays to a boolean expression
    Parameters
    ----------
    arrs: Length-k list of [ndarray(N)]
        List of arrays
    boolexpr: ndarray(N, dtype=np.boolean)
        Boolean expression
    Returns
    -------
    Length-k list of [ndarray(M)], where M is
    the number of true evaluations in boolexpr
    """
    return [arr[boolexpr] for arr in arrs]

def ray_intersect_curve(p0, v, X):
    """
    Parameters
    ----------
    p0: ndarray(2)
        Initial point on ray
    v: ndarray(2)
        Direction of ray
    X: ndarray(N, 2)
        Points on the curve

    Returns {'t': float,
                    The distance of the closest point of intersection, or
                    np.inf if there was no intersection
            'p': ndarray(2)
                    Point of intersection
            }
    """
    # vec{a} + s*vec{u} refers to line segments on curve
    # vec{b} + t*vec{v} refers to ray
    u = X[1::, :] - X[0:-1, :]
    u = np.concatenate((u, [u[-1, :]]), axis=0) # Continue the last slope
    # Use Cramer's rule
    # uy*vx - ux*vy
    denom = u[:, 1]*v[0] - u[:, 0]*v[1]
    # vx*(by - ay) - vy*(bx - ax)
    num_s = v[0]*(p0[1]-X[:, 1]) - v[1]*(p0[0]-X[:, 0])
    # ux*(by - ay) - uy*(bx - ax)
    num_t = u[:, 0]*(p0[1]-X[:, 1]) - u[:, 1]*(p0[0]-X[:, 0])

    idx = np.arange(X.shape[0])
    [num_s, num_t, denom, idx] = clamp([num_s, num_t, denom, idx], np.abs(denom) > 0)
    s = num_s/denom
    t = num_t/denom
    # Intersection is within the bounds of the segments
    [s, t, idx] = clamp([s, t, idx], (s >= 0)*(s <= 1))
    tmin = np.inf
    p = np.zeros(2)
    if s.size > 0:
        [s, t, idx] = clamp([s, t, idx], t >= 0)
        if t.size > 0:
            minidx = np.argmin(t)
            tmin = t[minidx]
            s = s[minidx]
            t = t[minidx]
            idx = idx[minidx]
            p = p0 + t*v
    return {'t':tmin, 'p':p}

class Environment(object):
    """
    Attributes
    ----------
    img: ndarray(M, N)
        The image for the maze

    """
    def __init__(self, path):
        """
        Setup an environment based on a file

        Parameters
        ----------
        path: string
            Path to environment file, where occupied cells are white
            and blocked cells are black
        """
        img = skimage.io.imread(path)
        if len(img.shape) > 2:
            img = np.sum(img, axis=2)
        self.img = img
        self.contours = [np.fliplr(x) for x in find_contours(img, 0.5)]
        self.setup_positions()

    def setup_positions(self):
        """
        A helper method for the constructor which sets up the open cells in
        the graph into a list of positions and neighbors
        """
        x, y = np.meshgrid(np.arange(self.img.shape[1]), np.arange(self.img.shape[0]))
        x = x[self.img > 0]
        y = y[self.img > 0]
        self.X = np.array([x, y]).T
        N = x.size
        pos2idx = {(x[i], y[i]):i for i in range(x.size)}
        neighbors = [[i] for i in range(N)]
        for i in range(N):
            xi = x[i]
            yi = y[i]
            for (dx, dy) in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                neighb = (xi+dx, yi+dy)
                if neighb in pos2idx:
                    neighbors[i].append(pos2idx[neighb])
        self.pos2idx = pos2idx
        self.neighbors = neighbors

    def plot(self, show_contours=False):
        """
        Plot the occupancy grid of the map with the extracted
        contours on top of it

        Parameters
        ----------
        show_contours: boolean
            If true, plot contours superimposed
        """
        plt.imshow(self.img, cmap='gray')
        if show_contours:
            for X in self.contours:
                plt.plot(X[:, 0], X[:, 1])
        plt.gca().invert_yaxis()

    def get_range_scan(self, pos, res, alpha=0):
        """
        Simulate a range scan
        Parameters
        ----------
        pos: ndarray(2)
            The x/y position of the camera
        res: int
            The number of samples to take
        alpha: float
            Noise amount

        Returns
        -------
        X: ndarray(N, 2)
            An 2D point cloud from the point of view of the camera
        XWorld: ndarray(N, 2)
            An 2D point cloud in world coordinates
        """
        scan = np.zeros(res)
        for i, theta in enumerate(np.linspace(0, 2*np.pi, res)):
            V = np.array([np.cos(theta), np.sin(theta)])
            # Find the closest point of intersection
            t = np.inf
            for curve in self.contours:
                res = ray_intersect_curve(pos, V, curve)
                if res['t'] < t:
                    t = res['t']
            if np.isfinite(t):
                scan[i] = max(0, t*(1+alpha*np.random.randn()))
        return scan

    def get_state_scans(self, res):
        """
        Compute perfect state scans for all states to use as a model

        Parameters
        ----------
        res: int
            Resolution of each scan
        """
        return [self.get_range_scan(self.X[i], res, alpha=0) for i in range(self.X.shape[0])]

    def plot_range_scan(self, x, scan, max_range=None):
        """
        Plot a range scan alongside a ground truth position on the map

        Parameters
        ----------
        x: ndarray(2)
            The x/y position of the camera
        scan: ndarray(N)
            The scan
        max_range: float
            Maximum range of scan (region to which to shrink plot)
        """
        N = len(scan)
        t = np.linspace(0, 2*np.pi, N)

        plt.subplot(121)
        self.plot()
        plt.title("Ground Truth Position")
        plt.scatter([x[0]], [x[1]])
        plt.plot([x[0], x[0]+2], [x[1], x[1]], 'C0')
        plt.plot([x[0], x[0]], [x[1], x[1]+2], 'C1')
        plt.subplot(122)
        plt.scatter([0], [0])
        plt.plot(scan*np.cos(t), scan*np.sin(t))
        if max_range:
            plt.xlim([-max_range, max_range])
            plt.ylim([-max_range, max_range])
        plt.title("Range Scan")

    def plot_probabilities(self, est, p=1, show_max=True):
        """
        Plot a set of probabilities

        Parameters
        ----------
        est: ndarray(N)
            (Log) Probabilities estimates of all points
        p: float
            Polynomial degree for compressing color range
        show_max: boolean
            Whether to show the location of the maximum probability
        """
        I = np.zeros(self.img.shape)
        idx = np.argmax(est)
        if np.min(est) < 0 or np.max(est) > 1:
            # Must be log probabilities or unnormalized probabilities,
            # so bring into the range [0, 1]
            est = est - np.min(est)
            est = est/np.max(est)
        est = est**p
        I[self.X[:, 1], self.X[:, 0]] += est+0.2
        res = []
        res.append(plt.imshow(I, cmap='magma'))
        plt.gca().invert_yaxis()
        plt.clim(0, 1.2) # Consistent coloring with 0.2 offset so all cells are visible
        if show_max:
            res.append(plt.scatter([self.X[idx, 0]], [self.X[idx, 1]], marker='X', c='C2'))
        return res

    def simulate_trajectory(self, xs):
        """
        Simulate a trajectory which is a shortest path between a set
        of points on the map

        Parameters
        ----------
        x: list of [float, float]
            Location of points

        """
        from collections import deque
        X = []
        for i in range(len(xs)-1):
            ## BFS
            x0 = tuple(xs[i])
            if not x0 in self.pos2idx:
                print("Error: ", x0, "is not a position on the grid")
                return
            x1 = tuple(xs[i+1])
            if not x1 in self.pos2idx:
                print("Error: ", x1, "is not a position on the grid")
                return
            i1 = self.pos2idx[x0]
            i2 = self.pos2idx[x1]
            visited = set([])
            queue = deque()
            queue.append((i1, None))
            found = False
            state = None
            while len(queue) > 0 and not found:
                (i, prev) = queue.popleft()
                if i == i2:
                    found = True
                    state = (i, prev)
                else:
                    for n in self.neighbors[i]:
                        if not n in visited:
                            visited.add(n)
                            queue.append((n, (i, prev)))
            Xi = []
            while state[-1]:
                Xi.append(self.X[state[0], :])
                state = state[-1]
            Xi.reverse()
            X += Xi
        return np.array(X)


text:
# -*- coding: utf-8 -*-
"""
optimizer for different models

Created on Sep 16 2018
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

__author__ = "yuhao"

from builtins import bytes, str
import json
from collections import OrderedDict

import numpy as np
from scipy.optimize import curve_fit

from pygeopressure.pressure.bowers import (
    virgin_curve, invert_virgin, power_bowers)
from pygeopressure.velocity.extrapolate import normal, normal_dt
from pygeopressure.pressure.obp import traugott
from pygeopressure.basic.well import Well
from pygeopressure.basic.well_log import Log
from pygeopressure.basic.utils import rmse, pick_sparse
from pygeopressure.pressure.eaton import power_eaton
from sklearn import linear_model

def optimize_bowers_virgin(well, vel_log, obp_log, upper, lower,
                           pres_log='loading', mode='nc', nnc=5):
    """
    Optimizer for Bowers loading curve

    Parameters
    ----------
    well : Well
    vel_log : Log or str
        Log object or well log name stored in well
    obp_log : Log or str
        Log object or well log name stored in well
    upper : float or str
        upper bound of nct, depth value or horizon name
    lower : float or str
        lower bound of nct, depth value or horizon name
    pres_log : Log or str
        Log object storing measured pressure value or Pressure name
        stored in well
    mode : {'nc', 'pres', 'both'}
        which pressure to use for optimization,
        - 'nc' : points on NCT
        - 'pres' : points in pres_log
        - 'both' : both of them
    nnc : int
        number of points to pick on NCT

    Returns
    -------
    a, b : tuple of floats
        optimized bowers loading curve coefficients
    rms_err : float
        root mean square error of pressure
    """
    if isinstance(upper, (bytes, str)):
        depth_upper = well.params['horizon'][upper]
    else:
        depth_upper = upper
    if isinstance(upper, (bytes, str)):
        depth_lower = well.params['horizon'][lower]
    else:
        depth_lower = lower
    if isinstance(vel_log, (bytes, str)):
        vel_log = well.get_log(vel_log)
    if isinstance(obp_log, (bytes, str)):
        obp_log = well.get_log(obp_log)
    if isinstance(pres_log, (bytes, str)):
        # pres_log = well.get_loading_pressure()
        pres_log = well.get_pressure(pres_log)
    depth = np.array(obp_log.depth)

    nct_vel_to_fit = []
    nct_es_to_fit = []
    pres_vel_to_fit = []
    pres_es_to_fit = []
    if mode == 'nct' or mode == 'both':
        nct_es_data = np.array(obp_log.data) - np.array(well.hydrostatic)
        nct_mask = depth < depth_lower
        nct_mask *= depth > depth_upper
        nct_mask *= depth < vel_log.stop

        nct_vel_interval = np.array(vel_log.data)[nct_mask]
        nct_es_interval = nct_es_data[nct_mask]

        nct_vel_to_fit = np.array(pick_sparse(nct_vel_interval, nnc))
        nct_es_to_fit = np.array(pick_sparse(nct_es_interval, nnc))
    if mode == 'pres' or mode == 'both':
        vel = list()
        obp = list()
        pres = list()
        for dp in pres_log.depth:
            idx = np.searchsorted(depth, dp)
            vel.append(vel_log.data[idx])
            obp.append(obp_log.data[idx])
        vel, obp, pres = np.array(vel), np.array(obp), np.array(pres_log.data)
        es = obp - pres

        pres_vel_to_fit = vel
        pres_es_to_fit = es
    vel_to_fit = np.append(nct_vel_to_fit, pres_vel_to_fit)
    es_to_fit = np.append(nct_es_to_fit, pres_es_to_fit)

    popt, _ = curve_fit(virgin_curve, es_to_fit, vel_to_fit)
    a, b = popt

    es_predicted = invert_virgin(vel_to_fit, a, b)
    rms_err = rmse(es_to_fit, es_predicted)

    return a, b, rms_err

def optimize_bowers_unloading(well, vel_log, obp_log, a, b,
                              vmax, pres_log='unloading'):
    """
    Optimize for Bowers Unloading curve parameter U

    Parameters
    ----------
    well : Well
    vel_log : Log or str
        Log object or well log name stored in well
    obp_log : Log or str
        Log object or well log name stored in well
    vmax : float
        vmax in bowers unloading curve
    pres_log : Log or str
        Log object storing measured pressure value or Pressure name
        stored in well

    Returns
    -------
    u : float
        unloading curve cofficient U
    error : float
        Relative RMS error
    """
    if isinstance(vel_log, (bytes, str)):
        vel_log = well.get_log(vel_log)
    if isinstance(obp_log, (bytes, str)):
        obp_log = well.get_log(obp_log)
    if isinstance(pres_log, (bytes, str)):
        pres_log = well.get_pressure(pres_log)

    depth = np.array(obp_log.depth)

    vel = list()
    obp = list()
    pres = list()
    for dp in pres_log.depth:
        idx = np.searchsorted(depth, dp)
        vel.append(vel_log.data[idx])
        obp.append(obp_log.data[idx])
    vel, obp, pres = np.array(vel), np.array(obp), np.array(pres_log.data)
    es = obp - pres

    sigma_max = invert_virgin(vmax, a, b)

    sigma_vc = invert_virgin(vel, a, b)

    popt, _ = curve_fit(power_bowers, sigma_vc/sigma_max, es/sigma_max)

    u, = popt

    return u

def optimize_bowers_trace(depth_tr, vel_tr, obp_tr, hydro_tr,
                          depth_upper, depth_lower):

    es_data = np.array(obp_tr) - np.array(hydro_tr)
    depth = np.array(depth_tr)

    mask = depth < depth_lower
    mask *= depth > depth_upper

    vel_interval = np.array(vel_tr)[mask]
    es_interval = es_data[mask]

    vel_to_fit = pick_sparse(vel_interval, 3)
    es_to_fit = pick_sparse(es_interval, 3)

    popt, _ = curve_fit(virgin_curve, es_to_fit, vel_to_fit)
    a, b = popt

    return a, b

def optimize_eaton(well, vel_log, obp_log, a, b, pres_log="loading"):
    """
    Optimizer for Eaton model

    Parameters
    ----------
    well : Well
    vel_log : Log or str
        Log object or well log name stored in well
    obp_log : Log or str
        Log object or well log name stored in well
    a, b : float
        coefficients of NCT
    pres_log : Log or str
        Log object storing measured pressure value or Pressure name
        stored in well

    Returns
    -------
    n : float
        optimized eaton exponential
    min_eer : float
        minimum error abtained by optimized n
    rms_err : array
        array of rms error of different n around minmum
    """
    if isinstance(vel_log, (bytes, str)):
        vel_log = well.get_log(vel_log)
    if isinstance(obp_log, (bytes, str)):
        obp_log = well.get_log(obp_log)
    if isinstance(pres_log, (bytes, str)):
        # pres_log = well.get_loading_pressure()
        pres_log = well.get_pressure(pres_log)

    depth = np.array(obp_log.depth)

    hydrostatic = np.array(well.hydrostatic)
    es_normal = np.array(obp_log.data) - hydrostatic
    v_normal = normal(depth, a, b)

    vel = list()
    vel_norm = list()
    es_norm = list()
    obp = list()
    pres = list()
    for dp in pres_log.depth:
        idx = np.searchsorted(depth, dp)
        vel.append(vel_log.data[idx])
        vel_norm.append(v_normal[idx])
        es_norm.append(es_normal[idx])
        obp.append(obp_log.data[idx])
    vel, vel_norm = np.array(vel), np.array(vel_norm)
    vel_ratio = vel / vel_norm

    obp, pres = np.array(obp), np.array(pres_log.data)
    es = obp - pres
    es_norm = np.array(es_norm)
    es_ratio = es / es_norm

    popt, _ = curve_fit(power_eaton, vel_ratio, es_ratio)
    n, = popt

    return n

def optimize_multivaraite(well, obp_log, vel_log, por_log, vsh_log, B,
                          upper, lower):
    if isinstance(vel_log, (bytes, str)):
        vel_log = well.get_log(vel_log)
    if isinstance(por_log, (bytes, str)):
        por_log = well.get_log(por_log)
    if isinstance(vsh_log, (bytes, str)):
        vsh_log = well.get_log(vsh_log)
    if isinstance(obp_log, (bytes, str)):
        obp_log = well.get_log(obp_log)
    # --------------------------------------------
    obp_data = np.array(obp_log.data)
    por_data = np.array(por_log.data)
    vp_data = np.array(vel_log.data)
    vsh_data = np.array(vsh_log.data)

    depth = well.depth
    hydrostatic = well.hydrostatic

    es = obp_data - hydrostatic

    B = well.params['bowers']['B'] if B is None else B  # 0.88
    es_data = es**B

    mask = np.isfinite(vp_data)
    mask *= np.isfinite(por_data)
    mask *= np.isfinite(vsh_data)
    mask *= depth < lower #3500
    mask *= depth > upper #1500

    por_data = por_data[mask]
    vsh_data = vsh_data[mask]
    es_data = es_data[mask]
    vp_data = vp_data[mask]
    depth = depth[mask]

    por_data = por_data.reshape((por_data.shape[0], 1))
    vsh_data = vsh_data.reshape((vsh_data.shape[0], 1))
    es_data = es_data.reshape((es_data.shape[0], 1))
    vp_data = vp_data.reshape((vp_data.shape[0], 1))

    #_____________________REGRESSION______________________
    data = np.hstack((por_data, vsh_data, es_data))
    # global REG
    reg = linear_model.LinearRegression()
    reg.fit(data, vp_data)

    r2 = reg.score(data, vp_data)

    a0 = reg.intercept_[0]
    a1 = -reg.coef_[0][0]
    a2 = -reg.coef_[0][1]
    a3 = reg.coef_[0][2]

    return a0, a1, a2, a3

def optimize_nct(vel_log, fit_start, fit_stop):
    """
    Fit velocity NCT

    Parameters
    ----------
    vel_log : Log
        Velocity log
    fit_start, fit_stop : float
        start and end depth for fitting

    Returns
    -------
    a, b : float
        NCT coefficients
    """
    fit_start = fit_start
    fit_stop = fit_stop

    if fit_start is None or fit_start < vel_log.top:
        fit_start = vel_log.top
    if fit_stop is None or fit_stop > vel_log.bottom:
        fit_stop = vel_log.bottom

    a, b = optimize_nct_trace(
        np.array(vel_log.depth), np.array(vel_log.data), fit_start, fit_stop,
        pick=False)

    return a, b

def optimize_nct_trace(depth, vel, fit_start, fit_stop, pick=True):

    mask = depth > fit_start
    mask *= depth < fit_stop
    mask *= np.isfinite(vel)

    depth_interval = np.array(depth)[mask]
    vel_interval = np.array(vel)[mask]

    if pick is True:
        depth_to_fit = pick_sparse(depth_interval, 3)
        vel_to_fit = pick_sparse(vel_interval, 3)
    else:
        depth_to_fit = depth_interval
        vel_to_fit = vel_interval

    dt = vel_to_fit**(-1)
    log_dt = np.log(dt)

    popt, _ = curve_fit(normal_dt, depth_to_fit, log_dt)
    a, b = popt

    return a, b

def optimize_traugott(den_log, fit_start, fit_stop, kb=0, wd=0):
    """
    Fit density variation against depth with Traugott equation

    Parameters
    ----------
    den_log : Log
        Density log
    fit_start, fit_stop : float
        start and end depth for fitting
    kb : float
        kelly bushing height in meters
    wd : float
        water depth in meters

    Returns
    -------
    a, b : float
        Traugott equation coefficients
    """
    start_idx = den_log.get_depth_idx(fit_start)
    stop_idx = den_log.get_depth_idx(fit_stop) + 1
    depth = np.array(den_log.depth[start_idx: stop_idx + 1])
    den = np.array(den_log.data[start_idx: stop_idx + 1])

    mask = np.isfinite(den)
    den_finite = den[mask]
    depth_finite = depth[mask]

    depth_finite_shift = depth_finite - kb - wd

    popt, _ = curve_fit(traugott, depth_finite_shift, den_finite)
    a, b = popt

    return a, b

max_stars_repo_name: yohanesnuwara/pyGeoPressure

reason: import numpy,from scipy


text:
# CRRN f1_score: 0.8112233240414933	 # UDLR f1_score: 0.7072364069321339 -- 500 / 2304
import os
import time
import torch
import argparse
from torch import optim
import torch.utils.data
import torchvision
# from model.ConvRRN import ST_EncDec as Model
from model.ConvRRN import Wrapper as Model
from model.CNN import CNN
import numpy as np
import cv2
from tqdm import tqdm

# os.environ["CUDA_VISIBLE_DEVICES"]='0'

class Dataset(torch.utils.data.Dataset):
    def __init__(self, mode='train'): ## mode must be train, val, test
        super(Dataset, self).__init__()
        print('Dataset initialize starts')
        self.mode = mode
        self.data_dir_path = '../Dataset/udlr+bicrrn/'
        self.size = (400,144)
        self.list_file = open(os.path.join(self.data_dir_path,"list", f"{mode}.txt"), 'r')
        self.data_list = []
        self.label_list = []
        prev_subdir_name = None
        set_counter = 0
        self.weight = torch.tensor([0.03905, 0.96095]) ## calculated by get_weight()
        while True:
            line = self.list_file.readline()
            while '90frame' in line:
                line = self.list_file.readline()
                if not line: break
            if not line: break

            ############################################################################################
            subdir_name = line[:-10]		####Make Data Multiple of 10
            if subdir_name ==  prev_subdir_name:
                set_counter += 1
            else:
                if set_counter%10 != 9 or set_counter!=0:
                    self.data_list = self.data_list[:-1*(set_counter%10+1)]
                    self.label_list = self.label_list[:-1*(set_counter%10+1)]

                set_counter = 0
            ############################################################################################

            data_name = f"{self.data_dir_path}/udlr_ENet_whiten{line[:-4]}jpg"
            label_name = f"{self.data_dir_path}/label{line[:-4]}png"
            self.data_list.append(data_name)
            self.label_list.append(label_name)
            prev_subdir_name = subdir_name
        print('Dataset initialize ends')

    def get_weight(self):
        ratio = 0
        count = 0
        for label in tqdm(self.label_list):
            label_img = cv2.imread(label, cv2.IMREAD_GRAYSCALE)
            label_size = np.size(label_img)
            num_of_one = int(label_img.sum()/255)
            count += 1
            if count == 0:
                ratio = num_of_one/label_size
            else:
                ratio = count/(count+1)*ratio + num_of_one/label_size/(count+1)
                if count%1000 == 0:
                    print(ratio)
        return ratio

    def __getitem__(self, index):
        data_list = self.data_list[index*10:(index+1)*10]
        label_list = self.label_list[index*10:(index+1)*10]

        inputs = []
        targets = []
        for i in range(10):
            input_numpy = torch.round(torch.Tensor(cv2.resize(cv2.imread(data_list[i], cv2.IMREAD_GRAYSCALE),self.size))/255)
            target_img = torch.Tensor(cv2.resize(cv2.imread(label_list[i], cv2.IMREAD_GRAYSCALE),self.size))/255
            h,w = input_numpy.size()
            h = h+1
            w = w+1
            padded_input = torch.zeros(h, w, dtype=torch.float)
            padded_input[:input_numpy.size(0), :input_numpy.size(1)] = input_numpy
            padded_target = torch.zeros(h, w, dtype=torch.float)
            padded_target[:input_numpy.size(0), :input_numpy.size(1)] = target_img
            inputs.append(padded_input)
            targets.append(padded_target)

        inputs = torch.stack(inputs,0)
        targets = torch.stack(targets,0)
        return inputs, targets, data_list

    def __len__(self):
        return int(len(self.data_list)/10)

def split(targets, inputs, outputs, i):
    target = targets.squeeze()[i]
    input = inputs.squeeze()[i]
    output = outputs.squeeze()[i]
    return target, input, output

def gt_lane_existance(gt_map, gt_th):
    if np.sum(gt_map) > gt_th:
        return True
    else:
        return False

def pred_lane_existance(pred_map, pred_th):
    if np.sum(pred_map) > pred_th:
        return True
    else:
        return False

def gt_pred_IOU(gt, pred, iou_th):
    assert(pred.shape == gt.shape)
    summation = pred+gt
    intersection = np.sum(summation>1)
    union = np.sum(summation>0)
    iou = intersection/union
    if iou > iou_th:
        return True
    else:
        return False

def f1_counter(targets, inputs, outputs, gt_th, pred_th, iou_th): #call after inference
    tp,fp,tn,fn = 0,0,0,0

    for i in range(10):
        if i < 9: ##last frame only
            continue
        target, input, output = split(targets, inputs, outputs, i)
        gt_existance = gt_lane_existance(target, gt_th)
        pred_existance = pred_lane_existance(output, pred_th)
        if gt_existance and pred_existance:
            if gt_pred_IOU(target, output, iou_th):
                tp+=1
            else:
                fp+=1
        elif not gt_existance and pred_existance:
            fp+=1
        elif not gt_existance and not pred_existance:
            tn+=1
        elif gt_existance and not pred_existance:
            fn+=1
    return tp,fp,tn,fn

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='PCB anomaly detection')
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--num_epochs', type=int, default=501)
    parser.add_argument('--n_layers', type=int, default=2)
    parser.add_argument('--lr', type=float, default=1e-4)
    #	parser.add_argument('--init_tr', type=float, default=0.25)
    #	parser.add_argument('--final_tr', type=float, default=0.0)
    parser.add_argument('--gpu_ids', nargs='+', type=int, default=[0])
    parser.add_argument('--is_attn', action='store_true')
    parser.add_argument('--aug', action='store_true')
    parser.add_argument('--mode', type=str, default='ConvConjLSTM', help='ConvLSTM, ConvConjLSTM, CNN')
    parser.add_argument('--test_mode', type=str, default='CRRN', help='CRRN or UDLR')

    args = parser.parse_args()

    if not os.path.isdir('visualization'):
        os.mkdir('visualization')

    if args.mode == 'CNN':
        vis_dir = 'visualization/CNN'
    else:
        vis_dir = 'visualization/CRRN'

    if not os.path.isdir(vis_dir):
        os.mkdir(vis_dir)

    # train_dataset = Dataset(mode='test')
    # train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True,
    #                                                num_workers=0)

    valid_dataset = Dataset(mode='test')
    valid_dataloader = torch.utils.data.DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False,
                                                   num_workers=6)

    if args.mode == 'CNN':
        model = CNN([3, 64, 64], [5, 5])
    else:
        model = Model(args.mode, [1, 64, 64], [5, 5], args.n_layers, args.is_attn)

    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
    if torch.cuda.device_count() > 1:
        if args.gpu_ids == None:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            device = torch.device('cuda:0')
        else:
            print("Let's use", len(args.gpu_ids), "GPUs!")
            device = torch.device('cuda:' + str(args.gpu_ids[0]))

    else:
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    print('args.gpu_ids', args.gpu_ids)

    model = torch.nn.DataParallel(model, device_ids=args.gpu_ids)
    model = model.to(device)

    start_time = time.time()
    state = torch.load('parameters/290_parameter')
    model.load_state_dict(state['state_dict'])
    model.eval()
    with torch.no_grad():
        threshold = 0.99 # [0.95, 0.99]

        precision = [0,0,0,0,0,0,0,0,0]
        recall = [0,0,0,0,0,0,0,0,0]
        cnt = 0
        out_dir = './output'
        outputs_list = []
        inputs_list = []
        targets_list = []
        total_tp = 0
        total_fp = 0
        total_tn = 0
        total_fn = 0
        for inputs, targets, data_list in tqdm(valid_dataloader):

            inputs = inputs.to(device)
            inputs = inputs.unsqueeze(2)
            targets = targets.to(device)

            if (len(inputs) != args.batch_size):
                break
            if args.test_mode == 'CRRN':
                outputs = model(inputs)  # batch by 10 by 2 by h by w
                outputs = torch.nn.Sigmoid()(outputs)
                outputs = outputs.squeeze()[:,1].detach().cpu().numpy()
                inputs = inputs.squeeze().detach().cpu().numpy()
                targets = targets.squeeze().detach().cpu().numpy()
            else:
                inputs = inputs.squeeze().detach().cpu().numpy()
                targets = targets.squeeze().detach().cpu().numpy()
                outputs = inputs

            outputs = (outputs>threshold)
            # outputs = torch.tensor(outputs)

        # for i in range(len(targets_list)):
        #     final_result = torch.cat((inputs_list[i].squeeze(), targets_list[i].float(), outputs_list[i][:,1]),0).unsqueeze(1)
        #     scaled_result = torchvision.utils.make_grid(final_result, nrow=10, padding=2, pad_value=1)
        #     os.makedirs('./testnoth', exist_ok=True)
        #     torchvision.utils.save_image(scaled_result, f'./testtest/{i}.png')
        # print('a')

            iou_th = 0.3
            # l1_th_list = [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]
            # for i, l1_th in enumerate(l1_th_list):
            tp,fp,tn,fn = f1_counter(targets, inputs, outputs, gt_th=0, pred_th=2304, iou_th=iou_th)
            total_tp += tp
            total_fp += fp
            total_tn += tn
            total_fn += fn
        precision = total_tp/(total_tp+total_fp)
        recall = total_tp/(total_tp+total_tn)
        print(f"tp: {total_tp} \nfp: {total_fp} \ntn: {total_tn}")
        print(f"precision: {precision}")
        print(f"recall: {recall}")
        print(f"f1_score: {2*precision*recall/(precision+recall)}")