text:
"""Background maps module."""

import os
import re

import numpy as np

import matplotlib.pyplot as plt
from matplotlib.image import imread
from matplotlib.path import Path
from matplotlib.patches import PathPatch
from matplotlib.collections import PatchCollection

from .vertices import path_lonlat
from ..projections.stereographic import r_stereo, xy as xy_stereo
from ..vars import ROOT_DATA
from ..vectors import deg180, deg360
from ..coordinates import slat, slon_e, slon_w

ROOT = ROOT_DATA / 'maps'

def parse(rexp, line):
    r"""Parse README line.

    Parameters
    ----------
    rexp: str
        Regular expression for parser.
    line: str
        Line to parse.

    Returns
    -------
    array
        Parsed value(s).

    Raises
    ------
    ValueError
        If not match was found.

    Example
    -------
    >>> parse(r'\d+', 'abc12def45')
    ['12', '45']

    """
    res = re.findall(rexp, line)
    if not res:
        raise ValueError(f'Invalid line: {line}')
    return res

def add(maps, keys, attr, value):
    """Add value(s) to maps dict.

    Parameters
    ----------
    maps: dict
        Global list of maps.
    keys: list
        List of file maps to append.
    attr: str
        Key value.
    value: any
        Value to append. If the value is a tuple, the values
        will loop according to ``keys`` index.

    Raises
    ------
    ValueError
        If the keys variable is ``None``.

    """
    if keys is None:
        raise ValueError('No map key(s) selected.')

    if isinstance(value, str) and value == 'None':
        return

    for i, key in enumerate(keys):
        if key not in maps.keys():
            maps[key] = {}
        maps[key][attr] = value[i] if isinstance(value, tuple) else value

def basename(fname):
    """Extract filename basename without extension.

    Parameters
    ----------
    fname: str
        Input filename

    Returns
    -------
    str
        File basename without extension.

    Example
    -------
    >>> basename('foo')
    'foo'
    >>> basename('foo.txt')
    'foo'
    >>> basename('test.foo.txt')
    'test.foo'
    >>> basename('test/foo.txt')
    'foo'
    >>> basename('/test/foo.txt')
    'foo'

    """
    return str(fname) if '.' not in str(fname) else \
        '.'.join(os.path.basename(str(fname)).split('.')[:-1])

def parse_readme(filename):
    """Parse README file.

    Parameters
    ----------
    filename: str
        README filename.

    Returns
    -------
    dict
        List of all the maps available in the README.

    """
    with open(filename, 'r') as f:
        lines = f.read().splitlines()

    maps = {}
    name = None
    keys = None
    for line in lines:
        if line.startswith('##'):
            name = line[2:].strip()
            keys = None

        elif line.startswith('* Filename:'):
            filenames = parse(r'`([\w\.\-\s/\\]+)`', line)

            files = tuple([os.path.basename(fname) for fname in filenames])
            roots = tuple([os.path.dirname(fname) for fname in filenames])

            keys = tuple([basename(f) for f in files])

            add(maps, keys, 'fname', files)
            add(maps, keys, 'root', roots)
            add(maps, keys, 'name', name)

        elif line.startswith('* Source:'):
            src, url = parse(r'\[([\w\.\-\s]*)\]\(([\w\.\-\s:/]*)\)', line)[0]

            add(maps, keys, 'src', src)
            add(maps, keys, 'url', url)

        elif line.startswith('* Extent:'):
            lon_1, lon_2, lat_1, lat_2 = parse(r'(-?\d+\.\d+|-?\d+)', line)
            extent = [float(lon_1), float(lon_2), float(lat_1), float(lat_2)]
            add(maps, keys, 'extent', extent)

        elif line.startswith('* Projection:'):
            projection = parse(r'`([\w\-_\s]+)`', line)[0]
            add(maps, keys, 'projection', projection.lower())

    return maps

def remove_from_readme(filename, fname):
    """Remove map from README file.

    Parameters
    ----------
    filename: str
        README filename.
    fname: str
        Filename to remove

    Returns
    -------
    dict
        List of all the maps available in the README.

    """
    with open(filename, 'r') as f:
        lines = f.read().splitlines()

    new = []
    buffer = []
    include = True
    for i, line in enumerate(lines):
        if line.startswith('##'):
            if buffer:
                new += buffer

            include = True
            buffer = []

        elif line.startswith('* Filename:'):
            filenames = parse(r'`([\w\.\-\s/\\]+)`', line)

            files = tuple([basename(fname) for fname in filenames])

            if fname in files:
                include = False
                buffer = []

        if include:
            buffer.append(line)

    if buffer:
        new += buffer

    with open(filename, 'w') as f:
        f.write('\n'.join(new))

class MapsDetails(type):
    """Map details parser."""

    __maps = {}

    def __repr__(cls):
        return '\n - '.join([
            f'<{cls.__name__}> Available: {len(cls)}',
            *cls.maps().keys()
        ])

    def __len__(cls):
        return len(cls.maps())

    def __contains__(cls, item):
        return item in cls.maps().keys()

    def __iter__(cls):
        return iter(cls.maps())

    def __getitem__(cls, item):
        try:
            return Map(**cls.maps()[item])
        except KeyError:
            raise KeyError(f'Unknown map `{item}`.')

    @classmethod
    def maps(cls):
        """Load all maps listed in content."""
        if not cls.__maps:
            cls.__maps = parse_readme(cls.filename())
        return cls.__maps

    @classmethod
    def filename(cls):
        """Get README filename.

        Raises
        ------
        FileNotFoundError
            If the file is missing.

        """
        filename = ROOT / 'README.md'

        if not filename.exists():
            with filename.open('w') as f:
                f.write('\n'.join([
                    'List of maps available',
                    '======================\n',
                    '## Titan VIMS/ISS\n',
                    '* Filename: `Titan_VIMS_ISS.jpg`',
                    '* Source: [Seignovert et al. 2019]'
                    '(https://doi.org/10.22002/D1.1173)',
                    '* Extent: `-180° 180° -90° 90°`',
                    '* Projection: `equirectangular`\n',
                ]))

        return str(filename)

    @classmethod
    def register(cls, m, update=False):
        """Register a new maps in README.

        Parameters
        ----------
        m: Map
            Background map object.
        update: bool, optional
            Enable map update with the same key.

        Raises
        ------
        TypeError
            If the map provided doest not have a
            ``markdown`` attribute.
        ValueError
            If the ``name`` attribute is already used.

        """
        if not hasattr(m, 'markdown'):
            raise TypeError(f'Map has an invalid type: `{type(m)}`.')

        key = basename(m)

        if key in cls.maps().keys():
            if update:
                cls.remove(key)
            else:
                raise ValueError(f'A map is already registered with the name: `{key}`.')

        with open(cls.filename(), 'a') as f:
            f.write(f'\n{m.markdown}')

        cls.__maps[key] = dict(m)

        print(f'Image `{key}` saved in MAPS registry.')

    @classmethod
    def remove(cls, fname):
        """Remove a map file from README."""
        key = basename(fname)

        if key in cls.maps().keys():
            remove_from_readme(cls.filename(), key)
            del cls.__maps[key]

            print(f'Image `{key}` removed from MAPS registry.')

    @classmethod
    def reload(cls):
        """Reload maps list from README."""
        cls.__maps = {}

class MAPS(metaclass=MapsDetails):
    """List of all the registered maps."""

class Map:
    """Background map object.

    Image filename format:

        `ROOT`/`Target`_`INSTR`.`ext`

    Parameters
    ----------
    fname: str
        Image filename.
    root: str, optional
        Data root folder containing the map.
        If ``None`` is provided (default), the
        package ``maps`` folder will be used.
    extent: list, required
        Map extent: [lon_w_0, lon_w_1, lat_1, lat_2]
    src: str, optional
        Image source name.
    url: str, optional
        Image URL source.
    projection:
        Image projection (default: ``equirectangular``).
    name: str, optional
        Name of the map (for the registry).

    """

    def __init__(self, fname, root=None, extent=None, src=None, url=None,
                 projection=None, name=None):
        self.root = root if root else str(ROOT)
        self.fname = fname

        self.data_extent = list(extent)
        self.src = src
        self.url = url
        self.proj = projection.lower() if isinstance(projection, str) else None

        self.name = name if name is not None else \
            '.'.join(self.fname.split('.')[:-1])

    def __str__(self):
        return self.fname if self.root == str(ROOT) else self.filename

    def __repr__(self):
        return '\n - '.join([
            f'<{self.__class__.__name__}> {self}',
            f'Extent: {"Undefined" if self.data_extent is None else self.data_extent}',
            f'Source: {"Undefined" if self.src is None else self.src}',
            f'URL: {"Undefined" if self.url is None else self.url}',
            f'Projection: {"Undefined" if self.proj is None else self.proj.title()}',
        ])

    def __iter__(self):
        for k, v in self.as_dict.items():
            yield k, v

    def __call__(self, *args):
        if len(args) == 1:
            if isinstance(args[0], PatchCollection):
                return self.xy_collection(args[0])

            if isinstance(args[0], PathPatch):
                return self.xy_patch(args[0])

            if isinstance(args[0], Path):
                return self.xy_path(args[0])

        if len(args) == 2:
            return self.xy(args[0], args[1])

        raise ValueError('A `PatchCollection`, `PathPatch`, `Patch` '
                         'or (lon_w, lat) attributes are required.')

    @property
    def fname(self):
        return self.__fname

    @fname.setter
    def fname(self, fname):
        if os.path.dirname(fname):
            self.root = os.path.dirname(fname)

        self.__fname = os.path.basename(fname)
        self.__img = None
        self.__n_pole = None

        if not os.path.exists(self.filename):
            raise FileNotFoundError(f'Map `{self.filename}` is not available.')

    @property
    def filename(self):
        """Image absolute path."""
        return os.path.join(self.root, self.fname)

    @property
    def extent_str(self):
        """Format data extent as string."""
        return ' '.join([f'{e}°' for e in self.data_extent])

    @property
    def markdown(self):
        """Convert Map object into markdown."""
        return '\n'.join([
            f'## {self.name}\n',
            f'* Filename: `{self}`',
            f'* Source: [{self.src}]({self.url})',
            f'* Extent: `{self.extent_str}`',
            f'* Projection: `{self.proj}`',
            f'\n![{self.name}]({self.filename})',
        ])

    @property
    def as_dict(self):
        """Convert as dict."""
        return {
            'name': self.name,
            'fname': self.fname,
            'root': self.root,
            'src': self.src,
            'url': self.url,
            'extent': self.data_extent,
            'projection': self.proj,
        }

    def register(self, update=False):
        """Register map in MAPS."""
        MAPS.register(self, update=update)

    @property
    def img(self):
        """Image data."""
        if self.__img is None:
            self.__img = imread(self.filename)
        return self.__img

    @property
    def shape(self):
        """Background image shape."""
        return self.img.shape

    @property
    def ndim(self):
        """Background image dimension."""
        return self.img.ndim

    @property
    def center_0(self):
        """Center equirectangular image background at 0°."""
        if self._proj != 'lonlat':
            raise ValueError('Recentering is only available for equirectangular maps.')

        if self.data_extent[0] == -180 and self.data_extent[1] == 180:
            return self

        if self.data_extent[0] != 360 or self.data_extent[1] != 0:
            raise ValueError('Longitude span invalid. Must be [360, 0], not '
                             f'[{self.data_extent[0]:.1f}, {self.data_extent[1]:.1f}].')

        if self.ndim == 2:
            _, ns = self.shape
            self.__img = np.hstack([self.img[:, ns // 2:], self.img[:, :ns // 2]])
        elif self.ndim == 3:
            _, ns, _ = self.shape
            self.__img = np.hstack([self.img[:, ns // 2:, :], self.img[:, :ns // 2, :]])
        else:
            raise ValueError(f'Image dimension invalid: {self.ndim}.')

        self.data_extent[:2] = [-180, 180]
        return self

    @property
    def center_180(self):
        """Center equirectangular image background at 180°."""
        if self._proj != 'lonlat':
            raise ValueError('Recentering is only available for equirectangular maps.')

        if self.data_extent[0] == 360 and self.data_extent[1] == 0:
            return self

        if self.data_extent[0] != -180 or self.data_extent[1] != 180:
            raise ValueError('Longitude span invalid. Must be [-180, 180], not '
                             f'[{self.data_extent[0]:.1f}, {self.data_extent[1]:.1f}].')

        if self.ndim == 2:
            _, ns = self.shape
            self.__img = np.hstack([self.img[:, ns // 2:], self.img[:, :ns // 2]])
        elif self.ndim == 3:
            _, ns, _ = self.shape
            self.__img = np.hstack([self.img[:, ns // 2:, :], self.img[:, :ns // 2, :]])
        else:
            raise ValueError(f'Image dimension invalid: {self.ndim}.')

        self.data_extent[:2] = [360, 0]
        return self

    @property
    def n_pole(self):
        """Pole observered for polar projection."""
        if self.__n_pole is None:
            if self._proj == 'stereo':
                self.__n_pole = self.data_extent[3] > 0
        return self.__n_pole

    @property
    def xright(self):
        """Check if x-axis is positive to the right."""
        return self.data_extent[1] < 0

    @property
    def _proj(self):
        if self.proj in [None, 'equi', 'equirectangular', 'plate carrée', 'lonlat']:
            return 'lonlat'

        if self.proj in ['stereo', 'stereographic']:
            return 'stereo'

        raise ValueError(f'Projection `{self.proj}` is not available.')

    @property
    def extent(self):
        """Projected data extent."""
        if self._proj in ['lon_w_lat', 'lon_e_lat']:
            return self.data_extent

        if self._proj == 'stereo':
            r = r_stereo(self.data_extent[2], n_pole=self.n_pole)
            return [-r, r, r, -r]

        return self.data_extent

    def xy(self, lon_w, lat):
        """Convert point coordinate in map coordinates.

        Parameters
        ----------
        lon_w: float
            Point longitude West (degree).
        lat: float
            Point latitude (degree).

        Returns
        -------
        float, float
            X-Y coordinates of the point on the surface.

        """
        if self._proj == 'lonlat':
            return deg180(np.asarray(lon_w)) if self.xright else deg360(lon_w), lat

        if self._proj == 'stereo':
            return xy_stereo(lon_w, lat, n_pole=self.n_pole)

        return self.data_extent

    def xy_path(self, path):
        """Convert path vertices in map coordinates.

        Parameters
        ----------
        path: matplotlib.path.Path
            Pyplot path.

        Returns
        -------
        matplotlib.path.Path
            Re-projected path.

        """
        if path is None:
            return None

        vertices = np.transpose(self.xy(*path.vertices.T))

        _path = Path(vertices, path.codes)

        if self._proj == 'lonlat':
            _path = path_lonlat(_path)

        return _path

    def xy_patch(self, patch):
        """Convert patch vertices in map coordinates.

        Parameters
        ----------
        patch: matplotlib.patches.PathPatch
            Pyplot patch.

        Returns
        -------
        matplotlib.patches.PathPatch
            Re-projected patch.

        """
        return PathPatch(
            self.xy_path(patch.get_path()),
            edgecolor=patch.get_ec(),
            facecolor=patch.get_fc(),
        )

    def xy_collection(self, collection):
        """Convert patch collection vertices in map coordinates.

        Parameters
        ----------
        collection: matplotlib.collections.PatchCollection
            Pyplot patch collection.

        Returns
        -------
        matplotlib.collections.PatchCollection
            Re-projected patch collection.

        """
        return PatchCollection(
            [
                PathPatch(self.xy_path(path))
                for path in collection.get_paths()
            ],
            edgecolor=collection.get_ec(),
            facecolor=collection.get_fc(),
        )

    def lons(self, lon_min=None, lon_max=None, npts=None):
        """Get longitude ticks.

        Parameters
        ----------
        lon_min: float, optional
            Minimal longitude. Use ``data_extent`` min value
            if ``None`` provided.
        lon_max: float, optional
            Minimal longitude. Use ``data_extent`` max value
            if ``None`` provided.
        npts: int, optional
            Number of points to output (default ``13``).

        Returns
        -------
        numpy.array
            List of longitudes.

        """
        if self._proj == 'stereo':
            return [0]

        lons = np.linspace(
            min(self.data_extent[:2]) if lon_min is None else lon_min,
            max(self.data_extent[:2]) if lon_max is None else lon_max,
            13 if npts is None else npts,
        )

        if self.data_extent[1] < 0:
            lons = lons[::-1]  # Revert x-axis

        return lons

    def lats(self, lat_min=None, lat_max=None, npts=None):
        """Get latitude ticks.

        Parameters
        ----------
        lat_min: float, optional
            Minimal latitude. Use ``data_extent`` min value
            if ``None`` provided.
        lat_max: float, optional
            Minimal latitude. Use ``data_extent`` max value
            if ``None`` provided.
        npts: int, optional
            Number of points to output (default ``7``).

        Returns
        -------
        numpy.array
            List of latitudes.

        """
        if self._proj == 'stereo':
            return [0]

        return np.linspace(
            min(self.data_extent[2:]) if lat_min is None else lat_min,
            max(self.data_extent[2:]) if lat_max is None else lat_max,
            7 if npts is None else npts,
        )

    def lonlabels(self, lon_min=None, lon_max=None, npts=None, precision=0):
        """Get longitude labels.

        Parameters
        ----------
        lon_min: float, optional
            Minimal longitude. Use ``data_extent`` min value
            if ``None`` provided.
        lon_max: float, optional
            Minimal longitude. Use ``data_extent`` max value
            if ``None`` provided.
        npts: int, optional
            Number of points to output (default ``13``).
        precision: int, optional
            Displayed float precision.

        Returns
        -------
        numpy.array
            List of longitude labels.

        """
        if self._proj == 'stereo':
            return [slon_w(0)]

        lons = self.lons(lon_min=lon_min, lon_max=lon_max, npts=npts)

        if self.data_extent[1] < 0:  # East longitude
            labels = [slon_e(-lon) for lon in lons]
        else:
            labels = [slon_w(lon) for lon in lons]

        return labels

    def latlabels(self, lat_min=None, lat_max=None, npts=None, precision=0):
        """Get latitude labels.

        Parameters
        ----------
        lat_min: float, optional
            Minimal latitude. Use ``data_extent`` min value
            if ``None`` provided.
        lat_max: float, optional
            Minimal latitude. Use ``data_extent`` max value
            if ``None`` provided.
        npts: int, optional
            Number of points to output (default ``13``).
        precision: int, optional
            Displayed float precision.

        Returns
        -------
        numpy.array
            List of latitude labels.

        """
        if self._proj == 'stereo':
            return [slon_w(90) if self.n_pole else slon_w(270)]

        lats = self.lats(lat_min=lat_min, lat_max=lat_max, npts=npts)

        return [slat(lat, precision) for lat in lats]

    @property
    def xlim(self):
        """X-axis limits based on image background."""
        return self.extent[0], self.extent[1]

    @property
    def ylim(self):
        """Y-axis limits based on image background."""
        return self.extent[2], self.extent[3]

    @staticmethod
    def paralleles(lat, nlons=181, lon_min=0, lon_max=360):
        """Paralleles coordinates.

        Parameters
        ----------
        lats: int, float or list
            Parallele latitude(s).
        nlons: int, optional
            Number of west longitudes in the parallele(s).
        lon_min: float, optional
            Minimum west longitude.
        lon_max: float, optional
            Maximum west longitude.

        Returns
        -------
        ([float, …], [float, …]) or [[float, float], …]
            West longitude and latitude of the parallele.

        """
        lons = np.linspace(lon_min, lon_max, nlons)
        lats = np.repeat(lat, nlons)

        if isinstance(lat, (list, np.ndarray)):
            nlats = len(lat)
            lats = np.repeat(lat, nlons).reshape(nlats, nlons).T
            lons = np.repeat(lons, nlats).reshape(nlons, nlats)

        return lons, lats

    @staticmethod
    def meridians(lon, nlats=91, lat_min=-90, lat_max=90):
        """Meridian coordinates.

        Parameters
        ----------
        lons: int, float or list
            Meridian longitude(s).
        nlats: int, optional
            Number of latitudes in the meridian(s).
        lat_min: float, optional
            Minimum latitude.
        lat_max: float, optional
            Maximum latitude.

        Returns
        -------
        ([float, …], [float, …]) or [[float, float], …]
            West longitude and latitude of the parallele.

        """
        lats = np.linspace(lat_min, lat_max, nlats)
        lons = np.repeat(lon, nlats)

        if isinstance(lon, (list, np.ndarray)):
            nlons = len(lon)
            lons = np.repeat(lon, nlats).reshape(nlons, nlats).T
            lats = np.repeat(lats, nlons).reshape(nlats, nlons)

        return lons, lats

    def subplots(self, *args, bgshow=True, ticks=True,
                 lim=True, grid=True, **kwargs):
        """Create subplots based on map background."""
        fig, ax = plt.subplots(*args, **kwargs)

        return fig, MapAxes(self, ax, bgshow=bgshow, ticks=ticks,
                            lim=lim, grid=grid)

    def figure(self, *args, **kwargs):
        """Create figure based on map background."""
        return self.subplots(*args, **kwargs)

class MapAxis:
    """Map axis object.

    Parameters
    ----------
    bg: Map
        Background Map object.
    ax: matplotlib.axes._subplots.AxesSubplot
        Matplotlib axis.
    bgshow: bool, optional
        Show basemap image (default: ``True``).
    ticks: bool, optional
        Show map X-Y ticks (default: ``True``).
    lim: bool, optional
        Set X-Y limits based on map extent (default: ``True``).
    grid: bool, optional
        Show default image grid (default: ``True``).

    """

    def __init__(self, bg, ax, bgshow=True, ticks=True,
                 lim=True, grid=True):
        self.bg = bg
        self.ax = ax

        if bgshow:
            self.set_bg()
        if ticks:
            self.set_xyticks()
            self.set_xyticklabels()
        if lim:
            self.set_xylim()
        if grid:
            self.grid()

    def __repr__(self):
        return f'<{self.__class__.__name__}> {self.bg}'

    def set_bg(self, cmap='gray', **kwargs):
        """Show background image scaled on the map."""
        self.ax.imshow(self.bg.img, extent=self.bg.extent, cmap=cmap, **kwargs)

    def set_xyticks(self, xticks=None, yticks=None):
        """Set axis X-Y ticks."""
        self.ax.set_xticks(self.bg.lons() if xticks is None else xticks)
        self.ax.set_yticks(self.bg.lats() if yticks is None else yticks)

    def set_xyticklabels(self, lats=[60, 70, 80], lon_w=45,
                         color='lightgray',
                         xticklabels=None, yticklabels=None):
        """Set axis X-Y ticklabels.

        Parameters
        ----------
        lats: list, optional
            Latitudes to display meridian values.
        lon_w: float, optional
            West longitude to display meridian values.

        """
        self.ax.set_xticklabels(
            self.bg.lonlabels() if xticklabels is None else xticklabels)
        self.ax.set_yticklabels(
            self.bg.latlabels() if yticklabels is None else yticklabels)

        if (self.bg._proj == 'stereo'
                and yticklabels is None
                and yticklabels is None):
            for lat in lats:
                self.ax.text(*self.bg(lon_w, lat), slat(lat),
                             rotation=-lon_w, color=color,
                             va='baseline', ha='center')

    def set_xylim(self, bl_lon_w=None, bl_lat=None, tr_lon_w=None, tr_lat=None):
        """Set X-Y axis limits based on map dimensions."""
        if None in [bl_lon_w, bl_lat, tr_lon_w, tr_lat]:
            self.ax.set_xlim(self.bg.xlim)
            self.ax.set_ylim(self.bg.ylim)
        else:
            xmin, ymin = self.bg(bl_lon_w, bl_lat)
            xmax, ymax = self.bg(tr_lon_w, tr_lat)
            self.ax.set_xlim(xmin, xmax)
            self.ax.set_ylim(ymin, ymax)

    def grid(self, lats=[60, 70, 80], lons=[30, 60, 120, 150, 210, 240, 310, 340],
             lat_min=60, lat_max=80, nlats=21,
             color='lightgray', lw=.5, **kwargs):
        """Set image ticks, limits and grid.

        Parameters
        ----------
        lats: list, optional
            List of parallele latitudes in stereographic projection.
        lons: list, optional
            List of meridians west longitudes in stereographic projection.
        lat_min: float, optional
            Minimum meridian latitude in stereographic projection.
        lat_max: float, optional
            Maximum meridian latitude in stereographic projection.
        nlats: int, optional
            Number of latitudes points per meridians in stereographic projection.
        color: str, optional
            Grid color.
        lw: float
            Grid line width.

        """
        self.ax.grid(color=color, lw=lw, **kwargs)

        if self.bg._proj == 'stereo':
            p = self.bg.paralleles(lats)
            m = self.bg.meridians(lons, nlats=nlats,
                                  lat_min=lat_min, lat_max=lat_max)

            self.plot(*p, color=color, lw=lw)
            self.plot(*m, color=color, lw=lw)

    def set_xlabel(self, *args, **kwargs):
        """Set x-label on map axis."""
        self.ax.set_xlabel(*args, **kwargs)

    def set_ylabel(self, *args, **kwargs):
        """Set y-label on map axis."""
        self.ax.set_ylabel(*args, **kwargs)

    def set_title(self, *args, **kwargs):
        """Set title on map axis."""
        self.ax.set_title(*args, **kwargs)

    def plot(self, lon_w, lat, *args, **kwargs):
        """Plot on the map."""
        return self.ax.plot(*self.bg(lon_w, lat), *args, **kwargs)

    def scatter(self, lon_w, lat, *args, **kwargs):
        """Scatter point(s) on the map."""
        return self.ax.scatter(*self.bg(lon_w, lat), *args, **kwargs)

    def text(self, lon_w, lat, txt, *args, **kwargs):
        """Text on the map."""
        self.ax.text(*self.bg(lon_w, lat), txt, *args, **kwargs)

    def add_path(self, path, **kwargs):
        """Add path on the map."""
        self.ax.add_patch(PathPatch(self.bg(path), **kwargs))

    def add_patch(self, patch, **kwargs):
        """Add patch on the map."""
        self.ax.add_patch(self.bg(patch), **kwargs)

    def add_collection(self, collection, **kwargs):
        """Add collection on the map."""
        self.ax.add_collection(self.bg(collection), **kwargs)

class MapAxes:
    """Multi axis map object.

    Parameters
    ----------
    bg: Map
        Background Map object.
    ax: matplotlib.axes._subplots.AxesSubplot or numpy.ndarray
        Matplotlib axis.
    bgshow: bool, optional
        Show basemap image (default: ``True``).
    ticks: bool, optional
        Show map X-Y ticks (default: ``True``).
    lim: bool, optional
        Set X-Y limits based on map extent (default: ``True``).
    grid: bool, optional
        Show default image grid (default: ``True``).

    """

    def __new__(cls, bg, ax, **kwargs):
        if np.ndim(ax) == 0:
            return MapAxis(bg, ax, **kwargs)

        if np.ndim(ax) == 1:
            return np.array([MapAxis(bg, axis, **kwargs) for axis in ax])

        if np.ndim(ax) == 2:
            return np.array([
                [MapAxis(bg, axis, **kwargs) for axis in _ax] for _ax in ax
            ])

        raise ValueError(f'Invalid axis dimension: {np.ndim(ax)}')

max_stars_repo_name: seignovert/pyvims

max_issues_repo_name: seignovert/pyvims

max_forks_repo_name: seignovert/pyvims


text:
from datetime import datetime
from functools import partial
from matplotlib import pyplot as plt
from matplotlib import gridspec as gs
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from operator import itemgetter

from XMC.GlasXC import GlasXC
from XMC.loaders import LibSVMLoader
from XMC.metrics import precision_at_k, ndcg_score_at_k

import math
import os
import yaml
import torch
import numpy as np
import torch.nn.functional as F
import torch.nn as nn

def weights_init(mdl, scheme):
    """
    Function to initialize weights

    Args:
        mdl : Module whose weights are going to modified
        scheme : Scheme to use for weight initialization
    """
    if isinstance(mdl, torch.nn.Linear):
        func = getattr(torch.nn.init, scheme + '_')  # without underscore is deprecated
        func(mdl.weight)

TIME_STAMP = datetime.utcnow().isoformat()

parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)

# data argument
parser.add_argument('--data_root', type=str, required=True,
                    help="""Root folder for dataset.
                            Note that the root folder should contain files either ending with
                            test / train""")
parser.add_argument('--dataset_info', type=str, required=True,
                    help='Dataset information in YAML format')

# architecture arguments
parser.add_argument('--input_encoder_cfg', type=str, required=True,
                    help='Input Encoder architecture configuration in YAML format')
parser.add_argument('--input_decoder_cfg', type=str, required=True,
                    help='Input Decoder architecture configuration in YAML format')
parser.add_argument('--output_encoder_cfg', type=str, required=True,
                    help='Output Encoder architecture configuration in YAML format')
parser.add_argument('--output_decoder_cfg', type=str, required=True,
                    help='Output Decoder architecture configuration in YAML format')
parser.add_argument('--regressor_cfg', type=str, required=True,
                    help='Regressor architecture configuration in YAML format')
parser.add_argument('--init_scheme', type=str, default='default',
                    choices=['xavier_uniform', 'kaiming_uniform', 'default'])

# training configuration arguments
parser.add_argument('--device', type=str, default='cpu',
                    help='PyTorch device string <device_name>:<device_id>')
parser.add_argument('--seed', type=int, default=None,
                    help='Manually set the seed for the experiments for reproducibility')
parser.add_argument('--batch_size', type=int, default=64,
                    help='Batch size for training')
parser.add_argument('--epochs', type=int, default=10,
                    help='Number of epochs to train the autoencoder for')
parser.add_argument('--interval', type=int, default=-1,
                    help='Interval between two status updates on training')
parser.add_argument('--input_ae_loss_weight', type=float, default=1.,
                    help='Weight to give the input autoencoder loss in the entire loss')
parser.add_argument('--output_ae_loss_weight', type=float, default=1.,
                    help='Weight to give the output autoencoder loss in the entire loss')
parser.add_argument('--plot', action='store_true',
                    help='Option to plot the loss variation over iterations')

# optimizer arguments
parser.add_argument('--optimizer_cfg', type=str, required=True,
                    help='Optimizer configuration in YAML format for GlasXC model')

# post training arguments
parser.add_argument('--save_model', type=str, default=None,
                    choices=['all', 'inputAE', 'outputAE', 'regressor'], nargs='+',
                    help='Options to save the model partially or completely')
parser.add_argument('--k', type=int, default=5,
                    help='k for Precision at k and NDCG at k')

# parse the arguments
args = parser.parse_args()

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ CUDA Capability ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

cur_device = torch.device(args.device)
USE_CUDA = cur_device.type == 'cuda'
if USE_CUDA and not torch.cuda.is_available():
    raise ValueError("You can't use CUDA if you don't have CUDA")

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~Check Num of CPU's~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
print("No of threads : ",torch.get_num_threads())

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Reproducibility ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

if args.seed is not None:
    torch.manual_seed(args.seed)
    if USE_CUDA:
        torch.cuda.manual_seed(args.seed)

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Model initialization ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

input_enc_cfg = yaml.load(open(args.input_encoder_cfg))
input_dec_cfg = yaml.load(open(args.input_decoder_cfg))
output_enc_cfg = yaml.load(open(args.output_encoder_cfg))
output_dec_cfg = yaml.load(open(args.output_decoder_cfg))
regress_cfg = yaml.load(open(args.regressor_cfg))

Glas_XC = GlasXC(input_enc_cfg, input_dec_cfg, output_enc_cfg, output_dec_cfg, regress_cfg)
if args.init_scheme != 'default':
    Glas_XC.apply(partial(weights_init, scheme=args.init_scheme))
Glas_XC = Glas_XC.to(cur_device)

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Optimizer initialization ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

opt_options = yaml.load(open(args.optimizer_cfg))
optimizer = getattr(torch.optim, opt_options['name'])(Glas_XC.parameters(),
                                                      **opt_options['args'])

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Dataloader initialization ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

loader_kwargs = {}
if USE_CUDA:
    loader_kwargs = {'num_workers': 1, 'pin_memory': True}

dset_opts = yaml.load(open(args.dataset_info))
USE_TEST_DSET = 'test_filename' in dset_opts.keys()

train_file = os.path.join(args.data_root, dset_opts['train_filename'])
train_loader = LibSVMLoader(train_file, dset_opts['train_opts'])
len_loader = len(train_loader)
train_data_loader = torch.utils.data.DataLoader(train_loader, batch_size=args.batch_size,
                                                shuffle=True, **loader_kwargs)

if USE_TEST_DSET:
    test_file = os.path.join(args.data_root, dset_opts['test_filename'])
    test_loader = LibSVMLoader(test_file, dset_opts['test_opts'])
    test_data_loader = torch.utils.data.DataLoader(test_loader, batch_size=1000,
                                                   shuffle=False)

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Train your model ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

all_iters = 0
ALPHA_INPUT = args.input_ae_loss_weight
ALPHA_OUTPUT = args.output_ae_loss_weight
K = args.k
INP_REC_LOSS = []
OTP_REC_LOSS = []
CLASS_LOSS = []
AVG_P_AT_K = []
AVG_NDCG_AT_K = []
LAMBDA = 10
mean = 0.5

#epsilon = 10e-5                # Use only for Eurlex because Z becomes singular

for epoch in range(args.epochs):
    #print("In epoch ",epoch)
    cur_no = 0
    for x, y in iter(train_data_loader):
        x = x.to(device=cur_device, dtype=torch.float)
        y = y.to(device=cur_device, dtype=torch.float)
        cur_no += x.size(0)
        div = 1/math.pow(y.size(1),2) # different for different datasets
        #print("In epoch ",epoch," size of x is ",x.size())
        #print("In epoch ",epoch," size of y is ",y.size())

        y_sum = np.sum(np.asmatrix(y.numpy()),axis=1)
        for s in range(y_sum.size):
            if y_sum[s] == 0:
                #print("idx : ",s, " sum : ", y_sum[s])
                #y[s] = torch.from_numpy(np.copy(y[s-1].numpy())) # if some sample has no labels, copy the previous sample
                y[s][np.random.randint(y.size(1))] = 1 # setting a random label as 1

        """
        y_sum = np.sum(np.asmatrix(y.numpy()),axis=1)
        for s in range(y_sum.size):
            if y_sum[s] == 0:
                print("idx : ",s, " sum : ", y_sum[s])
        """

        optimizer.zero_grad()

        inp_ae_fp, out_ae_fp, reg_fp, V = Glas_XC.forward(x, y)
        #print("Size of reg_fp is : ", reg_fp.size())
        #print("Size of decoder weight is : ", V.size())
        #print("The first two columns of decoder matrix is :", V[:, :2])
        #print(type(decoder_weight_mat[1,1]))

        # Build GLAS Regularizer

        # Sampling the Label Matrix - Start
        #print("Size of y is : ", y.size())
        count_labels = np.sum(np.asmatrix(y),axis=1)
        #print(count_labels)
        label_dict = {}

        for i in range(y.size(0)):
            label_dict[i] = int(count_labels[i]);

        #print(label_dict)

        label_dict_sorted = sorted(label_dict.items(), key = itemgetter(1))
        #{k: v for k, v in sorted(label_dict.items(), key=lambda item: item[1])}
        #print(label_dict_sorted)
        #print(label_dict_sorted[3][1])

        sorted_samples = torch.empty((y.size(0)),dtype=torch.int)
        for i in range(y.size(0)):
            sorted_samples[i] = label_dict_sorted[i][0]

        #print(sorted_samples)
        sampled_labels = []
        for i in range(y.size(0)):
            if len(sampled_labels) >= args.batch_size:
                break
            goto_next = False
            #while(not (goto_next)):
            ty = y[sorted_samples[i],:].numpy()
            #print(ty)
            valid_idx = np.nonzero(ty)  # get all the non zeros
            #print("Valid index is : ",valid_idx)
            #print("First value of valid index is : ",valid_idx[0])
            #t = torch.flatten(valid_idx)  # flatten the tuple into a list tensor
            #print("i : ", i ,"len valid_idx : ", len(valid_idx[0]), " len(sampled_labels) : ", len(sampled_labels) )
            s = 0
            while(not(goto_next) and s < len(valid_idx[0])) :
                rand_idx = np.random.randint(len(valid_idx[0]))
                #sampled_label = np.array([np.random.choice(valid_idx[0][idx],1) for idx in range(len(valid_idx[0]))])
                sampled_element = valid_idx[0][rand_idx]
                #print("i",i," Sampled element is : ",sampled_element)
                #sampled_labels.append()
                #print("sampled_element : ", sampled_element)
                s += 1;
                if sampled_element not in sampled_labels:
                    sampled_labels.append(sampled_element)
                    s = len(valid_idx[0]) # set while loop condition
                    goto_next = True # set while loopp condition
            #print("sampled_labels[",len(sampled_labels)-1,"] :", sampled_labels[len(sampled_labels)-1], " len(sampled_labels) : ", len(sampled_labels))

        #sampled_labels = np.asarray(sampled_labels)
        #print("sampled_labels : ", sampled_labels)
        #print("size of sampled_labels : ", len(sampled_labels))
        #print("size of unique sampled_labels : ", len(sampled_labels.unique()))

        sampled_labels = torch.from_numpy(np.asarray(sampled_labels))
        y_sampled = torch.index_select(y, 1, sampled_labels)  #indexes the input tensor along column using the entries in indices

        #sample_sum = np.sum(np.asmatrix(y_sampled.numpy()),axis=1)
        #for s in range(sample_sum.size):
            #if sample_sum[s] == 0:
            #print("s : ",s," sample_sum[s",s,"] = ", sample_sum[s])

        #print("Size of sampled y is : ", y_sampled.size())
        #print("Rank of sampled y is : ", torch.matrix_rank(y_sampled))
        #print(np.sum(np.asmatrix(y_sampled),axis=0))
        #print("Size of output decoder  matrix is : ", V.size())

        V_sampled = torch.index_select(V, 1, sampled_labels)  #indexes the input tensor along column using the entries in indices
        VtV  = torch.mm(V_sampled.t(), V_sampled)               # Label Embedding Matrix
        A  = torch.mm(y_sampled.t(), y_sampled)  			  # models co-occurence of labels
        #print(torch.nonzero(A[2,:]))
        #print("Size of yty is : ", A.size())
        #print(np.sum(np.asmatrix(A),axis=1))
        #print(torch.matrix_rank(A))
        #inp_ae_fp, out_ae_fp, reg_fp = Glas_XC.forward(x, y)
      #  reg_fip = torch.index_select(reg_fp, 1, sampled_labels)
        # Build GLAS Regularizer

        # Sampling the Label matrix per batch done!

        #v  = Glas_XC.encode_output(y)    # Label Embedding Matrix for mini-batch
        #V  = torch.mm(v.t(), v)               # co-occurence in the latent/embedded space
        #A  = torch.mm(y.t(), y)               # models co-occurence of labels

        Z  = torch.diag(A)  #+ epsilon        # returns the diagoan in vector form
        Z  = torch.diag(Z)                    # creates the diagonal from the vector
        #AZ = torch.add(torch.mm(A, torch.pinverse(Z)), torch.mm(torch.pinverse(Z), A)) # to be used for Eurlex4k
        AZ = torch.add(torch.mm(A, torch.inverse(Z)), torch.mm(torch.inverse(Z), A))
        M  = mean*AZ                          # Mean of conditional frequencies of label
        g  = torch.sub(VtV, M)
        gl = torch.norm(g, p='fro')
        loss_glas = div * gl*gl                  # final loss of glas regularizer

        #print("loss glas regularizer created")
        #print("Epoch : ", epoch)

       # print("The size of encoded output label is")
       # print(V.size())

       #	reg_fp = reg_fp.numpy()

       # for i in range(0, args.batch_size):
        #	reg_fp[i] = [1 if ele > 0.5 else 0 for ele in reg_fp[i]]

        # This is a custom that we will be using to backprop. It has three components:
        # 1. Reconstruction error of input
        # 2. Reconstruction error of output
        # 3. Classification (Binary cross entropy) of input-output
        # The first two are weighted using ALPHA_INPUT and ALPHA_OUTPUT.
        #print("The size of the reg_fp matrix is \n ", reg_fp.size())
        #print("The size of the y matrix is \n", y.size())

       # print("The first two rows of prediction score reg_fp matrix is :", reg_fp[:2, :])
       # loss_class = F.multilabel_margin_loss(reg_fp, y) + LAMBDA * loss_glas
       #	th = torch.Tensor([0.5])
       #	reg_fp1 = (reg_fp > th).float()*1
        loss_class = F.binary_cross_entropy(reg_fp, y) + LAMBDA * loss_glas
        net_loss = loss_class
        net_loss.backward()
        #print("Backprop for epoch ", epoch, " done")
        optimizer.step()
        all_iters += 1
        if all_iters % args.interval == 0:
            print("{} / {} :: {} / {} - CLASS_LOSS : {}"
                  .format(epoch, args.epochs, cur_no, len_loader,round(loss_class.item(), 5)))

        CLASS_LOSS.append(loss_class.item())

    pred_y = []
    actual_y = []
    for x, y in iter(train_data_loader):
        x = x.to(device=cur_device, dtype=torch.float)
        y = y.to(device=cur_device, dtype=torch.float)
        pred_y.append(Glas_XC.predict(x).detach().cpu().numpy())
        actual_y.append(y.numpy())

    pred_y = np.vstack(pred_y)
   # print("The first few entries of pred_y - ", pred_y[0:5])
    actual_y = np.vstack(actual_y)
   # print("The first few entries of actual_y - ", actual_y[0:5])
    p_at_k = [precision_at_k(actual_y[i], pred_y[i], K) for i in range(len(pred_y))]
    ndcg_at_k = [ndcg_score_at_k(actual_y[i], pred_y[i], K) for i in range(len(pred_y))]
    print("{0} / {1} :: Precision at {2}: {3}\tNDCG at {2}: {4}"
          .format(epoch, args.epochs, K, np.mean(p_at_k), np.mean(ndcg_at_k)))
    AVG_P_AT_K.append(np.mean(p_at_k))
    AVG_NDCG_AT_K.append(np.mean(ndcg_at_k))

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Plot graphs ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
if args.plot:
    fig = plt.figure(figsize=(9, 18))
    gridspec = gs.GridSpec(4, 6, figure=fig)
    gridspec.tight_layout(fig)
    ax1 = plt.subplot(gridspec[0, :2])
    ax2 = plt.subplot(gridspec[0, 2:4])
    ax3 = plt.subplot(gridspec[0, 4:])
    #ax4 = plt.subplot(gridspec[1:3, 1:5])
    #ax5 = plt.subplot(gridspec[3, :3])
    #ax6 = plt.subplot(gridspec[3, 3:])

    ax1.plot(list(range(1, all_iters + 1)), CLASS_LOSS, 'b', linewidth=2.0)
    ax1.set_title('Classification loss')
    ax2.plot(list(range(1, args.epochs + 1)), AVG_P_AT_K, 'g', linewidth=2.0)
    ax2.set_title('Average Precision at {} (over all datapoints) with epochs'.format(K))
    ax3.plot(list(range(1, args.epochs + 1)), AVG_NDCG_AT_K, 'b', linewidth=2.0)
    ax3.set_title('Average NDCG at {} (over all datapoints) with epochs'.format(K))
    #plt.show()
    plt.savefig('prec_plots.png')

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Save your model ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

if args.save_model is not None:
    if 'inputAE' in args.save_model or 'all' in args.save_model:
        torch.save(Glas_XC.input_encoder.to('cpu'),
                   'trained_input_encoder_{}.pt'.format(TIME_STAMP))
        torch.save(Glas_XC.input_decoder.to('cpu'),
                   'trained_input_decoder_{}.pt'.format(TIME_STAMP))

    if 'outputAE' in args.save_model or 'all' in args.save_model:
        torch.save(Glas_XC.output_encoder.to('cpu'),
                   'trained_output_encoder_{}.pt'.format(TIME_STAMP))
        torch.save(Glas_XC.output_decoder.to('cpu'),
                   'trained_output_decoder_{}.pt'.format(TIME_STAMP))

    if 'regressor' in args.save_model or 'all' in args.save_model:
        torch.save(Glas_XC.regressor.to('cpu'),
                   'trained_regressor_{}.pt'.format(TIME_STAMP))

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Prediction on test ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

if USE_TEST_DSET:
    print("Test set characteristics")
    pred_y = []
    actual_y = []
    for x, y in iter(test_data_loader):
        x = x.to(device=cur_device, dtype=torch.float)

        pred_y.append(Glas_XC.predict(x).detach().cpu().numpy())
        actual_y.append(y.numpy())

    pred_y = np.vstack(pred_y)
    actual_y = np.vstack(actual_y)
    p_at_k = [precision_at_k(actual_y[i], pred_y[i], K) for i in range(len(pred_y))]
    ndcg_at_k = [ndcg_score_at_k(actual_y[i], pred_y[i], K) for i in range(len(pred_y))]
    print("Precision at {2}: {3}\tNDCG at {2}: {4}"
          .format(epoch, args.epochs, K, np.mean(p_at_k), np.mean(ndcg_at_k)))

ext: py

max_stars_repo_name: siddhantkatyan/GlasXC

max_issues_repo_name: pyschedelicsid/GlasXC

max_forks_repo_name: pyschedelicsid/GlasXC

reason: import numpy


text:
# Licensed under a 3-clause BSD style license - see LICENSE.rst
import pytest
import numpy as np
from numpy.testing import assert_allclose, assert_equal

from astropy import units as u
from astropy.coordinates import Longitude, Latitude, Galactic, SkyCoord

from ..high_level import HEALPix

class TestHEALPix:

    def setup_class(self):
        self.pix = HEALPix(nside=256, order='nested')

    def test_pixel_area(self):
        pixel_area = self.pix.pixel_area
        assert_allclose(pixel_area.value, 1.5978966540475428e-05)
        assert pixel_area.unit == u.sr

    def test_pixel_resolution(self):
        pixel_resolution = self.pix.pixel_resolution
        assert_allclose(pixel_resolution.value, 13.741945647269624)
        assert pixel_resolution.unit == u.arcmin

    def test_level(self):
        assert self.pix.level == 8

    def test_npix(self):
        assert self.pix.npix == 12 * 256 ** 2

    # For the following tests, the numerical accuracy of this function is
    # already tested in test_cython_api.py, so we focus here on functionality
    # specific to the high-level functions.

    def test_healpix_to_lonlat(self):
        lon, lat = self.pix.healpix_to_lonlat([1, 2, 3])

        assert isinstance(lon, Longitude)
        assert isinstance(lat, Latitude)

        index = self.pix.lonlat_to_healpix(lon, lat)

        assert_equal(index, [1, 2, 3])

        lon, lat = self.pix.healpix_to_lonlat([1, 2, 3],
                                              dx=[0.1, 0.2, 0.3],
                                              dy=[0.5, 0.4, 0.7])

        assert isinstance(lon, Longitude)
        assert isinstance(lat, Latitude)

        index, dx, dy = self.pix.lonlat_to_healpix(lon, lat, return_offsets=True)

        assert_equal(index, [1, 2, 3])
        assert_allclose(dx, [0.1, 0.2, 0.3])
        assert_allclose(dy, [0.5, 0.4, 0.7])

    def test_healpix_to_xyz(self):
        x, y, z = self.pix.healpix_to_xyz([1, 2, 3])

        assert isinstance(x, np.ndarray)
        assert isinstance(y, np.ndarray)
        assert isinstance(z, np.ndarray)

        index = self.pix.xyz_to_healpix(x, y, z)

        assert_equal(index, [1, 2, 3])

        x, y, z = self.pix.healpix_to_xyz([1, 2, 3],
                                          dx=[0.1, 0.2, 0.3],
                                          dy=[0.5, 0.4, 0.7])

        assert isinstance(x, np.ndarray)
        assert isinstance(y, np.ndarray)
        assert isinstance(z, np.ndarray)

        index, dx, dy = self.pix.xyz_to_healpix(x, y, z, return_offsets=True)

        assert_equal(index, [1, 2, 3])
        assert_allclose(dx, [0.1, 0.2, 0.3])
        assert_allclose(dy, [0.5, 0.4, 0.7])

    def test_nested_to_ring(self):
        nested_index_1 = [1, 3, 22]
        ring_index = self.pix.nested_to_ring(nested_index_1)
        nested_index_2 = self.pix.ring_to_nested(ring_index)
        assert_equal(nested_index_1, nested_index_2)

    def test_bilinear_interpolation_weights(self):
        indices, weights = self.pix.bilinear_interpolation_weights([1, 3, 4] * u.deg,
                                                                   [3, 2, 6] * u.deg)
        assert indices.shape == (4, 3)
        assert weights.shape == (4, 3)

    def test_interpolate_bilinear_lonlat(self):
        values = np.ones(12 * 256 ** 2) * 3
        result = self.pix.interpolate_bilinear_lonlat([1, 3, 4] * u.deg,
                                                      [3, 2, 6] * u.deg, values)
        assert_allclose(result, [3, 3, 3])

    def test_interpolate_bilinear_lonlat_invalid(self):
        values = np.ones(222) * 3
        with pytest.raises(ValueError) as exc:
            self.pix.interpolate_bilinear_lonlat([1, 3, 4] * u.deg,
                                                 [3, 2, 6] * u.deg, values)
        assert exc.value.args[0] == 'values must be an array of length 786432 (got 222)'

    def test_cone_search_lonlat(self):
        lon, lat = 1 * u.deg, 4 * u.deg
        result = self.pix.cone_search_lonlat(lon, lat, 1 * u.deg)
        assert len(result) == 77

    def test_cone_search_lonlat_invalid(self):
        lon, lat = [1, 2] * u.deg, [3, 4] * u.deg
        with pytest.raises(ValueError) as exc:
            self.pix.cone_search_lonlat(lon, lat, 1 * u.deg)
        assert exc.value.args[0] == ('The longitude, latitude and radius must '
                                     'be scalar Quantity objects')

    def test_boundaries_lonlat(self):
        lon, lat = self.pix.boundaries_lonlat([10, 20, 30], 4)
        assert lon.shape == (3, 16)
        assert lat.shape == (3, 16)

    def test_neighbours(self):
        neigh = self.pix.neighbours([10, 20, 30])
        assert neigh.shape == (8, 3)

class TestCelestialHEALPix:

    def setup_class(self):
        self.pix = HEALPix(nside=256, order='nested', frame=Galactic())

    def test_healpix_from_header(self):
        """Test instantiation from a FITS header.

        Notes
        -----
        We don't need to test all possible options, because
        :meth:`~astropy_healpix.HEALPix.from_header` is just a wrapper around
        :meth:`~astropy_healpix.utils.parse_input_healpix_data`, which is
        tested exhaustively in :mod:`~astropy_healpix.tests.test_utils`.
        """

        pix = HEALPix.from_header(
            (np.empty(self.pix.npix), 'G'),
            nested=self.pix.order == 'nested')

        assert pix.nside == self.pix.nside
        assert type(pix.frame) == type(self.pix.frame)  # noqa
        assert pix.order == self.pix.order

    def test_healpix_to_skycoord(self):
        coord = self.pix.healpix_to_skycoord([1, 2, 3])

        assert isinstance(coord, SkyCoord)
        assert isinstance(coord.frame, Galactic)

        # Make sure that the skycoord_to_healpix method converts coordinates
        # to the frame of the HEALPix
        coord = coord.transform_to('fk5')

        index = self.pix.skycoord_to_healpix(coord)

        assert_equal(index, [1, 2, 3])

        coord = self.pix.healpix_to_skycoord([1, 2, 3],
                                             dx=[0.1, 0.2, 0.3],
                                             dy=[0.5, 0.4, 0.7])

        assert isinstance(coord, SkyCoord)
        assert isinstance(coord.frame, Galactic)

        # Make sure that the skycoord_to_healpix method converts coordinates
        # to the frame of the HEALPix
        coord = coord.transform_to('fk5')

        index, dx, dy = self.pix.skycoord_to_healpix(coord, return_offsets=True)

        assert_equal(index, [1, 2, 3])
        assert_allclose(dx, [0.1, 0.2, 0.3])
        assert_allclose(dy, [0.5, 0.4, 0.7])

    def test_interpolate_bilinear_skycoord(self):
        values = np.ones(12 * 256 ** 2) * 3
        coord = SkyCoord([1, 2, 3] * u.deg, [4, 3, 1] * u.deg, frame='fk4')
        result = self.pix.interpolate_bilinear_skycoord(coord, values)
        assert_allclose(result, [3, 3, 3])

        # Make sure that coordinate system is correctly taken into account

        values = np.arange(12 * 256 ** 2) * 3
        coord = SkyCoord([1, 2, 3] * u.deg, [4, 3, 1] * u.deg, frame='fk4')

        result1 = self.pix.interpolate_bilinear_skycoord(coord, values)
        result2 = self.pix.interpolate_bilinear_skycoord(coord.icrs, values)

        assert_allclose(result1, result2)

    def test_cone_search_skycoord(self):
        coord = SkyCoord(1 * u.deg, 4 * u.deg, frame='galactic')
        result1 = self.pix.cone_search_skycoord(coord, 1 * u.deg)
        assert len(result1) == 77
        result2 = self.pix.cone_search_skycoord(coord.icrs, 1 * u.deg)
        assert_allclose(result1, result2)

    def test_boundaries_skycoord(self):
        coord = self.pix.boundaries_skycoord([10, 20, 30], 4)
        assert coord.shape == (3, 16)

class TestCelestialHEALPixFrameAsClass(TestCelestialHEALPix):

    def setup_class(self):
        self.pix = HEALPix(nside=256, order='nested', frame=Galactic)

class TestCelestialHEALPixFrameAsString(TestCelestialHEALPix):

    def setup_class(self):
        self.pix = HEALPix(nside=256, order='nested', frame='galactic')

def test_invalid_frame_name():
    with pytest.raises(ValueError, match='Coordinate frame name "foobar"'):
        HEALPix(nside=256, frame='foobar')

def test_invalid_frame_type():
    with pytest.raises(ValueError, match='Coordinate frame must be a'):
        HEALPix(nside=256, frame=('obviously', 'not', 'a', 'frame'))

max_stars_repo_name: astropy/healpix

max_issues_repo_name: nstarman/astropy-healpix

max_forks_repo_name: nstarman/astropy-healpix

reason: import numpy,from numpy,from astropy


text:
import numpy as np
import tensorflow as tf
from sklearn.datasets import make_moons

# dataset
num_attributes = 2
num_classes = 2

examples, classes = make_moons(256)
examples += np.random.normal(scale=5e-2, size=np.shape(examples))

# hyperparameters
num_must = 12
num_cannot = 12
neighborhood_size = 64
auxiliary_weight = 1e-2
gamma = 1 # parameter for RBF
regularizer = tf.contrib.layers.l2_regularizer(1e-2)
batch_size = len(examples) # examples
min_batch_must = 0 # pairs of examples
min_batch_cannot = 0 # pairs of examples
learning_rate = 1e-4
num_episodes = 4096

class model:
    input = tf.placeholder(tf.float32, shape=(None, num_attributes))
    layers = [input]
    layers.append(tf.layers.dense(inputs=layers[-1], units=16, activation=tf.nn.relu, kernel_regularizer=regularizer))
    layers.append(tf.layers.dense(inputs=layers[-1], units=4, activation=tf.nn.relu, kernel_regularizer=regularizer))
    output = tf.layers.dense(inputs=layers[-1], units=num_classes, activation=tf.nn.softmax, kernel_regularizer=regularizer)
    layers.append(output)

def visualize(show=True, res=(64, 64)):
    import matplotlib.pyplot as plt
    plt.figure()
    plt.title('classification space')
    from d_graph_plus.manager import predict
    from d_graph_plus.data import must_link, cannot_link

    # calculate the box area to plot on
    left = np.min(examples[:, 0])
    right = np.max(examples[:, 0])
    bottom = np.min(examples[:, 1])
    top = np.max(examples[:, 1])
    half_pixel_x = (right - left) / res[1] / 2
    half_pixel_y = (top - bottom) / res[0] / 2
    extent = (left - half_pixel_x, right + half_pixel_x, bottom - half_pixel_y, top + half_pixel_y)

    # create background
    grid_x, grid_y = np.meshgrid(np.linspace(top, bottom, res[0]), np.linspace(left, right, res[1]))
    grid = np.array([grid_x.flatten(), grid_y.flatten()]).T
    predictions = predict(grid)
    bg = np.reshape([(pred[0], 1.0, pred[1]) for pred in predictions], (res[0], res[1], 3))

    # plot and show
    plt.imshow(bg, zorder=0, interpolation='gaussian', extent=extent)
    for pair in must_link:
        pair_examples = examples[np.array(list(pair))]
        plt.plot(pair_examples[:, 0], pair_examples[:, 1], c='k', zorder=1)
    for pair in cannot_link:
        pair_examples = examples[np.array(list(pair))]
        plt.plot(pair_examples[:, 0], pair_examples[:, 1], c='w', zorder=1)
    plt.scatter(examples[classes == 0][:, 0], examples[classes == 0][:, 1], c='r', zorder=2)
    plt.scatter(examples[classes == 1][:, 0], examples[classes == 1][:, 1], c='b', zorder=2)
    if show:
        plt.show()

max_stars_repo_name: malyvsen/d-graph-plus

max_issues_repo_name: malyvsen/d-graph-plus

max_forks_repo_name: malyvsen/d-graph-plus


text:
"""Physics-related helper functions."""

# external
import numpy as np

def snell_angle_2(angle_1, n_1, n_2):
    """Calculate the angle of refraction of a ray travelling between two mediums
    according to Snell's law.

    Args:     angle_1 (array_like[float]): angle of incidence with respect to surface
    normal         in radians.     n_1 (float): index of refraction in first medium.
    n_2 (float): index of refraction in second medium.  Returns:     float: angle of
    refraction in radians.

    """
    angle_2 = np.arcsin(n_1 / n_2 * np.sin(angle_1))

    return angle_2

max_stars_repo_name: spacesys-finch/payload-designer

max_issues_repo_name: spacesys-finch/payload-designer

max_forks_repo_name: spacesys-finch/payload-designer


text:
from typing import Any, Dict, Optional, Sequence, Tuple, Union

import numpy as np
import torch
from torch import nn

from tianshou.utils.net.discrete import NoisyLinear

class DQN(nn.Module):
    """Reference: Human-level control through deep reinforcement learning.
    For advanced usage (how to customize the network), please refer to
    :ref:`build_the_network`.
    """

    def __init__(
        self,
        c: int,
        h: int,
        w: int,
        action_shape: Sequence[int],
        device: Union[str, int, torch.device] = "cpu",
        features_only: bool = False,
        output_dim: Optional[int] = None,
    ) -> None:
        super().__init__()
        self.device = device
        self.net = nn.Sequential(
            nn.Conv2d(c, 32, kernel_size=8, stride=4), nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=4, stride=2), nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1), nn.ReLU(inplace=True),
            nn.Flatten()
        )
        with torch.no_grad():
            self.output_dim = np.prod(self.net(torch.zeros(1, c, h, w)).shape[1:])
        if not features_only:
            self.net = nn.Sequential(
                self.net, nn.Linear(self.output_dim, 512), nn.ReLU(inplace=True),
                nn.Linear(512, np.prod(action_shape))
            )
            self.output_dim = np.prod(action_shape)
        elif output_dim is not None:
            self.net = nn.Sequential(
                self.net, nn.Linear(self.output_dim, output_dim),
                nn.ReLU(inplace=True)
            )
            self.output_dim = output_dim

    def forward(
        self,
        obs: Union[np.ndarray, torch.Tensor],
        state: Optional[Any] = None,
        info: Dict[str, Any] = {},
    ) -> Tuple[torch.Tensor, Any]:
        r"""Mapping: s -> Q(s, \*)."""
        obs = torch.as_tensor(obs, device=self.device, dtype=torch.float32)
        return self.net(obs), state

class C51(DQN):
    """Reference: A distributional perspective on reinforcement learning.
    For advanced usage (how to customize the network), please refer to
    :ref:`build_the_network`.
    """

    def __init__(
        self,
        c: int,
        h: int,
        w: int,
        action_shape: Sequence[int],
        num_atoms: int = 51,
        device: Union[str, int, torch.device] = "cpu",
    ) -> None:
        self.action_num = np.prod(action_shape)
        super().__init__(c, h, w, [self.action_num * num_atoms], device)
        self.num_atoms = num_atoms

    def forward(
        self,
        obs: Union[np.ndarray, torch.Tensor],
        state: Optional[Any] = None,
        info: Dict[str, Any] = {},
    ) -> Tuple[torch.Tensor, Any]:
        r"""Mapping: x -> Z(x, \*)."""
        obs, state = super().forward(obs)
        obs = obs.view(-1, self.num_atoms).softmax(dim=-1)
        obs = obs.view(-1, self.action_num, self.num_atoms)
        return obs, state

class Rainbow(DQN):
    """Reference: Rainbow: Combining Improvements in Deep Reinforcement Learning.
    For advanced usage (how to customize the network), please refer to
    :ref:`build_the_network`.
    """

    def __init__(
        self,
        c: int,
        h: int,
        w: int,
        action_shape: Sequence[int],
        num_atoms: int = 51,
        noisy_std: float = 0.5,
        device: Union[str, int, torch.device] = "cpu",
        is_dueling: bool = True,
        is_noisy: bool = True,
    ) -> None:
        super().__init__(c, h, w, action_shape, device, features_only=True)
        self.action_num = np.prod(action_shape)
        self.num_atoms = num_atoms

        def linear(x, y):
            if is_noisy:
                return NoisyLinear(x, y, noisy_std)
            else:
                return nn.Linear(x, y)

        self.Q = nn.Sequential(
            linear(self.output_dim, 512), nn.ReLU(inplace=True),
            linear(512, self.action_num * self.num_atoms)
        )
        self._is_dueling = is_dueling
        if self._is_dueling:
            self.V = nn.Sequential(
                linear(self.output_dim, 512), nn.ReLU(inplace=True),
                linear(512, self.num_atoms)
            )
        self.output_dim = self.action_num * self.num_atoms

    def forward(
        self,
        obs: Union[np.ndarray, torch.Tensor],
        state: Optional[Any] = None,
        info: Dict[str, Any] = {},
    ) -> Tuple[torch.Tensor, Any]:
        r"""Mapping: x -> Z(x, \*)."""
        obs, state = super().forward(obs)
        q = self.Q(obs)
        q = q.view(-1, self.action_num, self.num_atoms)
        if self._is_dueling:
            v = self.V(obs)
            v = v.view(-1, 1, self.num_atoms)
            logits = q - q.mean(dim=1, keepdim=True) + v
        else:
            logits = q
        probs = logits.softmax(dim=2)
        return probs, state

class QRDQN(DQN):
    """Reference: Distributional Reinforcement Learning with Quantile \
    Regression.
    For advanced usage (how to customize the network), please refer to
    :ref:`build_the_network`.
    """

    def __init__(
        self,
        c: int,
        h: int,
        w: int,
        action_shape: Sequence[int],
        num_quantiles: int = 200,
        device: Union[str, int, torch.device] = "cpu",
    ) -> None:
        self.action_num = np.prod(action_shape)
        super().__init__(c, h, w, [self.action_num * num_quantiles], device)
        self.num_quantiles = num_quantiles

    def forward(
        self,
        obs: Union[np.ndarray, torch.Tensor],
        state: Optional[Any] = None,
        info: Dict[str, Any] = {},
    ) -> Tuple[torch.Tensor, Any]:
        r"""Mapping: x -> Z(x, \*)."""
        obs, state = super().forward(obs)
        obs = obs.view(-1, self.action_num, self.num_quantiles)
        return obs, state

max_stars_repo_name: NickKaparinos/OpenAI-Gym-Projects

max_issues_repo_name: NickKaparinos/OpenAI-Gym-Projects

max_forks_repo_name: NickKaparinos/OpenAI-Gym-Projects


text:
import numpy as np
import matplotlib.pyplot as plt
import sys

sys.path.append('../src')

import Steppers as Step
import Fluxes as Flux
from PyRsw import Simulation
from constants import minute, hour, day

def test():

    sim = Simulation()

    # Geometry and Model Equations
    sim.geomy       = 'periodic'
    sim.stepper     = Step.Euler
    sim.method      = 'Spectral'
    sim.dynamics    = 'Nonlinear'
    sim.flux_method = Flux.spectral_sw

    # Specify paramters
    sim.Ly  = 4000e3
    sim.Ny  = 128
    sim.cfl = 0.5
    sim.Hs  = [100.]
    sim.rho = [1025.]
    sim.end_time = 5.*minute

    # Plotting parameters
    sim.animate = 'None'
    sim.output = False
    sim.diagnose = False

    # Initialize the grid and zero solutions
    sim.initialize()

    for ii in range(sim.Nz):
        sim.soln.h[:,:,ii] = sim.Hs[ii]

    # Gaussian initial conditions
    x0 = 1.*sim.Lx/2.
    W  = 200.e3
    amp = 1.
    sim.soln.h[:,:,0] += amp*np.exp(-(sim.Y)**2/(W**2))

    sim.run()

max_stars_repo_name: nishantsule/PyRsw

max_issues_repo_name: nishantsule/PyRsw

max_forks_repo_name: nishantsule/PyRsw


text:
"""Utilities for neural networks."""
import copy
import os

import numpy as np
import torch
import torch.jit
import torch.nn as nn

def deep_copy_module(module):
    """Deep copy a module."""
    if isinstance(module, torch.jit.ScriptModule):
        module.save(module.original_name)
        out = torch.jit.load(module.original_name)
        os.system(f"rm {module.original_name}")
        return out
    return copy.deepcopy(module)

class Swish(nn.Module):
    """Swish activation function.

    The swish activation function is given by:
        f(x) = x * sigmoid(x)

    References
    ----------
    Ramachandran, P., Zoph, B., & Le, Q. V. (2017).
    Swish: a self-gated activation function. arXiv.
    """

    def forward(self, x):
        """Apply forward computation of the module."""
        return x * torch.sigmoid(x)

class Mish(nn.Module):
    """Mish activation function.

    The swish activation function is given by:
        f(x) = x * tanh(softplus(x))

    References
    ----------
    Misra, D. (2019).
    Mish: A self regularized non-monotonic neural activation function. arXiv.
    """

    def forward(self, x):
        """Apply forward computation of module."""
        return x * (torch.tanh(torch.nn.functional.softplus(x)))

class View(nn.Module):
    """View Layer.

    The  view layer flattens a tensor.
    It is useful for changing from convolutions to FC layers.
    """

    def forward(self, x):
        """Apply forward computation of module."""
        return x.view(x.size(0), -1)

def parse_nonlinearity(non_linearity):
    """Parse non-linearity."""
    if hasattr(nn, non_linearity):
        return getattr(nn, non_linearity)
    elif hasattr(nn, non_linearity.capitalize()):
        return getattr(nn, non_linearity.capitalize())
    elif hasattr(nn, non_linearity.upper()):
        return getattr(nn, non_linearity.upper())
    elif non_linearity.lower() == "swish":
        return Swish
    elif non_linearity.lower() == "mish":
        return Mish
    else:
        raise NotImplementedError(f"non-linearity {non_linearity} not implemented")

def parse_layers(layers, in_dim, non_linearity):
    """Parse layers of nn."""
    nonlinearity = parse_nonlinearity(non_linearity)
    layers_ = list()
    if len(in_dim) > 1:  # Convolutional Layers.
        w = min(in_dim[1], in_dim[2])
        k0 = max(2 * int(np.floor(np.sqrt(w) / 2)), 2)
        layers_.append(
            nn.Conv2d(
                in_channels=in_dim[0],
                out_channels=32,
                kernel_size=k0,
                stride=max(k0 // 2, 1),
            )
        )
        layers_.append(nn.BatchNorm2d(32))
        layers_.append(nn.ReLU())

        layers_.append(
            nn.Conv2d(
                in_channels=32,
                out_channels=64,
                kernel_size=max(k0 // 2, 2),
                stride=max(k0 // 4, 1),
            )
        )
        layers_.append(nn.BatchNorm2d(64))
        layers_.append(nn.ReLU())

        layers_.append(
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1)
        )
        layers_.append(nn.BatchNorm2d(64))
        layers_.append(nn.ReLU())
        layers_.append(View())

        x = torch.randn(in_dim).unsqueeze(0)
        for module in layers_:
            x = module(x)
        in_dim = x.shape[-1]
    else:
        in_dim = in_dim[0]

    for layer in layers:
        layers_.append(nn.Linear(in_dim, layer))
        layers_.append(nonlinearity())
        in_dim = layer

    return nn.Sequential(*layers_), in_dim

def update_parameters(target_module, new_module, tau=0.0):
    """Update the parameters of target_params by those of new_params (softly).

    The parameters of target_nn are replaced by:
        target_params <- (1-tau) * (target_params) + tau * (new_params)

    Parameters
    ----------
    target_module: nn.Module
    new_module: nn.Module
    tau: float, optional

    Returns
    -------
    None.
    """
    with torch.no_grad():
        target_state_dict = target_module.state_dict()
        new_state_dict = new_module.state_dict()

        for name in target_state_dict.keys():
            if target_state_dict[name] is new_state_dict[name]:
                continue
            else:
                if target_state_dict[name].data.ndim == 0:
                    target_state_dict[name].data = new_state_dict[name].data
                else:
                    target_state_dict[name].data[:] = (
                        tau * target_state_dict[name].data
                        + (1 - tau) * new_state_dict[name].data
                    )

        # It is not necessary to load the dict again as it modifies the pointer.
        # target_module.load_state_dict(target_state_dict)

def update_module(target_module, new_modules, tau=0.0):
    """Update the parameters of target_params by those of new_params (softly).

    The parameters of target_nn are replaced by:
        target_params <- (1-tau) * (target_params) + tau * (new_params)

    Parameters
    ----------
    target_module: nn.Module
    new_module: nn.Module
    tau: float, optional

    Returns
    -------
    None.
    """
    with torch.no_grad():
        target_state_dict = target_module.state_dict()
        new_state_dict = [new_module.state_dict() for new_module in new_modules]

        for name in target_state_dict.keys():
            if target_state_dict[name].data.ndim == 0:
                temp = [new_state_dict_head[name].data for new_state_dict_head in new_state_dict]
                target_state_dict[name].data = temp.mean()
            else:
                temp = [new_state_dict_head[name].data for new_state_dict_head in new_state_dict]
                temp = torch.stack(temp, dim=-1)
                target_state_dict[name].data[:] = torch.mean(temp, -1)

def count_vars(module):
    """Count the number of variables in a module."""
    return sum([np.prod(p.shape) for p in module.parameters()])

def zero_bias(module):
    """Zero all bias terms in the parameters.

    Parameters
    ----------
    module: module to zero the biases to.

    """
    for name, param in module.named_parameters():
        if name.endswith("bias"):
            nn.init.zeros_(param)

def init_head_bias(module, offset=0.0, delta=0.1):
    """Initialize the bias sampling u.a.r. from [offset - delta, offset + delta].

    This is useful for optimistic initialization of value functions.

    Parameters
    ----------
    module: nn.Module.
        Module to initialize head weights.
    offset: float.
        Mean of the bias.
    delta: float.
        Amplitude of bias.
    """
    for name, param in module.named_parameters():
        if name.endswith("head.bias"):
            torch.nn.init.uniform_(param, offset - delta, offset + delta)

def init_head_weight(module, mean_weight=0.1, scale_weight=0.01):
    """Initialize the head of a NN.

    The mean weights are sampled u.a.r from [-mean_weight, mean_weight].
    The scale weights are sampled u.a.r from [-scale_weight, scale_weight].

    Parameters
    ----------
    module: nn.Module.
        Module to initialize head weights.
    mean_weight: float.
        Amplitude of mean weights.
    scale_weight: float.
        Amplitude of scale weights.
    """
    for name, param in module.named_parameters():
        if name.endswith("head.weight"):
            torch.nn.init.uniform_(param, -mean_weight, mean_weight)
        elif name.endswith("scale.weight"):
            torch.nn.init.uniform_(param, -scale_weight, scale_weight)

def repeat_along_dimension(array, number, dim=0):
    """Return a view into the array with a new, repeated dimension.

    Parameters
    ----------
    array : torch.Tensor
    number : int
        The number of repeats.
    dim : int
        The dimension along which to repeat.

    Returns
    -------
    array : torch.Tensor
        A view into the input array. It has an additional dimension added at `dim` and
        the this new dimension is repeated `number` times.

    >>> import torch
    >>> array = torch.tensor([[1., 2.], [3., 4.]])

    >>> repeat_along_dimension(array, number=3, dim=0)
    tensor([[[1., 2.],
             [3., 4.]],
    <BLANKLINE>
            [[1., 2.],
             [3., 4.]],
    <BLANKLINE>
            [[1., 2.],
             [3., 4.]]])

    >>> repeat_along_dimension(array, number=3, dim=1)
    tensor([[[1., 2.],
             [1., 2.],
             [1., 2.]],
    <BLANKLINE>
            [[3., 4.],
             [3., 4.],
             [3., 4.]]])
    """
    expanded_array = array.unsqueeze(dim)
    shape = [-1] * expanded_array.dim()
    shape[dim] = number
    return expanded_array.expand(*shape)

def torch_quadratic(array, matrix):
    """Compute the quadratic form in pytorch.

    Parameters
    ----------
    array : torch.Tensor
    matrix : torch.Tensor

    Returns
    -------
    ndarray
        The quadratic form evaluated for each x.
    """
    squared_values = array * (array @ matrix)
    return torch.sum(squared_values, dim=-1, keepdim=True)

def inverse_softplus(x):
    """Inverse function to torch.functional.softplus.

    Parameters
    ----------
    x : torch.Tensor

    Returns
    -------
    output : torch.Tensor
    """
    if not isinstance(x, torch.Tensor):
        x = torch.tensor(x, dtype=torch.get_default_dtype())
    return torch.log(torch.exp(x) - 1.0)

class TileCode(nn.Module):
    """Tile coding implementation.

    A tile code discretizes the environment into bins.
    Given a continuous tensor, it encodes the output as the closest tile index per dim.
    The output can be either an integer or a one-hot encode vector.

    Parameters
    ----------
    low: array_like
        Array of lower value of bins (per dimension).
    high: array_like
        Array of higher value of bins (per dimesnion)
    bins: int
        Number of bins per dimension.
    one_hot: bool, optional (default = True).
        Flag that indicates if output is one-hot encoded or not.

    Notes
    -----
    Only a same number of bins is implemented as this makes compilation easier.
    """

    def __init__(self, low, high, bins, one_hot=True):
        super().__init__()
        dim = len(low)
        assert dim == len(high)

        self.tiles = torch.stack(
            [
                torch.linspace(low_, high_, bins + 1)[1:] - (high_ - low_) / (2 * bins)
                for low_, high_ in zip(low, high)
            ],
            dim=-1,
        )

        bins = torch.tensor([bins] * dim)
        self.bins = bins + 1
        self.num_outputs = self._tuple_to_int(bins).item() + 1

        if one_hot:
            self.extra_dims = -dim + self.num_outputs
        else:
            self.extra_dims = -dim + 1
        self.one_hot = one_hot

    def _tuple_to_int(self, tuple_):
        """Convert tuple of ints into a single integer."""
        out, tuple_ = tuple_[..., 0], tuple_[..., 1:]
        i = 1
        while tuple_.shape[-1] > 0:
            out = self.bins[i] * out + tuple_[..., 0]
            tuple_ = tuple_[..., 1:]
            i += 1
        return out

    def forward(self, x):
        """Encode a vector using tile-coding."""
        if x.dim() == 0 or (x.dim() == 1 and len(self.bins) == 1):
            x = x.unsqueeze(-1)
        code = self._tuple_to_int(digitize(x, self.tiles))
        if self.one_hot:
            return one_hot_encode(code, self.num_outputs)
        return code

def digitize(tensor, bin_boundaries):
    """Implement numpy digitize using torch."""
    result = torch.zeros(tensor.shape).long()
    for boundary in bin_boundaries:
        result += (tensor > boundary).long()
    return result

class OneHotEncode(nn.Module):
    """One Hot Encoder."""

    def __init__(self, num_classes: int):
        super().__init__()
        assert num_classes > 0
        self.num_classes = num_classes

    def forward(self, x):
        """One-Hot Encode Vector."""
        return one_hot_encode(x.long(), self.num_classes)

def one_hot_encode(tensor, num_classes: int):
    """Encode a tensor using one hot encoding.

    Parameters
    ----------
    tensor: torch.Tensor of dtype torch.long

    num_classes: int
        number of classes to encode.

    Returns
    -------
    tensor: torch.Tensor of dtype torch.float
        one-hot-encoded tensor.

    Raises
    ------
    TypeError:
        if tensor not of dtype long.

    """
    tensor = tensor.long()
    if tensor.dim() == 0:
        return torch.scatter(torch.zeros(num_classes), -1, tensor, 1)
    else:
        tensor_ = tensor.reshape(-1, 1)
        out = torch.scatter(torch.zeros(tensor_.shape[0], num_classes), -1, tensor_, 1)
        return out.reshape(tensor.shape + (num_classes,))

def reverse_cumsum(tensor, dim=-1):
    """Return reversed cumsum along dimensions."""
    return torch.flip(torch.cumsum(torch.flip(tensor, (dim,)), dim), (dim,))

def reverse_cumprod(tensor, dim=-1):
    """Return reversed cumprod along dimensions."""
    return torch.flip(torch.cumprod(torch.flip(tensor, (dim,)), dim), (dim,))

def get_batch_size(tensor, base_size):
    """Get the batch size of a tensor if it is a discrete or continuous tensor.

    Parameters
    ----------
    tensor: torch.tensor.
        Tensor to identify batch size
    base_size: Tuple.
        Base size of tensor.

    Returns
    -------
    batch_size: int or None

    """
    size = tensor.shape
    if len(base_size) == 0:  # Discrete
        return tuple(size)
    else:
        return tuple(size[: -len(base_size)])

def random_tensor(discrete, dim, batch_size=None):
    """Generate a random tensor with a given dimension and batch size.

    Parameters
    ----------
    discrete: bool
        Flag that indicates if random tensor is discrete.
    dim: int
        Dimensionality of random tensor
    batch_size: int, optional
        batch size of random tensor to generate.

    Returns
    -------
    tensor: torch.tensor

    """
    if discrete:
        if batch_size:
            return torch.randint(dim, (batch_size,))
        else:
            return torch.randint(dim, ())
    else:
        if batch_size:
            return torch.randn(batch_size, dim)
        else:
            return torch.randn(dim)

def freeze_parameters(module):
    """Freeze all module parameters.

    Can be used to exclude module parameters from the graph.

    Parameters
    ----------
    module : torch.nn.Module
    """
    for param in module.parameters():
        param.requires_grad = False

def unfreeze_parameters(module):
    """Unfreeze all module parameters.

    Can be used to include excluded module parameters in the graph.

    Parameters
    ----------
    module : torch.nn.Module
    """
    for param in module.parameters():
        param.requires_grad = True

def stop_learning(module):
    """Stop learning all module parameters.

    Can be used for early stopping parameters of the graph.

    Parameters
    ----------
    module : torch.nn.Module
    """
    for param in module.parameters():
        param.requires_grad = False
        param.grad = None

def resume_learning(module):
    """Resume learning all module parameters.

    Can be used to resume learning after early stopping parameters of the graph.

    Parameters
    ----------
    module : torch.nn.Module
    """
    for param in module.parameters():
        param.requires_grad = True
        param.grad = torch.zeros_like(param.data)

class DisableGradient(object):
    """Context manager to disable gradients temporarily.

    Gradients terms will be zero-ed, momentum terms will continue.

    Parameters
    ----------
    modules : sequence
        List of torch.nn.Module.
    """

    def __init__(self, *modules):
        self.modules = modules

    def __enter__(self):
        """Freeze the parameters."""
        for module in self.modules:
            if module is not None:
                freeze_parameters(module)

    def __exit__(self, *args):
        """Unfreeze the parameters."""
        for module in self.modules:
            if module is not None:
                unfreeze_parameters(module)

class EnableGradient(object):
    """Context manager to enables gradients temporarily.

    Gradients terms will be zero-ed, momentum terms will continue.

    Parameters
    ----------
    modules : sequence
        List of torch.nn.Module.
    """

    def __init__(self, *modules):
        self.modules = modules

    def __enter__(self):
        """Unfreeze the parameters."""
        for module in self.modules:
            if module is not None:
                unfreeze_parameters(module)

    def __exit__(self, *args):
        """Freeze the parameters."""
        for module in self.modules:
            if module is not None:
                freeze_parameters(module)

class DisableOptimizer(object):
    """Context manager to disable optimization steps temporarily.

    Gradients and momentum terms will be zero-ed.

    Parameters
    ----------
    modules : sequence
        List of torch.nn.Module.
    """

    def __init__(self, *modules):
        self.modules = modules

    def __enter__(self):
        """Freeze the parameters."""
        for module in self.modules:
            if module is not None:
                stop_learning(module)

    def __exit__(self, *args):
        """Unfreeze the parameters."""
        for module in self.modules:
            if module is not None:
                resume_learning(module)

max_stars_repo_name: shenao-zhang/DCPU

max_issues_repo_name: shenao-zhang/DCPU

max_forks_repo_name: shenao-zhang/DCPU


text:
import matplotlib.pyplot as plt
import numpy as np

from ifg import IfgCalculator

# temps = np.logspace(-2, 4, num=100)
thetas = np.linspace(0.5, 10.0, 1000)
vols = [0.1, 1, 100]
t = IfgCalculator(volumes=vols, thetas=thetas, output_in_si=False)
asymp = np.ones_like(thetas) * 5 / 3

plt.plot(thetas, t.C_P[:, 0] / t.C_V[:, 0], "k-", label=r"$v = 0.1$")
# NOTE: does not depend on volume
# plt.plot(thetas, t.C_P[:, 0] / t.C_V[:, 1], "k--", label=r"$v = 1$")
# plt.plot(thetas, t.C_P[:, 0] / t.C_V[:, 2], "k-.", label=r"$v = 10$")
plt.plot(thetas, asymp, "r--", label=r"High-$T$ asymp.")

plt.xscale("log")
plt.xlabel(r"$T$, Ha")
plt.ylabel(r"$C_P / C_V$")
# plt.xlim(1e-2, 1e4)
# plt.ylim(0, 2)
plt.legend()
plt.text(200, 1.7, r"$C_P/C_V = 5/3$")
plt.savefig("cp_by_cv.pdf", bbox_inches="tight")
plt.show()


text:
"""Extract information from the GeoLife dataset"""
import numpy as np
import os
import pandas as pd

from dateutil import parser
from time import time

from joblib import Parallel, delayed

def decode_str(s):
    return s.decode('utf-8')

class GeoLifeExtractor(object):

    def __init__(self, base_path='./geolife-gps-trajectory-dataset/'):
        self.base_path = base_path

    def extract_user_information(self, user):
        #
        # Get users dataset files
        #
        files = []
        user_dir = os.path.join(self.base_path, f'{user}/Trajectory/')

        for (_, _, filenames) in os.walk(user_dir):
            files.extend([f for f in filenames if f.endswith('.plt')])

        #
        # Build a pandas DataFrame gathering data for each file of the user
        #
        df = None

        columns = ['latitude', 'longitude', 'Unknown1', 'Unknown2', 'Unknown3', 'date', 'time']

        for file in files:
            data = np.genfromtxt(
                os.path.join(user_dir, file),
                delimiter=',',
                skip_header=6,
                converters={
                    0: float,
                    1: float,
                    2: float,
                    3: float,
                    4: float,
                    5: decode_str,
                    6: decode_str,
                }
            )

            if df is None:
                df = pd.DataFrame([list(l) for l in data], columns=columns)
            else:
                df_aux = pd.DataFrame([list(l) for l in data], columns=columns)
                df = pd.concat([df, df_aux])

        df['user'] = user

        print(f'{user}: Number of locations: {len(df)}')
        print(f'{user}: Trajectories from {df.date.min()} until {df.date.max()}')

        df.to_pickle(f'.tmp_user_tracjectories_{user}.pkl')

        return df

    def extract_all_users_information(self, n_jobs=1):
        self.users = [
            d
            for d in os.listdir(self.base_path)
            if os.path.isdir(os.path.join(self.base_path, d))
        ]

        t = time()
        dfs = Parallel(n_jobs=n_jobs)(
            delayed(self.extract_user_information)(user)
            for user in self.users
        )
        print(f'Took {time() - t} gather information from all users')

        return pd.concat(dfs)

max_stars_repo_name: leopiney/tscf

max_issues_repo_name: leopiney/tscf

max_forks_repo_name: leopiney/tscf


text:
#!/usr/bin/python
import argparse
import glob
import os
from UnsupervisedModels import AllMF, RandomGuess, MeanGuess
from SupervisedModels import AllZeroFeature, WordVector
from FreebaseFeatures import FreebaseFeatures
from Ensemble import Ensemble
import numpy as np
import cPickle
from utils import *

def check_type(file_path, professions, nationalities):
	hit_p = 0.0
	hit_n = 0.0
	with open(file_path, "r") as f:
		for line in f:
			value = line.strip().split("\t")[1]
			if value in professions:
				hit_p += 1
			if value in nationalities:
				hit_n += 1
	if hit_p >= hit_n:
		return "profession"
	else:
		return "nationality"

"""
This script is our submission software.
"""
if __name__ == "__main__":
	#parse arguments
	np.random.seed(0)
	parser = argparse.ArgumentParser(description="triple_scoring")
	parser.add_argument("-c", "--inputDataset", default = "../data/")
	parser.add_argument("-r", "--inputRun",  default = "../data/")
	parser.add_argument("-i", "--inputFiles", action = "append", nargs = "+")
	parser.add_argument("-o", "--outputDir", default = "../output/")
	parser.add_argument("-m", "--modelDir", default = "../models/")

	args = parser.parse_args()

	if not os.path.exists(args.outputDir):
		os.mkdir(args.outputDir)

	#load professions / nationality
	professions = read_one_column(os.path.join(args.inputDataset, "professions"))
	nationalities = read_one_column(os.path.join(args.inputDataset, "nationalities"))

	#load models
	"""
        modelP = FreebaseFeatures(freebase_features="../data/freebase_features/features_ipca.bin", labels= professions)
	modelP.load(os.path.join(args.modelDir, "freebase_profession.mod"))
        modelN = FreebaseFeatures(freebase_features="../data/freebase_features/features_ipca.bin", labels= nationalities)
	modelN.load(os.path.join(args.modelDir, "freebase_nationality.mod"))

	modelWP = WordVector(w2v_path = "../models/vectors.bin")
	modelFP = FreebaseFeatures(freebase_features="../data/freebase_features/features_ipca.bin", labels=professions)
        modelP = Ensemble(model_list = [modelWP, modelFP])
	modelP.load(os.path.join(args.modelDir, "ensemble_profession.mod"))

        modelWN = WordVector(w2v_path = "../models/vectors.bin")
        modelFN = FreebaseFeatures(freebase_features="../data/freebase_features/features_ipca.bin", labels=nationalities)
        modelN = Ensemble(model_list = [modelWN, modelFN])
	modelN.load(os.path.join(args.modelDir, "ensemble_nationality.mod"))
        """

        modelP = WordVector(w2v_path = "../models/vectors.bin")
	modelP.load(os.path.join(args.modelDir, "word2vec_profession.mod"))
	modelN = WordVector(w2v_path = "../models/vectors.bin")
	modelN.load(os.path.join(args.modelDir, "word2vec_nationality.mod"))

        if args.inputFiles is not None:
            for input_path in args.inputFiles:
                    input_path = input_path[0]
                    input_type = check_type(input_path, professions, nationalities)
                    file_name = os.path.split(input_path)[1]
                    if input_type == "profession":
                            modelP.test(input_path, os.path.join(args.outputDir, file_name))
                    elif input_type == "nationality":
                            modelN.test(input_path, os.path.join(args.outputDir, file_name))
        else:
            if os.path.isfile(args.inputDataset):
                    file_name = os.path.split(args.inputDataset)[1]
                    file_path = args.inputDataset
                    input_type = check_type(file_path, professions, nationalities)
                    if input_type == "profession":
                            modelP.test(file_path, os.path.join(args.outputDir, file_name))
                    elif input_type == "nationality":
                            modelN.test(file_path, os.path.join(args.outputDir, file_name))
            else:
                    for file_path in glob.glob(args.inputDataset + "/*.train"):
                            file_name = os.path.split(file_path)[1]
                            input_type = check_type(file_path, professions, nationalities)
                            if input_type == "profession":
                                    modelP.test(file_path, os.path.join(args.outputDir, file_name))
                            elif input_type == "nationality":
                                    modelN.test(file_path, os.path.join(args.outputDir, file_name))
                    for file_path in glob.glob(args.inputDataset + "/*.test"):
                            file_name = os.path.split(file_path)[1]
                            input_type = check_type(file_path, professions, nationalities)
                            if input_type == "profession":
                                    modelP.test(file_path, os.path.join(args.outputDir, file_name))
                            elif input_type == "nationality":
                                    modelN.test(file_path, os.path.join(args.outputDir, file_name))


text:
import numpy as np
from kitti_data_loader import dataloader
import matplotlib.pyplot as plt

class EKF(object):
	class State(object):
		def __init__(self,X=np.mat(np.zeros((6,1)))):
			self.x = X[0,0]
			self.y = X[1,0]
			self.v = X[2,0]
			self.a = X[3,0]
			self.phi = X[4,0]
			self.w = X[5,0]

		def state_to_matrix(self):
			return np.mat([[self.x],[self.y],[self.v],[self.a],[self.phi],[self.w]])

		def unpack(self):
			return self.x,self.y,self.v,self.a,self.phi,self.w

		def __add__(self,matrix):
			return EKF.State(self.state_to_matrix() + matrix)

		def __radd__(self,matrix):
			return EKF.State(matrix + self.state_to_matrix())

		def __mul__(self,matrix):
			return self.state_to_matrix()*matrix

		def __rmul__(self,matrix):
			return matrix*self.state_to_matrix()

	class dynamics(object):
		def __init__(self,X):
			self.X = X
			self.del_t = 0

		def step(self,del_t):
			self.del_t = del_t

			X_next = EKF.State()
			X_next.x = self.X.x + self.X.v*np.cos(self.X.phi+self.X.w*self.del_t/2)*self.del_t
			X_next.y = self.X.y + self.X.v*np.sin(self.X.phi+self.X.w*self.del_t/2)*self.del_t
			X_next.v = self.X.v + self.X.a*self.del_t
			X_next.a = self.X.a
			X_next.phi = self.X.phi + self.X.w*self.del_t
			X_next.w = self.X.w
			return X_next

		def measurement_update(self,measurement):
			self.X.x = measurement[0,0]
			self.X.y = measurement[1,0]
			self.X.a = measurement[2,0]
			self.X.w = measurement[3,0]

	def __init__(self,X_init,P_init,Q_init,R_init,del_t=0.001):
		self.X = X_init
		self.P = P_init
		self.Q = Q_init
		self.R = R_init
		self.del_t = del_t

	def state_predict(self):
		return EKF.dynamics(self.X).step(self.del_t)

	def get_transition_matrix(self):
		x,y,v,a,phi,w = self.X.unpack()
		Phi = np.mat([ \
		[1,0,0,0,0,0], \
		[0,1,0,0,0,0], \
		[np.cos(phi+w*self.del_t/2)*self.del_t,np.sin(phi+w*self.del_t/2)*self.del_t,1,0,0,0], \
		[0,0,self.del_t,1,0,0], \
		[-v*np.sin(phi+w*self.del_t/2)*self.del_t,v*np.cos(phi+w*self.del_t/2)*self.del_t,0,0,1,0], \
		[-v*np.sin(phi+w*self.del_t/2)*self.del_t**2 * 1/2,v*np.cos(phi+w*self.del_t/2)*self.del_t**2 * 1/2,0,0,self.del_t,1] \
		])
		return Phi.T

	def get_measurement_matrix(self):
		H = np.mat([ \
			[1,0,0,0,0,0], \
			[0,1,0,0,0,0], \
			[0,0,0,1,0,0], \
			[0,0,0,0,0,1] \
			])
		return H

	def covariance_predict(self):
		phi_transition = self.get_transition_matrix()
		return phi_transition*self.P*phi_transition.T + self.Q

	def step(self,measurement,del_t):
		self.del_t = del_t
		X_prediction = self.state_predict()
		P_prediction = self.covariance_predict()
		H_measurement = self.get_measurement_matrix()

		K_gain = P_prediction*H_measurement.T*(H_measurement*P_prediction*H_measurement.T+self.R).I
		self.X = X_prediction + K_gain*(measurement-H_measurement*X_prediction)
		self.P = (np.mat(np.eye(6))-K_gain*H_measurement)*P_prediction
		return self.X

def plot_pos_data(dataset,filtered_data,no_filtered_data):
	plt.figure(1)
	plt.plot(np.array(dataset)[:,0],np.array(dataset)[:,1],'b-',label='ground truth')
	plt.legend()

	plt.plot(np.array(filtered_data)[:,0],np.array(filtered_data)[:,1],'r-',label='kalman filter')
	plt.legend()

	plt.plot(np.array(no_filtered_data)[:,0],np.array(no_filtered_data)[:,1],'g-',label='no filter')
	plt.legend()

	plt.xlabel('X-Axis')
	plt.ylabel('Y-Axis')

	plt.show()

if __name__ == '__main__':
	X_init,dataset = dataloader()
	X_0 = np.mat([[X_init[0]],[X_init[1]],[X_init[2]],[X_init[3]],[X_init[4]],[X_init[5]]])
	X_0 = EKF.State(X_0)
	P_0 = np.mat(' \
		1 0 0 0 0 0; \
		0 1 0 0 0 0; \
		0 0 1 0 0 1; \
		0 0 0 1 0 0; \
		0 1 0 0 1 0; \
		0 0 0 0 0 1 \
		')
	Q = np.mat(' \
		0 0 0 0 0 0; \
		0 0 0 0 0 0; \
		0 0 0 0 0 0; \
		0 0 0 0 0 0; \
		0 0 0 0 0 0; \
		0 0 0 0 0 0 \
		')
	R = np.mat(' \
		10 0 0 0; \
		0 10 0 0; \
		0 0 10 0; \
		0 0 0 10 \
		')

	kalman_filter = EKF(X_0,P_0,Q,R)

	cur_t = 0
	filtered_data = [X_0.state_to_matrix()]
	no_filtered_data = [X_0.state_to_matrix()]
	dynamics = EKF.dynamics(X_0)

	for i in range(len(dataset)-1):
		random_noise = np.random.multivariate_normal([0,0,0,0],R)
		measurement = np.mat([[dataset[i][0] + random_noise[0]],[dataset[i][1] + random_noise[1]],[dataset[i][2] + random_noise[2]],[dataset[i][3] + random_noise[3]]])
		del_t = dataset[i][4] - cur_t

		X_filtered = kalman_filter.step(measurement,del_t)
		filtered_data.append(X_filtered.state_to_matrix())

		dynamics.measurement_update(measurement)
		X_no_filtered = dynamics.step(del_t)
		no_filtered_data.append(X_no_filtered.state_to_matrix())
		dynamics.X = X_no_filtered

		cur_t = dataset[i][4]

	plot_pos_data(dataset,filtered_data,no_filtered_data)

max_stars_repo_name: dohyeoklee/EKF-kitti-GPS-IMU

max_issues_repo_name: dohyeoklee/EKF-kitti-GPS-IMU

max_forks_repo_name: dohyeoklee/EKF-kitti-GPS-IMU