text:
# Copyright 2017 Battelle Energy Alliance, LLC
#
# 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.
from glob import glob
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.cm as cm

files = glob('dummy_*.csv')

fptr = open(files[0])
fptr.readline()
numTimeSteps = len([1 for line in fptr])
fptr.close()

numPoints = len(files)

xLabel = 'x1'
yLabel = 'x2'
label = 'labels'

X = np.zeros(shape=(numTimeSteps,numPoints))
Y = np.zeros(shape=(numTimeSteps,numPoints))
C = np.zeros(shape=(numTimeSteps,numPoints))

for filename in files:
  counter = int(filename.split('_')[1].split('.')[0])
  fptr = open(filename)
  labels = fptr.readline().split(',')
  for idx,line in enumerate(fptr):
    tokens = line.split(',')
    X[idx,counter] = float(tokens[labels.index(xLabel)])
    Y[idx,counter] = float(tokens[labels.index(yLabel)])
    C[idx,counter] = float(tokens[labels.index(label)])
  fptr.close()

mapper = cm.ScalarMappable(cmap=cm.get_cmap('Set1'))
mapper.set_clim(np.min(C),np.max(C))

def init_plot():
  pass

def update_plot(idx):
  fig.clf()
  if idx == numTimeSteps:
    for i in xrange(numPoints):
      ## The line below fails as the plot function cannot figure out how to take
      ## a list of rgba values
      # myScatter = plt.plot(X[:,i], Y[:,i], c=mapper.cmap(C[:,i]))
      myScatter = plt.plot(X[:,i], Y[:,i], c='k')
    for i in xrange(numTimeSteps):
      myScatter = plt.scatter(x=X[i,:], y=Y[i,:], c=C[i,:], cmap=mapper.cmap)
  else:
    myScatter = plt.scatter(x=X[idx,:], y=Y[idx,:], c=C[idx,:], cmap=mapper.cmap)
  plt.xlim(-2,2)
  plt.ylim(-0.5,3.5)
  return None

fig = plt.figure()
im_ani = animation.FuncAnimation(fig, update_plot, frames=range(numTimeSteps+1),
                                 interval=500, init_func=init_plot,
                                 repeat_delay=5000, repeat=True)

im_ani.save('test.gif', writer='imagemagick',fps=10)
im_ani.save('test.mp4', fps=10)
# plt.show()


text:
# encoding: utf-8
"""Pade analytic continuation for Green's functions and self-energies.

The main aim of this module is to provide analytic continuation based on
averaging over multiple Pade approximates (similar to [1]_).

In most cases the following high level function should be used:

`averaged`, `avg_no_neg_imag`
   Return one-shot analytic continuation evaluated at `z`.

`Averager`
   Returns a function for repeated evaluation of the continued function.

References
----------
.. [1] Schött et al. “Analytic Continuation by Averaging Pade Approximants”.
   Phys Rev B 93, no. 7 (2016): 075104.
   https://doi.org/10.1103/PhysRevB.93.075104.

"""
import logging

from abc import ABC, abstractmethod
from typing import Optional as Opt
from functools import partial
from itertools import islice

import numpy as np

from gftool import Result
from gftool.precision import PRECISE_TYPES as _PRECISE_TYPES

LOGGER = logging.getLogger(__name__)
LOGGER.addHandler(logging.NullHandler())

_nan_std = partial(np.nanstd, ddof=1, axis=0)

class KindSelector(ABC):
    """Abstract filter class to determine high-frequency behavior of Pade.

    We denote approximants with the corresponding high frequency behavior as
    *valid*.
    Considers all valid approximants including between `n_min` and `n_max`
    Matsubara frequencies.
    """

    @abstractmethod
    def __init__(self, n_min, n_max):
        """Consider approximants including between `n_min` and `n_max` Matsubara frequencies."""
        if n_min < 1:
            raise ValueError(f"`n_min` needs to be at least 1 (n_min: {n_min}).")
        if n_max <= n_min:
            raise ValueError(f"`n_max` ({n_max}) needs to be bigger than `n_min` ({n_min}).")
        self.start = n_min - 1  # indices start from 0
        self.stop = n_max - 1  # indices start from 0
        self.step = NotImplemented

    def islice(self, iterable):
        """Return an iterator whose next() method returns valid values from `iterable`."""
        return islice(iterable, self.start, self.stop, self.step)

    @property
    def slice(self):
        """Return slice selecting the valid approximants."""
        return slice(self.start, self.stop, self.step)

    def __getitem__(self, index):
        """Get indices."""
        return range(self.start, self.stop, self.step)[index]

    def __len__(self):
        """Get number of approximants."""
        return len(range(self.start, self.stop, self.step))

    def __repr__(self):
        """Return Meaningful string."""
        return (self.__class__.__name__
                + f'(start={self.start}, stop={self.stop}, step={self.step})')

class KindGf(KindSelector):
    """Filter approximants such that the high-frequency behavior is :math:`1/ω`.

    We denote approximants with the corresponding high frequency behavior as
    *valid*.
    Considers all valid approximants including between `n_min` and `n_max`
    Matsubara frequencies.
    """

    def __init__(self, n_min, n_max):
        """Consider approximants including between `n_min` and `n_max` Matsubara frequencies."""
        if not n_min % 2:
            n_min += 1  # odd number for 1/z required
        super().__init__(n_min, n_max)
        self.step = 2

class KindSelf(KindSelector):
    """Filter approximants such that the high-frequency behavior is a constant.

    We denote approximants with the corresponding high frequency behavior as
    *valid*.
    Considers all valid approximants including between `n_min` and `n_max`
    Matsubara frequencies.
    """

    def __init__(self, n_min, n_max):
        """Consider approximants including between `n_min` and `n_max` Matsubara frequencies."""
        if n_min % 2:
            n_min += 1  # even number for constant tail required
        super().__init__(n_min, n_max)
        self.step = 2

def FilterNegImag(threshold=1e-8):
    """Return function to check if imaginary part is smaller than `threshold`.

    This methods is designed to create `valid_pades` for `Averager`.
    The imaginary part of retarded Green's functions and self-energies must be
    negative, this is checked by this filter.
    A threshold is given as Pade overshoots when the function goes sharply to 0.
    See for example the semi-circular spectral function of the Bethe lattice
    with infinite Coordination number as example.
    """
    def filter_neg_imag(pade_iter):
        r"""Check which Pade approximants have a negative imaginary part.

        Parameters
        ----------
        pade_iter : iterable of (..., N_z) complex np.ndarray
            Iterator of analytic continuations as generated by `calc_iterator`.

        Returns
        -------
        is_valid : (...) bool np.ndarray
            True for all approximants that fulfill `apporximant.imag < threshold`.

        """
        is_valid = np.array([np.all(pade.imag < threshold, axis=-1) for pade in pade_iter])
        LOGGER.debug("Filter Pades with positive imaginary part (threshold: %s): %s",
                     threshold, np.count_nonzero(is_valid, axis=0))
        return is_valid

    return filter_neg_imag

def FilterNegImagNum(abs_num=None, rel_num=None):
    """Return function to check how bad the imaginary part gets.

    This methods is designed to create `valid_pades` for `Averager`.
    The imaginary part of retarded Green's functions and self-energies must be
    negative, this is checked by this filter.
    All continuations that are *valid* in this sense are kept, the worst invalid
    are dropped till only `abs_num` remain.

    Warnings
    --------
    Only checked for flat inputs.

    """
    assert abs_num is None or rel_num is None
    assert abs_num is not None or rel_num is not None
    def filter_neg_imag_num(pade_iter):
        badness = np.array([np.max(pade.imag, axis=-1) for pade in pade_iter])
        abs_num_ = rel_num * badness.shape[0] if abs_num is None else abs_num
        if abs_num_ >= badness.shape[0]:
            LOGGER.warning("Skipping filter, not enough Pades (#Pades = %s)",
                           badness.shape[0])
            return np.ones_like(badness, dtype=bool)

        keep = np.argsort(badness, axis=0)[:abs_num_]  # abs_num_ best results
        is_valid = np.zeros_like(badness, dtype=bool)
        is_valid[keep] = True
        is_valid[badness <= 0] = True
        LOGGER.debug("Filter Pades with positive imaginary part (keep best %s): %s",
                     abs_num_, np.count_nonzero(is_valid, axis=0))
        assert np.all(np.count_nonzero(is_valid, axis=0) >= abs_num_)
        return is_valid
    return filter_neg_imag_num

def FilterHighVariance(rel_num: Opt[float] = None, abs_num: Opt[int] = None):
    """Return function to filter continuations with highest variance.

    Parameters
    ----------
    rel_num : float, optional
        The relative number of continuations to keep.
    abs_num : int, optional
        The absolute number of continuations to keep.

    Returns
    -------
    filter_high_variance : callable
        The filter function (pade_iter) -> np.ndarray.

    """
    assert abs_num is None or rel_num is None
    assert abs_num is not None or rel_num is not None
    if rel_num is not None:
        assert 0. < rel_num < 1.
    if abs_num is not None:
        assert abs_num > 0
    def filter_high_variance(pade_iter):
        """Remove the continuations with highest variance.

        Parameters
        ----------
        pade_iter : iterable of (..., N_z) complex np.ndarray
            Iterator of analytic continuations as generated by `calc_iterator`.

        Returns
        -------
        is_valid : (...) bool np.ndarray
            Boolean array indicating which continuations to keep.

        """
        pade = np.array(list(pade_iter))
        pade_sum = np.sum(pade, axis=0)
        N_pades = pade.shape[0]
        # iteratively remove Pades with larges deviation
        # why iterative?
        # Awful Pades might give wrong features in average, so it should be corrected
        if abs_num is None:
            abs_num_ = int(rel_num*N_pades)
        bad = []  # isin needs list not set
        for nn in range(N_pades, abs_num_, -1):
            diff = nn*pade - pade_sum  # 50% of time
            distance = np.sum(diff.real**2 + diff.imag**2, axis=-1)  # 40% of time
            badness = np.argsort(distance, axis=0)[::-1]  # truncate what is not needed
            newbad = badness[np.isin(badness, bad, invert=True)][0]
            bad.append(newbad)
            pade_sum -= pade[newbad]
        try:
            is_valid = np.ones_like(distance, dtype=bool)
        except UnboundLocalError:
            LOGGER.warning("Not enough pades to filter (#Pades = %s)", N_pades)
            return np.ones_like(pade[..., 0], dtype=bool)
        is_valid[badness[:-abs_num_]] = False
        # assert set(badness[:-abs_num_]) == set(bad)  # FIXME
        return is_valid
    return filter_high_variance

def _contains_nan(array) -> bool:
    """Check if `array` contains any NaN."""
    flat = array.reshape(-1)
    return np.isnan(np.dot(flat, flat))

def coefficients(z, fct_z) -> np.ndarray:
    """Calculate the coefficients for the Pade continuation.

    Parameters
    ----------
    z : (N_z, ) complex ndarray
        Array of complex points
    fct_z : (..., N_z) complex ndarray
        Function at points `z`

    Returns
    -------
    coefficients : (..., N_z) complex ndarray
        Array of Pade coefficients, needed to perform Pade continuation.
        Has the same same shape as `fct_z`.

    Raises
    ------
    ValueError
        If the size of `z` and the last dimension of `fct_z` do not match.

    Notes
    -----
    The calculation is always done in quad precision (complex256), as it is
    very sensitive towards rounding errors. Afterwards the type of the result
    is cast back to double precision (complex128) unless the input data of
    `fct_z` was already quad precision {float128, complex256}, see
    `_PRECISE_TYPES`. This avoids giving the illusion that the results are more
    precise than the input.

    """
    if z.shape != fct_z.shape[-1:]:
        raise ValueError(f"Dimensions of `z` ({z.shape}) and `fct_z` ({fct_z.shape}) mismatch.")
    res = fct_z.astype(dtype=np.complex256, copy=True)
    for ii in range(z.size - 1):
        res[..., ii+1:] = (res[..., ii:ii+1]/res[..., ii+1:] - 1.)/(z[ii+1:] - z[ii])
    complex_pres = np.complex256 if fct_z.dtype in _PRECISE_TYPES else np.complex
    LOGGER.debug("Input type %s precise: %s -> result type: %s",
                 fct_z.dtype, fct_z.dtype in _PRECISE_TYPES, complex_pres)
    return res.astype(complex_pres, copy=False)

@partial(np.vectorize, otypes=[np.complex], signature='(n),(n)->(n)')
def masked_coefficients(z, fct_z):
    """Calculate coefficients but ignore extreme values.

    Like `coefficients` but probably better for noisy data.
    """
    mat = np.zeros((z.size, *fct_z.shape), dtype=np.complex256)
    mask = np.empty_like(z, dtype=bool)
    mask[:] = True
    # cutoff = 1e-6
    cutoff = 1e-4
    mat[0] = fct_z

    assert np.abs(fct_z[0]) > cutoff
    # last valid
    last_it = 0
    last_coeff = mat[last_it, last_it]

    def signi_diff(element) -> bool:
        """Return if the difference of `element` and `last_coeff` is larger `cutoff`."""
        return abs(last_coeff - element) > cutoff

    def comparable_mag(element) -> bool:
        """Return weather the magnitude of `element` is comparable to `last_coeff`."""
        return abs(last_coeff)/cutoff > abs(element) > abs(last_coeff)*cutoff

    for ii, mat_pi in enumerate(mat[1:], start=1):
        if signi_diff(mat[last_it, ii]) and comparable_mag(mat[last_it, ii]):
            # regular update
            mat_pi[ii] = (last_coeff/mat[last_it, ii] - 1.)/(z[ii] - z[last_it])
            for jj in range(ii+1, z.size):
                if not mask[jj]:
                    continue
                if comparable_mag(mat[last_it, jj]):
                    mat_pi[jj] = (last_coeff/mat[last_it, jj] - 1.)/(z[jj] - z[last_it])
                elif abs(last_coeff) < abs(mat[last_it, jj])*cutoff:  # tiny quotient
                    mat_pi[jj] = (-1)/(z[jj] - z[last_it])
                else:  # huge quotient
                    mat_pi[jj] = np.infty
            last_it = ii
            last_coeff = mat_pi[ii]
        else:
            mask[ii] = False
    # pylint: disable=invalid-unary-operand-type
    LOGGER.info("Number of eliminated coefficients: %s", np.count_nonzero(~mask))
    return mat.diagonal(axis1=0, axis2=-1)

def calc_iterator(z_out, z_in, coeff):
    r"""Calculate Pade continuation of function at points `z_out`.

    The continuation is calculated for different numbers of coefficients taken
    into account, where the number is in [n_min, n_max].
    The algorithm is take from [2]_.

    Parameters
    ----------
    z_out : complex ndarray
        points at with the functions will be evaluated
    z_in : (N_in,) complex ndarray
        complex mesh used to calculate `coeff`
    coeff : (..., N_in) complex ndarray
        coefficients for Pade, calculated from `pade.coefficients`

    Yields
    ------
    pade_calc : (..., N_in, z_out.shape) complex np.ndarray
        Function evaluated at points `z_out`.
        numbers of Matsubara frequencies between `n_min` and `n_max`.
        The shape of the elements is the same as `coeff.shape` with the last
        dimension corresponding to N_in replaced by the shape of `z_out`:
        (..., N_in, \*z_out.shape).

    References
    ----------
    .. [2] Vidberg, H. J., and J. W. Serene. “Solving the Eliashberg Equations
       by Means of N-Point Pade Approximants.” Journal of Low Temperature
       Physics 29, no. 3-4 (November 1, 1977): 179-92.
       https://doi.org/10.1007/BF00655090.

    """
    assert coeff.shape[-1] == z_in.size
    target_shape = coeff.shape[:-1] + z_out.shape

    z_out = z_out.reshape(-1)  # accept arbitrary shaped z_out
    id1 = np.ones_like(z_out, dtype=complex)

    pade_prev = 0.*id1
    pade = coeff[..., 0:1]*id1
    B2 = id1

    multiplier = (z_out - z_in[:-1, np.newaxis])*coeff[..., 1:, np.newaxis]
    # move N_in axis in front to iterate over it
    multiplier = np.moveaxis(multiplier, -2, 0).copy()

    for multiplier_im in multiplier:
        multiplier_im = multiplier_im / B2
        B2 = 1 + multiplier_im
        pade, pade_prev = (pade + multiplier_im*pade_prev)/B2, pade

        yield pade.reshape(target_shape)

def Averager(z_in, coeff, *, valid_pades, kind: KindSelector):
    """Create function for averaging Pade scheme.

    Parameters
    ----------
    z_in : (N_in,) complex ndarray
        complex mesh used to calculate `coeff`
    coeff : (..., N_in) complex ndarray
        coefficients for Pade, calculated from `pade.coefficients`
    valid_pades : list_like of bool
        Mask which continuations are correct, all Pades where `valid_pades`
        evaluates to false will be ignored for the average.
    kind : {KindGf, KindSelf}
        Defines the asymptotic of the continued function and the number of
        minimum and maximum input points used for Pade. For `KindGf` the
        function goes like :math:`1/z` for large `z`, for `KindSelf` the
        function behaves like a constant for large `z`.

    Returns
    -------
    average : function
        The continued function `f(z)` (`z`, ) -> Result. `f(z).x` contains the
        function values `f(z).err` the associated variance.

    Raises
    ------
    TypeError
        If `valid_pades` not of type `bool`
    RuntimeError
        If all there are none elements of `valid_pades` that evaluate to True.

    """
    valid_pades = np.array(valid_pades)
    if valid_pades.dtype != bool:
        raise TypeError(f"Invalid type of `valid_pades`: {valid_pades.dtype}\n"
                        "Expected `bool`.")
    if not valid_pades.any(axis=0).all():
        # for some axis no valid pade was found
        raise RuntimeError("No Pade fulfills is valid.\n"
                           f"No solution found for coefficient (shape: {coeff.shape[:-1]}) axes "
                           f"{np.argwhere(~valid_pades.any(axis=0))}")
    LOGGER.info("Number of valid Pade approximants: %s", np.count_nonzero(valid_pades, axis=0))

    def average(z) -> Result:
        """Calculate Pade continuation of function at points `z`.

        The continuation is calculated for different numbers of coefficients
        taken into account, where the number is in [n_min, n_max]. The function
        value es well as its variance is calculated. The variance should not be
        confused with an error estimate.

        Parameters
        ----------
        z : complex ndarray
            points at with the functions will be evaluated

        Returns
        -------
        pade.x : complex ndarray
            function evaluated at points `z`
        pade.err : complex ndarray
            variance associated with the function values `pade.x` at points `z`

        Raises
        ------
        RuntimeError
            If the calculated continuation contain any NaNs. This indicates
            invalid input in the coefficients and thus the original function.

        """
        z = np.asarray(z)

        pade_iter = kind.islice(calc_iterator(z, z_in, coeff=coeff))
        if valid_pades.ndim == 1:
            # validity determined for all dimensions -> drop invalid pades
            pades = np.array([pade for pade, valid in zip(pade_iter, valid_pades) if valid])
            if _contains_nan(pades):
                # check if fct_z already contained nans
                raise RuntimeError("Calculation of Pades failed, results contains NaNs")
        else:
            pades = np.array(list(pade_iter))
            if _contains_nan(pades):
                raise RuntimeError("Calculation of Pades failed, results contains NaNs")
            pades[~valid_pades] = np.nan + 1j*np.nan

        pade_avg = np.nanmean(pades, axis=0)
        std = _nan_std(pades.real) + 1j*_nan_std(pades.imag)

        return Result(x=pade_avg, err=std)
    return average

def Mod_Averager(z_in, coeff, mod_fct, *, valid_pades, kind: KindSelector, vectorized=True):
    r"""Create function for averaging Pade scheme using `mod_fct` before the average.

    This function behaves like `Averager` just that `mod_fct` is applied before
    taking the averages. This should be used, if not the analytic continuation
    but a mollification thereof is used.

    Parameters
    ----------
    z_in : (N_in,) complex ndarray
        complex mesh used to calculate `coeff`
    coeff : (..., N_in) complex ndarray
        coefficients for Pade, calculated from `pade.coefficients`
    mod_fct : callable
        Modification of the analytic continuation. The signature of the function
        should be `mod_fct` (z, pade_z, \*args, \*\*kwds), the tow first
        arguments are the point of evaluation `z` and the single Pade approximants.
    valid_pades : list_like of bool
        Mask which continuations are correct, all Pades where `valid_pades`
        evaluates to false will be ignored for the average.
    kind : {KindGf, KindSelf}
        Defines the asymptotic of the continued function and the number of
        minumum and maximum input points used for Pade. For `KindGf` the
        function goes like :math:`1/z` for large `z`, for `KindSelf` the
        function behaves like a constant for large `z`.
    vectorized : bool, optional
        If `vectorized`, all approximants are given to the function simultaniously
        where the first dimension corresponds to the approximants.
        If not `vectorized`, `mod_fct` will be called for every approximant
        seperately. (default: True)

    Returns
    -------
    mod_average : function
        The continued function `f(z)` (`z`, ) -> Result. `f(z).x` contains the
        function values `f(z).err` the associated variance.

    Raises
    ------
    TypeError
        If `valid_pades` not of type `bool`
    RuntimeError
        If all there are none elements of `valid_pades` that evaluate to True.

    """
    valid_pades = np.array(valid_pades)
    if valid_pades.dtype != bool:
        raise TypeError(f"Invalid type of `valid_pades`: {valid_pades.dtype}\n"
                        "Expected `bool`.")
    if not valid_pades.any(axis=0).all():
        # for some axis no valid pade was found
        raise RuntimeError("No Pade fulfills is valid.\n"
                           f"No solution found for coefficient (shape: {coeff.shape[:-1]}) axes "
                           f"{np.argwhere(~valid_pades.any(axis=0))}")
    LOGGER.info("Number of valid Pade approximants: %s", np.count_nonzero(valid_pades, axis=0))

    def mod_average(z, *args, **kwds) -> Result:
        """Calculate modified Pade continuation of function at points `z`.

        Calculate the averaged continuation of `mod_fct(f_z, *args, **kwds)`
        The continuation is calculated for different numbers of coefficients
        taken into account, where the number is in [n_min, n_max]. The function
        value es well as its variance is calculated. The variance should not be
        confused with an error estimate.

        Parameters
        ----------
        z : complex ndarray
            points at with the functions will be evaluated
        args, kwds :
            Passed to the `mod_fct` {mod_fct.__name__}.

        Returns
        -------
        pade.x : complex ndarray
            function evaluated at points `z`
        pade.err : complex ndarray
            variance associated with the function values `pade.x` at points `z`

        Raises
        ------
        RuntimeError
            If the calculated continuation contain any NaNs. This indicates
            invalid input in the coefficients and thus the original function.

        """
        z = np.asarray(z)

        pade_iter = kind.islice(calc_iterator(z, z_in, coeff=coeff))
        if valid_pades.ndim == 1:
            # validity determined for all dimensions -> drop invalid pades
            pades = np.array([pade for pade, valid in zip(pade_iter, valid_pades) if valid])
            if _contains_nan(pades):
                # check if fct_z already contained nans
                raise RuntimeError("Calculation of Pades failed, results contains NaNs")
        else:
            pades = np.array(list(pade_iter))
            if _contains_nan(pades):
                raise RuntimeError("Calculation of Pades failed, results contains NaNs")
            pades[~valid_pades] = np.nan + 1j*np.nan

        if vectorized:
            mod_pade = mod_fct(z, pades, *args, **kwds)
        else:
            mod_pade = np.array([mod_fct(z, pade_ii, *args, **kwds)
                                 for pade_ii in pades])
        pade_avg = np.nanmean(mod_pade, axis=0)
        # define helper pade_std np.nanstd( ,axis=0, ddof=1) if complex...
        if np.iscomplexobj(mod_pade):
            std = _nan_std(mod_pade.real) + 1j*_nan_std(mod_pade.imag)
        else:
            std = _nan_std(mod_pade)

        return Result(x=pade_avg, err=std)
    mod_average.__doc__ = mod_average.__doc__.format(mod_fct=mod_fct)
    return mod_average

def apply_filter(*filters, validity_iter):
    r"""Handle usage of filters for Pade.

    Parameters
    ----------
    filters : callable
        Functions to determine which continuations to keep.
    validity_iter : iterable of (..., N_z) complex np.ndarray
        The iterable of analytic continuations as generated by `calc_iterator`.

    Returns
    -------
    is_valid : (...) bool np.ndarray
        Array to index which continuations are good.

    """
    if len(filters) == 1:
        return filters[0](validity_iter)
    validity_iter = np.array(list(validity_iter))
    shape = validity_iter.shape
    validity_iter = np.moveaxis(validity_iter, 0, -2).reshape((-1, shape[0], shape[-1]))
    is_valid = np.ones(validity_iter.shape[0:2], dtype=bool)
    for i_valid, i_validity_iter in zip(is_valid, validity_iter):
        for filt in filters:
            is_valid_filt = filt(i_validity_iter[i_valid])
            i_valid[i_valid] = is_valid_filt
            if np.count_nonzero(i_valid) == 0:
                raise RuntimeError(
                    f"No Pade is valid due to filter {filt}.\n"
                    f"No solution found for coefficient (shape: {validity_iter.shape[1:-1]}) axes "
                    f"{np.argwhere(~is_valid.any(axis=0))}"
                )
    return np.moveaxis(is_valid, 0, -1).reshape(shape[:-1])

def averaged(z_out, z_in, *, valid_z=None, fct_z=None, coeff=None,
             filter_valid=None, kind: KindSelector):
    """Return the averaged Pade continuation with its variance.

    The output is checked to have an imaginary part smaller than `threshold`,
    as retarded Green's functions and self-energies have a negative imaginary
    part.
    This is a helper to conveniently get the continuation, it comes however with
    overhead.

    Parameters
    ----------
    z_out : (N_out,) complex ndarray
        points at with the functions will be evaluated
    z_in : (N_in,) complex ndarray
        complex mesh used to calculate `coeff`
    valid_z : (N_out,) complex ndarray, optional
        The output range according to which the Pade approximation is validated
        (compared to the `threshold`).
    fct_z : (N_z, ) complex ndarray, optional
        Function at points `z` from which the coefficients will be calculated.
        Can be omitted if `coeff` is directly given.
    coeff : (N_in,) complex ndarray, optional
        Coefficients for Pade, calculated from `pade.coefficients`. Can be given
        instead of `fct_z`.
    filter_valid : callable or iterable of callable
        Function determining which approximants to keep. The signature should
        be filter_valid(iterable) -> bool ndarray.
        Currently there are the functions {`FilterNegImag`, `FilterNegImagNum`,
        `FilterHighVariance`} implemented to generate filter functions.
        Look into the implemented for details to create new filters.
    kind : {KindGf, KindSelf}
        Defines the asymptotic of the continued function and the number of
        minimum and maximum input points used for Pade. For `KindGf` the
        function goes like :math:`1/z` for large `z`, for `KindSelf` the
        function behaves like a constant for large `z`.

    Returns
    -------
    averaged.x : (N_in, N_out) complex ndarray
        function evaluated at points `z`
    averaged.err : (N_in, N_out) complex ndarray
        variance associated with the function values `pade.x` at points `z`

    """
    assert fct_z is None or coeff is None
    z_in = z_in[:kind.stop]
    coeff = (coefficients(z_in, fct_z=fct_z[..., :kind.stop]) if coeff is None
             else coeff[..., :kind.stop])
    if valid_z is None:
        valid_z = z_out
    valid_z = valid_z.reshape(-1)
    if filter_valid is not None:
        validity_iter = kind.islice(calc_iterator(valid_z, z_in, coeff=coeff))
        try:
            filters = tuple(filter_valid)
        except TypeError:  # only one filter given
            filters = (filter_valid,)
        is_valid = apply_filter(*filters, validity_iter=validity_iter)
    else:
        is_valid = np.ones((len(kind), *coeff.shape[:-1]), dtype=bool)

    assert is_valid.shape[1:] == coeff.shape[:-1]

    _average = Averager(z_in, coeff=coeff, valid_pades=is_valid, kind=kind)
    return _average(z_out)

def avg_no_neg_imag(z_out, z_in, *, valid_z=None, fct_z=None, coeff=None,
                    threshold=1e-8, kind: KindSelector):
    """Average Pade filtering approximants with non-negative imaginary part.

    This function wraps `averaged`, see `averaged` for the parameters.

    Other Parameters
    ----------------
    threshold : float, optional
        The numerical threshold, how large of an positive imaginary part is
        tolerated (default: 1e-8). `np.infty` can be given to accept all.

    Returns
    -------
    averaged.x : (N_in, N_out) complex ndarray
        function evaluated at points `z`
    averaged.err : (N_in, N_out) complex ndarray
        variance associated with the function values `pade.x` at points `z`

    """
    filter_neg_imag = FilterNegImag(threshold)
    return averaged(z_out=z_out, z_in=z_in, valid_z=valid_z, fct_z=fct_z,
                    coeff=coeff, filter_valid=filter_neg_imag, kind=kind)

# def SelectiveAverage(object):
#     """Do not accept Matsubara frequencies, which make the result unphysical."""

ext: py

max_stars_repo_name: aostlin/gftools

max_issues_repo_name: aostlin/gftools

max_forks_repo_name: aostlin/gftools

reason: import numpy


text:
import tarfile
import pickle
import zipfile
from sklearn.datasets import load_files
from sklearn.datasets.twenty_newsgroups import strip_newsgroup_header, strip_newsgroup_footer, strip_newsgroup_quoting
from util.file import *
from scipy.io import arff
from scipy.sparse import csr_matrix
import numpy as np
import sklearn

def _proc_review(doc):
    parts = doc.split(' ')
    label = parts[-1].replace('#label#:', '').strip()
    assert label in ['positive','negative'], 'error parsing label {}'.format(label)
    label = 1 if label == 'positive' else 0
    repeat_word = lambda word, num: ' '.join([word] * int(num))
    text = ' '.join([repeat_word(*x.split(':')) for x in parts[:-1]])
    return text, label

def fetch_MDS(DOWNLOAD_URL = 'https://www.cs.jhu.edu/~mdredze/datasets/sentiment/processed_acl.tar.gz',
              dataset_home = '../datasets/MDS'):
    """
    Fetchs the processed version of the  Multi-Domain Sentiment Dataset (version 2.0) for cross-domain adaptation defined in:
    John Blitzer, Mark Dredze, Fernando Pereira.
    Biographies, Bollywood, Boom-boxes and Blenders: Domain Adaptation for Sentiment Classification.
    Association of Computational Linguistics (ACL), 2007.
    """
    dataset_path = join(dataset_home, 'processed_acl.tar.gz')
    create_if_not_exist(dataset_home)

    if not exists(dataset_path):
        print("downloading multidomain dataset (once and for all) into %s" % dataset_path)
        download_file(DOWNLOAD_URL, dataset_path)
        print("untarring dataset...")
        tarfile.open(dataset_path, 'r:gz').extractall(dataset_home)

    dataset_path = dataset_path.replace('.tar.gz','')

    documents = dict()
    for domain in list_dirs(dataset_path):
        documents[domain] = dict()
        for file in ['positive.review', 'negative.review', 'unlabeled.review']:
            documents[domain][file] = []
            for doc in open(join(dataset_path, domain, file), 'rt'):
                text, label = _proc_review(doc)
                documents[domain][file].append((text,label))
            print('{} documents read for domain {} in file {}'.format(len(documents[domain][file]), domain, file))

    return documents

def fetch_Webis_cls_10(DOWNLOAD_URL = 'https://zenodo.org/record/3251672/files/cls-acl10-processed.tar.gz?download=1',
                       #DOWNLOAD_URL = 'http://www.uni-weimar.de/medien/webis/corpora/corpus-webis-cls-10/cls-acl10-processed.tar.gz',
                       dataset_home = '../datasets/Webis-CLS-10',
                       max_documents=50000,
                       languages=['de','en','fr','jp'],
                       domains=None,
                       skip_translations=False,
                       dopickle=False):
    """
    Fetchs the processed version of the Webis-CLS-10 dataset for cross-lingual adaptation defined in:
    Prettenhofer, Peter, and Benno Stein.
    "Cross-language text classification using structural correspondence learning."
    Proceedings of the 48th annual meeting of the association for computational linguistics.
    Association for Computational Linguistics, 2010.
    """
    picklepath = join(dataset_home, 'webiscls10_processed.pkl')
    if dopickle and exists(picklepath):
        print('...loading pickle from {}'.format(picklepath))
        return pickle.load(open(picklepath, 'rb'))

    dataset_path = join(dataset_home, 'cls-acl10-processed.tar.gz')
    create_if_not_exist(dataset_home)

    if not exists(dataset_path):
        print("downloading Webis-CLS1-0 dataset (once and for all) into %s" % dataset_path)
        download_file(DOWNLOAD_URL, dataset_path)
        print("untarring dataset...")
        tarfile.open(dataset_path, 'r:gz').extractall(dataset_home)

    dataset_path = dataset_path.replace('.tar.gz','')

    documents = dict()
    for language in languages:
        documents[language] = dict()
        domain_list = domains if domains is not None else list_dirs(join(dataset_path,language))
        for domain in domain_list:
            documents[language][domain] = dict()
            for file in ['train.processed', 'test.processed', 'unlabeled.processed']:
                documents[language][domain][file] = []
                for doc in open(join(dataset_path, language, domain, file), 'rt'):
                    text, label = _proc_review(doc)
                    documents[language][domain][file].append((text,label))
                    if max_documents is not None and len(documents[language][domain][file]) >= max_documents:
                        break
                print('{} documents read for language {}, domain {}, in file {}'.format(
                    len(documents[language][domain][file]), language, domain, file))

    translations = dict()
    if not skip_translations:
        for language in ['de','fr','jp']:
            translations[language] = dict()
            for domain in list_dirs(join(dataset_path,language)):
                translations[language][domain] = dict()
                for file in ['test.processed']:
                    translations[language][domain][file] = []
                    for doc in open(join(dataset_path, language, domain, 'trans', 'en', domain, file), 'rt'):
                        text, label = _proc_review(doc)
                        translations[language][domain][file].append((text,label))
                    print('{} translations read for language {}, domain {}, in file {}'.format(
                        len(translations[language][domain][file]), language, domain, file))

    dictionaries = dict()
    split_t = lambda x: x.strip().replace(' ','').split('\t')
    for d in list_files(join(dataset_path,'dict')):
        dictionaries[d] = dict(map(split_t, open(join(dataset_path,'dict',d))))

    if dopickle:
        print('...pickling the dataset into {} to speed-up next calls'.format(picklepath))
        pickle.dump((documents, translations, dictionaries), open(picklepath, 'wb'), pickle.HIGHEST_PROTOCOL)

    return documents, translations, dictionaries

def fetch_reuters21579(DOWNLOAD_URL = 'http://www.cse.ust.hk/TL/dataset/Reuters.zip',
                       dataset_home='../datasets/Reuters21578', dopickle=False):
    """
    Fetchs a version of Reuters21578 for cross-domain adaptation, as defined in:
    Dai, W., Xue, G. R., Yang, Q., & Yu, Y. (2007, August).
    Co-clustering based classification for out-of-domain documents.
    In Proceedings of the 13th ACM SIGKDD international conference on Knowledge discovery and data mining (pp. 210-219). ACM.
    """
    picklepath = join(dataset_home, 'reuters.pkl')
    if dopickle and exists(picklepath):
        print('...loading pickle from {}'.format(picklepath))
        return pickle.load(open(picklepath, 'rb'))

    dataset_path = join(dataset_home, 'Reuters.zip')

    if not exists(dataset_path):
        create_if_not_exist(dataset_home)
        print("downloading Reuters dataset (once and for all) into %s" % dataset_path)
        download_file(DOWNLOAD_URL, dataset_path)
        print("unzipping dataset...")
        zip_ref = zipfile.ZipFile(dataset_path, 'r')
        zip_ref.extractall(dataset_home)
        zip_ref.close()

    corpus = {}
    for file in list_files(dataset_home):
        if not file.endswith('.arff'): continue
        task, domain, _ = file.split('.')
        data, meta = arff.loadarff(join(dataset_home, file))
        data = np.array([[int(x) for x in row] for row in data])
        X = csr_matrix(data[:, :-1])
        y = data[:, -1]
        if task not in corpus: corpus[task] = {}
        corpus[task][domain] = (X, y)
        print('loaded {} with shape {} and prevalence {}'.format(file, X.shape, y.mean()))

    if dopickle:
        print('...pickling the dataset into {} to speed-up next calls'.format(picklepath))
        pickle.dump(corpus, open(picklepath, 'wb'), pickle.HIGHEST_PROTOCOL)

    return corpus

def _index_by_label(news, min_words=1):
    corpus = {label:[] for label in news.target_names}
    removed=0
    for i,doc in enumerate(news.data):
        if len(doc.strip().split())>=min_words:
            label_id = news.target[i]
            label = news.target_names[label_id]
            corpus[label].append(doc)
        else: removed+=1
    print('removed {} documents'.format(removed))
    return corpus

def fetch_20newsgroups(dataset_home='../datasets/20news', dopickle=False):
    """
    Fetchs a version of 20Newsgroups for cross-domain adaptation, as defined in:
    Dai, W., Xue, G. R., Yang, Q., & Yu, Y. (2007, August).
    Co-clustering based classification for out-of-domain documents.
    In Proceedings of the 13th ACM SIGKDD international conference on Knowledge discovery and data mining (pp. 210-219). ACM.
    """
    picklepath = join(dataset_home, '20news.pkl')
    if dopickle and exists(picklepath):
        print('...loading pickle from {}'.format(picklepath))
        return pickle.load(open(picklepath, 'rb'))

    news = sklearn.datasets.fetch_20newsgroups(data_home=dataset_home, subset='all')
    news = _index_by_label(news)

    if dopickle:
        print('...pickling the dataset into {} to speed-up next calls'.format(picklepath))
        pickle.dump(news, open(picklepath, 'wb'), pickle.HIGHEST_PROTOCOL)

    return news

def fetch_sraa(DOWNLOAD_URL = 'http://people.cs.umass.edu/~mccallum/data/sraa.tar.gz',
                       dataset_home='../datasets/SRAA', dopickle=False):
    """
    Fetchs a version of Reuters21578 for cross-domain adaptation, as defined in:
    Dai, W., Xue, G. R., Yang, Q., & Yu, Y. (2007, August).
    Co-clustering based classification for out-of-domain documents.
    In Proceedings of the 13th ACM SIGKDD international conference on Knowledge discovery and data mining (pp. 210-219). ACM.
    """
    picklepath = join(dataset_home, 'sraa.pkl')
    if dopickle and exists(picklepath):
        print('...loading pickle from {}'.format(picklepath))
        return pickle.load(open(picklepath, 'rb'))

    dataset_path = join(dataset_home, 'sraa.tar.gz')

    if not exists(dataset_path):
        create_if_not_exist(dataset_home)
        print("downloading SRAA dataset (once and for all) into %s" % dataset_path)
        download_file(DOWNLOAD_URL, dataset_path)
        print("untarring dataset...")
        tarfile.open(dataset_path, 'r:gz').extractall(dataset_home)

    sraa = load_files(join(dataset_home,'sraa'), encoding='latin1')
    remove = ('headers', 'footers')#, 'quotes')
    if 'headers' in remove:
        sraa.data = [strip_newsgroup_header(text) for text in sraa.data]
    if 'footers' in remove:
        sraa.data = [strip_newsgroup_footer(text) for text in sraa.data]
    if 'quotes' in remove:
        sraa.data = [strip_newsgroup_quoting(text) for text in sraa.data]

    sraa = _index_by_label(sraa, min_words=10)

    if dopickle:
        print('...pickling the dataset into {} to speed-up next calls'.format(picklepath))
        pickle.dump(sraa, open(picklepath, 'wb'), pickle.HIGHEST_PROTOCOL)

    return sraa

max_stars_repo_name: AlexMoreo/pydci

max_issues_repo_name: AlexMoreo/pydci

max_forks_repo_name: AlexMoreo/pydci

reason: import numpy,from scipy


text:
import GPy
import numpy as np

# Before running this test, ensure andro.arff, downloaded from
# http://mulan.sourceforge.net/datasets-mtr.html is in this folder.

# Run this test from the command line using
# `python3 slfm_experiment.py`

if __name__ == "__main__":
    with open('andro.arff') as f:
        for row in range(0,40):
            f.readline()

        i = -1
        locs = []
        dat = np.zeros((54,6))
        for i in range(0,48):
            s = f.readline()
            vals = list(map(float,s.split(',')))
            dat[i,0] = vals[0]
            dat[i,1] = vals[1]
            dat[i,2] = vals[2]
            dat[i,3] = vals[3]
            dat[i,4] = vals[4]
            dat[i,5] = vals[5]

        s = f.readline()
        vals = list(map(float,s.split(',')))
        for i in range(0,6):
            for v in range(0,6):
                dat[i+48,v] = vals[v + 6*i]

    X = np.linspace(0,53,54)

    train = np.copy(dat)
    test = np.copy(dat)

    K1 = GPy.kern.RBF(1)
    K2 = GPy.kern.RBF(1)
    K3 = GPy.kern.RBF(1)
    K4 = GPy.kern.RBF(1)
    K5 = GPy.kern.RBF(1)
    K6 = GPy.kern.RBF(1)
    # Best performance achieved with 6 kernels
    LCM = GPy.util.multioutput.LCM(input_dim=1,num_outputs=6,kernels_list=[K1,K2,K3,K4,K5,K6])

    XTemp = X[:,None]
    YTemp = train[:,0][:,None]
    XPH = X[:,None]
    YPH = train[:,1][:,None]
    XCond = X[:,None]
    YCond = train[:,2][:,None]
    XSal = np.concatenate((X[0:20],X[30:]),axis=0)[:,None]
    YSal = np.concatenate((train[0:20,3],train[30:,3]),axis=0)[:,None]
    XOxy = np.concatenate((X[0:30],X[40:]),axis=0)[:,None]
    YOxy = np.concatenate((train[0:30:,4],train[40:,4]),axis=0)[:,None]
    XTurb = X[:,None]
    YTurb = train[:,5][:,None]

    # Normalize
    meanTemp = np.mean(YTemp)
    meanPH   = np.mean(YPH)
    meanCond = np.mean(YCond)
    meanSal  = np.mean(YSal)
    meanOxy  = np.mean(YOxy)
    meanTurb = np.mean(YTurb)
    stdTemp  = np.std(YTemp)
    stdPH    = np.std(YPH)
    stdCond  = np.std(YCond)
    stdSal   = np.std(YSal)
    stdOxy   = np.std(YOxy)
    stdTurb  = np.std(YTurb)

    YTemp = (YTemp - meanTemp) /stdTemp
    YPH   = (YPH   - meanPH)   /stdPH
    YCond = (YCond - meanCond) /stdCond
    YSal  = (YSal  - meanSal)  /stdSal
    YOxy  = (YOxy  - meanOxy)  /stdOxy
    YTurb = (YTurb - meanTurb) /stdTurb

    model = GPy.models.GPCoregionalizedRegression(X_list=[XTemp,XPH,XCond,XSal,XOxy,XTurb],Y_list=[YTemp,YPH,YCond,YSal,YOxy,YTurb],kernel=LCM)

    model.optimize('lbfgs',messages=True)

    pred_XSal = np.arange(20,30)[:,None]
    pred_XSal = np.hstack([pred_XSal,3*np.ones_like(pred_XSal)])
    YSal_meta = {'output_index':pred_XSal[:,1].astype(int)}

    pred_YSal = model.predict(pred_XSal,Y_metadata=YSal_meta)

    pred_XOxy = np.arange(30,40)[:,None]
    pred_XOxy = np.hstack([pred_XOxy,4*np.ones_like(pred_XOxy)])
    YOxy_meta = {'output_index':pred_XOxy[:,1].astype(int)}

    pred_YOxy = model.predict(pred_XOxy,Y_metadata=YOxy_meta)

    cseSal = 0
    cse_normSal = 0
    truthSal = dat[20:30,3]
    mean_truthSal = np.mean(truthSal)
    cnllSal = 0
    for i in range(10):
        pred_mean = meanSal + pred_YSal[0][i,0] * stdSal
        pred_var = pred_YSal[1][i,0] * stdSal**2
        cseSal += (pred_mean - truthSal[i])**2
        cse_normSal += (mean_truthSal - truthSal[i])**2
        cnllSal += 0.5 * ((pred_mean - truthSal[i])**2 / pred_var + np.log(2*np.pi*pred_var))

    cseOxy = 0
    cse_normOxy = 0
    truthOxy = dat[30:40,4]
    mean_truthOxy = np.mean(truthOxy)
    cnllOxy = 0
    for i in range(10):
        pred_mean = meanOxy + pred_YOxy[0][i,0] * stdOxy
        pred_var = pred_YOxy[1][i,0] * stdOxy**2
        cseOxy += (pred_mean - truthOxy[i])**2
        cse_normOxy += (mean_truthOxy - truthOxy[i])**2
        cnllOxy += 0.5 * ((pred_mean - truthOxy[i])**2 / pred_var + np.log(2*np.pi*pred_var))

    print("smseSal = {}".format(cseSal/cse_normSal))
    print("smseOxy = {}".format(cseOxy/cse_normOxy))
    print("average = {}".format((cseSal/cse_normSal + cseOxy/cse_normOxy)/2))
    print("mnllSal = {}".format(cnllSal/10))
    print("mnllOxy = {}".format(cnllOxy/10))
    print("average = {}".format((cnllSal/10 + cnllOxy/10)/2))

# Results:

# Trial 1:
# smseSal = 0.10239271002717211
# smseOxy = 0.04883199618682176
# average = 0.07561235310699693
# mnllSal = 3.1685164328490805
# mnllOxy = 2.127499480800946
# average = 2.648007956825013

# Trial 2:
# smseSal = 0.09027465247006487
# smseOxy = 0.2131194545074291
# average = 0.15169705348874699
# mnllSal = 2.226509518563331
# mnllOxy = 3.8774681000140054
# average = 3.0519888092886682

# Trial 3:
# smseSal = 0.07340337174621854
# smseOxy = 0.08262265987021256
# average = 0.07801301580821554
# mnllSal = 1.688124566168296
# mnllOxy = 2.578512673378579
# average = 2.1333186197734375

# Trial 4:
# smseSal = 0.0958180189204561
# smseOxy = 0.16720579983831446
# average = 0.13151190937938528
# mnllSal = 2.466207976952092
# mnllOxy = 3.530483346035022
# average = 2.998345661493557

# Trial 5:
# smseSal = 0.07097851158787918
# smseOxy = 0.0999299683436154
# average = 0.08545423996574729
# mnllSal = 1.3707329691067336
# mnllOxy = 3.0784285430652427
# average = 2.224580756085988

# Trial 6:
# smseSal = 0.07981208236541011
# smseOxy = 0.016262980694810875
# average = 0.04803753153011049
# mnllSal = 1.525926008886414
# mnllOxy = 1.5212872080292368
# average = 1.5236066084578255

# Trial 7:
# smseSal = 0.10990880933000702
# smseOxy = 0.16516706269958442
# average = 0.13753793601479572
# mnllSal = 3.34300651774684
# mnllOxy = 3.7230054549764127
# average = 3.5330059863616263

# Trial 8:
# smseSal = 0.09148169287083106
# smseOxy = 0.14412654264246946
# average = 0.11780411775665026
# mnllSal = 2.1971135055129114
# mnllOxy = 3.219618636827792
# average = 2.7083660711703517

# Trial 9:
# smseSal = 0.13458751481890768
# smseOxy = 0.18014223057216253
# average = 0.15736487269553512
# mnllSal = 3.852268557533268
# mnllOxy = 4.079252498486456
# average = 3.965760528009862

# Trial 10:
# smseSal = 0.08549209764236905
# smseOxy = 0.18066514621793375
# average = 0.1330786219301514
# mnllSal = 1.9544891570549534
# mnllOxy = 3.5644259909227856
# average = 2.7594575739888696

# Averages:
# smseSal = 0.09341494617793158 +/- 0.005673506840221591
# smseOxy = 0.12980738415733545 +/- 0.019360630064543354
# average = 0.11161116516763352 +/- 0.011143726200627247
# mnllSal = 2.379289521037392 +/- 0.2487762328718216
# mnllOxy = 3.1299981932536474 +/- 0.24634022636928674
# average = 2.7546438571455196 +/- 0.21015366888516424

max_stars_repo_name: aytonb/dag-gp

max_issues_repo_name: aytonb/dag-gp

max_forks_repo_name: aytonb/dag-gp


text:
import os
import tqdm
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

PREFIX = 'Hammer_'

PCD_ROOT = './SyntheticHammer/data/hammers_pcd'
PCD_SUFFIX = '.pcd'
KP_ROOT = './SyntheticHammer/data/hammers_kp'
KP_SUFFIX = '_kp.txt'

for i in range(2000):
    fig = plt.figure()
    ax = Axes3D(fig)
    pcd_file = os.path.join(PCD_ROOT, PREFIX + str(i) + PCD_SUFFIX)
    kp_file = os.path.join(KP_ROOT, PREFIX + str(i) + KP_SUFFIX)

    # Scatters
    x = []
    y = []
    z = []

    x_kp = []
    y_kp = []
    z_kp = []

    # Get point cloud
    pcd_contents = []
    with open(pcd_file) as f:
        pcd_contents = f.readlines()
    f.close()
    pcd_contents = pcd_contents[11:]
    for pcd_content in pcd_contents:
        point = pcd_content.split(' ')
        x.append(float(point[0]))
        y.append(float(point[1]))
        z.append(float(point[2]))

    # Get keypoints
    kp_contents = []
    with open(kp_file) as f:
        kp_contents = f.readlines()
    f.close()
    kp_contents = kp_contents[1:]
    for j, (kp_content) in enumerate(kp_contents):
        point = kp_content.split(' ')
        x_kp.append(float(point[0]))
        y_kp.append(float(point[1]))
        z_kp.append(float(point[2]))

    x = np.array(x)
    y = np.array(y)
    z = np.array(z)

    x_kp = np.array(x_kp)
    y_kp = np.array(y_kp)
    z_kp = np.array(z_kp)
    # print(x_kp.shape)

    ax.scatter(x, y, z, c='blue', marker='.', s=1, label='')
    ax.scatter(x_kp[0], y_kp[0], z_kp[0], c='red', marker='*', s=50, label='grasp point')
    ax.scatter(x_kp[1], y_kp[1], z_kp[1], c='green', marker='*', s=50, label='function point')
    ax.scatter(x_kp[2], y_kp[2], z_kp[2], c='yellow', marker='*', s=50, label='effect point')
    plt.legend()
    plt.show()
    plt.close()

max_stars_repo_name: Simon-Fuhaoyuan/SyntheticHammer

max_issues_repo_name: Simon-Fuhaoyuan/SyntheticHammer

max_forks_repo_name: Simon-Fuhaoyuan/SyntheticHammer


text:
import numpy as np

import keras
from keras.models import Model
from keras.layers import Input, Dense, Conv2D, Lambda, LeakyReLU, TimeDistributed, Flatten, BatchNormalization, concatenate
import keras.backend as K
import keras.callbacks as kc

def mk_dense(variant):
    assert variant in ['trim', 'full_separate', 'full_whole']
    def mk_model(**kwargs):
        wrap = lambda x: LeakyReLU(0.01)(x)
        input_lower = Input((None, 129), name="input_lf")
        layer = Lambda(lambda x: x / kwargs["MAX"])(input_lower)
        lyr = Lambda(K.expand_dims)(layer)
        windowing = Conv2D(15, (15, 1), padding='same', use_bias=False)
        lyr = windowing(lyr)
        windowing.set_weights([np.eye(15).reshape(15, 1, 1, 15)])
        windowing.trainable = False
        base = lyr = TimeDistributed(Flatten())(lyr)
        if variant == "trim":
            lyr = wrap(Dense(2048)(lyr))
            lyr = wrap(Dense(2048)(lyr))
            lf_and_hf = Dense(129)(lyr)
        if variant == "full_separate":
            lyr = wrap(Dense(1024)(base))
            lyr = wrap(Dense(1024)(lyr))
            out1 = Dense(129)(lyr)
            lyr = wrap(Dense(1024)(base))
            lyr = wrap(Dense(1024)(lyr))
            out2 = Dense(128)(lyr)
            lf_and_hf = concatenate([out1, out2])
        if variant == "full_whole":
            lyr = wrap(Dense(2048)(lyr))
            lyr = wrap(Dense(2048)(lyr))
            lf_and_hf = Dense(257)(lyr)
        lf_and_hf = Lambda(lambda x: x * kwargs["MAX"])(lf_and_hf)
        mdl = Model(input_lower, lf_and_hf)
        mdl.summary()
        mdl.compile('adam', 'mse')
        return mdl
    return mk_model

def mk_conv_dae(variant, complexity):
    assert variant in ['trim', 'full_separate', 'full_whole']
    def mk_model(**kwargs):
        input_lower = Input((None, 129), name="input_lf")
        layer = Lambda(lambda x: x / kwargs["MAX"])(input_lower)
        base = Lambda(K.expand_dims)(layer)
        if variant == "trim":
            layer = base
            for i in range(complexity):
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 20, kernel_size=(1,5), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 24, kernel_size=(9,1), padding='same', activation='linear')(layer))
            layer = LeakyReLU(0.01)(Conv2D(1, kernel_size=(1,1), padding='same', activation='linear')(layer))
            for i in range(complexity - 1, -1, -1):
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 24, kernel_size=(9,1), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 20, kernel_size=(1,5), padding='same', activation='linear')(layer))
            lf_and_hf = Dense(129)(TimeDistributed(Flatten())(layer))
        if variant == "full_separate":
            outs = []
            for sz in [129, 128]:
                layer = base
                for i in range(complexity):
                    layer = LeakyReLU(0.01)(Conv2D(i * 4 + 10, kernel_size=(1,5), padding='same', activation='linear')(layer))
                    layer = LeakyReLU(0.01)(Conv2D(i * 4 + 12, kernel_size=(9,1), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(1, kernel_size=(1,1), padding='same', activation='linear')(layer))
                for i in range(complexity - 1, -1, -1):
                    layer = LeakyReLU(0.01)(Conv2D(i * 4 + 12, kernel_size=(9,1), padding='same', activation='linear')(layer))
                    layer = LeakyReLU(0.01)(Conv2D(i * 4 + 10, kernel_size=(1,5), padding='same', activation='linear')(layer))
                outs.append(Dense(sz)(TimeDistributed(Flatten())(layer)))
            lf_and_hf = concatenate(outs)
        if variant == "full_whole":
            layer = base
            for i in range(complexity):
                layer = LeakyReLU(0.01)(Conv2D(i * 4 + 10, kernel_size=(1,5), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(i * 4 + 12, kernel_size=(9,1), padding='same', activation='linear')(layer))
            layer = LeakyReLU(0.01)(Conv2D(1, kernel_size=(1,1), padding='same', activation='linear')(layer))
            for i in range(complexity - 1, -1, -1):
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 24, kernel_size=(9,1), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 20, kernel_size=(1,5), padding='same', activation='linear')(layer))
            lf_and_hf = Dense(257)(TimeDistributed(Flatten())(layer))
        lf_and_hf = Lambda(lambda x: x * kwargs["MAX"])(lf_and_hf)
        mdl = Model(input_lower, lf_and_hf)
        mdl.summary()
        mdl.compile('adam', 'mse')
        return mdl
    return mk_model

def mk_conv_dense(variant, complexity):
    assert variant in ['trim', 'full_separate', 'full_whole']
    def mk_model(**kwargs):
        input_lower = Input((None, 129), name="input_lf")
        layer = Lambda(lambda x: x / kwargs["MAX"])(input_lower)
        base = Lambda(K.expand_dims)(layer)
        if variant == "trim":
            layer = base
            for i in range(complexity):
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 20, kernel_size=(1,5), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(i * 8 + 24, kernel_size=(9,1), padding='same', activation='linear')(layer))
            layer = TimeDistributed(Flatten())(layer)
            layer = LeakyReLU(0.01)(Dense(1024)(layer))
            hidden = layer = LeakyReLU(0.01, name='hidden')(Dense(96)(layer))
            for i in range(complexity - 1, -1, -1):
                layer = LeakyReLU(0.01)(Dense(512)(layer))
            lf_and_hf = Dense(129)(layer)
        if variant == "full_separate":
            outs = []
            for sz in [129, 128]:
                layer = base
                for i in range(complexity):
                    layer = LeakyReLU(0.01)(Conv2D(i * 4 + 10, kernel_size=(1,5), padding='same', activation='linear')(layer))
                    layer = LeakyReLU(0.01)(Conv2D(i * 4 + 12, kernel_size=(9,1), padding='same', activation='linear')(layer))
                layer = TimeDistributed(Flatten())(layer)
                layer = LeakyReLU(0.01)(Dense(512)(layer))
                layer = LeakyReLU(0.01)(Dense(96)(layer))
                for i in range(complexity - 1, -1, -1):
                    layer = LeakyReLU(0.01)(Dense(512)(layer))
                outs.append(Dense(sz)(layer))
            lf_and_hf = concatenate(outs)
        if variant == "full_whole":
            layer = base
            for i in range(complexity):
                layer = LeakyReLU(0.01)(Conv2D(i * 4 + 10, kernel_size=(1,5), padding='same', activation='linear')(layer))
                layer = LeakyReLU(0.01)(Conv2D(i * 4 + 12, kernel_size=(9,1), padding='same', activation='linear')(layer))
            layer = TimeDistributed(Flatten())(layer)
            layer = LeakyReLU(0.01)(Dense(1024)(layer))
            hidden = layer = LeakyReLU(0.01, name='hidden')(Dense(96)(layer))
            for i in range(complexity - 1, -1, -1):
                layer = LeakyReLU(0.01)(Dense(512)(layer))
            lf_and_hf = Dense(257)(layer)
        lf_and_hf = Lambda(lambda x: x * kwargs["MAX"])(lf_and_hf)
        mdl = Model(input_lower, lf_and_hf)
        mdl.summary()
        mdl.compile('adam', 'mse')
        return mdl
    return mk_model

def get_schema(model_type, variant):
    if model_type == "dense":
        return mk_dense(variant)
    if model_type.startswith("conv_dense"):
        return mk_conv_dense(variant, 2 if model_type == "conv_dense" else 4)
    if model_type.startswith("dae"):
        return mk_conv_dae(variant, 2 if model_type == "dae" else 4)

max_stars_repo_name: PICTEC/SEAR

max_issues_repo_name: PICTEC/SEAR

max_forks_repo_name: PICTEC/SEAR


text:
######################
#authors: Jens Hoeijmakers and Julia Seidel
#Description: Calculates the stellar spectrum
#
#
#####################
#Call like: python3 main.py 586.0 592.0 110000.0 90.0 90.0 0.0 0.0 500
#import statements
import sys
import numpy as np
import argparse
import lib.test as test
import lib.vgrid as vgrid
import lib.plotting as pl
import lib.operations as ops
import lib.stellar_spectrum as spectrum
import lib.integrate as integrate
import pdb
import time
import matplotlib.pyplot as plt
import math
from matplotlib.patches import Circle
import lib.planet_pos as ppos
import copy
#main body of code

class StarRotator(object):
    def __init__(self,wave_start,wave_end,grid_size,star_path='demo_star.txt',planet_path='demo_planet.txt',obs_path='demo_observations.txt'):
        """
            Welcome to StarRotator.
            ***********************
            The StarRotator object contains the main functionality of StarRotator.
            Upon initialization, the exoplanet system input is read from file, and the
            model is computed. The parameters needed to initialise are listed below:

            Parameters
            ----------
            wave_start : float
                Start of modelled wavelength range in nm in vacuum.
            wave_end : float
                Ending Wavelength range in nm in vacuum.
            grid_size: int
                Half-width of number of grid cells. Set to values greater than 200
                to limit numerical errors, or less if you are trying things out and just
                want speed.
            star_path : str
                Path to parameter file defining the star. This file should contain the following
                values on separate lines, and the default values are as follows:
                    50000.0     v_eq
                    90.0        stellar i
                    0.0         Differential rotation parameter (alpha)
                    5000.0      Stellar effective temperature (K)
                    0.0         Fe/H
                    4.5         logg
                    0.93        Limb-darkening coefficient u1
                    -0.23       Limb-darkening coefficient u2
                    0           Number of mu angles to consider. For values higher than zero, StarRotator switches to SPECTRUM rather than PHOENIX.
            planet_path: str
                Path to the parameter file defining the planet and its orbit. This file should contain
                the following values on separate lines, and the default values are as follows:
                    3.153       a/Rs
                    0.0         e
                    0.0         omega
                    86.79       Orbital inclination
                    85.0        Projected obliquity
                    0.08228     Rp/Rs
                    1.4811235   Orbital period
                    57095.68572 Transit center time - 2400000.
                    phases      mode, providing the interpretation of the timestamps of the observations:

            obs_path: str
                Path to the parameter file defining the timestamps of the observations.
                If mode (see previous) is set to 'phases', this file is assumed to contain
                a list of orbital phases. Otherwise it should be set to 'times', in which
                case the file should contain the times of the observations, in
                JD - 2400000. These times are assumed to mean the *start* times of
                each observation in the sequence. In addition, a second column should
                be provided giving the exposure time in seconds. This allows StarRotator
                to shift to the mid-times of the exposures. Note that for longer exposure times,
                the output of StarRotator will be less accurate because the signal of multiple
                exposures are convolved together. This effect is greater for aligned orbits.

                By default, the following phases are provided:
                    -0.06
                    -0.055
                    -0.05
                    -0.045
                    -0.04
                    -0.035
                    -0.03
                    -0.025
                    -0.02
                    -0.015
                    -0.01
                    -0.005
                    0.0
                    0.005
                    0.01
                    0.015
                    0.02
                    0.025
                    0.03
                    0.035
                    0.04
                    0.045
                    0.05
                    0.055
                    0.06

            Class methods
            -------------
            After initializing the class like a=StarRotator(588.0,592.0,200.0),
            the following methods are available to manipulate the simulation after
            it has been calculated the first time, or to plot the simulation output.
            star.read_system() Reads in (new) exoplanet system files (see above).
            star.compute_spectrum() recomputes the simulation.
            star.plot_residuals() Produces a 2D plot of the residuals.
            star.animate() Produces an animation of the transiting system. One frame
            corresponding to each phase provided.

            Output attributes
            -----------------
            After initializing the class like star = StarRotator(588.0,592.0,200.0),
            the primary simulation output can be accessed as follows:
            star.wl (wavelength array)
            star.spectra (matrix of spectra, 2D np.array)
            star.lightcurve (list of fluxes)
            star.residual (residual after dividing out of transit spectra, 2D np.array)

        """
        self.wave_start=float(wave_start)
        self.wave_end=float(wave_end)
        self.grid_size=int(grid_size)
        self.read_system(star_path=star_path,planet_path=planet_path,obs_path=obs_path)
        self.compute_spectrum()

        # return(self.wlF,F_out)#This is mad return statement. This whole function should be a class instead.

    def read_system(self,star_path='demo_star.txt',planet_path='demo_planet.txt',obs_path='demo_observations.txt'):
        """Reads in the stellar, planet and observation parameters from file; performing
        tests on the input and lifting the read variables to the class object.

        Parameters
        ----------
            star_path : str
                Path to parameter file defining the star.
            planet_path: str
                Path to the parameter file defining the planet and its orbit.
            obs_path: str
                Path to the parameter file defining the timestamps of the observations.
        """
        planetparams = open(planet_path,'r').read().splitlines()
        starparams = open(star_path,'r').read().splitlines()
        obsparams = open(obs_path,'r').read().splitlines()
        self.velStar = float(starparams[0].split()[0])
        self.stelinc = float(starparams[1].split()[0])
        self.drr = float(starparams[2].split()[0])
        self.T = float(starparams[3].split()[0])
        self.Z = float(starparams[4].split()[0])
        self.logg = float(starparams[5].split()[0])
        self.u1 = float(starparams[6].split()[0])
        self.u2 = float(starparams[7].split()[0])
        self.mus = float(starparams[8].split()[0])
        self.R = float(starparams[9].split()[0])
        self.sma_Rs = float(planetparams[0].split()[0])
        self.ecc = float(planetparams[1].split()[0])
        self.omega = float(planetparams[2].split()[0])
        self.orbinc = float(planetparams[3].split()[0])
        self.pob = float(planetparams[4].split()[0])#Obliquity.
        self.Rp_Rs = float(planetparams[5].split()[0])
        self.orb_p = float(planetparams[6].split()[0])
        self.transitC = float(planetparams[7].split()[0])
        self.mode = planetparams[8].split()[0]#This is a string.
        times = [] #These are in JD-24000000.0 or in orbital phase.
        self.exptimes = []
        for i in obsparams:
            times.append(float(i.split()[0]))
        self.times = np.array(times)
        self.Nexp = len(self.times)#Number of exposures.
        if self.mode == 'times':
            for i in obsparams:
                self.exptimes.append(float(i.split()[1]))
            self.exptimes = np.array(self.exptimes)
        try:
            test.typetest(self.wave_start,float,varname='wave_start in input')
            test.nantest(self.wave_start,varname='wave_start in input')
            test.notnegativetest(self.wave_start,varname='wave_start in input')
            test.notnegativetest(self.velStar,varname='velStar in input')
            test.notnegativetest(self.stelinc,varname='stelinc in input')
        #add all the other input parameters
        except ValueError as err:
            print("Parser: ",err.args)

        if self.mus != 0:
            self.mus = np.linspace(0.0,1.0,self.mus)#Uncomment this to run in CLV mode with SPECTRUM.

    def compute_spectrum(self):
        """This is where the main computation takes place. The simulation output
        and other variables are raised to class-wide attributes.

        Parameters
        ----------
            None
        """
        import math
        #Two arrays for the x and y axes
        self.x = np.linspace(-1,1,num=2*self.grid_size) #in units of stellar radius
        self.y = np.linspace(-1,1,num=2*self.grid_size) #in units of stellar radius
        #Calculate the velocity and flux grids
        print('--- Computing velocity/flux grids')
        self.vel_grid = vgrid.calc_vel_stellar(self.x,self.y, self.stelinc, self.velStar, self.drr,  self.pob)
        self.flux_grid = vgrid.calc_flux_stellar(self.x,self.y,self.u1,self.u2)

        if isinstance(self.mus,np.ndarray) != True:#SWITCH BETWEEN PHOENIX (mus=0) AND SPECTRUM
            print('--- Reading spectrum from PHOENIX')
            print('-----T=%sK, log(g)=%s, Z=%s.' % (self.T,self.logg,self.Z))
            wl,fx = spectrum.read_spectrum(self.T,self.logg,metallicity=self.Z)
            print('--- Integrating disk')
            if  self.drr == 0:
                print('------ Fast integration')
                wlF,F = integrate.build_spectrum_fast(wl,fx,self.wave_start,self.wave_end,self.x,self.y,self.vel_grid,self.flux_grid)
            else:
                print('------ Slow integration')
                wlF,F = integrate.build_spectrum_slow(wl,fx,self.wave_start,self.wave_end,self.x,self.y,self.vel_grid,self.flux_grid)
        else:#Meaning, if we have no mu's do:
            test.test_KURUCZ()
            print('--- Computing limb-resolved spectra with SPECTRUM')
            print('-----T=%sK, log(g)=%s, Z=%s.'% (self.T,self.logg,self.Z))

            # print(self.wave_start,math.floor(ops.vactoair(self.wave_start)*10.0)/10.0)
            # print(self.wave_end,math.ceil(ops.vactoair(self.wave_end)*10.0)/10.0)
            # sys.exit()
            maxvel = math.ceil(np.nanmax(np.abs(self.vel_grid)))
            wl,fx_list = spectrum.compute_spectrum(self.T,self.logg,self.Z,self.mus,math.floor(ops.vactoair(self.wave_start*ops.doppler((-1.0)*maxvel))*10.0)/10.0,math.ceil(ops.vactoair(self.wave_end*ops.doppler(maxvel))*10.0)/10.0,mode='anM')
            print('--- Integrating limb-resolved disk')
            # print(wl,self.wave_start)
            # sys.exit()
            wlF,F = integrate.build_spectrum_limb_resolved(wl,fx_list,self.mus, self.wave_start,self.wave_end,self.x,self.y,self.vel_grid)

        self.xp,self.yp,self.zp = ppos.calc_planet_pos(self.sma_Rs, self.ecc, self.omega, self.orbinc, self.pob, self.Rp_Rs, self.orb_p, self.transitC, self.mode, self.times, self.exptimes)

        # for i in range(len(self.xp)):
        #     print(self.xp[i],self.yp[i],self.zp[i])
        # pdb.set_trace()
        F_out = np.zeros((self.Nexp,len(F)))
        flux_out = []
        mask_out = []
        for i in range(self.Nexp):
            if isinstance(self.mus,np.ndarray) == True:
                wlp,Fp,flux,mask = integrate.build_local_spectrum_limb_resolved(self.xp[i],self.yp[i],self.zp[i],self.Rp_Rs,wl,fx_list,self.mus,self.wave_start,self.wave_end,self.x,self.y,self.vel_grid)
            else:
                wlp,Fp,flux,mask = integrate.build_local_spectrum_fast(self.xp[i],self.yp[i],self.zp[i],self.Rp_Rs,wl,fx,self.wave_start,self.wave_end,self.x,self.y,self.vel_grid,self.flux_grid)
            integrate.statusbar(i,self.Nexp)

            F_out[i,:]=F-Fp
            flux_out.append(flux)
            mask_out.append(mask)
        #This defines the output.
        self.wl = wlF
        self.stellar_spectrum = F
        self.spectra = copy.deepcopy(F_out)
        self.lightcurve = flux_out
        self.masks = mask_out
        self.compute_residuals()

    def compute_residuals(self):
        """Compute and return the residuals of the time_series of spectra.
        This may be the primary output of StarRotator.

        Parameters
        ----------
            None

        Returns
        -------
            residual : np.array
                2D matrix corresponding in dimensions to the length of the
                time-series times the number of wavelength points, containing
                the residuals of the spectra after dividing out the out-of-
                transit flux.
        """
        import pdb
        self.residual = self.spectra*0.0
        # self.residual_smooth = self.spectra*0.0
        for i in range(self.Nexp):
            self.residual[i,:]=self.spectra[i]/self.stellar_spectrum
        # pdb.set_trace()
        # self.apply_spectral_resolution(self.R)

    def plot_residuals(self):
        """Compute and plot the residuals.

        Parameters
        ----------
            None
        Returns
        -------
            None
        """
        import matplotlib.pyplot as plt
        # res = self.residuals()
        plt.pcolormesh(self.wl,self.times,self.residual)
        plt.xlabel('Wavelength (nm)')
        plt.ylabel('Phase')
        plt.show()

    def convolve_spectral_resolution(self):
        import lib.operations as ops
        import astropy.constants as const
        import copy
        import scipy.ndimage.filters as SNF
        import scipy.interpolate as interp
        from lib.integrate import statusbar as statusbar
        dv = const.c.value / self.R / 1000.0 #in km/s
        # wl_high = np.linspace(np.min(self.wl),np.max(self.wl),num=len(self.wl)*30.0)
        print('---Blurring')

        for i in range(len(self.spectra)):
            statusbar(i,len(self.spectra))
            # spec_i = interp.interp1d(self.wl,self.spectra[i])
            # self.residual_smooth[i] = ops.blur_spec(self.wl,copy.deepcopy(self.residual[i]),dv)
            # self.spectra_smooth[i] = ops.blur_spec(self.wl,copy.deepcopy(self.spectra[i]),dv)
            # self.spectra_smooth[i] = SNF.gaussian_filter(copy.deepcopy(self.spectra[i]),20.0,truncate=6.0)
            # self.spectra_smooth[i] = ops.smooth(copy.deepcopy(self.spectra[i]),40.0)
            # self.spectra_smooth[i] = interp.interp1d(wl_high,ops.blur_spec(wl_high,spec_i(wl_high),dv))(self.wl)

            res = self.residual[i]
            wl_constant_v,res_constant_v,a = ops.constant_velocity_wl_grid(self.wl,res,3.0)
            res_smooth = ops.smooth(res_constant_v,dv/a)
            self.residual[i] = interp.interp1d(wl_constant_v,res_smooth,fill_value='extrapolate',bounds_error=False)(self.wl)
            # self.residual[i] = ops.blur_spec(self.wl,copy.deepcopy(self.residual[i]),dv)
        # self.residual_save = copy.deepcopy(self.residual)
        # self.residual = copy.deepcopy(self.residual_smooth)
        # self.spectra_save = copy.deepcopy(self.spectra)
        # self.spectra = copy.deepcopy(self.spectra_smooth)

    # def undo_spectral_resolution(self):
    #     self.spectra = copy.deepcopy(self.spectra_save)
        # self.residual = copy.deepcopy(self.residual_save)

    def animate(self):
        """Plots an animation of the transit event, the stellar flux and velocity
        fields, and the resulting transit and line-shape variations. The animation
        is located in the subfolder `anim`. Anything located in this folder prior
        to running this function will be removed.

        Parameters
        ----------
            None
        Returns
        -------
            None
        """
        import matplotlib.pyplot as plt
        import lib.integrate as integrate
        import numpy as np
        from matplotlib.patches import Circle
        import shutil
        import os
        import os.path
        if os.path.isdir('anim/') == True:
            shutil.rmtree('anim/')#First delete the contents of the anim folder.
        os.mkdir('anim/')
        minflux = min(self.lightcurve)
        F=self.stellar_spectrum
        for i in range(self.Nexp):
            mask = self.masks[i]
            fig,ax = plt.subplots(nrows=2, ncols=2,figsize=(8,8))
            ax[0][0].pcolormesh(self.x,self.y,self.flux_grid*mask,vmin=0,vmax=1.0*np.nanmax(self.flux_grid),cmap='autumn')
            ax[1][0].pcolormesh(self.x,self.y,self.vel_grid*mask,cmap='bwr')
            ax[0][0].axes.set_aspect('equal')
            ax[1][0].axes.set_aspect('equal')
            if self.zp[i] > 0.0:
                planet1 = Circle((self.xp[i],self.yp[i]),self.Rp_Rs, facecolor='black', edgecolor='black', lw=1)
                planet2 = Circle((self.xp[i],self.yp[i]),self.Rp_Rs, facecolor='black', edgecolor='black', lw=1)
                ax[0][0].add_patch(planet1)
                ax[1][0].add_patch(planet2)
            ax[0][1].plot(self.times[0:i],self.lightcurve[0:i],'.',color='black')
            ax[0][1].set_xlim((min(self.times),max(self.times)))
            ax[0][1].set_ylim((minflux-0.1*self.Rp_Rs**2.0),1.0+0.1*self.Rp_Rs**2)
            ax[1][1].plot(self.wl,F/np.nanmax(F),color='black',alpha = 0.5)
            ymin = np.nanmin(F/np.nanmax(F))
            ymax = np.nanmax(F/np.nanmax(F))
            linedepth = ymax - ymin
            ax[1][1].plot(self.wl,self.spectra[i]/np.nanmax(self.spectra[i]),color='black')
            # ax[1][1].set_xlim((588.5,590.2))
            yl = (ymin-0.1*linedepth,ymax+0.3*linedepth)
            ax[1][1].set_ylim(yl)
            ax2 = ax[1][1].twinx()
            ax2.plot(self.wl,(self.spectra[i])*np.nanmax(F)/F/np.nanmax(self.spectra[i]),color='skyblue')
            sf = 30.0
            ax2.set_ylim((1.0-(1-yl[0])/sf,1.0+(yl[1]-1)/sf))
            ax2.set_ylabel('Ratio in transit / out of transit',fontsize = 7)
            ax2.tick_params(axis='both', which='major', labelsize=6)
            ax2.tick_params(axis='both', which='minor', labelsize=5)
            ax[0][0].set_ylabel('Y (Rs)',fontsize=7)
            ax[0][0].tick_params(axis='both', which='major', labelsize=6)
            ax[0][0].tick_params(axis='both', which='minor', labelsize=5)
            ax[1][0].set_ylabel('Y (Rs)',fontsize=7)
            ax[1][0].set_xlabel('X (Rs)',fontsize=7)
            ax[1][0].tick_params(axis='both', which='major', labelsize=6)
            ax[1][0].tick_params(axis='both', which='minor', labelsize=5)
            ax[0][1].set_ylabel('Normalised flux',fontsize=7)
            ax[0][1].set_xlabel('Orbital Phase',fontsize='small')
            ax[0][1].tick_params(axis='both', which='major', labelsize=6)
            ax[0][1].tick_params(axis='both', which='minor', labelsize=5)
            ax[1][1].set_ylabel('Normalised flux',fontsize=7)
            ax[1][1].set_xlabel('Wavelength (nm)',fontsize=7)
            ax[1][1].tick_params(axis='both', which='major', labelsize=6)
            ax[1][1].tick_params(axis='both', which='minor', labelsize=5)
            if len(str(i)) == 1:
                out = '000'+str(i)
            if len(str(i)) == 2:
                out = '00'+str(i)
            if len(str(i)) == 3:
                out = '0'+str(i)
            if len(str(i)) == 4:
                out = str(i)
            fig.savefig('anim/'+out+'.png', dpi=fig.dpi)
            integrate.statusbar(i,self.Nexp)
            plt.close()
        print('',end="\r")
        print('--- Saving to animation.gif')

        status = os.system('convert -delay 8 anim/*.png animation.gif')
        if status != 0:
            print('The conversion of the animation frames into a gif has')
            print('failed; probably because the Imagemagick convert command')
            print('was not found. The animation frames have been created in')
            print('the anim/ folder. If you want to convert these into a .gif,')
            print('please do it manually, or install Imagemagick, see')
            print('https://imagemagick.org')

        #This is for plotting random comparisons.
        # wl2,fx2 = spectrum.read_spectrum(T,logg)
        # wlF2,F2 = integrate.build_spectrum_fast(wl2,fx2, wave_start, wave_end,x,y,vel_grid,flux_grid)
        # wlp,Fp,flux,mask = integrate.build_local_spectrum_limb_resolved(-0.3,0.0,0.1,wl,fx_list,mus, wave_start, wave_end,x,y,vel_grid)
        # wlp2,Fp2,flux2,mask2 = integrate.build_local_spectrum_fast(-0.3,0.0,0.1,wl2,fx2, wave_start, wave_end,x,y,vel_grid,flux_grid)
        ##This overplots non-rotating SPECTRUM and PHOENIX spectra, normalised.
        ## plt.plot(wl,fx_list[-1]/max(fx_list[-1]),label='SPECTRUM')
        ## plt.plot(wl2,fx2/5e15,label='PHOENIX')
        ## plt.xlabel('Wavelength (nm)')
        ## plt.ylabel('Max-normalised flux')
        ## plt.title('T = %s K, log(g) = %s' % (T,logg))
        ## plt.legend()
        ## plt.show()

        # plt.plot(wlF,F/max(F),color='skyblue',alpha=0.5)
        # plt.plot(wlF,(F-Fp)/np.nanmax(F-Fp),color='skyblue',label='SPECTRUM')
        # plt.plot(wlF2,F2/max(F2),color='red',alpha=0.5)
        # plt.plot(wlF2,(F2-Fp2)/np.nanmax(F2-Fp2),color='red',label='PHOENIX')
        # plt.legend()
        # plt.xlabel('Wavelength (nm)')
        # plt.ylabel('Max-normalised flux')
        # plt.title('T = %s K, log(g) = %s, vsini = 110km/s' % (T,logg))
        # plt.show()

    # void1,void2,minflux,void3 = integrate.build_local_spectrum_fast(0,0,RpRs,wl,fx, wave_start, wave_end,x,y,vel_grid,flux_grid)

    #The following puts a large circular spot with a T 1000K less than the star in the center.
    # wl2,fx2 = spectrum.read_spectrum(T-1000,logg)
    # wlp,Fp,flux,mask = integrate.build_local_spectrum_fast(0,0,0.2,wl,fx, wave_start, wave_end,x,y,vel_grid,flux_grid)
    # wls,Fs,fluxs,masks = integrate.build_local_spectrum_fast(0,0,0.2,wl2,fx2, wave_start, wave_end,x,y,vel_grid,flux_grid)
    # Ft = F-Fp+Fs
    # plt.plot(wlF,F)
    # plt.plot(wlF,Ft)
    # plt.show()
    # sys.exit()

max_stars_repo_name: Hoeijmakers/StarRotator

max_issues_repo_name: Hoeijmakers/StarRotator

max_forks_repo_name: Hoeijmakers/StarRotator

reason: import numpy,import scipy,import astropy


text:
""" Tests the DiscreteGRF class.

"""
import numpy as np
import torch
from meslas.means import ConstantMean
from meslas.covariance.spatial_covariance_functions import Matern32
from meslas.covariance.cross_covariances import UniformMixing
from meslas.covariance.heterotopic import FactorCovariance
from meslas.geometry.grid import TriangularGrid, SquareGrid
from meslas.random_fields import GRF, DiscreteGRF
from meslas.excursion import coverage_fct_fixed_location
from meslas.plotting import plot_grid_values, plot_grid_probas

# Dimension of the response.
n_out = 2

# Spatial Covariance.
matern_cov = Matern32(lmbda=0.1, sigma=1.0)

# Cross covariance.
cross_cov = UniformMixing(gamma0=0.4, sigmas=[np.sqrt(1.0), np.sqrt(5.5)])

covariance = FactorCovariance(matern_cov, cross_cov, n_out=n_out)

# Specify mean function
mean = ConstantMean([0.0, 5.0])

# Create the GRF.
myGRF = GRF(mean, covariance)

# Create an equilateral triangular grid in 2 dims.
# Number of respones.
my_grid = TriangularGrid(40)
my_square_grid = SquareGrid(50, 2)

my_discrete_grf = DiscreteGRF.from_model(myGRF, my_grid)

# Sample and plot.
sample = my_discrete_grf.sample()
plot_grid_values(my_grid, sample)

# Sample and plot.
sample = my_discrete_grf.sample()
plot_grid_values(my_grid, sample)

# Sample the continuous version and compare.
sample_cont = myGRF.sample_isotopic(my_grid.points)
plot_grid_values(my_grid, sample_cont)

# Observe some data.
S_y = torch.tensor([[0.2, 0.1], [0.2, 0.2], [0.2, 0.3],
        [0.2, 0.4], [0.2, 0.5], [0.2, 0.6],
        [0.2, 0.7], [0.2, 0.8], [0.2, 0.9], [0.2, 1.0],
        [0.6, 0.5]])
L_y = torch.tensor([0, 0, 0, 0, 0, 1, 1, 0 ,0 ,0, 0])
y = torch.tensor(11*[-6]).float()

# Since we are working with a discrete GRF, we can only observa data at grid
# nodes. Hence get the ones corresponding to the measured data.
S_y_inds = my_grid.get_closest(S_y)

my_discrete_grf.update(S_y_inds, L_y, y, noise_std=0.05)
plot_grid_values(my_grid, my_discrete_grf.mean_vec, S_y, L_y)

# -----------------------------------------
# Now compare with the non-discrete version.
# -----------------------------------------
mu_cond, K_cond = myGRF.krig_isotopic(
        my_grid.points, S_y, L_y, y,
        noise_std=0.05,
        compute_post_cov=True)

plot_grid_values(my_grid, mu_cond, S_y, L_y)

# -----------------------------------------
# Try the variance part.
plot_grid_values(my_grid, my_discrete_grf.variance, S_y, L_y, cmap="proba")

max_stars_repo_name: CedricTravelletti/MESLAS

max_issues_repo_name: CedricTravelletti/MESLAS

max_forks_repo_name: CedricTravelletti/MESLAS


text:
from keras import backend as K
from keras.layers import Dense, Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
import numpy as np
import math

class QACAgent(object):
    def __init__(self, input_dims, n_actions, alpha, beta, gamma=0.99,
                critic_skew = 0.0, qscale = 1.0,
                layer1_size=1024, layer2_size=512):
        self.gamma = gamma
        self.alpha = alpha
        self.beta = beta
        self.input_dims = input_dims
        self.fc1_dims = layer1_size
        self.fc2_dims = layer2_size
        self.n_actions = n_actions
        self.critic_skew = critic_skew
        self.QScale = qscale

        self.value, self.advantage, self.policy = self.build_value_advantage_network()

        self.action_space = [i for i in range(n_actions)]
        self.all_layers = []

        self.q2_ma = None
        self.q2_ma_alpha = 0.01

    def save(self, path):
        weights = {}
        for i, l in enumerate(self.all_layers):
            for j, w in l.get_weights():
                weights["w_%d_%d" % (i, j)] = w
        np.savez(path, **weights)

    def load(self, path):
        data = np.load(path)
        for i, l in enumerate(self.all_layers):
            weights = [data["w_%d_%d" % (i, j)] for j, _ in enumerate(l.get_weights())]
            l.set_weights(weights)

    def build_value_advantage_network(self):
        input = Input(shape=(self.input_dims,))
        action_mask = Input(shape=(self.n_actions,))
        dense1 = Dense(self.fc1_dims, activation='relu')(input)
        dense2 = Dense(self.fc2_dims, activation='relu')(dense1)

        dense3 = Dense(self.fc2_dims//10, activation='relu')(dense2)
        advantage = Dense(self.n_actions, activation='linear')(dense3)

        #
        # Policy model
        #

        policy_model = Model(input=[input], output=[advantage])

        #
        # advantage training
        #

        def masked_advantage(args):
            mask, adv = args
            return K.reshape(K.sum(mask*adv, axis = -1), (-1,1))

        masked_advantage = Lambda(masked_advantage)([action_mask, advantage])

        advatnage_model = Model(inputs=[input, action_mask], outputs=[masked_advantage])
        advatnage_model.compile(optimizer=Adam(lr=self.alpha), loss="mse")

        #
        # value model
        #

        dense4 = Dense(self.fc2_dims//10, activation='relu')(dense2)
        values = Dense(1, activation='linear')(dense4)
        values_model = Model(input=[input], output=[values])

        def skewed_loss(y_, y):
            delta = y_ - y
            delta = (delta + K.abs(delta)*self.critic_skew)/(1.0+self.critic_skew)
            return K.mean(delta*delta, axis=-1)

        values_model.compile(optimizer=Adam(lr=self.beta), loss=skewed_loss)

        self.all_layers = [dense1, dense2, dense3, dense4, advantage, values]

        return values_model, advatnage_model, policy_model

    def choose_action(self, observation, test = False, epsilon=0.0, tau=None):
        if tau is None: tau = self.QScale
        weights = self.policy.predict(observation[np.newaxis,:])[0]
        q2 = np.mean(weights*weights)
        q = np.mean(weights)
        if self.q2_ma is None:
            self.q2_ma = q2
            self.q_ma = q
        self.q2_ma += self.q2_ma_alpha*(q2-self.q2_ma)
        self.q_ma += self.q2_ma_alpha*(q-self.q_ma)
        s = math.sqrt(self.q2_ma - self.q_ma**2)

        if test or tau < 0.0:
            action = np.argmax(weights)
        else:
            weights /= s
            weights = weights - np.max(weights)
            exp_values = np.exp(weights / tau)
            probs = exp_values / np.sum(exp_values)
            probs = (probs+epsilon)/(1.0+len(probs)*epsilon)
            action = np.random.choice(self.action_space, p=probs)
        return action

    def learn(self, state, action, reward, state_, done):
        state = state[np.newaxis,:]
        state_ = state_[np.newaxis,:]
        action = np.array([action])
        reward = np.array([reward])
        done = np.array([done])

        return self.learn_batch(state, action, reward, state_, done)

    def learn_batch(self, states, actions, rewards, states_, dones):
        n = len(states)
        values1 = self.value.predict(states_)
        values0 = self.value.predict(states)
        rewards = rewards.reshape((-1, 1))
        dones = dones.reshape((-1, 1))

        #print("learn_batch: states:", states.shape, " actions:", actions.shape, " rewards:", rewards.shape,
        #            " states_:", states_.shape, " dones:", dones.shape)

        #print("             critic_values_:", critic_values_.shape, " critic_values:", critic_values.shape)

        mask = np.zeros((n, self.n_actions))
        mask[np.arange(n), actions] = 1.0

        values_ = rewards + self.gamma*values1*(1.0-dones)
        advatages_ = values_ - values0
        advatages_ = advatages_.reshape((-1,1))

        #print("learn_batch: states:", states.shape, " deltas:", deltas.shape, " action_arrays:", action_array.shape,
        #            " tagets:", targets.shape)

        advantage_metrics = self.advantage.train_on_batch([states, mask], advatages_)
        value_metrics = self.value.train_on_batch(states, values_)

        return advantage_metrics, value_metrics

    def learn_batches(self, mb_size, state, action, reward, state_, done, tau=1.0, shuffle=True):
        # make sure data is np arrays
        state = np.array(state)
        #tau = np.ones((len(state),1))*tau
        action = np.array(action)
        reward = np.array(reward)
        state_ = np.array(state_)
        done = np.array(done, dtype=np.int)
        #print("learn_batches: inputs:", state.shape, action.shape, reward.shape, state_.shape, done.shape)

        n = len(state)

        critic_value = self.critic.predict(state)[:,0]
        critic_value_ = self.critic.predict(state_)[:,0]

        target = reward + self.gamma*critic_value_*(1-done)
        delta =  target - critic_value
        #print("learn_batches: delta:", delta.shape, delta)

        actions = np.zeros([n, self.n_actions])
        actions[np.arange(n), action] = 1

        target = target.reshape((-1,1))
        delta = delta.reshape((-1,1))
        #print("learn_batches: state:", state.shape, "  delta:", delta.shape, "  actions:", actions.shape, "  target:", target.shape)

        for i in range(0, n, mb_size):
            actor_metrics = self.actor.train_on_batch([state[i:i+mb_size], delta[i:i+mb_size]], actions[i:i+mb_size])
            critic_metrics = self.critic.train_on_batch(state[i:i+mb_size], target[i:i+mb_size])

        return actor_metrics, critic_metrics

    def run_episode(self, env, learn = False, test = False, render=False):
        done = False
        observation = env.reset()
        if render:
            env.render()
        score = 0.0
        record = []
        actor_metrics, critic_metrics = None, None
        while not done:
            action = self.choose_action(observation, test)
            observation_, reward, done, info = env.step(action)
            if render:
                env.render()
            record.append((observation, action, reward, observation_, 1.0 if done else 0.0, info))
            if learn:
                self.learn(observation, action, reward, observation_, done)
            observation = observation_
            score += reward
        return score, record

max_stars_repo_name: imandr/KeRLas

max_issues_repo_name: imandr/KeRLas

max_forks_repo_name: imandr/KeRLas