text:
from mnms import simio, tiled_ndmap, utils, soapack_utils as s_utils, tiled_noise, wav_noise, inpaint
from pixell import enmap, wcsutils
from enlib import bench
from optweight import wavtrans

import numpy as np

from abc import ABC, abstractmethod

# expose only concrete noise models, helpful for namespace management in client
# package development. NOTE: this design pattern inspired by the super-helpful
# registry trick here: https://numpy.org/doc/stable/user/basics.dispatch.html

REGISTERED_NOISE_MODELS = {}

def register(registry=REGISTERED_NOISE_MODELS):
    """Add a concrete NoiseModel implementation to the specified registry (dictionary)."""
    def decorator(noise_model_class):
        registry[noise_model_class.__name__] = noise_model_class
        return noise_model_class
    return decorator

# NoiseModel API and concrete NoiseModel classes.

class NoiseModel(ABC):

    def __init__(self, *qids, data_model=None, preload=True, ivar=None, mask_est=None,
                calibrated=True, downgrade=1, lmax=None, mask_version=None, mask_name=None,
                union_sources=None, notes=None, dtype=None, **kwargs):
        """Base class for all NoiseModel subclasses. Supports loading raw data necessary for all
        subclasses, such as masks and ivars. Also defines some class methods usable in subclasses.

        Parameters
        ----------
        qids : str
            One or more qids to incorporate in model.
        data_model : soapack.DataModel, optional
            DataModel instance to help load raw products, by default None.
            If None, will load the 'default_data_model' from the 'mnms' config.
        preload: bool, optional
            If ivar is None, load them ivar maps for each qid, by default True. Setting
            to False can expedite access to helper methods of this class without waiting
            for ivar to load from disk.
        ivar : array-like, optional
            Data inverse-variance maps, by default None. If provided, assumed properly
            downgraded into compatible wcs with internal NoiseModel operations. If None
            and preload, will be loaded via DataModel according to 'downgrade' and
            'calibrated' kwargs.
        mask_est : enmap.ndmap, optional
            Mask denoting data that will be used to determine the harmonic filter used
            to whiten the data before estimating its variance, by default None. If
            provided, assumed properly downgraded into compatible wcs with internal
            NoiseModel operations. If None and preload, will load a mask according to
            the 'mask_version' and 'mask_name' kwargs.
        calibrated : bool, optional
            Whether to load calibrated raw data, by default True.
        downgrade : int, optional
            The factor to downgrade map pixels by, by default 1.
        lmax : int, optional
            The bandlimit of the maps, by default None.
            If None, will be set to twice the theoretical CAR limit, ie 180/wcs.wcs.cdelt[1].
        mask_version : str, optional
           The mask version folder name, by default None.
           If None, will first look in config 'mnms' block, then block of default data model.
        mask_name : str, optional
            Name of mask file, by default None.
            If None, a default mask will be loaded from disk.
        union_sources : str, optional
            A soapack source catalog, by default None. If given, inpaint data and ivar maps.
        notes : str, optional
            A descriptor string to differentiate this instance from
            otherwise identical instances, by default None.
        dtype : np.dtype, optional
            The data type used in intermediate calculations and return types, by default None.
            If None, inferred from data_model.dtype.
        kwargs : dict, optional
            Optional keyword arguments to pass to simio.get_sim_mask_fn (currently just
            'galcut' and 'apod_deg'), by default None.
        """

        # if mask and ivar are provided, there is no way of checking
        # whether calibrated, what downgrade, etc; they are assumed to be 'correct'

        # store basic set of instance properties
        self._qids = qids
        if data_model is None:
            data_model = utils.get_default_data_model()
        self._data_model = data_model
        self._calibrated = calibrated
        self._downgrade = downgrade
        if mask_version is None:
            mask_version = utils.get_default_mask_version()
        self._mask_version = mask_version
        self._mask_name = mask_name
        self._notes = notes
        self._union_sources = union_sources
        self._kwargs = kwargs
        self._dtype = dtype if dtype is not None else self._data_model.dtype

        # get derived instance properties
        self._num_arrays = len(self._qids)
        self._num_splits = utils.get_nsplits_by_qid(self._qids[0], self._data_model)
        self._use_default_mask = mask_name is None

        # Get shape, wcs, ivars, and mask_observed -- need these for every operation.
        full_shape, full_wcs = self._check_geometry()
        self._shape, self._wcs = utils.downgrade_geometry_cc_quad(
            full_shape, full_wcs, self._downgrade
            )

        if ivar is not None:
            self._ivar = ivar
            self._mask_observed = utils.get_bool_mask_from_ivar(self._ivar)
        elif preload:
            self._ivar, self._mask_observed = self._get_ivar_and_mask_observed()
        else:
            raise ValueError('Models require ivar, please preload or supply manually')

        # Possibly store input data
        self._mask_est = mask_est

        # initialize unloaded noise model dictionary, holds noise model variables for each split
        self._nm_dict = {}

        # get lmax
        if lmax is None:
            lmax = utils.lmax_from_wcs(self._wcs)
        self._lmax = lmax

        # sanity checks
        if self._ivar is not None:
            assert self._num_splits == self._ivar.shape[-4], \
                'Num_splits inferred from ivar shape != num_splits from data model table'

    def _check_geometry(self, return_geometry=True):
        """Check that each qid in this instance's qids has compatible shape and wcs."""
        for i, qid in enumerate(self._qids):
            # Load up first geometry of first split's ivar. Only need pixel shape.
            shape, wcs = s_utils.read_map_geometry(self._data_model, qid, 0, ivar=True)
            shape = shape[-2:]
            assert len(shape) == 2, 'shape must have only 2 dimensions'

            # Check that we are using the geometry for each qid -- this is required!
            if i == 0:
                main_shape, main_wcs = shape, wcs
            else:
                with bench.show(f'Checking geometry compatibility between {qid} and {self._qids[0]}'):
                    assert(
                        shape == main_shape), 'qids do not share pixel shape -- this is required!'
                    assert wcsutils.is_compatible(
                        wcs, main_wcs), 'qids do not share a common wcs -- this is required!'

        if return_geometry:
            return main_shape, main_wcs
        else:
            return None

    def _get_ivar_and_mask_observed(self):
        """Load the inverse-variance maps according to instance attributes, and use it
        to construct and observed-by-all-splits pixel map.

        Returns
        -------
        ivars, mask_observed : (nmaps, nsplits, npol, ny, nx) enmap, (ny, nx) enmap
            Inverse-variance maps, possibly downgraded. Observed pixel map,
            possibly downgraded.
        """

        # first check for ivar compatibility and get map geometry
        full_shape, full_wcs = self._check_geometry()

        # load the first ivar map geometry so that we may allocate a buffer to accumulate
        # all ivar maps in -- this has shape (nmaps, nsplits, npol, ny, nx).
        ivars = self._empty(ivar=True)
        mask_observed = np.ones(full_shape, dtype=bool)

        for i, qid in enumerate(self._qids):
            with bench.show(self._action_str(qid, ivar=True)):
                if self._calibrated:
                    mul = s_utils.get_mult_fact(self._data_model, qid, ivar=True)
                else:
                    mul = 1

                # we want to do this split-by-split in case we can save
                # memory by downgrading one split at a time
                for j in range(self._num_splits):

                    ivar = s_utils.read_map(self._data_model, qid, j, ivar=True)
                    ivar *= mul

                    # iteratively build the mask_observed at full resolution,
                    # loop over leading dims
                    for idx in np.ndindex(*ivar.shape[:-2]):
                        mask_observed *= ivar[idx].astype(bool)

                    if self._downgrade != 1:
                        # use harmonic instead of interpolated downgrade because it is
                        # 10x faster
                        ivar = utils.harmonic_downgrade_cc_quad(
                            ivar, self._downgrade, area_pow=1
                            )

                    # zero-out any numerical negative ivar
                    ivar[ivar < 0] = 0
                    ivars[i,j] = ivar

        with bench.show('Generating observed-pixels mask'):
            mask_observed = enmap.enmap(mask_observed, wcs=full_wcs, copy=False)
            mask_observed = utils.interpol_downgrade_cc_quad(
                mask_observed, self._downgrade, dtype=self._dtype
                )

            # define downgraded mask_observed to be True only where the interpolated
            # downgrade is all 1 -- this is the most conservative route in terms of
            # excluding pixels that may not actually have nonzero ivar or data
            mask_observed = utils.get_mask_bool(mask_observed, threshold=1.)

            # finally, need to layer on any ivars that may still be 0 that aren't yet
            # masked
            mask_observed *= utils.get_bool_mask_from_ivar(ivars)

        return ivars*mask_observed, mask_observed

    def _empty(self, ivar=False, num_arrays=None, num_splits=None,
                shape=None, wcs=None):
        """Allocate an empty buffer that will broadcast against the Noise Model
        number of arrays, number of splits, and the map (or ivar) shape.

        Parameters
        ----------
        ivar : bool, optional
            If True, load the inverse-variance map shape for the qid and
            split. If False, load the source-free map shape for the same,
            by default False.
        num_arrays : int, optional
            The number of arrays (axis -5) in the empty ndmap, by default None.
            If None, inferred from the number of qids in the NoiseModel.
        num_splits : int, optional
            The number of splits (axis -4) in the empty ndmap, by default None.
            If None, inferred from the number of splits on disk.
        shape : tuple, optional
            A geometry footprint shape to use to build the empty ndmap, by
            default None. If None, will use the downgraded geometry inferred
            from the data on-disk.
        wcs : astropy.wcs.WCS
            A geometry wcs to use to build the empty ndmap, by default None.
            If None, will use the downgraded geometry inferred from the data
            on-disk.

        Returns
        -------
        enmap.ndmap
            An empty ndmap with shape (num_arrays, num_splits, num_pol, ny, nx),
            with dtype of the instance soapack.DataModel. If ivar is True, num_pol
            likely is 1. If ivar is False, num_pol likely is 3.
        """
        # read geometry from the map to be loaded. we really just need the first component,
        # a.k.a "npol", which varies depending on if ivar is True or False
        if shape is not None:
            footprint_shape = shape[-2:]
        else:
            footprint_shape = self._shape

        if wcs is not None:
            footprint_wcs = wcs
        else:
            footprint_wcs = self._wcs

        shape, _ = s_utils.read_map_geometry(self._data_model, self._qids[0], 0, ivar=ivar)
        shape = (shape[0], *footprint_shape)

        if num_arrays is None:
            num_arrays = self._num_arrays
        if num_splits is None:
            num_splits = self._num_splits

        shape = (num_arrays, num_splits, *shape)
        return enmap.empty(shape, wcs=footprint_wcs, dtype=self._dtype)

    def _action_str(self, qid, ivar=False):
        """Get a string for benchmarking the loading step of a map product.

        Parameters
        ----------
        qid : str
            Map identification string
        ivar : bool, optional
            If True, print 'ivar' where appropriate. If False, print 'imap'
            where appropriate, by default False.

        Returns
        -------
        str
            Benchmarking action string.

        Examples
        --------
        >>> from mnms import noise_models as nm
        >>> tnm = nm.TiledNoiseModel('s18_03', downgrade=2, union_sources='20210209_sncut_10_aggressive', notes='my_model')
        >>> tnm._action_str('s18_03')
        >>> 'Loading, inpainting, downgrading imap for s18_03'
        """
        if ivar:
            if self._downgrade != 1:
                return f'Loading, downgrading ivar for {qid}'
            else:
                return f'Loading ivar for {qid}'
        else:
            if self._downgrade != 1 and self._union_sources:
                return f'Loading, inpainting, downgrading imap for {qid}'
            elif self._downgrade != 1 and not self._union_sources:
                return f'Loading, downgrading imap for {qid}'
            elif self._downgrade == 1 and self._union_sources:
                return f'Loading, inpainting imap for {qid}'
            else:
                return f'Loading imap for {qid}'

    def get_model(self, check_on_disk=True, write=True, keep_model=False, keep_data=False, verbose=False, **kwargs):
        """Generate (or load) a sqrt-covariance matrix for this NoiseModel instance.

        Parameters
        ----------
        check_on_disk : bool, optional
            If True, first check if an identical model (including by 'notes') exists
            on-disk for each split. If it does, do nothing or store it in the object
            attributes, depending on the 'keep_model' kwarg. If it does not, generate the model
            for the missing splits instead. By default True.
        write : bool, optional
            Save the generated model to disk, by default True.
        keep_model : bool, optional
            Store the generated (or loaded) model in the instance attributes, by
            default False.
        keep_data: bool, optional
            Store the loaded, possibly downgraded, data split differences in the
            instance attributes, by default False.
        verbose : bool, optional
            Print possibly helpful messages, by default False.
        """
        if check_on_disk:
            # build a list of splits that don't have models on-disk
            does_not_exist = []
            for s in range(self._num_splits):
                res = self._check_model_on_disk(s, keep_model=keep_model)
                if not res:
                    does_not_exist.append(s)
            if not does_not_exist:
                # if all models exist on-disk, exit this function
                return
        else:
            does_not_exist = range(self._num_splits)

        # models need data split differences as inputs
        dmap = self._get_data_split_diffs()
        if keep_data:
            self._dmap = dmap

        # particular model subclasses may further modify data split differences
        dmap = self._get_dmap(dmap)

        # get the conservative mask for estimating the harmonic filter used to whiten
        # the difference maps
        self._mask_est = self._get_mask_est(min_threshold=1e-4)

        for s in does_not_exist:
            with bench.show(f'Generating noise model for split {s}'):
                nm_dict = self._get_model(s, dmap, verbose=verbose)

            if keep_model:
                self._keep_model(s, nm_dict)

            if write:
                fn = self._get_model_fn(s)
                self._write_model(fn, **nm_dict)

    def _get_data_split_diffs(self):
        """Load the raw data split differences according to instance attributes.

        Returns
        -------
        dmaps : (nmaps, nsplits, npol, ny, nx) enmap
            Data split difference maps, possibly downgraded.
        """
        # first check for mask compatibility and get map geometry
        full_shape, full_wcs = self._check_geometry()

        dmaps = self._empty(ivar=False)
        num_pol = dmaps.shape[-3]

        for i, qid in enumerate(self._qids):
            with bench.show(self._action_str(qid, ivar=False)):
                # load the first data map geometry so that we may allocate a buffer to accumulate
                # all data maps in -- this has shape (nmaps=1, nsplits, npol, ny, nx).
                if self._downgrade != 1:
                    # we do this one array at a time to save memory
                    imaps = self._empty(num_arrays=1, ivar=False, shape=full_shape, wcs=full_wcs)
                    ivars = self._empty(num_arrays=1, ivar=True, shape=full_shape, wcs=full_wcs)

                if self._calibrated:
                    mul_imap = s_utils.get_mult_fact(self._data_model, qid, ivar=False)
                    mul_ivar = s_utils.get_mult_fact(self._data_model, qid, ivar=True)
                else:
                    mul_imap = 1
                    mul_ivar = 1

                # we want to do this split-by-split in case we can save
                # memory by downgrading one split at a time
                for j in range(self._num_splits):
                    imap = s_utils.read_map(self._data_model, qid, j, ivar=False)
                    imap = enmap.extract(imap, full_shape, full_wcs)
                    imap *= mul_imap

                    if self._downgrade != 1:
                        # need to reload ivar at full res for differencing (and inpainting)
                        ivar = s_utils.read_map(self._data_model, qid, j, ivar=True)
                        ivar = enmap.extract(ivar, full_shape, full_wcs)
                        ivar *= mul_ivar
                    else:
                        ivar = self._ivar[i, j]

                    if self._union_sources:
                        # the boolean mask for this array, split, is non-zero ivar.
                        # iteratively build the boolean mask at full resolution,
                        # loop over leading dims
                        mask_bool = np.ones(full_shape, dtype=bool)
                        for idx in np.ndindex(*ivar.shape[:-2]):
                            mask_bool *= ivar[idx].astype(bool)

                        self._inpaint(imap, ivar, mask_bool, qid=qid, split_num=j)

                    if self._downgrade != 1:
                        imaps[0, j] = imap
                        ivars[0, j] = ivar
                    else:
                        dmaps[i, j] = imap

                # get difference map for this array. do 1 polarization at a time to save memory.
                # differences before harmonic downgrade avoids ringing around bright objects
                for k in range(num_pol):
                    sel = np.s_[..., k:k+1, :, :] # preserve pol dim to keep dim ordering

                    if self._downgrade != 1:
                        imaps[sel] = utils.get_noise_map(imaps[sel], ivars)
                    else:
                        dmaps[sel] = utils.get_noise_map(dmaps[sel], self._ivar)

                # downgrade each split separately to save memory
                if self._downgrade != 1:
                    for j in range(self._num_splits):
                        dmaps[i, j] = utils.harmonic_downgrade_cc_quad(
                            imaps[0, j], self._downgrade
                        )

        return dmaps*self._mask_observed

    def _inpaint(self, imap, ivar, mask, inplace=True, qid=None, split_num=None):
        """Inpaint point sources given by the union catalog in input map.

        Parameters
        ---------
        imap : (..., 3, Ny, Nx) enmap
            Maps to be inpainted.
        ivar : (Ny, Nx) or (..., 1, Ny, Nx) enmap
            Inverse variance map. If not 2d, shape[:-3] must match imap.
        mask : (Ny, Nx) bool array
            Mask, True in observed regions.
        inplace : bool, optional
            Modify input map.
        qid : str, optional
            Array identifier, used to determine seed for inpainting.
        split_num : int, optional
            The 0-based index of the split that is inpainted, used to get unique seeds
            per split if this function is called per split. Otherwise defaults to 0.
        """

        assert self._union_sources is not None, f'Inpainting needs union-sources, got {self._union_sources}'

        catalog = utils.get_catalog(self._union_sources)
        mask_bool = utils.get_mask_bool(mask)

        if qid:
            # This makes sure each qid gets a unique seed. The sim index is fixed.
            split_idx = 0 if split_num is None else split_num
            seed = utils.get_seed(*(split_idx, 999_999_999, self._data_model, qid))
        else:
            seed = None

        return inpaint.inpaint_noise_catalog(imap, ivar, mask_bool, catalog, inplace=inplace,
                                             seed=seed)

    def _get_mask_est(self, min_threshold=1e-3, max_threshold=1.):
        """Load the data mask from disk according to instance attributes.

        Returns
        -------
        mask : (ny, nx) enmap
            Sky mask. Dowgraded if requested.
        """
        with bench.show('Generating harmonic-filter-estimate mask'):

            # first check for ivar compatibility and get map geometry
            full_shape, full_wcs = self._check_geometry()

            for i, qid in enumerate(self._qids):
                fn = simio.get_sim_mask_fn(
                    qid, self._data_model, use_default_mask=self._use_default_mask,
                    mask_version=self._mask_version, mask_name=self._mask_name, **self._kwargs
                )
                mask = enmap.read_map(fn).astype(self._dtype, copy=False)

                # check that we are using the same mask for each qid -- this is required!
                if i == 0:
                    mask_est = mask
                else:
                    with bench.show(f'Checking mask compatibility between {qid} and {self._qids[0]}'):
                        assert np.allclose(
                            mask, mask_est), 'qids do not share a common mask -- this is required!'
                        assert wcsutils.is_compatible(
                            mask.wcs, mask_est.wcs), 'qids do not share a common mask wcs -- this is required!'

            # Extract mask onto geometry specified by the ivar map.
            mask_est = enmap.extract(mask, full_shape, full_wcs)

            if self._downgrade != 1:
                mask_est = utils.interpol_downgrade_cc_quad(mask_est, self._downgrade)

                # to prevent numerical error, cut below a threshold
                mask_est[mask_est < min_threshold] = 0.

                # to prevent numerical error, cut above a maximum
                mask_est[mask_est > max_threshold] = 1.

        return mask_est

    def _check_model_on_disk(self, split_num, keep_model=False, generate=True):
        """Check if this NoiseModel's model for a given split exists on disk.
        If it does, return True. Depending on the 'keep_model' kwarg, possibly store
        the model in memory. Depending on the 'generate' kwarg, return either
        False or raise a FileNotFoundError if it does not exist on-disk.

        Parameters
        ----------
        split_num : int
            The 0-based index of the split model to look for.
        keep_model : bool, optional
            Store the generated (or loaded) model in the instance attributes if it
            exists on-disk, by default False.
        generate : bool, optional
            If the model does not exist on-disk and 'generate' is True, then return
            False. If the model does not exist on-disk and 'generate' is False, then
            raise a FileNotFoundError. By default True.

        Returns
        -------
        bool
            If the model exists on-disk, return True. If 'generate' is True and the
            model does not exist on-disk, return False.

        Raises
        ------
        FileNotFoundError
            If 'generate' is False and the model does not exist on-disk.
        """
        try:
            self._get_model_from_disk(split_num, keep_model=keep_model)
            return True
        except (FileNotFoundError, OSError):
            fn = self._get_model_fn(split_num)
            if generate:
                print(f'Model for split {split_num} not found on-disk, generating instead')
                return False
            else:
                print(f'Model for split {split_num} not found on-disk, please generate it first')
                raise FileNotFoundError(fn)

    def _get_model_from_disk(self, split_num, keep_model=True):
        """Load a sqrt-covariance matrix from disk. If keep_model, store it in instance attributes."""
        fn = self._get_model_fn(split_num)
        nm_dict = self._read_model(fn)
        if keep_model:
            self._keep_model(split_num, nm_dict)

    @abstractmethod
    def _get_model_fn(self, split_num):
        """Get a noise model filename for split split_num; return as <str>"""
        return ''

    @abstractmethod
    def _read_model(self, fn):
        """Read a noise model with filename fn; return a dictionary of noise model variables"""
        return {}

    def _keep_model(self, split_num, nm_dict):
        """Store a dictionary of noise model variables in instance attributes under key split_num"""
        print(f'Loading model for split {split_num} into memory')
        self._nm_dict[split_num] = nm_dict

    # for memory management
    def _delete_data(self):
        """Remove data stored in memory under instance attribute self._imap"""
        self._imap = None

    @abstractmethod
    def _get_dmap(self, imap):
        """Return the required input difference map for a NoiseModel subclass, from split data imap"""
        return enmap.ndmap

    @abstractmethod
    def _get_model(self, split_num, dmap, verbose=False):
        """Return a dictionary of noise model variables for this NoiseModel subclass, for split split_num and from difference maps dmap"""
        return {}

    @abstractmethod
    def _write_model(self, fn, sqrt_cov_mat=None, sqrt_cov_ell=None, **kwargs):
        """Write sqrt_cov_mat, sqrt_cov_ell, and possibly more noise model variables to filename fn"""
        pass

    def get_sim(self, split_num, sim_num, alm=True, check_on_disk=True, write=False,
                keep_model=True, do_mask_observed=True, verbose=False):
        """Generate a sim from this NoiseModel.

        Parameters
        ----------
        split_num : int
            The 0-based index of the split to simulate.
        sim_num : int
            The map index, used in setting the random seed. Must be non-negative. If the sim
            is written to disk, this will be recorded in the filename. There is a maximum of
            9999, ie, one cannot have more than 10_000 of the same sim, of the same split,
            from the same noise model (including the 'notes').
        alm : bool, optional
            Generate simulated alms instead of a simulated map, by default True.
        check_on_disk : bool, optional
            If True, first check if the exact sim (including the noise model 'notes'),
            exists on disk, and if it does, load and return it. If it does not,
            generate the sim on-the-fly instead, by default True.
        write : bool, optional
            Save the generated sim to disk, by default False.
        keep_model : bool, optional
            Store the loaded model for this split in instance attributes, by default True.
            This spends memory to avoid spending time loading the model from disk
            for each call to this method.
        do_mask_observed : bool, optional
            Apply the mask determined by the observed patch to the sim, by default True.
            If not applied, the sim will bleed into unobserved pixels, but this can
            potentially avoid intermediate calculation if the user will be applying
            their own analysis mask to sims before processing them.
        verbose : bool, optional
            Print possibly helpful messages, by default False.

        Returns
        -------
        enmap.ndmap
            A sim of this noise model with the specified sim num, with shape
            (num_arrays, num_splits, num_pol, ny, nx), even if some of these
            axes have dimension 1. As implemented, num_splits is always 1.
        """

        assert sim_num <= 9999, 'Cannot use a map index greater than 9999'

        if check_on_disk:
            res = self._check_sim_on_disk(split_num, sim_num, alm=alm, mask_obs=do_mask_observed)
            if res is not None:
                return res

        if split_num not in self._nm_dict:
            self._check_model_on_disk(split_num, keep_model=True, generate=False)

        seed = self._get_seed(split_num, sim_num)

        with bench.show(f'Generating noise sim for split {split_num}, map {sim_num}'):

            mask = self._mask_observed if do_mask_observed else None
            if alm:
                sim = self._get_sim_alm(split_num, seed, verbose=verbose, mask=mask)
            else:
                sim = self._get_sim(split_num, seed, verbose=verbose, mask=mask)

        if not keep_model:
            self.delete_model(split_num)

        if write:
            fn = self._get_sim_fn(split_num, sim_num, alm=alm, mask_obs=do_mask_observed)
            if alm:
                utils.write_alm(fn, sim, dtype=self._dtype)
            else:
                enmap.write_map(fn, sim)
        return sim

    def _check_sim_on_disk(self, split_num, sim_num, alm=False, mask_obs=True):
        """Check if sim with split_num, sim_num exists on-disk; if so return it, else return None."""
        fn = self._get_sim_fn(split_num, sim_num, alm=alm, mask_obs=mask_obs)
        try:
            if alm:
                # we know the preshape is (num_arrays, num_splits=1)
                return utils.read_alm(fn, preshape=(self._num_arrays, 1))
            else:
                return enmap.read_map(fn)
        except FileNotFoundError:
            print(f'Sim for split {split_num}, map {sim_num} not found on disk, generating instead')
            return None

    @abstractmethod
    def _get_sim_fn(self, split_num, sim_num, alm=False, mask_obs=True):
        """Get a sim filename for split split_num, sim sim_num, and bool alm/mask_obs; return as <str>"""
        pass

    def _get_seed(self, split_num, sim_num):
        """Return seed for sim with split_num, sim_num."""
        return utils.get_seed(
            *(split_num, sim_num, self._data_model, *self._qids)
            )

    @abstractmethod
    def _get_sim(self, split_num, seed, mask=None, verbose=False):
        """Return a masked enmap.ndmap sim of split split_num, with seed <sequence of ints>"""
        return enmap.ndmap

    @abstractmethod
    def _get_sim_alm(self, split_num, seed, mask=None, verbose=False):
        """Return a masked alm sim of split split_num, with seed <sequence of ints>"""
        pass

    def delete_model(self, split_num):
        """Delete a dictionary entry of noise model variables from instance attributes under key split_num"""
        try:
            del self._nm_dict[split_num]
        except KeyError:
            print(f'Nothing to delete, no model in memory for split {split_num}')

    @property
    def num_splits(self):
        return self._num_splits

    @property
    def mask_est(self):
        return self._mask_est

    @property
    def mask_observed(self):
        return self._mask_observed

@register()
class TiledNoiseModel(NoiseModel):

    def __init__(self, *qids, data_model=None, preload=True, ivar=None, mask_est=None,
                calibrated=True, downgrade=1, lmax=None, mask_version=None, mask_name=None,
                union_sources=None, notes=None, dtype=None, width_deg=4., height_deg=4.,
                delta_ell_smooth=400, **kwargs):
        """A TiledNoiseModel object supports drawing simulations which capture spatially-varying
        noise correlation directions in map-domain data. They also capture the total noise power
        spectrum, spatially-varying map depth, and array-array correlations.

        Parameters
        ----------
        qids : str
            One or more qids to incorporate in model.
        data_model : soapack.DataModel, optional
            DataModel instance to help load raw products, by default None.
            If None, will load the 'default_data_model' from the 'mnms' config.
        preload: bool, optional
            If ivar is None, load them ivar maps for each qid, by default True. Setting
            to False can expedite access to helper methods of this class without waiting
            for ivar to load from disk.
        mask : enmap.ndmap, optional
            Mask denoting data that will be used to determine the harmonic filter used
            to whiten the data before estimating its variance, by default None.
            If None and preload, will load a mask according to the 'mask_version' and
            'mask_name' kwargs.
        ivar : array-like, optional
            Data inverse-variance maps, by default None. If provided, assumed properly
            downgraded into compatible wcs with internal NoiseModel operations. If None
            and preload, will be loaded via DataModel according to 'downgrade' and
            'calibrated' kwargs.
        mask_est : enmap.ndmap, optional
            Mask denoting data that will be used to determine the harmonic filter used
            to whiten the data before estimating its variance, by default None. If
            provided, assumed properly downgraded into compatible wcs with internal
            NoiseModel operations. If None and preload, will load a mask according to
            the 'mask_version' and 'mask_name' kwargs.
        calibrated : bool, optional
            Whether to load calibrated raw data, by default True.
        downgrade : int, optional
            The factor to downgrade map pixels by, by default 1.
        lmax : int, optional
            The bandlimit of the maps, by default None.
            If None, will be set to twice the theoretical CAR limit, ie 180/wcs.wcs.cdelt[1].
        mask_version : str, optional
           The mask version folder name, by default None.
           If None, will first look in config 'mnms' block, then block of default data model.
        mask_name : str, optional
            Name of mask file, by default None.
            If None, a default mask will be loaded from disk.
        union_sources : str, optional
            A soapack source catalog, by default None.
        notes : str, optional
            A descriptor string to differentiate this instance from
            otherwise identical instances, by default None.
        dtype : np.dtype, optional
            The data type used in intermediate calculations and return types, by default None.
            If None, inferred from data_model.dtype.
        width_deg : scalar, optional
            The characteristic tile width in degrees, by default 4.
        height_deg : scalar, optional
            The characteristic tile height in degrees,, by default 4.
        delta_ell_smooth : int, optional
            The smoothing scale in Fourier space to mitigate bias in the noise model
            from a small number of data splits, by default 400.
        kwargs : dict, optional
            Optional keyword arguments to pass to simio.get_sim_mask_fn (currently just
            'galcut' and 'apod_deg'), by default None.

        Notes
        -----
        Unless passed explicitly, the mask and ivar will be loaded at object instantiation time,
        and stored as instance attributes.

        Examples
        --------
        >>> from mnms import noise_models as nm
        >>> tnm = nm.TiledNoiseModel('s18_03', 's18_04', downgrade=2, notes='my_model')
        >>> tnm.get_model() # will take several minutes and require a lot of memory
                            # if running this exact model for the first time, otherwise
                            # will return None if model exists on-disk already
        >>> imap = tnm.get_sim(0, 123) # will get a sim of split 1 from the correlated arrays;
                                       # the map will have "index" 123, which is used in making
                                       # the random seed whether or not the sim is saved to disk,
                                       # and will be recorded in the filename if saved to disk.
        >>> print(imap.shape)
        >>> (2, 1, 3, 5600, 21600)
        """
        super().__init__(
            *qids, data_model=data_model, preload=preload, ivar=ivar, mask_est=mask_est,
            calibrated=calibrated, downgrade=downgrade, lmax=lmax, mask_version=mask_version,
            mask_name=mask_name, union_sources=union_sources, notes=notes, dtype=dtype, **kwargs
        )

        # save model-specific info
        self._width_deg = width_deg
        self._height_deg = height_deg
        self._delta_ell_smooth = delta_ell_smooth

    def _get_model_fn(self, split_num):
        """Get a noise model filename for split split_num; return as <str>"""
        return simio.get_tiled_model_fn(
            self._qids, split_num, self._width_deg, self._height_deg, self._delta_ell_smooth, self._lmax, notes=self._notes,
            data_model=self._data_model, mask_version=self._mask_version, bin_apod=self._use_default_mask, mask_name=self._mask_name,
            calibrated=self._calibrated, downgrade=self._downgrade, union_sources=self._union_sources, **self._kwargs
        )

    def _read_model(self, fn):
        """Read a noise model with filename fn; return a dictionary of noise model variables"""
        # read from disk
        sqrt_cov_mat, extra_hdu = tiled_ndmap.read_tiled_ndmap(
            fn, extra_hdu=['SQRT_COV_ELL']
        )
        sqrt_cov_ell = extra_hdu['SQRT_COV_ELL']

        return {
            'sqrt_cov_mat': sqrt_cov_mat,
            'sqrt_cov_ell': sqrt_cov_ell
            }

    def _get_dmap(self, imap):
        """Return the required input difference map for a NoiseModel subclass, from split data imap"""
        # model needs whitened noise maps as input
        return imap * np.sqrt(utils.get_ivar_eff(self._ivar))

    def _get_model(self, split_num, dmap, verbose=False):
        """Return a dictionary of noise model variables for this NoiseModel subclass, for split split_num and from difference maps dmap"""
        sqrt_cov_mat, sqrt_cov_ell = tiled_noise.get_tiled_noise_covsqrt(
            dmap, split_num, mask_observed=self._mask_observed, mask_est=self._mask_est,
            width_deg=self._width_deg, height_deg=self._height_deg,
            delta_ell_smooth=self._delta_ell_smooth, lmax=self._lmax, rfft=True, nthread=0, verbose=verbose
        )
        return {
            'sqrt_cov_mat': sqrt_cov_mat,
            'sqrt_cov_ell': sqrt_cov_ell
            }

    def _write_model(self, fn, sqrt_cov_mat=None, sqrt_cov_ell=None):
        """Write sqrt_cov_mat, sqrt_cov_ell, and possibly more noise model variables to filename fn"""
        tiled_ndmap.write_tiled_ndmap(
            fn, sqrt_cov_mat, extra_hdu={'SQRT_COV_ELL': sqrt_cov_ell}
        )

    def _get_sim_fn(self, split_num, sim_num, alm=False, mask_obs=True):
        """Get a sim filename for split split_num, sim sim_num; return as <str>"""
        return simio.get_tiled_sim_fn(
            self._qids, self._width_deg, self._height_deg, self._delta_ell_smooth, self._lmax, split_num, sim_num, alm=alm, notes=self._notes,
            data_model=self._data_model, mask_version=self._mask_version, bin_apod=self._use_default_mask, mask_name=self._mask_name,
            calibrated=self._calibrated, downgrade=self._downgrade, union_sources=self._union_sources, mask_obs=mask_obs, **self._kwargs
        )

    def _get_sim(self, split_num, seed, mask=None, verbose=False):
        """Return a masked enmap.ndmap sim of split split_num, with seed <sequence of ints>"""
        # get noise model variables
        sqrt_cov_mat = self._nm_dict[split_num]['sqrt_cov_mat']
        sqrt_cov_ell = self._nm_dict[split_num]['sqrt_cov_ell']

        sim = tiled_noise.get_tiled_noise_sim(
            sqrt_cov_mat, split_num, ivar=self._ivar,
            sqrt_cov_ell=sqrt_cov_ell, num_arrays=self._num_arrays,
            num_splits=self._num_splits, nthread=0, seed=seed, verbose=verbose
        )
        if mask is not None:
            sim *= mask
        return sim

    def _get_sim_alm(self, split_num, seed, mask=None, verbose=False):
        """Return a masked alm sim of split split_num, with seed <sequence of ints>"""

        sim = self._get_sim(split_num, seed, mask=mask, verbose=verbose)
        return utils.map2alm(sim, lmax=self._lmax)

@register()
class WaveletNoiseModel(NoiseModel):

    def __init__(self, *qids, data_model=None, preload=True, ivar=None, mask_est=None,
                 calibrated=True, downgrade=1, lmax=None, mask_version=None, mask_name=None,
                 union_sources=None, notes=None, dtype=None, lamb=1.3, smooth_loc=False,
                 fwhm_fact=2, **kwargs):
        """A WaveletNoiseModel object supports drawing simulations which capture scale-dependent,
        spatially-varying map depth. They also capture the total noise power spectrum, and
        array-array correlations.

        Parameters
        ----------
        qids : str
            One or more qids to incorporate in model.
        data_model : soapack.DataModel, optional
            DataModel instance to help load raw products, by default None.
            If None, will load the 'default_data_model' from the 'mnms' config.
        preload: bool, optional
            If mask and/or ivar are respectively None, load them, by default True. Setting
            to False can expedite access to helper methods of this class without waiting
            for data to load.
        mask : enmap.ndmap, optional
            Mask denoting data that will be used to determine the harmonic filter used
            to whiten the data before estimating its variance, by default None.
            If None and preload, will load a mask according to the 'mask_version' and
            'mask_name' kwargs.
        ivar : array-like, optional
            Data inverse-variance maps, by default None.
            If None, will be loaded via DataModel according to 'downgrade' and 'calibrated' kwargs.
        imap : array-like, optional
            Data maps, by default None.
            If None, will be loaded in call to NoiseModel.get_model(...), and may be retained in
            memory if keep_data=True is passed to that function.
        calibrated : bool, optional
            Whether to load calibrated raw data, by default True.
        downgrade : int, optional
            The factor to downgrade map pixels by, by default 1.
        lmax : int, optional
            The bandlimit of the maps, by default None.
            If None, will be set to twice the theoretical CAR limit, ie 180/wcs.wcs.cdelt[1].
        mask_version : str, optional
           The mask version folder name, by default None.
           If None, will first look in config 'mnms' block, then block of default data model.
        mask_name : str, optional
            Name of mask file, by default None.
            If None, a default mask will be loaded from disk.
        union_sources : str, optional
            A soapack source catalog, by default None.
        notes : str, optional
            A descriptor string to differentiate this instance from
            otherwise identical instances, by default None.
        dtype : np.dtype, optional
            The data type used in intermediate calculations and return types, by default None.
            If None, inferred from data_model.dtype.
        lamb : float, optional
            Parameter specifying width of wavelets kernels in log(ell), by default 1.3
        smooth_loc : bool, optional
            If passed, use smoothing kernel that varies over the map, smaller along edge of
            mask, by default False.
        fwhm_fact : float, optional
            Factor determining smoothing scale at each wavelet scale:
            FWHM = fact * pi / lmax, where lmax is the max wavelet ell.

        kwargs : dict, optional
            Optional keyword arguments to pass to simio.get_sim_mask_fn (currently just
            'galcut' and 'apod_deg'), by default None.

        Notes
        -----
        Unless passed explicitly, the mask and ivar will be loaded at object instantiation time,
        and stored as instance attributes.

        Examples
        --------
        >>> from mnms import noise_models as nm
        >>> wnm = nm.WaveletNoiseModel('s18_03', 's18_04', downgrade=2, notes='my_model')
        >>> wnm.get_model() # will take several minutes and require a lot of memory
                            # if running this exact model for the first time, otherwise
                            # will return None if model exists on-disk already
        >>> imap = wnm.get_sim(0, 123) # will get a sim of split 1 from the correlated arrays;
                                       # the map will have "index" 123, which is used in making
                                       # the random seed whether or not the sim is saved to disk,
                                       # and will be recorded in the filename if saved to disk.
        >>> print(imap.shape)
        >>> (2, 1, 3, 5600, 21600)
        """
        super().__init__(
            *qids, data_model=data_model, preload=preload, ivar=ivar, mask_est=mask_est,
            calibrated=calibrated, downgrade=downgrade, lmax=lmax, mask_version=mask_version,
            mask_name=mask_name, union_sources=union_sources, notes=notes, dtype=dtype, **kwargs
        )

        # save model-specific info
        self._lamb = lamb
        self._smooth_loc = smooth_loc
        self._fwhm_fact = fwhm_fact

    def _get_model_fn(self, split_num):
        """Get a noise model filename for split split_num; return as <str>"""
        return simio.get_wav_model_fn(
            self._qids, split_num, self._lamb, self._lmax, self._smooth_loc, self._fwhm_fact,
            notes=self._notes, data_model=self._data_model, mask_version=self._mask_version,
            bin_apod=self._use_default_mask, mask_name=self._mask_name, calibrated=self._calibrated,
            downgrade=self._downgrade, union_sources=self._union_sources, **self._kwargs
        )

    def _read_model(self, fn):
        """Read a noise model with filename fn; return a dictionary of noise model variables"""
        # read from disk
        sqrt_cov_mat, nm_dict = wavtrans.read_wav(
            fn, extra=['sqrt_cov_ell', 'w_ell']
        )
        nm_dict['sqrt_cov_mat'] = sqrt_cov_mat

        return nm_dict

    def _get_dmap(self, imap):
        """Return the required input difference map for a NoiseModel subclass, from split data imap"""

        # Correction factor turns split difference d_i to split noise n_i
        # (which is the quantity we want to simulate in the end).
        imap *= utils.get_corr_fact(self._ivar)

        return imap

    def _get_model(self, split_num, dmap, verbose=False):
        """
        Return a dictionary of noise model variables for this NoiseModel subclass,
        for split split_num and from difference maps dmap
        """

        sqrt_cov_mat, sqrt_cov_ell, w_ell = wav_noise.estimate_sqrt_cov_wav_from_enmap(
            dmap[:, split_num], self._mask_observed, self._lmax, self._mask_est, lamb=self._lamb,
            smooth_loc=self._smooth_loc, fwhm_fact=self._fwhm_fact
        )

        return {
            'sqrt_cov_mat': sqrt_cov_mat,
            'sqrt_cov_ell': sqrt_cov_ell,
            'w_ell': w_ell
            }

    def _write_model(self, fn, sqrt_cov_mat=None, sqrt_cov_ell=None, w_ell=None):
        """Write sqrt_cov_mat, sqrt_cov_ell, and possibly more noise model variables to filename fn"""
        wavtrans.write_wav(
            fn, sqrt_cov_mat, symm_axes=[[0, 1], [2, 3]],
            extra={'sqrt_cov_ell': sqrt_cov_ell, 'w_ell': w_ell}
        )

    def _get_sim_fn(self, split_num, sim_num, alm=False, mask_obs=True):
        """Get a sim filename for split split_num, sim sim_num; return as <str>"""
        return simio.get_wav_sim_fn(
            self._qids, split_num, self._lamb, self._lmax, self._smooth_loc, self._fwhm_fact,
            sim_num, alm=alm, notes=self._notes, data_model=self._data_model,
            mask_version=self._mask_version, bin_apod=self._use_default_mask, mask_name=self._mask_name,
            calibrated=self._calibrated, downgrade=self._downgrade, union_sources=self._union_sources,
            mask_obs=mask_obs, **self._kwargs
        )

    def _get_sim(self, split_num, seed, mask=None, verbose=False):
        """Return a masked enmap.ndmap sim of split split_num, with seed <sequence of ints>"""

        alm, ainfo = self._get_sim_alm(split_num, seed, mask=None, return_ainfo=True, verbose=verbose)
        sim = utils.alm2map(alm, shape=self._shape, wcs=self._wcs,
                            dtype=np.float32, ainfo=ainfo)
        if mask is not None:
            sim *= mask
        return sim

    def _get_sim_alm(self, split_num, seed, mask=None, return_ainfo=False, verbose=False):
        """Return a masked alm sim of split split_num, with seed <sequence of ints>"""

        # Get noise model variables.
        sqrt_cov_mat = self._nm_dict[split_num]['sqrt_cov_mat']
        sqrt_cov_ell = self._nm_dict[split_num]['sqrt_cov_ell']
        w_ell = self._nm_dict[split_num]['w_ell']

        alm, ainfo = wav_noise.rand_alm_from_sqrt_cov_wav(
            sqrt_cov_mat, sqrt_cov_ell, self._lmax,
            w_ell, dtype=np.complex64, seed=seed)

        # We always want shape (num_arrays, num_splits=1, num_pol, nelem).
        assert alm.ndim == 3, 'Alm must have shape (num_arrays, num_pol, nelem)'
        alm = alm.reshape(alm.shape[0], 1, *alm.shape[1:])

        if mask is not None:
            sim = utils.alm2map(alm, shape=self._shape, wcs=self._wcs,
                                dtype=np.float32, ainfo=ainfo)
            sim *= mask
            utils.map2alm(sim, alm=alm, ainfo=ainfo)
        if return_ainfo:
            return alm, ainfo
        else:
            return alm

class WavFiltTile(NoiseModel):

    def __init__(self):
        pass

    def _get_model(self):
        pass
        # 1. get wavelet noise model including sqrt_cov_ell and wavelet maps
        # 2. use sqrt_cov_ell to flatten noise maps
        # 3. do map->alm->wav on noise maps, flatten using wavelet maps
        # 4. do wav->alm->map to recover full-res flattened noise map
        # 5. make tiled noise model

    def _get_sim(self):
        pass
        # 1. draw tiled noise sim, stitch into flat map
        # 2. do map->alm->wav, unflatten using wavelet maps
        # 3. do wav->alm->map to recover full-res unflattened noise map
        # 4. use sqrt_cov_ell to unflatten
        # 5. use corr_fact to get sim

max_stars_repo_name: ACTCollaboration/mnms

max_issues_repo_name: ACTCollaboration/mnms

max_forks_repo_name: ACTCollaboration/mnms


text:
import numpy as np
import pandas as pd

from utils.random import scaled_inverse_chi_squared

class ThompsonSamplingGaussianSicqPriorWithStartup(object):
    def __init__(self, N, save_log=False):
        self.N = N
        self.threshold = 20
        self.ks = np.ones(N)
        self.mus = np.zeros(N)
        self.vs = np.ones(N)
        self.sigmas = np.ones(N)
        self.save_log = save_log
        self.variances = []
        self.thetas = []
        self.rewards = [[] for _ in range(self.N)]
        self.t = 0

    def initialize(self):
        self.ks = np.ones(self.N)
        self.mus = np.zeros(self.N)
        self.vs = np.ones(self.N)
        self.sigmas = np.ones(self.N)
        self.variances = []
        self.thetas = []
        self.rewards = [[] for _ in range(self.N)]
        self.t = 0

    def select_arm(self):
        if self.t <= self.threshold:
            return np.random.randint(self.N)
        return np.argmax(self.estimate_mean())

    def estimate_mean(self):
        # For each arm i=1,...,N, sample random value from sicq distribution
        variance = [scaled_inverse_chi_squared(self.vs[i], self.sigmas[i])
                    for i in range(self.N)]
        # For each arm i=1,...,N, sample random value from normal distribution
        theta = [np.random.normal(self.mus[i],
                                  np.math.sqrt(variance[i] / self.ks[i]))
                 for i in range(self.N)]
        if self.save_log:
            self.thetas.append(theta)
            self.variances.append(variance)
        return theta

    def update_param(self, arm_id, reward):
        if self.t <= self.threshold:
            self.rewards[arm_id].append(reward)
            self.vs[arm_id] = self.vs[arm_id] + 1.0
            self.ks[arm_id] = self.ks[arm_id] + 1.0

            if self.t == self.threshold:
                for i in range(self.N):
                    if len(self.rewards[i]) <= 1:
                        self.mus[i] = 0.0
                        self.sigmas[i] = 1.0
                    else:
                        self.mus[i] = np.mean(self.rewards[i])
                        self.sigmas[i] = np.var(self.rewards[i])
                self.mus = np.array(self.mus)
                self.sigmas = np.array(self.sigmas)
        else:
            sigma = self.sigmas[arm_id]
            v = self.vs[arm_id]
            mu = self.mus[arm_id]
            k = self.ks[arm_id]
            d = reward - mu
            # update parameter of sicq distribution
            self.sigmas[arm_id] = sigma * v / (v + 1.0) + d * d * k / (
                        (v + 1.0) * (k + 1.0))
            self.vs[arm_id] = v + 1.0
            # update parameter of normal distribution
            self.mus[arm_id] = (mu * k + reward) / (k + 1.0)
            self.ks[arm_id] = k + 1.0
        self.t += 1

    def save(self, folder_name):
        if self.save_log:
            self.thetas = pd.DataFrame(self.thetas)
            self.thetas.to_csv(folder_name + '/theta.csv')
            self.variances = pd.DataFrame(self.variances)
            self.variances.to_csv(folder_name + '/variance.csv')

ext: py

max_stars_repo_name: bakanaouji/ThompsonSamplingExp

max_issues_repo_name: bakanaouji/ThompsonSamplingExp

max_forks_repo_name: bakanaouji/ThompsonSamplingExp

reason: import numpy


text:
#import sys
from numpy import *

def quat_to_transmat(Q):
   w = Q[0,0]
   x = Q[0,1]
   y = Q[0,2]
   z = Q[0,3]
   tx = Q[0,4]
   ty = Q[0,5]
   tz = Q[0,6]
   # left handed or right? post/pre-multiply?
   T = matrix([[1 - 2*y*y - 2*z*z,  2*x*y - 2*z*w,	 2*x*z + 2*y*w,	      tx],
               [2*x*y + 2*z*w,	    1 - 2*x*x - 2*z*z,	 2*y*z - 2*x*w,	      ty],
               [2*x*z - 2*y*w,	    2*y*z + 2*x*w,	 1 - 2*x*x - 2*y*y,   tz]])
   #T[:,0:3] = T[:,0:3].T
   #T[:,3] = -T[:,3]
   return T

set_printoptions(precision=2)

#savetxt('K.txt', K, fmt='%.5f')
# only for array type txt file
#K = matrix(fromfile("K.txt", sep='\n').reshape(3,3))

# load camera matrix C
#C = eye(3)
#savetxt('C.txt', C, fmt='%.2f')
C = matrix(loadtxt('C.txt'))
print('C:\n', C, '\n')

# load 3D points
#pts3d = matrix([[0,0,6,1], [1,1,6,1], [-4,4,6,1], [2,-3,6,1], [3,3,7,1], [5,4,7,1], [-1,-2,7,1], [-2,5,7,1], [1,1,8,1], [2,2,8,1], [3,-3,8,1], [6,7,8,1]]).transpose()
#pts3d = matrix([[0,0,6,1], [5,5,6,1], [-4,10,6,1], [20,-3,6,1], [9,3,7,1], [25,11,7,1], [-10,-22,7,1],
#	        [-4,15,7,1], [10,11,8,1], [22,12,8,1], [23,-13,8,1], [16,7,8,1]]).transpose()
#savetxt('pts3d.txt', pts3d, fmt='%.2f')
pts3d = matrix(loadtxt('pts3d.txt'))
print('3D points:\n', pts3d, '\n')

# load camera movement transformations
#rt = matrix([[1,0,0,0,0,0,0], [1,0,0,0,0,0,1], [1,0,0,0,1,0,2], [1,0,0,0,2,0,3], [1,0,0,0,3,0,4]])
#savetxt('rt.txt', rt, fmt='%.2f')
rt = matrix(loadtxt('rt_rot.txt'))
print('camera rotation-translation params:\n', rt, '\n')

[r, c] = rt.shape
[d, npts] = pts3d.shape
nframes = r
projs = zeros([nframes, 3, npts])
for i in range(nframes):
   T = quat_to_transmat(rt[i,:])
   print("using transformation:\n", T, '\n')
   pts2d = C * T * pts3d
   pts2d = pts2d / pts2d[2,:]
   projs[i,:,:] = pts2d
   #append(projs, pts2d)
   print('Frame ', i, ':\n', pts2d, '\n')

print("All projections: \n")
print(projs)

fp = open("pts_data.txt", "w")
SPLIT_PTS = 2
for i in range(npts-SPLIT_PTS):
   fp.write(str(pts3d[0,i]) + " " + str(pts3d[1,i]) + " " + str(pts3d[2,i]) + " " + str(nframes))
   for f in range(nframes):
      fp.write(" " + str(f) + " " + str(projs[f,0,i]) + " " + str(projs[f,1,i]))
   fp.write("\n")

split_frames = 3
for i in range(npts-SPLIT_PTS, npts):
   fp.write(str(pts3d[0,i]) + " " + str(pts3d[1,i]) + " " + str(pts3d[2,i]) + " " + str(nframes-split_frames))
   for f in range(split_frames, nframes):
      fp.write(" " + str(f) + " " + str(projs[f,0,i]) + " " + str(projs[f,1,i]))
   fp.write("\n")

fp.close()

# now run: ./eucsbademo rt_pts.txt pts_data.txt C.txt

reason: from numpy


text:
from __future__ import annotations
from typing import *
import numpy as np #type: ignore
from enum import IntEnum, auto

FeatureExtractor = Callable[[np.ndarray], 'ImageFeatures']

#Features are always arranged row major
#This way accessing a single feature is faster
class ImageFeatures: #type: ignore
    __slots__=['FeatureDim', 'NumFeatureVecs','Features','IterIndex']
    def __init__(self, FeatureDim:int, NumFeatureVecs: int, Features:np.ndarray):
        self.FeatureDim: int = FeatureDim
        self.NumFeatureVecs: int = NumFeatureVecs
        self.Features: np.ndarray = Features
        self.IterIndex:int=0

    #Resets the iteration index of the feature vector
    def __iter__(self)->ImageFeatures:
        self.IterIndex = 0
        return self

    #Returns the next feature vector
    def __next__(self)->np.ndarray:
        if self.IterIndex >= self.NumFeatureVecs:
            raise StopIteration
        if len(self.Features.shape) <= 1:
            self.IterIndex+=1
            return self.Features
        res = self.Features[self.IterIndex,:]
        self.IterIndex+=1
        return res

    #Returns how many feature vectors are contained
    def __len__(self)->int:
        return self.NumFeatureVecs

    #Allows indexing into the features, always returns a ImageFeatures type
    def __getitem__(self, key: Union[int, slice])->np.ndarray:
        if not(isinstance(key, slice) or isinstance(key, int)):
            raise KeyError("ImageFeatures only accepts integers and slices as subscripts")
        features = self.Features[key,:]
        return ImageFeatures(self.FeatureDim, features.shape[0], features)

    def __setitem__(self, key, item):
        if not(isinstance(key, slice) or isinstance(key, int)):
            raise KeyError("ImageFeatures only accepts integers and slices as subscripts when setting values")
        self.Features[key, :] = item

    def __str__(self)->str:
        return f"ImageClass(FeatureDim={self.FeatureDim}, NumFeatureVecs={self.NumFeatureVecs}, IterIndex={self.IterIndex}, Features={self.Features})"


text:
"""
openPMD-specific fields

"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
# Copyright (c) 2015, Daniel Grassinger (HZDR)
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import numpy as np
from yt.funcs import mylog
import yt.utilities.physical_constants
from yt.fields.field_info_container import \
    FieldInfoContainer

# Function to calculate new fields out of other fields
# if a error occurs the field is not calculated

def _kinetic_energy(field, data):
    #calculate kinetic energy out of momentum
    c = 2.997e8 # velocity of light
    ke =       ( data["particle_momentum_x"]**2
               + data["particle_momentum_y"]**2
               + data["particle_momentum_z"]**2) * c
    return ke

# This function does no work !!
def setup_momentum_to_velocity(self, ptype):
        def _get_vel(axis):
            def velocity(field, data):
		c = 2.997e8
		moment = data[ptype,"particle_momentum_%" % axis]
                return moment / ((data[ptype,"particle_mass"]*data[ptype,"particle_weighting"])**2 + (moment**2)/(c**2))**0.5
            return velocity

        for ax in 'xyz':
            self.add_field((ptype, "particle_velocity_%s" % ax),
                           function=_get_vel(ax),
                           units="m/s")

def setup_poynting_vector(self):
        def _get_poyn(axis):
            def poynting(field, data):
		Efieldx = data["E_x"]
		Efieldy = data["E_y"]
		Efieldz = data["E_z"]
		Bfieldx = data["B_x"]
		Bfieldy = data["B_y"]
		Bfieldz = data["B_z"]

		u = 79577.4715459 # = 1/magnetic permeability

		if(axis == 'x'):
		    return u * (Efieldy *Bfieldz-Efieldz*Bfieldy)
		elif(axis == 'y'):
		    return u * (Efieldz *Bfieldx-Efieldx*Bfieldz)
		elif(axis == 'z'):
		    return u * (Efieldx *Bfieldy-Efieldy*Bfieldx)

            return poynting
        for ax in 'xyz':
            self.add_field(("openPMD", "poynting_vector_%s" % ax),
                           function=_get_poyn(ax),
                           units="T*V/m") #N/(m*s)

# We need to specify which fields we might have in our dataset.  The field info
# container subclass here will define which fields it knows about.  There are
# optionally methods on it that get called which can be subclassed.

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
#
# This class defines, which fields and particle fields could be in the HDF5-file
# The field names have to match the names in "OpenPMDHierarchy" in data_structures.py
# This also defines the units of the fields

class OpenPMDFieldInfo(FieldInfoContainer):
    known_other_fields = (
        # Each entry here is of the form
        # ( "name", ("units", ["fields", "to", "alias"], # "display_name")),
        ("B_x", ("T", [], None)),
	("B_y", ("T", [], None)),
	("B_z", ("T", [], None)),
	("E_x", ("V/m", [], None)),
	("E_y", ("V/m", [], None)),
	("E_z", ("V/m", [], None)),
	("J_x", ("A/m**2", [], None)),
	("J_y", ("A/m**2", [], None)),
	("J_z", ("A/m**2", [], None)),
	("rho", ("kg*s/m**3.0", [], None)),

   )

    known_particle_fields = (
        # Identical form to above
        # ( "name", ("units", ["fields", "to", "alias"], # "display_name")),
	("particle_charge", ("A*s", [], None)),
	("particle_mass", ("kg", [], None)),
	("particle_momentum_x", ("kg*m/s", [], None)),
        ("particle_momentum_y", ("kg*m/s", [], None)),
        ("particle_momentum_z", ("kg*m/s", [], None)),
        ("particle_position_x", ("m", [], None)),
        ("particle_position_y", ("m", [], None)),
        ("particle_position_z", ("m", [], None)),
	("particle_positionOffset_x", ("m", [], None)),
        ("particle_positionOffset_y", ("m", [], None)),
        ("particle_positionOffset_z", ("m", [], None)),
	("particle_weighting", ("", [], None)),
	("particle_kinetic_energy", ("dimensionless", [], None)),

   )

    def __init__(self, ds,field_list):
        super(OpenPMDFieldInfo, self).__init__(ds, field_list)
        # If you want, you can check field_list

    # Here you can create functions to calculate out of existing fields the
    # values of new fields e.g. calculate out of E-field and B-field the
    # Poynting vector
    def setup_fluid_fields(self):
        # Here we do anything that might need info about the dataset.
        # You can use self.alias, self.add_output_field and self.add_field .

	setup_poynting_vector(self)

    # !!! You have to call the function of the parent class to load particles
    def setup_particle_fields(self, ptype):
        # This will get called for every particle type.

	self.add_field((ptype,"particle_kinetic_energy"),function=_kinetic_energy,units="dimensionless")
	setup_momentum_to_velocity(self, ptype)

    # !!! Has to be called to load particles
	super(OpenPMDFieldInfo, self).setup_particle_fields(ptype)


text:
import numpy as np
from keras import layers, models
from tensorflow import keras

from preprocessing import value_preprocessing
from train_framework import train, calculate_test_mae

def create_vgg16_3d(dense=False, batch_norm=True, weights=None, input_shape=(260, 100, 15, 1)):
    """
    Creates slightly modified VGG16 model, ported to 3d with trainable BatchNormalization
    """

    def create_conv(filter, kernel_size, name):
        def conv_wrapper(inp):
            x = layers.Conv3D(filter, kernel_size, padding='same', name=name)(inp)
            if batch_norm:
                x = layers.BatchNormalization()(x)
            x = layers.Activation('relu')(x)
            return x

        return conv_wrapper

    img_input = layers.Input(shape=input_shape, name='input')

    # block 1
    x = create_conv(8, (3, 3, 3), name='block1_conv1')(img_input)
    x = create_conv(8, (3, 3, 3), name='block1_conv2')(x)
    x = layers.MaxPooling3D((2, 2, 2), strides=(2, 2, 2), name='block1_pool')(x)

    # block 2
    x = create_conv(16, (3, 3, 3), name='block2_conv1')(x)
    x = create_conv(16, (3, 3, 3), name='block2_conv2')(x)
    x = layers.MaxPooling3D((2, 2, 2), strides=(2, 2, 2), name='block2_pool')(x)

    # block 3
    x = create_conv(32, (3, 3, 3), name='block3_conv1')(x)
    x = create_conv(32, (3, 3, 3), name='block3_conv2')(x)
    x = create_conv(32, (3, 3, 3), name='block3_conv3')(x)
    x = layers.MaxPooling3D((2, 2, 2), strides=(2, 2, 2), name='block3_pool')(x)

    # block 4
    x = create_conv(32, (3, 3, 3), name='block4_conv1')(x)
    x = create_conv(32, (3, 3, 3), name='block4_conv2')(x)
    x = create_conv(32, (3, 3, 3), name='block4_conv3')(x)
    # delete MaxPooling to allow for 5th block

    # block 5
    x = create_conv(64, (3, 3, 3), name='block5_conv1')(x)
    x = create_conv(64, (3, 3, 3), name='block5_conv2')(x)
    x = create_conv(64, (3, 3, 3), name='block5_conv3')(x)

    if dense:
        x = layers.Flatten()(x)
        x = layers.Dense(4096)(x)
        x = layers.BatchNormalization()(x)
        x = layers.Dense(2048)(x)
        x = layers.BatchNormalization()(x)
        x = layers.Dense(1)(x)
    else:
        x = layers.GlobalMaxPooling3D()(x)

    model = models.Model(img_input, x, name='vgg16_3d')

    if weights is not None:
        model.load_weights(weights, by_name=True)
    return model

def preprocessing(ct):
    """wrapper function to pass to generators"""
    return value_preprocessing(ct, False)

def output_reshape(ct):
    """wrapper function to pass to generators"""
    return np.moveaxis(ct, 1, -1)

name = 'vgg16_3d_bn_better0.1'
patch_size = (260, 100, 15)
model = create_vgg16_3d(dense=True)
optimizer = keras.optimizers.Adam(0.01, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
train(model, optimizer, epochs=1000, batch_size=32, patch_size=patch_size, num_validation=32, name=name, loss='mae',
      preprocessing_func=preprocessing, output_reshape_func=output_reshape, training_generator_threads=6,
      training_sample_cache=16)
calculate_test_mae(model, optimizer, 'mae', 4, patch_size, preprocessing, output_reshape)

max_stars_repo_name: leonvol/bwki-boneage-experiments

max_issues_repo_name: leonvol/bwki-boneage-experiments

max_forks_repo_name: leonvol/bwki-boneage-experiments


text:
import numpy as np
import pickle
import sys

def parse_energies(filename):
    """
    Return all the energies in a molpro out-file
    """
    with open(filename) as f:
        lines = f.readlines()

    if "Molpro calculation terminated" not in lines[-1]:
        print("Error parsing %s" % filename)
        return

    energies = lines[-3].split()[:8]
    return energies

def parse_filename(filename):
    tokens = filename.split("/")[-1].split("_")
    tokens[-1] = tokens[-1][:-4]
    tokens[0] = int(tokens[0])
    return tokens

if __name__ == "__main__":
    filenames = sys.argv[1:]
    energies = {}
    for filename in filenames:
        energy = parse_energies(filename)
        if energy is None:
            continue
        atom, d1, d2 = parse_filename(filename)
        if atom not in energies: energies[atom]=[]

        energies[atom].append(np.asarray([d1,d2] + energy, dtype=float))

    for key in energies.keys():
        energies[key] = np.asarray(energies[key], dtype=float)

    with open("data.pkl", "wb") as f:
        pickle.dump(energies, f)

max_stars_repo_name: larsbratholm/plotting

max_issues_repo_name: larsbratholm/plotting

max_forks_repo_name: larsbratholm/plotting


text:
"""
Make models from a template
"""
import numpy as np
from scipy.optimize import nnls

from standard_glm import StandardGlm
from spatiotemporal_glm import SpatiotemporalGlm
from shared_tuningcurve_glm import SharedTuningCurveGlm
from simple_weighted_model import SimpleWeightedModel
from simple_sparse_model import SimpleSparseModel
from sparse_weighted_model import SparseWeightedModel
from sbm_weighted_model import SbmWeightedModel
from distance_weighted_model import DistanceWeightedModel

import copy

def make_model(template, N=None, dt=None):
    """ Construct a model from a template and update the specified parameters
    """
    if isinstance(template, str):
        # Create the specified model
        if template.lower() == 'standard_glm' or \
           template.lower() == 'standardglm':
            model = copy.deepcopy(StandardGlm)
        elif template.lower() == 'spatiotemporal_glm':
            model = copy.deepcopy(SpatiotemporalGlm)
        elif template.lower() == 'shared_tuning_curve':
            model = copy.deepcopy(SharedTuningCurveGlm)
        elif template.lower() == 'simple_weighted_model' or \
             template.lower() == 'simpleweightedmodel':
            model = copy.deepcopy(SimpleWeightedModel)
        elif template.lower() == 'simple_sparse_model' or \
             template.lower() == 'simplesparsemodel':
            model = copy.deepcopy(SimpleSparseModel)
        elif template.lower() == 'sparse_weighted_model' or \
             template.lower() == 'sparseweightedmodel':
            model = copy.deepcopy(SparseWeightedModel)
        elif template.lower() == 'sbm_weighted_model' or \
             template.lower() == 'sbmweightedmodel':
            model = copy.deepcopy(SbmWeightedModel)
        elif template.lower() == 'distance_weighted_model' or \
             template.lower() == 'distanceweightedmodel':
            model = copy.deepcopy(DistanceWeightedModel)
        else:
            raise Exception("Unrecognized template model: %s!" % template)
    elif isinstance(template, dict):
        model = copy.deepcopy(template)
    else:
        raise Exception("Unrecognized template model!")

    # Override template model parameters
    if N is not None:
        model['N'] = N

    if dt is not None:
        model['dt'] = dt

    # # Update other parameters as necessary
    # if template.lower() == 'distance_weighted_model' or \
    #    template.lower() == 'distanceweightedmodel':
    #     #model['network']['graph']['location_prior']['sigma'] = N/2.0/3.0
    #     model['network']['graph']['location_prior']['mu'] = \
    #         np.tile(np.arange(N).reshape((N,1)),
    #                 [1,model['network']['graph']['N_dims']]).ravel()

    return model

def stabilize_sparsity(model):
    """ Adjust the sparsity level for simple weighted models
        with Gaussian weight models and Bernoulli adjacency matrices.
        The variance of a Gaussian N(0,sigma) times a Bernoulli(rho) is
            E[B*N] = 0
            Var[B*N] = E[(B*N)**2] = rho*sigma^2 + (1-rho)*0 = rho*sigma^2

        Hence the eigenvalues will be distributed in a complex disk of radius
            \sqrt(N) * \sqrt(rho) * sigma
        For this to be less than (1-delta), we have
            \sqrt(rho) < (1-delta)/sqrt(N)/sigma
            rho < (1-delta)**2 / N / sigma**2
    """
    N = model['N']
    imp_model = model['impulse']
    weight_model = model['network']['weight']
    graph_model = model['network']['graph']

    delta = 0.3
    maxeig = 1.0-delta

    if graph_model['type'].lower() == 'erdos_renyi':
        if weight_model['prior']['type'].lower() == 'gaussian':

            # If we have a refractory bias on the diagonal weights then
            # we can afford slightly stronger weights
            if 'refractory_prior' in weight_model:
                maxeig -= weight_model['refractory_prior']['mu']

            sigma = weight_model['prior']['sigma']
            stable_rho = maxeig**2/N/sigma**2
            stable_rho = np.minimum(stable_rho, 1.0)
            print "Setting sparsity to %.2f for stability." % stable_rho
            graph_model['rho'] = stable_rho

    elif graph_model['type'].lower() == 'sbm':
        # Things are trickier in the SBM case because the entries in A
        # are not iid. But, we can still make some approximations by
        # thinking about the mean and variance within a single block, wherein
        # the entries really are i.i.d. Then we scale the eigs of
        # each block by N/R, as if the blocks were equal size and the
        # interdependencies between blocks did not matter. Obviously,
        # this is a hack.
        R = graph_model['R']
        if weight_model['type'].lower() == 'sbm':
            sig_mu = weight_model['sigma_mu']
            sig_w = weight_model['sigma_w']
        elif weight_model['type'].lower() == 'gaussian':
            sig_mu = 0.0
            sig_w = weight_model['sigma']
        else:
            raise Exception("Unrecognized weight model for SBM graph: %s" % weight_model['type'])

        # var_AW =  1./4. * (3.*sig_mu)**2 + sig_w**2
        # mean_lam_max = sig_mu * np.sqrt(R) * N/float(R) + 3*sig_w
        # sig_lam_max = np.sqrt(var_AW)
        # ub_lam_max = mean_lam_max + 3*sig_lam_max
        #
        # var_B = (((1.0-mean_lam_max)/3.0)**2 - sig_w**2) / (3*sig_mu)**2
        #
        # print "Setting b0 to %.2f to achive sparsity of %.2f." % (graph_model['b0'],stable_rho)

        #elif weight_model['type'].lower() == 'constant' and \
        #     imp_model['type'].lower() == 'basis':
        #    sigma = imp_model['sigma']
        #    maxeig = 0.7
        #
        #    # TODO Figure out how sigma actually relates to the eigenvalues
        #    stable_rho = maxeig/N/sigma**2
        #    stable_rho = np.minimum(stable_rho, 1.0)
        #    print "Setting sparsity to %.2f for stability." % stable_rho
        #    graph_model['rho'] = stable_rho\

    elif graph_model['type'].lower() == 'distance':
        # It's hard to characterize this exactly, but suppose that we had an exact
        # banded adjacency matrix where the -b-th lower diagonal to the +b-th upper diagonal
        # are normally distributed, and the rest of the matrix is zero.
        #
        # The max eigenvalue of this matrix tends to grow proportionally to log N
        # and b, and is scaled by the standard deviation of the weights, sigma.
        # The mean refractory weight adds a bias to the max eigenvalue.
        #
        # To estimate the bandwidth (i.e. 2b in this example), suppose that the
        # probability of connection p(A_n,n') decays exponentially with |n-n'| at
        # length scale d. That is, when |n-n'| > d, the probability has decreased to
        # about 1/3. Hence, at distance 2d, the probability of interaction is about 10%.
        # Thus, let's say the bandwidth of this matrix is roughly 2*(2d) = 4d,
        #
        # So, given the bandwidth of 4d and the matrix size N, we will scale the weight
        # distribution to hopefully achieve stability
        b = 8.0 * graph_model['delta']

        # Constant scaling factors
        # max eig ~= mu_ref + .25*sigma_W * (b+2*log(N))
        # So set sigma_W ~= 4*(1-mu_ref)/(b+2*log N)
        if weight_model['type'].lower() == 'gaussian':
            mu_ref = weight_model['refractory_prior']['mu']
            sig_w = 4.0*(1.0-mu_ref)/(b+2.0*np.log(model['N']))

            print "Setting sig_w to %.3f to ensure stability" % sig_w
            weight_model['prior']['sigma'] = sig_w
        else:
            raise Exception("Unrecognized weight model for distance graph: %s" % weight_model['type'])

def check_stability(model, x, N):
    """
    Check the stability of model parameters
    """

    if model['network']['weight']['type'].lower() == 'gaussian':
        Weff = x['net']['graph']['A'] * np.reshape(x['net']['weights']['W'], (N,N))
        maxeig = np.amax(np.real(np.linalg.eig(Weff)[0]))
        print "Max eigenvalue of Weff: %.2f" % maxeig
        return maxeig < 1
    else:
        print "Check stability: unrecognized model type. Defaulting to true."
        return True

def convert_model(from_popn, from_model, from_vars, to_popn, to_model, to_vars):
    """ Convert from one model to another model of a different type
        Generally this will involve projecting impulse responses, etc.
        It's hairy business.
    """

    # Idea: Get the state of the GLMs, e.g. the impulse responses, etc.
    #       Project those states onto the parameters of the to-model
    N = from_popn.N
    from_state = from_popn.eval_state(from_vars)
    to_state = to_popn.eval_state(to_vars)

    conv_vars = None
    if from_model['impulse']['type'].lower() == 'basis':
        if to_model['impulse']['type'].lower() == 'normalized' or \
           to_model['impulse']['type'].lower() == 'dirichlet':
            import copy
            conv_vars = copy.deepcopy(to_vars)

            # To convert from basis -> normalized, project the impulse
            # responses onto the normalized basis, divide by the area
            # under the curve to get the weight.
            W = np.zeros((N,N))
            for n2 in np.arange(N):
                B = to_state['glms'][n2]['imp']['basis'].shape[1]
                w_ir_n2 = np.zeros((N,B))
                for n1 in np.arange(N):
                    # Solve a nonnegative least squares problem
                    (w_ir_n1n2p, residp) = nnls(to_state['glms'][n2]['imp']['basis'],
                                                from_state['glms'][n2]['imp']['impulse'][n1,:])
                    (w_ir_n1n2n, residn) = nnls(to_state['glms'][n2]['imp']['basis'],
                                                -1.0*from_state['glms'][n2]['imp']['impulse'][n1,:])

                    # Take the better of the two solutions
                    if residp < residn:
                        Wsgn = 1.0
                        w_ir_n1n2 = w_ir_n1n2p
                    else:
                        Wsgn = -1.0
                        w_ir_n1n2 = w_ir_n1n2n

                    # Normalized weights must be > 0, sum to 1
                    w_ir_n1n2 = w_ir_n1n2
                    w_ir_n1n2 = np.clip(w_ir_n1n2,0.001,np.Inf)
                    # Normalize the impulse response to get a weight
                    W[n1,n2] = Wsgn*np.sum(w_ir_n1n2)

                    # Set impulse response to normalized impulse response
                    w_ir_n2[n1,:] = w_ir_n1n2 / np.sum(w_ir_n1n2)

                # Update to_vars
                if to_model['impulse']['type'].lower() == 'normalized':
                    conv_vars['glms'][n2]['imp']['w_lng'] = np.log(w_ir_n2.flatten())
                if to_model['impulse']['type'].lower() == 'dirichlet':
                    for n1 in range(N):
                        # Scale up the weights such that the average is preserved
                        alpha = to_popn.glm.imp_model.alpha
                        B = to_popn.glm.imp_model.B
                        conv_vars['glms'][n2]['imp']['g_%d' % n1] = alpha * B * w_ir_n2[n1,:]

            # Update to_vars
            conv_vars['net']['weights']['W'] = W.flatten()

            # Threshold the adjacency matrix to start with the right level of sparsity
            if 'rho' in to_model['network']['graph'].keys():
                W_sorted = np.sort(np.abs(W.ravel()))
                thresh = W_sorted[np.floor((1.0-2.0*to_model['network']['graph']['rho'])*(N**2-N)-N)]
                conv_vars['net']['graph']['A'] = (np.abs(W) >= thresh).astype(np.int8)
            else:
                conv_vars['net']['graph']['A'] = np.ones((N,N), dtype=np.int8)

    # Copy over the bias
    for n in np.arange(N):
        conv_vars['glms'][n]['bias']['bias'] = from_vars['glms'][n]['bias']['bias']

    # Copy over the background params
    if 'sharedtuningcurves' in to_model['latent'] and \
        from_model['bkgd']['type'] == 'spatiotemporal':
        convert_stimulus_filters_to_sharedtc(from_popn, from_model, from_vars,
                                             to_popn, to_model, conv_vars)

    return conv_vars

def convert_stimulus_filters_to_sharedtc(from_popn, from_model, from_vars, to_popn, to_model, to_vars):
    """
    Convert a set of stimulus filters to a shared set of tuning curves
    """
    # Get the spatial component of the stimulus filter for each neuron
    N = from_popn.N
    R = to_model['latent']['sharedtuningcurves']['R']
    from_state = from_popn.eval_state(from_vars)

    locs = np.zeros((N,2))
    local_stim_xs = []
    local_stim_ts = []
    for n in range(N):
        s_glm = from_state['glms'][n]

        # to_state = to_popn.eval_state(to_vars)
        assert 'stim_response_x' in s_glm['bkgd']

        # Get the stimulus responses
        stim_x = s_glm['bkgd']['stim_response_x']
        stim_t = s_glm['bkgd']['stim_response_t']
        loc_max = np.argmax(np.abs(stim_x))

        if stim_x.ndim == 2:
            locsi, locsj = np.unravel_index(loc_max, stim_x.shape)
            locs[n,0], locs[n,1] = locsi.ravel(), locsj.ravel()

            # # TODO: Test whether we have an issue with unraveling the data.
            # locs[n,0], locs[n,1] = locsj.ravel(), locsi.ravel()

        # Get the stimulus response in the vicinity of the mode
        # Create a meshgrid of the correct shape, centered around the max
        max_rb = to_model['latent']['latent_location']['location_prior']['max0']
        max_ub = to_model['latent']['latent_location']['location_prior']['max1']
        gsz = to_model['latent']['sharedtuningcurves']['spatial_shape']
        gwidth = (np.array(gsz) - 1)//2
        lb = max(0, locs[n,0]-gwidth[0])
        rb = min(locs[n,0]-gwidth[0]+gsz[0], max_rb)
        db = max(0, locs[n,1]-gwidth[1])
        ub = min(locs[n,1]-gwidth[1]+gsz[1], max_ub)
        grid = np.ix_(np.arange(lb, rb).astype(np.int),
                      np.arange(db, ub).astype(np.int))

        # grid = grid.astype(np.int)
        # Add this local filter to the list
        local_stim_xs.append(stim_x[grid])
        local_stim_ts.append(stim_t)

    # Cluster the local stimulus filters
    from sklearn.cluster import KMeans
    flattened_filters_x = np.array(map(lambda f: f.ravel(), local_stim_xs))
    flattened_filters_t = np.array(map(lambda f: f.ravel(), local_stim_ts))
    km = KMeans(n_clusters=R)
    km.fit(flattened_filters_x)
    Y = km.labels_
    print 'Filter cluster labels from kmeans: ',  Y

    # Initialize type based on stimulus filter
    to_vars['latent']['sharedtuningcurve_provider']['Y'] = Y

    # Initialize shared tuning curves (project onto the bases)
    from pyglm.utils.basis import project_onto_basis
    for r in range(R):
        mean_filter_xr = flattened_filters_x[Y==r].mean(axis=0)
        mean_filter_tr = flattened_filters_t[Y==r].mean(axis=0)

        # TODO: Make sure the filters are being normalized properly!

        # Project the mean filters onto the basis
        to_vars['latent']['sharedtuningcurve_provider']['w_x'][:,r] = \
            project_onto_basis(mean_filter_xr,
                               to_popn.glm.bkgd_model.spatial_basis).ravel()

        # Temporal part of the filter
        temporal_basis = to_popn.glm.bkgd_model.temporal_basis
        t_temporal_basis = np.arange(temporal_basis.shape[0])
        t_mean_filter_tr = np.linspace(0, temporal_basis.shape[0]-1, mean_filter_tr.shape[0])
        interp_mean_filter_tr = np.interp(t_temporal_basis, t_mean_filter_tr, mean_filter_tr)
        to_vars['latent']['sharedtuningcurve_provider']['w_t'][:,r] = \
            project_onto_basis(interp_mean_filter_tr, temporal_basis).ravel()

    # Initialize locations based on stimuls filters
    to_vars['latent']['location_provider']['L'] = locs.ravel().astype(np.int)

reason: import numpy,from scipy


text:
"""
Illustration of Oja's rule (Hebbian learning) in Spiking Neural Networks using LIF neuron.
Author: Ali Safa
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
import matplotlib.animation as animation

"""
Functions
"""
def sample_spherical(npoints, ndim): #sample a vector of dimension "ndim" from the unit sphere randomly
    vec = np.random.randn(ndim, npoints)
    vec /= np.linalg.norm(vec, axis=0)
    return vec[:,0]

def simulate_neuron_Oja(Tsim, dt, trc, tref, vrest, vth, Jbias, alpha, e, input_vec, tau, eta):
    N = int(np.round(Tsim/dt))
    Vprev = 0
    Jprev = 0
    spike_train = np.zeros(N)
    Vhist = np.zeros(N)
    psc = np.zeros(N)
    W_vec = np.zeros((N, len(e)))
    W = alpha*e
    W_vec[0,:] = W
    Jbias_vec = np.zeros(N)
    Jbias_vec[0] = Jbias
    mutex = 0
    for i in range(N):
        J = np.inner(W, input_vec[i,:]) + Jbias
        if mutex == 0:
            V = (J + Jprev - (1-2*trc/dt)*Vprev)/(1+2*trc/dt) #bilinear transform
            if V < vrest:
                V = vrest
            elif V > vth:
                spike_train[i] = 1
                V = vrest
                mutex = np.round(tref/dt)
            Vhist[i] = V
            Jprev = J
            Vprev = V
        else:
            mutex -= 1

        if i > 0:
            psc[i] = (spike_train[i] + spike_train[i-1])*(1/(1+2*tau/dt)) - ((1-2*tau/dt)/(1+2*tau/dt))*psc[i-1]
            #update weights following Oja's rule
            DELW = psc[i]*input_vec[i,:] - (psc[i]**2)*W
            W = W + eta*DELW
            W_vec[i,:] = W
            Jbias_vec[i] = Jbias

    return Vhist, spike_train, W_vec, psc, Jbias_vec

def PSC_filter(Tsim, dt, tau):
    t = np.linspace(0,Tsim,int(np.round(Tsim/dt)))
    h = np.exp(-(t-Tsim/2)/tau)
    h[0:len(h)//2] = 0
    h = (1/dt)*h/np.sum(h)
    return h

def normalize_imges(data): #normalize pixel values to [-1, 1]
    for i in range(data.shape[0]):
        img = data[i]
        data[i] = 2*(img - min(img))/(max(img) - min(img)) - 1
    return data

"""
Main
"""
np.random.seed(2) #to get reproducable results
plt.close('all')

D = 64 #data dimensions
F_max_l = 100 #100
F_max_h = 200
in_l = -1.0
in_h = 1.0
tref = 0.002 #2ms
trc = 0.02 #20ms
Tsim = 0.2
dt = 0.002
vrest = 0
vth = 1
Tlen = int(np.round(Tsim/dt))

digits = load_digits(n_class=1)
x = normalize_imges(digits.data)
input_vec = x[:Tlen,:]

amax = np.random.uniform(F_max_l,F_max_h,1) # maximum rate uniformly distributed between 100 and 200 HZ
#xi = np.random.uniform(in_l+0.05,in_h-0.05,1) # x-intercept
xi = np.random.uniform(in_l+0.05,in_h-0.05,1) # new idea x-intercept
alpha = (1/(1-np.exp((tref - 1/amax)/trc)) - 1)/(1-xi) #for LIF neuron
Jbias = 1-xi*alpha
e = sample_spherical(1, D)

Vhist, spike_train, W_vec, psc, Jbias_vec = simulate_neuron_Oja(Tsim, dt, trc, tref, vrest, vth, Jbias, alpha, e, input_vec, 0.05, 1)

plt.figure(1)
plt.gray()
plt.axis('off')
im = np.reshape(W_vec[-1,:], (8,8))
plt.imshow(im)

#if True:
plt.figure(2)
fig, ax = plt.subplots(1,1)
plt.gray()
plt.axis('off')
plt.title('On-line and Unsuppervised Pattern Learning', fontsize=15, color='black')

def img_anim(i):
    im = np.reshape(W_vec[i,:], (8,8))
    ax.matshow(im)
    print("Time step: " + str(i) + "/" + str(Tlen))

anim2 = animation.FuncAnimation(fig, img_anim, frames=Tlen, interval=1)
#plt.show()
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
#anim2.save('Hebbian.mp4', writer=writer)

plt.figure(4)
plt.subplot(2,2,1)
plt.gray()
plt.axis('off')
plt.title('Time step: 0ms', fontsize=15, color='black')
im = np.reshape(W_vec[0,:], (8,8))
plt.imshow(im)
plt.subplot(2,2,2)
plt.gray()
plt.axis('off')
plt.title('Time step: 20ms', fontsize=15, color='black')
im = np.reshape(W_vec[10,:], (8,8))
plt.imshow(im)
plt.subplot(2,2,3)
plt.gray()
plt.axis('off')
plt.title('Time step: 28ms', fontsize=15, color='black')
im = np.reshape(W_vec[14,:], (8,8))
plt.imshow(im)
plt.subplot(2,2,4)
plt.gray()
plt.axis('off')
plt.title('Time step: 60ms', fontsize=15, color='black')
im = np.reshape(W_vec[30,:], (8,8))
plt.imshow(im)


text:
import numpy as np

def xavier_normal_init(in_dim, out_dim):
    """
    Xavier normal initialization.

    Parameters
    ----------
    in_dim
    out_dim

    Returns
    -------

    """
    return np.random.randn(out_dim, in_dim) * 2. / (out_dim+in_dim), np.zeros((out_dim, 1))

def xavier_uniform_init(in_dim, out_dim):
    return np.random.uniform(low=-np.sqrt(6. / (in_dim + out_dim)),
                             high=np.sqrt(6. / (in_dim + out_dim)),
                             size=(out_dim, in_dim)), np.zeros((out_dim, 1))

def kai_ming_uniform_init(in_dim, out_dim):
    return np.random.uniform(low=-np.sqrt(6. / out_dim),
                             high=np.sqrt(6. / out_dim),
                             size=(out_dim, in_dim)), np.zeros((out_dim, 1))

def kai_ming_normal_init(in_dim, out_dim):
    return np.random.randn(out_dim, in_dim) * 2. / out_dim, np.zeros((out_dim, 1))


text:
import torch
import torch.nn as nn
import numpy as np
from transformers import AutoModel

class MyLSTM_Bert(nn.Module):
    def __init__(self,
                 lstm_layers,
                 hidden_dim,
                 target_size,
                 dropout_prob,
                 device,
                 seq_len=250,
                 bert_pretrained="distilbert-base-uncased",
                 train_bert=False):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_layers = lstm_layers
        self.device = device

        self.bert_pretrained = bert_pretrained
        self.train_bert = train_bert
        self.bert_model = AutoModel.from_pretrained(self.bert_pretrained)
        self.bert_emb_dim = 768

        self.lstm = nn.LSTM(self.bert_emb_dim, hidden_dim,
                            batch_first=True, num_layers=lstm_layers,
                            dropout=dropout_prob)

        self.dropout = nn.Dropout(0.1)
        self.linear1 = nn.Linear(hidden_dim*seq_len, 1024)
        self.hidden2tag = nn.Linear(hidden_dim*seq_len, target_size)
        self.sigmoid = nn.Sigmoid()
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, sentence, hidden):
        with torch.no_grad():
            embeds = self.bert_model(sentence)[0]
        x, (hid, out) = self.lstm(embeds, hidden)
        x = x.contiguous().view(x.shape[0], -1)
        x = self.dropout(x)
        # x = self.linear1(x)
        x = self.hidden2tag(x)
        # x = self.sigmoid(x)
        x = self.log_softmax(x)
        return x, (hid, out)

    def init_hidden(self, batch_size):
        weight = next(self.parameters()).data
        hidden = (weight.new(self.n_layers, batch_size, self.hidden_dim).zero_().to(self.device),
                  weight.new(self.n_layers, batch_size, self.hidden_dim).zero_().to(self.device))
        return hidden

    def switch_train(self):
        self.train()
        bert_model = next(self.children())
        bert_model.train(self.train_bert)

max_stars_repo_name: rubicco/music-genre-classification

max_issues_repo_name: rubicco/music-genre-classification

max_forks_repo_name: rubicco/music-genre-classification


text:
#!/sr/bin/env python3
import numpy as np
import pandas as pd
from PIL import Image

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import os
from shutil import copyfile
from tqdm import tqdm
import argparse

from LungDataset import LungDataset
from model.unet import UNet

import gc
import math

def generate_masks(net,dataloader,device,useful_size,current_size):

    """Iterate over data"""

    print("predict masks and croped images")

    predicted_masks=[]
    data_iter = tqdm(enumerate(dataloader), total=len(dataloader))
    for batch_idx, sample in data_iter:
        imgs, true_masks = sample['image'], sample['mask']
        imgs = imgs.to(device=device, dtype=torch.float32)
        # mask_type = torch.float32 if net.n_classes == 1 else torch.long

        with torch.set_grad_enabled(False):
            masks_pred = net(imgs)
            pred = torch.sigmoid(masks_pred) > 0.5
            #print(pred.size())
            pred = torch.squeeze(pred)
            #print(pred.size())

        masks = pred.detach().cpu().numpy().astype(np.uint8)

        if useful_size != current_size:

            resized_masks = np.zeros((masks.shape[0],useful_size,useful_size),dtype=np.uint8)

            for i,mask in enumerate(masks):

                mask = (mask*255).astype(np.uint8)

                mask_img = Image.fromarray(mask).resize((useful_size,useful_size),Image.LANCZOS)

                resized_masks[i,:,:] = (np.array(mask_img)/255).astype(np.uint8)

            predicted_masks.append(resized_masks)
        else:
            predicted_masks.append(masks)

    return np.concatenate(predicted_masks, axis=0)

def create_predict_data(path,img_list,out,predicted_masks_array):

    croped_out = os.path.join(out,'croped_lung')

    #print(img_list)

    for i,img_name in tqdm(enumerate(img_list)):

        img = Image.open(os.path.join(path,'train/'+img_name)).convert('L')

        mask = (predicted_masks_array[i,:,:]*255).astype(np.uint8)

        mask_img = Image.fromarray(mask).resize(img.size,Image.LANCZOS)

        #mask_img.save(os.path.join(masks_out,'mask_'+img_name))

        croped = np.where(np.array(mask_img) == 0, 0, np.array(img)).astype(np.uint8)

        Image.fromarray(croped).save(os.path.join(croped_out,img_name))

        #print(i)

        del img,mask,mask_img,croped
        gc.collect()

def get_args():

    parser = argparse.ArgumentParser(description = "Qata_Covid19 Segmentation" ,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # set your environment
    parser.add_argument('--path',type=str,default='./data/Qata_COV')
    parser.add_argument('--gpu', type=str, default = '0')
    # arguments for training
    parser.add_argument('--img_size_lung', type = int , default = 512)
    parser.add_argument('--img_size_qata', type = int , default = 224)

    parser.add_argument('--load_lung_model', type=str, default='best_checkpoint.pt', help='.pth file path to load model')
    parser.add_argument('--load_qata_model', type=str, default='best_checkpoint.pt', help='.pth file path to load model')

    parser.add_argument('--out', type=str, default='./dataset')
    return parser.parse_args()

def main():

    args = get_args()

    if ~ os.path.exists(args.out):
        print("path created")
        os.mkdir(args.out)
        #os.mkdir(os.path.join(args.out,'Images'))
        #os.mkdir(os.path.join(args.out,'Ground-truths'))
        #os.mkdir(os.path.join(args.out,'predict_Ground-truths'))
        #os.mkdir(os.path.join(args.out,'original_crop_images'))
        os.mkdir(os.path.join(args.out,'croped_lung'))

    # set GPU device
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # default: '0'
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # set model
    lung_model = UNet(n_channels=1, n_classes=1).to(device)

    checkpoint = torch.load(args.load_lung_model)
    lung_model.load_state_dict(checkpoint['model_state_dict'])

    """set img size
        - UNet type architecture require input img size be divisible by 2^N,
        - Where N is the number of the Max Pooling layers (in the Vanila UNet N = 5)
    """

    img_size = args.img_size_lung #default: 224

    # set transforms for dataset
    import torchvision.transforms as transforms
    from my_transforms import RandomHorizontalFlip,RandomVerticalFlip,ColorJitter,GrayScale,Resize,ToTensor
    eval_transforms = transforms.Compose([
        GrayScale(),
        Resize(img_size),
        ToTensor()
    ])

    img_path = os.path.join(args.path,'train')
    img_list = os.listdir(img_path)#[:1000]

    dataset = LungDataset(root_dir = args.path,split=img_list,transforms=eval_transforms,img_size=args.img_size_lung)
    dataloader = DataLoader(dataset = dataset , batch_size=16,shuffle=False)

    #create_original_data(args.path,args.out)

    masks_lung = generate_masks(lung_model,dataloader,device,args.img_size_lung,args.img_size_lung)

    #from numba import cuda
    #cuda.select_device(0)
    #cuda.close()
    #cuda.select_device(0)

    # set model
    qata_model = UNet(n_channels=1, n_classes=1).to(device)

    checkpoint = torch.load(args.load_qata_model)
    qata_model.load_state_dict(checkpoint['model_state_dict'])

    dataset = LungDataset(root_dir = args.path,split=img_list,transforms=eval_transforms,img_size=args.img_size_qata)
    dataloader = DataLoader(dataset = dataset , batch_size=16,shuffle=False)

    masks_qata = generate_masks(qata_model,dataloader,device,args.img_size_lung,args.img_size_qata)

    #print(masks_lung.shape)
    #print(masks_qata.shape)

    batch_size = 500

    num_batchs = math.ceil(len(img_list)/batch_size)

    for i in range(num_batchs):

        a = batch_size * i

        b = min(batch_size * (i+1),len(img_list))

        joined_masks = masks_lung[a:b,:,:]+masks_qata[a:b,:,:]

        masks = np.where(joined_masks==0,0,1)

        #print(masks.shape)

        create_predict_data(args.path,img_list[a:b],args.out,masks)

        del joined_masks, masks
        gc.collect()

    #df = create_annotation(args.path)

    #df.to_csv(os.path.join(args.out,'target.csv'),index=False)

if __name__ == '__main__':

    main()

reason: import numpy,from numba