text:
__authors__ = ["M Sanchez del Rio - ESRF ISDD Advanced Analysis and Modelling"]
__license__ = "MIT"
__date__ = "30-08-2018"

"""

Undulator code: computes undulator radiation distributions and samples rays according to them.

Fully replaces and upgrades the shadow3 undulator model.

The radiation is calculating using one of three cods: internal, pySRU and SRW. The internal
code is indeed based on pySRU.

The radiation divergences (far field) are computed in polar coordinates for a more efiicient sampling.

Usage:

su = SourceUndulator()        # use keywords to define parameters. It uses syned for electron beam and vertical undulator
rays = su.calculate_rays()    # sample rays. Result is a numpy.array of shape (NRAYS,18), exactly the same as in shadow3

"""

import numpy

from srxraylib.util.inverse_method_sampler import Sampler1D, Sampler2D, Sampler3D
import scipy.constants as codata
from scipy import interpolate

from syned.storage_ring.magnetic_structures.undulator import Undulator
# from syned.storage_ring.electron_beam import ElectronBeam
from shadow4.sources.s4_electron_beam import S4ElectronBeam

from shadow4.beam.beam import Beam

from shadow4.sources.undulator.source_undulator_factory import calculate_undulator_emission # SourceUndulatorFactory
from shadow4.sources.undulator.source_undulator_factory_srw import calculate_undulator_emission_SRW # SourceUndulatorFactorySrw
from shadow4.sources.undulator.source_undulator_factory_pysru import calculate_undulator_emission_pySRU # SourceUndulatorFactoryPysru

from shadow4.sources.undulator.s4_undulator import S4Undulator

from syned.storage_ring.light_source import LightSource
from shadow4.sources.s4_light_source import S4LightSource

from shadow4.sources.source_geometrical.source_gaussian import SourceGaussian
import scipy.constants as codata

INTEGRATION_METHOD = 1 # 0=sum, 1=trapz

class S4UndulatorLightSource(S4LightSource):

    def __init__(self, name="Undefined", electron_beam=None, magnetic_structure=None):
        super().__init__(name,
                         electron_beam=electron_beam if not electron_beam is None else S4ElectronBeam(),
                         magnetic_structure=magnetic_structure if not magnetic_structure is None else S4Undulator())

        # results of calculations

        self.__result_radiation = None
        self.__result_photon_size_distribution = None
        self.__result_photon_size_sigma = None

    def get_beam(self, NRAYS=5000, SEED=123456):
        user_unit_to_m = 1.0
        F_COHER = 0
        # use_gaussian_approximation = False

        if self.get_magnetic_structure().use_gaussian_approximation():
            return self.get_beam_in_gaussian_approximation(NRAYS=NRAYS, SEED=SEED)
        else:
            return Beam.initialize_from_array(self.__calculate_rays(
                user_unit_to_m=user_unit_to_m, F_COHER=F_COHER, NRAYS=NRAYS, SEED=SEED
                ))

    def get_resonance_ring(self,harmonic_number=1, ring_order=1):
        return 1.0/self.get_electron_beam().gamma()*numpy.sqrt( ring_order / harmonic_number * (1+0.5*self.get_magnetic_structure().K_vertical()**2) )

    def set_energy_monochromatic_at_resonance(self,harmonic_number):
        self.set_energy_monochromatic(self.syned_undulator.resonance_energy(
            self.syned_electron_beam.gamma(),harmonic=harmonic_number))
        # take 3*sigma - _MAXANGLE is in RAD
        self._MAXANGLE = 3 * 0.69 * self.syned_undulator.gaussian_central_cone_aperture(self.syned_electron_beam.gamma(),harmonic_number)

    # def set_energy_monochromatic(self,emin):
    #     """
    #     Sets a single energy line for the source (monochromatic)
    #     :param emin: the energy in eV
    #     :return:
    #     """
    #     self._EMIN = emin
    #     self._EMAX = emin
    #     self._NG_E = 1
    #
    #
    # def set_energy_box(self,emin,emax,npoints=None):
    #     """
    #     Sets a box for photon energy distribution for the source
    #     :param emin:  Photon energy scan from energy (in eV)
    #     :param emax:  Photon energy scan to energy (in eV)
    #     :param npoints:  Photon energy scan number of points (optinal, if not set no changes)
    #     :return:
    #     """
    #     self._EMIN = emin
    #     self._EMAX = emax
    #     if npoints != None:
    #         self._NG_E = npoints
    #
    # def get_energy_box(self):
    #     """
    #     Gets the limits of photon energy distribution for the source
    #     :return: emin,emax,number_of_points
    #     """
    #     return self._EMIN,self._EMAX,self._NG_E

    def __calculate_radiation(self):
        """
        Calculates the radiation (emission) as a function of theta (elevation angle) and phi (azimuthal angle)
        This radiation will be sampled to create the source

        It calls undul_phot* in SourceUndulatorFactory

        :param code_undul_phot: 'internal' (calls undul_phot), 'pysru' (calls undul_phot_pysru) or
                'srw' (calls undul_phot_srw)
        :return: a dictionary (the output from undul_phot*)
        """

        # h = self.to_dictionary()
        # print(self.info())
        # os.system("rm -f xshundul.plt xshundul.par xshundul.traj xshundul.info xshundul.sha")

        # if code_undul_phot != "internal" or code_undul_phot != "srw":
        #     dump_uphot_dot_dat = True

        syned_electron_beam = self.get_electron_beam()
        undulator = self.get_magnetic_structure()

        self.__result_radiation = None

        # undul_phot
        if undulator.code_undul_phot == 'internal':
            undul_phot_dict = calculate_undulator_emission(
                electron_energy= syned_electron_beam.energy(),
                electron_current= syned_electron_beam.current(),
                undulator_period= undulator.period_length(),
                undulator_nperiods= undulator.number_of_periods(),
                K         = undulator.K(),
                photon_energy= undulator._EMIN,
                EMAX      = undulator._EMAX,
                NG_E      = undulator._NG_E,
                MAXANGLE  = undulator._MAXANGLE,
                number_of_points= undulator._NG_T,
                NG_P      = undulator._NG_P,
                number_of_trajectory_points = undulator._NG_J,
                )

        elif undulator.code_undul_phot == 'pysru' or  undulator.code_undul_phot == 'pySRU':
            undul_phot_dict = calculate_undulator_emission_pySRU(
                electron_energy= syned_electron_beam.energy(),
                electron_current= syned_electron_beam.current(),
                undulator_period= undulator.period_length(),
                undulator_nperiods= undulator.number_of_periods(),
                K         = undulator.K(),
                photon_energy= undulator._EMIN,
                EMAX      = undulator._EMAX,
                NG_E      = undulator._NG_E,
                MAXANGLE  = undulator._MAXANGLE,
                number_of_points= undulator._NG_T,
                NG_P      = undulator._NG_P,
                )
        elif undulator.code_undul_phot == 'srw' or  undulator.code_undul_phot == 'SRW':
            undul_phot_dict = calculate_undulator_emission_SRW(
                electron_energy= syned_electron_beam.energy(),
                electron_current= syned_electron_beam.current(),
                undulator_period= undulator.period_length(),
                undulator_nperiods= undulator.number_of_periods(),
                K         = undulator.K(),
                photon_energy= undulator._EMIN,
                EMAX      = undulator._EMAX,
                NG_E      = undulator._NG_E,
                MAXANGLE  = undulator._MAXANGLE,
                number_of_points= undulator._NG_T,
                NG_P      = undulator._NG_P,
                )
        else:
            raise Exception("Not implemented undul_phot code: "+undulator.code_undul_phot)

        # add some info
        undul_phot_dict["code_undul_phot"] = undulator.code_undul_phot
        undul_phot_dict["info"] = self.info()

        self.__result_radiation = undul_phot_dict

    #
    # get from results
    #

    def get_result_photon_size_sigma(self):
        return self.__result_photon_size_sigma

    def get_result_dictionary(self):
        if self.__result_radiation is None:
            self.__calculate_radiation()
        return self.__result_radiation

    def get_result_radiation(self):
        return self.get_result_dictionary()["radiation"]

    def get_result_polarization(self):
        return self.get_result_dictionary()["polarization"]

    def get_result_polarisation(self):
        return self.get_result_polarization()

    def get_result_theta(self):
        return self.get_result_dictionary()["theta"]

    def get_result_phi(self):
        return self.get_result_dictionary()["phi"]

    def get_result_photon_energy(self):
        return self.get_result_dictionary()["photon_energy"]

    def get_radiation_polar(self):
        return self.get_result_radiation(),self.get_result_photon_energy(),self.get_result_theta(),self.get_result_phi()

    def get_radiation(self):
        return self.get_radiation_polar()

    def get_radiation_interpolated_cartesian(self,npointsx=100,npointsz=100,thetamax=None):

        radiation,photon_energy, thetabm,phi = self.get_radiation_polar()

        if thetamax is None:
            thetamax = thetabm.max()

        vx = numpy.linspace(-1.1*thetamax,1.1*thetamax,npointsx)
        vz = numpy.linspace(-1.1*thetamax,1.1*thetamax,npointsz)
        VX = numpy.outer(vx,numpy.ones_like(vz))
        VZ = numpy.outer(numpy.ones_like(vx),vz)
        VY = numpy.sqrt(1 - VX**2 - VZ**2)

        THETA = numpy.abs(numpy.arctan( numpy.sqrt(VX**2+VZ**2)/VY))
        PHI = numpy.arctan2( numpy.abs(VZ),numpy.abs(VX))

        radiation_interpolated = numpy.zeros((radiation.shape[0],npointsx,npointsz))

        for i in range(radiation.shape[0]):
            interpolator_value = interpolate.RectBivariateSpline(thetabm, phi, radiation[i])
            radiation_interpolated[i] = interpolator_value.ev(THETA, PHI)

        return radiation_interpolated,photon_energy,vx,vz

    def get_power_density(self):

        radiation = self.get_result_radiation().copy()
        theta = self.get_result_theta()
        phi = self.get_result_phi()
        photon_energy = self.get_result_photon_energy()

        step_e = photon_energy[1]-photon_energy[0]

        for i in range(radiation.shape[0]):
            radiation[i] *= 1e-3 * photon_energy[i] # photons/eV/rad2 -> photons/0.1%bw/rad2

        if INTEGRATION_METHOD == 0:
            power_density = radiation.sum(axis=0) * step_e * codata.e * 1e3 # W/rad2
        else:
            power_density = numpy.trapz(radiation,photon_energy,axis=0) * codata.e * 1e3 # W/rad2

        return power_density,theta,phi

    def get_power_density_interpolated_cartesian(self,npointsx=100,npointsz=100,thetamax=None):

        power_density_polar,theta,phi = self.get_power_density()

        if thetamax is None:
            thetamax = theta.max()

        vx = numpy.linspace(-1.1*thetamax,1.1*thetamax,npointsx)
        vz = numpy.linspace(-1.1*thetamax,1.1*thetamax,npointsz)
        VX = numpy.outer(vx,numpy.ones_like(vz))
        VZ = numpy.outer(numpy.ones_like(vx),vz)
        VY = numpy.sqrt(1 - VX**2 - VZ**2)

        THETA = numpy.abs(numpy.arctan( numpy.sqrt(VX**2+VZ**2)/VY))
        PHI = numpy.arctan2( numpy.abs(VZ),numpy.abs(VX))

        interpolator_value = interpolate.RectBivariateSpline(theta, phi, power_density_polar)
        power_density_cartesian = interpolator_value.ev(THETA, PHI)

        return power_density_cartesian,vx,vz

    def get_flux_and_spectral_power(self):

        radiation2 = self.get_result_radiation().copy()
        theta = self.get_result_theta()
        phi = self.get_result_phi()
        photon_energy = self.get_result_photon_energy()
        THETA = numpy.outer(theta,numpy.ones_like(phi))
        for i in range(radiation2.shape[0]):
            radiation2[i] *= THETA

        if INTEGRATION_METHOD == 0:
            flux = radiation2.sum(axis=2).sum(axis=1) * (1e-3*photon_energy) # photons/eV -> photons/0.1%bw
            flux *= 4 * (theta[1]-theta[0]) * (phi[1]-phi[0]) # adding the four quadrants!
        else:
            flux = 4 * numpy.trapz(numpy.trapz(radiation2,phi,axis=2),theta,axis=1) * (1e-3*photon_energy) # photons/eV -> photons/0.1%bw

        spectral_power = flux*codata.e*1e3

        return flux,spectral_power,photon_energy

    def get_flux(self):
        flux,spectral_power,photon_energy = self.get_flux_and_spectral_power()
        return flux,photon_energy

    def get_spectral_power(self):
        flux,spectral_power,photon_energy = self.get_flux_and_spectral_power()
        return spectral_power,photon_energy

    def get_photon_size_distribution(self):
        return self.__result_photon_size_distribution["x"], self.__result_photon_size_distribution["y"]

    def __calculate_rays(self,user_unit_to_m=1.0,F_COHER=0,NRAYS=5000,SEED=123456):
        """
        compute the rays in SHADOW matrix (shape (npoints,18) )
        :param F_COHER: set this flag for coherent beam
        :param user_unit_to_m: default 1.0 (m)
        :return: rays, a numpy.array((npoits,18))
        """

        if self.__result_radiation is None:
            self.__calculate_radiation()

        syned_electron_beam = self.get_electron_beam()
        undulator = self.get_magnetic_structure()

        sampled_photon_energy,sampled_theta,sampled_phi = self._sample_photon_energy_theta_and_phi(NRAYS)

        if SEED != 0:
            numpy.random.seed(SEED)

        sigmas = syned_electron_beam.get_sigmas_all()

        rays = numpy.zeros((NRAYS,18))

        #
        # sample sizes (cols 1-3)
        #

        if undulator._FLAG_EMITTANCE:
            x_electron = numpy.random.normal(loc=0.0,scale=sigmas[0],size=NRAYS)
            y_electron = 0.0
            z_electron = numpy.random.normal(loc=0.0,scale=sigmas[2],size=NRAYS)
        else:
            x_electron = 0.0
            y_electron = 0.0
            z_electron = 0.0

        # calculate (and stores) sizes of the photon undulator beam
        # see formulas 25 & 30 in Elleaume (Onaki & Elleaume)
        # sp_phot = 0.69*numpy.sqrt(lambda1/undulator_length)
        undulator_length = undulator.length()
        lambda1 = codata.h*codata.c/codata.e / numpy.array(sampled_photon_energy).mean()
        s_phot = 2.740/(4e0*numpy.pi)*numpy.sqrt(undulator_length*lambda1)
        self.__result_photon_size_sigma = s_phot

        if undulator._FLAG_SIZE == 0:
            x_photon = 0.0
            y_photon = 0.0
            z_photon = 0.0
            # for plot, a delta
            x = numpy.linspace(-1e-6,1e-6,101)
            y = numpy.zeros_like(x)
            y[y.size//2] = 1.0
            self.__result_photon_size_distribution = {"x":x, "y":y}
        elif undulator._FLAG_SIZE == 1:
            # TODO: I added this correction to obtain the sigma in the RADIAL coordinate, not in x and z.
            # RODO: TO be verified!
            s_phot_corrected = s_phot / numpy.sqrt(2)

            cov = [[s_phot_corrected**2, 0], [0, s_phot_corrected**2]]
            mean = [0.0,0.0]

            tmp = numpy.random.multivariate_normal(mean, cov, NRAYS)
            x_photon = tmp[:,0]
            y_photon = 0.0
            z_photon = tmp[:,1]

            # for plot, a Gaussian
            x = numpy.linspace(-5*s_phot,5*s_phot,101)
            y = numpy.exp(-x**2/2/s_phot**2)
            self.__result_photon_size_distribution = {"x":x, "y":y}

        elif undulator._FLAG_SIZE == 2:
            # we need to retrieve the emission as a function of the angle
            radiation,photon_energy, theta,phi = self.get_radiation_polar()

            mean_photon_energy = numpy.array(sampled_photon_energy).mean() # todo: use the weighted mean?
            shape_radiation = radiation.shape
            radial_flux = radiation.sum(axis=2) / shape_radiation[2]
            radial_flux = radial_flux.sum(axis=0) / shape_radiation[0]
            # doble the arrays for 1D propagation
            THETA = numpy.concatenate((-theta[::-1],theta[1::]),axis=None)
            RADIAL_FLUX = numpy.concatenate( (radial_flux[::-1],radial_flux[1::]),axis=None)

            #
            # we propagate the emission at a long distance back to the source plane
            #
            distance = 100.

            magnification = s_phot*10 / (theta[-1]*distance)

            # do the propagation; result is stored in self._photon_size_distribution
            self._back_propagation_for_size_calculation(THETA,RADIAL_FLUX,
                                mean_photon_energy,distance=distance,magnification=magnification)

            # we sample rays following the resulting radial distribution
            xx = self.__result_photon_size_distribution["x"]
            yy = self.__result_photon_size_distribution["y"]

            # #########################################################
            # # for plot, a Gaussian
            # xx = numpy.linspace(-5*s_phot,5*s_phot,101)
            # yy = numpy.exp(-xx**2/2/s_phot**2)
            # self._result_photon_size_distribution = {"x":xx,"y":yy}
            # #########################################################

            sampler_radial = Sampler1D(yy*numpy.abs(xx),xx)
            r,hy,hx = sampler_radial.get_n_sampled_points_and_histogram(NRAYS,bins=101)
            angle = numpy.random.random(NRAYS) * 2 * numpy.pi

            x_photon = r / numpy.sqrt(2.0) * numpy.sin(angle)
            y_photon = 0.0
            z_photon = r / numpy.sqrt(2.0) * numpy.cos(angle)

        rays[:,0] = x_photon + x_electron
        rays[:,1] = y_photon + y_electron
        rays[:,2] = z_photon + z_electron

        if user_unit_to_m != 1.0:
            rays[:,0] /= user_unit_to_m
            rays[:,1] /= user_unit_to_m
            rays[:,2] /= user_unit_to_m

        #
        # sample divergences (cols 4-6): the Shadow way
        #
        THETABM = sampled_theta
        PHI = sampled_phi
        A_Z = numpy.arcsin(numpy.sin(THETABM)*numpy.sin(PHI))
        A_X = numpy.arccos(numpy.cos(THETABM)/numpy.cos(A_Z))
        THETABM = A_Z
        PHI  = A_X
        # ! C Decide in which quadrant THETA and PHI are.
        myrand = numpy.random.random(NRAYS)
        THETABM[numpy.where(myrand < 0.5)] *= -1.0
        myrand = numpy.random.random(NRAYS)
        PHI[numpy.where(myrand < 0.5)] *= -1.0

        if undulator._FLAG_EMITTANCE:
            EBEAM1 = numpy.random.normal(loc=0.0,scale=sigmas[1],size=NRAYS)
            EBEAM3 = numpy.random.normal(loc=0.0,scale=sigmas[3],size=NRAYS)
            ANGLEX = EBEAM1 + PHI
            ANGLEV = EBEAM3 + THETABM
        else:
            ANGLEX = PHI # E_BEAM(1) + PHI
            ANGLEV =THETABM #  E_BEAM(3) + THETABM

        VX = numpy.tan(ANGLEX)
        VY = 1.0
        VZ = numpy.tan(ANGLEV)/numpy.cos(ANGLEX)
        VN = numpy.sqrt( VX*VX + VY*VY + VZ*VZ)
        VX /= VN
        VY /= VN
        VZ /= VN

        rays[:,3] = VX
        rays[:,4] = VY
        rays[:,5] = VZ

        #
        # electric field vectors (cols 7-9, 16-18) and phases (cols 14-15)
        #

        # beam.rays[:,6] =  1.0

        # ! C
        # ! C  ---------------------------------------------------------------------
        # ! C                 POLARIZATION
        # ! C
        # ! C   Generates the polarization of the ray. This is defined on the
        # ! C   source plane, so that A_VEC is along the X-axis and AP_VEC is along Z-axis.
        # ! C   Then care must be taken so that A will be perpendicular to the ray
        # ! C   direction.
        # ! C
        # ! C
        # A_VEC(1) = 1.0D0
        # A_VEC(2) = 0.0D0
        # A_VEC(3) = 0.0D0

        DIREC = rays[:,3:6].copy()
        A_VEC = numpy.zeros_like(DIREC)
        A_VEC[:,0] = 1.0

        # ! C
        # ! C   Rotate A_VEC so that it will be perpendicular to DIREC and with the
        # ! C   right components on the plane.
        # ! C
        # CALL CROSS (A_VEC,DIREC,A_TEMP)
        A_TEMP = self._cross(A_VEC,DIREC)
        # CALL CROSS (DIREC,A_TEMP,A_VEC)
        A_VEC = self._cross(DIREC,A_TEMP)
        # CALL NORM (A_VEC,A_VEC)
        A_VEC = self._norm(A_VEC)
        # CALL CROSS (A_VEC,DIREC,AP_VEC)
        AP_VEC = self._cross(A_VEC,DIREC)
        # CALL NORM (AP_VEC,AP_VEC)
        AP_VEC = self._norm(AP_VEC)

        #
        # obtain polarization for each ray (interpolation)
        #

        if undulator._NG_E == 1: # 2D interpolation
            sampled_photon_energy = numpy.array(sampled_photon_energy) # be sure is an array
            fn = interpolate.RegularGridInterpolator(
                (self.__result_radiation["theta"], self.__result_radiation["phi"]),
                self.__result_radiation["polarization"][0])

            pts = numpy.dstack( (sampled_theta,
                                 sampled_phi) )
            pts = pts[0]
            POL_DEG = fn(pts)
        else: # 3D interpolation
            fn = interpolate.RegularGridInterpolator(
                (self.__result_radiation["photon_energy"], self.__result_radiation["theta"], self.__result_radiation["phi"]),
                self.__result_radiation["polarization"])

            pts = numpy.dstack( (sampled_photon_energy,
                                 sampled_theta,
                                 sampled_phi) )
            pts = pts[0]
            POL_DEG = fn(pts)

        #     ! C
        #     ! C   WaNT A**2 = AX**2 + AZ**2 = 1 , instead of A_VEC**2 = 1 .
        #     ! C
        #     DENOM = SQRT(1.0D0 - 2.0D0*POL_DEG + 2.0D0*POL_DEG**2)
        #     AX = POL_DEG/DENOM
        #     CALL SCALAR (A_VEC,AX,A_VEC)
        #     ! C
        #     ! C   Same procedure for AP_VEC
        #     ! C
        #     AZ = (1-POL_DEG)/DENOM
        #     CALL SCALAR  (AP_VEC,AZ,AP_VEC)

        DENOM = numpy.sqrt(1.0 - 2.0 * POL_DEG + 2.0 * POL_DEG**2)
        AX = POL_DEG/DENOM
        for i in range(3):
            A_VEC[:,i] *= AX

        AZ = (1.0-POL_DEG)/DENOM
        for i in range(3):
            AP_VEC[:,i] *= AZ

        rays[:,6:9] =  A_VEC
        rays[:,15:18] = AP_VEC

        #
        # ! C
        # ! C Now the phases of A_VEC and AP_VEC.
        # ! C
        # IF (F_COHER.EQ.1) THEN
        #     PHASEX = 0.0D0
        # ELSE
        #     PHASEX = WRAN(ISTAR1) * TWOPI
        # END IF
        # PHASEZ = PHASEX + POL_ANGLE*I_CHANGE
        #
        POL_ANGLE = 0.5 * numpy.pi

        if F_COHER == 1:
            PHASEX = 0.0
        else:
            PHASEX = numpy.random.random(NRAYS) * 2 * numpy.pi

        PHASEZ = PHASEX + POL_ANGLE * numpy.sign(ANGLEV)

        rays[:,13] = PHASEX
        rays[:,14] = PHASEZ

        # set flag (col 10)
        rays[:,9] = 1.0

        #
        # photon energy (col 11)
        #

        A2EV = 2.0*numpy.pi/(codata.h*codata.c/codata.e*1e2)
        rays[:,10] =  sampled_photon_energy * A2EV

        # col 12 (ray index)
        rays[:,11] =  1 + numpy.arange(NRAYS)

        # col 13 (optical path)
        rays[:,11] = 0.0

        return rays

    def _back_propagation_for_size_calculation(self,theta,radiation_flux,photon_energy,
                                                distance=100.0,magnification=0.010000):
        """
        Calculate the radiation_flux vs theta at a "distance"
        Back propagate to -distance
        The result is the size distrubution

        :param theta:
        :param radiation_flux:
        :param photon_energy:
        :param distance:
        :param magnification:
        :return: None; stores results in self._photon_size_distribution
        """

        from wofry.propagator.wavefront1D.generic_wavefront import GenericWavefront1D
        from wofry.propagator.propagator import PropagationManager, PropagationElements, PropagationParameters
        from syned.beamline.beamline_element import BeamlineElement
        from shadow4.syned.element_coordinates import ElementCoordinates
        from wofryimpl.propagator.propagators1D.fresnel_zoom import FresnelZoom1D
        from wofryimpl.beamline.optical_elements.ideal_elements.screen import WOScreen1D

        input_wavefront = GenericWavefront1D().initialize_wavefront_from_arrays(theta*distance,numpy.sqrt(radiation_flux)+0j)
        input_wavefront.set_photon_energy(photon_energy)
        input_wavefront.set_spherical_wave(radius=distance,complex_amplitude=numpy.sqrt(radiation_flux)+0j)
        # input_wavefront.save_h5_file("tmp2.h5","wfr")

        optical_element = WOScreen1D()
        #
        # propagating
        #
        #
        propagation_elements = PropagationElements()
        beamline_element = BeamlineElement(optical_element=optical_element,
                        coordinates=ElementCoordinates(p=0.0,q=-distance,
                        angle_radial=numpy.radians(0.000000),
                        angle_azimuthal=numpy.radians(0.000000)))
        propagation_elements.add_beamline_element(beamline_element)
        propagation_parameters = PropagationParameters(wavefront=input_wavefront.duplicate(),propagation_elements = propagation_elements)
        propagation_parameters.set_additional_parameters('magnification_x', magnification)

        #
        propagator = PropagationManager.Instance()
        try:
            propagator.add_propagator(FresnelZoom1D())
        except:
            pass
        output_wavefront = propagator.do_propagation(propagation_parameters=propagation_parameters,handler_name='FRESNEL_ZOOM_1D')

        self.__result_photon_size_distribution = {"x":output_wavefront.get_abscissas(), "y":output_wavefront.get_intensity()}

    def _cross(self,u,v):
        # w = u X v
        # u = array (npoints,vector_index)

        w = numpy.zeros_like(u)
        w[:,0] = u[:,1] * v[:,2] - u[:,2] * v[:,1]
        w[:,1] = u[:,2] * v[:,0] - u[:,0] * v[:,2]
        w[:,2] = u[:,0] * v[:,1] - u[:,1] * v[:,0]

        return w

    def _norm(self,u):
        # w = u / |u|
        # u = array (npoints,vector_index)
        u_norm = numpy.zeros_like(u)
        uu = numpy.sqrt( u[:,0]**2 + u[:,1]**2 + u[:,2]**2)
        for i in range(3):
            u_norm[:,i] = uu
        return u / u_norm

    def _sample_photon_energy_theta_and_phi(self,NRAYS):

        #
        # sample divergences
        #

        theta = self.__result_radiation["theta"]
        phi = self.__result_radiation["phi"]
        photon_energy = self.__result_radiation["photon_energy"]

        photon_energy_spectrum = 'polychromatic' # 'monochromatic' #
        if self.get_magnetic_structure()._EMIN == self.get_magnetic_structure()._EMAX:
            photon_energy_spectrum = 'monochromatic'
        if self.get_magnetic_structure()._NG_E == 1:
            photon_energy_spectrum = 'monochromatic'

        if photon_energy_spectrum == 'monochromatic':

            #2D case
            tmp = self.__result_radiation["radiation"][0, :, :].copy()
            tmp /= tmp.max()

            # correct radiation for DxDz / DthetaDphi
            tmp_theta = numpy.outer(theta,numpy.ones_like(phi))
            tmp_theta /= tmp_theta.max()
            tmp_theta += 1e-6 # to avoid zeros
            tmp *= tmp_theta
            # plot_image(tmp_theta,theta,phi,aspect='auto')

            s2d = Sampler2D(tmp,theta,phi)
            sampled_theta,sampled_phi = s2d.get_n_sampled_points(NRAYS)

            sampled_photon_energy = self.get_magnetic_structure()._EMIN

        elif photon_energy_spectrum == "polychromatic":
            #3D case
            tmp = self.__result_radiation["radiation"].copy()
            tmp /= tmp.max()
            # correct radiation for DxDz / DthetaDphi
            tmp_theta = numpy.outer(theta,numpy.ones_like(phi))
            tmp_theta /= tmp_theta.max()
            tmp_theta += 1e-6 # to avoid zeros
            for i in range(tmp.shape[0]):
                tmp[i,:,:] *= tmp_theta

            s3d = Sampler3D(tmp,photon_energy,theta,phi)

            sampled_photon_energy,sampled_theta,sampled_phi = s3d.get_n_sampled_points(NRAYS)

        return sampled_photon_energy,sampled_theta,sampled_phi

    def get_beam_in_gaussian_approximation(self, NRAYS=1000, SEED=12345):
        emin, emax, npoints = self.get_magnetic_structure().get_energy_box()

        if self.get_magnetic_structure().get_flag_emittance():
            sigma_x, sigdi_x, sigma_z, sigdi_z = self.get_electron_beam().get_sigmas_all()
            Sx, Sz, Spx, Spz = self.get_undulator_photon_beam_sizes_by_convolution(
                sigma_x=sigma_x,
                sigma_z=sigma_z,
                sigdi_x=sigdi_x,
                sigdi_z=sigdi_z,
                undulator_E0=0.5*(emin+emax),
                undulator_length=self.get_magnetic_structure().length(),
            )
            a = SourceGaussian.initialize_from_keywords(number_of_rays=NRAYS,
                                                        sigmaX=Sx,
                                                        sigmaY=0,
                                                        sigmaZ=Sz,
                                                        sigmaXprime=Spx,
                                                        sigmaZprime=Spz,
                                                        real_space_center=[0.0, 0.0, 0.0],
                                                        direction_space_center=[0.0, 0.0])
        else:
            s, sp = self.get_undulator_photon_beam_sizes(
                undulator_E0=0.5*(emin+emax),
                undulator_length=self.get_magnetic_structure().length(),
            )
            a = SourceGaussian.initialize_from_keywords(number_of_rays=NRAYS,
                                                        sigmaX=s,
                                                        sigmaY=0,
                                                        sigmaZ=s,
                                                        sigmaXprime=sp,
                                                        sigmaZprime=sp,
                                                        real_space_center=[0.0, 0.0, 0.0],
                                                        direction_space_center=[0.0, 0.0])

        print(a.info())
        beam = a.get_beam()
        if emax != emin:
            e = numpy.random.random(NRAYS) * (emax - emin) + emin
            beam.set_photon_energy_eV(e)

        return beam

    @classmethod
    def get_undulator_photon_beam_sizes(cls,
                                        undulator_E0=15000.0,
                                        undulator_length=4.0):
        m2ev = codata.c * codata.h / codata.e
        lambda1 = m2ev / undulator_E0
        # calculate sizes of the photon undulator beam
        # see formulas 25 & 30 in Elleaume (Onaki & Elleaume)
        s_phot = 2.740/(4e0*numpy.pi)*numpy.sqrt(undulator_length*lambda1)
        sp_phot = 0.69*numpy.sqrt(lambda1/undulator_length)
        return s_phot, sp_phot

    @classmethod
    def get_undulator_photon_beam_sizes_by_convolution(cls,
                                                       sigma_x=1e-4,
                                                       sigma_z=1e-4,
                                                       sigdi_x=1e-6,
                                                       sigdi_z=1e-6,
                                                       undulator_E0=15000.0,
                                                       undulator_length=4.0):

        user_unit_to_m = 1.0

        s_phot, sp_phot = cls.get_undulator_photon_beam_sizes(undulator_E0=undulator_E0, undulator_length=undulator_length)

        photon_h = numpy.sqrt(numpy.power(sigma_x*user_unit_to_m,2) + numpy.power(s_phot,2) )
        photon_v = numpy.sqrt(numpy.power(sigma_z*user_unit_to_m,2) + numpy.power(s_phot,2) )
        photon_hp = numpy.sqrt(numpy.power(sigdi_x,2) + numpy.power(sp_phot,2) )
        photon_vp = numpy.sqrt(numpy.power(sigdi_z,2) + numpy.power(sp_phot,2) )

        m2ev = codata.c * codata.h / codata.e
        lambda1 = m2ev / undulator_E0
        txt = ""
        txt += "Photon single electron emission at wavelength %f A: \n"%(lambda1*1e10)
        txt += "    sigma_u:      %g um \n"%(1e6*s_phot)
        txt += "    sigma_uprime: %g urad \n" %(1e6*sp_phot)
        txt += "Electron sizes: \n"
        txt += "    sigma_x: %g um \n"%(1e6*sigma_x*user_unit_to_m)
        txt += "    sigma_z: %g um \n" %(1e6*sigma_z*user_unit_to_m)
        txt += "    sigma_x': %g urad \n"%(1e6*sigdi_x)
        txt += "    sigma_z': %g urad \n" %(1e6*sigdi_z)
        txt += "Photon source sizes (convolution): \n"
        txt += "    Sigma_x: %g um \n"%(1e6*photon_h)
        txt += "    Sigma_z: %g um \n" %(1e6*photon_v)
        txt += "    Sigma_x': %g urad \n"%(1e6*photon_hp)
        txt += "    Sigma_z': %g urad \n" %(1e6*photon_vp)

        print(txt)

        return (photon_h/user_unit_to_m, photon_v/user_unit_to_m, photon_hp, photon_vp)

    ########################################
    #
    ########################################
    def info(self,debug=False):
        syned_electron_beam = self.get_electron_beam()
        undulator = self.get_magnetic_structure()

        txt = ""

        txt += "-----------------------------------------------------\n"

        txt += "Input Electron parameters: \n"
        txt += "        Electron energy: %f geV\n"%syned_electron_beam._energy_in_GeV
        txt += "        Electron current: %f A\n"%syned_electron_beam._current
        if undulator._FLAG_EMITTANCE:
            sigmas = syned_electron_beam.get_sigmas_all()
            txt += "        Electron sigmaX: %g [um]\n"%(1e6*sigmas[0])
            txt += "        Electron sigmaZ: %g [um]\n"%(1e6*sigmas[2])
            txt += "        Electron sigmaX': %f urad\n"%(1e6*sigmas[1])
            txt += "        Electron sigmaZ': %f urad\n"%(1e6*sigmas[3])
        txt += "Input Undulator parameters: \n"
        txt += "        period: %f m\n"%undulator.period_length()
        txt += "        number of periods: %d\n"%undulator.number_of_periods()
        txt += "        K-value: %f\n"%undulator.K_vertical()

        txt += "-----------------------------------------------------\n"

        txt += "Lorentz factor (gamma): %f\n"%syned_electron_beam.gamma()
        txt += "Electron velocity: %.12f c units\n"%(numpy.sqrt(1.0 - 1.0 / syned_electron_beam.gamma() ** 2))
        txt += "Undulator length: %f m\n"%(undulator.period_length()*undulator.number_of_periods())
        K_to_B = (2.0 * numpy.pi / undulator.period_length()) * codata.m_e * codata.c / codata.e

        txt += "Undulator peak magnetic field: %f T\n"%(K_to_B*undulator.K_vertical())
        txt += "Resonances: \n"
        txt += "        harmonic number [n]                   %10d %10d %10d \n"%(1,3,5)
        txt += "        wavelength [A]:                       %10.6f %10.6f %10.6f   \n"%(\
                                                                1e10*undulator.resonance_wavelength(syned_electron_beam.gamma(),harmonic=1),
                                                                1e10*undulator.resonance_wavelength(syned_electron_beam.gamma(),harmonic=3),
                                                                1e10*undulator.resonance_wavelength(syned_electron_beam.gamma(),harmonic=5))
        txt += "        energy [eV]   :                       %10.3f %10.3f %10.3f   \n"%(\
                                                                undulator.resonance_energy(syned_electron_beam.gamma(),harmonic=1),
                                                                undulator.resonance_energy(syned_electron_beam.gamma(),harmonic=3),
                                                                undulator.resonance_energy(syned_electron_beam.gamma(),harmonic=5))
        txt += "        frequency [Hz]:                       %10.3g %10.3g %10.3g   \n"%(\
                                                                1e10*undulator.resonance_frequency(syned_electron_beam.gamma(),harmonic=1),
                                                                1e10*undulator.resonance_frequency(syned_electron_beam.gamma(),harmonic=3),
                                                                1e10*undulator.resonance_frequency(syned_electron_beam.gamma(),harmonic=5))
        txt += "        central cone 'half' width [urad]:     %10.6f %10.6f %10.6f   \n"%(\
                                                                1e6*undulator.gaussian_central_cone_aperture(syned_electron_beam.gamma(),1),
                                                                1e6*undulator.gaussian_central_cone_aperture(syned_electron_beam.gamma(),3),
                                                                1e6*undulator.gaussian_central_cone_aperture(syned_electron_beam.gamma(),5))
        txt += "        first ring at [urad]:                 %10.6f %10.6f %10.6f   \n"%(\
                                                                1e6*self.get_resonance_ring(1,1),
                                                                1e6*self.get_resonance_ring(3,1),
                                                                1e6*self.get_resonance_ring(5,1))

        txt += "-----------------------------------------------------\n"
        txt += "Grids: \n"
        if undulator._NG_E == 1:
            txt += "        photon energy %f eV\n"%(undulator._EMIN)
        else:
            txt += "        photon energy from %10.3f eV to %10.3f eV\n"%(undulator._EMIN,undulator._EMAX)
        txt += "        number of points for the trajectory: %d\n"%(undulator._NG_J)
        txt += "        number of energy points: %d\n"%(undulator._NG_E)
        txt += "        maximum elevation angle: %f urad\n"%(1e6*undulator._MAXANGLE)
        txt += "        number of angular elevation points: %d\n"%(undulator._NG_T)
        txt += "        number of angular azimuthal points: %d\n"%(undulator._NG_P)
        # txt += "        number of rays: %d\n"%(self.NRAYS)
        # txt += "        random seed: %d\n"%(self.SEED)
        txt += "-----------------------------------------------------\n"

        txt += "calculation code: %s\n"%undulator.code_undul_phot
        if self.__result_radiation is None:
            txt += "radiation: NOT YET CALCULATED\n"
        else:
            txt += "radiation: CALCULATED\n"
        txt += "Sampling: \n"
        if undulator._FLAG_SIZE == 0:
            flag = "point"
        elif undulator._FLAG_SIZE == 1:
            flag = "Gaussian"
        elif undulator._FLAG_SIZE == 2:
            flag = "Far field backpropagated"

        txt += "        Photon source size sampling flag: %d (%s)\n"%(undulator._FLAG_SIZE,flag)
        if undulator._FLAG_SIZE == 1:
            if undulator.__result_photon_size_sigma is not None:
                txt += "        Photon source size sigma (Gaussian): %6.3f um \n"%(1e6 * self.__result_photon_size_sigma)

        txt += "-----------------------------------------------------\n"
        return txt

    def to_python_code(self, data=None):
        script = ''
        try:
            script += self.get_electron_beam().to_python_code()
        except:
            script += "\n\n#Error retrieving electron_beam code"

        try:
            script += self.get_magnetic_structure().to_python_code()
        except:
            script += "\n\n#Error retrieving magnetic structure code"

        script += "\n\n\n# light source\nfrom shadow4.sources.undulator.s4_undulator_light_source import S4UndulatorLightSource"
        script += "\nlight_source = S4UndulatorLightSource(name='%s', electron_beam=electron_beam, magnetic_structure=source)" % \
                                                          (self.get_name())

        script += "\n\n\n#beamline\nfrom shadow4.beamline.s4_beamline import S4Beamline"
        script += "\nbeamline = S4Beamline(light_source=light_source)"
        return script

if __name__ == "__main__":

    import numpy

    do_plots = True

    su = S4Undulator(K_vertical=0.25,
                     period_length=0.032,
                     number_of_periods=50,
                     code_undul_phot='SRW',
                     use_gaussian_approximation=1,
                     )

    ebeam = S4ElectronBeam(energy_in_GeV=6.04,
                 energy_spread = 0.0,
                 current = 0.2,
                 number_of_bunches = 400,
                 moment_xx=(400e-6)**2,
                 moment_xxp=0.0,
                 moment_xpxp=(10e-6)**2,
                 moment_yy=(10e-6)**2,
                 moment_yyp=0.0,
                 moment_ypyp=(4e-6)**2 )

    ls = S4UndulatorLightSource(name="", electron_beam=ebeam, magnetic_structure=su)

    beam =  ls.get_beam()

    print(ls.info())
    print(ls.to_python_code())

max_stars_repo_name: srio/minishadow

max_issues_repo_name: srio/minishadow

reason: import numpy,import scipy,from scipy


text:
import keras
from keras import layers
from keras import backend as K
from keras.models import Model
from keras.datasets import mnist

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

# Global Variables
img_shape = (28, 28, 1)
batch_size = 16
latent_dim = 2

def image_input():
    input_img = keras.Input(shape=img_shape)

    x = layers.Conv2D(32, 3, padding="same", activation="relu")(input_img)
    x = layers.Conv2D(64, 3, padding="same", activation="relu", strides=(2, 2))(x)
    x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
    x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
    shape_before_flattening = K.int_shape(x)

    x = layers.Flatten()(x)
    x = layers.Dense(32, activation="relu")(x)

    z_mean = layers.Dense(latent_dim)(x)
    z_log_var = layers.Dense(latent_dim)(x)
    shape_after_flattening = K.int_shape(z_log_var)
    return z_mean, z_log_var, input_img, shape_before_flattening, shape_after_flattening

def sampling(args, latent_dim=2):
    mean, log_var = args
    epsilon = K.random_normal(
        shape=(K.shape(mean)[0], latent_dim), mean=0.0, stddev=1.0
    )
    return mean + K.exp(log_var) * epsilon

def decoder_input(z, shape_before_flattening):

    # This is the input where we will feed `z`.

    decoder_input = layers.Input(shape=K.int_shape(z)[1:])

    # Upsample to the correct number of units
    x = layers.Dense(np.prod(shape_before_flattening[1:]), activation="relu")(
        decoder_input
    )

    # Reshape into an image of the same shape as before our last `Flatten` layer
    x = layers.Reshape(shape_before_flattening[1:])(x)

    # We then apply then reverse operation to the initial
    # stack of convolution layers: a `Conv2DTranspose` layers
    # with corresponding parameters.
    x = layers.Conv2DTranspose(
        32, 3, padding="same", activation="relu", strides=(2, 2)
    )(x)
    x = layers.Conv2D(1, 3, padding="same", activation="sigmoid")(x)
    # We end up with a feature map of the same size as the original input.

    # This is our decoder model.
    decoder = Model(decoder_input, x)

    # We then apply it to `z` to recover the decoded `z`.
    z_decoded = decoder(z)

    return z_decoded, decoder

class CustomVariationalLayer(keras.layers.Layer):
    def vae_loss(self, x, z_decoded):
        x = K.flatten(x)
        z_decoded = K.flatten(z_decoded)
        xent_loss = keras.metrics.binary_crossentropy(x, z_decoded)
        kl_loss = -5e-4 * K.mean(
            1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1
        )
        return K.mean(xent_loss + kl_loss)

    def call(self, inputs):
        x = inputs[0]
        z_decoded = inputs[1]
        loss = self.vae_loss(x, z_decoded)
        self.add_loss(loss, inputs=inputs)
        # We don't use this output.
        return x

def autoencoder_train(input_img, y):
    vae = Model(input_img, y)
    vae.compile(optimizer="rmsprop", loss=None)
    vae.summary()

    # Train the VAE on MNIST digits
    (x_train, _), (x_test, y_test) = mnist.load_data()

    x_train = x_train.astype("float32") / 255.0
    x_train = x_train.reshape(x_train.shape + (1,))
    x_test = x_test.astype("float32") / 255.0
    x_test = x_test.reshape(x_test.shape + (1,))

    vae.fit(
        x=x_train,
        y=None,
        shuffle=True,
        epochs=10,
        batch_size=batch_size,
        validation_data=(x_test, None),
    )

def verify_plots(decoder):
    # Display a 2D manifold of the digits
    n = 15  # figure with 15x15 digits
    digit_size = 28
    figure = np.zeros((digit_size * n, digit_size * n))
    # Linearly spaced coordinates on the unit square were transformed
    # through the inverse CDF (ppf) of the Gaussian
    # to produce values of the latent variables z,
    # since the prior of the latent space is Gaussian
    grid_x = norm.ppf(np.linspace(0.05, 0.95, n))
    grid_y = norm.ppf(np.linspace(0.05, 0.95, n))

    for i, yi in enumerate(grid_x):
        for j, xi in enumerate(grid_y):
            z_sample = np.array([[xi, yi]])
            z_sample = np.tile(z_sample, batch_size).reshape(batch_size, 2)
            x_decoded = decoder.predict(z_sample, batch_size=batch_size)
            digit = x_decoded[0].reshape(digit_size, digit_size)
            figure[
                i * digit_size : (i + 1) * digit_size,
                j * digit_size : (j + 1) * digit_size,
            ] = digit

    plt.figure(figsize=(10, 10))
    plt.imshow(figure, cmap="Greys_r")
    plt.show()

if __name__ == "__main__":
    (
        z_mean,
        z_log_var,
        input_img,
        shape_before_flattening,
        shape_after_flattening,
    ) = image_input()
    z = layers.Lambda(sampling, output_shape=shape_after_flattening)(
        [z_mean, z_log_var]
    )
    z_decoded, decoder = decoder_input(z, shape_before_flattening)
    y = CustomVariationalLayer()([input_img, z_decoded])
    autoencoder_train(input_img, y)
    verify_plots(decoder=decoder)

ext: py

max_stars_repo_name: Anselmoo/DeepLearningwKeras

max_issues_repo_name: Anselmoo/DeepLearningwKeras

max_forks_repo_name: Anselmoo/DeepLearningwKeras

reason: import numpy,from scipy


text:
# Copyright (C) 2018 Paul King

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import csv
from collections import namedtuple

from pymedphys._imports import numpy as np

def read_mapcheck_txt(file_name):
    """
    Read native MapCheck data file and return dose array.

    Parameters
    ----------
    file_name : string
        | file name of MapCheck file including path

    Returns
    -------
    MapCheck : named tuple
        | MapCheck.x = np.array, float x-coords
        | MapCheck.y = np.array, float y-coords
        | MapCheck.dose = np.array (x,y), float dose
    """

    Mapcheck = namedtuple("Mapcheck", ["x", "y", "dose"])

    with open(file_name, "r") as mapcheck_file:
        m_chk = "\n".join(mapcheck_file.readlines())
        m_chk = m_chk.split("Dose Interpolated")[-1]
        m_chk = m_chk.split("\n")[2:]

    temp = [r for r in csv.reader(m_chk, delimiter="\t") if "Xcm" in r]
    x_coord = np.array(temp[0][2:]).astype(float)
    y_coord, dose = np.array([]), np.array([])

    for line in csv.reader(m_chk, delimiter="\t"):
        if len(line) > 1:
            try:
                line = [float(r) for r in line]
                y_coord = np.insert(y_coord, 0, float(line[0]))
                dose = np.append(dose, line[2:])
            except ValueError:
                pass
    dose = np.array(dose).flatten().astype(float)
    dose.shape = (len(x_coord), len(y_coord))
    return Mapcheck(x_coord, y_coord, dose)

max_forks_repo_name: cpbhatt/pymedphys

reason: import numpy


text:
from ast import Mult
import os
import random
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms

import torchvision.transforms as T
from main import ContrastiveLearningTransform2, Branch
# from main2 import ContrastiveLearningTransform, Branch

from functools import partial
from tqdm import tqdm
import pdb
st = pdb.set_trace

class Config(object):
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

def sample_data(loader):
    while True:
        for batch in loader:
            yield batch

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='cifar10')
    parser.add_argument('--data_path', type=str, default='data/c10_data2.pkl')
    parser.add_argument('--latent_path', type=str, default='data/c10_latent.pkl')
    parser.add_argument('--view_path', type=str, default='data/c10_obj=l2_eps=0.3_loss=l2_reg=0.01_eps2=0.5_n=8_seed=0.pkl')
    args = parser.parse_args()

    device = 'cuda'

    # model = Branch(args_simclr).to(device)
    # saved_dict = torch.load('../pretrained/simclr-cifar10-resnet18-800ep-1.pth')['state_dict']
    # model.load_state_dict(saved_dict, strict=True)
    opt = args
    if opt.dataset == 'cifar10':
        from main import Branch
        args_simclr = Config(dim_proj='2048,2048', dim_pred=512, loss='simclr')
        model = Branch(args_simclr).to(device)
        saved_dict = torch.load('../pretrained/simclr-cifar10-resnet18-800ep-1.pth')['state_dict']
        model.load_state_dict(saved_dict, strict=True)
        # expert_transform = ContrastiveLearningTransformCIFAR10()
        mean = (0.4914, 0.4822, 0.4465)
        std = (0.2023, 0.1994, 0.2010)
    elif opt.dataset == 'cifar100':
        from resnet_big import SupConResNet
        # import apex
        model = SupConResNet(name='resnet50')
        # model = apex.parallel.convert_syncbn_model(model)
        # model.encoder = torch.nn.DataParallel(model.encoder)
        state_dict = torch.load('../../SupContrast/logs/0425_c100_bs=1024_base/weights/last.pth', map_location='cpu')['model']
        new_state_dict = {}
        for k, v in state_dict.items():
            k = k.replace("module.", "")
            new_state_dict[k] = v
        state_dict = new_state_dict
        model.load_state_dict(state_dict, strict=True)
        model = model.to(device)
        # expert_transform = ContrastiveLearningTransformCIFAR100()
        mean = (0.5071, 0.4867, 0.4408)
        std = (0.2675, 0.2565, 0.2761)
    elif opt.dataset == 'tiny':
        from resnet_big import SupConResNet
        # supcon codebase
        model = SupConResNet(name='resnet50')
        state_dict = torch.load('../../SupContrast/logs/0430_t64_bs=512_base_2/weights/ckpt_epoch_1000.pth', map_location='cpu')['model']
        new_state_dict = {}
        for k, v in state_dict.items():
            k = k.replace("module.", "")
            new_state_dict[k] = v
        state_dict = new_state_dict
        model.load_state_dict(state_dict, strict=True)
        model = model.to(device)
        # expert_transform = ContrastiveLearningTransformTinyImageNet1()
        mean = (0.480, 0.448, 0.398)
        std = (0.277, 0.269, 0.282)

    normalize = T.Normalize(mean=mean, std=std)
    # single_transform = T.Compose([T.ToTensor(), normalize])
    single_transform = T.Compose([T.ToTensor(), normalize])

    model.eval()  # NOTE

    batch_size = 256

    from utils_data import MultiViewDataset
    train_dataset = MultiViewDataset(
        data_path=args.data_path,
        latent_path=args.latent_path,
        view_paths=[args.view_path],
        transform0=normalize,
        transform1=normalize,
        transform2=normalize,
        transform3=normalize,
        n_views=2,  # NOTE: we use 0, 3, 4
        train=True,
    )

    train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        shuffle=True,
        batch_size=batch_size,
        pin_memory=True,
        num_workers=0,
        drop_last=True
    )

    p = 2
    n_views = 3

    d01_list = list()
    d02_list = list()
    d12_list = list()

    s01_list = list()
    s02_list = list()
    s12_list = list()

    s_pos = list()
    s_neg = list()
    d_pos = list()
    d_neg = list()

    loader = sample_data(train_loader)

    with torch.no_grad():
        for i in tqdm(range(195)):
            imgs, _, _ = next(loader)
            h0 = model(imgs[0].cuda())
            h1 = model(imgs[3].cuda())
            h2 = model(imgs[4].cuda())
            h0 = F.normalize(h0, dim=-1, p=2)
            h1 = F.normalize(h1, dim=-1, p=2)
            h2 = F.normalize(h2, dim=-1, p=2)

            d01 = torch.norm(h0 - h1, p=p, dim=1)
            d02 = torch.norm(h0 - h2, p=p, dim=1)
            d12 = torch.norm(h1 - h2, p=p, dim=1)
            s01 = torch.sum(h0 * h1, dim=1)
            s02 = torch.sum(h0 * h2, dim=1)
            s12 = torch.sum(h1 * h2, dim=1)

            d01_list.append(d01)
            d02_list.append(d02)
            d12_list.append(d12)
            s01_list.append(s01)
            s02_list.append(s02)
            s12_list.append(s12)

            features = torch.cat([h0, h1, h2], dim=0)
            similarity_matrix = torch.matmul(features, features.T)
            distance_matrix = torch.cdist(features, features, p=p)

            labels = torch.cat([torch.arange(batch_size) for i in range(n_views)], dim=0)
            labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()
            mask = torch.eye(labels.shape[0], dtype=torch.bool)
            labels = labels[~mask].view(labels.shape[0], -1)
            similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)
            distance_matrix = distance_matrix[~mask].view(similarity_matrix.shape[0], -1)
            positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)
            negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)
            s_pos.append(positives)
            s_neg.append(negatives)

            dist_pos = distance_matrix[labels.bool()].view(labels.shape[0], -1)
            dist_neg = distance_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)
            d_pos.append(dist_pos)
            d_neg.append(dist_neg)

    d01 = torch.cat(d01_list, dim=0)
    d02 = torch.cat(d02_list, dim=0)
    d12 = torch.cat(d12_list, dim=0)
    s01 = torch.cat(s01_list, dim=0)
    s02 = torch.cat(s02_list, dim=0)
    s12 = torch.cat(s12_list, dim=0)
    s_pos = torch.cat(s_pos, dim=0)
    s_neg = torch.cat(s_neg, dim=0)
    d_pos = torch.cat(d_pos, dim=0)
    d_neg = torch.cat(d_neg, dim=0)
    print(f"d01: {d01.mean().item():.4f} +- {d01.std().item():.4f}")
    print(f"distance: (0, 1): {d01.mean().item():.4f}, (0, 2): {d02.mean().item():.4f}, (1, 2): {d12.mean().item():.4f}")
    print(f"cosine: (0, 1): {s01.mean().item():.4f}, (0, 2): {s02.mean().item():.4f}, (1, 2): {s12.mean().item():.4f}")
    print(f"cosine: positive: {s_pos.mean().item():.4f}, negative: {s_neg.mean().item():.4f}")
    print(f"distance: positive: {d_pos.mean().item():.4f}, negative: {d_neg.mean().item():.4f}")

    # save visuals

max_stars_repo_name: phymhan/essl

max_issues_repo_name: phymhan/essl

max_forks_repo_name: phymhan/essl


text:
# We need a normalized dump of the embeddings generated by TypedDM.
# Use this to normalize the relation and entity embeddings

import numpy as np
import pickle
import json
import os
import sys
import copy
import sklearn.preprocessing

MODEL_TYPE = '../dumps/fb15k237'

def normalize(X):
    return sklearn.preprocessing.normalize(X)

dump = {}
with open(MODEL_TYPE+'_distmult_dump.pkl', "rb") as f:
    dump = pickle.load(f)

def normalize_dump_distmult(dump):
    dump_norm = {}
    for key in dump:
        if(key == 'relation_to_id' or key == 'entity_to_id'):
            dump_norm[key] = dump[key]
        else:
            dump_norm[key] = normalize(dump[key])
    with open(MODEL_TYPE+"_distmult_dump_norm.pkl", "wb") as f:
        pickle.dump(dump_norm, f)

normalize_dump_distmult(dump)

max_stars_repo_name: dair-iitd/OxKBC

max_issues_repo_name: dair-iitd/OxKBC

max_forks_repo_name: dair-iitd/OxKBC


text:
# Copyright Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0

import numpy as np
from numpy.testing import assert_allclose, assert_raises
import pytest

from jax import jit, random, tree_multimap, value_and_grad, vmap
import jax.numpy as jnp

import numpyro
from numpyro import handlers
import numpyro.distributions as dist
from numpyro.distributions import constraints
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.util import log_density
import numpyro.optim as optim
from numpyro.util import not_jax_tracer, optional

@pytest.mark.parametrize('mask_last', [1, 5, 10])
@pytest.mark.parametrize('use_jit', [False, True])
def test_mask(mask_last, use_jit):
    N = 10
    mask = np.ones(N, dtype=np.bool)
    mask[-mask_last] = 0

    def model(data, mask):
        with numpyro.plate('N', N):
            x = numpyro.sample('x', dist.Normal(0, 1))
            with handlers.mask(mask=mask):
                numpyro.sample('y', dist.Delta(x, log_density=1.))
                with handlers.scale(scale=2):
                    numpyro.sample('obs', dist.Normal(x, 1), obs=data)

    data = random.normal(random.PRNGKey(0), (N,))
    x = random.normal(random.PRNGKey(1), (N,))
    if use_jit:
        log_joint = jit(lambda *args: log_density(*args)[0], static_argnums=(0,))(
            model, (data, mask), {}, {'x': x, 'y': x})
    else:
        log_joint = log_density(model, (data, mask), {}, {'x': x, 'y': x})[0]
    log_prob_x = dist.Normal(0, 1).log_prob(x)
    log_prob_y = mask
    log_prob_z = dist.Normal(x, 1).log_prob(data)
    expected = (log_prob_x + jnp.where(mask,  log_prob_y + 2 * log_prob_z, 0.)).sum()
    assert_allclose(log_joint, expected, atol=1e-4)

@pytest.mark.parametrize("num_particles", [1, 2])
@pytest.mark.parametrize("mask", [
    True,
    False,
    np.array([True]),
    np.array([False]),
    np.array([False, True, False]),
])
@pytest.mark.parametrize("Elbo", [
    Trace_ELBO,
])
def test_obs_mask_ok(Elbo, mask, num_particles):
    data = np.array([7., 7., 7.])

    def model():
        x = numpyro.sample("x", dist.Normal(0., 1.))
        with numpyro.plate("plate", len(data)):
            y = numpyro.sample("y", dist.Normal(x, 1.),
                               obs=data, obs_mask=mask)
            if not_jax_tracer(y):
                assert ((y == data) == mask).all()

    def guide():
        loc = numpyro.param("loc", np.zeros(()))
        scale = numpyro.param("scale", np.ones(()),
                              constraint=constraints.positive)
        x = numpyro.sample("x", dist.Normal(loc, scale))
        with numpyro.plate("plate", len(data)):
            with handlers.mask(mask=np.invert(mask)):
                numpyro.sample("y_unobserved", dist.Normal(x, 1.))

    elbo = Elbo(num_particles=num_particles)
    svi = SVI(model, guide, numpyro.optim.Adam(1), elbo)
    svi_state = svi.init(random.PRNGKey(0))
    svi.update(svi_state)

@pytest.mark.parametrize("num_particles", [1, 2])
@pytest.mark.parametrize("mask", [
    True,
    False,
    np.array([True]),
    np.array([False]),
    np.array([False, True, True, False]),
])
@pytest.mark.parametrize("Elbo", [
    Trace_ELBO,
])
def test_obs_mask_multivariate_ok(Elbo, mask, num_particles):
    data = np.full((4, 3), 7.0)

    def model():
        x = numpyro.sample("x", dist.MultivariateNormal(np.zeros(3), np.eye(3)))
        with numpyro.plate("plate", len(data)):
            y = numpyro.sample("y", dist.MultivariateNormal(x, np.eye(3)),
                               obs=data, obs_mask=mask)
            if not_jax_tracer(y):
                assert ((y == data).all(-1) == mask).all()

    def guide():
        loc = numpyro.param("loc", np.zeros(3))
        cov = numpyro.param("cov", np.eye(3),
                            constraint=constraints.positive_definite)
        x = numpyro.sample("x", dist.MultivariateNormal(loc, cov))
        with numpyro.plate("plate", len(data)):
            with handlers.mask(mask=np.invert(mask)):
                numpyro.sample("y_unobserved", dist.MultivariateNormal(x, np.eye(3)))

    elbo = Elbo(num_particles=num_particles)
    svi = SVI(model, guide, numpyro.optim.Adam(1), elbo)
    svi_state = svi.init(random.PRNGKey(0))
    svi.update(svi_state)

def test_mask_inf():
    def model():
        with handlers.mask(mask=jnp.zeros(10, dtype=bool)):
            numpyro.factor('inf', -jnp.inf)

    log_joint = log_density(model, (), {}, {})[0]
    assert_allclose(log_joint, 0.)

@pytest.mark.parametrize('use_context_manager', [True, False])
def test_scale(use_context_manager):
    def model(data):
        x = numpyro.sample('x', dist.Normal(0, 1))
        with optional(use_context_manager, handlers.scale(scale=10)):
            numpyro.sample('obs', dist.Normal(x, 1), obs=data)

    model = model if use_context_manager else handlers.scale(model, 10.)
    data = random.normal(random.PRNGKey(0), (3,))
    x = random.normal(random.PRNGKey(1))
    log_joint = log_density(model, (data,), {}, {'x': x})[0]
    log_prob1, log_prob2 = dist.Normal(0, 1).log_prob(x), dist.Normal(x, 1).log_prob(data).sum()
    expected = log_prob1 + 10 * log_prob2 if use_context_manager else 10 * (log_prob1 + log_prob2)
    assert_allclose(log_joint, expected)

def test_substitute():
    def model():
        x = numpyro.param('x', None)
        y = handlers.substitute(lambda: numpyro.param('y', None) * numpyro.param('x', None), {'y': x})()
        return x + y

    assert handlers.substitute(model, {'x': 3.})() == 12.

def test_seed():
    def _sample():
        x = numpyro.sample('x', dist.Normal(0., 1.))
        y = numpyro.sample('y', dist.Normal(1., 2.))
        return jnp.stack([x, y])

    xs = []
    for i in range(100):
        with handlers.seed(rng_seed=i):
            xs.append(_sample())
    xs = jnp.stack(xs)

    ys = vmap(lambda rng_key: handlers.seed(lambda: _sample(), rng_key)())(jnp.arange(100))
    assert_allclose(xs, ys, atol=1e-6)

def test_nested_seeding():
    def fn(rng_key_1, rng_key_2, rng_key_3):
        xs = []
        with handlers.seed(rng_seed=rng_key_1):
            with handlers.seed(rng_seed=rng_key_2):
                xs.append(numpyro.sample('x', dist.Normal(0., 1.)))
                with handlers.seed(rng_seed=rng_key_3):
                    xs.append(numpyro.sample('y', dist.Normal(0., 1.)))
        return jnp.stack(xs)

    s1, s2 = fn(0, 1, 2), fn(3, 1, 2)
    assert_allclose(s1, s2)
    s1, s2 = fn(0, 1, 2), fn(3, 1, 4)
    assert_allclose(s1[0], s2[0])
    assert_raises(AssertionError, assert_allclose, s1[1], s2[1])

def test_condition():
    def model():
        x = numpyro.sample('x', dist.Delta(0.))
        y = numpyro.sample('y', dist.Normal(0., 1.))
        return x + y

    model = handlers.condition(handlers.seed(model, random.PRNGKey(1)), {'y': 2.})
    model_trace = handlers.trace(model).get_trace()
    assert model_trace['y']['value'] == 2.
    assert model_trace['y']['is_observed']
    assert handlers.condition(model, {'y': 3.})() == 3.

def test_no_split_deterministic():
    def model():
        x = numpyro.sample('x', dist.Normal(0., 1.))
        y = numpyro.sample('y', dist.Normal(0., 1.))
        return x + y

    model = handlers.condition(model, {'x': 1., 'y': 2.})
    assert model() == 3.

def model_nested_plates_0():
    with numpyro.plate('outer', 10):
        x = numpyro.sample('y', dist.Normal(0., 1.))
        assert x.shape == (10,)
        with numpyro.plate('inner', 5):
            y = numpyro.sample('x', dist.Normal(0., 1.))
            assert y.shape == (5, 10)
            z = numpyro.deterministic('z', x ** 2)
            assert z.shape == (10,)

def model_nested_plates_1():
    with numpyro.plate('outer', 10, dim=-2):
        x = numpyro.sample('y', dist.Normal(0., 1.))
        assert x.shape == (10, 1)
        with numpyro.plate('inner', 5):
            y = numpyro.sample('x', dist.Normal(0., 1.))
            assert y.shape == (10, 5)
            z = numpyro.deterministic('z', x ** 2)
            assert z.shape == (10, 1)

def model_nested_plates_2():
    outer = numpyro.plate('outer', 10)
    inner = numpyro.plate('inner', 5, dim=-3)
    with outer:
        x = numpyro.sample('x', dist.Normal(0., 1.))
        assert x.shape == (10,)
    with inner:
        y = numpyro.sample('y', dist.Normal(0., 1.))
        assert y.shape == (5, 1, 1)
        z = numpyro.deterministic('z', x ** 2)
        assert z.shape == (10,)

    with outer, inner:
        xy = numpyro.sample('xy', dist.Normal(0., 1.), sample_shape=(10,))
        assert xy.shape == (5, 1, 10)

def model_nested_plates_3():
    outer = numpyro.plate('outer', 10, dim=-1)
    inner = numpyro.plate('inner', 5, dim=-2)
    numpyro.deterministic('z', 1.)

    with inner, outer:
        xy = numpyro.sample('xy', dist.Normal(jnp.zeros((5, 10)), 1.))
        assert xy.shape == (5, 10)

def model_dist_batch_shape():
    outer = numpyro.plate('outer', 10)
    inner = numpyro.plate('inner', 5, dim=-3)
    with outer:
        x = numpyro.sample('x', dist.Normal(jnp.zeros(10), 1.))
        assert x.shape == (10,)
    with inner:
        y = numpyro.sample('y', dist.Normal(0., jnp.ones(10)))
        assert y.shape == (5, 1, 10)
        z = numpyro.deterministic('z', x ** 2)
        assert z.shape == (10,)

    with outer, inner:
        xy = numpyro.sample('xy', dist.Normal(0., jnp.ones(10)), sample_shape=(10,))
        assert xy.shape == (5, 10, 10)

def model_subsample_1():
    outer = numpyro.plate('outer', 20, subsample_size=10)
    inner = numpyro.plate('inner', 10, subsample_size=5, dim=-3)
    with outer:
        x = numpyro.sample('x', dist.Normal(0., 1.))
        assert x.shape == (10,)
    with inner:
        y = numpyro.sample('y', dist.Normal(0., 1.))
        assert y.shape == (5, 1, 1)
        z = numpyro.deterministic('z', x ** 2)
        assert z.shape == (10,)

    with outer, inner:
        xy = numpyro.sample('xy', dist.Normal(0., 1.))
        assert xy.shape == (5, 1, 10)

def model_subsample_2():
    data = jnp.ones((10, 1, 20))
    outer = numpyro.plate('outer', data.shape[-1], subsample_size=10)
    inner = numpyro.plate('inner', data.shape[-3], subsample_size=5, dim=-3)
    with outer:
        x = numpyro.sample('x', dist.Normal(0., 1.))
        assert x.shape == (10,)
    with inner:
        y = numpyro.sample('y', dist.Normal(0., 1.))
        assert y.shape == (5, 1, 1)
        z = numpyro.deterministic('z', x ** 2)
        assert z.shape == (10,)

    with outer, inner:
        xy = numpyro.sample('xy', dist.Normal(0., 1.))
        assert xy.shape == (5, 1, 10)
        subsample_data = numpyro.subsample(data, event_dim=0)
        assert subsample_data.shape == (5, 1, 10)

@pytest.mark.parametrize('model', [
    model_nested_plates_0,
    model_nested_plates_1,
    model_nested_plates_2,
    model_nested_plates_3,
    model_dist_batch_shape,
    model_subsample_1,
    model_subsample_2,
])
def test_plate(model):
    trace = handlers.trace(handlers.seed(model, random.PRNGKey(1))).get_trace()
    jit_trace = handlers.trace(jit(handlers.seed(model, random.PRNGKey(1)))).get_trace()
    assert 'z' in trace
    for name, site in trace.items():
        if site['type'] == 'sample':
            assert_allclose(jit_trace[name]['value'].shape, site['value'].shape)

def test_subsample_data():
    data = jnp.arange(100.)
    subsample_size = 7
    with handlers.seed(rng_seed=0):
        with numpyro.plate("a", len(data), subsample_size=subsample_size) as idx:
            assert data[idx].shape == (subsample_size,)
            subsample_data = numpyro.subsample(data, event_dim=0)
            assert subsample_data.shape == (subsample_size,)

def test_subsample_param():
    data = jnp.arange(100.)
    subsample_size = 7
    with handlers.seed(rng_seed=0):
        with numpyro.plate("a", len(data), subsample_size=subsample_size):
            p0 = numpyro.param("p0", 0., event_dim=0)
            assert jnp.shape(p0) == ()
            p = numpyro.param("p", 0.5 * jnp.ones(len(data)), event_dim=0)
            assert len(p) == subsample_size

def test_subsample_substitute():
    data = jnp.arange(100.)
    subsample_size = 7
    subsample = jnp.array([13, 3, 30, 4, 1, 68, 5])
    with handlers.trace() as tr, handlers.seed(rng_seed=0), handlers.substitute(data={"a": subsample}):
        with numpyro.plate("a", len(data), subsample_size=subsample_size) as idx:
            assert data[idx].shape == (subsample_size,)
            assert_allclose(idx, subsample)
    assert tr["a"]["kwargs"]["rng_key"] is None

def test_subsample_replay():
    data = jnp.arange(100.)
    subsample_size = 7

    with handlers.trace() as guide_trace, handlers.seed(rng_seed=0):
        with numpyro.plate("a", len(data), subsample_size=subsample_size):
            pass

    with handlers.seed(rng_seed=1), handlers.replay(guide_trace=guide_trace):
        with numpyro.plate("a", len(data)):
            subsample_data = numpyro.subsample(data, event_dim=0)
            assert subsample_data.shape == (subsample_size,)

@pytest.mark.parametrize("scale", [1., 2.], ids=["unscaled", "scaled"])
@pytest.mark.parametrize("subsample", [False, True], ids=["full", "subsample"])
def test_subsample_gradient(scale, subsample):
    data = jnp.array([-0.5, 2.0])
    subsample_size = 1 if subsample else len(data)
    precision = 0.06 * scale

    def model(subsample):
        with handlers.substitute(data={"data": subsample}):
            with numpyro.plate("data", len(data), subsample_size) as ind:
                x = data[ind]
                z = numpyro.sample("z", dist.Normal(0, 1))
                numpyro.sample("x", dist.Normal(z, 1), obs=x)

    def guide(subsample):
        scale = numpyro.param("scale", 1.)
        with handlers.substitute(data={"data": subsample}):
            with numpyro.plate("data", len(data), subsample_size):
                loc = numpyro.param("loc", jnp.zeros(len(data)), event_dim=0)
                numpyro.sample("z", dist.Normal(loc, scale))

    if scale != 1.:
        model = handlers.scale(model, scale=scale)
        guide = handlers.scale(guide, scale=scale)

    num_particles = 50000
    optimizer = optim.Adam(0.1)
    elbo = Trace_ELBO(num_particles=num_particles)
    svi = SVI(model, guide, optimizer, loss=elbo)
    svi_state = svi.init(random.PRNGKey(0), None)
    params = svi.optim.get_params(svi_state.optim_state)
    normalizer = 2 if subsample else 1
    if subsample_size == 1:
        subsample = jnp.array([0])
        loss1, grads1 = value_and_grad(lambda x: svi.loss.loss(
            svi_state.rng_key, svi.constrain_fn(x), svi.model, svi.guide, subsample))(params)
        subsample = jnp.array([1])
        loss2, grads2 = value_and_grad(lambda x: svi.loss.loss(
            svi_state.rng_key, svi.constrain_fn(x), svi.model, svi.guide, subsample))(params)
        grads = tree_multimap(lambda *vals: vals[0] + vals[1], grads1, grads2)
        loss = loss1 + loss2
    else:
        subsample = jnp.array([0, 1])
        loss, grads = value_and_grad(lambda x: svi.loss.loss(
            svi_state.rng_key, svi.constrain_fn(x), svi.model, svi.guide, subsample))(params)

    actual_loss = loss / normalizer
    expected_loss, _ = value_and_grad(lambda x: svi.loss.loss(
        svi_state.rng_key, svi.constrain_fn(x), svi.model, svi.guide, None))(params)
    assert_allclose(actual_loss, expected_loss, rtol=precision, atol=precision)

    actual_grads = {name: grad / normalizer for name, grad in grads.items()}
    expected_grads = {'loc': scale * jnp.array([0.5, -2.0]), 'scale': scale * jnp.array([2.0])}
    assert actual_grads.keys() == expected_grads.keys()
    for name in expected_grads:
        assert_allclose(actual_grads[name], expected_grads[name], rtol=precision, atol=precision)

def test_messenger_fn_invalid():
    with pytest.raises(ValueError, match="to be a Python callable object"):
        with numpyro.handlers.mask(False):
            pass

@pytest.mark.parametrize('shape', [(), (5,), (2, 3)])
def test_plate_stack(shape):
    def guide():
        with numpyro.plate_stack("plates", shape):
            return numpyro.sample("x", dist.Normal(0, 1))

    x = handlers.seed(guide, 0)()
    assert x.shape == shape

@pytest.mark.parametrize('intervene,observe,flip', [
    (True, False, False),
    (False, True, False),
    (True, True, False),
    (True, True, True),
])
def test_counterfactual_query(intervene, observe, flip):
    # x -> y -> z -> w

    sites = ["x", "y", "z", "w"]
    observations = {"x": 1., "y": None, "z": 1., "w": 1.}
    interventions = {"x": None, "y": 0., "z": 2., "w": 1.}

    def model():
        with handlers.seed(rng_seed=0):
            x = numpyro.sample("x", dist.Normal(0, 1))
            y = numpyro.sample("y", dist.Normal(x, 1))
            z = numpyro.sample("z", dist.Normal(y, 1))
            w = numpyro.sample("w", dist.Normal(z, 1))
            return dict(x=x, y=y, z=z, w=w)

    if not flip:
        if intervene:
            model = handlers.do(model, data=interventions)
        if observe:
            model = handlers.condition(model, data=observations)
    elif flip and intervene and observe:
        model = handlers.do(
            handlers.condition(model, data=observations),
            data=interventions)

    with handlers.trace() as tr:
        actual_values = model()
    for name in sites:
        # case 1: purely observational query like handlers.condition
        if not intervene and observe:
            if observations[name] is not None:
                assert tr[name]['is_observed']
                assert_allclose(observations[name], actual_values[name])
                assert_allclose(observations[name], tr[name]['value'])
            if interventions[name] != observations[name]:
                if interventions[name] is not None:
                    assert_raises(AssertionError, assert_allclose, interventions[name], actual_values[name])
        # case 2: purely interventional query like old handlers.do
        elif intervene and not observe:
            assert not tr[name]['is_observed']
            if interventions[name] is not None:
                assert_allclose(interventions[name], actual_values[name])
            if observations[name] is not None:
                assert_raises(AssertionError, assert_allclose, observations[name], tr[name]['value'])
            if interventions[name] is not None:
                assert_raises(AssertionError, assert_allclose, interventions[name], tr[name]['value'])
        # case 3: counterfactual query mixing intervention and observation
        elif intervene and observe:
            if observations[name] is not None:
                assert tr[name]['is_observed']
                assert_allclose(observations[name], tr[name]['value'])
            if interventions[name] is not None:
                assert_allclose(interventions[name], actual_values[name])
            if interventions[name] != observations[name]:
                if interventions[name] is not None:
                    assert_raises(AssertionError, assert_allclose, interventions[name], tr[name]['value'])

def test_block():
    with handlers.trace() as trace:
        with handlers.block(hide=['x']):
            with handlers.seed(rng_seed=0):
                numpyro.sample('x', dist.Normal())
    assert 'x' not in trace

def test_scope():
    def fn():
        return numpyro.sample('x', dist.Normal())

    with handlers.trace() as trace:
        with handlers.seed(rng_seed=1):
            with handlers.scope(prefix='a'):
                fn()
            with handlers.scope(prefix='b'):
                with handlers.scope(prefix='a'):
                    fn()

    assert 'a/x' in trace
    assert 'b/a/x' in trace

def test_lift():
    def model():
        loc1 = numpyro.param("loc1", 0.)
        scale1 = numpyro.param("scale1", 1., constraint=constraints.positive)
        numpyro.sample("latent1", dist.Normal(loc1, scale1))

        loc2 = numpyro.param("loc2", 1.)
        scale2 = numpyro.param("scale2", 2., constraint=constraints.positive)
        latent2 = numpyro.sample("latent2", dist.Normal(loc2, scale2))
        return latent2

    loc1_prior = dist.Normal()
    scale1_prior = dist.LogNormal()
    prior = {"loc1": loc1_prior, "scale1": scale1_prior}

    with handlers.trace() as tr:
        with handlers.seed(rng_seed=1):
            model()

    with handlers.trace() as lifted_tr:
        with handlers.seed(rng_seed=2):
            with handlers.lift(prior=prior):
                model()

    for name in tr.keys():
        assert name in lifted_tr
        if name in prior:
            assert lifted_tr[name]['fn'] is prior[name]
            assert lifted_tr[name]['type'] == 'sample'
            assert lifted_tr[name]['value'] not in (0., 1.)
        elif name in ('loc2', 'scale2'):
            assert lifted_tr[name]['type'] == 'param'

def test_lift_memoize():
    def model():
        a = numpyro.param("loc")
        b = numpyro.param("loc")
        assert a == b

    with handlers.seed(rng_seed=1):
        with handlers.lift(prior=dist.Normal(0, 1)):
            model()

@pytest.mark.xfail(reason="Issue: https://github.com/pyro-ppl/numpyro/issues/964")
def test_collapse_beta_binomial():
    total_count = 10
    data = 3.

    def model1():
        c1 = numpyro.param("c1", 0.5, constraint=dist.constraints.positive)
        c0 = numpyro.param("c0", 1.5, constraint=dist.constraints.positive)
        with handlers.collapse():
            probs = numpyro.sample("probs", dist.Beta(c1, c0))
            numpyro.sample("obs", dist.Binomial(total_count, probs), obs=data)

    def model2():
        c1 = numpyro.param("c1", 0.5, constraint=dist.constraints.positive)
        c0 = numpyro.param("c0", 1.5, constraint=dist.constraints.positive)
        numpyro.sample("obs", dist.BetaBinomial(c1, c0, total_count),
                       obs=data)

    trace1 = handlers.trace(model1).get_trace()
    trace2 = handlers.trace(model2).get_trace()
    assert "probs" in trace1
    assert "obs" not in trace1
    assert "probs" not in trace2
    assert "obs" in trace2

    svi1 = SVI(model1, lambda: None, numpyro.optim.Adam(1), Trace_ELBO())
    svi2 = SVI(model2, lambda: None, numpyro.optim.Adam(1), Trace_ELBO())
    svi_state1 = svi1.init(random.PRNGKey(0))
    svi_state2 = svi2.init(random.PRNGKey(0))
    params1 = svi1.get_params(svi_state1)
    params2 = svi2.get_params(svi_state2)
    assert_allclose(params1["c1"], params2["c1"])
    assert_allclose(params1["c0"], params2["c0"])

    params1 = svi1.get_params(svi1.update(svi_state1)[0])
    params2 = svi2.get_params(svi2.update(svi_state2)[0])
    assert_allclose(params1["c1"], params2["c1"])
    assert_allclose(params1["c0"], params2["c0"])

def test_collapse_beta_bernoulli():
    data = 0.

    def model():
        c = numpyro.sample("c", dist.Gamma(1, 1))
        with handlers.collapse():
            probs = numpyro.sample("probs", dist.Beta(c, 2))
            numpyro.sample("obs", dist.Bernoulli(probs), obs=data)

    def guide():
        a = numpyro.param("a", 1., constraint=constraints.positive)
        b = numpyro.param("b", 1., constraint=constraints.positive)
        numpyro.sample("c", dist.Gamma(a, b))

    svi = SVI(model, guide, numpyro.optim.Adam(1), Trace_ELBO())
    svi_state = svi.init(random.PRNGKey(0))
    svi.update(svi_state)

@pytest.mark.xfail(reason="missing pattern in Funsor")
def test_collapse_beta_binomial_plate():
    data = np.array([0., 1., 5., 5.])

    def model():
        c = numpyro.sample("c", dist.Gamma(1, 1))
        with handlers.collapse():
            probs = numpyro.sample("probs", dist.Beta(c, 2))
            with numpyro.plate("plate", len(data)):
                numpyro.sample("obs", dist.Binomial(10, probs),
                               obs=data)

    def guide():
        a = numpyro.param("a", 1., constraint=constraints.positive)
        b = numpyro.param("b", 1., constraint=constraints.positive)
        numpyro.sample("c", dist.Gamma(a, b))

    svi = SVI(model, guide, numpyro.optim.Adam(1), Trace_ELBO())
    svi_state = svi.init(random.PRNGKey(0))
    svi.update(svi_state)

def test_prng_key():
    assert numpyro.prng_key() is None

    with handlers.seed(rng_seed=0):
        rng_key = numpyro.prng_key()

    assert rng_key.shape == (2,) and rng_key.dtype == "uint32"

def test_prng_key_with_vmap():
    def model(x=None):
        return numpyro.prng_key()

    x = handlers.seed(vmap(model), 0)(jnp.arange(10))
    assert (x == x[0]).all()
    y = vmap(handlers.seed(model, 0))(jnp.arange(10))
    assert (x == y).all()
    z = vmap(lambda i: handlers.seed(model, i)())(jnp.arange(10))
    z0 = handlers.seed(model, 0)()
    assert (z[1:] != z0).all()
    assert (z[0] == z0).all()

def test_subsample_fn():
    size = 20
    subsample_size = 11
    num_samples = 1000000

    @jit
    def subsample_fn(rng_key):
        return numpyro.primitives._subsample_fn(size, subsample_size, rng_key)

    rng_keys = random.split(random.PRNGKey(0), num_samples)
    subsamples = vmap(subsample_fn)(rng_keys)
    for k in range(1, 11):
        i = random.randint(random.PRNGKey(k), (), 0, size)
        assert_allclose(jnp.mean(subsamples == i, axis=0), jnp.full(subsample_size, 1 / size), atol=1e-3)

        # test that values are not duplicated
        assert len(set(subsamples[k])) == subsample_size

max_stars_repo_name: ahoho/numpyro

max_issues_repo_name: ahoho/numpyro

max_forks_repo_name: ahoho/numpyro

reason: import numpy,from numpy,import jax,from jax


text:
from flask import Flask, render_template, redirect, Response
from flask_frozen import Freezer

import os, sys, yaml, pytz, time, datetime, heapq  # 堆队列
import numpy as np
import click

add_datas = [
    "./_data/_effective_kaggle.yaml",
    "./_data/_effective_drivendata.yaml",
    "./_data/_effective_codalab.yaml",
    "./_data/_effective_challengedata.yaml",
    "./_data/_effective_crowdai.yaml",
    "./_data/_effective_signate.yaml",
    "./_data/_effective_crowdanalytix.yaml",
    "./_data/_effective_unearthed.yaml",
    "./_data/_effective_analyticsvidhya.yaml",
    "./_data/_effective_GECCO.yaml",
    "./_data/_effective_grand_challenge.yaml",
    "./_data/_effective_herox.yaml",
    "./_data/_effective_ICDAR.yaml",
    "./_data/_effective_kelvins.yaml",
    "./_data/_effective_天池.yaml",
    "./_data/_effective_iQIYI.yaml",
    "./_data/_effective_FUTURELAB_AI.yaml",
    "./_data/_effective_点石.yaml",
    "./_data/_effective_kesci.yaml",
    "./_data/_effective_dcjingsai.yaml",
    "./_data/_effective_datafountain.yaml",
    "./_data/_effective_biendata.yaml",
    "./_data/_effective_ICME.yaml",
    "./_data/_effective_NeurIPS.yaml",
    "./_data/_effective_EvalAI.yaml",
    "./_data/_effective_FlyAI.yaml",
    "./_data/_effective_tinymind.yaml",
    "./_data/_effective_MagicMirror.yaml",
    "./_data/_effective_others.yaml",
]

competitions = []
for add_data in add_datas:
    with open(add_data, "r", encoding="utf-8") as f:
        cfg = f.read()
        competitions.extend(yaml.load(cfg))
competitions = [
    c
    for c in competitions
    if c["deadtime"] > datetime.date.today().strftime("%Y-%m-%d")
]
competitions.sort(key=lambda x: x["deadtime"])

# Find CVPR competitions
# CVPR_comps = [comp for comp in competitions if 'CVPR' in  [ 'CVPR' for l in comp['hostby'] if 'CVPR' in [ 'CVPR' for key in l.keys() if 'CVPR' in key ] ]  ]

app = Flask(__name__)

app.config["FREEZER_DESTINATION"] = "docs"
# https://stackoverflow.com/questions/11577147/how-to-fix-page-404-on-github-page
app.config["FREEZER_DESTINATION_IGNORE"] = [".git*", "CNAME", "favicon.ico"]

freezer = Freezer(app)

id_type_checkboxs = [
    {"id": "PF-checkbox", "name": "Platform"}
    # ,{'id': 'IT-checkbox', 'name': 'Industry'}
    ,
    {"id": "AC-checkbox", "name": "Academia"},
]

id_type2_checkboxs = [
    {"id": "DM-checkbox", "name": "Data Mining"},
    {"id": "CV-checkbox", "name": "Computer Vision"},
    {"id": "NLP-checkbox", "name": "Natural Language Processing"},
    {"id": "RL-checkbox", "name": "Reinforcement Learning/Robotics"},
    {"id": "SP-checkbox", "name": "Speech/Signal Proccessing"},
]

@app.route("/")
def index(
    id_type_checkboxs=id_type_checkboxs,
    id_type2_checkboxs=id_type2_checkboxs,
    competitions=competitions,
):
    # 保证每个 comp 的 type 键的值是 list 类型
    [
        comp.update({"type1": [comp["type1"]]})
        for comp in competitions
        if not isinstance(comp["type1"], list)
    ]
    [
        comp.update({"type2": [comp["type2"]]})
        for comp in competitions
        if not isinstance(comp["type2"], list)
    ]
    # 判断 prize
    for comp in competitions:
        if (comp["prize"] == "NaN") or (comp["prize"] == "Kaggle Swag"):
            comp.update({"howprize": ["unrewarded"]})
        else:
            comp.update({"howprize": ["rewarded"]})
    # [ comp.update({'howprize': ['unrewarded']}) for comp in competitions if comp['prize'] == 'NaN']
    # [ comp.update({'howprize': ['rewarded']}) for comp in competitions if comp['prize'] ~= 'NaN']
    return render_template(
        "index.html",
        id_type_checkboxs=id_type_checkboxs,
        id_type2_checkboxs=id_type2_checkboxs,
        competitions=competitions,
    )

with open("./_data/_hosts.yaml", "r", encoding="utf-8") as f:
    cfg = f.read()
    hosts = yaml.load(cfg)

@app.route("/hostby.html")
def hostby(id_type_checkboxs=id_type_checkboxs, hosts=hosts):
    return render_template(
        "hostby.html", id_type_checkboxs=id_type_checkboxs, hosts=hosts
    )

# @app.route('/rss.xml')
# def log(competitions = competitions):
#     num_largest = 2
#     comps_pubtime = np.array([ int(i['pubtime'].replace('-', '')) for i in competitions ])
#     index_largest = [ np.where(comps_pubtime == largest_value)[0].tolist() for largest_value in heapq.nlargest(num_largest, np.unique(comps_pubtime)) ]
#     competitions = [ (competitions[largest_index], block_time) for block_time, largest_time in enumerate(index_largest) for largest_index in largest_time]

#     [ comp.update({'pubtime': datetime.datetime.fromtimestamp(int(datetime.datetime.strptime(comp['pubtime'], '%Y-%m-%d').timestamp()), pytz.timezone('Asia/Shanghai')).strftime("%a, %d %b %Y") }) for comp, _ in competitions ]

#     return render_template('rss.xml', competitions=competitions)

@app.route("/update_log.xml")
def products_xml(competitions=competitions):
    num_largest = 2
    comps_pubtime = np.array([int(i["pubtime"].replace("-", "")) for i in competitions])
    time_largest = heapq.nlargest(num_largest, np.unique(comps_pubtime))
    index_largest = [
        np.where(comps_pubtime == largest_value)[0].tolist()
        for largest_value in time_largest
    ]
    competitions = [
        (competitions[largest_index], block_time)
        for block_time, largest_time in enumerate(index_largest)
        for largest_index in largest_time
    ]

    [
        comp.update(
            {
                "pubtime": datetime.datetime.fromtimestamp(
                    int(
                        datetime.datetime.strptime(
                            comp["pubtime"], "%Y-%m-%d"
                        ).timestamp()
                    ),
                    pytz.timezone("Asia/Shanghai"),
                ).strftime("%a, %d %b %Y")
            }
        )
        for comp, _ in competitions
    ]

    output = '<?xml version="1.0" encoding="UTF-8" ?>'
    output += '<rss version="2.0">'
    output += "<channel>"
    output += "<title>Data Science Challenge / Competition</title>"
    output += "<link>https://iphysresearch.github.io/DataSciComp</link>"

    update_block = [
        datetime.datetime.fromtimestamp(
            int(
                datetime.datetime.strptime(
                    str(time_largest[block]), "%Y%m%d"
                ).timestamp()
            ),
            pytz.timezone("Asia/Shanghai"),
        ).strftime("%m/%d/%Y")
        for block in range(num_largest)
    ]
    output += "<description>Latest update at {} (GMT+0800).</description>".format(
        update_block[0]
    )

    for comp, _ in competitions:
        output += "<item>"
        output += "<title>{:s}</title>".format(comp["title"])
        output += "<link>{:s}</link>".format(comp["url"])
        output += "<category>{:s}</category>".format("/".join(comp["type1"]))
        output += "<category>{:s}</category>".format("/".join(comp["type2"]))
        output += "<pubDate>{:s}</pubDate>".format(comp["pubtime"])
        output += "<description>{:s}</description>".format(
            comp["note"].replace("<br>", "")
        )
        output += "</item>"
    output += "</channel>"
    output += "</rss>"
    return Response(output, mimetype="application/xml")

if __name__ == "__main__":
    # freezer.freeze()
    freezer.run(debug=True)

max_stars_repo_name: shuiyuejihua/DataSciComp

max_issues_repo_name: shuiyuejihua/DataSciComp

max_forks_repo_name: shuiyuejihua/DataSciComp